[TRTLLM-15177][chore] Kimi K3: inline MLA module, drop dead MoE comm plumbing

tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py

@@ -88,7 +88,6 @@ def create_strategy(
         alltoall_result_do_sum: bool = True,
         use_flashinfer: bool = False,
         hidden_size: Optional[int] = None,
-        communication_method: Optional[str] = None,
     ) -> Optional[Communication]:
         """
         Create the best communication method for the given configuration
@@ -114,8 +113,6 @@ def create_strategy(
             hidden_size: Actual MoE activation dimension (the A2A payload width).
                 For latent-MoE models this is moe_latent_size, not pretrained_config.hidden_size.
                 Falls back to pretrained_config.hidden_size when not provided.
-            communication_method: Optional model-selected communication method.
-                ``TRTLLM_FORCE_COMM_METHOD`` takes precedence when set.
             # TODO: Need a way to indicate whether EPLB is enabled.
 
         Returns:
@@ -151,15 +148,7 @@ def create_strategy(
         if mapping.moe_tp_size != 1:
             return AllGatherReduceScatter(mapping)
 
-        # A forced method comes either from the environment, which wins, or from the
-        # model-selected argument. Keep the source with the value so the log below can
-        # name the one the reader can actually go and change.
-        env_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD")
-        if env_method is not None:
-            force_method, force_source = env_method, "TRTLLM_FORCE_COMM_METHOD"
-        else:
-            force_method, force_source = communication_method, "communication_method"
-
+        force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD")
         if force_method is not None:
             strategy = CommunicationFactory._create_forced_method(
                 force_method,
@@ -175,7 +164,7 @@ def create_strategy(
             )
             logger.info(
                 f"Selected communication strategy: {strategy.__class__.__name__} "
-                f"({force_source}={force_method})"
+                f"(TRTLLM_FORCE_COMM_METHOD={force_method})"
             )
             return strategy
 

tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py

@@ -168,7 +168,6 @@ def __init__(
         trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None,
         trtllm_gen_activation_alpha: Optional[float] = None,
         trtllm_gen_activation_beta: Optional[float] = None,
-        communication_method: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(
@@ -190,7 +189,6 @@ def __init__(
         # Store model_config and aux_stream_dict for later use (e.g., backend setter)
         self.model_config = model_config
         self.aux_stream_dict = aux_stream_dict
-        self.communication_method = communication_method
 
         # If True, the router weight will be multiplied on the input rather than at the end of FC2
         self.apply_router_weight_on_input = apply_router_weight_on_input
@@ -640,7 +638,6 @@ def _create_comm_strategy_auto(self) -> Optional[Communication]:
             alltoall_result_do_sum=True,
             use_flashinfer=self.use_flashinfer,
             hidden_size=self.hidden_size,
-            communication_method=self.communication_method,
         )
 
     def forward_impl(

tensorrt_llm/_torch/moe/fused_moe/create_moe.py

@@ -383,7 +383,6 @@ def create_moe(
     trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None,
     trtllm_gen_activation_alpha: Optional[float] = None,
     trtllm_gen_activation_beta: Optional[float] = None,
-    communication_method: Optional[str] = None,
 ) -> MoE | VanillaMoE:
     """
     Create MoE instance with automatic parameter inference from model_config.
@@ -413,7 +412,6 @@ def create_moe(
         trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type
         trtllm_gen_activation_alpha: Optional backend-local activation alpha
         trtllm_gen_activation_beta: Optional backend-local activation beta
-        communication_method: Optional ConfigurableMoE communication method
 
     Returns:
         A complete MoE layer: a ``MoE`` (``ConfigurableMoE`` around an
@@ -508,13 +506,10 @@ def create_moe(
             trtllm_gen_activation_type=trtllm_gen_activation_type,
             trtllm_gen_activation_alpha=trtllm_gen_activation_alpha,
             trtllm_gen_activation_beta=trtllm_gen_activation_beta,
-            communication_method=communication_method,
         )
 
     # TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE
     # and own their communication and forward paths.
-    if communication_method is not None:
-        raise ValueError("communication_method requires ConfigurableMoE.")
     return create_moe_backend(
         moe_cls=moe_cls,
         routing_method=routing_method,

tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py

@@ -8,9 +8,9 @@
 import torch
 
 from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention
-from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import (
+from tensorrt_llm._torch.models.modeling_kimi_linear import (
     _KIMI_K3_MLA_GEN_BACKEND_ENV,
+    KimiK3MLAAttention,
     _kimi_k3_mla_decode_backend_policy,
     _select_mla_generation_backend,
     _validate_mla_generation_backend,

tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py

@@ -47,7 +47,6 @@
 import tensorrt_llm._torch.models.modeling_kimi_linear as modeling_kimi_linear
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3MoERuntime
-from tensorrt_llm._torch.moe.fused_moe.communication import CommunicationFactory
 from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import _MEGA_MOE_SYMM_BUFFER_CACHE
 from tensorrt_llm._torch.utils import ActType_TrtllmGen
 from tensorrt_llm._utils import get_free_port, get_sm_version
@@ -202,52 +201,6 @@ def test_padded_fused_shapes():
     assert padded_fused_shapes(2880, 96) == (3072, 2944, 128)
 
 
-def test_communication_factory_accepts_model_selected_method(monkeypatch):
-    mapping = SimpleNamespace(
-        enable_attention_dp=True,
-        dp_size=16,
-        moe_tp_size=1,
-        moe_ep_size=16,
-        has_cp_helix=lambda: False,
-    )
-    model_config = SimpleNamespace(
-        mapping=mapping,
-        pretrained_config=SimpleNamespace(hidden_size=3584),
-        torch_dtype=torch.bfloat16,
-        quant_config=None,
-        max_num_tokens=4096,
-        moe_max_num_tokens=65536,
-        use_cuda_graph=False,
-        use_low_precision_moe_combine=False,
-    )
-    selected = object()
-    method = None
-
-    def create_forced_method(force_method, *args, **kwargs):
-        nonlocal method
-        method = force_method
-        return selected
-
-    monkeypatch.delenv("TRTLLM_FORCE_COMM_METHOD", raising=False)
-    monkeypatch.setattr(
-        CommunicationFactory,
-        "_create_forced_method",
-        staticmethod(create_forced_method),
-    )
-    actual = CommunicationFactory.create_strategy(
-        model_config=model_config,
-        num_experts=896,
-        num_slots=896,
-        top_k=16,
-        expert_size_per_partition=56,
-        hidden_size=3584,
-        communication_method="ALLGATHER",
-    )
-
-    assert method == "ALLGATHER"
-    assert actual is selected
-
-
 @situ_supported
 def test_make_situ_alpha_beta_contract():
     alpha, beta = make_situ_alpha_beta(
@@ -712,7 +665,6 @@ def _make_routed_moe(
             else QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8)
         ),
         layer_idx=0,
-        communication_method=None,
     )
     if moe_backend == "TRTLLM":
         moe_kwargs.update(