Extend Double-Buffering to the Piecewise Cuda Graph Runner by NSagan271 · Pull Request #243 · mstar-project/mstar

NSagan271 · 2026-09-07T15:48:51Z

What does this PR do?

Extends double-buffering to the piecewise CUDA graph runner and fixes a host race in the sampler. Piecewise graphs previously ran single-buffered, so a resource that stages a step's layout into a reused pinned host buffer behind a
non-blocking H2D could have plan(N+1) overwrite the buffer before step N's DMA retired. The sampler's per-step gather
indices (_slot_idx_cpu) could hit this, e.g. in qwen3-tts prefill under a piecewise graph.

Adds an opt-in force_double_buffer resource flag (default off, set only by the sampler right now) that enables double-buffering off the pre-plan path; the piecewise runner honors it for any resource its declare-step touches. The main runner's existing supports_preplan-driven double-buffering is unchanged (but it also checks force_double_buffer).

How was it tested?

Ran perf and correctness on qwen3-omni, orpheus, and qwen3-tts.

Checklist

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

stephen-dwq · 2026-09-07T21:06:29Z

This key is still a 2-tuple, but we now use a 3-tuple (defined on L1115). This means that _resolve for batched regions would return None without an error, and packed regions would give a ValueError.

stephen-dwq · 2026-09-07T21:09:05Z

+                        "total_tokens=%d", self._label, shape.bs, shape.total_tokens,
+                        exc_info=True,
+                    )
+                    break


Because we now barrier on the number of run slots, but break out of the loop here, we can have a mismatch of barriers between ranks if one breaks out and the otheres do not. This can result in hang. I think we can continue, or try to hoist the barriers?

Fixed: the slot loop always walks every slot and configs with failures are dropped in a separate loop after

stephen-dwq · 2026-09-07T21:19:53Z

+        # path (e.g. in the piecewise runner). This is the slot the batch
+        # rotation cycles through; each runner narrows it to its own slot count.
+        double_buffered = any(
+            res.supports_preplan or res.force_double_buffer


Piecewise will do num_slots if double_buffer else 1, Main runner just does num_slots. If we have a resource with the staging hazard addressed in this PR but does not support preplanning, the main runner ends up with two slots for everything when we only need it for the piecewise because we have no general preplanning.

Technically, a hazard would also be present in the regular cuda graph case without double-buffering, so it's not piecewise-specific (it's just that the regular case is already double-buffered). It might be worthwhile to double-buffer per-slot based on whether any of the resources used need to be double-buffered, but I might leave that as a TODO comment.

garv901 · 2026-09-07T22:17:50Z

+    def force_double_buffer(self):
+        # Per-step gather indices are staged into a reused pinned buffer
+        # (``_slot_idx_cpu``) behind an async H2D, so a single-buffered region
+        # (a piecewise runner with no pre-plan) races when the host runs ahead.
+        return True


We should fix MSTAR_NUM_SLOTS to a minimum of 2 when forcing double buffers else we might reintroduce the race

Also we would set this property true for the flash infer ragged attention too right? It would have the same race issue without this

Fixed both, and also (1) set force_double_buffer for the non-ragged flashinfer wrappers, (2) added double-buffering for the eager flashinfer wrappers

…emains)

… on slot, add force double buffer for more resources

NSagan271 · 2026-09-10T19:45:50Z

@stephen-dwq @garv901 did I end up addressing your concerns?

stephen-dwq · 2026-09-10T22:26:54Z

MSTAR_NUM_SLOTS is now dead code, and we should also add MSTAR_PRE_PLAN_SPEC to environment variable documentation. Everything else looks fine.

NSagan271 · 2026-09-10T22:40:52Z

@stephen-dwq Technically, MSTAR_NUM_SLOTS is not dead code --- it is possible to bump up the number of slots above two. I guess this is arguably worse because speculation only plans one step ahead, so it's just adding more memory for nothing. I'm ok removing the environment variable if you want.

Also, MSTAR_PRE_PLAN_SPEC will be added to the documentation in #224, which I will get merged in soon.

stephen-dwq · 2026-09-10T23:50:14Z

I think it's okay to keep if we have any plan of bringing speculation to more than one step ahead. Otherwise, may as well get rid of it.

garv901

changes look good to me!

NSagan271 · 2026-09-13T05:45:40Z

@stephen-dwq if things look good can you approve the PR?

stephen-dwq · 2026-09-13T06:10:53Z

Apologies for delay, forgot that I have to overwrite my previous request change.

… all-reduce (#386) * 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. * 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. * 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. * 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. * 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. * 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. * 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. * 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. * 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). * 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. * 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. * 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. * 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. * 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. * test: CI runs the new CPU tests Each new test file with a CPU case, listed in test/cpu-core.txt. * 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. * 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. * 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. * cuda_graph_runner: empty compile mode env is unset Raising at import failed every model's boot on an empty export. * 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. * 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. * 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: #745's run-loop worker stubs the timeout restore #386's warmup ends with --------- Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com>

* 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. * 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. * 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. * 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. * 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. * 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. * 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. * 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. * 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. * 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 * 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 0dedee4, 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. * 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. * 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: 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. * 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 (caeea87, 1ded60c) 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. --------- Co-authored-by: Kirill Vasilev <290354918+vasilevklart@users.noreply.github.com>

NSagan271 marked this pull request as ready for review September 7, 2026 15:54

NSagan271 requested review from Gaurav-Shah05, garv901, merceod, stephen-dwq and vasilevklart September 7, 2026 16:03

stephen-dwq requested changes Sep 7, 2026

View reviewed changes

garv901 reviewed Sep 7, 2026

View reviewed changes

This was referenced Sep 8, 2026

worker: async scheduling for TP nodes (MSTAR_TP_ASYNC_SCHED) #224

Merged

sampling: write master rows with a scalar fill_, not through a reused pinned row #222

Merged

NSagan271 added 7 commits September 9, 2026 06:23

initial implementation (works, unprofiled)

4bb3914

revert changes to attn and kv manager for perf reasons (sampler fix r…

0ee899b

…emains)

ruff

130d65a

Fix more race conditions: sync on the single path, key eager wrappers…

b284db2

… on slot, add force double buffer for more resources

fix [untested]

50c85ee

modify previous fix

54a830a

double buffer for ragged

167dc46

NSagan271 force-pushed the extend-double-buffering branch from 67f7d0f to 167dc46 Compare September 9, 2026 06:24

double-buffer ragged + cleanup

6975390

NSagan271 requested review from garv901 and stephen-dwq September 9, 2026 06:56

garv901 approved these changes Sep 11, 2026

View reviewed changes

stephen-dwq approved these changes Sep 13, 2026

View reviewed changes

NSagan271 merged commit 7675236 into main Sep 13, 2026
3 checks passed

NSagan271 deleted the extend-double-buffering branch September 13, 2026 06:21

merceod mentioned this pull request Sep 24, 2026

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

Closed

dghosef mentioned this pull request Oct 14, 2026

Review prompt: minimal, prioritized, inline comments (into #370) #432

Merged

Conversation

NSagan271 commented Sep 7, 2026

What does this PR do?

How was it tested?

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NSagan271 commented Sep 10, 2026

Uh oh!

stephen-dwq commented Sep 10, 2026

Uh oh!

NSagan271 commented Sep 10, 2026

Uh oh!

stephen-dwq commented Sep 10, 2026

Uh oh!

garv901 left a comment

Choose a reason for hiding this comment

Uh oh!

NSagan271 commented Sep 13, 2026

Uh oh!

stephen-dwq commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants