A causal parallel draft head over fused hidden states, combining one-forward
drafting with branch-wise causal conditioning — up to
9.64× speedup on math and
4.58× on chat with Qwen3.
A concept-level memory that lets an LLM agent distill reasoning traces into
reusable abstractions and keep learning at test time without weight updates.
Hao AI Lab, UC San DiegoAdvised by Hao Zhang and Lanxiang Hu · Sep 2025–Present
Led JetSpec++, a causal correction for parallel drafting that
conditions each drafted token on its full branch history rather than on the preceding
token alone. At matched parameter count it outperformed DeepSeek’s DSpark, with
the margin widening at draft depths beyond those used in training.
Contributed to JetSpec, a draft head that generates a complete
speculative tree in one forward pass while preserving the target model’s
autoregressive factorization along every branch; built its standalone inference engine
(paged KV cache, Triton tree‑attention kernel, CUDA‑graph verification) and
its evaluation suite. Reaches up to 9.64×
end‑to‑end speedup over autoregressive decoding on Qwen3‑8B
(MATH‑500).
Ported DFlash speculative decoding to TPU in JAX as the primary
contributor, showing that verification cost stays flat in draft block size through
K=1024. Averaged 3.13× speedup
on TPU v5p; the draft model and proposer are merged into vLLM’s TPU backend,
with pipeline integration in review, and the work is featured on the Google Developers
Blog.
Built the training stack for a 1.8B‑parameter language
model on eight NVIDIA B200 GPUs, implementing the full
pre‑training and SFT and DPO post‑training system with DDP and
ZeRO‑1.
Q‑Lab, UC San DiegoAdvised by Lianhui Qin and Matthew Ho · Jan 2025–Jun 2026
Contributed to the design of ArcMemo’s program‑synthesis
memory ontology, a typed, parameterized representation in which reusable reasoning
concepts are stored and composed.
Worked on the reasoning‑guided concept selection that replaced
embedding‑based retrieval, after embedding search proved ineffective on the
ARC‑AGI‑1 abstract‑reasoning benchmark.
Built the pipeline that produces concept‑labeled helper puzzles: each
hand‑written concept is expanded by an LLM into puzzle‑generating code,
which is executed and tested so that every generated puzzle is verified to exercise
its concept.
Designed the framework’s adaptation to AIME competition mathematics and ran the
generalization experiments and memory‑format comparisons that informed the final
design.
