Skip to content

[Test] glm53/glm53_flash: Turn off crash dumps in the fault tests - #55

Merged
jhinpan merged 1 commit into
sgl-project:mainfrom
jhinpan:glm53-fault-coredump
Oct 25, 2026
Merged

[Test] glm53/glm53_flash: Turn off crash dumps in the fault tests#55
jhinpan merged 1 commit into
sgl-project:mainfrom
jhinpan:glm53-fault-coredump

Conversation

@jhinpan

@jhinpan jhinpan commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator

Motivation

The fault case of each GLM e2e file (test_a_stalled_rank_stops_the_server_without_tokens) makes one rank stall, so its launch stops with a GPU fault. The case then waits 300 s for the server to exit. On the MI355X images, two crash dumps can take longer than that:

Modifications

  • test/e2e/test_glm53_flash_server.py and test/e2e/test_glm53_server.py: the fault case's server environment adds HSA_DISABLE_COREDUMP_ON_EXCEPTION=1, SGLANG_PYSPY_DUMP_BEFORE_CRASH=0 and SGLANG_CUDA_COREDUMP_BEFORE_CRASH=0.

Accuracy Tests

Not applicable: test change only.

Speed Tests and Profiling

Not applicable: test change only.

Each fault case ran from this branch with none of the three variables set outside the test. Both passed. No core file was written to the working directory, and the server log no longer shows the crash diagnostics:

Case Setup Result
GLM-5.3-Flash, TP4 rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014, SGLang ab580b2e06, amd/GLM-5.3-Flash-Quark-MXFP4@b5688f25 passed in 2 min 29 s
GLM-5.3, TP8 lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019, SGLang e45cbce599, zai-org/GLM-5.3@ba32f8c1 passed in 2 min 50 s; failed after 13 min 47 s with only the core dumps off

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. (No kernel or forward code changes.)
  • A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)

Review and Merge Process

  1. The lint and Check PR title checks must pass.
  2. Ask a code owner from .github/CODEOWNERS to review.
  3. Merge with "Squash and merge": the PR title becomes the commit title on main.

@jhinpan
test(glm53): turn off crash dumps in the fault tests
79ac2a1 
The fault case of each GLM e2e file stops one rank with a GPU fault, then waits 300 s for the server to exit. The crash dumps can take longer: ROCr writes each faulting rank's VRAM as a core dump to the working directory, and SGLang's SIGQUIT handler runs py-spy on every scheduler, with no timeout, and then waits 60 s. Set HSA_DISABLE_COREDUMP_ON_EXCEPTION=1, SGLANG_PYSPY_DUMP_BEFORE_CRASH=0 and SGLANG_CUDA_COREDUMP_BEFORE_CRASH=0 in the fault case's server env.
@jhinpan jhinpan mentioned this pull request Oct 25, 2026
Merged
4 of 5 tasks
@jhinpan
jhinpan marked this pull request as ready for review October 25, 2026 00:12
@jhinpan
jhinpan merged commit 876ed7f into sgl-project:main Oct 25, 2026
3 checks passed
kevin-mii pushed a commit that referenced this pull request Oct 25, 2026
Kevin Mi @claude
[Test] minimax_m31: Serving tests for the FFN launches
4d75ee7 
Like every other model's: the DSpark verify server (2 requests of 8 rows at
the largest width) and the decode server (16 rows) show their tested batches
in the startup line and pass chat GSM8K past the largest width (0.975 /
0.985 and 0.975; the standard path scores 0.975; 8-shot completion prompts
score under 0.1 on either path, so the gate is chat). Check mode holds every
launch to NUMERICS, and a rank stalled by SGLANG_AMD_MONO_FAULT_INJECT stops
the server with faults at the null page only and no tokens sent; that test
turns crash dumps off, as #55 did for GLM's.

6 passed (MI350X, TP4, sgl-project/sglang#45936).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@jhinpan jhinpan mentioned this pull request Oct 25, 2026
Merged
4 of 5 tasks
jhinpan added a commit that referenced this pull request Oct 25, 2026
@jhinpan
[Perf] [MI355X] glm53: Select a long row's indexer top-k on up to 16 …
240d587 
…CTAs (#24)

## Motivation

**Summary.** With long contexts, the GLM-5.3 monokernels were slower
than the normal path, because the exact top-2048 selection of each
indexer row ran on one CTA. This PR selects the top-k of a row of more
than 32768 keys on up to 16 CTAs. The result does not change: it is the
exact top-2048, as before. On 8x MI355X at TP8, in InferenceX agentic
coding with MTP, the monokernels now decode faster than the normal path
with MTP: interactivity p50 is 15% higher at concurrency 1, 17% at 2 and
11% at 4. With #15 alone, it was 5%, 5% and 14% lower.

**On main.** This PR is two commits on main (`876ed7f`): `6bfceb6` is
`48ff491` rebased with the same changed lines (and `48ff491` is
`4f187f6` rebased), and `235bb05` is the fix from the review ("Review
fix" below). Apart from that fix, `gfx950/kernel.py`, `layout.py` and
the `common` modules they import have the same code as at `4f187f6`, so
the results below measured at `4f187f6` and `8068147` still apply. The
fix's own results are marked.

**Problem.**

- InferenceX agentic coding (AgentX) sends long requests. In the runs
below, the median input of a request is 89k to 285k tokens.
- GLM-5.3 runs the indexer on 21 of its 78 layers. The monokernels run
18 of them (layers 3 to 77 are covered; layers 0 to 2 are dense and use
the normal path).
- Without MTP, at sgl-amd-mono `e952b22`, the monokernels lost to the
normal path at concurrency 8: interactivity p50 -14%, p90 -42%. With MTP
(#15), they lost 5% to 14% at concurrency 1 to 4 (table in "Speed Tests
and Profiling").

**Cause.** A debug build (not in this PR) records the start and the end
of each stage task of each launch. In an indexer launch:

- The exact top-2048 selection of a row ran on one CTA. That CTA reads
all scores of the row at least two times, at about 0.8 ns a key: 211 us
a launch at 256k keys.
- The sparse MLA and the MoE of the layer wait for the selection. The
other CTAs wait idle.
- The launches without the indexer (57 of the 75 in a step) take the
same time at all contexts.

| Keys per row | Rows | Selection (us) | Indexer launch (us) | Launch
without indexer (us) | Estimated step (ms) |
|---|---|---|---|---|---|
| 8k | 1 | 31 | 97 | 53 | 4.79 |
| 256k | 1 | 211 | 325 | 54 | 8.91 |
| 8k | 4 | 27 | 136 | 72 | 6.55 |
| 256k | 4 | 237 | 500 | 72 | 13.11 |

The estimated step is 18 indexer launches and 57 launches without the
indexer. From 8k to 256k keys, it grows 4.1 ms at 1 row and 6.6 ms at 4
rows. The measured decode step grows 4.1 ms and 6.3 ms.

## Modifications

- `glm53/gfx950/layout.py`
- `INDEX_SEL_CHUNK = 4096` (the keys of a select chunk: 512 threads x
8), `INDEX_SEL_PAR_CHUNKS = 8`, `INDEX_SEL_SLICES = 16` and
`INDEX_CAND_STAGE_CAP = 16384`.
- `select_slices(bound)` gives the slices of a row and the chunks of
each slice. The kernel and the CPU test call the same function.
- Three stages before `index_select`: `index_hist` (16 tasks a row),
`index_thresh` (1 a row) and `index_gather` (16 a row), and their
scratch: about 2.4 MiB a row, 33 MiB a rank for widths 2, 4 and 8.
- `glm53/gfx950/kernel.py`: a row of more than 32768 keys (8 chunks of
4096) selects on up to 16 CTAs of the upper half of the grid (CTAs 128
to 255), each with a run of whole chunks. Before this change, these CTAs
waited idle for the selection.
1. `index_hist`: each slice (CTA) makes a 4096-bin histogram of its
chunks in LDS (the top 12 bits of each score key) and writes it to
scratch. A release fence orders the histogram before its ready flag.
2. `index_thresh` runs on the CTA of the row's `index_select`. It adds
the histograms of the slices, finds the threshold bin, and publishes the
bin, the keys still needed from the bin and the count in the bin.
3. `index_gather`: each slice writes, in position order, its positions
above the bin (winners) and its (key, position) pairs in the bin
(candidates) to scratch, and publishes the two counts.
4. `index_select` copies the winners and the candidates of the slices,
in slice order, into the lists that the one-CTA path makes. Then it runs
the same radix finish on the candidates.
- A row of 32768 keys or fewer keeps the one-CTA path: up to 8 chunks,
the hand-offs between the CTAs cost what the slices save. When the
threshold bin holds more candidates than LDS, the row uses the one-CTA
overflow path.
  - The score stage does not change.
- Review fix: `index_thresh` and `index_select` read the slices'
histograms and staged lists, plain stores released before their ready
pairs, after an agent-scope acquire, as the `indices` hand-off does. A
parallel row whose threshold bin overflows polls its scores in the
overflow pass, because its CTA has not read them before.
- `glm53/check.py`, `glm53/spec.py` (review fix): check mode compares
each indexer row's selection with the exact top-k of the launch's own
scores (score descending, then position): `layer.topk_exact.max_abs`,
the keys a row misses, bound 0. Every layer check case runs it.
- Tests
- `test/unit/test_glm53_select.py` (new): `select_slices` covers each
chunk of a row once, in position order, with at most 16 slices and no
empty slice (138 cases, 135 distinct bounds); the slice tasks of each
width fit in the upper half of the grid.
- `test/e2e/test_glm53_server.py`:
`test_each_covered_layer_past_the_parallel_top_k` runs the layer check
on rows of about 108k keys at widths 2 and 4 (27 chunks: 14 slices, 13
of two chunks) and of about 36k keys at width 8 (9 one-chunk slices). It
asserts that `past` (65536) is past one chunk a slice. `_check_layers`
gets `past` (an indexer check must reach this many keys) and requires
`layer.topk_exact` in the log.
- No test for the acquires: a missing acquire cannot be made to fail
deterministically on gfx950.

**No change.** Rows of 32768 keys or fewer, the decode steps of
GLM-5.3-Flash, DeepSeek-V4.1 and MiniMax-M3, and the normal path.

## Accuracy Tests

**Setup.** 8x MI355X (gfx950) at TP8; zai-org/GLM-5.3 @ `ba32f8c1`
(FP8); the image and the commits in "Speed Tests and Profiling"; the
server flags of `test/e2e/glm53_server.py`; a new server for each test.

**Test 1: exact top-k.** A debug build (not in this PR) dumps the scores
and the selected positions of each row of a launch. A script computes
the exact top-2048 of the dumped scores (score descending, then position
ascending) and compares the two sets. The prompts are disjoint windows
of real text (SGLang's source). 312 rows (8 ranks; 8k to 256k keys; 1 to
8 rows): 0 rows differ, and all the scores of each row come from the
same launch.

**Test 2: E2E tests at TP8 (review fix).** We ran
`test/e2e/test_glm53_server.py` as written at `235bb05` on main
`876ed7f`. All 7 tests passed in one run (1234 s):

- `TestMonokernelDecode`: average KL 0.015 to 0.025 for each batch
(bound 0.06); GSM8K (sgl-eval, thinking on) 0.995.
- `TestMonokernelMtp`: GSM8K 0.995, accept length 3.31.
- `TestMonokernelCheck`: no violations. `layer.topk_exact` and
`layer.topk_own` are 0 in all 918 indexer checks.
- Decode steps: 1425 checks, stream cosine 0.9925 or more,
`layer.topk.rel` 0.015 or less.
- Verify steps: 1200 checks, stream cosine 0.9977 or more,
`layer.topk.rel` 0.017 or less.
- Past the parallel top-k (rows of about 108k keys at widths 2 and 4,
about 36k at width 8): 1200 checks, stream cosine 0.9954 or more,
`layer.topk.rel` 0.051 or less.
- `layer.topk.rel` is the share of SGLang's selection that a row misses.
It has no bound: SGLang scores with FP8 queries and the kernel with bf16
queries, so keys near the threshold can swap. It is highest in the case
whose prompt repeats one paragraph thousands of times; we did not
measure why. `layer.topk_exact` checks the selection itself against the
launch's own scores.
- `TestMonokernelFault`: passed. Since #55, the test sets
`HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` and turns off SGLang's crash dumps
itself.

An earlier run of the same file at `8068147`, before the review fix,
also passed all 7 tests.

**Test 3: the exact check fails on a wrong selection.** A deliberate
error (not in this PR): the slices stage their candidates with the low
20 bits of each key flipped, so the radix finish takes the lowest keys
of the threshold bin instead of the highest. Every position stays valid.
On the parallel check case, `layer.topk_exact` reports 204 violations
(up to 1422 keys of a row's exact top-k missed), and the test fails. No
other bound catches it: the lowest stream cosine is 0.991 (bound 0.95),
and `layer.topk_own` stays 0. A first version of this control, which
dropped the last slice's winners, made the launch publish stale
positions and fault on memory instead.

The prompts of the GSM8K and KL tests have fewer than 32768 keys, so
these tests run the one-CTA path: they show that this path does not
change. Test 1 and the new layer check cover the new path.

## Speed Tests and Profiling

**Setup.**

- Hardware: one node with 8x MI355X (gfx950), TP8.
- Image: `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019`
(`sha256:8cb29fc407c2`).
- Software: SGLang `e45cbce599`, AITER `2bb549102b`, FlyDSL 0.3.4.1,
torch 2.9.1+rocm7.2.0.
- Model: zai-org/GLM-5.3 @ `ba32f8c1`.
- Server flags: the InferenceX GLM recipe for MI355X
(`--mem-fraction-static 0.80 --chunked-prefill-size 131072
--dsa-prefill-backend tilelang --dsa-decode-backend tilelang`), with
`--max-running-requests` and `--cuda-graph-max-bs-decode` at 2x the
concurrency. MTP: EAGLE, 3 draft steps, top-k 1, 4 draft tokens. The
arms differ only in `--monokernel-backend amd` and the checkout.
- A new server for each point, on an idle node (checked with `amd-smi`
before and after).

**Review fix (`235bb05`).** Fixed length without MTP: 131072-token
prompts, 512 new tokens, 4 prompts for each concurrent request, before
(`6bfceb6`) and after the fix, a new server for each, two rounds with
the second in reverse order. Median ITL (ms):

| Concurrency | Before (round 1 / 2) | After (round 1 / 2) | Change |
|---|---|---|---|
| 1 | 6.47 / 6.47 | 6.49 / 6.49 | +0.3% |
| 4 | 9.27 / 9.28 | 9.28 / 9.28 | 0% |

At concurrency 4 the TPOT median (30 to 40 ms) includes the other
requests' 128k-token prefills, so ITL is the decode measure. In a
compile-only build of the 22 GLM-5.3 keys, the 11 builds without the
indexer are unchanged, and the indexer builds have 0.3% to 1.4% more
instructions. At TP8 no build spills more VGPRs; at TP4, width 4 spills
10 instead of 9.

**Kernel timeline: #15 to this PR.** The debug build of the cause
section, on #15 (`91f0cc7`) and on this PR (`8068147`): a new server for
each, random prompts, 1 to 8 rows decoding together, the last launch of
each of 6 passes on each of the 8 ranks (medians). The estimated step is
18 indexer launches and 57 launches without the indexer.

| Keys per row | Rows | Selection (us) | Score stage (us) | Indexer
launch (us) | Estimated step (ms) |
|---|---|---|---|---|---|
| 8k | 1 | 31 -> 31 | 5 -> 5 | 95 -> 97 | 4.75 -> 4.84 (0.98x) |
| 8k | 2 | 26 -> 27 | 11 -> 12 | 105 -> 108 | 4.93 -> 4.98 (0.99x) |
| 8k | 4 | 27 -> 27 | 18 -> 18 | 135 -> 136 | 6.57 -> 6.61 (0.99x) |
| 8k | 8 | 27 -> 28 | 29 -> 29 | 205 -> 206 | 10.74 -> 10.65 (1.01x) |
| 32k | 1 | 38 -> 37 | 11 -> 11 | 118 -> 119 | 5.17 -> 5.20 (0.99x) |
| 32k | 4 | 42 -> 44 | 44 -> 45 | 177 -> 182 | 7.23 -> 7.45 (0.97x) |
| 32k | 8 | 44 -> 49 | 86 -> 89 | 277 -> 285 | 11.91 -> 12.10 (0.98x) |
| 64k | 1 | 63 -> 43 | 15 -> 15 | 150 -> 129 | 5.77 -> 5.33 (**1.08x**)
|
| 64k | 4 | 68 -> 50 | 63 -> 63 | 223 -> 206 | 8.22 -> 7.87 (1.05x) |
| 64k | 8 | 69 -> 58 | 124 -> 126 | 340 -> 334 | 12.88 -> 12.95 (0.99x)
|
| 128k | 1 | 109 -> 48 | 25 -> 25 | 204 -> 144 | 6.69 -> 5.60
(**1.20x**) |
| 128k | 2 | 117 -> 52 | 48 -> 50 | 238 -> 174 | 7.32 -> 6.19
(**1.18x**) |
| 128k | 4 | 119 -> 56 | 101 -> 102 | 314 -> 251 | 9.85 -> 8.71
(**1.13x**) |
| 128k | 8 | 123 -> 64 | 200 -> 203 | 476 -> 417 | 15.58 -> 14.56
(**1.07x**) |
| 256k | 1 | 207 -> 60 | 44 -> 44 | 320 -> 175 | 8.84 -> 6.20
(**1.42x**) |
| 256k | 2 | 216 -> 65 | 86 -> 88 | 375 -> 224 | 9.88 -> 7.05
(**1.40x**) |
| 256k | 4 | 243 -> 71 | 180 -> 180 | 514 -> 343 | 13.45 -> 10.37
(**1.30x**) |
| 256k | 5 | 224 -> 82 | 224 -> 227 | 593 -> 453 | 17.42 -> 14.94
(**1.17x**) |

The launches without the indexer run the same code in both builds;
between the two runs they differ by -3.2% to +3.1%. Thus the differences
at 8k and 32k keys are in the noise of this measurement. At 32k keys (9
chunks, just past the one-CTA path), the selection takes 2 to 5 us
longer at 4 and 8 rows. The score stage does not change.

**AgentX with MTP.** InferenceX agentic coding (`agentx-harness` 1.0.11,
corpus `23f152f6`), 3600 s a point. All arms use the same simulated
accept length (2.99), so the accept length does not differ between the
arms.

| Concurrency | Interactivity p50 (tok/s/user): normal path with MTP /
#15 / this PR | p90 | Total throughput (tok/s/GPU) | Requests done |
Input p50 (tokens) |
|---|---|---|---|---|---|
| 1 | 204.5 / 193.7 (0.95x) / **236.1 (1.15x)** | 178.5 / 168.0 (0.94x)
/ **203.8 (1.14x)** | 1921 / 1875 / **2135 (1.11x)** | 234 / 231 / 246 |
225k / 221k / 285k |
| 2 | 199.7 / 190.4 (0.95x) / **233.0 (1.17x)** | 149.6 / 135.4 (0.91x)
/ **154.8 (1.03x)** | 2044 / 2018 / **2348 (1.15x)** | 377 / 374 / 409 |
97k / 96k / 111k |
| 4 | 180.7 / 155.1 (0.86x) / **200.9 (1.11x)** | 128.6 / 106.5 (0.83x)
/ 123.7 (0.96x) | 2961 / 2880 / **3132 (1.06x)** | 501 / 493 / 536 | 91k
/ 91k / 89k |

In the same time, this PR completes more requests (for example 246
against 234 at concurrency 1), so its sessions go further; at
concurrency 1 and 2, its inputs are longer.

Time of a verify step for each generated token (ms), from the server log
intervals without prefill, all three concurrencies:

| Requests in the step (verify rows) | 0-64k | 64-128k | 128-192k |
192-320k | 320k or more |
|---|---|---|---|---|---|
| 1 (4) | 4.61 / 3.22 / **3.17** | 4.67 / 3.49 / **3.29** | 4.74 / 3.92
/ **3.49** | 4.86 / 4.86 / **3.84** | 4.97 / 5.72 / **4.46** |
| 2 (8) | 5.48 / 4.81 / **4.72** | 5.57 / 5.40 / **5.08** | 5.62 / 5.75
/ **5.32** | 5.86 / 7.59 / 6.30 | 6.00 / 9.09 / 7.03 |

(normal path with MTP / #15 / this PR). With 2 requests past 192k keys,
the monokernels are still slower than the normal path. At 5 to 8 rows,
the score stage takes about half of the indexer launch (timeline above);
a follow-up PR changes it.

**Fixed length without MTP.** 131072-token prompts, 1024 new tokens, 6
prompts for each concurrent request, the radix cache flushed before each
concurrency. TPOT median (ms):

| Concurrency | Normal path | Monokernels at `e952b22` | This PR |
|---|---|---|---|
| 1 | 11.03 | 7.54 | **6.47** |
| 4 | 23.32 | 25.85 | **19.69** |
| 8 | 46.82 | 52.08 | **46.13** |

**AgentX without MTP.** The same selection code on `e952b22` (without
#15), concurrency 4 and 8, 3600 s a point. Interactivity p50 / p90
(tok/s/user):

| Concurrency | Normal path | Monokernels at `e952b22` | With this
selection |
|---|---|---|---|
| 4 | 66.4 / 58.3 | 71.6 / 53.0 | **117.5 / 79.1** |
| 8 | 54.2 / 33.2 | 46.5 / 19.2 | **57.9 / 32.3** |

## Checklist

- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] 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.

## Review and Merge Process

1. The `lint` and `Check PR title` checks must pass.
2. Ask a code owner from `.github/CODEOWNERS` to review.
3. Merge with "Squash and merge": the PR title becomes the commit title
on `main`.
@jhinpan jhinpan mentioned this pull request Oct 25, 2026
Merged
4 of 5 tasks
jhinpan added a commit that referenced this pull request Oct 25, 2026
@jhinpan
[Perf] [MI355X] glm53: Score the paged indexer keys per wave (#42)
9484c60 
## Motivation

**Summary.** After #24, the score stage of the GLM-5.3 indexer read the
keys at about 0.7 TB/s, and at 5 to 8 rows past 128k keys it took about
half of an indexer launch. This PR scores each group of 16 keys on one
wave over all 32 heads, with the reads of the next groups in flight. The
scores do not change: in one launch, they are bitwise equal to the
scores of #24 on the same inputs. On 8x MI355X at TP8, in InferenceX
agentic coding with MTP, interactivity p50 is 1.43x, 1.46x and 1.23x
that of the normal path with MTP at concurrency 1, 2 and 4 (#24: 1.15x,
1.17x and 1.11x).

**On main after #24.** This PR is two commits on main (`240d587`, which
has #24 and #56): `e405d6b` is `e0d7655` rebased with the same changed
lines (and `e0d7655` is `9733ff4` rebased), and `3419a0e` is the fix
from the review. The results below were measured at `9733ff4` unless
marked.

**Cause.** In the paged score stage of #24:

- A CTA scored a tile of 64 keys with its 8 waves: 4 groups of 16 keys,
each for the two halves of the 32 heads. The two waves of a group read
the same keys, so each key was read two times.
- The two halves were added through LDS, with two block barriers a tile.
- The block-table read and the key reads of a tile started only after
the tile before it ended. Thus each tile waited for two memory round
trips: about 2.1 us a tile, 0.7 TB/s at 256k keys.
- Each key also compared its cache slot with the slots of all S rows of
the launch (the new-key lookup of #15). This cost grows with the width
(S = 2, 4 or 8).

## Modifications

- `glm53/gfx950/kernel.py`, paged score stage only:
- A wave scores a group of 16 keys over all 32 heads: two MFMA chains on
one read of the keys. The head sums stay in the lanes, so the stage has
no LDS reduction and no barrier between groups.
- A wave walks every 8th group of its CTA's tiles; the tiles of a CTA do
not change. The key reads of the next 2 groups are in flight, and their
block-table reads go one step further ahead. The first block-table reads
of a row overlap the staging of its query.
- The query fragments and the head weights of the row stay in registers.
The values of a row (bound, request, counts) are wave-uniform.
- The new-key lookup of #15 runs only for each wave's group that holds
any of the last S keys of a row: a launch writes only those keys (a
verify row's draft chain, a decode row's own key). A wave's groups are
`SCORE_STRIDE` (32768) keys apart, so a wave has at most one such group;
the walk asserts `S + 15 < SCORE_STRIDE` (review fix).
- Each score is computed in the same order as before: the same MFMA
chain, the same sums in the same order.
- `test/e2e/test_glm53_server.py` (review fix): the parallel check case
of #24 passes 98304 keys (rows of about 108k keys at widths 2 and 4),
where each wave's walk refills its buffers and looks blocks ahead.
Before, the longest case (about 36k keys) gave each wave at most 2
groups, so the refills never ran.
- No new check. With a wrong refill offset (a deliberate error, not in
this PR: a refilled buffer reads the keys 16 positions on), the parallel
check case fails on the existing stream-cosine bound: 22 violations,
lowest cosine 0.914 (bound 0.95). `layer.topk.rel` rises to 0.66 (0.052
without the error). So once a case reaches the walk's refills, the
existing bounds catch an error there. The bitwise check needs a debug
build (not in this PR).

## Accuracy Tests

**Setup.** 8x MI355X (gfx950) at TP8; zai-org/GLM-5.3 @ `ba32f8c1`
(FP8); the image and the commits of #24; the server flags of
`test/e2e/glm53_server.py`; a new server for each test.

**Test 1: bitwise scores in one launch.** A debug build (not in this PR)
also runs the per-tile score code of #24 in the same launch, on the same
inputs, and writes those scores to a second scratch buffer. A script
compares the two buffers.

| Steps | Rows | Keys | Keys bitwise equal |
|---|---|---|---|
| Decode: 1 to 8 rows, 8k to 256k keys | 256 | 29,961,984 | 29,961,984 |
| MTP verify: 1 and 2 requests, 8k to 256k keys | 192 | 29,631,744 |
29,631,744 |

The exact top-k check of #24 also found 0 wrong rows in both runs.

We did not compare two servers: on this image, two prefills of the same
prompt give different bits. After a prefill of 131072 tokens, the scores
of the first decode step correlate at 0.9988 between two passes (one
prefill, cached) and at 0.9457 between two prefills in the same server,
with the same build. Between a server with #24 and a server with this
PR, they correlate at 0.9460.

**Test 2: E2E tests at TP8 (review fix).** We ran
`test/e2e/test_glm53_server.py` as written at `3419a0e` on main
`240d587`. All 7 tests passed in one run (1214 s):

- `TestMonokernelDecode`: average KL 0.018 to 0.022 for each batch
(bound 0.06); GSM8K (sgl-eval, thinking on) 0.99.
- `TestMonokernelMtp`: GSM8K 0.99, accept length 3.28.
- `TestMonokernelCheck`: no violations. `layer.topk_exact` (#24) and
`layer.topk_own` are 0 in all 954 indexer checks.
- Decode steps: 1425 checks, stream cosine 0.9942 or more,
`layer.topk.rel` 0.014 or less.
- Verify steps: 1350 checks, stream cosine 0.9976 or more,
`layer.topk.rel` 0.017 or less.
- Past the parallel top-k, with rows past 98304 keys: 1200 checks,
stream cosine 0.9968 or more, `layer.topk.rel` 0.052 or less.
- `TestMonokernelFault`: passed. Since #55, the test sets
`HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` and turns off SGLang's crash dumps
itself.

Two earlier runs also passed all 7 tests: at `9733ff4`, before the
review fix, and at `9b48912`, the fix on #24's branch before #56 merged.

## Speed Tests and Profiling

**Setup.** The setup of #24: one node with 8x MI355X, TP8, image
`lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019`, SGLang
`e45cbce599`, AITER `2bb549102b`, the InferenceX GLM recipe flags, and a
new server for each point on an idle node. The arms differ only in
`--monokernel-backend amd` and the checkout.

**Kernel timeline: #24 to this PR.** The debug build of #24, random
prompts, 1 to 8 rows decoding together, the last launch of each of 6
passes on each of the 8 ranks (medians). The estimated step is 18
indexer launches and 57 launches without the indexer.

| Keys per row | Rows | Score stage (us) | Indexer launch (us) |
Estimated step (ms) |
|---|---|---|---|---|
| 8k | 1 | 5 -> 5 (0.9x) | 97 -> 96 | 4.84 -> 4.80 (1.01x) |
| 8k | 2 | 12 -> 9 (1.3x) | 108 -> 106 | 4.98 -> 4.96 (1.00x) |
| 8k | 4 | 18 -> 14 (1.3x) | 136 -> 134 | 6.61 -> 6.59 (1.00x) |
| 8k | 8 | 29 -> 24 (1.2x) | 206 -> 202 | 10.65 -> 10.61 (1.00x) |
| 32k | 1 | 11 -> 5 (2.0x) | 119 -> 114 | 5.20 -> 5.12 (1.01x) |
| 32k | 4 | 45 -> 29 (1.5x) | 182 -> 172 | 7.45 -> 7.26 (1.03x) |
| 32k | 8 | 89 -> 62 (1.4x) | 285 -> 261 | 12.10 -> 11.67 (1.04x) |
| 64k | 1 | 15 -> 6 (2.5x) | 129 -> 133 | 5.33 -> 5.49 (0.97x) |
| 64k | 4 | 63 -> 35 (1.8x) | 206 -> 182 | 7.87 -> 7.47 (**1.05x**) |
| 64k | 8 | 126 -> 74 (1.7x) | 334 -> 281 | 12.95 -> 11.95 (**1.08x**) |
| 128k | 1 | 25 -> 8 (3.1x) | 144 -> 126 | 5.60 -> 5.32 (**1.05x**) |
| 128k | 2 | 50 -> 19 (2.7x) | 174 -> 144 | 6.19 -> 5.59 (**1.11x**) |
| 128k | 4 | 102 -> 41 (2.5x) | 251 -> 191 | 8.71 -> 7.58 (**1.15x**) |
| 128k | 8 | 203 -> 87 (2.3x) | 417 -> 306 | 14.56 -> 12.44 (**1.17x**)
|
| 256k | 1 | 44 -> 12 (3.7x) | 175 -> 148 | 6.20 -> 5.73 (**1.08x**) |
| 256k | 2 | 88 -> 27 (3.2x) | 224 -> 166 | 7.05 -> 6.05 (**1.17x**) |
| 256k | 4 | 180 -> 60 (3.0x) | 343 -> 225 | 10.37 -> 8.33 (**1.25x**) |
| 256k | 5 | 227 -> 78 (2.9x) | 453 -> 305 | 14.94 -> 12.30 (**1.21x**)
|

The launches without the indexer run the same code in both builds;
between the two runs they differ by -1.8% to +3.4%. The differences at
8k keys and at 64k keys with 1 row are in this noise. The score stage's
time is its window, from the first start to the last end of its tasks.
With this PR, a tile's start and end stamps can come from two waves, so
the time of one task is not meaningful, but the window is.

**AgentX with MTP.** InferenceX agentic coding (`agentx-harness` 1.0.11,
corpus `23f152f6`), 3600 s a point, the same simulated accept length
(2.99) in all arms.

| Concurrency | Interactivity p50: normal path with MTP / #15 / #24 /
this PR | p90 | Total throughput (tok/s/GPU) | Requests done | Input p50
|
|---|---|---|---|---|---|
| 1 | 204.5 / 0.95x / 1.15x / **292.1 (1.43x)** | 178.5 / 0.94x / 1.14x
/ **216.1 (1.21x)** | 1921 / 0.98x / 1.11x / **2386 (1.24x)** | 234 /
231 / 246 / 261 | 225k / 221k / 285k / 300k |
| 2 | 199.7 / 0.95x / 1.17x / **292.0 (1.46x)** | 149.6 / 0.91x / 1.03x
/ **184.9 (1.24x)** | 2044 / 0.99x / 1.15x / **2418 (1.18x)** | 377 /
374 / 409 / 427 | 97k / 96k / 111k / 116k |
| 4 | 180.7 / 0.86x / 1.11x / **222.6 (1.23x)** | 128.6 / 0.83x / 0.96x
/ **146.3 (1.14x)** | 2961 / 0.97x / 1.06x / **3657 (1.24x)** | 501 /
493 / 536 / 582 | 91k / 91k / 89k / 112k |

Time of a verify step for each generated token (ms), from the server log
intervals without prefill, concurrency 1 and 2 (normal path with MTP /
#15 / #24 / this PR):

| Requests in the step (verify rows) | 0-64k | 64-128k | 128-192k |
192-320k | 320k or more |
|---|---|---|---|---|---|
| 1 (4) | 4.61 / 3.21 / 3.15 / **3.06** | 4.66 / 3.50 / 3.29 / **3.12**
| 4.74 / 3.93 / 3.49 / **3.21** | 4.86 / 4.92 / 4.00 / **3.33** | 4.95 /
5.65 / 4.39 / **3.48** |
| 2 (8) | 5.51 / 4.92 / 4.79 / **4.61** | 5.57 / 5.39 / 5.02 / **4.70**
| 5.61 / 5.73 / 5.31 / **4.85** | 5.82 / 7.51 / 6.20 / **5.04** | too
few steps |

With 2 requests past 192k keys, #24 was still slower than the normal
path; this PR is 1.15x faster there.

**Fixed length without MTP.** Prompts of 131072 and 262144 tokens, 1024
new tokens, 6 prompts for each concurrent request, the radix cache
flushed before each concurrency. Decode step without prefill (ms, median
of the server log intervals):

| Keys | Batch | Normal path | #24 | This PR |
|---|---|---|---|---|
| 128k | 1 | 11.04 | 6.47 | **6.25** |
| 128k | 4 | 12.94 | 9.28 | **8.16** |
| 128k | 8 | 15.52 | 14.78 | **12.52** |
| 256k | 1 | 11.36 | 7.01 | **6.49** |
| 256k | 4 | 13.55 | 10.93 | **8.75** |
| 256k | 5 | 15.92 | 15.40 | **12.61** |

At 256k keys, the #24 column is the same selection code on `e952b22`
(without #15); #24 itself was measured at 128k.

## Checklist

- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] 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.

## Review and Merge Process

1. The `lint` and `Check PR title` checks must pass.
2. Ask a code owner from `.github/CODEOWNERS` to review.
3. Merge with "Squash and merge": the PR title becomes the commit title
on `main`.
This was referenced Oct 26, 2026
Closed
Closed
Merged
Merged
jhinpan pushed a commit that referenced this pull request Oct 27, 2026
@heslami Hassan Eslami
[Perf] [MI355X] glm53_flash: Cheaper hand-offs: uncached scratch, few…
db11dcb 
…er vector loads ahead of the polls, tagged selected tokens (#62)

## 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 #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
#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 #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 #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

- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] 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.)

---------

Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan pushed a commit that referenced this pull request Oct 27, 2026
@heslami Hassan Eslami
[Perf] [MI355X] glm53_flash: Scalar-polled boxes for broadcast rows, …
f27e778 
…and o_proj before the state write-back at 2 rows (#63)

## 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):

- 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 #62, 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 #62 (36 / 108 / 216 files, 996 / 2,988 / 5,976 tensors at
TP2 / TP4 / TP8). So #23's accuracy results hold.
- **`SGLANG_AMD_MONO_CHECK=1`:** 874 checks, 0 violations at TP2, TP4
and TP8. As in #62, 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 #55).

## Speed Tests and Profiling

Setup as #62 (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 #62. 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:** #62'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

- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] 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.)

---------

Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jhinpan