Prhub
← 返回仓库详情

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

sgl-project/sglang · 作者 yuan-luo · 合并时间 2026-03-23 23:17

分析状态 已生成
文件变更 3提交数 3 · 评论 22
代码增减 +597 / -202
jit-kernel performance refactor quant

执行摘要

优化 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_qkvz 和 in_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

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

关键符号

fused_qkvzba_split_reshape_cat_kernel fused_qkvzba_split_reshape_cat_contiguous _make_packed_weight_loader create_qkvz_proj create_ba_proj

评论区精华

小模型性能回归 性能

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 链接，后续同步到相关引用后会出现在这里。

完整报告

Markdown PDF

执行摘要

本PR优化了Qwen3.5模型的Gated Delta Net投影层，通过引入Triton核融合split/reshape/cat操作，减少内核启动和内存分配，提升推理性能。尽管在小模型上存在性能讨论，但整体改进显著，需注意FP8量化兼容性和准确性验证。

功能与动机

作为PR #19321的后续，针对Qwen3.5检查点布局与Qwen3-Next不同的特点，将多个投影操作融合以优化性能。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.' 目标是减少内核启动和中间张量分配，提升预填充和解码阶段的效率。

实现拆解

  • 新增Triton核：在 python/sglang/jit_kernel/triton/gdn_fused_proj.py 中新增 fused_qkvzba_split_reshape_cat_contiguous 核函数，处理连续输入格式。
  • 重构投影层：修改 python/sglang/srt/models/qwen3_5.py，将原先的 in_proj_qkv、in_proj_z、in_proj_b、in_proj_a 合并为 in_proj_qkvz 和 in_proj_ba 两个融合层。
  • 增强权重加载器：实现 _make_packed_weight_loader 方法，支持融合和拆分检查点格式的权重加载，确保参数初始化正确。
  • 清理冗余代码：修改 python/sglang/srt/models/qwen3_next.py，移除旧Triton核，复用新核以保持代码一致性。

评论区精华

  • 性能争议：jasperjiaguo报告小模型性能下降，但作者yuan-luo验证后显示改进，引发对优化效果的讨论。引用yuan-luo：'I'll verify the small model's performance and do refactor to avoid the corresponding impact.'
  • FP8修复：yuan-luo在issue评论中修复FP8量化兼容性问题：'FP8 problem fixed.' 强调权重加载器需处理不同参数类型。
  • 准确性检查：cs-cat提到可能影响准确性：'This PR does indeed bring significant performance improvements, but it seems to affect the accuracy of the model? Please refer to #21696.' 需要进一步验证。

风险与影响

  • 技术风险：小模型性能可能出现波动；FP8量化模型需特殊处理权重加载器；变更可能引入准确性偏差，需加强测试。
  • 影响范围：对用户而言，大模型推理性能提升，但小模型需监控；系统层面减少内核启动和内存开销；团队需适应新的融合架构，并关注兼容性测试。

关联脉络

  • 与PR #19321直接关联，延续了投影层融合的优化路线，展示了跨模型版本的一致性优化策略。
  • 涉及jit-kernel和量化主题，与仓库中其他性能优化PR（如#21657的AMD优化）共享技术思路，反映了团队在核融合和量化适配上的持续演进。

参与讨论