Skip to content

[Feat] [MI355X] glm53_flash: Add GLM-5.3-Flash decode monokernels - #6

Merged
kevin-mii merged 23 commits into
sgl-project:mainfrom
jhinpan:glm53-flash
Oct 18, 2026
Merged

kevin-mii merged 23 commits into
sgl-project:mainfrom
jhinpan:glm53-flash

Conversation

@jhinpan

@jhinpan jhinpan commented Oct 16, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

On MI355X at TP4 with one request, a GLM-5.3-Flash decode step on SGLang's path runs 1,301 kernels (torch profiler). Its median TPOT is 7.75 ms. vllm-project/vllm#61662 runs each MoE decoder layer of this model as one persistent launch on gfx950.

This PR ports that kernel to the amd monokernel backend of SGLang (RFC sgl-project/sglang#43949). With --monokernel-backend amd, each covered layer of a decode step runs as one launch.

Result on MI355X:

  • At TP4 with one request, the median TPOT decreases from 7.75 ms to 4.40 ms (1.76x).
  • From 1 to 8 requests (1 to 4 at TP2), the speedup is 1.31x to 2.09x.
  • Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is 89.17% (off) and 91.67% (on). On GSM8K, it is 97.35% (off) and 97.27% (on).

Scope:

Modifications

How a decode step runs

flowchart TB
  F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP2)"}
  Q -- "no" --> N["SGLang's path"]
  Q -- "yes" --> R["glm53_flash spec forward"]
  R --> A["SGLang: embedding and<br/>dense layers 0 to 2"]
  A --> L["41 launches, layers 3 to 43<br/>(30 KDA, 11 sparse attention)"]
  L --> Z["SGLang: layer 44,<br/>final norm and logits"]
  L -.->|reads and writes in place| M[("SGLang: KV, index and KDA state pools,<br/>page table, loaded weights")]
  classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328
  classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328
  class F,Q,N,A,Z,M sgl
  class R,L amd
Loading

Grey boxes are SGLang code, and this PR changes none of it. Green boxes are new in this PR.

One launch does these operations for one layer, in this sequence:

  1. The attention-side mHC.
  2. The attention: KDA, or sparse MLA with its kpool indexer and its cache writes.
  3. o_proj and its TP all-reduce.
  4. The FFN-side mHC and norm.
  5. The router.
  6. The shared and routed MXFP4 experts.
  7. The MoE TP all-reduce.

Changes to SGLang (python/sglang/srt)

  • None. This PR adds and changes files in sgl-amd-mono only.
  • Without --monokernel-backend amd, SGLang runs as before. The spec registers for gfx950 only, so CUDA and other ROCm GPUs do not change.
  • With the flag, SGLang calls the spec's forward, not Glm5NextForConditionalGeneration.forward, for each decode step that the spec covers. SGLang still does the scheduling, the cache management, the sampling and all other steps.
  • The plugin reads SGLang internals by name. Each module lists these names in SGLANG_READS. test/unit/test_contract.py checks them on CPU against the installed SGLang.
  • At startup, the plugin uses more GPU memory for launch buffers and weight copies (see "Startup").

Files

  • sgl_amd_mono/glm53_flash/gfx950/: the decoder-layer kernel from [ROCm][Perf][GLM-5.3-Flash] Mono decode layer for GLM-5.3-Flash on MI355X vllm-project/vllm#61662 at ca328e297b. vLLM adapted it from [FlyDSL] GLM-5 fused decode-layer MonoKernel ROCm/aiter#6173. The first GLM-5.3-Flash commit copies the files verbatim. The next commit adapts them:
    • Every wait runs under common/guard.py (deadline, peer flags, fail-stop).
    • The kernel reads the KDA conv and recurrent state pool of SGLang in place. It also reads the DSA MLA and index caches, the compress tails and the 64-token page table in place.
    • The page table is a launch argument, so each captured graph reads its own table.
    • The kernel reads the FP8 attention projections in the AITER (16, 16) layout that SGLang's gfx950 block-FP8 linear method makes. It reads the MXFP4 experts in the layout that SGLang's Quark path makes.
    • Like the vLLM kernel, this kernel quantizes the expert inputs to MXFP8. SGLang's Quark path quantizes them to MXFP4 (W4A4, as the checkpoint declares). Thus the two paths are not bit-exact.
  • sgl_amd_mono/glm53_flash/spec.py: the spec and its forward. The forward is Glm5NextForConditionalGeneration.forward for a covered decode step.
    • A step with fewer rows than a kernel width (2, 4 or 8) runs at the next width. The added rows are pad rows, and pad rows write no cache.
    • At TP2, an 8-row sparse-attention launch does not fit in the LDS, so steps of 5 to 8 rows at TP2 use SGLang's path.
    • The framework sets DeviceLimits.widths for each ISA, not for each TP size. Thus at TP2, the framework still lists 5 to 8 rows as covered.
    • The spec refuses unsupported setups at startup. The config and model checks give their reasons: page size, KV dtype, MoE runner, --enable-linear-replayssm, rotary layout, and others.
    • If the kernel cannot read a weight layout (for example a bf16 shared expert), the build stops and gives the reason.
  • sgl_amd_mono/glm53_flash/check.py: with SGLANG_AMD_MONO_CHECK=1 (eager mode only), each covered layer also runs SGLang's layer on the same input and caches. The check compares the written streams, each state row and each cache row with the bounds in NUMERICS. It logs each value that is outside its bound.
  • sgl_amd_mono/common/: three fixes for problems that the GLM-5.3-Flash launches found:
    • The guard reads its peer flags in loads of 4 dwords or fewer. At TP8, the 8-dword load failed instruction selection.
    • Peer buffers use whole 2 MiB granules. Before this fix, a peer faulted on guard flags after the last whole granule.
    • The peer handshake fails on all ranks when the buffer sizes are different. Before this fix, ranks that built buffers in different orders waited until the deadline.
  • Registration in sgl_amd_mono/__init__.py, a README row, a NOTICE entry, and these tests:
    • test/unit/test_glm53_flash_spec.py (CPU)
    • test/kernels/test_glm53_flash_kernels.py (8 GPUs)
    • test/e2e/test_glm53_flash_server.py (TP4)

Not in this PR

  • An FP8 KV cache.
  • Speculative decoding (MTP).
  • Decode steps of more than 8 rows.

Accuracy Tests

Setup:

  • MI355X, TP4, amd/GLM-5.3-Flash-Quark-MXFP4 at b5688f25, sgl-eval 0.1.2.
  • The two arms ran on the same GPUs. Each arm used a fresh server with CUDA graphs and 8 or fewer requests at a time. Thus all decode steps of the "on" arm ran the launches.
Eval Settings Monokernels off Monokernels on
GSM8K (1,319) thinking, T 1.0, top-p 0.95, seed 42 97.35% 97.27% (13 vs 14 discordant, exact McNemar p = 1.00)
AIME25 (30 x 4 seeds) thinking, T 0.6, top-p 0.95, 65,536 tokens pass@1 82.50% pass@1 91.67%
RULER2 NIAH at 1M tokens (4 subtasks x 2) no thinking, greedy, 256 tokens 40.62% 44.12%
  • With the monokernels, AIME25 is 9.2 points higher. The paired bootstrap 95% CI over the 30 problems is [+1.7, +17.5] points. A sign test over the problems gives p = 0.09 (10 problems better, 3 worse).
  • 14 of 120 samples reached the token limit in the "off" arm, and 6 in the "on" arm. The median output was 7,348 tokens (off) and 5,051 tokens (on).
  • A possible cause of this difference is the MXFP8 expert inputs. I did not isolate it.

Per layer, against SGLang's own layer on the same input and caches:

  • Setup: SGLANG_AMD_MONO_CHECK=1, eager mode. Batches of 1, 3 and 8 rows (4, not 8, at TP2). Positions up to 2.4k. 779 checks for each TP size.
  • Result: no value is outside its bound.
Output TP2 TP4 TP8
written streams, cosine >= 0.974 >= 0.992 >= 0.980
MLA KV rows, cosine >= 0.9991 >= 0.9992 >= 0.9993
index cache values, cosine (scales bit-exact) >= 0.99998 >= 0.99998 >= 0.99998
KDA recurrent state, cosine >= 0.999999 >= 0.999999 >= 0.999998

The streams differ by approximately the error of the MoE: the kernel uses MXFP8 expert inputs, and SGLang's path uses MXFP4.

Decode against prefill (sglang.test.kl_test_utils, 8 prompts x 512 tokens, mean KL):

  • SGLang's path: 0.038.
  • Monokernels on: 0.009 to 0.042 for each batch, over two e2e runs.
  • The e2e test sets the bound to 0.06. It runs one graph for each batch of 2, 3, 7 and 8 rows.
  • With an injected fault (one page-table pointer for all graphs), the 3-row batch measured 0.146.

GPQA-Diamond did not run, because the dataset is gated on the Hub.

AIME26 on daf4739

GSM8K is near its ceiling for this model, so it finds only large errors. The SGLang contribution guide asks for a harder eval, and the GLM-5.3 cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on its final head.

  • Setup: MI355X, TP4, SGLang main 4cca7fcd3e, sgl-amd-mono daf4739, sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh server.
  • 8 or fewer requests ran at a time, so each decode step of the "on" arm ran the launches.
  • Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings of the model), 65,536 max tokens.
AIME26 (30 problems x 8) Monokernels off Monokernels on
pass@1 89.17% 91.67%
majority@8 96.67% 96.67%
samples at the token limit 14 of 240 12 of 240
median output tokens 8,463 6,516
  • Paired over the 30 problems, the "on" arm is 2.50 points higher (bootstrap 95% CI [+0.0, +5.4]). 9 problems are better, 3 are worse and 18 are tied (sign test p = 0.15).
  • The direction is the same as on AIME25: higher accuracy and shorter outputs with the monokernels. A possible cause is the MXFP8 expert inputs.

Speed Tests and Profiling

Setup:

  • Hardware: one node with 8 x MI355X (gfx950).
  • Image: rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014 (sha256:e4498f5d595f), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1 and AITER e7d2453f (= AITER_COMMIT_DEFAULT).
  • SGLang: [Feat] [AMD] monokernels: one runtime switch for vendor monokernel plugins sglang#44353 at 04fcb04e34. Its monokernel framework is the same as the merged ec7ab9d0c2.
  • sgl-amd-mono: this PR on 3970534. The rebase onto 815b104 changed no code. The later commits change no launched code:
    • One commit moves an import below the TP2 fallback.
    • daf4739 refuses --enable-linear-replayssm and corrects comments and docs.
    • The e2e tests ran again on daf4739 (see the end).
  • Server: the cookbook's MI355X high-throughput cell (no speculative decoding), with these differences:
    • This checkpoint. The cookbook's OneNexus/GLM-5.3-Flash-MXFP4 keeps the shared expert in bf16, and the spec refuses that.
    • --kv-cache-dtype bfloat16.
    • --ep-size 1.
    • TileLang DSA prefill and decode.
    • --moe-runner-backend aiter.
  • The only difference between the arms is --monokernel-backend amd.

Method:

  • Workload: sglang.bench_serving, random dataset with fixed lengths (--random-range-ratio 1.0), max(16, 4c) prompts after a warmup.
  • Each arm ran on a fresh server, in two rounds. The second round swapped the arm order and the GPU set (TP8: the order only).
  • At load, each server checked that no other process held memory on its GPUs.
  • The table gives the median TPOT per request in ms over both rounds, with a bootstrap 95% CI.
TP In/out Conc. Monokernels off Monokernels on Speedup
2 8192/256 1 9.31 [9.28, 9.31] 5.68 [5.68, 5.68] 1.64x
2 8192/256 2 9.58 [9.57, 9.60] 5.73 [5.72, 5.75] 1.67x
2 8192/256 4 10.25 [10.23, 10.27] 7.15 [7.13, 7.18] 1.43x
2 512/128 1 9.17 [9.14, 9.18] 5.48 [5.46, 5.48] 1.67x
2 512/128 4 10.31 [10.26, 10.33] 6.88 [6.86, 6.91] 1.50x
4 8192/256 1 7.75 [7.74, 7.75] 4.40 [4.40, 4.40] 1.76x
4 8192/256 2 7.98 [7.98, 8.00] 4.43 [4.42, 4.49] 1.80x
4 8192/256 4 8.39 [8.38, 8.40] 5.03 [5.02, 5.03] 1.67x
4 8192/256 8 9.23 [9.22, 9.49] 7.03 [7.01, 7.39] 1.31x
4 512/128 1 7.66 [7.64, 7.66] 4.16 [4.16, 4.17] 1.84x
4 512/128 2 7.91 [7.91, 7.92] 4.18 [4.18, 4.19] 1.89x
4 512/128 4 8.35 [8.33, 8.36] 4.83 [4.82, 4.83] 1.73x
4 512/128 8 9.39 [9.35, 9.41] 6.76 [6.74, 6.77] 1.39x
8 8192/256 1 7.30 [7.30, 7.30] 3.77 [3.77, 3.78] 1.93x
8 8192/256 2 7.51 [7.50, 7.53] 3.79 [3.79, 3.80] 1.98x
8 8192/256 4 7.90 [7.89, 7.90] 4.25 [4.24, 4.25] 1.86x
8 8192/256 8 8.43 [8.42, 8.62] 5.32 [5.31, 5.51] 1.58x
8 512/128 1 7.21 [7.20, 7.21] 3.52 [3.52, 3.53] 2.05x
8 512/128 2 7.41 [7.41, 7.42] 3.55 [3.54, 3.56] 2.09x
8 512/128 4 7.85 [7.84, 7.86] 4.04 [4.03, 4.04] 1.95x
8 512/128 8 8.40 [8.39, 8.41] 5.06 [5.05, 5.07] 1.66x

After the largest width, both arms use SGLang's path. The control rows stay within 1%:

  • each TP at 16 requests: 0.999x to 1.004x;
  • TP2 at 8 requests: 0.998x and 1.004x.

Cookbook shape on daf4739

The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This table uses that shape at TP4 on daf4739.

  • Each arm ran on a fresh server, in two rounds. The second round swapped the arm order and the GPU set.
  • TPOT is the median per request over both rounds, with a bootstrap 95% CI. TTFT is the median per request. Throughput is the mean over the rounds.
Conc. TPOT off (ms) TPOT on (ms) Speedup TTFT off / on (ms) Output tok/s per GPU off / on
1 7.75 [7.74, 7.75] 4.39 [4.39, 4.39] 1.76x 235 / 232 31.3 / 53.8
2 8.01 [8.00, 8.03] 4.42 [4.41, 4.43] 1.81x 321 / 317 59.6 / 104.3
4 8.40 [8.40, 8.41] 5.02 [5.01, 5.02] 1.67x 211 / 199 115.8 / 190.9
8 9.22 [9.21, 9.38] 7.02 [7.01, 7.17] 1.31x 1,116 / 1,112 176.6 / 220.3
16 10.08 [9.81, 10.30] 10.07 [9.80, 10.31] 1.00x 1,082 / 1,104 337.5 / 337.7

At 16 requests, a step has more than 8 rows, so both arms use SGLang's path.

Profile (TP4, one request, torch profiler, rank 0):

  • Monokernels on: a step runs 184 kernels, and the GPU is busy for 4.28 ms. Of this time, the 41 layer launches use 3.44 ms (84 us each). The trace shows each launch as down.
  • Monokernels off: a step runs 1,301 kernels.

Startup:

  • In a new container, the build of all widths takes 25 to 32 s (FlyDSL and comgr compilation). When the caches hold the kernels, the build takes approximately 2 s.
  • Extra memory for each GPU:
    • TP2: 490 MiB of launch buffers and 455 MiB of weight copies. 261 MiB of the copies is the packed shared expert, which SGLang keeps unfused at TP2.
    • TP4: 214 MiB of launch buffers and 150 MiB of weight copies.
    • TP8: 184 MiB of launch buffers and 128 MiB of weight copies.

Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e:

  • unit (CPU): 75 passed. They also passed on 04fcb04e34, and on daf4739 with SGLang main 4cca7fcd3e and 2948662026.
  • kernels: 24 passed on 5a8b90f:
    • Each GLM-5.3-Flash launch kind at TP 2, 4 and 8, on random weights in SGLang's layouts. All ranks return the same output and write only their valid rows. AITER-shuffled FP8 projections give the row-major result bitwise.
    • The guard tests.
    • The DeepSeek-V4.1 kernel tests with the common fixes.
  • e2e (TP4) has four tests: GSM8K, KL for each batch, the check mode, and an injected wait fault that must stop all ranks.
    • Before this review: 4 passed (GSM8K 0.99, KL 0.009 to 0.018 for each batch).
    • On daf4739, in one run: 4 passed. GSM8K-200 is 0.99. KL is 0.018, 0.015, 0.010 and 0.027 for batches of 2, 3, 7 and 8 rows. The 779 checks are within NUMERICS (lowest stream cosine 0.9904). After the injected stall, the faults are only at address 0x0.

The kernel comes from vllm-project/vllm#61662 (Apache-2.0). vLLM adapted it from ROCm/aiter#6173 (FlyDSL GLM monokernel).

