Skip to content

[Perf] [MI355X] glm53: Score a request's rows together and select the exact top-2048 on up to 64 slices - #74

Merged
jhinpan merged 4 commits into
sgl-project:mainfrom
MarloweAI:perf/glm53-index-parallel
Oct 29, 2026
Merged

jhinpan merged 4 commits into
sgl-project:mainfrom
MarloweAI:perf/glm53-index-parallel

Conversation

@heslami

@heslami heslami commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Stacked on #68, #69 and #73 ([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 #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:

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 (#73) → 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.56 → 4.44 (−2.8%) at two, 5.87 → 5.77 (−1.7%) 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.
  • 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 #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 ([Feat] [MI355X] glm53: Serve MTP verify steps on the monokernels #15, [Perf] [MI355X] glm53: Score the paged indexer keys per wave #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: Select a long row's indexer top-k on up to 16 CTAs #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) #69's tool did. With the base's own check_layer.py, which since [Feat] [MI355X] glm53: Cover layers 0-2 and the logits #73 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):
  • 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 #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.

Serving: as #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 10.760 10.539 −2.1%
128K 1 8.287 7.238 −12.7%
128K 2 12.655 11.483 −9.3%

The requests of a two-request wave prefill one after the other. The first-prefilled request waits out the other's prefill (380 ms at 8K, 7.9 s at 128K) before both decode, so its cycle is the one above plus that wait over its verify steps. The table gives the other request, which decodes beside the first from its first token to its last; the MTP-off rows below are measured the same way.

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.563 4.437 −2.8%
4 5.872 5.773 −1.7%

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

@jhinpan jhinpan left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I re-ran this head against #73 (3f5f995). The kernel and serving speed claims reproduce, and the exact top-k check passes. The bitwise claim holds for the common outcome; the run-to-run differences come from the uncached-buffer reuse that #76 fixes.

Kernel-only, time_layers.py, E4M3 KV, TP8, us a launch (mean of two passes)

Launch Context 2 rows 4 rows 8 rows
indexer 8K 72.36 -> 70.96 (-1.9%) 97.24 -> 87.25 (-10.3%) 145.90 -> 134.31 (-7.9%)
indexer 128K 97.62 -> 79.38 (-18.7%) 140.99 -> 101.63 (-27.9%) 225.18 -> 174.46 (-22.5%)
model order (75) 8K 55.07 -> 54.57 (-0.9%) 72.23 -> 70.78 (-2.0%) 116.61 -> 113.64 (-2.5%)
model order (75) 128K 61.42 -> 56.62 (-7.8%) 83.88 -> 73.63 (-12.2%) 136.30 -> 123.71 (-9.2%)
  • TP4 matches your TP4 table: indexer at 8K +1.7% (2 rows) and -6.4% (4 rows); at 128K -16.4% and -22.2%.
  • One TP4 pass at 128K / 4 rows had a median of 159.8 us (a max sample of 8.6 ms) against 139.5 us in the other pass. Neither TP8 pass repeated it.

Serving, TP8, E4M3 KV

Setting #73 This PR Change PR body
MTP, cycle, 8K, 1 request 7.151 ms 6.941 ms -2.9% -3.4%
MTP, cycle, 8K, 2 requests 11.670 ms 11.454 ms -1.9% -2.1% / -1.6%
MTP, cycle, 128K, 1 request 8.273 ms 7.254 ms -12.3% -12.7%
MTP off, TPOT, 8K, 1 / 2 / 4 4.529 / 4.902 / 6.331 4.492 / 4.777 / 6.240 -0.8% / -2.6% / -1.5% -0.5% / -2.6% / -1.5%

At 2 requests, our client mixes "both decoding" with "beside the other's prefill", so it lands between your two rows.

Accuracy

  • check_topk.py at TP4 (positions 8,447 and 131,071; the random, recent, ties and zeros caches; the serving, mixed and chain row sets): 96 checks, every row exact, every control caught, largest score error 5.0e-7.
  • Bitwise (check_layer.py --dump, TP4; kinds index, shared, dense and full; bf16 and E4M3; p0, p1 and v4; two runs per tree):
    • bf16 is equal everywhere.
    • On E4M3, every difference sits in a verify-row group where one of the two runs deviated (the uncached-buffer reuse that #76 fixes).
    • On the index kind, the common outcome is bitwise the same on #69 (3 runs), #73 and #74 (4 runs).
    • So "every output is bitwise the base's" holds for the common outcome. A single-run --compare can still report a difference that neither tree causes.
  • Kernel tests at TP8 (test_glm53_kernels.py): 1 failed, 7 passed (24 min). The failure is the same test_graph_replays_equal_the_eager_step GPU fault at TP4, in PyTorch's index_elementwise_kernel (the reuse #76 fixes). test_indexer_selects_the_exact_top_k passed at TP4 and TP8.
  • E2E at TP8: covered by #75's head (9 passed); not run on this head alone.

Nit: kernel.py adds a non-ASCII ≤ in a docstring ("The next bits (≤ 10) ..."). The glm53 package is otherwise ASCII-only.

Setup. MI355X x8 (gfx950), rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261026 (ROCm 10, AITER 285579b62, FlyDSL 0.4.0) with SGLang's python/sglang at 5044779f0f, zai-org/GLM-5.3 @ ba32f8c1, TP8, --kv-cache-dtype fp8_e4m3, the e2e server flags. Serving:

  • one fresh server per tree and setting;
  • in-tree sglang.benchmark.serving: random prompts of exactly 8,192 or 131,072 tokens from ShareGPT, 512 output tokens, a warm-up wave, then 4 waves of C requests (3 at 128K), flush_cache before each point;
  • MTP 3-1-4 at the golden accept length (SGLANG_SIMULATE_ACC_LEN=2.99, real-draft-token); the measured accept length was 2.99 on every server, and cycle = accept length x median TPOT.

Kernel-only numbers come from time_layers.py, two passes in opposite order. GPUs were verified idle over 30 samples.

@heslami

heslami commented Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto main, on the rebased #73; no code changes.

Hassan Eslami added 4 commits October 29, 2026 13:53
… index caches

test/kernels/glm53/check_topk.py runs indexer, dense or layer-0 launches at chosen
positions, on random, tied, zero and recency-heavy index caches, with the cached
scales of each step's own keys poisoned first. For every live row it checks the
launch's scores against an f64 reference and its top-k list, entry by entry,
against the radix select's list of the launch's own scores; negative controls (a
wrong slot, two entries swapped, a flipped score bit) must each fail. --dump writes
digests of every row's scores and list for check_layer.py --compare.
test_glm53_kernels.py runs it at TP4 and TP8.

check_layer.py's steps give each row its own req_ids entry, as the spec does (each
row has its request's page-table row); before, a second request's rows read the
first request's row.
…s first

The cache task normalized, roped and stored each row's index key after its KV row;
it is now an index_key task of its own, on the row's q_norm CTA after the CTA's
index-query tile. The q_norm CTAs run their own index-query tile right after
q_norm, and 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 (the index chain waits on it). Same
values in the same order.

The fused indexer for an unpaged index cache goes: every indexer build reads
SGLang's paged index cache.
Each indexer row re-read every key: a CTA scored tiles of 64 keys of one row. Now
up to four consecutive rows of one request score its keys in one pass. The
request's pages are split into runs, one a CTA; a CTA copies its run into LDS
eight pages at a time, the next batch's loads in flight, and each row's waves score
the batch with the row's query in registers. A key this launch writes is patched
into the batch from its row's index_k8 pairs (its FP8 values and scale as cached),
for every row of the request in the launch.

Each row's index query is rotated and Hadamard-transformed once, by index_qh tasks
(8 heads a task), instead of on every scoring CTA. Every score is computed as
before.
A row's selection ran on one CTA, or past 32768 keys on up to 16 slices whose
histograms one CTA added. Now every row selects on up to 64 slices (a slice per
1,024 keys), each a run of its keys in position order, by a radix select of 6-bit
digits from the top: per digit every slice publishes the histogram of its keys
under the prefix found so far and reads every slice's. Once the prefix has 12 or 24
bits and its keys fit the candidate room, or has all 32, each slice places its keys
above it and stages its keys under it, with their cache slots, for the row's first
slice, which resolves them in two 10-bit passes. The lists are the one-CTA radix
select's, in the same order. Short rows' slices interleave on one range of CTAs.

layout.select_slices gives a row's slices; the kernel and the CPU test call it.
@heslami
heslami force-pushed the perf/glm53-index-parallel branch from e9ea165 to f044534 Compare October 29, 2026 20:53
@jhinpan
jhinpan merged commit a0c3997 into sgl-project:main Oct 29, 2026
2 checks passed
jhinpan pushed a commit that referenced this pull request Oct 30, 2026
## Motivation

**Stacked on #73 and #74** (and #68, #69 under them). Merge those first;
this PR's diff is its last three commits. It uses #73's tail and token
builds and `test/kernels/glm53/` tools.

With MTP (EAGLE 3-1-4, SGLang's GLM-5.3 default), each cycle is one
verify step and three draft steps of the NextN layer. After #73 the
verify step runs entirely as launches. The draft steps still run
SGLang's NextN: the embedding, `enorm` / `hnorm` / `eh_proj`, one full
decoder layer (the indexer on the draft extend), `shared_head.norm`, the
LM head and the logits' all-gather. That is dozens of kernels and three
all-reduces a draft step, on 1 to 8 rows.

This PR runs each draft step as two launches: the NextN layer, starting
from the step's tokens through its `eh_proj` entry, and #73's tail with
the NextN's norm. The draft extend (each request's verify rows) runs the
indexer kind; the draft decode steps run the shared kind on the top-k
the draft extend selected (SGLang's MTP index share). SGLang needs no
change: Monokernel API v1 routes a draft model by the target's verify
key, which a draft never runs, so the draft spec puts a wrapper in front
of the draft's `model.forward`, as #29 does.

Result on MI355X, TP8, FP8 E4M3 KV cache, #74 → this PR:

- MTP 3-1-4 at the golden accept length 2.99, per cycle: 6.93 → **6.44
ms** at one request (−7.1%; 432 → 465 tok/s per user), 10.49 → **9.96
ms** at two (−5.0%).
- Real acceptance (SGLang's default rejection sampling, greedy
requests), paired over two rounds ABBA: accept length 3.29 → 3.30 at one
request (+0.01 ± 0.03), 3.36 → 3.33 at two (−0.03 ± 0.02), within the
servers' A/A spread (Accuracy Tests).

Scope:

- GPU: MI355X (gfx950), TP 4 or 8, as the target.
- Steps: draft extends of 1 or 2 requests (4 or 8 rows; 1 request at
TP4) and draft decode steps of 1 to 4 requests. Other draft steps,
including draft decode steps of 5 to 8 requests, run SGLang's path
(Reviewer notes).
- A draft decode step whose proposal is sampled (rejection sampling with
a batch that is not all greedy) runs SGLang's path; the draft extend's
proposal is SGLang's on the tail's logits.
- SGLang: no change.

## Modifications

Three commits: the NextN builds (kernel and runtime), the draft spec and
its forward, and the kernel checks.

### 1. The NextN entry, and a NextN launch from the step's tokens
(`kernel.py`, `layout.py`, `runner.py`, `weights.py`)

- **`Glm53Build.eh`** (with `tokens`): the layer starts from the NextN
entry, `eh_proj(enorm(e) | hnorm(h))`. It is off by default; with it
off, every existing build compiles to the same code object as the base's
(Speed Tests, Registers).
- e: the token ids' embedding rows, pushed by the rank whose vocabulary
shard holds each, as #73's layer 0 does, into their own symmetric buffer
(`ein`).
- The entry: each rank computes its `HIDDEN / tp` columns of x from its
rows of `eh_proj` (a row view of SGLang's bf16 weight), in 16-column
tiles of 2 tasks each (the e half and the h half, f32 sums; the h half's
partial goes to the e task's CTA through `ehp`), rounds them to bf16 and
pushes them to every rank as tagged pairs (`xin`). The layer then reads
x where layer 0 reads its embedding rows (the q_a / kv_a and index front
tiles, and the attention residual).
- The entry's tasks run on the last CTAs: at TP8 96 tasks on CTAs
160–255, at TP4 192 tasks on CTAs 64–255. The stages that share those
CTAs (qkv_a and q_b tiles) wait for x anyway; only their weight prefetch
starts later.
  - Everything after x is the indexer or shared kind's code, unchanged.
- **Arguments:** the h rows ride in `cur_pos`, which a layer launch does
not read; the entry's tensors (eh_proj's rows, the two norms, the
embedding rows) ride in the op's own copy of the pointer table
(`EH_ARGS`, after the indexer's entries), which the shared kind
otherwise does not have.
- **`Glm53Epochs`:** several runtimes' guard blocks in one tensor, so
one `add_` advances all their step epochs.
- **`head_weights(norm=)`:** #73's tail on another norm (the NextN's
`shared_head.norm`).

### 2. The MTP draft steps on the NextN launches (`draft.py`,
`tables.py`, `spec.py`, `check.py`, `__init__.py`)

- **The draft spec** (`glm53_nextn_hip`), registered for
`GlmMoeDsaForCausalLMNextN`.
- **Routing.** `_build` puts a wrapper in front of the draft's
`model.forward`; SGLang's runtime then installs its own forward over it,
which hands every draft step to the wrapper (they never match the
target's verify key). The wrapper decides on static properties of the
step only (forward mode, batch size, rows, and whether the step reuses a
top-k), so a captured graph replays the decision made at capture:
- a draft extend (`DRAFT_EXTEND_V2`) of whole requests at 4 or 8 rows (4
at TP4) runs the NextN indexer launch;
- a draft decode step (`DECODE`) of 1 to 4 rows that reuses the draft
extend's top-k runs the NextN shared launch at the next width (2 or 4);
  - every other step runs SGLang's forward.
- A draft-shaped step the launches cannot read (no hidden states, hidden
states of another shape or dtype, no reused top-k) is refused during
graph capture and runs SGLang's path eagerly, logged once a key. Each
key the launches take is logged at its first launch.
- The draft's attention backend is checked once per step mode, as the
runtime checks a routed step's.
- **A step.** The step tables (one launch, below), the NextN launch, the
seed, then the tail:
- the draft extend writes the NextN MLA and index-K rows of every row,
then SGLang's seed (`dsa_seed_topk_capture`: each row's top-k as KV
slots, then −1);
- a draft decode row attends exactly the seed's valid entries, compacted
in order (its count comes from the seed, not from its position);
- the tail runs on `shared_head.norm` and the LM head SGLang shares with
the target; the logits processor returns its f32 logits through the same
`_get_logits` binding as #73's target (the draft extend graph's selected
rows when `draft_extend_select_index` is set), and stores the tail's
normed rows as the step's hidden states, which the next draft step
reads.
- The embedding and the LM head are SGLang's, read in place; SGLang
shares them with the target after the draft is built, so a step binds
the ones current at that step (a pointer compare on the host).
- **Sampling.** The logits are the tail's, so SGLang's draft-extend
proposal on them is SGLang's, greedy or sampled. A draft decode step
whose proposal SGLang samples (rejection sampling, which this image
enables by default on HIP for EAGLE topk 1, with a batch that is not all
greedy) runs SGLang's path: its graph runner captures that variant
separately.
- **SGLang's path for the whole draft** (logged at startup): a
speculative token map, no MTP index share
(`index_share_for_mtp_iteration`), a seed width other than 2048, an
indexer that is not SGLang's gfx950 fused one, a TP size the target does
not run, or an `eh_proj` that is not a bf16 `nn.Linear`.
- **Step tables in one launch** (`tables.py`, Triton): `step_rows` fills
a width's `Glm53Step` (slots, positions clamped to the context,
`indptr`, request rows, token ids, and with a seed the compacted top-k)
in one launch, for the target's `_step` and the draft's; `write_seed`
writes the seed. The words equal the torch ops they replace (unit tests
under Triton's interpreter, kernel checks on the GPU). The target's step
epochs advance in one launch (`Glm53Epochs`).
- **Weights.** No new copies: eh_proj's rows are a view of SGLang's
weight; the NextN layer's tensors are read and packed as a target
layer's (`layer_weights`, with its W_UK / W_UV, router and indexer
copies logged at startup).
- **`SGLANG_AMD_MONO_CHECK`** (eager): each routed draft step against
SGLang's NextN on the same inputs and caches:
- the entry (`draft.eh`: x against SGLang's `eh_proj` on its `enorm |
hnorm` rows; `draft.eh_vs_fp32`: against an f32 entry);
- the NextN layer on the launch's own x, as a target layer's check but
reported as `nextn.*`;
- a draft decode row's keys against the seed's valid entries
(`draft.topk_extra`), and the seed SGLang gets against the draft
extend's top-k (`draft.seed`);
- the tail against SGLang's `shared_head.norm`, LM head and all-gather
and an f32 tail (#73's `head.*` bounds, logged as `draft tail`);
- the draft tokens: the tail's top-1 against SGLang's whole NextN (its
own entry, the layer routed as SGLang routes, its norm and LM head) on
the same inputs (`draft.token.agree`, reported, no bound), and SGLang's
logit gap where they differ (`draft.token_gap`).
-
`SGLANG_AMD_MONO_DRAFT_FAULT={swap_norms,count_by_position,seed_row0,rope_shift}`
is a test fault for the negative controls (Accuracy Tests).
- **`SGLANG_READS`** of `draft.py`: the NextN model's modules, the
forward batch's draft fields (`spec_info.hidden_states`,
`dsa_topk_indices`, `dsa_seed_topk_capture`, `reuse_dsa_topk_indices`,
`sampling_info.is_all_greedy`), the logits metadata's
`draft_extend_select_index`, the spec arguments read, and
`get_dsa_mtp_topk_width`.

### 3. Kernel checks of the NextN launches
(`test/kernels/glm53/check_nextn.py`, `test_glm53_kernels.py`)

- `check_nextn.py`, under torchrun at TP4 or TP8, on random weights in
SGLang's layouts, NextN index at 4 and 8 rows (draft extend rows, 4 a
request; 4 at TP4) and NextN shared at 2 and 4, each with and without a
pad row, bf16 or E4M3 MLA cache:
- **entry:** the NextN launch against the plain launch of its kind fed
the NextN launch's own x (its gathered `xin` rows): output rows, MLA
cache rows, index cache rows and top-k, bitwise; x within 1e-4 of a
torch entry and 7e-3 of an f32 entry; a pad row's x zero. The enorm and
hnorm weights swapped must fail both.
- **tables:** `step_rows` and `write_seed` against the torch ops, word
for word (SGLang's padding slot 0, pad rows, positions past the context,
−1 inside a seed); the seed sized by position must differ.
- **replay:** a draft step (NextN launch and tail) in a graph, 20
replays bitwise equal to the eager step, no error word.
- `test_glm53_kernels.py` runs it at TP4 and TP8 on both MLA caches.
- CPU unit tests (`test/unit/test_glm53_draft.py`): the registry entry;
the steps the wrapper routes, at TP8 and TP4, positional and keyword;
the steps it leaves to SGLang (other modes, more rows, a sampled draft
decode); the capture-time refusal and the eager fallback logged once;
the runtime's forward over the wrapper; the backend check once per mode;
every SGLang call site of `model.forward` the wrapper relies on; the
tail's logits through `_get_logits`; each setting that keeps the draft
standard; the entry, seed and decode-row metrics with their negative
controls (real `check.py` functions); the step tables against the torch
ops under Triton's interpreter, in a child process
(`test/unit/triton_interpreter.py`), as the target's `_step` tests now
run too.

## Accuracy Tests

The draft is not bitwise equal to SGLang's NextN, and is not meant to
be: it runs the target layers' numerics (bf16 attention against the FP8
cache, f32 routing). Everything else is bitwise or exact:

- **NextN launch = plain launch fed its own x, bitwise**
(`check_nextn.py`, entry): outputs, MLA rows, index rows and top-k, at
TP8 (NextN index 4 / 8 rows, NextN shared 2 / 4, each with and without a
pad row; 56 checks per cache) and TP4 (index 4, shared 2 / 4; 38 checks
per cache), bf16 and E4M3 MLA cache: all pass. The entry against a torch
entry: 3e-6 to 3e-5; against f32: 2.3e-3, the torch entry's own
distance; pad rows zero; the swapped-norm control gives 0.14 against
both bounds. The step tables equal the torch ops word for word; 20 graph
replays of a draft step are bitwise equal to the eager step.
- **Every existing build is unchanged:** the compile survey's 90
existing keys (indexer, shared, dense, token and tail; TP4 2 / 4, TP8 2
/ 4 / 8; both caches; `stop` on and off) give the same code objects as
the base. Seven of them compile nondeterministically on both trees (the
base alone gives two different objects across reruns, with equal
register counts); in reruns each head object equals one the base
produced.
- **`SGLANG_AMD_MONO_CHECK=1` with MTP on** (TP8, FP8 E4M3 KV, eager, 1
and 2 requests at 1K and 8K, 48 tokens each; 240 draft steps: draft
extends at 4 and 8 rows, draft decode steps at 1 and 2): **0
violations**, target layers included.

  | Metric | Measured | Bound | Negative control (must fail) |
  | --- | --- | --- | --- |
  | `draft.eh.rel` | max 4.5e-4 | 1.4e-3 | `swap_norms`: 0.57–0.79 |
| `draft.eh_vs_fp32.rel` | max 2.8e-3 (SGLang's own: 2.8e-3) | 8.8e-3 |
`swap_norms`: 0.57–0.79 |
  | `draft.topk_extra.max_abs` | 0 | 0 | `count_by_position`: 1–2 |
  | `draft.seed.max_abs` | 0 | 0 | `seed_row0`: 4–2048 |
| `nextn.streams.cosine` | min 0.963 | 0.89 | `rope_shift`: min 0.62
(110 of 264 steps) |
  | `nextn.kv.cosine` | min 0.99936 | 0.995 | `rope_shift`: 0.36–0.86 |
| `nextn.index.cosine` | min 0.99965 | 0.999 | `rope_shift`: from 0.85 |
| `nextn.topk_exact` / `nextn.topk_own` | 0 / 0 | 0 / 0 | the target's |
| draft tail `head.logits.rel` / `_vs_fp32` | max 4.4e-4 / 2.2e-3 |
#73's 1.7e-3 / 7.5e-3 | #73's |
| `draft.token.agree` (reported) | drafts 1 / 2 / 3: 95.9 / 95.5 / 99.1%
of rows | — | — |

Each control ran on its own server and failed only its own checks
(`rope_shift` also fails `nextn.kv` and `nextn.index`).
- **Why the stream bound is loose:** a target layer's stream is the
residual of every layer before it, so one layer's error moves it little.
The NextN stream is a fresh x plus one layer's attention and MLP. It
carries the full distance of SGLang's NextN attention, which runs FP8 q
and P on this cache where the launches run bf16. The tight checks on the
same path are the bitwise kernel check and the entry against f32.
- The draft tokens differ from SGLang's NextN on 1–5% of rows, at SGLang
logit gaps up to 1.25; these are the two attentions' numerics.
- **Accept length at real acceptance**, paired (TP8, FP8 E4M3 KV,
8,192-token ShareGPT prompts, 512 greedy tokens, SGLang's default
rejection sampling; two rounds ABBA, one server per arm and round, the
same prompts on every server):

| Requests | n | base | this PR | change (± s.e.) | A/A base / this PR |
  | --- | --- | --- | --- | --- | --- |
| 1 | 80 | 3.291 | 3.304 | +0.013 ± 0.032 | −0.032 ± 0.053 / −0.017 ±
0.047 |
| 2 | 160 | 3.363 | 3.328 | −0.035 ± 0.022 | +0.011 ± 0.027 / −0.030 ±
0.042 |

- **E2E** (`test/e2e/test_glm53_server.py`, TP8, this head):
- `TestMonokernelCheck`: 4 passed. Its MTP cases (bf16 and E4M3 KV) now
also check every draft step (`check.draft`) and assert drafts 1, 2 and 3
were checked; no violation.
- `TestMonokernelMtp`: passed. GSM8K 0.985, accept length 3.28, and the
draft extends of 1 and 2 requests ran the NextN launches (the test now
asserts their log lines; these requests sample, so under rejection
sampling their draft decode steps run SGLang's path).
  - `TestMonokernelDecode` (2) and `TestMonokernelHiddenStates`: passed.
- `TestMonokernelFault` fails here on `main` the same way, only at its
last check: at the 10 s deadline every rank stops with `HIP error: an
illegal memory access was encountered` and the server exits before any
token of the failed step reaches the client, but ROCr in this container
prints no `Memory access fault ... (nil)` line, so `assertTrue(faults)`
fails; in the reviewer's setup on #68 and #69 the line is printed and
the case passes.
- **GSM8K with MTP at two requests at a time** (greedy, thinking on, 200
questions, as the e2e test's GSM8K but at two requests, so every draft
step of 1 or 2 requests runs the launches): 0.990, accept length 3.38;
the draft extends of 1 and 2 requests and the draft decode steps of 1 to
4 ran the NextN launches.
- **TP4:** an MTP server (FP8 KV, golden AL 2.99, one request at 8K)
serves with the NextN launches: the draft extend at 4 rows and draft
decode steps of 1 to 4 rows ran them; cycle 8.93 ms, accept length 2.99;
NextN kernel checks as above.
- Unit tests on a CPU-only host: 335 passed (base: 289); the base's 2
failures and 3 collection errors (minimax, deepseek_v4).

## Speed Tests and Profiling

Setup as #68 (image
`lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs
`sha256:fb07e0ef78fe`) with SGLang's `python/sglang` overlaid at
`5044779f`, AITER 285579b6, FlyDSL 0.4.0, `zai-org/GLM-5.3` at
`ba32f8c1`, the e2e test's server flags with `--kv-cache-dtype
fp8_e4m3`). All numbers are from this head (a5bfe9d) against the base
(b6fdf8e, #74's head).

Serving as #68: one fresh server per arm, 8,192-token ShareGPT prompts,
512 greedy output tokens, MTP at the golden accept length 2.99 (measured
2.98–2.99 on every server), **out-of-tree streaming client**. One round.

At C≥2 the server prefills a wave's requests one after another, and a
request that has started decoding waits out the later prefills. The C≥2
numbers below are those of each wave's last-prefilled request, which
decodes beside the others from its first token to its last.

TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms:

| Requests | base | This PR | Change | tok/s per user at AL 2.99 |
| --- | --- | --- | --- | --- |
| 1 (draft extend 4 rows, draft decode 1 row) | 6.927 | 6.435 | −7.1% |
432 → 465 |
| 2 (8 rows, 2 rows) | 10.488 | 9.962 | −5.0% | 285 → 300 |

MTP off is unchanged: the target's launches are the base's code objects.
Its step tables are one launch instead of the torch ops' kernels, and
its epochs one launch.

At real acceptance (the ABBA servers above, median TPOT): 2.12 → 1.93 /
1.96 ms at one request, 3.04 / 3.03 → 2.96 / 2.95 ms at two (−2.6%).

### Registers

The compile survey (CPU-only, FlyDSL compile-only): every key the spec
and the tests build, TP4 (2, 4 rows) and TP8 (2, 4, 8), both caches,
`stop=True` (spec) and `stop=False` (tests). The existing keys are
unchanged (Accuracy Tests). The NextN keys, VGPR spills / private
segment bytes, the ones that have either:

| Build | `stop=True` (spec) | `stop=False` (tests) | the plain kind's
(`stop=True` / `False`) |
| --- | --- | --- | --- |
| TP8 NextN index, 8 rows | 12–16 / 52–68 | 12–16 / 52–68 | 20–24 /
84–100 |
| TP4 NextN index, 4 rows | 16 / 68 | 56 / 228 | 13–15 / 56–64, 48 / 196
|
| TP4 NextN shared, 4 rows | 24–26 / 100–108 | 40–48 / 164–196 | 0 / 0 |

Every other NextN key (TP8 index 4 rows, shared 2 / 4 rows; TP4 shared 2
rows) has no VGPR spill and no private segment. Every symbol is
`glm5_monokernel_0`, as on the base. The TP4 NextN shared spill at 4
rows is the entry's 16-column tile next to the 4 rows; it serves TP4's
draft decode steps of 3 and 4 requests.

## Reviewer notes

- **Draft decode steps of 5 to 8 requests run SGLang's path.** In
`check_nextn.py`, the NextN shared launch at 8 rows on the E4M3 cache
faulted (illegal address) as the ninth NextN model of one process, after
eight others were built, run, replayed and closed. Alone in a fresh
process it runs clean, both caches, NextN or plain first. The plain
shared launch at 8 rows never faulted (the kernel tests, this head and
the base). Until that is root-caused, the spec does not build the NextN
shared kind at 8 rows. MTP at 1 to 4 requests is unaffected.
- **The draft-step wrapper**, as in #29. It
relies on SGLang's draft call sites calling the instance attribute
`model.forward`; a unit test pins them. It can go when the monokernel
API routes a draft model's own step modes: then the spec declares them,
and the wrapper and its capture-time refusal go.
- **`_get_logits`** is the one SGLang method the draft spec replaces, as
#73's target does; the binding is active only inside a routed step. With
`draft_extend_select_index` set (the draft extend graph), it returns the
selected rows, as SGLang's pruning would hand them.
- **Eager and graph steps share the per-width step tables.** Every
routed step, eager or replayed, rewrites all of them in the same stream
before its launch (the step-table launch is inside the graph), so an
eager step between replays changes nothing a replay reads.
- **Kernel test harness:** `check_nextn.py` builds several models in one
process. For the runs above, both trees (base and head alike, not part
of this PR) carried a runner-only patch in which `Glm53Runtime.close()`
keeps its uncached buffers. In a process that closes runtimes and builds
new ones, freed uncached buffers corrupt later allocations; this is
root-caused separately. Serving never closes a runtime.
- **Copies:** the draft copies its own rotary table (128 MiB a rank at
the served context) and packs its layer's W_UK / W_UV, router and
indexer weights, as a target layer does (logged at startup).
- `kernel.py` grows by 152 lines (the entry, the `ein` buffer and the
pushed-x switch), a candidate for the stage-builder split.

## 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 Nov 3, 2026
…-split front, indexer launch, MoE expert dedup and weight loads, kernarg reload, LM-head loads (#85)

## Motivation

On main. One commit per optimization, 24 in all, then one cleanup commit
that defines the up/gate tile helpers once and folds `reload_args`'
guard into it (−202 lines; every compiled kernel is byte-identical to
the commit before it).

A GLM-5.3 decode step runs 75 MoE layer launches plus layer 0, the dense
layers 1–2 and the tail; with MTP each cycle also runs the NextN draft
launches. This PR shortens the hand-offs and the weight streams of those
launches:

- **Shared layer and MoE.** down read all of mid as tagged f32 on every
CTA; it now reads mid quantized to FP8 once. At 8 rows the split tiles'
KV rows waited behind q_b and uk; they now go out at the launch start.
At 4 rows every CTA routed all four rows before its first up/gate job; a
CTA now runs one row's jobs, its first after that row's first two picks.
TP4 now runs 8-row steps on the kernels.
- **NextN.** The entry runs as four K parts on 192 CTAs instead of two
halves on 96. Each stage reloads its pointer arguments from the kernarg
segment, which avoids a FlyDSL 0.4.0 LLVM miscompile (ROCm/FlyDSL#1344).
- **The front** (q_a | kv_a) ran as 164 whole-K tiles while 92 CTAs
waited; each tile's K is now split between its own CTA and a spare one.
Top-8 rounds reduce the expert ID only on ties, and the router's RMSNorm
moves into its logits.
- **The indexer launch.** Above 4 rows two rows score a wave; the index
projection tiles run alone and the scoring skips the q_norm CTAs; split
stages q and the new rows before the selection's wait; the selection
warms its code in L2, bounded at the code's end.
- **The MoE at 8 rows** runs each expert the step uses once in up/gate
and down. FP8 weights load 16 bytes a lane, weights read once a launch
load nontemporal, and at 4 rows a repeated expert streams once.
- **SGPR spills.** The indexer and dense-MLP launches reload their
pointer arguments at each stage, as the NextN launches do.
- **The tail** keeps its LM-head loads in flight across its tiles.

**Every output and every write stays bitwise equal to #75, except the
K-split front's** (commits 8 and 10, below). Every other commit is
bitwise against its parent at TP8 and TP4, every width, on the bf16 and
the E4M3 MLA cache. The NextN entry's partial-sum mailbox changes layout
(commit 4).

**The one non-bitwise part, the front:** each front part multiplies
bf16(x · gamma) rather than the normalized row, and the consumers scale
the sum by the rstd, so q_an, the KV rows, the router scores and the
outputs move by about one bf16 rounding; the router fold changes the
router logits only. Its accuracy evidence is GSM8K with MTP against an
A/A pair and KL against prefill at TP8 (Accuracy Tests).

Result on MI355X, #75 → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at the
golden accept length 2.99):

- **Per cycle at 8K: 6.467 → 6.017 ms at one request (−7.0%) and 9.967 →
8.560 ms at two (−14.1%)**, measured step by step from #75 through
commit 22 (Speed Tests). The tail's in-flight LM-head loads (commit 24)
and the bounded code warm (commit 15) were not in that serving tree, so
they are not in this number.
- **TP4, two requests, 8K:** −10.4% for commits 1–5, −5.0% for the MoE
dedup (16, 17), −4.3% for the 4-row expert-once and nontemporal loads
(19–21).
- Real acceptance (TP8, 8K, paired): TPOT −4.5% at one request and −7.8%
at two for commits 1–10 (Speed Tests).

## Modifications

24 commits, one per change, in the order below (#). Three only add test
options.

| # | Area | Commit | Change | Numerics |
|---|---|---|---|---|
| 1 | Shared layer and MoE | mid as FP8 once | Wave k of down task t
quantizes mid block t + k × 256 and publishes its FP8 bytes and scale
(`midq`, `midqs`). Every down task polls those instead of all of mid. At
TP4 and 8 rows (288 blocks) two waves a task quantize. MXFP4 experts
keep their path | bitwise |
| 2 | Shared layer and MoE | early split rows | Above 4 rows, in
launches without the indexer (TP8; at TP4 there are more split tiles
than CTAs), each split tile's keys and KV cache rows are loaded right
after the CTA's q_norm and cache tasks, ahead of its q_b or uk task |
bitwise |
| 3 | Shared layer and MoE | per-row up/gate jobs at 4 rows | With one
up/gate task round (TP8, 4 rows), the CTA holding tiles c of slot group
g runs row g % 4 at slots 1 + g/4, +2, +4, +6: the same tiles as before,
each on its XCD. Its routing waves run the first two top-8 rounds, park
their state in LDS, issue the first job's weights, and run the other six
rounds under them. A tied round hands the row to the existing top-8 |
bitwise |
| 5 | Shared layer and MoE | TP4 at 8 rows | The per-TP row cap goes.
TP4 routes target decode / verify steps and draft extends of up to 8
rows. Draft decode steps keep their 4-row limit at both TP sizes. No
code object changes | — |
| 4 | NextN entry and its kernarg reload | NextN entry in four parts |
Each 16-column entry tile runs as four K parts of `[enorm(e) |
hnorm(h)]` on 192 CTAs (TP8). The parts' f32 partials are added in part
order by the first part, which pushes the tile. With the indexer the
entry starts past the front's index projection tiles. TP4 keeps the two
halves (192 tasks already). The entry normalizes 4 rows at a time, and
at up to 4 rows a NextN launch polls its pushed rows 2 at a time in the
front tiles | bitwise, except the `ehp` mailbox layout |
| 6 | NextN entry and its kernarg reload | NextN kernarg reload |
`KERNARG_I64` lists the kernel's Int64 arguments in signature order
(argument i is the 8 bytes at 8·i of the kernarg segment). A NextN
launch takes the segment's address in its prologue
(`llvm.amdgcn.kernarg.segment.ptr`); `reload_args()` rebinds the 35
pointer arguments the stages read, and `xin` from the reloaded `sym`, as
invariant scalar loads behind an empty asm on the segment's address, so
LLVM cannot merge one stage's loads with another's. Each of the 14
stages after the NextN entry calls it at its start; LLVM drops the loads
a stage does not use. Other kinds trace exactly as before | bitwise |
| 7 | NextN entry and its kernarg reload | NextN kinds in
`time_layers.py` | `--kind nextn_index` and `--kind nextn_shared` time
one NextN launch a step, `check_nextn.py`'s draft layer, at the served
decode points' rows, in the same graph replays as the other kinds.
`check_nextn.Draft` takes the rows directly | test only |
| 8 | K-split front (the one non-bitwise part) | K-split front | Each
q_a / kv_a tile's first 4096 K run on its own CTA, and the remaining
2048 on one of the 92 CTAs past the tiles, two tiles' rests a CTA.
Neither part normalizes: each multiplies bf16(x · gamma). q_norm and the
cache tasks add the two partials (`q_a2`, `kv_a2` hold the rests) and
scale them by the row's rstd, which they compute from x before their
wait. The index projection tiles keep whole K, on the CTAs with one
rest, and the NextN entry starts past them. The MoE stages keep their
CTAs. A launch from pushed rows (layer 0, NextN) reads them in both
parts and in the rstd, with its weights issued before the rows' wait |
not bitwise: q_an, the KV rows, the router scores and the output move by
about one bf16 rounding (GSM8K and KL in Accuracy Tests) |
| 9 | K-split front (the one non-bitwise part) | top-8 tie shortcut | A
round takes the key maximum. When one lane holds it (a ballot), that
lane's head is the pick; only an exact tie runs the ID reduction. The
picks come out in the same order | bitwise |
| 10 | K-split front (the one non-bitwise part) | router norm in its
logits | The router task stores bf16(a · gamma) as its rows arrive, with
each wave's sum of squares beside them. One barrier serves both. The
rstd, summed in the same order, scales the FP8 activation blocks
(bitwise) and the 8 logits | not bitwise: router logits only |
| 11 | Indexer launch | two-row scoring above 4 rows | Above 4 rows a
scoring wave runs two rows' MFMAs on one LDS read and FP8 → bf16
conversion of each key, a run's keys load nontemporal, runs split evenly
once every scoring CTA has one, and a multi-batch run stages a batch's
scores in LDS and stores them after the next batch's keys land. The
index-weights tiles keep a quarter of their weights in flight before the
rows' staging instead of half. Timeline builds load the timeline's
address once a launch. At 2 and 4 rows the scoring code is unchanged |
bitwise |
| 12 | Indexer launch | index tiles and scoring placement | With the
indexer the 10 index K / weights projection tiles run alone on the
grid's last 10 CTAs and the K rests go two a CTA on the CTAs before
them; the NextN entry, which starts past the index tiles, starts at CTA
0. The scoring runs go to CTAs S .. G − 1 only | bitwise |
| 13 | Indexer launch | split before the selection, selection code warm
| With the indexer, split polls q and the rows this launch appends
before the selection's wait, staging the rows in its reduction and
output regions; each wave then polls its own 8 selected keys and issues
their cache rows as they land. Each selection slice, while its scores
are pending, loads the code after it into its XCD's L2 (one word a
64-byte line, from `s_getpc`). TP8 dense launches above 4 rows skip that
warm | bitwise |
| 14 | Indexer launch | the two-row scoring's queries in MFMA operand
order | Above 4 rows the scoring stages its rows' queries in LDS in MFMA
operand order (head half, 32-dim chunk, lane), so each wave reads a
query fragment as one conflict-free 16-byte load a lane; row-major, a
fragment's 16 lanes read 16 heads 64 words apart, all on one bank, twice
a wave | bitwise |
| 15 | Indexer launch | the code warm ends at the code's end |
`warm_code` takes its PC and the distance to `.Lfunc_end0` from one
inline asm (`s_getpc_b64`, a label, `s_mov_b32 sN,
.Lfunc_end0-.Lwarm_code`; the assembler resolves the distance). The
buffer's size is min(distance, 64 KB), and loads past it read nothing.
With a fixed 64 KB the dense and embed launches read up to about 8 KB
past their code, and TP4 dense launches faulted (illegal address). Where
64 KB fits, the read is unchanged | bitwise |
| 16 | MoE dedup and weight loads | each expert once in up/gate and down
| Above 4 rows (not dense layers) every wave finds the step's distinct
routed experts from the top-8s (lane p holds pick p, a ballot marks
first uses). up/gate jobs are the 8-intermediate tiles of the shared
expert and of each distinct expert, each with every using row as an MFMA
column, dealt round-robin over the CTAs; every CTA routes, issues the
activations' poll and its first job's weights, then checks the poll.
down's units take one expert's tile against every using row and store
the products in LDS; a row's output is summed over its units in the
one-row form's order and wave segments. The rows the multi-column MFMAs
read are padded by 4 words against bank conflicts | bitwise |
| 17 | MoE dedup and weight loads | dead-batch skip | A wave's down
weight batch whose first unit lies past the step's expert groups skips
its MFMAs | bitwise |
| 18 | MoE dedup and weight loads | 16-byte FP8 loads | `ld_fp8_chunk`,
shuffled layout: a lane of an even 16-lane row loads the whole 16-byte
segment of a row's first K32 block, and its partner 16 lanes up the same
segment of the second block, as one `dwordx4`. `_RowSwapped` holds the
raw 16 bytes; reading element `e` runs `v_permlane16_swap` on words (e %
2, e % 2 + 2), so the load stays in flight across the task's wait. The
packed layout and the callers are unchanged | bitwise |
| 19 | MoE dedup and weight loads | `--similar` | `check_layer.py` and
`time_layers.py --similar X`: the live rows of a request are one random
row plus X times their own, so they pick many of the same experts, as a
verify step's rows do. The default keeps the independent rows and their
digests | test only |
| 20 | MoE dedup and weight loads | 4-row repeated expert once | At 4
rows with one up/gate task round (TP8, not dense, not the NextN indexer
launch), each job takes every row as an MFMA column. The owner of an
expert is the row whose job of it comes first, the lowest such row on a
tie, found by a ballot over the step's picks; the owner writes mid for
every row that picked it, through the 8-row path's `job_emit`, and the
other rows' later jobs of it load nothing | bitwise |
| 21 | MoE dedup and weight loads | nontemporal once-read weights |
`ld_fp8_chunk` and the FP8 unit helpers take a cache policy (`CM_NT` in
`tiles.py`). At 2 and 8 rows: q_a \| kv_a, q_b, W_UK and o_proj, and the
experts' up/gate and down (routed, shared and the dense layers' slices).
At 4 rows: up/gate only, where commit 20 applies | bitwise |
| 22 | Kernarg reload for the indexer and dense launches | kernarg
reload for indexer and dense launches | `RELOAD_ARGS` (NextN, indexer or
dense-MLP launch, layer 0 included) selects the launches whose 14 stages
reload their pointer arguments from the kernarg segment, as the NextN
launches already do; `reload_args()` rebinds `xin` only where x is
pushed. Shared-layer and tail launches build exactly as before | bitwise
|
| 23 | LM-head tail | tail digest in `check_layer.py` | `--dump` digests
the tail's gathered logits | test only |
| 24 | LM-head tail | tail loads in flight | A tail wave keeps
`HD_AHEAD` units in flight across its tiles: each unit it multiplies is
replaced by the next one of its tile sequence, so the next tile's
weights load while the current tile reduces and pushes. The norm's row
loads go out before the first weights | bitwise |

- **Layout:** `midq` and `midqs` after `mid` (19 KB at 4 rows, 38 KB at
8); `ehp` holds three NextN entry parts' partials instead of one; `q_a2`
and `kv_a2` hold the front's K rests. The MoE group table and down's
products live in the LDS X region; no other new scratch.
- **Spec and draft:** `TP_MAX_ROWS` and `_Built.max_rows` go;
`DeviceLimits(widths=ROWS)` bounds every routed step at 8 rows.
`test_glm53_kernels.py` runs TP4 at widths 2, 4 and 8. The README's
GLM-5.3 row says up to 8 rows at both TP sizes.
- **Placement:** `qkv_a` has a task on every CTA; q_norm, q_b, split, uk
and uv follow at CTA 0 as before, and o starts at split's first CTA.
With the indexer the index projection tiles take the grid's last 10
CTAs; the NextN entry starts at CTA 0 there.
- **`poll`** takes an optional `between()` that runs after the first
loads are issued and before they are checked (up/gate issues its first
job's weights there).
- **What changes in #74's code:** the scoring keeps its per-row path at
2 and 4 rows and adds a two-row path above 4; a run's first CTA and page
range are counted over the scoring CTAs (S .. G − 1). The selection
gains one call, the code warm. Split's indexer path is reordered;
without the indexer split keeps its order.

## Accuracy Tests

All on MI355X in the image below, under torchrun on 8 GPUs (TP8) and 4
GPUs (TP4).

- **Bitwise, each commit against its parent** (every commit but 8 and
10). Commits 1–21 were measured with these tools on the same commits in
an earlier stack, whose kernel files are this branch's except commit
15's bounded code warm. Commits 15 and 22 were measured as this branch's
head against that stack's head. The dump set has no head kind, so commit
24's tail logits were not in this run. Runs at 2 rows need #76's
`peer.py`: without it, #75 run twice gave 1,693 differing digests at
TP8, and a few 2-row configs differ run to run at TP4 too; with it, 0.
- Run as, on each tree, for each kind:launches in `shared:3 index:3
model:8 dense:3 embed:3` and each kv in `bf16 fp8`: `torchrun
--nproc-per-node 8 test/kernels/glm53/check_layer.py --kind K --launches
N --rows 2,4,8 --rowsets p0,p1,v4 --kv KV --dump DIR/<tree>/tp8` (and
`--nproc-per-node 4` into `tp4`); then `python3
test/kernels/glm53/check_layer.py --compare DIR/<parent>/tp8
DIR/<tree>/tp8`. Commit 20 also with `--rowsets v4 --similar 1.0`.
Negative control: the same dump of #75 with `--perturb`.

  | Area (commits) | TP8 | TP4 |
  | --- | --- | --- |
| Shared layer, MoE and NextN entry (1–5) against #75 | 0 differ | 0
differ |
  | NextN reload (6) | 0 differ | 0 differ |
| Front (8–10) | differs by design (commits 8 and 10); every CHECK
passes | same |
  | Indexer launch (11–15) | 0 differ | 0 differ |
  | MoE dedup (16, 17) | 0 differ | 0 differ |
  | 16-byte loads (18) | 0 differ | 0 differ |
  | 4-row expert once, nontemporal (20, 21) | 0 differ | 0 differ |
| Bounded code warm (15), kernarg reload (22): head against the earlier
stack's head | 0 differ | 0 differ |
| Tail loads (24): the tail's logits (`--kind head`, commit 23) | not in
this run | not in this run |

- **The front, not bitwise (commits 8 and 10).** Measured at TP8 with
MTP on and real acceptance, the commit before the front (#23's head)
against the front's head, and #23's head again as an A/A pair:
- **GSM8K** (1,319 questions, 8-shot, two requests): 0.9234 (#23),
**0.9166 (front)**, 0.9136 (#23 again). McNemar, one-sided: front vs #23
p = 0.145; the A/A pair p = 0.065; all above 0.025.
- **KL decode ‖ prefill** (MTP off, 64 prompts × 128 tokens, top-20):
0.0848 (#23), **0.0766 (front)**, 0.0721 (#23 again). Top-1 agreement
0.931 / 0.932 / 0.936.
  - On both, the front's value lies within #23's own A/A spread.
- AIME 2025, 30 problems × 16 samples, with the model card's sampling
(temperature 1.0, top_p 0.95, thinking on, up to 163,840 tokens). TP8,
FP8 KV, MTP 3-1-4. Generations average about 23K tokens and some reach
the cap, so the indexer runs over long KV.

| | Accuracy | 95% CI | Accept length |
|---|---|---|---|
| SGLang without monokernels | 99.17% | 98.1–100 | 3.21 |
| sgl-amd-mono main | 99.17% | 98.3–99.8 | 3.23 |
| With this PR, #86, sgl-project/sglang#47324 and its follow-up | 98.33%
| 96.9–99.6 | 3.23 |

The stack has 8 wrong answers out of 480, against 4 for each of the
others (Fisher p = 0.39; against main, paired by problem, p = 0.36).
Samples are independent draws. The reference served 16 requests at a
time, the monokernel builds 2.
- Decode drift at long context: KL of decode against SGLang's stock
prefill of the same tokens (top-20), and top-1 agreement. TP8, FP8 KV,
MTP off; 128 prompts × 256 greedy tokens at each of 8K, 32K, 64K and
128K, cut from the AIME reasoning traces; two servers per build.

| | KL decode ‖ prefill | Top-1 agreement |
|---|---|---|
| SGLang without monokernels | 0.0249 | 95.94% |
| sgl-amd-mono main | 0.0199 | 96.34% |
| With this PR, #86, sgl-project/sglang#47324 and its follow-up | 0.0203
| 96.20% |

The stack against main: KL +1.9% (95% CI −1.0 to +5.0), top-1 −0.14
points, both within the spread between two servers of the same build (up
to 4% in KL and 0.25 points in top-1). The front alone (commit 10
against commit 7): KL −1.0% (−4.8 to +3.0). With MTP on at 64K: KL −5.7%
(−12.8 to +2.1).
- **Stress** (TP8, the front's head): 4 runs × 20,000 graph replays, 0
bad.
- **E2E** (`test/e2e/test_glm53_server.py`) at TP8 and at TP4
(`--tp-size 4`), on the same commits in an earlier stack: 8 passed.
`TestMonokernelFault` fails as it does on main.

## Speed Tests and Profiling

Setup as #75: image
`lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs
`sha256:fb07e0ef78fe`) with SGLang's `python/sglang` overlaid at
`5044779f`, AITER 285579b6, FlyDSL 0.4.0, `zai-org/GLM-5.3` at
`ba32f8c1`, the e2e test's server flags with `--kv-cache-dtype
fp8_e4m3`.

**Serving ladder**, TP8, 8K, golden AL 2.99, FP8 E4M3 KV, MTP 3-1-4; one
server per arm; the cycle at two requests is the last-prefilled
request's. Measured step by step on the same commits in an earlier
stack:

| Step (this PR's commits) | Cycle, 1 request (ms) | Δ | Cycle, 2
requests (ms) | Δ |
| --- | --- | --- | --- | --- |
| #75 | 6.467 (A/A 6.444) | | 9.967 (A/A 9.986) | |
| shared layer, MoE, NextN entry, TP4 at 8 rows (1–5) | 6.296 | −0.171
(−2.6%) | 9.585 | −0.382 (−3.8%) |
| NextN reload and the front (6–10) | 6.232 (first sweep 6.236) | −0.064
(−1.0%) | 9.257 (first sweep 9.320) | −0.328 (−3.4%) |
| indexer launch (11–15) | 6.108 | −0.124 (−2.0%) | 9.189 | −0.068
(−0.7%) |
| MoE dedup (16, 17) | 6.103 | −0.005 (−0.1%) | 8.780 | −0.409 (−4.5%) |
| 16-byte loads (18) | 6.126 | +0.023 (+0.4%) | 8.776 | −0.004 (0.0%) |
| 4-row expert once, nontemporal (19–21) | 6.049 | −0.077 (−1.3%) |
8.562 | −0.214 (−2.4%) |
| kernarg reload (22) | 6.017 | −0.032 (−0.5%) | 8.560 | −0.002 (0.0%) |
| **#75 → (22)** | | **−0.450 (−7.0%)** | | **−1.407 (−14.1%)** |

- 128K, the front's step → the indexer launch (11–15): 6.534 → 6.417 ms
(−1.8%) at one request, 10.195 → 10.070 ms (−1.2%) at two.
- The tail (23, 24) and the bounded code warm (15) are not in this
ladder.

**TP4, two requests, 8K** (cycle in ms):

| Step | Before | After | Δ |
| --- | --- | --- | --- |
| #75 → shared layer, MoE, NextN entry, TP4 at 8 rows (1–5) | 15.857 |
14.213 | −10.4% |
| indexer launch → MoE dedup (16, 17) | 14.043 | 13.346 | −5.0% |
| 16-byte loads → 4-row expert once, nontemporal (19–21) | 13.319 |
12.741 | −4.3% |

**Real acceptance**, #75 → the front's head (commits 1–10), TP8, 8K, 40
waves, paired:

| Requests | AL | TPOT (ms) | Δ TPOT | Paired Δ |
| --- | --- | --- | --- | --- |
| 1 | 3.265 → 3.269 | 1.936 → 1.836 | −5.2% | −4.5% |
| 2 | 3.290 → 3.310 | 2.979 → 2.671 | −10.3% | −7.8% |

**AgentX** (InferenceX `inferencex-agentx-mvp`, seed 42, traces in
InferenceX order) with AMD's InferenceX recipe: TP8, FP8 KV, MTP 3-1-4
at simulated accept length 2.99, one hour per point. Main and the stack
ran back to back on the same node.

| Concurrency | | P90 interactivity (tok/s/user) | Throughput
(tok/s/GPU) | ITL p50 (ms) | TTFT p50 (s) |
| --- | --- | --- | --- | --- | --- |
| 1 | SGLang without monokernels | 181.5 | 1,944 | 4.808 | 0.351 |
| | sgl-amd-mono main | 291.1 | 2,890 | 2.484 | 0.372 |
| | With this PR and #86 | 345.3 (+18.6%) | 3,135 (+8.5%) | 2.207
(−11.2%) | 0.375 |
| 2 | sgl-amd-mono main | 212.2 | 2,911 | 2.491 | 0.318 |
| | With this PR and #86 | 282.9 (+33.3%) | 3,160 (+8.5%) | 2.219
(−10.9%) | 0.318 |
| 4 | sgl-amd-mono main | 167.1 | 4,069 | 2.706 | 0.256 |
| | With this PR and #86 | 181.0 (+8.4%) | 4,369 (+7.4%) | 2.326
(−14.1%) | 0.263 |

- Paired by request (same conversation and turn), ITL falls 11.6% at
concurrency 1 (95% CI 9.0–12.3), 10.9% at 2 (8.9–11.6) and 10.0% at 4
(9.1–11.9); TTFT is unchanged at every point. Throughput is a fixed
one-hour window, so a faster server gets further into the traces.
- SGLang without monokernels ran at concurrency 1 only, on the same
node. Paired by request, ITL is 49.0% lower on main (95% CI 48.1–50.4)
and 54.7% lower with the stack (54.4–54.8).

### Registers and lint

From the code objects of the spec's builds (`stop=True`):
`/opt/rocm/llvm/bin/llvm-readelf --notes <key>.co`, fields
`.vgpr_count`, `.vgpr_spill_count`, `.sgpr_count`, `.sgpr_spill_count`
and `.private_segment_fixed_size`. Two figures a cell are bf16 / E4M3
MLA cache.

**#75 → this PR, at 8 rows:**
- TP8 index: VGPR spills 24 / 20 → 2 / 1, private segment 100 / 84 → 12
/ 8 B. TP8 dense: 20 → 1.
- TP8 NextN shared: 12 → 0; NextN index: 12 → 6 / 4.
- TP8 shared: SGPR spills 253 / 295 → 23 / 18.
- TP8 embed: VGPRs 228 → 256, no spills.
- TP4 index: VGPR spills 28 → 10; dense 24 / 20 → 4 / 5; NextN shared 20
→ 0.
- The one 8-row build that gains spills is TP4 embed: 0 → 3 VGPR spills
(16 B).
- SGPR spills fall on every 8-row build, from 300–1,700 to about 100.

**Lint:** #84's
`test/kernels/test_kernarg_lint.py` (the GLM-5.3 index and shared
launches at TP4 and TP8, widths 2, 4 and 8, both caches): 24 / 24 pass.

## Reviewer notes

- **The front:** the rstd is computed twice, in q_norm and in the cache
task, from x, before their waits; the extra x read overlaps the front.
With all 256 CTAs in the front, q_norm and the cache tasks sit behind a
front tile on CTAs 0..S−1, which the indexer commits then place around.
The pad-row skip the development tree carried at 2 rows is not part of
this PR.
- **Early split rows assume the selection is the previous launch's**;
the indexer launches select in-launch and use commit 13's path.
- **The code warm** reads code bytes through a buffer resource built
from `s_getpc`, sized by the assembler's distance to `.Lfunc_end0`
(capped at 64 KB), so it stays inside the kernel's code; the words are
kept live through `fresh` and not used.
- **Two scoring paths.** The per-row path stays for 2 and 4 rows; the
two-row path is used above 4 rows only.
- **The expert dedup is 8 rows only;** the dead-batch skip needs its
group count. The group table is per CTA: every CTA routes all rows, so
no cross-CTA hand-off is added and the tie order is unchanged. At 4
rows, `owns(e)` reads the step's picks after the routing barrier; a
row's first job is issued before it and always loads.
- **Nontemporal loads** return the same bits; W_UV, the indexer's
weights, the router, the NextN entry's eh projection and the LM head
keep the default policy, and TP4 at 4 rows keeps it everywhere.
- **The kernarg reload** only moves where the pointers come from. FlyDSL
0.4.0's code objects differ from build to build (ROCm/FlyDSL#1345);
#84 makes them reproducible and adds the lint.
- **The spec** still does not build the NextN shared launch at 8 rows
(#75's limit for draft decode steps).
- `kernel.py` grows by 1,397 lines (+1,843 / −446).

## Checklist

- [ ] `pre-commit run --all-files` passes. (ruff, ruff-format, isort and
codespell are clean on the touched files)
- [x] Tests under `test/` cover the change. (`check_layer.py --similar`,
the tail digest, the NextN kinds in `time_layers.py`; the kernel tests
compare every launch)
- [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 Nov 3, 2026
…ows, keys straight into registers, the scoring layout ahead, first digits from the scoring (#86)

## Motivation

In long-context decode the indexer launch is what grows with context. It
scores every cached key of each row and then selects each row's exact
top-2048, and it runs on 18 of a step's 75 MoE layer launches, on layer
0, on the dense layers 1–2 and in every NextN draft step. In
AgentX-style serving most decode steps sit at 256K keys and above. There
the scoring scan and the selection's first radix digit dominate the
launch:

- **The scan is bound by per-wave compute, not by bytes in flight.** One
8-page batch a CTA in flight already keeps the key loads streaming; the
scoring waves spend the time in address math, FP8 → bf16 conversion, the
head-half reduction and LDS traffic. At 4 rows (one request's MTP verify
rows) each key was read and converted once per row.
- **The scan copied every batch through LDS.** At 4 rows and more, each
8-page batch went through an LDS copy and barriers before any wave
scored it.
- **The scoring stage started with a chain of loads.** Each scoring CTA
first worked out which rows score together (each row's bound, and its
page-table entry against the previous row's slot) before it could poll
its query.
- **The selection's first digit re-read every score.** A selection slice
of up to 8,192 keys built its first 64-bin histogram in a local pass,
and most scores share the one or two bins of their common exponent, so
the LDS atomics serialized.

This PR shortens all four. **Every output, every write, every score and
every selection stays bitwise equal to the base PR**, at TP8 and TP4,
every width, on the bf16 and E4M3 MLA cache, at 8K to 512K keys. The one
addition is a new scratch mailbox, `score_hist`.

Result on MI355X, the base PR → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at
the golden accept length 2.99), measured on the same commits in an
earlier stack, whose kernel files are these PRs' except the base PR's
bounded code warm (Speed Tests):

- **8K:** 6.018 → 5.985 ms at one request (−0.5%), 8.479 → 8.443 ms at
two (−0.4%).
- **128K:** 6.313 → 6.134 ms at one request (−2.8%), 9.372 → 9.173 ms at
two (−2.1%).
- **512K, kernel-only on this branch:** the indexer launch at 4 / 8 rows
132.5 → 108.7 / 218.0 → 185.6 µs (−18% / −15%), the 75-launch model
order 76.8 → 70.8 / 117.5 → 110.2 µs; about 0.52 ms a step less at one
request (Speed Tests).

## Modifications

Four commits, one per change, `kernel.py` and `layout.py`. The head is
bitwise against the base PR.

| # | Commit | Change | Numerics |
|---|---|---|---|
| 1 | two-row scoring from 4 rows | A scoring wave's 16-key group
addresses come from its page and first key in scalar math. A row's two
head halves reduce across lane rows in three permlane swaps, and two
rows share one score store. From 4 rows (was: above 4) a wave scores two
rows, reading and converting each key once for both, with the rows'
queries staged in MFMA operand order; a run's keys load nontemporal,
long rows' runs split evenly, and every run stages its scores in LDS and
stores them as pairs after each batch. A lone live row (2 rows: a
request and its pad row) is scored by all 8 waves | bitwise |
| 2 | key groups straight into registers | From 4 rows each wave of a
row set takes every fourth 16-key group of the CTA's run and loads its
lane's words straight into registers as MFMA B operands, two groups
ahead (four 8-byte loads and the scale a group). The scores are stored
as pairs from the registers, the store's mask dropping keys at or past a
row's bound. No LDS copy of the run, no batch barriers, no score stage;
the direct loads are plain (commit 1's nontemporal loads go with the
batch path). A key this launch writes is taken from its row's `index_k8`
pairs, for every row of the request in the launch. Below 4 rows the
batch path reads all its block ids before its first load and its loads
are unconditional | bitwise |
| 3 | the scoring layout ahead | Above 4 rows, with the paged indexer,
the CTAs past the q_norm ones work out the scoring layout once their
front tile is done, from scalar loads of each row's slot, position and
page-table row: the row groups, their requests' first rows, and each
row's index bound and page-table row. They keep it in LDS ahead of the
arena, and the scoring stage reads it there. Up to 4 rows the scoring
stage still computes it | bitwise |
| 4 | first digits from the scoring | From 4 rows each scoring wave adds
the first radix digit (`radix_key >> 26`) of every score it stores to
LDS bins per row and slice segment. A request's rows join a group only
at the next position, so a run spans at most two of each row's slices.
The run publishes each live row's two 64-bin histograms in the new
`score_hist` mailbox. A selection slice polls its runs' segments (one
run lookup and a compare), publishes their sum as its first-digit
histogram before reading its scores, and its first exchange skips the
local pass; its first chunk of scores loads while the segments are
summed. Above 4 rows the q_norm CTAs, which score nothing, plan the
scoring layout at the scoring stage, since their selection slices read
it | bitwise; new `score_hist` mailbox |

- **Layout:** `score_hist` holds, per row and scoring CTA, two 64-bin
histograms as tagged pairs: 1 MB at 8 rows, 512 KB at 4. Above 4 rows
the scoring layout takes `(5 + 2 S + 3) // 4 * 4` LDS words ahead of the
arena.
- **What changes in #74's code:**
- The scoring (`score_group`): from 4 rows the paired path replaces the
batched one; below 4 rows the batch path stays, with its block ids read
first. The new-key patch keeps #74's rule: every new key of the
request's rows in the launch, from its first row on.
- `score_groups` moves to the launch level and takes the rows' bounds
and page-table lookups as arguments, so the front and the scoring stage
share it. From 4 rows a row joins its group only at the next position.
- The selection (`index_slice`): from 4 rows the first digit's local
pass is replaced by the sum of the slice's runs' segments. The later
digits, the counts and the placement are unchanged.
- **Unchanged:** the shared, NextN shared and tail launches' code
objects.

