minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] Add Reasoning Tokens Usage (#15562)

Browse Source 
Signed-off-by: Muqi Li <muqi1029@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Mufeez Amjad <mufeez.amjad@outlook.com>
Co-authored-by: cklxx <1293822641@qq.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
Muqi Li
2026-04-04 02:18:10 -07:00
committed by GitHub
co-authored by Xinyuan Tong Mufeez Amjad cklxx hnyls2002 Liangsheng Yin
parent bf984ae65d
commit 1ad6839659
12 changed files with 283 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/openai/serving_chat.py
+12 -1
View File
@@ -649,6 +649,7 @@ class OpenAIServingChat(OpenAIServingBase):
# Usage tracking
prompt_tokens = {}
reasoning_tokens = {}
completion_tokens = {}
cached_tokens = {}
hidden_states = {}
@@ -670,6 +671,9 @@ class OpenAIServingChat(OpenAIServingBase):
completion_tokens[index] = content["meta_info"].get(
"completion_tokens", 0
)
reasoning_tokens[index] = content["meta_info"].get(
"reasoning_tokens", 0
)
cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None)
@@ -752,6 +756,7 @@ class OpenAIServingChat(OpenAIServingBase):
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
reasoning_tokens=reasoning_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
)
@@ -805,6 +810,7 @@ class OpenAIServingChat(OpenAIServingBase):
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
reasoning_tokens=reasoning_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
)
@@ -885,8 +891,9 @@ class OpenAIServingChat(OpenAIServingBase):
if include_usage:
usage = UsageProcessor.calculate_streaming_usage(
prompt_tokens,
reasoning_tokens,
completion_tokens,
cached_tokens,
cached_tokens=cached_tokens,
n_choices=request.n,
enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report,
)
@@ -1356,9 +1363,11 @@ class OpenAIServingChat(OpenAIServingBase):
if continuous_usage_stats:
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
completion_tokens = content["meta_info"].get("completion_tokens", 0)
reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
reasoning_tokens=reasoning_tokens,
)
yield f"data: {chunk.model_dump_json()}\n\n"
@@ -1406,9 +1415,11 @@ class OpenAIServingChat(OpenAIServingBase):
if continuous_usage_stats:
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
completion_tokens = content["meta_info"].get("completion_tokens", 0)
reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
reasoning_tokens=reasoning_tokens,
)
yield f"data: {chunk.model_dump_json()}\n\n"
python/sglang/srt/entrypoints/openai/serving_completions.py
+7 -1
View File
@@ -220,6 +220,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
# Usage tracking
prompt_tokens = {}
completion_tokens = {}
reasoning_tokens = {}
cached_tokens = {}
hidden_states = {}
routed_experts = {}
@@ -241,6 +242,9 @@ class OpenAIServingCompletion(OpenAIServingBase):
completion_tokens[index] = content["meta_info"].get(
"completion_tokens", 0
)
reasoning_tokens[index] = content["meta_info"].get(
"reasoning_tokens", 0
)
cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None)
@@ -328,6 +332,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
reasoning_tokens=reasoning_tokens.get(index, 0),
)
yield f"data: {chunk.model_dump_json()}\n\n"
@@ -377,8 +382,9 @@ class OpenAIServingCompletion(OpenAIServingBase):
if include_usage:
usage = UsageProcessor.calculate_streaming_usage(
prompt_tokens,
reasoning_tokens,
completion_tokens,
cached_tokens,
cached_tokens=cached_tokens,
n_choices=request.n,
enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report,
)
python/sglang/srt/entrypoints/openai/usage_processor.py
+11 -1
View File
@@ -23,12 +23,16 @@ class UsageProcessor:
completion_tokens = sum(
r["meta_info"].get("completion_tokens", 0) for r in responses
)
prompt_tokens = sum(
responses[i]["meta_info"].get("prompt_tokens", 0)
for i in range(0, len(responses), n_choices)
)
# some API don't have reasoning_tokens semantics
reasoning_tokens = sum(
r["meta_info"].get("reasoning_tokens", 0) for r in responses
)
cached_details = None
if enable_cache_report:
cached_total = sum(
@@ -39,6 +43,7 @@ class UsageProcessor:
return UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens,
reasoning_tokens=reasoning_tokens,
completion_tokens=completion_tokens,
cached_tokens=cached_details,
)
@@ -46,6 +51,7 @@ class UsageProcessor:
@staticmethod
def calculate_streaming_usage(
prompt_tokens: Mapping[int, int],
reasoning_tokens: Mapping[int, int],
completion_tokens: Mapping[int, int],
cached_tokens: Mapping[int, int],
n_choices: int,
@@ -55,6 +61,7 @@ class UsageProcessor:
total_prompt_tokens = sum(
tok for idx, tok in prompt_tokens.items() if idx % n_choices == 0
)
total_reasoning_tokens = sum(reasoning_tokens.values())
total_completion_tokens = sum(completion_tokens.values())
cached_details = (
@@ -67,6 +74,7 @@ class UsageProcessor:
return UsageProcessor.calculate_token_usage(
prompt_tokens=total_prompt_tokens,
reasoning_tokens=total_reasoning_tokens,
completion_tokens=total_completion_tokens,
cached_tokens=cached_details,
)
@@ -75,6 +83,7 @@ class UsageProcessor:
def calculate_token_usage(
prompt_tokens: int,
completion_tokens: int,
reasoning_tokens: Optional[int] = 0,
cached_tokens: Optional[PromptTokensDetails] = None,
) -> UsageInfo:
"""Calculate token usage information"""
@@ -83,4 +92,5 @@ class UsageProcessor:
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=cached_tokens,
reasoning_tokens=reasoning_tokens,
)
python/sglang/srt/managers/detokenizer_manager.py
+1
View File
@@ -332,6 +332,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
output_strs=output_strs,
output_ids=recv_obj.output_ids,
prompt_tokens=recv_obj.prompt_tokens,
reasoning_tokens=recv_obj.reasoning_tokens,
completion_tokens=recv_obj.completion_tokens,
cached_tokens=recv_obj.cached_tokens,
cached_tokens_details=recv_obj.cached_tokens_details,
python/sglang/srt/managers/io_struct.py
+2
View File
@@ -974,6 +974,7 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
# Token counts
prompt_tokens: List[int]
reasoning_tokens: List[int]
completion_tokens: List[int]
cached_tokens: List[int]
@@ -1036,6 +1037,7 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
# Token counts
prompt_tokens: List[int]
completion_tokens: List[int]
reasoning_tokens: List[int]
cached_tokens: List[int]
# Logprobs
python/sglang/srt/managers/multi_tokenizer_mixin.py
+2
View File
@@ -146,6 +146,7 @@ def _handle_output_by_index(output, i):
no_stop_trim=_extract_field_by_index(output, "no_stop_trim", i),
prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
cached_tokens_details=_extract_field_by_index(
output, "cached_tokens_details", i
@@ -224,6 +225,7 @@ def _handle_output_by_index(output, i):
output_ids=_extract_field_by_index(output, "output_ids", i),
prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
input_token_logprobs_val=_extract_field_by_index(
output, "input_token_logprobs_val", i, check_length=False
python/sglang/srt/managers/schedule_batch.py
+19 -1
View File
@@ -634,9 +634,13 @@ class Req(ReqDllmMixin):
# For multi-http worker
self.http_worker_ipc = http_worker_ipc
# Require reasoning for the request (hybrid reasoning model only)
# Require reasoning for the request
self.require_reasoning = require_reasoning
# State indicating whether the reasoning phase has finished (only meaningful when require_reasoning is True)
self._is_reasoning_over = False
self.reasoning_tokens = 0
# Sampling info
if isinstance(sampling_params.custom_params, dict):
sampling_params = copy.copy(sampling_params)
@@ -1276,6 +1280,20 @@ class Req(ReqDllmMixin):
error_msg, HTTPStatus.BAD_REQUEST, "BadRequestError"
)
def update_reasoning_tokens(self, token_id, think_end_id):
if self._is_reasoning_over:
return
if not isinstance(token_id, list):
token_id = [token_id]
try:
end_pos = token_id.index(think_end_id)
self.reasoning_tokens += end_pos + 1
self._is_reasoning_over = True
except ValueError:
self.reasoning_tokens += len(token_id)
def __repr__(self):
return (
f"Req(rid={self.rid}, "
python/sglang/srt/managers/scheduler.py
+3
View File
@@ -552,6 +552,9 @@ class Scheduler(
self.tokenizer.think_end_id = self.tokenizer.encode(
reasoning_parser.detector.think_end_token, add_special_tokens=False
)[0]
self._think_end_id = self.tokenizer.think_end_id
else:
self._think_end_id = None
def init_mamba_backend(self) -> None:
initialize_mamba_selective_state_update_backend(self.server_args)
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+42 -11
View File
@@ -51,7 +51,7 @@ class SchedulerOutputProcessorMixin:
storage_backend_type = type(storage_backend).__name__
return storage_backend_type
def _get_cached_tokens_details(self, req: Req) -> Optional[dict]:
def _get_cached_tokens_details(self: Scheduler, req: Req) -> Optional[dict]:
"""Get detailed cache breakdown for a request, if available.
Returns:
@@ -182,8 +182,10 @@ class SchedulerOutputProcessorMixin:
# req output_ids are set here
req.output_ids.append(next_token_id)
req.check_finished()
self._maybe_update_reasoning_tokens(req, next_token_id)
req.check_finished()
if req.finished():
self.maybe_collect_routed_experts(req)
release_kv_cache(req, self.tree_cache)
@@ -404,6 +406,20 @@ class SchedulerOutputProcessorMixin:
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
else:
# for normal spec decoding: unify next_token_ids format
next_token_ids = []
cum_num_tokens = 0
next_token_ids_list = result.next_token_ids.tolist()
for i, req in enumerate(batch.reqs):
accept_length = result.accept_length_per_req_cpu[i]
next_token_ids.append(
next_token_ids_list[
cum_num_tokens : cum_num_tokens + accept_length + 1
]
)
cum_num_tokens += accept_length + 1
self.num_generated_tokens += len(batch.reqs)
if not batch.spec_algorithm.is_none():
@@ -436,6 +452,8 @@ class SchedulerOutputProcessorMixin:
req.output_ids.extend(next_token_id)
new_accepted_len = len(next_token_id)
self._maybe_update_reasoning_tokens(req, next_token_id)
# Update Mamba last track seqlen
self._mamba_prefix_cache_update(req, batch, result, i)
@@ -536,8 +554,18 @@ class SchedulerOutputProcessorMixin:
num_accepted_tokens=result.num_accepted_tokens,
)
def _maybe_update_reasoning_tokens(
self: Scheduler, req: Req, next_token_id: Union[int, List[int]]
):
if req.require_reasoning and self._think_end_id is not None:
req.update_reasoning_tokens(next_token_id, self._think_end_id)
def _mamba_prefix_cache_update(
self, req: Req, batch: ScheduleBatch, result: GenerationBatchResult, i: int
self: Scheduler,
req: Req,
batch: ScheduleBatch,
result: GenerationBatchResult,
i: int,
) -> None:
seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
if req.mamba_ping_pong_track_buffer is not None:
@@ -571,7 +599,7 @@ class SchedulerOutputProcessorMixin:
)
def _process_input_token_logprobs(
self, req: Req, input_token_logprobs: List
self: Scheduler, req: Req, input_token_logprobs: List
) -> None:
"""Process input token logprobs values and indices."""
is_multi_item_scoring = self._is_multi_item_scoring(req)
@@ -603,7 +631,7 @@ class SchedulerOutputProcessorMixin:
for x in input_token_logprobs_idx
]
def _process_input_top_logprobs(self, req: Req) -> None:
def _process_input_top_logprobs(self: Scheduler, req: Req) -> None:
"""Process input top logprobs."""
if req.top_logprobs_num <= 0:
return
@@ -632,7 +660,7 @@ class SchedulerOutputProcessorMixin:
req.temp_input_top_logprobs_idx = None
req.temp_input_top_logprobs_val = None
def _process_input_token_ids_logprobs(self, req: Req) -> None:
def _process_input_token_ids_logprobs(self: Scheduler, req: Req) -> None:
"""Process input token IDs logprobs."""
if req.token_ids_logprob is None:
return
@@ -664,7 +692,7 @@ class SchedulerOutputProcessorMixin:
req.temp_input_token_ids_logprobs_idx = None
req.temp_input_token_ids_logprobs_val = None
def _calculate_relevant_tokens_len(self, req: Req) -> int:
def _calculate_relevant_tokens_len(self: Scheduler, req: Req) -> int:
"""Calculate the expected length of logprob arrays based on whether multi-item scoring is enabled.
For multi-item scoring, only delimiter positions have logprobs.
@@ -685,7 +713,7 @@ class SchedulerOutputProcessorMixin:
return len(relevant_tokens)
def _calculate_num_input_logprobs(
self, req: Req, extend_input_len: int, extend_logprob_start_len: int
self: Scheduler, req: Req, extend_input_len: int, extend_logprob_start_len: int
) -> int:
"""Calculate the number of input logprobs based on whether multi-item scoring is enabled.
@@ -708,7 +736,7 @@ class SchedulerOutputProcessorMixin:
# Regular request: all tokens in the range
return extend_input_len - extend_logprob_start_len
def _is_multi_item_scoring(self, req: Req) -> bool:
def _is_multi_item_scoring(self: Scheduler, req: Req) -> bool:
"""Check if request uses multi-item scoring.
Multi-item scoring applies to prefill-only requests when a delimiter
@@ -845,7 +873,7 @@ class SchedulerOutputProcessorMixin:
return num_input_logprobs
def _initialize_empty_logprob_containers(self, req: Req) -> None:
def _initialize_empty_logprob_containers(self: Scheduler, req: Req) -> None:
"""
Initialize logprob fields to empty lists if unset.
@@ -882,7 +910,7 @@ class SchedulerOutputProcessorMixin:
envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get()
)
def _trigger_crash_for_tests(self, crash_threshold: int):
def _trigger_crash_for_tests(self: Scheduler, crash_threshold: int):
# Crash trigger: crash after stream_output is called N times
# This is used for testing purposes.
if not hasattr(self, "_test_stream_output_count"):
@@ -913,6 +941,7 @@ class SchedulerOutputProcessorMixin:
spaces_between_special_tokens = []
no_stop_trim = []
prompt_tokens = []
reasoning_tokens = []
completion_tokens = []
cached_tokens = []
cached_tokens_details = [] # Detailed breakdown by cache source
@@ -1013,6 +1042,7 @@ class SchedulerOutputProcessorMixin:
)
no_stop_trim.append(req.sampling_params.no_stop_trim)
prompt_tokens.append(len(req.origin_input_ids))
reasoning_tokens.append(req.reasoning_tokens)
completion_tokens.append(len(output_ids_))
cached_tokens.append(req.cached_tokens)
@@ -1142,6 +1172,7 @@ class SchedulerOutputProcessorMixin:
spaces_between_special_tokens=spaces_between_special_tokens,
no_stop_trim=no_stop_trim,
prompt_tokens=prompt_tokens,
reasoning_tokens=reasoning_tokens,
completion_tokens=completion_tokens,
cached_tokens=cached_tokens,
cached_tokens_details=cached_tokens_details,
python/sglang/srt/managers/tokenizer_manager.py
+1
View File
@@ -1574,6 +1574,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin):
if not isinstance(recv_obj, BatchEmbeddingOutput):
meta_info.update(
{
"reasoning_tokens": recv_obj.reasoning_tokens[i],
"completion_tokens": recv_obj.completion_tokens[i],
"cached_tokens": recv_obj.cached_tokens[i],
}
python/sglang/srt/speculative/ngram_worker.py
+3
View File
@@ -215,6 +215,7 @@ class NGRAMWorker:
spec_info = model_worker_batch.spec_info
num_accepted_tokens = 0
accept_lens = None
accept_length_per_req_cpu = None
if model_worker_batch.forward_mode.is_target_verify():
if batch.has_grammar:
@@ -256,6 +257,7 @@ class NGRAMWorker:
logits_output, next_token_ids, num_accepted_tokens = verify_input.verify(
batch, logits_output, self.page_size, vocab_mask
)
accept_length_per_req_cpu = verify_input.accept_length.cpu().tolist()
# Store accept_lens for per-request metrics
accept_lens = verify_input.accept_length
if batch.return_logprob:
@@ -277,6 +279,7 @@ class NGRAMWorker:
logits_output=logits_output,
next_token_ids=next_token_ids,
num_accepted_tokens=num_accepted_tokens,
accept_length_per_req_cpu=accept_length_per_req_cpu,
can_run_cuda_graph=can_run_cuda_graph,
accept_lens=accept_lens,
)