[None][fix] Synchronize KV cache V2 host fallback across ranks

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -18,6 +18,7 @@
 import sys
 from collections import OrderedDict, defaultdict, deque
 from dataclasses import dataclass, field, replace
+from enum import IntEnum
 from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -300,6 +301,25 @@ def _sync_host_tier_quota(host_quota: int, mapping: Mapping) -> int:
     return host_quota
 
 
+class _KVCacheManagerInitStatus(IntEnum):
+    # The numeric order is part of the allreduce(MAX) protocol: more severe
+    # outcomes must have larger values.
+    KEEP_HOST = 0
+    USE_NO_HOST = 1
+    ABORT = 2
+
+
+def _sync_kv_cache_manager_init_status(
+    local_status: _KVCacheManagerInitStatus, mapping: Mapping
+) -> _KVCacheManagerInitStatus:
+    """Return the most severe initialization status across all ranks."""
+    if mapping.world_size > 1:
+        local_status = _KVCacheManagerInitStatus(
+            Distributed.get(mapping).allreduce(int(local_status), op=ReduceOp.MAX)
+        )
+    return local_status
+
+
 def _estimate_swa_cache_size(
     layer_sizes: Sequence[int],
     attention_windows: Sequence[Optional[int]],
@@ -1089,25 +1109,73 @@ def append_to_kv_heads_per_layer(
             isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers
         )
 
-        self.kv_cache_manager_py_config = config
+        candidate: Optional[KVCacheManagerPy] = None
+        if not has_host_cache_tier:
+            candidate = KVCacheManagerPy(config, event_manager=self.event_manager)
+        else:
+            init_error: Optional[Exception] = None
+            local_init_status = _KVCacheManagerInitStatus.KEEP_HOST
+            try:
+                candidate = KVCacheManagerPy(config, event_manager=self.event_manager)
+            except Exception as error:
+                if isinstance(error, (CuError, KVCacheOutOfMemoryError)):
+                    local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST
+                else:
+                    init_error = error.with_traceback(None)
+                    local_init_status = _KVCacheManagerInitStatus.ABORT
 
-        try:
-            self.impl = KVCacheManagerPy(config, event_manager=self.event_manager)
-        except (CuError, KVCacheOutOfMemoryError):
-            if has_host_cache_tier:
+            init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping)
+
+            if init_status == _KVCacheManagerInitStatus.ABORT:
+                if candidate is not None:
+                    candidate.shutdown()
+                if init_error is not None:
+                    raise init_error
+                raise RuntimeError("KV cache manager initialization failed on another rank")
+
+            if init_status == _KVCacheManagerInitStatus.USE_NO_HOST:
                 logger.warning(
-                    "Failed to initialize KV cache manager with host cache "
-                    "tier (cuMemHostRegister may have failed). "
-                    "Retrying without host cache tier."
+                    "At least one rank could not use the KV cache manager host tier "
+                    "(cuMemHostRegister may have failed). Rebuilding without the "
+                    "host cache tier on all ranks."
                 )
