[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -615,8 +615,10 @@ void KvCache::_refreshStatsDirtyState()
 
 void KvCache::_recordDirectIterationStats(LifeCycleId lifeCycle, KVCacheIterationStatsDelta const& iterationStats)
 {
-    if (!_shouldRecordStats() || iterationStats.empty()
-        || !std::holds_alternative<AttnLifeCycle>(mManager->lifeCycles().getLifeCycle(lifeCycle)))
+    // Every lifecycle is reported, including SSM / recurrent ones: iteration
+    // statistics are keyed by lifecycle, so recurrent page movement stays
+    // distinguishable from attention movement downstream.
+    if (!_shouldRecordStats() || iterationStats.empty())
     {
         return;
     }
@@ -636,10 +638,7 @@ void KvCache::_recordMigratedSlots(
     for (auto const& page : pages)
     {
         LifeCycleId const lifeCycle = page->lifeCycle;
-        if (!std::holds_alternative<AttnLifeCycle>(mManager->lifeCycles().getLifeCycle(lifeCycle)))
-        {
-            continue;
-        }
+        bool const isAttention = std::holds_alternative<AttnLifeCycle>(mManager->lifeCycles().getLifeCycle(lifeCycle));
 
         PoolGroupIndex const poolGroup = mManager->storage().getPoolGroupIndex(lifeCycle);
         int64_t pageSize = 0;
@@ -657,8 +656,13 @@ void KvCache::_recordMigratedSlots(
         }
         else if (dstLevel == kGpuLevel)
         {
-            stats.allocTotalBlocks = 1;
-            stats.allocNewBlocks = 1;
+            // Global cache-hit accounting is attention-only. SSM movement is
+            // reported by lifecycle/pool-group iteration statistics instead.
+            if (isAttention)
+            {
+                stats.allocTotalBlocks = 1;
+                stats.allocNewBlocks = 1;
+            }
             iterationStats.iterAllocTotalBlocks = 1;
             iterationStats.iterAllocNewBlocks = 1;
             if (srcLevel > kGpuLevel)
@@ -692,10 +696,6 @@ void KvCache::_recordDroppedPages(std::vector<SharedPtr<Page>> const& pages, Cac
     for (auto const& page : pages)
     {
         LifeCycleId const lifeCycle = page->lifeCycle;
-        if (!std::holds_alternative<AttnLifeCycle>(mManager->lifeCycles().getLifeCycle(lifeCycle)))
-        {
-            continue;
-        }
         PoolGroupIndex const poolGroup = mManager->storage().getPoolGroupIndex(lifeCycle);
         int64_t pageSize = 0;
         for (size_t const size : mManager->storage().slotSize(poolGroup))

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2014,6 +2014,29 @@ def _get_mamba_cache_layer_masks(
     )
 
 
+# The V1 hybrid managers select the convolution-state layout by model_type;
+# MambaHybridCacheManagerV2 takes the layout by name and rejects model_type.
+_CONV_STATE_LAYOUT_BY_MODEL_TYPE = {
+    "nemotron_hybrid": "x_b_c",
+    "qwen3_next": "q_k_v",
+}
+
+
+def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type,
+                              model_type: str) -> dict:
+    """Constructor kwarg selecting the conv-state layout for a hybrid manager.
+
+    Keeps the V1-vs-V2 dispatch in one place: a manager branch that forgets it
+    would previously get V2's silent "x_b_c" default (the Kimi K3 bug fixed in
+    this change).
+    """
+    if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
+        return {
+            "conv_state_layout": _CONV_STATE_LAYOUT_BY_MODEL_TYPE[model_type]
+        }
+    return {"model_type": model_type}
+
+
 def _create_kv_cache_manager(
         model_engine: Optional[PyTorchModelEngine],
         kv_cache_manager_cls,
@@ -2221,6 +2244,10 @@ def _create_kv_cache_manager(
             if is_kda_mtp_verify_available():
                 kimi_extra_kwargs["kda_replay_num_spec"] = (
                     spec_config.tokens_per_gen_step - 1)
+        # KDA's conv state is a [Q | K | V] concatenation whose three sections
+        # have identical width, i.e. the qwen3_next section layout.
+        kimi_extra_kwargs.update(
+            _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next"))
         kv_cache_manager = kv_cache_manager_cls(
             # mamba (KDA) cache parameters
             mamba_params.state_size,
@@ -2248,9 +2275,6 @@ def _create_kv_cache_manager(
             spec_config=spec_config,
             is_estimating_kv_cache=estimating_kv_cache,
             execution_stream=execution_stream,
-            # Reuse the qwen3_next [Q | K | V] conv-state section layout;
-            # all three KDA sections have identical width.
-            model_type="qwen3_next",
             **kimi_extra_kwargs,
             **manager_extra_kwargs,
         )
@@ -2364,10 +2388,8 @@ def _create_kv_cache_manager(
                                          and mamba_params.mamba_ssm_cache_dtype
                                          == torch.float16)
         mamba_manager_extra_kwargs = dict(manager_extra_kwargs)
-        if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
-            mamba_manager_extra_kwargs["conv_state_layout"] = "x_b_c"
-        else:
-            mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid"
+        mamba_manager_extra_kwargs.update(
+            _mamba_conv_layout_kwargs(kv_cache_manager_cls, "nemotron_hybrid"))
         kv_cache_manager = kv_cache_manager_cls(
             # mamba cache parameters
             mamba_params.state_size,
@@ -2473,10 +2495,8 @@ def _create_kv_cache_manager(
                     ("ENABLED" if use_replay else "DISABLED"))
 
         mamba_manager_extra_kwargs = dict(manager_extra_kwargs)
-        if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
-            mamba_manager_extra_kwargs["conv_state_layout"] = "q_k_v"
-        else:
-            mamba_manager_extra_kwargs["model_type"] = "qwen3_next"
+        mamba_manager_extra_kwargs.update(
+            _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next"))
         kv_cache_manager = kv_cache_manager_cls(
             # mamba cache parameters
             mamba_params.state_size,

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2870,6 +2870,14 @@ def __init__(
         if conv_state_layout not in ("x_b_c", "q_k_v"):
             raise ValueError(
                 f"Unsupported convolution state layout: {conv_state_layout!r}")
+        if "model_type" in kwargs:
+            # The V1 managers select the conv-state layout by model_type; this
+            # class takes it explicitly. Silently absorbing model_type here
+            # means a caller's layout request would be dropped on the floor.
+            raise TypeError(
+                "MambaHybridCacheManagerV2 does not accept 'model_type' "
+                f"(got {kwargs['model_type']!r}); pass "
+                "conv_state_layout='x_b_c' or 'q_k_v' instead")
         total_layers = len(mamba_layer_mask)
         if layer_mask is None:
             full_attention_layer_mask = [False] * total_layers

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -427,11 +427,12 @@ def _refresh_stats_dirty_state(self) -> None:
         else:
             self.manager.clear_stats_dirty(self.id)
 
+    def _is_attention_life_cycle(self, life_cycle: LifeCycleId) -> bool:
+        return isinstance(self.manager._life_cycles.get_life_cycle(life_cycle), AttnLifeCycle)
+
     def _stats_life_cycle_key(self, life_cycle: LifeCycleId) -> LifeCycleId | None:
-        life_cycle_obj = self.manager._life_cycles.get_life_cycle(life_cycle)
-        if isinstance(life_cycle_obj, AttnLifeCycle):
-            return life_cycle
-        return None
+        """Key for the attention-only block-reuse (hit/miss range) accounting."""
+        return life_cycle if self._is_attention_life_cycle(life_cycle) else None
 
     def _refresh_generation_alloc_ready(self) -> None:
         expected_prompt_length = self._expected_prompt_length
@@ -498,10 +499,12 @@ def _subtract_pending_allocation_range(
     def _record_direct_iteration_stats(
         self, life_cycle: LifeCycleId, iteration_stats: KVCacheIterationStatsDelta
     ) -> None:
-        life_cycle_key = self._stats_life_cycle_key(life_cycle)
-        if life_cycle_key is None or iteration_stats.empty or not self._should_record_stats():
+        # Every life cycle is reported, including SSM / recurrent ones: iteration
+        # statistics are keyed by life cycle, so recurrent page movement stays
+        # distinguishable from attention movement downstream.
+        if iteration_stats.empty or not self._should_record_stats():
             return
-        self.manager.commit_stats(KVCacheStatsDelta(), {life_cycle_key: iteration_stats})
+        self.manager.commit_stats(KVCacheStatsDelta(), {life_cycle: iteration_stats})
 
     def _record_migrated_slots(
         self,
@@ -514,9 +517,7 @@ def _record_migrated_slots(
             return
         assert len(pages) == len(slots)
         for page in pages:
-            life_cycle_key = self._stats_life_cycle_key(page.life_cycle)
-            if life_cycle_key is None:
-                continue
+            is_attention = self._is_attention_life_cycle(page.life_cycle)
             pg_idx = self.manager._storage.get_pool_group_index(page.life_cycle)
             page_size = sum(self.manager._storage.slot_size(pg_idx))
             stats = KVCacheStatsDelta()
@@ -525,8 +526,11 @@ def _record_migrated_slots(
                 iteration_stats.iter_offload_blocks = 1
                 iteration_stats.iter_offload_bytes = page_size
             elif dst_level == GPU_LEVEL:
-                stats.alloc_total_blocks = 1
-                stats.alloc_new_blocks = 1
+                # Global cache-hit accounting is attention-only. SSM movement is
+                # reported by life-cycle/pool-group iteration statistics instead.
+                if is_attention:
+                    stats.alloc_total_blocks = 1
+                    stats.alloc_new_blocks = 1
                 iteration_stats.iter_alloc_total_blocks = 1
                 iteration_stats.iter_alloc_new_blocks = 1
                 if src_level > GPU_LEVEL:
@@ -536,7 +540,7 @@ def _record_migrated_slots(
                     iteration_stats.iter_intra_device_copy_blocks = 1
                     iteration_stats.iter_intra_device_copy_bytes = page_size
             if not stats.empty or not iteration_stats.empty:
-                self.manager.commit_stats(stats, {life_cycle_key: iteration_stats})
+                self.manager.commit_stats(stats, {page.life_cycle: iteration_stats})
 
     def _record_dropped_pages(
         self,
@@ -554,15 +558,12 @@ def _record_dropped_pages(
         if not self._should_record_stats() or not pages:
             return
         for page in pages:
-            life_cycle_key = self._stats_life_cycle_key(page.life_cycle)
-            if life_cycle_key is None:
-                continue
             pg_idx = self.manager._storage.get_pool_group_index(page.life_cycle)
             page_size = sum(self.manager._storage.slot_size(pg_idx))
             iteration_stats = KVCacheIterationStatsDelta()
             iteration_stats.iter_host_dropped_blocks = 1
             iteration_stats.iter_host_dropped_bytes = page_size
-            self.manager.commit_stats(KVCacheStatsDelta(), {life_cycle_key: iteration_stats})
+            self.manager.commit_stats(KVCacheStatsDelta(), {page.life_cycle: iteration_stats})
 
     # destroy ownership of memory blocks, so KV cache manager can decide to evict or drop them. After
     # close, uncommitted data in blocks for (beam_index >= beam_width) will be lost.
@@ -1292,13 +1293,16 @@ def resume(self, cuda_stream: CudaStream | None = None) -> bool:
                         )
                         if changed:
                             self.manager.mark_stats_dirty(self.id)
-                    self._record_direct_iteration_stats(
-                        lc_idx,
-                        KVCacheIterationStatsDelta(
-                            iter_intra_device_copy_blocks=1,
-                            iter_intra_device_copy_bytes=sum(storage.slot_size(pg_idx)),
-                        ),
-                    )
+                # Block-reuse accounting above is attention-only, but the copy
+                # itself is reported for every life cycle, SSM included —
+                # matching the C++ backend's deferred-copy loop (kvCache.cpp).
+                self._record_direct_iteration_stats(
+                    lc_idx,
+                    KVCacheIterationStatsDelta(
+                        iter_intra_device_copy_blocks=1,
+                        iter_intra_device_copy_bytes=sum(storage.slot_size(pg_idx)),
+                    ),
+                )
             # Unlock source pages — _record_event captures all prior cuda work
             # so the original pages know when we're done reading from them.
             if src_locks:

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -317,22 +317,80 @@ def test_kimi_explicit_v2_manager_geometry(monkeypatch: pytest.MonkeyPatch) -> N
     assert "kda_replay_num_spec" not in kwargs
 
 
-@pytest.mark.xfail(
-    reason="The Kimi route in _create_kv_cache_manager passes "
-    "model_type='qwen3_next' unconditionally; MambaHybridCacheManagerV2 "
-    "swallows it via **kwargs and falls back to the 'x_b_c' "
-    "conv_state_layout instead of the KDA [q|k|v] sectioning. Runtime-side "
-    "layout selection is a follow-up (TRTLLM-14813).",
-    strict=True,
-)
 def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    """TRTLLM-15216: MambaHybridCacheManagerV2 takes the KDA conv-state
+    sectioning by `conv_state_layout`, not by `model_type`. Passing
+    `model_type` instead is silently swallowed by **kwargs and leaves the
+    default 'x_b_c' layout, i.e. a wrong KDA conv state with no error."""
     _, kwargs = _capture_kimi_v2_manager_ctor(monkeypatch)
     assert kwargs["conv_state_layout"] == "q_k_v"
     assert "model_type" not in kwargs
 
 
+def test_kimi_v1_manager_still_selects_qwen3_next_model_type(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The V1 managers have no `conv_state_layout` parameter; they must keep
+    getting `model_type='qwen3_next'` (TRTLLM-15216 regression guard)."""
+    captured: dict[str, object] = {}
+
+    class RecordingV1Manager(CppMambaHybridCacheManager):
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            captured["kwargs"] = kwargs
+
+    model_config = _kimi_model_config()
+    _create_kv_cache_manager(
+        model_engine=None,
+        kv_cache_manager_cls=RecordingV1Manager,
+        mapping=Mapping(world_size=1, tp_size=1, pp_size=1),
+        kv_cache_config=KvCacheConfig(),
+        tokens_per_block=64,
+        max_seq_len=2048,
+        max_batch_size=4,
+        spec_config=None,
+        sparse_attention_config=None,
+        max_num_tokens=256,
+        max_beam_width=1,
+        kv_connector_manager=None,
+        model_config=model_config,
+        dtype=torch.bfloat16,
+        is_draft=False,
+    )
+    kwargs = captured["kwargs"]
+    assert kwargs["model_type"] == "qwen3_next"
+    assert "conv_state_layout" not in kwargs
+
+
+def test_v2_manager_rejects_model_type_kwarg() -> None:
+    """MambaHybridCacheManagerV2 must fail loudly when handed the V1 managers'
+    `model_type` instead of `conv_state_layout` — silently absorbing it into
+    **kwargs is how the TRTLLM-15216 wrong-layout bug went unnoticed."""
+    with pytest.raises(TypeError, match="conv_state_layout"):
+        MambaHybridCacheManagerV2(
+            16,  # mamba_d_state
+            4,  # mamba_d_conv
+            8,  # mamba_num_heads
+            1,  # mamba_n_groups
+            16,  # mamba_head_dim
+            2,  # mamba_num_layers
+            [True, True],  # mamba_layer_mask
+            torch.float16,
+            torch.float16,
+            KvCacheConfig(),
+            CacheTypeCpp.SELF,
+            num_layers=0,
+            num_kv_heads=1,
+            head_dim=16,
+            tokens_per_block=32,
+            max_seq_len=64,
+            max_batch_size=1,
+            mapping=Mapping(world_size=1, tp_size=1, pp_size=1),
+            model_type="qwen3_next",
+        )
+
+
 @pytest.mark.parametrize(
     ("use_v2", "enable_block_reuse", "expected"),
     [

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

@@ -2243,6 +2243,44 @@ def test_discard_ssm_snapshot_stats_clears_dirty_state(self) -> None:
         kv_cache.resume(cast(CudaStream, stream_holder.handle))
         kv_cache.close()
 
+    def test_ssm_resume_records_intra_device_copy(self) -> None:
+        """The SSM deferred copy on resume is counted in iteration stats.
+
+        First resume of a cache reusing an SSM snapshot copies the snapshot
+        into a private slot; the copy must appear in the SSM life cycle's
+        iteration stats (TRTLLM-15217). Runs against the selected backend, so
+        it checks the default C++ implementation and Python-backend parity.
+        """
+        tokens_per_block = 32
+        cfg = self._make_ssm_config(tokens_per_block=tokens_per_block)
+        self.manager = KVCacheManager(cfg)
+        stream_holder = CachedCudaStream()
+        stream = cast(CudaStream, stream_holder.handle)
+        prompt = [self.next_token() for _ in range(48)]
+
+        seed = self.manager.create_kv_cache()
+        seed.resume(stream)
+        seed.capacity = tokens_per_block
+        seed.history_length = tokens_per_block
+        seed.commit(prompt[:tokens_per_block], is_end=True)
+        seed.close()
+
+        reused = self.manager.create_kv_cache(input_tokens=prompt, id=101)
+        self.assertEqual(reused.num_committed_tokens, tokens_per_block)
+        reused.commit_pending_stats()
+        # Drop everything recorded so far; only the resume below should count.
+        self.manager.get_and_reset_ssm_snapshot_iteration_stats()
+        self.manager.get_and_reset_iteration_stats()
+
+        self.assertTrue(reused.resume(stream))
+        ssm_life_cycle_id = _introspection.ssm_life_cycle_id(self.manager)
+        assert ssm_life_cycle_id is not None
+        stats = self.manager.get_and_reset_iteration_stats()
+        self.assertIn(ssm_life_cycle_id, stats)
+        self.assertEqual(stats[ssm_life_cycle_id].iter_intra_device_copy_blocks, 1)
+        self.assertGreater(stats[ssm_life_cycle_id].iter_intra_device_copy_bytes, 0)
+        reused.close()
+
     def test_ssm(self) -> None:
         """Inference with SSM layer: prefill 63 tokens, decode 52 tokens."""
         cfg = self._make_ssm_config()