Skip to content

[None][perf] Reduce OpenAI stream postprocess overhead #14708

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
2ez4bz merged 3 commits into from
Jun 3, 2026
Merged

[None][perf] Reduce OpenAI stream postprocess overhead #14708

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b75f1e6
[None][perf] Reduce OpenAI stream postprocess overhead
2ez4bz May 27, 2026
ffdfef0
[None] Address review comments
2ez4bz May 29, 2026
8d05f99
[None][fix] Address review comments
2ez4bz May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 52 additions & 13 deletions tensorrt_llm/serve/harmony_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1427,7 +1427,9 @@ def create_openai_streaming_response(
tokens: list[int],
available_tools: list[dict[str, Any]] | None = None,
model_name: str = "harmony-model",
tool_choice: str | None = None) -> Tuple[list[str], bool]:
tool_choice: str | None = None,
stream_response_id: str | None = None,
stream_created: int | None = None) -> Tuple[list[str], bool]:
"""
Create properly formatted OpenAI streaming responses from harmony tokens.

Expand All @@ -1436,6 +1438,8 @@ def create_openai_streaming_response(
tokens: New tokens from this iteration
available_tools: Available tools for filtering
model_name: Model name for response
stream_response_id: Response ID shared by all chunks in the stream
stream_created: Creation timestamp shared by all chunks in the stream

Returns:
List of properly formatted streaming response strings
Expand Down Expand Up @@ -1536,9 +1540,11 @@ def create_openai_streaming_response(
finish_reason="stop" if should_stop else None,
stop_reason=None)

stream_response = ChatCompletionStreamResponse(model=model_name,
choices=[choice],
usage=None)
stream_response = _create_stream_response(
model=model_name,
choices=[choice],
stream_response_id=stream_response_id,
stream_created=stream_created)

# Convert to string
response_json = stream_response.model_dump_json(exclude_none=True)
Expand Down Expand Up @@ -1631,6 +1637,24 @@ def get_harmony_adapter() -> HarmonyAdapter:
return _SERVE_HARMONY_ADAPTER


def _create_stream_response(
model: str,
choices: List[ChatCompletionResponseStreamChoice],
usage: UsageInfo | None = None,
stream_response_id: str | None = None,
stream_created: int | None = None) -> ChatCompletionStreamResponse:
response_kwargs: dict[str, Any] = {
"model": model,
"choices": choices,
"usage": usage,
}
if stream_response_id is not None:
response_kwargs["id"] = stream_response_id
if stream_created is not None:
response_kwargs["created"] = stream_created
return ChatCompletionStreamResponse(**response_kwargs)


def handle_streaming_response(tools: List[ChatCompletionToolsParam],
tool_choice: str,
result: GenerationResult,
Expand All @@ -1640,7 +1664,9 @@ def handle_streaming_response(tools: List[ChatCompletionToolsParam],
num_prompt_tokens: int,
first_iteration: bool,
stream_options=None,
cached_tokens: int = 0) -> List[str]:
cached_tokens: int = 0,
stream_response_id: str | None = None,
stream_created: int | None = None) -> List[str]:
output = result.outputs[0]

# Convert tools to dictionary format for harmony adapter (standard pattern)
Expand Down Expand Up @@ -1670,9 +1696,12 @@ def end_streaming(res):
usage_info = _create_usage_info(num_prompt_tokens, result.outputs,
cached_tokens)

final_usage_chunk = ChatCompletionStreamResponse(choices=[],
model=model,
usage=usage_info)
final_usage_chunk = _create_stream_response(
model=model,
choices=[],
usage=usage_info,
stream_response_id=stream_response_id,
stream_created=stream_created)

final_usage_json = final_usage_chunk.model_dump_json(exclude_none=True)

Expand All @@ -1689,22 +1718,26 @@ def end_streaming(res):
tokens=output.token_ids_diff,
available_tools=tools_for_parser,
model_name=model,
tool_choice=tool_choice)
tool_choice=tool_choice,
stream_response_id=stream_response_id,
stream_created=stream_created)
if first_iteration and remaining_responses:
first_delta = DeltaMessage(role="assistant")
choice = ChatCompletionResponseStreamChoice(
index=0, delta=first_delta)
first_response = ChatCompletionStreamResponse(
first_response = _create_stream_response(
model=model,
choices=[choice],
stream_response_id=stream_response_id,
stream_created=stream_created,
)
response_json = first_response.model_dump_json(
exclude_none=True)
res.append(f"data: {response_json}\n\n")
res.extend(remaining_responses)

# Send final message with finish_reason
final_response = ChatCompletionStreamResponse(
final_response = _create_stream_response(
model=model,
choices=[
ChatCompletionResponseStreamChoice(
Expand All @@ -1713,6 +1746,8 @@ def end_streaming(res):
finish_reason=output.finish_reason,
stop_reason=output.stop_reason)
],
stream_response_id=stream_response_id,
stream_created=stream_created,
)

final_response_json = final_response.model_dump_json(
Expand All @@ -1725,7 +1760,9 @@ def end_streaming(res):
tokens=output.token_ids_diff,
available_tools=tools_for_parser,
model_name=model,
tool_choice=tool_choice)
tool_choice=tool_choice,
stream_response_id=stream_response_id,
stream_created=stream_created)
# Send first response after receiving the first output
if first_iteration:
first_iteration = False
Expand All @@ -1734,9 +1771,11 @@ def end_streaming(res):
choice = ChatCompletionResponseStreamChoice(index=0,
delta=first_delta)

first_response = ChatCompletionStreamResponse(
first_response = _create_stream_response(
model=model,
choices=[choice],
stream_response_id=stream_response_id,
stream_created=stream_created,
)

response_json = first_response.model_dump_json(
Expand Down
8 changes: 8 additions & 0 deletions tensorrt_llm/serve/openai_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1494,6 +1494,12 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
else:
prompts = request.prompt

stream_response_id = None
stream_created = None
if request.stream and len(prompts) > 1:
stream_response_id = f"cmpl-{uuid.uuid4().hex}"
stream_created = int(time.time())

promises: List[RequestOutput] = []
postproc_params_collection: List[Optional[PostprocParams]] = []
# Pass the model vocabulary size so ``logit_bias`` can be
Expand All @@ -1516,6 +1522,8 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
for idx, prompt in enumerate(prompts):
postproc_args = CompletionPostprocArgs.from_request(request)
postproc_args.prompt_idx = idx
postproc_args.stream_response_id = stream_response_id
postproc_args.stream_created = stream_created
if request.echo:
postproc_args.prompt = prompt
postproc_params = PostprocParams(
Expand Down
55 changes: 47 additions & 8 deletions tensorrt_llm/serve/postprocess_handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import time
from dataclasses import dataclass, field
from typing import Any, List, Literal, Optional, Tuple, Union

Expand Down Expand Up @@ -89,6 +90,10 @@ class ChatPostprocArgs(PostprocArgs):
tool_call_id_type: str = "random"
chat_template_kwargs: Optional[dict[str, Any]] = None
ctx_usage: Optional[UsageInfo] = None
# Cache per-request stream metadata so every chunk reuses the same response
# id and created timestamp instead of regenerating them for each chunk.
stream_response_id: Optional[str] = None
Comment thread
2ez4bz marked this conversation as resolved.
stream_created: Optional[int] = None

@classmethod
def from_request(cls, request: ChatCompletionRequest):
Expand All @@ -109,6 +114,15 @@ def from_request(cls, request: ChatCompletionRequest):
)


def _ensure_stream_metadata(args: Any, rsp: GenerationResultBase,
prefix: str) -> Tuple[str, int]:
if args.stream_response_id is None:
args.stream_response_id = f"{prefix}-{rsp.id}"
if args.stream_created is None:
args.stream_created = int(time.time())
return args.stream_response_id, args.stream_created


def create_logprobs(token_ids: List[int], tokenizer: TransformersTokenizer,
logprobs: List[float] | TokenLogprobs,
top_logprobs: bool) -> ChatCompletionLogProbs:
Expand Down Expand Up @@ -212,7 +226,9 @@ def yield_first_chat(num_tokens: int,
content=content),
finish_reason=None)
chunk = ChatCompletionStreamResponse(choices=[choice_data],
model=args.model)
model=args.model,
id=stream_response_id,
created=stream_created)
if include_continuous_usage:
chunk.usage = UsageInfo(
prompt_tokens=num_tokens,
Expand All @@ -229,6 +245,8 @@ def yield_first_chat(num_tokens: int,
finish_reason_sent = [False] * args.num_choices
prompt_tokens = args.num_prompt_tokens
ctx_usage = _ctx_usage_for_postproc(args, rsp.outputs)
stream_response_id, stream_created = _ensure_stream_metadata(
Comment thread
2ez4bz marked this conversation as resolved.
args, rsp, "chatcmpl")
if stream_option := args.stream_options:
include_usage = stream_option.include_usage
include_continuous_usage = include_usage and stream_option.continuous_usage_stats
Expand All @@ -253,7 +271,6 @@ def yield_first_chat(num_tokens: int,
continue

delta_text = output.text_diff

delta_text, reasoning_delta_text = apply_reasoning_parser(
args,
i,
Expand Down Expand Up @@ -326,7 +343,10 @@ def yield_first_chat(num_tokens: int,
choice.finish_reason = output.finish_reason
choice.stop_reason = output.stop_reason
finish_reason_sent[i] = True
chunk = ChatCompletionStreamResponse(choices=[choice], model=args.model)
chunk = ChatCompletionStreamResponse(choices=[choice],
model=args.model,
id=stream_response_id,
created=stream_created)
if include_continuous_usage:
chunk.usage = UsageInfo(prompt_tokens=prompt_tokens,
completion_tokens=output.length,
Expand All @@ -350,7 +370,9 @@ def yield_first_chat(num_tokens: int,

final_usage_chunk = ChatCompletionStreamResponse(choices=[],
model=args.model,
usage=final_usage)
usage=final_usage,
id=stream_response_id,
created=stream_created)
final_usage_data = final_usage_chunk.model_dump_json()
res.append(f"data: {final_usage_data}\n\n")
return res
Expand Down Expand Up @@ -446,6 +468,10 @@ class CompletionPostprocArgs(PostprocArgs):
return_logprobs: bool = False
stream_options: Optional[StreamOptions] = None
ctx_usage: Optional[UsageInfo] = None
# Cache per-request stream metadata so every chunk reuses the same response
# id and created timestamp instead of regenerating them for each chunk.
stream_response_id: Optional[str] = None
stream_created: Optional[int] = None

@classmethod
def from_request(cls, request: CompletionRequest):
Expand Down Expand Up @@ -500,6 +526,8 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase,
res: List[str] = []
prompt_tokens = args.num_prompt_tokens
ctx_usage = _ctx_usage_for_postproc(args, rsp.outputs)
stream_response_id, stream_created = _ensure_stream_metadata(
Comment thread
2ez4bz marked this conversation as resolved.
args, rsp, "cmpl")
if stream_option := args.stream_options:
include_usage = stream_option.include_usage
include_continuous_usage = include_usage and stream_option.continuous_usage_stats
Expand Down Expand Up @@ -527,7 +555,10 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase,
choice.logprobs = create_completion_logprobs(
token_ids, args.tokenizer, logprobs, output._last_text_len)

chunk = CompletionStreamResponse(model=args.model, choices=[choice])
chunk = CompletionStreamResponse(model=args.model,
choices=[choice],
id=stream_response_id,
created=stream_created)
if include_continuous_usage:
chunk.usage = UsageInfo(prompt_tokens=prompt_tokens,
completion_tokens=output.length,
Expand All @@ -549,9 +580,11 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase,
)
rewrite_usage_info_from_ctx(final_usage, ctx_usage)

final_usage_chunk = ChatCompletionStreamResponse(choices=[],
model=args.model,
usage=final_usage)
final_usage_chunk = CompletionStreamResponse(choices=[],
model=args.model,
usage=final_usage,
id=stream_response_id,
created=stream_created)
final_usage_data = final_usage_chunk.model_dump_json()
res.append(f"data: {final_usage_data}\n\n")
args.first_iteration = False
Expand Down Expand Up @@ -616,6 +649,8 @@ class ChatCompletionPostprocArgs(PostprocArgs):
stream_options: Optional[StreamOptions] = None
chat_template_kwargs: Optional[dict[str, Any]] = None
ctx_usage: Optional[UsageInfo] = None
stream_response_id: Optional[str] = None
stream_created: Optional[int] = None

@classmethod
def from_request(cls, request: ChatCompletionRequest):
Expand Down Expand Up @@ -662,6 +697,8 @@ def chat_harmony_streaming_post_processor(
if ctx_prompt_tokens is not None:
prompt_tokens = ctx_prompt_tokens
cached_tokens = ctx_cached_tokens
stream_response_id, stream_created = _ensure_stream_metadata(
args, rsp, "chatcmpl")
response = handle_streaming_response(
tools=args.tools,
tool_choice=args.tool_choice,
Expand All @@ -673,6 +710,8 @@ def chat_harmony_streaming_post_processor(
first_iteration=args.first_iteration,
stream_options=args.stream_options,
cached_tokens=cached_tokens,
stream_response_id=stream_response_id,
stream_created=stream_created,
)
args.first_iteration = False
return response
Expand Down
Loading
Loading