[OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter

.pre-commit-config.yaml

@@ -114,6 +114,7 @@ repos:
               modelopt/torch/fastgen/plugins/qwen_image_pdd.py|
               modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py|
               modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py|
+              modelopt/torch/models/kimi_k3/ep.py|
               modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py|
               modelopt/torch/quantization/export_onnx.py|
               modelopt/torch/quantization/plugins/attention.py|

LICENSE

@@ -224,6 +224,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0
   Copyright (c) 2024 Heming Xia
   Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
   Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team
+  Copyright 2026 the MiniMax AI Team and HuggingFace Team. All rights reserved.
   Copyright (c) OpenMMLab. All rights reserved.
   Copyright (c) 2022 megvii-model. All Rights Reserved.
 

modelopt/torch/models/__init__.py

@@ -90,6 +90,7 @@
     nemotron,
     nemotron_h,
     minimax_m3_vl,
+    kimi_k3,
     llama4_text,
     kimi_vl,
     kimi_k25,

modelopt/torch/models/kimi_k3/__init__.py

@@ -13,4 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Kimi-K3 (HF model type ``kimi_k3``)."""
+"""kimi_k3 specs (HF model type ``kimi_k3``)."""
+
+from . import specs

modelopt/torch/models/kimi_k3/specs.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""kimi_k3 specs (HF model type ``kimi_k3``).
+
+Kimi-K3: a remote-code VLM on the kimi_linear backbone, with packed mxfp4 experts that are
+decompressed at load.
+"""
+
+from ..specs import EPSpec, ModelSpec, register
+from .ep import BASE_MODEL_EP_PLAN, MODEL_TYPE, PATCHED_EVERY_LOAD, make_buildable, prepare
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type=MODEL_TYPE,
+        modeling_source="remote_code",
+        ep_spec=EPSpec(
+            prepare=prepare,
+            prepared_ep_plan=BASE_MODEL_EP_PLAN,
+            plan_config_attr="text_config",
+            make_buildable=make_buildable,
+            prepare_patches_every_load=PATCHED_EVERY_LOAD,
+            decompressed_formats=("compressed-tensors",),
+            # prepare's patches to the checkpoint's remote code were validated on this release.
+            requires_transformers="5.13",
+        ),
+    )
+)

modelopt/torch/models/specs.py

@@ -230,6 +230,13 @@ class EPSpec(SpecSection):
     the modules the code imports, and adapt a method to the arguments transformers now passes it,
     but patches no MoE or expert class, config, weight converter or weight-loading hook."""
 
+    prepare_patches_every_load: tuple[str, ...] = ()
+    """The model types ``prepare`` patches every later load of in the process, not only the
+    expert-parallel load it ran for: the model's own, and any other its patches reach (Kimi-K3's
+    swaps its remote MoE block's class, and registers weight converters under its ``kimi_linear``
+    backbone's type). Once such a ``prepare`` has run, a whole-model or FSDP2-sharded load of any of
+    them in the process is refused before it builds or loads one."""
+
     decompressed_formats: tuple[str, ...] = ()
     """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts
     ``prepare`` has decompress as they load, which the expert-parallel load then takes beside

modelopt/torch/utils/moe.py

@@ -220,6 +220,36 @@ def _spec_refuses_ep(model_type: str) -> bool:
     return False
 
 
