From cbed9d96c66e091205fc4631eb99acf5eff4fd51 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 01:44:18 -0700 Subject: [PATCH 1/4] [None][fix] Synchronize KV cache V2 host fallback across ranks Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 138 ++++++- .../executor/test_kv_cache_manager_v2.py | 364 +++++++++++++++++- 2 files changed, 490 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7291568fb70e..375aae1ef3f9 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,35 @@ 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 _shutdown_kv_cache_manager_candidate(candidate: Optional[KVCacheManagerPy]) -> None: + """Best-effort cleanup for an uncommitted KV cache manager candidate.""" + if candidate is None: + return + try: + candidate.shutdown() + except Exception as error: + logger.error("Failed to clean up an uncommitted KV cache manager:", error) + + def _estimate_swa_cache_size( layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]], @@ -1089,25 +1119,113 @@ def append_to_kv_heads_per_layer( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) - self.kv_cache_manager_py_config = config + candidate = None + init_error: Optional[Exception] = None + local_init_status = _KVCacheManagerInitStatus.KEEP_HOST + try: + candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + if not has_host_cache_tier: + local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST + except (CuError, KVCacheOutOfMemoryError) as error: + if has_host_cache_tier: + # Do not retain the traceback of a failed Python HostMem + # constructor: its frames can keep a partially allocated mmap + # alive while the hostless fallback is being built. + local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST + else: + init_error = error.with_traceback(None) + local_init_status = _KVCacheManagerInitStatus.ABORT + except Exception as error: + # Every rank must reach the consensus below even when only one + # sees an unexpected constructor error. The error is re-raised + # after peers have observed the fatal status. Drop its traceback + # so a partial constructor is not retained across the collective. + init_error = error.with_traceback(None) + local_init_status = _KVCacheManagerInitStatus.ABORT try: - self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) - except (CuError, KVCacheOutOfMemoryError): + init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) + except Exception: + _shutdown_kv_cache_manager_candidate(candidate) + raise + + if init_status == _KVCacheManagerInitStatus.ABORT: + _shutdown_kv_cache_manager_candidate(candidate) + if local_init_status == _KVCacheManagerInitStatus.ABORT: + assert 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: if has_host_cache_tier: 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." ) + fallback_config = None + fallback_error: Optional[Exception] = None + try: 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: + fallback_config = replace(config, cache_tiers=cache_tiers_without_host) + except Exception as error: + # Config derivation is rank-local too, so it must participate + # in the same completion vote as teardown and reconstruction. + fallback_error = error.with_traceback(None) + + stale_candidate = None + if has_host_cache_tier and candidate is not None: + stale_candidate = candidate + candidate = None + try: + stale_candidate.shutdown() + except Exception as error: + fallback_error = error.with_traceback(None) + else: + stale_candidate = None + + if fallback_error is None and candidate is None: + assert fallback_config is not None + try: + candidate = KVCacheManagerPy(fallback_config, event_manager=self.event_manager) + except Exception as error: + # A fallback-construction failure is fatal because there + # is no lower cache tier to retry. + fallback_error = error.with_traceback(None) + + local_fallback_status = ( + _KVCacheManagerInitStatus.USE_NO_HOST + if fallback_error is None + else _KVCacheManagerInitStatus.ABORT + ) + + try: + fallback_status = _sync_kv_cache_manager_init_status(local_fallback_status, mapping) + except Exception: + _shutdown_kv_cache_manager_candidate(candidate) + _shutdown_kv_cache_manager_candidate(stale_candidate) raise + + if fallback_status == _KVCacheManagerInitStatus.ABORT: + _shutdown_kv_cache_manager_candidate(candidate) + _shutdown_kv_cache_manager_candidate(stale_candidate) + 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 stale_candidate is None + assert candidate is not None + assert fallback_config is not None + config = fallback_config + + 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..cc17b4a10915 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -21,10 +21,18 @@ 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.bindings.internal.batch_manager import kv_cache_manager_v2 as kv2_bindings from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -107,8 +115,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, @@ -160,7 +171,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 +179,48 @@ def build_cache_config( return manager, impl_constructor +def _host_fallback_consensus_worker() -> dict[str, int | bool]: + """Exercise the real world collective from an attention-DP worker.""" + from tensorrt_llm._utils import mpi_rank, mpi_world_size + + rank = mpi_rank() + world_size = mpi_world_size() + fallback_impl = Mock() + initial_impl = Mock() if rank == 0 else None + impl_side_effect = ( + [initial_impl, fallback_impl] + if initial_impl is not None + 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, + enable_attention_dp=True, + ), + ) + + return { + "rank": rank, + "constructor_calls": impl_constructor.call_count, + "initial_shutdown_calls": ( + initial_impl.shutdown.call_count if initial_impl is not None else 0 + ), + "final_has_host": any( + isinstance(tier, HostCacheTierConfig) + for tier in manager.kv_cache_manager_py_config.cache_tiers + ), + "can_evict": manager.can_evict, + } + + @pytest.mark.parametrize( ("enable_block_reuse", "block_reuse_policy", "is_draft", "commit_min_snapshot"), [ @@ -311,6 +364,313 @@ 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 + ) + + +def test_kv_cache_manager_init_status_sync_is_noop_for_single_rank() -> None: + mapping = SimpleNamespace(world_size=1) + + with patch.object(Distributed, "get") as get_dist: + status = _sync_kv_cache_manager_init_status(_KVCacheManagerInitStatus.KEEP_HOST, mapping) + + assert status == _KVCacheManagerInitStatus.KEEP_HOST + get_dist.assert_not_called() + + +@pytest.mark.parametrize("derived_name", ["HostOOMError", "DiskOOMError", "CuOOMError"]) +def test_cpp_oom_exception_hierarchy_matches_python_backend(derived_name: str) -> None: + assert issubclass(getattr(kv2_bindings, derived_name), kv2_bindings.OutOfMemoryError) + + +@pytest.mark.cpu_only +@pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device (MPI) build required") +def test_attention_dp_ranks_converge_on_hostless_fallback() -> None: + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + session = MpiPoolSession(n_workers=2) + try: + results = session.submit_sync(_host_fallback_consensus_worker) + finally: + session.shutdown() + + results = sorted(results, key=lambda result: result["rank"]) + assert [result["constructor_calls"] for result in results] == [2, 2] + assert [result["initial_shutdown_calls"] for result in results] == [1, 0] + assert not any(result["final_has_host"] for result in results) + assert not any(result["can_evict"] for result in results) + + +def test_all_ranks_keep_successful_host_candidate() -> None: + initial_impl = Mock() + module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + + with patch( + f"{module}._sync_kv_cache_manager_init_status", + return_value=_KVCacheManagerInitStatus.KEEP_HOST, + ) as sync_status: + manager, impl_constructor = _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + ), + [initial_impl], + ) + + sync_status.assert_called_once() + initial_impl.shutdown.assert_not_called() + assert manager.impl is initial_impl + assert impl_constructor.call_count == 1 + assert any( + isinstance(tier, HostCacheTierConfig) + for tier in manager.kv_cache_manager_py_config.cache_tiers + ) + assert manager.can_evict + + +def test_peer_host_init_failure_rebuilds_successful_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.USE_NO_HOST, + ], + ) as sync_status: + manager, impl_constructor = _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + ), + [initial_impl, fallback_impl], + ) + + assert sync_status.call_count == 2 + initial_impl.shutdown.assert_called_once_with() + assert manager.impl is fallback_impl + assert impl_constructor.call_count == 2 + initial_tiers = impl_constructor.call_args_list[0].args[0].cache_tiers + fallback_tiers = impl_constructor.call_args_list[1].args[0].cache_tiers + assert any(isinstance(tier, HostCacheTierConfig) for tier in initial_tiers) + assert all(not isinstance(tier, HostCacheTierConfig) for tier in fallback_tiers) + assert all( + not isinstance(tier, HostCacheTierConfig) + for tier in manager.kv_cache_manager_py_config.cache_tiers + ) + assert not manager.can_evict + assert [call.args[0] for call in sync_status.call_args_list] == [ + _KVCacheManagerInitStatus.KEEP_HOST, + _KVCacheManagerInitStatus.USE_NO_HOST, + ] + + +def test_no_host_candidate_joins_peer_fallback_consensus_without_rebuild() -> None: + no_host_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.USE_NO_HOST, + ], + ) as sync_status: + manager, impl_constructor = _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + ), + [no_host_impl], + ) + + assert [call.args[0] for call in sync_status.call_args_list] == [ + _KVCacheManagerInitStatus.USE_NO_HOST, + _KVCacheManagerInitStatus.USE_NO_HOST, + ] + assert impl_constructor.call_count == 1 + no_host_impl.shutdown.assert_not_called() + assert manager.impl is no_host_impl + + +def test_peer_fatal_init_failure_cleans_successful_local_candidate() -> None: + initial_impl = Mock() + module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + + with ( + patch( + f"{module}._sync_kv_cache_manager_init_status", + return_value=_KVCacheManagerInitStatus.ABORT, + ), + pytest.raises(RuntimeError, match="initialization failed on another rank"), + ): + _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + ), + [initial_impl], + ) + + initial_impl.shutdown.assert_called_once_with() + + +def test_peer_fallback_failure_cleans_successful_local_fallback() -> 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() + + +def test_local_fallback_failure_is_shared_before_raising() -> None: + fallback_error = RuntimeError("fallback init failed") + unused_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, + ], + ) 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"), + fallback_error, + unused_impl, + ], + ) + + assert sync_status.call_count == 2 + assert [call.args[0] for call in sync_status.call_args_list] == [ + _KVCacheManagerInitStatus.USE_NO_HOST, + _KVCacheManagerInitStatus.ABORT, + ] + unused_impl.shutdown.assert_not_called() + + +def test_candidate_shutdown_failure_aborts_before_fallback_build() -> None: + initial_impl = Mock() + initial_impl.shutdown.side_effect = RuntimeError("candidate shutdown failed") + unused_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, + ], + ) as sync_status, + pytest.raises(RuntimeError, match="candidate shutdown failed"), + ): + _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + ), + [initial_impl, unused_impl], + ) + + assert sync_status.call_count == 2 + assert [call.args[0] for call in sync_status.call_args_list] == [ + _KVCacheManagerInitStatus.KEEP_HOST, + _KVCacheManagerInitStatus.ABORT, + ] + assert initial_impl.shutdown.call_count == 2 + unused_impl.shutdown.assert_not_called() + + +def test_consensus_failure_cleans_local_candidate() -> None: + initial_impl = Mock() + module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" + + with ( + patch( + f"{module}._sync_kv_cache_manager_init_status", + side_effect=RuntimeError("consensus failed"), + ), + pytest.raises(RuntimeError, match="consensus failed"), + ): + _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + ), + [initial_impl], + ) + + initial_impl.shutdown.assert_called_once_with() + + +def test_fallback_consensus_failure_cleans_local_fallback() -> 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, + RuntimeError("fallback consensus failed"), + ], + ), + pytest.raises(RuntimeError, match="fallback consensus failed"), + ): + _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)], From 140f8f0275ca43dd365844bc5ed6d9d4ed974217 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 02:43:07 -0700 Subject: [PATCH 2/4] [None][refactor] Trim host fallback synchronization scope Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 146 ++++------ .../executor/test_kv_cache_manager_v2.py | 269 ++---------------- 2 files changed, 81 insertions(+), 334 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 375aae1ef3f9..d29fc2dcb763 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -327,7 +327,7 @@ def _shutdown_kv_cache_manager_candidate(candidate: Optional[KVCacheManagerPy]) try: candidate.shutdown() except Exception as error: - logger.error("Failed to clean up an uncommitted KV cache manager:", error) + logger.error(f"Failed to clean up an uncommitted KV cache manager: {error}") def _estimate_swa_cache_size( @@ -1119,109 +1119,79 @@ def append_to_kv_heads_per_layer( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) - candidate = None - init_error: Optional[Exception] = None - local_init_status = _KVCacheManagerInitStatus.KEEP_HOST - try: + candidate: Optional[KVCacheManagerPy] = None + if not has_host_cache_tier: candidate = KVCacheManagerPy(config, event_manager=self.event_manager) - if not has_host_cache_tier: - local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST - except (CuError, KVCacheOutOfMemoryError) as error: - if has_host_cache_tier: - # Do not retain the traceback of a failed Python HostMem - # constructor: its frames can keep a partially allocated mmap - # alive while the hostless fallback is being built. + else: + init_error: Optional[Exception] = None + try: + candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + local_init_status = _KVCacheManagerInitStatus.KEEP_HOST + except (CuError, KVCacheOutOfMemoryError): local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST - else: + except Exception as error: init_error = error.with_traceback(None) local_init_status = _KVCacheManagerInitStatus.ABORT - except Exception as error: - # Every rank must reach the consensus below even when only one - # sees an unexpected constructor error. The error is re-raised - # after peers have observed the fatal status. Drop its traceback - # so a partial constructor is not retained across the collective. - init_error = error.with_traceback(None) - local_init_status = _KVCacheManagerInitStatus.ABORT - try: - init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) - except Exception: - _shutdown_kv_cache_manager_candidate(candidate) - raise - - if init_status == _KVCacheManagerInitStatus.ABORT: - _shutdown_kv_cache_manager_candidate(candidate) - if local_init_status == _KVCacheManagerInitStatus.ABORT: - assert 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: - if has_host_cache_tier: + try: + init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) + except Exception: + _shutdown_kv_cache_manager_candidate(candidate) + raise + + if init_status == _KVCacheManagerInitStatus.ABORT: + _shutdown_kv_cache_manager_candidate(candidate) + 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( "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." ) - fallback_config = None - fallback_error: Optional[Exception] = None - try: - cache_tiers_without_host = [ - tier for tier in config.cache_tiers if not isinstance(tier, HostCacheTierConfig) - ] - fallback_config = replace(config, cache_tiers=cache_tiers_without_host) - except Exception as error: - # Config derivation is rank-local too, so it must participate - # in the same completion vote as teardown and reconstruction. - fallback_error = error.with_traceback(None) - - stale_candidate = None - if has_host_cache_tier and candidate is not None: - stale_candidate = candidate + _shutdown_kv_cache_manager_candidate(candidate) candidate = None + fallback_error: Optional[Exception] = None try: - stale_candidate.shutdown() + config = replace( + config, + cache_tiers=[ + tier + for tier in config.cache_tiers + if not isinstance(tier, HostCacheTierConfig) + ], + ) except Exception as error: fallback_error = error.with_traceback(None) else: - stale_candidate = None - - if fallback_error is None and candidate is None: - assert fallback_config is not None - try: - candidate = KVCacheManagerPy(fallback_config, event_manager=self.event_manager) - except Exception as error: - # A fallback-construction failure is fatal because there - # is no lower cache tier to retry. - fallback_error = error.with_traceback(None) - - local_fallback_status = ( - _KVCacheManagerInitStatus.USE_NO_HOST - if fallback_error is None - else _KVCacheManagerInitStatus.ABORT - ) - - try: - fallback_status = _sync_kv_cache_manager_init_status(local_fallback_status, mapping) - except Exception: - _shutdown_kv_cache_manager_candidate(candidate) - _shutdown_kv_cache_manager_candidate(stale_candidate) - raise - - if fallback_status == _KVCacheManagerInitStatus.ABORT: - _shutdown_kv_cache_manager_candidate(candidate) - _shutdown_kv_cache_manager_candidate(stale_candidate) - if fallback_error is not None: - raise fallback_error - raise RuntimeError( - "KV cache manager initialization without the host cache tier " - "failed on another rank" + try: + 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 ) - - assert stale_candidate is None - assert candidate is not None - assert fallback_config is not None - config = fallback_config + try: + fallback_status = _sync_kv_cache_manager_init_status( + local_fallback_status, mapping + ) + except Exception: + _shutdown_kv_cache_manager_candidate(candidate) + raise + + if fallback_status == _KVCacheManagerInitStatus.ABORT: + _shutdown_kv_cache_manager_candidate(candidate) + 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 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 cc17b4a10915..bc59b1afc77b 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -32,7 +32,6 @@ 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.bindings.internal.batch_manager import kv_cache_manager_v2 as kv2_bindings from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -146,11 +145,14 @@ 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 ( @@ -179,17 +181,17 @@ def build_cache_config( return manager, impl_constructor -def _host_fallback_consensus_worker() -> dict[str, int | bool]: +def _host_fallback_consensus_worker() -> tuple[int, int, int, bool]: """Exercise the real world collective from an attention-DP 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() - initial_impl = Mock() if rank == 0 else None - impl_side_effect = ( + impl_side_effect: list[object] = ( [initial_impl, fallback_impl] - if initial_impl is not None + if rank == 0 else [_CacheTierInitError("rank-local host tier failure"), fallback_impl] ) @@ -207,18 +209,15 @@ def _host_fallback_consensus_worker() -> dict[str, int | bool]: ), ) - return { - "rank": rank, - "constructor_calls": impl_constructor.call_count, - "initial_shutdown_calls": ( - initial_impl.shutdown.call_count if initial_impl is not None else 0 - ), - "final_has_host": any( + 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 ), - "can_evict": manager.can_evict, - } + ) @pytest.mark.parametrize( @@ -378,21 +377,6 @@ def test_kv_cache_manager_init_status_sync_uses_world_max() -> None: ) -def test_kv_cache_manager_init_status_sync_is_noop_for_single_rank() -> None: - mapping = SimpleNamespace(world_size=1) - - with patch.object(Distributed, "get") as get_dist: - status = _sync_kv_cache_manager_init_status(_KVCacheManagerInitStatus.KEEP_HOST, mapping) - - assert status == _KVCacheManagerInitStatus.KEEP_HOST - get_dist.assert_not_called() - - -@pytest.mark.parametrize("derived_name", ["HostOOMError", "DiskOOMError", "CuOOMError"]) -def test_cpp_oom_exception_hierarchy_matches_python_backend(derived_name: str) -> None: - assert issubclass(getattr(kv2_bindings, derived_name), kv2_bindings.OutOfMemoryError) - - @pytest.mark.cpu_only @pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device (MPI) build required") def test_attention_dp_ranks_converge_on_hostless_fallback() -> None: @@ -404,159 +388,10 @@ def test_attention_dp_ranks_converge_on_hostless_fallback() -> None: finally: session.shutdown() - results = sorted(results, key=lambda result: result["rank"]) - assert [result["constructor_calls"] for result in results] == [2, 2] - assert [result["initial_shutdown_calls"] for result in results] == [1, 0] - assert not any(result["final_has_host"] for result in results) - assert not any(result["can_evict"] for result in results) - - -def test_all_ranks_keep_successful_host_candidate() -> None: - initial_impl = Mock() - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" - - with patch( - f"{module}._sync_kv_cache_manager_init_status", - return_value=_KVCacheManagerInitStatus.KEEP_HOST, - ) as sync_status: - manager, impl_constructor = _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=16 << 20, - ), - [initial_impl], - ) - - sync_status.assert_called_once() - initial_impl.shutdown.assert_not_called() - assert manager.impl is initial_impl - assert impl_constructor.call_count == 1 - assert any( - isinstance(tier, HostCacheTierConfig) - for tier in manager.kv_cache_manager_py_config.cache_tiers - ) - assert manager.can_evict - - -def test_peer_host_init_failure_rebuilds_successful_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.USE_NO_HOST, - ], - ) as sync_status: - manager, impl_constructor = _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=16 << 20, - ), - [initial_impl, fallback_impl], - ) - - assert sync_status.call_count == 2 - initial_impl.shutdown.assert_called_once_with() - assert manager.impl is fallback_impl - assert impl_constructor.call_count == 2 - initial_tiers = impl_constructor.call_args_list[0].args[0].cache_tiers - fallback_tiers = impl_constructor.call_args_list[1].args[0].cache_tiers - assert any(isinstance(tier, HostCacheTierConfig) for tier in initial_tiers) - assert all(not isinstance(tier, HostCacheTierConfig) for tier in fallback_tiers) - assert all( - not isinstance(tier, HostCacheTierConfig) - for tier in manager.kv_cache_manager_py_config.cache_tiers - ) - assert not manager.can_evict - assert [call.args[0] for call in sync_status.call_args_list] == [ - _KVCacheManagerInitStatus.KEEP_HOST, - _KVCacheManagerInitStatus.USE_NO_HOST, - ] - - -def test_no_host_candidate_joins_peer_fallback_consensus_without_rebuild() -> None: - no_host_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.USE_NO_HOST, - ], - ) as sync_status: - manager, impl_constructor = _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=0, - ), - [no_host_impl], - ) - - assert [call.args[0] for call in sync_status.call_args_list] == [ - _KVCacheManagerInitStatus.USE_NO_HOST, - _KVCacheManagerInitStatus.USE_NO_HOST, - ] - assert impl_constructor.call_count == 1 - no_host_impl.shutdown.assert_not_called() - assert manager.impl is no_host_impl - - -def test_peer_fatal_init_failure_cleans_successful_local_candidate() -> None: - initial_impl = Mock() - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" - - with ( - patch( - f"{module}._sync_kv_cache_manager_init_status", - return_value=_KVCacheManagerInitStatus.ABORT, - ), - pytest.raises(RuntimeError, match="initialization failed on another rank"), - ): - _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=16 << 20, - ), - [initial_impl], - ) - - initial_impl.shutdown.assert_called_once_with() - - -def test_peer_fallback_failure_cleans_successful_local_fallback() -> 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() + assert sorted(results) == [(0, 2, 1, False), (1, 2, 0, False)] def test_local_fallback_failure_is_shared_before_raising() -> None: - fallback_error = RuntimeError("fallback init failed") - unused_impl = Mock() module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" with ( @@ -576,75 +411,17 @@ def test_local_fallback_failure_is_shared_before_raising() -> None: ), [ _CacheTierInitError("host tier init failed"), - fallback_error, - unused_impl, + RuntimeError("fallback init failed"), ], ) - assert sync_status.call_count == 2 assert [call.args[0] for call in sync_status.call_args_list] == [ _KVCacheManagerInitStatus.USE_NO_HOST, _KVCacheManagerInitStatus.ABORT, ] - unused_impl.shutdown.assert_not_called() -def test_candidate_shutdown_failure_aborts_before_fallback_build() -> None: - initial_impl = Mock() - initial_impl.shutdown.side_effect = RuntimeError("candidate shutdown failed") - unused_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, - ], - ) as sync_status, - pytest.raises(RuntimeError, match="candidate shutdown failed"), - ): - _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=16 << 20, - ), - [initial_impl, unused_impl], - ) - - assert sync_status.call_count == 2 - assert [call.args[0] for call in sync_status.call_args_list] == [ - _KVCacheManagerInitStatus.KEEP_HOST, - _KVCacheManagerInitStatus.ABORT, - ] - assert initial_impl.shutdown.call_count == 2 - unused_impl.shutdown.assert_not_called() - - -def test_consensus_failure_cleans_local_candidate() -> None: - initial_impl = Mock() - module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" - - with ( - patch( - f"{module}._sync_kv_cache_manager_init_status", - side_effect=RuntimeError("consensus failed"), - ), - pytest.raises(RuntimeError, match="consensus failed"), - ): - _make_manager_for_cache_tier_test( - KvCacheConfig( - max_gpu_total_bytes=16 << 20, - host_cache_size=16 << 20, - ), - [initial_impl], - ) - - initial_impl.shutdown.assert_called_once_with() - - -def test_fallback_consensus_failure_cleans_local_fallback() -> None: +def test_peer_fallback_failure_discards_local_candidate() -> None: initial_impl = Mock() fallback_impl = Mock() module = "tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2" @@ -654,10 +431,10 @@ def test_fallback_consensus_failure_cleans_local_fallback() -> None: f"{module}._sync_kv_cache_manager_init_status", side_effect=[ _KVCacheManagerInitStatus.USE_NO_HOST, - RuntimeError("fallback consensus failed"), + _KVCacheManagerInitStatus.ABORT, ], ), - pytest.raises(RuntimeError, match="fallback consensus failed"), + pytest.raises(RuntimeError, match="failed on another rank"), ): _make_manager_for_cache_tier_test( KvCacheConfig( From 1f07080f2521d319dff43e12e716a661614949e4 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 22 Aug 2026 07:43:31 -0700 Subject: [PATCH 3/4] [None][refactor] Simplify host fallback error handling Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 52 ++++++------------- 1 file changed, 16 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index d29fc2dcb763..349f15e3f206 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -320,16 +320,6 @@ def _sync_kv_cache_manager_init_status( return local_status -def _shutdown_kv_cache_manager_candidate(candidate: Optional[KVCacheManagerPy]) -> None: - """Best-effort cleanup for an uncommitted KV cache manager candidate.""" - if candidate is None: - return - try: - candidate.shutdown() - except Exception as error: - logger.error(f"Failed to clean up an uncommitted KV cache manager: {error}") - - def _estimate_swa_cache_size( layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]], @@ -1124,23 +1114,21 @@ def append_to_kv_heads_per_layer( 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) - local_init_status = _KVCacheManagerInitStatus.KEEP_HOST - except (CuError, KVCacheOutOfMemoryError): - local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST except Exception as error: - init_error = error.with_traceback(None) - local_init_status = _KVCacheManagerInitStatus.ABORT + 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: - init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) - except Exception: - _shutdown_kv_cache_manager_candidate(candidate) - raise + init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) if init_status == _KVCacheManagerInitStatus.ABORT: - _shutdown_kv_cache_manager_candidate(candidate) + 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") @@ -1151,10 +1139,11 @@ def append_to_kv_heads_per_layer( "(cuMemHostRegister may have failed). Rebuilding without the " "host cache tier on all ranks." ) - _shutdown_kv_cache_manager_candidate(candidate) - candidate = None fallback_error: Optional[Exception] = None try: + if candidate is not None: + candidate.shutdown() + candidate = None config = replace( config, cache_tiers=[ @@ -1163,29 +1152,20 @@ def append_to_kv_heads_per_layer( if not isinstance(tier, HostCacheTierConfig) ], ) + candidate = KVCacheManagerPy(config, event_manager=self.event_manager) except Exception as error: fallback_error = error.with_traceback(None) - else: - try: - 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 ) - try: - fallback_status = _sync_kv_cache_manager_init_status( - local_fallback_status, mapping - ) - except Exception: - _shutdown_kv_cache_manager_candidate(candidate) - raise + fallback_status = _sync_kv_cache_manager_init_status(local_fallback_status, mapping) if fallback_status == _KVCacheManagerInitStatus.ABORT: - _shutdown_kv_cache_manager_candidate(candidate) + if candidate is not None: + candidate.shutdown() if fallback_error is not None: raise fallback_error raise RuntimeError( From 779202b96ec9aa624435e82d9ba08164b5f3640c Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:10:40 -0700 Subject: [PATCH 4/4] [None][test] Make fallback consensus test CPU-safe Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/executor/test_kv_cache_manager_v2.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 bc59b1afc77b..bea5ae250a58 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -157,6 +157,7 @@ def build_cache_config( 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), @@ -181,8 +182,8 @@ def build_cache_config( return manager, impl_constructor -def _host_fallback_consensus_worker() -> tuple[int, int, int, bool]: - """Exercise the real world collective from an attention-DP worker.""" +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() @@ -205,7 +206,6 @@ def _host_fallback_consensus_worker() -> tuple[int, int, int, bool]: world_size=world_size, rank=rank, tp_size=world_size, - enable_attention_dp=True, ), ) @@ -379,12 +379,12 @@ def test_kv_cache_manager_init_status_sync_uses_world_max() -> None: @pytest.mark.cpu_only @pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device (MPI) build required") -def test_attention_dp_ranks_converge_on_hostless_fallback() -> None: +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(_host_fallback_consensus_worker) + results = session.submit_sync(_multi_rank_host_fallback_consensus_worker) finally: session.shutdown()