A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 08:30 UTC
[Feature][V1]: suupports cached_tokens in response usage (#18149)
Co-authored-by: simon-mo <xmo@berkeley.edu>
Chauncey committed
May 23, 2025 at 16:41 UTC
b046cf792d49c36cabc5589501fcc1da31945f64
5 files changed
+27
-5
tests/v1/core/test_scheduler_e2e.py
+10
-1
@@ -19,7 +19,8 @@ def model() -> LLM:
19
enable_prefix_caching=True,
20
long_prefill_token_threshold=2,
21
max_num_batched_tokens=6,
22
- max_num_seqs=3)
22
+ max_num_seqs=3,
23
+ block_size=16)
24
25
26
def test_concurrent_partial_prefill(model):
@@ -27,3 +28,11 @@ def test_concurrent_partial_prefill(model):
28
assert len(outputs) == 3
29
for output in outputs:
30
assert len(output.outputs) == 1
31
+
32
+
33
+def test_prefix_cache_stats_is_recorded(model):
34
+ # 17 tokens will make sure first 16 tokens are cached in a block
35
+ input_tokens = {"prompt_token_ids": [101] * 17}
36
+ _ = model.generate([input_tokens])
37
+ outputs = model.generate([input_tokens])
38
+ assert outputs[0].num_cached_tokens == 16
vllm/v1/core/sched/scheduler.py
+4
-1
@@ -457,7 +457,9 @@ class Scheduler(SchedulerInterface):
457
token_budget -= num_new_tokens
458
request.status = RequestStatus.RUNNING
459
request.num_computed_tokens = num_computed_tokens
460
-
460
+ # Count the number of prifix cached tokens.
461
+ if request.num_cached_tokens < 0:
462
+ request.num_cached_tokens = num_computed_tokens
463
# Encoder-related.
464
if encoder_inputs_to_schedule:
465
scheduled_encoder_inputs[request.request_id] = (
@@ -798,6 +800,7 @@ class Scheduler(SchedulerInterface):
800
stop_reason=request.stop_reason,
801
events=request.take_events(),
802
kv_transfer_params=kv_transfer_params,
803
+ num_cached_tokens=request.num_cached_tokens,
804
))
805
806
else:
vllm/v1/engine/__init__.py
+3
@@ -107,6 +107,9 @@ class EngineCoreOutput(
107
events: Optional[list[EngineCoreEvent]] = None
108
kv_transfer_params: Optional[dict[str, Any]] = None
109
110
+ # The number of tokens with prefix cache hits.
111
+ num_cached_tokens: int = 0
112
+
113
@property
114
def finished(self) -> bool:
115
return self.finish_reason is not None
vllm/v1/engine/output_processor.py
+6
-3
@@ -147,6 +147,7 @@ class RequestState:
147
finish_reason: Optional[FinishReason],
148
stop_reason: Union[int, str, None],
149
kv_transfer_params: Optional[dict[str, Any]] = None,
150
+ num_cached_tokens: int = 0,
151
) -> Optional[RequestOutput]:
152
153
finished = finish_reason is not None
@@ -169,7 +170,7 @@ class RequestState:
170
return None
171
172
return self._new_request_output(request_id, outputs, finished,
172
- kv_transfer_params)
173
+ kv_transfer_params, num_cached_tokens)
174
175
def _new_request_output(
176
self,
@@ -177,6 +178,7 @@ class RequestState:
178
outputs: list[CompletionOutput],
179
finished: bool,
180
kv_transfer_params: Optional[dict[str, Any]] = None,
181
+ num_cached_tokens: int = 0,
182
) -> RequestOutput:
183
184
if self.output_kind == RequestOutputKind.DELTA:
@@ -193,6 +195,7 @@ class RequestState:
195
outputs=outputs,
196
finished=finished,
197
kv_transfer_params=kv_transfer_params,
198
+ num_cached_tokens=num_cached_tokens,
199
)
200
201
def _new_completion_output(
@@ -340,7 +343,7 @@ class OutputProcessor:
343
finish_reason = engine_core_output.finish_reason
344
stop_reason = engine_core_output.stop_reason
345
kv_transfer_params = engine_core_output.kv_transfer_params
343
-
346
+ num_cached_tokens = engine_core_output.num_cached_tokens
347
req_state.is_prefilling = False
348
349
# 2) Detokenize the token ids into text and perform stop checks.
@@ -356,7 +359,7 @@ class OutputProcessor:
359
# 4) Create and handle RequestOutput objects.
360
if request_output := req_state.make_request_output(
361
new_token_ids, finish_reason, stop_reason,
359
- kv_transfer_params):
362
+ kv_transfer_params, num_cached_tokens):
363
if req_state.queue is not None:
364
# AsyncLLM: put into queue for handling by generate().
365
req_state.queue.put(request_output)
vllm/v1/request.py
+4
@@ -77,6 +77,10 @@ class Request:
77
self.output_token_ids = ConstantList(self._output_token_ids)
78
self.all_token_ids = ConstantList(self._all_token_ids)
79
80
+ # State
81
+ # The number of tokens with prefix cache hits.
82
+ self.num_cached_tokens = -1
83
+
84
@classmethod
85
def from_engine_core_request(cls, request: EngineCoreRequest) -> "Request":
86
if request.mm_inputs is not None: