[None][test] Validate Kimi default KVCM V2 accuracy and disaggregated transfer

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -18,8 +18,11 @@
 
 Caching
 -------
-KDA states live on the mamba side of a ``MixedMambaHybridCacheManager``
-(wired in ``pyexecutor/_util.py``): per layer, a short-conv slot of
+KDA states use ``MambaHybridCacheManagerV2`` with the ``KvCacheConfig``
+default ``use_kv_cache_manager_v2=True`` (selected in ``pyexecutor/_util.py``).
+Setting this field to ``"auto"`` falls back to V1 for
+``KimiLinearForCausalLM``, which declares no V2 model preference.
+Per layer, the state consists of a short-conv slot of
 ``[3 * num_heads * head_dim, W]`` bf16 (the full FLA ``ShortConvolution``
 cache window, sections ``[q | k | v]``) and a delta-rule recurrent slot of
 ``[num_heads, head_dim, head_dim]`` fp32 (``[H, V, K]``, the
@@ -67,14 +70,13 @@
 conv/recurrent state back into the FLA kernels (``use_initial_states``)
 and the MLA prefill path natively attends over the cached latent prefix
 (``kv_len = cached + q_len``). KV-cache block reuse is supported as an
-opt-in via ``kv_cache_config.enable_block_reuse=true``, which routes to
-the unified-pool ``CppMambaHybridCacheManager`` (per-block KDA state
-snapshots every ``mamba_state_cache_interval`` tokens, FORCE_CHUNK
-context chunking).
+opt-in via ``kv_cache_config.enable_block_reuse=true`` on the default V2
+manager. Configuring ``mamba_state_config.periodic_snapshot_interval``
+adds periodic KDA state snapshots; reuse uses FORCE_CHUNK context chunking.
 
 Not supported: pipeline parallelism, draft-head spec-dec modes
 (MTP/Eagle — no draft-head checkpoint exists). SA speculative decoding
-is validated only without block reuse (Mixed cache manager).
+is validated only without block reuse.
 """
 
 from __future__ import annotations
@@ -2289,10 +2291,10 @@ def _setup_helix_mappings(
 
     @classmethod
     def get_model_defaults(cls, llm_args) -> dict:
-        # - enable_block_reuse defaults off: reuse is supported as an
-        #   explicit opt-in (routes to CppMambaHybridCacheManager with
-        #   per-block KDA state snapshots); the default stays on the
-        #   Mixed manager, which SA speculative decoding requires.
+        # - Keep block reuse opt-in because its combination with SA
+        #   speculative decoding is not validated. Both settings use V2
+        #   with the default use_kv_cache_manager_v2=True. The "auto"
+        #   setting follows the model preference and can resolve to V1.
         # - tokens_per_block=64: with 32, the flashinfer trtllm-gen FMHA lib
         #   rejects the MLA (576, 512) generation kernel (marked slower) and
         #   the fallback C++ path requires num_heads % 64 == 0, which K3's

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -274,22 +274,11 @@ def get_kv_cache_manager_cls(
                 Glm5NextCacheManager
             return Glm5NextCacheManager
 
-        # Kimi K3 (KDA + MLA hybrid): block reuse uses the unified C++ pool
-        # (CppMambaHybridCacheManager) like the other hybrid linear models —
-        # per-block KDA state snapshots every mamba_state_cache_interval
-        # tokens with FORCE_CHUNK context chunking. Without block reuse the
-        # Mixed manager (separate KV / recurrent-state pools) stays the
-        # default. SA speculative decoding is validated on the Mixed
-        # manager's SpeculativeState scratch path only; reuse + SA is
-        # unvalidated. Disaggregated serving (TRTLLM-14815) routes through
-        # the shared hybrid transceiver validation below: the Python NIXL
-        # transceiver selects the Mixed manager, whose KDA recurrent/conv
-        # states transfer through the bounce buffer.
         # Helix x speculation bookkeeping (per-token verify groups on the
         # superblock ledger, py_helix_decode_group_index advancement) exists
         # only in KVCacheManagerV2. The V1-family hybrid managers account
         # helix decode one token per iteration and have no helix-x-spec
-        # path, so a default (V1) resolution would run silently wrong.
+        # path, so selecting V1 would run silently wrong.
         if (model_config.mapping is not None
                 and model_config.mapping.has_cp_helix()
                 and model_config.spec_config is not None and not use_v2):
@@ -299,6 +288,12 @@ def get_kv_cache_manager_cls(
                 "hybrid managers do not implement per-token verify-group "
                 "bookkeeping.")
 
+        # Kimi uses MambaHybridCacheManagerV2 with or without reuse when
+        # use_kv_cache_manager_v2=True, the KvCacheConfig default. V1 can
+        # be selected explicitly or resolved from "auto". Aggregated V1
+        # uses Mixed without reuse and the C++ hybrid manager with reuse.
+        # Disaggregated serving requires Python NIXL below, using V2 for
+        # the configured default or Mixed for V1 without reuse.
         if is_kimi_linear(config) and not use_v2 and not is_disagg:
             if kv_cache_config.enable_block_reuse:
                 logger.info(

tests/integration/defs/accuracy/test_kimi3.py

@@ -190,15 +190,9 @@ def test_gpqa_diamond_w4a16_mxfp4(self) -> None:
                 max_num_tokens=33024,
                 use_low_precision_moe_combine=True,
             ),
-            # The qualified K3 GPQA configuration used KV cache manager V1.
-            # Like K3 MMMU, GPQA leaves very long generations after shorter
-            # requests finish. V2 has stalled or deadlocked at this shape even
-            # with a host tier and additional cache capacity. Pin V1 until the
-            # V2 x KDA-hybrid long-generation path is qualified.
             kv_cache_config=KvCacheConfig(
                 free_gpu_memory_fraction=0.25,
                 tokens_per_block=64,
-                use_kv_cache_manager_v2=False,
             ),
         ) as llm:
             # K3 stores its compressed-tensors quantization configuration in

tests/integration/test_lists/waives.txt

@@ -242,7 +242,7 @@ unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_hybrid_cache
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_hybrid_cache_manager_factory_rejects_unsupported_v2_disagg_route SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_hybrid_cache_manager_factory_routes_explicit_v2_disagg SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_hybrid_cache_manager_factory_v2_preference_does_not_select_v2 SKIP (https://nvbugs/6818710)
-unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_disagg_python_nixl_routes_to_mixed_manager SKIP (https://nvbugs/6818710)
+unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_disagg_explicit_v1_python_nixl_routes_to_mixed_manager SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_disagg_rejects_non_python_transceiver_route SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_explicit_v2_manager_enables_kda_replay SKIP (https://nvbugs/6818710)
 unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_explicit_v2_manager_geometry SKIP (https://nvbugs/6818710)

tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py

@@ -297,8 +297,10 @@ def _kimi_model_config() -> SimpleNamespace:
     )
     return SimpleNamespace(
         pretrained_config=config,
+        mapping=Mapping(),
         quant_config=None,
         sparse_attention_config=None,
+        spec_config=None,
         get_num_mamba_layers=lambda: 2,
     )
 
@@ -1027,7 +1029,34 @@ def test_kimi_disagg_rejects_non_python_transceiver_route(
         )
 
 
-def test_kimi_disagg_python_nixl_routes_to_mixed_manager(
+def test_kimi_disagg_defaults_use_v2_manager(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Kimi's default V2 manager supports the resolved Python NIXL route."""
+    from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM
+
+    monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    llm_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False, tokens_per_block=64),
+        cache_transceiver_config=CacheTransceiverConfig(backend="NIXL", transceiver_runtime="auto"),
+    )
+    _resolve_transceiver_runtime_auto(llm_args, KimiLinearForCausalLM)
+    _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM)
+
+    assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
+    assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON"
+    assert (
+        get_kv_cache_manager_cls(
+            _kimi_model_config(),
+            llm_args.kv_cache_config,
+            is_disagg=True,
+            cache_transceiver_config=llm_args.cache_transceiver_config,
+        )
+        is MambaHybridCacheManagerV2
+    )
+
+
+def test_kimi_disagg_explicit_v1_python_nixl_routes_to_mixed_manager(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)

tests/unittest/disaggregated/test_kda_mamba_transfer.py

@@ -34,8 +34,8 @@
   ``MambaPolicy`` descriptors tile each layer's slot bytes exactly.
 
 * ``test_kda_transfer`` performs a real single-node NIXL loopback transfer
-  and bitwise-compares both dtypes on the gen side: matched and
-  heterogeneous ctx/gen TP, attention-DP on and off.
+  with V1 and V2 cache managers and bitwise-compares both dtypes on the gen
+  side: matched and heterogeneous ctx/gen TP, attention-DP on and off.
 
 * ``test_kda_hetero_tp_sharded_accepted`` asserts that peer registration
   accepts heterogeneous ctx/gen TP with attention-DP off for what
@@ -49,7 +49,13 @@
 
 import pytest
 import torch
-from test_mamba_transfer import _create_transceivers, _run_concurrent
+from test_mamba_transfer import (
+    _create_transceivers,
+    _mamba_layer_ids,
+    _mamba_state_slot,
+    _run_concurrent,
+    _zero_mamba_states,
+)
 
 import tensorrt_llm
 import tensorrt_llm.bindings
@@ -67,7 +73,10 @@
 )
 from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup
 from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params
-from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MixedMambaHybridCacheManager
+from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import (
+    MambaHybridCacheManagerV2,
+    MixedMambaHybridCacheManager,
+)
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm.bindings import DataType
@@ -124,9 +133,12 @@ def _kimi_linear_config():
 
 
 def _create_kda_managers(
-    tp: int, enable_attention_dp: bool = False, max_batch_size: int = MAX_BATCH_SIZE
+    tp: int,
+    enable_attention_dp: bool = False,
+    max_batch_size: int = MAX_BATCH_SIZE,
+    use_v2: bool = False,
 ):
-    """Create MixedMambaHybridCacheManagers with K3-style KDA slots.
+    """Create V1 or V2 hybrid cache managers with K3-style KDA slots.
 
     Mirrors the ``is_kimi_linear`` route in ``_util.py``: the mamba params
     come from ``extract_mamba_kv_cache_params`` UNSCALED (global
@@ -150,7 +162,13 @@ def _create_kda_managers(
             pp_size=1,
             enable_attention_dp=enable_attention_dp,
         )
-        mgr = MixedMambaHybridCacheManager(
+        manager_cls = MambaHybridCacheManagerV2 if use_v2 else MixedMambaHybridCacheManager
+        manager_kwargs = (
+            {"is_disagg": True, "conv_state_layout": "q_k_v"}
+            if use_v2
+            else {"model_type": "qwen3_next"}
+        )
+        mgr = manager_cls(
             mamba_d_state=params.state_size,
             mamba_d_conv=params.conv_kernel,
             mamba_num_heads=params.num_heads,
@@ -165,6 +183,7 @@ def _create_kda_managers(
                 max_tokens=256 * max_batch_size,
                 enable_block_reuse=False,
                 event_buffer_max_size=0,
+                use_kv_cache_manager_v2=use_v2,
             ),
             kv_cache_type=CacheTypeCpp.SELF,
             num_layers=1,
@@ -176,7 +195,7 @@ def _create_kda_managers(
             max_batch_size=max_batch_size,
             mapping=mapping,
             dtype=DataType.FLOAT,
-            model_type="qwen3_next",
+            **manager_kwargs,
         )
         managers.append(mgr)
     return managers
@@ -541,7 +560,9 @@ def _generate_ground_truth(num_requests: int, seed: int = 20260722):
     return results
 
 
-def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = False):
+def run_kda_transfer_test(
+    ctx_tp: int, gen_tp: int, enable_attention_dp: bool = False, use_v2: bool = False
+):
     """Loopback KDA transfer, matched or heterogeneous ctx/gen TP.
 
     ADP off: per-rank state is a head shard (state_tp = tp) and every TP rank
@@ -560,13 +581,12 @@ def _ctx_ranks(req_idx: int) -> List[int]:
     def _gen_ranks(req_idx: int) -> List[int]:
         return [req_idx % gen_tp] if enable_attention_dp else list(range(gen_tp))
 
-    ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=enable_attention_dp)
-    gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=enable_attention_dp)
+    ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=enable_attention_dp, use_v2=use_v2)
+    gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=enable_attention_dp, use_v2=use_v2)
     ctx_tcs, gen_tcs = [], []
     try:
         for mgr in ctx_mgrs + gen_mgrs:
-            mgr._impl.mamba_cache.conv.zero_()
-            mgr._impl.mamba_cache.temporal.zero_()
+            _zero_mamba_states(mgr)
 
         config = CacheTransceiverConfig(
             backend="NIXL",
@@ -616,13 +636,28 @@ def _gen_ranks(req_idx: int) -> List[int]:
             gen_reqs.append(gen_req)
 
         ctx_batch = ScheduledRequests()
-        ctx_batch.reset_context_requests(ctx_reqs)
-        for mgr in ctx_mgrs:
-            mgr.prepare_resources(ctx_batch)
+        if use_v2:
+            ctx_batch.context_requests_last_chunk = ctx_reqs
+            for mgr in ctx_mgrs:
+                for req in ctx_reqs:
+                    assert mgr.prepare_context(req)
+                    assert mgr.resize_context(req, req.context_chunk_size)
+                mgr.prepare_resources(ctx_batch)
+        else:
+            ctx_batch.reset_context_requests(ctx_reqs)
+            for mgr in ctx_mgrs:
+                mgr.prepare_resources(ctx_batch)
         gen_batch = ScheduledRequests()
-        gen_batch.reset_context_requests(gen_reqs)
-        for mgr in gen_mgrs:
-            mgr.prepare_resources(gen_batch)
+        if use_v2:
+            gen_batch.context_requests_last_chunk = gen_reqs
+            for mgr in gen_mgrs:
+                for req in gen_reqs:
+                    assert mgr.prepare_disagg_gen_init(req)
+                mgr.prepare_resources(gen_batch)
+        else:
+            gen_batch.reset_context_requests(gen_reqs)
+            for mgr in gen_mgrs:
+                mgr.prepare_resources(gen_batch)
         for req in ctx_reqs + gen_reqs:
             req.context_current_position = req.prompt_len
             req.add_new_token(req.prompt_len, 0)
@@ -639,8 +674,8 @@ def _gen_ranks(req_idx: int) -> List[int]:
             # state); with ADP off rank r holds shard r.
             shard_rank = rank % ctx_state_tp
             for req_idx, rid in enumerate(ctx_rids):
-                slot = mgr.mamba_cache_index[rid]
-                for layer_idx in mgr._impl.mamba_layer_offsets:
+                slot = _mamba_state_slot(mgr, rid)
+                for layer_idx in _mamba_layer_ids(mgr):
                     full = ground_truth[req_idx][layer_idx]
                     mgr.get_conv_states(layer_idx)[slot] = _shard_kda_conv(
                         full["conv"], ctx_state_tp, shard_rank
@@ -681,8 +716,8 @@ def _gen_ranks(req_idx: int) -> List[int]:
             for req_idx, rid in enumerate(gen_rids):
                 if gen_rank not in _gen_ranks(req_idx):
                     continue
-                slot = mgr.mamba_cache_index[rid]
-                for layer_idx in mgr._impl.mamba_layer_offsets:
+                slot = _mamba_state_slot(mgr, rid)
+                for layer_idx in _mamba_layer_ids(mgr):
                     full = ground_truth[req_idx][layer_idx]
                     expected = {
                         "conv": _shard_kda_conv(full["conv"], gen_state_tp, shard_rank),
@@ -711,6 +746,7 @@ def _gen_ranks(req_idx: int) -> List[int]:
 
 
 @pytest.mark.timeout(180)
+@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"])
 @pytest.mark.parametrize(
     "ctx_tp,gen_tp,enable_attention_dp",
     [
@@ -730,11 +766,11 @@ def _gen_ranks(req_idx: int) -> List[int]:
         "hetero_ctx_tp4_gen_tp2_adp_off_sharded",
     ],
 )
-def test_kda_transfer(ctx_tp, gen_tp, enable_attention_dp):
+def test_kda_transfer(ctx_tp, gen_tp, enable_attention_dp, use_v2):
     """KDA two-dtype state transfer over real NIXL loopback.
 
     Covers matched TP with ADP on/off, hetero DEP with ADP on, and hetero
     ctx/gen TP with ADP off (head-sharded state re-tiled by the TP-mismatch
     mappers, both expand and contract directions).
     """
-    run_kda_transfer_test(ctx_tp, gen_tp, enable_attention_dp=enable_attention_dp)
+    run_kda_transfer_test(ctx_tp, gen_tp, enable_attention_dp=enable_attention_dp, use_v2=use_v2)