Skip to content

[ROCm][Perf] Qwen3.8-27B dense mono decode: persistent FlyDSL GDN and FFN launches (TP1 and TP2, MI355X/MI350X) - #43558

Closed
vorapolsiloai wants to merge 18 commits into
sgl-project:mainfrom
vorapolsiloai:rocm/qwen3-8-dense-mono-decode
Closed

vorapolsiloai wants to merge 18 commits into
sgl-project:mainfrom
vorapolsiloai:rocm/qwen3-8-dense-mono-decode

Conversation

@vorapolsiloai

@vorapolsiloai vorapolsiloai commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Draft. Stacked on #43516, which includes #43481. The diff below contains both PRs. Review only the last 15 commits (from "Add a FlyDSL dense FFN mono kernel for Qwen3.8-27B decode").

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:

  • at TP1, at most two persistent launches per decoder layer;
  • at TP2, one launch for each rank's share of each layer's GDN mixer and FFN, between the stock all-reduces.

Modifications

  • K2 (qwen3_5_mono_flydsl/dense_ffn.py), all 64 layers. One persistent launch runs:
    • at TP1: out_proj / o_proj and the residual add, post_attention_layernorm, then the MXFP4 gate_up, SiLU-mul and down projections;
    • at TP2: the rank's half of the intermediate (gate_up, SiLU-mul, down), from the stock post_attention_layernorm's output, up to down's partial.
  • K1 (qwen3_5_mono_flydsl/dense_gdn.py), the 48 GDN layers. One persistent launch runs:
    • at TP1: everything from the layer input to K2's input (the residual add and input_layernorm, in_proj_qkvz and in_proj_ba, the conv update, the gated delta rule, the gated norm);
    • at TP2: the rank's half of the GDN heads, from the stock input_layernorm's output, with out_proj folded in, up to the rank's out_proj partial. Nothing crosses ranks inside the kernel.
  • Runtime (srt/layers/qwen3_8_dense_mono.py) plus small hooks in qwen3_5.py.
    • When it runs: pure decode steps no wider than a kernel's row cap, on gfx950, at TP1 or TP2, with PP1 and no spec decode.
      • Every other step runs the stock layers, and attention layers keep the stock attention core.
      • At TP2 the stock layers run the step, and the kernels replace the rank's GDN mixer and FFN. The all-reduces, the norms, attention and its o_proj stay stock.
      • There is no fused all-reduce epilogue in this PR.
    • Row caps by TP: 4 rows for both K1 and K2 at TP1, and 8 rows for both at TP2.
      • SGLANG_QWEN3_8_DENSE_MONO_K1_MAX_ROWS / _K2_MAX_ROWS override them: up to 8 at TP1, 16 at TP2.
      • The chosen caps are logged once per rank.
      • A step wider than one kernel's cap runs that part stock.
      • The TP1 defaults come from an MI355X cap sweep. Set both overrides to 8 to take steps of up to 8 rows.
    • Weights: a (16,16)-shuffled copy of the MXFP4 weights, made before decode graph capture: 10.8 GiB at TP1 (8.9 GiB with K1 off). Prefill keeps the stock GEMMs on the checkpoint layout.
    • Guards. The path stays off and logs the cause in any of these cases:
      • an MXFP4 linear is not a plain Quark MXFP4 weight on the Triton gemm_afp4wfp4 path;
      • a weight is already preshuffled (is_shuffled / aiter_bpreshuffled), so nothing is shuffled twice;
      • a weight was dequantized to bf16 (AITER_FORCE_A8W4);
      • GemmaRMSNorm is not on its ROCm Triton kernel (needs SGLANG_USE_AITER=1);
      • fewer than 256 CUs;
      • hierarchical cache or LMCache.
    • Self-checks. After the shuffle, K1 and K2 are compared with the layers' own stock modules by torch.equal, at widths 1, min(8, cap) and the cap.
      • K2 checks its output and residual (TP1), or this rank's down partial (TP2).
      • K1 checks its core, residual (TP1), conv state and SSM state with a pad row, and with the fold, the out_proj partial.
      • Any mismatch turns the path off with a warning. At TP2 the ranks agree: the path is off on all ranks if any rank is off.
  • Flags:
    • 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_afp4wfp4 arithmetic, not only its result. For each width and TP:

  • it takes the stock config's split-K count: gate_up 2 splits up to 8 rows, else 1; down 2 at TP1 up to 8 rows, else 1; the TP1 o_proj fold 4;
  • it runs each split as one chain of scaled MFMAs over the K steps in order, from a zero accumulator;
  • it adds the split partials in the split-K reduce's order (tl.sum's pairwise tree);
  • it rounds to bf16 once.

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_afp4wfp4 configs choose.