-                cache_tiers_without_host = [
-                    tier for tier in config.cache_tiers if not isinstance(tier, HostCacheTierConfig)
-                ]
-                config = replace(config, cache_tiers=cache_tiers_without_host)
-                self.kv_cache_manager_py_config = config
-                self.impl = KVCacheManagerPy(config, event_manager=self.event_manager)
-            else:
-                raise
+                fallback_error: Optional[Exception] = None
+                try:
+                    if candidate is not None:
+                        candidate.shutdown()
+                    candidate = None
+                    config = replace(
+                        config,
+                        cache_tiers=[
+                            tier
+                            for tier in config.cache_tiers
+                            if not isinstance(tier, HostCacheTierConfig)
+                        ],
+                    )
+                    candidate = KVCacheManagerPy(config, event_manager=self.event_manager)
+                except Exception as error:
+                    fallback_error = error.with_traceback(None)
+
+                local_fallback_status = (
+                    _KVCacheManagerInitStatus.USE_NO_HOST
+                    if fallback_error is None
+                    else _KVCacheManagerInitStatus.ABORT
+                )
+                fallback_status = _sync_kv_cache_manager_init_status(local_fallback_status, mapping)
+
+                if fallback_status == _KVCacheManagerInitStatus.ABORT:
+                    if candidate is not None:
+                        candidate.shutdown()
+                    if fallback_error is not None:
+                        raise fallback_error
+                    raise RuntimeError(
+                        "KV cache manager initialization without the host cache tier "
+                        "failed on another rank"
+                    )
+
+        assert candidate is not None
+        self.kv_cache_manager_py_config = config
+        self.impl = candidate
         self.can_evict = len(config.cache_tiers) > 1
         if self.event_manager is not None:
             self.event_manager.set_layer_group_window_sizes(

tests/unittest/_torch/executor/test_kv_cache_manager_v2.py

@@ -21,9 +21,16 @@
 import pytest
 import torch
 
-from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2
+from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
+    BlockReusePolicy,
+    KVCacheManagerV2,
+    _KVCacheManagerInitStatus,
+    _sync_kv_cache_manager_init_status,
+)
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm.bindings import DataType
+from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE
 from tensorrt_llm.bindings.internal.batch_manager import CacheType
 from tensorrt_llm.conversation_params import ConversationParams
 from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, KvCacheConfig
@@ -107,8 +114,11 @@ def _make_manager_for_cache_tier_test(
     impl_side_effect: list[object],
     *,
     add_secondary_gpu_tier: bool = False,
+    mapping: Mapping | None = None,
 ) -> tuple[KVCacheManagerV2, Mock]:
     impl_constructor = Mock(side_effect=impl_side_effect)
