From 3bbd7b8f092d215c5067d323e571f071a6e9fbde Mon Sep 17 00:00:00 2001 From: YZLi Date: Fri, 11 Sep 2026 01:32:38 +0800 Subject: [PATCH 1/3] feat(vllm): hand the warm-up chain's recurrent state to hybrid-layout 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_ plus kvwarm_state_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 Signed-off-by: YZLi --- .../src/dynamo/vllm/instrumented_scheduler.py | 344 ++++++++++++++++-- .../observability/environment-variables.mdx | 6 +- 2 files changed, 312 insertions(+), 38 deletions(-) diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index 62aae5f15a4a..4c4d486ea95d 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -4710,25 +4710,118 @@ def _kvwarm_meta_init(self) -> dict: "points_fake_fallback": 0, "giant_kv_threshold": self._kvwarm_giant_threshold(), "giant_kv_repeats": self._kvwarm_giant_repeats(), + # Recurrent-state hand-off on hybrid layouts: the requested + # ``DYN_BENCH_KV_WARMUP_STATE`` and the mode the shadows + # actually used (None until a hybrid point is injected; + # ``mixed`` when auto settled differently across points). + "state_mode_requested": self._kvwarm_state_mode_requested(), + "state_mode": None, } self._kvwarm_meta = meta return meta def _kvwarm_state_layer_groups(self) -> list[str]: - """Names of KV-cache groups backed by recurrent state (Mamba/linear - attention). Their per-request state is updated in place, so a borrowed - shadow write would corrupt the chain; the warm-up must not run on them - until scratch state blocks exist.""" + """Names of recurrent-state KV-cache groups (Mamba/linear attention) + the warm-up cannot serve. Align mode keeps one running-state block per + request at a token-derived slot of a null-padded table, which a shadow + borrows (``_kvwarm_stage_state_shadow``); the other Mamba cache modes + lay the state out differently and stay unsupported.""" manager = getattr(self, "kv_cache_manager", None) config = getattr(manager, "kv_cache_config", None) groups = getattr(config, "kv_cache_groups", None) or [] names = [] for group in groups: - spec_name = type(getattr(group, "kv_cache_spec", None)).__name__ - if "Mamba" in spec_name: + spec = getattr(group, "kv_cache_spec", None) + spec_name = type(spec).__name__ + if ( + "Mamba" in spec_name + and getattr(spec, "mamba_cache_mode", None) != "align" + ): names.append(spec_name) return names + _KVWARM_STATE_MODES = ("auto", "copy", "inplace") + + def _kvwarm_state_mode_requested(self) -> str: + """``DYN_BENCH_KV_WARMUP_STATE``: how a shadow obtains its recurrent + state block on a hybrid layout. ``copy`` forks the chain's state block + (copy-on-write, needs a manager with ``_apply_cow``), ``inplace`` + hands the shadow the chain's own block, ``auto`` (default) copies when + the manager can and the pool holds the forks, else runs in place.""" + mode = os.environ.get("DYN_BENCH_KV_WARMUP_STATE", "auto").strip().lower() + if mode not in self._KVWARM_STATE_MODES: + raise ValueError( + "DYN_BENCH_KV_WARMUP_STATE must be one of " + f"{', '.join(self._KVWARM_STATE_MODES)}; got {mode!r}" + ) + return mode + + def _kvwarm_state_managers(self) -> list: + """Align-mode Mamba managers of the layout (empty for pure attention).""" + coordinator = getattr( + getattr(self, "kv_cache_manager", None), "coordinator", None + ) + return [ + mgr + for mgr in getattr(coordinator, "single_type_managers", ()) or () + if getattr(mgr, "mamba_cache_mode", None) == "align" + ] + + def _kvwarm_state_copy_supported(self) -> bool: + """Whether every state manager can fork a block (``_apply_cow``).""" + return all( + callable(getattr(mgr, "_apply_cow", None)) + for mgr in self._kvwarm_state_managers() + ) + + def _kvwarm_state_mode_floor(self) -> str | None: + """Least pool-hungry state mode an injection can settle on: the + requested mode when explicit, in-place under ``auto`` (its fallback); + None without state groups. The stage check bounds the shadows' pool + draw with it so a rung is failed only when even the fallback does not + fit.""" + if not self._kvwarm_state_managers(): + return None + requested = self._kvwarm_state_mode_requested() + return "inplace" if requested == "auto" else requested + + def _kvwarm_state_mode(self, context_lengths, headroom: int) -> str | None: + """Resolve the state mode of one injection: None without state groups, + the requested mode when explicit; under ``auto``, ``copy`` when every + state manager forks and the pool holds the forks on top of the tails + (``_kvwarm_shadow_pool_shortfall``), otherwise ``inplace``.""" + if not self._kvwarm_state_managers(): + return None + requested = self._kvwarm_state_mode_requested() + if requested != "auto": + return requested + if ( + self._kvwarm_state_copy_supported() + and self._kvwarm_shadow_pool_shortfall(context_lengths, headroom, "copy") + == 0 + ): + return "copy" + return "inplace" + + def _kvwarm_stamp_state(self, state_mode: str, depth: int) -> None: + """Record the state hand-off of the point being injected: the mode in + the KVWARM metadata and, per row, ``kvwarm_state_`` plus the + chain depth the borrowed state summarises (``kvwarm_state_depth=N``; + the recurrent layers see that depth, not the injected context).""" + meta = self._kvwarm_meta_init() + previous = meta.get("state_mode") + meta["state_mode"] = state_mode if previous in (None, state_mode) else "mixed" + point = getattr(self, "_bench_current_point", None) + if point is not None: + self._bench_current_point = replace( + point, + sample_reasons=[ + *point.sample_reasons, + f"kvwarm_state_{state_mode}", + f"kvwarm_state_depth={depth}", + ], + ) + def _kvwarm_seed_regime(self, point) -> str: """Row-level KV seed provenance for artifact consumers. @@ -4833,6 +4926,14 @@ def _kvwarm_warm_eligible(self) -> bool: reason = "prefix_caching_disabled" elif self._kvwarm_state_layer_groups(): reason = "hybrid_state_layers_unsupported" + elif ( + self._kvwarm_state_mode_requested() == "copy" + and not self._kvwarm_state_copy_supported() + ): + # An explicit copy request on a vLLM whose managers cannot + # fork a block (no ``_apply_cow``) is refused here, before + # any chain is built, instead of failing at injection. + reason = "state_copy_unsupported" else: reason = self._kvwarm_probe_content() eligible = reason is None @@ -5142,21 +5243,42 @@ def _kvwarm_prepare(self, mode: str) -> None: # ------- Warmup state machine (intercepts before phase dispatch) ------- def _kvwarm_shadow_tail_blocks(self, repeats: int) -> int: - """Worst-case private tail blocks one measurement shadow draws from - the pool on top of the chain prefix it shares (see - ``_kvwarm_register_shadow``): ``ceil((ctx + 1 + headroom) / bs) - - ctx // bs`` peaks at ``1 + ceil(headroom / bs)`` when ``ctx`` ends one - slot short of a block boundary; the headroom is the giant repeat - count (at least 2). Every KV-cache group draws its own tail.""" + """Worst-case private blocks one measurement shadow draws from the + pool on top of what it shares with its chain (see + ``_kvwarm_register_shadow``); the headroom is the giant repeat count + (at least 2) and every KV-cache group draws its own share. + + Attention groups: ``ceil((ctx + 1 + headroom) / bs) - ctx // bs`` + peaks at ``1 + ceil(headroom / bs)`` when ``ctx`` ends one slot short + of a block boundary. Recurrent-state groups: one fork of the chain's + state block in copy mode, a blank block for every state slot the + steady writes move into (``(ctx + headroom) // bs - (ctx - 1) // bs``, + at most ``ceil((headroom + 1) / bs)``) and the speculative scratch + blocks vLLM keeps after the state. The mode is settled per injection + (``_kvwarm_state_mode``), after the plan; the plan reserves the fork + whenever copy is possible so that ``auto`` can take it -- the price + is a slightly shallower chain when the pool is tight, never a shadow + the pool cannot serve. + """ coordinator = getattr( getattr(self, "kv_cache_manager", None), "coordinator", None ) - n_groups = max(1, len(getattr(coordinator, "single_type_managers", ()) or ())) - block_size = int( + managers = list(getattr(coordinator, "single_type_managers", ()) or ()) + default_bs = int( getattr(getattr(self, "cache_config", None), "block_size", 16) or 16 ) headroom = max(2, int(repeats)) - return n_groups * (1 + -(-headroom // block_size)) + fork = int(self._kvwarm_state_mode_requested() != "inplace") + total = 0 + for mgr in managers or [None]: + bs = int(getattr(mgr, "block_size", default_bs) or default_bs) + if getattr(mgr, "mamba_cache_mode", None) == "align": + spec = int(getattr(mgr, "num_speculative_blocks", 0) or 0) + can_fork = callable(getattr(mgr, "_apply_cow", None)) + total += fork * can_fork + -(-(headroom + 1) // bs) + spec + else: + total += 1 + -(-headroom // bs) + return total def _kvwarm_plan_covers(self, point) -> bool: """Plan-level coverage decision (independent of live chains): the shared @@ -5411,8 +5533,11 @@ def _kvwarm_stage_shadow_shortfall(self, batch: int | None) -> int: point this stage will serve, with its chains parked (0 when every point fits). Same per-point arithmetic as ``_kvwarm_inject_borrowed`` (``_kvwarm_shadow_pool_shortfall``), with the full repeat count as - headroom, so it bounds what injection will ask for.""" + headroom and the least pool-hungry state mode injection can fall + back to (``_kvwarm_state_mode_floor``), so it bounds what injection + will ask for.""" headroom = self._kvwarm_giant_repeats() + state_mode = self._kvwarm_state_mode_floor() worst = 0 for point in self._bench_grid: if ( @@ -5425,7 +5550,10 @@ def _kvwarm_stage_shadow_shortfall(self, batch: int | None) -> int: point.total_kv_read_tokens, point.batch_size ) injected = [max(1, ctx - 1) for ctx in ctxs] - worst = max(worst, self._kvwarm_shadow_pool_shortfall(injected, headroom)) + worst = max( + worst, + self._kvwarm_shadow_pool_shortfall(injected, headroom, state_mode), + ) return worst def _kvwarm_shed_chains(self) -> None: @@ -5471,23 +5599,31 @@ def _kvwarm_covers(self, point, injected_lengths) -> bool: ) def _kvwarm_register_shadow( - self, req_id: str, chain_id: str, ctx_len: int, headroom: int + self, + req_id: str, + chain_id: str, + ctx_len: int, + headroom: int, + state_mode: str | None = None, ) -> tuple[tuple[list[int], ...], list[int]]: """Register a measurement shadow with the KV-cache managers. - The shadow shares the chain's full prefix blocks (reference-counted, - never written) and owns private tail blocks for every position it - writes (the admission token plus ``headroom`` steady steps). The tail - is a copy-on-write fork of the chain's blocks when the manager offers - CoW; otherwise it is zero-filled (the few sub-block slots below - ``ctx_len`` then read zeros -- measurement-local, timing-neutral). + Attention groups: the shadow shares the chain's full prefix blocks + (reference-counted, never written) and owns private tail blocks for + every position it writes (the admission token plus ``headroom`` + steady steps). The tail is a copy-on-write fork of the chain's blocks + when the manager offers CoW; otherwise it is zero-filled (the few + sub-block slots below ``ctx_len`` then read zeros -- measurement- + local, timing-neutral). Recurrent-state groups (align-mode Mamba) + borrow the chain's single running-state block in ``state_mode`` + (``_kvwarm_state_mode``; see ``_kvwarm_stage_state_shadow``). Returns the shadow's block table per group and the block ids to zero. Registration is all-or-nothing across KV-cache groups. Every group's - geometry is checked and its private tail taken from the pool before + geometry is checked and its private blocks taken from the pool before any chain block is referenced or any table written, so a chain too shallow for a later group or a pool that cannot supply its tail - unwinds to a shadow that holds nothing: the tails already taken go + unwinds to a shadow that holds nothing: the blocks already taken go back to the pool (``free_blocks`` drops the single reference ``get_new_blocks`` gave them) and no group is left with a half-registered shadow or an over-referenced chain prefix. @@ -5498,9 +5634,23 @@ def _kvwarm_register_shadow( manager = self.kv_cache_manager block_pool = manager.block_pool managers = manager.coordinator.single_type_managers - staged: list[tuple[Any, int, list, list, list]] = [] + # Per manager: ``(mgr, plan)`` with the fresh pool blocks last in the + # plan, so a failure can return them whatever the manager kind. + staged: list[tuple[Any, tuple]] = [] try: for mgr in managers: + cache_mode = getattr(mgr, "mamba_cache_mode", None) + if cache_mode is not None: + if cache_mode != "align": + raise RuntimeError( + f"KVWARM: Mamba cache mode {cache_mode!r} unsupported " + f"for shadow {req_id}" + ) + plan = self._kvwarm_stage_state_shadow( + mgr, chain_id, req_id, ctx_len, headroom, state_mode + ) + staged.append((mgr, plan)) + continue chain_blocks = list(mgr.req_to_blocks[chain_id]) bs = int( getattr( @@ -5516,14 +5666,21 @@ def _kvwarm_register_shadow( ) tail_src = chain_blocks[n_shared:n_total] fresh = block_pool.get_new_blocks(len(tail_src)) - staged.append((mgr, n_shared, chain_blocks[:n_shared], tail_src, fresh)) + staged.append( + (mgr, (n_shared, chain_blocks[:n_shared], tail_src, fresh)) + ) except Exception: - for _, _, _, _, fresh in staged: - block_pool.free_blocks(fresh) + for _, plan in staged: + block_pool.free_blocks(plan[-1]) raise table: list[list[int]] = [] zero_ids: list[int] = [] - for mgr, n_shared, shared, tail_src, fresh in staged: + for mgr, plan in staged: + if getattr(mgr, "mamba_cache_mode", None) is not None: + blocks = self._kvwarm_commit_state_shadow(mgr, req_id, plan, state_mode) + table.append([b.block_id for b in blocks]) + continue + n_shared, shared, tail_src, fresh = plan block_pool.touch(shared) apply_cow = getattr(mgr, "_apply_cow", None) if callable(apply_cow): @@ -5548,10 +5705,103 @@ def _kvwarm_register_shadow( table.append([b.block_id for b in mgr.req_to_blocks[req_id]]) return tuple(table), zero_ids - def _kvwarm_shadow_pool_shortfall(self, context_lengths, headroom: int) -> int: - """Free blocks the pool lacks for the private tails of these shadows - (0 when they fit). Mirrors the per-group tail arithmetic of - ``_kvwarm_register_shadow`` so the check and the allocation agree.""" + def _kvwarm_stage_state_shadow( + self, + mgr, + chain_id: str, + req_id: str, + ctx_len: int, + headroom: int, + state_mode: str | None, + ) -> tuple: + """Plan a shadow's table for an align-mode recurrent-state group and + take its private blocks from the pool (the staging half of + ``_kvwarm_register_shadow``; ``_kvwarm_commit_state_shadow`` writes + the table). + + vLLM keeps one running-state block per request: the runner reads it + at table slot ``(num_computed - 1) // bs`` and writes the step's state + at ``(num_computed + scheduled - 1) // bs``, copying the state across + when the two differ (``preprocess_mamba``); every slot before the + state is a null block and the ``num_speculative_blocks`` slots after + it are draft-state scratch. The shadow's table follows that layout + for the span it writes: the chain's state block sits at + ``prev0 = (ctx_len - 1) // bs`` (in place, or forked by CoW so the + chain's own state is never written), a fresh blank block waits at + every later slot the ``headroom`` steady writes move into, and the + speculative scratch closes the table. The runner fills each blank + from the previous slot before reading it, so none is zeroed. + + The chain's state block is the last non-null block before its + speculative scratch: skipped slots are null and the chain's own + ``last_state_block_idx`` may already name one of them. The state + summarises the chain's full depth, not ``ctx_len`` (recorded per row + by ``_kvwarm_stamp_state``). + """ + if state_mode not in ("copy", "inplace"): + raise RuntimeError( + f"KVWARM: state mode {state_mode!r} unresolved for shadow {req_id}" + ) + if state_mode == "copy" and not callable(getattr(mgr, "_apply_cow", None)): + raise RuntimeError( + f"KVWARM: state copy for shadow {req_id} needs a manager with CoW" + ) + chain_blocks = list(mgr.req_to_blocks[chain_id]) + bs = int(mgr.block_size) + spec = int(getattr(mgr, "num_speculative_blocks", 0) or 0) + live = [index for index, block in enumerate(chain_blocks) if not block.is_null] + state_idx = live[-1] - spec if live else -1 + if state_idx < 0 or chain_blocks[state_idx].is_null: + raise RuntimeError( + f"KVWARM: chain {chain_id} holds no recurrent state block for " + f"shadow {req_id}" + ) + prev0 = (ctx_len - 1) // bs + last = (ctx_len + headroom) // bs + fresh = self.kv_cache_manager.block_pool.get_new_blocks( + int(state_mode == "copy") + (last - prev0) + spec + ) + return (chain_blocks[state_idx], prev0, last + 1 + spec, fresh) + + def _kvwarm_commit_state_shadow( + self, mgr, req_id: str, plan: tuple, state_mode: str | None + ) -> list: + """Write the table planned by ``_kvwarm_stage_state_shadow``. + + The chain's state block takes the shadow's reference + (``BlockPool.touch``); in place it stays the shadow's state slot and + the shadow's release returns that reference. In copy mode it is the + CoW source -- exactly like the attention tail, the reference is the + hit-ref the retained release after the copy consumes -- and + ``_apply_cow`` swaps the fork into the slot and queues the copy for + the admission step. ``num_cached_block`` covers the whole table so + vLLM never hashes the null slots or the state as a prefix. + """ + state, prev0, n_slots, fresh = plan + block_pool = self.kv_cache_manager.block_pool + blocks = [block_pool.null_block] * n_slots + blocks[prev0] = state + block_pool.touch([state]) + mgr.req_to_blocks[req_id] = blocks + if state_mode == "copy": + fork, fresh = fresh[0], fresh[1:] + mgr._apply_cow(req_id, prev0, state, fork) + for slot, block in enumerate(fresh, start=prev0 + 1): + blocks[slot] = block + cached = getattr(mgr, "num_cached_block", None) + if isinstance(cached, dict): + cached[req_id] = n_slots + return blocks + + def _kvwarm_shadow_pool_shortfall( + self, context_lengths, headroom: int, state_mode: str | None = None + ) -> int: + """Free blocks the pool lacks for the private blocks of these shadows + (0 when they fit). Mirrors the per-group arithmetic of + ``_kvwarm_register_shadow`` so the check and the allocation agree: + attention tails, and for recurrent-state groups the fork (``copy`` + only), the blank block of every state slot the steady writes move + into and the speculative scratch.""" manager = self.kv_cache_manager free_fn = getattr(manager.block_pool, "get_num_free_blocks", None) if not callable(free_fn): @@ -5561,6 +5811,13 @@ def _kvwarm_shadow_pool_shortfall(self, context_lengths, headroom: int) -> int: bs = int( getattr(mgr, "block_size", getattr(self.cache_config, "block_size", 16)) ) + if getattr(mgr, "mamba_cache_mode", None) == "align": + spec = int(getattr(mgr, "num_speculative_blocks", 0) or 0) + fork = int(state_mode == "copy") + for ctx_len in context_lengths: + blanks = (ctx_len + headroom) // bs - (ctx_len - 1) // bs + need += fork + blanks + spec + continue for ctx_len in context_lengths: need += -(-(ctx_len + 1 + headroom) // bs) - ctx_len // bs return max(0, need - int(free_fn())) @@ -5603,7 +5860,12 @@ def _kvwarm_inject_borrowed(self, context_lengths) -> "SchedulerOutput": # which is exactly what ``_kvwarm_point_need`` (1 + repeats) and the # plan margin reserve. headroom = max(1, int(getattr(self, "_bench_extra_steps_left", 1))) - shortfall = self._kvwarm_shadow_pool_shortfall(context_lengths, headroom) + # Recurrent-state hand-off (None on pure attention layouts), settled + # against this point's pool before the shortfall check reads it. + state_mode = self._kvwarm_state_mode(context_lengths, headroom) + shortfall = self._kvwarm_shadow_pool_shortfall( + context_lengths, headroom, state_mode + ) if shortfall > 0: # Not enough free blocks for the private tails: register nothing # and hand back an empty step; the caller's injection-shortfall @@ -5626,7 +5888,7 @@ def _kvwarm_inject_borrowed(self, context_lengths) -> "SchedulerOutput": req_id = f"__bench_{self._bench_seq}" self._bench_seq += 1 block_ids, req_zero_ids = self._kvwarm_register_shadow( - req_id, chain_id, ctx_len, headroom + req_id, chain_id, ctx_len, headroom, state_mode ) zero_ids.extend(req_zero_ids) prompt = list(chain_tokens[: ctx_len + 1]) @@ -5668,6 +5930,14 @@ def _kvwarm_inject_borrowed(self, context_lengths) -> "SchedulerOutput": # block left referenced fails the prefix-cache reset). self._kvwarm_take_cow_copies() raise + if state_mode is not None and context_lengths: + self._kvwarm_stamp_state( + state_mode, + max( + len(self._kvwarm_chain_prompts[self._kvwarm_chain_ids[index]]) + for index in range(len(context_lengths)) + ), + ) output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=CachedRequestData.make_empty(), diff --git a/docs/fern/pages/reference/observability/environment-variables.mdx b/docs/fern/pages/reference/observability/environment-variables.mdx index 364a608e6efc..03a86774bde2 100644 --- a/docs/fern/pages/reference/observability/environment-variables.mdx +++ b/docs/fern/pages/reference/observability/environment-variables.mdx @@ -275,7 +275,11 @@ See [Forward Pass Metrics Trace Reference](forward-pass-metrics-traces.mdx) for These variables tune the vLLM worker's self-benchmark collector (`--benchmark-mode`), which measures whole-model forward-pass latencies for the FPM performance database. They have no effect on serving. - KV warm-up master switch. When on, decode benchmark points on expert-parallel MoE models read KV produced by real prefill chains instead of self-injected synthetic KV. Set `off` to force the legacy synthetic-KV path. Ineligible configurations (dense models, no expert parallelism, prefix caching disabled, or a seeding dataset that is unavailable, empty, too shallow for the model length, or without a loadable tokenizer) skip the warm-up automatically and record the reason. Under attention data parallelism the decision is group-wide: if any rank is ineligible, every rank uses synthetic KV; if any rank's warm-up stage fails to build or cannot reserve its shadow blocks, every rank falls back to synthetic KV for that stage. + KV warm-up master switch. When on, decode benchmark points on expert-parallel MoE models read KV produced by real prefill chains instead of self-injected synthetic KV. Set `off` to force the legacy synthetic-KV path. Ineligible configurations (dense models, no expert parallelism, prefix caching disabled, or a seeding dataset that is unavailable, empty, too shallow for the model length, or without a loadable tokenizer) skip the warm-up automatically and record the reason. Under attention data parallelism the decision is group-wide: if any rank is ineligible, every rank uses synthetic KV; if any rank's warm-up stage fails to build or cannot reserve its shadow blocks, every rank falls back to synthetic KV for that stage. Hybrid layouts whose recurrent-state layers (Mamba / linear attention) use vLLM's `align` cache mode are supported; how a measured request obtains its state block is recorded as `state_mode` (see `DYN_BENCH_KV_WARMUP_STATE`). Other Mamba cache modes skip the warm-up with `hybrid_state_layers_unsupported`. + + + + Recurrent-state hand-off on hybrid layouts. Attention layers read the warm-up chain's KV by position; a Mamba / linear-attention layer keeps one running-state block per request, which the measured request borrows from its chain. `copy` forks that block (copy-on-write, one extra block per request and state group; needs a vLLM whose KV-cache managers support CoW), `inplace` hands the chain's own block to the measured request, `auto` copies when the manager can and the pool holds the forks, otherwise runs in place. The chosen mode is recorded as `state_mode` in the KV warm-up metadata and per row as `kvwarm_state_copy` / `kvwarm_state_inplace`, alongside `kvwarm_state_depth=N`, the chain depth the borrowed state summarises. An explicit `copy` on a vLLM without CoW skips the warm-up with `state_copy_unsupported`. From a5108ea94efdc7018b190de9d852d12bb4c7dea8 Mon Sep 17 00:00:00 2001 From: YZLi Date: Fri, 11 Sep 2026 01:39:30 +0800 Subject: [PATCH 2/3] test(vllm): cover the KV warm-up recurrent-state hand-off 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 Signed-off-by: YZLi --- .../tests/test_vllm_instrumented_scheduler.py | 541 +++++++++++++++++- 1 file changed, 528 insertions(+), 13 deletions(-) diff --git a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py index fee0c79edf4f..c11ecdfe1d51 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -37,6 +37,7 @@ def _isolate_synthetic_content_env(monkeypatch): monkeypatch.delenv("DYN_BENCH_PREFILL_REAL_SEED", raising=False) monkeypatch.delenv("DYN_BENCH_GIANT_KV_THRESHOLD", raising=False) monkeypatch.delenv("DYN_BENCH_GIANT_KV_REPEATS", raising=False) + monkeypatch.delenv("DYN_BENCH_KV_WARMUP_STATE", raising=False) # Module-level import: triggers real site-packages ``vllm`` to load before @@ -4810,7 +4811,18 @@ def _kvwarm_no_dataset_resolution(): raise AssertionError("the gate test must not resolve the real dataset") -def _kvwarm_gate_stub(*, state_groups=(), experts=8, ep=True, prefix=True): +def _kvwarm_gate_stub( + *, + state_groups=(), + experts=8, + ep=True, + prefix=True, + align_state=False, + managers=None, +): + """``align_state`` gives every Mamba spec vLLM's ``align`` cache mode; + ``managers`` optionally supplies the coordinator's single-type managers + (the gate consults them for the state-copy capability).""" stub = InstrumentedScheduler.__new__(InstrumentedScheduler) stub.vllm_config = SimpleNamespace( parallel_config=SimpleNamespace(enable_expert_parallel=ep), @@ -4819,12 +4831,17 @@ def _kvwarm_gate_stub(*, state_groups=(), experts=8, ep=True, prefix=True): ), ) stub.cache_config = SimpleNamespace(enable_prefix_caching=prefix) - groups = [ - SimpleNamespace(kv_cache_spec=type(name, (), {})()) for name in state_groups - ] + groups = [] + for name in state_groups: + attrs = {"mamba_cache_mode": "align"} if align_state and "Mamba" in name else {} + groups.append(SimpleNamespace(kv_cache_spec=type(name, (), attrs)())) stub.kv_cache_manager = SimpleNamespace( kv_cache_config=SimpleNamespace(kv_cache_groups=groups) ) + if managers is not None: + stub.kv_cache_manager.coordinator = SimpleNamespace( + single_type_managers=list(managers) + ) # Host-local inputs the gate proves: a pool that holds one chain at the # depth cap (max_model_len - 4 = 60 tokens) and a tokenizer. The dataset # itself is never resolved here (that would download it); tests that @@ -4863,12 +4880,49 @@ def _kvwarm_collection_bodies(count, length=80): def test_kvwarm_gate_rejects_recurrent_state_layers(monkeypatch): + """A Mamba group outside vLLM's ``align`` cache mode lays its state out + differently from the one-running-state-block table the shadow borrows.""" monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") stub = _kvwarm_gate_stub(state_groups=("FullAttentionSpec", "MambaSpec")) assert InstrumentedScheduler._kvwarm_warm_eligible(stub) is False assert stub._kvwarm_meta["skip_reason"] == "hybrid_state_layers_unsupported" +def test_kvwarm_gate_admits_align_mode_state_layers(monkeypatch): + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + stub = _kvwarm_gate_stub( + state_groups=("FullAttentionSpec", "MambaSpec"), align_state=True + ) + assert InstrumentedScheduler._kvwarm_state_layer_groups(stub) == [] + assert InstrumentedScheduler._kvwarm_warm_eligible(stub) is True + assert stub._kvwarm_meta["state_mode_requested"] == "auto" + assert stub._kvwarm_meta["state_mode"] is None + + +@pytest.mark.parametrize("requested", ["auto", "copy"]) +def test_kvwarm_gate_refuses_an_explicit_copy_without_cow(monkeypatch, requested): + """The frozen 0.25.1 image has no ``_apply_cow``: ``auto`` runs in place + there, an explicit ``copy`` is refused before any chain is built.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", requested) + pool = _FakePool() + stub = _kvwarm_gate_stub( + state_groups=("FullAttentionSpec", "MambaSpec"), + align_state=True, + managers=[ + _FakeManager([], cow=False), + _FakeMambaManager(pool, _FakeBlock(500), cow=False), + ], + ) + eligible = InstrumentedScheduler._kvwarm_warm_eligible(stub) + if requested == "copy": + assert eligible is False + assert stub._kvwarm_meta["skip_reason"] == "state_copy_unsupported" + else: + assert eligible is True + assert stub._kvwarm_meta["state_mode_requested"] == requested + + def test_kvwarm_gate_admits_pure_attention_moe_ep(monkeypatch): monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") stub = _kvwarm_gate_stub(state_groups=("FullAttentionSpec",)) @@ -5307,13 +5361,18 @@ def __init__(self, block_id): class _FakePool: """``BlockPool`` reference semantics: ``touch`` is +1, ``get_new_blocks`` hands out blocks at ref 1, ``free_blocks`` is -1 and a block joins the - free queue only when it reaches 0 (a free past 0 is a double free).""" + free queue only when it reaches 0 (a free past 0 is a double free). The + null block (id 0) pads skipped slots; vLLM does not maintain its count + and never queues it.""" def __init__(self, next_id=1000): self.next_id = next_id self.touched = [] self.freed = [] self.free_queue = [] + self.null_block = _FakeBlock(0) + self.null_block.is_null = True + self.null_block.ref_cnt = 0 def touch(self, blocks): self.touched.extend(blocks) @@ -5329,6 +5388,9 @@ def get_new_blocks(self, n): def free_blocks(self, ordered_blocks): for b in ordered_blocks: + if b.is_null: + b.ref_cnt -= 1 + continue assert b.ref_cnt > 0, f"double free of block {b.block_id}" b.ref_cnt -= 1 self.freed.append(b) @@ -5368,6 +5430,28 @@ def pop_blocks_for_free(self, req_id): return self.req_to_blocks.pop(req_id, []) +class _FakeMambaManager(_FakeManager): + """Align-mode ``MambaManager`` surface: one running-state block per request + at the end of a null-padded table (``[null, null, S]`` is a 192-token + chain at block size 64, its state block last), the align bookkeeping + ``pop_blocks_for_free`` drops like vLLM's, and ``_apply_cow`` as in + ``_FakeManager``.""" + + mamba_cache_mode = "align" + block_size = 64 + num_speculative_blocks = 0 + + def __init__(self, pool, state, cow=True): + super().__init__([pool.null_block, pool.null_block, state], cow=cow) + self.last_state_block_idx = {"chain": 1} + self._allocated_block_reqs = {"chain"} + + def pop_blocks_for_free(self, req_id): + self._allocated_block_reqs.discard(req_id) + self.last_state_block_idx.pop(req_id, None) + return super().pop_blocks_for_free(req_id) + + def _take_kv_cache_block_copies(manager): """``KVCacheManager.take_kv_cache_block_copies``: drain every manager's pending (source, cow) pairs into copy descriptors plus the retained @@ -5398,6 +5482,30 @@ def _shadow_stub(cow=True): return stub, mgr, pool, chain +def _hybrid_shadow_stub(cow=True): + """Two KV-cache groups over one pool: full attention (12 blocks of 16 = + 192 tokens) and an align-mode Mamba group whose 192-token chain holds its + running state in block 500 at slot 2 of ``[null, null, S]``.""" + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + pool = _FakePool() + chain = [_FakeBlock(i) for i in range(12)] + state = _FakeBlock(500) + attn = _FakeManager(chain, cow=cow) + mamba = _FakeMambaManager(pool, state, cow=cow) + manager = SimpleNamespace( + block_pool=pool, + coordinator=SimpleNamespace(single_type_managers=[attn, mamba]), + num_kv_cache_groups=2, + ) + if cow: + manager.take_kv_cache_block_copies = lambda: _take_kv_cache_block_copies( + manager + ) + stub.kv_cache_manager = manager + stub.cache_config = SimpleNamespace(block_size=16) + return stub, attn, mamba, pool, chain, state + + def test_kvwarm_shadow_registration_shares_prefix_and_forks_tail_with_cow(): stub, mgr, pool, chain = _shadow_stub(cow=True) # ctx=40 -> 2 shared full blocks, writes at 40..43 land in block 2 only @@ -5658,14 +5766,15 @@ def test_kvwarm_stage_shadow_shortfall_takes_the_rung_worst_case(monkeypatch): ) seen = [] - def shortfall(injected, headroom): - seen.append((list(injected), headroom)) + def shortfall(injected, headroom, state_mode): + seen.append((list(injected), headroom, state_mode)) return 5 if injected[0] == 46 else 0 stub._kvwarm_shadow_pool_shortfall = shortfall assert InstrumentedScheduler._kvwarm_stage_shadow_shortfall(stub, 4) == 5 - # Only this rung's covered points, with the full repeat count as headroom. - assert seen == [([63] * 4, 3), ([46] * 4, 3)] + # Only this rung's covered points, with the full repeat count as headroom + # and no state mode on a layout without recurrent-state groups. + assert seen == [([63] * 4, 3, None), ([46] * 4, 3, None)] def test_kvwarm_stage_sync_timeout_reaches_the_soft_deadline(): @@ -5974,9 +6083,347 @@ def test_kvwarm_shadow_lifecycle_returns_every_reference(cow): assert len(pool.free_queue) == len(chain) + 1 -def _kvwarm_injection_stub(chain_ids): +# --------------------------------------------------------------------------- +# Recurrent-state (align-mode Mamba) groups: the shadow borrows the chain's +# single running-state block instead of a positional prefix +# --------------------------------------------------------------------------- + + +def test_kvwarm_state_shadow_in_place_puts_the_chain_state_at_the_read_slot(): + """No CoW anywhere (the 0.25.1 image): the shadow's state table is null + up to ``prev0 = (ctx - 1) // 64``, where the chain's own state block sits + with the shadow's reference on it; nothing is zeroed for that group.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) + # ctx=100: attention shares 6 blocks and zero-fills one tail block; the + # state slot is 99 // 64 = 1 and writes 100..103 stay in that slot. + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, "inplace" + ) + assert table == ([*range(6), 1000], [0, 500]) + assert zero_ids == [1000] + assert mamba.req_to_blocks["shadow"][0] is pool.null_block + assert mamba.req_to_blocks["shadow"][1] is state + assert state.ref_cnt == 2 + assert [b.block_id for b in pool.touched] == [*range(6), 500] + assert mamba.cows == [] + # The whole table counts as cached: vLLM must not hash the null slots. + assert mamba.num_cached_block["shadow"] == 2 + assert attn.num_cached_block["shadow"] == 6 + # The chain's own table is untouched. + assert mamba.req_to_blocks["chain"] == [pool.null_block, pool.null_block, state] + + +def test_kvwarm_state_shadow_copy_forks_the_chain_state_like_an_attention_tail(): + """Copy mode: the state block takes the hit-ref first, the fork replaces + it in the read slot through ``_apply_cow`` (retention +1), and the + retained release after registration returns both to their owners.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=True) + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, "copy" + ) + fork = mamba.req_to_blocks["shadow"][1] + assert table == ([*range(6), 1000], [0, fork.block_id]) + assert fork.block_id == 1001 and zero_ids == [] + assert mamba.cows == [(500, 1001)] + assert (state.ref_cnt, fork.ref_cnt) == (2, 2) + assert [b.block_id for b in pool.touched] == [*range(6), 6, 500] + copies = InstrumentedScheduler._kvwarm_take_cow_copies(stub) + assert copies == [(6, 1000), (500, 1001)] + assert (state.ref_cnt, fork.ref_cnt) == (1, 1) + assert pool.free_queue == [] + + +@pytest.mark.parametrize( + ("ctx", "expected"), + [ + # 125 // 64 = 1, 129 // 64 = 2: the third steady write crosses. + (126, [0, 500, 1002]), + # 127 // 64 = 1, 128 // 64 = 2: the admission write itself crosses. + (128, [0, 500, 1001]), + # 63 // 64 = 0, 67 // 64 = 1: a slot-0 state moving into slot 1. + (64, [500, 1001]), + ], +) +def test_kvwarm_state_shadow_pre_places_a_blank_block_at_the_crossing(ctx, expected): + """Every state slot a write within the headroom moves into gets a fresh + private block at registration; the runner copies the previous slot's + state into it before reading, so the blank is not zeroed and the steady + steps allocate nothing.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", ctx, 3, "inplace" + ) + assert table[1] == expected + blank = mamba.req_to_blocks["shadow"][-1] + assert blank.ref_cnt == 1 and blank.block_id not in zero_ids + # The attention group's zero list is what it always was. + assert zero_ids == table[0][ctx // 16 :] + + +def test_kvwarm_state_shadow_keeps_the_speculative_scratch_after_the_state(): + """vLLM reads ``1 + num_speculative_blocks`` slots from the state slot on; + the shadow's table ends with that many private blocks.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) + mamba.num_speculative_blocks = 1 + mamba.req_to_blocks["chain"] = [pool.null_block, state, _FakeBlock(501)] + table, _ = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, "inplace" + ) + assert table[1] == [0, 500, 1001] + pool.get_num_free_blocks = lambda: 0 + assert ( + InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [100], 3, "inplace") + == 2 + ) + + +def test_kvwarm_state_shadow_finds_the_state_behind_a_nulled_last_state_idx(): + """The chain's ``last_state_block_idx`` may name a slot + ``remove_skipped_blocks`` has already nulled; the state is the last + non-null block, and a chain without one cannot serve a shadow.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) + mamba.req_to_blocks["chain"] = [pool.null_block, state, pool.null_block] + mamba.last_state_block_idx["chain"] = 0 + table, _ = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, "inplace" + ) + assert table[1] == [0, 500] + mamba.req_to_blocks["chain"] = [pool.null_block, pool.null_block] + with pytest.raises(RuntimeError, match="no recurrent state"): + InstrumentedScheduler._kvwarm_register_shadow( + stub, "other", "chain", 100, 3, "inplace" + ) + + +@pytest.mark.parametrize( + ("state_mode", "cow", "match"), + [ + (None, True, "unresolved"), + ("auto", True, "unresolved"), + ("copy", False, "needs a manager with CoW"), + ], +) +def test_kvwarm_state_shadow_rejects_an_unusable_mode_and_unwinds( + state_mode, cow, match +): + """The attention group is staged first; a state group that cannot be + served returns its tail and leaves no reference behind.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=cow) + with pytest.raises(RuntimeError, match=match): + InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, state_mode + ) + assert [b.block_id for b in pool.freed] == [1000] + assert pool.free_queue == pool.freed and pool.touched == [] + assert all(b.ref_cnt == 1 for b in [*chain, state]) + for mgr in (attn, mamba): + assert "shadow" not in mgr.req_to_blocks + assert mgr.take_pending_cow_copies() == [] + + +def test_kvwarm_state_shadow_rejects_a_non_align_mamba_group(): + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) + mamba.mamba_cache_mode = "all" + with pytest.raises(RuntimeError, match="unsupported"): + InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 100, 3, "inplace" + ) + assert pool.touched == [] and all(b.ref_cnt == 1 for b in chain) + + +def test_kvwarm_shadow_tail_blocks_count_the_state_group_per_manager(monkeypatch): + """Planning reserves, per shadow and state group, the fork (whenever copy + is possible and not ruled out), one blank for a crossing within the + headroom and the speculative scratch, on top of the attention tail.""" + attention = SimpleNamespace(block_size=16) + state_cow = SimpleNamespace( + block_size=64, + mamba_cache_mode="align", + num_speculative_blocks=0, + _apply_cow=lambda *args: None, + ) + state_plain = SimpleNamespace( + block_size=64, mamba_cache_mode="align", num_speculative_blocks=0 + ) + + def planner(managers): + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + stub.cache_config = SimpleNamespace(block_size=16) + stub.kv_cache_manager = SimpleNamespace( + coordinator=SimpleNamespace(single_type_managers=managers) + ) + return stub + + # attention 1 + ceil(3 / 16) = 2; state: fork 1 + ceil(4 / 64) = 2 + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks( + planner([attention, state_cow]), 3 + ) + == 4 + ) + # No fork without CoW, or when in-place is requested. + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks( + planner([attention, state_plain]), 3 + ) + == 3 + ) + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", "inplace") + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks( + planner([attention, state_cow]), 3 + ) + == 3 + ) + monkeypatch.delenv("DYN_BENCH_KV_WARMUP_STATE") + # A headroom past one state block can cross twice. + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks(planner([state_plain]), 70) + == 2 + ) + state_cow.num_speculative_blocks = 2 + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks(planner([state_cow]), 3) == 4 + ) + + +def test_kvwarm_shadow_pool_shortfall_adds_the_state_group_blocks(): + """Pre-check arithmetic per group and request: attention tails as before; + the state group draws the fork only in copy mode plus one blank per + crossed slot (ctx=126 crosses within a headroom of 3, ctx=100 does not).""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=True) + # attention: ctx=100 -> 7 - 6 = 1, ctx=126 -> 9 - 7 = 2; state (copy): + # ctx=100 -> 1 + 0, ctx=126 -> 1 + 1: six blocks in total. + pool.get_num_free_blocks = lambda: 6 + assert ( + InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [100, 126], 3, "copy") + == 0 + ) + pool.get_num_free_blocks = lambda: 5 + assert ( + InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [100, 126], 3, "copy") + == 1 + ) + assert ( + InstrumentedScheduler._kvwarm_shadow_pool_shortfall( + stub, [100, 126], 3, "inplace" + ) + == 0 + ) + pool.get_num_free_blocks = lambda: 3 + assert ( + InstrumentedScheduler._kvwarm_shadow_pool_shortfall( + stub, [100, 126], 3, "inplace" + ) + == 1 + ) + + +def test_kvwarm_state_mode_resolution(monkeypatch): + """``DYN_BENCH_KV_WARMUP_STATE`` parsing and the ``auto`` rule: copy when + every state manager forks and the pool holds the forks, else in place; + the stage check's floor is the fallback; pure attention has no mode.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=True) + pool.get_num_free_blocks = lambda: 10 + assert InstrumentedScheduler._kvwarm_state_mode_requested(stub) == "auto" + assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "copy" + assert InstrumentedScheduler._kvwarm_state_mode_floor(stub) == "inplace" + pool.get_num_free_blocks = lambda: 1 # the attention tail alone fits + assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "inplace" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", " COPY ") + assert InstrumentedScheduler._kvwarm_state_mode_requested(stub) == "copy" + assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "copy" + assert InstrumentedScheduler._kvwarm_state_mode_floor(stub) == "copy" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", "bogus") + with pytest.raises(ValueError, match="bogus"): + InstrumentedScheduler._kvwarm_state_mode_requested(stub) + monkeypatch.delenv("DYN_BENCH_KV_WARMUP_STATE") + plain, *_ = _hybrid_shadow_stub(cow=False) + plain.kv_cache_manager.block_pool.get_num_free_blocks = lambda: 10 + assert InstrumentedScheduler._kvwarm_state_mode(plain, [100], 3) == "inplace" + pure, *_ = _shadow_stub(cow=True) + assert InstrumentedScheduler._kvwarm_state_mode(pure, [100], 3) is None + assert InstrumentedScheduler._kvwarm_state_mode_floor(pure) is None + + +def test_kvwarm_stamp_state_marks_the_row_and_the_metadata(): + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + stub._bench_current_point = BenchmarkPoint( + point_type="decode", + total_kv_read_tokens=101, + batch_size=1, + sample_reasons=["kvwarm_real_kv"], + ) + InstrumentedScheduler._kvwarm_stamp_state(stub, "copy", 192) + assert stub._bench_current_point.sample_reasons == [ + "kvwarm_real_kv", + "kvwarm_state_copy", + "kvwarm_state_depth=192", + ] + assert stub._kvwarm_meta["state_mode"] == "copy" + InstrumentedScheduler._kvwarm_stamp_state(stub, "copy", 192) + assert stub._kvwarm_meta["state_mode"] == "copy" + InstrumentedScheduler._kvwarm_stamp_state(stub, "inplace", 192) + assert stub._kvwarm_meta["state_mode"] == "mixed" + + +@pytest.mark.parametrize("cow", [True, False]) +def test_kvwarm_state_shadow_lifecycle_returns_every_reference(cow): + """``test_kvwarm_shadow_lifecycle_returns_every_reference`` over a hybrid + layout, in copy mode (CoW) and in place (no CoW): after the shadow and + the chain are released through ``free_blocks(reversed( + pop_blocks_for_free(req_id)))`` per manager, every real block sits at + ref 0 exactly once, the null block was never queued, and the Mamba + manager keeps no align bookkeeping for either request.""" + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=cow) + mode = "copy" if cow else "inplace" + ctx, headroom = 126, 3 # crosses into a second state slot + table, _ = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", ctx, headroom, mode + ) + if cow: + # Two attention tail forks and the state fork ride the admission + # step; their retentions are released at once. + assert len(InstrumentedScheduler._kvwarm_take_cow_copies(stub)) == 3 + assert [b.ref_cnt for b in chain] == [2] * 7 + [1] * 5 + assert state.ref_cnt == (1 if cow else 2) + private = [ + b + for mgr in (attn, mamba) + for b in mgr.req_to_blocks["shadow"] + if not b.is_null and b not in chain and b is not state + ] + # Two attention tail blocks (forked or zero-filled), the crossing blank + # and, in copy mode, the state fork. + assert len(private) == (4 if cow else 3) + assert all(b.ref_cnt == 1 for b in private) + + for mgr in (attn, mamba): + pool.free_blocks(reversed(mgr.pop_blocks_for_free("shadow"))) + assert "shadow" not in mgr.num_cached_block + assert all(b.ref_cnt == 1 for b in [*chain, state]) + assert all(b.ref_cnt == 0 for b in private) + assert sorted(b.block_id for b in pool.free_queue) == sorted( + b.block_id for b in private + ) + assert pool.null_block not in pool.free_queue + + for mgr in (attn, mamba): + pool.free_blocks(reversed(mgr.pop_blocks_for_free("chain"))) + assert all(b.ref_cnt == 0 for b in [*chain, state]) + assert sorted(b.block_id for b in pool.free_queue) == sorted( + b.block_id for b in [*private, *chain, state] + ) + assert len(pool.free_queue) == len(private) + len(chain) + 1 + assert mamba.req_to_blocks == {} and mamba.last_state_block_idx == {} + assert mamba._allocated_block_reqs == set() + + +def _kvwarm_injection_stub(chain_ids, hybrid=False): """Everything ``_kvwarm_inject_borrowed`` reads, over the fake pool and a - CoW manager; every chain owns ten blocks (160 tokens) at ref 1, parked.""" + CoW manager; every chain owns ten blocks (160 tokens) at ref 1, parked. + ``hybrid`` adds an align-mode Mamba group (block 500 + 100 * index holds + each chain's state) and the point being injected.""" stub = InstrumentedScheduler.__new__(InstrumentedScheduler) chains = { chain_id: [_FakeBlock(100 * index + i) for i in range(10)] @@ -5986,10 +6433,28 @@ def _kvwarm_injection_stub(chain_ids): mgr.req_to_blocks = dict(chains) pool = _FakePool() pool.get_num_free_blocks = lambda: 10 + managers = [mgr] + if hybrid: + states = { + chain_id: _FakeBlock(500 + 100 * index) + for index, chain_id in enumerate(chain_ids) + } + mamba = _FakeMambaManager(pool, states[chain_ids[0]], cow=True) + mamba.req_to_blocks = { + chain_id: [pool.null_block, pool.null_block, state] + for chain_id, state in states.items() + } + managers.append(mamba) + stub._bench_current_point = BenchmarkPoint( + point_type="decode", + total_kv_read_tokens=101 * len(chain_ids), + batch_size=len(chain_ids), + sample_reasons=["kvwarm_real_kv"], + ) manager = SimpleNamespace( block_pool=pool, - coordinator=SimpleNamespace(single_type_managers=[mgr]), - num_kv_cache_groups=1, + coordinator=SimpleNamespace(single_type_managers=managers), + num_kv_cache_groups=len(managers), ) manager.take_kv_cache_block_copies = lambda: _take_kv_cache_block_copies(manager) stub.kv_cache_manager = manager @@ -6046,6 +6511,56 @@ def test_kvwarm_inject_borrowed_releases_cow_retentions_before_the_step_runs(): assert stub.requests["__bench_0"].status == RequestStatus.RUNNING +@pytest.mark.parametrize("requested", ["auto", "inplace"]) +def test_kvwarm_inject_borrowed_hands_the_state_over_and_stamps_the_point( + monkeypatch, requested +): + """On a hybrid layout the injection settles the state mode against this + point's pool, registers every group, rides the state fork's copy on the + admission step like the attention forks, and stamps the point with the + mode and the chain depth the borrowed state summarises.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", requested) + stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"], hybrid=True) + mamba = stub.kv_cache_manager.coordinator.single_type_managers[1] + state = mamba.req_to_blocks["chain-a"][2] + + output = InstrumentedScheduler._kvwarm_inject_borrowed(stub, [100]) + + shadow_state = mamba.req_to_blocks["__bench_0"] + assert shadow_state[0] is pool.null_block + if requested == "auto": + # Copy: the fork sits in the read slot, the copy rides the step. + fork = shadow_state[1] + assert fork is not state and fork.ref_cnt == 1 + assert output.kv_cache_block_copies == [(6, 1000), (state.block_id, 1001)] + assert output.scheduled_new_reqs[0].block_ids == ([*range(6), 1000], [0, 1001]) + assert state.ref_cnt == 1 + mode = "copy" + else: + assert shadow_state[1] is state and state.ref_cnt == 2 + assert output.kv_cache_block_copies == [(6, 1000)] + assert output.scheduled_new_reqs[0].block_ids == ([*range(6), 1000], [0, 500]) + mode = "inplace" + assert mamba.take_pending_cow_copies() == [] + assert stub._bench_current_point.sample_reasons == [ + "kvwarm_real_kv", + f"kvwarm_state_{mode}", + "kvwarm_state_depth=160", + ] + assert stub._kvwarm_meta["state_mode"] == mode + assert stub._kvwarm_meta["state_mode_requested"] == requested + + +def test_kvwarm_inject_borrowed_leaves_pure_attention_points_unstamped(): + stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"]) + stub._bench_current_point = BenchmarkPoint( + point_type="decode", total_kv_read_tokens=41, sample_reasons=["kvwarm_real_kv"] + ) + InstrumentedScheduler._kvwarm_inject_borrowed(stub, [40]) + assert stub._bench_current_point.sample_reasons == ["kvwarm_real_kv"] + assert getattr(stub, "_kvwarm_meta", None) is None + + def test_kvwarm_inject_borrowed_drops_queued_copies_when_a_later_shadow_fails(): """A failure while registering the second shadow leaves the first one's fork queued in the manager. The step it would have ridden on is never From c53d9ffedbdeb80bc29873661f44ab14c7cf98a9 Mon Sep 17 00:00:00 2001 From: YZLi Date: Fri, 11 Sep 2026 02:19:05 +0800 Subject: [PATCH 3/3] fix(vllm): address review findings on the hybrid state hand-off 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 --- .../src/dynamo/vllm/instrumented_scheduler.py | 219 +++++++++----- .../tests/test_vllm_instrumented_scheduler.py | 278 +++++++++++++++--- .../observability/environment-variables.mdx | 2 +- 3 files changed, 389 insertions(+), 110 deletions(-) diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index 4c4d486ea95d..1e456e7d6535 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -4710,11 +4710,12 @@ def _kvwarm_meta_init(self) -> dict: "points_fake_fallback": 0, "giant_kv_threshold": self._kvwarm_giant_threshold(), "giant_kv_repeats": self._kvwarm_giant_repeats(), - # Recurrent-state hand-off on hybrid layouts: the requested - # ``DYN_BENCH_KV_WARMUP_STATE`` and the mode the shadows - # actually used (None until a hybrid point is injected; - # ``mixed`` when auto settled differently across points). - "state_mode_requested": self._kvwarm_state_mode_requested(), + # Recurrent-state hand-off on hybrid layouts: the value of + # ``DYN_BENCH_KV_WARMUP_STATE`` as set (the gate validates it + # only where it matters, ``_kvwarm_state_gate_reason``) and + # the mode the shadows use (None until a hybrid point is + # injected; one mode per run, ``_kvwarm_state_mode``). + "state_mode_requested": self._kvwarm_state_mode_env(), "state_mode": None, } self._kvwarm_meta = meta @@ -4725,7 +4726,13 @@ def _kvwarm_state_layer_groups(self) -> list[str]: the warm-up cannot serve. Align mode keeps one running-state block per request at a token-derived slot of a null-padded table, which a shadow borrows (``_kvwarm_stage_state_shadow``); the other Mamba cache modes - lay the state out differently and stay unsupported.""" + lay the state out differently and stay unsupported. + + The cache mode is read from the spec, the key the managers are built + from (``_kvwarm_state_managers`` selects them by the same attribute), + so the gate and the registration agree whatever the spec is called; a + spec named after Mamba that carries no cache mode is refused too. + """ manager = getattr(self, "kv_cache_manager", None) config = getattr(manager, "kv_cache_config", None) groups = getattr(config, "kv_cache_groups", None) or [] @@ -4733,22 +4740,31 @@ def _kvwarm_state_layer_groups(self) -> list[str]: for group in groups: spec = getattr(group, "kv_cache_spec", None) spec_name = type(spec).__name__ - if ( - "Mamba" in spec_name - and getattr(spec, "mamba_cache_mode", None) != "align" - ): + mode = getattr(spec, "mamba_cache_mode", None) + if mode is None and "Mamba" not in spec_name: + continue + if mode != "align": names.append(spec_name) return names _KVWARM_STATE_MODES = ("auto", "copy", "inplace") + def _kvwarm_state_mode_env(self) -> str: + """``DYN_BENCH_KV_WARMUP_STATE`` as set (``auto`` when unset), not + validated: recorded in the metadata for every run, parsed only where + it matters (``_kvwarm_state_mode_requested``).""" + return os.environ.get("DYN_BENCH_KV_WARMUP_STATE", "auto") + def _kvwarm_state_mode_requested(self) -> str: """``DYN_BENCH_KV_WARMUP_STATE``: how a shadow obtains its recurrent state block on a hybrid layout. ``copy`` forks the chain's state block (copy-on-write, needs a manager with ``_apply_cow``), ``inplace`` hands the shadow the chain's own block, ``auto`` (default) copies when - the manager can and the pool holds the forks, else runs in place.""" - mode = os.environ.get("DYN_BENCH_KV_WARMUP_STATE", "auto").strip().lower() + the managers can, else runs in place (``_kvwarm_state_mode``). Only + layouts with recurrent-state groups consult it; the gate turns an + invalid value into a skip there (``_kvwarm_state_gate_reason``), so + a typo cannot fail a run the variable does not concern.""" + mode = self._kvwarm_state_mode_env().strip().lower() if mode not in self._KVWARM_STATE_MODES: raise ValueError( "DYN_BENCH_KV_WARMUP_STATE must be one of " @@ -4774,34 +4790,80 @@ def _kvwarm_state_copy_supported(self) -> bool: for mgr in self._kvwarm_state_managers() ) - def _kvwarm_state_mode_floor(self) -> str | None: - """Least pool-hungry state mode an injection can settle on: the - requested mode when explicit, in-place under ``auto`` (its fallback); - None without state groups. The stage check bounds the shadows' pool - draw with it so a rung is failed only when even the fallback does not - fit.""" - if not self._kvwarm_state_managers(): - return None - requested = self._kvwarm_state_mode_requested() - return "inplace" if requested == "auto" else requested + def _kvwarm_state_inplace_ok(self, mgr) -> bool: + """Whether a shadow may hold the chain's own state block in ``mgr``'s + table. + + A vLLM whose Mamba manager registers partial-tail prefix entries + (``_cache_partial_tail_block``) hashes the block at a request's state + slot under the request's own key on every output update once the + group's block size exceeds the hash block size (a hybrid layout's + default). Done to the chain's block, that stacks the shadow's key on + the chain's, cuts the shadow's ``num_cached_block`` back to the state + slot and, at the first steady write that reaches the next slot, fails + the promotion of the chain's entry to a full one -- so such builds + fork instead (``copy``). A manager without the method never hashes a + slot the shadow marks cached (``_kvwarm_commit_state_shadow``). + """ + if not callable(getattr(mgr, "_cache_partial_tail_block", None)): + return True + block_pool = getattr( + getattr(self, "kv_cache_manager", None), "block_pool", None + ) + return int(mgr.block_size) == getattr(block_pool, "hash_block_size", None) + + def _kvwarm_state_inplace_supported(self) -> bool: + """Whether every state manager admits in-place state + (``_kvwarm_state_inplace_ok``).""" + return all( + self._kvwarm_state_inplace_ok(mgr) for mgr in self._kvwarm_state_managers() + ) - def _kvwarm_state_mode(self, context_lengths, headroom: int) -> str | None: - """Resolve the state mode of one injection: None without state groups, - the requested mode when explicit; under ``auto``, ``copy`` when every - state manager forks and the pool holds the forks on top of the tails - (``_kvwarm_shadow_pool_shortfall``), otherwise ``inplace``.""" + def _kvwarm_state_mode(self) -> str | None: + """The state mode every shadow of the run uses: None without state + groups, the requested mode when explicit; under ``auto``, ``copy`` + when every state manager forks, else ``inplace`` when every manager + admits it, else None (the gate refuses the layout, + ``_kvwarm_state_gate_reason``). + + The mode follows from the layout and the vLLM build alone, never from + the pool: every attention-DP rank resolves the same mode (a + rank-local choice would fork the point digest the ranks compare), and + the plan, the stage check and the injection reserve the blocks of one + and the same mode. A pool too small for it skips the point + (``_kvwarm_shadow_pool_shortfall``) instead of switching modes. + """ if not self._kvwarm_state_managers(): return None requested = self._kvwarm_state_mode_requested() if requested != "auto": return requested - if ( - self._kvwarm_state_copy_supported() - and self._kvwarm_shadow_pool_shortfall(context_lengths, headroom, "copy") - == 0 - ): + if self._kvwarm_state_copy_supported(): return "copy" - return "inplace" + if self._kvwarm_state_inplace_supported(): + return "inplace" + return None + + def _kvwarm_state_gate_reason(self) -> str | None: + """Skip reason of the recurrent-state hand-off, None when the layout + has no align-mode state groups or serves the requested mode: + ``state_mode_invalid`` for an unparseable ``DYN_BENCH_KV_WARMUP_STATE``, + ``state_copy_unsupported`` / ``state_inplace_unsupported`` for an + explicit mode the managers cannot serve, ``state_hand_off_unsupported`` + when ``auto`` finds neither. Decided before any chain is built.""" + if not self._kvwarm_state_managers(): + return None + try: + requested = self._kvwarm_state_mode_requested() + except ValueError: + return "state_mode_invalid" + supported = { + "copy": self._kvwarm_state_copy_supported(), + "inplace": self._kvwarm_state_inplace_supported(), + } + if requested == "auto": + return None if any(supported.values()) else "state_hand_off_unsupported" + return None if supported[requested] else f"state_{requested}_unsupported" def _kvwarm_stamp_state(self, state_mode: str, depth: int) -> None: """Record the state hand-off of the point being injected: the mode in @@ -4809,8 +4871,7 @@ def _kvwarm_stamp_state(self, state_mode: str, depth: int) -> None: chain depth the borrowed state summarises (``kvwarm_state_depth=N``; the recurrent layers see that depth, not the injected context).""" meta = self._kvwarm_meta_init() - previous = meta.get("state_mode") - meta["state_mode"] = state_mode if previous in (None, state_mode) else "mixed" + meta["state_mode"] = state_mode point = getattr(self, "_bench_current_point", None) if point is not None: self._bench_current_point = replace( @@ -4926,14 +4987,11 @@ def _kvwarm_warm_eligible(self) -> bool: reason = "prefix_caching_disabled" elif self._kvwarm_state_layer_groups(): reason = "hybrid_state_layers_unsupported" - elif ( - self._kvwarm_state_mode_requested() == "copy" - and not self._kvwarm_state_copy_supported() - ): - # An explicit copy request on a vLLM whose managers cannot - # fork a block (no ``_apply_cow``) is refused here, before - # any chain is built, instead of failing at injection. - reason = "state_copy_unsupported" + elif (state_reason := self._kvwarm_state_gate_reason()) is not None: + # A state hand-off this vLLM's managers cannot serve (or an + # unparseable mode) is refused here, before any chain is + # built, instead of failing at injection. + reason = state_reason else: reason = self._kvwarm_probe_content() eligible = reason is None @@ -5254,11 +5312,9 @@ def _kvwarm_shadow_tail_blocks(self, repeats: int) -> int: state block in copy mode, a blank block for every state slot the steady writes move into (``(ctx + headroom) // bs - (ctx - 1) // bs``, at most ``ceil((headroom + 1) / bs)``) and the speculative scratch - blocks vLLM keeps after the state. The mode is settled per injection - (``_kvwarm_state_mode``), after the plan; the plan reserves the fork - whenever copy is possible so that ``auto`` can take it -- the price - is a slightly shallower chain when the pool is tight, never a shadow - the pool cannot serve. + blocks vLLM keeps after the state. The mode is the run's + (``_kvwarm_state_mode``, fixed by the layout), so the plan, the + stage check and the injection reserve the same blocks. """ coordinator = getattr( getattr(self, "kv_cache_manager", None), "coordinator", None @@ -5268,14 +5324,13 @@ def _kvwarm_shadow_tail_blocks(self, repeats: int) -> int: getattr(getattr(self, "cache_config", None), "block_size", 16) or 16 ) headroom = max(2, int(repeats)) - fork = int(self._kvwarm_state_mode_requested() != "inplace") + fork = int(self._kvwarm_state_mode() == "copy") total = 0 for mgr in managers or [None]: bs = int(getattr(mgr, "block_size", default_bs) or default_bs) if getattr(mgr, "mamba_cache_mode", None) == "align": spec = int(getattr(mgr, "num_speculative_blocks", 0) or 0) - can_fork = callable(getattr(mgr, "_apply_cow", None)) - total += fork * can_fork + -(-(headroom + 1) // bs) + spec + total += fork + -(-(headroom + 1) // bs) + spec else: total += 1 + -(-headroom // bs) return total @@ -5533,11 +5588,10 @@ def _kvwarm_stage_shadow_shortfall(self, batch: int | None) -> int: point this stage will serve, with its chains parked (0 when every point fits). Same per-point arithmetic as ``_kvwarm_inject_borrowed`` (``_kvwarm_shadow_pool_shortfall``), with the full repeat count as - headroom and the least pool-hungry state mode injection can fall - back to (``_kvwarm_state_mode_floor``), so it bounds what injection - will ask for.""" + headroom and the run's state mode (``_kvwarm_state_mode``), so it + bounds what injection will ask for.""" headroom = self._kvwarm_giant_repeats() - state_mode = self._kvwarm_state_mode_floor() + state_mode = self._kvwarm_state_mode() worst = 0 for point in self._bench_grid: if ( @@ -5732,11 +5786,19 @@ def _kvwarm_stage_state_shadow( speculative scratch closes the table. The runner fills each blank from the previous slot before reading it, so none is zeroed. - The chain's state block is the last non-null block before its - speculative scratch: skipped slots are null and the chain's own - ``last_state_block_idx`` may already name one of them. The state + The chain's state sits at ``(depth - 1) // bs`` for its computed + depth: the slot its last step wrote and the one the runner would read + next, taken from the request rather than from the table's shape. + Lower slots are null or, after a chunked prefill, the still-live + earlier state block the chain's own ``last_state_block_idx`` names + until its next allocation nulls it (a parked chain never allocates + again). Without draft tokens every allocation is exact and the state + is the table's last block, which is checked; draft tokens can leave + the table over-allocated past the state and its scratch (a rejected + draft's blocks stay), so there the tail is only reported. The state summarises the chain's full depth, not ``ctx_len`` (recorded per row - by ``_kvwarm_stamp_state``). + by ``_kvwarm_stamp_state``). In place is refused on a manager that + hashes partial tails (``_kvwarm_state_inplace_ok``). """ if state_mode not in ("copy", "inplace"): raise RuntimeError( @@ -5746,16 +5808,37 @@ def _kvwarm_stage_state_shadow( raise RuntimeError( f"KVWARM: state copy for shadow {req_id} needs a manager with CoW" ) + if state_mode == "inplace" and not self._kvwarm_state_inplace_ok(mgr): + raise RuntimeError( + f"KVWARM: in-place state for shadow {req_id} unsupported: the " + "manager hashes partial tails at the state slot" + ) chain_blocks = list(mgr.req_to_blocks[chain_id]) bs = int(mgr.block_size) spec = int(getattr(mgr, "num_speculative_blocks", 0) or 0) - live = [index for index, block in enumerate(chain_blocks) if not block.is_null] - state_idx = live[-1] - spec if live else -1 - if state_idx < 0 or chain_blocks[state_idx].is_null: + state_idx = (int(self.requests[chain_id].num_computed_tokens) - 1) // bs + if not 0 <= state_idx < len(chain_blocks) or chain_blocks[state_idx].is_null: raise RuntimeError( f"KVWARM: chain {chain_id} holds no recurrent state block for " f"shadow {req_id}" ) + live = [index for index, block in enumerate(chain_blocks) if not block.is_null] + extra = live[-1] - spec - state_idx + if extra < 0 or (extra and spec == 0): + raise RuntimeError( + f"KVWARM: chain {chain_id} table ends {live[-1] - state_idx} " + f"block(s) past its state slot {state_idx}, not the {spec} " + "speculative block(s) expected" + ) + if extra: + logger.warning( + "KVWARM: chain %s table ends %d block(s) past its state slot %d " + "and %d speculative block(s); trusting the state slot", + chain_id, + extra, + state_idx, + spec, + ) prev0 = (ctx_len - 1) // bs last = (ctx_len + headroom) // bs fresh = self.kv_cache_manager.block_pool.get_new_blocks( @@ -5775,7 +5858,11 @@ def _kvwarm_commit_state_shadow( hit-ref the retained release after the copy consumes -- and ``_apply_cow`` swaps the fork into the slot and queues the copy for the admission step. ``num_cached_block`` covers the whole table so - vLLM never hashes the null slots or the state as a prefix. + vLLM never hashes the null slots or the state as a prefix; a manager + that hashes partial tails still registers the block at the state slot + under the shadow's key on the admission step's output update, which + is why that slot holds a private fork there + (``_kvwarm_state_inplace_ok``). """ state, prev0, n_slots, fresh = plan block_pool = self.kv_cache_manager.block_pool @@ -5860,9 +5947,9 @@ def _kvwarm_inject_borrowed(self, context_lengths) -> "SchedulerOutput": # which is exactly what ``_kvwarm_point_need`` (1 + repeats) and the # plan margin reserve. headroom = max(1, int(getattr(self, "_bench_extra_steps_left", 1))) - # Recurrent-state hand-off (None on pure attention layouts), settled - # against this point's pool before the shortfall check reads it. - state_mode = self._kvwarm_state_mode(context_lengths, headroom) + # Recurrent-state hand-off (None on pure attention layouts): the run's + # mode, the same on every rank and in every pool check. + state_mode = self._kvwarm_state_mode() shortfall = self._kvwarm_shadow_pool_shortfall( context_lengths, headroom, state_mode ) diff --git a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py index c11ecdfe1d51..3fb2dbaa5680 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -4920,9 +4920,92 @@ def test_kvwarm_gate_refuses_an_explicit_copy_without_cow(monkeypatch, requested assert stub._kvwarm_meta["skip_reason"] == "state_copy_unsupported" else: assert eligible is True + assert InstrumentedScheduler._kvwarm_state_mode(stub) == "inplace" assert stub._kvwarm_meta["state_mode_requested"] == requested +@pytest.mark.parametrize("requested", ["auto", "inplace"]) +def test_kvwarm_gate_refuses_in_place_state_where_partial_tails_are_hashed( + monkeypatch, requested +): + """A vLLM whose Mamba manager hashes partial tails (block size above the + hash block size) would register the shadow's state slot under the + shadow's key on every output update; with the chain's own block there + that corrupts the chain's prefix entry. ``auto`` forks on such a build, + an explicit ``inplace`` is refused before any chain is built.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", requested) + pool = _FakePool() + stub = _kvwarm_gate_stub( + state_groups=("FullAttentionSpec", "MambaSpec"), + align_state=True, + managers=[ + _FakeManager([], cow=True), + _FakeMambaManager(pool, _FakeBlock(500), cow=True, partial_tail=True), + ], + ) + stub.kv_cache_manager.block_pool = pool + eligible = InstrumentedScheduler._kvwarm_warm_eligible(stub) + if requested == "inplace": + assert eligible is False + assert stub._kvwarm_meta["skip_reason"] == "state_inplace_unsupported" + else: + assert eligible is True + assert InstrumentedScheduler._kvwarm_state_mode(stub) == "copy" + # A state block as large as the hash block is never partially hashed. + pool.hash_block_size = 64 + assert InstrumentedScheduler._kvwarm_state_inplace_supported(stub) is True + + +def test_kvwarm_gate_keys_state_groups_on_the_cache_mode_not_the_name(monkeypatch): + """The gate and the registration must refuse the same groups: a spec that + carries a non-align ``mamba_cache_mode`` is refused whatever it is + called, and a Mamba-named spec without a cache mode is refused too.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + stub = _kvwarm_gate_stub(state_groups=("FullAttentionSpec",)) + groups = stub.kv_cache_manager.kv_cache_config.kv_cache_groups + groups.append( + SimpleNamespace( + kv_cache_spec=type("LinearStateSpec", (), {"mamba_cache_mode": "all"})() + ) + ) + assert InstrumentedScheduler._kvwarm_state_layer_groups(stub) == ["LinearStateSpec"] + assert InstrumentedScheduler._kvwarm_warm_eligible(stub) is False + assert stub._kvwarm_meta["skip_reason"] == "hybrid_state_layers_unsupported" + groups[-1] = SimpleNamespace(kv_cache_spec=type("MambaSpec", (), {})()) + assert InstrumentedScheduler._kvwarm_state_layer_groups(stub) == ["MambaSpec"] + groups[-1] = SimpleNamespace( + kv_cache_spec=type("LinearStateSpec", (), {"mamba_cache_mode": "align"})() + ) + assert InstrumentedScheduler._kvwarm_state_layer_groups(stub) == [] + + +def test_kvwarm_gate_validates_the_state_mode_only_where_it_matters(monkeypatch): + """``DYN_BENCH_KV_WARMUP_STATE`` concerns hybrid layouts with the warm-up + on; anywhere else an unparseable value is recorded and ignored rather + than failing a run it has no bearing on.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", "bogus") + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "off") + off = _kvwarm_gate_stub(state_groups=("FullAttentionSpec", "MambaSpec")) + assert InstrumentedScheduler._kvwarm_warm_eligible(off) is False + assert off._kvwarm_meta["skip_reason"] == "flag_off" + assert off._kvwarm_meta["state_mode_requested"] == "bogus" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + pure = _kvwarm_gate_stub(state_groups=("FullAttentionSpec",)) + assert InstrumentedScheduler._kvwarm_warm_eligible(pure) is True + assert InstrumentedScheduler._kvwarm_state_mode(pure) is None + pool = _FakePool() + hybrid = _kvwarm_gate_stub( + state_groups=("FullAttentionSpec", "MambaSpec"), + align_state=True, + managers=[_FakeManager([], cow=True), _FakeMambaManager(pool, _FakeBlock(500))], + ) + assert InstrumentedScheduler._kvwarm_warm_eligible(hybrid) is False + assert hybrid._kvwarm_meta["skip_reason"] == "state_mode_invalid" + with pytest.raises(ValueError, match="bogus"): + InstrumentedScheduler._kvwarm_state_mode_requested(hybrid) + + def test_kvwarm_gate_admits_pure_attention_moe_ep(monkeypatch): monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") stub = _kvwarm_gate_stub(state_groups=("FullAttentionSpec",)) @@ -5363,7 +5446,8 @@ class _FakePool: hands out blocks at ref 1, ``free_blocks`` is -1 and a block joins the free queue only when it reaches 0 (a free past 0 is a double free). The null block (id 0) pads skipped slots; vLLM does not maintain its count - and never queues it.""" + and never queues it. ``hash_block_size`` is the prefix-hash granularity + of a hybrid layout: the attention block size, below the state block.""" def __init__(self, next_id=1000): self.next_id = next_id @@ -5373,6 +5457,7 @@ def __init__(self, next_id=1000): self.null_block = _FakeBlock(0) self.null_block.is_null = True self.null_block.ref_cnt = 0 + self.hash_block_size = 16 def touch(self, blocks): self.touched.extend(blocks) @@ -5435,16 +5520,21 @@ class _FakeMambaManager(_FakeManager): at the end of a null-padded table (``[null, null, S]`` is a 192-token chain at block size 64, its state block last), the align bookkeeping ``pop_blocks_for_free`` drops like vLLM's, and ``_apply_cow`` as in - ``_FakeManager``.""" + ``_FakeManager``. ``partial_tail`` marks a vLLM whose manager registers + partial-tail prefix entries (``_cache_partial_tail_block``): the shadow's + state slot gets hashed on every output update there, so the scheduler + must not put the chain's own block in it.""" mamba_cache_mode = "align" block_size = 64 num_speculative_blocks = 0 - def __init__(self, pool, state, cow=True): + def __init__(self, pool, state, cow=True, partial_tail=False): super().__init__([pool.null_block, pool.null_block, state], cow=cow) self.last_state_block_idx = {"chain": 1} self._allocated_block_reqs = {"chain"} + if partial_tail: + self._cache_partial_tail_block = lambda request, num_tokens: None def pop_blocks_for_free(self, req_id): self._allocated_block_reqs.discard(req_id) @@ -5482,16 +5572,17 @@ def _shadow_stub(cow=True): return stub, mgr, pool, chain -def _hybrid_shadow_stub(cow=True): +def _hybrid_shadow_stub(cow=True, partial_tail=False): """Two KV-cache groups over one pool: full attention (12 blocks of 16 = - 192 tokens) and an align-mode Mamba group whose 192-token chain holds its - running state in block 500 at slot 2 of ``[null, null, S]``.""" + 192 tokens) and an align-mode Mamba group whose 192-token chain (parked + with every token computed) holds its running state in block 500 at slot + 2 of ``[null, null, S]``.""" stub = InstrumentedScheduler.__new__(InstrumentedScheduler) pool = _FakePool() chain = [_FakeBlock(i) for i in range(12)] state = _FakeBlock(500) attn = _FakeManager(chain, cow=cow) - mamba = _FakeMambaManager(pool, state, cow=cow) + mamba = _FakeMambaManager(pool, state, cow=cow, partial_tail=partial_tail) manager = SimpleNamespace( block_pool=pool, coordinator=SimpleNamespace(single_type_managers=[attn, mamba]), @@ -5503,6 +5594,7 @@ def _hybrid_shadow_stub(cow=True): ) stub.kv_cache_manager = manager stub.cache_config = SimpleNamespace(block_size=16) + stub.requests = {"chain": SimpleNamespace(num_computed_tokens=192)} return stub, attn, mamba, pool, chain, state @@ -6165,6 +6257,8 @@ def test_kvwarm_state_shadow_keeps_the_speculative_scratch_after_the_state(): the shadow's table ends with that many private blocks.""" stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) mamba.num_speculative_blocks = 1 + # A 128-token chain: its state at slot 1, one speculative block after it. + stub.requests["chain"].num_computed_tokens = 128 mamba.req_to_blocks["chain"] = [pool.null_block, state, _FakeBlock(501)] table, _ = InstrumentedScheduler._kvwarm_register_shadow( stub, "shadow", "chain", 100, 3, "inplace" @@ -6177,38 +6271,82 @@ def test_kvwarm_state_shadow_keeps_the_speculative_scratch_after_the_state(): ) -def test_kvwarm_state_shadow_finds_the_state_behind_a_nulled_last_state_idx(): - """The chain's ``last_state_block_idx`` may name a slot - ``remove_skipped_blocks`` has already nulled; the state is the last - non-null block, and a chain without one cannot serve a shadow.""" +def test_kvwarm_state_shadow_reads_the_state_slot_from_the_chain_depth(caplog): + """The state slot is ``(computed depth - 1) // bs`` of the chain request, + not a property of the table's shape: a chunked prefill leaves the earlier + state block live below it (the chain's ``last_state_block_idx`` still + names it), ``remove_skipped_blocks`` may have nulled it instead, and + draft tokens can over-allocate the table past the state and its + speculative scratch. Without drafts the state must be the last block.""" stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=False) - mamba.req_to_blocks["chain"] = [pool.null_block, state, pool.null_block] + stub.requests["chain"].num_computed_tokens = 128 + earlier = _FakeBlock(499) + mamba.req_to_blocks["chain"] = [earlier, state] mamba.last_state_block_idx["chain"] = 0 table, _ = InstrumentedScheduler._kvwarm_register_shadow( stub, "shadow", "chain", 100, 3, "inplace" ) assert table[1] == [0, 500] + assert (earlier.ref_cnt, state.ref_cnt) == (1, 2) + mamba.req_to_blocks["chain"] = [pool.null_block, state] + table, _ = InstrumentedScheduler._kvwarm_register_shadow( + stub, "nulled", "chain", 100, 3, "inplace" + ) + assert table[1] == [0, 500] + # Over-allocated by a rejected draft: the state slot is trusted, the + # extra block only reported. + mamba.num_speculative_blocks = 1 + mamba.req_to_blocks["chain"] = [ + pool.null_block, + state, + _FakeBlock(501), + _FakeBlock(502), + ] + with caplog.at_level("WARNING"): + table, _ = InstrumentedScheduler._kvwarm_register_shadow( + stub, "drafted", "chain", 100, 3, "inplace" + ) + assert table[1] == [0, 500, 1003] + assert "past its state slot 1" in caplog.text + # Without drafts a block past the state is a layout this code does not + # model; a chain whose state slot is null or beyond its table holds no + # state at all. + mamba.num_speculative_blocks = 0 + mamba.req_to_blocks["chain"] = [pool.null_block, state, _FakeBlock(503)] + with pytest.raises(RuntimeError, match="past its state slot"): + InstrumentedScheduler._kvwarm_register_shadow( + stub, "long", "chain", 100, 3, "inplace" + ) mamba.req_to_blocks["chain"] = [pool.null_block, pool.null_block] with pytest.raises(RuntimeError, match="no recurrent state"): InstrumentedScheduler._kvwarm_register_shadow( stub, "other", "chain", 100, 3, "inplace" ) + mamba.req_to_blocks["chain"] = [pool.null_block, state] + stub.requests["chain"].num_computed_tokens = 192 + with pytest.raises(RuntimeError, match="no recurrent state"): + InstrumentedScheduler._kvwarm_register_shadow( + stub, "short", "chain", 100, 3, "inplace" + ) @pytest.mark.parametrize( - ("state_mode", "cow", "match"), + ("state_mode", "cow", "partial_tail", "match"), [ - (None, True, "unresolved"), - ("auto", True, "unresolved"), - ("copy", False, "needs a manager with CoW"), + (None, True, False, "unresolved"), + ("auto", True, False, "unresolved"), + ("copy", False, False, "needs a manager with CoW"), + ("inplace", True, True, "hashes partial tails"), ], ) def test_kvwarm_state_shadow_rejects_an_unusable_mode_and_unwinds( - state_mode, cow, match + state_mode, cow, partial_tail, match ): """The attention group is staged first; a state group that cannot be served returns its tail and leaves no reference behind.""" - stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=cow) + stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub( + cow=cow, partial_tail=partial_tail + ) with pytest.raises(RuntimeError, match=match): InstrumentedScheduler._kvwarm_register_shadow( stub, "shadow", "chain", 100, 3, state_mode @@ -6232,9 +6370,9 @@ def test_kvwarm_state_shadow_rejects_a_non_align_mamba_group(): def test_kvwarm_shadow_tail_blocks_count_the_state_group_per_manager(monkeypatch): - """Planning reserves, per shadow and state group, the fork (whenever copy - is possible and not ruled out), one blank for a crossing within the - headroom and the speculative scratch, on top of the attention tail.""" + """Planning reserves, per shadow and state group, the fork (when the + run's state mode is copy), one blank for a crossing within the headroom + and the speculative scratch, on top of the attention tail.""" attention = SimpleNamespace(block_size=16) state_cow = SimpleNamespace( block_size=64, @@ -6245,12 +6383,20 @@ def test_kvwarm_shadow_tail_blocks_count_the_state_group_per_manager(monkeypatch state_plain = SimpleNamespace( block_size=64, mamba_cache_mode="align", num_speculative_blocks=0 ) + state_partial = SimpleNamespace( + block_size=64, + mamba_cache_mode="align", + num_speculative_blocks=0, + _apply_cow=lambda *args: None, + _cache_partial_tail_block=lambda *args: None, + ) def planner(managers): stub = InstrumentedScheduler.__new__(InstrumentedScheduler) stub.cache_config = SimpleNamespace(block_size=16) stub.kv_cache_manager = SimpleNamespace( - coordinator=SimpleNamespace(single_type_managers=managers) + block_pool=SimpleNamespace(hash_block_size=16), + coordinator=SimpleNamespace(single_type_managers=managers), ) return stub @@ -6261,7 +6407,13 @@ def planner(managers): ) == 4 ) - # No fork without CoW, or when in-place is requested. + assert ( + InstrumentedScheduler._kvwarm_shadow_tail_blocks( + planner([attention, state_partial]), 3 + ) + == 4 + ) + # No fork without CoW (auto runs in place), or when in-place is requested. assert ( InstrumentedScheduler._kvwarm_shadow_tail_blocks( planner([attention, state_plain]), 3 @@ -6321,29 +6473,37 @@ def test_kvwarm_shadow_pool_shortfall_adds_the_state_group_blocks(): def test_kvwarm_state_mode_resolution(monkeypatch): """``DYN_BENCH_KV_WARMUP_STATE`` parsing and the ``auto`` rule: copy when - every state manager forks and the pool holds the forks, else in place; - the stage check's floor is the fallback; pure attention has no mode.""" + every state manager forks, else in place where the managers admit it -- + a function of the layout, never of the pool, so every attention-DP rank + and every pool check (plan, stage, injection) name the same mode. Pure + attention has no mode.""" stub, attn, mamba, pool, chain, state = _hybrid_shadow_stub(cow=True) - pool.get_num_free_blocks = lambda: 10 assert InstrumentedScheduler._kvwarm_state_mode_requested(stub) == "auto" - assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "copy" - assert InstrumentedScheduler._kvwarm_state_mode_floor(stub) == "inplace" - pool.get_num_free_blocks = lambda: 1 # the attention tail alone fits - assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "inplace" + for free in (10, 1, 0): + pool.get_num_free_blocks = lambda free=free: free + assert InstrumentedScheduler._kvwarm_state_mode(stub) == "copy" monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", " COPY ") assert InstrumentedScheduler._kvwarm_state_mode_requested(stub) == "copy" - assert InstrumentedScheduler._kvwarm_state_mode(stub, [100], 3) == "copy" - assert InstrumentedScheduler._kvwarm_state_mode_floor(stub) == "copy" + assert InstrumentedScheduler._kvwarm_state_mode(stub) == "copy" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", "inplace") + assert InstrumentedScheduler._kvwarm_state_mode(stub) == "inplace" monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", "bogus") with pytest.raises(ValueError, match="bogus"): InstrumentedScheduler._kvwarm_state_mode_requested(stub) monkeypatch.delenv("DYN_BENCH_KV_WARMUP_STATE") plain, *_ = _hybrid_shadow_stub(cow=False) - plain.kv_cache_manager.block_pool.get_num_free_blocks = lambda: 10 - assert InstrumentedScheduler._kvwarm_state_mode(plain, [100], 3) == "inplace" + assert InstrumentedScheduler._kvwarm_state_mode(plain) == "inplace" + hashed, *_ = _hybrid_shadow_stub(cow=True, partial_tail=True) + assert InstrumentedScheduler._kvwarm_state_mode(hashed) == "copy" + stuck, *_ = _hybrid_shadow_stub(cow=False, partial_tail=True) + assert InstrumentedScheduler._kvwarm_state_mode(stuck) is None + assert ( + InstrumentedScheduler._kvwarm_state_gate_reason(stuck) + == "state_hand_off_unsupported" + ) pure, *_ = _shadow_stub(cow=True) - assert InstrumentedScheduler._kvwarm_state_mode(pure, [100], 3) is None - assert InstrumentedScheduler._kvwarm_state_mode_floor(pure) is None + assert InstrumentedScheduler._kvwarm_state_mode(pure) is None + assert InstrumentedScheduler._kvwarm_state_gate_reason(pure) is None def test_kvwarm_stamp_state_marks_the_row_and_the_metadata(): @@ -6363,8 +6523,6 @@ def test_kvwarm_stamp_state_marks_the_row_and_the_metadata(): assert stub._kvwarm_meta["state_mode"] == "copy" InstrumentedScheduler._kvwarm_stamp_state(stub, "copy", 192) assert stub._kvwarm_meta["state_mode"] == "copy" - InstrumentedScheduler._kvwarm_stamp_state(stub, "inplace", 192) - assert stub._kvwarm_meta["state_mode"] == "mixed" @pytest.mark.parametrize("cow", [True, False]) @@ -6462,7 +6620,8 @@ def _kvwarm_injection_stub(chain_ids, hybrid=False): stub._kvwarm_chain_ids = list(chain_ids) stub._kvwarm_chain_prompts = {chain_id: list(range(160)) for chain_id in chain_ids} stub.requests = { - chain_id: SimpleNamespace(request_id=chain_id) for chain_id in chain_ids + chain_id: SimpleNamespace(request_id=chain_id, num_computed_tokens=160) + for chain_id in chain_ids } stub.running = [] stub.finished_req_ids = set() @@ -6515,10 +6674,10 @@ def test_kvwarm_inject_borrowed_releases_cow_retentions_before_the_step_runs(): def test_kvwarm_inject_borrowed_hands_the_state_over_and_stamps_the_point( monkeypatch, requested ): - """On a hybrid layout the injection settles the state mode against this - point's pool, registers every group, rides the state fork's copy on the - admission step like the attention forks, and stamps the point with the - mode and the chain depth the borrowed state summarises.""" + """On a hybrid layout the injection uses the run's state mode, registers + every group, rides the state fork's copy on the admission step like the + attention forks, and stamps the point with the mode and the chain depth + the borrowed state summarises.""" monkeypatch.setenv("DYN_BENCH_KV_WARMUP_STATE", requested) stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"], hybrid=True) mamba = stub.kv_cache_manager.coordinator.single_type_managers[1] @@ -6551,6 +6710,39 @@ def test_kvwarm_inject_borrowed_hands_the_state_over_and_stamps_the_point( assert stub._kvwarm_meta["state_mode_requested"] == requested +def test_kvwarm_inject_borrowed_stamps_the_same_state_mode_on_every_rank(): + """The stamp enters the point digest attention-DP ranks compare, so the + mode must not depend on the rank's pool: two ranks with different free + counts stamp the same mode and digest, and a pool short of the fork + skips the point (the shortfall path) instead of switching modes.""" + digests = [] + for free in (10, 2): + stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"], hybrid=True) + pool.get_num_free_blocks = lambda free=free: free + InstrumentedScheduler._kvwarm_inject_borrowed(stub, [100]) + assert stub._bench_current_point.sample_reasons == [ + "kvwarm_real_kv", + "kvwarm_state_copy", + "kvwarm_state_depth=160", + ] + digests.append( + instrumented_scheduler_module._benchmark_point_digest( + stub._bench_current_point + ) + ) + assert digests[0] == digests[1] + stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"], hybrid=True) + mamba = stub.kv_cache_manager.coordinator.single_type_managers[1] + pool.get_num_free_blocks = lambda: 1 # the attention tail alone fits + output = InstrumentedScheduler._kvwarm_inject_borrowed(stub, [100]) + assert output.total_num_scheduled_tokens == 0 + assert ( + "__bench_0" not in mgr.req_to_blocks and "__bench_0" not in mamba.req_to_blocks + ) + assert stub._bench_current_point.sample_reasons == ["kvwarm_real_kv"] + assert getattr(stub, "_kvwarm_meta", None) is None + + def test_kvwarm_inject_borrowed_leaves_pure_attention_points_unstamped(): stub, mgr, pool, chains = _kvwarm_injection_stub(["chain-a"]) stub._bench_current_point = BenchmarkPoint( diff --git a/docs/fern/pages/reference/observability/environment-variables.mdx b/docs/fern/pages/reference/observability/environment-variables.mdx index 03a86774bde2..957fbed73895 100644 --- a/docs/fern/pages/reference/observability/environment-variables.mdx +++ b/docs/fern/pages/reference/observability/environment-variables.mdx @@ -279,7 +279,7 @@ These variables tune the vLLM worker's self-benchmark collector (`--benchmark-mo - Recurrent-state hand-off on hybrid layouts. Attention layers read the warm-up chain's KV by position; a Mamba / linear-attention layer keeps one running-state block per request, which the measured request borrows from its chain. `copy` forks that block (copy-on-write, one extra block per request and state group; needs a vLLM whose KV-cache managers support CoW), `inplace` hands the chain's own block to the measured request, `auto` copies when the manager can and the pool holds the forks, otherwise runs in place. The chosen mode is recorded as `state_mode` in the KV warm-up metadata and per row as `kvwarm_state_copy` / `kvwarm_state_inplace`, alongside `kvwarm_state_depth=N`, the chain depth the borrowed state summarises. An explicit `copy` on a vLLM without CoW skips the warm-up with `state_copy_unsupported`. + Recurrent-state hand-off on hybrid layouts. Attention layers read the warm-up chain's KV by position; a Mamba / linear-attention layer keeps one running-state block per request, which the measured request borrows from its chain. `copy` forks that block (copy-on-write, one extra block per request and state group; needs a vLLM whose KV-cache managers support CoW), `inplace` hands the chain's own block to the measured request, `auto` copies when the managers can and otherwise runs in place. The mode is fixed by the layout and the vLLM build for the whole run, never by the free pool: a pool too small for it skips the point like any other shadow-block shortfall. It is recorded as `state_mode` in the KV warm-up metadata and per row as `kvwarm_state_copy` / `kvwarm_state_inplace`, alongside `kvwarm_state_depth=N`, the chain depth the borrowed state summarises. `inplace` lets the measured request write into the chain's state block: later points that borrow the same chain in that stage, and the next generation's deeper chain that resumes from it through the prefix cache, start from the advanced state (timing-neutral, but no longer the pure chain state). Which modes a vLLM supports depends on its KV-cache managers: without CoW only `inplace` is possible; a build that hashes partial tails at a finer granularity than the state block (the vLLM 0.28.0 line) would corrupt the chain's prefix entry from a shared state block, so only `copy` is used there. An explicit mode the build cannot serve skips the warm-up with `state_copy_unsupported` or `state_inplace_unsupported`; an unparseable value skips it with `state_mode_invalid` (on layouts without recurrent-state groups the variable is ignored).