[Shan2L]

Motivation

sgl-kernel cannot be built for Radeon RDNA4 GPUs (gfx1201: Radeon AI PRO R9700 / R9600D, RX 9070 series): the target is not in the setup_rocm.py allowlist. Adding it to the list is not enough. The MoE TopK routing kernels then fail at launch:

test (gfx1201, allowlist-only build)	result
test_moe_topk_softmax.py	Fatal Python error: Aborted (core dump)
test_moe_topk_sigmoid.py	torch.AcceleratorError: HIP error: unspecified launch failure

Root cause: on ROCm, WARP_SIZE is the wave size of the code object being compiled. It is 32 in the device pass for gfx1250 / gfx1201, but it is always 64 in the host pass. The TopK host launchers derive the per-thread work (VPT), the template instantiation, and block_dim from WARP_SIZE. For wave32 targets the host therefore launches the wave64 configuration of a kernel that was compiled for wave32. gfx1250, which is already in the allowlist, has the same host/device mismatch. kvcacheio/transfer.cu also sizes its blocks with the host-side WARP_SIZE.

Modifications

• include/utils.h:
  • Documents what WARP_SIZE means on ROCm.
  • Adds sgl_launch_warp_size(), the runtime wave size of the current device (from the cached at::cuda device properties).
  • Adds SGL_DISPATCH_WARP_SIZE(...), which runs a lambda with constexpr int kWarpSize set to 32 or 64.
  • Adds sgl_wave_matches_target<WAVE>, which is true when WAVE matches the code object being compiled.
  • On CUDA these resolve to WARP_SIZE / 32, with no behavior change.
• moe_topk_softmax_kernels.cu, moe_topk_sigmoid_kernels.cu:
  • topkGating{Softmax,Sigmoid} and TopkConstants are templated on the wave size, and the launchers dispatch on the device's wave size.
  • Each code object gets both instantiations. Only the one matching its wave size has a body; the other compiles to a 2–5 instruction trap stub and is never launched.
• kvcacheio/transfer.cu: on ROCm, threads_per_block uses sgl_launch_warp_size().
• setup_rocm.py: adds gfx1201 to supported_targets, with comments on wave size, FP8 format (OCP E4M3), and LDS budget.
• docs/.../amd_gpu.mdx: adds a "Radeon RDNA4 / gfx1201 (experimental)" section covering build, wave-size rules for kernel authors, and serving notes.

Other wave-size uses were checked and need no change:

• moe_align: the host only over-provisions threads, and the ROCm path is independent of wave size.
• topk.hip and deepseek_v4_topk.cu: wave handling is device-side only.
• dsv4_norm_rope: already writes E4M3FN on gfx12.

Wave64 targets are unchanged

sgl-kernel was built from main and from this PR for each target (AMDGPU_TARGET=<arch>, GPUs hidden), and the disassembled code objects were compared kernel by kernel, ignoring only addresses and symbol names:

target	kernels compared	identical	extra (other-wave trap stubs)
gfx942	215	215	57 (2 instructions each)
gfx950	215	215	57 (2 instructions each)
gfx1250	215	215	57 (5 instructions each)
gfx1201	215	215 (vs. allowlist-only build)	57

On gfx942 / gfx950 the device code is therefore unchanged. The runtime cost is one cached device-property lookup per TopK / KV-transfer launch. On gfx1250 the device code is also unchanged, but the host now launches the wave32 configuration it was compiled for; I have no gfx1250 hardware to test this on.

Accuracy Tests

Kernel unit tests on R9700 (gfx1201), PyTorch 2.11.0+rocm10.0.0:

test	allowlist only	this PR
test_moe_topk_softmax.py	abort (core dump)	744 passed
test_moe_topk_sigmoid.py	hipErrorLaunchFailure	774 passed
test_moe_align.py	4368 passed, 72 skipped	4368 passed, 72 skipped
test_topk.py	162 passed, 1 skipped	162 passed, 1 skipped
test_kvcacheio.py	144 passed, 50 skipped	144 passed, 50 skipped
test_activation.py	945 passed	945 passed
test_dsv4_norm_rope.py	16 passed	16 passed

