ROCm: gated delta net on AITER and FlyDSL, and capture buckets a resource can refuse by merceod · Pull Request #686 · mstar-project/mstar

merceod · 2026-10-19T11:01:08Z

Stacked on #684 (all-reduce fusion). The refusal check lands in the capture loops that PR reworks (the graph_capture() scope on both graph runners) and runs before a bucket's scope is entered.

GDN kernels. The GDN wrappers now call their kernels through a GDNKernels strategy. FlashInferGDNKernels holds today's calls unchanged. AiterGDNKernels runs on AMD GPUs and works the K-last state pool in place by slot.

eager prefill on AITER's chunk_gated_delta_rule_opt_vk, with its schedule built once per step
captured prefill on AITER's sequential kernel, up to linear_attn.max_captured_prefill_tokens (512 by default, a YAML override)
decode on FlyDSL, or the sequential kernel where FlyDSL's fixed scale or tiling doesn't fit

LinearAttnBackend gains AUTO (now the default, AITER on ROCm and FlashInfer elsewhere) and AITER. Mamba-2 keeps its Triton kernels on every platform.

Capture refusal. AITER's chunked prefill builds its schedule on the host, so it can't be captured. Rather than failing inside capture, both graph runners now ask Resource.capture_refusal(spec) before each capture. A refused bucket runs eagerly. The answer comes from the spec and the resource's config only, so every TP rank refuses the same buckets and skips them together, barrier and capture scope included. The runner logs one warning with the reason, and the bucket stays in dropped_buckets, so MSTAR_REQUIRE_ACCELERATOR_GRAPHS=1 still fails startup for it. GDN on AITER refuses prefill buckets above max_captured_prefill_tokens. CGSlotSpec.max_row_tokens is the longest row a bucket can hold. adding_models.rst gets the LinearAttnSpec rows and the capture bound here; the hook itself is written up in the docs PR (#687).

NVIDIA: AUTO resolves to FlashInfer and nothing refuses a bucket, so the captured set doesn't change.

Tests

What	Result
`ruff check .`	clean
`test_gdn_kernels_dispatch.py` (new), refusal cases in `test_accelerator_graph_capture.py`, plus `test_gdn_`, `test_mamba2_`, `test_recurrent_*`, `test_resource_runner.py`, `test_ar_fusion.py`	pass in the run of the whole integration branch on MI300X below
`test/rocm/test_gdn_aiter.py` on MI300X	14 passed (in the r8 `test/rocm` run below). Kernel parity vs exact fp32: decode 9.6e-4, chunked up to 1.25e-3 max abs error, bitwise repeatable over 5 runs, capture and replay on unseen layouts
Qwen3.5 greedy vs HF, 8 prompts x 64 tokens, MI300X	identical on 7/8 (0.8B), 7/8 (2B), 5/8 (4B), 8/8 (9B). Every divergence starts at an HF bf16 logit tie
`benchmark/correctness` (harness PR) with the norm and attention PRs, MI300X	Qwen3.5 0.8B to 9B pass on text and image prompts (2026-10-20 cut). Image requests on the current cut need #685's `max_segments_per_request` fix for the AITER ragged managers; with it the 0.8B image row passes again on MI300X
captured buckets, Qwen3.5 on MI300X	decode 10/10, prefill captured up to 512 tokens, larger ones refused and run eagerly
`test/modular`, `test/distributed` and `test/sampling_test` on an MI300X node with GPUs hidden, on the integration branch that carries this whole stack, rebased on main	only the 7 environment failures main has there too (rsync, ftfy)
`test/rust` there	359 passed
`test/rocm` on MI300X, same branch	674 passed, plus the fused all-reduce harness at 2, 4 and 8 GPUs
Rebase onto main `54f29da` (2026-10-26, the whole stack)	CPU on a devel node, GPUs hidden: `test/modular` only fails the two `test_dump_stacks` environment cases main fails too (7174 passed before the last three test fixes, targeted rerun on the tip 398 passed), other CPU suites 248 passed, TP lockstep sim 7 passed, `test/rust` 368 passed. One MI300X: `test/rocm` 675 passed, fused-MoE and sampler GPU tests 228 passed, served smoke rows for Qwen3.5 0.8B, Kokoro, Whisper large-v3 and Cosmos3 DROID match the earlier runs

Numbers

FlyDSL decode vs the sequential kernel at 9B: +2.3-3.2% tok/s (benchmark.runner, one MI300X VM slice, indicative).

Prompts over the 512-token bound are where this costs the most. One MI325X of an 8-GPU node with nothing else running, benchmark/amd/long_prompt.py (benchmark.runner under benchmark.worker_phases), 4 concurrent requests:

	exec per step	TTFT p50 / p95	tok/s
0.8B, ~1k-token prompts	13.6 ms	82 / 158 ms	1133
0.8B, short prompts	1.4 ms	12 / 18 ms	1838
9B, ~1k-token prompts	17.7 ms	129 / 184 ms	453
9B, short prompts	1.3 ms	20 / 27 ms	627

event_sync stays small on the long prompts (0.01 ms at 0.8B, 1.45 ms at 9B), so those steps are host-bound, from the 29 prefill buckets that run eagerly. Raising max_captured_prefill_tokens doesn't help: at 1024 and 2048 on one MI300X (indicative, another table running beside it) the bursts lose 5-14% tok/s, 9B's long-prompt TTFT p50 goes from 61 to 137 ms at 2048 and startup grows by 25-30 s, so the default stays. The exclusive-node text tables follow.

Not in this PR

Eager prefill above 512 tokens is host-bound (13.6 ms of exec per step at 0.8B and 17.7 ms at 9B over the ~1k-token window above, against 1.4 / 1.3 ms with short prompts). A capturable chunked prefill would fix it; a higher bound doesn't (A/B above). Chunked prefill ([Chunked Prefill #4a] Chunked prefill core #365) isn't on main yet.
AITER's metadata builder does one pageable H2D per eager prefill step.
FlyDSL compiles each tiling on first use (~5 s), which lands in cold start.

Since r8

Since the r8 cut: AITER's GDN prefill kernels keep no state checkpoints, and main's #877 now feeds prefix snapshots through the kernel interface, so on ROCm the recurrent pool's prefix cache stays shut: the GDN manager tells the pool, the pool gives its snapshot slots back (about 2.3 GiB on 0.8B and 2B, 6 GiB on 4B and 9B that would otherwise sit idle), and the engine keeps the node's caches shut with it. docs/models.rst says so. Qwen3.5 serves on ROCm as it did, without prefix reuse; the sequential kernel's state map is the follow-up path.

Checks on this rebase

The stack now sits on main f7cb5ca (2026-10-29), every branch cut from one integration tree.

Check	Result
CPU on a devel node with the GPUs hidden: ruff, `test/ci`, the TP lockstep sim, `test/rust`, CPU Core	ruff clean, sim 14 passed, `test/rust` 436 passed, CPU Core fails nothing main doesn't (main's own 34: `test_dump_stacks` x2, `loop_output_ownership` x32)
`test/rocm` on one MI325X (a VM slice, torch 2.12.1+rocm7.2)	715 passed, 9 skipped, 0 failed
MoE and sampler tests there	MoE EP 53 passed, sampler 145 passed. The 75 `test_moe` failures are main's GLM-5.2 kernel tests, the same set on the previous head
Served rows there through the correctness runner	all PASS: Qwen3.5 0.8B and 9B text (no divergence away from near ties, 4/4 needles), Whisper large-v3 (WER 1.64%, 99/100 transcripts equal to the reference), Kokoro (WER delta 0), Qwen3-Omni colocated (50/50, spoken WER 0), pi0.5 and Cosmos3 DROID (two instances bit-same)
Qwen3.5 9B TP2 on two GPUs there	PASS (21/32 exact, 11 near ties, 0 unexplained), the fused all-reduce on both ranks

Nothing fails on this tree that the previous head doesn't: the failing-set diff against it is empty for test/rocm, MoE, the sampler and the rows.

For this PR: Qwen3.5 9B holds 38.58 GiB of VRAM at ready on this tree against 44.78 GiB with the snapshot slots kept (the pool logs giving back 6288 MiB, 0.8B 2385 MiB), and the 9B row passes the runner (32/32, 4/4 needles).

Resource.capture_refusal(spec) is asked by both graph runners. A refused bucket is warmed up when captures are sized, so its eager peak counts toward the floor, but it is not captured to size it, so the capture plan never takes it. Every rank refuses the same buckets and skips them at capture, scope and all, as it skips an unplanned one. The runner logs one warning with the reason, and the bucket stays in dropped_buckets (or dropped_shapes), so MSTAR_REQUIRE_ACCELERATOR_GRAPHS still fails startup. CGSlotSpec.max_row_tokens is the longest row a bucket can hold.

