Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for fused AllReduce + RMSNorm + quantization (including MXFP4, per-group FP8, and per-token FP8) on AMD GPUs using AITER. It adds benchmarking scripts, updates communication operations, integrates the fused paths into layer communicators and layer normalization modules, and updates the Qwen 3.5 model implementation to leverage these optimizations. The feedback highlights several critical issues: the platform check helper _is_gfx95_supported is evaluated as a bare variable instead of being called as a function, which bypasses the check; the benchmark timing function should disable gradient tracking to prevent memory overhead; and potential runtime crashes should be mitigated by handling 3D tensors in _should_use_1stage_mxfp4_ar and adding defensive checks for 1D tensors in _maybe_transpose_aiter_bpreshuffle_scale.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
57713e8 to
2c9514f
Compare
950810f to
fd23e99
Compare
Conflicts: python/sglang/srt/distributed/parallel_state.py
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
@amd-bot ci-status |
CI Status for PR #29723Merge verdict: ❌ Not ready to merge. PR CI is incomplete (the AMD ROCm 7.2 run still has a queued Caution This PR's changed code is not verified by the CI that ran.
Changed files: Executed CI failure attribution: AMD: 3 executed test failures (0 clearly related) + 1 infra install-fail · Others: NVIDIA 3 executed (0 related) + 2 fast-fail-skipped · NPU 4 (0 related). AMD ROCm 7.2 run still has 1 job pending — not counted as passed. AMD Executed Failures
Other Executed Failures
Gate/rollup jobs ( Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
@amd-bot ci-status |
CI Status for PR #29723Merge verdict: ❌ Not ready to merge — CI is still running and incomplete. A new commit ( Caution This PR's core value (per-token FP8 / MXFP4 fused all-reduce→RMSNorm→quant) is exercised only by two
The third changed test, Changed source: Executed CI failure attribution: AMD: 0 real failures (main gate still running) · Others: 0 real failures (all "Extra" reds are benign label opt-outs). Real test jobs still pending/in-progress — do not read as passed. Benign gate "failures" (not caused by this PR)
Still running (decisive signal)
What to do before merge
Note: a prior ci-status comment (2026-08-13 03:08 UTC) was for an earlier commit where
|
Conflicts: python/sglang/srt/models/qwen3_5.py
|
@amd-bot ci-status |
CI Status for PR #29723Merge verdict: 🚫 Not ready — PR CI is incomplete, and the changed code has not been exercised yet. The only completed "failures" are by-design gate cascades in the Extra workflows (missing Caution This PR's core value — the AMD aiter fused AllReduce→RMSNorm→per-token-FP8/MXFP4 quant kernels — is not yet verified by any completed PR-CI test:
Changed files: mostly AMD paths — Executed CI failure attribution: AMD: 2 failures (0 related) · Others: 2 failures (0 related). All 4 are Extra-workflow gate cascades. Main NVIDIA + AMD + NPU/XPU/Xeon/Arm64/MLX pipelines are still pending — not counted as passed. Other Executed Failures
AMD Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
sogalin
left a comment
sogalin
left a comment
There was a problem hiding this comment.
Looks good to me now, we cover test cases.
|
I found several blockers on the current head
Additional issues:
The branch is also ~176 main commits behind and overlaps semantically with #34502 in five shared infrastructure files. Please rebase and reconcile a single fused-AR/tuple API before merge. Finally, the E2E benchmark toggle is confounded: |
kkHuang-amd
left a comment
kkHuang-amd
left a comment
There was a problem hiding this comment.
Inline findings for current head 4def6486, expanding the summary in #29723 (comment). Each comment identifies the concrete changed line, reachable impact, and requested fix.
| if len(hidden_states) == 2 and _linear_accepts_fp8_tuple(linear): | ||
| if len(hidden_states) == 2 and _linear_accepts_quant_tuple(linear): | ||
| return hidden_states | ||
| raise TypeError( |
There was a problem hiding this comment.
[P0] The current-head MI35X production test still reaches this raise. During Qwen3.5 CUDA-graph capture the fused producer returns a 2-tuple, while merged in_proj_ba cannot consume that exact format and has no BF16 sidecar to select. The cached boolean capability therefore does not keep producer and both consumers in agreement. Please define an explicit emitted-format/consumer-format contract and always request/select a BF16 sidecar when either GDN projection cannot consume the exact quant tuple. Add a regression using the failing Qwen3.5-397B-A17B-FP8 merged in_proj_ba path.
| if hidden_size == 7168: | ||
| # CUDA-graph microbench: direct MXFP4 epilogue is faster through 56 | ||
| # tokens, while fallback wins from 64 tokens onward. | ||
| return tokens <= 56 |
There was a problem hiding this comment.
[P0] This can select the 1-stage kernel above its documented 80-token hard limit. For example BF16 [128, 512] is exactly 128 KiB and returns true here despite having 128 tokens. Please require tokens <= 80 for every default 1-stage decision; keep the measured K=7168 cutoff as an additional restriction, not a replacement for the hard limit. Add boundary tests for 80/81 tokens.
| use_1stage_ar = _should_use_1stage_mxfp4_ar(input_) | ||
|
|
||
| try: | ||
| return ca_comm.custom_fused_ar_rms_mxfp4_quant( |
There was a problem hiding this comment.
[P0] The new quantized collectives omit the existing TC-piecewise CUDA-graph guard. AITER's custom communicator can return dummy zero outputs when its global capture state is active but the current stream is not capturing; this non-None tuple is then treated as real activations/residuals/scales. Mirror fused_allreduce_rmsnorm's capture-state handling before calling the MXFP4 and per-token wrappers, and add piecewise capture/replay correctness tests for both formats.
| use_1stage_ar, | ||
| emit_bf16=emit_bf16, | ||
| ) | ||
| except Exception: |
There was a problem hiding this comment.
[P0] An arbitrary collective runtime failure cannot safely become a per-rank local fallback. Once one or more ranks have entered the fused collective, independently returning None can make peers hang, run a second all-reduce, or consume partially mutated communicator/tensor state. Return None only from deterministic preflight checks before collective entry. Once the backend call begins, propagate the exception or implement a rank-consistent failure protocol. The same issue exists in the per-token wrapper's broad catch.
| WQ=weight.T, | ||
| x_scale=x_scale, | ||
| w_scale=weight_scale, | ||
| dtype=torch.bfloat16, |
There was a problem hiding this comment.
[P1] This tuple path loses the producer's original activation dtype. The normal path emits input.dtype, but a prequantized FP16 activation is silently promoted to BF16 here. Carry the intended output dtype in the tuple/API (as the static prequantized path does), pass it to the GEMM, and add FP16 plus BF16 tuple-path coverage.
| enable_fused_ar_quant = ( | ||
| _enable_qwen35_fused_ar_quant() | ||
| and _linear_accepts_fp8_tuple(self.linear_attn.in_proj_qkvz) | ||
| and _linear_accepts_quant_tuple(self.linear_attn.in_proj_qkvz) |
There was a problem hiding this comment.
[P1] InternS2 computes capability but never completes the new communicator wiring. Both decoder forwards still call prepare_attn_and_capture_last_layer_outputs without the required quant_format, emit_bf16, and fuse_quant=True, so the communicator does not select the per-token/MXFP4 path (and the prior opt-in is effectively disabled). Please store the consumer format and pass the full fusion contract from both forwards, or remove this ineffective opt-in.
| register_amd_ci( | ||
| est_time=480, | ||
| suite="stage-c-test-large-8-gpu-amd-mi35x", | ||
| nightly=True, |
There was a problem hiding this comment.
[P1] This registration makes the direct fused-op coverage effectively orphaned. PR stage-C invokes the suite without --nightly, while no nightly workflow invokes this exact suite; current-head logs contain no execution of this file. Either remove nightly=True for the PR suite or register a real nightly suite that is actually dispatched, then provide a non-skipped current-head run.
| ) | ||
| QWEN35_MXFP4_MODEL_PATH = os.environ.get( | ||
| "QWEN35_MXFP4_MODEL_PATH", | ||
| "amd/Qwen3.5-397B-A17B-MXFP4", |
There was a problem hiding this comment.
[P1] This all-MXFP4 checkpoint does not establish the claimed per-token FP8 path. Setting SGLANG_USE_AITER_FP8_PER_TOKEN=1 cannot create FP8 attention consumers when the checkpoint itself is all MXFP4, and current CI aborts in the preceding per-group case before this method starts. Use the MXFP4-AttnFP8 checkpoint, assert that both per-token-FP8 and MXFP4 epilogues actually fire, and calibrate the acceptance threshold against that checkpoint's baseline.
9 participants
Motivation
The existing fused AllReduce → RMSNorm → quant path only supports per-group FP8 scales.
Checkpoints that use per-token FP8 activation scaling (and MXFP4) therefore fall back to the
unfused AllReduce → RMSNorm → quantize sequence — three kernel launches plus intermediate bf16
round-trips — before every tensor-parallel linear.
Modifications
Adds a per-token FP8 variant (and a 1-stage MXFP4 variant) of the fused kernel, collapsing those
three launches into a single aiter kernel. The change is additive and gated: when the per-token
fast path is not applicable it falls back to the existing per-group / plain AR+RMSNorm path with
no behavior change.
Key changes
Collective entry points (
distributed/parallel_state.py,distributed/communication_op.py)tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_tokenreturns(fp8_output, residual_out, per_token_scale)withper_token_scaleshaped(M, 1), orNonewhen the backend cannot fuse (fallback signal).
tensor_model_parallel_fused_allreduce_rmsnorm_mxfp4_quantplus the_should_use_1stage_mxfp4_arheuristic for the MXFP4 1-stage path.fused_allreduce_rmsnorm_quant_per_token/fused_allreduce_rmsnorm_mxfp4_quant, calling aitercustom_fused_ar_rms_quant(post_per_token_quant=True).LayerNorm hook (
layers/layernorm.py)forward_with_allreduce_fusion_quant_per_token(plus the shared_forward_with_allreduce_fusion_quant_per_tokenimpl), with a cached_aiter_per_token_quantfunctor and a non-aiter /
residual is Nonefallback.Communicator gating (
layers/communicator.py)_try_fused_allreduce_rmsnorm_quantdispatches to the per-token path whenquant_format == "fp8_per_token"and the LayerNorm exposes the per-token hook; otherwise itfalls through to per-group / unfused.
Quant fast path (
layers/quantization/fp8_utils.py)(fp8, per_token_scale)inputs directly, skipping the linear's internalquantization.
Qwen3.5 wiring (
models/qwen3_5.py)_detect_fused_ar_quant_formatreturns"fp8_per_token"underSGLANG_USE_AITER_FP8_PER_TOKEN; consumes the fused 2-tuple/3-tuple handoff in both the GDNlinear-attention (
in_proj_qkvz/in_proj_ba) path and the full-attention (qkv_proj) path.Enablement
SGLANG_USE_AITER=1+--enable-aiter-allreduce-fusion(existing), plusSGLANG_USE_AITER_FP8_PER_TOKEN=1to select the per-token format. MXFP4 checkpoints auto-selectthe MXFP4 variant. ROCm/aiter/gfx95-gated; other backends and non-eligible configs are unchanged.
Accuracy Tests
The fused kernel is numerically equivalent to the reference AR → RMSNorm → per-token quant
sequence; when the fast path is unavailable it defers to the existing per-group / plain path.
Verified on Qwen3.5-397B MXFP4-AttnFP8 with GSM8K (1319q): fusion ON = 0.932, and toggling the
per-token FP8 fusion off (
SGLANG_DISABLE_FUSED_AR_QUANT) leaves accuracy unchanged at 0.932,confirming the fallback is numerically faithful.
Unit tests:
test/registered/ops/test_aiter_allreduce_fusion_amd.pyadds per-token FP8 and MXFP4cases (fused vs unfused reference) and the
None-fallback contract.Speed Tests and Profiling
Kernel benchmark:
benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.pyaddsper-token FP8 and MXFP4 coverage against the 3-launch baseline.
End-to-end
Setup: Qwen3.5-397B MXFP4-AttnFP8, TP2 on MI355X, fp8 KV cache,
--attention-backend aiter,random dataset, output len 1024, range ratio 0.8, num-prompts = 10 x cc. (1k,1k) uses
--enable-mixed-chunk; (8k,1k) does not. The arms differ only bySGLANG_USE_AITER_FP8_PER_TOKEN(OFF = per-token fusion disabled → per-group/plain fallback;ON = per-token fused path). Accuracy is identical between arms (GSM8K 1319q = 0.932 both, see
Accuracy Tests), so this is a like-for-like numerical comparison.
Total token throughput (tok/s), higher is better:
Mean TPOT (ms/token), lower is better:
The fused path is faster or equal in 9 of 10 cells, and TPOT moves in lockstep with throughput,
which is the signature of removing launches from a decode-bound step rather than of noise. The
largest gains are ~2.4% at (1k,1k) cc=16/32; the rest is neutral. That is the expected magnitude:
the fused region is a small share of a decode step, so collapsing three launches into one removes
real work without dominating end-to-end serving throughput.
Caveat on provenance: these are single samples per cell, and the ON and OFF arms were collected in
separate sessions rather than interleaved, so deltas below ~1% should be read as neutral.
Applicability
The saving is largest when the consumer of the fused output accepts the quantized tensor directly.
If the downstream linear is not quantized in a matching format, the kernel must still produce a
bf16 side output, and the benefit reduces to the removed launches alone.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #31792406179
Latest PR Test (Extra): ❌ Run #31792406163