## Accuracy Tests

The checks, on MI355X under torchrun on 8 GPUs (TP8) and 4 GPUs (TP4),
each against the base PR:

- **Dumps, bitwise:** on each tree, for each kind:launches in `shared:3
index:3 model:8 dense:3 embed:3` and each kv in `bf16 fp8`, `torchrun
--nproc-per-node 8 test/kernels/glm53/check_layer.py --kind K --launches
N --rows 2,4,8 --rowsets p0,p1,v4 --kv KV --dump DIR/<tree>/tp8` (and
`--nproc-per-node 4` into `tp4`); then `python3
test/kernels/glm53/check_layer.py --compare DIR/base/tp8 DIR/pr/tp8`.
With #76's `peer.py` applied, 0 of 33,552 digests differ at TP8 and 0 of
1,944 at TP4. Without it, only a few 2-row configs differ, which is the
known 2-row run-to-run variation that #76 fixes (TP8: bf16 model and fp8
dense, index and embed; TP4: bf16 model and fp8 embed).
- **Top-k exact at 8K to 512K, scores and selections bitwise:**
`torchrun --nproc-per-node 8 test/kernels/glm53/check_topk.py --pos
8447,131071,262143,524287 --rows 2,4,8 --dump DIR/topk/<tree>/tp8` on
the base PR and on this PR (and `--nproc-per-node 4`), then
`check_layer.py --compare` of the two digest sets. The selections are
exact on both trees, on both caches and every rowset. Scores and
selections are bitwise: 0 of 41,984 digests differ at TP8 and 0 of
20,992 at TP4. The `check_layer.py` dump set stays under 9K keys, so
this is the check of the multi-batch runs, the even split, the direct
loads and the first digits.
- **NextN launch = plain launch fed its own x, bitwise:** `torchrun
--nproc-per-node 8 test/kernels/glm53/check_nextn.py --rows 2,4,8 --kv
fp8 --only entry` (and `--kv bf16`, and `--nproc-per-node 4`). Passes at
TP8 and TP4 on both caches.
- AIME 2025, 30 problems × 16 samples, with the model card's sampling
(temperature 1.0, top_p 0.95, thinking on, up to 163,840 tokens). TP8,
FP8 KV, MTP 3-1-4. Generations average about 23K tokens and some reach
the cap, so the indexer runs over long KV.

