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

modelopt/torch/models/__init__.py

@@ -92,6 +92,7 @@
     minimax_m3_vl,
     llama4_text,
     kimi_vl,
+    kimi_k3,
     kimi_k25,
     nemotron_h_omni,
     qwen2_moe,

modelopt/torch/models/kimi_k3/__init__.py

@@ -0,0 +1,18 @@
+# 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``)."""
+
+from . import specs

modelopt/torch/models/kimi_k3/specs.py

@@ -0,0 +1,38 @@
+# 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, prepare
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="kimi_k3",
+        modeling_source="remote_code",
+        ep_spec=EPSpec(
+            prepare=prepare,
+            prepared_ep_plan=BASE_MODEL_EP_PLAN,
+            # prepare's patches to the checkpoint's remote code were validated on this release.
+            requires_transformers="5.13",
+        ),
+    )
+)

modelopt/torch/utils/moe.py

@@ -514,8 +514,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_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/unit/recipe/test_presets.py

@@ -22,9 +22,11 @@
 """
 
 import argparse
+import copy
 
 import pytest
 import torch
+import torch.nn as nn
 
 import modelopt.torch.quantization as mtq
 from modelopt.recipe import load_recipe, presets
@@ -116,6 +118,48 @@ def test_default_exclusions_disable_mtp_at_any_depth():
     assert not model.language_model.mtp.proj.input_quantizer.is_enabled
 
 
+class _KimiK3Layer(nn.Module):
+    """Kimi-K3's decoder-layer projections that vLLM builds unquantized, beside one it quantizes."""
+
+    def __init__(self, hidden=32, latent=16):
+        super().__init__()
+        self.self_attn = nn.Module()
+        self.self_attn.o_proj = nn.Linear(hidden, hidden, bias=False)
+        self.self_attention_res_proj = nn.Linear(hidden, 1, bias=False)
+        self.mlp_res_proj = nn.Linear(hidden, 1, bias=False)
+        self.block_sparse_moe = nn.Module()
+        self.block_sparse_moe.routed_expert_down_proj = nn.Linear(hidden, latent, bias=False)
+        self.block_sparse_moe.routed_expert_up_proj = nn.Linear(latent, hidden, bias=False)
+
+
+@pytest.mark.parametrize("qformat", ["fp8", "nvfp4"])
+def test_default_presets_leave_the_kimi_k3_projections_vllm_builds_unquantized(qformat):
+    # vLLM's Kimi-K3 model passes quant_config=None for these, so it cannot load them quantized.
+    # Kimi-K2.5 and Kimi-K3 also name their multimodal projector ``mm_projector``; like the other
+    # projector spellings it stays BF16 unless a recipe re-enables it.
+    model = nn.Module()
+    model.layers = nn.ModuleList([_KimiK3Layer()])
+    model.output_attn_res_proj = nn.Linear(32, 1, bias=False)
+    model.mm_projector = nn.Sequential(nn.Linear(32, 32), nn.GELU(), nn.Linear(32, 32))
+    config = copy.deepcopy(presets.QUANT_CFG_CHOICES[qformat])
+    config["algorithm"] = None  # placement only, no calibration forward
+    mtq.quantize(model, config)
+
+    layer = model.layers[0]
+    for module in (
+        layer.block_sparse_moe.routed_expert_down_proj,
+        layer.block_sparse_moe.routed_expert_up_proj,
+        layer.self_attention_res_proj,
+        layer.mlp_res_proj,
+        model.output_attn_res_proj,
+        model.mm_projector[0],
+        model.mm_projector[2],
+    ):
+        assert not module.weight_quantizer.is_enabled
+        assert not module.input_quantizer.is_enabled
+    assert layer.self_attn.o_proj.weight_quantizer.is_enabled
+
+
 def test_local_hessian_layerwise_is_scoped_to_the_model_recipe():
     # Layerwise calibration needs identifiable decoder layers, so it belongs to the model
     # recipe the published numbers come from -- not to the shared preset, which also backs