LinearAttnBackend gains AUTO (the default) and AITER. AUTO resolves once at build: AITER on an AMD GPU, FlashInfer elsewhere. FlashInfer's calls move unchanged into FlashInferGDNKernels, and its prefill state-dtype table now asks the platform for an NVIDIA capability. On AITER every kernel works the K-last pool in place by slot: eager prefill on chunk_gated_delta_rule_opt_vk with its schedule built once per step, captured prefill buckets on the sequential kernel up to max_captured_prefill_tokens (larger ones are refused and run eagerly), and decode on FlyDSL, or the sequential kernel where FlyDSL's fixed scale or tiling does not fit. A split walk's mixed step (GDNMixedWrapper) takes both halves from the same GDNKernels, through without_qk_l2norm since the caller normalises q and k for both. On NVIDIA that makes the calls it made before; on AITER the decode rows run on FlyDSL beside the prefill kernel. A kernel that can normalise reports it from its kernels' qk_l2norm. test_gdn_mixed_split.py's FlashInfer parity runs on NVIDIA only, and test/rocm/test_gdn_aiter.py checks a mixed step on AITER.

…checkpoints AITER's GDN kernels update the pool in place and keep no state checkpoints, so a snapshot slot would hold whatever it held before. The GDN manager tells the pool, whose prefix cache then stays shut, and the engine keeps the node's caches shut with it, as it does above one rank.

The engine's process-wide 'high' matmul precision runs the oracle's fp32 einsums as XF32 on gfx942, about 1e-2 off fp64, which failed the kernel check there; the Triton kernel itself is within 4e-6 of fp64.

The pool sizes its snapshot slots before the GDN manager can say its kernels write none, so on AMD GPUs the slots sat allocated and idle: about 2.3 GiB on Qwen3.5 0.8B and 2B and 6 GiB on 4B and 9B with the fp32 state and 128 slots. refuse_snapshots now reallocates the blocks without them and hands the memory back before anything is captured. The docs say prefix caching is off there.

* engine: MLA, fp8 MoE, step budget, symm all-reduce 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 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 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 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 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 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 The conductor now builds it only for a model that sets prebuild_fused_moe. * fused_moe: prefill experts as one custom op 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 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 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 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 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 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 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 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 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 Each new test file with a CPU case, listed in test/cpu-core.txt. * distributed: symm all-reduce on every rank or none 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 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 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 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. * fused_moe: check block scales and inputs 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. * glm52: a bad sampling kwarg is a 400 None now means unset. A string or null reached the conductor's ingest and raised there, outside any request; process_prompt reads the same kwargs first, where a bad one is the client's 400. * kv: SHM transfer reads an MLA cache by token The reader sliced NHD's K/V axis first; on an MLA cache that took latent columns, not tokens. * glm52: refuse min_p as a 400 The sampler is built without min-p, so the worker refused it at ingest as a 500, after every rank had the request. * glm52: fp8 kernels, fused AR norm, one-graph MTP Decode kernels and fused all-reduce + add + norm, on in the shipped configs; the MTP decode step as one captured graph; small steps' kv write plans built on the host. * conductor: pre-build the MoE op only when asked 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. * api_server: TextChatAdapter refuses what it drops n, stop, logprobs, logit_bias and the presence and frequency penalties returned 200 and were dropped. Left off they pass; set, they are a 400. * glm52: MTP's prompt limit leaves two verify blocks An MTP step writes k + 1 rows, so a prompt leaves 2(k + 1) under the limit, as glm5_next's does. * attn: refuse what only the MLA backend reads 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. * glm52: no max_tokens runs to the window The default budget was 1024 on every window, so a thinking reply was cut with length; vLLM runs it to what the prompt leaves, as the submodule's cap now does. * glm52: MTP's no-compile beside the env hatch disable_torch_compile became a property (MSTAR_GLM52_GRAPH_COMPILE=0 covers uncaptured steps); MTP's own choice now sits beside it. * test: CI runs the new CPU tests Each new test file with a CPU case, listed in test/cpu-core.txt. * cli: mstar serve glm52 It answered unknown model; the default is the TP8 yaml. * glm52: probe cuBLAS's block-scaled fp8 GEMM _W8A8_OK checked only that the API exists; torch 2.10-2.12 raise off SM90 or with cuBLASLt < 12.9 (cu128 builds). Probed at load, a failure takes the w8a16 kernel. * engine: KDA on the recurrent pool, pread loader recurrent: roll back a refused admit; per-slot and speculative blocks. linear_attn: KDA manager and Triton kernels. engine: raise on a TP follower's divergent refusal. worker: node metadata. loader: pread. * engine: select TP async scheduling from config 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. * pyproject: glm extra Every other model family has one; GLM's runtime came only with all. * glm52: MTP review fixes Refuse a negative mtp_num_draft_tokens; raise on a second next-token layer; fix the pairing log's "default" label and the long-context yaml's bit-identity claim. * kda: a captured bucket's walk picks its kernels A prefill bucket of as many tokens as rows captured one-token spans as a decode, then replayed uneven spans on the decode kernels, each token on another request's slot. * api: warmup_requests can repeat their text A long prompt for a path only long prompts take, such as a long-context prefill, without pasting it into the yaml. * glm52: an oversized top_k or seed is a 400 The sampler keeps them in int32 / int64 rows. Past that its step raised on every rank, failing the whole batch; under TP async the followers then hung. * test: MTP comment no longer claims bit-identity * loader: pread edge cases A slice dim past the tensor wrapped onto another axis instead of raising; the safe_open fallback queued its tensors as 0 bytes, outside the read-ahead; one preadv of 2 GiB or more fails on macOS (EINVAL). * worker: tp_async_sched reads boolean words A quoted yaml true or false arrives as a string and was taken as a node name, so the switch silently did nothing. * glm52: an oversized max_tokens is a 400 Past 2**64 it failed the wire to the conductor as a 500. The bound is int64, what the Rust server's max_tokens holds. * configs: glm52 yamls turn on TP async * recurrent: _release runs before any forward Its callers are a refused admit and a dropped pre-plan. * cuda_graph_runner: empty compile mode env is unset Raising at import failed every model's boot on an empty export. * glm52: an out-of-range sampling value is a 400 A negative temperature or a top_p of 0 returned 200 with garbage: GLM-5.2's eager prefill inverted the logits or sampled token 0. * glm52: MTP ends the reply on a prefill stop The MTP check_stop still named decode_loop on the prefill walk, and a verify step returned before postprocess could drop its tokens. * test: drop an unrelated stop-check test It covers main's batched stop check, not the KDA pool, and passes without this branch. * distributed: load timeout ends after warmup dist_timeout_s held for every serving collective, so a wedged rank hung its peers for the whole load timeout. Serving gets torch's default. * api_server: three more dropped fields are a 400 best_of, response_format and stop_token_ids returned 200 and were dropped, as n and stop were. * glm52: refuse sampling under MTP as a 400 prepare_inputs refused a temperature or penalty under MTP as a 500, after every rank had the request; process_prompt refuses it first, as GLM-5.3's does. * test: the worker tells the runtime its TP nodes Runs Worker.__init__ over stubs and checks the runtime then refuses a transition into the parallel node, which the static check cannot see. * micro_scheduler: drop the step token budget Chunked prefill (#365) brings a per-(node, walk) max_batch_tokens, and the two would fight over the same batch; GLM moves to it once it lands. The round-robin pass over walks with no room stays, and so do its tests. * attn: refuse context_only on the backends that don't read it Every backend but MLA planned it as an ordinary step, so the queries also attended to the tokens the step appends. * attn: drop FlashMLA Off by default and measured no faster than FlashInfer's MLA kernel; it needs 64-token pages, which no GLM config uses, and served eager plans only. * cuda_graph_runner: drop MSTAR_GRAPH_COMPILE_MODE It set the mode of every capture of every model at once. A capture config's compile_mode, declared per (node, walk), is now checked when the config is built, so a bad one fails warmup instead of serving its buckets eager. * distributed: tp_allreduce and dist_timeout_s from the config only MSTAR_TP_ALLREDUCE, MSTAR_TP_SYMM_AR_MAX_KB and MSTAR_DIST_TIMEOUT_S were a second way to set three serving-config keys, and nothing sets them. * distributed: all-reduce backends as classes AllReduce is NCCL; SymmMemAllReduce takes the small messages it can and hands the rest to it; AllReduce.new picks one from tp_allreduce. CommGroup holds one. resolve_dist_timeout moves to distributed/utils.py. * distributed: one multimem kernel on every rank The one-shot kernel ran only where the caller's tensor started 16-byte aligned, which each rank decided for itself, and ranks on different kernels hang the collective. A misaligned tensor gets the sum in an aligned one and a copy. * docs: the MLA attention resource in adding_models The MLA cache layout, the FLASHINFER_MLA backend and AttentionStep's context_only, beside the resources they extend. * distributed: a class per symmetric-memory kernel SymmMemAllReduce keeps the buffer and the routing; OneShotAllReduce and MultimemAllReduce are its two kernels, in place of a mode string checked on every call. AllReduce.new refuses a tp_allreduce it doesn't know, which used to serve on NCCL behind a warning. * glm52: the accelerator graph names from #262, and capture takes #341's planned * test: the GLM-5.2 tests install the CPU FlashInfer per test, not at import * test: GLM-5.2's attention spec takes #603's AUTO backend * pyproject: the glm extra takes main's Mooncake floor * test: the piecewise dummy-row test fakes #262's graph backend and accelerator memory calls, and passes planned * glm52: the accelerator graph names from #262 * test: the GLM-5.2 MTP, dense fp8 and moe_decode tests install the CPU FlashInfer per test * kda: the accelerator graph hook name from #262 * test: #561's hand-built engine has the TP follower set #389's admit reads * test: the node-metadata test calls #362's per-row speculate_node, its worker graph names its walks * distributed: compiled one-shot all-reduce stays on NCCL Inductor gives one_shot_all_reduce's input a persistent symmetric buffer sized by the first call, so BAGEL at TP2 with symm_oneshot failed decode capture at its second batch size. * test: #755's looping worker has the timeout restore warmup ends on * test: rust-transport lists TextChatAdapter's chat as not ported mstar-server does not serve GLM's flattened text chat yet (#798). * test: the GLM-5.2 serve harnesses stop as the runtime does They ran decode on the prefill walk's info, whose check_stop ends no loop since 0dedee4c8, so the loop never stopped. They now check the prefill's stop, then decode on the decode walk, as the capture test does. * glm52: request kwargs through normalize_generation_kwargs The model's own sampling_kwargs duplicated main's check. As main has it, a number sent as a string is a 400, and a top_k past int32 is served as off. * glm52: checkpoints resolve through hf_snapshot.resolve_snapshot_dir The model's copy is gone; a repo id the Hub cannot fetch still raises. * glm52: a request's token budget and prefill stop in its request_state The engine drops the state with the request, so cleanup_request only evicts the DSA k-store. * glm52: positions from a PositionSpec The step's tokens land where the position resource's counter says, not at the KV stream's stored length, as every other model does. The layers still rotate q_pe / k_pe themselves. KVManager.stored_len had no other user. * fp8 block-scale checkpoints in mstar/utils/fp8_block.py Moved from mstar/model/glm52/quantization.py: nothing in it is GLM's. quant_fp8 takes its FP8_DTYPE from there, and the post-load hook walk is in mstar/model/utils.py. * moe: the noaux_tc router and fp8 expert loaders are shared Moved from mstar/model/glm52/components/moe.py into components/moe.py, nothing in them is GLM's. The router keeps main's router contract (a third value, None: it is stateless). * rope: the interleaved fp32 RoPE is shared Moved from mstar/model/glm52/components/rope.py into components/rope.py as InterleavedRotaryEmbedding, same math: one fp32 cos/sin table per forward. * mla: the shared MLAAttention; GLM-5.2's layer adds DSA MLA with weight absorption (projections, the w_kc / w_vc fold, the naive and absorbed paths) moves to components/mla.py. Glm52MLAAttention is that plus its DSA, through two hooks: _select and _attend_selected. * fp8 linears, their kernels and MLA prep are shared The shared MLA builds on them, and nothing in them is GLM's: components/fp8_linear.py, utils/dense_fp8.py, utils/mla_prep.py. The router tests take the shared router's third value. * glm52: stage host tensors through main's PinnedStager mstar/utils/pinned_staging.py duplicated mstar.utils.h2d.PinnedStager with a fresh pinned allocation per call. GLM-5.2's preprocess now uses a PinnedStager like Qwen3-TTS; PinnedStager.to_device returns a new device tensor of the values' shape, for the callers that build one. * dsa: GLM-5.2's DSA as an engine resource mstar/engine/resources/dsa holds what glm52/dsa.py and the attention layer did: the index key store, each step's rows off the KV plan, the indexer's scores and top-k, IndexShare's shared pick and sparse MLA over the picks. The model keeps the indexer projections and declares a DSASpec under dsa_long_context. Still eager only; the pages are staged pinned and the selection no longer syncs on positions.max(). * glm52: drop the unmeasured MoE fused all-reduce switch MSTAR_GLM52_MOE_FUSED_ALLREDUCE came with no measurement, which AGENTS.md invariant 9 asks for. The block reduces the routed and the shared expert's outputs separately, as before the switch; ParallelGatedMLP is main's again. * docs: DSASpec in adding_models.rst What the DSA resource builds, what it depends on and what a step declares, with its limits in this PR: eager only, prompts within topk, keys outside the paged pool. * glm52: moe_fused_allreduce from the serving config only MSTAR_GLM52_MOE_FUSED_ALLREDUCE overrode the config's moe_fused_allreduce. A tuning value belongs in the serving config (AGENTS.md invariant 9), and the shipped TP8 configs set it. * position: refuse rope_table_dtype widths the CUDA kernel can't rotate FlashInfer's cos/sin-cache kernel moves 16 elements a thread; at rotary_dim 8 it returned garbage, different from run to run, instead of failing. Build now refuses a CUDA table whose width is not a multiple of 16. The comment no longer calls the in-kernel angles exact: they come from a fast sin/cos, about a bf16 ulp from a float32 table. * glm52: positions and RoPE from the position resource, softmax_scale from config The forward reads PositionManager.plan's ids (pos_ids("main")) instead of building its own in preprocess. The layers and the DSA indexer rotate with apply_qk, which takes explicit positions (pos_ids=) and rotates q alone (k=None); GLM-5.2's PositionConfig is interleaved over qk_rope_head_dim from an fp32 table, the angles as before. components/rope.py is gone. MLAAttention reads config.softmax_scale. The reduced test config rotates 16 dims (8 before), the narrowest the CUDA table kernel takes. * dsa: a selecting step's page tables go to the device once, in plan attend copied each request's page list to the device in every layer, and the stager's ring waited on a copy eight calls back, so the CPU ran at most eight layers ahead of the GPU. plan() now stages the tables once per step and attend gathers every row's picks in one batch. The decode-only check moves to plan, before any layer runs. The docs row says what the paged follow-up changes. * api_server: TextChatAdapter refuses tool calls A chat with tools reached the model as plain text and came back with a 200, as if the model had declined every tool; _refuse_tool_calls makes it a 400, as for the other chat adapters. A test pins that each turn keeps its role. * glm52: chat turns in their roles, main's intake checks, batched check_stop - process_prompt rendered every message's text as one user turn, the system prompt and earlier replies included; it now builds one message per turn from prompt_parts, as Qwen3.5 does, and a chat the template refuses is a 400. - The sampling-kwarg and min_p checks in process_prompt repeated what the conductor already does at intake since #402 (SamplingReqConfig.validate, and a ValueError before dispatch is a 400); they are gone, and the tests drive the conductor's own methods. - get_max_output_tokens_limit publishes the context window, so a larger max_tokens is the conductor's 400 rather than a silent clamp. - check_stop_batched reads the step's tokens in one tolist; the sampler returns them as row outputs, one clone and one D2H per step. - MSTAR_GLM52_GRAPH_COMPILE becomes model_kwargs.graph_compile. * glm52: every rank keeps rank 0's greedy MTP verdict Each rank took its own argmax of the verify logits and trimmed its KV and seeded its next step by it, which agrees across ranks only while every all-reduce path (NCCL, symmetric memory, the one-shot kernel) sums bitwise alike; a near-tie read apart leaves the ranks' streams apart and the instance stuck at its next collective. The verdict is broadcast from rank 0 before anything reads it, as the sampled verify's draw is. * recurrent: gather and scatter_ a step's rows through its addressing A model that keeps per-request state in the pool (GLM-5.3's MTP seed) read and wrote it with index_select / index_copy_ on block() by slot ids it took from the plan. The pool now does that access, addressed by the step's own addressing, so eager and captured forwards reach their rows the same way. The signature follows #383's; a per-slot block takes no layer_idx. * glm52: GLM-5.2 declares its KV layout Since #832 a paged pool needs its labels declared, and each request's layout returned from process_prompt; without them every GLM-5.2 request was refused. - kv_labels: the one "main" label, growing on "decode" under decode_loop. - process_prompt returns the prompt's ids span and one slot per decode iteration under "kv_layout". - get_loop_caps: max_output_tokens counts the token the prefill emits, so decode_loop runs max_output_tokens - 1 iterations, where check_stop ends it. #808 refuses a cap past the loop's max_iters (context_limit - 1), which is the bound get_max_output_tokens_limit published; #808 removed that hook, so GLM-5.2's override goes too. - The prefill's prompt-aware clamp stays: the conductor holds prompt plus max_tokens to the KV pool's max_seq_len, the clamp to GLM-5.2's own limit, one row short for a step still in flight when the stop lands. - The walk harness passes the model's loop caps to the graph runtime. Claude-Session: https://claude.ai/code/session_01FeseQ9qqtZRk2CqVpUJ6MH * glm52: no token-budget clamp; the KV pool is the context window 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 * loader: pread into private buffers, an index for stacked rules (from #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 (caeea8771, 1ded60cf4) so #387 boots without it. * configs: glm52_tp8_longctx points its DSA limits at the resource * test: the GLM-5.2 layout test asks main's validation checks #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. * test: GLM-5.2 MTP and last-layer-rows tests read the forward through 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. * test: the MTP window test asks what the conductor admits #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. * test: the node-metadata test passes the runtime its loop caps Since #808 PythonGraphRuntime.add_request takes each request's loop_caps; the test came from before it and raised TypeError. * test: the recompile-limit tests set dynamo's flag through torch_flag (pytest 9.2) pytest 9.2's monkeypatch undo deletes an attribute it didn't find in the instance __dict__. On torch._dynamo.config that removed recompile_limit for the rest of the worker, so a GLM-5.2 capture-config test that ran later on it failed with AttributeError. On a submodule built with object.__new__, the delete itself raised in nn.Module.__delattr__ at teardown. The flag now goes through #917's torch_flag fixture, and the throwaway submodule takes its stub by plain assignment. * kda: KDACallable resolves a layer's pool blocks; the torch reference moves to the engine A KDA layer named its blocks of the recurrent pool (pool.block("conv" / "state"), SpecBlocks.of), and the GLM-5.3 package carried the torch reference bundle that indexed them at the plan's slot ids. KDACallable, beside LinearAttnCallable, resolves a layer's conv, state and speculative blocks once and hands them to KDAManager.run, and settles a verify step (set_prefix_len, moved off the manager unchanged). A model holds the callable and never names a block. kda_torch.py takes the reference math and TorchKDAKernels from the model package. The manager defaults to it where the fused kernels don't run, and its prefill no longer reads the slot ids back with .tolist(): each row's slot stays a device index. The new CUDA test fails on the old prefill under sync debugging. * mla: one MLA layer for GLM-5.2, GLM-5.3 and Kimi K3 MLAAttention (components/mla.py) reads three optional config fields, and two constructor arguments, so GLM-5.3's MLA and Kimi K3's are this layer rather than copies of it: - mla_use_nope: no rotation and no position resource (pos_key=None). With qk_rope_head_dim > 0 the decoupled dims are kept unrotated (Kimi K3); with 0 there are none (GLM-5.3): the cache row is the normed latent and a zero rope slot mla_cache_kpe wide, and its query part is kpe_query_fill. The two latent norms and the zero slot run as one Triton kernel on CUDA (mstar/utils/mla_latent_norm.py, GLM-5.3's mla_kernels.py moved here). - mla_use_output_gate: o_proj reads the attention output times sigmoid(g_proj(x)) (Kimi K3). - mla_latent_norm_eps: the q and kv latent norms' eps. Kimi K3's reference builds them as KimiRMSNorm(q_lora_rank) / KimiRMSNorm(kv_lora_rank), eps 1e-6 against rms_norm_eps 1e-5 (modeling_kimi_linear.py:376,384); DeepSeek-V3's does the same. - kv_label: the cache stream every write and attention runs under (label=, passed only when given, so a model without one makes the same calls as before). - cache_layer: the KV layer, for a model that packs its attention layers' planes. The absorbed output goes through _out: eager, bmm writes (T, H, Dv) in place of einsum's (H, T, Dv), so the reshape is a view (GLM-5.3's form); compiled or under autograd, the einsum (GLM-5.2's compiled path unchanged). On H100 the two give the same bits at every shape tried (405 eager and captured, plus the MTP block row's strided rows). w_kc / w_vc are built detached, as GLM-5.3's were: derived inference buffers need no autograd graph. _attend_selected, the absorbed sparse hook, takes the step's cache rows (T, L + kpe) in place of kv_c and k_pe: GLM-5.2's DSA wrote their concatenation, and a latent with no decoupled dims has one row already. The concatenation moves from the hook into the forward; the bits are the same. forward_block_row (the MTP block step) covers a latent with no decoupled dims, and applies the output gate. test_mla_variants: Kimi K3's KimiMLAAttention transcribed (now with its 1e-6 latent eps), naive and absorbed on CPU; GLM-5.3's layout (no decoupled dims, a zero rope slot of 0 or 4) against plain attention from the same weights; the label and packed plane reaching the resources. * kda: the KDA layer in components, beside the resource KimiDeltaAttention (components/kda.py) and ParallelKimiDeltaAttention (components/distributed/kda.py), over a plain KDAConfig, are GLM-5.3's KDA layer moved out of glm5_next/ with generic names, so a model gets the resource and the layer from this PR: - KimiDeltaAttention holds the weights, runs a step through the KDA resource (forward, with the layer's index among the KDA layers or the KDACallable cursor) and keeps the dense reference (prefill / decode_step on caller-held state) as the parity oracle. No distributed import. - ParallelKimiDeltaAttention builds one rank's heads, attaches the checkpoint loaders (whole or pre-sliced tensors; the three q/k/v convs by shard id) and sums o_proj's partials, as ParallelGatedDeltaNet does for GDN. - KDAConfig.full_rank_output_gate: Kimi K3's output gate is one hidden -> heads * head_dim projection (g_proj) rather than the low-rank g_a_proj / g_b_proj pair. The layer then applies the gated norm itself (params().g_b is None): TorchKDAKernels returns the un-normed output for that, and the fused bundle refuses it (kda_triton._merged_row), so K3 on the fused kernels needs a kernel mode. Resource keys default to "kda" and "kda_state" and are constructor arguments, as GDN's. test_kda_layer: on CPU through the resource with the torch bundle, a prefill and decodes equal the dense reference for both output gates; TP2 ranks loaded from one checkpoint sum to the full layer. * test: the piecewise dummy-row test's capture fake takes node_name (#810) Main's #810 passes the node's name to capture_into_graph, which this test replaces with a fake that runs the region once. The fake's signature had no node_name, so every capture raised TypeError and the runner fell back to eager: three cases failed on this branch in CPU Core (any_graphs False, no captured graph for bs=1). The fake takes node_name and ignores it, as it does the dtype. * kv: correct_len's sealed-tail drop goes through main's page accounting (#869) Main's #869 counts each admitted request's pages against an allotment: a page leaves a stream through _release, which takes it out of the stream's lent set and the allotment, and _recount brings the pool's count of unused allotted pages up to date. _drop_sealed_tail, which a rewind over a sealed page runs, still handed the dropped pages straight to the arena, so a page lent to the stream stayed in its lent set, the allotment kept counting it, and assert_pages_conserved failed. Only a deciding pool (world size 1) keeps allotments, so TP8 serving is unaffected. The tail now goes through _release and the request is recounted, and correct_len touches the plan caches, as every change to a stream outside a commit does. test_kv_correct_len.py: a reserved request rewinds over a sealed page it lends; the lent set empties and the pages stay conserved (fails before: "counts pages it no longer holds as lent"). * glm52: the MoE block keeps the all-reduce's result (#957) Glm52SparseMoeBlock reduced its partials with comm_group.all_reduce and dropped the return value at six places: the decode chain's partial, the fused prefill and runner outputs, the bf16 dispatch, the summed partials under moe_fused_allreduce, and the fp8 reference loop. NCCL and symmetric memory reduce in place, so the shipped yamls never saw it, but FlashInfer's all-reduce (tp_allreduce: flashinfer) writes a new tensor and leaves the input as the rank's partial: each such layer then summed nothing. Main's GLM-5.2 block had three of the six; they are fixed here too. test_glm52_moe.py: a collective that returns its sum in a new tensor gives the block twice one rank's partials, as the in-place one gives one, on the bf16 dispatch, the fp8 reference loop, both reduction orders and the three fused paths (stand-in kernels). Each case fails before the fix. * glm52: MTP pairs the trunk's post-final-norm state, without an env var MSTAR_GLM52_MTP_PAIR_POSTNORM chose which trunk stream the MTP plane pairs drafts against: the post-final-norm state by default (the reference's pairing, which accepts better), the pre-norm state at 0. It was an A/B switch whose answer is settled (AGENTS.md invariant 9), so the post-norm pairing is now the only one: the env read, _mtp_pair_rows, the trunk's return_prenorm path, the acceptance log's pairing label and the docs row go. Glm52MTPModule.forward's docstring said the module expects the un-normalized stream; it now says what it is handed. The default changes nothing: the same state was paired before. * glm52: drop prefill_max_step_tokens, which nothing reads GLM-5.2 declared a per-step prefill token budget (max_step_tokens, set "auto" in both TP8 yamls) for a scheduler hook #386 removed: neither the engine nor the worker calls max_step_tokens, here or on main, and chunked prefill (#916) brings max_batch_tokens instead. The yamls' comment that every prefill step replays a graph therefore didn't hold. The config field, its model_kwargs parse, the submodule method, the two yaml keys and the tests that pinned the budget go. * utils: the fp8 GEMM probe and the MLA prep's PDL ask current_platform dense_fp8._probe_fi and mla_prep._pdl read torch.cuda's compute capability, which an AMD GPU reports in CUDA's form (an MI300X is 9.4), so the MLA prep launched with PDL there (AGENTS.md invariant 8). Both now go through current_platform: cuda_capability for the SM90-only GEMM and has_cuda_capability(9) for PDL, neither true off NVIDIA. test_glm52_kernel_platform.py: with the platform answering as ROCm and torch reporting 9.0, neither PDL nor the SM90 GEMM is taken (PDL was). * kv: host write plans only for MTP-sized rows #388 sent every unstaged KV step of at most 512 tokens through the host write plan (a Python loop per token and one H2D) in place of main's packed path, so every model's small prefills changed path, unmeasured on any of them (AGENTS.md invariants 4, 6, 10). The host plan is for MTP verify steps, whose rows write k + 1 tokens, so it now takes a step only when each row writes at most HOST_PLAN_MAX_SPAN (8) tokens, still within the 512-token total; prefills keep the packed path. Decode steps (one token a row) took the host plan before #388 and still do. test_kv_host_plan.py: a step with a 12-token row takes the packed path (it took the host plan before). * glm52: the MTP draft chain and the seed read the post-norm stream Each chained draft fed the next MTP pass the layer output under its shared_head norm, while the trunk pairing feeds the trunk's post-final-norm stream, so the MTP layer's hnorm saw one stream for d1 and another for drafts 2..k, and those were accepted less: GLM-5.2-FP8 at TP8 on 8xH200, GSM8K-200 under k=3, conditional acceptance per position 0.92 / 0.77 / 0.59 against vLLM 0.31's 0.92 / 0.88 / 0.85 on the same questions (vLLM chains the post-norm hidden since vllm#45895; SGLang's NextN returns it too). The chain and the stored seed now carry the shared_head-normed state, the one the head reads, so _mtp_seed takes it alone. The verdict is the same greedy verify against the target's argmax, so outputs don't change; only which drafts are accepted does. Found and fixed on the 8xH200 node (91a9ef4e2); this is that fix without the env var the pairing no longer has. test_glm52_mtp_configs.py::test_the_draft_chain_and_the_seed_carry_the_post_norm_stream fails without it: the second pass reads the first pass's output under the norm. * test: the slot-block test's pool declares the label its fork opens (#869) Main's #869 has a recurrent pool admit only the labels its config declares (RecurrentStateConfig.labels, ("main",) by default) and fail a step that opens another. This test forks "main" into "copy", so its admit failed after the merge. The pool declares ("main", "copy"), as main's own forking tests do (test_admit_rollback.py, test_recurrent_preplan.py). * test: a dropped pre-plan gives back the reservation its slot took (#869) The merge with main's #869 has the pool's _release, which a refused admit and a dropped pre-plan both run, hand back the reservation a slot was taken from, as main's rollback_admit does. Without that, a pre-plan that took a reserved request's first slot and was dropped left the pool keeping one slot fewer for admitted requests. The test pins it at world size 1, where the pool decides admission. rid_table.py lists the pool's admit records (_admit_taken, _preplan_taken, _reserved) among the handle-keyed state that remove_request purges. * test: the fused KDA kernels' parity tests, beside the kernels kda_triton.py's prefill, decode and verify kernels came with this PR, but their parity tests were in #390's GLM-5.3 files, so none of this PR's tests ran a kernel. test_kda_triton.py holds them now, unchanged but for their skip (kda_triton._HAS_TRITON, not GLM-5.3's fused_decode): the prefill against the per-span fp32 reference, the decode against the layer's decode_step, the verify kernel against TorchKDAKernels.run_verify. CUDA only; #390 drops its copies. * engine: a TP follower that refuses a step rank 0 cannot have refused exits _check_follower_refusal raised RuntimeError, which the worker's main loop turns into failed requests while the worker keeps serving; rank 0 then waits in the refused step's collectives until the NCCL watchdog (~600 s). A hang costs more than an error (AGENTS.md invariant 11), so it raises SystemExit("TP ranks diverged: ..."), main's convention for a diverged page trail (micro_scheduler.py), and the worker stops. The docstrings say the guard relies on a cap that follows free state, GLM-5.3's prefill cap. * test: a TP follower admits against rank 0's free recurrent slots GLM-5.3 takes a KDA slot on prefill only into free slots, so a follower that admits a step against fewer free slots than rank 0, or hands out another slot, has diverged. test_tp_async_teardown_order.py drives Worker.run over stubs and pinned the teardown order with a page allocator; this case gives each rank a real RecurrentStatePool (sharded over two ranks, so neither decides alone) and a new request in N+1. Rank 0 freed two requests' slots before it sent N+1; the follower, still running N, tears them down once N is done. At each of its admits its free slots equal rank 0's, in rank 0's order, and N+1's new request gets the slot it got on rank 0. * kda: the decode kernel issues its independent loads first _kda_decode_kernel loaded the pooled state, dt_bias, A_log, beta's input and the norm weight where their values are first used: behind the gate matvecs and the conv for the state, at the very end for the norm weight. At decode sizes the kernel is bound by that chain of load latencies. Every load that depends on no other result now issues at the top, with the gate weights; the arithmetic and its order are unchanged, so the outputs, the state and the conv windows are bitwise the same. One TP8 rank of GLM-5.3 (8 heads of 128), CUDA-graph replay, us a call, before -> after: tokens 1x H100 1x H200 1 7.94 -> 7.04 7.98 -> 7.07 2 8.11 -> 7.05 8.21 -> 7.14 4 8.43 -> 7.35 8.54 -> 7.44 8 8.67 -> 7.67 8.58 -> 7.55 32 13.13 -> 13.01 12.56 -> 12.27 64 27.02 -> 26.54 23.64 -> 23.54 GLM-5.3 runs it in 34 of its 45 layers. Tests: test_glm5next_fused_decode.py (the KDA decode against the reference), test_kda_resource.py, test_kda_layer.py, test_glm5next_kda*.py on 1x H200 and 1x H100: passed. kirill/fuse/bench_kda_decode.py compares a copy of the old kernel bit for bit. * mla: the fused prep can serve a selecting path too components/mla.py ran its fused prep (mla_prep: q_a's norm, then kv_a's norm, both RoPEs and the latent's cache write, in two kernels) only for a forward that selects no rows. A selecting subclass that attends rows the prep already wrote (_attends_written, _attend_written) now takes it too: q_norm first, the selection from its q_c, then kv_write_q_rope and attention over the selected slots. The kernels reproduce the unfused path bit for bit. Nothing here sets _attends_written, so every model keeps its path on this branch; GLM-5.2's selecting attention takes it with #392's paged DSA. * fused_moe: GLM-5.2's decode tiles on the H200 decode.py's H200 table was swept at GLM-5.3's hidden 4096, and at GLM-5.2's experts (hidden 6144, 256 experts of 256 at TP8) it is the slow one at 1-6 and 9-10 tokens. A table swept at that shape on an H200, with every call's weights read from HBM (the per-pair winner per token count up to 8, the grouped kernels from 9), now serves that shape on every GPU but the H100; any other shape keeps the H200 table, and the H100 keeps its own. One layer's experts (routed and shared), CUDA-graph replay, us a call, H200 table -> GLM-5.2's: tokens 1 2 4 5 6 8 9 10 29.85 33.38 53.52 70.87 81.02 90.45 113.13 116.85 -> 24.15 31.12 51.94 63.94 75.37 89.78 102.08 107.80 GLM-5.2-FP8 on 8x H200 (TP8, glm52_tp8.yaml), random 1024 in / 512 out, ignore_eos, greedy, a fresh server per arm: c=1 131.3 -> 138.7 tok/s (TPOT 7.48 -> 7.06 ms), c=8 555.1 -> 556.6. GSM8K-200 183 -> 184. Greedy parity (8 prompts x 128 tokens): 3 of 8 identical, the others first differing at tokens 7 to 71, since the new tiles sum in another order; a second server whose decode MoE path was unchanged matched 8 of 8. Tests: test_fused_moe_decode.py (the table pick by GPU name and shape). * dense_fp8: GLM-5.2's o_proj and dense down on FlashInfer's GEMM from 3 tokens on the H200 On an H200, FlashInfer's SM90 block-scale GEMM beats the w8a16 kernel from 3 tokens at GLM-5.2's o_proj (N 6144, K 2048) and dense-layer down (N 6144, K 1536), where the shape table kept w8a16 up to 32 tokens. Cold weights on one H200, us a call at 4 / 8 / 16 / 32 tokens: o_proj w8a16 10.32 / 11.05 / 9.84 / 10.59 FlashInfer 8.01 / 7.99 / 8.09 / 9.28 dense down w8a16 9.27 / 9.28 / 8.74 / 9.55 FlashInfer 7.51 / 7.56 / 7.63 / 7.97 The new rows are keyed by GPU model, since the H100 and the H200 are both SM90. reserve() resolves a device's limits once, at load and outside capture. They apply only where FlashInfer's GEMM runs, so a GPU without it keeps w8a16 at these sizes instead of falling through to cuBLAS's slower GEMM. Every other GPU and shape keeps its path. A 1x H100 shows the same crossover (o_proj 10.53 -> 9.04 us at 3 tokens, 11.65 -> 9.09 at 8), but GLM-5.2-FP8 does not fit 8x H100, so the H100 gets no row until it is measured end to end. FlashInfer quantizes the activations to e4m3 per 128 columns, as dense_fp8 already does from 33 tokens, so outputs move at 3-32 tokens. GLM-5.2-FP8 on 8x H200 (TP8, configs/glm52_tp8.yaml; random 1024 in / 512 out, greedy, closed loop; a fresh server per arm, in the order base, this, this, base), tok/s: c base this change TPOT p50 (ms), mean of the rounds 1 138.7 / 138.7 138.8 / 138.6 0.0% 7.06 -> 7.06 8 556.6 / 556.8 564.2 / 563.1 +1.3% 13.76 -> 13.50 32 1177.6 / 1270.8 1260.7 / 1280.6 +0.8% 23.42 -> 23.36 (second round) 64 1924.7 / 1963.6 1919.1 / 1968.5 0.0% 29.82 -> 29.65 At c=32 both first-round servers stalled for up to 2.4 s, the base's more, so only the second round is compared there. c=1 and c=64 are controls: at 1 token nothing changes, and from 33 tokens both arms already run FlashInfer's GEMM. With MTP (one server each), c=1 goes 197.3 -> 199.9 tok/s (+1.3%) and c=8 785.2 -> 792.1 (+0.9%). GSM8K-200 is 185 against 183, with 0 errors on both. Tests: test_glm52_kernel_platform.py checks that the H200 takes FlashInfer from 3 tokens on these shapes, and that an H100, an H200 NVL and an H200 without FlashInfer keep w8a16 up to 32 tokens. test_glm52_dense_fp8.py checks reserve()'s limits, and that FlashInfer stays within the rounding bound at 3 and 8 tokens. * mla: drop the options only Kimi K3 uses components/mla.py carried three options for Kimi K3, which no GLM model sets: an output gate (mla_use_output_gate: g_proj and _gate), the latent norms' own eps (mla_latent_norm_eps), and unrotated decoupled dims (mla_use_nope with qk_rope_head_dim > 0). GLM-5.2 rotates its decoupled dims and GLM-5.3 has none; K3 is not in this launch, so the options and their tests move to the K3 stack (refs/prep/k3/00-mla-k3-options). mla_use_nope stays for GLM-5.3 and now means no decoupled dims: with qk_rope_head_dim above 0 it is refused. Every GLM path is unchanged: with the gate off _gate returned its input, the eps fell back to rms_norm_eps, and GLM-5.3's config already refuses mla_use_nope with rope dims. Tests: test_mla_variants.py keeps GLM-5.3's layout and pins the new refusal; K3's reference and gate tests leave with the options. * glm52: one kernels option instead of nine switches; precision follows the checkpoint GLM-5.2 had nine kernel switches (moe_quant_kernel, moe_decode_kernel, moe_router_kernel, moe_prefill_kernel, moe_fused_allreduce, dense_fp8, mla_fused_prep, fused_add_rmsnorm, prefill_last_layer_rows), off by default and on in every serving yaml. Now: - The MoE prefill kernels, which keep the W8A8 runner's tiles and K order (the same bits, test_glm52_moe_decode.py), run on every fused block, with no switch. - The fp8 linears follow the checkpoint: the non-expert linears a block fp8 checkpoint stores in fp8 stay fp8 (dense_fp8 is now a property of the quantization config), and a bf16 config, as the reduced test models are, takes the bf16 path. - What changes numerics is one option, model_kwargs.kernels. "auto" (the default) is what the yamls turned on: the fused fp8 MoE with fused_moe's decode chain and router kernel (their per-GPU tables live in fused_moe/decode.py and dense_fp8.py), the MLA latent prep as two kernels, each all-reduce fused with the next add + RMSNorm and the MoE's partials reduced once, and the last layer's FFN on the sampled rows at prefill. "reference" is the numerics anchor, served eager: the per-expert dequant loop, the fp8 linears dequantized per call, the latent prep op by op, unfused all-reduces and norms, every prefill row through the last layer. "triton" (the fused MoE or an error) is gone: auto serves the reference loop, with a warning, where the fused kernel can't run. The fused latent prep sits under auto, not beside the prefill kernels: its kernels have the ops' bits, and plain decode keeps them (a reduced GLM-5.2 on an H200, round 9 without it against this with it: 3/3 plain cases identical), but 3 of 7 MTP cases differ, in KV pages (two captured cases) and in one eager case's LM-head bits, with every token the same. A config that still names a retired switch is refused at startup, and a fp8 linear the checkpoint doesn't store is named by the loader's missing-weight check. The serving yamls drop the switches and serve what they served: GLM-5.2 bitwise, round 9 with the yamls' switches against this with auto, 10/10 identical; round 9's defaults against this with reference, 10/10. Every other model is unchanged: only GLM-5.2's config names kernels. Tests: the GLM-5.2 kernel, MoE, fused-norm, MLA-prep and dense fp8 files ask for kernels instead of the switches; the integration checkpoint stores the attention, shared-expert and dense-MLP linears in fp8 as the real one does. * glm52: shorter comments; cut tests that pin flags, constants and logs Comments and docstrings keep the why and lose the history. Tests that re-checked a switch, a default or a log line, or copied production logic (the unused host verify helpers, the loader's key inventory) are cut or merged into parametrized cases. * glm52: name the shared fp8 and MLA prep ops under mstar::; drop unused verify helpers The two custom ops live in mstar/utils and serve GLM-5.3 too. mtp_greedy_verify and its host twin had no caller outside tests: verify runs on the device in the step. The models doc names glm52_tp8_mtp.yaml. * engine: a submodule's dynamo recompile limit holds around its own compiled calls GLM-5.2's submodule raised torch._dynamo.config.recompile_limit for the whole thread (the whole process on torch < 2.12), so every other node on a worker compiled under it. A submodule now asks through NodeSubmodule.dynamo_recompile_limit(), and the engine and the graph runner patch the limit around that submodule's compiled callables only. One that doesn't override the hook gets its callables back unwrapped. * kda: shorter comments; cut tests that re-check a call or a constant Comments and docstrings keep the why. The static AST check of the worker's metadata call goes (test_worker_node_metadata runs the real __init__), the recurrent-pool teardown variant goes (the page-allocator test pins the order), and near-duplicate resource and pool tests merge. * kda: the torch kernels on ROCm, where the fused kernels' PTX can't run * recurrent pool: a refused admit or a dropped pre-plan gives back the slots it took A step refused after the pool's admit kept the slots that admit took when the step had been pre-planned: the live admit is a no-op then, and it cleared the pre-plan pass's record of them. A pre-plan the worker discarded (its speculation went unused) kept its slots too. Either way the slots stayed with requests that never ran the step. Rollback now frees the pre-plan pass's slots as well, clear_preplan gives them back, and both free newest first so the free list returns to its old order on every TP rank. * worker: tell the graph runtime which nodes are lockstep-parallel again #333 dropped Worker.__init__'s set_node_metadata call, so the runtime saw no parallel node and its speculation filter could speculate a transition into a TP node, which only the leader may decide. The call is back, for the Python and Rust runtimes alike. * kda: the resource picks its kernels from what the spec and pool declare The fused Triton kernels run where they cover the declaration (CUDA, not ROCm, a validated head dim, a width-4 conv, the low-rank output gate); the torch reference runs otherwise. LinearAttnConfig declares the output gate (`full_rank_output_gate`, Kimi K3's), so a model no longer has to swap the bundle. `set_kernels` goes, and `prefill_capturable` says whether the picked kernels replay a captured prefill. * kda: the decode kernels follow force_prefill and the spans, as GDN's do A captured bucket keyed its kernels on the walk being named "decode", so a model whose decode walk has another name would replay the prefill kernels. The packed prefill bucket that check was for is what `ctx.force_prefill` marks, so a step decodes when no packed capture forces prefill and every row is one token, captured or not. * kda: KDACallable.current_plan runs outside torch.compile, like its other methods * recurrent: scatter_ takes its value in the block's dtype again The cast turned a dtype mismatch that raised into a silent one for every pool user; the slot-block test was the one caller that relied on it. * glm52: the MTP verdict wait has no timeout; an overwritten verdict stops the worker A rank that timed out waiting for a verdict, or found its mailbox reused, raised RuntimeError on its own: its batch failed while the other ranks went on into the next all-reduce and hung. Every read follows the step's launch (a pre-plan waits on the step's commit), so only a hung device or peer delays a verdict, as with any collective. The overwrite check is a broken invariant rather than a request's failure, so it raises SystemExit, as micro_scheduler does for diverged TP ranks. * mla: the fused prep writes and attends under the layer's kv_label _write_fused_and_attend called write_slots() and the attention's run() without the label, so a RoPE layer with kv_label set wrote and read the default stream on the fused path. Both paths take it from one property now. * mla: the fused prep only for interleaved RoPE mla_prep's _rope turns the pairs (2i, 2i + 1), while the op-by-op path follows PositionConfig.interleave, so a checkpoint with "rope_interleave": false under kernels: auto had q_pe and k_pe rotated on the wrong pairs with no error. RopeManager exposes interleave, and bind_resources keeps the ops for NeoX pairing, as it does without a cos/sin table. * glm52: the kernels comment says reference runs the routed experts in a loop The MoE prefill kernels run only on a fused block, and only kernels: auto builds one (Glm52SparseMoeBlock.process_weights_after_loading). * engine: say why the recompile limit stays a submodule's before torch 2.12 dynamo's config.patch is process-wide there, and pyproject still allows torch 2.10; the limit stays per submodule only because a worker runs one compiled call at a time. * docs: KVManager.correct_len and the position resource's * moe: the router takes its top-k kernel in finalize_weights; fewer probes Glm52SparseMoeBlock set NoAuxTcRouter._topk_op from outside. Whether the kernel runs is known only on the loaded device, so the router's own load hook takes it as an argument and imports the op. Glm52LLMSubmodule reads its config, and an attribute __init__ sets, instead of getattr fallbacks, and the GLM-5.2 MoE block reads config.kernels. * kda: the torch reference builds its decay mask one chunk at a time chunk_kda built exp(g_i - g_j) for every chunk of a span at once, [B, H, N, C, C, D] in fp32: ~2 GB for one 8k-token span on a TP8 rank of GLM-5.3, twice that with the product after it, allocated after the KV pool is sized. Per chunk, both reductions that read it are built in one pass; the output and state are bit-identical, and a 4k-token span on 8 heads of 128 peaks at 238 MB instead of 3.2 GB on CPU. * kda: the layer's input projections view the fused weight, one copy process_weights_after_loading concatenated the six input projections and kept the originals beside the copy: 26 MiB per layer at GLM-5.3's TP8 shape, ~0.9 GiB per GPU over 34 layers. The projections now view their rows of the fused weight, so the dense reference and the loaders' names stay as they are. The fused weight is a non-persistent buffer, so .to() moves it, and _apply views the moved copy again. * kda: the layers declare their gates; the resource checks them before it plans LinearAttnConfig declared the gate twice over: gate_lower_bound, which nothing read (the kernels take the layer's KDAParams.lower_bound) and whose comment described a softplus path that doesn't exist, and full_rank_output_gate, which had to match the layer's KDAConfig by hand. Both go. A KDA layer hands the resource its bound and output gate as it binds, and the pool's conv dtype comes in at build, so the fused kernels are not picked for a bound past what their decay blocks hold or for activations other than bf16, which passed build and then asserted in the first prefill or verify. Those fall back to the torch reference with a warning, as a head dim the fused kernels aren't validated for already did. * kda: the resource checks that each verify step settles its count once A verify step's forward has to hand set_prefix_len its accepted counts exactly once, and a request that verified has to keep verifying, since run_paged never replays the accepted drafts; nothing checked either, so a slip left the state silently wrong. The resource now refuses a second set_prefix_len for one step or one after a step that didn't verify, an eager verify step that commits unsettled, a replay of a verify capture that never settled, and a step that doesn't verify for a request that did. The requests that verified are handle-keyed, purged by remove_request. Applying the count from the step itself, rather than by the model's call, waits on the engine handing a forward's outputs to its resources. * resources: a refused admit drops the stage of the step it refused When a pre-planned step's live admit was refused, the unwind gave back the slots its pre-admit took, but the runner kept the stage and its record of the pre-admit. The same step retried then skipped the pool's admit as covered and promoted the staged addressing: slots the requests no longer held, which the pool could have handed to others. The runner clears the stage when it refuses the step it was for, so a retry admits and plans afresh. * glm52: pass a request's penalize_prompt to the sampler get_request_resource_configs built the sampler config from a fixed key list without penalize_prompt, so a request that set it false with a repetition_penalty still had its prompt penalized. It now passes with top_k and min_p as a sampler knob with no config default (as #390 does for GLM-5.3). test_glm52_request_kwargs: every knob the sampler declares (REAL_KNOB_RULES, INT_KNOB_RULES, BOOL_KNOBS) reaches the sampler config or is refused at intake. * glm52: refuse capture ladders, MTP k and graph_compile it can't serve Each of these loaded and then failed or served something else: - prefill_capture_batch_sizes [0, ...] or a negative entry failed the capture build at startup (ZeroDivisionError, IndexError), and an empty list quietly meant the default; - prefill_batched_token_buckets took 0 and negative buckets into capture; - graph_compile: "false" compiled (bool of a string); - mtp_num_draft_tokens took True as 1 and 2.5 as 2, and a k whose two verify blocks fill the context left max_prompt_tokens negative, so every prompt was refused at the API. Now each is a ValueError at load naming the knob. * test: KDA TP ranks sum to the full layer through their own sharded pools The TP test summed each rank's dense reference prefill, a path the engine never runs. It now serves a prefill, a decode and three verify steps through each rank's own pool and KDA resource (speculative blocks sharded with the heads, o_proj into the all-reduce buffer) and checks that the partials sum to the unsharded layer at every step. * glm52: refuse a prefill_token_buckets ladder of floats, bools or a string int() read [64.5, 128] as [64, 128], [True, 64] as [1, 64] and "64" as [6, 4], so each loaded as some other ladder. The ladder now reaches the submodule as given, where the base class's check (Qwen3.5's) refuses them. Test: test_glm52_prefill::test_a_bad_prefill_ladder_is_refused. * engine: a submodule's recompile limit lifts dynamo's per-code-object cap too with_recompile_limit raised recompile_limit only. Dynamo also stops at accumulated_recompile_limit (256) entries per code object, so a submodule asking for more than 256 still ran the captures past it uncompiled. GLM-5.2's shipped ladders need about 145, so nothing served today hits it. Test: test_torch_config::test_a_submodules_recompile_limit_holds_inside_its_calls_only. * test: a packed replay after the release pads with a zero-length row The release tests replayed one row in a bs-1 bucket, so no replay had a padding row. The new case replays one row in a bs-2 bucket and checks that the padding row takes no pages. DummyRowPool's docstring now says which regions keep their padding rows' pages. * rid_table: list GLM-5.2's MTP maps among the handle-keyed state Glm52LLMSubmodule keys its MTP state by handle and clears it in cleanup_request; the list is where a reader checks what a recycled handle could inherit. * kda: the resource refuses a captured prefill its kernels can't replay The torch bundle's prefill slices each step's spans on the host, so a captured prefill would replay the capture's spans on every later split, with no error. GLM-5.3 leaves those buckets out when prefill_capturable is False, but a model that forgets got wrong rows. A prefill planned under a slot lease on such a bundle now raises, which fails the bucket's warmup at load. Decode and verify buckets still capture. #1200 moves this to a capture refusal once #686 lands. Test: test_kda_resource::test_a_captured_prefill_on_the_torch_bundle_is_refused * dense_fp8: the kernel per token count comes from the registry's fp8_linear table dense_fp8 picked w8a16, FlashInfer's block-scale GEMM or cuBLAS's per call from token limits in code, plus a per-GPU dict that only the H200 had (GLM-5.2's o_proj and dense down on FlashInfer from 3 tokens). That dict is now tables/fp8_linear/NVIDIA_H200.json, an fp8_linear op in #986's registry, declared next to the kernels as moe_decode is. A shape's rows are looked up once per device; a row whose implementation doesn't run on the GPU takes the built-in rule, which is the old dispatch without the per-GPU dict, so every GPU picks what it picked before, an H200 without FlashInfer included. Ported from the kernels stack's fp8_linear (ea9f5b0df, aff4d564f) onto main's registry API, minus its sweep harness and GLM-5.3's q_a + kv_a row (GLM-5.3 builds no fp8 linear here). Test: test_glm52_kernel_platform::test_each_gpu_picks_what_dense_fp8_picked_before_the_table checks the kernel and tile for every GLM-5.2 TP8/TP4 shape (and GLM-5.3's MLA and MLP shapes) at 24 token counts on 7 GPU and library combinations against a table recorded from the old dispatch. * model utils: one capture-ladder check for every model's knobs GLM-5.2's `_capture_ladder` moves to `mstar.model.utils.capture_ladder`, so GLM-5.3 (#390) refuses its ladders by the same rule instead of a copy. No behaviour change; test_glm52_prefill's knob cases cover it. * docs: correct_len's wording, and _decode_locations keeps main's docstring Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * kda: a dtype or gate bound the fused kernels can't take is refused at load Where the fused kernels would run (CUDA, a supported head dim), an fp16 pool or a forget-gate bound past -5 used to fall back to the torch reference with a warning, which also left prefill uncaptured. Both are now a ValueError at build or bind that says what to change. A full-rank output gate and a head dim the kernels aren't validated for still run the torch reference, and off CUDA the torch reference takes any dtype and bound. Test: test_kda_resource::test_what_the_fused_kernels_cant_run_is_refused_at_load * glm52: a verdict read gives up after 5 minutes and exits the process A verify step that failed after it was numbered never stages its verdict, and the plan thread waited for it with no deadline: the worker hung with its requests open. The read now raises SystemExit after _VERDICT_DEADLINE_S, as a read spanning two steps does instead of an assert: one rank failing its batch alone would leave the others in the step's collective, so the process exits where everyone sees it. Tests: test_glm52_mtp_engine_cycle::test_a_verdict_that_never_lands_exits_the_process, ::test_a_read_across_two_steps_exits_the_process. * worker: start the GPU and plan threads before the loop Before torch 2.12 dynamo's config is process-wide. The plan thread started at the first pre-plan, and its initializer (apply_dynamo_config) set the engine's recompile limit while the GPU thread was inside a submodule's raised one (with_recompile_limit), so the shapes after it compiled under the lower limit. Both one-thread pools now run their initializer when they are built. Test: test_torch_config::test_the_engine_threads_start_before_a_compiled_call_can_raise_the_limit. * glm52: the MTP stop check asks TokenStop for the stop token _mtp_stops repeated TokenStop's stop-id and ignore_eos check by hand. It now calls _token_stops on the step's last token and _stop_names for what to stop, and keeps only the budget count, which a verify step's several tokens need. --------- Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com>

This was referenced Oct 19, 2026

Sampler: keep the eager path out of dynamo, and take rows from the planned step #682

Draft

Docs: running M* on AMD GPUs #687

Draft

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from cc9f6dc to 78e3a1d Compare October 19, 2026 22:31

merceod force-pushed the users/atindra/pr-04-gdn branch from 5f6a041 to c0b5d1a Compare October 19, 2026 22:31

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 78e3a1d to 1f1a5c8 Compare October 19, 2026 22:46

merceod force-pushed the users/atindra/pr-04-gdn branch from c0b5d1a to 36f38be Compare October 19, 2026 22:46

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 1f1a5c8 to 15cdeb4 Compare October 19, 2026 22:50

merceod force-pushed the users/atindra/pr-04-gdn branch from 36f38be to 1a07215 Compare October 19, 2026 22:50

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 15cdeb4 to ae856e1 Compare October 20, 2026 05:24

merceod force-pushed the users/atindra/pr-04-gdn branch from 1a07215 to 76b8c96 Compare October 20, 2026 05:24

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from ae856e1 to 15c978f Compare October 20, 2026 05:59

merceod force-pushed the users/atindra/pr-04-gdn branch from 76b8c96 to 0f7484a Compare October 20, 2026 05:59

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 15c978f to 4e0c11a Compare October 20, 2026 06:02

merceod force-pushed the users/atindra/pr-04-gdn branch from 0f7484a to 087999c Compare October 20, 2026 06:02

github-actions Bot added area/engine area/resources size/L labels Oct 20, 2026

dghosef added type/feature type/perf labels Oct 20, 2026

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 4e0c11a to 1eef02f Compare October 21, 2026 02:30

merceod force-pushed the users/atindra/pr-04-gdn branch from 087999c to 937cc89 Compare October 21, 2026 02:30

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 1eef02f to 2753cfc Compare October 21, 2026 08:47

merceod force-pushed the users/atindra/pr-04-gdn branch from 937cc89 to 6e1ca18 Compare October 21, 2026 08:47

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 2753cfc to 7832713 Compare October 21, 2026 10:22

merceod force-pushed the users/atindra/pr-04-gdn branch from 6e1ca18 to c457558 Compare October 21, 2026 10:23

merceod force-pushed the users/atindra/pr-05-ar-fusion branch from 7832713 to e4ab232 Compare October 21, 2026 10:32

merceod force-pushed the users/atindra/pr-04-gdn branch from c457558 to f838869 Compare October 21, 2026 10:32

dghosef assigned dghosef and unassigned dghosef Oct 29, 2026

merceod added 5 commits October 31, 2026 23:36

test: keep the mamba2 oracle in exact fp32

054b31d

The engine's process-wide 'high' matmul precision runs the oracle's fp32 einsums as XF32 on gfx942, about 1e-2 off fp64, which failed the kernel check there; the Triton kernel itself is within 4e-6 of fp64.

gdn: document the linear-attention specs and the AITER capture bound

cfe6abf

merceod force-pushed the users/atindra/pr-04-gdn branch from f838869 to acb54aa Compare November 1, 2026 07:49

merceod mentioned this pull request Nov 1, 2026

Docs: AMD media packages, compile caches and model status #691

Draft

dghosef mentioned this pull request Nov 3, 2026

Add Kimi Delta Attention on the recurrent pool #389

Merged

vasilevklart mentioned this pull request Nov 4, 2026

KDAManager refuses the prefill buckets its kernels can't capture #1200

Open

Navigation Menu

Repository navigation

ROCm: gated delta net on AITER and FlyDSL, and capture buckets a resource can refuse - #686

ROCm: gated delta net on AITER and FlyDSL, and capture buckets a resource can refuse#686
merceod wants to merge 6 commits into
users/atindra/pr-05-ar-fusionfrom
users/atindra/pr-04-gdn

merceod commented Oct 19, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Since r8

Checks on this rebase

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

merceod commented Oct 19, 2026 •

edited

Loading