[TRTLLM-14813][feat] Add Kimi K3 (KimiLinear) model

cpp/tensorrt_llm/thop/kdaDecodeOp.cpp

@@ -205,14 +205,21 @@ at::Tensor kda_decode_fusion_forward(at::Tensor x_q, at::Tensor x_k, at::Tensor
     at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, at::Tensor onorm_g,
     at::Tensor onorm_weight, std::optional<at::Tensor> ssm_state_indices, at::Tensor cu_seqlens, at::Tensor state,
     bool apply_onorm, bool update_conv_cache, bool use_lower_bound, bool apply_beta_sigmoid, double lower_bound,
-    double scale, double onorm_eps)
+    double scale, double onorm_eps, std::optional<at::Tensor> output)
 {
     validate_kda_decode_fusion_inputs(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v,
         a_log, g, dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, apply_onorm,
         update_conv_cache);
     int const B = static_cast<int>(x_q.size(1));
     int const HV = static_cast<int>(x_v.size(2));
-    auto out = at::empty({B, 1, HV, kDimV}, x_q.options());
+    auto out = output.has_value() ? *output : at::empty({B, 1, HV, kDimV}, x_q.options());
+    if (output.has_value())
+    {
+        TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kBFloat16, "out must be a CUDA bfloat16 tensor");
+        TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
+        TORCH_CHECK(out.dim() == 4 && out.size(0) == B && out.size(1) == 1 && out.size(2) == HV && out.size(3) == kDimV,
+            "out must have shape [B, 1, HV, 128]");
+    }
     launch_selected_kernel(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, a_log, g,
         dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, out, apply_onorm, update_conv_cache,
         use_lower_bound, apply_beta_sigmoid, lower_bound, scale, onorm_eps);
@@ -236,7 +243,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "Tensor? ssm_state_indices, Tensor cu_seqlens, Tensor(d!) state, "
         "bool apply_onorm, bool update_conv_cache, bool use_lower_bound, "
         "bool apply_beta_sigmoid, float lower_bound, float scale, "
-        "float onorm_eps) -> Tensor");
+        "float onorm_eps, Tensor(e!)? output=None) -> Tensor(e!)");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)

tensorrt_llm/_torch/configs/__init__.py

@@ -23,6 +23,7 @@
     Gemma4UnifiedTextConfig,
     Gemma4UnifiedVisionConfig,
 )
+from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig
 from tensorrt_llm._torch.configs.laguna import LagunaConfig
 from tensorrt_llm._torch.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig
 
@@ -54,6 +55,12 @@ def _register_custom_configs_with_transformers() -> None:
         "kimi_k2": DeepseekV3Config,
         "deepseek_v4": DeepseekV4Config,
         "gemma4_assistant": Gemma4AssistantConfig,
+        # Kimi K3 text config ("kimi_linear"). The composite "kimi_k3"
+        # model_type is flattened to the text config by
+        # pyexecutor.config_utils.load_pretrained_config; registering the
+        # text config here lets AutoConfig / AutoTokenizer resolve
+        # "kimi_linear" without trust_remote_code.
+        "kimi_linear": KimiLinearConfig,
         "laguna": LagunaConfig,
         # minicpmv4_6 is only registered in transformers>=5.7.0; register our
         # own composite config so AutoTokenizer.from_pretrained works on older
@@ -86,6 +93,7 @@ def _register_custom_configs_with_transformers() -> None:
     "Gemma4UnifiedConfig",
     "Gemma4UnifiedTextConfig",
     "Gemma4UnifiedVisionConfig",
+    "KimiLinearConfig",
     "LagunaConfig",
     "MiniCPMV4_6Config",
     "MiniCPMV4_6VisionConfig",

tensorrt_llm/_torch/configs/kimi_linear.py

