[#12116][fix] prevent KVBM hang when requests abort during KV cache transfer

tensorrt_llm/_torch/pyexecutor/kv_cache_connector.py

@@ -519,6 +519,23 @@ def request_finished(self, req: LlmRequest,
 
         return saving_async
 
+    def abort_request(self, req: LlmRequest) -> None:
+        """Remove a request from all async tracking structures.
+
+        Must be called when a request is cancelled or fails to prevent
+        worker.get_finished() from blocking on transfers that will never complete.
+
+        Args:
+            req: The request to abort.
+        """
+        req_id = req.request_id
+        for store in (self.new_async_requests, self.pending_async_requests,
+                      self.local_finished_async_requests):
+            store.loading.pop(req_id, None)
+            store.saving.pop(req_id, None)
+        self.scheduler_output_manager.requests.pop(req_id, None)
+        self.scheduler_output_manager.external_loads.pop(req_id, None)
+
     def get_finished(self) -> List[LlmRequest]:
         """
         Process requests that have finished loading and saving.
@@ -533,9 +550,18 @@ def get_finished(self) -> List[LlmRequest]:
         self.pending_async_requests.add_from(self.new_async_requests)
 
         # Pass these newly finished requests into get_finished, and get the list of requests that have finished saving and loading.
-        (finished_saving,
-         finished_loading) = self.worker.get_finished(finished_gen_req_ids,
-                                                      started_loading_req_ids)
+        # Wrap in try/except so that a worker failure still reaches the mpi_allgather barrier,
+        # preventing a collective stall where other ranks wait indefinitely for this rank.
+        try:
+            (finished_saving, finished_loading) = self.worker.get_finished(
+                finished_gen_req_ids, started_loading_req_ids)
+        except Exception as e:
+            logger.error(
+                f"KV connector worker get_finished() raised an exception: {e}. "
+                "Reporting no completions this iteration to unblock mpi_allgather."
+            )
+            finished_saving = []
+            finished_loading = []
 
         # Remove the requests from our pending list that have finished locally.
         new_local_finished_async_requests = self.pending_async_requests.extract_by_id(

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3175,6 +3175,8 @@ def _handle_errors(self,
             ]
         self._enqueue_responses(list(error_responses.items()))
         for request in failed_requests:
+            if self.kv_connector_manager is not None:
+                self.kv_connector_manager.abort_request(request)
             self._terminate_request(request)
 
     def _terminate_request(self, request: LlmRequest):
@@ -3202,6 +3204,10 @@ def _try_cancel_request(self, request) -> bool:
         Returns:
             bool: True if the request can be canceled (either successfully cancelled or doesn't need cancellation).
         """
+        if self.kv_connector_manager is not None:
+            self.kv_connector_manager.abort_request(request)
+            return True
+
         if self.kv_cache_transceiver is None:
             return True
 

tests/unittest/_torch/test_connector.py

@@ -175,6 +175,66 @@ def test():
     run_across_mpi(mpi_pool_executor, test, 2)
 
 
+def test_abort_request_cleans_async_state():
+    """Verify abort_request removes a request from all async tracking structures."""
+    worker = MagicMock()
+    scheduler = MagicMock()
+
+    manager = KvCacheConnectorManager(worker, scheduler=scheduler)
+
+    req = MagicMock()
+    req.request_id = 42
+
+    # Manually populate every tracking structure as if a transfer was in flight.
+    for store in (manager.new_async_requests, manager.pending_async_requests,
+                  manager.local_finished_async_requests):
+        store.loading[42] = req
+        store.saving[42] = req
+    manager.scheduler_output_manager.requests[42] = req
+    manager.scheduler_output_manager.external_loads[42] = req
+
+    manager.abort_request(req)
+
+    for store in (manager.new_async_requests, manager.pending_async_requests,
+                  manager.local_finished_async_requests):
+        assert 42 not in store.loading
+        assert 42 not in store.saving
+    assert 42 not in manager.scheduler_output_manager.requests
+    assert 42 not in manager.scheduler_output_manager.external_loads
+
+
+@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
+@pytest.mark.threadleak(enabled=False)
+def test_connector_manager_get_finished_exception_unblocks_allgather(
+        mpi_pool_executor):
+    """Verify that a worker.get_finished() exception on rank 0 does not cause a collective stall."""
+
+    def test():
+        worker = MagicMock()
+
+        if mpi_rank() == 0:
+            scheduler = MagicMock()
+            scheduler.request_finished.return_value = True
+            # Simulate a Rust/KVBM panic on rank 0.
+            worker.get_finished.side_effect = RuntimeError(
+                "simulated KVBM worker failure")
+        else:
+            scheduler = None
+            worker.get_finished.return_value = ([], [])
+
+        manager = KvCacheConnectorManager(worker, scheduler=scheduler)
+
+        req = MagicMock()
+        req.request_id = 42
+        manager.request_finished(req, [])
+
+        # Should not raise, return empty, and both ranks must reach this point.
+        result = manager.get_finished()
+        assert result == []
+
+    run_across_mpi(mpi_pool_executor, test, 2)
+
+
 def test_scheduler_output_num_scheduled_tokens_with_mtp():
     """Test that num_scheduled_tokens is correctly set for MTP (multi-token prediction)."""
     NUM_DRAFT_TOKENS = 3