[OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter

modelopt/torch/models/__init__.py

@@ -91,6 +91,7 @@
     nemotron_h,
     minimax_m3_vl,
     llama4_text,
+    kimi_vl,
     nemotron_h_omni,
     qwen2_moe,
     qwen3,

modelopt/torch/models/_kimi.py

@@ -0,0 +1,206 @@
+# 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.
+
+"""What the Kimi-K2.5 and Kimi-VL adapters change in the modeling code their checkpoints ship.
+
+Both checkpoints' own VLM class builds its language model as ``DeepseekV3ForCausalLM(
+config.text_config)``, a DeepSeek-V3 whose MoE is an ``nn.ModuleList`` of per-expert Linears that
+expert parallelism cannot shard. The adapters point that name at transformers' native class, and
+keep the vision tower and projector from the checkpoint's code.
+"""
+
+# modelopt.torch.utils and transformers are imported inside the functions: the model registry imports
+# this module, and it must import without transformers, the optional [hf] extra.
+from ._remote_code import keep_patched
+
+
+def checkpoint_to_patch(
+    label: str, hf_config, ckpt_path: str | None, trust_remote_code: bool
+) -> str:
+    """Return the checkpoint whose own modeling code the adapter ``label`` patches.
+
+    Refuses first, before its ``prepare`` changes anything, a config the adapter cannot patch.
+    """
+    if type(hf_config).__module__.startswith("transformers."):
+        # transformers' own port (Kimi-K2.5's, from 5.13), which the patch never reaches: a load
+        # without trust_remote_code builds it, and so does one of a checkpoint naming no config class.
+        raise ValueError(
+            f"{label}: expert parallelism runs on the checkpoint's own modeling code, but this "
+            f"config is transformers' own port ({type(hf_config).__module__}). Load it with "
+            "trust_remote_code=True, from a checkpoint whose auto_map names its own AutoConfig."
+        )
+    if not hasattr(hf_config, "text_config"):
+        raise ValueError(f"{label}: hf_config has no .text_config (expected a VLM config)")
+    if ckpt_path is None:
+        raise ValueError(
+            f"{label}: expert parallelism patches the checkpoint's own modeling code, so it needs "
+            "the checkpoint's path (apply_distributed_support's ckpt_path)."
+        )
+    if not trust_remote_code:
+        # Patching imports the checkpoint's own modeling code, and transformers'
+        # get_class_from_dynamic_module runs it without checking trust_remote_code itself.
+        raise ValueError(
+            f"{label}: expert parallelism imports the checkpoint's own modeling code; "
+            "load it with trust_remote_code=True."
+        )
+    return ckpt_path
+
+
+def refuse_quantization_only_on_text_config(label: str, hf_config) -> None:
+    """Refuse a checkpoint whose ``quantization_config`` only its ``text_config`` declares.
+
+    The adapter replaces ``text_config`` with a native config without it, and the expert-parallel
+    load would then read the quantized experts as stored. It loads unquantized checkpoints only, so
+    the way in is an unquantized copy of the model.
+    """
+    if (
+        getattr(hf_config.text_config, "quantization_config", None) is not None
+        and getattr(hf_config, "quantization_config", None) is None
+    ):
+        raise ValueError(
+            f"{label}: the checkpoint declares quantization_config on text_config only, which the "
+            "native deepseek_v3 config replacing it drops, so the expert-parallel load would read "
+            "its experts as stored. Start from an unquantized copy of the model."
+        )
+
+
+def install_compat_import_shims() -> None:
+    """Add stubs for symbols the Kimi modeling imports that newer transformers removed, so it imports.
+
+    Each guards a torch.fx / legacy-attn-mask code path the native LM + MoonViT forward never hit, so a benign stub is
+    safe. Idempotent (only sets missing attrs).
+    """
+    from modelopt.torch.utils import print_rank_0
+
+    try:
+        import transformers.utils.import_utils as iu
+
+        if not hasattr(iu, "is_torch_fx_available"):
+            iu.is_torch_fx_available = lambda: False
+    except Exception as e:
+        print_rank_0(f"[kimi] compat-shim warning (import_utils): {type(e).__name__}: {e}")
+    try:
+        import transformers.pytorch_utils as pu
+
+        if not hasattr(pu, "is_torch_greater_or_equal_than_1_13"):
+            pu.is_torch_greater_or_equal_than_1_13 = True
+    except Exception as e:
+        print_rank_0(f"[kimi] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}")
+    try:
+        import transformers.modeling_attn_mask_utils as amu
+
+        if not hasattr(amu, "_prepare_4d_causal_attention_mask"):
+
+            def _prepare_4d_causal_attention_mask(*args, **kwargs):  # pragma: no cover
+                raise RuntimeError(
+                    "kimi: remote _prepare_4d_causal_attention_mask stub called -- the native LM "
+                    "builds its own mask, so this path should be unreachable."
+                )
+
+            amu._prepare_4d_causal_attention_mask = _prepare_4d_causal_attention_mask
+    except Exception as e:
+        print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}")
+
+
+def _drop_tie_weights_kwargs(orig_tie):
+    """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it.
+
+    Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote
+    ``tie_weights`` only calls its language model's, and the Kimi checkpoints set
+    ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. A
+    tied one would still load its embedding into both, with ``lm_head.weight`` reported missing.
+    """
+
+    def _tie_weights(self, *args, **kwargs):  # drop recompute_mapping etc.
+        return orig_tie(self)
+
+    setattr(_tie_weights, "_kimi_tie_weights_shim", True)
+    return _tie_weights
+
+
+def _native_for_the_prepared_config(own_class, native_class):
+    """The ``DeepseekV3ForCausalLM`` the patched module builds its language model with.
+
+    transformers' native class for the native ``DeepseekV3Config`` ``prepare`` puts on
+    ``text_config``, and the checkpoint's own class for its own config. Deciding the plan prepares a
+    copy of the config before the plan is known, and the patch outlives that copy: a load the plan
+    keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own
+    language model, as it would with no patch.
+    """
+    from transformers import DeepseekV3Config
+
+    def build(config, *args, **kwargs):
+        built = native_class if isinstance(config, DeepseekV3Config) else own_class
+        return built(config, *args, **kwargs)
+
+    setattr(build, "_kimi_native_swap", True)
+    return build
+
+
+def patch_language_model_to_native(
+    label: str, modeling: str, vlm_class: str, ckpt_path: str, trust_remote_code: bool
+) -> None:
+    """Point the checkpoint's ``modeling`` module's ``DeepseekV3ForCausalLM`` at the native class.
+
+    Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM``, so
+    ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 for ``self.language_model``
+    instead of the per-expert Linears version, from the config ``prepare`` made
+    (:func:`_native_for_the_prepared_config`).
+    """
+    from transformers.dynamic_module_utils import get_class_from_dynamic_module
+    from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM
+
+    from modelopt.torch.utils import print_rank_0
+
+    # Force the dynamic-module loader to stage + import the module into sys.modules.
+    get_class_from_dynamic_module(
+        f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code
+    )
+
+    def to_native(module):
+        current = getattr(module, "DeepseekV3ForCausalLM", None)
+        if current is None:
+            # Not "already native": the VLM class then builds its language model from a class this
+            # swap does not reach (renamed, or imported inside __init__), with per-expert Linears.
+            raise RuntimeError(
+                f"{label}: {module.__name__} defines no module-level DeepseekV3ForCausalLM to swap "
+                "for the transformers-native class; the checkpoint's modeling code changed."
+            )
+        if getattr(current, "_kimi_native_swap", False):
+            return False
+        module.DeepseekV3ForCausalLM = _native_for_the_prepared_config(
+            current, NativeDeepseekV3ForCausalLM
+        )
+        # The remote tie_weights() takes no kwargs, but transformers' post_init and load call
+        # tie_weights(recompute_mapping=False) -> TypeError.
+        kimi_cls = getattr(module, vlm_class, None)
+        if kimi_cls is not None:
+            own_tie = vars(kimi_cls).get("tie_weights")
+            if own_tie is not None and not getattr(own_tie, "_kimi_tie_weights_shim", False):
+                kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie)
+        return True
+
+    # The dynamic-module loader names the module like
+    # ``transformers_modules.<short-repo>.<hash>.<modeling>``. It stays patched when from_pretrained
+    # resolves the class again, even from a re-executed module.
+    patched = keep_patched(f".{modeling}", to_native)
+    if not patched:
+        raise RuntimeError(
+            f"{label}: could not find a loaded {modeling} in sys.modules to monkey-patch"
+        )
+    print_rank_0(
+        f"[{label}] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) "
+        f"in: {patched}"
+    )

