[heslami]

Motivation

Stacked on sgl-project#68, sgl-project#69 and the coverage PR ([Feat] [MI355X] glm53: Cover layers 0-2 and the logits). Merge those first. This PR's diff is its last four commits, and it uses sgl-project#68's test/kernels/glm53/ tools.

On an indexer launch the attention waits for the index chain: index query, key scoring, then the exact top-2048. That covers 18 of the 75 MoE layer launches of a step, plus layer 0 and the dense layers 1–2, which reuse the indexer stages. Two parts of the chain grow with context:

• Scoring reads each key once per row. [Perf] [MI355X] glm53: Score the paged indexer keys per wave sgl-project/sgl-amd-mono#42's per-wave scoring scores one row at a time, so at MTP's 4 verify rows every key of the request is read 4 times.
• The selection is mostly serial. A row of up to 32,768 keys selects on one CTA. Past that, [Perf] [MI355X] glm53: K-split front, top-8 tie shortcut, router norm in its logits (preview) #24 spreads it over at most 16 slices, and the row's CTA adds their 4,096-bin histograms and, once the slices have gathered, copies their lists.

This PR scores a request's rows together, reading each key once from LDS batches of whole pages. It then selects every row's top-2048 on up to 64 slices with all-gathered 6-bit digit histograms. The selection stays exact: the same keys, in the same order and with the same ties, as the base's radix select. Every launch's outputs are bitwise equal to the base's (Accuracy Tests).

Result on MI355X, base (the coverage PR) → this PR:

• TP8, MTP 3-1-4 at the golden accept length 2.99, cycle at one request:
  • 8K context: 7.16 → 6.92 ms (−3.4%);
  • 128K context: 8.29 → 7.24 ms (−12.7%).
• At two requests, while both decode:
  • 8K: 10.76 → 10.54 ms (−2.1%);
  • 128K: 12.66 → 11.48 ms (−9.3%).
• TP8, MTP off, median TPOT at 8K: 4.45 → 4.42 ms (−0.5%) at one request, 4.93 → 4.80 (−2.6%) at two, 6.60 → 6.50 (−1.5%) at four.
• Kernel-only, time_layers.py, E4M3 MLA cache, TP8, µs a launch at 4 rows:
  • indexer launch: 96.42 → 86.33 at 8K (−10.5%) and 140.00 → 100.00 at 128K (−28.6%);
  • the 75-launch model order: 71.12 → 68.78 (−3.3%) and 82.88 → 72.20 (−12.9%).

The gain grows with context because both rewritten stages scale with the keys a row reads.

Modifications

The first commit adds a check. The next three are the changes, each measured on top of the one before (Speed Tests, "Each change on its own").

0. An exact top-k check (test/kernels/glm53/check_topk.py, check_layer.py, test_glm53_kernels.py)

• check_topk.py runs indexer launches under torchrun at TP4 or TP8, at chosen positions (--pos). --kind dense and --kind embed run layer 0 and the dense layers.
• It uses four index caches:
  • check_layer.py's random keys;
  • four keys repeated over every position, so the 2048th score falls in a tie (ties);
  • all-zero keys (zeros);
  • each row's last 4,096 keys scaled 8x, so the selection crowds into them (recent).
• Row sets:
  • the served decode points (serving, as time_layers.py);
  • one request a row, at --pos and below, some under 2,048 keys (mixed);
  • one request's rows, a draft chain as long as the width (chain).
