[AMD] Add fused all-reduce RMSNorm per-group quant for Qwen3.5 FP8

[hubertlu-tw]

Co-author: @mqhc2020, @yichiche

Motivation

This PR adds a fused aiter path for the Qwen3.5 FP8 attention/GDN input normalization path on AMD. The baseline path around prepare_attn runs all-reduce, RMSNorm, and per-group activation quant as separate work before the FP8 projection consumes (fp8, scale).

The new path lets LayerCommunicator.prepare_attn request fused all-reduce + RMSNorm + per-group FP8 quant. Standard attention receives (fp8, scale) directly for qkv_proj. GDN receives (bf16, fp8, scale) so in_proj_qkvz can skip its internal activation quant while in_proj_ba still reads the bf16 activation it needs.

The fused-quant path is scoped to ROCm/aiter. The aiter single-kernel dispatch is gated to gfx95 / gfx950-class GPUs through is_gfx95_supported(), and plain --enable-aiter-allreduce-fusion AR+RMSNorm behavior remains separate. If the fused-quant path is unavailable, callers fall back to fused AR+RMSNorm or the existing generic path.

Modifications

• Added a tensor-parallel API for fused all-reduce + RMSNorm + per-group FP8
  quant:
  • tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group(...)
  • GroupCoordinator.fused_allreduce_rmsnorm_quant_per_group(...)
• Added layernorm helpers that return tuple activations for FP8 consumers:
  • ((fp8, scale), residual) for standard attention.
  • ((bf16, fp8, scale), residual) for GDN, where bf16 is required by in_proj_ba.
• Wired LayerCommunicator.prepare_attn to prefer the fused-quant helper when _sglang_needs_allreduce_fusion is set and the layer opts in with enable_fused_ar_quant=True.
• Wired Qwen3.5 standard attention with keep_bf16=False and GDN with keep_bf16=True.
• Updated Qwen3.5 GDN input projection to consume the (bf16, fp8, scale) tuple without dequantizing fp8 back to bf16.
• Added SGLANG_DISABLE_FUSED_AR_QUANT=1 as the operator opt-out for this quantized handoff while keeping the existing AR+RMSNorm fusion eligible.

Fallback order:

1. aiter single-kernel AR + RMSNorm + per-group FP8 quant, optionally with bf16 side-output.
2. aiter fused AR + RMSNorm plus separate aiter per-1x128 quant.
3. Generic caller fallback.

The single-kernel dispatch is intentionally limited to ROCm + aiter + gfx95.

Accuracy Tests

GSM8K was run on Qwen3.5-397B-A17B-FP8 with TP=8 and
--enable-aiter-allreduce-fusion.

Config	GSM8K command	Accuracy
Baseline	python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319 --num-shots 5 --port 9000	0.946
Fused AR+RMSNorm+per-group quant	python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319 --num-shots 5 --port 9000	0.955

The fused path is above the 0.94 accuracy gate used for this model.

Speed Tests and Profiling

Server command

SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 \
python3 -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-397B-A17B-FP8 \
  --trust-remote-code \
  --model-loader-extra-config '{"enable_multithread_load": true}' \
  --host 0.0.0.0 --port 30001 --tensor-parallel-size 4 \
  --attention-backend aiter --kv-cache-dtype fp8_e4m3 --page-size 16 \
  --chunked-prefill-size 8192 --mem-fraction-static 0.8 --watchdog-timeout 1200 \
  --disable-radix-cache \
  --enable-aiter-allreduce-fusion

Serving benchmark command

for cc in 2 8 32; do
  python3 -m sglang.bench_serving --host 0.0.0.0 --port 30001 \
    --model /data2/Qwen/Qwen3.5-397B-A17B-FP8 \
    --dataset-name random --random-input 8192 --random-output 1024 \
    --random-range-ratio 0.8 --max-concurrency $cc \
    --num-prompts $((cc*10)) --seed 0
done

Case A: baseline without --enable-aiter-allreduce-fusion
Case B: with --enable-aiter-allreduce-fusion

cc	Output tok/s A -> B	Delta	Mean TPOT A -> B	Delta	Mean TTFT A -> B
2	186.42 -> 197.21	+5.8%	10.46 ms -> 9.88 ms	+5.5%	222.27 ms -> 225.39 ms
8	598.54 -> 625.89	+4.6%	12.82 ms -> 12.23 ms	+4.6%	292.07 ms -> 305.96 ms
32	1332.17 -> 1328.84	-0.2%	22.81 ms -> 22.90 ms	-0.4%	571.93 ms -> 560.36 ms

Improvement formula:

• Throughput: (candidate - baseline) / baseline * 100.
• TPOT: (baseline - candidate) / baseline * 100.

Kernel-count profile, per decode pass on Qwen3.5-397B-A17B-FP8 / TP=8 with 60
decoder layers:

Kernel	Baseline	Fused path	Delta
cross_device_reduce_*stage (plain AR)	60	0	-60
add_rmsnorm_quant_kernel (RMSNorm+add)	60	0	-60
dynamic_per_group_scaled_quant_kernel	60	0	-60
allreduce_fusion_kernel_1stage (AR+RMSNorm)	0	0	0
allreduce_fusion_kernel_1stage_per_group	0	60	+60
Comm-side separate quant for GDN keep-bf16 path	45	0	-45
Total kernels saved on prepare_attn			-165

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.

---

CI States

Latest PR Test (Base): ✅ Run #29728384591
Latest PR Test (Extra): ❌ Run #29728384518

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yctseng0211]

@HaiShaw
#30940 is ready now (verified on ci mi355 runners)

[fused-ar-rms-per-group-quant] metrics={'accuracy': 0.953, 'invalid': 0.01, 'latency': 73.984, 'output_throughput': 2923.251}

Qwen3.5-FP8 aiter AR-fusion (MI35x, parallel TP4)

Variant	GPUs	Accuracy	Invalid	Latency (s)	Output tok/s	Threshold	Status
disable-fused-ar-quant-opt-out	4,5,6,7	0.942	0.012	74.16	2947.32	0.94	PASS
fused-ar-rms-per-group-quant	0,1,2,3	0.953	0.010	73.98	2923.25	0.94	PASS

[HaiShaw]

#30940 merged

[yctseng0211]

4b23443 fixed broken cuda CI: https://github.com/sgl-project/sglang/actions/runs/29680669860/job/88269276341?pr=24651

[yctseng0211]

cuda PR Test is clean
AMD PR Test:
stage-c all green

, waiting for the last one AMD PR-Test : https://github.com/sgl-project/sglang/actions/runs/29728384583/job/88794722663?pr=24651

[yctseng0211]

@HaiShaw
Cuda CI/ AMD CI all green, safe to merge