modelopt/torch/models/deepseek_v3/specs.py

@@ -23,6 +23,17 @@
 
 __all__: list[str] = []
 
+# The plan for transformers releases that ship none for deepseek_v3 (before 5.13). From 5.13 its own
+# plan wins (an existing plan always does, see apply_distributed_support): the same experts, plus an
+# ep_router entry for the gate, which the load swaps for torch_a2a_router (global ids, unmasked). No
+# gate entry here on purpose: the router is logits-only and the MoE block computes GLOBAL top-k ids,
+# so the experts already receive unmasked global ids -- what the all-to-all dispatch wants.
+BASE_MODEL_EP_PLAN = {
+    "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+    "layers.*.mlp.experts.down_proj": "grouped_gemm",
+    "layers.*.mlp.experts": "moe_tp_experts",
+}
+
 # DeepseekV3MoE is invisible to the generic MoE detection: its class name does not end
 # in "SparseMoeBlock" and it calls its router ``gate``, so neither the name test nor the
 # structural router+experts test in is_moe matches. block_names below is what puts it on
@@ -37,17 +48,7 @@
         model_type="deepseek_v3",
         min_transformers_version="5.5",
         export_spec=ExportSpec(grouped_expert_export=True),
-        ep_spec=EPSpec(
-            # Stock transformers ships no plan for deepseek_v3 before 5.13; from 5.13 the config's
-            # own plan wins (see apply_distributed_support) and this one goes unused. No gate entry
-            # on purpose: the router is logits-only and the MoE block computes GLOBAL top-k ids, so
-            # the experts already receive unmasked global ids -- what the all-to-all dispatch wants.
-            base_model_ep_plan={
-                "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
-                "layers.*.mlp.experts.down_proj": "grouped_gemm",
-                "layers.*.mlp.experts": "moe_tp_experts",
-            },
-        ),
+        ep_spec=EPSpec(base_model_ep_plan=BASE_MODEL_EP_PLAN),
         moe_spec=MoESpec(
             block_names=("DeepseekV3MoE",),
             expert_linear_names=("gate_proj", "down_proj", "up_proj"),

modelopt/torch/models/kimi_vl/__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_vl specs (HF model type ``kimi_vl``)."""
+
+from . import specs

modelopt/torch/models/kimi_vl/ep.py

@@ -0,0 +1,56 @@
+# 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.
+
+"""Expert-parallel preparation for remote-code Kimi-VL (``kimi_vl``).
+
+``prepare`` replaces the per-expert DeepSeek-V3 language model the checkpoint's own code builds with
+transformers' native DeepSeek-V3, whose experts the expert-parallel load can shard, and keeps the
+vision tower and projector from that code. It needs trusted remote code and transformers 5.7.
+"""
+
+from .. import _kimi, _onfly_decompress
+from .._remote_code import stage_full_remote_closure
+
+
+def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_):
+    """Load Kimi-VL's language model as the transformers-native deepseek_v3.
+
+    Swap text_config to native deepseek_v3 + apply the EP plan + monkey-patch Kimi-VL's modeling to
+    build the native EP-capable DeepseekV3ForCausalLM for the language sub-model. No decompressor: the
+    Kimi-VL checkpoint is plain bf16.
+    """
+    # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe
+    # imports the registry.
+    from modelopt.torch.utils.moe import apply_distributed_support
+
+    ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code)
+    _kimi.refuse_quantization_only_on_text_config("kimi_vl", hf_config)
+
+    _kimi.install_compat_import_shims()
+    stage_full_remote_closure()
+
+    # Native deepseek_v3 config for the LM (drops the remote auto_map so the native class is used).
+    hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config)
+    # EP plan onto the native text_config (bubbles up to the VLM root under language_model.*).
+    hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config)
+
+    _kimi.patch_language_model_to_native(
+        "kimi_vl",
+        "modeling_kimi_vl",
+        "KimiVLForConditionalGeneration",
+        ckpt_path,
+        trust_remote_code,
+    )
+    return hf_config

modelopt/torch/models/kimi_vl/specs.py

@@ -0,0 +1,41 @@
+# 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_vl specs (HF model type ``kimi_vl``).
+
+Kimi-VL-A3B: a remote-code VLM whose text backbone is a Moonlight MoE. Its bundled
+modeling builds the language model from per-expert Linears, which the EP machinery cannot shard, so
+``ep.py`` swaps in the transformers-native deepseek_v3 language model.
+
+Only the expert-parallel section is registered under ``kimi_vl``: the MoE and export sections stay
+``deepseek_v3``'s, so lookups by ``kimi_vl`` find neither yet. Nor do ``is_moe``'s name and structure
+fallbacks recognize the native ``DeepseekV3MoE`` blocks ``prepare`` installs, so the HF export's
+MoE-block handling skips them. Sharing the sections is left to a follow-up.
+"""
+
+from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN
+from ..specs import EPSpec, ModelSpec, register
+from .ep import prepare
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="kimi_vl",
+        modeling_source="remote_code",
+        # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan.
+        ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN),
+    )
+)