[TRTLLM-17125][fix] Address PR review on agg Helix preemption

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1845,7 +1845,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
             if isinstance(self.event_manager, StreamingKVCacheEventManager)
             else self.event_manager
         )
-        if not has_host_cache_tier:
+        # Aggregated helix runs GPU-only but its CP ranks still need a common
+        # verdict: a rank that fails to allocate would otherwise leave its peers
+        # waiting in the first collective.
+        sync_init_status = has_host_cache_tier or (self._has_cp_helix and not is_disagg)
+        if not sync_init_status:
             candidate = KVCacheManagerPy(
                 config,
                 event_manager=event_sink,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -880,11 +880,13 @@ def __init__(
                                          "0") == "1"
         self._check_helix_schedule = (check_requested
                                       and self.dist.mapping.has_cp_helix()
+                                      and self.dist.mapping.pp_size == 1
                                       and self.kv_cache_transceiver is None)
         if check_requested and not self._check_helix_schedule:
             logger.warning(
                 "TLLM_HELIX_AGG_CHECK_SCHEDULE=1 is ignored: the schedule check "
-                "only runs with aggregated helix context parallelism.")
+                "only runs with aggregated helix context parallelism "
+                "and pipeline parallelism disabled.")
         self.sample_stream = torch.cuda.Stream()
         self.finish_sample_event = torch.cuda.Event()
         # Set of request IDs that are currently in flight across all micro batches

tests/integration/test_lists/test-db/l0_a10.yml

@@ -42,6 +42,7 @@ l0_a10:
   - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py
   - unittest/_torch/executor/kv_cache/sharing
   - unittest/_torch/executor/kv_cache/test_kv_cache_v2_branch_snapshot.py
+  - unittest/_torch/executor/kv_cache/test_helix_agg_preemption.py
   - unittest/_torch/executor/test_kv_cache_layout.py
   - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py
   - unittest/_torch/executor/test_error_classification.py

tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py

@@ -53,7 +53,7 @@
     MTPDecodingConfig,
     PARDDecodingConfig,
 )
-from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.mapping import CpType, Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     BAD_PAGE_INDEX,
     DEFAULT_BEAM_INDEX,
@@ -1191,6 +1191,38 @@ def test_peer_init_failure_discards_local_candidate() -> None:
     local_impl.shutdown.assert_called_once_with()
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("peer_failed", [False, True])
+def test_agg_helix_gpu_only_init_syncs_status_across_ranks(peer_failed: bool) -> None:
+    """Aggregated helix has no host tier but must still agree on init success.
+
+    Without the sync, a rank that failed to allocate raises alone and its peers
+    carry on into collectives that never complete.
+    """
+    local_impl = Mock()
+    helix_mapping = Mapping(
+        world_size=2,
+        rank=0,
+        cp_size=2,
+        cp_config={"cp_type": CpType.HELIX},
+    )
+    module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2"
+    status = _KVCacheManagerInitStatus.ABORT if peer_failed else _KVCacheManagerInitStatus.OK
+
+    with patch(f"{module}._sync_kv_cache_manager_init_status", return_value=status) as sync:
+        config = KvCacheConfig(
+            max_gpu_total_bytes=16 << 20, host_cache_size=0, enable_block_reuse=False
+        )
+        if peer_failed:
+            with pytest.raises(RuntimeError, match="failed on another rank"):
+                _make_manager_for_cache_tier_test(config, [local_impl], mapping=helix_mapping)
+            local_impl.shutdown.assert_called_once_with()
+        else:
+            _make_manager_for_cache_tier_test(config, [local_impl], mapping=helix_mapping)
+
+    sync.assert_called_once_with(_KVCacheManagerInitStatus.OK, helix_mapping)
+
+
 @pytest.mark.parametrize("add_secondary_gpu_tier", [False, True])
 def test_host_init_failure_propagates_without_rebuilding(add_secondary_gpu_tier: bool) -> None:
     """A configured host tier is honoured or the manager does not come up.