fix(mtp): preserve speculative groups in TP-SP batch alignment - #1587
Merged
Merged
Conversation
shihaobai
approved these changes
Sep 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
固定布局 MTP 与 TP/SP 同时启用时,只按 TP 对齐 decode batch 会破坏验证分组。例如 TP=8、
mtp_step=2时,3 行补成 8 行,无法满足固定布局的整除约束。本改动按 MTP 基础步长与 TP 的最小公倍数对齐普通 decode、overlap decode、CUDA Graph 档位及模型内部 graph 容量,上述情况补成 24 行。复用
MtpManager的每请求 token 数和基础对齐粒度:动态主模型容量仍按最大 verify 宽度预留,对齐粒度为 1;草稿模型沿用自身布局,并保留 chained/recurrent 模式常量及显式分支。graph 步长仅保留动态主模型特例,其余复用每请求 token 数。模型只增加共用的 batch 对齐函数,无额外布局初始化函数或缓存字段。保留下游断言和调度器请求上限。验证:
python -m pytest unit_tests/common/basemodel/test_model_output.py unit_tests/common/basemodel/test_cuda_graph_layout.py unit_tests/common/basemodel/test_cuda_graph_autotune.py unit_tests/common/basemodel/test_mtp_manager.py unit_tests/common/basemodel/test_overlap_utils.py unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py -q:179 passed,在 H200 环境执行,包含真实 CUDA 算子测试。mtp_manager的构造遗漏;真实 CUDA 上验证动态压缩及 FA3 元数据算子。git diff --check通过。decode 调用链测试使用替代 forward/graph 对象,不能替代真实 attention 或 CUDA Graph capture/replay 验证。尚未运行真实多卡模型及性能评测;检查时 8 张 H200 均处于高负载,未启动多卡服务或压测。TP=8、verify 宽度=3 时最小 batch 从 8 行增加到 24 行,实际延迟及显存影响待测。