Add --stream-response-default-include-usage server flag (#16711)

This commit is contained in:
faceless void
2026-04-03 21:36:00 -07:00
committed by GitHub
parent 31c9d8e885
commit de9859073f
7 changed files with 47 additions and 19 deletions
@@ -341,7 +341,6 @@ async def lifespan(fast_api_app: FastAPI):
_global_state.tokenizer_manager,
_global_state.template_manager,
enable_prompt_tokens_details=True,
enable_force_include_usage=True,
tool_server=tool_server,
)
except Exception:
@@ -41,6 +41,7 @@ from sglang.srt.entrypoints.openai.utils import (
process_cached_tokens_details_from_ret,
process_hidden_states_from_ret,
process_routed_experts_from_ret,
should_include_usage,
to_openai_style_logprobs,
)
from sglang.srt.function_call.core_types import ToolCallItem
@@ -655,6 +656,11 @@ class OpenAIServingChat(OpenAIServingBase):
stream_started = False
try:
include_usage, continuous_usage_stats = should_include_usage(
request.stream_options,
self.tokenizer_manager.server_args.stream_response_default_include_usage,
)
async for content in self.tokenizer_manager.generate_request(
adapted_request, raw_request
):
@@ -743,10 +749,7 @@ class OpenAIServingChat(OpenAIServingBase):
)
# Add usage stats if continuous_usage_stats is enabled
if (
request.stream_options
and request.stream_options.continuous_usage_stats
):
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
@@ -767,6 +770,7 @@ class OpenAIServingChat(OpenAIServingBase):
content,
request,
has_tool_calls,
continuous_usage_stats,
):
if chunk:
yield chunk
@@ -798,10 +802,7 @@ class OpenAIServingChat(OpenAIServingBase):
)
# Add usage stats if continuous_usage_stats is enabled
if (
request.stream_options
and request.stream_options.continuous_usage_stats
):
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
@@ -881,7 +882,7 @@ class OpenAIServingChat(OpenAIServingBase):
yield f"data: {routed_experts_chunk.model_dump_json()}\n\n"
# Additional usage chunk
if request.stream_options and request.stream_options.include_usage:
if include_usage:
usage = UsageProcessor.calculate_streaming_usage(
prompt_tokens,
completion_tokens,
@@ -1313,6 +1314,7 @@ class OpenAIServingChat(OpenAIServingBase):
content: Dict[str, Any],
request: ChatCompletionRequest,
has_tool_calls: Dict[int, bool],
continuous_usage_stats: bool = False,
):
"""Process tool calls in streaming response"""
if index not in parser_dict:
@@ -1351,7 +1353,7 @@ class OpenAIServingChat(OpenAIServingBase):
)
# Add usage stats if continuous_usage_stats is enabled
if request.stream_options and request.stream_options.continuous_usage_stats:
if continuous_usage_stats:
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
completion_tokens = content["meta_info"].get("completion_tokens", 0)
chunk.usage = UsageProcessor.calculate_token_usage(
@@ -1401,7 +1403,7 @@ class OpenAIServingChat(OpenAIServingBase):
)
# Add usage stats if continuous_usage_stats is enabled
if request.stream_options and request.stream_options.continuous_usage_stats:
if continuous_usage_stats:
prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
completion_tokens = content["meta_info"].get("completion_tokens", 0)
chunk.usage = UsageProcessor.calculate_token_usage(
@@ -23,6 +23,7 @@ from sglang.srt.entrypoints.openai.utils import (
process_cached_tokens_details_from_ret,
process_hidden_states_from_ret,
process_routed_experts_from_ret,
should_include_usage,
to_openai_style_logprobs,
)
from sglang.srt.managers.io_struct import GenerateReqInput
@@ -225,6 +226,11 @@ class OpenAIServingCompletion(OpenAIServingBase):
stream_started = False
try:
include_usage, continuous_usage_stats = should_include_usage(
request.stream_options,
self.tokenizer_manager.server_args.stream_response_default_include_usage,
)
async for content in self.tokenizer_manager.generate_request(
adapted_request, raw_request
):
@@ -318,10 +324,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
)
# Add usage stats if continuous_usage_stats is enabled
if (
request.stream_options
and request.stream_options.continuous_usage_stats
):
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
@@ -371,7 +374,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
yield f"data: {routed_experts_chunk.model_dump_json()}\n\n"
# Handle final usage chunk
if request.stream_options and request.stream_options.include_usage:
if include_usage:
usage = UsageProcessor.calculate_streaming_usage(
prompt_tokens,
completion_tokens,
@@ -76,7 +76,6 @@ class OpenAIServingResponses(OpenAIServingChat):
template_manager: TemplateManager,
*,
enable_prompt_tokens_details: bool = False,
enable_force_include_usage: bool = False,
tool_server: Optional[ToolServer] = None,
) -> None:
super().__init__(tokenizer_manager, template_manager)
@@ -84,7 +83,6 @@ class OpenAIServingResponses(OpenAIServingChat):
# template_manager is already set by parent class
self.reasoning_parser = self.tokenizer_manager.server_args.reasoning_parser
self.enable_prompt_tokens_details = enable_prompt_tokens_details
self.enable_force_include_usage = enable_force_include_usage
# Get default sampling params from model config if available
self.default_sampling_params = {}
@@ -6,6 +6,7 @@ from sglang.srt.entrypoints.openai.protocol import (
ChatCompletionRequest,
CompletionRequest,
LogProbs,
StreamOptions,
)
logger = logging.getLogger(__name__)
@@ -73,6 +74,23 @@ def process_hidden_states_from_ret(
return hidden_states
def should_include_usage(
stream_options: StreamOptions | None, stream_response_default_include_usage: bool
) -> tuple[bool, bool]:
# When stream_options are specified in the request
if stream_options:
include_usage = (
stream_options.include_usage or stream_response_default_include_usage
)
continuous_usage_stats = bool(stream_options.continuous_usage_stats)
else:
include_usage, continuous_usage_stats = (
stream_response_default_include_usage,
False,
)
return include_usage, continuous_usage_stats
def process_routed_experts_from_ret(
ret_item: Dict[str, Any],
request: Union[
+7
View File
@@ -374,6 +374,7 @@ class ServerArgs:
pp_max_micro_batch_size: Optional[int] = None
pp_async_batch_depth: int = 0
stream_interval: int = 1
stream_response_default_include_usage: bool = False
incremental_streaming_output: bool = False
enable_streaming_session: bool = False
random_seed: Optional[int] = None
@@ -4155,6 +4156,12 @@ class ServerArgs:
action="store_true",
help="Whether to output as a sequence of disjoint segments.",
)
parser.add_argument(
"--stream-response-default-include-usage",
action="store_true",
help="Include usage in every streaming response "
"(even when stream_options is not specified).",
)
parser.add_argument(
"--stream-output",
action=DeprecatedStoreTrueAction,