[https://nvbugs/5961736][fix] Prebuild disagg ctx response to avoid ctx_request_id race (NVIDIA#12466)

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

Author: peihu-nv

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1717,7 +1717,8 @@ class GenericLlmRequest
 
     [[nodiscard]] bool isFinished() const noexcept
     {
-        return isGenerationCompleteState() || mState == LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS;
+        return isGenerationCompleteState() || mState == LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS
+            || isDisaggContextCompleteState();
     }
 
     /// Returns true if finished_reason is length for all beams

cpp/tensorrt_llm/batch_manager/llmRequest.cpp

@@ -89,7 +89,7 @@ std::optional<executor::Result> LlmRequest::createResult(bool useFastLogits, int
 
     auto const maxNbTokens = getMaxBeamNumTokens();
 
-    if (isDisaggContextTransmissionState() && isContextOnlyRequest())
+    if ((isDisaggContextTransmissionState() || isDisaggContextCompleteState()) && isContextOnlyRequest())
     {
         auto const reqBeamWidth = mSamplingConfig.beamWidth;
         std::vector<TokenIdType> firstGenTokens;

cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp

@@ -751,6 +751,44 @@ TEST_P(ParamTest, createResponse)
     }
 }
 
+// Regression test for nvbug/5961736: createResult() must produce a valid
+// response with contextPhaseParams when the request is in
+// kDISAGG_CONTEXT_COMPLETE, not just kDISAGG_CONTEXT_TRANS_IN_PROGRESS.
+// Without the fix, createResult() returns nullopt for CONTEXT_COMPLETE,
+// causing ctx_request_id=None in the disaggregated serving response.
+TEST_F(LlmRequestTest, createResultDisaggContextComplete)
+{
+    VecTokens inputTokens{1, 2, 3, 4, 5};
+    SizeType32 maxNewTokens{10};
+    texec::IdType requestId{42};
+
+    // Build an executor::Request and configure it as context-only with ContextPhaseParams.
+    texec::Request execReq(inputTokens, maxNewTokens);
+    execReq.setRequestType(texec::RequestType::REQUEST_TYPE_CONTEXT_ONLY);
+    texec::ContextPhaseParams ctxParams({100}, requestId, static_cast<void*>(nullptr), std::nullopt);
+    execReq.setContextPhaseParams(std::move(ctxParams));
+
+    tb::LlmRequest llmReq(requestId, execReq);
+    EXPECT_TRUE(llmReq.isContextOnlyRequest());
+
+    // Add a generated token (required by createResult's firstGenTokens extraction).
+    llmReq.addNewTokens(VecTokens{42});
+
+    // Verify isFinished() covers DISAGG_CONTEXT_COMPLETE.
+    llmReq.setState(tb::LlmRequestState::kDISAGG_CONTEXT_COMPLETE);
+    EXPECT_TRUE(llmReq.isFinished());
+
+    // This is the regression case — without the fix, createResult() returns nullopt
+    // because DISAGG_CONTEXT_COMPLETE was not handled by createResult's early guard
+    // or its context-phase branch.
+    auto response = llmReq.createResult(/*useFastLogits=*/false, /*mpiWorldRank=*/0);
+    ASSERT_TRUE(response.has_value()) << "createResult() must not return nullopt for DISAGG_CONTEXT_COMPLETE";
+    EXPECT_TRUE(response->contextPhaseParams.has_value())
+        << "contextPhaseParams must be populated for context-only DISAGG_CONTEXT_COMPLETE requests";
+    EXPECT_EQ(response->contextPhaseParams->getReqId(), requestId);
+    EXPECT_TRUE(response->isSequenceFinal);
+}
+
 INSTANTIATE_TEST_SUITE_P(LlmRequestTest, ParamTest,
     testing::Combine(
         // TODO: Support and add coverage for streamLLM

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3530,6 +3530,9 @@ def _handle_first_token_response(self, scheduled_batch):
     def _handle_responses(self):
         new_responses = []
         requests_to_terminate = []
+        # Requests terminated by _check_disagg_ctx_cache_transfer_status (DISAGG_CONTEXT_COMPLETE);
+        # included in the return value for stats but not re-terminated here.
+        requests_finished_by_transfer = []
         new_active_requests = []
         logger.debug(
             f'------before _handle_responses, rank = {self.dist.rank}, output = {self.active_requests}'
@@ -3618,11 +3621,18 @@ def _handle_responses(self):
                 # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1,
                 # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse
                 # path is fixed.
-                if self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa and self.dist.pp_size == 1:
+                force_terminate_for_partial_reuse = (
+                    self.enable_partial_reuse_for_disagg
+                    and not self.kv_cache_manager.is_vswa
+                    and self.dist.pp_size == 1)
+                if request.is_disagg_context_complete_state:
+                    # Already terminated by _check_disagg_ctx_cache_transfer_status;
+                    # track for stats only to avoid double-free (nvbug/5961736).
+                    requests_finished_by_transfer.append(request)
+                elif force_terminate_for_partial_reuse:
+                    requests_to_terminate.append(request)
+                elif not request.is_disagg_context_transmission_state:
                     requests_to_terminate.append(request)
-                else:
-                    if not request.is_disagg_context_transmission_state:
-                        requests_to_terminate.append(request)
             else:
                 new_active_requests.append(request)
 
@@ -3632,7 +3642,7 @@ def _handle_responses(self):
         self._enqueue_responses(new_responses)
         for request in requests_to_terminate:
             self._terminate_request(request)
-        return requests_to_terminate
+        return requests_to_terminate + requests_finished_by_transfer
 
     def _await_any_response(self,
                             timeout: Optional[float] = None

tests/unittest/_torch/executor/test_py_executor.py

@@ -178,3 +178,70 @@ def test_getter_methods(mock_executor):
     assert mock_executor.get_expected_num_active_requests() == 5
     assert mock_executor._get_new_active_requests_queue_latency() == 10.5
     assert mock_executor.get_waiting_queue_size() == 1
+
+
+def _classify_termination(request, enable_partial_reuse_for_disagg, is_vswa, pp_size):
+    """Reproduce the termination logic from _handle_responses (py_executor.py).
+
+    Returns:
+        "terminate" | "stats_only" | "skip"
+    """
+    force_terminate_for_partial_reuse = (
+        enable_partial_reuse_for_disagg and not is_vswa and pp_size == 1
+    )
+    if request.is_disagg_context_complete_state:
+        return "stats_only"
+    elif force_terminate_for_partial_reuse:
+        return "terminate"
+    elif not request.is_disagg_context_transmission_state:
+        return "terminate"
+    return "skip"
+
+
+def _make_request(complete_state, transmission_state):
+    req = Mock()
+    req.is_disagg_context_complete_state = complete_state
+    req.is_disagg_context_transmission_state = transmission_state
+    return req
+
+
+class TestDisaggTerminationGuard:
+    """Verify _handle_responses does not double-terminate DISAGG_CONTEXT_COMPLETE
+    requests that were already cleaned up by _check_disagg_ctx_cache_transfer_status
+    (nvbug/5961736)."""
+
+    def test_normal_path_skips_context_complete(self):
+        """Without partial reuse, CONTEXT_COMPLETE goes to stats only."""
+        req = _make_request(complete_state=True, transmission_state=False)
+        assert _classify_termination(req, False, False, 1) == "stats_only"
+
+    def test_normal_path_skips_transmission_in_progress(self):
+        """Without partial reuse, TRANS_IN_PROGRESS is skipped (still in flight)."""
+        req = _make_request(complete_state=False, transmission_state=True)
+        assert _classify_termination(req, False, False, 1) == "skip"
+
+    def test_normal_path_terminates_regular_request(self):
+        """Without partial reuse, a normal finished request is terminated."""
+        req = _make_request(complete_state=False, transmission_state=False)
+        assert _classify_termination(req, False, False, 1) == "terminate"
+
+    def test_partial_reuse_terminates_non_complete(self):
+        """With partial reuse, non-CONTEXT_COMPLETE requests are terminated."""
+        for complete, transmission in [(False, True), (False, False)]:
+            req = _make_request(complete, transmission)
+            assert _classify_termination(req, True, False, 1) == "terminate"
+
+    def test_partial_reuse_skips_context_complete(self):
+        """With partial reuse, CONTEXT_COMPLETE still goes to stats only."""
+        req = _make_request(complete_state=True, transmission_state=False)
+        assert _classify_termination(req, True, False, 1) == "stats_only"
+
+    def test_partial_reuse_disabled_by_vswa(self):
+        """VSWA disables partial reuse path, falling back to normal logic."""
+        req = _make_request(complete_state=True, transmission_state=False)
+        assert _classify_termination(req, True, True, 1) == "stats_only"
+
+    def test_partial_reuse_disabled_by_pp(self):
+        """PP > 1 disables partial reuse path, falling back to normal logic."""
+        req = _make_request(complete_state=True, transmission_state=False)
+        assert _classify_termination(req, True, False, 2) == "stats_only"