执行摘要

• 一句话：修复 AMD gfx950 上的 Triton 编译断言错误，使用标量分支替换指针级 tl.where。
• 推荐动作：该 PR 值得精读，特别是对于关注 AMD Triton 兼容性或 JIT kernel 优化的工程师。关键设计决策包括：如何在不增加加载次数的前提下避免指针级选择，以及利用标量均匀性消除分支成本。建议结合历史 PR 如 #21691 和 #20974，了解跨硬件的性能修复模式。

功能与动机

根据 PR body，AMD Triton 的 TritonAMDGPUCanonicalizePointers pass 在 arith.select 用于指针张量时触发断言（ConvertArithSelectOp::matchAndRewrite_），导致 RuntimeError: PassManager::run failed，使得 CI 测试 test_diffusion_generation[qwen_image_edit_2511_ti2i] 在 gfx950 (MI350X) 硬件上失败。

实现拆解

变更集中在文件 python/sglang/jit_kernel/diffusion/triton/scale_shift.py。修改了两个 kernel 函数：_fused_layernorm_scale_shift_gate_select01_kernel 和 _fused_residual_layernorm_scale_shift_gate_select01_kernel。关键改动是将 scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs) 等三行替换为 if idx: 分支结构，直接加载对应的指针。注释解释了避免指针级 tl.where 以绕过 AMD Triton 编译器 bug，同时保持加载次数不变且 idx 为标量均匀，无线程分歧。

关键文件：

• python/sglang/jit_kernel/diffusion/triton/scale_shift.py（模块 jit_kernel/diffusion）: 修复了 AMD Triton 编译断言，修改了两个关键 kernel 函数，是扩散模型 JIT 编译的核心部分。

关键符号：_fused_layernorm_scale_shift_gate_select01_kernel, _fused_residual_layernorm_scale_shift_gate_select01_kernel

评论区精华

review 过程中无具体讨论评论，reviewer gemini-code-assist[bot] 指出 'no feedback'，yctseng0211 和 HaiShaw 快速批准。Issue 评论中，bingxche 请求 review 和 CI 状态检查，amd-bot 回复 CI 状态显示可能相关的错误，但修复应有助于而非损害。整体讨论聚焦于 CI 验证和修复有效性，无争议点。

• CI 测试验证与修复有效性 (testing): 修复被确认有效，CI 测试应通过。

风险与影响

• 风险：风险较低：变更逻辑简单，直接修复编译错误。潜在风险包括：
  1) 可能引入新的编译问题在其他硬件或 Triton 版本，但作者声明无性能影响且通过 CI 测试；
  2) 对非 AMD 硬件的兼容性未显式测试，但基于标量分支的代码在 NVIDIA GPU 上应无问题；
  3) 文件 scale_shift.py 是扩散模型核心 JIT kernel，修改需确保准确性，但变更仅限于指针选择逻辑。
  • 影响：影响范围较小：主要影响 AMD gfx950 (MI350X) GPU 上的扩散模型生成测试，修复了 CI 失败。对用户而言，确保 SGLang 在 AMD 硬件上的稳定运行；对系统性能无负面影响，因为分支无发散成本且加载次数不变。影响程度为中等，限于特定硬件和模块。
  • 风险标记：硬件特定依赖, 编译兼容性风险

关联脉络

• PR #21691 [AMD] fix performance regression issue when run gpt-oss with "--context-length 13824": 同为 AMD 相关的 bugfix，涉及性能优化，展示跨硬件兼容性问题。
• PR #20974 [NPU][Diffusion] fix sp modulate for qwen-image-edit: 同为扩散模型的 bugfix，涉及 JIT kernel 修改，展示类似修复模式。
• PR #21383 [diffusion] [NPU] support ring attention on NPU with FA: 扩散模型相关 PR，涉及新功能支持，与本 PR 的模块相同。