Prhub

#21019 [Qwen3.5] Fuse split/reshape/cat ops in GDN projection with Triton kernel

原始 PR 作者 yuan-luo 合并时间 2026-03-23 23:17 文件变更 3 提交数 3 评论 22 代码增减 +597 / -202

执行摘要

优化 Qwen3.5 GDN 投影层,通过 Triton 核融合操作减少内核启动和内存分配。

PR body中说明:'In PR https://github.com/sgl-project/sglang/pull/19321 we fused Qwen3-Next GDN's qkvz_proj and ba_proj. This PR is a follow up.' 由于Qwen3.5检查点布局与Qwen3-Next不同,需要新的融合方案以减少内核启动和中间张量分配,提升性能。

该PR值得精读,特别关注Triton核的设计实现、权重加载器的兼容性处理,以及性能优化策略。工程师可从中学习内核融合技术和量化模型适配方法。

讨论亮点

review评论中核心讨论:

  1. jasperjiaguo报告Qwen3.5-0.8B小模型性能下降约4.5%,但作者yuan-luo验证后显示性能提升5.1%,引发对性能影响的争议。
  2. yuan-luo修复FP8量化兼容性问题,因权重加载器在FP8模型上设置只读属性导致错误。
  3. cs-cat在issue评论中提到可能影响模型准确性,需参考#21696进行验证。

实现拆解

  1. 新增Triton核文件 python/sglang/jit_kernel/triton/gdn_fused_proj.py,引入 fused_qkvzba_split_reshape_cat_contiguous 核函数。
  2. 修改 python/sglang/srt/models/qwen3_5.py,重构投影层为 in_proj_qkvzin_proj_ba,实现 _make_packed_weight_loader 处理权重加载。
  3. 修改 python/sglang/srt/models/qwen3_next.py,移除冗余Triton核,复用新增核以保持代码一致性。
文件 模块 状态 重要度
python/sglang/jit_kernel/triton/gdn_fused_proj.py jit_kernel added 8.0
python/sglang/srt/models/qwen3_5.py models modified 8.0
python/sglang/srt/models/qwen3_next.py models modified 5.0

关键符号

fused_qkvzba_split_reshape_cat_kernel fused_qkvzba_split_reshape_cat_contiguous _make_packed_weight_loader create_qkvz_proj create_ba_proj

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

评论区精华

小模型性能回归 性能

jasperjiaguo 报告 Qwen3.5-0.8B 模型在 H200 上性能下降 4.5%,但 yuan-luo 验证后显示提升 5.1%,双方讨论性能影响。

结论:作者验证后显示性能改进,但需持续监控小模型场景。 · 已解决

FP8 量化兼容性问题 正确性

yuan-luo 发现 FP8 量化模型因权重加载器属性设置错误导致启动失败,随后修复了权重加载器逻辑。

结论:已修复,确保 FP8 模型兼容性。 · 已解决

准确性影响 正确性

cs-cat 在 issue 评论中提到变更可能影响模型准确性,需参考 #21696 进行验证。

结论:未在 PR 中完全解决,需要额外测试。 · 待处理

风险与影响

技术风险包括:

  1. 性能风险:小模型上可能出现性能回归,需进一步验证和优化。
  2. 兼容性风险:FP8量化模型需要特殊处理权重加载器,否则可能导致初始化错误。
  3. 准确性风险:变更可能影响模型输出准确性,需加强测试覆盖。

影响范围:

  1. 对用户:在大模型上提升推理性能,但小模型需监控性能波动。
  2. 对系统:减少内核启动次数和内存分配,优化资源使用。
  3. 对团队:代码复杂度增加,但提供了统一的融合方案,便于后续维护和扩展。
小模型性能风险 FP8 兼容性问题 准确性验证不足

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论