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.

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

Speed Tests and Profiling

SGLANG_USE_AITER_UNIFIED_ATTN=1 \
SGLANG_USE_AITER=1 \
SGLANG_AITER_UNIFIED_VERIFY=1 \
python3 -m sglang.launch_server \
  --model-path /data2/Qwen/Qwen3.5-397B-A17B-FP8/ \
  --tp 8 \
  --attention-backend aiter \
  --trust-remote-code \
  --model-loader-extra-config '{"enable_multithread_load": true}' \
  --watchdog-timeout 1200 \
  --mem-fraction-static 0.9 \
  --host 0.0.0.0 \
  --port 9000 \
  --disable-radix-cache \
  --enable-aiter-allreduce-fusion \
  --max-running-requests 128 \
  --page-size 16

python3 benchmark_serving.py \
  --model /data2/Qwen/Qwen3.5-397B-A17B-FP8 \
  --base-url http://0.0.0.0:9000 \
  --backend sglang \
  --dataset-name random \
  --random-input-len 8192 \
  --random-output-len 8192 \
  --random-range-ratio 1.0 \
  --num-prompts 32 \
  --max-concurrency 16 \
  --request-rate inf \
  --ignore-eos \
  --num-warmups 16 \
  --percentile-metrics ttft,tpot,itl,e2el

End-to-end serving result, 8192 input / 8192 output, concurrency 16, 32 prompts,
ignore-eos:

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:

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.