@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d58da79de7

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

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>

Zijing Liu committed Oct 1, 2026 at 18:24 UTC d58da79de76d2c3dfd6799f3355dec9eb8af9b19

5 files changed +116

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+28

@@ -5,6 +5,7 @@ from unittest.mock import patch

6 import pytest

8 +from vllm.v1.metrics.stats import KV_FETCH_STAGES

9 from vllm.v1.outputs import (

10 EMPTY_MODEL_RUNNER_OUTPUT,

11 KVConnectorOutput,

@@ -28,6 +29,11 @@ def _num_waiting_requests(scheduler) -> int:

29 return len(scheduler.waiting) + len(scheduler.kv_holding_waiting)

32 +def _kv_fetch_stages(scheduler) -> tuple[int, ...]:

33 + """(waiting_to_start, in_progress, completed_waiting) request counts."""

34 + return tuple(scheduler._kv_fetch_counts[s] for s in KV_FETCH_STAGES)

35 +

36 +

37 def test_basic_lifecycle():

38 """Test lifecycle of a remote prefill."""

39 vllm_config = create_vllm_config()

@@ -50,10 +56,12 @@ def test_basic_lifecycle():

57 scheduler.add_request(request)

58 request_id = request.request_id

59 + assert _kv_fetch_stages(scheduler) == (1, 0, 0)

61 # STEP (1):

62 # (1a): schedule()

63 scheduler_output = scheduler.schedule()

64 + assert _kv_fetch_stages(scheduler) == (0, 1, 0)

66 # Nothing running and empty scheduler output.

67 assert len(scheduler.running) == 0

@@ -105,11 +113,13 @@ def test_basic_lifecycle():

113 )

114 assert _num_waiting_requests(scheduler) == 1

115 assert request_id in scheduler.finished_recving_kv_req_ids

116 + assert _kv_fetch_stages(scheduler) == (0, 0, 1)

117

118 # STEP (3):

119 # (3a): schedule(): this should actually schedule.

120 scheduler_output = scheduler.schedule()

121 assert len(scheduler.running) == 1

122 + assert _kv_fetch_stages(scheduler) == (0, 0, 0)

123

124 # Confirm the block are actually allocated.

125 num_hashed_blocks = 0

@@ -582,6 +592,24 @@ def test_cannot_recv():

592 assert_scheduler_empty(scheduler)

593

594

595 +def test_kv_fetch_stages_cleared_on_abort():

596 + vllm_config = create_vllm_config()

597 + scheduler = create_scheduler(vllm_config)

598 + BLOCK_SIZE = vllm_config.cache_config.block_size

599 + request = create_request(

600 + request_id=1,

601 + block_size=BLOCK_SIZE,

602 + num_tokens=int(BLOCK_SIZE * 2.5),

603 + do_remote_prefill=True,

604 + )

605 + scheduler.add_request(request)

606 + scheduler.schedule()

607 + assert _kv_fetch_stages(scheduler) == (0, 1, 0)

608 +

609 + scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)

610 + assert _kv_fetch_stages(scheduler) == (0, 0, 0)

611 +

612 +

613 @patch(

614 "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform"

615 )

tests/v1/kv_connector/unit/utils.py

@@ -60,6 +60,8 @@ def assert_scheduler_empty(scheduler: Scheduler):

60 assert len(scheduler.finished_req_ids) == 0

61 assert len(scheduler.finished_recving_kv_req_ids) == 0

62 assert len(scheduler._inflight_prefills) == 0

63 + assert not scheduler._kv_fetch_stages

64 + assert not any(scheduler._kv_fetch_counts.values())

66 # EncoderCacheManager.

67 assert len(scheduler.encoder_cache_manager.freed) == 0

vllm/v1/core/sched/scheduler.py

+34

@@ -62,6 +62,10 @@ from vllm.v1.kv_cache_interface import (

62 )

63 from vllm.v1.metrics.perf import ModelMetrics, PerfStats

64 from vllm.v1.metrics.stats import (

65 + KV_FETCH_COMPLETED_WAITING,

66 + KV_FETCH_IN_PROGRESS,

67 + KV_FETCH_STAGES,

68 + KV_FETCH_WAITING_TO_START,

69 PrefixCacheStats,

70 RequestSpecDecodeMetrics,

71 SchedulerStats,

@@ -412,6 +416,13 @@ class Scheduler(SchedulerInterface):

416 # async KV loads). Their remaining-block reservation gates async loads.

417 self._inflight_prefills: set[Request] = set()

418

419 + # Requests with an async KV load, by stage (KV_FETCH_STAGES).

420 + # None when not tracked (stats disabled or no KV connector).

421 + self._kv_fetch_stages: dict[Request, str] | None = (

422 + {} if self.log_stats and self.connector is not None else None

423 + )

424 + self._kv_fetch_counts: dict[str, int] = dict.fromkeys(KV_FETCH_STAGES, 0)

425 +

426 def _mamba_block_aligned_split(

427 self,

428 request: Request,

@@ -1039,6 +1050,9 @@ class Scheduler(SchedulerInterface):

1050 num_local_cached_tokens=num_new_local_computed_tokens,

1051 num_external_cached_tokens=num_external_computed_tokens,

1052 )

1053 + self._set_kv_fetch_stage(

1054 + request, KV_FETCH_WAITING_TO_START if load_kv_async else None

1055 + )

1056 else:

1057 # KVTransfer: WAITING reqs have num_computed_tokens > 0

1058 # after async KV recvs are completed. A streaming-input

@@ -1276,6 +1290,7 @@ class Scheduler(SchedulerInterface):

1290 # only the successfully loaded tokens.

1291 request.num_computed_tokens = num_computed_tokens

1292 self._inflight_prefills.add(request)

1293 + self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS)

1294 if self.needs_kv_cache_zeroing:

1295 # Skip zeroing of the blocks the async load will

1296 # overwrite; the zeroing could race the write.

@@ -1316,6 +1331,7 @@ class Scheduler(SchedulerInterface):

1331 input_budget -= num_new_tokens + draft_slots

1332 request.status = RequestStatus.RUNNING

1333 request.num_computed_tokens = num_computed_tokens

1334 + self._set_kv_fetch_stage(request, None)

1335 if pad_spec_decode:

1336 assert num_new_tokens == 1 + self.num_spec_tokens

1337 scheduled_spec_decode_tokens[request_id] = [

@@ -2588,6 +2604,11 @@ class Scheduler(SchedulerInterface):

2604 )

2605 if self.connector is not None:

2606 self.connector.on_new_request(request)

2607 + kv_transfer_params = request.kv_transfer_params

2608 + if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"):

2609 + # Count P/D requests on arrival: connector matching is

2610 + # skipped while all run slots are taken.

2611 + self._set_kv_fetch_stage(request, KV_FETCH_WAITING_TO_START)

2612 if self.log_stats:

2613 request.record_event(EngineCoreEventType.QUEUED)

2614

@@ -2664,6 +2685,7 @@ class Scheduler(SchedulerInterface):

2685 if self.aux_output_connector is not None:

2686 self.aux_output_connector.request_finished(request)

2687 self._inflight_prefills.discard(request)

2688 + self._set_kv_fetch_stage(request, None)

2689 connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)

2690

2691 # EC Connector: mirror the KV hook. The contract requires firing

@@ -2889,6 +2911,7 @@ class Scheduler(SchedulerInterface):

2911 num_running_reqs=num_running,

2912 num_waiting_reqs=num_waiting - num_deferred,

2913 num_skipped_waiting_reqs=num_deferred,

2914 + num_kv_fetch_reqs_by_stage=self._kv_fetch_counts.copy(),

2915 kv_cache_usage=self.kv_cache_manager.usage,

2916 prefix_cache_stats=prefix_cache_stats,

2917 connector_prefix_cache_stats=connector_prefix_cache_stats,

@@ -3028,6 +3051,16 @@ class Scheduler(SchedulerInterface):

3051 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size

3052 )

3053

3054 + def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None:

3055 + kv_fetch_stages = self._kv_fetch_stages

3056 + if kv_fetch_stages is not None:

3057 + prev_stage = kv_fetch_stages.pop(request, None)

3058 + if prev_stage is not None:

3059 + self._kv_fetch_counts[prev_stage] -= 1

3060 + if stage is not None:

3061 + kv_fetch_stages[request] = stage

3062 + self._kv_fetch_counts[stage] += 1

3063 +

3064 def _inflight_prefill_reserved_blocks(self) -> int:

3065 """Num blocks in-flight prefills still need to finish (their reservation)."""

