From a922c8e8848206a0b6cc5874b9e7944fa07dad19 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:52:24 -0700 Subject: [PATCH 1/2] [https://nvbugs/6693989][fix] Apply the prompt-token offset only once a count exists chat_stream_post_processor eagerly computed num_prompt_tokens - num_prompt_tokens_offset, but num_prompt_tokens is Optional and only populated later by the executor or the server, so any args reaching the handler before that raised TypeError. Only the usage-reporting branches consume the value, so gate the subtraction on a known count and drop the now-passing test's waiver. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/serve/postprocess_handlers.py | 6 +++++- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 4f16f55693d3..cf998ee860ec 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -354,7 +354,11 @@ def yield_first_chat(num_tokens: int, res: List[str] = [] finish_reason_sent = [False] * args.num_choices - prompt_tokens = args.num_prompt_tokens - args.num_prompt_tokens_offset + # num_prompt_tokens stays None until a prompt length is recorded, and only + # the usage branches below consume it, so offset it only once it exists. + prompt_tokens = args.num_prompt_tokens + if prompt_tokens is not None: + prompt_tokens -= args.num_prompt_tokens_offset ctx_usage = _ctx_usage_for_postproc(args, rsp.outputs) stream_response_id, stream_created = _ensure_stream_metadata( args, rsp, "chatcmpl") diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 88768513f396..aea4f915c6c0 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -348,7 +348,6 @@ unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfe unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) -unittest/llmapi/test_llm.py::test_chat_stream_post_processor_reuses_stream_metadata SKIP (https://nvbugs/6693989) unittest/llmapi/test_llm.py::test_generate_with_detokenization_stop_words_streaming[/scratch.trt_llm_data/llm-models/gemma/gemma-3-1b-it] SKIP (https://nvbugs/6566772) unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" SKIP (https://nvbugs/6428092) unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-False-True] SKIP (https://nvbugs/6432826) From eea8ef354d85a951403baf6bdb1fe0d2f7c5b287 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 2 Sep 2026 09:42:46 -0700 Subject: [PATCH 2/2] [https://nvbugs/6693989][fix] Fail clearly when streaming usage is requested without a prompt token count Address the CodeRabbit review on #18509: keeping prompt_tokens as None when the count is unknown still let None reach the usage arithmetic and UsageInfo's int fields once include_usage or continuous_usage_stats is on. Real requests never hit this (the server and the executor both record num_prompt_tokens before the first chunk is post-processed), so rather than silently dropping usage chunks, raise a ValueError that names the missing field. Add regression tests for the offset applied to streaming usage and for the new error. Signed-off-by: Michal Guzek --- tensorrt_llm/serve/postprocess_handlers.py | 11 +++++++ tests/unittest/llmapi/test_llm.py | 37 +++++++++++++++++++++- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index cf998ee860ec..0a4c9709ec07 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -368,6 +368,17 @@ def yield_first_chat(num_tokens: int, else: include_usage = False include_continuous_usage = False + if include_usage and prompt_tokens is None: + # The usage chunks below feed prompt_tokens into UsageInfo (int fields) + # and into the total_tokens arithmetic. The server records the prompt + # length before the first chunk is post-processed (the executor does so + # on the postproc-worker path), so a missing count here means the + # caller wired PostprocArgs without one; fail with a clear message + # instead of a TypeError from the usage math. + raise ValueError( + "Streaming usage was requested, but PostprocArgs.num_prompt_tokens " + "is not set; record the prompt token count before " + "chat_stream_post_processor reports usage.") if args.first_iteration: for i in range(args.num_choices): res.append( diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 5b62972ec8c2..2ff1d9d8171e 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -40,7 +40,7 @@ from tensorrt_llm.llmapi.tokenizer import (TokenizerBase, TransformersTokenizer, load_hf_tokenizer) from tensorrt_llm.sampling_params import LogitsProcessor, SamplingParams -from tensorrt_llm.serve.openai_protocol import CompletionRequest +from tensorrt_llm.serve.openai_protocol import CompletionRequest, StreamOptions from tensorrt_llm.serve.openai_server import OpenAIServer from tensorrt_llm.serve.postprocess_handlers import (ChatPostprocArgs, chat_stream_post_processor) @@ -1246,6 +1246,41 @@ def test_chat_stream_post_processor_reuses_stream_metadata() -> None: assert payloads[-1]["choices"][0]["delta"]["content"] == "y" +def test_chat_stream_post_processor_usage_applies_prompt_token_offset() -> None: + result = GenerationResultBase(123, SamplingParams()) + output = result._outputs[0] + output.text = "x" + output.token_ids = [1] + output.finish_reason = "stop" + result._done = True + + args = ChatPostprocArgs(role="assistant", + model="test-model", + num_prompt_tokens=5, + num_prompt_tokens_offset=3, + stream_options=StreamOptions(include_usage=True)) + payloads = _stream_payloads_from_chunks( + chat_stream_post_processor(result, args)) + + final_chunk = payloads[-1] + assert final_chunk["choices"] == [] + assert final_chunk["usage"]["prompt_tokens"] == 2 + assert final_chunk["usage"]["completion_tokens"] == 1 + assert final_chunk["usage"]["total_tokens"] == 3 + + +def test_chat_stream_post_processor_usage_requires_prompt_token_count() -> None: + # Usage arithmetic needs a concrete count: a missing one must surface as a + # clear error, never as None leaking into UsageInfo or a TypeError. + result = GenerationResultBase(123, SamplingParams()) + args = ChatPostprocArgs(role="assistant", + model="test-model", + stream_options=StreamOptions(include_usage=True)) + + with pytest.raises(ValueError, match="num_prompt_tokens"): + chat_stream_post_processor(result, args) + + class _FakeCompletionGeneratorArgs: backend = "pytorch" gather_generation_logits = False