Repository navigation
[Perf] [AMD] qwen3: Fuse the 1x128 FP8 input quant into RMSNorm and SiLU-and-mul for block-FP8 linears on gfx950 - #48043
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
a300473 to
46b550d
Compare
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
[BBuf Review Agent] One thing to fix before this can go in.
| hidden_states = self.ffn_boundary.prepare( | ||
| hidden_states, | ||
| forward_batch, | ||
| quant_format=ffn_quant, |
There was a problem hiding this comment.
[BBuf Review Agent] With Qwen3 --tp 2 --attn-dp-size 2, even CUDA/BF16 prefill now fails because _reduce_update_read_dp_gather rejects quant_format, including the empty string. Could we handle this keyword on the DP boundary paths and cover the real decoder-to-boundary call?
46b550d to
151ce19
Compare
This comment has been minimized.
This comment has been minimized.
|
/tag-and-rerun-ci |
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
LGTM. The attention-DP regression is addressed.
|
[BBuf Review Agent] Could you rebase this onto the latest main? It conflicts with main now. |
151ce19 to
c52e332
Compare
|
/tag-and-rerun-ci |
|
/tag-and-rerun-ci |
PR Pundit review(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)
1. 🟠 Should fix: fp16 Qwen3 block-FP8 model decode emits bf16 qkv/MLP outputs: this PR's fp8 norm fusion lacks an activation-dtype guard
Who is affected: Only Qwen3 block-FP8 checkpoints run with Checked against the code by an independent pass. Related past reviews: #44266 on 2. 🟠 Should fix: Under torch.compile the SiLU fusion falls back while norms stay fused, losing the SiLU speedup: this PR gave the gates different condition lists
Who is affected: Qwen3 block-FP8 on gfx950 with SGLANG_USE_AITER=1 and --enable-torch-compile. Checked against the code by an independent pass. 3. 🟠 Should fix: Llama/Qwen2 block-FP8 decode on gfx950 silently switches to the fused quant kernel with no model-level test: this PR added the dispatch to the shared forward
Who is affected: Qwen2/Llama-family block-FP8 models decoding on gfx950 with SGLANG_USE_AITER=1 and ROCm >= 7.2. Checked against the code by an independent pass. Related past reviews: #19746 on 4. 🟠 Should fix: Qwen3 decoder DP test would pass even if the decode fp8 fusion were disabled: this PR's attention stub swallows and never records
|
c52e332 to
5a5c815
Compare
|
@peymanr Thanks for the review. Point by point, fixed in 5a5c815 (also rebased onto the latest main):
|
|
[BBuf Review Agent] @zengleixin-amd Checked 5a5c815: points 1, 2, 4 and 5 are addressed. I also accept the shared Qwen2/Llama SiLU path with the bitwise kernel checks, so all five points can be closed. |
PR Pundit review(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)
1. 🟠 Should fix: Two Qwen3 qkv_proj shapes skip the gfx950 small-M decode kernel: this PR pre-quantizes qkv_proj input
On gfx950 decode, the fused norm-quant hands Who is affected: Block-FP8 Qwen3-0.6B/1.7B served at TP1 on gfx950 with SGLANG_USE_AITER=1, decode batches of at most 32 tokens. Checked against the code by an independent pass. Related past reviews: #46908 on 2. 🟠 Should fix: Dense Qwen3-FP8 decode regression on MI35x passes CI: this PR's decoder tests stub the predicate and collaborators
The new unit tests for Who is affected: Dense Qwen3 block-FP8 models served on gfx950 with SGLANG_USE_AITER=1 in decode, if the predicate or tuple plumbing regresses. Checked against the code by an independent pass. Related past reviews: #46908 on 3. 🟠 Should fix: Fused 1x128 quant silently skipped under TP>1 all-reduce fusion: this PR's reads_quantized ignores quant_format
Who is affected: Qwen3 block-FP8 decode on gfx950 with SGLANG_USE_AITER and TP>1 where the aiter/custom all-reduce+RMSNorm fusion takes the batch. Checked against the code by an independent pass. Related past reviews: #32779 on Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at |
5a5c815 to
47e214b
Compare
|
@peymanr Thanks for the second pass. Point by point, against 47e214b:
|
|
[BBuf Review Agent] @zengleixin-amd I accept the small-M fix, existing MI355X coverage and bf16 fallback for all-reduce fusion. But qwen3.py:654 still reads hidden_states.shape before any gates or boundary prepare, so a previous FFN's OwedOutput crashes ordinary TP2, including CUDA/BF16 prefill; could we get the row count from batch/boundary metadata and cover the two-layer handoff? |
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
[BBuf Review Agent] One thing to fix before this can go in.
| capture_output=None, | ||
| ) -> torch.Tensor: | ||
| attn_quant, ffn_quant = self._quant_formats( | ||
| forward_batch, num_tokens=hidden_states.shape[0] |
There was a problem hiding this comment.
[BBuf Review Agent] With Qwen3 --tp 2, the previous FFN can return an OwedOutput, so the next decoder layer now raises AttributeError on hidden_states.shape. This runs before the decode/ROCm checks, so BF16 prefill on CUDA is affected too. Could we obtain the row count through batch/boundary metadata and add a two-layer deferred-output regression test?
…or block-FP8 linears on gfx950 On gfx950 with aiter, block-FP8 dense linears take a 1x128 group-quantized FP8 input, which Qwen3 produced with a separate quant kernel after each input norm and after SiLU-and-mul. In decode, the attention and FFN input norms now request the fp8 quant format (aiter fused_rms_fp8_group_quant through the layer boundary's NormQuantReadout) when the next linear runs the bpreshuffle block-FP8 GEMM, and DenseSiluAndMulQuant runs a new Triton kernel that fuses SiLU-and-mul with the quant. The SiLU-and-mul kernel rounds like SiluAndMul, so its output equals the unfused path. The attention norm stays unfused when post_residual_addition is passed. A read takes the quant format only where the boundary's entry reads last, on the rows the stage computes (StageBoundary.reads_quantized); the FFN entry under attention DP gathers after its read and stays unfused. The quant format is passed to prepare only when non-empty, and the entry steps that cannot quantize accept it and reject a non-empty one. The fused norm needs bf16 activations, since the GEMM returns bf16 for a pre-quantized input, and both fused producers (norm and SiLU-and-mul) now share one runtime gate, fp8_group_quant_fusion_allowed(), so they turn off together (torch.compile tracing, deterministic inference / RL on-policy, batch-invariant mode, LoRA). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
47e214b to
3b748e1
Compare
|
@BBuf Thanks, confirmed and fixed in 3b748e1. With
Verified on MI355X: Qwen3-8B BF16 TP2, Qwen3-14B-FP8 TP2 and TP1 serve and generate; |
|
[BBuf Review Agent] @zengleixin-amd Checked 3b748e1: the row count now comes from positions after the fusion gates, and hidden_states is inspected only after boundary prepare. The two-layer test covers the OwedOutput handoff in prefill and decode, so this addresses my concern. |
|
/tag-and-rerun-ci |
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
LGTM. The deferred-output issue is fixed.
Motivation
On gfx950 with aiter, block-FP8 dense linears run the bpreshuffle blockscale GEMM, which takes a 1x128 group-quantized FP8 input. In a Qwen3 decoder layer each such input is quantized by a separate kernel after the op that produces it: after the input RMSNorm (
qkv_proj), after the post-attention RMSNorm (gate_up_proj) and after SiLU-and-mul (down_proj). Each quant reads back a bf16 tensor that was just written and adds a launch. Qwen3.5 already reads its attention input through the layer boundary'sNormQuantReadout(aiterfused_rms_fp8_group_quant); Qwen3 does not.Modifications
models/qwen3.py: in decode, the attention and FFN input norms request the fp8 quant format when the following linear runs the aiter bpreshuffle block-FP8 GEMM (aiter_takes_fp8_group_input), on gfx95 with aiter, with bf16 activations (the GEMM returns bf16 for a pre-quantized input), and under the same runtime gate as the SiLU-and-mul fusion (fp8_group_quant_fusion_allowed(): not while torch.compile traces, nor under deterministic inference, an RL on-policy target, batch-invariant mode or LoRA). The FFN boundary declaresread=NORM_QUANT_READOUT. The attention norm stays unfused whenpost_residual_additionis passed (Qwen3-VL deepstack), because the fused path does not apply it. The token count and the aiter fused KV-write dtype check accept the(fp8, scale)tuple. A norm requests the quant only where its stage's entry reads last, on the rows the stage computes (StageBoundary.reads_quantized): under attention DP the FFN entry gathers after its read, so the FFN norm stays unfused there.quant_formatis passed toprepareonly when non-empty, so every other configuration calls the boundary exactly as before.layers/quantization/fp8_activation.py: on the same gfx95 bpreshuffle path,DenseSiluAndMulQuant(used by the Qwen2, Qwen3 and Llama dense MLPs) runs SiLU-and-mul and the 1x128 quant in one kernel. Its runtime conditions moved intofp8_group_quant_fusion_allowed(), shared with the fused norm, without changing them.kernels/ops/gemm/smallm_fp8_blockscale_gfx950: the gfx950 small-M block-FP8 GEMM quantizes its bf16 input itself and is skipped for a pre-quantized one.prefers_bf16_input(M, N, K)reports, per listed shape, the row counts at which it on a bf16 input beats the bpreshuffle GEMM on a pre-quantized input (measured on MI355X, below), and both producers (the fused norm, throughaiter_prefers_bf16_inputinfp8_activation.py, and the fused SiLU-and-mul) keep their output in bf16 there. For Qwen3 this only touchesqkv_projof Qwen3-0.6B and Qwen3-1.7B at TP1, up to 16 rows. The decoder takes that row count frompositions(one entry per token; the layer input can be an opaqueOwedOutputunder TP > 1) and reads it only past the decode, gfx95 aiter and runtime gates.kernels/ops/quantization/silu_mul_fp8_group_quant.py. It rounds likeSiluAndMul(silu to bf16, then the product to bf16) and quantizes like aiter's per-1x128 quant, so its output is identical to the unfused path.layers/layer_boundary/stage.py: newStageBoundary.reads_quantized(forward_batch), and thepreparedocstring states which reads takequant_format, and that it is a request: a fusion that completes the sum with the read (e.g. the aiter / custom all-reduce + RMSNorm fusion under TP > 1) returns the bf16 input, which the GEMM then quantizes itself.layers/layer_boundary/prepare.py:reads_quantized(entry)(only_reduce_update_read/_update_readsteps, with no input move after them, hand on a quantized read). The six entry steps that cannot quantize (_reduce_update_read_dp_gather,_dp_gather_sum_read,_attn_tp_reduce_scatter_update_read,_attn_tp_slice_update_read,_tp_reduce_scatter_update_read_gather,_tp_sum_with_residual_read) acceptquant_formatand raiseNotImplementedErrorfor a non-empty one.test/registered/kernels/ops/quantization/test_fused_norm_act_fp8_group_quant_gfx95.py(stage-b-test-1-gpu-small-amd-mi35x);test_boundary_output_contracts.pygets the new_quant_formatshook on its layer stub, a regression test that runsQwen3DecoderLayer.forwardinto the FFN entry bound under attention DP (prefill, and a decode that would otherwise fuse; it failed withTypeError: _reduce_update_read_dp_gather() got an unexpected keyword argument 'quant_format'before the fix), and a check that the decode reads stay quantized without attention DP and the FFN's turns off with it; the attention-DP test also asserts that decode still sendsquant_format="fp8"to the attention read; a gate test checks that bf16 fuses while fp16 and torch.compile tracing do not (both failed before the gate change); the kernel test adds the TP8/TP4 per-rank intermediate sizes 2176 and 1536, whose group counts leave a partial last program. A test checks that the Qwen3-0.6Bqkv_projnorm stays unfused at 16 rows and fuses at 17 rows and withSGLANG_ROCM_SMALLM_FP8_BS=0, withgate_up_projfused throughout. A two-layer test runsQwen3DecoderLayer.forwardon real attention / FFN stages under attention TP 2, where each layer hands on anOwedOutput, in prefill and in a decode that checks the fusion.Accuracy Tests
MI355X (gfx950), image
lmsysorg/sglang-rocm:v0.5.22-rocm10-mi35x-20261025(torch 2.11.0+rocm10, aiter285579b6),SGLANG_USE_AITER=1, Qwen3-14B-FP8, TP1.Kernel outputs. The fused SiLU-and-mul quant equals the unfused path exactly (fp8 values and scales) for the Qwen3-14B and Qwen3-1.7B intermediate sizes at M = 1 to 512; the test asserts equality. The fused norm is aiter's
fused_rms_fp8_group_quant, which quantizes the fp32 normalized value directly while the unfused path rounds it to bf16 first. On 112 batch sizes (M = 1 to 64, every CUDA-graph decode size up to 512, a few odd sizes) and on 1,600 real decode inputs dumped from the server, about 2.2% of the fp8 values differ by one e4m3 step, scales differ within the bf16 rounding of the group amax, and the residual output is identical. The test bounds this (at most one step, at most 5% of values).GSM8K (
python -m sglang.test.few_shot_gsm8k --num-questions 1319, greedy):68e57f4005, 12 runs eachd74749b0b6, 3 runs eachAccuracies are the eval's own; the per-question outcomes are parsed from its dumped outputs. Runs of the same build differ only by nondeterminism, so the per-question test is the one that matters; neither difference is significant. The change in the norm path flips about 50 borderline questions, and which way the net goes is not tied to accuracy: a torch reimplementation of the same arithmetic as the aiter kernel (at most 0.04% of fp8 values differ from it) scored 0.9018 against main's 0.9012 on the same build (12 runs each), and the bf16-rounded variant of that reimplementation scored 0.9003. The SiLU-and-mul fusion alone (norm fusion disabled) scored 0.9012 against 0.9007.
Distance to the BF16 model. Teacher-forced on the gold answers of 400 GSM8K questions (48,692 tokens), compared with
Qwen/Qwen3-14B(BF16). For this measurement the fused norm was also applied in prefill (an experiment-only switch; the PR fuses decode only).The fused norm is not further from the BF16 model than main (the gap in |Δ logprob| is significant in its favor, 95% bootstrap interval over questions [0.00025, 0.00165]).
Attention DP (MI355X,
--tp 2 --attn-dp-size 2): Qwen3-8B (BF16) serves, with the same greedy outputs as main on the check prompts; Qwen3-14B-FP8 serves. Without attention DP, a TP1 Qwen3-14B-FP8 eager profile of 30 steps launches the same fused kernels as before this change (fused_rms_fp8_group_quant2240,silu_mul_fp8_group_quant1120).Speed Tests and Profiling
python -m sglang.bench_serving --backend sglang --dataset-name random-ids --random-input-len 1024 --random-output-len 1024 --random-range-ratio 1.0, 80 prompts at concurrency 16 and 320 at concurrency 64, servers in the order main, PR, PR, main on one GPU with the node otherwise idle, 4 runs per server with the first dropped:68e57f400568e57f4005d74749b0b6d74749b0b668e57f400568e57f4005The two runs of each build agree within 0.4%. The SiLU-and-mul fusion alone gives +3.1% / +2.0% (concurrency 16 / 64) on Qwen3-14B-FP8.
A torch-profiler kernel census of eager decode on Qwen3-14B-FP8 shows
_fused_rms_fp8_group_quant_kerneland_silu_mul_fp8_group_quant_kernelon this PR (2240 and 1120 calls in 30 steps, batch 8) and neither on main.Small-M GEMM shapes (MI355X, CUDA graph, per call, norm or SiLU-and-mul plus GEMM): the bf16 input to the small-M GEMM is faster up to the listed row count, the pre-quantized input above it.
python -m sglang.benchmark.one_batchon Qwen3-1.7B-FP8 TP1 (input 512, output 128, two rounds with the GPUs swapped), median decode step of this PR with vs without the small-M rule: 0.973 / 0.980 / 0.981 / 0.981 at batch 1 / 4 / 8 / 16, 1.001 / 1.001 at batch 24 / 32.Scope: block-FP8 checkpoints in compressed-tensors format (for example the published Llama-3.1-8B and Qwen2.5 FP8-block ones) load through
CompressedTensorsLinearMethodand do not take these paths (census onRedHatAI/Llama-3.1-8B-Instruct-FP8-block: no fused kernels). Llama/Qwen2 checkpoints in the nativefp8format take the SiLU-and-mul fusion only, which is bit-exact; their speed was not measured.Checklist
Found by Hyperloom.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ⏳ Run #42193308108
Latest PR Test (Extra): ❌ Run #42193308010
Latest PR Test (AMD ROCm 10): ⏳ Run #42193307957
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: