@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0512c04aee

Fork 0

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 14, 05:30 UTC

[frontend][gptoss] Add per turn stats into Harmony Context (#25061)

Signed-off-by: lacora <hyelacora@gmail.com> Co-authored-by: Ye Hu <yehu@fb.com>

Ye Hu committed Oct 14, 2025 at 16:48 UTC 0512c04aee408367a068b5960e7857c722ed204d

4 files changed +188 -62

tests/entrypoints/test_context.py

+84 -9

@@ -6,7 +6,11 @@ from unittest.mock import MagicMock, patch

6 import pytest

7 from openai_harmony import Author, Message, Role, StreamState, TextContent

9 -from vllm.entrypoints.context import HarmonyContext, StreamingHarmonyContext

9 +from vllm.entrypoints.context import (

10 + HarmonyContext,

11 + StreamingHarmonyContext,

12 + TurnMetrics,

13 +)

14 from vllm.outputs import CompletionOutput, RequestOutput

@@ -101,8 +105,12 @@ def test_single_turn_token_counting():

105

106 # Verify internal state tracking

107 assert not context.is_first_turn

104 - assert context.previous_turn.input_tokens == 5

105 - assert context.previous_turn.output_tokens == 3

108 + assert len(context.all_turn_metrics) == 1

109 + previous_turn = context.all_turn_metrics[0]

110 + assert previous_turn.input_tokens == 5

111 + assert previous_turn.output_tokens == 3

112 + assert previous_turn.cached_input_tokens == 2

113 + assert previous_turn.tool_output_tokens == 0

114

115

116 @pytest.mark.asyncio

@@ -156,6 +164,15 @@ async def test_multi_turn_token_counting():

164 assert context.num_tool_output_tokens == expected_tool_output

165 assert context.num_cached_tokens == 5 + 15

166

167 + # Validate all turn metrics

168 + assert len(context.all_turn_metrics) == 3

169 + for i, turn in enumerate(context.all_turn_metrics):

170 + assert turn.input_tokens == prompt_token_counts[i]

171 + assert turn.output_tokens == output_token_counts[i]

172 + assert turn.cached_input_tokens == cached_token_counts[i]

173 + assert context.all_turn_metrics[1].tool_output_tokens == 7

174 + assert context.all_turn_metrics[2].tool_output_tokens == 1

175 +

176

177 def test_empty_output_tokens():

178 """Test behavior when RequestOutput has empty output tokens."""

@@ -314,6 +331,10 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

331 # Create a streaming context

332 context = StreamingHarmonyContext(messages=[], available_tools=["browser"])

333

334 + num_prompt_tokens = [3, 8, 13]

335 + num_output_tokens = [3, 3, 2]

336 + num_cached_tokens = [0, 3, 8]

337 +

338 # Simulate three turns of conversation:

339 # Turn 1: stream tokens one by one, then finish the message

340 # Turn 2: new prompt, stream more tokens with a reasoning segment

@@ -325,7 +346,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

346 create_mock_request_output(

347 prompt_token_ids=[1, 2, 3], # 3 prompt tokens

348 output_token_ids=[101], # Single token

328 - num_cached_tokens=0,

349 + num_cached_tokens=num_cached_tokens[0],

350 finished=False, # Not end of message yet

351 )

352 )

@@ -370,7 +391,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

391 5,

392 ], # 8 tokens (includes previous)

393 output_token_ids=[201],

373 - num_cached_tokens=3, # Some tokens cached

394 + num_cached_tokens=num_cached_tokens[1], # Some tokens cached

395 finished=False,

396 )

397 )

@@ -422,7 +443,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

443 7,

444 ], # 13 tokens

445 output_token_ids=[301],

425 - num_cached_tokens=8, # More cached tokens

446 + num_cached_tokens=num_cached_tokens[2], # More cached tokens

447 finished=False,

448 )

449 )

@@ -435,10 +456,12 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

456 )

457

458 # Final token counts check

438 - assert context.num_prompt_tokens == 3 + 8 + 13 # All prompts

439 - assert context.num_output_tokens == 3 + 3 + 2 # All outputs

459 + assert context.num_prompt_tokens == sum(num_prompt_tokens) # All prompts

460 + assert context.num_output_tokens == sum(num_output_tokens) # All outputs

461 assert context.num_reasoning_tokens == 3 # Unchanged from second turn

441 - assert context.num_cached_tokens == 3 + 8 # Accumulated cached tokens

462 + assert context.num_cached_tokens == sum(

463 + num_cached_tokens

464 + ) # Accumulated cached tokens

465

466 # Additional tool tokens from third turn

467 # Formula: this turn prompt - last turn prompt - last turn output

@@ -447,6 +470,15 @@ async def test_streaming_multi_turn_token_counting(mock_parser):

470 context.num_tool_output_tokens == expected_tool_tokens + additional_tool_tokens

471 )

472

473 + # Validate all turn metrics

474 + assert len(context.all_turn_metrics) == 3

475 + for i, turn in enumerate(context.all_turn_metrics):

476 + assert turn.input_tokens == num_prompt_tokens[i]

477 + assert turn.output_tokens == num_output_tokens[i]

478 + assert turn.cached_input_tokens == num_cached_tokens[i]

479 + assert context.all_turn_metrics[1].tool_output_tokens == 2

480 + assert context.all_turn_metrics[2].tool_output_tokens == 2

481 +

482

483 @pytest.mark.asyncio

484 async def test_streaming_message_synchronization(mock_parser):

@@ -522,3 +554,46 @@ async def test_streaming_message_synchronization(mock_parser):

554 assert len(context._messages) == 3

555 assert context.num_init_messages == 1

556 assert context._messages[2].content[0].text == "Response 4"

557 +

558 +

559 +def test_turn_metrics_copy_and_reset():

560 + """Test TurnMetrics copy and reset methods work correctly."""

561 + # Create a TurnMetrics with specific values

562 + original_metrics = TurnMetrics(

563 + input_tokens=10,

564 + output_tokens=20,

565 + cached_input_tokens=5,

566 + tool_output_tokens=3,

567 + )

568 +

569 + # Test copy functionality

570 + copied_metrics = original_metrics.copy()

571 +

572 + # Verify copy has same values

573 + assert copied_metrics.input_tokens == 10

574 + assert copied_metrics.output_tokens == 20

575 + assert copied_metrics.cached_input_tokens == 5

576 + assert copied_metrics.tool_output_tokens == 3

577 +

578 + # Verify they are separate objects

579 + assert copied_metrics is not original_metrics

580 +

581 + # Modify copy to ensure independence

582 + copied_metrics.input_tokens = 999

583 + assert original_metrics.input_tokens == 10 # Original unchanged

584 + assert copied_metrics.input_tokens == 999

585 +

586 + # Test reset functionality

587 + original_metrics.reset()

588 +

589 + # Verify all fields are reset to zero

590 + assert original_metrics.input_tokens == 0

591 + assert original_metrics.output_tokens == 0

592 + assert original_metrics.cached_input_tokens == 0

593 + assert original_metrics.tool_output_tokens == 0

594 +

595 + # Verify copied metrics are unaffected by reset

596 + assert copied_metrics.input_tokens == 999

597 + assert copied_metrics.output_tokens == 20

598 + assert copied_metrics.cached_input_tokens == 5

599 + assert copied_metrics.tool_output_tokens == 3

vllm/entrypoints/context.py

+43 -22

@@ -45,21 +45,36 @@ def _map_tool_name_to_tool_type(tool_name: str) -> str:

45 return _TOOL_NAME_TO_TYPE_MAP[tool_name]

48 -class TurnTokens:

49 - """Tracks token counts for a single conversation turn."""

48 +class TurnMetrics:

49 + """Tracks token and toolcall details for a single conversation turn."""

51 - def __init__(self, input_tokens=0, output_tokens=0):

