[heslami]

Motivation

Stacked on #PERF1 ([Perf] [MI355X] glm53_flash: Cheaper hand-offs: uncached scratch, ...). Merge that one first; this PR's diff is its last three commits. The boxes need #PERF1's uncached scratch: scalar loads of another XCD's device-scope stores were tested only on uncached memory.

Some hand-offs still wait behind vector loads the waiting task already has in flight: the gates' rows, which every CTA reads, and the KDA heads' outputs, which every o_proj CTA reads. This PR polls them with scalar loads, which don't queue behind vector loads; reads the attention gate's row once, as plain words; and at 2 rows on TP4 and TP8 runs a KDA CTA's o_proj task before its state write-back. Every output and every write stays bitwise equal to the base, at TP2, TP4 and TP8 and every width.

Result on MI355X, median TPOT, base (#PERF1) → this PR (two rounds each, same allocation):

• TP4, 2048 / 256: 2.433 → 2.370 ms at one request (−2.6%), 2.458 → 2.392 at 2 (−2.7%), 2.962 → 2.903 at 4 (−2.0%), 4.169 → 4.049 at 8 (−2.9%).
• TP4, 8192 / 256: 2.465 → 2.394 ms at one request (−2.9%), 4.152 → 4.043 at 8 (−2.6%).
• TP2, 2048 / 256: 3.339 → 3.285 ms at one request (−1.6%), 3.403 → 3.350 at 2 (−1.6%), 4.405 → 4.303 at 4 (−2.3%).
• TP8, 2048 / 256: 2.140 → 2.078 ms at one request (−2.9%), 2.155 → 2.094 at 2 (−2.8%), 2.613 → 2.566 at 4 (−1.8%), 3.377 → 3.330 at 8 (−1.4%).
• With #PERF1, against main: 2.534 → 2.370 ms at TP4 and one request (−6.5%), 3.446 → 3.285 at TP2 (−4.7%), 2.253 → 2.078 at TP8 (−7.8%).
• Kernel-only (time_layers.py --kind model, µs a launch): 48.0 → 46.5 / 59.1 → 57.6 / 87.3 → 84.5 at TP4, 2 / 4 / 8 rows.

Modifications

Three changes, one commit each, each measured on its own (Speed Tests).

1. Scalar-polled boxes for the broadcast rows (kernel.py, layout.py)

A box gives each consumer CTA its own 128-byte lines, one word a producer, 16 producers a line (one s_load_dwordx16).

• Producer: after its payload stores, one wave stores the launch tag into its word of every CTA's line (signal: 4 stores a lane for 256 CTAs).
• Consumer: waits on its line with one scalar glc load (await_box), inside the guarded spin like every other wait (common.epoch.spin_until under guard.scope(launch_guard)). Scalar loads count in lgkmcnt and don't queue behind the CU's vector loads. The polling wave is chosen wave-uniformly, since scalar code runs whatever the EXEC mask.
• The loads are inline assembly (sloads: s_nop 4; s_load_dword{xN} ... glc; s_waitcnt lgkmcnt(0)) through llvm.inline_asm, as common.ops.fresh does. They are the package's first scalar memory loads.
• Where:
  • the attention gate's row (box_a_in, replacing the a_in_rdy flag) for inproj and qkv_a; the FFN gate's row (box_a, replacing a_rdy) for the dense MLP, the router and the expert inputs;
  • the KDA heads' outputs (box_o_in) for o_proj, replacing wave 0's sampled poll of o_in (one pair in 32). There are S × heads producers, and wave w polls line w: 4 / 8 lines at TP2, 2 / 4 / 8 at TP4, 1 / 2 / 4 at TP8 (asserted at build).
• Why it's correct: the boxed payloads stay tagged pairs, so a box word seen early only costs a re-poll of the payload, never a wrong value. Box words are launch tags and are never cleared, like the mailboxes.
• Memory: 32 KB a line of every CTA (256 × 128 bytes): 32 KB for each gate row's box, 32 to 256 KB for box_o_in, in the scratch.

2. The attention gate's row as plain words (kernel.py)

The attention gate publishes its row (a_in) as plain packed bf16 pairs, drains its stores (s_waitcnt vmcnt(0)), then sets its box word. inproj and qkv_a read the row once after the box, in 16-byte loads (stage_plain): half the bytes of the tagged pairs, and no tags to check. The FFN gate's row (a) stays tagged, since check.py reads it as pairs.

This hand-off's release is the drain: device-scope (sc1) stores, vmcnt(0), then the box word; the reader loads with sc1.

3. o_proj before the state write-back at 2 rows, on TP4 and TP8 (kernel.py)

At 2 rows a KDA task publishes the head's output, runs its CTA's o_proj task, then stores the recurrent and conv states; the conv inputs wait in the top LDS words, past o_proj's rows. The o_proj weights are then issued right behind the output, and the 64 KB state write-back leaves the path. A CTA that ran its o_proj task inside its KDA task (at most one KDA task a CTA at 2 rows) skips the o_proj stage.

• It changes the KDA, dense and embed launches at 2 rows.
• Not at TP2: there, without the exclusion, the 2-row KDA launch was 0.5 µs slower and TPOT at one request 0.5% higher, so TP2 keeps change 2's code (every TP2 code object equals change 2's).
• The 4- and 8-row and the sparse-attention code objects are byte-identical to change 2's.

