diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index a78adfa36c02..56b6b1301536 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -90,6 +90,7 @@ Models that select the V2 manager by default: | GPT-OSS | Sliding window on every other layer (VSWA), so the sliding-window and full-attention pools are sized independently | | Gemma3 / Gemma4 (text and multimodal) | Alternating sliding-window and full-attention layers (VSWA); same independent pool sizing | | Llama / Llama4 | Uniform KV pool layout; chunked attention does not partition the pools | +| KimiLinear | Hybrid KDA recurrent state and paged MLA cache | Separately, Gemma4 hybrid attention and sparse-attention models are routed to V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 60d840506301..dc409bcb18ed 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -18,7 +18,7 @@ Caching ------- -KDA states live on the mamba side of a ``MixedMambaHybridCacheManager`` +KDA states live on the mamba side of a ``MambaHybridCacheManagerV2`` (wired 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 @@ -68,13 +68,12 @@ 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). +``MambaHybridCacheManagerV2`` (per-block KDA state snapshots every +``mamba_state_cache_interval`` tokens, 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 @@ -2103,9 +2102,7 @@ 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. + # explicit opt-in with per-block KDA state snapshots. # - 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 @@ -2117,6 +2114,14 @@ def get_model_defaults(cls, llm_args) -> dict: } } + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for KimiLinear.""" + return "V2" + @classmethod def get_preferred_transceiver_runtime( cls, 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 d9902377405b..c2b18556b718 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 @@ -844,10 +844,10 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) -def test_kimi_without_v2_preference_uses_mixed_manager( +def test_kimi_model_preference_uses_v2_manager( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Kimi K3 uses separate KV and recurrent-state pools for SA decoding.""" + """Kimi's model preference selects V2 for its KDA and MLA cache.""" from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) @@ -862,13 +862,13 @@ def test_kimi_without_v2_preference_uses_mixed_manager( ) resolved = _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM) - assert resolved is False - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + assert resolved is True + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.kv_cache_config.tokens_per_block == 64 assert ( get_kv_cache_manager_cls(_kimi_model_config(), llm_args.kv_cache_config) - is MixedMambaHybridCacheManager + is MambaHybridCacheManagerV2 ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index c1e09bd89930..152efe3b11e6 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -827,6 +827,7 @@ def test_registered_models_prefer_v2(self) -> None: "MistralLarge3ForCausalLM", "DeepseekV4ForCausalLM", "KimiK25ForConditionalGeneration", + "KimiLinearForCausalLM", "MiniMaxM2ForCausalLM", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", @@ -871,6 +872,7 @@ def test_registered_models_keep_v2_on_nixl(self) -> None: "MistralLarge3ForCausalLM", "GptOssForCausalLM", "KimiK25ForConditionalGeneration", + "KimiLinearForCausalLM", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", "Qwen3NextForCausalLM",