LLM end to end (exercises sgl_kernel.topk_softmax):

• Model and setup: Qwen3-30B-A3B BF16 (128 experts, top-8), TP4 on 4x R9700.
• Flags: SGLANG_USE_AITER=0 --attention-backend triton --disable-custom-all-reduce.
• GSM8K (200 questions, 5-shot): accuracy 0.930, invalid 0.000.
• With the allowlist-only build, the same server fails at startup during CUDA graph capture with RuntimeError: topk_softmax launch error: unspecified launch failure.

Diffusion smoke: MiniMax-H3 (Ref2AV) with TP2 on 2x R9700 and --attention-backend torch_sdpa, using the custom all-reduce fix from #43676, completes 2 denoising steps and writes the video.

Roadmap: gfx1201 SageAttention / Sol-Attn in SGLang Diffusion

This PR is the first step toward fast MiniMax-H3 (and other DiT) serving on Radeon gfx1201 with SGLang Diffusion. Today, on gfx1201:

• SAGE_ATTN and SOL_ATTN are rejected on ROCm, because they call CUDA-only packages.
• AITER_SAGE runs, but it uses the aiter Triton fav3_sage kernel. That kernel takes 12.4 s per MiniMax-H3 attention call, which makes it 3.25x slower end to end than BF16 SDPA (688.3 vs 211.6 s per DiT step).

Two gfx1201 attention ops are being upstreamed to AITER as a four-part series:

• gfx1201_sage_attention: INT8 QK^T, FP8 PV, FP32 accumulation; a HIP kernel plus a hand-scheduled code object. It is [HIP] [JIT] [Build] gfx1201: SageAttention forward (INT8 QK, FP8 PV) ROCm/aiter#6390 plus a follow-up.
• gfx1201_sol_attention: Sol-Attn block routing with the same arguments as NVlabs sol_attn; also HIP plus a code object.
• The overall integration plan for LightX2V, diffusers, vLLM-Omni and SGLang Diffusion is in this comment on ROCm/aiter#6292.

With these ops plugged into stock LightX2V, MiniMax-H3 ref2av (TP2, 2x R9700, 768x1344, 362 frames, 29 steps) measures:

attention	DiT s/step	vs BF16 SDPA
BF16 SDPA	211.6	1.00x
aiter Triton fav3_sage (today's AITER_SAGE)	688.3	0.31x
gfx1201_sage_attention	89.1	2.38x
gfx1201_sol_attention (Sol config)	60.8	3.48x

Planned SGLang changes, in order:

1. This PR: sgl-kernel builds and runs correctly on gfx1201 (wave32).
2. [Fix] [AMD] custom all-reduce: skip IPC setup on ROCm when it cannot be selected #43676: custom all-reduce IPC setup no longer breaks multi-GPU runs on PCIe Radeon with expandable_segments.
3. SGLang Diffusion AITER_SAGE: dispatch to gfx1201_sage_attention on gfx1201, keeping the existing fav3_sage path on other archs. This lands after the AITER Sage parts are merged.
4. SGLang Diffusion SOL_ATTN: accept it on ROCm through gfx1201_sol_attention, with AITER_SAGE as the dense guard. This lands after the AITER Sol parts are merged.
5. A gfx1201 CI runner, if maintainers are open to it, so this target stays covered.

Steps 3 and 4 are opt-in and only change behavior on gfx1201. The CUDA and gfx94x/gfx95x paths stay unchanged.

Checklist

• Format code with pre-commit.
• Add documentation (amd_gpu.mdx).
• Provide accuracy results (above).

---

CI States

Latest PR Test (Base): ❌ Run #38135876973
Latest PR Test (Extra): ❌ Run #38135876782
Latest PR Test (AMD ROCm 10): ❌ Run #38135876885