3066 return sum(

@@ -3148,6 +3181,7 @@ class Scheduler(SchedulerInterface):

3181 req = self.requests[req_id]

3182 if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:

3183 self.finished_recving_kv_req_ids.add(req_id)

3184 + self._set_kv_fetch_stage(req, KV_FETCH_COMPLETED_WAITING)

3185 else:

3186 assert RequestStatus.is_finished(req.status)

3187 self._free_blocks(self.requests[req_id])

vllm/v1/metrics/loggers.py

+40

@@ -26,6 +26,7 @@ from vllm.v1.metrics.buckets import histogram_buckets

26 from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm

27 from vllm.v1.metrics.prometheus import unregister_vllm_metrics

28 from vllm.v1.metrics.stats import (

29 + KV_FETCH_STAGES,

30 CachingMetrics,

31 IterationStats,

32 MultiModalCacheStats,

@@ -292,6 +293,13 @@ class LoggingStatLogger(StatLoggerBase):

293 log_parts.append("Deferred: %d reqs")

294 log_args.append(self.last_scheduler_stats.num_skipped_waiting_reqs)

295

296 + kv_fetch = self.last_scheduler_stats.num_kv_fetch_reqs_by_stage

297 + if any(kv_fetch.values()):

298 + log_parts.append(

299 + "KV fetch: %d waiting to start, %d in progress, %d completed waiting"

300 + )

301 + log_args.extend(kv_fetch.get(stage, 0) for stage in KV_FETCH_STAGES)

302 +

303 if self.num_preemptions > 0:

304 log_parts.append("Preemptions: %d")

305 log_args.append(self.num_preemptions)

@@ -393,6 +401,10 @@ class AggregatedLoggingStatLogger(LoggingStatLogger, AggregateStatLoggerBase):

401 self.last_scheduler_stats.num_skipped_waiting_reqs += (

402 last_scheduler_stats.num_skipped_waiting_reqs

403 )

404 + engine_kv_fetch = last_scheduler_stats.num_kv_fetch_reqs_by_stage

405 + total_kv_fetch = self.last_scheduler_stats.num_kv_fetch_reqs_by_stage

406 + for stage, num_reqs in engine_kv_fetch.items():

407 + total_kv_fetch[stage] = total_kv_fetch.get(stage, 0) + num_reqs

408 self.last_scheduler_stats.kv_cache_usage += (

409 last_scheduler_stats.kv_cache_usage

410 )

@@ -547,6 +559,30 @@ class PrometheusStatLogger(AggregateStatLoggerBase):

559 gauge_waiting_by_reason, per_engine_labelvalues_with_reason

560 )

561

562 + self.gauge_kv_fetch_by_stage: dict[str, dict[int, Gauge]] = {}

563 + kv_transfer_config = vllm_config.kv_transfer_config

564 + if kv_transfer_config is not None and kv_transfer_config.is_kv_consumer:

565 + gauge_kv_fetch_by_stage = self._gauge_cls(

566 + name="vllm:num_requests_kv_fetch_by_stage",

567 + documentation=(

568 + "Number of waiting requests by async KV load stage. "

569 + "Stage labels: 'waiting_to_start' = needs an async KV load "

570 + "that has not started; 'in_progress' = load started and not "

571 + "yet reported finished; 'completed_waiting' = load finished, "

572 + "request not running yet."

573 + ),

574 + multiprocess_mode="mostrecent",

575 + labelnames=labelnames + ["stage"],

576 + )

577 + for stage in KV_FETCH_STAGES:

578 + per_engine_labelvalues_with_stage = {

579 + idx: labelvalues + [stage]

580 + for idx, labelvalues in per_engine_labelvalues.items()

581 + }

582 + self.gauge_kv_fetch_by_stage[stage] = create_metric_per_engine(

583 + gauge_kv_fetch_by_stage, per_engine_labelvalues_with_stage

584 + )

585 +

586 gauge_engine_sleep_state = self._gauge_cls(

587 name="vllm:engine_sleep_state",

588 documentation=(

@@ -1062,6 +1098,10 @@ class PrometheusStatLogger(AggregateStatLoggerBase):

1098 self.gauge_waiting_by_reason[WAITING_REASON_DEFERRED][engine_idx].set(

1099 scheduler_stats.num_skipped_waiting_reqs

1100 )

1101 + for stage, gauges in self.gauge_kv_fetch_by_stage.items():

1102 + gauges[engine_idx].set(

1103 + scheduler_stats.num_kv_fetch_reqs_by_stage.get(stage, 0)

1104 + )

1105 self.gauge_kv_cache_usage[engine_idx].set(scheduler_stats.kv_cache_usage)

1106

1107 self.counter_prefix_cache_queries[engine_idx].inc(

vllm/v1/metrics/stats.py

+12

@@ -182,6 +182,16 @@ class SchedulerIterationDetails:

182 is_dummy: bool = False

183

184

185 +KV_FETCH_WAITING_TO_START = "waiting_to_start"

186 +KV_FETCH_IN_PROGRESS = "in_progress"

187 +KV_FETCH_COMPLETED_WAITING = "completed_waiting"

188 +KV_FETCH_STAGES = (

189 + KV_FETCH_WAITING_TO_START,

190 + KV_FETCH_IN_PROGRESS,

191 + KV_FETCH_COMPLETED_WAITING,

192 +)

193 +

194 +

195 @dataclass

196 class SchedulerStats:

197 """Stats associated with the scheduler."""

@@ -192,6 +202,8 @@ class SchedulerStats:

202 # waiting requests deferred by transient constraints or a blocked status

203 num_skipped_waiting_reqs: int = 0

204

205 + num_kv_fetch_reqs_by_stage: dict[str, int] = field(default_factory=dict)

206 +

207 # These are used for internal DP load-balancing.

208 step_counter: int = 0

209 current_wave: int = 0