51 + def __init__(

52 + self,

53 + input_tokens=0,

54 + output_tokens=0,

55 + cached_input_tokens=0,

56 + tool_output_tokens=0,

57 + ):

58 self.input_tokens = input_tokens

59 self.output_tokens = output_tokens

60 + self.cached_input_tokens = cached_input_tokens

61 + self.tool_output_tokens = tool_output_tokens

63 def reset(self):

64 """Reset counters for a new turn."""

65 self.input_tokens = 0

66 self.output_tokens = 0

67 + self.cached_input_tokens = 0

68 + self.tool_output_tokens = 0

70 def copy(self):

71 """Create a copy of this turn's token counts."""

62 - return TurnTokens(self.input_tokens, self.output_tokens)

72 + return TurnMetrics(

73 + self.input_tokens,

74 + self.output_tokens,

75 + self.cached_input_tokens,

76 + self.tool_output_tokens,

77 + )

80 class ConversationContext(ABC):

@@ -102,6 +117,8 @@ class SimpleContext(ConversationContext):

117 self.num_cached_tokens = 0

118 # todo num_reasoning_tokens is not implemented yet.

119 self.num_reasoning_tokens = 0

120 + # not implemented yet for SimpleContext

121 + self.all_turn_metrics = []

122

123 def append_output(self, output) -> None:

124 self.last_output = output

@@ -154,8 +171,9 @@ class HarmonyContext(ConversationContext):

171 self.num_tool_output_tokens = 0

172

173 # Turn tracking - replaces multiple individual tracking variables

157 - self.current_turn = TurnTokens()

158 - self.previous_turn = TurnTokens()

174 + self.current_turn_metrics = TurnMetrics()

175 + # Track metrics for all turns

176 + self.all_turn_metrics: list[TurnMetrics] = []

177 self.is_first_turn = True

178 self.first_tok_of_message = True # For streaming support

179

@@ -173,11 +191,10 @@ class HarmonyContext(ConversationContext):

191 # Check if the current token is part of reasoning content

192 self._update_num_reasoning_tokens()

193 self._update_prefill_token_usage(output)

176 - # Reset current turn output tokens for this turn

177 - self.current_turn.output_tokens = 0

194 self._update_decode_token_usage(output)

179 - # Move current turn to previous turn for next turn's calculations

180 - self.previous_turn = self.current_turn.copy()

195 + # Append current turn to all turn list for next turn's calculations

196 + self.all_turn_metrics.append(self.current_turn_metrics.copy())

197 + self.current_turn_metrics.reset()

198 # append_output is called only once before tool calling

199 # in non-streaming case

200 # so we can append all the parser messages to _messages

@@ -213,20 +230,21 @@ class HarmonyContext(ConversationContext):

230 logger.error("RequestOutput appended contains no prompt_token_ids.")

231

232 # Update current turn input tokens

216 - self.current_turn.input_tokens = this_turn_input_tokens

233 + self.current_turn_metrics.input_tokens = this_turn_input_tokens

234 self.num_prompt_tokens += this_turn_input_tokens

235

236 # Calculate tool tokens (except on first turn)

237 if self.is_first_turn:

238 self.is_first_turn = False

239 else:

240 + previous_turn = self.all_turn_metrics[-1]

241 # start counting tool after first turn

242 # tool tokens = this turn prefill - last turn prefill -

243 # last turn decode

244 this_turn_tool_tokens = (

227 - self.current_turn.input_tokens

228 - - self.previous_turn.input_tokens

229 - - self.previous_turn.output_tokens

245 + self.current_turn_metrics.input_tokens

246 + - previous_turn.input_tokens

247 + - previous_turn.output_tokens

248 )

249

250 # Handle negative tool token counts (shouldn't happen in normal

@@ -237,17 +255,20 @@ class HarmonyContext(ConversationContext):

255 "(current_input=%d, previous_input=%d, "

256 "previous_output=%d). Setting to 0.",

257 this_turn_tool_tokens,

240 - self.current_turn.input_tokens,

241 - self.previous_turn.input_tokens,

242 - self.previous_turn.output_tokens,

258 + self.current_turn_metrics.input_tokens,

259 + previous_turn.input_tokens,

260 + previous_turn.output_tokens,

261 )

262 this_turn_tool_tokens = 0

263

264 self.num_tool_output_tokens += this_turn_tool_tokens

265 + self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens

266

267 # Update cached tokens

249 - if output.num_cached_tokens is not None:

250 - self.num_cached_tokens += output.num_cached_tokens

268 + num_cached_token = output.num_cached_tokens

269 + if num_cached_token is not None:

270 + self.num_cached_tokens += num_cached_token

271 + self.current_turn_metrics.cached_input_tokens = num_cached_token

272

273 def _update_decode_token_usage(self, output: RequestOutput) -> int:

274 """Update token usage statistics for the decode phase of generation.

@@ -272,7 +293,7 @@ class HarmonyContext(ConversationContext):

293 # only keep last round

294 updated_output_token_count += len(completion_output.token_ids)

295 self.num_output_tokens += updated_output_token_count

275 - self.current_turn.output_tokens += updated_output_token_count

296 + self.current_turn_metrics.output_tokens += updated_output_token_count

297 return updated_output_token_count

298

299 @property

@@ -452,7 +473,6 @@ class StreamingHarmonyContext(HarmonyContext):

473 # so we only want to add the prompt tokens once for each message.

474 if self.first_tok_of_message:

475 self._update_prefill_token_usage(output)

455 - self.current_turn.output_tokens = 0

476 # Reset self.first_tok_of_message if needed:

477 # if the current token is the last one of the current message

478 # (finished=True), then the next token processed will mark the

@@ -464,7 +484,8 @@ class StreamingHarmonyContext(HarmonyContext):

484

485 # For streaming, update previous turn when message is complete

486 if output.finished:

467 - self.previous_turn = self.current_turn.copy()

487 + self.all_turn_metrics.append(self.current_turn_metrics.copy())

488 + self.current_turn_metrics.reset()

489 # Check if the current token is part of reasoning content

490 self._update_num_reasoning_tokens()

491 self.last_tok = tok

vllm/entrypoints/openai/protocol.py

@@ -2103,11 +2103,15 @@ class TranscriptionStreamResponse(OpenAIBaseModel):

2103

2104 class InputTokensDetails(OpenAIBaseModel):

2105 cached_tokens: int

2106 + input_tokens_per_turn: list[int] = Field(default_factory=list)

2107 + cached_tokens_per_turn: list[int] = Field(default_factory=list)

2108

2109

2110 class OutputTokensDetails(OpenAIBaseModel):

2111 reasoning_tokens: int = 0

2112 tool_output_tokens: int = 0

2113 + output_tokens_per_turn: list[int] = Field(default_factory=list)

2114 + tool_output_tokens_per_turn: list[int] = Field(default_factory=list)

2115

2116

2117 class ResponseUsage(OpenAIBaseModel):

vllm/entrypoints/openai/serving_responses.py

+57 -31

@@ -589,10 +589,24 @@ class OpenAIServingResponses(OpenAIServing):

589 input_tokens=num_prompt_tokens,

590 output_tokens=num_generated_tokens,

591 total_tokens=num_prompt_tokens + num_generated_tokens,

592 - input_tokens_details=InputTokensDetails(cached_tokens=num_cached_tokens),

592 + input_tokens_details=InputTokensDetails(

593 + cached_tokens=num_cached_tokens,

594 + input_tokens_per_turn=[

595 + turn.input_tokens for turn in context.all_turn_metrics

596 + ],

597 + cached_tokens_per_turn=[

598 + turn.cached_input_tokens for turn in context.all_turn_metrics

599 + ],

600 + ),

601 output_tokens_details=OutputTokensDetails(

602 reasoning_tokens=num_reasoning_tokens,

603 tool_output_tokens=num_tool_output_tokens,

604 + output_tokens_per_turn=[

605 + turn.output_tokens for turn in context.all_turn_metrics

606 + ],

607 + tool_output_tokens_per_turn=[

608 + turn.tool_output_tokens for turn in context.all_turn_metrics

609 + ],

610 ),

611 )

612 response = ResponsesResponse.from_request(

@@ -665,11 +679,13 @@ class OpenAIServingResponses(OpenAIServing):

679 token=text,

680 logprob=max(token_logprob.logprob, -9999.0),

681 bytes=list(text.encode("utf-8", errors="replace")),

668 - top_logprobs=self._topk_logprobs(

669 - logprob, top_logprobs=top_logprobs, tokenizer=tokenizer

670 - )

671 - if top_logprobs

672 - else [],

682 + top_logprobs=(

683 + self._topk_logprobs(

684 + logprob, top_logprobs=top_logprobs, tokenizer=tokenizer

685 + )

686 + if top_logprobs

687 + else []

688 + ),

689 )

690 )

691 return out

@@ -758,14 +774,16 @@ class OpenAIServingResponses(OpenAIServing):

774 text=content,

775 annotations=[], # TODO

776 type="output_text",

761 - logprobs=self._create_response_logprobs(

762 - token_ids=final_output.token_ids,

763 - logprobs=final_output.logprobs,

764 - tokenizer=tokenizer,

765 - top_logprobs=request.top_logprobs,

766 - )

767 - if request.is_include_output_logprobs()

768 - else None,

777 + logprobs=(

778 + self._create_response_logprobs(

779 + token_ids=final_output.token_ids,

780 + logprobs=final_output.logprobs,

781 + tokenizer=tokenizer,

782 + top_logprobs=request.top_logprobs,

783 + )

784 + if request.is_include_output_logprobs()

785 + else None

786 + ),

787 )

788 message = ResponseOutputMessage(

789 id=f"msg_{random_uuid()}",

@@ -870,15 +888,21 @@ class OpenAIServingResponses(OpenAIServing):

888 with_custom_tools = has_custom_tools(tool_types)

889 sys_msg = get_system_message(

890 reasoning_effort=reasoning_effort,

873 - browser_description=self.tool_server.get_tool_description("browser")

874 - if enable_browser and self.tool_server is not None

875 - else None,

876 - python_description=self.tool_server.get_tool_description("python")

877 - if enable_code_interpreter and self.tool_server is not None

878 - else None,

879 - container_description=self.tool_server.get_tool_description("container")

880 - if enable_container and self.tool_server is not None

881 - else None,

891 + browser_description=(

892 + self.tool_server.get_tool_description("browser")

893 + if enable_browser and self.tool_server is not None

894 + else None

895 + ),

896 + python_description=(

897 + self.tool_server.get_tool_description("python")

898 + if enable_code_interpreter and self.tool_server is not None

899 + else None

900 + ),

901 + container_description=(

902 + self.tool_server.get_tool_description("container")

903 + if enable_container and self.tool_server is not None

904 + else None

905 + ),

906 instructions=request.instructions,

907 with_custom_tools=with_custom_tools,

908 )

@@ -1283,14 +1307,16 @@ class OpenAIServingResponses(OpenAIServing):

1307 output_index=current_output_index,

1308 item_id=current_item_id,

1309 delta=delta_message.content,

1286 - logprobs=self._create_stream_response_logprobs(

1287 - token_ids=output.token_ids,

1288 - logprobs=output.logprobs,

1289 - tokenizer=tokenizer,

1290 - top_logprobs=request.top_logprobs,

1291 - )

1292 - if request.is_include_output_logprobs()

1293 - else [],

1310 + logprobs=(

1311 + self._create_stream_response_logprobs(

1312 + token_ids=output.token_ids,

1313 + logprobs=output.logprobs,

1314 + tokenizer=tokenizer,

1315 + top_logprobs=request.top_logprobs,

1316 + )

1317 + if request.is_include_output_logprobs()

1318 + else []

1319 + ),

1320 )

1321 )

1322 current_content_index += 1