diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7291568fb70e..349f15e3f206 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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( diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index dc51d461a77d..bea5ae250a58 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/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,7 +174,7 @@ 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(), @@ -168,6 +182,44 @@ def build_cache_config( 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)],