+    if mapping is None:
+        mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1)
 
     def build_base_config(
         self: KVCacheManagerV2,
@@ -135,15 +145,19 @@ def build_cache_config(
             )
         return config
 
-    fake_impl = impl_side_effect[-1]
-    assert not isinstance(fake_impl, BaseException)
-    fake_impl.layer_grouping = [[0]]
-    fake_impl.pool_group_descs = []
-    fake_impl.get_layer_group_id.side_effect = lambda _: 0
+    fake_impl = next(
+        (item for item in reversed(impl_side_effect) if not isinstance(item, BaseException)),
+        None,
+    )
+    if fake_impl is not None:
+        fake_impl.layer_grouping = [[0]]
+        fake_impl.pool_group_descs = []
+        fake_impl.get_layer_group_id.side_effect = lambda _: 0
 
     module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2"
     with (
         patch(f"{module}.CuError", _CacheTierInitError),
+        patch(f"{module}.IndexMapper"),
         patch(f"{module}.KVCacheManagerPy", impl_constructor),
         patch.object(KVCacheManagerV2, "_build_base_config", build_base_config),
         patch.object(KVCacheManagerV2, "_build_cache_config", build_cache_config),
@@ -160,14 +174,52 @@ def build_cache_config(
             tokens_per_block=TOKENS_PER_BLOCK,
             max_seq_len=MAX_SEQ_LEN,
             max_batch_size=1,
-            mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+            mapping=mapping,
             dtype=DataType.HALF,
             vocab_size=16,
             execution_stream=Mock(),
         )
     return manager, impl_constructor
 
 
+def _multi_rank_host_fallback_consensus_worker() -> tuple[int, int, int, bool]:
+    """Exercise the real world collective from an MPI worker."""
+    from tensorrt_llm._utils import mpi_rank, mpi_world_size
+
+    rank = mpi_rank()
+    world_size = mpi_world_size()
+    initial_impl = Mock()
+    fallback_impl = Mock()
+    impl_side_effect: list[object] = (
+        [initial_impl, fallback_impl]
+        if rank == 0
+        else [_CacheTierInitError("rank-local host tier failure"), fallback_impl]
+    )
+
+    manager, impl_constructor = _make_manager_for_cache_tier_test(
+        KvCacheConfig(
+            max_gpu_total_bytes=16 << 20,
+            host_cache_size=16 << 20,
+        ),
+        impl_side_effect,
+        mapping=Mapping(
+            world_size=world_size,
+            rank=rank,
+            tp_size=world_size,
+        ),
+    )
+
+    return (
+        rank,
+        impl_constructor.call_count,
+        initial_impl.shutdown.call_count,
+        any(
+            isinstance(tier, HostCacheTierConfig)
+            for tier in manager.kv_cache_manager_py_config.cache_tiers
+        ),
+    )
+
+
 @pytest.mark.parametrize(
     ("enable_block_reuse", "block_reuse_policy", "is_draft", "commit_min_snapshot"),
     [
@@ -311,6 +363,91 @@ def test_disk_init_failure_does_not_use_host_fallback(tmp_path) -> None:
         )
 
 
+def test_kv_cache_manager_init_status_sync_uses_world_max() -> None:
+    mapping = SimpleNamespace(world_size=2)
+    dist = Mock()
+    dist.allreduce.return_value = int(_KVCacheManagerInitStatus.USE_NO_HOST)
+
+    with patch.object(Distributed, "get", return_value=dist):
+        status = _sync_kv_cache_manager_init_status(_KVCacheManagerInitStatus.KEEP_HOST, mapping)
+
+    assert status == _KVCacheManagerInitStatus.USE_NO_HOST
+    dist.allreduce.assert_called_once_with(
+        int(_KVCacheManagerInitStatus.KEEP_HOST), op=ReduceOp.MAX
+    )
+
+
+@pytest.mark.cpu_only
+@pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device (MPI) build required")
+def test_world_ranks_converge_on_hostless_fallback() -> None:
+    from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
+
+    session = MpiPoolSession(n_workers=2)
+    try:
+        results = session.submit_sync(_multi_rank_host_fallback_consensus_worker)
+    finally:
+        session.shutdown()
+
+    assert sorted(results) == [(0, 2, 1, False), (1, 2, 0, False)]
+
+
+def test_local_fallback_failure_is_shared_before_raising() -> None:
+    module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2"
+
+    with (
+        patch(
+            f"{module}._sync_kv_cache_manager_init_status",
+            side_effect=[
+                _KVCacheManagerInitStatus.USE_NO_HOST,
+                _KVCacheManagerInitStatus.ABORT,
+            ],
+        ) as sync_status,
+        pytest.raises(RuntimeError, match="fallback init failed"),
+    ):
+        _make_manager_for_cache_tier_test(
+            KvCacheConfig(
+                max_gpu_total_bytes=16 << 20,
+                host_cache_size=16 << 20,
+            ),
+            [
+                _CacheTierInitError("host tier init failed"),
+                RuntimeError("fallback init failed"),
+            ],
+        )
+
+    assert [call.args[0] for call in sync_status.call_args_list] == [
+        _KVCacheManagerInitStatus.USE_NO_HOST,
+        _KVCacheManagerInitStatus.ABORT,
+    ]
+
+
+def test_peer_fallback_failure_discards_local_candidate() -> None:
+    initial_impl = Mock()
+    fallback_impl = Mock()
+    module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2"
+
+    with (
+        patch(
+            f"{module}._sync_kv_cache_manager_init_status",
+            side_effect=[
+                _KVCacheManagerInitStatus.USE_NO_HOST,
+                _KVCacheManagerInitStatus.ABORT,
+            ],
+        ),
+        pytest.raises(RuntimeError, match="failed on another rank"),
+    ):
+        _make_manager_for_cache_tier_test(
+            KvCacheConfig(
+                max_gpu_total_bytes=16 << 20,
+                host_cache_size=16 << 20,
+            ),
+            [initial_impl, fallback_impl],
+        )
+
+    initial_impl.shutdown.assert_called_once_with()
+    fallback_impl.shutdown.assert_called_once_with()
+
+
 @pytest.mark.parametrize(
     ("add_secondary_gpu_tier", "expected_can_evict"),
     [(False, False), (True, True)],