Prhub

#29576 Fix DSA indexer fusion bug causing excessive memory consumption.

原始 PR 作者 b8zhong 合并时间 2026-06-29 05:11 文件变更 1 提交数 1 评论 4 代码增减 +11 / -6

执行摘要

修复 DSA indexer 融合中 rope buffer 每层重复创建导致内存暴增

由 PR #29564 发现 bug:DSA indexer 融合中 rope buffer 按每层 x 上下文长度创建,导致显存暴增。PR body 量化了修复前的严重后果:TP=4 的 GLM-5.2 模型推理容量从正常的 ~2,400k tok/pool 下降至 ~2,020k tok/pool,说明融合 bug 显著降低了吞吐量。

值得精读:这是一个典型的高内存泄漏修复案例,展示了将重复计算提取为跨实例共享缓存的模式。代码简洁,逻辑清晰,对理解 DSA indexer 的内存管理有参考价值。

讨论亮点

讨论简洁,Fridge003 触发 e2e 回归测试(test_dsa_glm5_dp_mtptest_dsa_glm5_tp_mtp)和内核单元测试(test_dsa_indexer)。mmangkad 回复“Nice catch”确认 bug 发现。审查者 Fridge003 直接 approve,无实质性争议。

实现拆解

  1. 提取共享函数:在 dsa_indexer.py 新增 _shared_indexer_freqs_cis(rotary_emb) 函数,该函数使用 getattr+setattrfreqs_cis 缓存到 rotary_emb 模块的自定义属性 _dsa_indexer_freqs_cis 上,确保同一 rotary_emb 实例只构建一次。
  2. 修改 Indexer.init:移除原来每层独立构建 cos_sin_cachefreqs_cis 的 6 行内联代码,改为调用新的共享函数,并删除了原注释行。
  3. 仅 1 个文件变更python/sglang/srt/layers/attention/dsa/dsa_indexer.py,+11/-6,无其他文件或测试配套改动。
文件 模块 状态 重要度
python/sglang/srt/layers/attention/dsa/dsa_indexer.py DSA 注意力 modified 6.77

关键符号

_shared_indexer_freqs_cis

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

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险极低:变更仅限于提取一个纯函数并替换内联代码,逻辑等价性由缓存机制保证——首次调用与旧行为完全一致,后续复用缓存。唯一风险是如果不同 Indexer 实例希望持有不同的 freqs_cis(例如不同 layer_id 影响 rope),但根据设计所有层共享相同配置,此风险不存在。回归测试已覆盖 GLM-5.2 模型的 DP 和 TP 场景。

正向影响:修复后每 GPU 显存释放约 0.25 GB(具体取决于模型隐藏层数),在 TP=4 的 GLM-5.2 上推理容量提升约 18%。零负面影响;无 API 变动,无配置更改。

核心路径变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论