Prhub

#24319 [AMD] fix tbo specv2 seq_lens_cpu NoneType error

原始 PR 作者 billishyahao 合并时间 2026-05-05 16:54 文件变更 1 提交数 2 评论 4 代码增减 +5 / -0

执行摘要

修复 EAGLE SpecV2 + TBO 下 seq_lens_cpu 空指针问题

CI 任务中 TestMTPwithTBOLowLatency 测试因 AttributeError: 'NoneType' object has no attribute 'shape' 崩溃,根因是 EagleVerifyInputseq_lens_cpuNone。关联 issue #24212 追踪此兼容性问题,并参考了 PR #24205 的 workaround。

值得快速合入,修复明确,影响范围小。可作为 AMD 平台 SpecV2 兼容性修复的参考模式。

讨论亮点

Review 中 gemini-code-assist[bot] 提出了代码风格建议:将新增的两行赋值移到已有的 if not batch.forward_mode.is_idle(): 条件块内部,以避免冗余的条件判断。作者在第二次提交中采纳了该建议。HaiShaw 和 hubertlu-tw 均批准了 PR。

实现拆解

仅在 python/sglang/srt/speculative/eagle_info_v2.pyprepare_for_v2_verify 方法中,在 ForwardBatch.init_new 调用之前,增加了两行赋值语句:

  1. self.seq_lens_cpu = batch.seq_lens_cpu
  2. self.seq_lens_sum = batch.seq_lens_sum

这些赋值被放置在一个 if not batch.forward_mode.is_idle(): 条件块内,与之前的 draft token 处理逻辑条件一致,确保仅在非空闲模式下填充。该修复同时覆盖了 multi_layer_eagle_worker_v2.py 中的路径,因为它们共用 EagleVerifyInputV2Mixin

文件 模块 状态 重要度
python/sglang/srt/speculative/eagle_info_v2.py 推测解码 modified 5.53

关键符号

prepare_for_v2_verify

关键源码片段

python/sglang/srt/speculative/eagle_info_v2.py core-logic

唯一修改的文件,在 `prepare_for_v2_verify` 方法中填充 `seq_lens_cpu`/`seq_lens_sum` 以修复 TBO 路径下的空指针崩溃。

# python/sglang/srt/speculative/eagle_info_v2.py
# 在 prepare_for_v2_verify 中,于 ForwardBatch.init_new 之前填充
if not batch.forward_mode.is_idle():
    # ... 之前已有的 draft token 处理逻辑 ...
    batch.mamba_track_mask = None
    batch.mamba_track_seqlens = None
​
    # ---- 新增:从 batch 复制 seq_lens_cpu / seq_lens_sum 到 self ----
    # 保证 TBO 的 split_spec_info 在切片 custom_mask 时能读取到有效值
    self.seq_lens_cpu = batch.seq_lens_cpu
    self.seq_lens_sum = batch.seq_lens_sum# 后续创建 ForwardBatch 并返回
batch.forward_mode = (
    ForwardMode.IDLE
    if batch.forward_mode.is_idle()
    else ForwardMode.TARGET_VERIFY
)
batch.capture_hidden_mode = CaptureHiddenMode.FULL
verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
# ... 后续返回 verify_forward_batch, can_run_cuda_graph

评论区精华

将新增赋值移入已有条件块 style

gemini-code-assist[bot] 建议将新增的两行赋值移到已有的 `if not batch.forward_mode.is_idle():` 块内部,以避免冗余的条件判断。

结论:作者在 commit `5553c474` 中采纳了该建议。 · 已解决

风险与影响

风险极低:变更仅在两处条件分支内增加了 GPU 张量引用赋值,不涉及新依赖或计算逻辑,且 CI 测试 TestMTPwithTBOLowLatency 已通过。回归风险主要在于 TBO 未启用时该赋值不会执行,无影响。

仅影响 AMD 平台启用 EAGLE SpecV2 且开启 TBO 优化时的工作流,修复了 SpecV2 默认启用后的崩溃问题。对非 AMD 或非 SpecV2 场景无影响。

核心路径变更

关联 Issue

#24212 [Feature][AMD] Address the specv2 + moriep tbo compatibility

完整报告

参与讨论