Review fixes (on top of daf4739)

  • New startup refusals:
    • an index cache that SGLang keeps row-major (no AITER preshuffle paged MQA);
    • --enable-hisparse;
    • --expert-distribution-recorder-mode;
    • SGLang's benchmark routing overrides;
    • moe_layer_freq other than 1;
    • a TP2 server whose decode graphs all exceed the 4-row cap.
  • SGLANG_READS declares every SGLang name that the spec and the weight readers read.
  • The stage loop no longer overwrites the kernel's symbol. Before, every variant was named down.
  • Cleanup:
    • startup errors say what they need;
    • the docs match the code (the README row, the copy lists, NOTICE's provenance);
    • common/peer.py's 2 MiB unit is named _GRANULE.
  • The README has one GPU test command per model.
  • CODEOWNERS adds **/*glm53_flash* @jhinpan.
  • Tests on 7944894:

Checklist

  • pre-commit run --all-files passes.
  • Tests under test/ cover the change.
  • The SGLang commit or PR this was tested against is named above.
  • Accuracy and speed results are attached when kernel or forward code changes.

Review and Merge Process

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

A buffer load carries at most four dwords, so at TP8 the single 8-wide load
of every rank's error flags failed instruction selection.
A peer maps an exported uncached allocation in 2 MiB granules. With 4 KiB
rounding, the guard flags past the last whole granule (2 MiB + 256 bytes into
a 4-row GLM-5.3-Flash buffer) faulted on the peer that raised them.
Each buffer's handshake is a collective. Ranks that build several buffers in
different orders map another buffer's memory and wait until the deadline. The
handshake now exchanges each rank's size with its handle and fails on every
rank when they differ.
Verbatim copies of vllm-project/vllm#61662 at ca328e297b (vllm/models/
glm5next/amd/ops/mono and mono/weights.py), themselves adapted from
ROCm/aiter#6173 at c39b56c36 (Apache-2.0). The next commit adapts them to
SGLang and this package.
With --monokernel-backend amd, each MoE decoder layer of GLM-5.3-Flash but the
last runs as one persistent launch on decode steps of 1 to 8 rows (1 to 4 at
TP2), on TP 2, 4 or 8: the attention-side mHC, the attention (KDA, or sparse
MLA with its kpool indexer and cache writes), o_proj and its TP reduce, the FFN
mHC and norm, the router, the shared and routed MXFP4 experts and the MoE TP
reduce. A step runs at the next kernel width (2, 4 or 8), the rest pad rows
that write no cache. The spec's forward runs SGLang's embedding, dense layers,
last layer, final norm and logits processor around the launches.

The kernels read SGLang's state in place: the KDA conv and recurrent state
pool, the DSA MLA and index caches and compress tails, and the 64-token page
table (a launch argument, so each captured graph reads its own). They read
the FP8 projections in the AITER (16, 16) layout SGLang's gfx950 block-FP8
linear method leaves them in, and the MXFP4 experts as SGLang's Quark path
shuffles them.

SGLANG_AMD_MONO_CHECK=1 (eager) runs the normal layer next to each covered
layer and checks the written streams and every state or cache row it writes
against the spec's NUMERICS.
- unit (CPU): the covered layers, the registry entry, the row routing and the
  TP2 row cap, and the config and model refusal reasons
- kernels (8 GPUs): each launch kind on random weights in SGLang's layouts at
  TP 2, 4 and 8 finishes on every rank with the same output and writes exactly
  its valid rows' state or cache rows; AITER-shuffled FP8 projections equal
  row-major ones bitwise
- e2e (TP4): each covered layer within NUMERICS against the normal layer (1, 3
  and 8 rows, positions past 2048); GSM8K while requests join and leave a batch
  that steps past 8 rows; decode against prefill (KL) on one graph per batch of
  2, 3, 7 and 8 rows; and a rank stalled by fault injection stops the server
  before any token of the failed step is sent
On SGLang main, importing layer_boundary reaches fp8_kernel's gfx1250
probe, so the CPU unit test of the TP2 fallback failed without a GPU.
The fallback path does not use the import.
The launch updates the KDA state in place, while --enable-linear-replayssm
defers SGLang's state updates to a ring, so the spec now requires it off.
SGLANG_AMD_MONO_CHECK is read once at build instead of per layer and step.
Two comments, the README design summary and NOTICE now describe the
GLM-5.3-Flash code as it is.
@jhinpan
jhinpan marked this pull request as ready for review October 16, 2026 23:04
@jhinpan

jhinpan commented Oct 16, 2026

Copy link
Copy Markdown
Collaborator Author

@kevin-mii This PR is ready for review, and #7 (GLM-5.3) is stacked on it.

  • CI: the Lint run for this fork PR needs a maintainer to approve it: https://github.com/sgl-project/sgl-amd-mono/actions/runs/38979382220. The repo's pre-commit config passes locally on the head (daf4739).
  • The first three commits change common/ (the guard's peer-flag loads, 2 MiB peer granules and the handshake size check), which DeepSeek-V4.1 runs too.
  • GPU tests: test/kernels needs 8 x MI355X, and test/e2e needs GLM_53_FLASH_MODEL_PATH at TP4. The e2e tests passed again on daf4739; the description has the numbers.

Kevin Mi and others added 9 commits October 18, 2026 09:24
Without AITER's preshuffle paged MQA (Triton < 3.5 and no AOT gluon kernels,
SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1, or SGLANG_USE_AITER=0), SGLang's k-pool
writer keeps the index cache row-major, while the kernels write and read it as
AITER's 16-pool tiles. A routed step then scores a misread cache; SGLang's own
k-pool decode refuses that runtime, and the spec now does too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The contract test pins only declared names, so these reads (the TP group's CPU
group, BuildContext's fields, forward_mode, the decoder layer's norms and
modules, the shared expert's projections and the experts' w13 / w2 weights)
would fail at a GPU server's startup instead of on CPU after an SGLang rename.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The README row counted the covered layers ambiguously, the copy lists missed
the f32 router bias, weights.py said the caller records the copies, NOTICE
covered only the kernel of the files adapted from vLLM, the e2e docstring
skipped the check-mode test, and the kernel docstring still described vLLM's
FFN-only launch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The cache-binding reasons, the o_norm check and the page-table check named
only what they found (or misreported a wrong shape as a missing weight), so a
layout change in SGLang left the reader to find the expected layout in the code.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The constant changed from a 4 KiB page to the 2 MiB IPC granule its comment
describes, so _PAGE no longer named what it holds.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
In weights.py, idx held the layer number next to w_idx and idx_mod, where idx
meant the indexer; spec.py already calls the layer number i.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
spec.py declared the decoder layer's mlp, which only weights.py reads (and
declares).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Kevin Mi and others added 6 commits October 18, 2026 10:05
SGLang accepts --enable-hisparse for GLM-5.3-Flash (a DSA model): its pool hands
out logical KV slots and translates them to device rows on every write, while the
kernels write and read the raw slots, past the device buffer or into other
requests' rows. With --expert-distribution-recorder-mode (also set for EPLB and
expert balancedness reports) the covered layers record nothing, and layer 44
records outside SGLang's per-layer context, which fails decode graph capture.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The forward hands each covered layer's owed FFN post to the next covered layer,
so it needs the MoE layers back to back; a config with moe_layer_freq 2 passed
the checks and gave a dense layer the owed struct at the first routed step.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The stage-base loop rebound name, so every variant (KDA or DSA, each width)
was emitted as the last stage's name, "down": profiles could not tell the
launches apart and IR dumps of all variants shared one directory.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ow cap

At TP2 only steps of up to 4 rows run on the kernels; when every captured
width was wider (say --cuda-graph-bs-decode 8 16), no runtime was built and
every step ran SGLang's forward while startup reported the widths as built.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The serving tests launch their model from the hub id when its path variable is
unset, so the combined command on a DeepSeek-V4.1 box also launched (and
downloaded) GLM-5.3-Flash at TP4.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SGLANG_SIMULATE_UNIFORM_EXPERTS and SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS replace
the router's choice in SGLang's top-k, which routed steps never call: the
kernel route with the model's router, so a benchmark with either set compared
prefill under the override against decode under real routing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@kevin-mii
kevin-mii merged commit 2038b05 into sgl-project:main Oct 18, 2026
3 checks passed
kevin-mii pushed a commit to jhinpan/sgl-amd-mono that referenced this pull request Oct 18, 2026
main has sgl-project#6 (GLM-5.3-Flash with its review fixes) and sgl-project#10. Every GLM-5.3-Flash
and common/ file takes main's version, since this branch carried sgl-project#6's pre-review
copies; NOTICE keeps both provenance paragraphs, the README both model rows (the
GLM-5.3 row in the Flash row's wording) and one GPU test command per model, and
the registry both GLM specs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii pushed a commit that referenced this pull request Oct 18, 2026
main now has GLM-5.3-Flash (#6) and the DeepSeek-V4.1 fixes (#10). CODEOWNERS,
the registry and the README keep both models; the README's launch-count bullet
takes main's wording with MiniMax-M3's one-launch sentence appended, and the GPU
test commands are one per model.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii pushed a commit that referenced this pull request Oct 18, 2026
## Motivation

On MI355X at TP8 with one request, a GLM-5.3 decode step on SGLang's
path has a median TPOT of 11.2 ms. The model is the official
`zai-org/GLM-5.3` (FP8 block-128). ROCm/aiter#6173 runs each GLM-5
decoder layer as one persistent FlyDSL launch. vllm-project/vllm#60831
integrates a variant of this kernel that keeps only MXFP4 experts.

This PR ports the aiter kernel to the `amd` monokernel backend of SGLang
(RFC sgl-project/sglang#43949). It keeps the FP8 block-128 expert path
that the official checkpoint needs. With `--monokernel-backend amd`,
each MoE decoder layer of a covered decode step runs as one launch.

Result on MI355X:

- At TP8 with one request, the median TPOT decreases from 11.24 ms to
5.60 ms (2.01x).
- From 1 to 8 requests at TP8, the speedup is 1.38x to 2.18x. From 1 to
4 requests at TP4, it is 1.45x to 1.82x.
- Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is
82.92% (off) and 86.67% (on). On GSM8K, it is 97.35% (off) and 97.27%
(on).

Scope:

- GPU: MI355X (gfx950), at TP 4 or 8.
- Steps: decode steps of 1 to 8 rows (1 to 4 at TP4). All other steps
use SGLang's path.
- Layers: the 75 MoE decoder layers 3 to 77. 18 launches run the
indexer, and 57 reuse the previous top-k. The dense layers 0 to 2 use
SGLang's path.
- SGLang: this PR needs sgl-project/sglang#44353 (merged as ec7ab9d0c2).
It changes no SGLang file.
- Stack: this PR uses the three `common` fixes of #6, which has merged.
Its diff now shows only GLM-5.3.

## Modifications

### How a decode step runs

```mermaid
flowchart TB
  F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP4)"}
  Q -- "no" --> N["SGLang's path"]
  Q -- "yes" --> R["glm53 spec forward"]
  R --> A["SGLang: embedding and<br/>dense layers 0 to 2"]
  A -- "top-k of layer 2" --> L["75 launches, layers 3 to 77<br/>(18 run the indexer,<br/>57 reuse the previous top-k)"]
  L --> Z["SGLang: final norm and logits"]
  L -.->|reads and writes in place| M[("SGLang: MLA cache, paged FP8 index cache,<br/>page table, FP8 weights")]
  classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328
  classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328
  class F,Q,N,A,Z,M sgl
  class R,L amd
```

Grey boxes are SGLang code, and this PR changes none of it. Green boxes
are new in this PR.

One launch does these operations for one layer, in this sequence:

1. The input RMSNorm, the q_a and kv_a projections, the q_a RMSNorm, and
q_b with RoPE.
2. The KV RMSNorm, the k_pe RoPE and the MLA cache write.
3. The indexer, on an indexer layer only:
   - the index K, Q and weight projections;
   - the index K norm, RoPE, Hadamard transform and cache write;
   - the index scores and the exact top-2048.

   A layer without an indexer reads the top-k of the layer before it.
4. The absorbed query (W_UK), the sparse MLA, W_UV and `o_proj`.
5. The attention TP all-reduce and the residual add.
6. The post-attention RMSNorm, the router in f32, and the FP8
quantization of the activations.
7. The top-8 selection and the experts: 1 shared and 8 routed, in FP8
block-128.
8. The MoE TP all-reduce and the residual add.

### Changes to SGLang (`python/sglang/srt`)

- None. This PR adds and changes files in sgl-amd-mono only.
- Without `--monokernel-backend amd`, SGLang runs as before. The spec
registers for gfx950 only, so CUDA and other ROCm GPUs do not change.
- With the flag, SGLang calls the spec's forward, not
`DeepseekV2ForCausalLM.forward`, for each decode step that the spec
covers. SGLang still does the scheduling, the cache management, the
sampling and all other steps.
- The plugin reads SGLang internals by name. Each module lists these
names in `SGLANG_READS`. `test/unit/test_contract.py` checks them on CPU
against the installed SGLang.
- At startup, the plugin uses more GPU memory. Thus the KV pool is 0.7%
smaller at TP8 and 3.4% smaller at TP4 (see "Startup").

### Files

- `sgl_amd_mono/glm53/gfx950/`: the decoder-layer kernel from
ROCm/aiter#6173 at 56d6390def. The first GLM-5.3 commit copies the files
verbatim. The next commits adapt them:
- Every wait runs under `common/guard.py` (deadline, peer flags,
fail-stop).
- The kernel reads SGLang's FP8 weights in place, in AITER's (16, 16)
layout. These are the attention projections, the indexer projections and
the experts. They include the two gate/up paths that load expert weights
outside the shared unit.
- The kernel reads SGLang's bf16 MLA cache, paged FP8 index cache and
64-token page table in place.
- The page table and its stride are launch arguments, so each captured
graph reads its own table.
- Index keys and queries get the 128-point Hadamard transform after
RoPE. SGLang's gfx950 fused indexer caches rotated keys, so the cache
keeps one format. The spec refuses all other indexers, because they
cache unrotated keys.
- `sgl_amd_mono/glm53/spec.py`: the spec and its forward. The forward is
`DeepseekV2ForCausalLM.forward` for a covered decode step.
- Layer 3 reuses the top-k of layer 2. The forward copies SGLang's
indices into the compact buffer that the launch reads.
- A step with fewer rows than a kernel width (2, 4 or 8) runs at the
next width. One row runs at width 2, because a one-row launch does not
complete on gfx950. The added rows are pad rows, and pad rows write no
cache.
- At TP4, a step of more than 4 rows uses SGLang's path. At TP4, an
8-row launch is slower than SGLang's path: TPOT 18.15 ms against 16.95
ms at 8192 / 256 and c8.
- The spec refuses unsupported setups at startup and gives the reasons.
Some examples:
    - another checkpoint format, page size, KV dtype or MoE runner;
    - an unfused shared expert;
- FP8 projections outside AITER's (16, 16) layout, for example with
`--fp8-gemm-backend triton`;
    - experts that do not use the AITER MoE runner;
    - an indexer other than SGLang's gfx950 fused indexer.
- Weight copies:
- W_UK and W_UV, packed from SGLang's per-tensor FP8 `w_kc` and `w_vc`
(the same values that SGLang multiplies by);
- the router and the indexer's `weights_proj`, packed for the MFMA
tiles;
  - the indexer's k LayerNorm in f32;
  - one bf16 rotary table.
- `sgl_amd_mono/glm53/check.py`: with `SGLANG_AMD_MONO_CHECK=1` (eager
mode only), each covered layer also runs SGLang's layer on the same
input and caches. The check compares the stream, the MLA KV row and the
index keys (after dequantization) with the bounds in `NUMERICS`.
- Registration in `sgl_amd_mono/__init__.py`, a README row, a NOTICE
entry, and these tests:
  - `test/unit/test_glm53_spec.py` (CPU)
  - `test/e2e/test_glm53_server.py` (TP8)
- There is no random-weight launch test like the one for glm53_flash.
The e2e check compares each covered layer with SGLang's layer on the
real weights. The cost is that the test needs the checkpoint.

### Routing numerics

- The kernel routes in f32, as the model defines it: sigmoid plus
correction bias for the choice, and sigmoid for the weights.
- SGLang's AITER top-k casts the correction bias of GLM-5.3
(approximately 34) to bf16. Thus SGLang's choice and weights are
different from f32 routing.
- For this reason, the check mode routes SGLang's reference layer in
f32.

### Not in this PR

- An FP8 KV cache.
- Speculative decoding (MTP).
- MXFP4 checkpoints.
- Decode steps of more than 8 rows (4 at TP4).

## Accuracy Tests

Setup (also for the speed tests):

- Hardware: one node with 8 x MI355X (gfx950).
- Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`
(`sha256:e4498f5d595f`), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1
and AITER e7d2453f.
- SGLang: main at 4cca7fcd3e.
- sgl-amd-mono: this PR. The TP8 runs are older than two commits, and
neither commit changes what TP8 launches:
  - cbc8f15 adds the TP4 row cap.
- 02d42fa adds the review fixes: startup refusals, the rotary check over
each indexer, msgspec config containers, comments and notices.
  - The e2e tests ran again on 02d42fa (see the end).
- Model: `zai-org/GLM-5.3` at ba32f8c19c (FP8 block-128).
- Server: the cookbook's AMD recipe (TP8, TileLang DSA prefill and
decode, `--chunked-prefill-size 131072`, no MTP), with
`--mem-fraction-static 0.80`, `--max-running-requests 96` and
`--cuda-graph-max-bs-decode 96`. SGLang selects the bf16 KV cache,
64-token pages and the AITER MoE runner. TP4 is not a cookbook cell, and
it uses the same flags.
- The only difference between the arms is `--monokernel-backend amd`.
Each arm runs on a fresh server.
- Evals: sgl-eval 0.1.2 at TP8, with 8 or fewer requests at a time.

| Eval | Settings | Monokernels off | Monokernels on |
| --- | --- | --- | --- |
| GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.35% | 97.27%
(10 vs 11 discordant, exact McNemar p = 1.00) |
| AIME25 (30 x 4 seeds) | thinking, T 0.6, top-p 0.95, 65,536 tokens |
90.00% | 89.17% (-0.83 points, bootstrap 95% CI [-5.8, +4.2] over
problems; 2 better, 4 worse; sign test p = 0.69) |
| RULER2 NIAH at 1M tokens (4 subtasks x 2) | no thinking, greedy, 256
tokens | 31.56% | 25.60% |

- AIME25: 11 answers (off) and 13 answers (on) reached the token limit.
The median length was 7,630 tokens (off) and 9,351 tokens (on). For each
problem, the monokernels ran longer on 17 problems and shorter on 12
(sign test p = 0.46).
- NIAH: `mk_niah_basic` scores 100 in both arms. Most other answers
reached the 256-token limit (4 of 8 samples off, 6 of 8 on). The
reasoning lists the needles, and the limit stops the answer. Thus the
multi-value subtasks mostly measure where the answer stops.
- The arms also use different routing precision: f32 (on) and SGLang's
bf16 correction bias (off). See "Routing numerics".

Per layer, against SGLang's own layer on the same input and caches:

- Setup: `SGLANG_AMD_MONO_CHECK=1`, eager mode. Batches of 1, 3 and 8
rows. Positions up to 2.4k. 1,425 checks for each TP size.
- Result: no value is outside its bound.
- The TP4 run is older than the TP4 row cap, so it also covers 8-row
launches.

| Output | TP4 | TP8 |
| --- | --- | --- |
| written stream, cosine | >= 0.9916 | >= 0.9938 |
| MLA KV row, cosine | >= 0.9996 | >= 0.9995 |
| index keys after dequantization, cosine | >= 0.99993 | >= 0.99994 |

Decode against prefill (`sglang.test.kl_test_utils`, 512 tokens, mean KL
for each batch):

- 0.009, 0.017, 0.016 and 0.030 for batches of 2, 3, 7 and 8 rows.
- The bound is 0.06.

### AIME26 on 02d42fa

GSM8K is near its ceiling for this model, so it finds only large errors.
The SGLang contribution guide asks for a harder eval, and the GLM-5.3
cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on
its final head.

- Setup: MI355X, TP8, SGLang main 4cca7fcd3e, sgl-amd-mono 02d42fa,
sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh
server.
- 8 or fewer requests ran at a time, so each decode step of the "on" arm
ran the launches.
- Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings
of the model), 65,536 max tokens.

| AIME26 (30 problems x 8) | Monokernels off | Monokernels on |
| --- | --- | --- |
| pass@1 | 82.92% | 86.67% |
| majority@8 | 96.67% | 96.67% |
| samples at the token limit | 34 of 240 | 29 of 240 |
| median output tokens | 12,197 | 11,492 |

- Paired over the 30 problems, the "on" arm is 3.75 points higher
(bootstrap 95% CI [-1.2, +9.6]). 7 problems are better, 5 are worse and
18 are tied (sign test p = 0.77). The difference is within noise.
- 12% to 14% of the samples reach the 65,536-token limit. Thus these
scores are lower than scores with a larger limit, and you cannot compare
them with the cookbook numbers. Both arms use the same limit.
- The arms also use different routing precision (see "Routing
numerics").

## Speed Tests and Profiling

Method:

- The setup is the same as for the accuracy tests.
- Workload: `sglang.bench_serving` with the random dataset at fixed
lengths (8192 in and 256 out, 512 in and 128 out). Each concurrency c
uses max(16, 4c) prompts after a warmup.
- Each arm runs on a fresh server, in two rounds. The second round swaps
the arm order. TP4 also swaps the GPU set.
- At load, each server checks that no other process holds memory on its
GPUs.
- The tables give the median TPOT per request over both rounds, with a
bootstrap 95% CI.

TP8 (GPUs 0-7):

| Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup |
| --- | --- | --- | --- | --- |
| 8192 / 256 | 1 | 11.24 [11.23, 11.24] | 5.60 [5.60, 5.60] | 2.01x |
| 8192 / 256 | 2 | 12.36 [12.35, 12.37] | 5.78 [5.77, 5.78] | 2.14x |
| 8192 / 256 | 4 | 13.30 [13.29, 13.31] | 7.33 [7.32, 7.34] | 1.81x |
| 8192 / 256 | 8 | 15.85 [15.81, 17.80] | 11.51 [11.49, 13.47] | 1.38x |
| 8192 / 256 | 16 | 18.65 [18.65, 22.51] | 18.94 [18.91, 22.53] | 0.98x
|
| 512 / 128 | 1 | 10.79 [10.79, 10.79] | 5.43 [5.43, 5.44] | 1.99x |
| 512 / 128 | 2 | 11.95 [11.93, 11.97] | 5.48 [5.47, 5.48] | 2.18x |
| 512 / 128 | 4 | 12.81 [12.81, 12.83] | 6.86 [6.85, 6.86] | 1.87x |
| 512 / 128 | 8 | 15.41 [15.40, 15.41] | 10.93 [10.91, 10.95] | 1.41x |
| 512 / 128 | 16 | 17.71 [17.68, 17.73] | 17.68 [17.66, 17.70] | 1.00x |

- At c16, a step has more than 8 rows, so both arms use SGLang's path.
The 16-row decode steps take the same time in both arms (median ITL 18.6
ms at 8192 / 256).
- The 8192 / 256 c16 TPOT is different because of the scheduler. It
admitted one burst of prefix-cached prompts as 1 + 10 + 5 requests, not
as 1 + 15. This stops approximately 10 running requests once, for
approximately 0.1 s.

TP4 (round 1 on GPUs 0-3, round 2 on GPUs 4-7):

| Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup |
| --- | --- | --- | --- | --- |
| 8192 / 256 | 1 | 12.02 [11.99, 12.04] | 7.09 [7.09, 7.10] | 1.70x |
| 8192 / 256 | 2 | 13.05 [13.04, 13.06] | 7.25 [7.24, 7.26] | 1.80x |
| 8192 / 256 | 4 | 14.68 [14.65, 14.70] | 10.10 [10.08, 10.11] | 1.45x |
| 8192 / 256 | 8 | 16.96 [16.93, 19.52] | 16.95 [16.93, 19.51] | 1.00x |
| 8192 / 256 | 16 | 23.19 [23.17, 27.96] | 23.17 [23.06, 27.95] | 1.00x
|
| 512 / 128 | 1 | 11.61 [11.58, 11.64] | 6.95 [6.95, 6.95] | 1.67x |
| 512 / 128 | 2 | 12.63 [12.61, 12.65] | 6.94 [6.94, 6.95] | 1.82x |
| 512 / 128 | 4 | 14.33 [14.31, 14.34] | 9.70 [9.68, 9.71] | 1.48x |
| 512 / 128 | 8 | 16.74 [16.71, 16.75] | 16.73 [16.68, 16.75] | 1.00x |
| 512 / 128 | 16 | 22.26 [22.20, 22.29] | 22.22 [22.18, 22.23] | 1.00x |

- At c8 and c16, a TP4 step has more than 4 rows, so both arms use
SGLang's path.
- Before the TP4 row cap, the 8-row launches made c8 slower:
  - 8192 / 256: 18.15 ms against 16.95 ms (0.93x).
  - 512 / 128: 17.44 ms against 16.76 ms (0.96x).

### Cookbook shape on 02d42fa

The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This
table uses that shape at TP8 on 02d42fa.

- Each arm ran on a fresh server, in two rounds. The second round
swapped the arm order.
- TPOT is the median per request over both rounds, with a bootstrap 95%
CI. TTFT is the median per request. Throughput is the mean over the
rounds.

| Conc. | TPOT off (ms) | TPOT on (ms) | Speedup | TTFT off / on (ms) |
Output tok/s per GPU off / on |
| --- | --- | --- | --- | --- | --- |
| 1 | 11.24 [11.24, 11.24] | 5.58 [5.58, 5.59] | 2.01x | 541 / 540 |
10.6 / 20.2 |
| 2 | 12.34 [12.33, 12.35] | 5.75 [5.75, 5.76] | 2.15x | 287 / 285 |
19.7 / 40.9 |
| 4 | 13.27 [13.25, 13.28] | 7.28 [7.28, 7.29] | 1.82x | 365 / 353 |
31.3 / 49.7 |
| 8 | 15.80 [15.79, 16.26] | 11.44 [11.42, 11.89] | 1.38x | 3,067 /
3,059 | 49.1 / 62.5 |
| 16 | 18.62 [18.60, 19.51] | 18.60 [18.60, 19.51] | 1.00x | 1,080 /
1,037 | 86.7 / 87.0 |

At 16 requests, a step has more than 8 rows, so both arms use SGLang's
path.

Profile (torch profiler, rank 0, TP8, one request, mean of 10 decode
steps):

- SGLang's path: a decode step runs 1,821 GPU kernels in 11.70 ms of GPU
time.
- Monokernels on: a decode step runs 212 kernels in 5.59 ms. The 75
launches use 4.76 ms (63.5 us each). The other 137 kernels (0.83 ms) are
the embedding, the dense layers 0 to 2, the final norm and the logits.

Startup:

- In a new container, the kernels build in approximately 40 s at TP8 (6
kernels) and 30 s at TP4 (4 kernels).
- A later server in the same container builds them in approximately 2 s,
because the container keeps the compiled code objects
(`~/.cache/comgr`).
- During the build, PyTorch memory increases by 511.8 MiB for each GPU
at TP8 and by 573.1 MiB at TP4.
- The weight copies are 491 MiB at TP8 and 622 MiB at TP4:
  - W_UK / W_UV: 131 or 263 MiB;
  - the router: 225 MiB;
  - the rotary tables: 128 MiB;
  - `weights_proj`: 7 MiB.
- The peer buffers are outside PyTorch.
- In total, the KV pool decreases by 0.7% at TP8 (from 1,533,824 to
1,522,624 tokens). At TP4, it decreases by 3.4% (from 539,584 to 521,024
tokens).

Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e:

- unit (CPU): 88 passed on 02d42fa, with SGLang main 4cca7fcd3e and
2948662026. Two of the cases fail on cbc8f15: the layout and indexer
refusals, and the rotary check over each indexer.
- e2e (TP8): 4 passed on 02d42fa in one run:
  - GSM8K-200 through the launches: 0.995.
  - KL: 0.036, 0.025, 0.018 and 0.021 for batches of 2, 3, 7 and 8 rows.
  - The check mode: all values are within `NUMERICS`.
- An injected wait fault stops all ranks, and the faults are only at
address 0x0.
- The lowest stream cosine of the check mode changes between runs, on
one row. It was 0.9938 and 0.9939 in the earlier runs, and 0.9895 (layer
31, position 30) on 02d42fa. All other checks are 0.994 or more.

The kernel comes from ROCm/aiter#6173 at 56d6390def (Apache-2.0, FlyDSL
Project Contributors). The file headers are kept.
vllm-project/vllm#60831 integrates the MXFP4 variant.

### Review fixes (on top of 02d42fa)

- main is merged in (#6 with its review fixes, and #10).
- New startup refusals:
- unfused DSA top-k (`SGLANG_DSA_FUSE_TOPK=0`). With it, layer 2 hands
layer 3 request-local positions instead of KV slots, and layers 3 to 5
would attend other requests' rows.
  - `--enable-hisparse`, whose logical KV slots the kernels write raw.
  - `--expert-distribution-recorder-mode`.
  - SGLang's benchmark routing overrides.
- a scaled RoPE such as YaRN. SGLang scales the softmax by mscale
squared, while the kernel's scale is fixed.
  - an f32 router gate that bf16 cannot hold exactly.
  - a served context longer than the copied rotary table.
  - a TP4 server whose decode graphs all exceed the 4-row cap.
- The kernel cache key holds the sources digest, as GLM-5.3-Flash's and
DeepSeek-V4.1's do. So an edit to `common/`, or to an aiter helper the
kernel imports, rebuilds it.
- `SGLANG_READS` declares every SGLang name that the spec reads.
- Cleanup:
  - the lineage comments are gone;
  - the stale kernel and runner docstrings match the code;
  - the guard flags size uses `common.guard.FLAG_BYTES`;
  - a few names are clearer, and `gfx950/` has a package `__init__`.
- CODEOWNERS adds `**/*glm53* @jhinpan`, above the Flash line, because
the last match wins.
- The README has the GLM-5.3 row and its own test command.
- Tests on a470273:
- `test/unit`: 89 passed, with SGLang main 4eb606b2ae and
sgl-project/sglang#44601.
  - pre-commit passes.
- Not executed: the TP8 e2e. It needs 8 x MI355X and the zai-org/GLM-5.3
checkpoint.

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

## Review and Merge Process

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

---------

Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii pushed a commit that referenced this pull request Oct 18, 2026
## Motivation

Adds MiniMax-M3 to the AMD monokernel backend: each sparse MoE layer of
a decode or EAGLE3 verify step runs as one launch.

## Modifications

- `sgl_amd_mono/minimax_m3`: the spec and its forward. The embedding,
the 3 dense layers and the final norm run SGLang's modules; each of the
57 sparse MoE layers is one launch. TP4, MXFP4 experts, 1 to 16 rows a
step (EAGLE3: up to 4 requests of 4 tokens).
- One launch runs K1 (input norm, fused qkv / index projection, q / k
norms, RoPE, K / V and index-K writes, index scores) and then block
selection, sparse attention, o_proj, the router, the MXFP4 MoE and both
all-reduces. Every wait goes through `common/guard`.
- Weight copies made at startup and logged: per-channel FP8 qkv / index
and o_proj (the checkpoint keeps them in bf16), the bf16 router gate,
the fp32 routing bias and one bf16 RoPE table, 1.7 GiB a rank.
- `SGLANG_AMD_MONO_CHECK=1` compares every launch with SGLang's layer on
the same input in eager steps; the spec's bounds come from it.
- Tests: `test/unit` (spec, refusals on meta-device layers),
`test/kernels` (row tables, every width compiles), `test/e2e` (EAGLE3
verify server, HiCache decode server, check mode, fault injection).
README.

Needs SGLang's page16 SHUFFLE K / V, sgl-project/sglang#41998 on #41708.
Until both merge, `--monokernel-backend amd` refuses MiniMax-M3 and
names them; their three switches are read with a fallback and move into
`SGLANG_READS` after the merge. SGLang's `sgl-amd-mono` pin needs a bump
for M3 to ship.

## Accuracy Tests

MI355X, TP4, SGLang main a038a0dca0 + sgl-project/sglang#41998
(544bfad891).

`amd/MiniMax-M3-MXFP4`, full 5-shot GSM8K (1,319 questions), chat
completions, 16384-token budget, temperature 0, real EAGLE3
verification, 4 concurrent requests so every verify step runs the
monokernels:

| | Run 1 | Run 2 |
|---|---:|---:|
| Monokernels off | 96.36% | 96.89% |
| Monokernels on | 96.51% | 96.82% |


- `test/e2e`: 9 passed.
- `test/unit`: the MiniMax-M3 tests and contract items pass on SGLang
main and on main + #41998. On SGLang main alone, `--monokernel-backend
amd` refuses MiniMax-M3 at startup, naming #41998 and #41708.

## Speed Tests and Profiling

AgentX on MI355X, TP4, EAGLE3, 3600 seconds per run.

| Concurrency | p90 interactivity (tok/s/user), off → on | Total
throughput (tok/s/GPU), off → on |
|---|---:|---:|
| C1 | 289.2 → 409.8 (+41.7%) | 4,802 → 5,771 (+20.2%) |
| C4 | 246.4 → 330.3 (+34.1%) | 7,711 → 8,601 (+11.5%) |
| C38, HiCache | 50.3 → 54.4 (+8.0%) | 47,040 → 47,720 (+1.4%) |

### Review fixes (on top of 374860b)

- main is merged in (GLM-5.3-Flash #6, the DeepSeek-V4.1 fixes #10 and
GLM-5.3 #7). The README, the registry and CODEOWNERS keep every model.
- The kernel cache key holds the sources digest, as the other models'
do. So an edit to `common/`, to `minimax_m3/gfx950` (including
`layout.py`'s dict-valued stage bases) or to an imported aiter helper
rebuilds the kernels.
- New startup refusals:
- experts whose gate and up rows SGLang interleaves (a SiTU activation,
A16W4). They are marked shuffled and have the separated layout's sizes,
while the kernels read separated rows.
  - `--enable-hisparse`;
  - `--expert-distribution-recorder-mode`;
  - SGLang's benchmark routing overrides.
- `spec.py`'s `SGLANG_READS` lists what `spec.py` reads: the parallel
sizes, the shape key fields, the `"target_verify"` literal, the layer
modules, `forward_native`, `out_cache_loc` and `text_config`. It no
longer lists names that only `weights.py` reads.
- Code that nothing called is removed: `post_attn.py`'s `_f32`,
`index_score.py`'s `q_frag_from`, and two runner fields.
- Cleanup:
  - two refusals say what they need;
- the docs match the code: the MoE is MXFP8 x MXFP4 (a8w4), and the
attention projections run in FP8 (now also in the numerics note and the
README);
  - CODEOWNERS adds `**/*minimax_m3* @zijiecode`.
- Tests:
- Setup: MI350X, TP4, SGLang main + sgl-project/sglang#44601 + #41998
(which carries #41708).
- `test/unit`: 119 passed with and without #41998, on main with every
model.
- `test/e2e`: the decode server's 3 tests passed. On 2c6213a, the final
M3 code (8b37f17 adds only main's other models), mean KL was 0.081
(bound 0.12) and GSM8K-200 was 0.885 (bound 0.87), and the tested
batches ran the monokernels. Before the fixes (f6ec322) they were 0.074
and 0.895, 2 questions apart.
- Not executed here: the EAGLE3 verify, check-mode and fault tests.
Their EAGLE3 draft's rotary fails on this host's sgl_kernel, which lacks
main's f32 cos/sin cache, before the monokernels run.
- Not in this PR: the unreachable context-parallel index path (about 220
lines across five kernel files) and the runner helpers duplicated from
DeepSeek-V4.1.

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

---------

Co-authored-by: Zijie Chen <300606707+zijiecode@users.noreply.github.com>
Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
jhinpan pushed a commit that referenced this pull request Oct 23, 2026
## Motivation

The GLM-5.3-Flash layer launch (#6) follows vLLM in two places where
SGLang's own decode does something else, so on some inputs it picks
other experts, or attends to other keys, than SGLang's path would on the
same input:

- **Router ties.** SGLang's decode under Quark with AITER
(`biased_grouped_topk`, the path `biased_grouped_topk_gpu` takes while
`SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK` is off, its default) rounds the
router logits and the correction bias to bf16, and keeps equal keys in
AITER's grouped top-k lane order. The launch kept f32 keys and broke
ties by the lower expert id.
- **Keys of long rows.** SGLang's kpool top-k
(`fast_kpool_topk_transform_fused`) gives a row with more than 512
complete pools the selected pools' 2,048 tokens and then its tail (the 1
to 3 tokens past its last complete pool), up to 2,051 keys; any shorter
row (up to 2,051 tokens) attends to every token. The launch scored rows
from 2,049 tokens on (`pos >= 2048`) and kept 2,048 keys, the tail
replacing the last selected pool's last tokens.

This PR makes the launch select as SGLang does. It was the first commit
of #22 and is split out so that the routing change lands, and can be
bisected or reverted, on its own. **Merge it before #22**, which is
rebased on it and already contains these commits.

## Modifications

- `gfx950/kernel.py`:
- Router: the GEMM's logits and the bias are rounded to bf16 before the
sigmoid and the add, and equal keys go in AITER's order (`tie_rank`, the
lane priority of AITER's one-wave grouped top-k).
- Sparse attention: `row_keys(pos)` gives each row's complete pools,
whether it is scored, and its key count. A scored row's selection is the
512 pools' tokens followed by its tail; the split attention takes a 33rd
task for the tail's keys, and the merge adds it to the second half.
- `gfx950/layout.py`: `MAX_KEYS`, `BASE_SPLITS`, `N_SPLIT = 33`,
`SEL_STRIDE`; the selection mailbox holds `SEL_STRIDE` tokens a row.
- `spec.py`: `requires.env` adds
`SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK: False` next to the two
`SGLANG_SIMULATE_*` overrides, and `SGLANG_READS` declares it. With it
on, ROCm routes through the Triton gate, where the lower expert id wins
ties, so the launch would no longer match SGLang there.
- Tests:
- `check_layer.py --check-routing`: router weights whose keys tie
exactly in bf16 and differ in f32. The launch's top-8, in order, must
equal AITER's `biased_grouped_topk` on the launch's own FFN input, every
draw. The same run counts the rows where lower-id ties or f32 keys would
pick differently (both must be nonzero, so the inputs exercise the
change). The host-side order with the kernel's `tie_rank` is printed as
a diagnostic only.
- `check_layer.py --check-keys`: rows of 4,102, 5,003 and 7,001 tokens
(tails of 2, 3 and 1) and one of 2,051. A scored row's selection must be
512 whole, distinct, complete pools followed by its tail tokens in
order; perturbing the MLA rows of the tokens a 2,048-key selection would
drop for the tail must change the row's output; perturbing an unselected
pool must not. The 2,051-token row attends to every token.
  - `test_glm53_flash_kernels.py` runs both at TP 2, 4 and 8.
- `test_glm53_flash_spec.py`: the spec declares
`SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK: False` in `requires.env`.
- Registers, on main's kernel: at TP4, 8 rows the KDA build spills 12
VGPRs (main 10) and the sparse-attention build 13 (main 17); TP2 is
unchanged.

## Accuracy Tests

All results in this section were measured on this PR alone (its head
over main's kernel), on `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`,
SGLang ab580b2e06, `amd/GLM-5.3-Flash-Quark-MXFP4@b5688f25`. Reruns by
@jhinpan on 21309bc (this head) unless marked.

- **Kernel checks at TP4 and TP2:** `--writes` (FFN, KDA, sparse
attention), `--compare-shuffled`, `--check-routing`, `--check-keys` all
pass.
- Routing: top-8 == AITER `biased_grouped_topk` in 18/18 rank-draws; the
negative controls (lower-id ties, f32 keys) differ in 12/12 rows each.
- Keys: every scored row is 512 whole distinct pools followed by its
tail; the tokens a 2,048-key selection drops change the output, and an
unselected pool does not.
- **Both checks fail on main** (bug regressions): main's top-8 differs
from AITER's in every draw (TP8: 24/24 rank-draws), and its scored rows
fail the whole-pool and tail checks.
- **Same code as reviewed on #22:** outputs and every state and cache
write are bitwise equal to `a524ebf`.
- **Env gate:** a server started with
`SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK=1` is refused at startup:
`--monokernel-backend amd cannot run this setup:
SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK True is not False`.
- **Unit tests:** 119 passed; the 2 failures are deepseek_v41's contract
cases, as on main (they need an SGLang newer than ab580b2e06 since #17).
- **E2E, `test_glm53_flash_server.py` at TP4:** 4/4 pass. GSM8K past the
largest width 0.985; KL decode against prefill 0.005 to 0.014; check
mode 779 checks, 0 violations (minimum cosine 0.979 sparse-attention
streams, 0.992 KDA streams, 0.9993 MLA KV rows, 0.999999 KDA state); the
injected fault stops the server with no token returned.
- **TP8:** `test_glm53_flash_kernels.py` and `test_guard_hip.py` on 8
GPUs (TP 2, 4 and 8): 8 passed. Check mode at TP8: 779 checks, 0
violations (minimum cosine 0.989 sparse-attention streams, 0.990 KDA
streams, 0.9992 MLA KV rows, 0.999998 KDA state).
- **Check mode at TP2** (ours, this PR's kernel over main): 779 checks,
0 violations; sparse attention past 2,048 tokens exercised (minimum
cosine 0.993 sparse-attention streams, 0.992 KDA streams, 0.9991 MLA KV
rows, 0.999999 KDA state).
- **GSM8K** (sgl-eval 0.1.3; 1,319 questions, thinking, T 1.0, top-p
0.95, seed 42, 8 at a time, TP4): 97.04% with this PR, 96.97% off; 10
questions right only with this PR and 9 only with off (exact McNemar p =
1.0).

## Speed Tests and Profiling

On main's kernel this change costs a little time. Measured by @jhinpan
on 21309bc against main, TP4:

- Serving TPOT, 2048 / 256, fresh servers, two rounds each: +0.6% at one
request (4.37 → 4.40 ms), +0.7% at 2, +1.6% at 4 (5.03 → 5.11 ms), +1.1%
at 8, and −0.2% at 16 (SGLang's path).
- Kernel-only, model order, per launch: +0.7 / +0.8 / +1.5% at 2 / 4 / 8
rows with the sparse-attention rows at position 2,047 (no tail), and
+1.1 / +1.9 / +1.4% at position 8,190 (a tail).
- A single sparse-attention launch at position 8,190: +2.5 / +5.0 /
+3.9%.
- At position 2,049 the launch is 15 / 10 / 9% faster, because rows of
2,049 to 2,051 tokens are no longer scored.

The cost comes from the 33rd split task for the tail's keys: it does
work on 3 of every 4 positions past 2,051 tokens, and even idle it adds
a task to every (row, head group) of the split stage and moves the task
placement of the stages after it, which is why the step costs time
without a tail too. The selection also writes the tail after the 512
pools, the merge adds the tail's partial, and the router computes tie
ranks for its keys. #22 replaces this kernel, so these numbers describe
main only between the two merges.

## 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
(main ab580b2e06).
- [x] Accuracy and speed results are attached when kernel or forward
code changes.

---------

Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan pushed a commit that referenced this pull request Oct 23, 2026
<!-- Branch fix/glm53-flash-index-check, one commit on main (992a3b2).
-->

## Motivation

Under `SGLANG_AMD_MONO_CHECK`, the check compares a sparse-attention
layer's completed kpool index pages against SGLang's own. Today it
compares the FP8 values and the per-pool power-of-two scales separately
(`dsa.index.cosine` 0.999, `dsa.index_scale.rel` 0.01).

The indexer reads neither the values nor the scales on their own. It
reads `FP8 value × pool scale`. A pool whose amax sits at a power of two
can legitimately take the next scale in the launch and the previous one
in SGLang, with its values requantized. The keys are then the same to
FP8 precision, but the values and the scale each differ.

This happened once in #23's TP2 serving check: layer 15, a pool
completed at position 31, 3 rows. The check read `dsa.index.cosine`
0.986 and `dsa.index_scale.rel` 0.10, two violations on a correct
launch. On the dequantized keys the same case reads 0.999996.

@jhinpan agreed in the review of #23 that comparing dequantized keys is
the right invariant, and asked for it as its own PR because it changes
#6's check for every sparse-attention layer.

## Modifications

- `check.py`: the index-cache output is now compared on its dequantized
keys (each tiled FP8 value times its pool's f32 scale) against
`dsa.index.cosine`.
- `spec.py`: `dsa.index_scale.rel` goes from `NUMERICS`.
- `test/unit/test_glm53_flash_check.py` (CPU) covers three cases:
  - the tiled page decode (16-pool × 16-byte tiles);
  - a pool requantized at the next scale (passes);
- a pool's scale doubled with its values kept (fails). This negative
control stays: for pools of similar norm, one doubled scale drops the
page's cosine to about 0.993, against the 0.999 bound.

No launch, spec forward or weight changes. Only the check's numerics
change.

## Accuracy Tests

No serving numerics change (the check runs only under
`SGLANG_AMD_MONO_CHECK=1`).

- Check mode at TP4 on this branch: 779 checks, 0 violations; the
minimum `dsa.index.cosine` (dequantized) is 0.999984. Unit tests pass,
apart from the two DeepSeek-V4.1 contract cases that main's #17 needs a
newer SGLang for.
- The same check on #23's head: 874 checks, 0 violations; the minimum
`dsa.index.cosine` is 0.999983 (TP4) and 0.999981 (TP2, including the
layer-15 case).

## Speed Tests and Profiling

None: check mode only.

## 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
(SGLang main ab580b2e06).
- [x] Accuracy and speed results are attached when kernel or forward
code changes (no kernel or forward change).

Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan pushed a commit that referenced this pull request Oct 23, 2026
## Motivation

The GLM-5.3-Flash monokernels (#6) run each MoE decoder layer 3 to 43 as
one persistent launch. At TP4 with one request, a decode step takes 4.40
ms (8K/256), of which the 41 launches take about 3.44 ms. This PR makes
those launches about 1.6 to 1.7 times faster without changing what they
cover or how SGLang drives them.

Result on MI355X, 2048 in / 256 out, median TPOT, SGLang's path → main's
monokernels → this PR:

- TP4, one request: 7.63 → 4.34 → 2.87 ms (1.52x over main, 2.66x over
SGLang's path).
- TP4, 8 requests: 9.33 → 7.08 → 4.59 ms (1.54x over main).
- TP2, one request: 9.22 → 5.62 → 3.87 ms (1.45x over main).
- TP8, one request: 7.21 → 3.73 → 2.62 ms (1.43x over main, 2.75x over
SGLang's path).
- Past the largest width all arms run SGLang's path and match within
0.3%.
- Kernel-only, the 41 launches in model order: 1.6 to 1.7 times faster
at TP2 and TP4, 1.46 to 1.59 times at TP8.

Accuracy does not decrease: GSM8K (1,319, thinking) 97.12% off, 97.19%
on (11 vs 12 discordant); AIME25 (30) 83.3% off, 86.7% on.

**On top of #27** (`[Fix] glm53_flash: Route ties and long rows' keys as
SGLang does`, merged), which aligns the router's ties and the keys of
rows past 2,048 tokens with SGLang's own decode. This PR's diff is the
kernel and its tests; the measurements below include #27.

Scope (unchanged from #6):

- GPU: MI355X (gfx950), TP 2, 4 or 8.
- Steps: decode steps of 1 to 8 rows (1 to 4 at TP2).
- Layers: 3 to 43 (41 launches: 30 KDA, 11 sparse attention).
- Spec, weights, `check.py`, `NUMERICS`, `SGLANG_READS`, `common/`:
unchanged (beyond #27).
- SGLang: no change; tested on main ab580b2e06.

## Modifications

### 1. Faster decode-layer launches (`c6ee629`, with the review's
`e27297d`, `ebe1443`, `d6192f4`)

The layer kernel is replaced with a faster one, on the same interfaces:
`FlashBuild`, `runner.py`, the 64-token page table as a launch argument,
`common/guard.py`, `common/peer.py`, `common/ops.py`, FP8 projections
read in place. `kernel.py` and `layout.py` change; nothing else does.

The changes, by kind (the commit message lists them):

- **Placement.** Every CU gets a tile in the projection stages; tail
tasks (the tail split; the beta, f_a and g_a rows) run first; a stage's
tasks stay off CTAs whose earlier role ends late.
- **Prefetch.** A task issues its weights, scales, states and table
entries before its wait, sized to the registers the wait leaves free,
and the next task's before the current one reduces. The per-step words
(epoch, peer, the argument and row tables) load at launch start with the
first stages' wait-free loads.
- **Hand-offs.** Broadcast rows carry a ready flag polled from one lane;
many-producer broadcasts are sampled from one wave; releases drain
device-scope stores instead of fencing; inputs are polled in the order
their producers finish.
- **Fusion and reuse.** The gate publishes the layer input row before
its Sinkhorn (16 lanes, DPP) and quantizes the FFN row to MXFP8 for the
experts. An up/gate task covers 32 intermediate values for half the rows
and publishes mid as MXFP8, so down no longer quantizes it. The shared
expert runs before the router scores. A split task covers 16 heads, so
each key row is gathered once.
- **Registers.** Late stages recompute thread indices; the guard reads
its deadline and its stop words (`guard.ERROR_WORDS`) wave-uniform. No
build spills VGPRs at TP 2, 4 or 8, any width, with FP8 preshuffled
(what serving uses) or row-major. (Before the review round the row-major
sparse-attention build at TP2, 4 rows spilled one VGPR; removing the
code warming freed it.) Main spills at TP2, 4 rows (KDA 4; sparse
attention 4 preshuffled, 8 row-major) and at TP4, 8 rows (KDA 10, sparse
attention 17).

The guard, its error words, `SGLANG_AMD_MONO_FAULT_INJECT` (sequential
sites, armed by steps), the kernel symbols (`kernel_symbol`, one per
build) and the two output sets by layer parity are as on main.

**Numerics.** This commit also changes numerics slightly, through the
rewritten gate (Sinkhorn on 16 lanes with DPP, the seam's mix partials
summed as per-stream chains: another fp32 summation order). Measured on
review, one launch of each kind on the same weights and inputs, before
and after this commit (TP4, 4 and 8 rows): the state and cache writes
are bitwise equal; the FFN-only launch differs in post / comb by up to
4e-7 absolute; through the attention-side mixes the difference grows to
about 6e-5, and 24 to 30% of the KDA launch's bf16 `x` values differ
(cosine 0.999991). The per-layer checks against SGLang below are
unchanged within noise.

### 2. Tests and tools (`a749c35`, `e102e1c`)

`check_layer.py --replays N --layers K`: a step of K launches of one
kind on one runtime (each layer its own state or caches, inputs
chained), captured after an eager step; N replays from the same state
give the eager outputs and writes bitwise, with clean error words and
every rank the same output. `test_glm53_flash_kernels.py` runs it for
KDA and sparse attention at every width each TP size builds.

`test/kernels/glm53_flash/time_layers.py`: the kernel-only timing driver
behind the table below (one graph of a whole step, each layer its own
random weights, state and caches; model order or one kind; any width and
TP size). The README's testing section gives the command.

## Accuracy Tests

Setup: MI355X, TP4, `amd/GLM-5.3-Flash-Quark-MXFP4` at b5688f25,
sgl-eval 0.1.3 (the version SGLang pins), the e2e test's server flags
(`test/e2e/glm53_flash_server.py`). Each arm a fresh server on the same
GPUs, 8 or fewer requests at a time (every decode step of the "on" arm
runs the launches).

| Eval | Settings | Off | On (this PR) | Paired |
| --- | --- | --- | --- | --- |
| GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.12% | 97.19%
| 11 vs 12 discordant; one-sided McNemar p = 0.66 (on worse), two-sided
1.00 |
| AIME25 (30 x 1) | thinking, T 0.6, top-p 0.95, 65,536 tokens | 83.3% |
86.7% | 1 vs 2 discordant; 5 / 4 samples at the token limit |
| RULER2 NIAH at 128K (4 subtasks x 2) | thinking off in the request,
greedy, 4,096 tokens | 96.9% | 96.9% | identical per subtask (mk basic /
easy 1.0, mv basic 1.0, mv easy 0.875) |

- RULER2: #6 used 256 output tokens at 1M. GLM-5.3-Flash reasons even
with thinking off in the request (sgl-eval's `--no-thinking` sets
`thinking`; the template reads `enable_thinking`, and the model still
writes reasoning with both false), so at 256 tokens the multi-value
subtasks hit the token limit in both arms (both 50.0%, identical). The
table uses 4,096 tokens; 128K keeps the run short.
- GSM8K: one-sided McNemar is the no-regression test (fails only when
"on" is worse, p < 0.025).

Per layer, against SGLang's own layer on the same input and caches
(`SGLANG_AMD_MONO_CHECK=1`, eager; batches of 1, 3 and 8 rows, 4 at TP2;
positions up to about 2.6K): no value outside its bound.

| Output | TP2 | TP4 | TP8 |
| --- | --- | --- | --- |
| written streams, cosine: min / p05 / median, this PR | 0.970 / 0.994 /
0.9986 | 0.983 / 0.994 / 0.9986 | 0.991 / 0.994 / 0.9987 |
| the same, main, same requests | 0.988 / 0.994 / 0.9987 | 0.992 / 0.994
/ 0.9987 | 0.985 / 0.994 / 0.9987 |
| MLA KV rows, cosine (min), this PR / main | 0.9991 / 0.9992 | 0.9987 /
0.9992 | 0.9991 / 0.9991 |
| KDA recurrent state, cosine (min) | 0.999999 | 0.999999 | 0.999998 |
| checks, violations | 779, 0 | 779, 0 | 779, 0 |

The minimum is one row of one layer (TP2: layer 23, sparse attention; at
TP8 main's is the lower); the 5th percentile and the median match main's
at every TP size. The streams differ from SGLang's mostly through the
MoE (MXFP8 expert inputs against MXFP4, as on main).

## Speed Tests and Profiling

Setup:

- Hardware: one node with 8 x MI355X.
- Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`
(`sha256:e4498f5d595f`, #6's image): ROCm 10.0.0, torch 2.11.0, FlyDSL
0.3.4.1, AITER e7d2453f.
- SGLang main ab580b2e06 (contains #44353 and #44601). sgl-amd-mono main
0bce188.
- Server: the e2e test's flags; the arms differ only in
`--monokernel-backend amd` and the plugin tree (main or this PR).
- Method: `sglang.benchmark.serving`, random dataset with fixed lengths,
max(16, 4c) prompts after c warmup requests; two rounds, the second in
reverse arm order; the tables give the mean of the two rounds' median
TPOT (ITL in parentheses).

TP4, 2048 in / 256 out (8192 / 256 in the second table): median TPOT in
ms, ITL in parentheses.

| Conc. | Off | Main | This PR | vs main | vs off |
| --- | --- | --- | --- | --- | --- |
| 1 | 7.63 (7.63) | 4.34 (4.34) | 2.87 (2.87) | 1.52x | 2.66x |
| 2 | 7.96 (7.95) | 4.38 (4.37) | 2.91 (2.90) | 1.51x | 2.74x |
| 4 | 8.36 (8.35) | 5.02 (5.01) | 3.43 (3.42) | 1.47x | 2.44x |
| 8 | 9.33 (9.23) | 7.08 (6.98) | 4.59 (4.51) | 1.54x | 2.03x |
| 16 (SGLang's path in every arm) | 9.80 (9.59) | 9.77 (9.57) | 9.78
(9.58) | 1.00x | 1.00x |

| Conc. | Off | Main | This PR | vs main | vs off |
| --- | --- | --- | --- | --- | --- |
| 1 | 7.65 (7.65) | 4.35 (4.35) | 2.90 (2.90) | 1.50x | 2.64x |
| 2 | 7.98 (7.96) | 4.38 (4.37) | 2.95 (2.93) | 1.49x | 2.71x |
| 4 | 8.36 (8.34) | 5.01 (5.01) | 3.45 (3.44) | 1.45x | 2.42x |
| 8 | 9.30 (9.17) | 7.05 (6.98) | 4.59 (4.53) | 1.54x | 2.03x |
| 16 (SGLang's path) | 10.34 (9.40) | 10.32 (9.37) | 10.34 (9.39) |
1.00x | 1.00x |

TP2, 2048 / 256 (4-row cap: 8 and 16 requests run SGLang's path in every
arm):

| Conc. | Off | Main | This PR | vs main | vs off |
| --- | --- | --- | --- | --- | --- |
| 1 | 9.22 (9.22) | 5.62 (5.62) | 3.87 (3.88) | 1.45x | 2.38x |
| 2 | 9.53 (9.52) | 5.72 (5.71) | 3.95 (3.94) | 1.45x | 2.41x |
| 4 | 10.35 (10.33) | 7.23 (7.21) | 4.94 (4.93) | 1.46x | 2.10x |
| 8 | 11.98 (11.78) | 11.98 (11.78) | 11.99 (11.80) | 1.00x | 1.00x |
| 16 | 13.06 (12.63) | 13.06 (12.62) | 13.09 (12.62) | 1.00x | 1.00x |

TP8, 2048 / 256:

| Conc. | Off | Main | This PR | vs main | vs off |
| --- | --- | --- | --- | --- | --- |
| 1 | 7.21 (7.21) | 3.73 (3.73) | 2.62 (2.62) | 1.43x | 2.75x |
| 2 | 7.42 (7.41) | 3.76 (3.75) | 2.64 (2.64) | 1.42x | 2.81x |
| 4 | 7.87 (7.85) | 4.21 (4.21) | 3.13 (3.13) | 1.35x | 2.51x |
| 8 | 8.36 (8.32) | 5.26 (5.22) | 3.86 (3.84) | 1.36x | 2.17x |
| 16 (SGLang's path) | 8.04 (7.95) | 8.04 (7.95) | 8.04 (7.95) | 1.00x |
1.00x |

Both rounds agree within 0.6% on every "on" point. The TP2 runs used an
earlier build of this PR whose code objects are byte-identical to the
head's (two comments differ).

### Kernel-only

`test/kernels/glm53_flash/time_layers.py`: one HIP graph of a whole step
of launches, each layer its own random weights (so they stream from
HBM), no profiler; µs per launch. Model order is layers 3 to 43 (30 KDA,
11 sparse attention at position 2,047) on two runtimes, as the spec
builds them. The table was measured with the same method before the
review round; after it (warming removed), the committed driver gives
50.1 / 61.8 / 88.7 µs per launch in model order at TP4, 2 / 4 / 8 rows
(table: 50.0 / 61.6 / 88.3).

| TP | Rows | KDA | Sparse attention, pos 640 | Sparse attention, pos
8191 | Model order (pos 2047) |
| --- | --- | --- | --- | --- | --- |
| 2 | 2 | 107.0 → 69.0 | 109.4 → 66.7 | 131.4 → 83.6 | 111.2 → 69.5
(1.60x) |
| 2 | 4 | 135.7 → 91.9 | 155.4 → 93.5 | 186.4 → 109.0 | 144.0 → 93.0
(1.55x) |
| 4 | 2 | 78.4 → 48.1 | 86.8 → 53.2 | 108.7 → 68.0 | 84.3 → 50.0 (1.69x)
|
| 4 | 4 | 92.0 → 59.1 | 109.8 → 66.2 | 129.9 → 81.6 | 99.3 → 61.6
(1.61x) |
| 4 | 8 | 133.5 → 86.5 | 163.4 → 92.0 | 198.1 → 110.4 | 145.3 → 88.3
(1.65x) |
| 8 | 2 | 64.3 → 41.4 | 74.7 → 48.0 | 97.6 → 63.4 | 69.4 → 43.8 (1.59x)
|
| 8 | 4 | 73.3 → 52.7 | 89.7 → 58.1 | 111.3 → 73.7 | 80.0 → 55.0 (1.46x)
|
| 8 | 8 | 94.6 → 68.7 | 120.1 → 78.9 | 144.3 → 94.2 | 105.3 → 72.2
(1.46x) |

Main → this PR, µs per layer, medians of 7 trials of 200 replays, two
passes in alternating order (they agree within 0.3 µs).

- FP8 projections read in place (main's choice, kept) against row-major
copies, the sparse-attention launch: in place is 1.4 to 2.4 µs faster at
TP2, 0.4 to 0.9 µs faster at TP4, and within 0.5 µs either way at TP8.
- TP8 gains less than TP4 at 4 and 8 rows: the kernel was tuned at TP4,
and at TP8 the per-rank work is smaller, so the launch's chain of
cross-CU and cross-GPU hand-offs is a larger share.

### Tests

On the head after the review, rebased onto main after #27 merged
(`e102e1c`). The PR's diff is unchanged by the rebase (same patch as
before it), and all 24 code objects (TP 2, 4 and 8; KDA and sparse
attention; every width; preshuffled and row-major FP8) are
byte-identical to the pre-rebase head's; main's `common/` changes since
are comments only. SGLang ab580b2e06:

- unit (CPU): 121 passed on SGLang main f91b81b82a (main: 120). On
ab580b2e06, 119 passed; the 2 failures are deepseek_v41's contract
cases, which need an SGLang newer than ab580b2e06 since #17, and fail
the same way on main.
- kernels: `check_layer.py` in every mode (`--writes` with a pad row,
`--compare-shuffled`, `--check-routing`, `--check-keys`, `--replays 20
--layers 4` at every width) at TP2 and TP4: all pass, on this PR's head
and on #27 alone. Every build's kernel symbol is distinct
(`glm53_flash_layer_s<rows>_tp<tp>_k<kda>_d<dsa>_q1_h<shuffled>`, TP 2,
4 and 8, KDA and sparse attention, every width).
- `SGLANG_AMD_MONO_CHECK` at TP4: 779 checks, 0 violations (minimum
cosine 0.989 sparse-attention streams, 0.986 KDA streams, 0.9992 MLA KV
rows, 0.999999 KDA state).
- Speed, the head against the previous head (`720c4dc`, with the code
warming), same server flags, two rounds, TP4 2048 / 256, median TPOT:
2.87 → 2.88 ms at one request, 3.43 → 3.45 at 4, 4.59 → 4.61 at 8 (+0.3
to +0.6%; ITL likewise). The speed tables above were measured on the
previous head; with the warming removed, the head is that much slower.
Kernel-only, model order at TP4 (3 alternating passes, µs per launch,
previous head → head): sparse attention at position 640, 49.4 → 49.7 /
61.2 → 61.4 / 88.1 → 88.1 at 2 / 4 / 8 rows; at 8191, 53.5 → 54.0 / 65.4
→ 65.7 / 93.0 → 93.4. Sparse attention alone at 8191: 68.1 → 67.6 / 81.5
→ 80.9 / 110.4 → 110.1.

On the previous head (`720c4dc`, the same kernel with the code warming):
`test_glm53_flash_kernels.py` over TP 2, 4 and 8 (8 GPUs) 5 passed,
`test_guard_hip.py` 4 passed, and `test_glm53_flash_server.py` at TP4 4
passed (GSM8K past the largest width, decode against prefill KL, the
check mode, the injected fault).

### Startup

- Launch buffers: 490.2 MiB at TP2 (main 490), 214.0 MiB at TP4 (main
213.7), 184.0 MiB at TP8 (main 184); weight copies unchanged
(`weights.py` is not touched).
- Build with the FlyDSL cache warm: 1.8 s at TP2, 2.3 to 2.5 s at TP4,
2.6 s at TP8.
- Server ready after 156 to 166 s in every arm (weights from the page
cache).

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

---------

Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan pushed a commit that referenced this pull request Oct 25, 2026
## Motivation

After #22, a routed decode step of GLM-5.3-Flash runs layers 3 to 43 as
launches (41 of them). Everything else in the step is still SGLang's:
the embedding and its all-reduce, the dense layers 0 to 2, layer 44, the
final norm, the LM head and the logits all-gather. Together these take
about 0.35 ms of the 2.86 ms step at TP4 with one request.

This PR runs every layer and the tail as launches. That makes 46
launches per step: layer 0 from the tokens, layers 1 and 2 (dense),
layers 3 to 44, and one tail launch.

Result on MI355X, 2048 in / 256 out, median TPOT, #22 → this PR:

- TP4, one request: #22 2.88 → **2.53 ms** (1.14x).
- TP4, 8 requests: #22 4.61 → **4.25 ms** (1.08x).
- On the heads before #22's review round, TP2 and TP8 at one request
gained 1.12x and 1.17x over #22, and the TP4 ratios were the same as
above.
- Past the largest width, every arm runs SGLang's path and they match
within 0.4%.

Accuracy does not decrease (TP4, same server flags, on vs off): GSM8K
(1,319, thinking) 96.82% off, 97.12% on; AIME25 (30) 86.7% in both arms;
RULER2 NIAH at 128K 87.5% off, 93.75% on.

Scope:

- GPU: MI355X (gfx950), TP 2, 4 or 8.
- Steps: decode steps of 1 to 8 rows (1 to 4 at TP2), as in #22.
- Coverage: all 45 layers, the embedding, the final norm, the LM head
and the logits all-gather. Output assembly and sampling stay SGLang's.
- SGLang: no change; tested on main ab580b2e06.

## Modifications

### 1. Dense, token and tail launches (`kernel.py`, `layout.py`,
`config.py`, `runner.py`)

There are three new builds of the layer launch, each selected by a
`FlashBuild` field that is off by default. With those fields off, the
KDA, sparse-attention and FFN launches compile to the same code objects
as #22, byte for byte, at every TP size and width.

Each build has its own kernel symbol: `glm53_flash_tail_s<rows>_tp<tp>`
for the tail, and `kernel_symbol`'s layer name with `n1` (dense), `e1`
(tokens) and `h1` (AITER-shuffled dense weights) for the others. A
compile-only build of every key the spec and the tests build (TP 2, 4
and 8; every width; preshuffled and row-major FP8) gives 72 distinct
symbols. The new builds use #22's guard as is (`_Guard`, whose stop
words come from `guard.ERROR_WORDS`), and none of #22's removed warming.

- **`dense`** (with `kda`): a dense layer (0 to 2).
  - KDA attention, as the KDA launch runs it.
- The dense SwiGLU MLP on its FP8 gate | up and down weights with 128 x
128 block scales.
- The weights are read in place, row-major or AITER-shuffled
(`dense_shuffled`; the two layouts are tested bitwise equal).
  - bf16 activations, then the MLP's TP reduce.
- **`tokens`** (with `dense`): layer 0, the step's first launch.
- The attention seam looks up this rank's rows of the vocab-parallel
embedding for the step's tokens.
- It pushes them to every rank and sums the ranks' rows in rank order,
as SGLang's embedding all-reduce does (one rank holds each token, so the
sum is the row).
- One CTA writes the step's KDA state indices and sparse-attention row
table for the later launches, in exactly the spec's own format.
  - The width's rows past the step's are pads.
- **`head`**: the tail.
- The last layer's owed mHC post, then the mean of the four streams, the
final RMSNorm, this rank's LM-head rows, and the logits' all-gather into
f32 [rows, vocab].
- The gathered logits use two symmetric slots by launch parity, so a
rank one tail ahead never overwrites logits a slower peer is still
reading.

Off TP4, the dense launch needed two geometry fixes. At TP2 it stages
its intermediate in a larger LDS region. At TP8 it runs the down
projection's 12 K blocks two per wave, with the spare blocks scaled by
0.

### 2. Kernel tests and timing (`check_layer.py`,
`test_glm53_flash_kernels.py`, `time_layers.py`)

`time_layers.py` (#22's kernel-only driver) times the new kinds too:
`--kind dense`, `embed` and `head`, and `--kind full` for the whole
served step (layer 0 from the tokens, the dense layers, layers 3 to 44
and the tail).

The kernel tests all run at TP 2, 4 and 8:

- **Each new kind:** rank agreement, clean error words, and only its own
rows written.
- **Graph replays:** 20 replays of each new kind equal the eager step
bitwise, including 4 tails back to back.
- **Dense layout:** shuffled FP8 equals row-major bitwise.
- **`--check-dense`:** the MLP output against an f32 MLP on the launch's
own FFN input stays within the spec's bound; a negative control must
fail it.
- **`--check-embed`:** the token launch equals the dense launch fed the
gathered rows (outputs and state writes, bitwise), at 2 of 2, 3 of 4 and
5 of 8 rows, and writes the spec's tables.
- **`--check-head`:** the logits match a torch tail and an f32 tail; a
negative control must fail.

### 3. The spec (`spec.py`, `weights.py`, `check.py`)

- **Layer 0** takes the step's tokens. SGLang's embedding is not run,
and the spec's torch row and index tables go away (layer 0's launch
writes them). Layers 1 and 2 are dense launches.
- **Layer 44** is covered.
- After each KDA launch, the spec runs SGLang's decode track-state copy
(`_track_mamba_state_decode`). At the last mamba layer, that copy saves
every KDA layer's state for the radix cache, exactly as SGLang's own
decode does.
  - The refusal of a covered last KDA layer goes.
- **The tail** runs after layer 44.
  - SGLang's final norm is not run.
- The logits processor returns the tail's logits through `_get_logits`,
so the output assembly and sampling stay SGLang's (see Reviewer notes).
  - A routed step with input embeddings runs the model's own forward.
- **Refusals:**
  - a dense MLP that is not block FP8 (e.g. bf16);
- an embedding or LM head that is quantized, padded, has added rows, or
is not sharded over the TP group;
  - an f32 LM head;
- logits gathered other than over the TP group (DP attention, DP LM
head);
  - a logit scale or soft-capping.
- **Hidden states.** When the step captures hidden states (fixed per
decode graph), the logits processor is handed the final norm's rows
(`normed_rows`), which it stores. Otherwise the last layer's x stands
in. The launches stay on either way.
- **Step tables.** The token launch's step-input tables are kept by
content. A graph reads the table its eager warm-up step made, and an
eager step on other tensors takes its own table, so it never rewrites a
captured one.
- **Weights:** `weights.py` reads the dense MLP, the embedding and the
LM head in place. There are no new copies.
- **`SGLANG_READS`** gains the logits processor's `_get_logits` and the
flags above, `vocab_parallel_embedding`'s shard indices,
`_track_mamba_state_decode`, `residual.batch.clear` and
`capture_hidden_mode`.
- **`SGLANG_AMD_MONO_CHECK`:**
  - layer 0 is checked against SGLang's embedding plus layer 0;
- the dense layers are checked, including their MLP against an f32 MLP;
  - layer 44 is checked on its contracted output;
- the tail's logits are checked against SGLang's final norm and logits,
and against an f32 tail.

### 4. Bounds (`NUMERICS`)

Bounds are relative, at 3x the largest value `SGLANG_AMD_MONO_CHECK`
measured in serving at TP2, TP4 and TP8 (the server check's prompts, 1,
3 and 8 or 4 rows, positions past 2048):

| Bound | Value | Serving check, max (TP2 / TP4 / TP8) | What it catches
|
| --- | --- | --- | --- |
| `dense.x_vs_fp32.rel` | 1.2e-2 | 3.1e-3 / 3.0e-3 / 2.9e-3 | Errors in
the dense MLP: against an f32 MLP on the launch's own FFN input.
`check_layer.py --check-dense`: one rank's down scales x2 gives 0.33 to
0.70 |
| `dense.mlp_in.rel` | 2e-2 | 6.2e-3 / 5.3e-3 / – | The launch's FFN
input against the input of SGLang's MLP (a forward pre-hook): attention,
o_proj, the mixes and `post_attention_layernorm`. Control, the dense
launches' `post_attention_layernorm` x1.1: 0.099 to 0.101, failing on
all three dense layers (57 of 57 checks) |
| `dense.streams.rel` | 0.11 | 3.4e-2 / 3.4e-2 / 3.4e-2 | Gross errors
only; see below |
| `dense.conv.rel` | 3e-3 | 9.3e-4 / 9.4e-4 / 8.8e-4 | The conv state
rows the dense layers write |
| `dense.state.rel` | 3.5e-3 | 9.1e-4 / 7.7e-4 / 1.1e-3 | The recurrent
state rows |
| `head.logits.rel` | 1.8e-3 | 2.4e-4 / 3.7e-4 / 3.1e-4 | Against
SGLang's own norm, LM head and gather. `check_layer.py --check-head`:
final norm x1.01 gives 9e-3 |
| `head.logits_vs_fp32.rel` | 0.06 | 9e-5 / 1.3e-4 / 2.2e-4 | Errors in
the tail: each row's distance to an f32 tail in excess of SGLang's |

- **Why `dense.streams.rel` is loose.** On layers 0 to 2 the streams are
the embedding's, which are small. There, the difference from SGLang is
SGLang's own FP8-activation MLP error: 3.6e-2 from an f32 MLP, against
the launch's 3.7e-3. So this bound measures stock's error. The bounds
referenced to f32 (`dense.x_vs_fp32`, `head.logits_vs_fp32`) and to
SGLang's MLP input (`dense.mlp_in`) catch real errors, together with the
state bounds.
- **Negative controls.** The kernel tests' controls above are caught at
TP 2, 4 and 8.
- **Every bound is checked.** `check.REPORTED` lists each kind's
outputs. `_report` rejects any other output, and a unit test keeps every
bound to a listed output.
- **No `head.logits.max_abs`.** A single bf16 ulp of a logit of 8 or
more (0.0625) exceeded a fixed absolute bound. The two relative logits
bounds do not depend on the logits' scale.

## Accuracy Tests

Setup: MI355X, TP4, `amd/GLM-5.3-Flash-Quark-MXFP4` at b5688f25,
sgl-eval 0.1.3, the e2e test's server flags. Each arm is a fresh server
on the same GPUs, at most 8 requests at a time. Measured on this PR's
head before #22's review round, whose launches are unchanged here.
@jhinpan's rerun on `d9cd3c4` gives GSM8K 97.19% on against 97.04% off
(11 vs 9 discordant, McNemar two-sided p = 0.82).

| Eval | Settings | Off | On (this PR) | Paired |
| --- | --- | --- | --- | --- |
| GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 96.82% | 97.12%
| 9 vs 13 discordant; one-sided McNemar p = 0.86 (on worse), two-sided
0.52 |
| AIME25 (30 x 1) | thinking, T 0.6, top-p 0.95, 65,536 tokens | 86.7% |
86.7% | 2 vs 2 discordant |
| RULER2 NIAH at 128K (4 subtasks x 2) | thinking off in the request,
greedy, 4,096 tokens | 87.5% | 93.75% | mk basic / easy 1.0 both, mv
basic 1.0 both, mv easy 0.5 / 0.75 |
| KL decode vs prefill (2, 3, 7, 8 prompts, 512 tokens) | greedy |
0.0256 | 0.0162 | #22: 0.0133 |

All five cases of upstream's e2e test pass at TP4 on this PR's head
(`88e1aed`), with `HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` for the fault
case: GSM8K past the largest width; decode against prefill (KL); the
check test (874 checks, 0 violations; also 0 at TP2); the returned
hidden states (each decode row times the LM head gives the sampled
token: 21 of 21 rows); and the fault stop.

## Speed Tests and Profiling

Setup as #22: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`, SGLang main
ab580b2e06, a fresh server per arm with the e2e test's flags,
`sglang.benchmark.serving` random 2048 / 256, max(16, 4c) prompts after
c warm-up requests. Two rounds, the second in reverse order; the table
gives the mean of the two rounds' median TPOT, with ITL in parentheses.
The rounds agree within 0.2%.

TP4, 2048 / 256, #22 at `e102e1c` against this PR at `d6ecd5e`. Rebased
on main, this PR's code objects are byte-identical to `d6ecd5e`'s, and
main's launches are #22's:

| Conc. | #22 | This PR | vs #22 |
| --- | --- | --- | --- |
| 1 | 2.88 (2.88) | 2.53 (2.53) | 1.14x |
| 2 | 2.92 (2.92) | 2.56 (2.55) | 1.14x |
| 4 | 3.44 (3.43) | 3.09 (3.08) | 1.11x |
| 8 | 4.61 (4.52) | 4.25 (4.16) | 1.08x |
| 16 (SGLang's path) | 9.79 (9.58) | 9.79 (9.56) | 1.00x |

@jhinpan's rerun on the previous heads (#22 `34898f4`, this PR
`d9cd3c4`, TP4) gives the same ratios: 2.92 → 2.57, 2.96 → 2.60, 3.47 →
3.12 and 4.64 → 4.29 ms at 1 / 2 / 4 / 8 requests, and 7.72 / 8.00 /
8.41 / 9.36 ms with the monokernels off.

### Kernel-only (`time_layers.py`)

`torchrun --nproc-per-node 4 test/kernels/glm53_flash/time_layers.py
--kind <kind>`: µs per launch at 2 / 4 / 8 rows (random weights, one HIP
graph a step). TP4 is measured on this PR's head (`88e1aed`); TP2 on
`d9cd3c4`, whose code objects are the same.

| TP | `model` (#22's 41 launches) | `full` (this PR's 46) | `dense` |
`embed` | `head` |
| --- | --- | --- | --- | --- | --- |
| 4 | 50.2 / 61.5 / 88.7 | 49.9 / 61.0 / 86.6 | 39.6 / 45.2 / 55.3 |
40.9 / 47.0 / 59.0 | 58.4 / 61.7 / 66.6 |
| 2 | 69.7 / 93.0 | 70.1 / 92.0 | | | |

- `model` runs on code objects that are byte-identical to #22's.
- **Spills:** no build that the spec or the kernel tests make spills
VGPRs or uses scratch, at any TP size, width or FP8 layout. #22's 32
builds are byte-identical in this PR, and all 72 kernel symbols are
distinct.

### Startup and memory

- **Launch buffers:** 508.3 MiB at TP2 (#22 490.2), 236.7 MiB at TP4
(#22 214.0) and 206.5 MiB at TP8 (#22 184.0). The new runtimes are layer
0, the dense layers and the tail, with the tail's logits slots.
- **KV pool:**
- TP4: 9,010,176 tokens against #22's 9,011,456 (-0.01%) and off's
9,016,320 (-0.07%).
  - TP8: 9,970,688 against 9,973,760 (-0.03%).
  - TP2: 7,022,848 against 7,018,240 (page granularity).
- **Build time:** 3.7 s (TP2), 5.6 s (TP4) and 6.3 s (TP8) with the
FlyDSL cache warm. A tree's first server boot builds in about 24 s.
- **Server ready:** 131 s at TP4 in every arm; 146 to 261 s at TP2 and
161 to 196 s at TP8 across arms (SGLang's radix init wait varies).

## Reviewer notes

- **`_get_logits` is the one SGLang method the spec replaces** (a hook
in SGLang's monokernel interface could remove it later).
- The spec binds the logits processor's `_get_logits` to its own at
build time. The binding is active only inside a routed step and
otherwise calls SGLang's.
- It is needed because the tail launch already produces the gathered f32
logits [rows, vocab], exactly what `_get_logits` returns after its LM
head and all-gather.
- Replacing just that method keeps `LogitsProcessor.forward`'s pruning,
output assembly, logprob paths and SGLang's sampler unchanged. The
alternative would re-implement the processor's forward in the plugin.
- The stored hidden states are the final norm's rows whenever the step
captures them, as SGLang's own forward stores them (see Hidden states).
- Everything else the spec does is its own forward (as in #6), calling
SGLang only through declared reads.
- **The dense and tail bounds** come from the serving check (section 4).
- **Numerics.** Only the new launches change numerics. Commit 1 leaves
#22's KDA, sparse-attention and FFN code objects byte-identical, and the
token launch writes the same tables as the spec's torch ops did
(tested). Commit 4 changes only the bounds.

## Checklist

- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] Accuracy and speed results are attached when kernel or forward
code changes.
- [ ] A new model adds one line to `.github/CODEOWNERS` that names its
owner. (No new model.)

---------

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants