[heslami]

Motivation

Stacked on the previous PR (shared layer, MoE and NextN entry, itself on sgl-project#75). Merge that first; this PR's diff is its last three commits.

The front of a GLM-5.3 layer launch (q_a | kv_a, 2,624 rows over the 6,144-wide input) ran as 164 whole-K tiles. 164 CTAs each streamed 96 KB of weights while the other 92 waited. Then two parts of the routing window run long. Each top-8 round takes two wave reductions, though the second matters only on exact ties. And the router task normalizes its row before the GEMV, which costs a block reduction and two barriers.

This PR:

• splits each front tile's K between its own CTA and a spare one, so all 256 CTAs stream at most 64 KB;
• reduces a top-8 round's expert ID only when its maximum is tied;
• moves the router's RMSNorm into its logits.

Numerics change (not bitwise):

• The K-split front: each part multiplies bf16(x · gamma) rather than the normalized row, and the consumers scale the sum by the rstd. q_b's input moves by about 3.1e-3 relative, inside the 6.4e-3 that one bf16 ulp on the layer input gives.
• The router fold: only the router logits change.
• The top-8 shortcut is bitwise.

Result on MI355X, the previous PR → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at the golden accept length 2.99, 8K):

• Per cycle: TBD ms at one request, TBD ms at two.
• Kernel-only estimates from development (Speed Tests): about −0.10 ms a cycle at one request and −0.2 ms at two.
• Accuracy: GSM8K with MTP, with an A/A pair, and KL against prefill at TP8: TBD.

Modifications

Three commits, one per change. Each has its own numerics check against its parent (Accuracy Tests) and its own kernel-only estimate (Speed Tests).

#	Commit	Change	Numerics	Kernel-only estimate (TP8, µs a launch, model order, 2 / 4 / 8 rows)
1	K-split front	Each q_a / kv_a tile's first 4096 K run on its own CTA, and the remaining 2048 on one of the 92 CTAs past the tiles, two tiles' rests a CTA. Neither part normalizes: each multiplies bf16(x · gamma). q_norm and the cache tasks add the two partials (q_a2, kv_a2 hold the rests) and scale them by the row's rstd, which they compute from x before their wait. The index projection tiles keep whole K, on the CTAs with one rest, and the NextN entry starts past them. The MoE stages keep their CTAs. A launch from pushed rows (layer 0, NextN) reads them in both parts and in the rstd, with its weights issued before the rows' wait	not bitwise: q_an within 3.3e-3 relative	−0.33 / −1.37 / −2.84
2	top-8 tie shortcut	A round takes the key maximum. When one lane holds it (a ballot), that lane's head is the pick; only an exact tie runs the ID reduction. The picks come out in the same order	bitwise	−0.19 / −0.01 / +0.11
3	router norm in its logits	The router task stores bf16(a · gamma) as its rows arrive, with each wave's sum of squares beside them. One barrier serves both. The rstd, summed in the same order, scales the FP8 activation blocks (bitwise) and the 8 logits	not bitwise: router logits only	−0.27 / −0.09 / −0.30

• The kernel-only estimates come from the development tree's TP8 ladder (E4M3 KV, 8K, 3 rounds in alternating order). They were measured on an older base and with a pad-row skip at 2 rows that is not upstream. The table on this head is TBD.
• Layout: qkv_a has a task on every CTA. q_norm, q_b, split, uk and uv follow at CTA 0, in the order they had before, and o starts at split's first CTA as before. router, ug and down keep the CTAs they had with whole tiles. New scratch: q_a2 and kv_a2.
• The front asserts FP8 block-128 q_a / kv_a weights, which is what GLM-5.3 ships.

Accuracy Tests

The kernel checks ran on one GPU: one process runs rank 0 of a TP8 (or TP4) launch alone, with its peers' partial sums zeros tagged with the launch's epoch. The out-of-tree harness wraps check_layer.py and check_nextn.py unchanged.

• Numerics per commit. A shared and an indexer launch from the same input rows on both trees: widths 2, 4 and 8, one request a row, bf16 and E4M3 MLA cache. The table gives relative L2 distances against the parent commit (q_an, KV latent rows written, router scores, output) and the routed experts that differ, out of 8 a row. Ranges are over 2 / 4 / 8 rows.

  Commit	TP	q_an	KV rows (bf16 / E4M3 cache)	router scores	output	routed experts differing
  1 K-split front	8	3.0–3.3e-3	3.1–3.4e-3 / 0.9–1.4e-2	1.0–2.4e-3	6.7–10.8e-3	0 in 23 of 24 launches; 1 of 64 (bf16 indexer, 8 rows)
  1 K-split front	4	3.1–3.3e-3	3.1–3.5e-3 / 0.9–1.4e-2	1.0–3.1e-3	1.0–1.5e-2	0 in all 24
  2 top-8 shortcut	8, 4	0	0	0	0	0
  3 router fold	8	0	0	1.2–1.4e-3	4.1–5.7e-4	0 in 23 of 24; 1 of 64 (bf16 indexer, 8 rows)
  3 router fold	4	0	0	1.3–1.5e-3	3.5–6.6e-4	0 in 23 of 24; 1 of 64 (bf16 indexer, 8 rows)

  • For scale: one bf16 ulp on the layer input moves q_an by 6.4e-3 and the output by about 4e-2. (Measured at TP8 on the same random layers.)
  • The one routed-expert flip moves that launch's output by 3.7e-2 (TP8) and 4.4e-2 (TP4).
  • The E4M3 KV rows' larger distance is the cache's own rounding: a 3e-3 move of the latent changes some E4M3 values by one step.