StockChainTest builds 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_afp4wfp4 with dynamic_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 raw v_sqrt / v_exp / v_log.

  • Unit tests (test/registered/amd/test_qwen3_8_dense_mono.py, AMD gfx950 CI, skipped elsewhere), on random MXFP4 weights:
    • widths 1–8 at TP1, and 1–16 on both ranks' slices at TP2;
    • K2: output and residual (TP1), down partial (TP2);
    • K1: core, residual, conv state and SSM state, on the first-layer path and with pad rows (slot -1), plus the folded out_proj partial at TP2;
    • StockChainTest and StockSplitTest (stock split counts and chains), StockAccumulationTest (stock GEMM == bf16(fp32(exact sum)) at K 5120 / 17408 / 8704), RowCapsTest (default caps by TP and overrides);
    • a norm reduced in a different order fails these tests: it changes an MXFP4 code on about 1 row in 640.
  • End to end (GSM8K, 5-shot, greedy):
    • TP1: 0.940 with the flag on and 0.940 with it off, and all 200 generated texts are identical between on and off.
    • TP2 (cap 8): 0.935 on and off, and 200/200 greedy completions are identical to stock. With caps 16: 0.940 (stock 0.940). After the exact-by-construction K2: 0.940 at TP1 and 0.935 at TP2, unchanged.
    • Batched greedy texts at batch 16 cannot be compared directly. Stock is not self-deterministic there (115/128), because step widths depend on prefill admission and stock's GEMM partition changes between M 8 and 9. Every mono configuration agrees with every other at 123–128/128.

Benchmark (MI350X, amd/Qwen3.8-27B-Quark-AWQ-MXFP4, random 8k in / 1k out, median TPOT)

TP1:

flag off flag on (K1 + K2)
TPOT c1 8.87 ms 4.94 ms (−44%)
TPOT c4 10.06 ms 6.01 ms (−40%)
Output tok/s c1 / c4 109 / 367 192 / 582
TTFT median c1 / c4 290 / 888 ms 284 / 913 ms

TP2:

stock mono (cap 8) mono (caps 16)
c1 7.36 ms 4.33 ms (−41%)
c8 9.45 ms 6.01 ms (−36%)
c12 9.13 ms 9.13 ms 7.10 ms (−22%)
c16 9.58 ms 9.57 ms 8.38 ms (−13%)
c32 13.08 ms 13.13 ms 13.09 ms

Notes on these tables:

  • The TP1 table was measured with K1 and K2 at 8 rows.
  • The TP2 c1 / c8 rows are from the fold commit's run, and c12–c32 from the cap-16 screen (one session).
  • Both tables were measured before the exact-by-construction K2. Against the previous K2, it adds +2.0% / +1.0% TPOT at TP1 c1 / c4 and +1.8% / +1.6% / +1.2% at TP2 c1 / c4 / c8 (one round).

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:

c1 c2 c4 c8
flag off 7.95 ms 8.65 ms 8.90 ms 9.36 ms
K2 only 5.90 ms (−25.8%) 6.54 ms (−24.4%) 6.84 ms (−23.1%) 7.57 ms (−19.1%)
K1 + K2 4.75 ms (−40.3%) 5.07 ms (−41.4%) 5.55 ms (−37.6%) 6.63 ms (−29.2%)

Serving sweep rows (2 rounds, GPUs swapped):

Row Throughput TTFT TPOT
8k/1k c4 +52.7% −2.5% −35.8%
8k/1k c8 +35.2% −0.5% −26.8%
58k c4 +14.8% +0.2% −16.7%
125k c4 +9.6% +0.6% −11.3%

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

  • Make K1 bit-exact and add its load-time self-check.
  • Detect weights that are already preshuffled and never shuffle them twice.
  • Add unit tests (bit-exact at widths 1–8 at TP1 and 1–16 at TP2) and register them in AMD CI.
  • TP2: K1 and K2 per rank, out_proj folded into K1, row caps by TP.
  • Make K2 exact by construction (stock split counts and chains).
  • Add MI355X TP1 benchmarks.
  • Add MI355X TP2 benchmarks, and TP1 at the default caps.

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

billishyahao and others added 3 commits October 10, 2026 03:26
…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.
@vorapolsiloai
vorapolsiloai force-pushed the rocm/qwen3-8-dense-mono-decode branch from af1e93e to d90a42b Compare October 10, 2026 13:21
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.
@vorapolsiloai
vorapolsiloai force-pushed the rocm/qwen3-8-dense-mono-decode branch from d90a42b to fa95336 Compare October 10, 2026 20:19
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.
@vorapolsiloai vorapolsiloai changed the title [ROCm][Perf] Qwen3.8-27B dense mono decode: persistent FlyDSL GDN and FFN launches (TP1, gfx950) [ROCm][Perf] Qwen3.8-27B dense mono decode: persistent FlyDSL GDN and FFN launches (TP1 and TP2, MI355X/MI350X) Oct 12, 2026
@jhinpan

