[heslami]

Motivation

After #23, a decode step of GLM-5.3-Flash runs 46 launches. Inside a launch, much of the time goes to hand-offs between CTAs: a consumer polls a mailbox, and the poll waits behind whatever its CU has already queued. This PR makes the hand-offs cheaper in three ways. It doesn't change any value: every output and every write is bitwise equal to main at TP2, TP4 and TP8, at every width.

Result on MI355X, 2048 in / 256 out, median TPOT, main → this PR (two rounds each, same allocation):

• TP4: 2.534 → 2.433 ms at one request (−4.0%), 2.558 → 2.458 at 2 (−3.9%), 3.082 → 2.962 at 4 (−3.9%), 4.251 → 4.169 at 8 (−1.9%).
• TP4, 8192 / 256: 2.567 → 2.465 ms at one request (−4.0%), 4.231 → 4.152 at 8 (−1.9%).
• TP2: 3.446 → 3.339 ms at one request (−3.1%), 3.516 → 3.403 at 2 (−3.2%), 4.519 → 4.405 at 4 (−2.5%).
• TP8: 2.253 → 2.140 ms at one request (−5.0%), 2.275 → 2.155 at 2 (−5.3%), 2.773 → 2.613 at 4 (−5.8%), 3.526 → 3.377 at 8 (−4.2%).
• Past the largest width every arm runs SGLang's path; they match within 0.5%.
• Kernel-only (time_layers.py --kind model, µs a launch): 50.1 → 48.0 / 61.7 → 59.1 / 88.7 → 87.3 at TP4, 2 / 4 / 8 rows; 69.8 → 67.1 / 93.1 → 90.7 at TP2; 43.8 → 41.0 / 55.0 → 51.8 / 72.6 → 69.2 at TP8.

Merge order: independent of sgl-project#61 (the first launch's wait deadline), which changes only spec.py. A second perf PR (scalar-polled boxes, plain rows, o_proj first at 2 rows) is stacked on this one.

Modifications

Three changes, one commit each, each with its own measured gain (Speed Tests, "Each change on its own").

1. The launch scratch is uncached (runner.py)

The scratch holds every in-launch hand-off: mailboxes, ready flags, published rows, scores. FlashRuntime now allocates it as a one-rank common.peer.PeerBuffer, the uncached allocation (hipExtMallocWithFlags(..., hipDeviceMallocUncached), 2 MiB granules) that the symmetric peer buffer already uses. common/ doesn't change, and close() frees the buffer.

• The kernel is unchanged: the same device-scope stores and loads, the same code objects.
• A consumer's load is now served at the line's home memory instead of from the producer XCD's L2.
• Memory: the scratch is now outside PyTorch's allocator, so the monokernel's launch-buffer line, which counts PyTorch's allocations, drops from 236.7 to 173.9 MiB at TP4 (508.3 → 473.9 MiB at TP2). The KV pool, which is sized from free device memory, sees it: 9,010,176 → 9,006,336 tokens at TP4 (−0.04%), 7,022,848 → 7,018,752 at TP2 (−0.06%).
• check.py and the kernel tests read rt.scratch through torch, as before.

2. Fewer vector loads ahead of the polls (kernel.py)

Every vector memory instruction holds the CU's in-order queue for about the same time, so loads issued ahead of a poll delay it, whatever their size. Two places now issue fewer:

• Norm weights: a thread's 8 contiguous bf16 weights (4 for the q_a norm) in one load (ld_bf16s), in the gates, the router's own row, the q_a norm and the tail.
• Gate mix partials: a row's mix partials go to LDS two pairs a load (25 loads a CTA instead of 56) and are summed there in the same order. The router that computes its own row keeps its mix sums past them (xs + PM_WORDS).

3. The selected tokens are tagged pairs (kernel.py, layout.py)

The index selection used to drain its tokens and then set the row's tok_ready flag; a split polled the flag and then read its 64 tokens: two round trips in series. Now:

