[avjves]

Motivation

sgl-project#35573 enabled the fused QK-norm and QK-norm+RoPE paths on ROCm, but the kernels behind them still don't build on all models. elementwise/qknorm.cuh fails outright, and diffusion/qknorm_rope.cuh fails for neox instantiations. Non-neox rope is the one case that already worked, which is why FLUX.2 saw the benefit of sgl-project#35573 and MiniMax H3 did not.

This PR removes two dead CUDA includes in qknorm.cuh, and on ROCm pins the shuffle width to the kernel's 32-lane sub-group instead of letting HIP default it to the physical wave64. With both fixed, MiniMax-H3's QK-norm paths accept ROCm.

Without these changes neither of the kernels work on ROCm. Theser are necessary for eager / compile mode parity.

CUDA is unaffected. The shuffle change is behind #ifndef USE_ROCM and the dropped includes were dead.

Modifications

• elementwise/qknorm.cuh carried cuda_bf16.h / cuda_fp16.h as dead includes, which hipcc cannot find ("fatal error: 'cuda_bf16.h' file not found"). The file reaches the GPU only through impl/norm.cuh and the runtime.cuh macro shims, both of which hipcc builds, so the includes can just go.

• __shfl_sync(kActiveMask, ...) passes a uint32_t mask, and HIP static_asserts sizeof(MaskT) == 8. On ROCm this now calls __shfl(..., partner_lane, device::kWarpThreads) instead. device::kWarpThreads is 32u unconditionally, so the kernel indexes lanes in 32-wide logical warps, while HIP defaults the shuffle width to the physical wave64. partner_lane is an index inside that 32-lane sub-group, so the width has to be explicit.

• minimax_h3.py gated _apply_qk_norm and _use_fused_qknorm_rope on current_platform.is_cuda(). Both now accept ROCm, which is what reaches the two kernels above. MiniMax uses is_neox=True, the instantiation that did not
  build.

Accuracy Tests

For testing, using MiniMax-H3. The output changes slightly (exactly 1 BF16 ULP) as expected by the rope kernel. Enabling torch.compile would change it as well. No artifacts present.

Command:

sglang generate \
    --model-path MiniMaxAI/MiniMax-H3 \
    --ulysses-degree 8 --ring-degree 1 --num-gpus 8 \
    --num-inference-steps 50 \
    --dit-cpu-offload False --dit-layerwise-offload False \
    --text-encoder-cpu-offload False --image-encoder-cpu-offload False \
    --offload-during-compile False --vae-cpu-offload False \
    --warmup-mode request --warmup-steps 50 \
    --vae-precision bf16 --image-encoder-precision bf16 \
    --output-path /outputs/minimax_h3.pr1 \
    --prompt 'A cat flying through the universe on a broom, playing rap music from his boombox' \
    --seed 42 \
    --config /outputs/minimax_h3.pr1/config.json \
    --perf-dump-path /outputs/minimax_h3.pr1/perf/pr1.json"

Before

v0.mp4

After

pr1.mp4

Speed Tests and Profiling

1. High-level Summary

Metric	Baseline	New	Diff	Status
E2E Latency	29317.56 ms	28142.04 ms	-1175.52 ms (-4.0%)	✅
Throughput	0.03 req/s	0.04 req/s	-	-

2. Stage Breakdown

Stage Name	Baseline (ms)	New (ms)	Diff (ms)	Diff (%)	Status
InputValidationStage	0.03	0.04	+0.01	+23.8%	⚪️
MiniMaxH3PartitionAdmissionStage	0.02	0.02	+0.01	+33.6%	⚪️
MiniMaxH3TextEncodingStage	29.64	29.23	-0.41	-1.4%	⚪️
MiniMaxH3VisualEncodingStage	0.01	0.02	+0.00	+8.7%	⚪️
MiniMaxH3AudioEncodingStage	0.01	0.01	+0.00	+3.8%	⚪️
MiniMaxH3LatentPreparationStage	11.70	11.23	-0.47	-4.0%	⚪️
MiniMaxH3TimestepPreparationStage	0.13	0.15	+0.02	+14.2%	⚪️
MiniMaxH3DenoisingStage	24650.74	23212.86	-1437.88	-5.8%	⚪️
MiniMaxH3DecodingStage	1590.55	1582.83	-7.71	-0.5%	⚪️
Scheduler.return_result.spill_arrays	0.06	0.06	+0.00	+0.7%	⚪️
SchedulerClient.materialize_file_refs	0.01	0.01	+0.00	+1.2%	⚪️

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.