iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Metrics][P/D] Add KV-fetch stage gauges for async KV loads (#58874)

Browse Source 
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Zijing Liu
2026-10-02 01:24:33 +00:00
committed by GitHub
co-authored by Nick Hill Claude Opus 5.5
parent eec3c36a47
commit d58da79de7
5 changed files with 116 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+28
View File
@@ -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"
)
tests/v1/kv_connector/unit/utils.py
+2
View File
@@ -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
vllm/v1/core/sched/scheduler.py
+34
View File
@@ -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])
vllm/v1/metrics/loggers.py
+40
View File
@@ -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(
vllm/v1/metrics/stats.py
+12
View File
@@ -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