jhinpan commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator

MI355X verification of c8285f9189 (TP1 and TP2)

Both kernels engage at TP1 and TP2 on MI355X. Greedy outputs with the flag on are byte-identical to the flag-off server's, and median TPOT drops 35–41% on every row within the row caps. The TP2 rows and the TP1 rows at the default caps (4/4) are the ones still open in the TODO list.

Setup

  • Image rocm/sgl-dev:v0.5.21-rocm720-mi35x-20261009 (sha256:3d9129547126…), with /sgl-workspace/sglang (editable install) checked out at c8285f9189. ROCm 7.2.0, torch 2.9.1+rocm7.2.0, Triton 3.7.0 (git89002410), aiter e7d2453f25, flydsl 0.3.4.1.
  • One 8x MI355X node (gfx950, 256 CUs per GPU), host amdgpu 6.16.13.
  • amd/Qwen3.8-27B-Quark-AWQ-MXFP4 with --tp-size 1 or --tp-size 2 and --trust-remote-code; every other server argument is the default (SGLANG_USE_AITER=1 comes from the image). Only SGLANG_QWEN3_8_DENSE_MONO=1 differs between arms, at the default row caps.
  • Timing runs: one server on the node at a time, 8/8 GPUs idle over 30 one-second amd-smi samples before each server, two rounds with the GPU and the arm order swapped.

Unit tests: test/registered/amd/test_qwen3_8_dense_mono.py passes 9/9 on one MI355X (91 s).

Engagement (flag-on server logs; the load-time self-check passed on every server, and no off or skips a decode step message appeared):

  • TP1: Qwen3.8 dense mono on for decode steps of <= 4 rows at TP1, o_proj folded, FFN K2 (<= 4 rows), GDN K1 on (<= 4 rows) (10.8 GiB of shuffled weights), then step at width 1, 2 and 4 during decode graph capture.
  • TP2, both ranks: Qwen3.8 dense mono on for decode steps of <= 8 rows at TP2, o_proj in K1, FFN K2 (<= 8 rows), GDN K1 on (<= 8 rows) (5.3 GiB of shuffled weights).

Accuracy: GSM8K, 5-shot, 200 questions, greedy (sglang.test.few_shot_gsm8k). Each pair is a flag-off and a flag-on server with the same traffic before the run.

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.

@vorapolsiloai
vorapolsiloai force-pushed the rocm/qwen3-8-dense-mono-decode branch from c8285f9 to b25e92f Compare October 15, 2026 08:25
@kevin-mii

Copy link
Copy Markdown
Collaborator

Thanks! Since #44353 landed, model monokernels live in sgl-project/sgl-amd-mono (--monokernel-backend amd). Could you port the dense K1 (GDN) and K2 (FFN) launches there as a spec, with sgl-project/sgl-amd-mono#6 (GLM-5.3-Flash) as the template? srt/layers/qwen3_8_dense_mono.py and the hooks in qwen3_5.py and kimi_k3.py would go away; the framework handles routing, graph-capture coverage and startup refusals.

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.

@vorapolsiloai

Copy link
Copy Markdown
Contributor Author

@kevin-mii Thanks, done: the port is up as a draft at sgl-project/sgl-amd-mono#28 (spec qwen3_8_dense_hip, following #6). I'm re-running the GPU tests on its rebased head, then I'll mark it ready and close this PR in favour of it.

vorapolsiloai added a commit to vorapolsiloai/sgl-amd-mono that referenced this pull request Oct 22, 2026
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.
@vorapolsiloai

Copy link
Copy Markdown
Contributor Author

@kevin-mii The port is ready for review as sgl-project/sgl-amd-mono#28 (spec qwen3_8_dense_hip, following sgl-amd-mono#6). On MI355X at TP1 and TP2, accuracy matches with the monokernels on and off, per-user decode is about 1.5x up to the row cap, and serving throughput at c4 is up 14-27% with no loss past the cap. Numbers are in that PR. Closing this one in favour of it.

vorapolsiloai added a commit to vorapolsiloai/sgl-amd-mono that referenced this pull request Oct 26, 2026
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.
vorapolsiloai added a commit to vorapolsiloai/sgl-amd-mono that referenced this pull request Oct 30, 2026
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants