[TRTLLM-13308][feat] Make the KVCacheManagerV2 KV pool rebalance safe under TP, CP, PP and attention DP

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -560,7 +560,11 @@ void KvCache::close()
     stopCommitting();
     TLLM_CHECK_DEBUG(_checkSanity());
 
-    if (mCapacity > 0)
+    // Dummy/warmup caches are reserved at the model's full declared context, not at a realistic
+    // sequence length, and mAvgSqrCapacity is an RMS -- so a handful of them dominates the
+    // statistic outright and the tuner sizes pools for sequences that never arrive. They are
+    // already tracked as stats-excluded at creation; honour that here too.
+    if (mCapacity > 0 && !mManager->isStatsExcluded(id))
     {
         mAvgCapacity.update(static_cast<double>(mCapacity));
         mManager->updateAvgSqrCapacity(mAvgCapacity.value() * mAvgCapacity.value());

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -583,7 +583,12 @@ def close(self) -> None:
         self.stop_committing()
         assert NDEBUG or self._check_sanity()
         manager = self.manager
-        if self.capacity > 0:
+        # Dummy/warmup caches are reserved at the model's full declared context,
+        # not at a realistic sequence length, and _avg_sqr_capacity is an RMS --
+        # so a handful of them dominates the statistic outright and the tuner
+        # sizes pools for sequences that never arrive. They are already tracked
+        # as stats-excluded at creation; honour that here too.
+        if self.capacity > 0 and not manager.is_stats_excluded(self.id):
             self._avg_capacity.update(self.capacity)
             manager._avg_sqr_capacity.update(self._avg_capacity.value**2)
             manager._avg_sqr_history_length.update(self._avg_history_length.value**2)

tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py

@@ -43,10 +43,18 @@ def _inject_pool_ratio_mismatch(llm: LLM, *, skew: float = 2.0) -> None:
     past the auto-tuner's adjustment threshold. The hook requires a model with
     >=2 pool groups (e.g. Gemma-3-1B with VSWA) and raises otherwise, so a future
     model change can't silently turn this test into a no-op.
+
+    Also drops the executor's rebalance-check throttle to every iteration, so
+    the test does not depend on how ``KV_POOL_REBALANCE_CHECK_INTERVAL`` compares
+    to the number of iterations this short prompt set happens to run.  Raise that
+    interval above the iteration count and the hook would never fire, leaving the
+    token comparison below to pass vacuously; the ratio assertion in
+    ``_generate_tokens`` is the backstop that would catch it.
     """
     from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection
 
     executor = llm._executor.engine
+    executor._rebalance_check_interval = 1
     kv_cache_manager = executor.kv_cache_manager
     _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew)
 
@@ -95,14 +103,38 @@ def _generate_tokens(*, model_path: str, disable_overlap: bool, enable_rebalance
     responsible for setting that env var (via monkeypatch or otherwise)
     before invoking this helper.
     """
+    from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection
+
     with LLM(
         model_path,
         disable_overlap_scheduler=disable_overlap,
         kv_cache_config=_vswa_kv_cache_config(enable_rebalance=enable_rebalance),
     ) as llm:
+        impl = llm._executor.engine.kv_cache_manager.impl
         if enable_rebalance:
             _inject_pool_ratio_mismatch(llm)
+        ratio_before = list(_introspection.current_gpu_ratio(impl))
         outputs = llm.generate(_PROMPTS, _SAMPLING)
+        ratio_after = list(_introspection.current_gpu_ratio(impl))
+
+        # Guard against a vacuous pass.  Token equality between the rebalance
+        # and no-rebalance arms proves nothing if adjust() never ran, and
+        # nothing in the run logs at info level to tell us it did.  The pool
+        # ratio moving is the observable signature that it happened.
+        if enable_rebalance:
+            assert ratio_after != ratio_before, (
+                "rebalance never fired: GPU pool ratio unchanged at "
+                f"{ratio_before}. The token comparison would pass vacuously. "
+                "Check the executor's rebalance-check throttle and the V2 "
+                "auto-tuner's sample-count / cooldown gates."
+            )
+        else:
+            assert ratio_after == ratio_before, (
+                "pool ratio moved with enable_kv_pool_rebalance=False "
+                f"({ratio_before} -> {ratio_after}); the baseline arm is "
+                "supposed to hold pool ratios fixed."
+            )
+
         return [list(o.outputs[0].token_ids) for o in outputs]
 
 

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -98,6 +98,11 @@ l0_dgx_h100:
   tests:
   # ------------- PyTorch tests ---------------
   - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu4"
+  # unittest/_torch/multi_gpu is otherwise only collected by the 2-GPU stage, so
+  # the 4-rank cases of the KV-pool rebalance agreement test (tp4, pp4, CP x TP,
+  # TP x PP, attention-DP over CP) always skip with "need 4 GPUs, have 2". Name
+  # the file here so those compositions actually run.
+  - unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py
   # ------------- Model specific tests ---------------
   - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=False]
   - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=True]