disagg-specific check

Signed-off-by: raayandhar <[email protected]>

Author: raayandhar

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -874,8 +874,8 @@ def _executor_loop_pp(self):
 
                 if self.kv_cache_transceiver and self.ctx_in_transmission_requests:
                     self._terminate_ctx_finished_requests()
-                
-                if self.dist.pp_size > 1 and self.enable_kv_cache_reuse:
+
+                if self.dist.pp_size > 1 and self.enable_kv_cache_reuse and self.kv_cache_transceiver:
                     self._sync_termination(prev_microbatch_id)
 
                 # march forward in microbatch slots
@@ -1598,7 +1598,7 @@ def _handle_errors(self,
         self._enqueue_responses(error_responses.items())
 
     def _terminate_request(self, request: LlmRequest):
-        if self.dist.pp_size > 1 and self.enable_kv_cache_reuse:
+        if self.dist.pp_size > 1 and self.enable_kv_cache_reuse and self.kv_cache_transceiver:
             # If pp_size > 1 and enable_kv_cache_reuse, we need to sync termination across PP ranks
             # otherwise, different ranks may have different KV cache blocks and a request may
             # have different PrepopulatedPromptLen which leads to a NCCL hang.
@@ -1612,7 +1612,7 @@ def _terminate_request(self, request: LlmRequest):
                 state['ready_to_terminate'].add(self.dist.rank)
         else:
             self._free_resources_for_request(request)
-    
+
     def _free_resources_for_request(self, request: LlmRequest):
         logger.debug(f"free resources for request {request.py_request_id}")
         self.resource_manager.free_resources(request)
@@ -1773,17 +1773,18 @@ def _sync_termination(self, microbatch_id: int):
             src=self.dist.prev_pp_rank,
             tag=microbatch_id,
         )
-        logger.debug(f"received remote state for microbatch {microbatch_id}, prev pp rank: {self.dist.prev_pp_rank} state {remote_state}")
+        logger.debug(
+            f"received remote state for microbatch {microbatch_id}, prev pp rank: {self.dist.prev_pp_rank} state {remote_state}"
+        )
 
         if remote_state:
             for req_id, state in remote_state.items():
                 local = self.pending_termination.get(req_id)
                 if local is None:
                     self.pending_termination[req_id] = {
-                        'ready_to_terminate':
-                        state.get('ready_to_terminate', set()),
-                        'terminated':
-                        state.get('terminated', set()),
+                        'ready_to_terminate': state.get('ready_to_terminate',
+                                                        set()),
+                        'terminated': state.get('terminated', set()),
                     }
                 else:
                     for key in ('ready_to_terminate', 'terminated'):