diff --git a/components/src/dynamo/vllm/args.py b/components/src/dynamo/vllm/args.py index 7fb9a0a61932..4347de90f1c0 100644 --- a/components/src/dynamo/vllm/args.py +++ b/components/src/dynamo/vllm/args.py @@ -437,9 +437,11 @@ def update_engine_config_with_dynamo( benchmark_config: Dict[str, Any] = { "mode": dynamo_config.benchmark_mode, "randomize_kda_state": dynamo_config.benchmark_randomize_kda_state, + "hybrid_live_state": dynamo_config.benchmark_hybrid_live_state, "warmup_iterations": dynamo_config.benchmark_warmup_iterations, "output_path": dynamo_config.benchmark_output_path, "timeout": dynamo_config.benchmark_timeout, + "max_batch_size": dynamo_config.benchmark_max_batch_size, "collect_imbalanced": dynamo_config.benchmark_collect_imbalanced, } explicit_points = dynamo_config._benchmark_points diff --git a/components/src/dynamo/vllm/backend_args.py b/components/src/dynamo/vllm/backend_args.py index 89fe0e2c0fe8..470aa33a82bc 100644 --- a/components/src/dynamo/vllm/backend_args.py +++ b/components/src/dynamo/vllm/backend_args.py @@ -480,6 +480,20 @@ def add_arguments(self, parser) -> None: "vLLM GPU worker. These are synthetic states, not real context history." ), ) + add_negatable_bool_argument( + g, + flag_name="--benchmark-hybrid-live-state", + env_var="DYN_BENCHMARK_HYBRID_LIVE_STATE", + default=False, + help=( + "Hybrid (KDA/Mamba) models: let the real-KV decode warm-up run with " + "recurrent-state groups forked from the parked chain's live state block " + "instead of skipping the warm-up. Attention KV is the chain's real prefix; " + "the recurrent state is a valid but deeper-context state (shallow points " + "read a few percent fast). Mutually exclusive with " + "--benchmark-randomize-kda-state." + ), + ) add_argument( g, flag_name="--benchmark-warmup-iterations", @@ -524,6 +538,21 @@ def add_arguments(self, parser) -> None: "A bounded cleanup grace still fails closed if no result is written." ), ) + add_argument( + g, + flag_name="--benchmark-max-batch-size", + env_var="DYN_BENCHMARK_MAX_BATCH_SIZE", + default=None, + arg_type=int, + help=( + "Cap the decode benchmark batch-size axis without changing the " + "engine's own limits. Points above the cap are never generated " + "(not measured, not skipped). Decoupled from --max-num-seqs: keep " + "the engine at its deployment concurrency while sweeping only the " + "batch range the performance model needs; setting --max-num-seqs " + "to at least twice this cap keeps every swept rung warmable." + ), + ) # @dataclass() @@ -571,9 +600,11 @@ class DynamoVllmConfig(ConfigBase): benchmark_mode: Optional[BenchmarkMode] = None benchmark_points_file: Optional[str] = None benchmark_randomize_kda_state: bool = False + benchmark_hybrid_live_state: bool = False benchmark_warmup_iterations: int = 5 benchmark_output_path: str = "/tmp/benchmark_results.json" benchmark_timeout: int = 900 + benchmark_max_batch_size: Optional[int] = None prefill_max_new_token_samples: int = 64 prefill_max_kv_read_token_samples: int = 16 decode_max_kv_read_token_samples: int = 128 @@ -693,6 +724,18 @@ def _resolve_legacy_benchmark_sampling(self) -> None: setattr(self, replacement_name, mapped_value) def _validate_benchmark_sampling(self) -> None: + if self.benchmark_hybrid_live_state and self.benchmark_randomize_kda_state: + raise ValueError( + "--benchmark-hybrid-live-state and --benchmark-randomize-kda-state " + "are mutually exclusive" + ) + if self.benchmark_hybrid_live_state and self.benchmark_mode not in ( + "decode", + "agg", + ): + raise ValueError( + "--benchmark-hybrid-live-state requires --benchmark-mode decode or agg" + ) if self.benchmark_randomize_kda_state and self.benchmark_mode not in ( "decode", "agg", diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index cbc2a12947ce..b6c44213a04e 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -171,6 +171,7 @@ def recurrent_shadow_range( class BenchmarkConfig: mode: BenchmarkMode = "agg" randomize_kda_state: bool = False + hybrid_live_state: bool = False warmup_iterations: int = 5 output_path: str = "/tmp/benchmark_results.json" timeout: int = 900 @@ -179,6 +180,10 @@ class BenchmarkConfig: decode_max_kv_read_token_samples: int = 128 decode_max_batch_size_samples: int = 128 prefix_max_batch_size_samples: int = 3 + # Cap on the decode batch-size axis, independent of the engine's own + # max_num_running_reqs: points above the cap are never generated. None + # keeps the historical behavior (axis runs to the engine limit). + max_batch_size: int | None = None # Measure the manifest's imbalanced prefill points (explicit rows, or a # partition) as well as its uniform ones. Those points come from an # explicit --benchmark-points-file; see the flag's comment in backend_args @@ -2210,6 +2215,7 @@ def _cleanup_finished(self, output: SchedulerOutput) -> None: # ------------------------------------------------------------------ _bench_random_kda: bool = False + _bench_hybrid_live_state: bool = False def _bench_init(self, vllm_config: "VllmConfig") -> None: """Parse benchmark config and initialise state machine.""" @@ -2244,6 +2250,10 @@ def _bench_init(self, vllm_config: "VllmConfig") -> None: for k in _INT_FIELDS: if k in cfg and not isinstance(cfg[k], int): cfg[k] = int(cfg[k]) + if cfg.get("max_batch_size") is not None and not isinstance( + cfg["max_batch_size"], int + ): + cfg["max_batch_size"] = int(cfg["max_batch_size"]) # A bool that arrives as JSON text: "false" is a non-empty string and # would otherwise turn the collection on. if "collect_imbalanced" in cfg and isinstance(cfg["collect_imbalanced"], str): @@ -2260,6 +2270,13 @@ def _bench_init(self, vllm_config: "VllmConfig") -> None: if not isinstance(self._bench_config.randomize_kda_state, bool): raise ValueError("benchmark randomize_kda_state must be a boolean") self._bench_random_kda = self._bench_config.randomize_kda_state + if not isinstance(self._bench_config.hybrid_live_state, bool): + raise ValueError("benchmark hybrid_live_state must be a boolean") + self._bench_hybrid_live_state = self._bench_config.hybrid_live_state + if self._bench_hybrid_live_state and self._bench_random_kda: + raise ValueError( + "benchmark hybrid_live_state and randomize_kda_state are mutually exclusive" + ) if ( self._bench_random_kda and vllm_config.parallel_config.worker_cls != RANDOM_KDA_WORKER @@ -2269,6 +2286,11 @@ def _bench_init(self, vllm_config: "VllmConfig") -> None: ) if self._bench_config.timeout <= 0: raise ValueError("benchmark timeout must be positive") + if ( + self._bench_config.max_batch_size is not None + and self._bench_config.max_batch_size < 1 + ): + raise ValueError("benchmark max_batch_size must be positive") uniform_sample_limits = { "prefill_max_new_token_samples": ( self._bench_config.prefill_max_new_token_samples @@ -2931,11 +2953,24 @@ def _bench_generate_prefill_grid(self) -> None: ) return + new_token_candidates = _cudagraph_axis_points( + self._bench_prefill_capture_sizes, + max_tokens, + ) + if getattr(self, "need_mamba_block_aligned_split", False): + # Hybrid align mode splits any chunk that crosses a cache-block boundary at + # that boundary, so a point whose per-request new-token count is above one + # block and not block-aligned can never run as planned and fails the + # feasibility check below. The cudagraph axis (powers of two) has no such + # values, which leaves small-batch prefill with more than one block of new + # tokens uncollected -- exactly the steps a served long prompt produces + # (block-multiple chunks). Add block-multiple totals so those shapes are + # candidates; the configured sample limit applies to the combined axis. + new_token_candidates = self._bench_block_aligned_prefill_axis( + new_token_candidates, max_tokens + ) total_prefill_tokens = _limit_cudagraph_axis( - _cudagraph_axis_points( - self._bench_prefill_capture_sizes, - max_tokens, - ), + new_token_candidates, self._bench_prefill_capture_sizes, self._bench_config.prefill_max_new_token_samples, ) @@ -2980,6 +3015,25 @@ def _bench_generate_prefill_grid(self) -> None: # decode ordering and the aggregate prefill-before-decode boundary intact. self._bench_grid.extend(reversed(prefill_points)) + def _bench_block_aligned_prefill_axis( + self, axis: Sequence[int], max_tokens: int + ) -> list[int]: + """Union of ``axis`` with the whole-block totals ``k * block_size`` (k >= 2) + that fit ``max_tokens``: per-request chunks of whole cache blocks are the only + multi-block chunks a hybrid align-mode scheduler runs unsplit, and every + small-batch multiple of the block size is itself such a total. + """ + block_size = int( + getattr(getattr(self, "cache_config", None), "block_size", 0) + or self.block_size + or 0 + ) + if block_size <= 0: + return list(axis) + totals = set(int(t) for t in axis) + totals.update(range(2 * block_size, int(max_tokens) + 1, block_size)) + return sorted(totals) + def _bench_prefill_batch_sizes(self, total_tokens: int) -> list[int]: """Return the smallest configured presets from the legal batch axis.""" upper_bound = min( @@ -3224,8 +3278,17 @@ def _bench_blocks_per_req( *, has_cache_hit: bool = False, apply_admission_cap: bool = False, + resident_chain: bool = False, ) -> int: - """Predict the shared-pool block footprint of one request.""" + """Predict the shared-pool block footprint of one request. + + ``resident_chain``: a prefilled kvwarm chain that stays resident and keeps + growing across stages holds TWO Mamba 'align' state blocks per group (current + + previous aligned boundary; vLLM MambaSpec sizes align mode as page_size * (2 + + num_speculative_blocks)). A freshly injected decode request holds one (+1 when + it pins a prefix-cache hit). Under-counting the chain footprint let the stage + plan admit more chains than the pool holds (stage-build deadlock at the pool + edge).""" coordinator = getattr( getattr(self, "kv_cache_manager", None), "coordinator", None ) @@ -3267,6 +3330,11 @@ def _bench_blocks_per_req( 0, ) blocks = 2 + speculative_blocks + checkpoint_blocks + if resident_chain and num_tokens < block_size: + # A parked kvwarm chain below one cache block has not crossed a + # block boundary yet and holds only its live state block; past the + # first boundary the allocator bound above applies. + blocks = 1 + speculative_blocks elif mamba_cache_mode is not None: blocks += speculative_blocks @@ -3457,6 +3525,10 @@ def _bench_generate_decode_grid(self) -> None: return feasible_max_batch = self._bench_decode_feasible_max_batch_size() + if self._bench_config.max_batch_size is not None: + feasible_max_batch = min( + feasible_max_batch, self._bench_config.max_batch_size + ) self._bench_feasible_max_decode_batch_size = feasible_max_batch if feasible_max_batch < 1: logger.warning("KV cache too small for decode grid, skipping") @@ -4888,7 +4960,11 @@ def _kvwarm_warm_eligible(self) -> bool: # to deepen incrementally; with prefix cache off a full chain # rebuild is prohibitively expensive -- prefer skipping. reason = "prefix_caching_disabled" - elif self._kvwarm_state_layer_groups() and not self._bench_random_kda: + elif ( + self._kvwarm_state_layer_groups() + and not self._bench_random_kda + and not self._bench_hybrid_live_state + ): reason = "hybrid_state_layers_unsupported" else: reason = self._kvwarm_probe_content() @@ -5150,12 +5226,14 @@ def _kvwarm_prepare(self, mode: str) -> None: repeats = self._kvwarm_giant_repeats() margin = 1 + repeats plan: dict = {} + rung_ctxs: dict = {} for p in decode_pts: ctxs = self._bench_decode_context_lengths( p.total_kv_read_tokens, p.batch_size ) want = min(max(ctxs) + margin, self._kvwarm_depth_cap()) plan[p.batch_size] = max(plan.get(p.batch_size, 0), want) + rung_ctxs.setdefault(p.batch_size, set()).update(int(c) for c in ctxs) # Shadows own private tail blocks (the admission write plus the steady # headroom) on top of the shared chain prefix, drawn from the same pool # while the chains are parked. Reserve them per request and per KV @@ -5166,15 +5244,56 @@ def _kvwarm_prepare(self, mode: str) -> None: # which rung every rank builds and which points it warms, and the # stage round (``_kvwarm_stage_round``) relies on every rank # agreeing on both. - shadow_tail_blocks = self._kvwarm_shadow_tail_blocks(repeats) + worst_case_tail = self._kvwarm_shadow_tail_blocks(repeats) for batch, depth in list(plan.items()): + # Reserve the tails the rung's OWN points take (exact per-group arithmetic + # at each measured context), bounded by the worst case; the worst case alone + # (two blocks per group) over-reserves on hybrids. + shadow_tail_blocks = min( + worst_case_tail, + max( + ( + # the shadow is admitted at ctx - 1 with ``repeats`` steady + # steps (``_bench_step_decode`` / ``_kvwarm_inject_borrowed``); + # the recurrent read slot can cross a block boundary between + # ctx and ctx - 1, so reserve at the admission geometry + self._kvwarm_shadow_tail_blocks_for( + max(1, ctx - 1), max(1, repeats) + ) + for ctx in rung_ctxs.get(batch, ()) + ), + default=worst_case_tail, + ), + ) usable = self._bench_grid_usable_blocks(batch, reserve_watermark=True) + # Plan against a margin of the pool: the per-request footprint estimate is a + # lower bound (block-boundary rounding, transient Mamba boundary blocks), + # and a stage whose chains do not ALL fit loses its real coverage, so a + # small margin buys full stages. Under attention-DP every rank must build + # the same stages: leave more headroom (a per-rank stall would desynchronize + # the ranks) -- 15% for DP>1, 5% otherwise. + pool_margin = 0.85 if getattr(self, "_bench_dp_size", 1) > 1 else 0.95 + pool = int(usable * pool_margin) while depth > 8 and ( - (self._bench_blocks_per_req(depth) + shadow_tail_blocks) * batch - > usable + ( + self._bench_blocks_per_req( + depth, apply_admission_cap=True, resident_chain=True + ) + + shadow_tail_blocks + ) + * batch + > pool ): depth -= 1 - required = (self._bench_blocks_per_req(depth) + shadow_tail_blocks) * batch + required = ( + self._bench_blocks_per_req( + depth, apply_admission_cap=True, resident_chain=True + ) + + shadow_tail_blocks + ) * batch + # The margin only steers the depth trim above; whether a stage is built at + # all is decided against the full pool, as upstream does (a rung already at + # the depth floor is not demoted by the margin). if required > usable: # Reaching the depth floor does not prove the fleet fits. # This also covers an initially short chain below the floor. @@ -5198,6 +5317,20 @@ def _kvwarm_prepare(self, mode: str) -> None: ) depth = 0 plan[batch] = depth + # Slot budget: a stage parks ``batch`` chains on the worker and measures each of + # its points by injecting ``batch`` shadow requests on top, so ``2 * batch`` + # request slots must exist (worker asserts "No free indices" otherwise). Rungs + # above that fall back to fake injection (measured after the chains are + # released). On models whose max_num_seqs is memory-capped (Mamba/KDA state + # blocks) this bites at batch > max_num_seqs / 2. + try: + slots = int(self._bench_capacity_limit("max_num_running_reqs")) + except (AttributeError, TypeError, ValueError): + # capacity unknown (e.g. partially constructed scheduler): leave the plan alone + slots = 0 + if slots > 0: + for batch in [b for b in plan if 2 * b > slots]: + plan.pop(batch) self._kvwarm_plan = plan # Second reordering: all warmed points first, fake fallbacks last -- # fake injection fills the whole pool and evicts the chains' cached @@ -5205,6 +5338,56 @@ def _kvwarm_prepare(self, mode: str) -> None: # deepening back to full rebuilds. warmed_pts = [p for p in decode_pts if self._kvwarm_plan_covers(p)] fake_pts = [p for p in decode_pts if not self._kvwarm_plan_covers(p)] + if getattr(self, "_bench_dp_size", 1) > 1 and fake_pts: + # Attention-DP: fake injection is not rank-consistent on the small per-rank + # pools (TP1): points truncate or OOM on some ranks, get skipped, and the + # group_prepare barrier times out (observed dep4: dozens of + # measured_decode_context_mismatch skips then "timed out waiting for + # attention-DP benchmark group_prepare"). Keep the grid real-KV only under + # DP; renumber so published ids stay contiguous and 1-based. Explicit + # manifest points carry the contract "measured or rejected with an error": + # never drop them silently. + explicit = [pt for pt in fake_pts if "explicit" in pt.sample_reasons] + if explicit: + coords = ", ".join( + f"(batch={pt.batch_size}, total_kv_read_tokens={pt.total_kv_read_tokens})" + for pt in explicit + ) + raise RuntimeError( + "KVWARM: attention-DP runs measure real-KV points only, and the " + f"warm-up plan cannot cover explicit decode point(s) {coords} " + "(rung depth trimmed by the per-rank pool / slot budget); lower " + "the requested batch or context, or run without attention-DP" + ) + logger.warning( + "KVWARM: attention-DP (dp_size=%d): dropping %d fake-fallback decode " + "points; grid is real-KV only", + self._bench_dp_size, + len(fake_pts), + ) + fake_pts = [] + # Eager warm-up replicas are executed but their results are discarded + # (``_bench_save_current_point``): count only real points, as + # ``_bench_build_grid`` does; it assigns the final IDs afterwards. + self._bench_expected_points = sum( + EAGER_WARMUP_REASON not in pt.sample_reasons + for pt in other_pts + warmed_pts + ) + if not warmed_pts: + # The filter emptied the decode phase (e.g. a per-rank slot budget + # below two requests): the phase was generated, so + # ``_bench_build_grid`` did not record it as missing. Record it here + # so the artifact cannot report a complete, usable run with no decode + # measurements. + missing = getattr(self, "_bench_missing_phases", None) + if missing is None: + missing = self._bench_missing_phases = [] + if "decode" not in missing: + missing.append("decode") + logger.warning( + "KVWARM: attention-DP warm-up plan covers no decode point; " + "decode phase recorded as missing" + ) self._bench_grid = deque(other_pts + warmed_pts + fake_pts) self._kvwarm_chain_ids: list = [] self._kvwarm_chain_prompts: dict = {} @@ -5222,11 +5405,86 @@ def _kvwarm_prepare(self, mode: str) -> None: # ------- Warmup state machine (intercepts before phase dispatch) ------- + @staticmethod + def _kvwarm_admission_cap(manager) -> int | None: + """Per-request block cap of an admission-capped KV-cache group (vLLM + ``_max_admission_blocks_per_request``, e.g. sliding-window and k-pool-tail + managers); None for groups whose tables grow with the context. + """ + cap = getattr(manager, "_max_admission_blocks_per_request", None) + return cap if isinstance(cap, int) and cap > 0 else None + + def _kvwarm_circular_table_manager(self, manager) -> bool: + """True for admission-capped groups whose block table is a fixed-length circular + buffer rather than a position-indexed table (GLM5-Next's k-pool tail: one block + per request, excluded from prefix caching). Sliding-window / chunked-local + groups are admission-capped too, but keep positional tables with null + placeholders for expired positions, so they take the positional path. + """ + if self._kvwarm_admission_cap(manager) is None: + return False + spec = getattr(manager, "kv_cache_spec", None) + return getattr(spec, "participates_in_prefix_caching", True) is False + + def _kvwarm_live_state_manager(self, manager) -> bool: + """Recurrent-state groups served by live-state borrowing (hybrid live-state mode, + random mode off). + """ + # Keyed on the KV-cache spec like the eligibility gate ("Mamba" in the spec + # name) and the random-state path (``isinstance(spec, MambaSpec)``), not on + # manager class names, so a vLLM rename or subclass cannot silently route a + # recurrent group to the positional path. The k-pool tail is not a recurrent + # state table; it is served by the circular-table predicate. + spec = getattr(manager, "kv_cache_spec", None) + return ( + self._bench_hybrid_live_state + and not self._bench_random_kda + and (isinstance(spec, MambaSpec) or "Mamba" in type(spec).__name__) + ) + def _kvwarm_random_state_manager(self, manager) -> bool: return self._bench_random_kda and isinstance( getattr(manager, "kv_cache_spec", None), MambaSpec ) + def _kvwarm_shadow_tail_blocks_for(self, ctx_len: int, headroom: int) -> int: + """Private tail blocks one shadow at ``ctx_len`` takes across the KV-cache groups + (the exact arithmetic of ``_kvwarm_register_shadow`` / + ``_kvwarm_shadow_pool_shortfall``): the write positions ctx .. ctx+headroom + minus the shared full prefix blocks; random-state groups their state-table span; + admission-capped groups their whole (capped) table. The worst case + (``_kvwarm_shadow_tail_blocks``) charges two blocks per group for any context; + on a hybrid with many small groups that alone pushes the large rungs over the + pool. + """ + coordinator = getattr( + getattr(self, "kv_cache_manager", None), "coordinator", None + ) + managers = getattr(coordinator, "single_type_managers", ()) + if not managers: + block_size = int(getattr(self.cache_config, "block_size", 16) or 16) + return -(-(ctx_len + 1 + headroom) // block_size) - ctx_len // block_size + need = 0 + for manager in managers: + bs = int( + getattr( + manager, "block_size", getattr(self.cache_config, "block_size", 16) + ) + ) + if self._kvwarm_circular_table_manager(manager): + need += ( + self._kvwarm_admission_cap(manager) or 0 + ) # fixed ring, whatever the mode + continue + if self._kvwarm_random_state_manager( + manager + ) or self._kvwarm_live_state_manager(manager): + first, end = recurrent_shadow_range(ctx_len, headroom, bs) + need += end - first + continue + need += -(-(ctx_len + 1 + headroom) // bs) - ctx_len // bs + return need + def _kvwarm_shadow_tail_blocks(self, repeats: int) -> int: """Bound private attention tails and recurrent-state write positions.""" coordinator = getattr( @@ -5240,7 +5498,16 @@ def _kvwarm_shadow_tail_blocks(self, repeats: int) -> int: return sum( 1 + -( - -(headroom + int(self._kvwarm_random_state_manager(manager))) + -( + headroom + + int( + self._kvwarm_random_state_manager(manager) + or ( + self._kvwarm_live_state_manager(manager) + and not self._kvwarm_circular_table_manager(manager) + ) + ) + ) // manager.block_size ) for manager in managers @@ -5639,12 +5906,50 @@ def _kvwarm_register_shadow( n_shared = ctx_len // bs n_total = -(-(ctx_len + 1 + headroom) // bs) chain_blocks = list(mgr.req_to_blocks[chain_id]) - if n_total > len(chain_blocks): - raise RuntimeError( - f"KVWARM: chain {chain_id} too shallow for shadow {req_id}: " - f"needs {n_total} blocks, has {len(chain_blocks)}" - ) - tail_src = chain_blocks[n_shared:n_total] + is_circular = self._kvwarm_circular_table_manager(mgr) and len( + chain_blocks + ) <= (self._kvwarm_admission_cap(mgr) or 0) + # A circular table (k-pool tail) is a fixed ring, never a positional + # state table: it keeps its own geometry even when live-state mode is on + # (the manager matches both predicates). + is_live_state = ( + not is_circular + and self._kvwarm_live_state_manager(mgr) + and bool(chain_blocks) + ) + if is_live_state: + # Recurrent state is read at ceil(ctx/bs)-1 and written at the + # admission and steady positions (``recurrent_shadow_range``); at an + # exact block boundary the read slot is the last SHARED entry, which + # may be a pruned/null checkpoint. Fork the whole read/write span + # privately. + n_shared, n_total = recurrent_shadow_range(ctx_len, headroom, bs) + if is_circular: + # Circular fixed-length table (k-pool tail): the chain holds ``cap`` + # blocks whatever its depth and the runner's table for the group has + # ``cap`` columns, so the shadow shares nothing and forks the + # chain's block(s) instead of walking positional indices. + n_shared = 0 + n_total = len(chain_blocks) + tail_src = list(chain_blocks) + else: + if n_total > len(chain_blocks): + raise RuntimeError( + f"KVWARM: chain {chain_id} too shallow for shadow {req_id}: " + f"needs {n_total} blocks, has {len(chain_blocks)}" + ) + tail_src = chain_blocks[n_shared:n_total] + if is_live_state: + # Live-state mode for recurrent groups: the state has no + # per-position history and the only valid state a chain holds is in + # its LIVE (last) block. Entries below it are align-mode retention + # checkpoints that may already be evicted and recycled; a shadow at + # ctx << chain depth forked from chain_blocks[ctx // bs] copies a + # stale block, its hidden state degenerates and MoE routing + # collapses (2-3x too-fast decode steps measured on GLM-5.3-Flash). + # Fork the private span from the live block instead: a valid, + # deeper-context state. + tail_src = [chain_blocks[-1]] * len(tail_src) fresh = block_pool.get_new_blocks(len(tail_src)) staged.append((mgr, n_shared, chain_blocks[:n_shared], tail_src, fresh)) except Exception: @@ -5682,23 +5987,17 @@ def _kvwarm_register_shadow( 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.""" + (0 when they fit). Uses the per-context reserve of + ``_kvwarm_shadow_tail_blocks_for`` so planning, this check and the + allocation in ``_kvwarm_register_shadow`` agree.""" manager = self.kv_cache_manager free_fn = getattr(manager.block_pool, "get_num_free_blocks", None) if not callable(free_fn): return 0 - need = 0 - for mgr in manager.coordinator.single_type_managers: - bs = int( - getattr(mgr, "block_size", getattr(self.cache_config, "block_size", 16)) - ) - for ctx_len in context_lengths: - if self._kvwarm_random_state_manager(mgr): - first, end = recurrent_shadow_range(ctx_len, headroom, bs) - need += end - first - else: - need += -(-(ctx_len + 1 + headroom) // bs) - ctx_len // bs + need = sum( + self._kvwarm_shadow_tail_blocks_for(ctx_len, headroom) + for ctx_len in context_lengths + ) return max(0, need - int(free_fn())) def _kvwarm_take_cow_copies(self) -> list: @@ -5978,6 +6277,26 @@ def _bench_step_decode(self) -> SchedulerOutput | None: sample_reasons=[*point.sample_reasons, "context_clamped"], ) kvwarm_real = self._kvwarm_covers(point, injected_lengths) + if ( + not kvwarm_real + and self._kvwarm_flag_on() + and getattr(self, "_bench_dp_size", 1) > 1 + and self._kvwarm_warm_eligible() + ): + # Attention-DP grids are real-KV only (``_kvwarm_prepare``); a point reaches + # here without coverage only when its rung's stage failed afterwards (group + # verdict, so every rank sees the same point). Fake injection is not + # rank-consistent, so record a skip instead -- explicit points raise in the + # skip path. + logger.warning( + "KVWARM: attention-DP point without real-KV coverage (stage failed): " + "skipping %s", + point, + ) + self._bench_skip_point(point, "stage_failed_under_attention_dp") + self._bench_current_point = None + self._bench_extra_steps_left = 0 + return None if self._kvwarm_flag_on(): meta = self._kvwarm_meta_init() if kvwarm_real: @@ -6014,7 +6333,9 @@ def _bench_step_decode(self) -> SchedulerOutput | None: repeats = min(repeats, max(1, self.max_model_len - 1 - max_ctx)) if not kvwarm_real: multi = sum( - self._bench_blocks_per_req(max(c, 2) + repeats) + self._bench_blocks_per_req( + max(c, 2) + repeats, apply_admission_cap=True + ) for c in injected_lengths ) if multi > self._bench_usable_blocks( @@ -6080,6 +6401,7 @@ def _bench_save_current_point(self) -> None: chosen = dict(steadies[0]) chosen["wall_time"] = walls[len(walls) // 2] chosen["kvwarm_giant_median_of"] = len(steadies) + chosen["kvwarm_steady_sample"] = True local_fpms = [chosen] elif expected_fpms > 1 and len(local_fpms) >= expected_fpms: # Keep only the steady-state sample; the admission step is @@ -6087,8 +6409,12 @@ def _bench_save_current_point(self) -> None: # admission FPM alone sends it through collect_result # unchanged: the shape validator rejects it # (sum_decode_kv_tokens mismatch) and every rank skips the - # point together -- no rank ever bypasses the barrier. - local_fpms = local_fpms[-1:] + # point together -- no rank ever bypasses the barrier. Both + # save paths mark the sample they keep as a recorded steady + # step, which the giant off-by-batch correction requires. + chosen = dict(local_fpms[-1]) + chosen["kvwarm_steady_sample"] = True + local_fpms = [chosen] if self._bench_synchronizer is not None: group_result = self._bench_synchronizer.collect_result( point, @@ -6112,9 +6438,21 @@ def _bench_save_current_point(self) -> None: wall_times: list[float] = [] validation_failure: tuple[int, str] | None = None + measured_decode_totals: dict[int, int] = {} for result in rank_results: dp_rank = result["dp_rank"] fpms = result.get("fpms") + if isinstance(fpms, list) and not fpms: + # The point deadline elapsed before this rank recorded a + # single FPM (the first pass at a fresh giant shape can + # outlast the deadline where a kernel JIT-compiles or the + # backend is simply slower than the validation platform). + # The deadline contract is a group-synchronized skip -- + # the same funnel an admission-only shape mismatch takes + # -- not a sweep abort. + if validation_failure is None: + validation_failure = (dp_rank, "no_fpm_before_deadline") + continue if not isinstance(fpms, list) or len(fpms) != 1: raise RuntimeError( "each self-benchmark point must produce exactly one FPM: " @@ -6137,6 +6475,26 @@ def _bench_save_current_point(self) -> None: reason = self._bench_fpm_validation_failure(point, fpm) if reason is not None and validation_failure is None: validation_failure = (dp_rank, reason) + if point.point_type == "decode": + measured_decode_totals[dp_rank] = int( + fpm.get("scheduled_requests", {}).get( + "sum_decode_kv_tokens", -1 + ) + ) + if ( + point.point_type == "decode" + and validation_failure is None + and len(set(measured_decode_totals.values())) > 1 + ): + # Giant fake-KV points may run one batch short of the plan; the + # normalization below must move every rank to the SAME coordinate, + # otherwise the rank artifacts disagree and the merge rejects them. + odd = next( + r + for r, m in measured_decode_totals.items() + if m != measured_decode_totals[rank_results[0]["dp_rank"]] + ) + validation_failure = (odd, "giant_measured_coordinate_mismatch") if EAGER_WARMUP_REASON in point.sample_reasons: # Warmup replicas are best-effort scaffolding and must be @@ -6178,6 +6536,25 @@ def _bench_save_current_point(self) -> None: self._bench_request_timeout_stop(point) return + if point.point_type == "decode": + # Benign giant-KV fake off-by-batch (accepted by + # _bench_fpm_validation_failure): the steady step measured one token per + # request short of the declared coordinate. Record the point AT THE + # MEASURED coordinate so the artifact stays self-consistent -- the + # collector re-checks scheduled.sum_decode_kv_tokens == + # point.total_kv_read_tokens exactly and fails the whole cell otherwise. + measured = set(measured_decode_totals.values()) + if len(measured) == 1: + m = measured.pop() + if m > 0 and m != point.total_kv_read_tokens: + point = replace( + point, + total_kv_read_tokens=m, + sample_reasons=[ + *point.sample_reasons, + "giant_fake_off_by_batch", + ], + ) self._bench_results.append( BenchmarkPointResult( point=point, @@ -6203,8 +6580,9 @@ def _bench_save_current_point(self) -> None: self._bench_current_fpms = [] self._bench_point_deadline = 0.0 - @staticmethod - def _bench_fpm_validation_failure(point: BenchmarkPoint, fpm: dict) -> str | None: + def _bench_fpm_validation_failure( + self, point: BenchmarkPoint, fpm: dict + ) -> str | None: scheduled = fpm.get("scheduled_requests", {}) batch_size_key = ( "num_prefill_requests" @@ -6219,6 +6597,23 @@ def _bench_fpm_validation_failure(point: BenchmarkPoint, fpm: dict) -> str | Non if scheduled.get("sum_prefill_kv_tokens") != point.total_kv_read_tokens: return "measured_kv_read_mismatch" elif scheduled.get("sum_decode_kv_tokens") != point.total_kv_read_tokens: + # Giant-KV fake points: the repeated-steady-step fake path measures exactly + # one token per request short of the declared coordinate (measured == + # declared - batch, all requests admitted). That is a <0.1% context shift on + # a giant point; ``_bench_save_current_point`` records it at the MEASURED + # coordinate instead of skipping the point and failing the strict + # all-or-nothing publish gate. The admission step has the same total, so + # the correction requires a recorded steady sample (``kvwarm_steady_sample``, + # set by both save paths): a point that hit its deadline with the admission + # FPM alone stays a validation skip. + measured = scheduled.get("sum_decode_kv_tokens") + if ( + "kvwarm_fake_fallback" in (point.sample_reasons or ()) + and point.total_kv_read_tokens >= self._kvwarm_giant_threshold() + and measured == point.total_kv_read_tokens - point.batch_size + and bool(fpm.get("kvwarm_steady_sample")) + ): + return None return "measured_decode_context_mismatch" return None 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 8e11c957c1ca..e56e2825585d 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -2516,6 +2516,33 @@ def test_prefill_grid_uses_total_tokens_and_piecewise_boundaries(): assert engine_limit.sample_reasons == ["eager_tail", "engine_limit"] +def test_prefill_grid_applies_the_sample_limit_after_block_alignment(): + """Hybrid align mode adds whole-block totals to the new-token axis; the + configured sample limit bounds the combined axis, so the aligned candidates + cannot grow the grid past ``prefill_max_new_token_samples``.""" + stub = _prefill_grid_stub(num_gpu_blocks=512) + stub.max_num_scheduled_tokens = 40 + stub._bench_prefill_capture_sizes = list(range(1, 41)) + stub._bench_config.prefill_max_new_token_samples = 4 + stub._bench_config.prefill_max_kv_read_token_samples = 3 + stub._bench_config.prefix_max_batch_size_samples = 1 + stub.need_mamba_block_aligned_split = True + stub.cache_config.block_size = stub.block_size + # whole-block chunks pass the align-mode split unchanged + stub._mamba_block_aligned_split = lambda request, new_tokens, **_: new_tokens + block_size = int(stub.block_size) + aligned = list(range(2 * block_size, 41, block_size)) + assert aligned, "fixture must leave room for whole-block totals" + assert InstrumentedScheduler._bench_block_aligned_prefill_axis( + stub, [1, 40], 40 + ) == sorted({1, 40, *aligned}) + + InstrumentedScheduler._bench_generate_prefill_grid(stub) + + totals = sorted({point.total_prefill_tokens for point in stub._bench_grid}) + assert len(totals) <= 4 + + def test_prefill_grid_uniformly_limits_new_tokens_batch_and_kv_axes(): stub = _prefill_grid_stub(num_gpu_blocks=512) stub.max_num_scheduled_tokens = 40 @@ -4100,8 +4127,7 @@ def test_prefill_point_with_exact_batch_shape_is_saved(): ] -@pytest.mark.parametrize("fpm_count", [0, 2]) -def test_benchmark_point_rejects_non_single_fpm_count(fpm_count): +def _decode_point_and_fpm(): point = BenchmarkPoint( point_type="decode", benchmark_id=4, @@ -4114,12 +4140,32 @@ def test_benchmark_point_rejects_non_single_fpm_count(fpm_count): "sum_decode_kv_tokens": 48, } } - stub = _benchmark_save_stub(point, [fpm.copy() for _ in range(fpm_count)]) + return point, fpm + + +def test_benchmark_point_rejects_multiple_fpms(): + # Two FPMs under one benchmark_id means the point isolation broke: abort. + point, fpm = _decode_point_and_fpm() + stub = _benchmark_save_stub(point, [fpm.copy() for _ in range(2)]) with pytest.raises(RuntimeError, match="exactly one FPM"): InstrumentedScheduler._bench_save_current_point(stub) +def test_benchmark_point_with_no_fpm_is_skipped_not_aborted(): + # The deadline passed before a single FPM was recorded (first pass at a + # fresh giant shape, kernel JIT): a group-synchronized skip, not a sweep abort. + point, _ = _decode_point_and_fpm() + stub = _benchmark_save_stub(point, []) + + InstrumentedScheduler._bench_save_current_point(stub) + + assert stub._bench_results == [] + assert stub._bench_skipped_points == [ + SkippedBenchmarkPoint(point=point, reason="no_fpm_before_deadline") + ] + + def test_decode_point_with_no_fpm_stops_waiting_at_deadline(monkeypatch): point = BenchmarkPoint( point_type="decode", @@ -4468,7 +4514,8 @@ def test_save_records_only_the_steady_fpm(): InstrumentedScheduler._bench_save_current_point(stub) assert len(stub._bench_results) == 1 - assert stub._bench_results[0].fpms == [steady] + # the kept sample is marked as a recorded steady step + assert stub._bench_results[0].fpms == [{**steady, "kvwarm_steady_sample": True}] assert stub._bench_skipped_points == [] @@ -5285,11 +5332,12 @@ def test_kvwarm_prepare_reserves_shadow_tail_blocks(monkeypatch): @pytest.mark.core @pytest.mark.parametrize("ctx", [2, 1000]) def test_kvwarm_does_not_build_a_stage_over_budget_at_the_depth_floor(ctx, monkeypatch): - # Four one-block chains fit, but their two private tail blocks per - # request raise the warmup bound to twelve. Reaching depth 8 (or starting - # below it) must not mark the short point as covered by real KV. + # Four one-block chains fit, but the private tail block each shadow takes + # at these contexts (exact per-rung reserve, one block per group here) + # raises the warmup bound to eight. Reaching depth 8 (or starting below + # it) must not mark the short point as covered by real KV. monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") - stub = _kvwarm_planner_stub(usable_blocks=11) + stub = _kvwarm_planner_stub(usable_blocks=7) short = BenchmarkPoint(point_type="decode", batch_size=4, total_kv_read_tokens=8) deepest = replace(short, total_kv_read_tokens=4 * ctx) stub._bench_grid = deque([deepest, short] if ctx > 2 else [short]) @@ -5304,8 +5352,8 @@ def test_kvwarm_does_not_build_a_stage_over_budget_at_the_depth_floor(ctx, monke { "batch": 4, "depth": min(ctx + 4, 8), - "required_blocks": 12, - "usable_blocks": 11, + "required_blocks": 8, + "usable_blocks": 7, } ] stub._bench_active_req_ids = set() @@ -5662,6 +5710,329 @@ def test_kvwarm_shadow_pool_shortfall_matches_tail_arithmetic(): assert InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [40, 47], 3) == 0 +def test_kvwarm_shadow_registration_keeps_positional_table_for_sliding_window(): + """An admission cap bounds resident blocks, not the positional table: a + sliding-window group keeps absolute positions (null placeholders included), + so the shadow at ctx=152 forks table index 9, never the last ``cap`` entries.""" + stub, mgr, pool, chain = _shadow_stub(cow=True) + mgr._max_admission_blocks_per_request = 4 + mgr.kv_cache_spec = SimpleNamespace(participates_in_prefix_caching=True) + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 152, 3 + ) + assert table == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 1000],) + assert mgr.cows == [(9, 1000)] + assert zero_ids == [] + assert mgr.num_cached_block["shadow"] == 9 + + +def test_kvwarm_shadow_registration_forks_circular_tail_table(): + """GLM5-Next's k-pool tail: one circularly reused block per request + (admission cap 1, excluded from prefix caching). The chain holds a single + block whatever its depth; the shadow shares nothing and forks that block.""" + stub, mgr, pool, chain = _shadow_stub(cow=True) + mgr.req_to_blocks["chain"] = chain[:1] + mgr.block_size = 4 + mgr._max_admission_blocks_per_request = 1 + mgr.kv_cache_spec = SimpleNamespace(participates_in_prefix_caching=False) + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 152, 3 + ) + assert table == ([1000],) + assert mgr.cows == [(0, 1000)] + assert zero_ids == [] + assert mgr.num_cached_block["shadow"] == 0 + + +class _FakeMambaSpec: + """Spec whose class name carries "Mamba", the key of the recurrent-group checks.""" + + +class MambaManager(_FakeManager): + def __init__(self, chain_blocks, cow=True): + super().__init__(chain_blocks, cow=cow) + self.kv_cache_spec = _FakeMambaSpec() + + +def test_kvwarm_live_state_shadow_forks_the_recurrent_read_slot_at_a_boundary(): + """At an exact block boundary vLLM reads the previous state at + ceil(ctx/bs)-1, one position below ctx//bs. The live-state fork must cover + that read slot too (``recurrent_shadow_range``), or the shadow starts from a + pruned checkpoint: ctx=32, bs=16 -> positions 1..2 forked from the live block 9.""" + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + chain = [_FakeBlock(i) for i in range(10)] + mgr = MambaManager(chain, cow=True) + pool = _FakePool() + stub.kv_cache_manager = SimpleNamespace( + block_pool=pool, coordinator=SimpleNamespace(single_type_managers=[mgr]) + ) + stub.cache_config = SimpleNamespace(block_size=16) + stub._bench_hybrid_live_state = True + stub._bench_random_kda = False + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 32, 3 + ) + assert table == ([0, 1000, 1001],) + assert mgr.cows == [(9, 1000), (9, 1001)] + assert zero_ids == [] + assert mgr.num_cached_block["shadow"] == 1 + # the per-context reserve and the pool-shortfall mirror agree on the span + assert InstrumentedScheduler._kvwarm_shadow_tail_blocks_for(stub, 32, 3) == 2 + assert InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [32], 3) == 0 + + +class KpoolTailManager(_FakeManager): + """Matches the circular-table predicate (admission cap, spec excluded from + prefix caching); its ``SimpleNamespace`` spec is not a Mamba spec, so the + live-state predicate does not match.""" + + +def test_kvwarm_live_state_keeps_circular_kpool_geometry(): + """With live-state on, the k-pool tail (one circular block, admission cap 1, + excluded from prefix caching) must keep its ring geometry: fork the single + block, never walk ``recurrent_shadow_range`` positions (65 entries at ctx 255).""" + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + chain = [_FakeBlock(0)] + mgr = KpoolTailManager(chain, cow=True) + mgr.block_size = 4 + mgr._max_admission_blocks_per_request = 1 + mgr.kv_cache_spec = SimpleNamespace(participates_in_prefix_caching=False) + pool = _FakePool() + stub.kv_cache_manager = SimpleNamespace( + block_pool=pool, coordinator=SimpleNamespace(single_type_managers=[mgr]) + ) + stub.cache_config = SimpleNamespace(block_size=16) + stub._bench_hybrid_live_state = True + stub._bench_random_kda = False + table, zero_ids = InstrumentedScheduler._kvwarm_register_shadow( + stub, "shadow", "chain", 255, 3 + ) + assert table == ([1000],) + assert mgr.cows == [(0, 1000)] + assert zero_ids == [] + assert InstrumentedScheduler._kvwarm_shadow_tail_blocks_for(stub, 255, 3) == 1 + assert InstrumentedScheduler._kvwarm_shadow_pool_shortfall(stub, [255], 3) == 0 + assert InstrumentedScheduler._kvwarm_shadow_tail_blocks(stub, 3) == 2 + + +def test_kvwarm_plan_keeps_kimi_long_context_real_kv_coverage(monkeypatch): + """Kimi K3 TP8/DCP8 layout (three align-mode recurrent groups, block 1536; + one attention group, effective block 12288): a 1M-token batch-8 point must + stay on real-KV warm-up. The resident recurrent estimate is the allocator's + (align pair + prefill checkpoint), not a token-proportional retention term.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + monkeypatch.setenv("DYN_BENCH_GIANT_KV_REPEATS", "3") + stub = _kvwarm_planner_stub(usable_blocks=2144) + del stub._bench_blocks_per_req # exercise the real footprint arithmetic + stub.max_model_len = 1048576 + stub.block_size = 12288 + stub.cache_config = SimpleNamespace(block_size=12288) + managers = [ + SimpleNamespace( + block_size=1536, + mamba_cache_mode="align", + num_speculative_blocks=0, + kv_cache_spec=SimpleNamespace(num_prefill_checkpoint_blocks=1), + ) + for _ in range(3) + ] + managers.append(SimpleNamespace(block_size=12288)) + stub.kv_cache_manager = SimpleNamespace( + coordinator=SimpleNamespace(single_type_managers=managers) + ) + context = 1_000_000 + point = BenchmarkPoint( + point_type="decode", + benchmark_id=1, + batch_size=8, + total_kv_read_tokens=8 * context, + ) + stub._bench_grid = deque([point]) + stub._kvwarm_prepare("decode") + assert stub._bench_blocks_per_req(context + 4, resident_chain=True) == 91 + assert stub._kvwarm_plan[8] == context + 4 + assert stub._kvwarm_plan_covers(point) + + +def test_kvwarm_reserve_uses_the_admission_context_at_a_block_boundary(monkeypatch): + """The shadow is admitted at ctx-1 with ``repeats`` steady steps. At a block + boundary the recurrent read slot moves one position down, so the reserve + must be taken at that geometry: ctx 33 -> admission 32 needs 3 private + blocks (1 attention + recurrent positions 1..2), not the 2 that + ``(33, 1 + repeats)`` suggests. A pool that fits only the smaller reserve + must NOT admit the warm stage at the measured depth.""" + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + monkeypatch.setenv("DYN_BENCH_GIANT_KV_REPEATS", "3") + stub = _kvwarm_planner_stub(usable_blocks=37, groups=2) + stub._bench_hybrid_live_state = True + stub._bench_random_kda = False + stub.kv_cache_manager.coordinator.single_type_managers[ + 1 + ].kv_cache_spec = _FakeMambaSpec() + assert InstrumentedScheduler._kvwarm_shadow_tail_blocks_for(stub, 32, 3) == 3 + assert InstrumentedScheduler._kvwarm_shadow_tail_blocks_for(stub, 33, 4) == 2 + point = BenchmarkPoint( + point_type="decode", benchmark_id=1, batch_size=4, total_kv_read_tokens=4 * 33 + ) + stub._bench_grid = deque([point]) + stub._kvwarm_prepare("decode") + # depth 37 with the 3-block reserve needs (6 + 3) * 4 = 36 > 35 (0.95 pool): + # the rung is trimmed below the point instead of admitting an underfunded stage + assert stub._kvwarm_plan[4] < 37 + assert not stub._kvwarm_plan_covers(point) + + +def _dp_planner_stub(monkeypatch, points, usable_blocks=100): + monkeypatch.setenv("DYN_BENCH_KV_WARMUP", "on") + stub = _kvwarm_planner_stub(usable_blocks=usable_blocks) + stub._bench_dp_size = 2 + stub._bench_expected_points = sum( + EAGER_WARMUP_REASON not in p.sample_reasons for p in points + ) + stub._bench_grid = deque(points) + return stub + + +def test_kvwarm_dp_filter_rejects_uncovered_explicit_points(monkeypatch): + covered = BenchmarkPoint( + point_type="decode", + benchmark_id=1, + batch_size=1, + total_kv_read_tokens=16, + sample_reasons=["explicit"], + ) + deep = BenchmarkPoint( + point_type="decode", + benchmark_id=2, + batch_size=3, + total_kv_read_tokens=1500, + sample_reasons=["explicit"], + ) + stub = _dp_planner_stub(monkeypatch, [covered, deep]) + with pytest.raises( + RuntimeError, match=r"explicit decode point.*batch=3, total_kv_read_tokens=1500" + ): + stub._kvwarm_prepare("decode") + + +def test_kvwarm_dp_filter_counts_only_real_points(monkeypatch): + covered = BenchmarkPoint( + point_type="decode", benchmark_id=1, batch_size=1, total_kv_read_tokens=16 + ) + replica = replace(covered, benchmark_id=3, sample_reasons=[EAGER_WARMUP_REASON]) + deep = BenchmarkPoint( + point_type="decode", benchmark_id=2, batch_size=3, total_kv_read_tokens=1500 + ) + stub = _dp_planner_stub(monkeypatch, [covered, deep, replica]) + stub._kvwarm_prepare("decode") + kept = list(stub._bench_grid) + assert [(p.batch_size, p.total_kv_read_tokens) for p in kept] == [(1, 16), (1, 16)] + assert stub._bench_expected_points == 1 + # IDs are left to ``_bench_build_grid``, which numbers the final order. + assert {p.benchmark_id for p in kept} == {1, 3} + + +def test_kvwarm_dp_filter_marks_decode_missing_when_nothing_is_covered(monkeypatch): + """Attention-DP keeps real-KV decode points only. When the plan covers none of + them the decode phase is gone, and the artifact must not report a complete, + usable run with zero decode measurements.""" + deep = BenchmarkPoint( + point_type="decode", benchmark_id=1, batch_size=3, total_kv_read_tokens=1500 + ) + stub = _dp_planner_stub(monkeypatch, [deep]) + stub._bench_missing_phases = [] + stub._kvwarm_prepare("decode") + assert list(stub._bench_grid) == [] + assert stub._bench_expected_points == 0 + assert stub._bench_missing_phases == ["decode"] + + +def test_giant_fake_off_by_batch_correction_requires_a_steady_sample(): + """The admission step also measures ``declared - batch``; only a recorded steady + sample (``kvwarm_steady_sample``, set by both save paths) may be accepted at the + measured coordinate. A giant fake point that reached its deadline with the + admission FPM alone is a validation skip, not a decode measurement.""" + stub = SimpleNamespace(_kvwarm_giant_threshold=lambda: 1000) + point = BenchmarkPoint( + point_type="decode", + benchmark_id=1, + batch_size=2, + total_kv_read_tokens=2000, + sample_reasons=["kvwarm_fake_fallback"], + ) + scheduled = {"num_decode_requests": 2, "sum_decode_kv_tokens": 1998} + admission_only = {"scheduled_requests": scheduled} + assert ( + InstrumentedScheduler._bench_fpm_validation_failure(stub, point, admission_only) + == "measured_decode_context_mismatch" + ) + steady = {"scheduled_requests": scheduled, "kvwarm_steady_sample": True} + assert ( + InstrumentedScheduler._bench_fpm_validation_failure(stub, point, steady) is None + ) + # the median marker alone is not the steady evidence + median_only = {"scheduled_requests": scheduled, "kvwarm_giant_median_of": 3} + assert ( + InstrumentedScheduler._bench_fpm_validation_failure(stub, point, median_only) + == "measured_decode_context_mismatch" + ) + + +def _giant_fake_point(): + return BenchmarkPoint( + point_type="decode", + benchmark_id=1, + batch_size=2, + total_kv_read_tokens=2000, + sample_reasons=["kvwarm_fake_fallback"], + ) + + +def test_two_fpm_save_path_records_the_steady_sample_at_the_measured_coordinate( + monkeypatch, +): + """`DYN_BENCH_GIANT_KV_REPEATS=1` (or a pool/model-length limit) reduces a giant + fake point to admission plus one steady step. That path keeps the steady FPM + and must mark it as recorded, so the off-by-batch correction accepts it and the + point is saved at the measured coordinate.""" + monkeypatch.setenv("DYN_BENCH_GIANT_KV_THRESHOLD", "1000") + scheduled = {"num_decode_requests": 2, "sum_decode_kv_tokens": 1998} + stub = _benchmark_save_stub( + _giant_fake_point(), + [ + {"scheduled_requests": dict(scheduled), "wall_time": 0.5}, + {"scheduled_requests": dict(scheduled), "wall_time": 0.02}, + ], + ) + stub._bench_expected_fpms = 2 + + InstrumentedScheduler._bench_save_current_point(stub) + + assert stub._bench_skipped_points == [] + (result,) = stub._bench_results + assert result.point.total_kv_read_tokens == 1998 + assert "giant_fake_off_by_batch" in result.point.sample_reasons + (fpm,) = result.fpms + assert fpm["wall_time"] == 0.02, "the admission step is scaffolding" + assert fpm["kvwarm_steady_sample"] is True + + +def test_two_fpm_save_path_skips_an_admission_only_giant_sample(monkeypatch): + monkeypatch.setenv("DYN_BENCH_GIANT_KV_THRESHOLD", "1000") + scheduled = {"num_decode_requests": 2, "sum_decode_kv_tokens": 1998} + stub = _benchmark_save_stub( + _giant_fake_point(), [{"scheduled_requests": dict(scheduled)}] + ) + stub._bench_expected_fpms = 2 + + InstrumentedScheduler._bench_save_current_point(stub) + + assert stub._bench_results == [] + assert [s.reason for s in stub._bench_skipped_points] == [ + "measured_decode_context_mismatch" + ] + + def test_kvwarm_shadow_registration_rejects_too_shallow_chain(): stub, mgr, pool, chain = _shadow_stub() with pytest.raises(RuntimeError, match="too shallow"): diff --git a/components/src/dynamo/vllm/tests/test_vllm_unit.py b/components/src/dynamo/vllm/tests/test_vllm_unit.py index 882813faa7a6..5c16aa52ce26 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unit.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unit.py @@ -971,9 +971,11 @@ def test_benchmark_operational_controls_reach_scheduler_config( assert config._benchmark_additional_config == { "mode": "prefill", "randomize_kda_state": False, + "hybrid_live_state": False, "warmup_iterations": 2, "output_path": str(output), "timeout": 900, + "max_batch_size": None, "prefill_max_new_token_samples": 64, "prefill_max_kv_read_token_samples": 16, "decode_max_kv_read_token_samples": 128, @@ -1696,6 +1698,8 @@ def _make_dynamo_config(**overrides): "fpm_trace": False, "benchmark_mode": None, "benchmark_randomize_kda_state": False, + "benchmark_hybrid_live_state": False, + "benchmark_max_batch_size": None, "benchmark_warmup_iterations": 5, "benchmark_output_path": "/tmp/benchmark_results.json", "benchmark_timeout": 900,