From b75f1e677e4d07a542be379af59d3a8447468d5f Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Wed, 27 May 2026 15:47:33 -0700 Subject: [PATCH 1/3] [None][perf] Reduce OpenAI stream postprocess overhead Streaming responses created fresh metadata for each chunk when callers did not pass it explicitly. High-concurrency workloads can emit hundreds of thousands of chunks, which makes UUID generation, and time lookups part of the CPU hot path. Reuse stream metadata per request. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- tensorrt_llm/serve/postprocess_handlers.py | 45 ++++++++++++++++++---- tests/unittest/llmapi/test_llm.py | 34 +++++++++++++++- 2 files changed, 69 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 6dc6d8a77d8e..52cc0984c1c5 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -1,3 +1,4 @@ +import time from dataclasses import dataclass, field from typing import Any, List, Literal, Optional, Tuple, Union @@ -89,6 +90,8 @@ class ChatPostprocArgs(PostprocArgs): tool_call_id_type: str = "random" 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): @@ -109,6 +112,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: @@ -212,7 +224,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, @@ -229,6 +243,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( + 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 @@ -253,7 +269,6 @@ def yield_first_chat(num_tokens: int, continue delta_text = output.text_diff - delta_text, reasoning_delta_text = apply_reasoning_parser( args, i, @@ -326,7 +341,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, @@ -350,7 +368,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 @@ -446,6 +466,8 @@ class CompletionPostprocArgs(PostprocArgs): return_logprobs: bool = False stream_options: Optional[StreamOptions] = None ctx_usage: Optional[UsageInfo] = None + stream_response_id: Optional[str] = None + stream_created: Optional[int] = None @classmethod def from_request(cls, request: CompletionRequest): @@ -500,6 +522,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( + 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 @@ -527,7 +551,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, @@ -549,9 +576,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 diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index fb727826cc57..41c897f1620d 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -30,8 +30,9 @@ from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm.bindings import executor as tllm from tensorrt_llm.executor import (GenerationExecutorWorker, GenerationRequest, - GenerationResult, LoRARequest, - PromptAdapterRequest, RequestError) + GenerationResult, GenerationResultBase, + LoRARequest, PromptAdapterRequest, + RequestError) from tensorrt_llm.llmapi import (BuildCacheConfig, EagleDecodingConfig, ExtendedRuntimePerfKnobConfig, KvCacheConfig, KvCacheRetentionConfig, @@ -50,6 +51,8 @@ from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode from tensorrt_llm.sampling_params import (BatchedLogitsProcessor, LogitsProcessor, SamplingParams) +from tensorrt_llm.serve.postprocess_handlers import (ChatPostprocArgs, + chat_stream_post_processor) # isort: off sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..") @@ -2575,6 +2578,33 @@ def test_llm_with_postprocess_parallel(): run_llm_with_postprocess_parallel(tp_size=1) +def test_chat_stream_post_processor_reuses_stream_metadata(): + result = GenerationResultBase(123, SamplingParams()) + output = result._outputs[0] + output.text = "x" + output.token_ids = [1] + + args = ChatPostprocArgs(role="assistant", model="test-model") + chunks = chat_stream_post_processor(result, args) + + output._last_text_len = len(output.text) + output._last_token_ids_len = len(output.token_ids) + output.text = "xy" + output.token_ids.append(2) + chunks += chat_stream_post_processor(result, args) + + payloads = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: "): + payloads.append(json.loads(line[len("data: "):].strip())) + + assert {payload["id"] for payload in payloads} == {"chatcmpl-123"} + assert len({payload["created"] for payload in payloads}) == 1 + assert payloads[0]["choices"][0]["delta"]["role"] == "assistant" + assert payloads[-1]["choices"][0]["delta"]["content"] == "y" + + def run_llm_with_postprocess_parallel_and_result_handler( streaming, backend, tp_size: int = 1): # avoid import error when running in CI From ffdfef0c726f7a49e5f834dcf420f8c7a23622c6 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Thu, 28 May 2026 21:44:30 -0700 Subject: [PATCH 2/3] [None] Address review comments Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- tensorrt_llm/serve/postprocess_handlers.py | 4 ++++ tests/unittest/llmapi/test_llm.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 52cc0984c1c5..322259dbe392 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -90,6 +90,8 @@ 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 stream_created: Optional[int] = None @@ -466,6 +468,8 @@ 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 diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 41c897f1620d..935a84644c77 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -2578,7 +2578,7 @@ def test_llm_with_postprocess_parallel(): run_llm_with_postprocess_parallel(tp_size=1) -def test_chat_stream_post_processor_reuses_stream_metadata(): +def test_chat_stream_post_processor_reuses_stream_metadata() -> None: result = GenerationResultBase(123, SamplingParams()) output = result._outputs[0] output.text = "x" From 8d05f992bd7e0ed4364a14fbdf70c685c8c6e5c4 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Fri, 29 May 2026 11:08:37 -0700 Subject: [PATCH 3/3] [None][fix] Address review comments Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- tensorrt_llm/serve/harmony_adapter.py | 65 ++++++-- tensorrt_llm/serve/openai_server.py | 8 + tensorrt_llm/serve/postprocess_handlers.py | 6 + .../llmapi/apps/test_harmony_parsing.py | 122 ++++++++++++++ tests/unittest/llmapi/test_llm.py | 150 +++++++++++++++++- 5 files changed, 333 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/serve/harmony_adapter.py b/tensorrt_llm/serve/harmony_adapter.py index 6e9949802c45..f326fc74360e 100644 --- a/tensorrt_llm/serve/harmony_adapter.py +++ b/tensorrt_llm/serve/harmony_adapter.py @@ -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. @@ -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 @@ -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) @@ -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, @@ -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) @@ -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) @@ -1689,14 +1718,18 @@ 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) @@ -1704,7 +1737,7 @@ def end_streaming(res): res.extend(remaining_responses) # Send final message with finish_reason - final_response = ChatCompletionStreamResponse( + final_response = _create_stream_response( model=model, choices=[ ChatCompletionResponseStreamChoice( @@ -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( @@ -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 @@ -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( diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 55e624d7bf8a..6d10cd6ac116 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -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 @@ -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( diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 322259dbe392..9b74ae328415 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -649,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): @@ -695,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, @@ -706,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 diff --git a/tests/unittest/llmapi/apps/test_harmony_parsing.py b/tests/unittest/llmapi/apps/test_harmony_parsing.py index 7636e0edb90d..8238bc3643f8 100644 --- a/tests/unittest/llmapi/apps/test_harmony_parsing.py +++ b/tests/unittest/llmapi/apps/test_harmony_parsing.py @@ -38,6 +38,10 @@ ) from tensorrt_llm.serve.openai_protocol import StreamOptions, _logit_bias_to_embedding_bias from tensorrt_llm.serve.openai_server import OpenAIServer + from tensorrt_llm.serve.postprocess_handlers import ( + ChatCompletionPostprocArgs, + chat_harmony_streaming_post_processor, + ) _harmony_available = True except (ImportError, ModuleNotFoundError): @@ -97,6 +101,7 @@ def _make_mock_result( output.token_ids = token_ids or [1, 2, 3] output.finish_reason = finish_reason output.stop_reason = stop_reason + output.disaggregated_params = None result = Mock() result.outputs = [output] @@ -105,6 +110,15 @@ def _make_mock_result( return result +def _stream_payloads(responses): + payloads = [] + for response in responses: + for line in response.splitlines(): + if line.startswith("data: "): + payloads.append(json.loads(line[len("data: ") :].strip())) + return payloads + + # =========================================================================== # Fix 1 — fe9e1a33: Agentic coding use cases # =========================================================================== @@ -618,6 +632,114 @@ def test_create_response_message_without_reasoning(self): assert "reasoning_content" not in result +class TestStreamMetadata: + """Verify Harmony chat streams reuse one response ID and timestamp.""" + + def test_postprocessor_passes_cached_stream_metadata(self): + mock_result = _make_mock_result(token_ids_diff=[1, 2, 3]) + mock_result.id = 456 + mock_result._done = False + args = ChatCompletionPostprocArgs( + model="test-model", + tools=[], + tool_choice=None, + num_prompt_tokens=5, + ) + + with patch( + "tensorrt_llm.serve.postprocess_handlers.handle_streaming_response", return_value=[] + ) as mock_handle: + chat_harmony_streaming_post_processor(mock_result, args) + first_created = args.stream_created + chat_harmony_streaming_post_processor(mock_result, args) + + assert mock_handle.call_count == 2 + call_kwargs = mock_handle.call_args_list[0].kwargs + second_call_kwargs = mock_handle.call_args_list[1].kwargs + assert call_kwargs["stream_response_id"] == "chatcmpl-456" + assert call_kwargs["stream_created"] == first_created + assert second_call_kwargs["stream_response_id"] == "chatcmpl-456" + assert second_call_kwargs["stream_created"] == first_created + assert args.stream_response_id == "chatcmpl-456" + assert isinstance(args.stream_created, int) + + def test_adapter_streaming_chunks_use_supplied_metadata(self): + adapter = HarmonyAdapter(harmony_input=False, harmony_output=False) + request_id = "test-stream-metadata" + adapter.create_stream_state(request_id=request_id, available_tools=None, tool_choice=None) + try: + with patch.object( + adapter, + "stateful_stream_harmony_tokens_to_openai_deltas", + return_value=[{"content": "hello"}, {"reasoning": "thinking"}], + ): + responses, _ = adapter.create_openai_streaming_response( + request_id=request_id, + tokens=[1, 2, 3], + available_tools=None, + model_name="test-model", + tool_choice=None, + stream_response_id="chatcmpl-fixed", + stream_created=123456, + ) + + payloads = _stream_payloads(responses) + assert len(payloads) == 2 + assert {payload["id"] for payload in payloads} == {"chatcmpl-fixed"} + assert {payload["created"] for payload in payloads} == {123456} + finally: + adapter.cleanup_stream_state(request_id) + + def test_handle_streaming_response_uses_supplied_metadata(self): + mock_result = _make_mock_result(token_ids_diff=[10, 20]) + harmony_adapter = get_harmony_adapter() + request_id = "test-handle-stream-metadata" + harmony_adapter.create_stream_state( + request_id=request_id, available_tools=None, tool_choice=None + ) + streamed_chunk = json.dumps( + { + "id": "chatcmpl-fixed", + "object": "chat.completion.chunk", + "created": 123456, + "model": "test-model", + "choices": [{"index": 0, "delta": {"content": "hi"}}], + } + ) + try: + with patch.object( + harmony_adapter, + "create_openai_streaming_response", + return_value=([f"data: {streamed_chunk}\n\n"], False), + ) as mock_create: + responses = handle_streaming_response( + tools=[], + tool_choice=None, + result=mock_result, + model="test-model", + request_id=request_id, + done=True, + num_prompt_tokens=5, + first_iteration=True, + stream_options=StreamOptions(include_usage=True), + stream_response_id="chatcmpl-fixed", + stream_created=123456, + ) + + mock_create.assert_called_once() + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs["stream_response_id"] == "chatcmpl-fixed" + assert call_kwargs["stream_created"] == 123456 + + payloads = _stream_payloads(responses) + assert len(payloads) == 4 + assert {payload["id"] for payload in payloads} == {"chatcmpl-fixed"} + assert {payload["created"] for payload in payloads} == {123456} + finally: + if request_id in harmony_adapter._stream_states: + harmony_adapter.cleanup_stream_state(request_id) + + class TestRemainingTokensOnDone: """Verify handle_streaming_response processes leftover tokens on done=True.""" diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 935a84644c77..fb174a99152d 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -51,6 +51,8 @@ from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode from tensorrt_llm.sampling_params import (BatchedLogitsProcessor, LogitsProcessor, SamplingParams) +from tensorrt_llm.serve.openai_protocol import CompletionRequest +from tensorrt_llm.serve.openai_server import OpenAIServer from tensorrt_llm.serve.postprocess_handlers import (ChatPostprocArgs, chat_stream_post_processor) @@ -2578,6 +2580,19 @@ def test_llm_with_postprocess_parallel(): run_llm_with_postprocess_parallel(tp_size=1) +def _stream_payloads_from_chunks(chunks): + payloads = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in chunk.splitlines(): + if line.startswith("data: "): + data = line[len("data: "):].strip() + if data != "[DONE]": + payloads.append(json.loads(data)) + return payloads + + def test_chat_stream_post_processor_reuses_stream_metadata() -> None: result = GenerationResultBase(123, SamplingParams()) output = result._outputs[0] @@ -2593,11 +2608,7 @@ def test_chat_stream_post_processor_reuses_stream_metadata() -> None: output.token_ids.append(2) chunks += chat_stream_post_processor(result, args) - payloads = [] - for chunk in chunks: - for line in chunk.splitlines(): - if line.startswith("data: "): - payloads.append(json.loads(line[len("data: "):].strip())) + payloads = _stream_payloads_from_chunks(chunks) assert {payload["id"] for payload in payloads} == {"chatcmpl-123"} assert len({payload["created"] for payload in payloads}) == 1 @@ -2605,6 +2616,135 @@ def test_chat_stream_post_processor_reuses_stream_metadata() -> None: assert payloads[-1]["choices"][0]["delta"]["content"] == "y" +class _FakeCompletionGeneratorArgs: + backend = "pytorch" + gather_generation_logits = False + num_postprocess_workers = 0 + return_perf_metrics = False + + +class _FakeModelConfig: + vocab_size = 32000 + + +class _FakeCompletionStreamResult(GenerationResultBase): + + @property + def finished(self): + return self._done + + @property + def request_id(self): + return self.id + + +class _FakeCompletionPromise: + + def __init__(self, result, prompt_token_ids): + self._result = result + self._yielded = False + self.prompt_token_ids = prompt_token_ids + self.aborted = False + + @property + def finished(self): + return True + + @property + def request_id(self): + return self._result.id + + def abort(self): + self.aborted = True + + def __aiter__(self): + return self + + async def __anext__(self): + if self._yielded: + raise StopAsyncIteration + self._yielded = True + return self._result + + +class _FakeCompletionGenerator: + + def __init__(self): + self.args = _FakeCompletionGeneratorArgs() + self.postproc_args = [] + + def input_processor(self, prompt, _sampling_params): + token_id = ord(prompt["prompt"]) + return [token_id], {} + + def generate_async(self, inputs, sampling_params, _postproc_params, + streaming, **_kwargs): + assert streaming + result_id = 100 + len(self.postproc_args) + result = _FakeCompletionStreamResult(result_id, sampling_params) + result._done = True + + output = result._outputs[0] + output.text = f"text-{result_id}" + output.token_ids = [result_id] + output.finish_reason = "stop" + + self.postproc_args.append(_postproc_params.postproc_args) + return _FakeCompletionPromise(result, inputs["prompt_token_ids"]) + + +class _FakeRawRequestState: + pass + + +class _FakeRawRequest: + + def __init__(self): + self.headers = {} + self.state = _FakeRawRequestState() + self.client = "test-client" + + async def is_disconnected(self): + return True + + +def test_openai_completion_list_prompt_stream_reuses_stream_metadata() -> None: + + async def run_request(): + generator = _FakeCompletionGenerator() + server = object.__new__(OpenAIServer) + server.generator = generator + server.model = "test-model" + server.model_config = _FakeModelConfig() + server.tokenizer = None + server.metrics_collector = None + server.perf_metrics = None + + request = CompletionRequest(model="test-model", + prompt=["A", "B"], + stream=True) + response = await server.openai_completion(request, _FakeRawRequest()) + chunks = [chunk async for chunk in response.body_iterator] + return generator, _stream_payloads_from_chunks(chunks) + + generator, payloads = asyncio.run(run_request()) + + ids = {payload["id"] for payload in payloads} + created = {payload["created"] for payload in payloads} + choice_indexes = { + payload["choices"][0]["index"] + for payload in payloads if payload["choices"] + } + + assert len(payloads) == 2 + assert len(ids) == 1 + assert ids.isdisjoint({"cmpl-100", "cmpl-101"}) + assert len(created) == 1 + assert choice_indexes == {0, 1} + assert {args.stream_response_id for args in generator.postproc_args} == ids + assert {args.stream_created for args in generator.postproc_args} == created + + def run_llm_with_postprocess_parallel_and_result_handler( streaming, backend, tp_size: int = 1): # avoid import error when running in CI