[TRTLLM-15179][test] Add multi-rank test for KVCM v2 host-tier quota sync

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(

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)

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]