[Bug]: GLM-5.x (DSA) under pipeline parallelism: reuse layers at a stage start read stale top-k indices

问题场景

使用 vllm serve 加载 GLM-5.x（如 zai-org/GLM-5.3 FP8，架构 GlmMoeDsaForCausalLM），并开启 pipeline parallelism（--pipeline-parallel-size 8 或任意 >1 的值），配合 --moe-backend deep_gemm、--no-enable-prefix-caching 等参数。典型复现还包括 2×8 RTX PRO 6000（SM120）上 PP2×TP8 跑 nvidia/GLM-5.3-NVFP4。

GLM-5.x 的 DSA indexer 只在部分层运行（GLM-5.3 为 0、1、2 层，之后从第 6 层起每 4 层一次，即 index_topk_freq=4、index_skip_topk_offset=3），其余层复用上一个 indexer 层算出的 top-k 索引。当 stage 边界恰好落在“indexer 层”和“复用其索引的层”之间时，问题触发。

报错原文

[Bug]: GLM-5.x (DSA) under pipeline parallelism: reuse layers at a stage start read stale top-k indices

注意：Issue 描述的并不是崩溃型异常，不会打印 Python traceback；表现为输出仍流畅但质量下降（复制任务出错、跳过条目、丢字符、循环到 token 上限等）。

原因分析

在 vLLM 0.30 的 PP 实现中，stage 之间只传递 IntermediateTensors({"hidden_states": ..., "residual": ...})。而 topk_indices_buffer 是每个 rank 自己的 torch.empty(...)，不会跨 stage 传递。

当某个 rank 的第一层不是 indexer 层（has_indexer=False）时，SparseMLAIndexGroupBuilder.register_layer 会为该层建组，但没有代码去从上一个 rank 取回索引。于是该层的 skip_topk 逻辑会读取本地 buffer 中残留的内容：要么是上一批自己最后一个 indexer 层算出的旧 top-k，要么在第一步时是未初始化显存。调度相关的 attention metadata（block table、slot mapping）在每个 rank 上都重建正确，缺的只是这些计算出来的索引。

按默认切分 [9,10,10,10,10,10,10,9]，PP8 下 78 层中有 13 层受影响（9；19-21；29；39-41；49；59-61；69）；PP2 的 [39,39] 下是 39-41 层。Issue 作者在接收 stage 上实测：每个真实 prefill batch 中 16,384/16,384 行（B200）与 8,192/8,192 行（H200）全部不同，只有 warm-up / CUDA graph capture 的 dummy batch 一致。

环境排查

• 确认 vLLM 版本：Issue 基于 0.30.0，其他版本代码路径可能不同。
• 确认 --pipeline-parallel-size 是否 > 1；PP=1 时不会走 stage 间 hand-off，不受影响。
• 确认模型架构为 GlmMoeDsaForCausalLM，即 GLM-5.x DSA 系列。
• 确认模型配置中的 index_topk_freq、index_skip_topk_offset，并对照当前 PP 切分，找出哪些 rank 的第一层是复用层。
• 确认注意力后端：B200 为 FLASHMLA_SPARSE，H200 为 FLASH_ATTN_MLA_SPARSE，SM120 + NVFP4 为 FLASHINFER_MLA_SPARSE_SM120。三者均已见复现报告。
• 确认 torch（2.13.0）、CUDA（13.0）、triton（3.7.1）、transformers（5.20.2）等依赖与运行环境是否与 Issue 接近；硬件不同不构成排除理由。

解决步骤

1. 先确认现象：用固定 prompt 的复制任务（如从约 250 token 的 prompt 中逐字复制 50 项列表）或 prompt log-likelihood 测试对比 PP 与 TP 输出。温度设为 0，复制必须完全一致才算通过；Issue 复现者 6 次中仅通过 2 次。输出流畅并不代表正常。
2. 如不需要 PP，可优先尝试 --pipeline-parallel-size 1（改为纯 TP8），以绕开 stage 间 hand-off。这是规避手段，不是修复。
3. 若必须使用 PP，可采用 Issue 中验证过的 out-of-tree 插件方案：在 DeepseekV32ForCausalLM.forward 中，非最后一个 stage 把 topk_indices_buffer[:n_tokens]（int32，[tokens, 2048]）追加到发出的 IntermediateTensors；非首个 stage 收到后、在第一层执行前把这些行拷贝进自己的 topk_indices_buffer；同时在 make_empty_intermediate_tensors 中加入该 key，保证常驻接收缓冲区带它。isend/irecv_tensor_dict 本身已支持额外的 int32 key。
4. 若所在分支已有相关修复（如 Issue 提到的 #59961 草案），需要先确认它已合入并包含 stage send/receive 路径，再作为正式方案使用。
5. 在不改代码的前提下，也可把 indexer 层与复用层尽量放在同一 stage 以避开边界；但这依赖具体切分，Issue 未给出通用切分方案，属于可尝试方向。

验证方法

• 在接收 stage（如 PP8 的 stage 1-7）上，对比进入 forward 时本地 buffer 与从上一 stage 收到的索引，真实 prefill batch 的行应一致，而不是像 stock 版那样整批不同。
• 做 prompt log-likelihood 对比（如 8 段 WikiText × 6,144 token，echo + logprobs），PP 修正后均值 NLL 应接近 TP8 参考值：B200 上约 0.817 对 0.819，H200 上约 0.816 对 0.817；stock PP8 明显更差。
• 跑贪心解码对比 PP 与 TP8 的最长公共前缀（修正后可达 128-161 字符，stock 仅 1-35 字符）。
• 确认吞吐无明显回退：转发开销约 8 KB/token，相对 hidden_states+residual 的 24 KB 较小。

参考来源

vllm-project/vllm #63170

AI 工具推荐

想把多个 AI 模型放在一个入口？

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型，适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接：通过此链接购买，我可能获得佣金，不影响你的价格。