From e2e6f2228e21ecc99b54b220e9464176cff50b21 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 06:46:59 -0700 Subject: [PATCH 1/6] [TRTLLM-15216][fix] pass conv_state_layout for Kimi K3 on KV cache manager V2 The kimi_linear branch of _create_kv_cache_manager passed model_type="qwen3_next" unconditionally. MambaHybridCacheManagerV2 has no model_type parameter: it absorbs it into **kwargs and selects the KDA convolution-state sectioning from conv_state_layout, which defaults to "x_b_c". An explicit use_kv_cache_manager_v2=True opt-in therefore built the KDA conv state with the wrong section layout, silently, with no error and no warning. Select the kwarg from the manager class, matching what the qwen3_hybrid branch in the same function already does: conv_state_layout="q_k_v" for MambaHybridCacheManagerV2, model_type="qwen3_next" for the V1 managers. This path is opt-in only today (use_kv_cache_manager_v2 must be set to True; "auto" does not select V2 for kimi_linear), so no default configuration changes behavior. Test: test_kimi_explicit_v2_manager_uses_qkv_convolution_layout already existed as a strict xfail describing this bug; the marker is dropped. Adds test_kimi_v1_manager_still_selects_qwen3_next_model_type to guard the V1 branch. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/pyexecutor/_util.py | 13 ++++-- .../executor/test_mamba_cache_manager.py | 46 +++++++++++++++---- 2 files changed, 48 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..12d4646dcfcb 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2221,6 +2221,16 @@ 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. The V1 + # managers select that layout by `model_type`; MambaHybridCacheManagerV2 + # takes an explicit `conv_state_layout` and silently defaults to + # "x_b_c" (and absorbs `model_type` into **kwargs), so V2 must be given + # the layout by name. Mirrors the qwen3_hybrid branch below. + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + kimi_extra_kwargs["conv_state_layout"] = "q_k_v" + else: + kimi_extra_kwargs["model_type"] = "qwen3_next" kv_cache_manager = kv_cache_manager_cls( # mamba (KDA) cache parameters mamba_params.state_size, @@ -2248,9 +2258,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, ) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 81db2977398d..ff37ee7d92ca 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -317,22 +317,52 @@ 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 + + @pytest.mark.parametrize( ("use_v2", "enable_block_reuse", "expected"), [ From 3f6c398fd67db542847d8255b32e9109ab3fcc36 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 06:58:53 -0700 Subject: [PATCH 2/6] [TRTLLM-15217][chore] report SSM life cycles in KV cache manager V2 iteration stats The V2 page-movement recorders (_recordDirectIterationStats, _recordMigratedSlots, _recordDroppedPages, and their Python mirrors) returned early or skipped the page whenever the life cycle was not an AttnLifeCycle. Offload, onboard, intra-device copy and host-tier drop of recurrent (SSM/KDA) state were therefore never recorded, and any recurrent-cache iteration statistics read back as zeros. Iteration statistics are already keyed by life cycle, so recurrent movement stays distinguishable from attention movement without the filter. Keep the filter only where it is semantically required: the global cache-hit counters (allocTotalBlocks / allocNewBlocks) and the block-reuse hit/miss range accounting stay attention-only. Observability only, no change to allocation or reuse behavior, and reachable only under use_kv_cache_manager_v2=True. Test: tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py drives the Python recorders directly (no GPU needed) and asserts offload and host-drop are reported for both life-cycle kinds while the global cache-hit counters remain attention-only. The C++ recorders have no equivalent CPU-only seam. Signed-off-by: Brian Nguyen --- .../kv_cache_manager_v2/kvCache.cpp | 24 ++-- .../kv_cache_manager_v2/_core/_kv_cache.py | 35 ++--- .../test_kv_cache_stats_life_cycles.py | 124 ++++++++++++++++++ 3 files changed, 154 insertions(+), 29 deletions(-) create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 6beb2bd4f9eb..e77d104abc35 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/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(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(mManager->lifeCycles().getLifeCycle(lifeCycle))) - { - continue; - } + bool const isAttention = std::holds_alternative(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> const& pages, Cac for (auto const& page : pages) { LifeCycleId const lifeCycle = page->lifeCycle; - if (!std::holds_alternative(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)) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index d71bf3db6acb..9992e3d0bbff 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/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. diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py new file mode 100644 index 000000000000..fd75fefe9f41 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TRTLLM-15217: SSM/recurrent life cycles must appear in V2 iteration stats. + +The page-movement recorders used to drop every non-attention life cycle, which +made KDA (Kimi K3) recurrent-state offload / onboard / drop invisible in +iteration statistics. These tests drive the recorders directly with a +duck-typed stand-in so they run without a GPU or an allocated cache. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.runtime.kv_cache_manager_v2._common import GPU_LEVEL, CacheLevel +from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache import KvCache +from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import ( + AttnLifeCycle, + LifeCycleId, + SsmLifeCycle, +) + +ATTN_LC = LifeCycleId(0) +SSM_LC = LifeCycleId(1) +PAGE_BYTES = 16 +HOST_LEVEL = CacheLevel(GPU_LEVEL + 1) + + +def _make_recorder(): + """Duck-typed KvCache exposing only what the stats recorders touch. + + The recording methods are bound off the real class, so the life-cycle + filtering under test is the production implementation. + """ + committed = [] + life_cycles = {ATTN_LC: AttnLifeCycle(None, 0), SSM_LC: SsmLifeCycle()} + manager = SimpleNamespace( + _life_cycles=SimpleNamespace(get_life_cycle=life_cycles.__getitem__), + _storage=SimpleNamespace( + get_pool_group_index=lambda life_cycle: life_cycle, + slot_size=lambda _pool_group: [PAGE_BYTES], + ), + commit_stats=lambda stats, by_life_cycle: committed.append((stats, by_life_cycle)), + ) + recorder = SimpleNamespace(manager=manager) + recorder._should_record_stats = lambda: True + for name in ( + "_is_attention_life_cycle", + "_record_direct_iteration_stats", + "_record_migrated_slots", + "_record_dropped_pages", + ): + setattr(recorder, name, getattr(KvCache, name).__get__(recorder)) + return recorder, committed + + +@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) +def test_offload_is_recorded_for_every_life_cycle(life_cycle: LifeCycleId) -> None: + recorder, committed = _make_recorder() + page = SimpleNamespace(life_cycle=life_cycle) + + recorder._record_migrated_slots([page], [object()], GPU_LEVEL, HOST_LEVEL) + + assert len(committed) == 1 + _, by_life_cycle = committed[0] + assert set(by_life_cycle) == {life_cycle} + assert by_life_cycle[life_cycle].iter_offload_blocks == 1 + assert by_life_cycle[life_cycle].iter_offload_bytes == PAGE_BYTES + + +@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) +def test_host_drop_is_recorded_for_every_life_cycle(life_cycle: LifeCycleId) -> None: + recorder, committed = _make_recorder() + page = SimpleNamespace(life_cycle=life_cycle) + + recorder._record_dropped_pages([page], HOST_LEVEL) + + assert len(committed) == 1 + _, by_life_cycle = committed[0] + assert set(by_life_cycle) == {life_cycle} + assert by_life_cycle[life_cycle].iter_host_dropped_blocks == 1 + assert by_life_cycle[life_cycle].iter_host_dropped_bytes == PAGE_BYTES + + +def test_onboard_counts_globally_only_for_attention() -> None: + """Onboard is per-life-cycle; global cache-hit counters are attention-only. + + alloc_total_blocks / alloc_new_blocks feed the global cache-hit rate, which + is defined over attention blocks only. + """ + recorder, committed = _make_recorder() + + recorder._record_migrated_slots( + [SimpleNamespace(life_cycle=SSM_LC)], [object()], HOST_LEVEL, GPU_LEVEL + ) + recorder._record_migrated_slots( + [SimpleNamespace(life_cycle=ATTN_LC)], [object()], HOST_LEVEL, GPU_LEVEL + ) + + assert len(committed) == 2 + ssm_stats, ssm_by_life_cycle = committed[0] + attn_stats, attn_by_life_cycle = committed[1] + + for life_cycle, by_life_cycle in ((SSM_LC, ssm_by_life_cycle), (ATTN_LC, attn_by_life_cycle)): + assert by_life_cycle[life_cycle].iter_onboard_blocks == 1 + assert by_life_cycle[life_cycle].iter_onboard_bytes == PAGE_BYTES + + assert ssm_stats.alloc_total_blocks == 0 + assert ssm_stats.alloc_new_blocks == 0 + assert attn_stats.alloc_total_blocks == 1 + assert attn_stats.alloc_new_blocks == 1 From 9a11822f1b7cfa5a1cc03c7822760abd6cd22c07 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 20:07:32 -0700 Subject: [PATCH 3/6] Address trivial review comments Signed-off-by: Brian Nguyen --- .../runtime/kv_cache_manager_v2/_core/_kv_cache.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 9992e3d0bbff..44944e00c953 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1293,13 +1293,13 @@ 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)), - ), - ) + 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: From 9e75e20c53810c28811b2c1710b25374cea8c206 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 20:40:51 -0700 Subject: [PATCH 4/6] [TRTLLM-15217][test] cover SSM deferred-copy iteration stats on both backends Review follow-up. The Python backend used to record iter_intra_device_copy_{blocks,bytes} only for non-SSM life cycles in the resume() deferred-copy loop, unlike the default C++ backend; the previous commit aligned the call sites. Add the coverage that was missing: - test_kv_cache_stats_life_cycles.py now drives _record_direct_iteration_stats (it was bound into the stand-in but never exercised) for both attention and SSM life cycles. - TestSSMSupport.test_ssm_resume_records_intra_device_copy asserts the SSM intra-device-copy delta end to end against the selected backend, so the default C++ implementation is covered and a Python-backend parity break fails the same test (verified: it fails against the pre-fix Python backend). Also comment why the recorder call sits outside the ssm_lc_id guard. Signed-off-by: Brian Nguyen --- .../kv_cache_manager_v2/_core/_kv_cache.py | 3 ++ .../test_kv_cache_manager_v2.py | 38 +++++++++++++++++++ .../test_kv_cache_stats_life_cycles.py | 27 +++++++++++++ 3 files changed, 68 insertions(+) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 44944e00c953..2fe111b1bd46 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1293,6 +1293,9 @@ def resume(self, cuda_stream: CudaStream | None = None) -> bool: ) if changed: self.manager.mark_stats_dirty(self.id) + # 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( diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 48c10a006fc0..e0519b3fd2f2 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/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() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py index fd75fefe9f41..661f65693548 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py @@ -32,6 +32,7 @@ LifeCycleId, SsmLifeCycle, ) +from tensorrt_llm.runtime.kv_cache_manager_v2._stats import KVCacheIterationStatsDelta ATTN_LC = LifeCycleId(0) SSM_LC = LifeCycleId(1) @@ -95,6 +96,32 @@ def test_host_drop_is_recorded_for_every_life_cycle(life_cycle: LifeCycleId) -> assert by_life_cycle[life_cycle].iter_host_dropped_bytes == PAGE_BYTES +@pytest.mark.parametrize("life_cycle", [ATTN_LC, SSM_LC]) +def test_direct_iteration_stats_are_recorded_for_every_life_cycle( + life_cycle: LifeCycleId, +) -> None: + """SSM deferred copies must reach iteration stats. + + The resume() deferred copy reports iter_intra_device_copy_* through this + recorder for SSM life cycles too, matching the C++ backend. + """ + recorder, committed = _make_recorder() + + recorder._record_direct_iteration_stats( + life_cycle, + KVCacheIterationStatsDelta( + iter_intra_device_copy_blocks=1, + iter_intra_device_copy_bytes=PAGE_BYTES, + ), + ) + + assert len(committed) == 1 + _, by_life_cycle = committed[0] + assert set(by_life_cycle) == {life_cycle} + assert by_life_cycle[life_cycle].iter_intra_device_copy_blocks == 1 + assert by_life_cycle[life_cycle].iter_intra_device_copy_bytes == PAGE_BYTES + + def test_onboard_counts_globally_only_for_attention() -> None: """Onboard is per-life-cycle; global cache-hit counters are attention-only. From 6209e47a5a611cc0d5911e572c9f138c6c6cd75c Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 20:19:28 -0700 Subject: [PATCH 5/6] [TRTLLM-15216][fix] reject model_type in MambaHybridCacheManagerV2 and dedupe layout dispatch Review follow-up, two hardenings against the bug class fixed earlier in this PR: - MambaHybridCacheManagerV2.__init__ now raises TypeError when handed the V1 managers' model_type kwarg instead of conv_state_layout. Silently absorbing it into **kwargs is exactly how the wrong-layout bug went unnoticed. - The three copies of the "V2 takes conv_state_layout, V1 takes model_type" dispatch in _create_kv_cache_manager (kimi, nemotron, qwen3_hybrid branches) collapse into one helper with a single model_type-to-layout mapping, so a future hybrid branch cannot pick one convention and drop the other. Test: test_v2_manager_rejects_model_type_kwarg; existing kimi/qwen3 layout-selection tests cover the helper refactor. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/pyexecutor/_util.py | 47 ++++++++++++------- .../_torch/pyexecutor/mamba_cache_manager.py | 8 ++++ .../executor/test_mamba_cache_manager.py | 28 +++++++++++ 3 files changed, 66 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 12d4646dcfcb..f193e270c614 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, @@ -2222,15 +2245,9 @@ def _create_kv_cache_manager( 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. The V1 - # managers select that layout by `model_type`; MambaHybridCacheManagerV2 - # takes an explicit `conv_state_layout` and silently defaults to - # "x_b_c" (and absorbs `model_type` into **kwargs), so V2 must be given - # the layout by name. Mirrors the qwen3_hybrid branch below. - if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - kimi_extra_kwargs["conv_state_layout"] = "q_k_v" - else: - kimi_extra_kwargs["model_type"] = "qwen3_next" + # 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, @@ -2371,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, @@ -2480,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, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 12df9681e6ab..10a7dba55540 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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 diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index ff37ee7d92ca..21317da11c4c 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -363,6 +363,34 @@ def __init__(self, *args: object, **kwargs: object) -> None: 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"), [ From 99b20ceb2729191fb3bc9136f61a8d3fa7610f4a Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 20:37:33 -0700 Subject: [PATCH 6/6] [TRTLLM-15217][fix] repair test_kv_cache_stats_life_cycles import The module imported KvCache but the class is named _KVCache, so the file failed pytest collection and none of its tests ever ran. Review follow-up discovered while adding the _record_direct_iteration_stats case. Signed-off-by: Brian Nguyen --- .../test_kv_cache_stats_life_cycles.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py index 661f65693548..a5c841e7ed2f 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py @@ -26,7 +26,7 @@ import pytest from tensorrt_llm.runtime.kv_cache_manager_v2._common import GPU_LEVEL, CacheLevel -from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache import KvCache +from tensorrt_llm.runtime.kv_cache_manager_v2._core._kv_cache import _KVCache from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import ( AttnLifeCycle, LifeCycleId, @@ -41,7 +41,7 @@ def _make_recorder(): - """Duck-typed KvCache exposing only what the stats recorders touch. + """Duck-typed _KVCache exposing only what the stats recorders touch. The recording methods are bound off the real class, so the life-cycle filtering under test is the production implementation. @@ -64,7 +64,7 @@ def _make_recorder(): "_record_migrated_slots", "_record_dropped_pages", ): - setattr(recorder, name, getattr(KvCache, name).__get__(recorder)) + setattr(recorder, name, getattr(_KVCache, name).__get__(recorder)) return recorder, committed