• The top-8 shortcut is bitwise: dumps of commit 2 against commit 1 (check_layer.py --dump, the previous PR's set: shared, index, model order, dense; 2 / 4 / 8 rows; p0, p1, v4; both caches; 3 steps): TP8 6,588 equal, 0 differ; TP4 6,588 equal, 0 differ; no error word.

• NextN launch = plain launch fed its own x, bitwise on this head (check_nextn.py --only entry: NextN index at 4 and 8 rows, NextN shared at 2, 4 and 8, with and without pad rows; TP8 and TP4, bf16 and E4M3): 20 / 20 checks pass in each of the 4 runs. Entry against the torch entry: up to 2.8e-5. Against f32: 2.4e-3. The swapped-norm control gives 0.14.

• Serving accuracy: TBD at TP8 (GSM8K with MTP, A/A pair; KL decode ‖ prefill; SGLANG_AMD_MONO_CHECK=1).

  • During development, both non-bitwise changes passed pre-registered rules at TP4 on an older base:
    • GSM8K at two requests: 0.9242 / 0.9196 on the A/A pair of the base, 0.9166 with the K-split, 0.9158 with both. One-sided McNemar p = 0.13 and 0.08 against the base.
    • KL against prefill: 0.0704 / 0.0741 on the base, 0.0702 / 0.0726 with the changes.
    • The serving check logged no violation.
  • These runs were made on a different tree, so they are not this PR's evidence.

• Unit tests on a CPU-only host, in the image: 330 passed, with the same 7 failures and 3 collection errors as [Feat] [MI355X] glm53: The NextN draft steps on their own launches sgl-project/sgl-amd-mono#75 (spec-registration and contract tests that need the SGLang overlay; deepseek_v4, minimax).

Speed Tests and Profiling

Setup as sgl-project#75 (image lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025, squashfs sha256:fb07e0ef78fe, SGLang python/sglang at 5044779f, zai-org/GLM-5.3 at ba32f8c1, FP8 E4M3 KV).

TP8, MTP 3-1-4 at AL 2.99, cycle in ms: TBD (previous PR → this PR, 1 and 2 requests, 8K and 128K).

Kernel-only (time_layers.py), each commit against its parent: TBD.

Registers

The compile survey covers every kind at TP4 and TP8, widths 2 / 4 / 8, both caches, stop on and off: 148 keys a commit. Every key builds on every commit.

• Spills against the previous PR's head:
  • The front removes the 8-row spills of the TP8 indexer and dense builds (20–24 VGPR spills → 0) and of the TP4 8-row indexer and dense builds (stop=True: 20–28 → 0).
  • It adds spills to the NextN indexer builds at 8 rows: TP8 0 → 3 (12 B), TP4 0 → 16–17 (52–56 B).
  • It adds a spill to TP8 embed w4: 0 → 2 (12 B).
• Lint: the head's objects are 148 / 148 clean, and so are commit 2's. Commit 1 has one hit in TP4 nextn_index w4 (E4M3, stop=True): six kernel arguments (s_q_b … w_o) are loaded and never read, the miscompile class of SGPR-spilled kernarg restores. Its SGPR spills go 240 → 756 with the front. The builds are deterministic for this key: a rebuild gives the same object on commit 1, the head and the previous PR's head. The head's object is clean.

Reviewer notes

• The rstd is computed twice: in q_norm and in the cache task, from x, before their waits. The extra x read (12 KB a row) overlaps the front.
• Placement: with all 256 CTAs in the front, q_norm and the cache tasks sit behind a front tile on CTAs 0..S−1. During development this cost the indexer and dense launches 1.4–1.9 µs at 2–4 rows at 8K, while the model order still gained at every width.
• The top-8 shortcut applies in route_top8, which the 4-row per-row up/gate path also uses for tied rows.
• The pad-row skip that the development tree carried at 2 rows is not part of this PR.

Checklist

• pre-commit run --all-files passes. (ruff, ruff-format and isort clean on the touched files; full run TBD)
• Tests under test/ cover the change. (The existing kernel tests; numerics are bounded by the serving check)
• The SGLang commit or PR this was tested against is named above.
• Accuracy and speed results are attached when kernel or forward code changes. (TBD)
• A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)

[heslami]

Superseded by #24 (rebased onto the kernarg-reload fix, #23).