Repository navigation
[Perf] [MI355X] glm53: Score a request's rows together and select the exact top-2048 on up to 64 slices - #74
Conversation
jhinpan
left a comment
There was a problem hiding this comment.
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.pyat 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
--comparecan 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 sametest_graph_replays_equal_the_eager_stepGPU fault at TP4, in PyTorch'sindex_elementwise_kernel(the reuse #76 fixes).test_indexer_selects_the_exact_top_kpassed 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_cachebefore 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.
b6fdf8e to
e9ea165
Compare
|
Rebased onto |
… 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.
e9ea165 to
f044534
Compare
## 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>
…-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>
…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>
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'stest/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:
time_layers.py, E4M3 MLA cache, TP8, µs a launch at 4 rows: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.pyruns indexer launches under torchrun at TP4 or TP8, at chosen positions (--pos).--kind denseand--kind embedrun layer 0 and the dense layers.check_layer.py's random keys;ties);zeros);recent).serving, astime_layers.py);--posand below, some under 2,048 keys (mixed);chain).check.poison_new_index_scales). A launch that scores a key this launch writes from the cache therefore fails.check.launch_scores) against an f64 reference. The reference uses the launch's index query (itsindex_qmailbox, 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.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.missedalso counts the keys of the exact top-2048 (score descending, then position) that the list lacks, ascheck.py'slayer.topk_exactdoes.--dump DIRwrites SHA-256 digests of every row's scores and list incheck_layer.py's format, socheck_layer.py --comparecan match two trees.test_glm53_kernels.pyruns it at TP4 (widths 2, 4) and TP8 (2, 4, 8):check_layer.py's steps gave each row its request'sreq_idsentry, 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_keytask of its own, on the row's q_norm CTA.runner.pypassesindex_paged=key.with_indexer).index_keytask's pairs.2. A request's rows scored together, each key read once (
kernel.py,layout.py)index_keytask also publishes its row's key as cached, FP8 value words and scale, inindex_k8pairs. A batch holding any of those keys patches them in from the pairs, so the cache store need not be visible yet.check_topk.py'schainrowset at 8 rows checks it. With each group patching only its own rows, that check fails (Accuracy Tests).index_qhtasks rope and Hadamard-transform each row's index query once and publish it in bf16; the scoring CTAs poll it into LDS.index_readyflag and theindex_k_newpairs.3. Every row's exact top-2048 on up to 64 slices (
kernel.py,layout.py)layout.select_slicesgives 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.sel_histandsel_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).test/unit/test_glm53_select.pycoversselect_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
check_topk.py, E4M3 MLA cache, 2 indexer launches × 2 steps).chainset 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.check_layer.py --dump,--compare).shared,index,model(the first 8 launches),denseandfull(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.check_layer.pywith 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 owncheck_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.layer.topk_owncompares 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:layer.topk_exactandlayer.topk_ownare 0 in every indexer check;layer.topk.relis at most 0.048 andlayer.index.cosineat least 0.99997.test/e2e/test_glm53_server.py, TP8, on this head):PYTHONPATHwith a pure-Python stand-in for its compilededitdistance, as for [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) #68;mainthe same way, only at its last check: every rank stops at the wait deadline withHIP error: an illegal memory access was encounteredand the server exits without sending a token, but ROCr in our container prints noMemory access fault ... on address (nil)line (not even for a bare one-GPU kernel that only stores to 0x0), so the test'sassertTrue(faults)fails. In the reviewer's setup on [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) #68 and [Perf] [MI355X] glm53: Shorter attention chain, one-pass 8-row front, cheaper routing hand-off (2/2) #69 the line is printed and the case passes.Speed Tests and Profiling
Setup: as #68: one node with 8 × MI355X; image
lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025(squashfssha256:fb07e0ef78fe), AITER 285579b6, FlyDSL 0.4.0;zai-org/GLM-5.3atba32f8c1(FP8); the e2e test's server flags with--kv-cache-dtype fp8_e4m3.python/sglangat5044779f.--enable-amd-monokerneland the plugin's spec registration from before [Feat] [MI355X] dsv4: Add whole-layer DeepSeek-V4-Pro mono kernels (TP8) #48.check.pyandspec.py.Serving: as #68.
SGLANG_SIMULATE_ACC_LEN=2.99; measured 2.99 on every server). A cycle is a request's decode time over its verify steps.sglang.benchmark.servingdoes not import on this image, so a small streaming client against/generatemeasures TTFT, TPOT and the verify count frommeta_info.TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms (median over requests):
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:
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:
TP4, E4M3 MLA cache, indexer launch:
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.
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.glm5_monokernel_0, as on the base.stop=True(spec)stop=False(tests)Tests
test_glm53_kernels.pywith 8 GPUs (TP4 and TP8): 8 passed (17 min 32 s), including the new exact top-k test.Checklist
pre-commit run --all-filespasses.test/cover the change..github/CODEOWNERSthat names its owner. (No new model.)