diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 4f16f55693d3..0a4c9709ec07 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") @@ -364,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/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) 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