#21649 fix: TRT-LLM MHA CUDA illegal address with EAGLE v2 + DP attention

原始 PR 作者 Kangyan-Zhou 合并时间 2026-04-06 00:41 文件变更 1 提交数 4 评论 11 代码增减 +1 / -1

执行摘要

修复 TRT-LLM MHA 在 EAGLE v2 推测解码 +DP 注意力下因批次大小不一致导致的 CUDA 非法地址错误。

根据PR body描述，在Qwen3.5-397B模型上运行MMMU-Pro VLM评估时，启用DP注意力、EAGLE v2推测解码和TRT-LLM MHA后端后，约388-405/500个问题处会一致触发CUDA非法地址错误。根本原因是prepare_mlp_sync_batch为MLP同步将forward_batch.batch_size填充至DP组最大批次大小，但init_forward_metadata已基于原始批次大小计算元数据张量，导致TRT-LLM FMHA内核访问越界。

该PR值得精读，尤其关注：

1) DP注意力下批次大小不一致的根本原因分析；
2) 从forward_batch.batch_size到元数据推导的设计决策，体现了与其他后端行为对齐的架构一致性；
3) review中关于填充目的和注意力独立性的讨论，有助于理解分布式推理中的数据流设计。

讨论亮点

review中主要讨论了修复方案的正确性：Qiaolin-Yu质疑为何选择元数据而非forward_batch的填充状态；ispobock澄清填充仅用于MLP通信，注意力应基于真实批次大小独立计算；gemini-code-assist[bot]指出初始方案中CUDA图捕获路径缺失batch_size设置，可能导致相同错误。最终结论是采用元数据推导方案，既解决非法地址问题，又避免维护额外字段。

实现拆解

核心改动在python/sglang/srt/layers/attention/trtllm_mha_backend.py的forward_extend函数中，将batch_size参数从forward_batch.batch_size改为self.forward_metadata.cu_seqlens_q.shape[0] - 1，从而从元数据张量形状推导真实批次大小，与其他注意力后端（FlashInfer、Triton）行为保持一致。提交历史显示最初尝试在TRTLLMMHAMetadata中存储batch_size字段，但最终采用更简洁的推导方案。

文件	模块	状态	重要度
`python/sglang/srt/layers/attention/trtllm_mha_backend.py`	attention_backend	modified	10.0

关键符号

forward_extend init_forward_metadata prepare_mlp_sync_batch

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

修复方案正确性：为何使用元数据而非 forward_batch 的填充批次大小 正确性

Qiaolin-Yu 质疑当 DP 注意力引入填充时，forward_batch 应反映填充状态，为何代码选择使用不一致的元数据；ispobock 澄清填充仅用于 MLP 同步，注意力各 DP rank 独立计算，元数据张量应始终反映真实批次大小

结论：采用从 cu_seqlens_q 推导真实批次大小的方案，确保注意力计算与元数据边界一致 · 已解决

初始方案中 CUDA 图捕获路径的完整性 正确性

gemini-code-assist[bot] 指出初始方案在 init_forward_metadata_capture_cuda_graph 中未设置 metadata.batch_size，可能导致 CUDA 图上下文中相同错误

结论：最终方案放弃存储 batch_size 字段，直接推导，避免该问题 · 已解决

风险与影响

风险较低：

1) 变更仅影响TRT-LLM MHA后端在DP注意力+推测解码场景，其他后端或配置不受影响；
2) 从cu_seqlens_q推导批次大小与其他后端逻辑一致，降低不一致风险；
3) 已通过MMMU-Pro完整评估验证（1730/1730问题无崩溃）。潜在风险：若cu_seqlens_q张量形状异常，可能推导错误批次大小，但该张量由同一初始化逻辑生成，风险可控。

影响范围：

1) 用户：修复了多模态评估中的稳定性问题，确保启用DP注意力、EAGLE v2和TRT-LLM MHA后端的场景可稳定运行；
2) 系统：消除CUDA非法地址崩溃，提升系统可靠性；
3) 团队：明确了DP填充仅用于MLP通信、注意力应基于真实批次大小的设计原则，为类似问题提供参考。影响程度中等，针对特定配置的崩溃修复。

核心路径变更分布式推理边界条件

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

执行摘要

一句话：修复TRT-LLM MHA在EAGLE v2推测解码+DP注意力下因批次大小不一致导致的CUDA非法地址错误。
推荐动作：该PR值得精读，尤其关注：
1) DP注意力下批次大小不一致的根本原因分析；
2) 从forward_batch.batch_size到元数据推导的设计决策，体现了与其他后端行为对齐的架构一致性；
3) review中关于填充目的和注意力独立性的讨论，有助于理解分布式推理中的数据流设计。

功能与动机

实现拆解

关键文件：

python/sglang/srt/layers/attention/trtllm_mha_backend.py（模块 attention_backend）: 唯一修改文件，包含forward_extend函数的关键修复，将batch_size参数从forward_batch.batch_size改为从cu_seqlens_q推导

关键符号：forward_extend, init_forward_metadata, prepare_mlp_sync_batch

评论区精华

修复方案正确性：为何使用元数据而非forward_batch的填充批次大小 (correctness): 采用从cu_seqlens_q推导真实批次大小的方案，确保注意力计算与元数据边界一致
初始方案中CUDA图捕获路径的完整性 (correctness): 最终方案放弃存储batch_size字段，直接推导，避免该问题

风险与影响

风险：风险较低：
1) 变更仅影响TRT-LLM MHA后端在DP注意力+推测解码场景，其他后端或配置不受影响；
2) 从cu_seqlens_q推导批次大小与其他后端逻辑一致，降低不一致风险；
3) 已通过MMMU-Pro完整评估验证（1730/1730问题无崩溃）。潜在风险：若cu_seqlens_q张量形状异常，可能推导错误批次大小，但该张量由同一初始化逻辑生成，风险可控。
影响：影响范围：
1) 用户：修复了多模态评估中的稳定性问题，确保启用DP注意力、EAGLE v2和TRT-LLM MHA后端的场景可稳定运行；
2) 系统：消除CUDA非法地址崩溃，提升系统可靠性；
3) 团队：明确了DP填充仅用于MLP通信、注意力应基于真实批次大小的设计原则，为类似问题提供参考。影响程度中等，针对特定配置的崩溃修复。
- 风险标记：核心路径变更, 分布式推理边界条件

关联脉络

PR #22146 Isolate spec V1 path in decode post-processing: 同涉及推测解码（speculative decoding）路径的修改，本PR修复EAGLE v2下TRT-LLM MHA问题，22146隔离Spec V1后处理路径，显示推测解码模块的持续演进
PR #22148 Unify think_end_id to model_config as single source of truth: 同属一致性（consistency）改进，本PR统一批次大小推导逻辑与其他后端一致，22148统一think_end_id存储，体现代码库消除冗余、提升一致性的趋势
PR #22104 [SpecV2]: Reopen kl accuracy test for qwen3 + SpecV2: 同涉及推测解码测试，本PR修复实际运行问题，22104重新启用SpecV2测试，反映推测解码功能的测试与修复并行推进

#21649 fix: TRT-LLM MHA CUDA illegal address with EAGLE v2 + DP attention

执行摘要

修复 TRT-LLM MHA 在 EAGLE v2 推测解码 +DP 注意力下因批次大小不一致导致的 CUDA 非法地址错误。

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论