[P/D] Report prefill worker cache hits in prompt_tokens_details (#54222)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
Robert Shaw committed
Sep 16, 2026 at 11:23 UTC
42919b49c573aa0b06710616a3b9785674f36ef5
4 files changed
+133
-1
tests/v1/engine/test_output_processor.py
+53
-1
@@ -23,6 +23,7 @@ from vllm.tokenizers import TokenizerLike
23
from vllm.v1.engine import (
24
EngineCoreEvent,
25
EngineCoreEventType,
26
+ EngineCoreOutput,
27
EngineCoreOutputs,
28
EngineCoreRequest,
29
FinishReason,
@@ -32,7 +33,7 @@ from vllm.v1.engine.output_processor import (
33
RequestOutputCollector,
34
RequestState,
35
)
35
-from vllm.v1.metrics.stats import IterationStats, SchedulerStats
36
+from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats
37
38
39
@pytest.mark.parametrize("flat_logprobs", [False, True])
@@ -171,6 +172,57 @@ def test_incremental_detokenization(
172
assert not output_processor.has_unfinished_requests()
173
174
175
+@pytest.mark.parametrize("do_remote_prefill", [True, False])
176
+def test_remote_prefill_cached_tokens_override(do_remote_prefill: bool):
177
+ """P/D disaggregation: num_cached_tokens should report the P worker's
178
+ cache hits (passed via kv_transfer_params) instead of the local count,
179
+ which sees the KVs pulled from the remote prefill as a ~100% hit."""
180
+ output_processor = OutputProcessor(tokenizer=None, log_stats=False)
181
+
182
+ prompt_tokens = [1, 2, 3, 4, 5, 6, 7, 8]
183
+ kv_transfer_params = {
184
+ "do_remote_prefill": do_remote_prefill,
185
+ "remote_prefill_cached_tokens": 5,
186
+ }
187
+ request = EngineCoreRequest(
188
+ request_id="request-0-int",
189
+ external_req_id="request-0",
190
+ prompt_token_ids=prompt_tokens,
191
+ mm_features=None,
192
+ arrival_time=0,
193
+ lora_request=None,
194
+ cache_salt=None,
195
+ data_parallel_rank=None,
196
+ sampling_params=SamplingParams(
197
+ detokenize=False,
198
+ extra_args={"kv_transfer_params": kv_transfer_params},
199
+ ),
200
+ pooling_params=None,
201
+ )
202
+ output_processor.add_request(request, prompt=None)
203
+
204
+ prefill_stats = PrefillStats()
205
+ prefill_stats.set(
206
+ num_prompt_tokens=len(prompt_tokens),
207
+ num_local_cached_tokens=0,
208
+ num_external_cached_tokens=len(prompt_tokens) - 1,
209
+ )
210
+ processed = output_processor.process_outputs(
211
+ [
212
+ EngineCoreOutput(
213
+ request_id="request-0-int",
214
+ new_token_ids=[42],
215
+ prefill_stats=prefill_stats,
216
+ )
217
+ ]
218
+ )
219
+ request_output = processed.request_outputs[0]
220
+ if do_remote_prefill:
221
+ assert request_output.num_cached_tokens == 5
222
+ else:
223
+ assert request_output.num_cached_tokens == len(prompt_tokens) - 1
224
+
225
+
226
def test_request_stream_interval_raises_but_not_below_engine_default(
227
dummy_test_vectors,
228
):
tests/v1/kv_connector/unit/test_nixl_connector.py
+53
@@ -74,6 +74,7 @@ from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
74
from vllm.v1.request import RequestStatus
75
76
from .utils import (
77
+ create_model_runner_output,
78
create_request,
79
create_scheduler,
80
create_vllm_config,
@@ -348,6 +349,58 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
349
assert req_meta.awaiting_kvs is False
350
351
352
+def test_prefill_exports_cached_tokens_in_kv_transfer_params():
353
+ """The P worker reports its own prefix-cache hits in the returned
354
+ kv_transfer_params so the D worker can surface them in
355
+ prompt_tokens_details instead of the ~100% local hit it measures
356
+ when pulling the KVs from the remote."""
357
+ vllm_config = create_vllm_config()
358
+ scheduler = create_scheduler(vllm_config)
359
+
360
+ BLOCK_SIZE = vllm_config.cache_config.block_size
361
+ NUM_TOKENS = BLOCK_SIZE * 3
362
+
363
+ # Warm the prefix cache with a request sharing the full prompt.
364
+ warmup = create_request(
365
+ request_id=1,
366
+ num_tokens=NUM_TOKENS,
367
+ common_prefix_len=NUM_TOKENS,
368
+ block_size=BLOCK_SIZE,
369
+ )
370
+ scheduler.add_request(warmup)
371
+ scheduler_output = scheduler.schedule()
372
+ model_runner_output = create_model_runner_output([warmup], use_eos=True)
373
+ scheduler.update_from_output(scheduler_output, model_runner_output)
374
+
375
+ # P-side request with the same prompt hits the local cache on all but
376
+ # the last block (recomputed to obtain logits).
377
+ request = create_request(
378
+ request_id=2,
379
+ num_tokens=NUM_TOKENS,
380
+ common_prefix_len=NUM_TOKENS,
381
+ block_size=BLOCK_SIZE,
382
+ do_remote_decode=True,
383
+ )
384
+ scheduler.add_request(request)
385
+ scheduler_output = scheduler.schedule()
386
+ assert request.prefill_stats is not None
387
+ assert request.prefill_stats.num_cached_tokens == NUM_TOKENS - BLOCK_SIZE
388
+
389
+ # max_tokens=1, so the request finishes and returns kv_transfer_params.
390
+ model_runner_output = create_model_runner_output([request])
391
+ engine_core_outputs = scheduler.update_from_output(
392
+ scheduler_output, model_runner_output
393
+ )
394
+
395
+ output = engine_core_outputs[0].outputs[0]
396
+ assert output.finish_reason is not None
397
+ assert output.kv_transfer_params is not None
398
+ assert (
399
+ output.kv_transfer_params["remote_prefill_cached_tokens"]
400
+ == NUM_TOKENS - BLOCK_SIZE
401
+ )
402
+
403
+
404
@patch(
405
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
406
FakeNixlWrapper,
vllm/v1/core/sched/scheduler.py
+11
@@ -2110,6 +2110,17 @@ class Scheduler(SchedulerInterface):
2110
finished = self._handle_stopped_request(request)
2111
if finished:
2112
kv_transfer_params, ec_transfer_params = self._free_request(request)
2113
+ if (
2114
+ prefill_stats is not None
2115
+ and kv_transfer_params is not None
2116
+ and kv_transfer_params.get("do_remote_prefill")
2117
+ ):
2118
+ # P-side cache hits, so D can report them in
2119
+ # prompt_tokens_details instead of its own (~100%)
2120
+ # hit rate from the KV transfer.
2121
+ kv_transfer_params["remote_prefill_cached_tokens"] = (
2122
+ prefill_stats.num_cached_tokens
2123
+ )
2124
2125
if status_before_stop == RequestStatus.RUNNING:
2126
stopped_running_reqs.add(request)
vllm/v1/engine/output_processor.py
+16
@@ -155,6 +155,7 @@ class RequestState:
155
n: int | None = None,
156
temperature: float | None = None,
157
stream_input: bool = False,
158
+ remote_prefill_cached_tokens: int | None = None,
159
):
160
self.request_id = request_id
161
self.external_req_id = external_req_id
@@ -179,6 +180,7 @@ class RequestState:
180
self.queue = queue
181
self.num_cached_tokens = 0
182
self.num_cache_creation_tokens = 0
183
+ self.remote_prefill_cached_tokens = remote_prefill_cached_tokens
184
# Per-sequence spec-decode accumulator; arrives once (on finish) via
185
# EngineCoreOutput, then attached to this sequence's CompletionOutput.
186
self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None
@@ -230,7 +232,18 @@ class RequestState:
232
log_stats: bool,
233
stream_interval: int,
234
) -> "RequestState":
235
+ remote_prefill_cached_tokens = None
236
if sampling_params := request.sampling_params:
237
+ # In a remote prefill scenario, report cached tokens
238
+ # as cache hit rate on the remote P worker.
239
+ if sampling_params.extra_args:
240
+ kv_transfer_params = sampling_params.extra_args.get(
241
+ "kv_transfer_params"
242
+ )
243
+ if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"):
244
+ cached = kv_transfer_params.get("remote_prefill_cached_tokens")
245
+ if isinstance(cached, int):
246
+ remote_prefill_cached_tokens = cached
247
if not sampling_params.detokenize:
248
tokenizer = None
249
output_kind = sampling_params.output_kind
@@ -281,6 +294,7 @@ class RequestState:
294
log_stats=log_stats,
295
stream_interval=stream_interval,
296
stream_input=request.resumable,
297
+ remote_prefill_cached_tokens=remote_prefill_cached_tokens,
298
)
299
300
def make_request_output(
@@ -687,6 +701,8 @@ class OutputProcessor:
701
req_state.num_cache_creation_tokens = (
702
engine_core_output.prefill_stats.num_cache_creation_tokens
703
)
704
+ if req_state.remote_prefill_cached_tokens is not None:
705
+ req_state.num_cached_tokens = req_state.remote_prefill_cached_tokens
706
req_state.is_prefilling = False
707
708
if engine_core_output.spec_decode_metrics is not None: