diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index daddd20b4a..59fcc678bf 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest +from vllm.v1.metrics.stats import KV_FETCH_STAGES from vllm.v1.outputs import ( EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput, @@ -28,6 +29,11 @@ def _num_waiting_requests(scheduler) -> int: return len(scheduler.waiting) + len(scheduler.kv_holding_waiting) +def _kv_fetch_stages(scheduler) -> tuple[int, ...]: + """(waiting_to_start, in_progress, completed_waiting) request counts.""" + return tuple(scheduler._kv_fetch_counts[s] for s in KV_FETCH_STAGES) + + def test_basic_lifecycle(): """Test lifecycle of a remote prefill.""" vllm_config = create_vllm_config() @@ -50,10 +56,12 @@ def test_basic_lifecycle(): scheduler.add_request(request) request_id = request.request_id + assert _kv_fetch_stages(scheduler) == (1, 0, 0) # STEP (1): # (1a): schedule() scheduler_output = scheduler.schedule() + assert _kv_fetch_stages(scheduler) == (0, 1, 0) # Nothing running and empty scheduler output. assert len(scheduler.running) == 0 @@ -105,11 +113,13 @@ def test_basic_lifecycle(): ) assert _num_waiting_requests(scheduler) == 1 assert request_id in scheduler.finished_recving_kv_req_ids + assert _kv_fetch_stages(scheduler) == (0, 0, 1) # STEP (3): # (3a): schedule(): this should actually schedule. scheduler_output = scheduler.schedule() assert len(scheduler.running) == 1 + assert _kv_fetch_stages(scheduler) == (0, 0, 0) # Confirm the block are actually allocated. num_hashed_blocks = 0 @@ -582,6 +592,24 @@ def test_cannot_recv(): assert_scheduler_empty(scheduler) +def test_kv_fetch_stages_cleared_on_abort(): + vllm_config = create_vllm_config() + scheduler = create_scheduler(vllm_config) + BLOCK_SIZE = vllm_config.cache_config.block_size + request = create_request( + request_id=1, + block_size=BLOCK_SIZE, + num_tokens=int(BLOCK_SIZE * 2.5), + do_remote_prefill=True, + ) + scheduler.add_request(request) + scheduler.schedule() + assert _kv_fetch_stages(scheduler) == (0, 1, 0) + + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + assert _kv_fetch_stages(scheduler) == (0, 0, 0) + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform" ) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 2febe3109f..6f27212281 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -60,6 +60,8 @@ def assert_scheduler_empty(scheduler: Scheduler): assert len(scheduler.finished_req_ids) == 0 assert len(scheduler.finished_recving_kv_req_ids) == 0 assert len(scheduler._inflight_prefills) == 0 + assert not scheduler._kv_fetch_stages + assert not any(scheduler._kv_fetch_counts.values()) # EncoderCacheManager. assert len(scheduler.encoder_cache_manager.freed) == 0 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fcb1421b10..5585c9c70c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -62,6 +62,10 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import ( + KV_FETCH_COMPLETED_WAITING, + KV_FETCH_IN_PROGRESS, + KV_FETCH_STAGES, + KV_FETCH_WAITING_TO_START, PrefixCacheStats, RequestSpecDecodeMetrics, SchedulerStats, @@ -412,6 +416,13 @@ class Scheduler(SchedulerInterface): # async KV loads). Their remaining-block reservation gates async loads. self._inflight_prefills: set[Request] = set() + # Requests with an async KV load, by stage (KV_FETCH_STAGES). + # None when not tracked (stats disabled or no KV connector). + self._kv_fetch_stages: dict[Request, str] | None = ( + {} if self.log_stats and self.connector is not None else None + ) + self._kv_fetch_counts: dict[str, int] = dict.fromkeys(KV_FETCH_STAGES, 0) + def _mamba_block_aligned_split( self, request: Request, @@ -1039,6 +1050,9 @@ class Scheduler(SchedulerInterface): num_local_cached_tokens=num_new_local_computed_tokens, num_external_cached_tokens=num_external_computed_tokens, ) + self._set_kv_fetch_stage( + request, KV_FETCH_WAITING_TO_START if load_kv_async else None + ) else: # KVTransfer: WAITING reqs have num_computed_tokens > 0 # after async KV recvs are completed. A streaming-input @@ -1276,6 +1290,7 @@ class Scheduler(SchedulerInterface): # only the successfully loaded tokens. request.num_computed_tokens = num_computed_tokens self._inflight_prefills.add(request) + self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS) if self.needs_kv_cache_zeroing: # Skip zeroing of the blocks the async load will # overwrite; the zeroing could race the write. @@ -1316,6 +1331,7 @@ class Scheduler(SchedulerInterface): input_budget -= num_new_tokens + draft_slots request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens + self._set_kv_fetch_stage(request, None) if pad_spec_decode: assert num_new_tokens == 1 + self.num_spec_tokens scheduled_spec_decode_tokens[request_id] = [ @@ -2588,6 +2604,11 @@ class Scheduler(SchedulerInterface): ) if self.connector is not None: self.connector.on_new_request(request) + kv_transfer_params = request.kv_transfer_params + if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"): + # Count P/D requests on arrival: connector matching is + # skipped while all run slots are taken. + self._set_kv_fetch_stage(request, KV_FETCH_WAITING_TO_START) if self.log_stats: request.record_event(EngineCoreEventType.QUEUED) @@ -2664,6 +2685,7 @@ class Scheduler(SchedulerInterface): if self.aux_output_connector is not None: self.aux_output_connector.request_finished(request) self._inflight_prefills.discard(request) + self._set_kv_fetch_stage(request, None) connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) # EC Connector: mirror the KV hook. The contract requires firing @@ -2889,6 +2911,7 @@ class Scheduler(SchedulerInterface): num_running_reqs=num_running, num_waiting_reqs=num_waiting - num_deferred, num_skipped_waiting_reqs=num_deferred, + num_kv_fetch_reqs_by_stage=self._kv_fetch_counts.copy(), kv_cache_usage=self.kv_cache_manager.usage, prefix_cache_stats=prefix_cache_stats, connector_prefix_cache_stats=connector_prefix_cache_stats, @@ -3028,6 +3051,16 @@ class Scheduler(SchedulerInterface): 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size ) + def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None: + kv_fetch_stages = self._kv_fetch_stages + if kv_fetch_stages is not None: + prev_stage = kv_fetch_stages.pop(request, None) + if prev_stage is not None: + self._kv_fetch_counts[prev_stage] -= 1 + if stage is not None: + kv_fetch_stages[request] = stage + self._kv_fetch_counts[stage] += 1 + def _inflight_prefill_reserved_blocks(self) -> int: """Num blocks in-flight prefills still need to finish (their reservation).""" return sum( @@ -3148,6 +3181,7 @@ class Scheduler(SchedulerInterface): req = self.requests[req_id] if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: self.finished_recving_kv_req_ids.add(req_id) + self._set_kv_fetch_stage(req, KV_FETCH_COMPLETED_WAITING) else: assert RequestStatus.is_finished(req.status) self._free_blocks(self.requests[req_id]) diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index ac93b26a8f..681f80ce29 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -26,6 +26,7 @@ from vllm.v1.metrics.buckets import histogram_buckets from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import ( + KV_FETCH_STAGES, CachingMetrics, IterationStats, MultiModalCacheStats, @@ -292,6 +293,13 @@ class LoggingStatLogger(StatLoggerBase): log_parts.append("Deferred: %d reqs") log_args.append(self.last_scheduler_stats.num_skipped_waiting_reqs) + kv_fetch = self.last_scheduler_stats.num_kv_fetch_reqs_by_stage + if any(kv_fetch.values()): + log_parts.append( + "KV fetch: %d waiting to start, %d in progress, %d completed waiting" + ) + log_args.extend(kv_fetch.get(stage, 0) for stage in KV_FETCH_STAGES) + if self.num_preemptions > 0: log_parts.append("Preemptions: %d") log_args.append(self.num_preemptions) @@ -393,6 +401,10 @@ class AggregatedLoggingStatLogger(LoggingStatLogger, AggregateStatLoggerBase): self.last_scheduler_stats.num_skipped_waiting_reqs += ( last_scheduler_stats.num_skipped_waiting_reqs ) + engine_kv_fetch = last_scheduler_stats.num_kv_fetch_reqs_by_stage + total_kv_fetch = self.last_scheduler_stats.num_kv_fetch_reqs_by_stage + for stage, num_reqs in engine_kv_fetch.items(): + total_kv_fetch[stage] = total_kv_fetch.get(stage, 0) + num_reqs self.last_scheduler_stats.kv_cache_usage += ( last_scheduler_stats.kv_cache_usage ) @@ -547,6 +559,30 @@ class PrometheusStatLogger(AggregateStatLoggerBase): gauge_waiting_by_reason, per_engine_labelvalues_with_reason ) + self.gauge_kv_fetch_by_stage: dict[str, dict[int, Gauge]] = {} + kv_transfer_config = vllm_config.kv_transfer_config + if kv_transfer_config is not None and kv_transfer_config.is_kv_consumer: + gauge_kv_fetch_by_stage = self._gauge_cls( + name="vllm:num_requests_kv_fetch_by_stage", + documentation=( + "Number of waiting requests by async KV load stage. " + "Stage labels: 'waiting_to_start' = needs an async KV load " + "that has not started; 'in_progress' = load started and not " + "yet reported finished; 'completed_waiting' = load finished, " + "request not running yet." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames + ["stage"], + ) + for stage in KV_FETCH_STAGES: + per_engine_labelvalues_with_stage = { + idx: labelvalues + [stage] + for idx, labelvalues in per_engine_labelvalues.items() + } + self.gauge_kv_fetch_by_stage[stage] = create_metric_per_engine( + gauge_kv_fetch_by_stage, per_engine_labelvalues_with_stage + ) + gauge_engine_sleep_state = self._gauge_cls( name="vllm:engine_sleep_state", documentation=( @@ -1062,6 +1098,10 @@ class PrometheusStatLogger(AggregateStatLoggerBase): self.gauge_waiting_by_reason[WAITING_REASON_DEFERRED][engine_idx].set( scheduler_stats.num_skipped_waiting_reqs ) + for stage, gauges in self.gauge_kv_fetch_by_stage.items(): + gauges[engine_idx].set( + scheduler_stats.num_kv_fetch_reqs_by_stage.get(stage, 0) + ) self.gauge_kv_cache_usage[engine_idx].set(scheduler_stats.kv_cache_usage) self.counter_prefix_cache_queries[engine_idx].inc( diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index da88c9379a..4db98884fe 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -182,6 +182,16 @@ class SchedulerIterationDetails: is_dummy: bool = False +KV_FETCH_WAITING_TO_START = "waiting_to_start" +KV_FETCH_IN_PROGRESS = "in_progress" +KV_FETCH_COMPLETED_WAITING = "completed_waiting" +KV_FETCH_STAGES = ( + KV_FETCH_WAITING_TO_START, + KV_FETCH_IN_PROGRESS, + KV_FETCH_COMPLETED_WAITING, +) + + @dataclass class SchedulerStats: """Stats associated with the scheduler.""" @@ -192,6 +202,8 @@ class SchedulerStats: # waiting requests deferred by transient constraints or a blocked status num_skipped_waiting_reqs: int = 0 + num_kv_fetch_reqs_by_stage: dict[str, int] = field(default_factory=dict) + # These are used for internal DP load-balancing. step_counter: int = 0 current_wave: int = 0