Repository navigation
feat(vllm): hand the KV warm-up chain's recurrent state to hybrid-layout shadows - #14628
Draft
liyuanzhe1991 wants to merge 3 commits into
Draft
liyuanzhe1991 wants to merge 3 commits into
liyuanzhe1991 wants to merge 3 commits into
Conversation
github-actions
Bot
added
feat
documentation
Contributor
… shadows KV warm-up shadows borrowed a chain's KV by position, which is right for attention groups but wrong for align-mode Mamba/linear-attention groups: vLLM keeps one running-state block per request at the slot derived from its token count and pads the slots before it with the shared null block. Cutting the chain's table at the shadow's token count left the shadow's state slot on the null block (a zero state read, a shared block written), so the gate rejected every model with a Mamba group. Register such groups by their own layout instead. The chain's state block (the last non-null block before its speculative scratch; the manager's own last_state_block_idx may already name a nulled slot) goes to the slot the runner reads at admission, either in place or forked by copy-on-write through the manager's _apply_cow like the attention tail, with a fresh blank block pre-placed at every later slot the steady writes move into (preprocess_mamba copies the state across before reading it) and the speculative scratch closing the table. DYN_BENCH_KV_WARMUP_STATE selects copy / inplace / auto (copy when the manager can fork and the pool holds the forks, else in place; the frozen 0.25.1 image has no _apply_cow and runs in place). The chosen mode lands in the KVWARM metadata and per row as kvwarm_state_<mode> plus kvwarm_state_depth=<chain depth>, since the recurrent layers see the chain's full depth rather than the injected context. Planning and the pool pre-check count the fork, the blank and the speculative blocks per state group; the gate now admits align-mode groups and refuses an explicit copy request on a vLLM without CoW up front. Validation: python3 -m py_compile; black 23.1.0 --check; the changed methods lifted with ast and run against the unit-test fake pool/manager stubs for in-place, copy, block-crossing, speculative-scratch, unwind, gate, planning, mode-resolution and full release lifecycles (every reference back to baseline, null block never queued). Real pytest and a GPU smoke run remain to be done in the container. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: YZLi <yuanli@nvidia.com>
Unit tests for the align-mode Mamba path of the KV warm-up shadows, in the existing fake pool/manager style: a _FakeMambaManager (align mode, block size 64, a null-padded table ending in the running-state block, CoW optional) and a hybrid two-group stub over one pool. Covered: in-place registration (chain state at the read slot, the rest null, +1 on the state block), copy registration (hit-ref first, fork through _apply_cow, the retained release), a pre-placed blank block at every crossed state slot, the speculative scratch after the state, locating the state behind a nulled last_state_block_idx, unwinding on an unresolved mode / copy without CoW / a non-align group, the gate admitting align mode and refusing an explicit copy without CoW, per-manager planning and pool pre-check arithmetic, DYN_BENCH_KV_WARMUP_STATE parsing and the auto rule, the row and metadata stamps, a full release lifecycle in both modes with every reference back to baseline and the null block never queued, and the injection path handing the state over and stamping the point. The pure attention layout keeps its previous registration results; only the stage-check stub gains the state-mode argument it is now passed. The fake pool grows a null block and mirrors vLLM's free_blocks on it. Validation: python3 -m py_compile; black 23.1.0 --check; the test functions and fakes lifted with ast and run against the lifted scheduler methods (43 cases, all passing, including the pre-existing shadow registration, lifecycle and planning tests). Real pytest in the container is still pending. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: YZLi <yuanli@nvidia.com>
Why:
- A vLLM whose MambaManager registers partial-tail prefix entries
(_cache_partial_tail_block, state block size above the hash block size)
hashes the block at a request's state slot on every output update. An
in-place shadow put the chain's own state block there: the shadow's key
stacked on the chain's, num_cached_block dropped back to the state slot
and the first steady write reaching the next slot failed the
partial-to-full promotion assertion in cache_full_blocks. In-place is
now refused on such managers (gate: state_inplace_unsupported;
registration raises) and auto forks instead. Builds without the method
(the frozen 0.25.1 image) keep in-place.
- auto settled the mode per injection against the rank-local pool, so two
attention-DP ranks could stamp different kvwarm_state_* reasons for one
point and abort on the point-digest mismatch; the plan, the stage check
and the injection also priced three different modes ("mixed"). The mode
is now a function of the layout and the vLLM build alone -- copy when
every state manager forks, else in place where admitted -- and a pool too
small for it skips the point through the existing shortfall path.
- The chain's state slot is read from the chain request's computed depth
((depth - 1) // bs, the slot preprocess_mamba reads next) instead of the
table's shape; without draft tokens the state must be the last block
(checked), with drafts an over-allocated tail is only reported.
- The gate keys state groups on mamba_cache_mode, the attribute the
registration uses, so a non-align spec is skipped whatever its name.
- DYN_BENCH_KV_WARMUP_STATE is parsed only on layouts with align-mode state
groups (skip reason state_mode_invalid); the metadata records the raw
value, so a typo cannot fail a run the variable does not concern.
- Docs: in-place side effects, per-build mode availability, new reasons.
Validation:
- python3 -m py_compile on both Python files; black 23.1.0 --check clean.
- All 263 tests of test_vllm_instrumented_scheduler.py pass in a local
harness that stubs the vLLM imports (no vLLM on this host), including
six new cases for the partial-tail gate, the cache-mode key, the
env-value scoping, the depth-derived state slot, the pool-independent
mode and the rank-invariant stamp. The container pytest run is pending.
Signed-off-by: YZLi <yuanli@nvidia.com>
liyuanzhe1991
force-pushed
the
yuanli/fpm-kvwarm-kda-state
branch
from
September 11, 2026 06:24
36c2215 to
c53d9ff
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #14029 (merged as 2d975ac); rebased onto
main, three commits.Hybrid layouts whose recurrent-state layers (Mamba / linear attention such as KDA) keep one in-place-updated state block per request could not take part in the KV warm-up: the eligibility gate rejected them wholesale (
hybrid_state_layers_unsupported, added in #14029 review), so every decode point of such a model fell back to fake injection. Before that gate, the borrowed shadow's state slot pointed at the shared null block (the chain's running state sits in the last slot of the chain's table and the shadow only borrowed the leading slots), so the recurrent layers ran from a zero state and wrote into the null block.This PR hands the chain's recurrent state to the shadow explicitly:
e9fdff2155):_kvwarm_register_shadowdispatches per KV-cache group. Attention/MLA groups are unchanged. For an align-mode Mamba group the shadow's table is filled with null blocks, the chain's running state block goes into the slotpreprocess_mambawill read ((ctx-1)//block_size), private blank blocks cover any block-boundary crossing during the steady steps (the runner copies the state forward itself) and any speculative slots, andnum_cached_blockis set so vLLM does not re-hash or cache the shadow's blocks.DYN_BENCH_KV_WARMUP_STATE=auto|copy|inplace:copyforks the state block through the manager's CoW path (same sequence as the attention tail fork; retentions are settled inside the injection window);inplacelets the shadow read and update the chain's state block directly (reference-counted, stamped on the row askvwarm_state_inplace), which is the only option on vLLM builds without_apply_cowand the only memory-neutral option for models whose state page is large.autois resolved from the layout and the vLLM build only — never from the free pool — so attention-DP ranks agree; a pool shortfall goes through the existing shadow-shortfall skip.copyon a build without CoW, or an explicitinplaceon a build whose Mamba manager registers partial-tail hashes on the state block (vLLM 0.28.0) are skipped with named reasons (state_copy_unsupported,state_inplace_unsupported,state_mode_invalid).kvwarm_state_<mode>andkvwarm_state_depth=<chain depth>; the artifact'skvwarmblock recordsstate_mode.Known approximation, unchanged by this PR: the state handed over is the chain's state at its full depth, not at the shadow's context length (recurrent layers keep no per-position history). The recurrent layers' compute does not depend on state content; the effect is confined to expert routing and is to be quantified by the pilot below.
Where should the reviewer start?
components/src/dynamo/vllm/instrumented_scheduler.py:_kvwarm_stage_state_shadow/_kvwarm_commit_state_shadow(the hand-off),_kvwarm_state_mode/_kvwarm_state_gate_reason(mode resolution and gate), then the per-group arithmetic in_kvwarm_shadow_tail_blocksand_kvwarm_shadow_pool_shortfall. Tests:test_kvwarm_state_*and the hybrid variants of the shadow lifecycle tests.Related Issues
Follow-up to #14029 (review thread on recurrent-state layouts by @tedzhouhk: "reject those layouts until scratch/COW state blocks are implemented").
Validation
blackclean;mypyadds no diagnostics on the scheduler module.inplacemode (2026-09-11, Nemotron-3.5-Lightning-30B-A3B NVFP4 = mamba/moe/attention hybrid, TP2+EP, 2×B200, frozen collection image with vLLM 0.25.1,--mamba-cache-mode align,--enforce-eagerbecause this checkpoint does not compile on either image): the gate admits the layout, three rungs build their chains (batch 8/32/128, depth 32772/32772/8196), 9 real-KV points, 0 fake, 0 skipped, every row stampedkvwarm_state_inplace+kvwarm_state_depth,reset_prefix_cachesucceeds at the end (no leaked references). Two boots per arm: warm-up vs no-warm-up medians differ by +0.3..+2.5% while boot-to-boot spread is 4..10% — in eager mode the step time (≈40 ms, independent of batch and KV) is dominated by host overhead, so this pilot establishes functional correctness, not the timing effect of real state.copymode on GPU: not yet exercised — the vLLM 0.28.0 runtime image fails engine warm-up on this checkpoint (aten::emptydispatcher error) on the node available today; covered by unit tests (test_kvwarm_state_*, hybrid lifecycle) pending a working 0.28 GPU environment.🤖 Generated with Claude Code