[Bugfix] Prefix-cache metrics double-counted when a KV connector defers requests (#48860)
Signed-off-by: Seiji Eicher <58963096+eicherseiji@users.noreply.github.com> Signed-off-by: Seiji Eicher <seiji@anyscale.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Lixi Zhang <muxixibbb@gmail.com> Co-authored-by: harjothkhara <harjoth.khara@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Seiji Eicher committed
Jul 21, 2026 at 07:02 UTC
8950394e0a40a781c61a9ef1f099ef847af03891
5 files changed
+204
-41
tests/v1/core/test_scheduler.py
+135
@@ -1068,6 +1068,141 @@ def test_preempt_during_execution():
1068
assert requests[1].output_token_ids[0] == 42
1069
1070
1071
+def test_prefix_cache_query_not_inflated_by_connector_defer():
1072
+ """The GPU prefix-cache query is recorded at admission, so a request the
1073
+ connector defers several times is counted once, not once per retry."""
1074
+ num_defers_before_matching = 3
1075
+ scheduler = create_scheduler(
1076
+ enable_prefix_caching=True,
1077
+ use_kv_connector=mock_kv(
1078
+ matched_tokens=0,
1079
+ is_async=False,
1080
+ num_defers_before_matching=num_defers_before_matching,
1081
+ ),
1082
+ )
1083
+ request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0]
1084
+ scheduler.add_request(request)
1085
+
1086
+ # Each deferred step re-runs the lookup but records nothing.
1087
+ for _ in range(num_defers_before_matching):
1088
+ assert not scheduler.schedule().scheduled_new_reqs
1089
+
1090
+ output = scheduler.schedule()
1091
+ assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs)
1092
+
1093
+ stats = scheduler.kv_cache_manager.prefix_cache_stats
1094
+ assert stats is not None
1095
+ assert stats.requests == 1
1096
+ assert stats.queries == request.num_tokens
1097
+
1098
+
1099
+def test_preemption_re_records_prefix_cache_query():
1100
+ """A preempted request re-enters the lookup on resume, so its recomputation
1101
+ is counted again into the preempted stats."""
1102
+ scheduler = create_scheduler(enable_prefix_caching=True)
1103
+ request = create_requests(num_requests=1)[0]
1104
+ scheduler.add_request(request)
1105
+
1106
+ scheduler.schedule()
1107
+ stats = scheduler.kv_cache_manager.prefix_cache_stats
1108
+ assert stats is not None
1109
+ assert (stats.requests, stats.preempted_requests) == (1, 0)
1110
+
1111
+ scheduler.running.remove(request)
1112
+ scheduler._preempt_request(request, 0.0)
1113
+ assert request.status == RequestStatus.PREEMPTED
1114
+
1115
+ scheduler.schedule()
1116
+ assert request.status == RequestStatus.RUNNING
1117
+ assert stats.preempted_requests == 1
1118
+
1119
+
1120
+def test_prefix_cache_stats_not_recorded_when_caching_disabled():
1121
+ """With prefix caching off there is no local lookup, so admitting a request
1122
+ records no phantom miss."""
1123
+ scheduler = create_scheduler(enable_prefix_caching=False)
1124
+ for request in create_requests(num_requests=2):
1125
+ scheduler.add_request(request)
1126
+
1127
+ scheduler.schedule()
1128
+
1129
+ stats = scheduler.kv_cache_manager.prefix_cache_stats
1130
+ assert stats is not None
1131
+ assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1132
+
1133
+
1134
+def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
1135
+ """A real cache hit rejected once by allocate_slots and admitted on the next
1136
+ step is counted exactly once, hits included."""
1137
+ block_size = 16
1138
+ scheduler = create_scheduler(
1139
+ enable_prefix_caching=True,
1140
+ enable_chunked_prefill=False,
1141
+ block_size=block_size,
1142
+ )
1143
+
1144
+ # Seed the cache so the next request with the same prompt hits it.
1145
+ seed = create_requests(
1146
+ num_requests=1,
1147
+ num_tokens=block_size * 2,
1148
+ max_tokens=2,
1149
+ same_prompt=True,
1150
+ block_size=block_size,
1151
+ req_ids=["seed"],
1152
+ )[0]
1153
+ scheduler.add_request(seed)
1154
+ _step_until_done(
1155
+ scheduler,
1156
+ scheduler.schedule(),
1157
+ ModelRunnerOutput(
1158
+ req_ids=["seed"],
1159
+ req_id_to_index={"seed": 0},
1160
+ sampled_token_ids=[[1000]],
1161
+ logprobs=None,
1162
+ prompt_logprobs_dict={},
1163
+ pooler_output=[],
1164
+ ),
1165
+ )
1166
+
1167
+ # The seeding step swapped in a fresh accumulator, so re-read it; the
1168
+ # retried request must be the only thing recorded from here on.
1169
+ stats = scheduler.kv_cache_manager.prefix_cache_stats
1170
+ assert stats is not None
1171
+ assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1172
+
1173
+ retried = create_requests(
1174
+ num_requests=1,
1175
+ num_tokens=block_size * 3,
1176
+ max_tokens=1,
1177
+ same_prompt=True,
1178
+ block_size=block_size,
1179
+ req_ids=["retried"],
1180
+ )[0]
1181
+ scheduler.add_request(retried)
1182
+
1183
+ # Reject the first allocation attempt, then delegate to the real one.
1184
+ orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots
1185
+ allocate_results: list = []
1186
+
1187
+ def spy_allocate_slots(*args, **kwargs):
1188
+ result = None if not allocate_results else orig_allocate_slots(*args, **kwargs)
1189
+ allocate_results.append(result)
1190
+ return result
1191
+
1192
+ scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots
1193
+
1194
+ assert not scheduler.schedule().scheduled_new_reqs
1195
+ assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
1196
+
1197
+ assert "retried" in scheduler.schedule().num_scheduled_tokens
1198
+ assert allocate_results[0] is None and allocate_results[1] is not None
1199
+ assert (stats.requests, stats.queries, stats.hits) == (
1200
+ 1,
1201
+ retried.num_tokens,
1202
+ block_size * 2,
1203
+ )
1204
+
1205
+
1206
def test_scheduler_reset_prefix_cache():
1207
scheduler = create_scheduler(enable_prefix_caching=True)
1208
requests = create_requests(num_requests=10)
tests/v1/core/utils.py
+9
-2
@@ -38,8 +38,12 @@ from vllm.v1.structured_output import StructuredOutputManager
38
EOS_TOKEN_ID = 50256
39
40
41
-def mock_kv(matched_tokens: int, is_async: bool):
42
- return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
41
+def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0):
42
+ return MockKVConfig(
43
+ matched_tokens=matched_tokens,
44
+ is_async=is_async,
45
+ num_defers_before_matching=num_defers_before_matching,
46
+ )
47
48
49
def create_scheduler(
@@ -115,6 +119,9 @@ def create_scheduler(
119
kv_connector_extra_config={
120
"matched_tokens": use_kv_connector.matched_tokens,
121
"is_async": use_kv_connector.is_async,
122
+ "num_defers_before_matching": (
123
+ use_kv_connector.num_defers_before_matching
124
+ ),
125
},
126
)
127
elif isinstance(use_kv_connector, str):
tests/v1/kv_connector/unit/utils.py
+10
@@ -362,6 +362,7 @@ class TestExampleConnector(ExampleConnector):
362
class MockKVConfig:
363
matched_tokens: int = 0
364
is_async: bool = False
365
+ num_defers_before_matching: int = 0
366
367
368
class MockKVConnectorMetadata(KVConnectorMetadata):
@@ -384,6 +385,12 @@ class MockKVConnector(KVConnectorBase_V1):
385
self.config = MockKVConfig(
386
matched_tokens=extra_config["matched_tokens"],
387
is_async=extra_config["is_async"],
388
+ num_defers_before_matching=extra_config.get(
389
+ "num_defers_before_matching", 0
390
+ ),
391
+ )
392
+ self._defers_left: defaultdict[str, int] = defaultdict(
393
+ lambda: self.config.num_defers_before_matching
394
)
395
396
def get_num_new_matched_tokens(
@@ -391,6 +398,9 @@ class MockKVConnector(KVConnectorBase_V1):
398
request: Request,
399
num_computed_tokens: int,
400
) -> tuple[int | None, bool]:
401
+ if self._defers_left[request.request_id] > 0:
402
+ self._defers_left[request.request_id] -= 1
403
+ return (None, False)
404
return (self.config.matched_tokens, self.config.is_async)
405
406
def update_state_after_alloc(
vllm/v1/core/kv_cache_manager.py
+16
-9
@@ -204,6 +204,21 @@ class KVCacheManager:
204
self.prefix_cache_stats = PrefixCacheStats()
205
return stats
206
207
+ def prefix_cache_lookup_enabled(self, request: Request) -> bool:
208
+ """Whether a local prefix cache lookup may be run for this request."""
209
+ return self.enable_caching and not request.skip_reading_prefix_cache
210
+
211
+ def record_prefix_cache_stats(self, request: Request, num_hits: int) -> None:
212
+ # Don't count a request that skipped the cache lookup.
213
+ if not self.log_stats or not self.prefix_cache_lookup_enabled(request):
214
+ return
215
+ assert self.prefix_cache_stats is not None
216
+ self.prefix_cache_stats.record(
217
+ num_tokens=request.num_tokens,
218
+ num_hits=num_hits,
219
+ preempted=request.num_preemptions > 0,
220
+ )
221
+
222
def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int]:
223
"""Get the computed (cached) blocks for the request.
224
Note that the computed blocks must be full.
@@ -225,7 +240,7 @@ class KVCacheManager:
240
# disabled or the request is marked as skipping kv cache read
241
# (which happens when the request requires prompt logprobs
242
# or calls a pooling model with all pooling).
228
- if not self.enable_caching or request.skip_reading_prefix_cache:
243
+ if not self.prefix_cache_lookup_enabled(request):
244
return self.empty_kv_cache_blocks, 0, 0
245
246
# NOTE: When all tokens hit the cache, we must recompute the last token
@@ -269,14 +284,6 @@ class KVCacheManager:
284
num_new_computed_tokens + num_uncached if num_uncached else 0
285
)
286
272
- if self.log_stats:
273
- assert self.prefix_cache_stats is not None
274
- self.prefix_cache_stats.record(
275
- num_tokens=request.num_tokens,
276
- num_hits=num_new_computed_tokens,
277
- preempted=request.num_preemptions > 0,
278
- )
279
-
287
blocks = self.create_kv_cache_blocks(computed_blocks)
288
return blocks, num_new_computed_tokens, shared_prefix_boundary
289
vllm/v1/core/sched/scheduler.py
+34
-30
@@ -706,46 +706,47 @@ class Scheduler(SchedulerInterface):
706
num_external_computed_tokens = 0
707
load_kv_async = False
708
connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0
709
+ did_prefix_cache_lookup = False
710
711
# Get already-cached tokens.
712
if request.num_computed_tokens == 0:
713
+ did_prefix_cache_lookup = True
714
# Get locally-cached tokens.
715
if (
716
self.connector is not None
717
and self.has_mamba_layers
718
and isinstance(
717
- self.kv_cache_manager.coordinator,
718
- HybridKVCacheCoordinator,
719
+ self.kv_cache_manager.coordinator, HybridKVCacheCoordinator
720
)
721
):
721
- computed, per_group_hits = (
722
- self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group(
723
- request.block_hashes, request.num_tokens - 1
724
- )
725
- )
726
- new_computed_blocks = (
727
- self.kv_cache_manager.create_kv_cache_blocks(computed)
728
- )
729
- # NOTE(ZhanqiuHu): For Mamba hybrid models,
730
- # num_new_local_computed_tokens should be the FA hit
731
- # length. This value is passed to the connector's
732
- # get_num_new_matched_tokens which computes:
733
- # external = total - local_computed.
734
- # Using the FA hit skips re-transferring FA blocks
735
- # already cached on D-side. The Mamba state (always
736
- # the last block) is transferred unconditionally by
737
- # _apply_prefix_caching in nixl/worker.py.
738
- num_new_local_computed_tokens = max(per_group_hits)
722
# The per-group lookup does not detect an uncached shared
723
# prefix, so there is no junction to pin in this path.
724
request.shared_prefix_boundary = 0
742
- if self.kv_cache_manager.log_stats:
743
- assert self.kv_cache_manager.prefix_cache_stats is not None
744
- self.kv_cache_manager.prefix_cache_stats.record(
745
- num_tokens=request.num_tokens,
746
- num_hits=num_new_local_computed_tokens,
747
- preempted=request.num_preemptions > 0,
725
+ kv_cache_manager = self.kv_cache_manager
726
+ if not kv_cache_manager.prefix_cache_lookup_enabled(request):
727
+ # Mirror the get_computed_blocks() early-out: the
728
+ # request must recompute its prompt.
729
+ new_computed_blocks = kv_cache_manager.empty_kv_cache_blocks
730
+ num_new_local_computed_tokens = 0
731
+ else:
732
+ computed, per_group_hits = (
733
+ self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group(
734
+ request.block_hashes, request.num_tokens - 1
735
+ )
736
+ )
737
+ new_computed_blocks = (
738
+ self.kv_cache_manager.create_kv_cache_blocks(computed)
739
)
740
+ # NOTE(ZhanqiuHu): For Mamba hybrid models,
741
+ # num_new_local_computed_tokens should be the FA hit
742
+ # length. This value is passed to the connector's
743
+ # get_num_new_matched_tokens which computes:
744
+ # external = total - local_computed.
745
+ # Using the FA hit skips re-transferring FA blocks
746
+ # already cached on D-side. The Mamba state (always
747
+ # the last block) is transferred unconditionally by
748
+ # _apply_prefix_caching in nixl/worker.py.
749
+ num_new_local_computed_tokens = max(per_group_hits)
750
else:
751
(
752
new_computed_blocks,
@@ -799,10 +800,7 @@ class Scheduler(SchedulerInterface):
800
continue
801
802
# Track first scheduled prefill, not post-preemption repeat prefills
802
- if (
803
- request.prefill_stats is not None
804
- and request.num_preemptions <= 0
805
- ):
803
+ if request.prefill_stats and request.num_preemptions <= 0:
804
assert num_computed_tokens <= request.num_prompt_tokens
805
request.prefill_stats.set(
806
num_prompt_tokens=request.num_prompt_tokens,
@@ -971,6 +969,12 @@ class Scheduler(SchedulerInterface):
969
preempted=request.num_preemptions > 0,
970
)
971
972
+ # Record at admission so unscheduled lookups are not counted.
973
+ if did_prefix_cache_lookup:
974
+ self.kv_cache_manager.record_prefix_cache_stats(
975
+ request, num_new_local_computed_tokens
976
+ )
977
+
978
request = request_queue.pop_request()
979
if load_kv_async:
980
# If loading async, allocate memory and put request