Repository navigation
[Perf] [MI355X] glm53_flash: Scalar-polled boxes for broadcast rows, and o_proj before the state write-back at 2 rows - #63
Merged
Conversation
added 6 commits
October 25, 2026 15:32
The scratch holds every in-launch hand-off (mailboxes, ready flags, the published rows). It is now a one-rank PeerBuffer, the uncached allocation the symmetric peer buffer already uses, so a consumer's load is served at the line's home memory instead of through the producer XCD's L2. close() frees it.
Every vector memory instruction takes its turn in the CU's in-order queue, so loads issued ahead of a poll delay it whatever their size: - a thread's contiguous bf16 norm weights in one load (the gates, the router's own row, the q_a norm, the tail); - a gate's mix partials two pairs a load into LDS, summed there in the same order.
The index selection drained its tokens and set the row's tok_ready flag, and a split task polled the flag, then read its 64 tokens: two round trips in series. The selected and tail tokens are now tagged pairs, with no drain and no flag, and a split's wave 0 polls its 64 tokens together with the cache's kv_ready before its acquire. check_layer.py --check-keys reads tok_sel as pairs and checks that every selected and tail token of a scored row carries the launch's tag.
A box gives each consumer CTA its own 128-byte lines, a word a producer.
After its payload stores, a producer's wave stores the launch tag into its
word of every CTA's line; the consumer waits on its line with one scalar
glc load (s_load_dword{xN}), inside the guarded spin. Scalar loads count in
lgkmcnt, so the wait does not queue behind the CU's vector loads, which the
waiting tasks have in flight.
The gates' rows (box_a_in, box_a) replace the a_in_rdy and a_rdy flags, and
the kda heads' box (box_o_in) replaces the KDA o_proj's sampled poll of
o_in. The payloads stay tagged, so a word seen early costs a re-poll, not a
value.
The attention gate publishes its row (a_in) as plain packed bf16 pairs and drains its stores before it sets the row's box word, so inproj and qkv_a read the row once after the box, half the bytes of the tagged pairs, in 16-byte loads.
…at 2 rows At 2 rows on TP4 and TP8 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. The o_proj weights are then issued right behind the output, and the 64 KB state write-back leaves the path. TP2, where it measured slower, other widths and the sparse-attention launch keep their code.
This was referenced Oct 27, 2026
jhinpan
approved these changes
Oct 27, 2026
Collaborator
There was a problem hiding this comment.
Approving, with one text correction. I reproduced the claims against #62 and on current main (with #19).
-
Bitwise: every output and every write equals #62's: 72 of 72 configs (
stop=False), and 48 of 48 on main (stop=True). -
Code objects: against commit 2 (
e50ee7e), only 12 of 72 builds change: KDA, dense and embed at 2 rows on TP4 and TP8. This agrees with the PR text. -
Tests: on main + #61 + this PR, kernel tests: 13 of 13 pass. E2E file at TP4: 5 of 5 pass, on this head and on that merge.
-
Speed: median TPOT at TP4, 2048 / 256, two rounds:
Conc. #62 → this PR main + #62 → main + this PR 1 2.475 → 2.406 ms (−2.8%) 2.426 → 2.354 ms (−3.0%) 4 3.015 → 2.951 ms (−2.1%) 2.950 → 2.903 ms (−1.6%) 8 4.242 → 4.129 ms (−2.7%) 4.185 → 4.089 ms (−2.3%) 16 (SGLang's path) no change no change
Please correct
- The Registers section says that the extra SGPR spills do not go to memory. On this head, the TP8 4-row sparse-attention builds report a 36-byte private segment (#62: 0). On main + this PR, it is 0. Please state this.
Nits
- Shorten the 4-line
OPROJ_IN_KDAcomment. The PR body gives the reason. - The second PR of a stack can end its title with
(2/2).
4 tasks done
4 tasks done
4 of 5 tasks
This was referenced Nov 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Stacked on #62 (
[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 #62'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 (#62) → this PR (two rounds each, same allocation):
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%).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).signal: 4 stores a lane for 256 CTAs).glcload (await_box), inside the guarded spin like every other wait (common.epoch.spin_untilunderguard.scope(launch_guard)). Scalar loads count inlgkmcntand don't queue behind the CU's vector loads. The polling wave is chosen wave-uniformly, since scalar code runs whatever the EXEC mask.sloads:s_nop 4; s_load_dword{xN} ... glc; s_waitcnt lgkmcnt(0)) throughllvm.inline_asm, ascommon.ops.freshdoes. They are the package's first scalar memory loads.box_a_in, replacing thea_in_rdyflag) for inproj and qkv_a; the FFN gate's row (box_a, replacinga_rdy) for the dense MLP, the router and the expert inputs;box_o_in) for o_proj, replacing wave 0's sampled poll ofo_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).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, sincecheck.pyreads it as pairs.This hand-off's release is the drain: device-scope (
sc1) stores,vmcnt(0), then the box word; the reader loads withsc1.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.
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
SGLANG_AMD_MONO_CHECK=1: 874 checks, 0 violations at TP2, TP4 and TP8. As in [Perf] [MI355X] glm53_flash: Cheaper hand-offs: uncached scratch, fewer vector loads ahead of the polls, tagged selected tokens #62, the values move between runs asmain's do (this head differs from the base in 32 of 43 metrics, three runs ofmainfrom each other in 31 to 34); the launches are bitwise equal.test_glm53_flash_server.pypass 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 (withHSA_DISABLE_COREDUMP_ON_EXCEPTION=1from [Test] glm53/glm53_flash: Turn off crash dumps in the fault tests #55).Speed Tests and Profiling
Setup as #62 (image
sha256:e4498f5d595f, SGLangab580b2e06, checkpointb5688f25, the e2e server flags, two rounds with the second reversed). All numbers are from this head (f79e3cc) against the base (9bad2e8).TP4, 2048 / 256:
TP4, 8192 / 256:
TP2, 2048 / 256:
TP8, 2048 / 256:
Kernel-only (
time_layers.py)µs a launch, base → this PR, the mean of two passes in alternating order in one allocation.
TP4:
TP2:
TP8:
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.a_inCommit 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 #62. SGPR spills (to VGPR lanes, not memory) rise on some sparse-attention builds, e.g. 184 → 375 at TP4, 8 rows.
Tests
main.test_glm53_flash_kernels.pywith 8 GPUs (TP2, TP4 and TP8): 8 passed (18 min 9 s).Checklist
pre-commit run --all-filespasses.test/cover the change..github/CODEOWNERSthat names its owner. (No new model.)