[TRTLLM-13639][test] Migrate Kimi dis-agg tests to Transceiver v2 and trim tests

tensorrt_llm/_torch/models/modeling_kimi_k25.py

@@ -34,7 +34,7 @@
 import os
 import tempfile
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -1517,6 +1517,21 @@ class KimiK25ForConditionalGeneration(PreTrainedModel):
 
     _LANG_PREFIX = "language_model."
 
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Literal["PYTHON"]:
+        """Kimi-K2.5 defaults to the Python (v2) KV-cache transceiver.
+
+        The DeepSeek-V3 MLA backbone transfers a large latent KV, which the
+        Python transceiver handles better in disaggregated serving. This is
+        only adopted when the user leaves
+        ``cache_transceiver_config.transceiver_runtime`` at 'auto' and the
+        effective backend is NIXL; otherwise the C++ transceiver is used.
+        """
+        return "PYTHON"
+
     def __init__(
         self,
         model_config: ModelConfig[PretrainedConfig],

tests/integration/defs/accuracy/test_disaggregated_serving.py

@@ -2081,59 +2081,6 @@ def test_mixed_ctx_gen_model(self, ctx_pp, gen_tp):
                                  gen_instances=1)
 
 
-@pytest.mark.timeout(10800)
-@skip_pre_blackwell
-class TestKimiK2(LlmapiAccuracyTestHarness):
-    MODEL_NAME = "moonshotai/Kimi-K2-Thinking"
-    MODEL_PATH = f"{llm_models_root()}/Kimi-K2-Thinking-NVFP4"
-
-    @pytest.mark.skip_less_device(8)
-    @pytest.mark.skip_less_device_memory(200000)
-    def test_nvfp4(self):
-        ctx_server_config = {
-            "max_batch_size": 16,
-            "disable_overlap_scheduler": True,
-            "cache_transceiver_config": {
-                "backend": "DEFAULT",
-                "max_tokens_in_buffer": 4096
-            },
-            "tensor_parallel_size": 4,
-            "enable_attention_dp": True,
-            "trust_remote_code": True,
-            "kv_cache_config": {
-                "free_gpu_memory_fraction": 0.8,
-            },
-        }
-        gen_server_config = {
-            "max_batch_size": 16,
-            "disable_overlap_scheduler": True,
-            "cache_transceiver_config": {
-                "backend": "DEFAULT",
-                "max_tokens_in_buffer": 4096
-            },
-            "tensor_parallel_size": 4,
-            "enable_attention_dp": True,
-            "trust_remote_code": True,
-            "kv_cache_config": {
-                "free_gpu_memory_fraction": 0.8,
-            },
-        }
-        disaggregated_server_config = {
-            "hostname": "localhost",
-            "backend": "pytorch",
-            "context_servers": {
-                "num_instances": 1
-            },
-            "generation_servers": {
-                "num_instances": 1
-            }
-        }
-        with launch_disaggregated_llm(disaggregated_server_config,
-                                      ctx_server_config, gen_server_config,
-                                      self.MODEL_PATH) as llm:
-            run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"])
-
-
 @pytest.mark.timeout(10800)
 @skip_pre_blackwell
 class TestKimiK25(LlmapiAccuracyTestHarness):
@@ -2145,9 +2092,13 @@ class TestKimiK25(LlmapiAccuracyTestHarness):
     def test_nvfp4(self):
         """Disaggregated GSM8K accuracy for Kimi-K2.5 (NVFP4).
 
-        ctx and gen servers are each TP4 (8 GPUs total) over the default cache
+        ctx and gen servers are each TP4 (8 GPUs total). The cache transceiver
+        uses backend=NIXL + transceiver_runtime=PYTHON: NIXL is required so the
+        disagg test harness skips its TRTLLM_USE_UCX_KVCACHE=1 fallback (which
+        would make the effective backend UCX and force the C++ transceiver),
+        letting the ctx->gen MLA-latent KV transfer run over the Python (v2)
         transceiver. GSM8K is text-only, so requests run through the DeepSeek-V3
-        MLA backbone (no vision) and the ctx->gen KV transfer is the MLA latent.
+        MLA backbone (no vision).
         Kimi-K2.5 ships custom HF modeling code (auto_map in config.json), so
         trust_remote_code must be set on both servers or executor init fails at
         config parse time.
@@ -2165,7 +2116,8 @@ def test_nvfp4(self):
             "max_num_tokens": 8192,
             "disable_overlap_scheduler": True,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
+                "transceiver_runtime": "PYTHON",
                 "max_tokens_in_buffer": 4096
             },
             "tensor_parallel_size": 4,
@@ -2181,7 +2133,8 @@ def test_nvfp4(self):
             "max_num_tokens": 8192,
             "disable_overlap_scheduler": True,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
+                "transceiver_runtime": "PYTHON",
                 "max_tokens_in_buffer": 4096
             },
             "tensor_parallel_size": 4,

tests/integration/test_lists/qa/llm_function_core.txt

@@ -21,7 +21,6 @@ accuracy/test_disaggregated_serving.py::TestGPTOSS::test_auto_dtype[True]
 accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python[cache_mgr_v1]
 accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python[cache_mgr_v2]
 accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl_python[cache_mgr_v1]
-accuracy/test_disaggregated_serving.py::TestKimiK2::test_nvfp4
 accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4
 accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[False-False-False-False]
 accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[False-False-False-True]

tests/unittest/_torch/modeling/test_modeling_kimi_k25.py

@@ -605,6 +605,12 @@ def test_auto_model_registered(self):
         self.assertIsNotNone(cls, "KimiK25ForConditionalGeneration not in MODEL_CLASS_MAPPING")
         self.assertIs(cls, KimiK25ForConditionalGeneration)
 
+    def test_prefers_python_transceiver(self):
+        """Kimi-K2.5 defaults to the Python KV-cache transceiver in disagg."""
+        self.assertEqual(
+            KimiK25ForConditionalGeneration.get_preferred_transceiver_runtime(), "PYTHON"
+        )
+
 
 # ---------------------------------------------------------------------------
 # E2E Smoke Test — requires GPU + model checkpoint