diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index 10c58cf84739..17e6d91a42ef 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -196,6 +196,88 @@ class _BenchmarkGroupResult: stop_requested: bool +@dataclass(frozen=True) +class _BenchmarkCapacityEnvelope: + """Rank-local limits that affect synthetic benchmark grid feasibility.""" + + max_model_len: int + max_num_scheduled_tokens: int + max_num_running_reqs: int + usable_blocks_without_watermark: int + usable_blocks_with_watermark: int + grid_invariants_digest: str + + @classmethod + def from_dict(cls, payload: object) -> _BenchmarkCapacityEnvelope: + if not isinstance(payload, dict): + raise RuntimeError("attention-DP benchmark capacity must be an object") + positive_fields = ( + "max_model_len", + "max_num_scheduled_tokens", + "max_num_running_reqs", + ) + nonnegative_fields = ( + "usable_blocks_without_watermark", + "usable_blocks_with_watermark", + ) + values: dict[str, int] = {} + for name in positive_fields: + value = payload.get(name) + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise RuntimeError( + f"attention-DP benchmark capacity has invalid {name}={value!r}" + ) + values[name] = value + for name in nonnegative_fields: + value = payload.get(name) + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise RuntimeError( + f"attention-DP benchmark capacity has invalid {name}={value!r}" + ) + values[name] = value + digest = payload.get("grid_invariants_digest") + if not isinstance(digest, str) or len(digest) != 64: + raise RuntimeError( + "attention-DP benchmark capacity has invalid grid invariants digest" + ) + return cls( + max_model_len=values["max_model_len"], + max_num_scheduled_tokens=values["max_num_scheduled_tokens"], + max_num_running_reqs=values["max_num_running_reqs"], + usable_blocks_without_watermark=values["usable_blocks_without_watermark"], + usable_blocks_with_watermark=values["usable_blocks_with_watermark"], + grid_invariants_digest=digest, + ) + + @classmethod + def common( + cls, capacities: Sequence[_BenchmarkCapacityEnvelope] + ) -> _BenchmarkCapacityEnvelope: + if not capacities: + raise RuntimeError("attention-DP benchmark has no capacity reports") + invariant_digests = {capacity.grid_invariants_digest for capacity in capacities} + if len(invariant_digests) != 1: + raise RuntimeError( + "attention-DP benchmark grid invariants differ across ranks" + ) + return cls( + max_model_len=min(capacity.max_model_len for capacity in capacities), + max_num_scheduled_tokens=min( + capacity.max_num_scheduled_tokens for capacity in capacities + ), + max_num_running_reqs=min( + capacity.max_num_running_reqs for capacity in capacities + ), + usable_blocks_without_watermark=min( + capacity.usable_blocks_without_watermark for capacity in capacities + ), + usable_blocks_with_watermark=min( + capacity.usable_blocks_with_watermark for capacity in capacities + ), + grid_invariants_digest=capacities[0].grid_invariants_digest, + ) + + def _benchmark_point_digest(point: BenchmarkPoint) -> str: payload = json.dumps( asdict(point), @@ -406,6 +488,67 @@ def close(self) -> None: ) self._socket.close(linger=linger) + def negotiate_capacity( + self, local_capacity: _BenchmarkCapacityEnvelope + ) -> _BenchmarkCapacityEnvelope: + """Agree on the minimum capacity that every attention-DP rank can run.""" + message = { + "type": "capacity", + "benchmark_id": 0, + "dp_rank": self.dp_rank, + "capacity": asdict(local_capacity), + } + if self.dp_rank == 0: + return self._coordinate_capacity(message) + + deadline = time.monotonic() + self.timeout_seconds + self._socket.send_json(message) + reply = self._recv_follower(deadline, 0, "capacity_result") + common_capacity = _BenchmarkCapacityEnvelope.from_dict(reply.get("capacity")) + self._socket.send_json( + { + "type": "capacity_ack", + "benchmark_id": 0, + "dp_rank": self.dp_rank, + } + ) + deadline = time.monotonic() + self.timeout_seconds + self._recv_follower(deadline, 0, "capacity_commit") + return common_capacity + + def synchronize_grid( + self, + *, + grid_digest: str, + expected_points: int, + missing_phases: Sequence[str], + ) -> None: + """Verify that all ranks built the same complete grid before warmup.""" + message = { + "type": "grid", + "benchmark_id": 0, + "dp_rank": self.dp_rank, + "grid_digest": grid_digest, + "expected_points": expected_points, + "missing_phases": list(missing_phases), + } + if self.dp_rank == 0: + self._coordinate_grid(message) + return + + deadline = time.monotonic() + self.timeout_seconds + self._socket.send_json(message) + self._recv_follower(deadline, 0, "grid_prepare") + self._socket.send_json( + { + "type": "grid_ack", + "benchmark_id": 0, + "dp_rank": self.dp_rank, + } + ) + deadline = time.monotonic() + self.timeout_seconds + self._recv_follower(deadline, 0, "grid_commit") + def synchronize( self, point: BenchmarkPoint, @@ -443,6 +586,114 @@ def synchronize( self._run_id = run_id return run_id + def _coordinate_capacity(self, local_message: dict) -> _BenchmarkCapacityEnvelope: + capacities = [ + _BenchmarkCapacityEnvelope.from_dict(local_message.get("capacity")) + ] + identities: dict[int, bytes] = {} + seen_identities: set[bytes] = set() + deadline = time.monotonic() + self.timeout_seconds + try: + while len(identities) < self.dp_size - 1: + identity, message = self._recv_router(deadline, 0) + seen_identities.add(identity) + rank = message.get("dp_rank") + if ( + message.get("type") != "capacity" + or not isinstance(rank, int) + or not 1 <= rank < self.dp_size + or rank in identities + or identity != str(rank).encode() + ): + raise RuntimeError( + f"invalid attention-DP benchmark capacity: {message}" + ) + capacities.append( + _BenchmarkCapacityEnvelope.from_dict(message.get("capacity")) + ) + identities[rank] = identity + + common_capacity = _BenchmarkCapacityEnvelope.common(capacities) + self._send_to_all( + identities, + { + "type": "capacity_result", + "benchmark_id": 0, + "capacity": asdict(common_capacity), + }, + ) + self._coordinate_phase( + identities, + time.monotonic() + self.timeout_seconds, + benchmark_id=0, + expected_type="capacity_ack", + ) + self._send_to_all( + identities, + {"type": "capacity_commit", "benchmark_id": 0}, + ) + return common_capacity + except Exception as error: + self._notify_error(seen_identities, str(error)) + raise + + def _coordinate_grid(self, local_message: dict) -> None: + expected = { + "grid_digest": local_message.get("grid_digest"), + "expected_points": local_message.get("expected_points"), + "missing_phases": local_message.get("missing_phases"), + } + identities: dict[int, bytes] = {} + seen_identities: set[bytes] = set() + mismatches: list[str] = [] + deadline = time.monotonic() + self.timeout_seconds + try: + while len(identities) < self.dp_size - 1: + identity, message = self._recv_router(deadline, 0) + seen_identities.add(identity) + rank = message.get("dp_rank") + if ( + message.get("type") != "grid" + or not isinstance(rank, int) + or not 1 <= rank < self.dp_size + or rank in identities + or identity != str(rank).encode() + ): + raise RuntimeError( + f"invalid attention-DP benchmark grid report: {message}" + ) + actual = { + "grid_digest": message.get("grid_digest"), + "expected_points": message.get("expected_points"), + "missing_phases": message.get("missing_phases"), + } + if actual != expected: + mismatches.append( + "attention-DP benchmark grid mismatch on " + f"rank {rank}: rank0={expected} rank{rank}={actual}" + ) + identities[rank] = identity + + if mismatches: + raise RuntimeError("; ".join(mismatches)) + self._send_to_all( + identities, + {"type": "grid_prepare", "benchmark_id": 0}, + ) + self._coordinate_phase( + identities, + time.monotonic() + self.timeout_seconds, + benchmark_id=0, + expected_type="grid_ack", + ) + self._send_to_all( + identities, + {"type": "grid_commit", "benchmark_id": 0}, + ) + except Exception as error: + self._notify_error(seen_identities, str(error)) + raise + def collect_result( self, point: BenchmarkPoint, @@ -813,8 +1064,7 @@ def _coordinate_results( rank_stop_requested = message.get("stop_requested") if not isinstance(rank_stop_requested, bool): raise RuntimeError( - f"attention-DP benchmark rank {rank} sent invalid " - "stop decision" + f"attention-DP benchmark rank {rank} sent invalid stop decision" ) stop_requested = stop_requested or rank_stop_requested identities[rank] = identity @@ -1811,6 +2061,8 @@ def _bench_init(self, vllm_config: "VllmConfig") -> None: self._bench_prefix_cache_cleared = False self._bench_grid_error: str | None = None self._bench_grid_digest: str | None = None + self._bench_local_capacity: _BenchmarkCapacityEnvelope | None = None + self._bench_negotiated_capacity: _BenchmarkCapacityEnvelope | None = None self._bench_started_at: str | None = None self._bench_completed_at: str | None = None self._bench_start_monotonic: float | None = None @@ -1878,10 +2130,134 @@ def _bench_init(self, vllm_config: "VllmConfig") -> None: # -- Grid generation ------------------------------------------------ + def _bench_grid_invariants_digest(self) -> str: + coordinator = getattr( + getattr(self, "kv_cache_manager", None), "coordinator", None + ) + managers = getattr(coordinator, "single_type_managers", ()) + manager_layout = [] + for manager in managers: + manager_layout.append( + { + "type": ( + f"{type(manager).__module__}.{type(manager).__qualname__}" + ), + "block_size": getattr(manager, "block_size", None), + "admission_cap": getattr( + manager, "_max_admission_blocks_per_request", None + ), + "mamba_cache_mode": getattr(manager, "mamba_cache_mode", None), + "num_speculative_blocks": getattr( + manager, "num_speculative_blocks", 0 + ), + "cross_attention": isinstance(manager, CrossAttentionManager), + } + ) + benchmark_config = { + name: getattr(self._bench_config, name) + for name in self._bench_config.__dataclass_fields__ + if name != "output_path" + } + scheduler_config = getattr(self, "scheduler_config", None) + payload = { + "benchmark_config": benchmark_config, + "block_size": self.block_size, + "hash_block_size": self._bench_hash_block_size, + "cache_block_size": getattr(self.cache_config, "block_size", None), + "enable_prefix_caching": getattr( + self.cache_config, "enable_prefix_caching", True + ), + "num_lookahead_tokens": getattr(self, "num_lookahead_tokens", 0), + "long_prefill_token_threshold": getattr( + scheduler_config, "long_prefill_token_threshold", 0 + ), + "need_mamba_block_aligned_split": getattr( + self, "need_mamba_block_aligned_split", False + ), + "use_eagle": getattr( + getattr(self, "kv_cache_manager", None), "use_eagle", False + ), + "uses_per_group_cache_lookup": self._bench_uses_per_group_cache_lookup(), + "manager_layout": manager_layout, + "prefill_cudagraph_mode": self._bench_prefill_cudagraph_mode, + "decode_cudagraph_mode": self._bench_decode_cudagraph_mode, + "prefill_capture_sizes": self._bench_prefill_capture_sizes, + # ``_bench_decode_capture_sizes`` is filtered by + # ``max_num_running_reqs`` — a negotiable capacity value that may + # legitimately differ across ranks — so hash the unfiltered + # configuration and re-filter after negotiation. + "decode_capture_sizes": ( + list(self._bench_cudagraph_capture_sizes) + if self._bench_decode_cudagraph_mode != "NONE" + else [] + ), + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + def _bench_make_local_capacity(self) -> _BenchmarkCapacityEnvelope: + available_blocks = self._bench_available_blocks() + watermark_blocks = max( + 0, + int( + getattr(getattr(self, "kv_cache_manager", None), "watermark_blocks", 0) + ), + ) + return _BenchmarkCapacityEnvelope( + max_model_len=int(self.max_model_len), + max_num_scheduled_tokens=int(self.max_num_scheduled_tokens), + max_num_running_reqs=int(self.max_num_running_reqs), + usable_blocks_without_watermark=available_blocks, + usable_blocks_with_watermark=max(0, available_blocks - watermark_blocks), + grid_invariants_digest=self._bench_grid_invariants_digest(), + ) + + def _bench_capacity_limit(self, name: str) -> int: + capacity = getattr(self, "_bench_negotiated_capacity", None) + if capacity is not None: + return int(getattr(capacity, name)) + return int(getattr(self, name)) + + def _bench_grid_usable_blocks( + self, batch_size: int, *, reserve_watermark: bool = False + ) -> int: + capacity = getattr(self, "_bench_negotiated_capacity", None) + if capacity is None: + return self._bench_usable_blocks( + batch_size, reserve_watermark=reserve_watermark + ) + if reserve_watermark or batch_size > 1: + return capacity.usable_blocks_with_watermark + return capacity.usable_blocks_without_watermark + def _bench_build_grid(self) -> None: """Generate the sweep grid once scheduler limits are known.""" if self._bench_grid_built: return + + local_capacity = self._bench_make_local_capacity() + self._bench_local_capacity = local_capacity + synchronizer = getattr(self, "_bench_synchronizer", None) + if synchronizer is not None: + common_capacity = synchronizer.negotiate_capacity(local_capacity) + else: + common_capacity = local_capacity + self._bench_negotiated_capacity = common_capacity + logger.info( + "Benchmark capacity: rank=%d local=%s common=%s", + getattr(self, "_fpm_dp_rank", 0), + asdict(local_capacity), + asdict(common_capacity), + ) + # The activation-time filter used the local request limit; re-filter + # with the negotiated one so every rank builds the decode grid from + # the same capture list. + self._bench_decode_capture_sizes = [ + size + for size in self._bench_decode_capture_sizes + if size <= common_capacity.max_num_running_reqs + ] + self._bench_grid_built = True mode = self._bench_config.mode explicit_points = self._bench_explicit_points @@ -1909,6 +2285,12 @@ def _bench_build_grid(self) -> None: separators=(",", ":"), ).encode() self._bench_grid_digest = hashlib.sha256(grid_payload).hexdigest() + if synchronizer is not None: + synchronizer.synchronize_grid( + grid_digest=self._bench_grid_digest, + expected_points=self._bench_expected_points, + missing_phases=self._bench_missing_phases, + ) logger.info("Benchmark grid: %d points (%s mode)", len(self._bench_grid), mode) def _bench_build_explicit_grid(self, points: BenchmarkPoints) -> None: @@ -1947,7 +2329,7 @@ def _bench_materialize_prefill_candidate( capture_size, padding_tokens, reasons = self._bench_cudagraph_metadata( candidate.total_prefill_tokens, self._bench_prefill_capture_sizes, - self.max_num_scheduled_tokens, + self._bench_capacity_limit("max_num_scheduled_tokens"), ) return BenchmarkPoint( point_type="prefill", @@ -2011,7 +2393,7 @@ def _bench_raise_explicit_infeasible( ) def _bench_generate_prefill_grid(self) -> None: - max_tokens = self.max_num_scheduled_tokens + max_tokens = self._bench_capacity_limit("max_num_scheduled_tokens") if max_tokens < 1: logger.warning( "max_num_scheduled_tokens=%d too small, skipping prefill grid", @@ -2072,8 +2454,8 @@ def _bench_prefill_batch_sizes(self, total_tokens: int) -> list[int]: """Return the smallest configured presets from the legal batch axis.""" upper_bound = min( total_tokens, - self.max_num_running_reqs, - self.max_num_scheduled_tokens, + self._bench_capacity_limit("max_num_running_reqs"), + self._bench_capacity_limit("max_num_scheduled_tokens"), ) legal_batches = [ batch_size @@ -2179,9 +2561,10 @@ def _bench_prefill_point_feasible( ) -> bool: if ( total_prefill_tokens < 1 - or total_prefill_tokens > self.max_num_scheduled_tokens + or total_prefill_tokens + > self._bench_capacity_limit("max_num_scheduled_tokens") or batch_size < 1 - or batch_size > self.max_num_running_reqs + or batch_size > self._bench_capacity_limit("max_num_running_reqs") ): return False try: @@ -2198,7 +2581,8 @@ def _bench_prefill_point_feasible( new_tokens + kv_read_tokens for new_tokens, kv_read_tokens in zip(new_token_lengths, kv_read_lengths) ] - if any(prompt_len + 1 > self.max_model_len for prompt_len in prompt_lengths): + max_model_len = self._bench_capacity_limit("max_model_len") + if any(prompt_len + 1 > max_model_len for prompt_len in prompt_lengths): return False if any( self._bench_prefill_scheduled_tokens_per_req(prompt_len, kv_read_tokens) @@ -2219,8 +2603,7 @@ def _bench_prefill_point_feasible( for kv_read_tokens in kv_read_lengths ] if any( - prompt_len + 1 > self.max_model_len - for prompt_len in seed_prompt_lengths + prompt_len + 1 > max_model_len for prompt_len in seed_prompt_lengths ): return False seed_required_blocks = sum( @@ -2232,7 +2615,7 @@ def _bench_prefill_point_feasible( for prompt_len in seed_prompt_lengths ) required_blocks = max(required_blocks, seed_required_blocks) - return required_blocks <= self._bench_usable_blocks(batch_size) + return required_blocks <= self._bench_grid_usable_blocks(batch_size) def _bench_prefill_blocks_per_req(self, isl: int, kv_read_tokens: int) -> int: tokens_with_lookahead = isl + getattr(self, "num_lookahead_tokens", 0) @@ -2354,7 +2737,11 @@ def _bench_max_prefill_kv_read_tokens( hash_block_size = max(1, self._bench_hash_block_size) max_blocks = sum( - max(0, self.max_model_len - new_tokens - 1) // hash_block_size + max( + 0, + self._bench_capacity_limit("max_model_len") - new_tokens - 1, + ) + // hash_block_size for new_tokens in new_token_lengths ) if max_blocks < batch_size: @@ -2410,7 +2797,8 @@ def _bench_cached_kv_read_tokens(self, req: Request) -> int: return cached_tokens def _bench_generate_decode_grid(self) -> None: - if self.max_model_len < 3: + max_model_len = self._bench_capacity_limit("max_model_len") + if max_model_len < 3: logger.warning("max_model_len too small for decode grid, skipping") return @@ -2464,22 +2852,30 @@ def _bench_generate_decode_grid(self) -> None: ) def _bench_decode_feasible_max_batch_size(self) -> int: - if self.max_model_len < 3: + max_model_len = self._bench_capacity_limit("max_model_len") + max_num_running_reqs = self._bench_capacity_limit("max_num_running_reqs") + max_num_scheduled_tokens = self._bench_capacity_limit( + "max_num_scheduled_tokens" + ) + if max_model_len < 3: return 0 min_blocks_per_request = self._bench_blocks_per_req(2) if min_blocks_per_request < 1: - feasible_max_batch = self.max_num_running_reqs + feasible_max_batch = max_num_running_reqs else: feasible_max_batch = ( - self._bench_usable_blocks( - self.max_num_running_reqs, reserve_watermark=True + self._bench_grid_usable_blocks( + max_num_running_reqs, reserve_watermark=True ) // min_blocks_per_request ) - return min( - self.max_num_running_reqs, - self.max_num_scheduled_tokens, - feasible_max_batch, + return max( + 0, + min( + max_num_running_reqs, + max_num_scheduled_tokens, + feasible_max_batch, + ), ) @staticmethod @@ -2511,7 +2907,9 @@ def _bench_decode_steady_kv_tokens( def _bench_decode_point_feasible( self, batch_size: int, total_kv_read_tokens: int ) -> bool: - if batch_size < 1 or batch_size > self.max_num_running_reqs: + if batch_size < 1 or batch_size > self._bench_capacity_limit( + "max_num_running_reqs" + ): return False try: context_lengths = self._bench_decode_context_lengths( @@ -2522,25 +2920,26 @@ def _bench_decode_point_feasible( # A decode point admits at max(1, ctx-1) and measures its steady step # one position later, so a request occupies max(ctx, 2) + 1 slots and # the runner's post-step bookkeeping writes through slot - # max(ctx, 2) + 2, which must stay within max_model_len (ctx + 2 for - # the ordinary ctx >= 2 case; one extra slot for clamped ctx = 1 - # entries, which run one token deeper than their nominal coordinate). + # max(ctx, 2) + 2, which must stay within the negotiated + # max_model_len (ctx + 2 for the ordinary ctx >= 2 case; one extra + # slot for clamped ctx = 1 entries, which run one token deeper than + # their nominal coordinate). + max_model_len = self._bench_capacity_limit("max_model_len") if any( - max(context_len, 2) + 2 > self.max_model_len - for context_len in context_lengths + max(context_len, 2) + 2 > max_model_len for context_len in context_lengths ): return False required_blocks = sum( self._bench_blocks_per_req(max(context_len, 2) + 1) for context_len in context_lengths ) - return required_blocks <= self._bench_usable_blocks( + return required_blocks <= self._bench_grid_usable_blocks( batch_size, reserve_watermark=True ) def _bench_max_decode_kv_read_tokens(self, batch_size: int) -> int: low = batch_size - high = batch_size * (self.max_model_len - 2) + high = batch_size * (self._bench_capacity_limit("max_model_len") - 2) best = 0 while low <= high: mid = (low + high) // 2 @@ -2948,8 +3347,7 @@ def _bench_abort(self, error: Exception) -> None: except Exception as prefix_error: cleanup_error = prefix_error self._bench_grid_error = ( - f"{self._bench_grid_error}; prefix-cache cleanup failed: " - f"{prefix_error}" + f"{self._bench_grid_error}; prefix-cache cleanup failed: {prefix_error}" ) try: self._bench_write_results() @@ -3622,6 +4020,18 @@ def _bench_write_results(self) -> None: "decode": "steady_state_second_step", "prefill": "single_step", }, + "capacity": { + "common": ( + asdict(negotiated_capacity) + if ( + negotiated_capacity := getattr( + self, "_bench_negotiated_capacity", None + ) + ) + is not None + else None + ), + }, "cudagraph": { "mode": getattr(self, "_bench_cudagraph_mode", "NONE"), "prefill_mode": getattr(self, "_bench_prefill_cudagraph_mode", "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 de7217b6bee3..7042bbfa778d 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -32,7 +32,10 @@ # not be resolvable and ``instrumented_scheduler`` will fail to load with # ``ModuleNotFoundError: No module named 'vllm.sampling_params'``. import dynamo.vllm.instrumented_scheduler as instrumented_scheduler_module # noqa: E402 -from dynamo.vllm.benchmark_points import BenchmarkPoints # noqa: E402 +from dynamo.vllm.benchmark_points import ( # noqa: E402 + BenchmarkPoints, + PrefillPointCandidate, +) from dynamo.vllm.instrumented_scheduler import ( # noqa: E402 BenchmarkConfig, BenchmarkPoint, @@ -53,6 +56,38 @@ ) or getattr(RequestStatus, "WAITING_FOR_FSM") +def _benchmark_capacity( + *, + max_model_len: int = 256, + max_num_scheduled_tokens: int = 10_000, + max_num_running_reqs: int = 10_000, + usable_blocks_without_watermark: int = 1_000, + usable_blocks_with_watermark: int | None = None, + grid_invariants_digest: str = "a" * 64, +): + if usable_blocks_with_watermark is None: + usable_blocks_with_watermark = usable_blocks_without_watermark + return instrumented_scheduler_module._BenchmarkCapacityEnvelope( + max_model_len=max_model_len, + max_num_scheduled_tokens=max_num_scheduled_tokens, + max_num_running_reqs=max_num_running_reqs, + usable_blocks_without_watermark=usable_blocks_without_watermark, + usable_blocks_with_watermark=usable_blocks_with_watermark, + grid_invariants_digest=grid_invariants_digest, + ) + + +def _install_test_capacity_preflight(stub, capacity=None): + capacity = capacity or _benchmark_capacity() + stub._bench_make_local_capacity = lambda: capacity + stub._bench_synchronizer = None + # ``_bench_build_grid`` re-filters the decode capture list against the + # negotiated request limit before generating the grid; stubs that don't + # model captures still need the attribute to exist. + if not hasattr(stub, "_bench_decode_capture_sizes"): + stub._bench_decode_capture_sizes = [] + + def _make_request(status, num_tokens: int, num_computed_tokens: int = 0): """Build a minimal stand-in for ``vllm.v1.request.Request``. @@ -522,6 +557,203 @@ def synchronize_follower(): rank0.close() +def test_benchmark_synchronizer_negotiates_minimum_capacity_and_grid(): + endpoint = f"inproc://benchmark-sync-{uuid.uuid4().hex}" + rank0 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=0, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + rank1 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=1, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + rank0_capacity = _benchmark_capacity( + max_model_len=383_168, + usable_blocks_without_watermark=23_944, + usable_blocks_with_watermark=23_940, + ) + rank1_capacity = _benchmark_capacity( + max_model_len=351_104, + max_num_scheduled_tokens=8_192, + usable_blocks_without_watermark=21_940, + usable_blocks_with_watermark=21_936, + ) + follower_result = {} + + def run_follower(): + follower_result["capacity"] = rank1.negotiate_capacity(rank1_capacity) + rank1.synchronize_grid( + grid_digest="b" * 64, + expected_points=1_368, + missing_phases=[], + ) + follower_result["grid_synchronized"] = True + + follower = threading.Thread(target=run_follower) + follower.start() + try: + common = rank0.negotiate_capacity(rank0_capacity) + rank0.synchronize_grid( + grid_digest="b" * 64, + expected_points=1_368, + missing_phases=[], + ) + follower.join(timeout=2) + assert not follower.is_alive() + assert common == follower_result["capacity"] + assert common.max_model_len == 351_104 + assert common.max_num_scheduled_tokens == 8_192 + assert common.usable_blocks_without_watermark == 21_940 + assert common.usable_blocks_with_watermark == 21_936 + assert follower_result["grid_synchronized"] is True + finally: + rank1.close() + rank0.close() + + +def test_benchmark_synchronizer_rejects_capacity_invariant_mismatch(): + endpoint = f"inproc://benchmark-sync-{uuid.uuid4().hex}" + rank0 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=0, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + rank1 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=1, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + follower_error = {} + + def run_follower(): + try: + rank1.negotiate_capacity( + _benchmark_capacity(grid_invariants_digest="b" * 64) + ) + except RuntimeError as error: + follower_error["error"] = error + + follower = threading.Thread(target=run_follower) + follower.start() + try: + with pytest.raises(RuntimeError, match="grid invariants differ"): + rank0.negotiate_capacity(_benchmark_capacity()) + follower.join(timeout=2) + assert not follower.is_alive() + assert "grid invariants differ" in str(follower_error["error"]) + finally: + rank1.close() + rank0.close() + + +def _digest_stub(max_num_running_reqs: int): + """Populate only the attributes ``_bench_grid_invariants_digest`` reads, + mirroring the activation-time filtering of the decode capture sizes.""" + stub = InstrumentedScheduler.__new__(InstrumentedScheduler) + stub._bench_config = BenchmarkConfig() + stub.block_size = 16 + stub._bench_hash_block_size = 16 + stub.cache_config = SimpleNamespace(block_size=16, enable_prefix_caching=True) + stub.max_num_running_reqs = max_num_running_reqs + stub._bench_prefill_cudagraph_mode = "PIECEWISE" + stub._bench_decode_cudagraph_mode = "FULL" + stub._bench_cudagraph_capture_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256] + stub._bench_prefill_capture_sizes = list(stub._bench_cudagraph_capture_sizes) + stub._bench_decode_capture_sizes = [ + size + for size in stub._bench_cudagraph_capture_sizes + if size <= max_num_running_reqs + ] + return stub + + +def test_capacity_digest_ignores_request_limit_filtered_capture_sizes(): + """Ranks that differ only in ``max_num_running_reqs`` filter different + decode capture lists at activation. The invariants digest must hash the + unfiltered configuration so ``common()`` negotiates the minimum instead + of rejecting the ranks as structurally different.""" + small = _digest_stub(max_num_running_reqs=128) + large = _digest_stub(max_num_running_reqs=256) + assert small._bench_decode_capture_sizes != large._bench_decode_capture_sizes + + small_digest = InstrumentedScheduler._bench_grid_invariants_digest(small) + large_digest = InstrumentedScheduler._bench_grid_invariants_digest(large) + assert small_digest == large_digest + + common = instrumented_scheduler_module._BenchmarkCapacityEnvelope.common( + [ + _benchmark_capacity( + max_num_running_reqs=128, grid_invariants_digest=small_digest + ), + _benchmark_capacity( + max_num_running_reqs=256, grid_invariants_digest=large_digest + ), + ] + ) + assert common.max_num_running_reqs == 128 + + +def test_benchmark_synchronizer_rejects_grid_mismatch_before_warmup(): + endpoint = f"inproc://benchmark-sync-{uuid.uuid4().hex}" + rank0 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=0, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + rank1 = instrumented_scheduler_module._BenchmarkSynchronizer( + dp_rank=1, + dp_size=2, + master_ip="unused", + port=0, + timeout=1, + endpoint=endpoint, + ) + follower_error = {} + + def run_follower(): + try: + rank1.synchronize_grid( + grid_digest="b" * 64, + expected_points=1_368, + missing_phases=[], + ) + except RuntimeError as error: + follower_error["error"] = error + + follower = threading.Thread(target=run_follower) + follower.start() + try: + with pytest.raises(RuntimeError, match="grid mismatch"): + rank0.synchronize_grid( + grid_digest="a" * 64, + expected_points=1_368, + missing_phases=[], + ) + follower.join(timeout=2) + assert not follower.is_alive() + assert "grid mismatch" in str(follower_error["error"]) + finally: + rank1.close() + rank0.close() + + def test_benchmark_synchronizer_shares_timeout_stop_decision(): endpoint = f"inproc://benchmark-sync-{uuid.uuid4().hex}" rank0 = instrumented_scheduler_module._BenchmarkSynchronizer( @@ -1299,6 +1531,60 @@ def test_decode_grid_uses_live_free_block_count_after_manager_reservations(): ) +def test_decode_grid_uses_common_attention_dp_capacity(): + # Regression for a DEP4 GLM-5.2 run where vLLM auto-fit two ranks to + # 5,987 blocks / 383,168 tokens and two ranks to 5,486 blocks / 351,104 + # tokens. Independent grids had the same count but diverged at point 20: + # 383,103 versus 351,039 total KV-read tokens. + larger_rank = _grid_stub_with_kv_capacity(num_gpu_blocks=5_987, block_size=64) + smaller_rank = _grid_stub_with_kv_capacity(num_gpu_blocks=5_486, block_size=64) + larger_rank.max_model_len = 383_168 + smaller_rank.max_model_len = 351_104 + capture_sizes = [ + 1, + 2, + 4, + *range(8, 257, 8), + *range(272, 513, 16), + ] + for stub in (larger_rank, smaller_rank): + stub.max_num_scheduled_tokens = 8_192 + stub.max_num_running_reqs = 1_024 + stub._bench_decode_capture_sizes = capture_sizes + common = _benchmark_capacity( + max_model_len=351_104, + max_num_scheduled_tokens=8_192, + max_num_running_reqs=1_024, + usable_blocks_without_watermark=5_485, + usable_blocks_with_watermark=5_485, + ) + larger_rank._bench_negotiated_capacity = common + smaller_rank._bench_negotiated_capacity = common + + InstrumentedScheduler._bench_generate_decode_grid(larger_rank) + InstrumentedScheduler._bench_generate_decode_grid(smaller_rank) + + larger_grid = [point.__dict__ for point in larger_rank._bench_grid] + smaller_grid = [point.__dict__ for point in smaller_rank._bench_grid] + assert larger_grid == smaller_grid + # Steady-coordinate normalization merges each batch's sub-2B presets into + # one point, so batch=1 keeps 19 ladder entries and its feasibility + # boundary sits at index 18. + assert len(larger_grid) == 1_266 + assert larger_rank._bench_grid[18].total_kv_read_tokens == 351_039 + + # The common grid must remain feasible under each rank's original local + # capacity when the conservative shared envelope is removed. + for stub in (larger_rank, smaller_rank): + stub._bench_negotiated_capacity = None + assert all( + InstrumentedScheduler._bench_decode_point_feasible( + stub, point.batch_size, point.total_kv_read_tokens + ) + for point in stub._bench_grid + ) + + @pytest.mark.parametrize( ("mode", "prefill_points", "decode_points", "expected_missing_phases"), [ @@ -1319,6 +1605,7 @@ def test_benchmark_grid_tracks_each_requested_empty_phase( stub._bench_grid = deque() stub._bench_grid_built = False stub._bench_missing_phases = [] + _install_test_capacity_preflight(stub) def generate_prefill_grid(): stub._bench_grid.extend( @@ -1347,6 +1634,7 @@ def test_benchmark_grid_has_no_point_cap(): stub._bench_grid_built = False stub._bench_missing_phases = [] stub._bench_grid_error = None + _install_test_capacity_preflight(stub) def generate_prefill_grid(): stub._bench_grid.extend( @@ -1372,6 +1660,7 @@ def test_benchmark_grid_assigns_stable_contiguous_ids_and_digest(): stub._bench_grid_built = False stub._bench_missing_phases = [] stub._bench_grid_error = None + _install_test_capacity_preflight(stub) def generate_prefill_grid(): stub._bench_grid.extend( @@ -1414,6 +1703,7 @@ def _explicit_grid_stub(mode="agg", points=None): "decode": [{"total_kv_read_tokens": 16, "batch_size": 1}], } ) + _install_test_capacity_preflight(stub) return stub @@ -1487,7 +1777,11 @@ def test_explicit_decode_respects_scheduled_token_limit(): "decode": [{"total_kv_read_tokens": 2, "batch_size": 2}], }, ) - stub.max_num_scheduled_tokens = 1 + # The limit is read through the negotiated capacity envelope, so the + # constraint must be installed there rather than on the stub attribute. + _install_test_capacity_preflight( + stub, _benchmark_capacity(max_num_scheduled_tokens=1) + ) with pytest.raises(ValueError, match=r"decode\[0\].*infeasible"): InstrumentedScheduler._bench_build_grid(stub) @@ -1705,6 +1999,15 @@ def _prefill_grid_stub( stub._bench_prefill_capture_sizes = [8, 16] stub._bench_prefill_cudagraph_mode = "PIECEWISE" stub.num_lookahead_tokens = 0 + _install_test_capacity_preflight( + stub, + _benchmark_capacity( + max_model_len=stub.max_model_len, + max_num_scheduled_tokens=stub.max_num_scheduled_tokens, + max_num_running_reqs=stub.max_num_running_reqs, + usable_blocks_without_watermark=num_gpu_blocks - 1, + ), + ) return stub @@ -1884,6 +2187,71 @@ def test_prefill_grid_runs_larger_workload_coordinates_first(): assert coordinates == sorted(coordinates, reverse=True) +def test_prefill_grid_uses_common_attention_dp_capacity(): + larger_rank = _prefill_grid_stub(num_gpu_blocks=64) + smaller_rank = _prefill_grid_stub(num_gpu_blocks=48) + common = _benchmark_capacity( + max_model_len=96, + max_num_scheduled_tokens=40, + max_num_running_reqs=8, + usable_blocks_without_watermark=47, + usable_blocks_with_watermark=47, + ) + larger_rank._bench_negotiated_capacity = common + smaller_rank._bench_negotiated_capacity = common + + InstrumentedScheduler._bench_generate_prefill_grid(larger_rank) + InstrumentedScheduler._bench_generate_prefill_grid(smaller_rank) + + larger_grid = [point.__dict__ for point in larger_rank._bench_grid] + smaller_grid = [point.__dict__ for point in smaller_rank._bench_grid] + assert larger_grid == smaller_grid + + for stub in (larger_rank, smaller_rank): + stub._bench_negotiated_capacity = None + assert all( + InstrumentedScheduler._bench_prefill_point_feasible( + stub, + point.total_prefill_tokens, + point.batch_size, + point.total_kv_read_tokens, + ) + for point in stub._bench_grid + ) + + +def test_explicit_prefill_point_uses_negotiated_scheduled_token_limit(): + """An explicit point at the negotiated ``max_num_scheduled_tokens`` must + get identical cudagraph metadata on every rank regardless of the rank's + local limit — otherwise ``sample_reasons`` (engine_limit vs + geometric_tail) and therefore the per-point digests diverge.""" + at_limit_rank = _prefill_grid_stub() + above_limit_rank = _prefill_grid_stub() + above_limit_rank.max_num_scheduled_tokens = 48 + common = _benchmark_capacity( + max_model_len=128, + max_num_scheduled_tokens=40, + max_num_running_reqs=8, + usable_blocks_without_watermark=63, + ) + candidate = PrefillPointCandidate( + total_prefill_tokens=40, batch_size=1, total_kv_read_tokens=0 + ) + + points = [] + for stub in (at_limit_rank, above_limit_rank): + stub._bench_negotiated_capacity = common + points.append( + InstrumentedScheduler._bench_materialize_prefill_candidate( + stub, candidate, "points[0]" + ) + ) + + assert points[0] == points[1] + assert "engine_limit" in points[0].sample_reasons + assert "geometric_tail" not in points[1].sample_reasons + + def test_agg_grid_contains_piecewise_prefill_then_full_decode_points(): stub = _prefill_grid_stub() stub._bench_grid = deque()