diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index f522ecf90ac6..8606626fa398 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -275,6 +275,31 @@ def _compute_auto_host_tier_quota( return host_quota +def _sync_host_tier_quota(host_quota: int, mapping: Mapping) -> int: + """Reduce the auto-provisioned host cache tier quota to the fleet minimum. + + ``_compute_auto_host_tier_quota`` reads rank-local host state + (``SC_AVPHYS_PAGES``, ``RLIMIT_MEMLOCK``), so co-scheduled ranks can arrive + at divergent host quotas (observed up to 10x). Divergent host-tier + retention makes per-rank MAX_UTILIZATION schedulers disagree about which + suspended requests can resume, which wedges collectives on + non-attention-DP TP. Reducing to the fleet minimum makes the + most-constrained rank set the value for everyone, mirroring the device + quota sync. A single-rank job needs no collective and is returned as-is. + + Args: + host_quota: This rank's locally-computed host tier quota in bytes. + mapping: Parallelism mapping; ``world_size`` selects whether a + collective is needed. + + Returns: + The globally-agreed host tier quota in bytes (the fleet ``MIN``). + """ + if mapping.world_size > 1: + host_quota = Distributed.get(mapping).allreduce(host_quota, op=ReduceOp.MIN) + return host_quota + + def _estimate_swa_cache_size( layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]], @@ -1030,13 +1055,9 @@ def append_to_kv_heads_per_layer( f"{local_ranks} co-located rank(s) on this node, " f"available host memory {mem_available / (1 << 30):.2f}GiB" ) - # Auto sizing reads rank-local host state, so co-scheduled ranks - # can get divergent host quotas (observed up to 10x). Divergent - # host-tier retention makes per-rank schedulers disagree, which - # wedges collectives on non-attention-DP TP. Sync to the fleet - # minimum, like the device quota above. - if mapping.world_size > 1: - host_quota = Distributed.get(mapping).allreduce(host_quota, op=ReduceOp.MIN) + # Reduce the rank-local auto quota to the fleet minimum; see + # _sync_host_tier_quota for why divergence wedges collectives. + host_quota = _sync_host_tier_quota(host_quota, mapping) if host_quota > 0: cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) logger.info( diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index dbf67ce1e568..cb0c137b98ce 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -111,6 +111,7 @@ disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-C disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) +examples/test_ad_speculative_decoding.py::test_nemotron_mtp_model_with_weights SKIP (https://nvbugs/6630699) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) diff --git a/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py b/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py index f4c3c237592d..9604ee7b797f 100644 --- a/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py +++ b/tests/unittest/_torch/executor/test_kvv2_host_tier_sizing.py @@ -3,11 +3,19 @@ The auto-provisioned host tier is computed per rank but drawn from a node-level memory budget, so it must be divided by the number of ranks co-located on the same physical node to avoid host OOM. + +The per-rank computation reads rank-local host state, so co-scheduled ranks +can arrive at divergent quotas; ``TestSyncHostTierQuota`` covers the +cross-rank ``allreduce(MIN)`` that reconciles them to the fleet minimum. """ import pytest -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import _compute_auto_host_tier_quota +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + _compute_auto_host_tier_quota, + _sync_host_tier_quota, +) +from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE GiB = 1 << 30 @@ -91,3 +99,83 @@ def test_result_is_always_positive(self): ) > 0 ) + + +def _host_tier_sync_worker(per_rank_mem_gib): + """Run on every MPI rank by the MpiPoolSession harness below. + + Each rank simulates reading a different amount of available host memory + (as ``os.sysconf("SC_AVPHYS_PAGES")`` would differ across co-scheduled + ranks), computes its own auto host-tier quota, then runs the real + cross-rank sync. Returns this rank's pre- and post-sync quota so the parent + process can assert on convergence. + """ + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + _compute_auto_host_tier_quota, + _sync_host_tier_quota, + ) + from tensorrt_llm._utils import mpi_rank, mpi_world_size + from tensorrt_llm.mapping import Mapping + + rank = mpi_rank() + world_size = mpi_world_size() + + # A large device quota so the (rank-local) available-memory budget is the + # binding constraint, producing a different local quota on each rank. + local_quota = _compute_auto_host_tier_quota( + quota=1024 * GiB, + local_ranks=1, + mem_available=float(per_rank_mem_gib[rank] * GiB), + memlock_limit=float("inf"), + ) + mapping = Mapping(world_size=world_size, rank=rank, tp_size=world_size) + synced_quota = _sync_host_tier_quota(local_quota, mapping) + return {"rank": rank, "local_quota": local_quota, "synced_quota": synced_quota} + + +class TestSyncHostTierQuota: + def test_single_rank_is_a_noop(self): + # world_size == 1 must not touch the collective layer (no MPI needed); + # the local quota is returned unchanged. + class _FakeMapping: + world_size = 1 + + quota = 173 * GiB + assert _sync_host_tier_quota(quota, _FakeMapping()) == quota + + @pytest.mark.cpu_only + @pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device (MPI) build required") + def test_multi_rank_syncs_to_fleet_min(self): + """Every rank must end up with the same host-tier quota after the sync. + + Regression guard for the hang fixed in PR #17380 (TRTLLM-15179): when + co-scheduled ranks auto-compute divergent host quotas, per-rank + MAX_UTILIZATION schedulers disagree about which suspended requests can + resume and wedge collectives on non-attention-DP TP. The + ``allreduce(MIN)`` in ``_sync_host_tier_quota`` makes the + most-constrained rank set the fleet value. + """ + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + world_size = 2 + # rank 0 sees ~440 GiB available (-> 220 GiB local quota), rank 1 sees + # ~880 GiB (-> 440 GiB). rank 0 is the most-constrained rank. + per_rank_mem_gib = [440, 880] + + session = MpiPoolSession(n_workers=world_size) + try: + results = session.submit_sync(_host_tier_sync_worker, per_rank_mem_gib) + finally: + session.shutdown() + + results = sorted(results, key=lambda r: r["rank"]) + local_quotas = [r["local_quota"] for r in results] + synced_quotas = [r["synced_quota"] for r in results] + + # Pre-sync the ranks genuinely disagreed (else the test proves nothing). + assert local_quotas[0] != local_quotas[1], local_quotas + # Post-sync every rank agrees... + assert len(set(synced_quotas)) == 1, synced_quotas + # ...on the global MIN (the most-constrained rank sets the fleet value). + assert synced_quotas[0] == min(local_quotas) + assert synced_quotas[0] == local_quotas[0]