• the selection writes every selected token, and the tail tokens of [Perf] [MI355X] glm53: Above 4 rows run each expert the step uses once in up/gate and down (preview, tables TBD) #27's long-row path (at SEL_STRIDE + TOPK_TOKENS), as tagged pairs (put), with no drain and no flag;
• a split's wave 0 polls its 64 tokens together with the cache's kv_ready, then acquires, as before;
• tok_sel doubles in size and tok_ready goes.

Test change, needed by the new layout: check_layer.py --check-keys reads tok_sel as pairs and also checks that every selected and tail token of a scored row carries the launch's tag (the one kv_ready carries). Negative control: with the tail tokens stored untagged, the check fails on every scored row and rank; it passes on this head.

Every build keeps main's number of guarded wait sites (the split polls its tokens where it polled the flag), so SGLANG_AMD_MONO_FAULT_INJECT's site numbers don't move.

Accuracy Tests

• Bitwise: every output and every write of every launch kind equals main's, byte for byte, at TP2, TP4 and TP8 and every width each TP builds. Kinds: FFN-only, KDA, sparse attention (scored rows past 2,048 tokens with tails, and short rows), dense, embed and tail; FP8 preshuffled and row-major; 3 launches a step (inputs chained), 3 steps, the last row a pad row from 4 rows on. That is 36 / 108 / 216 files and 996 / 2,988 / 5,976 tensors (3.4 / 11.8 / 20.8 GB) at TP2 / TP4 / TP8, compared by SHA-256. So [Fix] [MI355X] glm53: Load a NextN launch's pointer arguments from the kernarg segment at each stage (preview) #23's accuracy results hold.
• SGLANG_AMD_MONO_CHECK=1 (the e2e check case's requests): 874 checks, 0 violations at TP2, TP4 and TP8, on main and on this head. The values aren't equal across runs, because SGLang's reference forward isn't run-to-run reproducible: three runs of main at TP4 differ from each other in 31 to 34 of the 43 reported metrics (min / max over the checks), and this head differs from main in 34. The launches themselves are bitwise equal (above).
• E2E: all five cases of test_glm53_flash_server.py pass at TP4 on this head: GSM8K past the largest width, decode against prefill (KL), the check case, the returned hidden states and the injected fault (with HSA_DISABLE_COREDUMP_ON_EXCEPTION=1 from [Test] glm53/glm53_flash: Turn off crash dumps in the fault tests sgl-project/sgl-amd-mono#55).

Speed Tests and Profiling

Setup as #22 and #23: one node with 8 × MI355X, rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014 (sha256:e4498f5d595f, FlyDSL 0.3.4.1), SGLang ab580b2e06, amd/GLM-5.3-Flash-Quark-MXFP4 at b5688f25, the e2e test's server flags, sglang.benchmark.serving random with fixed lengths, max(16, 4c) prompts after c warm-up requests, two rounds with the second in reverse order. The tables give the mean of the rounds' median TPOT in ms (ITL in parentheses); the rounds agree within 0.5%.

All numbers are from this head (9bad2e8) against main (9484c60).

TP4, 2048 / 256:

Conc.	main	This PR	Change
1	2.534 (2.537)	2.433 (2.436)	−4.0%
2	2.558 (2.555)	2.458 (2.455)	−3.9%
4	3.082 (3.074)	2.962 (2.953)	−3.9%
8	4.251 (4.163)	4.169 (4.071)	−1.9%
16 (SGLang's path)	9.784 (9.565)	9.799 (9.566)	+0.2%

TP4, 8192 / 256:

Conc.	main	This PR	Change
1	2.567 (2.567)	2.465 (2.465)	−4.0%
2	2.592 (2.584)	2.495 (2.486)	−3.7%
4	3.095 (3.088)	2.977 (2.971)	−3.8%
8	4.231 (4.176)	4.152 (4.092)	−1.9%

TP2, 2048 / 256 (4-row cap: 8 and 16 requests run SGLang's path):

Conc.	main	This PR	Change
1	3.446 (3.450)	3.339 (3.343)	−3.1%
2	3.516 (3.513)	3.403 (3.400)	−3.2%
4	4.519 (4.513)	4.405 (4.398)	−2.5%
8	11.988 (11.755)	11.968 (11.772)	−0.2%
16	13.075 (12.628)	13.072 (12.641)	+0.0%

TP8, 2048 / 256:

Conc.	main	This PR	Change
1	2.253 (2.257)	2.140 (2.142)	−5.0%
2	2.275 (2.267)	2.155 (2.151)	−5.3%
4	2.773 (2.768)	2.613 (2.609)	−5.8%
8	3.526 (3.502)	3.377 (3.351)	−4.2%
16 (SGLang's path)	8.050 (7.954)	8.037 (7.953)	−0.2%

Kernel-only (time_layers.py)

torchrun --nproc-per-node <TP> test/kernels/glm53_flash/time_layers.py --kind <kind> --rows 2,4,8 (--pos 640 and --pos 8191 for sparse attention), µs a launch, main → this PR; the mean of two passes in alternating order in one allocation (they agree within 1%). model is #22's 41 launches with sparse-attention rows at position 2,047, full #23's whole step of 46.

TP4:

Launch	2 rows	4 rows	8 rows
model order (41)	50.1 → 48.0 (−4.3%)	61.7 → 59.1 (−4.3%)	88.7 → 87.3 (−1.6%)
full step (46)	49.7 → 48.1 (−3.3%)	60.8 → 58.4 (−4.0%)	86.3 → 85.0 (−1.6%)
KDA	48.1 → 46.5 (−3.1%)	59.0 → 56.5 (−4.3%)	86.0 → 84.1 (−2.2%)
sparse attention, pos 640	53.0 → 50.1 (−5.5%)	65.7 → 62.5 (−4.9%)	91.0 → 87.9 (−3.4%)
sparse attention, pos 8191	67.7 → 63.9 (−5.7%)	80.8 → 76.6 (−5.2%)	110.1 → 106.7 (−3.1%)
dense	39.6 → 38.5 (−2.6%)	45.2 → 43.5 (−3.7%)	55.2 → 53.9 (−2.4%)
embed (layer 0)	40.9 → 39.8 (−2.8%)	46.7 → 45.1 (−3.4%)	59.0 → 57.4 (−2.8%)
tail	57.7 → 57.5 (−0.4%)	61.1 → 60.8 (−0.4%)	64.8 → 64.7 (−0.1%)

TP2:

Launch	2 rows	4 rows
model order (41)	69.8 → 67.1 (−3.8%)	93.1 → 90.7 (−2.7%)
full step (46)	70.3 → 67.9 (−3.3%)	92.1 → 89.7 (−2.6%)
KDA	69.0 → 67.0 (−2.9%)	92.3 → 90.0 (−2.5%)
sparse attention, pos 640	66.2 → 63.6 (−3.9%)	94.3 → 90.4 (−4.1%)
sparse attention, pos 8191	83.0 → 78.6 (−5.3%)	109.1 → 105.9 (−2.9%)
dense	60.7 → 59.2 (−2.6%)	69.1 → 67.0 (−3.0%)
embed (layer 0)	61.7 → 60.4 (−2.1%)	70.4 → 68.4 (−2.8%)
tail	109.0 → 108.8 (−0.2%)	112.4 → 112.1 (−0.3%)

TP8:

Launch	2 rows	4 rows	8 rows
model order (41)	43.8 → 41.0 (−6.5%)	55.0 → 51.8 (−5.8%)	72.6 → 69.2 (−4.6%)
full step (46)	43.0 → 40.6 (−5.6%)	54.0 → 51.1 (−5.3%)	70.6 → 67.5 (−4.5%)
KDA	41.4 → 39.0 (−5.8%)	52.8 → 49.6 (−6.0%)	68.6 → 65.7 (−4.2%)
sparse attention, pos 640	48.2 → 45.1 (−6.4%)	58.0 → 53.9 (−7.1%)	78.8 → 74.6 (−5.3%)
sparse attention, pos 8191	63.4 → 59.8 (−5.7%)	73.9 → 68.3 (−7.5%)	93.9 → 89.0 (−5.2%)
dense	33.7 → 32.8 (−2.6%)	40.1 → 38.8 (−3.2%)	45.9 → 44.2 (−3.8%)
embed (layer 0)	34.6 → 33.9 (−2.2%)	41.7 → 40.5 (−2.8%)	50.9 → 48.7 (−4.3%)
tail	32.5 → 32.5 (−0.2%)	34.1 → 33.9 (−0.7%)	38.1 → 37.5 (−1.5%)

Each change on its own

Each commit against the one before it, in the same allocations as above. Kernel-only --kind model, µs a launch; the token change also at --kind dsa --pos 8191, because the model order's sparse-attention rows sit at 2,047 and aren't scored. Serving: TP4, median TPOT at 1 / 2 / 4 / 8 requests, two rounds.

Change	TP2, 2 / 4 rows	TP4, 2 / 4 / 8 rows	TP8, 2 / 4 / 8 rows	TP4 serving, 2048 / 256	TP4 serving, 8192 / 256
1. uncached scratch	−1.87 / −1.76	−1.35 / −1.90 / −1.27	−2.11 / −2.25 / −2.61	−2.0 / −1.9 / −2.3 / −1.0%	–
2. fewer vector loads	−0.71 / −0.75	−0.63 / −0.62 / −0.40	−0.85 / −1.01 / −0.77	−1.2 / −1.2 / −1.2 / −0.8%	−3.3 / −3.2 / −3.5 / −1.5% (1 and 2 together)
3. tagged tokens, model order	−0.09 / +0.03	−0.16 / −0.12 / +0.25	+0.12 / +0.10 / +0.06	−0.9 / −0.8 / −0.5 / −0.2%	−0.7 / −0.5 / −0.3 / −0.3%
3. tagged tokens, DSA at 8191	−0.86 / −1.36	−0.94 / −1.17 / −1.10	−0.23 / −1.16 / −1.47

Registers

A compile-only build of every key the spec and the tests build (TP2, TP4 and TP8; FFN-only, KDA, sparse attention, dense, embed, tail; every width; FP8 preshuffled, and row-major for sparse attention, dense and embed) gives 72 distinct symbols, as on main.

• No preshuffled build (what serving uses) spills VGPRs.
• The row-major sparse-attention build at TP2, 4 rows spills 1 VGPR (8 bytes of scratch); on main it uses all 256 VGPRs and spills none.
• The FFN-only build at TP4, 8 rows spills 8 VGPRs here and on main alike when built with stop=True, which neither the spec nor the tests use.

Tests

• Unit tests on a CPU-only host: the same results as main on SGLang ab580b2e06 (257 passed; the deepseek_v41 and minimax_m31 contract cases and the two minimax spec modules fail there, as on main).
• test_glm53_flash_kernels.py with 8 GPUs (TP2, TP4 and TP8): 8 passed (17 min 54 s).
• check_layer.py --check-keys with the tag check passes at TP4; its negative control fails (above).
• Out of tree, not part of this PR: a mixed-width stress. One graph holds a step of every width the TP builds, 12 launches each with its own weights and state or caches, replayed 30,000 times; after every step every output, state and cache write and scratch word is compared with the first run's. KDA, sparse attention at 640 and 8191, dense, embed and tail, at TP2 (2 and 4 rows), TP4 and TP8 (2, 4 and 8): 0 bad words and 0 error words, 0.72 / 1.08 / 1.08 M launches a rank per kind. With one input bit flipped for a batch, the stress reports it (positive control).

Checklist

• pre-commit run --all-files passes.
• Tests under test/ cover the change.
• The SGLang commit or PR this was tested against is named above.
• Accuracy and speed results are attached when kernel or forward code changes.
• A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)