GLM-5.2/GLM-5.3 PR 2: GLM-5.2 with MLA, fp8 MoE and DSA by vasilevklart · Pull Request #387 · mstar-project/mstar

vasilevklart · 2026-10-09T20:39:00Z

GLM-5.2-FP8 (GlmMoeDsaForCausalLM, 753B total / 40B active, 78 layers) served with mstar-serve --config configs/glm52_tp8.yaml, TP8 on the shared engine.

Default (glm52_tp8.yaml): contexts up to index_topk (2048 tokens), where dense MLA attends exactly what the DSA indexer would pick. Prefill and decode are captured.
Opt-in (glm52_tp8_longctx.yaml, dsa_long_context: true): contexts up to 8192 through the DSA engine resource (mstar/engine/resources/dsa/), eager only. Its index keys sit in a growing IndexKeyStore, a stopgap until they move onto the MLA cache's pages (DSA: paged index keys (#392) and what it leaves out #880, done in Extend GLM-5.2 to 256k-token contexts #392).
Loading: Add Kimi Delta Attention on the recurrent pool #389's pread loader. TP8 on 8×H200 with a warm page cache, weights load in 323 s instead of 961 s and the server is ready in 635 s instead of 1,148 s, with the same greedy outputs.

Later in the stack: fp8 kernels and one-graph MTP (#388), paged index keys and 256k context (#392).

Shared code

	where
MLA with weight absorption	`mstar/model/components/mla.py` (`MLAAttention`; DSA hooks in through `_select` / `_attend_selected`)
DeepSeek-V3 router, fp8 expert loaders	`mstar/model/components/moe.py` (`NoAuxTcRouter`)
fp8 block-scale checkpoints	`mstar/utils/fp8_block.py`
RoPE	`PositionManager.apply_qk`, plus `pos_ids=` and `k=None` (for the indexer)

Tests

Parity: absorbed MLA against dense (1e-5) and against the resource path; the indexer's batched selection against the per-row loop, and sparse equal to dense bitwise within index_topk; DSA attention against a masked reference with per-request page tables; the fp8 MoE dispatch against bf16.
Real checkpoint, TP8 on 8×H200: GSM8K-200 190/200 (Joseph's run, an earlier head). On the same set vLLM 0.31 scores 186 and SGLang 0.5.21 183 (Speed up GLM-5.2 with fp8 kernels and one-graph MTP #388's table).
CPU Core, CI's command and environment, on 8945dcf76 (this branch with main 73c6c17c4): only the tests main fails on that machine, by name. ruff check . clean.
GPU, 1×H100 (torch 2.12.1+cu129), the GLM, MLA, KV and capture files: 936 passed. Every failure also fails on main, except three compile-recompile checks that hit an import-order error in that machine's torch build; they fail the same on 202900ea2 and pass in coriander's CPU runs.

Known gaps

Long context runs eager, its keys sit outside the paged pool, and once one row of a batch passes index_topk the whole batch attends its picks (fp32, not bitwise the dense path).
A stop token sampled by the prefill costs one extra decode step, since the prefill walk has no loop to stop. Qwen3.5 has the same gap.
A stacked expert counts as loaded once any of its shards lands, and config.json geometry is read only from a local checkpoint path.
Hopper only; the shipped config needs 8 × 141 GB.

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.

GLM-5.2 at k=0: MLA with the DSA indexer, fp8 MoE, weight loader, configs. Replaces #272. Open: the configs' new max_concurrent_requests caps.

A prompt of limit - 1 tokens failed its whole decode batch (the speculative step overran the window), and a long-context prompt past index_topk failed its prefill batch. process_prompt now refuses both; prepare_inputs agrees.

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

Per-token decode turned a character split across byte-level BPE tokens into U+FFFD; the 10-07 GSM8K answers have them at ÷ and √.

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.

The conductor now builds it only for a model that sets prebuild_fused_moe.

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.

An fp8 checkpoint whose quant config went undetected loaded clean with unscaled e4m3 weights; a failed Hub fetch came back as a local path.

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.

A non-numeric max_output_tokens from extra_body reached the conductor's min() and took its message batch down; image or audio parts were decoded and dropped. Both are 400s now.

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.

top_k and min_p reach the sampler; a fused-kernel import failure is logged; MSTAR_GLM52_GRAPH_COMPILE=0 also leaves uncaptured steps uncompiled; the single-rank test config is dummy mode, not a 750 GB pull.

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

from_hf_config read only rope_theta, so a YaRN or dynamic variant ran as plain RoPE, its frequencies and softmax scale wrong.

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

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

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

Built from tokenizer.json alone (transformers 4.x), the tokenizer had no special tokens, so the detokenizer printed <|user|> at the end of replies.

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 prefill walk has no decode_loop: the runtime dropped (and logged) its stop, and decode wrote on past the stop token. The first decode step now stops and emits nothing; the CPU stream harness runs through the runtime.

dghosef · 2026-10-23T21:05:09Z

+        path, else index_topk, where dense MLA is exactly GLM-5.2's DSA."""
+        return self.config.max_seq_len if self.config.dsa_long_context else self.config.index_topk
+
+    def _clamp_token_budget(self, engine_inputs: ModelInputsFromEngine, rid: str, room: int) -> None:


don't need this check anymore since #808 merged

Since #808 and #832 the conductor holds every request to the room its prompt leaves in the KV pool (it refuses a max_tokens past it, and fits a request that names none), and the decode loop's cap ends decode where check_stop does, so no step runs past it. GLM-5.2's own prompt-aware clamp duplicated that, one row short, so a request sized exactly to the window got one token fewer than it asked for. check_stop now reads the request's max_tokens. The KV pool's max_seq_len is now Glm52Model.context_limit(), so the conductor sizes requests against the window the preprocess guard holds contexts to. Before, the two agreed only by convention (index_topk equals max_seq_len in the checkpoint's config); the reduced variant's did not (64 against 512). Claude-Session: https://claude.ai/code/session_01FeseQ9qqtZRk2CqVpUJ6MH

NSagan271

Thanks @vasilevklart, this addresses most of my comment on making DSA a resource. I'm fine landing the growing IndexKeyStore as an eager-only stopgap since it's isolated, and once keys are paged select collapses into the same gather attend already does.

Can you open the follow-up issue the module docstring describes (keys on the MLA cache's pages, under its page table, allocator, eviction and prefix reuse)?

What I want to see before merge: can you move the fast loading from #389 into this PR? I'm seeing it take over half an hour for weight loading when I try to test this PR (which gets the job automatically killed for idleness on coriander), which would make GLM 5.2 very annoying to run between #387 and #389 being merged.

Nits: glm52_tp8_longctx.yaml still points the follow-ups at components/attention.py. And is bf16 storage for the index keys close enough on parity? It would halve the ~2.1 GB/rank.

main 54f29da (15 commits). docs/adding_models.rst: main's sliding_window note on AttentionSpec, then the DSASpec row.

…389) Tensors are read with pread by a small thread pool a bounded window ahead of the consumer instead of faulting through safetensors' mmap, which moves a few pages per fault on ZFS and network filesystems. Per-expert stacked rule lists are matched through a token index instead of trying every rule on every key. Moved here from #389 (caeea87, 1ded60c) so #387 boots without it.

vasilevklart · 2026-10-26T08:29:00Z

@NSagan271 moved #389's pread loader in and fixed the longctx yaml pointer (1c6ac1f)

On bf16 index keys: ~1.2 of 2048 picks change per query (layer 0, 8k tokens) same as bf16 compute already does
index keys follow-up is in #392 (keys on the MLA cache's pages, device top-k, captured decode) + opened an issue for now while there are still gaps #880

vasilevklart · 2026-10-26T23:07:36Z

TP8 boot on 8xH200 with a warm page cache: weights load 961 ->323 s, ready 1148 -> 635 s

#640 moved the paged-pool and prefix-walk refusals out of mstar/worker/engine_manager.py into mstar/validation/resources.py, as generators of problems, so test_glm52_kv_layout.py failed to import after the merge of main.

NSagan271

LGTM, though I haven't run live tests myself (coriander was too contended)

Main since 73c6c17: #883 mstar_rust in the Linux wheels, #901 output spans (BatchedModelOutput.output_spans, NodeSubmodule.spanned_outputs, the worker's crop), #904 DiT seed noise drawn at ingest (on_request_added, NoiseStager deleted, to_device_nonblocking as one memcpy into pinned memory), #903 API-server torch pool from the CPU quota and pooled image encodes, #924 robotics harness, #819 Inductor cache under ~/.cache, #881 CommGroup.agree / JointGroups.agree (agree_across_ranks and flagged_on_any_rank deleted), #724 CUDA IPC KV in one 2D memcpy a page, #387 GLM-5.2 (MLA, fp8 MoE with NoAuxTcRouter, the DSA resource), #921 vjepa2 AC crop, #922 Cosmos3 contiguous positions, #923 dense attention falls back when fa3-fwd has no kernel for the GPU. Conflicts: - mstar/utils/h2d.py, to_device_nonblocking: the stack's signature (data: Tensor | Sequence, device, dtype=None) and DeviceCapability pin_memory gate, with main's tensor path (#904): a CPU tensor is written into a fresh pinned buffer with one memcpy rather than pin_memory()'s copy_ over the intra-op pool. Every CPU tensor now goes through a fresh buffer, so the stack's "an already-pinned source is read in place" caveat (R3.3) no longer applies and is dropped from the docstring. test_accelerator_graph_backend's pinning test checks the pinned torch.empty and its non-blocking copy instead of pin_memory(). - mstar/model/components/diffusion/denoise_loop.py: imports, both (main's copy, the stack's functools). Main code on a superseded path, ported here: - GLM-5.2 (glm52/submodules.py): its own check_stop / check_stop_batched (four arguments, step_walk, its own row-map loop) / _stops -> token_stop = TokenStop(decode_loop, new_token, eos_token_ids, sampler) with a _token_stops override for its two differences (a stop token at prefill ends the reply at the first decode step; max_tokens counts the prefill's token), and a check_stop override for that closing step, whose token postprocess drops. test_glm52_model's helper and test_model_kv_layouts' GLM cap test build the real submodule (over a stub language model) instead of object.__new__, which would leave token_stop unset. - PinnedStager.to_device (#387), a second "host values -> new device tensor" helper beside to_device_nonblocking: deleted with its test (test_pinned_stager_to_device.py); its callers, GLM-5.2's prefill seq_lens and the DSA resource's page tables, call to_device_nonblocking(list, device, dtype=torch.long). test_dsa_resource spies on that instead. The other files both sides touched (engine.py, accelerator_graph_runner.py, kv/manager.py, attn/base.py, ragged/base.py, communication.py, model/base.py, submodule_base.py, entrypoint.py, conductor.py, moe.py, image_io.py, vjepa2_model.py, docs, tests) auto-merged into disjoint hunks: per file the merged delta equals main's own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…per_request_outputs #387 now returns a plain prefill or decode step as a BatchedModelOutput (row outputs, one D2H for check_stop), so the MTP engine-cycle driver and the last-layer-rows test, which indexed the forward's dict by request, raised TypeError on the k=0 baseline. They now go through _testing.per_request_outputs, as the plain engine-cycle test does.

#387 dropped GLM-5.2's own prompt-aware clamp without MTP: the conductor holds prompt plus max_tokens to the window (#808, #832), so the k=0 arm asked for 40 tokens in a 24-row window and tripped the context guard. Both arms now ask for the most the conductor admits (limit - prompt + 1): the k=0 arm fills the window exactly, and under MTP the clamp still leaves two verify blocks.

Since #832 a paged pool needs its labels declared and each request its layout from process_prompt, and since #640 validate_deployment refuses a paged pool with no labels before any worker starts, so GLM-5.3 could not start once main was merged in. As GLM-5.2 does (#387): - kv_labels: the one "main" label, growing on "decode" under decode_loop. - process_prompt returns the prompt's ids span and, per decode iteration, one row, or under MTP the whole k + 1 block a verify step keeps (its rejected drafts stay as holes against kv_rows), so the pool sizes a request at its worst case. - get_loop_caps: decode_loop runs max_output_tokens - 1 iterations; the prefill emits the first token. - A request that names no max_tokens starts from what the decode loop can run (under MTP the cache rows bind before index_topk), which the conductor then fits to its prompt. Asking for index_topk, every such MTP request would need more iterations than the loop has, a 400.

…efusal #387 has TextChatAdapter refuse a chat that asks for a tool call, since no parser there splits one out of the reply. GlmChatAdapter, which does (glm_chat.GlmReplyParser), runs through TextChatAdapter.chat_to_request, so after the merge every GLM chat with tools was a 400 and tool_choice "required" got that message instead of its own. An adapter that parses tool calls now says so (serves_tool_calls); the others still refuse them.

…ake fails no reads - process_prompt returns a ProcessPromptOutput since GLM-5.2 (#387) and GLM-5.3 declare their KV layouts, so the rendered-chat test reads .new_input_tensors. - main's data worker asks the tensor manager for reads that failed (#419, take_failed_reads); the usage test's fake has none.

GLM-5.2 (#387) and GLM-5.3 return a ProcessPromptOutput, with their KV layouts, since the merge of main; the test indexed it as a dict.

vasilevklart added this pull request to stack #393 October 9, 2026 20:39

vasilevklart mentioned this pull request Oct 9, 2026

GLM-5.2/GLM-5.3 PR 1: engine: MLA attention, fp8 fused MoE, symmetric all-reduce #386

Merged

2 tasks done

vasilevklart changed the title ~~glm/02 glm52~~ GLM-5.2/GLM-5.3 PR 2: GLM-5.2 with MLA, fp8 MoE and DSA Oct 9, 2026

This was referenced Oct 9, 2026

Speed up GLM-5.2 with fp8 kernels and one-graph MTP #388

Open

glm52: GLM-5.2 with MLA, fp8 MoE and MTP #272

Closed

vasilevklart marked this pull request as ready for review October 9, 2026 21:54

vasilevklart force-pushed the glm/02-glm52 branch from 3924640 to 14a6094 Compare October 10, 2026 22:20

vasilevklart force-pushed the glm/02-glm52 branch from 14a6094 to 351c178 Compare October 13, 2026 15:16

vasilevklart added 21 commits October 14, 2026 23:01

engine: MLA, fp8 MoE, step budget, symm all-reduce

a74c56a

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.

glm52: GLM-5.2 with MLA and fp8 MoE

e443ade

GLM-5.2 at k=0: MLA with the DSA indexer, fp8 MoE, weight loader, configs. Replaces #272. Open: the configs' new max_concurrent_requests caps.

glm52: refuse a prompt without room, as a 400

d5867ac

A prompt of limit - 1 tokens failed its whole decode batch (the speculative step overran the window), and a long-context prompt past index_topk failed its prefill batch. process_prompt now refuses both; prepare_inputs agrees.

serving_chat: decode text across chunk boundaries

fe290bf

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

glm52: emit each token's raw bytes

324ae35

Per-token decode turned a character split across byte-level BPE tokens into U+FFFD; the 10-07 GSM8K answers have them at ÷ and √.

attn: MLA wrappers per slot, graph indices per row

74eb6a5

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.

glm52: ask for the MoE op pre-build

b3b31bb

The conductor now builds it only for a model that sets prebuild_fused_moe.

fused_moe: prefill experts as one custom op

7978911

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.

glm52: fail loading loudly, not serve garbage

53448e7

An fp8 checkpoint whose quant config went undetected loaded clean with unscaled e4m3 weights; a failed Hub fetch came back as a local path.

micro_scheduler: step budget takes the oldest

2db5679

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.

api_server: TextChatAdapter checks what it passes

e34dfac

A non-numeric max_output_tokens from extra_body reached the conductor's min() and took its message batch down; image or audio parts were decoded and dropped. Both are 400s now.

fused_moe: decode router takes at most 64 tokens

49423a4

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.

glm52: smaller review fixes

c3b4b87

top_k and min_p reach the sampler; a fused-kernel import failure is logged; MSTAR_GLM52_GRAPH_COMPILE=0 also leaves uncaptured steps uncompiled; the single-rank test config is dummy mode, not a 750 GB pull.

attn: FlashMLA only for eager MLA plans

421c683

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

glm52: refuse a scaled RoPE

4b2efd2

from_hf_config read only rope_theta, so a YaRN or dynamic variant ran as plain RoPE, its frequencies and softmax scale wrong.

fused_moe: an empty step skips the align op

94251af

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

test: CI runs the new CPU tests

f97bcda

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

distributed: symm all-reduce on every rank or none

42bd26a

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

glm52: tokenizer registers special tokens

2720d19

Built from tokenizer.json alone (transformers 4.x), the tokenizer had no special tokens, so the detokenizer printed <|user|> at the end of replies.

cuda_graph_runner: refuse empty/cudagraph modes

8cf5b51

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.

glm52: a prefill stop token ends the reply

0dedee4

The prefill walk has no decode_loop: the runtime dropped (and logged) its stop, and decode wrote on past the stop token. The first decode step now stops and emits nothing; the CPU stream harness runs through the runtime.

dghosef approved these changes Oct 23, 2026

View reviewed changes

NSagan271 requested changes Oct 24, 2026

View reviewed changes

vasilevklart added 4 commits October 25, 2026 02:35

Merge main into glm/02-glm52

3488aea

main 54f29da (15 commits). docs/adding_models.rst: main's sliding_window note on AttentionSpec, then the DSASpec row.

configs: glm52_tp8_longctx points its DSA limits at the resource

1c6ac1f

Merge main into glm/02-glm52

202900e

vasilevklart mentioned this pull request Oct 26, 2026

DSA: paged index keys (#392) and what it leaves out #880

Open

vasilevklart requested a review from NSagan271 October 26, 2026 20:25

vasilevklart added 5 commits October 26, 2026 19:12

Merge main into glm/02-glm52

7aa9eed

test: the GLM-5.2 layout test asks main's validation checks

5c0b75d

#640 moved the paged-pool and prefix-walk refusals out of mstar/worker/engine_manager.py into mstar/validation/resources.py, as generators of problems, so test_glm52_kv_layout.py failed to import after the merge of main.

Merge main into glm/02-glm52

798de5c

Merge main into glm/02-glm52

5315bf0

Merge main into glm/02-glm52

1557b04

NSagan271 approved these changes Oct 27, 2026

View reviewed changes

Merge main into glm/02-glm52

8945dcf

vasilevklart merged commit f417b90 into main Oct 28, 2026
6 checks passed

vasilevklart deleted the glm/02-glm52 branch October 28, 2026 01:26

dghosef mentioned this pull request Oct 28, 2026

Performance 1/7: shared engine, conductor and component changes #858

Merged

vasilevklart added a commit that referenced this pull request Oct 28, 2026

test: the reply-prefix test reads process_prompt's tensors

0a40f0b

GLM-5.2 (#387) and GLM-5.3 return a ProcessPromptOutput, with their KV layouts, since the merge of main; the test indexed it as a dict.

This was referenced Oct 29, 2026

FusedExperts: one owner for the routed experts and their dispatch (8/9) #945

Open

Per-request LoRA on MoE experts, in the fused-MoE grouped GEMM (9/9) #946

Open

Navigation Menu

Repository navigation

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 2: GLM-5.2 with MLA, fp8 MoE and DSA#387
vasilevklart merged 98 commits into
mainfrom
glm/02-glm52

vasilevklart commented Oct 9, 2026 •

edited

Loading

Uh oh!

dghosef Oct 23, 2026

Uh oh!

NSagan271 left a comment

Uh oh!

vasilevklart commented Oct 26, 2026

Uh oh!

vasilevklart commented Oct 26, 2026

Uh oh!

NSagan271 left a comment •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

vasilevklart commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Shared code

Tests

Known gaps

Uh oh!

dghosef Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

vasilevklart commented Oct 26, 2026

Uh oh!

vasilevklart commented Oct 26, 2026

Uh oh!

NSagan271 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

vasilevklart commented Oct 9, 2026 •

edited

Loading

NSagan271 left a comment •

edited

Loading