From 1876a7e69e5998c7ca18be4fcca4f7b0e367e1ca Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:11:06 -0700 Subject: [PATCH 1/5] [None][test] Exercise default KVCM V2 in Kimi K3 GPQA Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/integration/defs/accuracy/test_kimi3.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index 9fc72555efd5..b9680e2ab59e 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/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 From 0cc4d14a2963ca3ffc08b57471cbbfcceaf5cf57 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 02:49:04 -0700 Subject: [PATCH 2/5] [None][test] Cover Kimi V2 disaggregated cache transfer Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 +- .../kv_cache/test_mamba_cache_manager.py | 29 ++++++++- .../disaggregated/test_kda_mamba_transfer.py | 59 +++++++++++++------ 3 files changed, 69 insertions(+), 21 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index cf5d0e06e5f8..4ae1a418040d 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 2bdea0b58426..fb9a35242907 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -1027,7 +1027,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) diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 9494849f23fd..84551de70776 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/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", @@ -639,8 +659,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 +701,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 +731,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 +751,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) From 1a6a6c3de59934ffcc683667fdfea61808002835 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:05:24 -0700 Subject: [PATCH 3/5] [None][test] Fix Kimi disaggregated cache test initialization Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache/test_mamba_cache_manager.py | 2 ++ .../disaggregated/test_kda_mamba_transfer.py | 27 ++++++++++++++----- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index fb9a35242907..d31c46767fee 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/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, ) diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 84551de70776..5606fb1a558b 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -636,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) From 7838bb89266ac85ee602e2ef7c645d8eabdddcf7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:42:36 -0700 Subject: [PATCH 4/5] [None][docs] Clarify Kimi default V2 cache routing Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 20 +++++++++---------- tensorrt_llm/_torch/pyexecutor/_util.py | 18 ++++++----------- 2 files changed, 15 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c726cb6c290f..05a33f13ee9e 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -18,8 +18,8 @@ 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`` by default +(selected in ``pyexecutor/_util.py``): per layer, 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 +67,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 same 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 +2288,9 @@ 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 + # unless the user explicitly selects a V1 manager. # - 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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7aedaf8a81e6..e7f2ae5eee28 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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,11 @@ def get_kv_cache_manager_cls( "hybrid managers do not implement per-token verify-group " "bookkeeping.") + # Kimi defaults to MambaHybridCacheManagerV2 with or without reuse. + # Explicit V1 aggregated serving uses Mixed without reuse and the + # C++ hybrid manager with reuse. Disaggregated serving requires + # Python NIXL below, using V2 by default or Mixed for explicit 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( From 70047d853c3aa4a42baffc8587943a961404b3c7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:56:16 -0700 Subject: [PATCH 5/5] [None][docs] Distinguish Kimi V2 defaults from auto resolution Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 12 ++++++++---- tensorrt_llm/_torch/pyexecutor/_util.py | 11 ++++++----- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 05a33f13ee9e..07d232f62afe 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -18,8 +18,11 @@ Caching ------- -KDA states use ``MambaHybridCacheManagerV2`` by default -(selected 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,7 +70,7 @@ 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`` on the same V2 +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. @@ -2290,7 +2293,8 @@ def _setup_helix_mappings( def get_model_defaults(cls, llm_args) -> dict: # - Keep block reuse opt-in because its combination with SA # speculative decoding is not validated. Both settings use V2 - # unless the user explicitly selects a V1 manager. + # 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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e7f2ae5eee28..bc3a62f7d171 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -288,11 +288,12 @@ def get_kv_cache_manager_cls( "hybrid managers do not implement per-token verify-group " "bookkeeping.") - # Kimi defaults to MambaHybridCacheManagerV2 with or without reuse. - # Explicit V1 aggregated serving uses Mixed without reuse and the - # C++ hybrid manager with reuse. Disaggregated serving requires - # Python NIXL below, using V2 by default or Mixed for explicit V1 - # without reuse. + # 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(