Skip to content

[Perf] [MI355X] glm53: FP8 mid hand-off, early split rows, per-row up/gate jobs, NextN entry in four parts, TP4 at 8 rows (preview) - #21

Closed
heslami wants to merge 5 commits into
review/feat-glm53-nextn-draftsfrom
review/perf-glm53-shared-moe-nextn
Closed

[Perf] [MI355X] glm53: FP8 mid hand-off, early split rows, per-row up/gate jobs, NextN entry in four parts, TP4 at 8 rows (preview)#21
heslami wants to merge 5 commits into
review/feat-glm53-nextn-draftsfrom
review/perf-glm53-shared-moe-nextn

Conversation

@heslami

@heslami heslami commented Oct 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Motivation

Stacked on sgl-project#75 (and sgl-project#68, sgl-project#69, sgl-project#73, sgl-project#74 under it). Merge those first; this PR's diff is its last five commits.

A GLM-5.3 decode step runs 75 MoE layer launches, and with MTP each cycle also runs the NextN draft launches. This PR shortens three hand-offs inside the layer launch and the NextN entry, and lets TP4 run 8-row steps on the kernels:

  • down read all of mid as tagged f32. Every down task polled the whole mid tensor from the scratch and quantized it itself: 256 CTAs read the same 74 KB (4 rows) or 147 KB (8 rows). Each mid block is now quantized to FP8 once and every task reads the FP8 blocks, a quarter of the bytes.
  • At 8 rows the split tiles' KV rows waited behind q_b and uk. In a launch without the indexer the selection is the previous launch's, so the rows go out at the launch start, ahead of the CTA's q_b or uk task.
  • At 4 rows every CTA routed all four rows before its first up/gate job. A CTA now runs one row's jobs, and its first job goes out after that row's first two top-8 picks.
  • The NextN entry ran as two halves on 96 CTAs. It now runs as four K parts on 192 CTAs, starting past the front's index projection tiles.
  • TP4 capped steps at 4 rows. At TP4, MTP at two requests is one 8-row launch, which is about 14% a cycle faster than SGLang's forward. TP4 now routes up to 8 rows, as TP8 does.

Every output and every write stays bitwise equal to sgl-project#75, at TP8 and TP4, every width, on the bf16 and the E4M3 MLA cache (Accuracy Tests). The one exception is the layout of the NextN entry's partial-sum mailbox.

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

  • Per cycle at one request: TBD ms (TBD tok/s per user). At two requests: TBD ms.
  • TP4, two requests: TBD ms a cycle.
  • Kernel-only estimates from development, measured on a 1-GPU loopback of a TP8 rank (Speed Tests): about −0.24 ms a cycle at one request and −0.7 ms at two (8K).

Modifications

Five commits, one per change. Each has its own check against its parent (Accuracy Tests) and its own kernel-only estimate (Speed Tests).

# Commit Change Numerics Kernel-only estimate
1 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 shared launch −1.0 / −1.9 / −7.4 µs at 2 / 4 / 8 rows; model order −0.8 / −1.7 / −7.5
2 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 8 rows: shared −1.3 µs, model order −0.65
3 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. These are the same tiles as before, and each keeps 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 4 rows: shared −1.0 µs, model order −1.45, indexer −0.6
4 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
5 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 — TP4, MTP at two requests: about −14% a cycle against SGLang's forward (serving, during development)
  • The kernel-only estimates come from development trees, measured on the 1-GPU loopback with E4M3 KV at 8K, µs a launch. Rows 1 and 2 were measured on a tree that also had the K-split front and an 8-row expert dedup in down; at 8 rows this PR's down path is the one without the dedup, so the 8-row numbers there are upper bounds. The table on this head is TBD (Speed Tests).
  • Layout: midq and midqs go in the scratch after mid: 19 KB at 4 rows, 38 KB at 8. ehp holds three parts' partials instead of one.
  • Spec and draft: TP_MAX_ROWS and _Built.max_rows go. The framework's DeviceLimits(widths=ROWS) already bounds every routed step at 8 rows. test_glm53_kernels.py runs TP4 at widths 2, 4 and 8.
  • README: the GLM-5.3 row says up to 8 rows at both TP sizes.

Accuracy Tests

