diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py index 710cd2d04c..f3927c9d71 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py +++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py @@ -169,6 +169,7 @@ def _make_request_output( text: str = "Hello, world!", token_ids: Sequence[int] = (1, 2, 3), finish_reason: str = "stop", + num_cache_creation_tokens: int | None = None, ) -> RequestOutput: return RequestOutput( request_id="test", @@ -186,6 +187,7 @@ def _make_request_output( ) ], finished=True, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -425,6 +427,14 @@ def test_reasoning_tokens_counted_per_round(): assert ctx.num_reasoning_tokens == 3 +def test_cache_creation_tokens_counted_per_round(): + ctx = _make_context(None) + + for count, expected in [(None, 0), (0, 0), (3, 3), (4, 7)]: + ctx.append_output(_make_request_output(num_cache_creation_tokens=count)) + assert ctx.num_cache_creation_tokens == expected + + def test_reasoning_tokens_zero_without_parser(): ctx = _make_context(None) diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index a5b1afe8c5..f0af86231a 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -20,6 +20,7 @@ def create_mock_request_output( output_token_ids=None, num_cached_tokens=0, finished=True, + num_cache_creation_tokens=None, ): """Helper function to create a mock RequestOutput object for testing.""" outputs = [] @@ -44,6 +45,7 @@ def create_mock_request_output( outputs=outputs, finished=finished, num_cached_tokens=num_cached_tokens, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -117,7 +119,8 @@ def make_harmony_context( return context, fake_parser -def test_single_turn_token_counting(): +@pytest.mark.parametrize("num_cache_creation_tokens", [None, 0, 3]) +def test_single_turn_token_counting(num_cache_creation_tokens): """Test token counting behavior for a single turn.""" # Create a context context, _ = make_harmony_context() @@ -127,6 +130,7 @@ def test_single_turn_token_counting(): prompt_token_ids=[1, 2, 3, 4, 5], # 5 prompt tokens output_token_ids=[6, 7, 8], # 3 output tokens num_cached_tokens=2, # 2 cached tokens + num_cache_creation_tokens=num_cache_creation_tokens, ) # Append the output to the context @@ -136,6 +140,7 @@ def test_single_turn_token_counting(): assert context.num_prompt_tokens == 5 assert context.num_output_tokens == 3 assert context.num_cached_tokens == 2 + assert context.num_cache_creation_tokens == (num_cache_creation_tokens or 0) assert context.num_tool_output_tokens == 0 # No tool tokens in first turn # Verify internal state tracking @@ -398,6 +403,7 @@ async def test_streaming_multi_turn_token_counting(): num_prompt_tokens = [3, 8, 13] num_output_tokens = [3, 3, 2] num_cached_tokens = [0, 3, 8] + num_cache_creation_tokens = [3, 4, 2] # Simulate three turns of conversation: # Turn 1: stream tokens one by one, then finish the message @@ -411,6 +417,7 @@ async def test_streaming_multi_turn_token_counting(): prompt_token_ids=[1, 2, 3], # 3 prompt tokens output_token_ids=[101], # Single token num_cached_tokens=num_cached_tokens[0], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, # Not end of message yet ) ) @@ -419,6 +426,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[102], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, ) ) @@ -427,6 +435,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[103], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=True, # End of message ) ) @@ -435,6 +444,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_prompt_tokens == 3 # Initial prompt tokens assert context.num_output_tokens == 3 # Three output tokens assert context.num_cached_tokens == 0 + assert context.num_cache_creation_tokens == 3 assert context.num_tool_output_tokens == 0 # No tool output in first turn assert context.first_tok_of_message is True # Ready for next message @@ -454,6 +464,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 8 tokens (includes previous) output_token_ids=[201], num_cached_tokens=num_cached_tokens[1], # Some tokens cached + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -463,6 +474,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[202], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -471,6 +483,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[203], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=True, # End of reasoning message ) ) @@ -480,6 +493,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_output_tokens == 3 + 3 # First turn + second turn assert context.num_reasoning_tokens == 3 # All tokens in analysis channel assert context.num_cached_tokens == 3 # Cached tokens from second turn + assert context.num_cache_creation_tokens == 3 + 4 # Formula: this turn prompt tokens - last turn prompt - last turn output expected_tool_tokens = 8 - 3 - 3 # = 2 @@ -505,6 +519,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 13 tokens output_token_ids=[301], num_cached_tokens=num_cached_tokens[2], # More cached tokens + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=False, ) ) @@ -512,6 +527,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[302], + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=True, ) ) @@ -523,6 +539,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_cached_tokens == sum( num_cached_tokens ) # Accumulated cached tokens + assert context.num_cache_creation_tokens == sum(num_cache_creation_tokens) # Additional tool tokens from third turn # Formula: this turn prompt - last turn prompt - last turn output diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 334e8032e2..9f5a281163 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -315,6 +315,7 @@ class ParsableContext(ConversationContext): self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 # not implemented yet for ParsableContext self.all_turn_metrics: list[TurnMetrics] = [] @@ -347,6 +348,7 @@ class ParsableContext(ConversationContext): def append_output(self, output: RequestOutput) -> None: self.num_prompt_tokens = len(output.prompt_token_ids or []) self.num_cached_tokens = output.num_cached_tokens or 0 + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params @@ -651,6 +653,7 @@ class HarmonyContext(ConversationContext): self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 self.num_tool_output_tokens = 0 @@ -773,6 +776,8 @@ class HarmonyContext(ConversationContext): self.num_cached_tokens += num_cached_token self.current_turn_metrics.cached_input_tokens = num_cached_token + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 + def _update_decode_token_usage(self, output: RequestOutput) -> int: """Update token usage statistics for the decode phase of generation. diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index ba4aa04347..78a3f1015f 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -881,7 +881,7 @@ class OpenAIServingResponses(GenerateBaseServing): output_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, input_tokens_details=InputTokensDetails( - cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0), + cache_write_tokens=context.num_cache_creation_tokens, cached_tokens=num_cached_tokens, input_tokens_per_turn=[ turn.input_tokens for turn in context.all_turn_metrics