[AMD] Fix CI multimodal-gen-test-1-gpu-amd for gen model

[yichiche]

Motivation

Fix Triton CanonicalizePointers assertion on AMD gfx950 by branching on scalar idx instead of tl.where on pointers.

AMD Triton's TritonAMDGPUCanonicalizePointers pass hits an assertion in ConvertArithSelectOp::matchAndRewrite_ when arith.select is used on pointer tensors whose two branches have different "narrowability" properties. The _fused_layernorm_scale_shift_gate_select01_kernel and _fused_residual_layernorm_scale_shift_gate_select01_kernel kernels use tl.where(idx, ptr1, ptr0) to conditionally select between two sets of scale/shift/gate pointers, which triggers this bug on gfx950 (MI350X).

Modifications

• python/sglang/jit_kernel/diffusion/triton/scale_shift.py: In both _fused_layernorm_scale_shift_gate_select01_kernel and _fused_residual_layernorm_scale_shift_gate_select01_kernel, replaced pointer-level tl.where with if/else branching on the scalar idx:
  • Before: tl.where(idx, scale1_ptrs, scale0_ptrs) then tl.load(scale_ptrs, ...)
  • After: if idx: scale = tl.load(scale1_ptrs, ...) else: scale = tl.load(scale0_ptrs, ...)
  • This keeps it at 3 loads (not 6), avoids the pointer-level tl.where entirely, and since idx is uniform across all threads the branch has no divergence cost.

Accuracy Tests

python -m pytest -xvs     "sglang/multimodal_gen/test/server/test_server_a.py::TestDiffusionServerOneGpu::test_diffusion_generation[qwen_image_edit_2511_ti2i]"

CI test test_diffusion_generation[qwen_image_edit_2511_ti2i] was failing with RuntimeError: PassManager::run failed due to the Triton compilation assertion. This fix resolves the compilation failure.

Speed Tests and Profiling

No performance impact. The change replaces a pointer-select + 3 loads with an if/else branch + 3 loads. The idx value is a scalar uniform across all threads in the block, so the branch has zero divergence cost.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request modifies the Triton kernels in scale_shift.py to replace tl.where pointer selection with explicit if/else branching. This change addresses an assertion error in the AMD Triton compiler's CanonicalizePointers pass on gfx950 hardware while ensuring that only three loads are performed without thread divergence. I have no feedback to provide as there were no review comments.

[bingxche]

@amd-bot review

[bingxche]

@bingxche requested a review

Claude Code Review

PR #21621: [AMD] Fix CI multimodal-gen-test-1-gpu-amd for gen model
Reviewed at 2026-03-30 05:35 UTC

Code Review: [AMD] Fix CI multimodal-gen-test-1-gpu-amd for gen model

1. Summary

This PR fixes a Triton compilation failure on AMD gfx950 (MI350X) by replacing tl.where(idx, ptr1, ptr0) pointer-level selects with if/else branching on the scalar idx in two Triton kernels. The AMD Triton CanonicalizePointers pass asserts when arith.select is used on pointer tensors with differing "narrowability" properties. The fix is applied identically to both _fused_layernorm_scale_shift_gate_select01_kernel and _fused_residual_layernorm_scale_shift_gate_select01_kernel.

2. Code Quality

Correctness

• The transformation is semantically correct. idx is a scalar (loaded from idx_ptr + batch_idx, a single element per kernel instance), so it is uniform across all threads in the block. The if/else branch produces identical results to the tl.where on pointers — both select between the {scale,shift,gate}0 and {scale,shift,gate}1 pointer sets based on idx.
• No edge cases missed. The idx value is used as a boolean (truthy/falsy), which matches the original tl.where semantics. If idx is 0, the else branch (set 0) is taken; any nonzero value takes the if branch (set 1), consistent with the kernel name suffix select01.

Code Style and Readability

• The comments are clear, well-written, and explain the why behind the change (AMD Triton bug, no divergence cost, same number of loads). This is excellent documentation practice.
• The duplicated comment blocks in both kernels are appropriate since the kernels are independent and the explanation should be co-located with each fix.
• Code structure is clean and easy to follow.

Minor Observation

• The gate variable is loaded without .to(tl.float32) in both branches, which is consistent with the original code. No issue here — just confirming the cast pattern was preserved faithfully.

3. Performance

No performance concerns. The analysis in the PR description is correct:

• idx is a scalar uniform across all threads in the block (it's a per-batch-element index, not per-column). Therefore, the branch has zero divergence cost on both NVIDIA and AMD GPUs.
• The change maintains exactly 3 loads per kernel invocation (not 6), same as before.
• In Triton, if/else on a compile-time-uniform scalar will compile to a simple branch, which is essentially free compared to the memory access latency of the loads themselves.
• On NVIDIA GPUs, the original tl.where on pointers compiles fine and is also efficient, so this change is neutral there.

4. Security

No security concerns. This is a Triton kernel-level change with no user-facing input handling implications.

5. Testing

• The PR mentions running the specific failing CI test (test_diffusion_generation[qwen_image_edit_2511_ti2i]), which validates the fix resolves the compilation failure.
• The existing test suite should cover functional correctness since the semantic behavior is unchanged.
• Suggestion: It would be ideal to confirm the test also passes on NVIDIA GPUs (or that CI covers this), to ensure the if/else pattern doesn't cause any Triton codegen issues on the NVIDIA backend. Given that if/else on scalars is a standard Triton pattern, this is low risk.

6. Suggestions

1. Consider a TODO/upstream reference: If there's an upstream AMD Triton issue tracking the CanonicalizePointers assertion bug, linking it in the comment would help future maintainers know when this workaround can be reverted.

# TODO: Revert to tl.where once upstream AMD Triton fixes
# CanonicalizePointers assertion on pointer-select.
# See: https://github.com/triton-lang/triton/issues/XXXX (if applicable)

2. No other changes needed. The fix is minimal, targeted, and correct. Refactoring beyond this (e.g., extracting a helper function for the duplicated load patterns) would be out of scope for a bug fix PR.

7. Overall Assessment

Approve ✅

This is a clean, well-explained, minimal fix for a platform-specific compiler bug. The transformation is semantically equivalent, has no performance impact, and the PR description thoroughly documents the root cause and rationale. The code comments are helpful for future maintainability. The only minor suggestion is to add an upstream issue reference if one exists.

---

Automated review by amd-bot using Claude. This is an AI-generated review — please use your judgment.

[HaiShaw]

@mickqian please review

[bingxche]

@amd-bot ci-status

[amd-bot]

@bingxche

CI Status for PR #21621

1 possibly related · 6 unrelated

Workflow	Job	Error	Related?	Log
PR Test (AMD)	stage-c-test-large-8-gpu-amd-mi35x (linux-mi35x-gpu-8, 0)	RuntimeError: invalid argument for batch_prefill	🟡 Possibly -- This runs on AMD MI35x GPUs and the PR specifically changes Triton kernel behavior to fix an AMD Triton issue; while the fix should help rather than hurt, the changed code path is exercised on this hardware.	View
PR Test	stage-c-test-4-gpu-b200 (2)	Fast-fail: skipping â�� root cause job(s): stage-c-test-deepep-8-gpu-h200	🟢 Unlikely -- The failure is in 'check-stage-health', an infrastructure health check step unrelated to the Triton kernel code changes.	View
PR Test	stage-c-test-4-gpu-b200 (3)	Fast-fail: skipping â�� root cause job(s): stage-c-test-deepep-8-gpu-h200	🟢 Unlikely -- The failure is in 'check-stage-health', an infrastructure health check step unrelated to the Triton kernel code changes.	View
PR Test	stage-c-test-deepep-8-gpu-h200	RuntimeError: ('The detected CUDA version (%s) mismatches the version that was used to compilePyTorch (%s). Please make sure to use the same CUDA versions.', '13.0', '12.8')	🟢 Unlikely -- The failure is in 'Install dependencies', a setup/infrastructure step unrelated to the Triton kernel modification.	View
PR Test (NPU)	per-commit-1-npu-a2 (1)	Error: failed to run script step: Error: command terminated with non-zero exit code: command terminated with exit code 255	🟢 Unlikely -- This is an NPU (Ascend) test on aarch64 infrastructure, completely unrelated to the Triton kernel changes for AMD GPU pointer handling.	View
PR Test (XPU)	pr-gate / pr-gate	Cleaning up orphan processes	🟢 Unlikely -- The failure is in 'Require run-ci label (optional)', which is a CI gating/label check unrelated to any code changes.	View
PR Test (Xeon)	pr-gate / pr-gate	Cleaning up orphan processes	🟢 Unlikely -- The failure is in 'Require run-ci label (optional)', which is a CI gating/label check unrelated to any code changes.	View

---

Generated by amd-bot