+# The model types whose spec's prepare ran in this process and patched every later load of them.
+_PATCHED_EVERY_LOAD: set[str] = set()
+
+
+def mark_patched_every_load(model_type: str) -> None:
+    """Record that a ``prepare`` patched ``model_type`` for every later load in this process.
+
+    :func:`apply_distributed_support` records it once such a ``prepare`` returns; a ``prepare`` calls
+    this itself at its first patch that lasts, so a raise after it still leaves the record.
+    """
+    _PATCHED_EVERY_LOAD.add(model_type)
+
+
+def check_not_patched_by_prepare(model_type: str) -> None:
+    """Refuse a load other than an expert-parallel one of a model whose ``prepare`` patched it for good.
+
+    A spec whose ``prepare`` patches every later load of its model
+    (:attr:`~modelopt.torch.models.EPSpec.prepare_patches_every_load`) leaves code only an
+    expert-parallel load builds from: once it ran in this process, a whole-model or FSDP2-sharded
+    load of the model is refused before it builds or loads one. The plan decision's throwaway meta
+    build runs earlier, so under torchrun an EP degree of 1 meets the decision's own refusal first.
+    """
+    if model_type in _PATCHED_EVERY_LOAD:
+        raise NotImplementedError(
+            f"An expert-parallel load in this process, of {model_type} or of a model built on it, "
+            "patched its modeling code for every later load of it, which a whole-model or "
+            "FSDP2-sharded load cannot build from. Load it in a new process."
+        )
+
+
 def apply_distributed_support(
     model_type: str,
     hf_config,
@@ -289,6 +319,8 @@ def apply_distributed_support(
 
     if prepare and ep.prepare is not None:
         hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code)
+        for patched in ep.prepare_patches_every_load:
+            mark_patched_every_load(patched)
     elif not prepare and ep.make_buildable is not None:
         ep.make_buildable(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code)
     return hf_config
@@ -560,8 +592,9 @@ def expert_fn(sorted_tok, offsets):
 def register_torch_a2a():
     """Register everything the torch_a2a path needs into stock transformers (idempotent).
 
-    Done lazily from ``load_and_prepare_ep``, not at import time, so importing this module has no
-    transformers side effect:
+    Done lazily, not at import time, so importing this module has no transformers side effect: by
+    ``load_and_prepare_ep``, and by a spec's ``prepare`` that names ``torch_a2a`` on the config a
+    model is built from before that load (Kimi-K3's). It registers:
 
     - the ``torch_a2a`` experts forward into ``ExpertsInterface`` (``ALL_EXPERTS_FUNCTIONS``);
     - two EP parallel styles into ``ALL_PARALLEL_STYLES``:

modelopt/torch/utils/plugins/model_load_utils.py

@@ -82,6 +82,7 @@
     _spec_refuses_ep,
     apply_distributed_support,
     check_distributed_support,
+    check_not_patched_by_prepare,
     config_levels,
     register_torch_a2a,
 )
@@ -995,6 +996,8 @@ def handed() -> DistributionPlan | None:
             model_type = getattr(chosen, "model_type", None) or (
                 _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type
             )
+        if route != _EXPERT_PARALLEL:
+            check_not_patched_by_prepare(model_type)
         load_kwargs = (
             {}
             if route is None

modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml

@@ -50,12 +50,25 @@
     enable: false
   - quantizer_name: 'output.*'
     enable: false
+  # Kimi-K3: vLLM builds the latent-expert projections around the routed experts and the
+  # scalar residual gates unquantized (quant_config=None), so it cannot load a checkpoint
+  # that quantizes them. The [1, hidden] gates gain nothing from quantization either.
+  - quantizer_name: '*routed_expert_down_proj*'
+    enable: false
+  - quantizer_name: '*routed_expert_up_proj*'
+    enable: false
+  - quantizer_name: '*self_attention_res_proj*'
+    enable: false
+  - quantizer_name: '*mlp_res_proj*'
+    enable: false
+  - quantizer_name: '*output_attn_res_proj*'
+    enable: false
   # Multimodal vision branch: keep the vision encoder (SigLIP / ViT) and any
   # multimodal embedding projection in BF16 by default. Recipes that enable bare
   # `*weight_quantizer` / `*input_quantizer` or `*mlp*` wildcards otherwise also
   # match the vision tower (`model.vision_tower.*`, `model.visual.*`,
   # Llama-4 `vision_model.*`) and the embedding projection (`model.embed_vision.*`,
-  # Llama-4 `multi_modal_projector.*`); quantizing the vision branch crashes
+  # Llama-4 `multi_modal_projector.*`, Kimi `mm_projector.*`); quantizing the vision branch crashes
   # export / produces garbage image embeddings on VL models (gemma-4,
   # Qwen3.5-VL — NVBugs 6293731, 6293762, 6294017; Llama-4-Scout — NVBugs
   # 6359097, where `vision_model.patch_embedding.linear` has in_features=588,
@@ -71,6 +84,8 @@
     enable: false
   - quantizer_name: '*multi_modal_projector*'
     enable: false
+  - quantizer_name: '*mm_projector*'
+    enable: false
   - parent_class: 'nn.BatchNorm1d'
     quantizer_name: '*'
     enable: false

tests/_test_utils/torch/kimi_k3_remote.py

@@ -13,21 +13,111 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A stand-in for Kimi-K3's own remote modeling code: the classes and functions the fixes for it
-patch, for tests on tiny CPU tensors.
+"""A stand-in for Kimi-K3's own remote modeling code: the classes and functions its adapter patches,
+with the per-expert ``nn.ModuleList`` MoE the checkpoint's code ships, for tests on tiny CPU tensors.
 """
 
 import types
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 
+def situ(gate_up, beta, linear_beta):
+    """The stand-in ``situ`` gated activation over a full [gate; up] row."""
+    gate, up = gate_up.chunk(2, dim=-1)
+    return gate * torch.sigmoid(beta * gate) * (up + linear_beta)
+
+
 def remote_module(name="kimi_k3_remote_stub"):
-    """A stand-in for the remote modeling namespace: what the fixes patch, as shipped."""
+    """A stand-in for the remote modeling namespace: everything the adapter patches, as shipped."""
     remote = types.ModuleType(name)
     remote.calls = []
 
+    def _get_situ_activation_params(config):
+        return config.situ_beta, config.situ_linear_beta
+
+    class SituAndMul(nn.Module):
+        def __init__(self, beta, linear_beta):
+            super().__init__()
+            self.beta, self.linear_beta = beta, linear_beta
+
+        def forward(self, gate_up):
+            return situ(gate_up, self.beta, self.linear_beta)
+
+    class KimiBlockSparseMLP(nn.Module):
+        """One expert as the checkpoint stores it: w1 = gate, w3 = up, w2 = down."""
+
+        def __init__(self, config, hidden):
+            super().__init__()
+            inter = config.moe_intermediate_size
+            self.w1 = nn.Linear(hidden, inter, bias=False)
+            self.w3 = nn.Linear(hidden, inter, bias=False)
+            self.w2 = nn.Linear(inter, hidden, bias=False)
+            situ = config.hidden_act == "situ"
+            self.situ_params = _get_situ_activation_params(config) if situ else None
+
+        def forward(self, x):
+            gate, up = self.w1(x), self.w3(x)
+            if self.situ_params is None:
+                return self.w2(F.silu(gate) * up)
+            return self.w2(situ(torch.cat([gate, up], dim=-1), *self.situ_params))
+
+    class KimiMoEGate(nn.Module):
+        def __init__(self, config):
+            super().__init__()
+            self.top_k = config.num_experts_per_token
+            self.weight = nn.Parameter(torch.randn(config.num_experts, config.hidden_size))
+
+        def forward(self, hidden_states):
+            flat = hidden_states.reshape(-1, hidden_states.shape[-1])
+            weights, idx = torch.topk(torch.sigmoid(flat @ self.weight.t()), self.top_k, dim=-1)
+            return idx, weights
+
+    class KimiSparseMoeBlock(nn.Module):
+        """Latent projections around a per-expert Python loop, plus always-on shared experts."""
+
+        def __init__(self, config):
+            super().__init__()
+            self.config = config
+            # Kimi-K3's config has no use_latent_moe field: the block derives it, as vLLM does.
+            self.use_latent_moe = config.routed_expert_hidden_size is not None
+            self.latent_moe_use_norm = config.latent_moe_use_norm
+            hidden = config.hidden_size
+            width = config.routed_expert_hidden_size if self.use_latent_moe else hidden
+            self.experts = nn.ModuleList(
+                KimiBlockSparseMLP(config, width) for _ in range(config.num_experts)
+            )
+            self.gate = KimiMoEGate(config)
+            if self.use_latent_moe:
+                self.routed_expert_down_proj = nn.Linear(hidden, width, bias=False)
+                self.routed_expert_up_proj = nn.Linear(width, hidden, bias=False)
+                if config.latent_moe_use_norm:
+                    self.routed_expert_norm = nn.LayerNorm(width)
+            if config.num_shared_experts is not None:
+                self.shared_experts = KimiBlockSparseMLP(config, hidden)
+
+        def forward(self, hidden_states):
+            idx, weights = self.gate(hidden_states)
+            x = hidden_states.reshape(-1, hidden_states.shape[-1])
+            if self.use_latent_moe:
+                x = self.routed_expert_down_proj(x)
+            y = torch.stack(
+                [
+                    sum(w * self.experts[int(e)](token) for e, w in zip(top_e, top_w))
+                    for token, top_e, top_w in zip(x, idx, weights)
+                ]
+            )
+            if self.use_latent_moe:
+                if self.latent_moe_use_norm:
+                    y = self.routed_expert_norm(y)
+                y = self.routed_expert_up_proj(y)
+            y = y.reshape(hidden_states.shape)
+            if self.config.num_shared_experts is not None:
+                y = y + self.shared_experts(hidden_states)
+            return y
+
     class KimiDeltaAttention(nn.Module):
         def __init__(self, num_heads, head_dim):
             super().__init__()
@@ -70,6 +160,11 @@ def tie_weights(self):  # the remote override, with the old no-kwarg signature
             return "tied"
 
     for obj in (
+        _get_situ_activation_params,
+        SituAndMul,
+        KimiBlockSparseMLP,
+        KimiMoEGate,
+        KimiSparseMoeBlock,
         KimiDeltaAttention,
         chunk_kda,
         fused_recurrent_kda,