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/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..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, @@ -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, 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/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..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 @@ -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: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 81db2977398d..21317da11c4c 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/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"), [ 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 new file mode 100644 index 000000000000..a5c841e7ed2f --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py @@ -0,0 +1,151 @@ +# 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, +) +from tensorrt_llm.runtime.kv_cache_manager_v2._stats import KVCacheIterationStatsDelta + +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 + + +@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. + + 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