| | Accuracy | 95% CI | Accept length |
|---|---|---|---|
| SGLang without monokernels | 99.17% | 98.1–100 | 3.21 |
| sgl-amd-mono before #85 | 99.17% | 98.3–99.8 | 3.23 |
| With #85, this PR, sgl-project/sglang#47324 and
sgl-project/sglang#48224 | 98.33% | 96.9–99.6 | 3.23 |

The stack has 8 wrong answers out of 480, against 4 for each of the
others (Fisher p = 0.39; against the build before #85, paired by
problem, p = 0.36). Samples are independent draws. The reference served
16 requests at a time, the monokernel builds 2.
- Decode drift at long context: KL of decode against SGLang's stock
prefill of the same tokens (top-20), and top-1 agreement. TP8, FP8 KV,
MTP off; 128 prompts × 256 greedy tokens at each of 8K, 32K, 64K and
128K, cut from the AIME reasoning traces; two servers per build.

| | KL decode ‖ prefill | Top-1 agreement |
|---|---|---|
| SGLang without monokernels | 0.0249 | 95.94% |
| sgl-amd-mono before #85 | 0.0199 | 96.34% |
| With #85, this PR, sgl-project/sglang#47324 and
sgl-project/sglang#48224 | 0.0203 | 96.20% |

