[https://nvbugs/6223556][fix] Propagate gen-first ctx usage via aux buffer to postproc (#15246)

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>

Author: reasonsolo

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -578,6 +578,9 @@ def __init__(self,
         self._py_result = py_result
         self.is_final = is_final
         self.cached_tokens = 0
+        # Context-worker usage for gen-first disagg, delivered via the
+        # KV-transfer aux buffer (see _maybe_attach_ctx_usage).
+        self.ctx_usage = None
         # Time breakdown metrics for performance analysis
         # Contains: step_metrics (list), ctx_gpu_forward_time (float), ctx_gpu_sample_time (float)
         self.time_breakdown_metrics = time_breakdown_metrics

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -789,6 +789,7 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest):
             response = request.create_response(False, self.dist.rank)
             if response:
                 response.result.cached_tokens = request.cached_tokens
+                self._maybe_attach_ctx_usage(request, response)
                 # Buffer the response instead of enqueueing immediately.
                 # With ADP, _enqueue_responses does a tp_gather collective.
                 # Calling it here would deadlock because only the owning DP
@@ -4321,6 +4322,15 @@ def fail_request(message: str) -> bool:
         cum_log_probs[seq_slot, :beam_width].copy_(values)
         return True
 
+    @staticmethod
+    def _maybe_attach_ctx_usage(request: LlmRequest, response):
+        """Surface gen-first ctx usage (delivered via the KV-transfer aux
+        buffer in RxSession.unpack_aux) onto the response so the postprocessor
+        adopts the context-side prompt/cached token accounting."""
+        disagg_params = request.py_disaggregated_params
+        if disagg_params is not None and disagg_params.ctx_usage is not None:
+            response.result.ctx_usage = disagg_params.ctx_usage
+
     def _maybe_prepend_logprobs_and_logits(self, req, beam_width):
         """Prepend logprobs and generation logits for first_gen_tokens
         if transferred from prefill."""
@@ -4980,6 +4990,7 @@ def _emit_first_token_responses(self, prev_scheduled_requests):
             if response is None:
                 continue
             response.result.cached_tokens = request.cached_tokens
+            self._maybe_attach_ctx_usage(request, response)
             if logits_snapshot is not None:
                 response.result.generation_logits = logits_snapshot
             new_responses.append((request.py_request_id, response))
@@ -5067,6 +5078,7 @@ def _handle_responses(self, emit_first_iter: bool = True):
                 if response:
                     request_done = request.is_finished
                     response.result.cached_tokens = request.cached_tokens
+                    self._maybe_attach_ctx_usage(request, response)
                     response.result.per_pos_drafted = request.py_per_pos_drafted
                     response.result.per_pos_accepted = request.py_per_pos_accepted
                     new_responses.append((req_id, response))

tensorrt_llm/executor/result.py

@@ -498,6 +498,17 @@ def _handle_response(self,
             self.per_pos_accepted = getattr(response_result, 'per_pos_accepted',
                                             None)
             self.avg_decoded_tokens_per_iter = response_result.avg_decoded_tokens_per_iter
+            # Expose gen-first ctx usage so the postprocessor
+            # (_ctx_usage_from_outputs) can adopt the context-side accounting.
+            # ctx_usage only exists on the Python LlmResult wrapper; the raw C++
+            # bindings.executor.Result (non-disagg / benchmark path) does not
+            # have it, so fall back to None as with cached_tokens above.
+            ctx_usage = getattr(response_result, 'ctx_usage', None)
+            if ctx_usage is not None:
+                self._disaggregated_params = dataclasses.replace(
+                    self._disaggregated_params or DisaggregatedParams(),
+                    ctx_usage=ctx_usage,
+                )
             if context_phase_params is not None:
                 existing_disagg_params = self.disaggregated_params
                 # Use `replace` to preserve things like `mrope_position_ids_handle` and

tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml

@@ -12,9 +12,9 @@ context_servers:
   tensor_parallel_size: 1
   pipeline_parallel_size: 1
   kv_cache_config:
-    enable_block_reuse: False
+    enable_block_reuse: True
     free_gpu_memory_fraction: 0.2
-    enable_partial_reuse: False
+    enable_partial_reuse: True
 
   cache_transceiver_config:
     backend: DEFAULT
@@ -29,9 +29,9 @@ generation_servers:
   max_num_tokens: 4096
   max_seq_len: 4096
   kv_cache_config:
-    enable_block_reuse: False
+    enable_block_reuse: True
     free_gpu_memory_fraction: 0.2
-    enable_partial_reuse: False
+    enable_partial_reuse: True
   cache_transceiver_config:
     backend: DEFAULT
     transceiver_runtime: PYTHON

tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml

@@ -12,9 +12,9 @@ context_servers:
   tensor_parallel_size: 1
   pipeline_parallel_size: 4
   kv_cache_config:
-    enable_block_reuse: False
+    enable_block_reuse: True
     free_gpu_memory_fraction: 0.2
-    enable_partial_reuse: False
+    enable_partial_reuse: True
 
   cache_transceiver_config:
     backend: DEFAULT
@@ -29,9 +29,9 @@ generation_servers:
   max_num_tokens: 4096
   max_seq_len: 4096
   kv_cache_config:
-    enable_block_reuse: False
+    enable_block_reuse: True
     free_gpu_memory_fraction: 0.2
-    enable_partial_reuse: False
+    enable_partial_reuse: True
   cache_transceiver_config:
     backend: DEFAULT
     transceiver_runtime: PYTHON

tests/integration/test_lists/waives.txt

@@ -128,8 +128,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_two_mtp[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322)
 disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] SKIP (https://nvbugs/6245317)
-disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6223556)
-disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6223556)
 disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] SKIP (https://nvbugs/6266302)
 disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322)
 disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322)