The guarded wait sites change with the waits (fault injection numbers them in order), e.g. 64 → 60 for TP4's 8-row sparse-attention build. The e2e fault case passes.

Accuracy Tests

• Bitwise: every output and every write of every launch kind equals the base's, byte for byte, at TP2, TP4 and TP8, every width, with the dump set of #PERF1 (36 / 108 / 216 files, 996 / 2,988 / 5,976 tensors at TP2 / TP4 / TP8). 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: 874 checks, 0 violations at TP2, TP4 and TP8. As in #PERF1, the values move between runs as main's do (this head differs from the base in 32 of 43 metrics, three runs of main from each other in 31 to 34); the launches are bitwise equal.
• 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 #PERF1 (image sha256:e4498f5d595f, SGLang ab580b2e06, checkpoint b5688f25, the e2e server flags, two rounds with the second reversed). All numbers are from this head (f79e3cc) against the base (9bad2e8).

TP4, 2048 / 256:

Conc.	base	This PR	Change
1	2.433 (2.436)	2.370 (2.371)	−2.6%
2	2.458 (2.455)	2.392 (2.388)	−2.7%
4	2.962 (2.953)	2.903 (2.898)	−2.0%
8	4.169 (4.071)	4.049 (3.963)	−2.9%
16 (SGLang's path)	9.799 (9.566)	9.783 (9.568)	−0.2%

TP4, 8192 / 256:

Conc.	base	This PR	Change
1	2.465 (2.465)	2.394 (2.394)	−2.9%
2	2.495 (2.486)	2.423 (2.412)	−2.9%
4	2.977 (2.971)	2.915 (2.910)	−2.1%
8	4.152 (4.092)	4.043 (3.986)	−2.6%

TP2, 2048 / 256:

Conc.	base	This PR	Change
1	3.339 (3.343)	3.285 (3.288)	−1.6%
2	3.403 (3.400)	3.350 (3.347)	−1.6%
4	4.405 (4.398)	4.303 (4.293)	−2.3%
8 (SGLang's path)	11.968 (11.772)	11.976 (11.771)	+0.1%
16	13.072 (12.641)	13.055 (12.618)	−0.1%

TP8, 2048 / 256:

Conc.	base	This PR	Change
1	2.140 (2.142)	2.078 (2.081)	−2.9%
2	2.155 (2.151)	2.094 (2.086)	−2.8%
4	2.613 (2.609)	2.566 (2.559)	−1.8%
8	3.377 (3.351)	3.330 (3.298)	−1.4%
16 (SGLang's path)	8.037 (7.953)	8.044 (7.942)	+0.1%

Kernel-only (time_layers.py)

µs a launch, base → this PR, the mean of two passes in alternating order in one allocation.

TP4:

Launch	2 rows	4 rows	8 rows
model order (41)	48.0 → 46.5 (−3.2%)	59.1 → 57.6 (−2.4%)	87.3 → 84.5 (−3.1%)
full step (46)	48.1 → 46.3 (−3.7%)	58.4 → 57.0 (−2.3%)	85.0 → 82.2 (−3.2%)
KDA	46.5 → 44.3 (−4.9%)	56.5 → 55.0 (−2.7%)	84.1 → 81.7 (−2.8%)
sparse attention, pos 640	50.1 → 49.8 (−0.6%)	62.5 → 62.1 (−0.7%)	87.9 → 86.2 (−2.0%)
sparse attention, pos 8191	63.9 → 63.6 (−0.4%)	76.6 → 76.1 (−0.7%)	106.7 → 104.9 (−1.6%)
dense	38.5 → 35.6 (−7.7%)	43.5 → 41.7 (−4.1%)	53.9 → 51.2 (−5.0%)
embed (layer 0)	39.8 → 37.2 (−6.4%)	45.1 → 43.4 (−3.8%)	57.4 → 54.9 (−4.4%)
tail (same code objects)	57.5 → 57.8 (+0.5%)	60.8 → 61.2 (+0.7%)	64.7 → 65.4 (+1.1%)

TP2:

Launch	2 rows	4 rows
model order (41)	67.1 → 66.2 (−1.4%)	90.7 → 88.9 (−1.9%)
full step (46)	67.9 → 67.1 (−1.3%)	89.7 → 88.0 (−1.9%)
KDA	67.0 → 66.0 (−1.6%)	90.0 → 87.1 (−3.1%)
sparse attention, pos 640	63.6 → 63.3 (−0.4%)	90.4 → 89.2 (−1.3%)
sparse attention, pos 8191	78.6 → 78.1 (−0.6%)	105.9 → 105.5 (−0.3%)
dense	59.2 → 57.9 (−2.1%)	67.0 → 64.2 (−4.2%)
embed (layer 0)	60.4 → 59.0 (−2.3%)	68.4 → 65.7 (−3.9%)
tail (same code objects)	108.8 → 109.0 (+0.3%)	112.1 → 112.6 (+0.5%)

TP8:

Launch	2 rows	4 rows	8 rows
model order (41)	41.0 → 39.5 (−3.5%)	51.8 → 50.8 (−2.0%)	69.2 → 68.2 (−1.5%)
full step (46)	40.6 → 39.2 (−3.6%)	51.1 → 49.8 (−2.5%)	67.5 → 66.2 (−1.9%)
KDA	39.0 → 37.3 (−4.3%)	49.6 → 48.3 (−2.8%)	65.7 → 64.6 (−1.7%)
sparse attention, pos 640	45.1 → 44.2 (−1.9%)	53.9 → 53.6 (−0.6%)	74.6 → 73.5 (−1.5%)
sparse attention, pos 8191	59.8 → 59.0 (−1.4%)	68.3 → 67.8 (−0.8%)	89.0 → 87.6 (−1.7%)
dense	32.8 → 30.9 (−5.7%)	38.8 → 37.0 (−4.6%)	44.2 → 43.0 (−2.6%)
embed (layer 0)	33.9 → 32.1 (−5.2%)	40.5 → 38.6 (−4.9%)	48.7 → 47.3 (−2.7%)
tail (same code objects)	32.5 → 32.2 (−0.9%)	33.9 → 33.7 (−0.6%)	37.5 → 37.3 (−0.5%)

The tail's code objects are the same in both trees, so its column shows the run-to-run spread (up to 1.1% at TP4 and 8 rows).

Each change on its own

Each commit against the one before it, same allocations. Kernel-only --kind model, µs a launch; 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
1. boxes	−0.42 / −1.09	−0.43 / −0.65 / −0.94	+0.11 / −0.79 / −0.39	−0.5 / −0.5 / −0.9 / −1.1%
2. plain a_in	−0.70 / −0.56	−0.34 / −0.88 / −1.77	−0.43 / −0.30 / −0.63	−0.5 / −0.4 / −1.0 / −1.7%
3. o_proj first, 2 rows	off (same code objects; +0.15 / −0.06 measured)	−0.76 / – / –	−1.12 / – / –	−1.7 / −1.8 / – / –%

Commit 3 per kind at 2 rows: KDA −1.36 (TP4) / −1.47 (TP8), dense −2.05 / −1.75 µs a launch.

Registers

72 distinct symbols over every key the spec and the tests build. No preshuffled build spills VGPRs. The row-major sparse-attention build at TP2, 4 rows spills 1 VGPR, as in #PERF1. SGPR spills (to VGPR lanes, not memory) rise on some sparse-attention builds, e.g. 184 → 375 at TP4, 8 rows.

Tests

• Unit tests on a CPU-only host: as main.
• test_glm53_flash_kernels.py with 8 GPUs (TP2, TP4 and TP8): 8 passed (18 min 9 s).
• Out of tree, not part of this PR: #PERF1's mixed-width stress on this PR's launches (KDA, sparse attention at 640 and 8191, dense, embed and tail; 30,000 replays of a graph with every width the TP builds, every word compared after every step): 0 bad words and 0 error words at TP2, TP4 and TP8. This covers the o_proj-first 2-row KDA, dense and embed objects. With one input bit flipped for a batch, the stress reports it (positive control, KDA and dense).

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.)