feat: support adaptive speculative decoding.#1876
Conversation
weizhehuang0827
requested review from
Clement-Wang26,
DongheJin,
DragonFive,
JimHsiung,
Kang-Meng,
RobbieLeung,
XuZhang99,
liujinguang0125,
liutongxuan,
walsonyang,
xiao-yu-chen,
yingxudeng,
yq33victor and
zhang-minchao
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Adaptive Speculative Decoding for MTP, introducing an AdaptiveSpeculativeController to dynamically prune draft tokens before target validation based on path probabilities and a validate-time predictor. It also adds a SpeculativeProfileRegistry and integrates profiling within the ProfileManager. The review feedback is highly constructive, identifying an opportunity to avoid wasteful GPU kernel launches on the hot path when destination indices are empty, and recommending the cleanup of the unused target_validate_time_by_sl parameter across several controller and worker files.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
weizhehuang0827
force-pushed
the
adaptive-speculatuve
branch
2 times, most recently
from
a50c467 to
7973a32
Compare
…ning. Implement adaptive speculative decode that prunes low-confidence draft tokens before target validation, reducing validate compute via ATB variable-length q_len support. All ranks compute pruning independently (no cross-rank broadcast needed) since draft probs are deterministic.
weizhehuang0827
force-pushed
the
adaptive-speculatuve
branch
from
7973a32 to
30565e1
Compare
Move AdaptiveSpeculativeController and SpeculativeProfileRegistry from runtime/ to framework/speculative/ to keep runtime/ focused on worker/executor concerns. Update all include paths and CMake targets.
There was a problem hiding this comment.
remove and add design docs to xllm-docs.
2 participants
Description
实现自适应投机推理的 validate 前剪枝优化。在每个 decode step 的 draft 阶段结束后、validate 阶段开始前，根据 draft model 产出的 token 概率动态决定每个 sequence 保留多少 speculative tokens 参与 validate，裁剪掉置信度低、预期收益为负的 draft tokens。
核心算法：将 batch 内所有请求所有位置的候选 token 按 path probability（从第一个 token 到当前位置的累积概率）从大到小排列，从最高概率开始依次尝试加入验证集。每加入一个 token，基于 profile 拟合的时间预测模型更新估计吞吐（expected_emitted_tokens / estimated_time）。若加入后估计吞吐不增反降则停止加入该 token。最终每个 sequence 可能有不同的验证长度。
validate 阶段利用 ATB attention kernel 的 variable q_len 能力，只对每个 sequence 实际保留的 tokens 做计算，真正减少 validate 的 attention 和 FFN 计算量。
此方案参考 DSPark 自适应投机推理算法（区别主要在于计算接受概率时DSpark用的是模型本身的置信度头，我们这里先使用draft模型输出的logits），但还未专门针对graph和ZOS做适配。后续还可以进一步支持:(1)对draft的early stop;(2)根据近期各位置接受率动态调整num_speculative_tokens
Related Issues
N/A
Change Type
Pull Request Checklist
PR Title and Commit Messages
<type>: <subject>.Pre-commit Checks
pre-commitby runningpip install pre-commitor an equivalent command.pre-commit install.pre-commit run --all-filesand fixed any reported issues.Self Review
.agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.mainbranch.Build and Test Coverage
python setup.py build testhas passed on a CUDA machine.python setup.py build testhas passed on an NPU machine.python setup.py build testhas passed on an MLU machine.Reviewer Notes
测试数据（ShareGPT + DeepSeek-V3.2 MTP，TP=16 NPU， batch_size=32）：
SL=5/7 时 adaptive 显著优于 static（吞吐 +18
30%，acceptance_rate +5073%）。SL=3 时裁剪空间小，两者持平。关闭enable_adaptive_speculative_decode后完全不影响 static 路径。