The stack against the build before #85: KL +1.9% (95% CI −1.0 to +5.0),
top-1 −0.14 points, both within the spread between two servers of the
same build (up to 4% in KL and 0.25 points in top-1). With MTP on at
64K: KL −5.7% (−12.8 to +2.1).

## Speed Tests and Profiling

Setup as #75: image
`lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs
`sha256:fb07e0ef78fe`) with SGLang's `python/sglang` overlaid at
`5044779f`, AITER 285579b6, FlyDSL 0.4.0, `zai-org/GLM-5.3` at
`ba32f8c1`, the e2e test's server flags with `--kv-cache-dtype
fp8_e4m3`.

TP8, MTP 3-1-4 at golden AL 2.99, FP8 E4M3 KV, cycle in ms:

| Context | Requests | the base PR | This PR | Change |
| --- | --- | --- | --- | --- |
| 8K | 1 / 2 | 6.018 / 8.479 | 5.985 / 8.443 | −0.5% / −0.4% |
| 128K | 1 / 2 | 6.313 / 9.372 | 6.134 / 9.173 | −2.8% / −2.1% |

- Measured on the same commits in an earlier stack (the base PR's head
as `1b93ec8`, this PR's as `522a7ed`): golden AL 2.99, one server per
arm, the cycle at two requests the last-prefilled request's. The gain
grows with context: the scoring scan and the first digit scale with the
keys a row reads.

**Kernel-only, on this branch** (`time_layers.py --kv fp8 --similar 1.0`
under `torchrun --nproc-per-node 8`, E4M3 cache, µs a launch, the mean
of 2 rounds in alternating order, each round the median of 7 trials of
50 replays; the rounds agree within 0.3 µs). `model` is the mean over
the 75 MoE launches in model order, run one width a process:

| Launch | Rows | 8K | 128K | 512K |
| --- | --- | --- | --- | --- |
| index | 2 | | | 96.85 → 90.97 (−5.9) |
| index | 4 | 79.56 → 77.94 (−1.6) | 92.09 → 85.59 (−6.5) | 132.48 →
108.70 (−23.8) |
| index | 8 | 111.78 → 108.41 (−3.4) | 147.81 → 140.51 (−7.3) | 217.98 →
185.56 (−32.4) |
| dense | 2 | | | 89.83 → 83.41 (−6.4) |
| dense | 4 | 63.08 → 62.09 (−1.0) | 75.53 → 69.66 (−5.9) | 115.72 →
92.02 (−23.7) |
| dense | 8 | 79.52 → 75.80 (−3.7) | 113.53 → 105.53 (−8.0) | 181.71 →
149.72 (−32.0) |
| model | 4 | 63.88 → 63.57 (−0.3) | 67.06 → 65.55 (−1.5) | 76.84 →
70.81 (−6.0) |
| model | 8 | 92.10 → 91.80 (−0.3) | 101.22 → 99.67 (−1.6) | 117.50 →
110.19 (−7.3) |

- From these, the 75 MoE launches and the 3 dense ones of a step take
about 0.52 ms less at 512K and one request (4 rows), 0.13 ms at 128K and
0.03 ms at 8K.
- Commands: `torchrun --nproc-per-node 8
test/kernels/glm53/time_layers.py --kind {index,dense} --rows 2,4,8
--pos 524287 --kv fp8 --similar 1.0`, `--kind model --rows 4` and
`--rows 8` (one width a process), and the same at `--pos 131071` and
`--pos 8447` with `--rows 4,8`.

**AgentX** (InferenceX `inferencex-agentx-mvp`, seed 42, traces in
InferenceX order) with AMD's InferenceX recipe: TP8, FP8 KV, MTP 3-1-4
at simulated accept length 2.99, one hour per point. The build before
#85 and the stack ran back to back on the same node.

| Concurrency | | P90 interactivity (tok/s/user) | Throughput
(tok/s/GPU) | ITL p50 (ms) | TTFT p50 (s) |
| --- | --- | --- | --- | --- | --- |
| 1 | SGLang without monokernels | 181.5 | 1,944 | 4.808 | 0.351 |
| | sgl-amd-mono before #85 | 291.1 | 2,890 | 2.484 | 0.372 |
| | With #85 and this PR | 345.3 (+18.6%) | 3,135 (+8.5%) | 2.207
(−11.2%) | 0.375 |
| 2 | sgl-amd-mono before #85 | 212.2 | 2,911 | 2.491 | 0.318 |
| | With #85 and this PR | 282.9 (+33.3%) | 3,160 (+8.5%) | 2.219
(−10.9%) | 0.318 |
| 4 | sgl-amd-mono before #85 | 167.1 | 4,069 | 2.706 | 0.256 |
| | With #85 and this PR | 181.0 (+8.4%) | 4,369 (+7.4%) | 2.326
(−14.1%) | 0.263 |

- Paired by request (same conversation and turn), ITL falls 11.6% at
concurrency 1 (95% CI 9.0–12.3), 10.9% at 2 (8.9–11.6) and 10.0% at 4
(9.1–11.9); TTFT is unchanged at every point. Throughput is a fixed
one-hour window, so a faster server gets further into the traces.
- SGLang without monokernels ran at concurrency 1 only, on the same
node. Paired by request, ITL is 49.0% lower before #85 (95% CI
48.1–50.4) and 54.7% lower with the stack (54.4–54.8).

### Registers and lint

From the code objects of the spec's builds (`stop=True`):
`/opt/rocm/llvm/bin/llvm-readelf --notes <key>.co`, fields
`.vgpr_count`, `.vgpr_spill_count`, `.sgpr_count`, `.sgpr_spill_count`
and `.private_segment_fixed_size`. Two figures a cell are bf16 / E4M3
MLA cache.

**the base PR → this PR:**
- TP8 index, dense and NextN index at 8 rows lose their last VGPR spills
(1–6 → 0).
- TP4 at 8 rows: index 10 → 1 / 2, dense and embed 3–5 → 1.
- Two builds gain one: TP4 index at 4 rows 0 → 1 (8 B), and TP8 NextN
index at 4 rows, bf16 cache, 0 → 1.
- The shared, NextN shared and tail objects are unchanged.

**Lint:** #84's
`test/kernels/test_kernarg_lint.py` (the GLM-5.3 index and shared
launches at TP4 and TP8, widths 2, 4 and 8, both caches): 24 / 24 pass.

## Reviewer notes

- **New keys, as #74 has them.** A key this launch writes is scored from
its row's `index_k8` pairs for every row of the request in the launch
from its first row on (`c0`), not only the rows of the scoring group.
That matters when one request has more rows than a group (an 8-row draft
chain: two groups of 4). The direct path keeps the batch path's rule,
over the S rows (`new_bounds`), and the scoring layout of commit 3
carries the requests' first rows (`firsts`) for it. `check_topk.py
--rowsets chain` at 8 rows covers it.
- **The join rule.** From 4 rows a row joins its request's group only at
the next position (`bounds[r] == bounds[r - 1] + 1`), so a run spans at
most two of each row's selection slices. A request's rows in an MTP
verify step or a draft chain are always consecutive; other rows start a
new group, which keeps `firsts` for the new keys.
- **`score_hist`.** A scoring CTA publishes, per live row of its group,
the 64 first-digit counts of each of its two slice segments as tagged
pairs. A slice sums the segments of the runs over its key range; the
runs' layout it recomputes is the scoring's own (`score_groups`,
`select_slices`, the even split).
- **Two scoring paths remain.** Below 4 rows the batch path (one row a
wave, LDS batches) stays; from 4 rows the paired direct path. A runtime
switch by key count was not taken: both paths in one object grew it and
cost at short context.
- **The scoring layout ahead (commit 3)** applies above 4 rows only; at
4 rows it does not pay. It assumes the atom KV cache layout (asserted),
as the paged indexer does.
- **Commit 4's trade.** It moves the first digit's counting into the
scoring CTAs: a small cost per scoring CTA at short context (8K), and
the slice's local pass saved at long context. Commit 4 can be dropped
alone.
- `kernel.py` grows by 317 lines (+587 / −270); `layout.py` +3.

## Checklist

- [ ] `pre-commit run --all-files` passes. (ruff, ruff-format, isort and
codespell are clean on the touched files)
- [x] Tests under `test/` cover the change. (`check_topk.py` and the
kernel tests)
- [x] The SGLang commit or PR this was tested against is named above.
- [ ] Accuracy and speed results are attached when kernel or forward
code changes. (Speed: serving at 8K and 128K, above.)
- [ ] 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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants