Jikai Jason Li

I am an undergraduate student studying Computer Science at the University of Toronto. I am currently an undergraduate research assistant in the Department of Computer Science, advised by Prof. Fan Long.

My current research focuses on LLM-based program optimization, from multi-file C and C++ systems to high-performance GPU kernels.

Research interests: machine learning, reinforcement learning, and ML systems and infrastructure.

Email: jikaijasonli [AT] gmail.com


Research

LLM Compiler Superoptimization

I am working on making LLM-based code optimization practical beyond isolated functions. My work extends a compiler-superoptimization pipeline to multi-file C and C++ programs and develops the search, verification, and evaluation infrastructure needed to test changes safely at scale. The benchmark currently covers 139 cases from 49 LLVM benchmark families, including the 87-file SPASS theorem prover.

In an 11-workload study, I orchestrated 400 optimization attempts with 16 parallel search workers and four verifier agents. By sending verifiers only new evidence and using their counterexamples to refine optimization constraints, I reduced later prompt sizes by about 75% and improved aggregate speedup on the targeted workloads from 1.07× to 1.24× while preserving correctness.


Preprints

CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe
Tara Saba, Zhiyang Chen, Jikai Jason Li, Anne Ouyang, Xujie Si, Fan Long
arXiv preprint; submitted to NeurIPS 2026

CuTeGen studies how an agentic LLM workflow can generate and optimize CUDA/CuTe kernels. My work centered on evaluation and debugging: I ran generated kernels across 209 KernelBench tasks on remote NVIDIA GPUs and built tooling for correctness validation, performance timing, token accounting, and failed-run recovery. I also traced compilation failures, GPU crashes, and incorrect outputs to precision, tensor-shape, and kernel implementation errors.


Projects

TinyGPT

I built a decoder-only Transformer language model from scratch in PyTorch to understand the full training stack, including causal multi-head attention, positional embeddings, autoregressive generation, supervised fine-tuning, checkpointing, and device-aware training.

I also implemented an end-to-end RLHF pipeline with a learned reward model, frozen reference policy, value critic, batched rollouts, GAE, token-level KL regularization, and clipped PPO updates. A 500-token BPE tokenizer compressed a 1.1-million-character corpus to 574,000 tokens (1.94×), and the evaluation tools help surface reward-model failures, reward hacking, and exploration behavior.