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
[Feature][gpt-oss] Add support for num_cached_tokens and num_reasoning_tokens tracking (#23460)
Signed-off-by: George Nagy II <george.nagy0969@gmail.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com>
George Nagy II committed
Sep 3, 2025 at 15:08 UTC
36c260dad604ccc845150753f2530b5b2ba9d7e6
1 file changed
+21
vllm/entrypoints/context.py
+21
@@ -93,17 +93,35 @@ class HarmonyContext(ConversationContext):
93
# as new prompt each time. Hence the sum.
94
self.num_prompt_tokens += len(output.prompt_token_ids)
95
96
+ def _update_num_cached_tokens(self, output: RequestOutput):
97
+ if output.num_cached_tokens is not None:
98
+ #Similar to num_prompt_tokens
99
+ self.num_cached_tokens += output.num_cached_tokens
100
+
101
def _update_num_output_tokens(self, token_ids: Sequence[int]):
102
self.num_output_tokens += len(token_ids)
103
104
+ def _update_num_reasoning_tokens(self, token_ids: Sequence[int]):
105
+ # Count tokens that are part of reasoning content (analysis channel
106
+ # or tool-directed messages like python/browser calls)
107
+ is_analysis = self.parser.current_channel == "analysis"
108
+ is_tool_call = (self.parser.current_recipient is not None and
109
+ (self.parser.current_recipient.startswith("python") or
110
+ self.parser.current_recipient.startswith("browser.")))
111
+ if is_analysis or is_tool_call:
112
+ self.num_reasoning_tokens += len(token_ids)
113
+
114
def append_output(self, output) -> None:
115
if isinstance(output, RequestOutput):
116
self._update_num_prompt_tokens(output)
117
+ self._update_num_cached_tokens(output)
118
output_token_ids = output.outputs[0].token_ids
119
self._update_num_output_tokens(output_token_ids)
120
self.parser = get_streamable_parser_for_assistant()
121
for token_id in output_token_ids:
122
self.parser.process(token_id)
123
+ # Check if the current token is part of reasoning content
124
+ self._update_num_reasoning_tokens([token_id])
125
output_msgs = self.parser.messages
126
else:
127
# Tool output.
@@ -204,6 +222,7 @@ class StreamingHarmonyContext(HarmonyContext):
222
# so we only want to add the prompt tokens once for each message.
223
if self.first_tok_of_message:
224
self._update_num_prompt_tokens(output)
225
+ self._update_num_cached_tokens(output)
226
# Reset self.first_tok_of_message if needed:
227
# if the current token is the last one of the current message
228
# (finished=True), then the next token processed will mark the
@@ -212,6 +231,8 @@ class StreamingHarmonyContext(HarmonyContext):
231
tok = output.outputs[0].token_ids[0]
232
self.parser.process(tok)
233
self._update_num_output_tokens(output.outputs[0].token_ids)
234
+ # Check if the current token is part of reasoning content
235
+ self._update_num_reasoning_tokens([tok])
236
self.last_tok = tok
237
else:
238
# Handle the case of tool output in direct message format