[https://nvbugs/6693989][fix] Apply the K3 prompt-token offset only when a prompt count is actually known…

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(

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)

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