From 8b17461eb52089ed61b53994c564c702a6e64952 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:32:32 -0700 Subject: [PATCH 1/3] [None][fix] Enable KVCM V2 by default for KimiLinear Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 28 +++--- tensorrt_llm/_torch/pyexecutor/_util.py | 14 +-- .../references/acceptance_length.yaml | 4 + tests/integration/defs/accuracy/test_kimi3.py | 84 +++++++++++++++- .../test_lists/test-db/l0_b200.yml | 1 + .../kv_cache/test_mamba_cache_manager.py | 97 +++++++++++++++++-- tests/unittest/llmapi/test_llm_args.py | 4 + 7 files changed, 201 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 60d840506301..ff7a276a341a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -18,10 +18,10 @@ 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 +``[3 * num_heads * head_dim, W - 1]`` bf16 (the convolution history, +sections ``[q | k | v]``) and a delta-rule recurrent slot of ``[num_heads, head_dim, head_dim]`` fp32 (``[H, V, K]``, the ``state_v_first`` FLA layout). MLA layers use the paged-KV side with ``num_kv_heads=1`` and ``head_dim = kv_lora_rank + qk_rope_head_dim`` (576), @@ -67,14 +67,12 @@ 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``; the hybrid manager +tracks recurrent-state snapshots alongside the MLA cache blocks. 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 @@ -2102,10 +2100,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. + # - enable_block_reuse defaults off: keep reuse an explicit opt-in. + # V2 manages both KDA state and the paged MLA cache, including + # speculative-decoding state. # - 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,13 @@ def get_model_defaults(cls, llm_args) -> dict: } } + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + return "V2" + @classmethod def get_preferred_transceiver_runtime( cls, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 6bc67d8eb825..24f736449be7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -177,17 +177,9 @@ def get_kv_cache_manager_cls( "use_kv_cache_manager_v2=True; V1 supports only " "periodic_snapshot_interval.") - # 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. + # Kimi K3's V1 compatibility route uses the unified C++ pool for + # block reuse and separate KV / recurrent-state pools otherwise. + # V2 and disaggregated serving use the shared hybrid routing below. if is_kimi_linear(config) and not use_v2 and not is_disagg: if kv_cache_config.enable_block_reuse: logger.info( diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 46afc55d3e90..47135d878409 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -24,6 +24,10 @@ TestQwen3_5_4B::test_dflash: TestKimiK3::test_w4a16_mxfp4[sa]: ref_al: 1.318 min_al: 1.15 +TestKimiK3::test_text_only_gsm8k_tep16[sa]: + # Full GSM8K with the text model, TEP16 on 16 B200 GPUs. + ref_al: 1.337462 + min_al: 1.15 TestKimiK3::test_w4a16_mxfp4[dspark]: ref_al: 4.354 min_al: 4.1363 diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index 52bb992a70f4..97f26d03ecb5 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +from pathlib import Path + import pytest from tensorrt_llm import LLM from tensorrt_llm._torch.configs import KimiK3Config, KimiLinearConfig +from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config from tensorrt_llm.llmapi import ( CudaGraphConfig, @@ -161,6 +165,82 @@ def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: llm, ) + @skip_pre_blackwell + @pytest.mark.skip_less_mpi_world_size(16) + @pytest.mark.skip_less_device_memory(140000) + @pytest.mark.parametrize("mode", ["baseline", "sa"]) + def test_text_only_gsm8k_tep16( + self, mode: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Validate the complete text model and automatic V2 selection. + + For a multinode run, set pytest's --basetemp to a shared directory + mounted at the same path on every worker. The checkpoint overlay must + remain visible until all workers have finished loading the model. + """ + source = Path(self.MODEL_PATH) + config = json.loads((source / "config.json").read_text()) + config["language_model_only"] = True + checkpoint = tmp_path / "kimi-linear" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps(config)) + for entry in source.iterdir(): + if entry.name != "config.json": + (checkpoint / entry.name).symlink_to(entry) + + resolved = load_pretrained_config(str(checkpoint), trust_remote_code=True) + assert isinstance(resolved, KimiLinearConfig) + assert resolved.architectures == ["KimiLinearForCausalLM"] + model_class = get_registered_model_class(resolved.architectures[0]) + assert model_class.get_preferred_kv_cache_manager_version(resolved) == "V2" + assert resolved.num_hidden_layers == 93 + assert len(resolved.linear_attn_config["kda_layers"]) == 69 + assert len(resolved.linear_attn_config["full_attn_layers"]) == 24 + + if mode == "sa": + monkeypatch.setenv("TLLM_EVAL_SPEC_STATS", "1") + batch_size = 8 if mode == "sa" else 32 + with LLM( + str(checkpoint), + # Apply the reference precision on every MPI worker, including + # workers that were started before this test process. + env_overrides={"KIMI_K3_ROUTER_BF16": "0", "KIMI_K3_FP8_WEIGHT_READ": "0"}, + tensor_parallel_size=16, + moe_expert_parallel_size=16, + enable_attention_dp=False, + allreduce_strategy="NCCL", + max_batch_size=batch_size, + max_num_tokens=8192, + max_seq_len=8192, + trust_remote_code=True, + enable_chunked_prefill=mode == "baseline", + disable_overlap_scheduler=mode == "sa", + cuda_graph_config=CudaGraphConfig( + enable_padding=mode == "baseline", max_batch_size=batch_size + ), + moe_config=MoeConfig( + backend="TRTLLM", + max_num_tokens=33024, + use_low_precision_moe_combine=True, + ), + speculative_config=SADecodingConfig(max_draft_len=2) if mode == "sa" else None, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=0.25, + tokens_per_block=64, + ), + enable_iter_perf_stats=True, + max_stats_len=-1 if mode == "sa" else 256, + return_perf_metrics=True, + ) as llm: + self._assert_resolved_args(llm, mode, attention_dp=False) + # Resolution happens in the MPI workers; keep the frontend in + # auto mode so this exercises the registered model preference. + assert llm.args.kv_cache_config.use_kv_cache_manager_v2 == "auto" + GSM8K(self.MODEL_NAME).evaluate(llm) + if mode == "sa": + assert_acceptance_length_for_llm("TestKimiK3::test_text_only_gsm8k_tep16[sa]", llm) + @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(16) @pytest.mark.skip_less_device_memory(200000) @@ -220,8 +300,8 @@ def _assert_checkpoint_routing(self) -> None: assert getattr(config.text_config, "sliding_window", None) is None @staticmethod - def _assert_resolved_args(llm: LLM, mode: str) -> None: - assert llm.args.enable_attention_dp is True + def _assert_resolved_args(llm: LLM, mode: str, attention_dp: bool = True) -> None: + assert llm.args.enable_attention_dp is attention_dp assert llm.args.kv_cache_config.enable_block_reuse is (mode == "reuse") # K3's routed-expert quantization is nested in the composite checkpoint # and is not represented by the modelopt-style args quant_algo field. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 55c319f1d36b..a654a0f23f37 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -119,6 +119,7 @@ l0_b200: - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py # ------------- Kimi K3 (KimiLinear) unit tests --------------- + - unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_auto_v2_allocates_kda_and_mla_cache - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py - unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py 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..992ca2793a6e 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 @@ -68,6 +68,7 @@ KvCacheConfig, MambaStateConfig, MTPDecodingConfig, + SADecodingConfig, TorchLlmArgs, ) from tensorrt_llm.llmapi.llm_utils import ( @@ -844,10 +845,15 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) -def test_kimi_without_v2_preference_uses_mixed_manager( +@pytest.mark.cpu_only +@pytest.mark.parametrize("manager_setting", ["auto", True, False]) +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_kimi_model_preference_selects_cache_manager( monkeypatch: pytest.MonkeyPatch, + manager_setting: str | bool, + is_disagg: bool, ) -> None: - """Kimi K3 uses separate KV and recurrent-state pools for SA decoding.""" + """Resolve the registered text model before selecting its hybrid manager.""" from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) @@ -856,19 +862,31 @@ def test_kimi_without_v2_preference_uses_mixed_manager( llm_args = TorchLlmArgs( model="/tmp/dummy_model", kv_cache_config=KvCacheConfig( + use_kv_cache_manager_v2=manager_setting, enable_block_reuse=False, tokens_per_block=64, ), + cache_transceiver_config=CacheTransceiverConfig(backend="NIXL") if is_disagg else None, ) + _resolve_transceiver_runtime_auto(llm_args, KimiLinearForCausalLM) 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 + expected_v2 = manager_setting is not False + assert resolved is expected_v2 + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.kv_cache_config.tokens_per_block == 64 + if is_disagg: + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + expected_manager = MambaHybridCacheManagerV2 if expected_v2 else MixedMambaHybridCacheManager assert ( - get_kv_cache_manager_cls(_kimi_model_config(), llm_args.kv_cache_config) - is MixedMambaHybridCacheManager + get_kv_cache_manager_cls( + _kimi_model_config(), + llm_args.kv_cache_config, + is_disagg=is_disagg, + cache_transceiver_config=llm_args.cache_transceiver_config, + ) + is expected_manager ) @@ -880,6 +898,73 @@ def test_kimi_preferred_transceiver_runtime() -> None: assert KimiLinearForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" +@skip_no_cuda +@pytest.mark.parametrize("use_sa", [False, True]) +def test_kimi_auto_v2_allocates_kda_and_mla_cache( + monkeypatch: pytest.MonkeyPatch, use_sa: bool +) -> None: + """Construct the resolved Kimi manager with real KDA and MLA storage.""" + 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( + max_tokens=1024, + enable_block_reuse=False, + tokens_per_block=64, + ), + speculative_config=SADecodingConfig(max_draft_len=2) if use_sa else None, + ) + _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM) + model_config = _kimi_model_config() + manager_cls = get_kv_cache_manager_cls(model_config, llm_args.kv_cache_config) + manager = _create_kv_cache_manager( + model_engine=None, + kv_cache_manager_cls=manager_cls, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + kv_cache_config=llm_args.kv_cache_config, + tokens_per_block=64, + max_seq_len=256, + max_batch_size=4, + spec_config=llm_args.speculative_config, + 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, + ) + try: + assert isinstance(manager, MambaHybridCacheManagerV2) + assert manager.kv_cache_type == CacheTypeCpp.SELFKONLY + assert manager.local_num_mamba_layers == 2 + assert manager.blocks_in_primary_pool > 0 + for layer_idx in (1, 3): + buffers = manager.get_buffers(layer_idx) + assert buffers.shape[1:] == (1, 64, 1, 40) + assert buffers.dtype is torch.bfloat16 + assert manager.mamba_pp_layers == [0, 2] + for layer_idx in manager.mamba_pp_layers: + cache = manager.mamba_layer_cache(layer_idx) + assert cache.conv.shape[1:] == (96, 3) + assert cache.conv.dtype is torch.bfloat16 + assert cache.temporal.shape[1:] == (4, 8, 8) + assert cache.temporal.dtype is torch.float32 + + manager.add_dummy_requests([123], token_nums=[8], is_gen=False) + indices = manager.get_state_indices([123], [False]) + assert indices[0] >= 0 + assert manager.cuda_state_indices[0].item() == indices[0] + if use_sa: + cache = manager.mamba_layer_cache(0) + assert cache.kda_qkg_cache is not None or cache.intermediate_ssm is not None + finally: + manager.shutdown() + + @pytest.mark.parametrize( "cache_transceiver_config", [ diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index c1e09bd89930..b5d1a53b24f1 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -827,6 +827,8 @@ def test_registered_models_prefer_v2(self) -> None: "MistralLarge3ForCausalLM", "DeepseekV4ForCausalLM", "KimiK25ForConditionalGeneration", + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", "MiniMaxM2ForCausalLM", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", @@ -871,6 +873,8 @@ def test_registered_models_keep_v2_on_nixl(self) -> None: "MistralLarge3ForCausalLM", "GptOssForCausalLM", "KimiK25ForConditionalGeneration", + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", "Qwen3NextForCausalLM", From b7b068a853d74440746046834e09e6f20eee2096 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:41:24 -0700 Subject: [PATCH 2/3] [None][fix] Reuse existing Kimi tests for the V2 default Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- docs/source/features/kvcache.md | 1 + .../_torch/models/modeling_kimi_linear.py | 1 + .../references/acceptance_length.yaml | 4 - tests/integration/defs/accuracy/test_kimi3.py | 84 +--------------- .../test_lists/test-db/l0_b200.yml | 1 - .../kv_cache/test_mamba_cache_manager.py | 97 ++----------------- tests/unittest/llmapi/test_llm_args.py | 4 - 7 files changed, 10 insertions(+), 182 deletions(-) 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 ff7a276a341a..3c1e9809628d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2119,6 +2119,7 @@ def get_preferred_kv_cache_manager_version( cls, pretrained_config: Any = None, ) -> Literal["V2"]: + """Prefer KV cache manager V2 for KimiLinear.""" return "V2" @classmethod diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 47135d878409..46afc55d3e90 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -24,10 +24,6 @@ TestQwen3_5_4B::test_dflash: TestKimiK3::test_w4a16_mxfp4[sa]: ref_al: 1.318 min_al: 1.15 -TestKimiK3::test_text_only_gsm8k_tep16[sa]: - # Full GSM8K with the text model, TEP16 on 16 B200 GPUs. - ref_al: 1.337462 - min_al: 1.15 TestKimiK3::test_w4a16_mxfp4[dspark]: ref_al: 4.354 min_al: 4.1363 diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index 97f26d03ecb5..52bb992a70f4 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -13,14 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -from pathlib import Path - import pytest from tensorrt_llm import LLM from tensorrt_llm._torch.configs import KimiK3Config, KimiLinearConfig -from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config from tensorrt_llm.llmapi import ( CudaGraphConfig, @@ -165,82 +161,6 @@ def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: llm, ) - @skip_pre_blackwell - @pytest.mark.skip_less_mpi_world_size(16) - @pytest.mark.skip_less_device_memory(140000) - @pytest.mark.parametrize("mode", ["baseline", "sa"]) - def test_text_only_gsm8k_tep16( - self, mode: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Validate the complete text model and automatic V2 selection. - - For a multinode run, set pytest's --basetemp to a shared directory - mounted at the same path on every worker. The checkpoint overlay must - remain visible until all workers have finished loading the model. - """ - source = Path(self.MODEL_PATH) - config = json.loads((source / "config.json").read_text()) - config["language_model_only"] = True - checkpoint = tmp_path / "kimi-linear" - checkpoint.mkdir() - (checkpoint / "config.json").write_text(json.dumps(config)) - for entry in source.iterdir(): - if entry.name != "config.json": - (checkpoint / entry.name).symlink_to(entry) - - resolved = load_pretrained_config(str(checkpoint), trust_remote_code=True) - assert isinstance(resolved, KimiLinearConfig) - assert resolved.architectures == ["KimiLinearForCausalLM"] - model_class = get_registered_model_class(resolved.architectures[0]) - assert model_class.get_preferred_kv_cache_manager_version(resolved) == "V2" - assert resolved.num_hidden_layers == 93 - assert len(resolved.linear_attn_config["kda_layers"]) == 69 - assert len(resolved.linear_attn_config["full_attn_layers"]) == 24 - - if mode == "sa": - monkeypatch.setenv("TLLM_EVAL_SPEC_STATS", "1") - batch_size = 8 if mode == "sa" else 32 - with LLM( - str(checkpoint), - # Apply the reference precision on every MPI worker, including - # workers that were started before this test process. - env_overrides={"KIMI_K3_ROUTER_BF16": "0", "KIMI_K3_FP8_WEIGHT_READ": "0"}, - tensor_parallel_size=16, - moe_expert_parallel_size=16, - enable_attention_dp=False, - allreduce_strategy="NCCL", - max_batch_size=batch_size, - max_num_tokens=8192, - max_seq_len=8192, - trust_remote_code=True, - enable_chunked_prefill=mode == "baseline", - disable_overlap_scheduler=mode == "sa", - cuda_graph_config=CudaGraphConfig( - enable_padding=mode == "baseline", max_batch_size=batch_size - ), - moe_config=MoeConfig( - backend="TRTLLM", - max_num_tokens=33024, - use_low_precision_moe_combine=True, - ), - speculative_config=SADecodingConfig(max_draft_len=2) if mode == "sa" else None, - kv_cache_config=KvCacheConfig( - enable_block_reuse=False, - free_gpu_memory_fraction=0.25, - tokens_per_block=64, - ), - enable_iter_perf_stats=True, - max_stats_len=-1 if mode == "sa" else 256, - return_perf_metrics=True, - ) as llm: - self._assert_resolved_args(llm, mode, attention_dp=False) - # Resolution happens in the MPI workers; keep the frontend in - # auto mode so this exercises the registered model preference. - assert llm.args.kv_cache_config.use_kv_cache_manager_v2 == "auto" - GSM8K(self.MODEL_NAME).evaluate(llm) - if mode == "sa": - assert_acceptance_length_for_llm("TestKimiK3::test_text_only_gsm8k_tep16[sa]", llm) - @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(16) @pytest.mark.skip_less_device_memory(200000) @@ -300,8 +220,8 @@ def _assert_checkpoint_routing(self) -> None: assert getattr(config.text_config, "sliding_window", None) is None @staticmethod - def _assert_resolved_args(llm: LLM, mode: str, attention_dp: bool = True) -> None: - assert llm.args.enable_attention_dp is attention_dp + def _assert_resolved_args(llm: LLM, mode: str) -> None: + assert llm.args.enable_attention_dp is True assert llm.args.kv_cache_config.enable_block_reuse is (mode == "reuse") # K3's routed-expert quantization is nested in the composite checkpoint # and is not represented by the modelopt-style args quant_algo field. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a654a0f23f37..55c319f1d36b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -119,7 +119,6 @@ l0_b200: - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py # ------------- Kimi K3 (KimiLinear) unit tests --------------- - - unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_kimi_auto_v2_allocates_kda_and_mla_cache - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py - unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py 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 992ca2793a6e..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 @@ -68,7 +68,6 @@ KvCacheConfig, MambaStateConfig, MTPDecodingConfig, - SADecodingConfig, TorchLlmArgs, ) from tensorrt_llm.llmapi.llm_utils import ( @@ -845,15 +844,10 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) -@pytest.mark.cpu_only -@pytest.mark.parametrize("manager_setting", ["auto", True, False]) -@pytest.mark.parametrize("is_disagg", [False, True]) -def test_kimi_model_preference_selects_cache_manager( +def test_kimi_model_preference_uses_v2_manager( monkeypatch: pytest.MonkeyPatch, - manager_setting: str | bool, - is_disagg: bool, ) -> None: - """Resolve the registered text model before selecting its hybrid manager.""" + """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,31 +856,19 @@ def test_kimi_model_preference_selects_cache_manager( llm_args = TorchLlmArgs( model="/tmp/dummy_model", kv_cache_config=KvCacheConfig( - use_kv_cache_manager_v2=manager_setting, enable_block_reuse=False, tokens_per_block=64, ), - cache_transceiver_config=CacheTransceiverConfig(backend="NIXL") if is_disagg else None, ) - _resolve_transceiver_runtime_auto(llm_args, KimiLinearForCausalLM) resolved = _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM) - expected_v2 = manager_setting is not False - assert resolved is expected_v2 - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 + 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 - if is_disagg: - assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" - expected_manager = MambaHybridCacheManagerV2 if expected_v2 else MixedMambaHybridCacheManager assert ( - get_kv_cache_manager_cls( - _kimi_model_config(), - llm_args.kv_cache_config, - is_disagg=is_disagg, - cache_transceiver_config=llm_args.cache_transceiver_config, - ) - is expected_manager + get_kv_cache_manager_cls(_kimi_model_config(), llm_args.kv_cache_config) + is MambaHybridCacheManagerV2 ) @@ -898,73 +880,6 @@ def test_kimi_preferred_transceiver_runtime() -> None: assert KimiLinearForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" -@skip_no_cuda -@pytest.mark.parametrize("use_sa", [False, True]) -def test_kimi_auto_v2_allocates_kda_and_mla_cache( - monkeypatch: pytest.MonkeyPatch, use_sa: bool -) -> None: - """Construct the resolved Kimi manager with real KDA and MLA storage.""" - 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( - max_tokens=1024, - enable_block_reuse=False, - tokens_per_block=64, - ), - speculative_config=SADecodingConfig(max_draft_len=2) if use_sa else None, - ) - _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM) - model_config = _kimi_model_config() - manager_cls = get_kv_cache_manager_cls(model_config, llm_args.kv_cache_config) - manager = _create_kv_cache_manager( - model_engine=None, - kv_cache_manager_cls=manager_cls, - mapping=Mapping(world_size=1, tp_size=1, pp_size=1), - kv_cache_config=llm_args.kv_cache_config, - tokens_per_block=64, - max_seq_len=256, - max_batch_size=4, - spec_config=llm_args.speculative_config, - 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, - ) - try: - assert isinstance(manager, MambaHybridCacheManagerV2) - assert manager.kv_cache_type == CacheTypeCpp.SELFKONLY - assert manager.local_num_mamba_layers == 2 - assert manager.blocks_in_primary_pool > 0 - for layer_idx in (1, 3): - buffers = manager.get_buffers(layer_idx) - assert buffers.shape[1:] == (1, 64, 1, 40) - assert buffers.dtype is torch.bfloat16 - assert manager.mamba_pp_layers == [0, 2] - for layer_idx in manager.mamba_pp_layers: - cache = manager.mamba_layer_cache(layer_idx) - assert cache.conv.shape[1:] == (96, 3) - assert cache.conv.dtype is torch.bfloat16 - assert cache.temporal.shape[1:] == (4, 8, 8) - assert cache.temporal.dtype is torch.float32 - - manager.add_dummy_requests([123], token_nums=[8], is_gen=False) - indices = manager.get_state_indices([123], [False]) - assert indices[0] >= 0 - assert manager.cuda_state_indices[0].item() == indices[0] - if use_sa: - cache = manager.mamba_layer_cache(0) - assert cache.kda_qkg_cache is not None or cache.intermediate_ssm is not None - finally: - manager.shutdown() - - @pytest.mark.parametrize( "cache_transceiver_config", [ diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b5d1a53b24f1..c1e09bd89930 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -827,8 +827,6 @@ def test_registered_models_prefer_v2(self) -> None: "MistralLarge3ForCausalLM", "DeepseekV4ForCausalLM", "KimiK25ForConditionalGeneration", - "KimiLinearForCausalLM", - "KimiK3ForConditionalGeneration", "MiniMaxM2ForCausalLM", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", @@ -873,8 +871,6 @@ def test_registered_models_keep_v2_on_nixl(self) -> None: "MistralLarge3ForCausalLM", "GptOssForCausalLM", "KimiK25ForConditionalGeneration", - "KimiLinearForCausalLM", - "KimiK3ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", "Qwen3NextForCausalLM", From 682a0d770eed34ff1d302e3eff6395a3b78a7f15 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:47:58 -0700 Subject: [PATCH 3/3] [None][test] Cover registered KimiLinear V2 defaults Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 14 +++++++------- tensorrt_llm/_torch/pyexecutor/_util.py | 14 +++++++++++--- tests/unittest/llmapi/test_llm_args.py | 2 ++ 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 3c1e9809628d..dc409bcb18ed 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -20,8 +20,8 @@ ------- 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 - 1]`` bf16 (the convolution history, -sections ``[q | k | v]``) and a delta-rule recurrent 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 ``state_v_first`` FLA layout). MLA layers use the paged-KV side with ``num_kv_heads=1`` and ``head_dim = kv_lora_rank + qk_rope_head_dim`` (576), @@ -67,8 +67,9 @@ 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``; the hybrid manager -tracks recurrent-state snapshots alongside the MLA cache blocks. +opt-in via ``kv_cache_config.enable_block_reuse=true``, which routes to +``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 @@ -2100,9 +2101,8 @@ def _setup_helix_mappings( @classmethod def get_model_defaults(cls, llm_args) -> dict: - # - enable_block_reuse defaults off: keep reuse an explicit opt-in. - # V2 manages both KDA state and the paged MLA cache, including - # speculative-decoding state. + # - enable_block_reuse defaults off: reuse is supported as an + # 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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 24f736449be7..6bc67d8eb825 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -177,9 +177,17 @@ def get_kv_cache_manager_cls( "use_kv_cache_manager_v2=True; V1 supports only " "periodic_snapshot_interval.") - # Kimi K3's V1 compatibility route uses the unified C++ pool for - # block reuse and separate KV / recurrent-state pools otherwise. - # V2 and disaggregated serving use the shared hybrid routing below. + # 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. if is_kimi_linear(config) and not use_v2 and not is_disagg: if kv_cache_config.enable_block_reuse: logger.info( 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",