diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index dd73abbb42aa..d0d05e3a483e 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -649,6 +649,7 @@ async def _generate_chat_stream( # Usage tracking prompt_tokens = {} + reasoning_tokens = {} completion_tokens = {} cached_tokens = {} hidden_states = {} @@ -670,6 +671,9 @@ async def _generate_chat_stream( 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 @@ async def _generate_chat_stream( 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 @@ async def _generate_chat_stream( 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 @@ async def _generate_chat_stream( 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 @@ async def _process_tool_call_stream( 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 @@ async def _process_tool_call_stream( 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" diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index d99678fa8fa3..9d089fc3219e 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -220,6 +220,7 @@ async def _generate_completion_stream( # Usage tracking prompt_tokens = {} completion_tokens = {} + reasoning_tokens = {} cached_tokens = {} hidden_states = {} routed_experts = {} @@ -241,6 +242,9 @@ async def _generate_completion_stream( 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 @@ async def _generate_completion_stream( 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 @@ async def _generate_completion_stream( 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, ) diff --git a/python/sglang/srt/entrypoints/openai/usage_processor.py b/python/sglang/srt/entrypoints/openai/usage_processor.py index de88e9eb001c..8a6c9d7a29cc 100644 --- a/python/sglang/srt/entrypoints/openai/usage_processor.py +++ b/python/sglang/srt/entrypoints/openai/usage_processor.py @@ -23,12 +23,16 @@ def calculate_response_usage( 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 @@ def calculate_response_usage( 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 @@ def calculate_response_usage( @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 @@ def calculate_streaming_usage( 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 @@ def calculate_streaming_usage( 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 @@ def calculate_streaming_usage( 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 @@ def calculate_token_usage( completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, prompt_tokens_details=cached_tokens, + reasoning_tokens=reasoning_tokens, ) diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index ce27113845c7..568699c4ddcd 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -332,6 +332,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput): 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, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 53a6fc90280e..bd979653458a 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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 diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 4da03863068b..e0a1669fb3e6 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index b44c75a5dbee..35858c705a2a 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -634,9 +634,13 @@ def __init__( # 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 @@ def set_finish_with_abort(self, error_msg: str): 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}, " diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fe3379c95f84..67af2d0de943 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -552,6 +552,9 @@ def init_tokenizer(self): 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) diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index e3b8dc87e26d..496cd96656e5 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -51,7 +51,7 @@ def _get_storage_backend_type(self) -> str: 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 @@ def process_batch_result_prefill( # 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 @@ def process_batch_result_decode( 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 @@ def process_batch_result_decode( 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 @@ def process_batch_result_decode( 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 @@ def _mamba_prefix_cache_update( ) 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 @@ def _process_input_token_logprobs( 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 @@ def _process_input_top_logprobs(self, req: Req) -> None: 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 @@ def _process_input_token_ids_logprobs(self, req: Req) -> None: 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 @@ def _calculate_relevant_tokens_len(self, req: Req) -> int: 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 @@ def _calculate_num_input_logprobs( # 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 @@ def add_logprob_return_values( 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 @@ def stream_output( 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 @@ def stream_output_generation( 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 @@ def stream_output_generation( ) 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 @@ def stream_output_generation( 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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index eca904d8c153..81424329a0c1 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1574,6 +1574,7 @@ def _handle_batch_output( 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], } diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 8c108915c939..098338c5ee5b 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -215,6 +215,7 @@ def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResul 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 @@ def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResul 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 @@ def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResul 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, ) diff --git a/test/registered/openai_server/features/test_reasoning_usage_tokens.py b/test/registered/openai_server/features/test_reasoning_usage_tokens.py new file mode 100644 index 000000000000..d53ef88c440e --- /dev/null +++ b/test/registered/openai_server/features/test_reasoning_usage_tokens.py @@ -0,0 +1,180 @@ +"""Usage: +python3 -m unittest openai_server.features.test_reasoning_usage_tokens.TestNormalReasoningTokenUsage +python3 -m unittest openai_server.features.test_reasoning_usage_tokens.TestSpecReasoningTokenUsage +python3 -m unittest openai_server.features.test_reasoning_usage_tokens.TestSpecV2ReasoningTokenUsage +""" + +import json +import os +import unittest + +import requests +from openai import OpenAI + +from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.srt.utils import kill_process_tree +from sglang.srt.utils.hf_transformers_utils import get_tokenizer +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_REASONING_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=90, suite="stage-b-test-1-gpu-large") + + +def remove_prefix(text: str, prefix: str) -> str: + return text[len(prefix) :] if text.startswith(prefix) else text + + +class ReasoningTokenUsageMixin: + model = "" + reasoning_parser_name = "" + extra_server_args = [] + extra_env_vars = {} + max_new_tokens = 1024 + + @classmethod + def setUpClass(cls): + for k, v in cls.extra_env_vars.items(): + os.environ[k] = v + + assert cls.model + cls.base_url = DEFAULT_URL_FOR_TEST + cls.api_key = "sk-1234" + + # get think_end_token_id + cls.tokenizer = get_tokenizer(cls.model) + reasoning_parser = ReasoningParser(cls.reasoning_parser_name) + cls.think_end_token_id = cls.tokenizer.convert_tokens_to_ids( + reasoning_parser.detector.think_end_token + ) + assert ( + cls.think_end_token_id + ), f"think_end_token_id for {cls.reasoning_parser_name} shouldn't be None" + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + api_key=cls.api_key, + other_args=[ + "--reasoning-parser", + cls.reasoning_parser_name, + ] + + cls.extra_server_args, + ) + cls.client = OpenAI(base_url=f"{cls.base_url}/v1", api_key=cls.api_key) + cls.messages = [{"role": "user", "content": "What is 1+3?"}] + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process"): + kill_process_tree(cls.process.pid) + + def test_generate_api_non_streaming(self): + response = requests.post( + url=f"{self.base_url}/generate", + headers={"Authorization": f"Bearer {self.api_key}"}, + json={ + "text": self.tokenizer.apply_chat_template( + self.messages, add_generation_prompt=True, tokenize=False + ), + "model": self.model, + "require_reasoning": True, + "sampling_params": {"max_new_tokens": self.max_new_tokens}, + }, + ) + response.raise_for_status() + res_json = response.json() + report_reasoning_tokens = res_json["meta_info"]["reasoning_tokens"] + actual_reasoning_tokens = ( + res_json["output_ids"].index(self.think_end_token_id) + 1 + ) + assert ( + report_reasoning_tokens == actual_reasoning_tokens + ), f"Expected {actual_reasoning_tokens}, got {report_reasoning_tokens}" + + def test_generate_api_streaming(self): + response = requests.post( + url=f"{self.base_url}/generate", + headers={"Authorization": f"Bearer {self.api_key}"}, + json={ + "text": self.tokenizer.apply_chat_template( + self.messages, add_generation_prompt=True, tokenize=False + ), + "model": self.model, + "require_reasoning": True, + "sampling_params": {"max_new_tokens": 1024}, + "stream": True, + }, + stream=True, + ) + response.raise_for_status() + for chunk in response.iter_lines(): + if not chunk: + continue + decoded_str = remove_prefix(chunk.decode("utf-8"), "data: ") + if decoded_str != "[DONE]": + data = json.loads(decoded_str) + report_reasoning_tokens = data["meta_info"]["reasoning_tokens"] + if self.think_end_token_id in data["output_ids"]: + actual_reasoning_tokens = ( + data["output_ids"].index(self.think_end_token_id) + 1 + ) + else: + actual_reasoning_tokens = len(data["output_ids"]) + assert report_reasoning_tokens == actual_reasoning_tokens + + def test_chat_api_non_streaming(self): + response = self.client.chat.completions.create( + model=self.model, messages=self.messages, max_tokens=1024 + ) + assert response.usage is not None + assert response.usage.reasoning_tokens > 0 + + def test_chat_api_streaming(self): + response = self.client.chat.completions.create( + model=self.model, + messages=self.messages, + max_tokens=1024, + stream=True, + stream_options={"include_usage": True, "continuous_usage_stats": True}, + ) + for chunk in response: + if chunk.usage: + assert chunk.usage.reasoning_tokens > 0 + + +class TestNormalReasoningTokenUsage(ReasoningTokenUsageMixin, CustomTestCase): + model = DEFAULT_REASONING_MODEL_NAME_FOR_TEST + reasoning_parser_name = "deepseek-r1" + extra_server_args = ["--cuda-graph-max-bs", "2"] + + +class TestSpecReasoningTokenUsage(ReasoningTokenUsageMixin, CustomTestCase): + model = "Qwen/Qwen3-30B-A3B" # select this model due to its suitable eagle model + reasoning_parser_name = "qwen3" + extra_env_vars = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"} + extra_server_args = [ + "--speculative-algorithm", + "EAGLE3", + "--speculative-draft-model-path", + "nex-agi/SGLANG-EAGLE3-Qwen3-30B-A3B-Nex-N1", + "--cuda-graph-max-bs", + "2", + ] + + +class TestSpecV2ReasoningTokenUsage(TestSpecReasoningTokenUsage): + extra_env_vars = { + "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", + "SGLANG_ENABLE_SPEC_V2": "1", + } + + +if __name__ == "__main__": + unittest.main()