GLM-5.2/GLM-5.3 PR 1: engine: MLA attention, fp8 fused MoE, symmetric all-reduce by vasilevklart · Pull Request #386 · mstar-project/mstar

vasilevklart · 2026-10-09T20:38:56Z

What does this PR do?

GLM-5.2/GLM-5.3 PR 1 of 11: the engine pieces both GLM models share. PRs 2–11 (#387–#392, #397, #398, #405, #406) build on it, each based on the one below, so each diff is that layer only. Merge bottom-up; I rebase and retarget the next after each.

MLA attention resource (from attn: MLA attention resource (split from #257) #321, AttnBackend.FLASHINFER_MLA). Absorbed MLA over a latent KV layout (KVLayout.MLA: ckv 512 plus a rope slot of 64 or 0), on FlashInfer's MLA kernel with one wrapper per graph slot and graph indices per row. The other backends refuse what only this one reads: sm_scale at build, and AttentionStep(context_only=True) (which Kimi K3 support with DSpark speculative decoding #257's DSpark draft plans) at plan; and MLA attention refuses a cache with no query heads. The SHM KV transfer reads an MLA cache by token. Documented in docs/adding_models.rst.
fp8 block-scale fused MoE (from fused_moe: fp8 block-scale experts, optional SwiGLU clamp #270). Triton grouped GEMMs over 128×128 block-scaled fp8 experts, per-token-group activation quant (16 groups per program), an optional SwiGLU clamp. Decode kernels take a fused router (at most 64 tokens); the prefill experts run as one custom op, opaque to dynamo. The grid covers the padding the active experts can produce (_grid_rows), not all E experts. moe_align retries a lost JIT lock, and the conductor builds the op once at load for models that ask, so a TP worker can't fall back to the uncapturable path.
Symmetric-memory all-reduce from the serving config (tp_allreduce: symm_oneshot or symm_multimem, tp_allreduce_max_kb). mstar/distributed/all_reduce.py holds the backends: AllReduce is NCCL and the fallback, SymmMemAllReduce sends the small messages it can take through one persistent buffer per group, and OneShotAllReduce and MultimemAllReduce are its two kernels. AllReduce.new refuses an unknown mode and stays on NCCL unless every rank of the group set symmetric memory up. Which kernel a call takes depends only on what every rank shares (mode, message size, eager or compiled), never on one rank's pointer alignment.
Process-group timeout (dist_timeout_s, which a 750 GB load needs). It covers the load and warmup; serving collectives get torch's default, so a wedged rank doesn't hang its peers for the whole load timeout.
compile_mode per capture config, checked when the config is built: an unknown mode, or one that turns on Inductor's own CUDA graphs, fails warmup instead of serving its buckets eager.
TP async scheduling from the config (tp_async_sched: true, false or a list of nodes), so a shipped yaml can turn it on; it stays off by default and MSTAR_TP_ASYNC_SCHED still overrides it. The GLM yamls in Speed up GLM-5.2 with fp8 kernels and one-graph MTP #388 and Add GLM-5.3-Flash on 8 GPUs, plain and with MTP #390–Extend GLM-5.2 to 256k-token contexts #392 set it.
Streaming chat decodes text across chunk boundaries, so a multi-byte character split over two tokens arrives whole. It keeps worker: device loop-back, rows-only steps, leaner API delivery (round 2) #380's one frame per batch of tokens; the role rides on the first token.
Warmup requests can repeat their text (repeat: N), a long prompt for a path only long prompts take; Extend GLM-5.3-Flash to 1M-token contexts #391 and Extend GLM-5.2 to 256k-token contexts #392's yamls use it.

The new settings are serving-config keys; this PR adds no environment variables. Everything the stack sends over the control wire is plain data, as #407 requires.

Numbers

fp8 MoE at decode sizes (M ≤ 16) uses tiles from a sweep; above that, vLLM's block-fp8 defaults, which took the fp8 MoE at M = 8192 from 2067 to 1425 µs on an H200 (microbenchmark, 2026-10-05).
End to end: the tables in Speed up GLM-5.2 with fp8 kernels and one-graph MTP #388 (GLM-5.2) and Add GLM-5.3-Flash on 8 GPUs, plain and with MTP #390 (GLM-5.3), against vLLM 0.31 measured the same day on 8×H200. They were measured with the per-step prefill token budget this PR no longer has (see Known gaps), so they get measured again when the stack is rebased onto this head.

Known gaps

MoE tiles above 16 rows are vLLM's defaults (hand-set in the prefill kernel), not tuned per shape or GPU. fused_moe: tuned per-device configs, tighter grid, precomputed alignment #377's tuner feeds only the bf16 fused_experts; its grid bound equals _grid_rows, so the two conflict textually, not in behaviour.
No step token budget: GLM's prefill cap (prefill_max_step_tokens, from Speed up GLM-5.2 with fp8 kernels and one-graph MTP #388 on) moves to ChunkingPolicy.max_batch_tokens ([Chunked Prefill #4a] Chunked prefill core #365) once the chunked-prefill stack lands. Until GLM-5.2/GLM-5.3 PR 2: GLM-5.2 with MLA, fp8 MoE and DSA #387–GLM-5.2/GLM-5.3 PR 11: prefix cache and MTP with temperature #406 are rebased onto this head, they still carry the old budget.
api, model: keep each chat message's role through to the prompt #340 (on main) makes a GLM tool result a 400; Serve GLM chat with roles, tool calls, reasoning and usage #397 fixes it.

How was it tested?

CPU tests for every piece are in test/cpu-core.txt (test_fused_moe_*, test_mla_kv_layout, test_flashinfer_mla_gates, test_context_only_attention, test_symm_allreduce_routing, test_dist_timeout, test_cuda_graph_compile_mode, test_scheduler_room, test_conductor_moe_prebuild, …); GPU tests in test/integration/test_fused_moe_fp8.py, test_quant_fp8_gpu.py and test/modular/test_gpu_mla_attention.py.

This head (402f3ff, with main 1493521 merged in, 2026-10-18): ruff check . passes. @test/cpu-core.txt minus test/cpu-mac-exclude.txt on macOS (torch 2.12.1, no triton or einops): 1512 passed and 187 skipped, and its 68 failures and 34 errors are the same tests that fail on main 1493521 on that machine. CommGroup.all_reduce traces into one graph under torch.compile(fullgraph=True), on NCCL and through the symmetric backend's NCCL fallback. GPU gates on this head are pending; the tables below are from earlier heads.

Gates for every layer of the stack against main 3c5f0dd on the same machines (Laude, 2026-10-09):

gate	result
`pytest test/modular` (CPU)	main's 17 failures and no others on all 7 layers (3 collection errors from the runner's older diffusers, 14 that fail on main)
CI replica (torch 2.9.1 CPU, `@test/cpu-core.txt`)	passes on all 7 layers
1×H100: the GLM, MLA, KDA, KV, fused-MoE, fp8, loader and TP test files	main's 6 failures (`test_kv_prefix_stage`) and no others
`ruff check .` (0.16.10)	passes on all 7 layers

Gates for every layer of the stack on main 4d01740 (coriander, 8×H200, 2026-10-12/13):

gate	result
`pytest test/modular test/ci test/rust` (CPU, torch 2.12.1)	main's failures and no others, on all 11 layers
CI replica (`@test/cpu-core.txt` on 4 CPUs, `ruff check .`, dynamo smoke on torch 2.9.1)	passes on all 11 layers
1×H200: the GLM, MLA, KDA, KV, fused-MoE, fp8, loader and TP test files, #407's CUDA tensor transport	pass
TP8 at the top of the stack (#406, long-context MTP yamls)	GSM8K-200 186 (GLM-5.2) and 197 (GLM-5.3); passkey 3/3 to 252k and 503k; an opencode session with tools

Again after the rebase onto main 641162a (2026-10-14):

gate	result
CI replica (`@test/cpu-core.txt` on 4 CPUs, `ruff check .` 0.16 and 0.17, `test/ci`)	passes on all 11 layers
The router, wire-security and tensor-transport tests the replica's env skips, and the API tests, in the full env	main's results and no new failures
1×H200: the GLM, MLA, KDA, KV, fused-MoE, fp8, loader, TP and speculative-sampling test files	503 passed
TP8, the stack on main `4d01740` vs this one, same box and day, 1024 in / 512 out, c1 / c8 / c32 / c64 tok/s	GLM-5.3 MTP (8×H100): 363 / 1436 / 2932 / 4167 before, 363–364 / 1409–1459 / 2942–2965 / 4155–4201 after (two runs). GLM-5.2 long-context MTP (8×H200): 131 / 476 / 862 / 1226 before, 131 / 482 / 895 / 1238 after
TP8 at the top of the stack (#406): GLM-5.3 long-context MTP at a 256k window (8×H100)	GSM8K-200 197; passkey 2/2 and multikey 4/4 at 32k and 131k; agent sessions 77/77 turns; a repeated 131k prompt answers identically in 0.56 s (4.48 s first); tool calls and `finish_reason: length`

Again after the rebase onto main 7fa8b8e (#340, 2026-10-14):

gate	result
CI replica, as above	passes on all 11 layers
The router, wire-security, tensor-transport and API tests in the full env	main's results and no new failures
TP8 at the top of the stack (#406): GLM-5.3 long-context MTP at a 256k window (8×H100)	GSM8K-200 197; passkey 2/2 and multikey 4/4 at 32k and 131k; agent sessions 77/77 turns; a repeated 131k prompt answers identically in 0.47 s (4.50 s first); a tool result fed back gets its answer. Two of three runs, and one of two on the previous head, stopped at #405's prefix divergence (its known gaps)
TP8 at the top of the stack (#406): GLM-5.2 long-context MTP (8×H200)	queued; coriander's 8 GPUs are busy

Checklist

ruff check . passes
Added or updated tests / docs where relevant

attn: MLA on a latent KV layout (#321). fused_moe: fp8 block-scale experts (#270). worker: per-step token budget; api_server: role on the first token. distributed: group timeout, compile_mode, symmetric-memory all-reduce.

Byte-level BPE models emit each token's raw bytes, which can end inside a character; decoding per chunk turned those into U+FFFD.

As FlashInferManager since #243/#280: an eager plan staged into the one wrapper per label while the last step's copy could still be queued, and a captured step whose rows share a cached prefix overran the kv_indices buffer.

As fused_experts_fp8: dynamo traced the launch path inline, which breaks the graph and fails on a symbolic token count under torch 2.9.1.

The runtime's ready order is handle order in Rust, and handles are recycled: a large prefill could be passed over indefinitely. Backlogged rows count too, since the batch builder tops a backlog up with fresh rows.

Its logits kernel holds every token in one tile; past 128 the first launch raised OutOfResources. The models never send more; now it says so.

Chosen per plan, a captured bucket could re-plan into the other kernel while the graph replayed the captured one on stale buffers.

The CUDA op launched a 0-block grid unchecked; the latched error surfaced in the next unrelated kernel.

A rank whose setup failed fell back to NCCL alone, and its peers' first all-reduce waited on it forever. tp_allreduce_max_kb: 0 now turns it off instead of reading as unset (512).

An empty MSTAR_GRAPH_COMPILE_MODE compiled in Inductor's default mode silently; reduce-overhead and max-autotune capture inside the runner's own graph, fail it, and the bucket serves eager.

The decode and prefill kernels index scales, x and shared as dense rows; a strided view was mis-scaled or mis-added silently. The prefill test now covers GLM-5.2's hidden 6144, the fused sum-and-add branch.

The reader sliced NHD's K/V axis first; on an MLA cache that took latent columns, not tokens.

Every deployment JIT-built the fused-MoE align op before spawning its workers: non-MoE models paid the nvcc build, or hung on a stuck build lock. A model now asks with prebuild_fused_moe.

sm_scale on another backend was dropped silently; an MLA cache without num_qo_heads defaulted to the one latent KV head, and the kernel planned 1.

Each new test file with a CPU case, listed in test/cpu-core.txt.

The deployment config's tp_async_sched key (true, false or a node list) turns it on so a shipped yaml can opt in; MSTAR_TP_ASYNC_SCHED still overrides it. Default stays off.

A long prompt for a path only long prompts take, such as a long-context prefill, without pasting it into the yaml.

A quoted yaml true or false arrives as a string and was taken as a node name, so the switch silently did nothing.