执行摘要
本 PR 在 BlackWell 架构上添加了 DeepSeekV3 router gemm 操作的基准测试脚本，并集成了 flashinfer 内核以优化性能。通过准确性验证和性能对比，显示两内核性能基本持平，决策默认启用 PDL。影响包括潜在性能提升和硬件兼容性考虑，为未来优化提供数据支持。
功能与动机
动机源于 Issue #14453，旨在将 flashinfer 的优化集成到 DeepSeekV3 模型中，以提升推理效率。PR body 明确指出：“Comparing dsv3_router_gemm performance between flashinfer and the current sglang kernel to optimize performance”，这反映了团队对性能调优的持续关注。
实现拆解
实现分为两个关键部分：
-
基准测试脚本 (benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py)
- 新增脚本，支持在多种配置（m=1-16, n=256, k=7168, tp=1-8）下运行准确性测试和性能基准。
- 使用 triton.testing.do_bench 测量执行时间，对比 sglang 内核和 flashinfer 内核。
-
模型代码集成 (python/sglang/srt/models/deepseek_v2.py)
- 修改 DeepSeekV2DecoderLayer.forward 方法，在条件满足时（SM>=100 且 weight.shape[0]==256）使用 flashinfer 的 mm_M1_16_K7168_N256 内核。
- 添加自定义操作 flashinfer_dsv3_router_gemm，通过 register_custom_op 注册，确保兼容性和可维护性。
- 关键代码片段：
if _device_sm >= 100 and self.weight.shape[0] == 256:
logits = torch.empty(...)
flashinfer_dsv3_router_gemm(logits, hidden_states, self.weight)
else:
logits = dsv3_router_gemm(...)
评论区精华
Review 讨论中突出了几个关键交锋：
-
内核选择争议：Fridge003 质疑导入特定 flashinfer 内核是否合理，leejnau 回应：“actually I believe this is the correct kernel in that it is specific to dsv3 and optimized for that architecture.” 这确认了设计决策的正确性。
-
PDL 设置权衡：nv-yunzheq 和 leejnau 讨论 PDL 启用方式，最终决定默认启用 PDL。leejnau 指出：“Previous default was PDL on...”，Fridge003 总结：“We can turn it on by default and optionally turn it off with an extra environ”，避免了环境变量泛滥。
-
兼容性问题解决：b8zhong 报告 SM103 崩溃，nvpohanh 询问是否为 flashinfer bug，最终通过 PR #22134 解决，展示了团队对稳定性的重视。
风险与影响
-
技术风险：
- 兼容性：flashinfer 内核在 SM103 GPU 上初始失败，需额外修复。
- 性能：基准测试显示性能持平，但需监控实际部署中是否出现回归。
- 维护：新增条件逻辑和自定义操作可能增加代码复杂度，需确保测试覆盖。
-
影响分析：
- 用户：可能获得轻微性能提升，但需确保硬件支持。
- 系统：引入 flashinfer 依赖，可能影响部署流程和 CI 测试。
- 团队：提供了标准化基准工具，有助于未来性能优化和决策。
关联脉络
本 PR 是 Issue #14453 的一部分，该 Issue 列出了多项 flashinfer 优化任务。历史 PR 中，如 #22006 和 #22143 也涉及 DeepSeek 模型优化，显示了团队在该领域的持续投入。解决 SM103 问题的 PR #22134 直接关联，确保了本变更的广泛兼容性。整体上，这反映了 sglang 项目在集成第三方内核以提升性能方面的演进趋势。
参与讨论