Repository navigation
[2/3] kv: cross-request prefix caching - #280
Conversation
There was a problem hiding this comment.
Three issues I found:
-
The implementation has wrong-position keying. Keys assume text_inputs[0] is stream position 0. A Bagel edit sent as image,text writes the VAE and ViT walks into main first, and their pages get indexed under the text's keys. A later T2I with cfg off and the same prompt matched "2 of 2 pages" and produced a garbage image. Please try to reproduce yourself too and see the output yourself because I might be mistaken. Potential fix: only let commits the engine offered to the probe advance the chain, or declare the keyed walk on PrefixStream.
-
Noticed an unconverted hit after a refused batch admit. The stream keeps pages with stored_len 0, the retry goes again, admit skips the conversion, the trimmed step runs from position 0, and its pages get filed under the prefix keys so later requests read garbage. A one-line guard might fix it. Also, your PR #281 fixes it by probing at ingest.
-
The CUDA-graph index buffers overflow. The attention wrappers size paged_kv_indices by max_num_pages which is a bound that is broken by sharing. A small-pool hitFlashInfer's plan check and failed two whole batches. Also, the default pool has the same problem once batch size times pages per request passes 2048.
Note (should-fix): the leaves heap grows one entry per and nothing drains it until the pool fills (10k lookups left 10k entries).
Adding the following part later (a few details to go with the points above):
On the first one, to reproduce it the edit has to be the first writer of that prompt. If a t2i or a text-first edit with the same text ran earlier, the correct text pages are already in the index and "first writer wins" refuses the VAE pages, so nothing looks wrong. It should be a fresh prompt of at least 128 tokens, send the edit with input_modalities=image,text and a picture, then a t2i with the same prompt and cfg_text_scale=1.0, cfg_img_scale=1.0 (with cfg on the t2i is never tested). The log then says matched N of its N pages and the image is garbage. Unit version: ingest keys for a 40 token prompt on a 16 page manager, commit a 50 token step on main first, and the index has two entries. The check at manager.py line 453 only looks backward, it passes because 0 is not past covered_len.
On the second, the guard is the eligibility test in resolve_cached_prefix (manager.py line 399): add or stream.page_indices there, or convert anyway at admit (manager.py line744) and release the private pages.
On the third, it's the _paged_kv_indices_buf = torch.zeros(max_num_pages, ...) in wrappers.py at 84 and 229. The error is flashinfer's "The size of indices should be less than or equal to the allocated buffer", raised from the decode preplan. Sizing these as capture batch size times ceil(max_seq_len / page_size) should fix it, it's only int32s.
On the heap, it's the push in lookup (prefix_index.py line 46), one entry per hit on a leaf. Lazy re-push keeps it bounded so don't push on lookup, and in _pop_leaf when an entry's stamp is stale but the page is still a leaf, push it back with the current stamp instead of dropping it.
| for page in pages: | ||
| self._stamp[page] = self._clock | ||
| if self._children[page] == 0: | ||
| heapq.heappush(self._leaves, (self._clock, page)) |
There was a problem hiding this comment.
Every lookup that ends on a leaf pushes a new heap entry, and nothing pops them until the pool is full. A system prompt gets re-stamped on every request, so on a server whose pool never fills this list grows by one tuple per hit (10000 lookups in a test left 10000 entries), and the first eviction then walks through all of them under the lock. Lazy re-push would keep it bounded (so don't push here, and in _pop_leaf when an entry's stamp is stale but the page is still a leaf, push it back with the current stamp instead of dropping it).
There was a problem hiding this comment.
Good catch! I dropped the push in lookup, and _pop_leaf now puts a stale entry back when the page is still a leaf rather than losing it. Ten thousand reads of one leaf leave one entry
| or inputs.custom_pos_ids is not None | ||
| or inputs.tensor_inputs | ||
| or inputs.kwargs | ||
| or inputs.resource_step_info |
There was a problem hiding this comment.
This seems fragile. It works for Bagel because requires_cfg is False for chat and True for t2i, but a model that puts 0 or an empty container here would and one that puts anything else would never, and nothing in the model says so. Letting the model name the keyed walk (say, on PrefixStream) makes this clear, and it also becomes the store-side signal for the position-zero problem in _index_filled_pages.
There was a problem hiding this comment.
The walk is named on PrefixStream now, and that's the half that fixes the indexing. I kept this check beside it though, because requires_cfg is per request: the same walk is sliceable for a chat and not for a guided t2i, so the model can't say up front which one it is.
| model.preprocess_fingerprint() if model is not None else "", | ||
| ] | ||
| for key, resource in self._resources.items(): | ||
| if not isinstance(resource, KVManager): |
There was a problem hiding this comment.
Small layering thing, the engine picking out KVManager by type. If enable_prefix_cache were a Resource with a no-op default the engine wouldn't need to know which resource kind keeps an index.
There was a problem hiding this comment.
Fair, the engine shouldn't know which resource kind keeps an index. It's a Resource method with a no-op default now.
| if not matched: | ||
| return | ||
| self._arena.retain(matched) | ||
| stream.page_indices = list(matched) |
There was a problem hiding this comment.
If the stream already held pages here (a refused admit leaves some), they're overwritten and never released. An assert that it's empty or a release would do.
There was a problem hiding this comment.
Right, those were leaked. It releases them first now, in the same commit as the admit conversion.
| return False | ||
| self._clock += 1 | ||
| self._by_key[key] = page | ||
| self._key[page] = key |
There was a problem hiding this comment.
Nothing checks the page isn't already indexed under another key. Shouldn't happen today (I think), but if it ever does the old _by_key entry just hangs, and an assert here would be good.
There was a problem hiding this comment.
Added, with the page and both keys in the message.
| @@ -340,8 +372,19 @@ def _process_input( | |||
| prompt_parts=input.prompt_parts, | |||
| **(input.model_kwargs or {}), | |||
There was a problem hiding this comment.
process_prompt still gets the client's original kwargs here, prefix_keys included if they sent some. The stripped model_kwargs is right above and passing that instead would keep the "only this worker keys a prompt" idea alive all the way through.
There was a problem hiding this comment.
You're right, that one slipped through. It gets the stripped copy now.
|
@merceod I've pushed the fixes and updated the docs. Position-zero keying: the model names the walk that writes the keyed text, and a commit from any other walk ends that stream's chain. Index buffers are sized per row, so the 256-page pool that failed 29 of 32 requests completes all of them. Two I did differently from what you suggested. The buffer bound is Let me know if these make sense to you and if you'd like me to change anything. |
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Approving this since everything lgtm, but had one important note.
I ran a 96 page pool with 32 requests sharing a 59 page prefix. They all were allowed/admitted, since, with the prefix cached, each prefill only needed itstail page. Then every decode step wanted one more page, nothing was evictable because every page in the index was live, and with no cpu offload configured the worker just held all of them at the 50ms backoff until the clients timed out. It recovered fine once they did. The hold path is main brnach's, but a prefill no longer reserves the prompt's pages, so it no longer bounds concurrency by what the pool can hold. The sizing rule from your RFC fixes it I think, so maybe a small comment in serving.rst for now, and a headroom check at admission or a preemption path when a decode step finds nothing to evict would be a good follow up. Also that path logs a WARNING per retry per batch, I got 19k lines in six minutes, so rate limiting it would help.
|
Good catch! It's worse than a sizing note: prompt reservation was doing admission control by accident. Before this PR a request reserved its whole prompt, so concurrency was bounded by pool over prompt pages; with a hit the prefill reserves a tail page and that bound is gone. The cache is on by default, so a pool sized for the old behaviour over-admits on upgrade. Worth knowing, the victim path already exists: The new commits add the sizing rule in serving.rst, the rate limit on that warning, and two lines that would have made this visible: one at load naming which nodes are keyed and warning when a cached node has no host pages, one the first time requests reach a keyed node with no keys. There were additional fixes I added as well. Reading the manager for #281 I found @merceod Do you want I'm running a suite of benchmarks and will report the numbers here. |
|
@merceod Okay I have also filed the follow-ups. #307 covers the stall you found, #308 covers TP, and #309 covers guided prefill. None of them block this PR. The stall has a workaround and a warning at load, TP stays off above one rank, and admission needs the benchmark numbers which I will add soon. |
@Gaurav-Shah05 Keep it on I guess. The hold only happens once the pool is full of live pages, and main doesn't deal with that that either, it just gets there later because a prompt reserves its pages up front. What changed is the threshold, not the failure. Two things worth naming in the docs sentence and the load warning: max_concurrent_requests in the YAML bounds admission directly today (pool over expected pages per request is a fine setting), and cpu_offload_pages gives the worker a victim. With those two named, #307 can take its time. Merging. |
stephen-dwq
left a comment
stephen-dwq
left a comment
There was a problem hiding this comment.
fuzzing of enable_prefix_cache, ingest_request, resolve_cached_prefix, apply_cached_prefix, admit, plan, commit, extend_prefix_chain, reset_request, remove_request comes clean
lgtm
A colocated decode reads its own publish at every step, so clearing there stopped the chain at the first sampled token and none of the generated pages reached the index.
A stream that stops keying now drops its cursor and reported flag with its keys; nothing reads them again except a later ingest that reseeds that stream, and a reset of a keyed stream, which only capture rows reach.
eaeb336 to
5d51741
Compare
|
@merceod @NSagan271 @stephen-dwq @vasilevklart These are the benchmark numbers for prefix caching for Bagel on 1 H100. Head is #280 (
Page size, kvshare c4, 262,144-token pool:
|
The adapter predates #280 and joined every message into one user turn, so a chat through the Rust frontend rendered differently from the same chat through the default server.
* api: keep each message's role on its prompt part flatten_messages drops the role and joins text across messages, so a reply reaches the model inside the next user turn. A part now carries its role, text merges only while the role stays the same, and an attachment outside a user message is refused. * test: roles survive the chat adapter * model: bagel renders one turn per role change A reply now comes back as its own turn, so a chat's next prompt starts with the last prompt and its reply, token for token. Text chat keeps the bare format, one block per turn; understanding writes each turn in its own role block. * test: bagel renders a chat one turn at a time * model: qwen3-omni renders one turn per role change The chat template gets one entry per turn with its role, and the default system message only when the client sent none. * test: qwen3-omni's chat matches its own template * api: the dynamo bridge keeps the prompt parts _submit dropped prompt_parts, so a chat over Dynamo reached the model as one user turn, and text written after an attachment moved ahead of it. * test: the dynamo bridge keeps the prompt parts * model: build a chat template's messages from the prompt parts in one place Qwen3-Omni built its messages list inline. Kimi K2.7, VoiceChat and Ming will take the same list instead of flattening, so it moves to multimodal.py as messages_from_parts; Qwen3-Omni renders the same ids. * test: messages_from_parts gives one message per turn * model: a chat reply leaves out its end token Bagel and Qwen3-Omni decode every emitted token with its special tokens, so a chat reply ended with a literal <|im_end|>; a client that sends the reply back sent the end token twice, once as text. * test: a chat reply leaves out its end token * ci: run the rust frontend's tests The job built rust/server and ran none of its tests, so the adapter's and protocol's tests ran nowhere. * api: the submit message carries the parts The Rust frontend's submit had no field for a chat's ordered parts, so the bridge could only call submit_request without them. SubmitArgs gains parts, the message carries them as prompt_parts (left out when empty), and the bridge builds a PromptPart from each; nothing fills them yet. * test: the submit message carries the parts * api: the rust frontend builds the prompt parts The adapter predates #280 and joined every message into one user turn, so a chat through the Rust frontend rendered differently from the same chat through the default server. * test: a chat through the rust frontend keeps its roles and attachment order * api: articles in the refused attachment's message It read "a assistant message cannot carry a image attachment", in Python and in the Rust adapter that mirrors it. * api: drop the realtime end-token strip A reply no longer carries its end token, which now decodes to an empty chunk that the empty-delta check already skips. * model: bagel keys its replies again * model: qwen3.5 and cosmos3-edge render one turn per role change * test: their chats match their own templates * api: developer reads as system, and a role a chat can't render is a 400 * test: roles at the chat adapter * api: an attachment in any message renders in a user turn * test: a reply's image comes back * model: bagel renders text turns in role blocks * test: bagel's text chat renders one role block per turn * model: keep end tokens under ignore_eos * test: ignore_eos keeps every end token * model: qwen3.5's reply leaves out the stop token that ended it * model: cosmos3-edge keeps its end tokens under ignore_eos * model: bagel, qwen3-omni and qwen3.5 drop stop tokens through one helper * test: one chat test for every model with a chat template * api: a chat that asks for a tool call is a 400 * api: the older functions and function_call get the same 400 * api: the dynamo bridge passes tools and tool_choice to the adapter's 400 * api: a 400 from the python server is an invalid_request_error * api: the rust frontend types a backend's 400 invalid_request_error * test: a tool-call chat is a 400, and a 400 is an invalid_request_error * model: bagel generation separates the system text * model: a template error is a 400 on every model * model: comment on qwen3-omni's persona * api: tools with tool_choice none pass * test: tool_choice none and template errors * model: an empty cosmos3-edge chat is a 400
… 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>
What does this PR do?
Part of #210, stacked on #279 . Reuses KV pages across requests that start the same way: chained SHA-256 page keys under a fingerprint root, computed on the preprocess worke. On by default for any model that declares a token stream for a KV resource (Bagel's text LLM and Orpheus here), inert for everything else.
A hit
Engine._prepare_inputsasks the node's resources how much of the request is cached, takes the minimum, tells each resource, and trims that many tokens before the step is declared; bucket selection, attention, positions and commit see a shorter request. The sampler is told too, so its repetition-penalty mask still covers the skipped tokens. Keys come only from the preprocess worker; anything a client sends under those names is discarded.Storage
The arena's own pages under a second owner, no copies. A page is offered to the index by the commit that fills it on the keyed walk, prompt or generated. An allocation or reload that comes up short evicts leaves from the index before failing, skipping any a live request still reads.
What is cached
Bagel text chat and Orpheus prompts, generated tokens included. The model names the walk that writes the keyed text, so that walk alone is probed and its commits alone extend the chain; another walk writing the same stream ends it, as does a write the keys do not cover.
Not cached: anything after an image, since there is no digest stream yet, which the follow up adds; and Bagel's guided image prefill, whose inputs the default cut refuses because
requires_cfgtravels with the request rather than the walk. The follow up does not change that second one, which needs a Bagelsplit_inputsthat cuts thecfg_imgspan and the fork. A windowed stream keys only its protected prefix. TP is #281.Declaring a model
A model opts in with
prefix_key_streams(), naming the tensor that holds the token ids and the walks that write and decode it, pluscheckpoint_path(), since the weights are part of the root and a model that names none needs a salt to use the cache at all. It is about four lines each for Bagel and Orpheus, and a model that declares nothing is untouched. A declaration the node cannot honour is refused at load. I have updated docs/adding_models.rst w/ the contract.How was it tested?
178 new tests.
pytest test/modular -qatd172d293gives1 failed, 1112 passed, 20 skipped(the failure is pre-existing on main). Cache off is bit-identical to main at424a112cwithin one allocation. A hit changes the output without changing what it says: of six two-turn chats, three answered exactly as the uncached run and three paraphrased it, and on a text-to-image about 1% of pixels move by more than 32 of 255 with the scene unchanged, soprefix_cache=Falseis there when an output has to be reproducible. Turn-two TTFT went from 44.9 ms to 26.4 ms. An image-first edit followed by the same prompt as text-to-image matches nothing and its image is identical to the uncached run, and 32 same-prefix requests on a 256-page pool all complete.Checklist
ruff check .passes