All on MI355X in the image below. The kernel checks ran on one GPU: one process runs rank 0 of a TP8 (or TP4) launch alone. Its peers' partial sums are zeros tagged with the launch's epoch, and its pushes go to scratch sinks. A NextN launch's tokens are the rank's own, and the peers' columns of its entry are the torch entry's. The harness wraps check_layer.py and check_nextn.py unchanged, out of tree. For check_nextn.py it adds the NextN shared launch at 8 rows and takes the entry's bounds on the rank's own columns. A run of the in-tree tools under torchrun at TP8 and TP4 is TBD.

  • Bitwise, per commit (check_layer.py --dump, --compare):

    • Kinds shared (3 launches), index (3), model (the first 8 launches) and dense (3).
    • Widths 2, 4 and 8 at TP8 and TP4.
    • Row sets p0, p1 and v4. bf16 and E4M3 MLA cache. 3 steps.
    • 6,588 digests a TP a tree.
    Commit against its parent TP8 TP4
    1 mid as FP8 once 6,588 equal, 0 differ 6,588 equal, 0 differ
    2 early split rows 6,588 equal, 0 differ 6,588 equal, 0 differ
    3 per-row up/gate jobs 6,588 equal, 0 differ 6,588 equal, 0 differ
    4 NextN entry (the kinds above) 6,588 equal, 0 differ 6,588 equal, 0 differ
    4 against [Feat] [MI355X] glm53: The NextN draft steps on their own launches sgl-project/sgl-amd-mono#75 (commit 5 changes no code object) 6,588 equal, 0 differ 6,588 equal, 0 differ

  • NextN launch = plain launch fed its own x, bitwise (check_nextn.py --only entry):

    • NextN index at 4 and 8 rows, with and without a pad row; NextN shared at 2, 4 and 8 rows (8: 8 and 5 live rows).
    • TP8 and TP4, bf16 and E4M3: 20 / 20 checks pass in each of the 4 runs, on this head and on commit 3.
    • Entry against the torch entry: up to 2.8e-5 (bound 1e-4). Against f32: 2.3e-3, the torch entry's own distance (bound 7e-3).
    • The swapped-norm control gives 0.14 against both bounds.
    • The TP4 NextN shared launch at 8 rows on the bf16 cache runs clean on both trees.

  • The new entry gives the old one's values:

    • x (every column), the launch's outputs and its MLA, index and top-k writes are bitwise equal between commit 3 and commit 4 in all 40 entry configurations above.
    • The partial sums are added in a new order, but no bf16 rounding of x changed.

  • E2E (test/e2e/test_glm53_server.py, TP8): TBD.

  • Unit tests on a CPU-only host, in the image without the SGLang overlay: 330 passed, with the same 7 failures and 3 collection errors as [Feat] [MI355X] glm53: The NextN draft steps on their own launches sgl-project/sgl-amd-mono#75. These are the spec-registration and contract tests, which need the overlay's SGLang, plus deepseek_v4 and minimax. The spec and draft tests this PR changes (the row routing, the draft build at TP4 with 8 rows) pass on a CPU host with SGLang 5044779f on the path.

Speed Tests and Profiling

Setup as sgl-project#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, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms:

Requests sgl-project#75 this PR Change
1 (verify 4 rows, draft extend 4) TBD TBD TBD
2 (8 rows) TBD TBD TBD

TP4, FP8 E4M3 KV, MTP at AL 2.99, two requests: TBD → TBD ms.

Kernel-only (time_layers.py), each commit against its parent: TBD.

Registers

The compile survey (CPU only) covers every kind (index, shared, dense, embed, head, nextn_index, nextn_shared) at TP4 and TP8, widths 2 / 4 / 8, both caches, stop on and off: 148 keys a tree.

  • Every key builds on every commit.
  • From [Feat] [MI355X] glm53: The NextN draft steps on their own launches sgl-project/sgl-amd-mono#75 to commit 3, no key gains a VGPR spill or private segment. At TP4, 4 rows, the indexer and NextN builds lose theirs (TP4 index w4: 13–15 VGPR spills / 56–64 B → 0; nextn_shared w4: 24–26 → 0–1).
  • Commit 4 changes only the NextN objects (and TP4 compile jitter elsewhere, with equal counts). It removes their spills:
    • TP8 nextn_index and nextn_shared at 8 rows: 12 → 0;
    • TP4 nextn_index and nextn_shared at 8 rows, stop=True: 16–20 → 0.
    • The exception is the tests' TP4 nextn_shared w4 stop=False build: 1 → 3 VGPR spills (8 → 16 B).
  • Commit 5 changes no code object.

The kernel-argument lint (a scan for kernel arguments that are loaded but never read, or whose address halves mix two arguments) finds 0 hits in every object of every commit: 148 / 148 clean on sgl-project#75 and on commits 1 to 4.

Reviewer notes

  • The expert dedup's dead-batch skip is not here. It skips down's weight batches past the step's expert groups, which exist only with the 8-row expert dedup. The dedup is not in this stack.
  • Early split rows assume the selection is the previous launch's. The indexer launches select in-launch and keep the old order. Anything that moved the selection into a shared launch would have to turn this off.
  • Per-row up/gate jobs apply at TP8, 4 rows, outside the NextN indexer launch, where the first job's weights next to the remaining rounds add VGPR spills. Ties in a round go to the existing top-8 code, so tie order is unchanged.
  • TP4 at 8 rows, and the NextN shared launch at 8 rows. An earlier development tree faulted in this launch at TP4 on the bf16 cache, through a compiler miscompile of kernarg restores (gfx950: SGPR spill restores lanes that were never stored; a kernel then addresses a buffer with two arguments' words (bundled LLVM 24, greedy live-range splitting) ROCm/FlyDSL#1344). That tree had the old two-half entry. On this stack the TP4 bf16 NextN shared launch at 8 rows runs clean with either entry, and the lint is clean. The spec still does not build the NextN shared launch at 8 rows at either TP size ([Feat] [MI355X] glm53: The NextN draft steps on their own launches sgl-project/sgl-amd-mono#75's limit for draft decode steps).
  • kernel.py grows by about 320 lines, a candidate for the stage-builder split.

Checklist

  • pre-commit run --all-files passes. (ruff, ruff-format and isort clean on the touched files; full run TBD)
  • 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. (Serving: TBD)
  • A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)

Hassan Eslami added 5 commits October 28, 2026 15:41
Hassan Eslami
glm53: Quantize each mid block to FP8 once, for every down task to read
494f579 
Every down task polled all of mid as tagged f32 pairs from the scratch and
quantized it itself: 256 CTAs reading the same 74 KB (4 rows) or 147 KB (8
rows). Now wave k of down task t quantizes block t + k x 256 (one wave at
TP8; two at TP4 and 8 rows, where mid has 288 blocks) and publishes its FP8
bytes and scale (midq, midqs); every down task polls those, a quarter of
mid's bytes. The quantization is the same, so every value is unchanged. MXFP4
experts keep their own path.
Hassan Eslami
glm53: Above 4 rows load the split tiles' KV cache rows ahead of q_b …
69b544b 
…and uk

In a launch without the indexer the selection is the previous launch's, so a
split tile's keys and their KV latent and k_pe cache rows are known at the
launch start. Above 4 rows most split tiles share a CTA with a q_norm, q_b or
uk task; their rows now go out after the CTA's q_norm and cache tasks, ahead
of its q_b or uk task, and the tile reads them where it read the rows it
loaded itself. The indexer launches, TP4 (more split tiles than CTAs) and 2
and 4 rows keep the old order. Every value is unchanged.
Hassan Eslami
glm53: At 4 rows run one row's up/gate jobs a CTA, the first after th…
9cc3392 
…at row's first picks

Every CTA routed all four rows (a wave a row, 8 top-8 rounds each) before it
issued its first up/gate job. With one up/gate task round (TP8), the CTA with
tiles c of slot group g now runs row g % 4 at slots 1 + g / 4, 3 + g / 4,
5 + g / 4 and 7 + g / 4: the same (row, slot, c) tiles as before, on other
CTAs, with each tile's chunk (and so its XCD) unchanged. Routing wave w takes
row (row + w) % 4, so wave 0 routes this CTA's own row; the routing waves run
two rounds and keep their state in LDS, the CTA issues the first job's
weights, and the other six rounds, the route weights and the key stores run
under them.

These rounds take the wave maximum of the lane heads into every lane (DPP and
cross-row swaps), and the lane holding it records its expert ID; a round
whose maximum is in more than one lane hands the row to route_top8, which
breaks the tie by ID as before. The picks, the route weights, mid and down
are unchanged. The router CTAs issue their shared tile's weights where they
run it. Other widths, TP4, dense layers and the NextN indexer launch keep
their code.
Hassan Eslami
glm53: Run the NextN entry in four K parts, past the index projection…
1d5e48d 
… tiles

Each of a NextN launch's 48 entry tiles (TP8) runs as four K parts of
[enorm(e) | hnorm(h)] on 192 CTAs instead of two halves on 96: a part reads
98 KB of eh_proj, stages its half's normalized rows and sends its f32
partial to the first part, which adds them in part order and pushes the bf16
tile. With the indexer the entry starts past the front's index projection
tiles (wrapping around the grid), so their weights go out at the launch
start. TP4 keeps the two halves (192 tasks already).

Registers: the entry stages 4 rows at a time, and at up to 4 rows a NextN
layer polls the pushed rows two at a time in its front tiles. The other
kinds' code is unchanged.
Hassan Eslami
glm53: Run steps of up to 8 rows on the kernels at TP4, as at TP8
0b8b657 
On MI355X, MTP at two requests at TP4 served as one 8-row launch is 14% a
cycle faster than SGLang's forward, so TP4 now routes up to 8 rows as TP8
does: the target's decode and verify steps and the draft extend of two
requests. The row cap per TP size goes; no code object changes. The kernel
tests run TP4 at 8 rows too.
@heslami

heslami commented Oct 29, 2026

Copy link
Copy Markdown
Collaborator Author

Superseded by the regrouped preview #34 (and #35 for the long-context indexer).

@heslami heslami closed this Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@heslami