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

[Front-end] [Messages] Populate `num_cache_creation_tokens` (#48535)

Signed-off-by: Yifan Zong <yzong@redhat.com>

yzong-rh committed Jul 18, 2026 at 13:04 UTC a287eb163fb6f8f007a4a78411fb54c8dde64cc7
15 files changed +201 -127
rust/src/engine-core-client/src/protocol/stats.rs
+3
@@ -103,6 +103,9 @@ pub struct PrefillStats {
103 /// Tokens to be prefilled from external KV transfer.
104 #[serde(default)]
105 pub num_external_cached_tokens: u32,
106 + /// Prompt tokens newly admitted into the local prefix cache.
107 + #[serde(default)]
108 + pub num_cache_creation_tokens: u32,
109 }
110
111 /// Stats for debugging the metrics calculation.
rust/src/llm/src/request_metrics.rs
+1
@@ -375,6 +375,7 @@ mod tests {
375 num_cached_tokens: 4,
376 num_local_cached_tokens: 4,
377 num_external_cached_tokens: 0,
378 + ..Default::default()
379 }),
380 ..Default::default()
381 },
tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
+25 -60
@@ -23,7 +23,6 @@ from vllm.entrypoints.anthropic.protocol import (
23 from vllm.entrypoints.anthropic.serving import (
24 AnthropicServingMessages,
25 _build_anthropic_usage,
26 - _get_cached_tokens,
26 )
27 from vllm.entrypoints.openai.chat_completion.protocol import (
28 ChatCompletionResponse,
@@ -668,42 +667,6 @@ class TestThinkingBlockConversion:
667 # ======================================================================
668
669
671 -class TestGetCachedTokens:
672 - """Tests for _get_cached_tokens helper."""
673 -
674 - def test_none_usage(self):
675 - assert _get_cached_tokens(None) is None
676 -
677 - def test_no_prompt_tokens_details(self):
678 - usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
679 - assert _get_cached_tokens(usage) is None
680 -
681 - def test_cached_tokens_present(self):
682 - usage = UsageInfo(
683 - prompt_tokens=100,
684 - completion_tokens=10,
685 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
686 - )
687 - assert _get_cached_tokens(usage) == 80
688 -
689 - def test_cached_tokens_zero(self):
690 - """Zero cached tokens should return 0, not None."""
691 - usage = UsageInfo(
692 - prompt_tokens=100,
693 - completion_tokens=10,
694 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
695 - )
696 - assert _get_cached_tokens(usage) == 0
697 -
698 - def test_cached_tokens_none_in_details(self):
699 - usage = UsageInfo(
700 - prompt_tokens=100,
701 - completion_tokens=10,
702 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
703 - )
704 - assert _get_cached_tokens(usage) is None
705 -
706 -
670 class TestBuildAnthropicUsage:
671 """Tests for _build_anthropic_usage helper.
672
@@ -711,36 +674,32 @@ class TestBuildAnthropicUsage:
674 vLLM's prompt_tokens is the total.
675 """
676
714 - def test_no_cache_info(self):
715 - """When cache info is unavailable, return raw prompt_tokens."""
716 - result = _build_anthropic_usage(100, 10, None)
717 - assert result.input_tokens == 100
718 - assert result.output_tokens == 10
719 - assert result.cache_read_input_tokens is None
720 - assert result.cache_creation_input_tokens is None
721 -
677 def test_cache_hit(self):
678 """When cache is hit, input_tokens excludes cached tokens."""
679 usage = UsageInfo(
680 prompt_tokens=100,
681 completion_tokens=10,
727 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
682 + prompt_tokens_details=PromptTokenUsageInfo(
683 + cached_tokens=80, created_cache_tokens=10
684 + ),
685 )
729 - result = _build_anthropic_usage(100, 10, usage)
730 - assert result.input_tokens == 20 # 100 - 80
686 + result = _build_anthropic_usage(usage)
687 + assert result.input_tokens == 10 # 100 - 80 - 10
688 assert result.output_tokens == 10
689 assert result.cache_read_input_tokens == 80
733 - assert result.cache_creation_input_tokens == 0
690 + assert result.cache_creation_input_tokens == 10
691
692 def test_zero_cached_tokens(self):
693 """Zero cached tokens should still set cache_creation to 0."""
694 usage = UsageInfo(
695 prompt_tokens=100,
696 completion_tokens=10,
740 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
697 + prompt_tokens_details=PromptTokenUsageInfo(
698 + cached_tokens=0, created_cache_tokens=0
699 + ),
700 )
742 - result = _build_anthropic_usage(100, 10, usage)
743 - assert result.input_tokens == 100 # 100 - 0
701 + result = _build_anthropic_usage(usage)
702 + assert result.input_tokens == 100 # 100 - 0 - 0
703 assert result.cache_read_input_tokens == 0
704 assert result.cache_creation_input_tokens == 0
705
@@ -749,9 +708,11 @@ class TestBuildAnthropicUsage:
708 usage = UsageInfo(
709 prompt_tokens=100,
710 completion_tokens=10,
752 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
711 + prompt_tokens_details=PromptTokenUsageInfo(
712 + cached_tokens=100, created_cache_tokens=0
713 + ),
714 )
754 - result = _build_anthropic_usage(100, 10, usage)
715 + result = _build_anthropic_usage(usage)
716 assert result.input_tokens == 0
717 assert result.cache_read_input_tokens == 100
718 assert result.cache_creation_input_tokens == 0
@@ -759,7 +720,7 @@ class TestBuildAnthropicUsage:
720 def test_no_prompt_tokens_details(self):
721 """UsageInfo without prompt_tokens_details returns no cache info."""
722 usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
762 - result = _build_anthropic_usage(100, 10, usage)
723 + result = _build_anthropic_usage(usage)
724 assert result.input_tokens == 100
725 assert result.cache_read_input_tokens is None
726 assert result.cache_creation_input_tokens is None
@@ -1241,7 +1202,9 @@ class TestStreamingCacheUsageSemantics:
1202 prompt_tokens=100,
1203 completion_tokens=5,
1204 total_tokens=105,
1244 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
1205 + prompt_tokens_details=PromptTokenUsageInfo(
1206 + cached_tokens=80, created_cache_tokens=10
1207 + ),
1208 ),
1209 )
1210 yield "data: [DONE]"
@@ -1263,9 +1226,9 @@ class TestStreamingCacheUsageSemantics:
1226 delta_usage = next(
1227 data["usage"] for ev, data in events if ev == "message_delta"
1228 )
1266 - assert delta_usage["input_tokens"] == 20 # 100 - 80
1229 + assert delta_usage["input_tokens"] == 10 # 100 - 80 - 10
1230 assert delta_usage["cache_read_input_tokens"] == 80
1268 - assert delta_usage["cache_creation_input_tokens"] == 0
1231 + assert delta_usage["cache_creation_input_tokens"] == 10
1232
1233 @pytest.mark.asyncio
1234 async def test_streaming_no_cache_hit(self):
@@ -1284,7 +1247,9 @@ class TestStreamingCacheUsageSemantics:
1247 prompt_tokens=50,
1248 completion_tokens=5,
1249 total_tokens=55,
1287 - prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
1250 + prompt_tokens_details=PromptTokenUsageInfo(
1251 + cached_tokens=0, created_cache_tokens=0
1252 + ),
1253 ),
1254 )
1255 yield "data: [DONE]"
@@ -1302,7 +1267,7 @@ class TestStreamingCacheUsageSemantics:
1267 assert start_usage["input_tokens"] == 50
1268 assert "cache_read_input_tokens" not in start_usage
1269 assert "cache_creation_input_tokens" not in start_usage
1305 - assert delta_usage["input_tokens"] == 50 # 50 - 0
1270 + assert delta_usage["input_tokens"] == 50 # 50 - 0 - 0
1271 assert delta_usage["cache_read_input_tokens"] == 0
1272 assert delta_usage["cache_creation_input_tokens"] == 0
1273
tests/entrypoints/anthropic/test_messages.py
+52
@@ -18,6 +18,7 @@ def server():
18 "--max-model-len",
19 "2048",
20 "--enforce-eager",
21 + "--enable-prompt-tokens-details",
22 "--enable-auto-tool-choice",
23 "--tool-call-parser",
24 "hermes",
@@ -191,3 +192,54 @@ async def test_anthropic_structured_output(client: anthropic.AsyncAnthropic):
192 json_obj = json.loads(response.content[0].text)
193 for key in ["name", "email", "plan_interest", "demo_requested"]:
194 assert key in json_obj, f"Missing key in output: {key}"
195 +
196 +
197 +@pytest.mark.asyncio
198 +async def test_anthropic_streaming_cache_usage(client: anthropic.AsyncAnthropic):
199 + async def get_stream_usage(resp):
200 + prompt_tokens = None
201 + usage = None
202 + async for chunk in resp:
203 + if (
204 + chunk.type == "message_start"
205 + and chunk.message is not None
206 + and chunk.message.usage is not None
207 + ):
208 + prompt_tokens = chunk.message.usage.input_tokens
209 + elif chunk.type == "message_delta" and chunk.usage is not None:
210 + usage = chunk.usage
211 +
212 + assert usage is not None
213 + assert usage.input_tokens >= 0
214 + assert usage.output_tokens >= 0
215 + cache_created = usage.cache_creation_input_tokens
216 + cache_read = usage.cache_read_input_tokens
217 + assert cache_read is not None
218 + assert cache_created is not None
219 + assert cache_created >= 0
220 + assert cache_read >= 0
221 + assert prompt_tokens == usage.input_tokens + cache_created + cache_read
222 + return usage
223 +
224 + request = dict(
225 + model="claude-3-7-sonnet-latest",
226 + max_tokens=1,
227 + temperature=0.0,
228 + messages=[
229 + {
230 + "role": "user",
231 + "content": "Cache coverage sentinel. " * 256
232 + + "Answer with exactly one word: ok.",
233 + }
234 + ],
235 + stream=True,
236 + )
237 +
238 + cold_usage = await get_stream_usage(await client.messages.create(**request))
239 + assert cold_usage.cache_read_input_tokens == 0
240 + assert cold_usage.cache_creation_input_tokens is not None
241 + assert cold_usage.cache_creation_input_tokens > 0
242 +
243 + warm_usage = await get_stream_usage(await client.messages.create(**request))
244 + assert warm_usage.cache_read_input_tokens is not None
245 + assert warm_usage.cache_read_input_tokens > 0
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+12 -5
@@ -831,16 +831,23 @@ def test_mm_prompt_tokens_details():
831 assert counts == {"image": 600, "video": 1200}
832
833 # Gated off, or nothing to report -> no details.
834 - assert _make_prompt_tokens_details(False, 5, counts) is None
835 - assert _make_prompt_tokens_details(True, None, None) is None
834 + assert _make_prompt_tokens_details(False, 5, 0, counts) is None
835 + assert _make_prompt_tokens_details(True, None, None, None) is None
836
837 # Zero cached_tokens is still reported (not None), matching the cached-only
838 # behavior; multimodal counts ride alongside even when cached_tokens is None.
839 - assert _make_prompt_tokens_details(True, 0, None).cached_tokens == 0
840 - details = _make_prompt_tokens_details(True, None, counts)
839 + details = _make_prompt_tokens_details(True, 0, 0, None)
840 + assert details.cached_tokens == 0
841 + assert details.created_cache_tokens == 0
842 + assert details.multimodal_tokens is None
843 + details = _make_prompt_tokens_details(True, None, None, counts)
844 assert details.cached_tokens is None
845 + assert details.created_cache_tokens is None
846 + assert details.multimodal_tokens == {"image": 600, "video": 1200}
847 + details = _make_prompt_tokens_details(True, 3, 0, counts)
848 + assert details.cached_tokens == 3
849 + assert details.created_cache_tokens == 0
850 assert details.multimodal_tokens == {"image": 600, "video": 1200}
843 - assert _make_prompt_tokens_details(True, 3, counts).cached_tokens == 3
851
852
853 @pytest.mark.asyncio
tests/v1/core/test_async_scheduler.py
+1
@@ -280,6 +280,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
280 scheduler.waiting = Mock()
281 scheduler.kv_cache_manager = Mock()
282 scheduler.kv_cache_manager.take_events.return_value = None
283 + scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
284 scheduler.kv_event_publisher = Mock()
285 scheduler.finished_req_ids = set()
286 scheduler.finished_req_ids_dict = None
tests/v1/core/test_scheduler.py
+1
@@ -3036,6 +3036,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3036 scheduler.waiting = Mock()
3037 scheduler.kv_cache_manager = Mock()
3038 scheduler.kv_cache_manager.take_events.return_value = None
3039 + scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
3040 scheduler.kv_event_publisher = Mock()
3041 scheduler.finished_req_ids = set()
3042 scheduler.finished_req_ids_dict = None
vllm/entrypoints/anthropic/serving.py
+29 -51
@@ -53,47 +53,43 @@ from vllm.renderers.online_renderer import OnlineRenderer
53 logger = logging.getLogger(__name__)
54
55
56 -def _get_cached_tokens(usage: UsageInfo | None) -> int | None:
57 - """Extract cached token count from OpenAI UsageInfo."""
58 - if usage is None or usage.prompt_tokens_details is None:
59 - return None
60 - return usage.prompt_tokens_details.cached_tokens
61 -
62 -
56 def _build_anthropic_usage(
64 - prompt_tokens: int,
65 - completion_tokens: int | None,
57 usage: UsageInfo | None,
58 ) -> AnthropicUsage:
68 - """Build an AnthropicUsage from OpenAI-style token counts.
59 + """Build an AnthropicUsage from UsageInfo.
60
61 Anthropic defines ``total_input == input_tokens + cache_read +
62 cache_creation``. vLLM's ``prompt_tokens`` is the total, so
72 - ``input_tokens = prompt_tokens - cached_tokens``.
63 + ``input_tokens = prompt_tokens - cache_read - cache_creation``.
64
74 - OpenAI usage only exposes ``cached_tokens`` (hits); there is no
75 - cache-creation analog, so ``cache_creation_input_tokens`` is ``0``
76 - when cache info is present. When cache info is absent (e.g.
77 - ``--enable-prompt-tokens-details`` off, or a streaming chunk that
78 - hasn't carried it yet), cache fields are left **unset** so
79 - ``exclude_unset=True`` serialization omits them entirely.
65 + Cache fields are taken from ``UsageInfo.prompt_tokens_details``.
66 + When cache info is absent (e.g. ``--enable-prompt-tokens-details``
67 + off, or a streaming chunk that hasn't carried it yet), cache fields
68 + are left **unset** so ``exclude_unset=True`` serialization omits them
69 + entirely.
70
71 ``completion_tokens`` follows ``UsageInfo`` and may be ``None`` on
72 intermediate stream chunks; we coerce to ``0`` for the wire format.
73 """
84 - output_tokens = completion_tokens or 0
85 - cached = _get_cached_tokens(usage)
86 - if cached is not None:
87 - return AnthropicUsage(
88 - input_tokens=prompt_tokens - cached,
89 - output_tokens=output_tokens,
90 - cache_read_input_tokens=cached,
91 - cache_creation_input_tokens=0,
92 - )
93 - return AnthropicUsage(
94 - input_tokens=prompt_tokens,
95 - output_tokens=output_tokens,
96 - )
74 + kwargs = {}
75 + if usage is None:
76 + kwargs["input_tokens"] = 0
77 + kwargs["output_tokens"] = 0
78 + else:
79 + kwargs["output_tokens"] = usage.completion_tokens or 0
80 + input_tokens = usage.prompt_tokens
81 +
82 + if (details := usage.prompt_tokens_details) is not None:
83 + if (cache_read := details.cached_tokens) is not None:
84 + input_tokens -= cache_read
85 + kwargs["cache_read_input_tokens"] = cache_read
86 +
87 + if (cache_creation := details.created_cache_tokens) is not None:
88 + input_tokens -= cache_creation
89 + kwargs["cache_creation_input_tokens"] = cache_creation
90 +
91 + kwargs["input_tokens"] = max(0, input_tokens)
92 + return AnthropicUsage(**kwargs)
93
94
95 def wrap_data_with_event(data: str, event: str):
@@ -625,11 +621,7 @@ class AnthropicServingMessages(OpenAIServingChat):
621 id=generator.id,
622 content=[],
623 model=generator.model,
628 - usage=_build_anthropic_usage(
629 - generator.usage.prompt_tokens,
630 - generator.usage.completion_tokens,
631 - generator.usage,
632 - ),
624 + usage=_build_anthropic_usage(generator.usage),
625 kv_transfer_params=generator.kv_transfer_params,
626 ec_transfer_params=generator.ec_transfer_params,
627 )
@@ -816,13 +808,7 @@ class AnthropicServingMessages(OpenAIServingChat):
808 model=origin_chunk.model,
809 stop_reason=None,
810 stop_sequence=None,
819 - usage=_build_anthropic_usage(
820 - origin_chunk.usage.prompt_tokens
821 - if origin_chunk.usage
822 - else 0,
823 - 0,
824 - origin_chunk.usage,
825 - ),
811 + usage=_build_anthropic_usage(origin_chunk.usage),
812 ),
813 )
814 first_item = False
@@ -840,15 +826,7 @@ class AnthropicServingMessages(OpenAIServingChat):
826 chunk = AnthropicStreamEvent(
827 type="message_delta",
828 delta=AnthropicDelta(stop_reason=stop_reason),
843 - usage=_build_anthropic_usage(
844 - origin_chunk.usage.prompt_tokens
845 - if origin_chunk.usage
846 - else 0,
847 - origin_chunk.usage.completion_tokens
848 - if origin_chunk.usage
849 - else 0,
850 - origin_chunk.usage,
851 - ),
829 + usage=_build_anthropic_usage(origin_chunk.usage),
830 )
831 data = chunk.model_dump_json(exclude_unset=True)
832 yield wrap_data_with_event(data, "message_delta")
vllm/entrypoints/openai/chat_completion/serving.py
+11 -1
@@ -90,15 +90,21 @@ def _get_mm_token_counts(engine_input: EngineInput) -> dict[str, int]:
90 def _make_prompt_tokens_details(
91 enable_prompt_tokens_details: bool,
92 num_cached_tokens: int | None,
93 + num_cache_creation_tokens: int | None,
94 mm_token_counts: dict[str, int] | None,
95 ) -> PromptTokenUsageInfo | None:
96 """Build ``prompt_tokens_details`` from cached + multimodal token counts."""
97 if not enable_prompt_tokens_details:
98 return None
98 - if num_cached_tokens is None and not mm_token_counts:
99 + if (
100 + num_cached_tokens is None
101 + and num_cache_creation_tokens is None
102 + and not mm_token_counts
103 + ):
104 return None
105 return PromptTokenUsageInfo(
106 cached_tokens=num_cached_tokens,
107 + created_cache_tokens=num_cache_creation_tokens,
108 multimodal_tokens=mm_token_counts or None,
109 )
110
@@ -427,6 +433,7 @@ class OpenAIServingChat(GenerateBaseServing):
433 finish_reason_sent = [False] * num_choices
434 num_prompt_tokens = 0
435 num_cached_tokens = None
436 + num_cache_creation_tokens = None
437 tools_streamed = [False] * num_choices
438
439 if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
@@ -479,6 +486,7 @@ class OpenAIServingChat(GenerateBaseServing):
486 # response (by the try...catch).
487 if first_iteration:
488 num_cached_tokens = res.num_cached_tokens
489 + num_cache_creation_tokens = res.num_cache_creation_tokens
490 # Send first response for each request.n (index) with
491 # the role
492 role = self.get_chat_request_role(request)
@@ -756,6 +764,7 @@ class OpenAIServingChat(GenerateBaseServing):
764 final_usage.prompt_tokens_details = _make_prompt_tokens_details(
765 self.enable_prompt_tokens_details,
766 num_cached_tokens,
767 + num_cache_creation_tokens,
768 mm_token_counts,
769 )
770
@@ -1030,6 +1039,7 @@ class OpenAIServingChat(GenerateBaseServing):
1039 usage.prompt_tokens_details = _make_prompt_tokens_details(
1040 self.enable_prompt_tokens_details,
1041 final_res.num_cached_tokens,
1042 + final_res.num_cache_creation_tokens,
1043 mm_token_counts,
1044 )
1045
vllm/entrypoints/openai/engine/protocol.py
+1
@@ -101,6 +101,7 @@ class ModelList(OpenAIBaseModel):
101
102 class PromptTokenUsageInfo(OpenAIBaseModel):
103 cached_tokens: int | None = None
104 + created_cache_tokens: int | None = None
105 multimodal_tokens: dict[str, int] | None = None
106 """Prompt tokens contributed by each input modality, keyed by modality name
107 (e.g. `image`, `audio`, `video`). A breakdown of the multimodal
vllm/outputs.py
+6 -1
@@ -103,6 +103,8 @@ class RequestOutput:
103 encoder_prompt_token_ids: The token IDs of the encoder prompt.
104 None if decoder-only.
105 num_cached_tokens: The number of tokens with prefix cache hit.
106 + num_cache_creation_tokens: Prompt tokens currently counted as local
107 + prefix-cache writes for this request.
108 kv_transfer_params: The params for remote K/V transfer.
109 ec_transfer_params: The params for remote encoder-cache transfer.
110 """
@@ -120,6 +122,7 @@ class RequestOutput:
122 encoder_prompt: str | None = None,
123 encoder_prompt_token_ids: list[int] | None = None,
124 num_cached_tokens: int | None = None,
125 + num_cache_creation_tokens: int | None = None,
126 *,
127 kv_transfer_params: dict[str, Any] | None = None,
128 ec_transfer_params: dict[str, Any] | None = None,
@@ -142,6 +145,7 @@ class RequestOutput:
145 self.encoder_prompt = encoder_prompt
146 self.encoder_prompt_token_ids = encoder_prompt_token_ids
147 self.num_cached_tokens = num_cached_tokens
148 + self.num_cache_creation_tokens = num_cache_creation_tokens
149 self.kv_transfer_params = kv_transfer_params
150 self.ec_transfer_params = ec_transfer_params
151
@@ -188,7 +192,8 @@ class RequestOutput:
192 f"finished={self.finished}, "
193 f"metrics={self.metrics}, "
194 f"lora_request={self.lora_request}, "
191 - f"num_cached_tokens={self.num_cached_tokens})"
195 + f"num_cached_tokens={self.num_cached_tokens}, "
196 + f"num_cache_creation_tokens={self.num_cache_creation_tokens})"
197 )
198
199
vllm/v1/core/kv_cache_manager.py
+29
@@ -657,6 +657,35 @@ class KVCacheManager:
657 clipped_block_ids.append(ids[:num_valid_blocks])
658 return tuple(clipped_block_ids)
659
660 + def estimate_cached_tokens(self, request: Request) -> int:
661 + """Estimate the number of tokens cached by the request."""
662 + cached_tokens: int | None = None
663 + for group, blocks in zip(
664 + self.kv_cache_config.kv_cache_groups,
665 + self.get_blocks(request.request_id).blocks,
666 + ):
667 + if isinstance(
668 + group.kv_cache_spec,
669 + (CrossAttentionSpec, EncoderOnlyAttentionSpec),
670 + ):
671 + # Cross-attention and encoder-only groups are not prefix cached.
672 + continue
673 +
674 + group_cached_tokens = 0
675 + for block in blocks:
676 + group_cached_tokens = max(
677 + group_cached_tokens,
678 + block.block_hash_num_tokens or 0,
679 + )
680 +
681 + cached_tokens = (
682 + group_cached_tokens
683 + if cached_tokens is None
684 + else min(cached_tokens, group_cached_tokens)
685 + )
686 +
687 + return cached_tokens or 0
688 +
689 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
690 """Cache the blocks for the request, if enabled.
691
vllm/v1/core/sched/scheduler.py
+17 -9
@@ -799,7 +799,10 @@ class Scheduler(SchedulerInterface):
799 continue
800
801 # Track first scheduled prefill, not post-preemption repeat prefills
802 - if request.prefill_stats is not None:
802 + if (
803 + request.prefill_stats is not None
804 + and request.num_preemptions <= 0
805 + ):
806 assert num_computed_tokens <= request.num_prompt_tokens
807 request.prefill_stats.set(
808 num_prompt_tokens=request.num_prompt_tokens,
@@ -1708,6 +1711,7 @@ class Scheduler(SchedulerInterface):
1711 pooler_output = pooler_outputs[req_index] if pooler_outputs else None
1712 kv_transfer_params = None
1713 ec_transfer_params = None
1714 + prefill_stats = None
1715 status_before_stop = request.status
1716 num_output_tokens_before = len(request._output_token_ids)
1717
@@ -1787,6 +1791,16 @@ class Scheduler(SchedulerInterface):
1791 # Normal decode / re-prefill: token(s) at the END.
1792 routed_experts = routing_data[end - len(new_token_ids) : end]
1793
1794 + should_emit_output = bool(
1795 + new_token_ids or pooler_output is not None or stopped
1796 + )
1797 + if should_emit_output:
1798 + prefill_stats = request.take_prefill_stats()
1799 + if prefill_stats is not None:
1800 + prefill_stats.finalize(
1801 + self.kv_cache_manager.estimate_cached_tokens(request)
1802 + )
1803 +
1804 finish_reason = None
1805 if stopped:
1806 # Capture finish_reason BEFORE _handle_stopped_request, which may
@@ -1814,13 +1828,7 @@ class Scheduler(SchedulerInterface):
1828
1829 # Get prompt logprobs for this request.
1830 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
1817 - if (
1818 - new_token_ids
1819 - or pooler_output is not None
1820 - or kv_transfer_params
1821 - or ec_transfer_params
1822 - or stopped
1823 - ):
1831 + if should_emit_output:
1832 # Add EngineCoreOutput for this Request.
1833 outputs[request.client_index].append(
1834 EngineCoreOutput(
@@ -1832,7 +1840,7 @@ class Scheduler(SchedulerInterface):
1840 pooling_output=pooler_output,
1841 stop_reason=request.stop_reason,
1842 events=request.take_events(),
1835 - prefill_stats=request.take_prefill_stats(),
1843 + prefill_stats=prefill_stats,
1844 kv_transfer_params=kv_transfer_params,
1845 ec_transfer_params=ec_transfer_params,
1846 trace_headers=request.trace_headers,
vllm/v1/engine/output_processor.py
+5
@@ -172,6 +172,7 @@ class RequestState:
172 self.is_prefilling = True
173 self.queue = queue
174 self.num_cached_tokens = 0
175 + self.num_cache_creation_tokens = 0
176
177 self.stats = RequestStateStats(arrival_time=arrival_time) if log_stats else None
178
@@ -377,6 +378,7 @@ class RequestState:
378 kv_transfer_params=kv_transfer_params,
379 ec_transfer_params=ec_transfer_params,
380 num_cached_tokens=self.num_cached_tokens,
381 + num_cache_creation_tokens=self.num_cache_creation_tokens,
382 metrics=self.stats,
383 )
384
@@ -639,6 +641,9 @@ class OutputProcessor:
641 req_state.num_cached_tokens = (
642 engine_core_output.prefill_stats.num_cached_tokens
643 )
644 + req_state.num_cache_creation_tokens = (
645 + engine_core_output.prefill_stats.num_cache_creation_tokens
646 + )
647 req_state.is_prefilling = False
648
649 if pooling_output is None:
vllm/v1/metrics/stats.py
+8
@@ -265,6 +265,7 @@ class PrefillStats:
265 num_cached_tokens: Tokens to be prefilled without actual compute work.
266 num_local_cached_tokens: Tokens to be prefilled from local prefix cache.
267 num_external_cached_tokens: Tokens to be prefilled from external KV transfer.
268 + num_cache_creation_tokens: Tokens computed and written to the prefix cache.
269 """
270
271 num_prompt_tokens: int = 0
@@ -272,6 +273,7 @@ class PrefillStats:
273 num_cached_tokens: int = 0
274 num_local_cached_tokens: int = 0
275 num_external_cached_tokens: int = 0
276 + num_cache_creation_tokens: int = 0
277
278 def set(
279 self,
@@ -288,6 +290,12 @@ class PrefillStats:
290 self.num_local_cached_tokens = num_local_cached_tokens
291 self.num_external_cached_tokens = num_external_cached_tokens
292
293 + def finalize(self, num_cached_tokens: int) -> None:
294 + assert num_cached_tokens >= 0
295 + self.num_cache_creation_tokens = max(
296 + 0, min(num_cached_tokens, self.num_prompt_tokens) - self.num_cached_tokens
297 + )
298 +
299
300 @dataclass
301 class PromptTokenStats: