Repository navigation
[ROCm][Perf] Qwen3.8-27B dense mono decode: persistent FlyDSL GDN and FFN launches (TP1 and TP2, MI355X/MI350X) - #43558
vorapolsiloai wants to merge 18 commits into
Conversation
…layer on MI355X Port the vllm-project/vllm#60833 kernels. A decode step of at most 8 rows runs each layer as K1 (GDN front) and K2 (o_proj/out_proj, in-kernel TP all-reduces, norm, router, MXFP4 experts, shared expert). The full-attention layers keep SGLang's fused HIP front through a new attention_core(). Builds on the Kimi-K3 mono framework (k3_mono_flydsl common/ and stages/). Off by default: SGLANG_QWEN3_5_MONO_DECODE=1.
Change the launch asserts to need() checks. python -O does not remove them. Refuse GPUs with fewer than 256 compute units. Refuse the hierarchical cache and LMCache. Their copy kernels can hold CUs.
af1e93e to
d90a42b
Compare
One persistent launch (256 CTAs x 512 threads, gfx950) runs a decoder layer's out_proj / o_proj, the residual add, post_attention_layernorm and the MXFP4 gate_up, SiLU-mul and down projections for decode steps of 1 to 8 rows at TP1. It reads (16, 16)-shuffled copies of the MXFP4 weights and matches the stock path (dynamic_mxfp4_quant + Triton gemm_afp4wfp4, GemmaRMSNorm's ROCm Triton kernel, SiluAndMul) bit for bit. The norm sums its squares in the stock kernel's order (per-thread FMA pairs, its warp and cross-warp butterflies) and takes rstd as 1 / sqrt(sum / n + eps) with IEEE divisions and the raw v_sqrt; any other order changes an MXFP4 code on about one row in 640.
One persistent launch runs a GDN layer from its input to out_proj's input: the residual add and input_layernorm, the MXFP4 in_proj_qkvz and in_proj_ba, the conv update, the gated delta rule and the gated norm. It is bit-exact with the stock decode kernels on gfx950: it uses their lane layouts and reduction orders, the FMA contractions Triton's fp-contract=fast produces, and the same exp / log / sqrt lowerings. Pad rows (slot -1) write a zero core row and no state, as stock.
SGLANG_QWEN3_8_DENSE_MONO=1 (default off) runs pure decode steps of up to 8 rows at TP1 on gfx950 through the two kernels: K1 for the GDN layers' front, K2 (with out_proj / o_proj folded) for every layer's FFN half. Attention layers keep the stock attention core; any other step, and anything the checks reject, takes the stock layers. Before decode graph capture the runtime checks that every MXFP4 linear is a plain Quark MXFP4 weight on the Triton path and not already preshuffled, makes the shuffled copies, then compares both kernels with the layers' own stock path on random inputs with torch.equal. Any mismatch turns the whole path off with a warning that names the cause. It also refuses partitioned GPUs (fewer than 256 CUs), hierarchical cache and LMCache, and a GemmaRMSNorm not on its ROCm Triton kernel (needs SGLANG_USE_AITER=1), and keeps the stock prefix-cache state tracking.
K1 and K2 must match the stock decode path bit for bit at widths 1 to 8 on random MXFP4 weights, 512 input draws a width: K2's output and residual; K1's core, residual, conv state and SSM state, on the first layer and with pad rows. The reference norm is GemmaRMSNorm's ROCm Triton kernel; a norm reduced in another order fails a few draws. Runs in the AMD gfx950 CI suite; skipped elsewhere.
d90a42b to
fa95336
Compare
Add gemma_rms_norm_uses_rocm_triton() to layernorm.py and use it in GemmaRMSNorm.forward_hip and in the dense mono guard, instead of reading the module's private flags. The guard also names MXFP4 linears that were dequantized to bf16 at load (AITER_FORCE_A8W4) instead of reporting a shape mismatch. No numerics change.
At TP2 each rank runs K1 on its half of the GDN heads (8 key / 24 value heads). Nothing crosses ranks inside the kernel. Everything else stays stock: the all-reduces, input_layernorm, out_proj, the attention layers and the FFN. K2 stays TP1-only. dense_gdn: - Shapes per rank via Geo(tp). - in_proj_qkvz tiles at 2 groups a CTA, in_proj_ba on 24 CTAs. - GDN tasks numbered over the S x tpt tasks of the step, ceil(1.5 S) a CTA; a dead task never writes. TP1 keeps its index math and timing. - K1Build.normed: start from input_layernorm's output, with every CTA quantizing the rows itself. At TP2 this measured faster than K1's own add + norm prologue, and it is bit-exact with stock either way. Runtime: - The path accepts TP1 or TP2. - At TP2, on a pure decode step of <= 8 rows, each GDN layer's linear_attn runs K1 then the stock out_proj, inside the stock layers. - The K1 self-check runs per rank on the rank's slices, and the ranks agree (all off if any rank is off). - Per-rank shape and state-layout guards. Test: K1 at TP2, ranks 0 and 1 (the column-parallel slices of one layer) from the stock norm's output, widths 1-8, 512 draws each, bit-exact. MI350X TP2, random 8k in / 1k out, TPOT off -> on: c1 7.33 -> 6.25 ms, c2 8.15 -> 6.56, c4 8.57 -> 6.95, c8 9.43 -> 7.88. GSM8K: 0.935 on and off; 200/200 greedy completions identical.
At TP2 each rank runs K2 on its half of the intermediate. It starts from the stock post_attention_layernorm's output and writes down's bf16 partial; the stock all-reduce sums the partials. Nothing crosses ranks inside the kernel. dense_ffn: - Shapes per rank via FGeo(tp). gate_up has 544 groups of 16 rows: 2 whole groups a CTA, plus K quarters of the 32 split groups on CTAs 0..127; the quarter task of CTAs 128..255 reads past the weight and never writes. down's K quarters are 17 steps, split 3 / 2 over the waves. - FfnBuild.normed: every CTA quantizes the normed rows itself (the block qdq is shared with the o_proj core path). - TP1 is unchanged: back-to-back A/B within noise. Runtime: - SGLANG_QWEN3_8_DENSE_MONO_K2 (default on): at TP2, on a pure decode step of <= 8 rows, each layer's mlp runs K2 inside the stock layers. A tuple input (a fused AR + norm + quant) falls back to stock. - set_gdn_step becomes set_step: it gates K1 and K2. - Per-rank weight guards, and a self-check of the rank's partial against the stock mlp at widths 1 and 8; the ranks agree on the outcome (MIN). Test: K2 at TP2, ranks 0 and 1 (the column-parallel slices of one layer), widths 1-8, 256 draws each, bit-exact. Per layer and rank (MI350X, real weights): stock FFN 48-52 us, K2 21-24 us. MI350X TP2, random 8k in / 1k out, TPOT with K1 -> K1 + K2: c1 6.30 -> 4.65 ms, c2 6.56 -> 4.94, c4 6.98 -> 5.31, c8 7.88 -> 6.26. GSM8K 0.935 on and off; 200/200 greedy completions identical to stock.
With SGLANG_QWEN3_8_DENSE_MONO_OPROJ (default on), each GDN layer's out_proj runs at the end of K1 at TP2. K1 then writes the rank's out_proj partial for the stock all-reduce. dense_gdn (K1Build.oproj, from normed rows): - Every CTA stores its core rows past the CU cache and publishes a flag. - CTAs 0..159 poll the 256 flags and qdq the S core rows. - Each then runs its 32 hidden columns of out_proj's K share (dense_ffn's o_proj stage). - The out_proj weights are issued after the publish: issuing them earlier holds back the GDN hand-off polls (vmcnt drains in order). Splits: the fold keeps the stock GEMM's split-K partition at K 3072, i.e. 3 splits of 8 steps, each chained in order and the splits summed in order. fp32 sums are not associative; K quarters mismatched stock by 1 ulp on real weights. dense_ffn's o_loads / o_stage take the width and the split count. Runtime: the GDN mixer returns K1's partial instead of calling out_proj. The self-check compares it with the stock out_proj on the rank's core. prepare's log line no longer tests tensors for truth. Test: K1 at TP2 with the fold, ranks 0 and 1, widths 1-8, 512 draws each, core, states and partial bit-exact; without the fold, rank 0, 128 draws. Per GDN layer and rank (MI350X, real weights): K1 + stock out_proj 26-32 us -> K1 with the fold 20-28 us. MI350X TP2, random 8k in / 1k out, TPOT: - K1 + K2 -> with the fold: c1 4.65 -> 4.33 ms, c2 4.94 -> 4.62, c4 5.31 -> 5.05, c8 6.26 -> 6.01. - Stock -> all on: c1 7.36 -> 4.33 ms, c2 8.16 -> 4.62, c4 8.57 -> 5.05, c8 9.45 -> 6.01. GSM8K 0.935 on and off; 200/200 greedy completions identical to stock.
K2 sums gate_up and down over K in its own partition: K quarters, wave runs and fp32 adds. The stock Triton gemm_afp4wfp4 uses its split-K instead: gate_up 2 splits, down 2 at TP1 and 1 at TP2. The bits still match. - An fp4 x fp4 product scaled by e8m0 codes is exact in fp32. - On real weights, the stock output equals bf16(fp32(exact sum)): 0 of 4.4 M elements differ at TP2 shapes, including heavy-tailed activations. - So any partition whose fp32 partials are exact gives the stock bf16. It would fail only if a row's exact sum spanned more than 24 bits, i.e. products whose block scales differ by more than about 2^14. The partitions would then round differently. That was not seen in 2048 TP2 cases on real weights or in the end-to-end greedy checks, and the load-time self-check guards it per rank. The out_proj fold in dense_gdn reads core rows, whose range is wider, and a partition other than stock's did mismatch there by 1 ulp. So the fold keeps stock's 3 splits. - dense_ffn: the argument as a comment at the K split. - Test: StockAccumulationTest checks the stock GEMM against bf16(fp32(exact fp64 sum)) at K2's K sizes (5120, 17408, 8704), with random MXFP4 weights and randn and heavy-tailed rows. No kernel change.
…o 16 rows At TP2, K1 (GDN front + out_proj fold) and K2 (FFN) ran only on decode steps of at most 8 rows. Both kernels can take up to 16: a token is lane % 16 of the MFMA tile, so 16 rows still fit one N-tile. Each cap is now set on its own, so either kernel can run wider while the other stays at 8: - SGLANG_QWEN3_8_DENSE_MONO_K1_MAX_ROWS - SGLANG_QWEN3_8_DENSE_MONO_K2_MAX_ROWS Both take 1..16 and default to 8, so the default behaviour is unchanged. TP1 stays at 8. Kernel changes: - dense_ffn.max_tokens(tp) is 16 at TP2. MAX_TOKENS (8) still sizes TP1. - K2: down's cross-wave reduce has D_RG * S * 8 units, which is more than the 512 threads above 12 rows. It now runs in rounds. Without that, rows 13..16 come out wrong. - K1: stock's out_proj GEMM (gemm_afp4wfp4, N 5120, K 3072) uses NUM_KSPLIT 1 at M 9..16, against 3 splits at M <= 8. The fold follows it (o_splits(s)). With 3 splits, some output rows differ from stock above 8 rows. Runtime changes: - A step runs mono when it is no wider than the cap of an enabled kernel. Each kernel then runs only on steps within its own cap, and the stock layers run the rest of the step. - The load-time self-check also runs each kernel at its cap (widths 1, 8 and the cap). The ranks still agree on the outcome. - set_step takes the step's row count. Tests: - The TP2 cases of DenseFfnTest and DenseGdnTest cover widths 1..16. - The GDN test pool has 20 slots, so width 16 stays in range. - With the reduce in one round, every width 13..16 case fails on both ranks (256 of 256 draws each).
K1 takes decode steps of up to 8 rows at TP1 and TP2. K2 takes up to 4 rows at TP1 and 8 at TP2. One image serves both TPs without a cap env var. SGLANG_QWEN3_8_DENSE_MONO_K1_MAX_ROWS / _K2_MAX_ROWS still override either cap (1..8 at TP1, 1..16 at TP2); unset means the TP default. Each rank logs its caps at startup. At TP1 a step wider than K2's cap but within K1's runs K1, then the stock out_proj, post_attention_layernorm and MLP. o_proj is folded into K2 only on K2 steps. The defaults were chosen on MI355X with the AITER ASM FP4 GEMM as the stock path. Against the Triton GEMM (MI350X, TP1, 8k/1k, median TPOT at c5..c8): stock 10.24-10.74 ms, K2 cap 4 8.96-9.39 ms, K2 cap 8 6.81-7.24 ms.
The row caps test landed between the skip decorator and StockAccumulationTest. Give each class its own skip.
K2 summed gate_up and down in K quarters and relied on the stock MXFP4 GEMM returning bf16(fp32(exact sum)). That holds on real weights, but a row whose exact sum needs more than 24 bits rounds differently in another K grouping. K2 now follows the stock gemm_afp4wfp4 arithmetic at every width and TP: - the stock split counts (gate_up 2 splits up to 8 rows, else 1; down 2 at TP1 up to 8 rows, else 1; o_proj 4 at TP1), - each split one chain of scaled MFMAs over the 128-K steps in K order from a zero accumulator; a wave's K quarter continues the previous quarter's accumulator through LDS, - the split partials added as tl.sum's pairwise tree, then one bf16 rounding. gate_up's split column group and down's split row group are no longer summed across CTAs: each CTA takes a sliver of them whole (bf16 H for the gate_up sliver columns, re-quantized when down loads H). Tests: StockSplitTest pins the split counts to the stock config's partials. StockChainTest crafts rows with a 2^30-scaled block and its negation next to small blocks, across a quarter boundary and across the split boundary; there the GEMMs summed in K quarters differ from stock at every width (TP1 1..8, TP2 1..16), and K2 matches stock (the previous kernel fails all 24 widths).
The MI355X cap matrix on the a74aa01f stack settled the TP1 defaults at 4 rows for both kernels; TP2 stays K1 8 / K2 8.
MI355X verification of
|
| TP | Parallel | Pairs | Accuracy off / on | Byte-identical generations |
|---|---|---|---|---|
| 1 | 1 | 1 | 0.940 / 0.940 | 200/200 |
| 1 | 4 | 3 | 0.935 / 0.935, 0.940 / 0.940, 0.940 / 0.940 | 199/200, 200/200, 200/200 |
| 2 | 1 | 1 | 0.915 / 0.915 | 200/200 |
| 2 | 8 | 3 | 0.935 / 0.935, 0.915 / 0.915, 0.915 / 0.915 | 200/200 in each pair |
- The one differing question (TP1, parallel 4, first pair) came from servers that shared the node with other servers. Both answers are correct. The four later parallel-4 runs (two off, two on, node otherwise idle) are byte-identical to one another.
- Accuracy moves between pairs for stock and mono alike. Two flag-off servers whose GSM8K runs followed different traffic agree on only 20/200 generations at parallel 4.
Speed: sglang.benchmark.serving, random dataset, range ratio 1.0, seed 1, one warm-up row first. Each value is the mean of the two rounds' medians. Median TPOT differs by at most 0.4% between rounds.
TP1, default caps K1 4 / K2 4:
| Cell (prompts) | Median TPOT off / on (ms) | Change | Output tok/s off / on | Change | Median TTFT off / on (ms) |
|---|---|---|---|---|---|
| 1k/512 c1 (8) | 7.89 / 4.84 | −38.7% | 123.8 / 200.2 | +61.7% | 85 / 84 |
| 1k/512 c2 (16) | 8.63 / 5.22 | −39.6% | 223.2 / 360.5 | +61.5% | 159 / 157 |
| 1k/512 c4 (32) | 8.88 / 5.55 | −37.5% | 425.7 / 659.4 | +54.9% | 173 / 172 |
| 1k/512 c8 (32), wider than the caps | 9.44 / 9.42 | −0.1% | 786.9 / 788.9 | +0.3% | 201 / 200 |
| 8k/1k c4 (16) | 9.43 / 6.08 | −35.5% | 388.3 / 559.0 | +44.0% | 852 / 854 |
| 8k/1k c8 (32), wider than the caps | 10.81 / 10.79 | −0.2% | 667.0 / 667.7 | +0.1% | 1238 / 1245 |
TP2, default caps K1 8 / K2 8:
| Cell (prompts) | Median TPOT off / on (ms) | Change | Output tok/s off / on | Change | Median TTFT off / on (ms) |
|---|---|---|---|---|---|
| 1k/512 c1 (8) | 6.85 / 4.22 | −38.4% | 142.3 / 227.3 | +59.7% | 95 / 95 |
| 1k/512 c2 (16) | 7.45 / 4.43 | −40.5% | 255.3 / 416.0 | +62.9% | 177 / 177 |
| 1k/512 c4 (32) | 7.89 / 4.66 | −40.9% | 482.2 / 787.5 | +63.3% | 198 / 197 |
| 1k/512 c8 (32) | 8.33 / 5.20 | −37.6% | 871.8 / 1324.1 | +51.9% | 305 / 305 |
| 8k/1k c4 (16) | 8.26 / 5.03 | −39.1% | 436.1 / 672.7 | +54.3% | 907 / 908 |
| 8k/1k c8 (32) | 9.07 / 5.94 | −34.6% | 764.2 / 1090.5 | +42.7% | 1296 / 1296 |
Not covered: TP1 with the caps raised to 8/8, MI350X, and the ROCm 10 image.
c8285f9 to
b25e92f
Compare
|
Thanks! Since #44353 landed, model monokernels live in sgl-project/sgl-amd-mono ( This stacks on #43516 and #43481, so it's best ported after or together with them, so the shared kernels land once. If the largest step differs between TP1 and TP2: kernel widths are declared per GPU type today, not per TP size. GLM handles that inside its forward for now, and making widths TP-dependent in the framework is on our list. |
|
@kevin-mii Thanks, done: the port is up as a draft at sgl-project/sgl-amd-mono#28 (spec |
Port the dense K1 (GDN) and K2 (FFN) decode launches from sgl-project/sglang#43558 as spec qwen3_8_dense_hip: TP1 and TP2 on gfx950, decode steps up to 4 (TP1) / 8 (TP2) rows, wider ones on the model's forward. Kernel bodies are unchanged; imports move to sgl_amd_mono.common.
|
@kevin-mii The port is ready for review as sgl-project/sgl-amd-mono#28 (spec |
Port the dense K1 (GDN) and K2 (FFN) decode launches from sgl-project/sglang#43558 as spec qwen3_8_dense_hip: TP1 and TP2 on gfx950, decode steps up to 4 (TP1) / 8 (TP2) rows, wider ones on the model's forward. Kernel bodies are unchanged; imports move to sgl_amd_mono.common.
Port the dense K1 (GDN) and K2 (FFN) decode launches from sgl-project/sglang#43558 as spec qwen3_8_dense_hip: TP1 and TP2 on gfx950, decode steps up to 4 (TP1) / 8 (TP2) rows, wider ones on the model's forward. Kernel bodies are unchanged; imports move to sgl_amd_mono.common.
Motivation
At small batch, each Qwen3.8-27B decode step runs about 21 small launches per GDN layer. On gfx950 that leaves weight streaming well below HBM bandwidth. This PR brings the mono-kernel approach from #43516 (Qwen3.5 MoE) to the dense 27B model:
Modifications
qwen3_5_mono_flydsl/dense_ffn.py), all 64 layers. One persistent launch runs:qwen3_5_mono_flydsl/dense_gdn.py), the 48 GDN layers. One persistent launch runs:srt/layers/qwen3_8_dense_mono.py) plus small hooks inqwen3_5.py.SGLANG_QWEN3_8_DENSE_MONO_K1_MAX_ROWS/_K2_MAX_ROWSoverride them: up to 8 at TP1, 16 at TP2.gemm_afp4wfp4path;is_shuffled/aiter_bpreshuffled), so nothing is shuffled twice;AITER_FORCE_A8W4);SGLANG_USE_AITER=1);torch.equal, at widths 1, min(8, cap) and the cap.SGLANG_QWEN3_8_DENSE_MONO(default off) turns the path on.SGLANG_QWEN3_8_DENSE_MONO_OPROJ(default on): at TP1 it folds out_proj / o_proj into K2; at TP2 it folds each GDN layer's out_proj into K1.SGLANG_QWEN3_8_DENSE_MONO_K1(default on) enables K1.SGLANG_QWEN3_8_DENSE_MONO_K2(default on) runs K2 at TP2.SGLANG_QWEN3_8_DENSE_MONO_K1_MAX_ROWS/_K2_MAX_ROWS(unset: the TP's default caps above).Exact by construction
Every GEMM in K2 reproduces aiter's
gemm_afp4wfp4arithmetic, not only its result. For each width and TP:tl.sum's pairwise tree);The MFMA shape (16x16x128 vs stock's 32x32x64) does not change the bits. The out_proj fold in K1 also uses the stock split at each width: 3 splits up to 8 rows, 1 at 9..16, as aiter's
gemm_afp4wfp4configs choose.StockChainTestbuilds rows whose exact sums need more than 24 bits: a large block and its negation next to small ones, across quarter and split boundaries. The same GEMMs summed in another K grouping differ from stock there, and K2 matches stock at every width.Accuracy
Both kernels are bit-exact with the stock decode path: Triton
gemm_afp4wfp4withdynamic_mxfp4_quant, GemmaRMSNorm's ROCm Triton kernel, SiluAndMul, the Triton conv update, the packed gated delta rule kernel and RMSNormGated. They use the stock kernels' lane layouts, reduction orders and math lowerings on gfx950, including the FMAs that Triton contracts and the rawv_sqrt/v_exp/v_log.test/registered/amd/test_qwen3_8_dense_mono.py, AMD gfx950 CI, skipped elsewhere), on random MXFP4 weights:StockChainTestandStockSplitTest(stock split counts and chains),StockAccumulationTest(stock GEMM == bf16(fp32(exact sum)) at K 5120 / 17408 / 8704),RowCapsTest(default caps by TP and overrides);Benchmark (MI350X,
amd/Qwen3.8-27B-Quark-AWQ-MXFP4, random 8k in / 1k out, median TPOT)TP1:
TP2:
Notes on these tables:
MI355X
TP1, SGLang nightly 20261008 with this PR at
fa95336368(caps 8/8, before the exact-by-construction K2 and the 4/4 default), stock server arguments. Flag on vs off.Median TPOT, random 1024 in / 512 out, 3 runs x 2 rounds with GPUs swapped:
Serving sweep rows (2 rounds, GPUs swapped):
Rows wider than the caps (c32 and above) do not take the path. An image-input row at c64 came out at −2.3% throughput; the weight copy shrinks the KV pool, which is the likely cause (not yet confirmed).
TODO: MI355X TP2, and TP1 at the default caps on this head.
TODO before review
CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.