@@ -0,0 +1,175 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""In-tree config for Kimi K3 ("kimi_linear") text checkpoints.
+
+Mirrors the checkpoint-shipped ``configuration_kimi_k3.KimiLinearConfig`` so
+TRT-LLM can parse Kimi K3 checkpoints without ``trust_remote_code`` for the
+config. The top-level Kimi K3 checkpoints use a composite VLM config
+(``model_type: kimi_k3``) whose ``text_config`` is this class;
+``load_pretrained_config`` flattens the composite config to this text config
+(TRT-LLM runs the text model only).
+"""
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class KimiLinearConfig(PretrainedConfig):
+    model_type = "kimi_linear"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=163840,
+        hidden_size=4096,
+        head_dim=None,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        tie_word_embeddings=False,
+        moe_intermediate_size: Optional[int] = None,
+        moe_renormalize: bool = True,
+        moe_router_activation_func: str = "sigmoid",
+        num_experts: Optional[int] = None,
+        num_experts_per_token: Optional[int] = None,
+        num_shared_experts: int = 0,
+        routed_scaling_factor: float = 1.0,
+        first_k_dense_replace: int = 0,
+        moe_layer_freq: int = 1,
+        use_grouped_topk: bool = True,
+        num_expert_group: int = 1,
+        topk_group: int = 1,
+        q_lora_rank: Optional[int] = None,
+        kv_lora_rank: Optional[int] = None,
+        qk_nope_head_dim: Optional[int] = None,
+        qk_rope_head_dim: Optional[int] = None,
+        v_head_dim: Optional[int] = None,
+        mla_use_nope: Optional[bool] = False,
+        mla_use_output_gate: Optional[bool] = False,
+        num_nextn_predict_layers: int = 0,
+        linear_attn_config: Optional[dict] = None,
+        attn_res_block_size: Optional[int] = None,
+        latent_moe_use_norm: bool = False,
+        activation_situ_beta: Optional[float] = None,
+        activation_situ_linear_beta: Optional[float] = None,
+        max_position_embeddings: int = 4096,
+        routed_expert_hidden_size: Optional[int] = None,
+        topk_method: str = "noaux_tc",
+        **kwargs,
+    ):
+        # NOTE: unlike the checkpoint-shipped config class, do not accept a
+        # ``model_type`` kwarg that shadows the class attribute; transformers
+        # keys registry lookups off the class attribute.
+        kwargs.pop("model_type", None)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.mla_use_nope = mla_use_nope
+        self.mla_use_output_gate = mla_use_output_gate
+        # moe config
+        self.num_experts = num_experts
+        self.num_experts_per_token = num_experts_per_token
+        self.moe_renormalize = moe_renormalize
+        self.num_shared_experts = num_shared_experts
+        self.routed_scaling_factor = routed_scaling_factor
+        self.moe_router_activation_func = moe_router_activation_func
+        assert self.moe_router_activation_func in ("softmax", "sigmoid")
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.moe_layer_freq = moe_layer_freq
+        self.use_grouped_topk = use_grouped_topk
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        self.attn_res_block_size = attn_res_block_size
+        self.latent_moe_use_norm = latent_moe_use_norm
+        self.activation_situ_beta = activation_situ_beta
+        self.activation_situ_linear_beta = activation_situ_linear_beta
+        self.max_position_embeddings = max_position_embeddings
+        self.routed_expert_hidden_size = routed_expert_hidden_size
+        self.topk_method = topk_method
+
+        if linear_attn_config is not None:
+            assert linear_attn_config["kda_layers"] is not None
+            assert linear_attn_config["full_attn_layers"] is not None
+        self.linear_attn_config = linear_attn_config
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def is_mla(self):
+        return (
+            self.q_lora_rank is not None
+            or self.kv_lora_rank is not None
+            or self.qk_nope_head_dim is not None
+            or self.qk_rope_head_dim is not None
+            or self.v_head_dim is not None
+            or self.mla_use_nope is True
+        )
+
+    @property
+    def is_moe(self):
+        return self.num_experts is not None
+
+    @property
+    def is_linear_attn(self) -> bool:
+        return not (
+            self.linear_attn_config is None
+            or (
+                isinstance(self.linear_attn_config, dict)
+                and self.linear_attn_config["kda_layers"] is not None
+                and len(self.linear_attn_config["kda_layers"]) == 0
+            )
+        )
+
+    def is_kda_layer(self, layer_idx: int) -> bool:
+        """0-indexed layer check; ``kda_layers`` in the config is 1-indexed."""
+        return (
+            self.linear_attn_config is not None
+            and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
+        )
+
+    def is_full_attn_layer(self, layer_idx: int) -> bool:
+        """0-indexed layer check; ``full_attn_layers`` is 1-indexed."""
+        return (
+            self.linear_attn_config is not None
+            and (layer_idx + 1) in self.linear_attn_config["full_attn_layers"]
+        )

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -294,11 +294,39 @@ def _(logits,
         pass
 
     @torch.library.register_fake("trtllm::kda_decode")
-    def _(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v,
-          conv_state_q, conv_state_k, conv_state_v, a_log, g, dt_bias, beta,
-          onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state,
-          apply_onorm, update_conv_cache, use_lower_bound, apply_beta_sigmoid,
-          lower_bound, scale, onorm_eps):
+    def _(x_q: torch.Tensor,
+          x_k: torch.Tensor,
+          x_v: torch.Tensor,
+          w_q_t: torch.Tensor,
+          w_k_t: torch.Tensor,
+          w_v_t: torch.Tensor,
+          bias_q: torch.Tensor,
+          bias_k: torch.Tensor,
+          bias_v: torch.Tensor,
+          conv_state_q: torch.Tensor,
+          conv_state_k: torch.Tensor,
+          conv_state_v: torch.Tensor,
+          a_log: torch.Tensor,
+          g: torch.Tensor,
+          dt_bias: torch.Tensor,
+          beta: torch.Tensor,
+          onorm_g: torch.Tensor,
+          onorm_weight: torch.Tensor,
+          ssm_state_indices: Optional[torch.Tensor],
+          cu_seqlens: torch.Tensor,
+          state: torch.Tensor,
+          apply_onorm: bool,
+          update_conv_cache: bool,
+          use_lower_bound: bool,
+          apply_beta_sigmoid: bool,
+          lower_bound: float,
+          scale: float,
+          onorm_eps: float,
+          output: Optional[torch.Tensor] = None) -> torch.Tensor:
+        # Mirror the CUDA impl: write into the caller-provided output when
+        # given (schema returns Tensor(e!)), else allocate.
+        if output is not None:
+            return output
         # x_q is [1, tokens, H, 128]; the kernel emits one row per token.
         return x_q.new_empty((x_q.size(1), 1, x_v.size(2), x_v.size(3)))
 

tensorrt_llm/_torch/model_config.py

@@ -29,8 +29,8 @@
 from transformers.utils import HF_MODULES_CACHE
 
 from tensorrt_llm._torch.pyexecutor.config_utils import (
-    get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid,
-    load_pretrained_config)
+    get_kimi_linear_num_attention_layers, get_qwen3_hybrid_num_attention_layers,
+    is_kimi_linear, is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config)
 from tensorrt_llm._utils import (get_sm_version, is_sm_100f,
                                  torch_dtype_to_binding)
 from tensorrt_llm.bindings import LayerType as LayerTypeCpp
@@ -72,11 +72,31 @@ def _is_lock_infra_error(exc: BaseException) -> bool:
     if isinstance(exc, PermissionError):
         return True
     if isinstance(exc, OSError):
+        # EEXIST: filelock's ensure_directory_exists() can lose the
+        # mkdir(exist_ok=True) race on NFS when many ranks start at once
+        # (the post-EEXIST is_dir() recheck sees a stale attribute cache).
+        # An un-creatable lock dir is broken infra, not contention.
         return exc.errno in (errno.EACCES, errno.EPERM, errno.ENOLCK,
-                             errno.ESTALE)
+                             errno.ESTALE, errno.EEXIST)
     return False
 
 
+def _release_lock_ignoring_infra_errors(lock: "filelock.BaseFileLock") -> None:
+    """Release ``lock``, downgrading broken-lock-infra errors to a warning.
+
+    NFS can return ENOLCK/ESTALE from the unlock ``flock`` call itself (e.g.
+    lock-daemon exhaustion when many ranks start simultaneously). The config
+    load the lock protected has already completed at release time, so
+    crashing the process here would fail an otherwise healthy executor.
+    """
+    try:
+        lock.release()
+    except (PermissionError, OSError) as e:
+        if not _is_lock_infra_error(e):
+            raise
+        logger.warning(f"config lock release failed ({e}), continuing")
+
+
 @contextlib.contextmanager
 def config_file_lock(timeout: int = 10):
     """
@@ -120,12 +140,12 @@ def config_file_lock(timeout: int = 10):
             try:
                 yield
             finally:
-                tmp_lock.release()
+                _release_lock_ignoring_infra_errors(tmp_lock)
     else:
         try:
             yield
         finally:
-            lock.release()
+            _release_lock_ignoring_infra_errors(lock)
 
 
 @dataclass(kw_only=True)
@@ -1391,6 +1411,8 @@ def get_num_attention_layers(self) -> int:
             return cfg.hybrid_override_pattern.count("*")
         if is_qwen3_hybrid(cfg):
             return get_qwen3_hybrid_num_attention_layers(cfg)
+        if is_kimi_linear(cfg):
+            return get_kimi_linear_num_attention_layers(cfg)
         return cfg.num_hidden_layers
 
     def get_num_mamba_layers(self) -> int:
@@ -1401,6 +1423,9 @@ def get_num_mamba_layers(self) -> int:
         if is_qwen3_hybrid(cfg):
             return cfg.num_hidden_layers - get_qwen3_hybrid_num_attention_layers(
                 cfg)
+        if is_kimi_linear(cfg):
+            return cfg.num_hidden_layers - get_kimi_linear_num_attention_layers(
+                cfg)
         return 0
 
 

tensorrt_llm/_torch/models/__init__.py

@@ -30,6 +30,7 @@
 from .modeling_hunyuan_moe import HunYuanMoEV1ForCausalLM
 from .modeling_hyperclovax import HCXVisionForCausalLM
 from .modeling_kimi_k25 import KimiK25ForConditionalGeneration
+from .modeling_kimi_linear import KimiLinearForCausalLM
 from .modeling_laguna import LagunaForCausalLM
 from .modeling_llama import LlamaForCausalLM
 from .modeling_llava_next import LlavaNextModel
@@ -89,6 +90,7 @@
     "HunYuanDenseV1ForCausalLM",
     "HunYuanMoEV1ForCausalLM",
     "KimiK25ForConditionalGeneration",
+    "KimiLinearForCausalLM",
     "LlamaForCausalLM",
     "LlavaNextModel",
     "MiniCPMV4_6Model",