Tsinghua University · Computer Science

Huajun Bai 白华君

Ph.D. student · AI systems researcher

Efficient inference for
language models & agents.

I work on speculative decoding, efficient reasoning, and agentic inference—connecting algorithm design with the systems that run it.

At Tsinghua, I am advised by Prof. Jiwu Shu in the Storage Research Group.

Explore my work Get in touch

Google Scholar ↗GitHub ↗LinkedIn ↗

Tsinghua UniversityBeijing, China

Selected workSPORKarXiv ’26 SpecExitICML ’26 PARDICLR ’26

01 / Research

A systems view of inference.

From individual tokens to multi-step agents, I study where inference spends time and how to reduce that cost. Three complementary lines of work connect this agenda.

01
Token generation
Draft more efficiently
Parallel drafting with a reusable, low-cost draft model.
PARD
02
Model reasoning
Avoid redundant reasoning
Use the draft model’s own signals to guide when reasoning can stop.
SpecExit
03
Agent execution
Overlap reasoning and tools
Anticipate tool calls so useful work can begin before reasoning ends.
SPORK

02 / Selected publications

Ideas, implementations, evidence.

Google Scholar ↗

SPORK

Preprint · 2026

First author

Self-Speculative Forking to Accelerate Agentic LLM Inference

Huajun Bai, Weiwei Lv, Huichuan Zheng, Youyou Lu, Jiwu Shu

A training-free controller that anticipates an agent’s tool call and starts execution while reasoning continues. Shared prefixes keep speculation lightweight; verification determines whether to reuse the result or fall back.

Paper Code

18%lower P95 latency

Qwen3-32B · GAIA
131.9 → 108.1 seconds

SpecExit

ICML 2026

Co-first author

Accelerating Large Reasoning Model via Speculative Exit

Rubing Yang^*, Huajun Bai^*, Song Liu, et al.

Uses draft-model hidden states to predict an early-exit signal alongside future tokens, reducing redundant reasoning without a separate probe. The evaluation measures the resulting latency–accuracy tradeoff.

Paper Code

2.5×end-to-end speedup, up to

Relative to the speculative
decoding baseline

PARD

ICLR 2026

Co-first author

Accelerating LLM Inference with Low-Cost PARallel Draft Model Adaptation

Zihao An^*, Huajun Bai^*, Ziqiong Liu, Dong Li, Emad Barsoum

Adapts a small model to draft multiple future tokens in parallel. A target-independent draft can be reused across a family of target models, reducing the cost of adaptation and token generation.

Paper Code

264.88tokens / second

LLaMA3.1-8B · HumanEval
vLLM · A100-40GB

Results are reported in the linked papers under the configurations shown; they are not cross-paper comparisons. * Equal contribution.

More publications 2019–2025

MobiCom 2025
EdgeFM Workshop
Hierarchical N-Gram speculative decoding for AI PCs First author
NeurIPS 2023
FACE: Evaluating Natural Language Generation with Fourier Analysis of Cross-Entropy
WWW 2023
Expressive User Embedding from Churn and Recommendation Multi-Task Learning First author
ACL 2019
A Corpus for Reasoning about Natural Language Grounded in Photographs

03 / Education

Academic background

2024–Present
Tsinghua University
Ph.D. student, Computer Science and Technology
2017–2018
Cornell Tech
M.Eng., Computer Science
2013–2017
Cornell University
B.A., Mathematics and Computer Science

04 / Beyond papers

MLSYS FeedFollowing the field, out loud.

I host a Chinese-language podcast on machine learning systems. An agentic pipeline helps retrieve, rank, and narrate research papers—a way to explore ideas beyond my own projects.

Browse episodes ↗Talk: AI agents for research ↗