Skip to content

fix(mtp): preserve speculative groups in TP-SP batch alignment - #1587

Merged
shihaobai merged 11 commits into
ModelTC:mainfrom
sufubao:fix-mtp-tpsp-alignment
Sep 22, 2026
Merged

shihaobai merged 11 commits into
ModelTC:mainfrom
sufubao:fix-mtp-tpsp-alignment

Conversation

@sufubao

@sufubao sufubao commented Sep 21, 2026

Copy link
Copy Markdown
Collaborator

固定布局 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 算子测试。
  • 容量测试执行真实构造函数,覆盖 overlap 折半、MTP 倍率与向上对齐,在加载模型配置和权重前停止。
  • 覆盖固定/动态验证、TP=1/2/4/8、草稿模式、空 batch、graph 边界和 eager 回退;动态 decode 使用每请求 1/2/3 个 token 的变长输入;补充动态配置下 DSpark/DFlash 草稿仍保留块宽度的回归。
  • 修复已有 overlap 测试缺少 mtp_manager 的构造遗漏;真实 CUDA 上验证动态压缩及 FA3 元数据算子。
  • Black、flake8 和 git diff --check 通过。

decode 调用链测试使用替代 forward/graph 对象,不能替代真实 attention 或 CUDA Graph capture/replay 验证。尚未运行真实多卡模型及性能评测;检查时 8 张 H200 均处于高负载,未启动多卡服务或压测。TP=8、verify 宽度=3 时最小 batch 从 8 行增加到 24 行,实际延迟及显存影响待测。

@shihaobai
shihaobai merged commit 49acd05 into ModelTC:main Sep 22, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants