Resource pools: shared engine resources with a staged migration

[merceod]

What does this PR do?

This PR introduces resource pools per the design doc i.e. positional embedding, attention, and KV storage become independently implemented resources with a uniform admit/plan/commit/publish lifecycle, so a new implementation of any one of them touches that resource and nothing else. The migration runs in stages on this branch; every stage keeps all model families serving with unchanged behavior (I verified against baselines captured at the branch point).

The scope of the redesign is quite large so I have divided it up into smaller stages.

Done:

Stage 1, boundary types and the KV pool.New package mstar/engine/resources: the step boundary values (Segment, Reservation, SequenceView, PositionPlan, all immutable) and the storage split (PageArena holds the tensor and page allocator; KVCachePool is per-request accounting with admit/view/commit). The cache manager's five planning paths (self-attention, batched CFG, cross-attention write and plan, dense-gen) now reserve capacity through pool.admit and build FlashInfer index tensors from views instead of reading KVRequestState fields and allocating mid-plan. PagedAllocationManager stays the storage owner for now. Also, allocation order, arithmetic, and the failure surface are unchanged.

Stage 2: positional embedder resource; plan_rope/apply_rope delegate to it; advance_seq_lens and the custom_pos_advance side channel become pool commit.

Stage 3: attention managers and a step runner on the eager path; BatchedCacheManager becomes a stateless facade essentially; cross-attention becomes a keyed pool and manager instead of a label naming convention.

Stage 4: captured path; replay plans the real segment list into slot-indexed buffers instead of aliasing request state onto dummy slots; gated by graph-vs-eager parity and decode phase timings.

Stage 5: engine surface; per-node resource dicts built from model-declared specs; offload/LRU move onto the pool; a shared scratch KV pool replaces the private code-predictor tensors in qwen3_omni and qwen3_tts.

Stage 6: cosmos3, qwen3_omni, and bagel move from the facade to the resource surface, with a windowed retention policy, a second positional scheme, the keyed scratch cache. The remaining families follow in later PRs, after which get_kv_cache_config is deprecated.

Underway: None

Remaining: All stages done

How was it tested?

The stale allocator thread-safety test is rewritten against the pool and 18 pool unit tests are added. The CPU and GPU failure sets identical to baseline except the intentionally rewritten test; serving A/B against the branch point is byte-identical for cosmos3 (all hashes, cold and warm t2v), bagel, and whisper (which exercises the reworked cross-attention every decode step); qwen3omni matches a base boot exactly (the base tree itself varies across boots on one greedy near-tie); orpheus reproduces the identical token stream with PCM inside the vocoder's own run-to-run noise.

Checklist

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

[stephen-dwq]

We can probably use build_paged_indptrs here rather than writing it in-line.

[merceod]

Done in my commit 6e7b1e2 where plan_batched_cfg now builds its index tensors through build_paged_indptrs. The combined seq_lens list stays local since the helper doesn't produce it. plan() keeps its inline loop on purpose as it also accumulates kv_cache_locations for the decode fast path in the same pass, which the helper doesn't cover.

[stephen-dwq]

Storage is done through PagedAllocationManager rather than KVCachePool. I think we would need to add a reset_label to KVCachePool, a public accessor for the arena tensor which is a bit iffy, and a flush for KVCachePool to take place of PagedAllocationManager.

[merceod]

Agreed on the end state; I want to fold this into the "facade retirement" that I was mentioning on slack rather than do it now. The runner has to keep holding the manager either way today, because create_cache_manager takes the raw kv_cache tensor and manager, so adding pool.reset/pool.flush now leaves both surfaces alive in the same file and removes no field. The arena accessor you call iffy turns out to be avoidable entirely I think as the tensor is only needed for facade and attention-manager construction, which is engine-side wiring, so once the facade goes the runner's remaining needs (dummy reset between captures, padding-slot release, flush) become pool methods without exposing the arena. And flush_to_store is a no-op placeholder right now, so a pool flush would wrap dead code. Capture/replay is the riskiest area so I'd rather churn it once, with the real migration.

[stephen-dwq]

pause_request and resume_request are directly fiddling with the PagedAllocationManager albeit in a minimal manner and with no entanglement.

add_request and remove_request are looping over PagedAllocationManager and adding requests directly. We can use KVCachePool::add_request instead (although this just wraps the manager function).

[merceod]

Fixed in my new commit in 2725156.

[stephen-dwq]

Will we deprecate managers presence here eventually?

[merceod]

Yes commit 2725156 drops cross_alloc_managers from KVManagement.

[NSagan271]

Maybe a TODO for later, but I was envisioning allowing multiple KV cache configs for a node, and also possibly pulling the cross-attention config and RoPE config out of the KV cache config (making it its own separate entity)

[merceod]

Agreed on both and I put your TODO on NodeResourceSpec in commit 9ddc9cb.

[merceod]

Btw the one-config today is deliberate because the default get_node_resources wraps get_kv_cache_config unchanged so the unmigrated families need no change, and the cross-attention and rope settings stay on KVCacheConfig because the facade and those families still read them there. When get_kv_cache_config goes, the spec grows named KV configs and the cross/rope split happens in the same pass, without churning every model config twice.

[NSagan271]

Now that I think about it, it might be more flexible for get_node_resources to replace get_kv_cache_config (and later the functions for retrieving sampling configs), allowing every node to have an arbitrary number of KV caches / samplers / etc. I believe this implementation also ties get_node_resources to KV cache node groups, which would limit flexibility.

[merceod]

That's my intention. get_node_resources replaces get_kv_cache_config once every in-tree model migrates, and the per-config pairing becomes per-node declarations with arbitrary cardinality (the per-node resource dict already keys arbitrary resources, so named KV pools and, later, samplers slot in).

Doing the swap now would churn all ten families mid-review for no behavior change, so it will come with the deprecation.

[NSagan271]

I think the plan is for a node to potentially access several different resources (e.g., KV caches, cross-attention, samplers); I think it'll be cleanest in the end to have this function return a step declaration for each resource that it uses (right now, the StepDeclaration definition seems to be mainly hardcoded to attention and RoPE for a single KV cache)

[merceod]

Hmm. Partly agreed. The end state needs multi-resource steps, but I'd keep one declaration per step rather than one per resource. The step is the unit the runner sequences (one drive, one commit around one forward), and ordering across resources has to live somewhere and it seems to me that the plans tuple is that place. The extension is incremental change when when the first multi-resource family arrives: PlanSpec grows a resource key defaulting to "kv" (whisper's cross-attention migration is the concrete consumer), drive routes each plan to the named resource, and commit already walks per plan so per-resource commits fall out. Samplers stay engine-owned this workstream, so they don't enter the declaration. I'd rather design the multi-resource contract against whisper than guess it now. What do you think?