• Before each step, the cached scales of the step's own keys are poisoned (check.poison_new_index_scales). A launch that scores a key this launch writes from the cache therefore fails.
• After each launch it checks every live row:
  • Scores (check.launch_scores) against an f64 reference. The reference uses the launch's index query (its index_q mailbox, rotated and transformed in f32 as the kernel does, then bf16), its head weights, and the row's keys as the launch leaves the index cache. The error is taken over the row's largest reference score; the bound is 1e-4.
  • The list (the top-k at the row's indptr), entry by entry, against the radix select's list of the launch's own scores. That list is the keys above the 2048th key's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the 2048th, then its ties in position order. missed also counts the keys of the exact top-2048 (score descending, then position) that the list lacks, as check.py's layer.topk_exact does.
  • Negative controls on copies of each row's result: a selected slot replaced by the best unselected key, two entries swapped, and one score's bits flipped must each fail.
• --dump DIR writes SHA-256 digests of every row's scores and list in check_layer.py's format, so check_layer.py --compare can match two trees.
• test_glm53_kernels.py runs it at TP4 (widths 2, 4) and TP8 (2, 4, 8):
  • the indexer launch at about 8K and 300K keys (at 300K a scoring CTA's run takes three batches and a slice two LDS chunks);
  • layer 0 and a dense layer at 8K and 128K.
  • It passes on the base too: it checks what this PR keeps, the base's list.
• Test fix, in the same commit: check_layer.py's steps gave each row its request's req_ids entry, while the page table has a row per step row. At 8 verify rows, the second request's rows read the first request's page-table row. Each row now gets its own entry, as the spec passes them (req_ids = row index, page_table = the DSA backend's row per step row).

1. Each row's index key on its own task, and the index-query tiles first (kernel.py, layout.py)

• Index key: the cache task used to normalize, rope and store a row's index key after the row's KV row. The key is now an index_key task of its own, on the row's q_norm CTA.
• q_norm CTAs: they run their own index-query tile right after q_norm, ahead of the index key and the KV row.
  • This replaces [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) sgl-project/sgl-amd-mono#68's change 2, where those tiles ran second, on the CTAs of the first extra tiles.
  • That change's shuffle and its TP4 guard go.
• q_b CTAs: a q_b CTA issues its index-query tile's weights with q_b's, before the wait, and runs that tile before its q_b tile.
• Values: the same values, computed in the same order.
• Unpaged fused indexer: it goes. Every indexer build reads SGLang's paged index cache (runner.py passes index_paged=key.with_indexer).
• On its own this change costs 0.6 to 3.3 µs (Speed Tests). It stays because changes 2 and 3 were built and measured on it: the scoring patches this launch's keys from the index_key task's pairs.

2. A request's rows scored together, each key read once (kernel.py, layout.py)

• Before ([Perf] [MI355X] glm53: Score the paged indexer keys per wave sgl-project/sgl-amd-mono#42): a CTA scored tiles of 64 keys of one row, a wave every 8th group of 16 keys. A request's 4 verify rows read each key 4 times, and every scoring CTA polled, roped and Hadamard-transformed its row's whole query.
• Groups of rows: up to 4 consecutive rows of one request score its keys in one pass.
  • A row belongs to the previous row's request when its page table maps that row's position to the slot that row writes; a page being written belongs to its request alone.
  • Each row's 2 waves (4 rows) or 4 waves (2 rows) hold its query and head weights in registers.
• Runs of pages: the request's pages are split into runs of at least 2 pages, one run a CTA, up to 256 runs. They are counted from a CTA that the group's first row offsets.
  • A CTA copies its run into LDS 8 pages at a time with 16-byte loads. The next batch's loads are in flight while its waves score the current one.
  • A run's block ids load with the query's poll.
• Keys this launch writes:
  • Each index_key task also publishes its row's key as cached, FP8 value words and scale, in index_k8 pairs. A batch holding any of those keys patches them in from the pairs, so the cache store need not be visible yet.
  • The patch covers every row of the request in this launch, not only the group's. A draft chain of more than 4 rows (MTP with 4 or more steps) spans two groups, and the second group reads the first group's new keys.
  • The base handles this case by comparing every key's slot with every row's slot ([Perf] [MI355X] glm53: Shorter attention chain, one-pass 8-row front, cheaper routing hand-off (2/2) (preview) #15, [Perf] [MI355X] glm53: Score the paged indexer keys per wave sgl-project/sgl-amd-mono#42). Here, check_topk.py's chain rowset at 8 rows checks it. With each group patching only its own rows, that check fails (Accuracy Tests).
• The query once a row: index_qh tasks rope and Hadamard-transform each row's index query once and publish it in bf16; the scoring CTAs poll it into LDS.
  • A task covers a row's 8 heads, a wave a head.
  • At TP4, q_b runs on every CTA and produces every index-query tile, so these tasks run after q_b.
• Scores: every score is computed as before, with the same MFMA chains over the same bf16 query and keys and the same head sums in the same order. The scores are bitwise equal to the base's (Accuracy Tests).
• Removed: the index_ready flag and the index_k_new pairs.

3. Every row's exact top-2048 on up to 64 slices (kernel.py, layout.py)

• Before ([Perf] [MI355X] glm53: K-split front, top-8 tie shortcut, router norm in its logits (preview) #24): a row of up to 32,768 keys selected on one CTA, reading every score at least twice. Past that, up to 16 slices built 4,096-bin histograms of their chunks; one CTA added them and published the threshold bin; the slices gathered their winners and candidates; and the row's CTA copied them and finished the radix select.
• Now every row selects on up to 64 slices, a slice per 1,024 keys. layout.select_slices gives the split, and the kernel and the CPU test both call it. Each slice is a run of the row's keys in position order, held as radix keys in LDS.
  • Digit histograms: a radix select by 6-bit digits from the top. For each digit, every slice publishes the 64-bin histogram of its keys under the prefix found so far (32 tagged pairs of 16-bit counts) and reads every other slice's. It then finds the digit holding the remaining k-th key and counts the keys above it in the slices before its own.
  • Placing and resolving: the prefix stops at 12 or 24 bits once its keys fit the candidate room (as the base's does), or runs to all 32 bits. Each slice then places its keys above the prefix, whose places are known, and stages its keys with the prefix, with their cache slots, for the row's first slice. That slice resolves them in two 10-bit passes, or one 8-bit pass at 24 bits. At 32 bits the ties take their places directly.
  • Order: the list is the base's one-CTA radix select's, entry by entry. That is the keys above the threshold's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the threshold, then its ties in position order.
• Placement: the slices start past uk's CTAs.
  • When every row has at most 256 / S slices, slice c of row s is task c·S + s, so short rows' slices share one range of CTAs.
  • Otherwise each row takes a run of 64 tasks.
• Scratch: the histogram, threshold and gather stages' buffers give way to sel_hist and sel_cand. The indexer runtimes of widths 2, 4 and 8 together go from 163.4 to 135.2 MiB a rank at TP8, and from 76.5 to 64.4 MiB at TP4 (widths 2, 4).
• Tests, in the same commit:
  • test/unit/test_glm53_select.py covers select_slices: the slices tile a row's keys once, in position order, none empty, at most 64, a slice's keys within its LDS room and its 16-bit counts. It also covers the stage order.
  • test_glm53_server.py's parallel-top-k case asserts that its rows reach all 64 slices.

Accuracy Tests

• Exact top-2048 (check_topk.py, E4M3 MLA cache, 2 indexer launches × 2 steps).
  • Coverage: positions 8,447, 131,071, 262,143 and 524,287, the four index caches, and the serving, mixed and chain row sets.
  • TP8, widths 2, 4, 8: 288 checks on this head and 288 on the base, every row exact (missed 0, order exact). The largest score error is 1.8e-6, and every negative control is caught.
  • Head vs base: the digests of every row's scores and list are equal (41,984 digests, 0 differ). The scores are bitwise the base's, and so is the list, order and ties included.
  • TP4 (widths 2, 4) on this head: every row exact.
  • Negative control for the new-key patch: with each scoring group patching only its own rows, the 8-row chain set fails at 8K and 128K on every index cache. The launch scores the earlier rows' new keys from their poisoned cached scales, giving a score error of about 1e30.
  • The kernel tests (below) also run it on layer 0 and a dense layer at 8K and 128K, at TP4 and TP8.
• Bitwise (check_layer.py --dump, --compare).
  • Coverage: kinds shared, index, model (the first 8 launches), dense and full (layer 0, the dense layers 1–2 and the tail); bf16 and E4M3 MLA cache; every width the TP builds; one request a row, the same with a pad row, and MTP's verify rows; 3 steps.
  • TP8: 63,648 digests, all equal. TP4: 16,416 digests, all equal except one run of one configuration (bf16 dense, 2 rows with a pad row: 127 digests, the first dense launch's output with its top-k equal). That configuration, rerun 4 times, was equal base vs head and run to run. The negative control (one input value negated) changes 7,056 of 10,944 digests.
  • Every output and cache row is therefore bitwise equal to the base's, and so is the top-k, at every TP, width and cache dtype the spec builds.
  • These dumps used check_layer.py with its runtimes built before the layers' weights, as [Perf] [MI355X] glm53: Shorter attention chain, one-pass 8-row front, cheaper routing hand-off (2/2) sgl-project/sgl-amd-mono#69's tool did. With the base's own check_layer.py, which since the coverage PR draws every layer's weights first, two TP8 width-2 configurations (E4M3 index with verify rows, bf16 model with a pad row) vary from run to run on the base itself, with the same kernel code objects. This is under investigation separately and is not from this PR.
• What can differ from SGLang: nothing new. The serving check's layer.topk_own compares with SGLang's selection from FP8-q scores, while the kernel scores the bf16 query. At a threshold tie, the two can pick a different key, as on the base.
• SGLANG_AMD_MONO_CHECK=1:
  • The e2e file's four check cases pass on this head, each asserting no violation:
    • decode rows 1 / 3 / 8;
    • MTP verify rows 4 / 8;
    • verify rows on the E4M3 MLA cache;
    • rows past the parallel top-k (about 108K keys at widths 2 and 4, on 64 slices; about 36K at width 8, on 36).
  • Past 128K keys (two requests of 140K tokens, E4M3 KV, eager), on this head, 0 violations:
    • decode at 1 and 2 rows: 750 layer checks, 54 of them indexer checks past 131,072 keys;
    • MTP verify at 4 and 8 rows: 525 checks, 36 past 131,072 keys.
    • layer.topk_exact and layer.topk_own are 0 in every indexer check; layer.topk.rel is at most 0.048 and layer.index.cosine at least 0.99997.
• E2E (test/e2e/test_glm53_server.py, TP8, on this head):
  • decode against prefill (KL), at 2 / 3 / 7 / 8 rows: passed;
  • the four check cases: passed (above);
  • GSM8K past the largest width, without and with MTP, and the MTP accept-length bound: passed. sgl-eval 0.1.3 was put on PYTHONPATH with a pure-Python stand-in for its compiled editdistance, as for [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) sgl-project/sgl-amd-mono#68;
  • the injected fault fails on main the same way, as in [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) sgl-project/sgl-amd-mono#68 and [Perf] [MI355X] glm53: Shorter attention chain, one-pass 8-row front, cheaper routing hand-off (2/2) sgl-project/sgl-amd-mono#69: no fault at address 0x0 reaches the log (the test's assertTrue(faults) fails), and the server log shows SGLang's watchdog stopping it first.
• Long context, decode against prefill at 128K (an out-of-tree client: 4 prompts of 131,072 ShareGPT tokens with a GSM8K question, 64 greedy tokens, the decode's top-20 log-probabilities against a prefill of prompt and output; E4M3 KV, MTP off):
  • base: KL mean 0.081, top-1 agreement 0.941;
  • this PR: KL mean 0.063, top-1 agreement 0.956;
  • both arms have the same largest KL (1.87) over 252 positions.

Speed Tests and Profiling

Setup: as sgl-project#68: one node with 8 × MI355X; image lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025 (squashfs sha256:fb07e0ef78fe), AITER 285579b6, FlyDSL 0.4.0; zai-org/GLM-5.3 at ba32f8c1 (FP8); the e2e test's server flags with --kv-cache-dtype fp8_e4m3.

• The e2e and check runs overlay SGLang's python/sglang at 5044779f.
• The serving runs use the image's SGLang with --enable-amd-monokernel and the plugin's spec registration from before [Feat] [MI355X] dsv4: Add whole-layer DeepSeek-V4-Pro mono kernels (TP8) sgl-project/sgl-amd-mono#48.
• Kernel-only numbers don't depend on SGLang.
• All numbers are from this head against its base:
  • serving, e2e, the check runs, the long-context runs and the bitwise dumps: this head (b6fdf8e) against the base (3f5f995);
  • kernel-only, compile survey, exact top-k and kernel tests: the same four commits on the coverage PR's previous head (809f6ec), whose kernel code is the same; the coverage PR's last change touched only check.py and spec.py.

Serving: as sgl-project#68.

• One server per arm and setting: MTP 3-1-4 at 8K, MTP 3-1-4 at 128K, MTP off at 8K. All of a setting's points ran on the same server.
• Prompts of 8,192 or 131,072 ShareGPT tokens, no shared prefix, radix cache flushed before each wave. 512 greedy output tokens with EOS ignored. 4 waves (3 at 128K) of C concurrent requests after a warm-up wave.
• MTP runs at the golden accept length (SGLANG_SIMULATE_ACC_LEN=2.99; measured 2.99 on every server). A cycle is a request's decode time over its verify steps.
• One round. The client is out of tree: sglang.benchmark.serving does not import on this image, so a small streaming client against /generate measures TTFT, TPOT and the verify count from meta_info.

TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms (median over requests):

Context	Requests	base	This PR	Change
8K	1	7.164	6.920	−3.4%
8K	2, both decoding	10.760	10.539	−2.1%
8K	2, beside the other's prefill	12.978	12.765	−1.6%
128K	1	8.287	7.238	−12.7%
128K	2, both decoding	12.655	11.483	−9.3%
128K	2, beside the other's prefill	58.54	57.36	−2.0%

The requests of a two-request wave fall into two groups. Each prompt's prefill is its own chunk, so one request decodes for part of its output beside the other's prefill. Both groups are given here; the median over all of a wave's requests depends on how many fall in each.

TP8, FP8 E4M3 KV, MTP off, 8K, median TPOT in ms:

Requests	base	This PR	Change
1	4.446	4.424	−0.5%
2	4.926	4.799	−2.6%
4	6.597	6.497	−1.5%

Kernel-only (time_layers.py)

torchrun --nproc-per-node <TP> test/kernels/glm53/time_layers.py --kind <kind> --rows <rows> --kv fp8 --pos <8447 | 131071>, µs a launch, base → this PR. Each number is the mean of two passes in alternating order in one allocation; each pass is the median of 7 trials of 50 replays. The passes agree within 0.8 µs (1.3 µs on TP4's 4-row indexer launch).

TP8, E4M3 MLA cache:

Launch	Context	2 rows	4 rows	8 rows
indexer	8K	71.54 → 69.99 (−2.2%)	96.42 → 86.33 (−10.5%)	145.52 → 132.98 (−8.6%)
indexer	128K	96.66 → 78.65 (−18.6%)	140.00 → 100.00 (−28.6%)	223.84 → 173.65 (−22.4%)
model order (75)	8K	54.08 → 53.68 (−0.7%)	71.12 → 68.78 (−3.3%)	115.86 → 112.66 (−2.8%)
model order (75)	128K	60.46 → 55.52 (−8.2%)	82.88 → 72.20 (−12.9%)	135.27 → 122.45 (−9.5%)

TP4, E4M3 MLA cache, indexer launch:

Context	2 rows	4 rows
8K	87.88 → 88.77 (+1.0%)	134.13 → 125.26 (−6.6%)
128K	113.92 → 95.55 (−16.1%)	178.05 → 138.61 (−22.2%)

Each change on its own

Each commit against the one before it, TP8, E4M3 MLA cache, indexer launch, µs at 2 / 4 / 8 rows. Two passes in alternating order, one allocation.

Change	8K	128K
1. index key on its own task, index-query tiles first	+0.65 / +0.93 / +2.00	+1.68 / +1.13 / +3.24
2. a request's rows scored together	+0.51 / −6.04 / −11.64	+2.98 / −15.13 / −35.29
3. top-2048 on up to 64 slices	−2.71 / −4.98 / −2.90	−22.67 / −26.00 / −18.14
total	−1.55 / −10.09 / −12.54	−18.01 / −40.00 / −50.19

• Change 1 on its own is slower by up to 3.2 µs. It is kept because changes 2 and 3 were built and measured on it.
• Change 2 costs the 2-row launch 0.5 µs at 8K and 3.0 µs at 128K: a 2-row launch has one live row, which gains nothing from reading a key once for several rows.
• Change 3 takes most of the 128K gain.

Registers

The coverage PR's compile survey: every key the spec builds (stop=True) and the same keys as the tests build them (stop=False). That is TP4 (2, 4 rows) and TP8 (2, 4, 8); indexer, shared, dense, layer-0 and tail builds; bf16 and E4M3 MLA cache; 90 builds.

• Every symbol is glm5_monokernel_0, as on the base.
• The shared and tail code objects are byte-identical to the base's.
• VGPR spills / private segment bytes, base → this PR, for the builds that have either:

Build	stop=True (spec)	stop=False (tests)
TP8 indexer, 8 rows, bf16	4 / 20 → 24 / 100	5 / 24 → 20 / 84
TP8 indexer, 8 rows, E4M3	4 / 20 → 20 / 84	5 / 24 → 24 / 100
TP8 dense, 8 rows, bf16	1 / 8 → 20 / 84	0 / 0 → 20 / 84
TP8 dense, 8 rows, E4M3	0 / 0 → 20 / 84	0 / 0 → 20 / 84
TP4 indexer, 4 rows, bf16	11 / 48 → 15 / 64	40 / 164 → 48 / 196
TP4 indexer, 4 rows, E4M3	16 / 68 → 13 / 56	34 / 140 → 48 / 196

• The 8-row indexer and dense builds spill about 20 VGPRs (84–100 bytes), against 0–5 on the base. They come with change 2. Even so, the 8-row indexer launch is 8.6% faster at 8K and 22% at 128K.
• Every other build has no VGPR spill and no private segment on either head.
• The layer-0 builds have none at any width.

Tests

• Unit tests on a CPU-only host: the base's results (this head 288 passed, the base 276; the base's 2 contract failures and 3 collection errors).
• test_glm53_kernels.py with 8 GPUs (TP4 and TP8): 8 passed (17 min 32 s), including the new exact top-k test.

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

[heslami]

Sent upstream as sgl-project#74.