A high-throughput and memory-efficient inference and serving engine for LLMs
[Core][KV Connector] fix scheduler KV connector stats aggregation (#43877)
Fixes scheduler-side KV connector stats collection so that: 1. update_connector_output() runs before scheduler-side stats are collected. 2. worker-side and scheduler-side KV connector stats are aggregated when both are present. 3. scheduler-only KV connector stats are still emitted when no worker-side stats exist. Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>
Srinivas Krovvidi committed
Jun 12, 2026 at 09:51 UTC
9ff278b1d2304ae606a13e8eebab75fcde2d2281
3 files changed
+102
-8
tests/v1/core/test_scheduler.py
+82
@@ -15,6 +15,7 @@ from vllm.config import (
15
SpeculativeConfig,
16
VllmConfig,
17
)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
19
from vllm.multimodal.inputs import (
20
MultiModalFeatureSpec,
21
MultiModalKwargsItem,
@@ -3990,6 +3991,87 @@ def test_delayed_kv_connector_free_keeps_scheduler_active():
3991
assert not scheduler.has_finished_requests()
3992
3993
3994
+def test_scheduler_kv_connector_stats():
3995
+ """Test worker-side, scheduler-side, and combined KV connector stats."""
3996
+
3997
+ class GenericKVConnectorStats(KVConnectorStats):
3998
+ def reset(self):
3999
+ self.data = {}
4000
+
4001
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
4002
+ self.data.update(other.data)
4003
+ return self
4004
+
4005
+ def reduce(self) -> dict[str, int | float]:
4006
+ return {}
4007
+
4008
+ def is_empty(self) -> bool:
4009
+ return not self.data
4010
+
4011
+ test_cases = (
4012
+ ({"worker": 1}, None, {"worker": 1}),
4013
+ (None, {"scheduler": 2}, {"scheduler": 2}),
4014
+ ({"worker": 1}, {"scheduler": 2}, {"worker": 1, "scheduler": 2}),
4015
+ )
4016
+
4017
+ for worker_data, scheduler_data, expected_data in test_cases:
4018
+ scheduler = create_scheduler()
4019
+ worker_stats = (
4020
+ GenericKVConnectorStats(data=worker_data) if worker_data else None
4021
+ )
4022
+ scheduler_stats = (
4023
+ GenericKVConnectorStats(data=scheduler_data) if scheduler_data else None
4024
+ )
4025
+ scheduler.connector = Mock()
4026
+ scheduler.connector.get_kv_connector_stats.return_value = (
4027
+ scheduler_stats if worker_stats is None else None
4028
+ )
4029
+ scheduler.connector.take_events.return_value = []
4030
+
4031
+ def update_connector_output(
4032
+ kv_connector_output: KVConnectorOutput,
4033
+ scheduler=scheduler,
4034
+ scheduler_stats=scheduler_stats,
4035
+ ):
4036
+ scheduler.connector.get_kv_connector_stats.return_value = scheduler_stats
4037
+
4038
+ scheduler.connector.update_connector_output.side_effect = (
4039
+ update_connector_output
4040
+ )
4041
+
4042
+ model_output = ModelRunnerOutput(
4043
+ req_ids=["req_0"],
4044
+ req_id_to_index={"req_0": 0},
4045
+ sampled_token_ids=[[123]],
4046
+ logprobs=None,
4047
+ prompt_logprobs_dict={},
4048
+ pooler_output=[None],
4049
+ kv_connector_output=KVConnectorOutput(kv_connector_stats=worker_stats)
4050
+ if worker_stats
4051
+ else None,
4052
+ )
4053
+ scheduler_output = SchedulerOutput(
4054
+ scheduled_new_reqs=[],
4055
+ scheduled_cached_reqs=None,
4056
+ num_scheduled_tokens={"req_0": 1},
4057
+ total_num_scheduled_tokens=1,
4058
+ scheduled_spec_decode_tokens={},
4059
+ scheduled_encoder_inputs={},
4060
+ num_common_prefix_blocks=[0],
4061
+ finished_req_ids=set(),
4062
+ free_encoder_mm_hashes=[],
4063
+ )
4064
+
4065
+ engine_core_outputs = scheduler.update_from_output(
4066
+ scheduler_output, model_output
4067
+ )
4068
+
4069
+ final_stats = next(
4070
+ iter(engine_core_outputs.values())
4071
+ ).scheduler_stats.kv_connector_stats
4072
+ assert final_stats == expected_data
4073
+
4074
+
4075
# ==============================================================================
4076
# Variable-length encoder cross-attention block allocation tests
4077
# ==============================================================================
tests/v1/kv_connector/unit/test_multi_connector.py
+3
-1
@@ -58,6 +58,7 @@ class MockConnector(KVConnectorBase_V1):
58
mock = MagicMock(spec_set=KVConnectorBase_V1)
59
# Override just build_kv_connector_stats
60
mock.build_kv_connector_stats = cls.build_kv_connector_stats
61
+ mock.get_kv_connector_stats.return_value = None
62
return mock
63
64
@classmethod
@@ -93,6 +94,7 @@ class MockHMAConnector(KVConnectorBase_V1, SupportsHMA):
94
95
def __new__(cls, *args, **kwargs):
96
mock = MagicMock(spec_set=cls)
97
+ mock.get_kv_connector_stats.return_value = None
98
return mock
99
100
def start_load_kv(self, forward_context, **kwargs):
@@ -368,7 +370,7 @@ def test_multi_example_connector_consistency():
370
def _ignore_event_collection(events: list[str]) -> list[str]:
371
# Filter out per-step polling hooks that the scheduler calls repeatedly
372
# and which are not meaningful state transitions for these assertions.
371
- ignored = {"take_events", "has_pending_push_work"}
373
+ ignored = {"get_kv_connector_stats", "has_pending_push_work", "take_events"}
374
return [event for event in events if event not in ignored]
375
376
vllm/v1/core/sched/scheduler.py
+17
-7
@@ -1412,13 +1412,6 @@ class Scheduler(SchedulerInterface):
1412
1413
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
1414
spec_decoding_stats: SpecDecodingStats | None = None
1415
- kv_connector_stats: KVConnectorStats | None = (
1416
- kv_connector_output.kv_connector_stats if kv_connector_output else None
1417
- )
1418
- if kv_connector_stats and self.connector:
1419
- kv_stats = self.connector.get_kv_connector_stats()
1420
- if kv_stats:
1421
- kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
1415
1416
failed_kv_load_req_ids = None
1417
if kv_connector_output and kv_connector_output.invalid_block_ids:
@@ -1665,6 +1658,23 @@ class Scheduler(SchedulerInterface):
1658
if kv_connector_output:
1659
self._update_from_kv_xfer_finished(kv_connector_output)
1660
1661
+ # Worker-side KV connector stats from the model runner output.
1662
+ kv_connector_stats: KVConnectorStats | None = (
1663
+ kv_connector_output.kv_connector_stats if kv_connector_output else None
1664
+ )
1665
+ if self.connector:
1666
+ # Scheduler-side KV connector stats collected after connector update.
1667
+ scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
1668
+ if (
1669
+ scheduler_kv_connector_stats is not None
1670
+ and not scheduler_kv_connector_stats.is_empty()
1671
+ ):
1672
+ kv_connector_stats = (
1673
+ kv_connector_stats.aggregate(scheduler_kv_connector_stats)
1674
+ if kv_connector_stats is not None
1675
+ else scheduler_kv_connector_stats
1676
+ )
1677
+
1678
# collect KV cache events from KV cache manager
1679
events = self.kv_cache_manager.take_events()
1680