From 030726f8301f2436658893db0ee1713aa98b82da Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:53:57 +0000 Subject: [PATCH 01/18] Add the Kimi-K3 expert-parallel adapter Kimi-K3's text backbone, kimi_linear, has no native transformers modeling to swap in, so the adapter patches the checkpoint's remote code in place: - it replaces the MoE block's per-expert ModuleList with fused experts that dispatch through the experts interface, keeping K3's latent expert space and its custom SwiGLU; - it registers converters that fuse the checkpoint's per-expert w1/w3 into gate_up_proj and stack w2 into down_proj, and has the MXFP4 experts decompressed as they load; - it fixes the parts of the remote model that the expert-parallel load and calibration trip over: KDA's A_log shape and its sigmoid(beta), which fla-core silently skips, the causal-mask and cache signatures, and the kwargs newer transformers pass to tie_weights. It needs transformers 5.8, for the decompression. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 1 + modelopt/torch/models/kimi_k3/__init__.py | 18 + modelopt/torch/models/kimi_k3/ep.py | 625 ++++++++++++++++++ modelopt/torch/models/kimi_k3/specs.py | 33 + tests/unit/torch/models/test_kimi_k3_ep.py | 720 +++++++++++++++++++++ 5 files changed, 1397 insertions(+) create mode 100644 modelopt/torch/models/kimi_k3/__init__.py create mode 100644 modelopt/torch/models/kimi_k3/ep.py create mode 100644 modelopt/torch/models/kimi_k3/specs.py create mode 100644 tests/unit/torch/models/test_kimi_k3_ep.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index e2f5aa2425b..37c4fd463c4 100644 --- a/modelopt/torch/models/__init__.py +++ b/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, diff --git a/modelopt/torch/models/kimi_k3/__init__.py b/modelopt/torch/models/kimi_k3/__init__.py new file mode 100644 index 00000000000..05b323e388c --- /dev/null +++ b/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 diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py new file mode 100644 index 00000000000..4ba29efcfdd --- /dev/null +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -0,0 +1,625 @@ +# 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. + +"""DP=EP all-to-all support for Kimi-K3 (``kimi_k3``). + +a remote-code VLM whose text backbone is ``kimi_linear`` (KDA/linear attention + DeepSeek-V3-style routed MoE with a +LATENT expert space). + +Kimi-K3 has no transformers-native modeling (``CONFIG_MAPPING_NAMES`` has no ``kimi_linear``/``kimi_k3``), +so -- unlike ``kimi_k25``, which swaps to native ``deepseek_v3`` -- we cannot get EP-capable fused experts +by swapping in a native backbone. We must adapt the checkpoint's REMOTE code in place. Two structural +obstacles and how this plugin removes them: + +1. **Per-expert ``nn.ModuleList``.** The remote ``KimiSparseMoeBlock`` holds ``self.experts`` as an + ``nn.ModuleList`` of ``num_experts`` ``KimiBlockSparseMLP`` (each a separate ``w1``=gate / ``w3``=up / + ``w2``=down Linear), and its ``forward`` runs a hardcoded per-expert Python loop (``moe_infer``). The + transformers EP framework (and ModelOpt's ``torch_a2a`` transport) instead needs the FUSED 3-D grouped + params ``gate_up_proj [E, 2*I, H]`` / ``down_proj [E, H, I]`` behind ``ExpertsInterface``. This plugin + replaces ``self.experts`` with a fused ``KimiK3Experts`` module (decorated with + ``use_experts_implementation``) and rewrites ``KimiSparseMoeBlock.forward`` to dispatch through it -- so + the generic ``torch_a2a`` experts forward (NCCL all-to-all -> quantizable ``_grouped_linear``) takes over. + +2. **No expert fusion for the checkpoint's names.** The checkpoint stores per-expert ``w1``/``w3``/``w2`` + (gate/up/down), and transformers registers no expert fusion for custom-code models. This plugin + registers ``kimi_linear`` WeightConverters that fuse ``...experts..w1|w3`` into ``gate_up_proj`` and + stack ``w2`` into ``down_proj``, which the ``grouped_gemm`` plan entries then EP-shard. + +3. **Compressed (mxfp4) experts.** In this checkpoint the 896 routed experts are compressed-tensors + ``mxfp4-pack-quantized`` (4-bit, group 32; stored as ``w1/w3/w2.weight_packed`` + ``.weight_scale``) -- + everything else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, + vision) is plain bf16. This plugin installs the format-agnostic on-the-fly decompressor + (``_onfly_decompress``) so from_pretrained dequantizes each per-expert packed weight -> bf16 INSIDE + the fused-expert conversion (EP-shard-aware -- only this rank's local experts), then drops the + ``quantization_config`` so transformers does not build a compressed-tensors HF quantizer. ModelOpt PTQ + then inserts fp8/nvfp4 quantizers on the resulting bf16 fused params. (Re-quantizing an already-4-bit + model to fp8/nvfp4 is lossy-on-lossy; that is the meaning of PTQ-ing this checkpoint.) + +Kimi-K3 specifics preserved as DENSE / replicated (NOT quantized through the EP experts path): + * the LATENT MoE wrapper -- ``routed_expert_down_proj`` (H->H_latent) / ``routed_expert_up_proj`` + (H_latent->H) and the optional ``routed_expert_norm`` -- the experts run in the H_latent + (``routed_expert_hidden_size``) space between them; the block forward keeps these around the fused call. + * the always-on ``shared_experts`` (a dense ``KimiMLP``) and the inline ``gate`` (``KimiMoEGate``, + sigmoid noaux_tc). The gate runs over the REPLICATED router weight and emits GLOBAL top-k ids, so it + stays a plain replicated module (no ``ep_router`` plan entry -- as for deepseek_v3's global-id + gate). The custom ``situ`` SwiGLU (``SituAndMul``) is preserved via the fused module's + ``_apply_gate``. + +VLM: the plan is set on ``text_config`` and bubbles up under ``language_model.model.*`` (the VLM builds +``self.language_model = KimiLinearForCausalLM(config.text_config)``); vision tower + mm projector stay bf16. + +Also shims ``OutputRecorder`` back into ``transformers.utils.generic`` (the remote code imports it there, +but it moved to ``transformers.modeling_utils``) so the modeling imports at all. + +Version floor: transformers >= 5.8, for the on-the-fly decompression. Developed and validated against +transformers 5.13 + flash-linear-attention (FLA) 0.5.0, which the remote code's KDA layers need at runtime. +""" + +import contextlib +import sys + +import torch +import torch.nn as nn + +from .. import _onfly_decompress +from .._remote_code import stage_full_remote_closure + +# base-relative EP plan (bubbles to `language_model.model.` in the VLM). MoE lives at +# `layers.*.block_sparse_moe` (NOT `.mlp`; layer 0 is dense `.mlp`). Fused expert params shard dim-0 via +# grouped_gemm; the experts container gets the EP group via moe_tp_experts (-> torch_a2a_experts for a2a). +# No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. +BASE_MODEL_EP_PLAN = { + "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", + "layers.*.block_sparse_moe.experts.down_proj": "grouped_gemm", + "layers.*.block_sparse_moe.experts": "moe_tp_experts", +} + +_PATCHED = False + + +def _install_output_recorder_shim() -> None: + """Restore the generic-utils symbol the remote modeling imports. + + The remote ``modeling_kimi`` does ``from transformers.utils.generic import OutputRecorder`` (used + once, as output-recording metadata), but that symbol moved to ``transformers.modeling_utils`` in + recent transformers. Re-export it into ``utils.generic`` (or inject a benign stub) BEFORE the remote + modeling imports so the import succeeds. Idempotent. + """ + import transformers.utils.generic as _g + + if hasattr(_g, "OutputRecorder"): + return + found = None + for _p in ("transformers.modeling_utils", "transformers.utils"): + try: + _m = __import__(_p, fromlist=["OutputRecorder"]) + if hasattr(_m, "OutputRecorder"): + found = _m.OutputRecorder + break + except Exception: + pass + if found is None: + + class OutputRecorder: # minimal stub: satisfies the import + the metadata use only + def __init__(self, *a, **k): + self.args, self.kwargs = a, k + + found = OutputRecorder + _g.OutputRecorder = found + + +def _find_module_defining(class_name: str): + """Return the loaded module that DEFINES a class named ``class_name``, or None. + + Identify the remote module by the CLASS it defines, NOT by its file name -- the plugin depends only on + the class names it patches, never on the remote code's file layout, so this survives remote-file + renames (e.g. ``modeling_kimi`` vs ``modeling_kimi_linear``). transformers + owns parsing + importing the remote files (via ``trust_remote_code`` + ``auto_map``); we only locate + the already-imported module. ``__module__ == module.__name__`` requires the class to be DEFINED here + (not merely imported), which uniquely picks the defining module (e.g. the backbone that defines + ``KimiSparseMoeBlock``, not the VLM wrapper that may import it). Guarded with try/except because + iterating ``sys.modules`` can touch lazy/magic modules -- notably ``torch.classes``, whose + ``__getattr__`` returns a namespace proxy for ANY name and raises on use (the proxy is not a ``type``, + so it is rejected regardless). + """ + for _name, mod in list(sys.modules.items()): + if mod is None: + continue + try: + obj = getattr(mod, class_name, None) + except Exception: + continue + if isinstance(obj, type) and getattr(obj, "__module__", None) == getattr( + mod, "__name__", None + ): + return mod + return None + + +def _find_remote_module(): + """Return the loaded remote backbone module. + + the one that DEFINES ``KimiSparseMoeBlock`` (the Kimi MoE block this plugin patches) -- or None. Name-independent; + see ``_find_module_defining``. + """ + return _find_module_defining("KimiSparseMoeBlock") + + +def _build_fused_experts_class(kmod): + """Build the fused experts class the expert-parallel path needs. + + Build the fused ``KimiK3Experts`` class (bound to the remote module's ``SituAndMul``), decorated + with ``use_experts_implementation`` so its forward dispatches through transformers' ``ExpertsInterface`` + and honors ``config._experts_implementation`` (``torch_a2a`` at EP-PTQ time). Stores the standard fused + grouped params ``gate_up_proj [E, 2*I, H_latent]`` / ``down_proj [E, H_latent, I]`` (``is_transposed= + False``, ``[E, out, in]``). ``_apply_gate`` preserves Kimi's ``situ`` SwiGLU. + """ + from transformers.integrations.moe import use_experts_implementation + + SituAndMul = kmod.SituAndMul # noqa: N806 (a class, keeping its own name) + _get_situ = kmod._get_situ_activation_params + + class KimiK3Experts(nn.Module): + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + # experts run in the LATENT space when use_latent_moe (routed_expert_hidden_size), else hidden. + self.hidden_dim = ( + config.routed_expert_hidden_size + if getattr(config, "use_latent_moe", False) + else config.hidden_size + ) + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_proj = nn.Parameter( + torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim) + ) + self.down_proj = nn.Parameter( + torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim) + ) + # explicit attrs the torch_a2a experts forward reads (also set by use_experts_implementation). + self.has_gate = True + self.has_bias = False + self.is_transposed = False + if getattr(config, "hidden_act", None) == "situ": + beta, linear_beta = _get_situ(config) + self._situ = SituAndMul(beta=beta, linear_beta=linear_beta) + else: + self._situ = None + + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: + # SituAndMul takes the FULL concat [gate; up] and splits internally -- matches the fused + # gate_up_proj row layout (rows [0:I]=gate from w1, [I:2I]=up from w3). + if self._situ is not None: + return self._situ(gate_up) + gate, up = gate_up.chunk(2, dim=-1) + return nn.functional.silu(gate) * up + + def forward(self, hidden_states, top_k_index, top_k_weights): + # Fallback (non-EP / eager) dense per-expert loop; the EP-PTQ path replaces this with the + # torch_a2a experts forward via config._experts_implementation. Mirrors the native + # DeepseekV3Experts/MiniMaxM3VLExperts reference loop, with situ gating. + final = torch.zeros_like(hidden_states) + with torch.no_grad(): + mask = nn.functional.one_hot(top_k_index, num_classes=self.num_experts).permute( + 2, 1, 0 + ) + hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(mask[expert_idx]) + cur = self._apply_gate( + nn.functional.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]) + ) + cur = nn.functional.linear(cur, self.down_proj[expert_idx]) + cur = cur * top_k_weights[token_idx, top_k_pos, None] + final.index_add_(0, token_idx, cur.to(final.dtype)) + return final + + # is_transposed=False: [E, out, in] (like deepseek/qwen). has_gate=True: gate_up is FUSED (2*I) and + # split by _apply_gate. has_bias=False: bias-free experts. + use_experts_implementation(KimiK3Experts, is_transposed=False, has_bias=False, has_gate=True) + return KimiK3Experts + + +def _patch_kimi_moe_block(kmod) -> None: + """Replace the per-expert ModuleList with the fused experts module. + + Swap ``KimiSparseMoeBlock.experts`` (per-expert ModuleList) for the fused experts module and rewrite + its ``forward`` to dispatch through it (keeping the latent down/up projections, optional norm, and + shared experts around the fused call). Idempotent. + """ + Block = kmod.KimiSparseMoeBlock # noqa: N806 (a class, keeping its own name) + if getattr(Block, "_kimi_k3_ep_patched", False): + return + FusedExperts = _build_fused_experts_class(kmod) # noqa: N806 (a class, keeping its own name) + + _orig_init = Block.__init__ + + def _init(self, config): + _orig_init(self, config) + # Replace the ModuleList of per-expert MLPs with the fused grouped-param module (EP-capable). + del self.experts + self.experts = FusedExperts(config) + + def _forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + # Fused / a2a experts: returns the top-k-combined output in the latent space. + y = self.experts(hidden_states, topk_idx, topk_weight) + 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.view(*orig_shape) + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + Block.__init__ = _init + Block.forward = _forward + Block._kimi_k3_ep_patched = True + print( + "[kimi_k3] patched KimiSparseMoeBlock: fused experts + ExpertsInterface forward", flush=True + ) + + +def _patch_kda_a_log(kmod) -> None: + """Fix a stale-HF-port bug in the KDA (linear-attention) layers. + + ``KimiDeltaAttention.__init__`` allocates ``A_log`` as ``empty(num_heads)`` (96), but the checkpoint stores it as + ``[head_dim]`` (128 = ``linear_attn_config['head_dim']``). ``from_pretrained`` therefore raises a size mismatch on + ``layers.*.self_attn.A_log`` (and ``ignore_mismatched_sizes=True`` would REINIT the trained decay -> garbage). The + forward passes ``self.A_log`` straight to FLA's ``chunk_kda`` / ``fused_recurrent_kda`` (co-designed with this + checkpoint), so the ONLY bug is the allocation shape. Re-allocate ``A_log`` as ``[head_dim]`` right after the + original __init__ so the trained values load into the right shape; ``torch.empty`` is fine (overwritten by the + checkpoint on load). Idempotent. + """ + KDA = getattr(kmod, "KimiDeltaAttention", None) # noqa: N806 (a class, keeping its own name) + if KDA is None or getattr(KDA, "_kimi_k3_a_log_patched", False): + return + _orig_init = KDA.__init__ + + def _init(self, *args, **kwargs): + _orig_init(self, *args, **kwargs) + a_log = getattr(self, "A_log", None) + if a_log is not None and a_log.shape[0] != self.head_dim: + self.A_log = nn.Parameter( + torch.empty(self.head_dim, dtype=a_log.dtype, device=a_log.device) + ) + + KDA.__init__ = _init + KDA._kimi_k3_a_log_patched = True + print( + "[kimi_k3] patched KimiDeltaAttention.A_log alloc: [num_heads] -> [head_dim] (checkpoint shape)", + flush=True, + ) + + +def _patch_kda_beta_sigmoid(kmod) -> None: + """Fix the KDA calibration NaN. + + The remote ``KimiDeltaAttention.forward`` passes the RAW gate scalar ``beta = b_proj(hidden_states)`` to FLA's + ``chunk_kda`` / ``fused_recurrent_kda`` with ``use_beta_sigmoid_in_kernel=True``, expecting the kernel to apply + ``sigmoid`` (the delta-rule learning rate must be in ``(0,1)``). In fla-core 0.5.x that flag is a **silent no-op** + for the main recurrence -- the kernel uses the raw beta. Real trained weights produce ``beta`` up to ~8.4, so the + per-step delta-rule eigenvalue ``|1 - beta*||k||^2|`` exceeds 1 (``beta*||k||^2 ~ 7.7 >> 2``) and the recurrent + state grows exponentially (~1e11 over 128 tokens -> ``inf`` -> ``o_norm`` turns it into NaN -> ``detected nan values + in amax`` during calibration). Random-weight smokes never tripped it because their beta is small (``beta*||k||^2 < + 2``). Verified offline on the real layer-1 inputs: pre-applying ``sigmoid(beta)`` and passing + ``use_beta_sigmoid_in_kernel=False`` bounds the output (o~0.1, state~4.6, matching the pure-torch + ``naive_recurrent_kda`` reference), vs 8e9 unfixed. + + Fix: wrap the module-level ``chunk_kda`` / ``fused_recurrent_kda`` in the remote module; when invoked + with ``use_beta_sigmoid_in_kernel=True``, replace ``beta`` with ``beta.sigmoid()`` and set the flag + False so the transform happens (correctly) in torch. Idempotent. + """ + import inspect + + import torch + + if getattr(kmod, "_kimi_k3_beta_sigmoid_patched", False): + return + patched_any = False + for fname in ("chunk_kda", "fused_recurrent_kda"): + orig = getattr(kmod, fname, None) + if orig is None or getattr(orig, "_kimi_k3_beta_wrapped", False): + continue + + def _squash(arguments): + if arguments.get("use_beta_sigmoid_in_kernel"): + beta = arguments.get("beta") + if isinstance(beta, torch.Tensor): + arguments["beta"] = torch.sigmoid(beta.float()).to(beta.dtype) + arguments["use_beta_sigmoid_in_kernel"] = False + + def _make(orig_fn): + # Bind to the kernel's signature so beta and the flag are found however they are passed. + try: + signature = inspect.signature(orig_fn) + except (TypeError, ValueError): + signature = None + + def _wrapped(*args, **kwargs): + try: + bound = signature.bind(*args, **kwargs) # type: ignore[union-attr] + except (AttributeError, TypeError): # no signature, or a call it rejects + _squash(kwargs) + return orig_fn(*args, **kwargs) + bound.apply_defaults() + _squash(bound.arguments) + return orig_fn(*bound.args, **bound.kwargs) + + setattr(_wrapped, "_kimi_k3_beta_wrapped", True) + return _wrapped + + setattr(kmod, fname, _make(orig)) + patched_any = True + kmod._kimi_k3_beta_sigmoid_patched = True + if patched_any: + print( + "[kimi_k3] patched chunk_kda/fused_recurrent_kda: apply sigmoid(beta) in torch " + "(fla-core use_beta_sigmoid_in_kernel is a no-op -> KDA state explodes on real weights)", + flush=True, + ) + + +def _shim_create_causal_mask(kmod) -> None: + """tf-version compat for the backbone forward's ``create_causal_mask`` call. + + The remote code (written for an older transformers) calls it as ``create_causal_mask(config=, input_embeds=, + attention_mask=, cache_position=, past_key_values=, position_ids=)``, but tf 5.13's ``create_causal_mask`` renamed + ``input_embeds`` -> ``inputs_embeds`` and no longer accepts ``cache_position`` -> ``TypeError`` at the first + calibration forward. The remote module binds ``create_causal_mask`` at import (``from transformers.masking_utils + import ...``), so wrap that module-level reference: translate ``input_embeds`` -> ``inputs_embeds`` and drop any + kwargs the installed signature does not accept (robust to further drift). Idempotent. + """ + import functools + import inspect + + orig = getattr(kmod, "create_causal_mask", None) + if orig is None or getattr(orig, "_kimi_k3_ccm_shim", False): + return + try: + accepted = set(inspect.signature(orig).parameters) + except (TypeError, ValueError): + accepted = None + + @functools.wraps(orig) + def _wrapped(*args, **kwargs): + if "input_embeds" in kwargs and "inputs_embeds" not in kwargs: + kwargs["inputs_embeds"] = kwargs.pop("input_embeds") + if accepted is not None: + kwargs = {k: v for k, v in kwargs.items() if k in accepted} + return orig(*args, **kwargs) + + setattr(_wrapped, "_kimi_k3_ccm_shim", True) + kmod.create_causal_mask = _wrapped + print( + "[kimi_k3] shimmed create_causal_mask (input_embeds->inputs_embeds; drop unaccepted kwargs; tf>=5.13)", + flush=True, + ) + + +def _patch_kimi_cache_get_mask_sizes(kmod) -> None: + """tf-version compat for the KDA hybrid cache. + + tf 5.13's mask preprocessing calls ``cache.get_mask_sizes(q_length: int, layer_idx)`` (q_length = + inputs_embeds.shape[1], an int), but the remote ``KimiDynamicCache.get_mask_sizes(cache_position, layer_idx)`` does + ``cache_position.shape[0]`` -- it expects a TENSOR (older tf) -> ``AttributeError: 'int' object has no attribute + 'shape'`` in the calibration forward. Re-implement to accept either an int query_length or a tensor cache_position + (the body is otherwise trivial: kv_length = query_length + past_seen_tokens, kv_offset = 0). Idempotent. + """ + Cache = getattr(kmod, "KimiDynamicCache", None) # noqa: N806 (a class, keeping its own name) + if Cache is None or getattr(Cache, "_kimi_k3_gms_patched", False): + return + + def _get_mask_sizes(self, query_length, layer_idx): + ql = query_length.shape[0] if torch.is_tensor(query_length) else int(query_length) + past_seen_tokens = self.get_seq_length(layer_idx) + return ql + past_seen_tokens, 0 + + Cache.get_mask_sizes = _get_mask_sizes + Cache._kimi_k3_gms_patched = True + print( + "[kimi_k3] patched KimiDynamicCache.get_mask_sizes (accept int query_length; tf>=5.13)", + flush=True, + ) + + +def _shim_tie_weights_kwargs() -> None: + """Absorb the ``recompute_mapping`` kwarg newer transformers passes to ``tie_weights``. + + Transformers 5.13's ``from_pretrained`` calls ``model.tie_weights(recompute_mapping=...)``, but + Kimi-K3's remote ``KimiK3ForConditionalGeneration.tie_weights(self)`` has the older no-kwarg signature + (it just delegates to ``self.language_model.tie_weights()``) -> ``TypeError``. Wrap it to swallow the + extra kwargs (they only concern lm_head/embeddings tying, which the delegate handles). Idempotent. + (Same class of fix as kimi_vl.) + + Locate the VLM wrapper by the class it DEFINES (``KimiK3ForConditionalGeneration``) via + ``_find_module_defining`` -- name-independent and safe against the ``torch.classes`` magic module + (see that helper). The LM ``KimiLinearForCausalLM.tie_weights`` is the native ``PreTrainedModel`` one + (accepts the kwargs) so it needs no wrap. + """ + + def _make_tie(orig): + def _tie(self, *args, **kwargs): + return orig(self) + + setattr(_tie, "_kimi_k3_kwargs_shim", True) + return _tie + + mod = _find_module_defining("KimiK3ForConditionalGeneration") + if mod is None: + return + cls = getattr(mod, "KimiK3ForConditionalGeneration", None) + if cls is None: + return + tw = cls.__dict__.get("tie_weights") # only if defined ON the class + if tw is not None and not getattr(tw, "_kimi_k3_kwargs_shim", False): + cls.tie_weights = _make_tie(tw) + + +def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> None: + """Import the remote modeling now, so its classes can be patched before the model is built. + + Force the remote modeling into ``sys.modules`` (so we can monkey-patch its classes BEFORE + ``from_pretrained`` instantiates the model) by resolving the model class the checkpoint's config + declares in its ``auto_map`` -- no hardcoded remote file/module name. Runs the OutputRecorder + closure + shims first. + """ + _install_output_recorder_shim() + stage_full_remote_closure() + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + auto_map = getattr(hf_config, "auto_map", None) or {} + class_ref = auto_map.get("AutoModelForCausalLM") or auto_map.get("AutoModel") + if not class_ref: + raise RuntimeError( + "kimi_k3: config.auto_map has no AutoModel/AutoModelForCausalLM entry to force-load the remote " + "modeling" + ) + get_class_from_dynamic_module(class_ref, ckpt_path, trust_remote_code=trust_remote_code) + + +def _register_kimi_k3_conversions() -> None: + """Register the per-expert to fused-expert weight converters. + + Register the per-expert ``w1/w3/w2 -> gate_up_proj/down_proj`` fusion WeightConverters for Kimi-K3's + ``kimi_linear`` backbone. + + transformers 5.13's ``get_model_conversion_mapping`` SKIPS building any weight conversions for + custom-code (remote) models unless their class/model_type is in ``USER_REGISTERED_MAPPINGS`` + (``if submodule.is_custom_code() and class_name/model_type not in USER_REGISTERED_MAPPINGS: continue``). + Kimi-K3 is remote code, so the per-expert -> fused-expert fusion is NEVER built: the checkpoint's + ``experts..w1/w3/w2.weight_packed`` land as UNEXPECTED and the model's fused ``gate_up_proj`` / + ``down_proj`` as MISSING (newly-initialized garbage, and allocated at FULL 896-expert size on every + rank -> OOM at scale). Register the fusion converters ourselves (mirrors transformers' own native + ``mlp.experts.*.w1/w3/w2.weight`` fusion, adapted to K3's ``block_sparse_moe.experts``): w1=gate, + w3=up (Concatenate gate-then-up on dim 1), w2=down. The on-the-fly decompressor then rewrites these + ``.weight`` sources to the checkpoint's ``.weight_packed`` + scales (mxfp4 -> bf16). Idempotent. + """ + try: + from transformers.conversion_mapping import register_checkpoint_conversion_mapping + from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter + except Exception as e: # pragma: no cover - version guard + print( + f"[kimi_k3] WARN: transformers conversion API unavailable ({type(e).__name__}: {e}); " + "per-expert fusion NOT registered -- experts will load as MISSING", + flush=True, + ) + return + + def _mapping(): + return [ + WeightConverter( + source_patterns=[ + "block_sparse_moe.experts.*.w1.weight", + "block_sparse_moe.experts.*.w3.weight", + ], + target_patterns="block_sparse_moe.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="block_sparse_moe.experts.*.w2.weight", + target_patterns="block_sparse_moe.experts.down_proj", + operations=[MergeModulelist(dim=0)], + ), + ] + + # Register for the backbone MODEL_TYPE "kimi_linear" ONLY (not the class names). The VLM nests + # KimiLinearForCausalLM (`language_model`) -> KimiLinearModel (`language_model.model`), both model_type + # kimi_linear; registering by model_type applies the fusion ONCE (at the outer `language_model.` scope) + # and the inner submodule is de-duplicated (ancestor already claimed the model_type). Registering the + # CLASS names too would apply it a second time at `language_model.model.` scope -> the same expert keys + # converted twice -> conflict at load. + register_checkpoint_conversion_mapping("kimi_linear", _mapping(), overwrite=True) + print( + "[kimi_k3] registered per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion for kimi_linear " + "(custom-code conversion; tf>=5.13 skips these otherwise)", + flush=True, + ) + + +def prepare( + hf_config, + *, + ckpt_path: str | None = None, + trust_remote_code: bool = False, + **_, +): + """Adapt Kimi-K3's remote MoE to the EP/a2a framework and set its EP plan. + + Steps: shim OutputRecorder -> force-load + monkey-patch the remote KimiSparseMoeBlock (fused experts + + ExpertsInterface forward) -> set base_model_ep_plan on text_config -> register the w1/w3/w2 fusion + converters and the decompressor. Returns the (unchanged) config. + + The 896 ``mxfp4-pack-quantized`` routed experts are DECOMPRESSED at load: ``_onfly_decompress`` + dequantizes each local expert's packed weight -> bf16 INSIDE the fused-expert conversion, stored as + bf16 fused params, and the compressed-tensors ``quantization_config`` is dropped. ModelOpt then + inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM + (the memory bottleneck), but PTQ needs the bf16 weights. + """ + global _PATCHED + from modelopt.torch.utils.moe import require_transformers + + require_transformers("5.8", "kimi_k3 expert-parallel PTQ") + + text_cfg = getattr(hf_config, "text_config", None) or hf_config + + # Monkey-patch the remote modeling (needs ckpt_path to force-load it). If ckpt_path is missing we can + # only set the config plan; the load would then fall back to the non-EP ModuleList layout -> callers + # must thread ckpt_path through apply_distributed_support. + if ckpt_path is not None and not _PATCHED: + _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code) + kmod = _find_remote_module() + if kmod is None: + raise RuntimeError( + "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" + ) + _patch_kimi_moe_block(kmod) + _patch_kda_a_log(kmod) + _patch_kda_beta_sigmoid(kmod) + _shim_create_causal_mask(kmod) + _patch_kimi_cache_get_mask_sizes(kmod) + _shim_tie_weights_kwargs() + _PATCHED = True + + if not getattr(text_cfg, "base_model_ep_plan", None): + # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts + # (-> torch_a2a_experts). + text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) + text_cfg._experts_implementation = "torch_a2a" + + # The 896 routed experts are compressed-tensors `mxfp4-pack-quantized` (4-bit, group 32); everything + # else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, vision) is + # plain bf16. Register the per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion converters (tf 5.13 + # skips these for custom-code models) BEFORE the decompressor patches get_model_conversion_mapping, so + # the decompressor has converters to rewrite into compressed-aware (weight_packed -> bf16) ones. Drop + # quantization_config afterwards so transformers builds no competing HF quantizer. + _register_kimi_k3_conversions() + # ValueError == no compressed-tensors quantization_config, i.e. plain bf16 experts and nothing to + # decompress. + with contextlib.suppress(ValueError): + _onfly_decompress.install_onfly_decompress(hf_config) + for cfg in (hf_config, text_cfg): + if getattr(cfg, "quantization_config", None) is not None: + cfg.quantization_config = None + return hf_config diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py new file mode 100644 index 00000000000..6a9bfd90542 --- /dev/null +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -0,0 +1,33 @@ +# 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 prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k3", + modeling_source="remote_code", + ep_spec=EPSpec(prepare=prepare), + ) +) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py new file mode 100644 index 00000000000..56b300aa7b2 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -0,0 +1,720 @@ +# 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 adapter for Kimi-K3 (``kimi_k3``), whose kimi_linear backbone is remote code. + +The checkpoint's modeling is not available here, so the tests build a stand-in remote module -- the +classes and functions the adapter patches, with the per-expert ``nn.ModuleList`` MoE the remote code +ships -- and check what the adapter does to it on tiny CPU tensors. Every process-wide patch is +undone after each test. +""" + +import copy +import sys +import types +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.6"): + pytest.skip("Kimi-K3 expert parallelism needs transformers >= 5.6", allow_module_level=True) + +import transformers.conversion_mapping as cm +import transformers.core_model_loading as cml +import transformers.dynamic_module_utils as dmu +import transformers.modeling_utils as mu +import transformers.utils.generic as hf_generic +from transformers.core_model_loading import WeightConverter, build_glob_alternation +from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS + +from modelopt.torch.models import _onfly_decompress, get_spec +from modelopt.torch.models.kimi_k3 import ep as kimi_k3 + +ROUTES = torch.tensor([[0, 3], [1, 2], [3, 1], [2, 0], [1, 3]]) # top-2 experts of 5 tokens + + +def _restore(monkeypatch, owner, name): + """Have ``monkeypatch`` put ``owner.`` back after the test, or remove it if it was absent.""" + if hasattr(owner, name): + monkeypatch.setattr(owner, name, getattr(owner, name)) + else: + monkeypatch.setattr(owner, name, None, raising=False) + monkeypatch.delattr(owner, name) + + +@pytest.fixture +def private_conversion_registry(monkeypatch): + """Let a test register conversion mappings without leaking them into transformers' registry.""" + monkeypatch.setattr(cm, "_checkpoint_conversion_mapping_cache", None) # rebuilt on first use + if hasattr(cm, "USER_REGISTERED_MAPPINGS"): # transformers >= 5.13 + monkeypatch.setattr(cm, "USER_REGISTERED_MAPPINGS", set(cm.USER_REGISTERED_MAPPINGS)) + + +@pytest.fixture +def restore_loader_hooks(monkeypatch): + """Undo the process-wide transformers hooks the on-the-fly decompressor installs.""" + for owner, name in ( + (WeightConverter, "add_tensor"), + (WeightConverter, "materialize_tensors"), + (cml, "_materialize_copy"), + (mu, "get_model_conversion_mapping"), + ): + monkeypatch.setattr(owner, name, getattr(owner, name)) + + +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 _config(**overrides): + """A tiny kimi_linear MoE config: 4 routed experts, top-2, run in a 6-wide latent space.""" + fields = { + "num_experts": 4, + "num_experts_per_token": 2, + "hidden_size": 8, + "routed_expert_hidden_size": 6, + "moe_intermediate_size": 5, + "hidden_act": "silu", + "situ_beta": 1.5, + "situ_linear_beta": 0.25, + "use_latent_moe": True, + "latent_moe_use_norm": True, + "num_shared_experts": 1, + "_experts_implementation": "eager", + } + return SimpleNamespace(**{**fields, **overrides}) + + +def _randomize(module, seed=0): + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for param in module.parameters(): + param.copy_(torch.randn(param.shape, generator=gen)) + + +def _remote_module(name="kimi_k3_remote_stub"): + """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 + self.use_latent_moe = config.use_latent_moe + self.latent_moe_use_norm = config.latent_moe_use_norm + hidden = config.hidden_size + width = config.routed_expert_hidden_size if config.use_latent_moe else hidden + self.experts = nn.ModuleList( + KimiBlockSparseMLP(config, width) for _ in range(config.num_experts) + ) + self.gate = KimiMoEGate(config) + if config.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__() + self.head_dim = head_dim + self.A_log = nn.Parameter(torch.zeros(num_heads, dtype=torch.bfloat16)) # stale shape + + def chunk_kda(q, k, v, g, beta=None, use_beta_sigmoid_in_kernel=False): + remote.calls.append(("chunk_kda", beta, use_beta_sigmoid_in_kernel)) + return "chunk_kda" + + def fused_recurrent_kda(q, k, v, g, beta=None, use_beta_sigmoid_in_kernel=False): + remote.calls.append(("fused_recurrent_kda", beta, use_beta_sigmoid_in_kernel)) + return "fused_recurrent_kda" + + def create_causal_mask( + config, inputs_embeds, attention_mask, past_key_values, position_ids=None + ): + # current transformers' signature: ``inputs_embeds``, and no ``cache_position`` + args = (config, inputs_embeds, attention_mask, past_key_values, position_ids) + remote.calls.append(("create_causal_mask", *args)) + return "mask" + + class KimiDynamicCache: + def __init__(self, seen_tokens): + self.seen_tokens = seen_tokens + + def get_seq_length(self, layer_idx=0): + remote.calls.append(("get_seq_length", layer_idx)) + return self.seen_tokens + + def get_mask_sizes(self, cache_position, layer_idx): # wants a tensor, as older callers + return cache_position.shape[0] + self.get_seq_length(layer_idx), 0 + + class KimiK3ForConditionalGeneration: + def tie_weights(self): # the remote override, with the old no-kwarg signature + remote.calls.append(("tie_weights", self)) + return "tied" + + for obj in ( + _get_situ_activation_params, + SituAndMul, + KimiBlockSparseMLP, + KimiMoEGate, + KimiSparseMoeBlock, + KimiDeltaAttention, + chunk_kda, + fused_recurrent_kda, + create_causal_mask, + KimiDynamicCache, + KimiK3ForConditionalGeneration, + ): + obj.__module__ = name # defined in the remote module, as far as any lookup can tell + setattr(remote, obj.__name__, obj) + return remote + + +def _load(converters, checkpoint): + """Route, collect and fuse ``checkpoint`` the way transformers' loader does.""" + per_target = {} + for key, tensor in checkpoint.items(): + for conv in converters: + target, pattern = conv.rename_source_key(key) + if pattern is not None: + break + else: + raise AssertionError(f"no converter claims {key}") + future = Future() + future.set_result(tensor) + per_target.setdefault(target, copy.deepcopy(conv)).add_tensor(target, key, pattern, future) + fused = {} + for target, conv in per_target.items(): + fused.update(conv.convert(target)) + return fused + + +# --------------------------------------------------------------------------- the fused experts + + +@pytest.mark.parametrize("hidden_act", ["silu", "situ"]) +def test_fused_experts_match_a_per_expert_reference(hidden_act): + """Per token: the routing-weighted sum over its top-k experts of down(act(gate x) * up x), with + gate read from rows [0:I] and up from rows [I:2I] of ``gate_up_proj``. A ``situ`` model hands + the full [gate; up] row to the remote ``SituAndMul``. Experts live in the latent width.""" + config = _config(hidden_act=hidden_act) + experts = kimi_k3._build_fused_experts_class(_remote_module())(config) + n, inter, latent = 4, config.moe_intermediate_size, config.routed_expert_hidden_size + assert experts.gate_up_proj.shape == (n, 2 * inter, latent) + assert experts.down_proj.shape == (n, latent, inter) + _randomize(experts) + gen = torch.Generator().manual_seed(1) + tokens = torch.randn(len(ROUTES), latent, generator=gen) + route_weights = torch.rand(ROUTES.shape, generator=gen) + + def act(gate_up): + if hidden_act == "situ": + return _situ(gate_up, config.situ_beta, config.situ_linear_beta) + return F.silu(gate_up[:inter]) * gate_up[inter:] + + with torch.no_grad(): + expected = torch.stack( + [ + sum( + w * (experts.down_proj[e] @ act(experts.gate_up_proj[e] @ x)) + for e, w in zip(top_e, top_w) + ) + for x, top_e, top_w in zip(tokens, ROUTES, route_weights) + ] + ) + torch.testing.assert_close(experts(tokens, ROUTES, route_weights), expected) + + +def test_fused_experts_dispatch_through_the_experts_interface(monkeypatch): + """``config._experts_implementation`` picks the forward (``torch_a2a`` under EP), so the module + must route through transformers' ``ExpertsInterface``; and a stock implementation reading the + fused params agrees with the eager loop, so the declared layout ([E, out, in], gated) is right.""" + routed = [] + batched_mm = ALL_EXPERTS_FUNCTIONS["batched_mm"] + + def probe(module, *args): + routed.append(module) + return batched_mm(module, *args) + + monkeypatch.setitem(ALL_EXPERTS_FUNCTIONS, "modelopt_test_probe", probe) + config = _config(hidden_act="situ") + experts = kimi_k3._build_fused_experts_class(_remote_module())(config) + _randomize(experts) + gen = torch.Generator().manual_seed(1) + tokens = torch.randn(len(ROUTES), config.routed_expert_hidden_size, generator=gen) + route_weights = torch.rand(ROUTES.shape, generator=gen) + + with torch.no_grad(): + eager = experts(tokens, ROUTES, route_weights) + config._experts_implementation = "modelopt_test_probe" + dispatched = experts(tokens, ROUTES, route_weights) + + assert routed == [experts] + torch.testing.assert_close(dispatched, eager) + + +@pytest.mark.parametrize( + "overrides", + [ + {}, + { + "use_latent_moe": False, + "latent_moe_use_norm": False, + "routed_expert_hidden_size": None, + "num_shared_experts": None, + }, + ], + ids=["latent-norm-shared", "plain"], +) +def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): + """Swapping the per-expert ModuleList for the fused module keeps the block's output: the latent + projections, the latent norm and the shared experts stay around the fused call, and w1 / w3 / + w2 are the gate rows, the up rows and the down of the fused params.""" + remote = _remote_module() + config = _config(**overrides) + reference = remote.KimiSparseMoeBlock(config) + _randomize(reference) + gen = torch.Generator().manual_seed(1) + hidden_states = torch.randn(2, 3, config.hidden_size, generator=gen) + with torch.no_grad(): + expected = reference(hidden_states) + + kimi_k3._patch_kimi_moe_block(remote) + block = remote.KimiSparseMoeBlock(config) + assert not isinstance(block.experts, nn.ModuleList) + dense = {k: v for k, v in reference.state_dict().items() if not k.startswith("experts.")} + result = block.load_state_dict(dense, strict=False) + assert sorted(result.missing_keys) == ["experts.down_proj", "experts.gate_up_proj"] + assert not result.unexpected_keys + with torch.no_grad(): + gate_up = [torch.cat([m.w1.weight, m.w3.weight]) for m in reference.experts] + block.experts.gate_up_proj.copy_(torch.stack(gate_up)) + block.experts.down_proj.copy_(torch.stack([m.w2.weight for m in reference.experts])) + torch.testing.assert_close(block(hidden_states), expected) + + +def test_patching_the_block_twice_wraps_it_once(): + """A second ``prepare`` must not stack another fused-experts allocation onto ``__init__``.""" + remote = _remote_module() + kimi_k3._patch_kimi_moe_block(remote) + patched = (remote.KimiSparseMoeBlock.__init__, remote.KimiSparseMoeBlock.forward) + + kimi_k3._patch_kimi_moe_block(remote) + + assert (remote.KimiSparseMoeBlock.__init__, remote.KimiSparseMoeBlock.forward) == patched + + +# ------------------------------------------------------------------- the linear-attention fixes + + +def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(): + """fla-core ignores ``use_beta_sigmoid_in_kernel``, so both KDA entry points apply + ``sigmoid(beta)`` themselves (in beta's dtype) and clear the flag; calls without the flag pass + beta through untouched. Each entry point is wrapped once.""" + remote = _remote_module() + kimi_k3._patch_kda_beta_sigmoid(remote) + wrapped = (remote.chunk_kda, remote.fused_recurrent_kda) + kimi_k3._patch_kda_beta_sigmoid(remote) + assert (remote.chunk_kda, remote.fused_recurrent_kda) == wrapped + beta = torch.tensor([-8.0, 0.0, 8.4], dtype=torch.bfloat16) + + assert remote.chunk_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=True) == "chunk_kda" + remote.fused_recurrent_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=True) + + squashed = torch.sigmoid(beta.float()).to(torch.bfloat16) + assert [name for name, *_ in remote.calls] == ["chunk_kda", "fused_recurrent_kda"] + for _, seen, flag in remote.calls: + assert flag is False + assert seen.dtype == torch.bfloat16 + assert torch.equal(seen, squashed) + + remote.calls.clear() + remote.chunk_kda(1, 2, 3, 4, beta=beta) + remote.fused_recurrent_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=False) + assert [(seen is beta, flag) for _, seen, flag in remote.calls] == [(True, False)] * 2 + + +def test_kda_beta_passed_positionally_is_squashed_too(): + """The fix must not depend on how the remote code spells the call.""" + remote = _remote_module() + kimi_k3._patch_kda_beta_sigmoid(remote) + beta = torch.tensor([-8.0, 8.4], dtype=torch.bfloat16) + + remote.chunk_kda(1, 2, 3, 4, beta, True) + + ((_, seen, flag),) = remote.calls + assert flag is False + assert torch.equal(seen, torch.sigmoid(beta.float()).to(torch.bfloat16)) + + +def test_kda_a_log_takes_the_checkpoints_head_dim_shape(): + """The stale port allocates ``A_log`` as [num_heads]; the checkpoint stores [head_dim], which + would fail to load (or be re-initialized). The reallocation keeps the dtype.""" + remote = _remote_module() + kimi_k3._patch_kda_a_log(remote) + + attention = remote.KimiDeltaAttention(num_heads=3, head_dim=4) + + assert attention.A_log.shape == (4,) + assert attention.A_log.dtype == torch.bfloat16 + + +# ------------------------------------------------------------- transformers-version compatibility + + +def test_create_causal_mask_shim_adapts_the_remote_call_once(): + """The remote code calls ``create_causal_mask(input_embeds=..., cache_position=...)``; current + transformers takes ``inputs_embeds`` and no ``cache_position``. The shim renames the one, drops + what the installed signature does not accept, and wraps only once.""" + remote = _remote_module() + kimi_k3._shim_create_causal_mask(remote) + shim = remote.create_causal_mask + kimi_k3._shim_create_causal_mask(remote) + assert remote.create_causal_mask is shim + + mask = remote.create_causal_mask( + config="config", + input_embeds="embeds", + attention_mask="mask_2d", + cache_position="positions", + past_key_values="cache", + position_ids="ids", + ) + + assert mask == "mask" + assert remote.calls == [("create_causal_mask", "config", "embeds", "mask_2d", "cache", "ids")] + + +def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): + """transformers >= 5.13 passes the query length as an int; the remote cache expects a + ``cache_position`` tensor. Either way kv_length is the query length plus the cached tokens.""" + remote = _remote_module() + kimi_k3._patch_kimi_cache_get_mask_sizes(remote) + cache = remote.KimiDynamicCache(seen_tokens=5) + + assert cache.get_mask_sizes(3, 2) == (8, 0) + assert cache.get_mask_sizes(torch.arange(4), 1) == (9, 0) + assert remote.calls == [("get_seq_length", 2), ("get_seq_length", 1)] + + +def test_tie_weights_shim_drops_the_new_kwargs_once(monkeypatch): + """transformers >= 5.13 calls ``tie_weights(recompute_mapping=...)``; the remote VLM's override + takes no kwargs. The shim drops them, delegates once per call, and wraps only once.""" + remote = _remote_module() + monkeypatch.setitem(sys.modules, remote.__name__, remote) + vlm = remote.KimiK3ForConditionalGeneration + kimi_k3._shim_tie_weights_kwargs() + shim = vlm.__dict__["tie_weights"] + kimi_k3._shim_tie_weights_kwargs() + assert vlm.__dict__["tie_weights"] is shim + + model = vlm() + + assert model.tie_weights(missing_keys=set(), recompute_mapping=False) == "tied" + assert remote.calls == [("tie_weights", model)] + + +def test_tie_weights_shim_leaves_an_inherited_tie_weights_alone(monkeypatch): + """Only an override defined on the VLM class has the old signature; an inherited + ``tie_weights`` already takes the kwargs and must keep receiving them.""" + + class PreTrainedStub: + def tie_weights(self, **kwargs): + return kwargs + + remote = types.ModuleType("kimi_k3_remote_inheriting_stub") + vlm = type("KimiK3ForConditionalGeneration", (PreTrainedStub,), {"__module__": remote.__name__}) + remote.KimiK3ForConditionalGeneration = vlm + monkeypatch.setitem(sys.modules, remote.__name__, remote) + + kimi_k3._shim_tie_weights_kwargs() + + assert vlm().tie_weights(recompute_mapping=False) == {"recompute_mapping": False} + + +# ------------------------------------------------------------------- weight conversion and plan + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_registered_conversions_fuse_w1_w3_into_gate_up_and_stack_w2(): + """transformers builds no weight conversions for remote code unless they are registered. The + registered pair fuses the checkpoint's per-expert w1 (gate) and w3 (up) into + ``gate_up_proj = [w1; w3]`` and stacks w2 into ``down_proj``, under the VLM's prefix.""" + kimi_k3._register_kimi_k3_conversions() + converters = cm.get_checkpoint_conversion_mapping("kimi_linear") + n_experts, inter, hidden = 3, 2, 4 + shapes = {"w1": (inter, hidden), "w3": (inter, hidden), "w2": (hidden, inter)} + gen = torch.Generator().manual_seed(0) + weights = { + (name, e): torch.randn(shape, generator=gen) + for e in range(n_experts) + for name, shape in shapes.items() + } + prefix = "language_model.model.layers.1.block_sparse_moe.experts." + + fused = _load(converters, {f"{prefix}{e}.{name}.weight": w for (name, e), w in weights.items()}) + + gate_up = [torch.cat([weights["w1", e], weights["w3", e]]) for e in range(n_experts)] + assert torch.equal(fused[prefix + "gate_up_proj"], torch.stack(gate_up)) + down = [weights["w2", e] for e in range(n_experts)] + assert torch.equal(fused[prefix + "down_proj"], torch.stack(down)) + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_the_plan_the_conversions_and_the_fused_params_name_the_same_tensors(): + """Three separately written names must agree, or EP silently shards nothing: the fused module's + parameters, the conversion targets that fill them, and the EP plan globs (matched the way + transformers matches them). The router, latent projections and shared experts stay replicated.""" + remote = _remote_module() + kimi_k3._patch_kimi_moe_block(remote) + layer = nn.Module() + layer.block_sparse_moe = remote.KimiSparseMoeBlock(_config()) + model = nn.Module() + model.layers = nn.ModuleList([layer]) + kimi_k3._register_kimi_k3_conversions() + converters = cm.get_checkpoint_conversion_mapping("kimi_linear") + plan = kimi_k3.BASE_MODEL_EP_PLAN + matcher, glob_of_group, _ = build_glob_alternation(list(plan)) + + def style(name): + match = matcher.search(name) + return plan[glob_of_group[match.lastgroup]] if match else None + + sharded = {name for name, _ in model.named_parameters() if style(name) is not None} + + assert sharded == { + "layers.0.block_sparse_moe.experts.gate_up_proj", + "layers.0.block_sparse_moe.experts.down_proj", + } + assert {style(name) for name in sharded} == {"grouped_gemm"} + assert {name.removeprefix("layers.0.") for name in sharded} == { + target for conv in converters for target in conv.target_patterns + } + assert style("layers.0.block_sparse_moe.experts") == "moe_tp_experts" + + +# ------------------------------------------------------------------------------------- prepare + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the_decompressor( + monkeypatch, +): + """The plan goes on the text config (its globs are written against the language model), the + experts forward is ``torch_a2a``, the kimi_linear fusion is registered, and the decompressor + sees the compressed-tensors config before it is dropped (so transformers builds no competing + quantizer).""" + seen = [] + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: seen.append((cfg.quantization_config, cfg.text_config.quantization_config)), + ) + quant = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + text = SimpleNamespace(quantization_config=quant.copy()) + config = SimpleNamespace(text_config=text, quantization_config=quant.copy()) + + assert kimi_k3.prepare(config) is config + + assert text.base_model_ep_plan == kimi_k3.BASE_MODEL_EP_PLAN + assert text.base_model_ep_plan is not kimi_k3.BASE_MODEL_EP_PLAN + assert not hasattr(config, "base_model_ep_plan") + assert text._experts_implementation == "torch_a2a" + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") + assert any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) + assert seen == [(quant, quant)] + assert config.quantization_config is None + assert text.quantization_config is None + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_keeps_a_plan_the_config_already_has(monkeypatch): + """A plan the checkpoint's config already carries describes its own layout and wins.""" + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + text = SimpleNamespace(base_model_ep_plan={"layers.*.custom": "grouped_gemm"}) + + kimi_k3.prepare(SimpleNamespace(text_config=text)) + + assert text.base_model_ep_plan == {"layers.*.custom": "grouped_gemm"} + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_with_a_checkpoint_patches_the_loaded_remote_module_once(monkeypatch): + """Given a checkpoint, prepare imports the remote modeling through the class the config's + ``auto_map`` names, then patches the module that *defines* the MoE block (found by class, not + by file name, and not a VLM wrapper that merely imports it) -- once per process -- and leaves + the remote-code staging shim installed.""" + remote = _remote_module() + wrapper = types.ModuleType("kimi_k3_remote_stub_vlm") # imports the block, defines nothing + wrapper.KimiSparseMoeBlock = remote.KimiSparseMoeBlock + loads = [] + + def import_remote_class(class_ref, ckpt_path, **kwargs): + loads.append((class_ref, ckpt_path, kwargs)) + for module in (wrapper, remote): # what importing the remote code does + monkeypatch.setitem(sys.modules, module.__name__, module) + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + staging = dmu.get_cached_module_file + monkeypatch.setattr(dmu, "get_cached_module_file", staging) + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.setattr(kimi_k3, "_PATCHED", False) + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + auto_map = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} + + for _ in range(2): + config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) + kimi_k3.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + assert loads == [(auto_map["AutoModelForCausalLM"], "/ckpt", {"trust_remote_code": True})] + assert dmu.get_cached_module_file is not staging + assert hasattr(hf_generic, "OutputRecorder") # the remote code imports it from there + # every incompatibility is patched on the module the checkpoint's code defined + assert not isinstance(remote.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + assert remote.KimiDeltaAttention(num_heads=3, head_dim=4).A_log.shape == (4,) + remote.chunk_kda(1, 2, 3, 4, beta=torch.zeros(1), use_beta_sigmoid_in_kernel=True) + _, beta, flag = remote.calls.pop() + assert (beta.item(), flag) == (0.5, False) + remote.create_causal_mask( + config=None, input_embeds=None, attention_mask=None, past_key_values=None, cache_position=0 + ) + assert remote.KimiDynamicCache(seen_tokens=1).get_mask_sizes(2, 0) == (3, 0) + assert remote.KimiK3ForConditionalGeneration().tie_weights(recompute_mapping=False) == "tied" + + +def test_prepare_fails_loudly_when_no_loaded_module_defines_the_moe_block(monkeypatch): + """Skipping the patch silently would load the per-expert layout that EP cannot shard.""" + monkeypatch.setattr(kimi_k3, "_PATCHED", False) + monkeypatch.setattr(kimi_k3, "_force_load_remote_modeling", lambda *args: None) + + with pytest.raises(RuntimeError, match="KimiSparseMoeBlock"): + kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace()), ckpt_path="/ckpt") + + +def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares(monkeypatch): + """``AutoModelForCausalLM`` is preferred, then ``AutoModel``; a config declaring neither is an + error rather than a load that silently skips the patches.""" + loads = [] + monkeypatch.setattr( + dmu, "get_class_from_dynamic_module", lambda ref, path, **kwargs: loads.append(ref) + ) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + + for auto_map in ( + {"AutoModel": "remote.Base", "AutoModelForCausalLM": "remote.ForCausalLM"}, + {"AutoModel": "remote.Base"}, + ): + kimi_k3._force_load_remote_modeling(SimpleNamespace(auto_map=auto_map), "/ckpt", False) + + assert loads == ["remote.ForCausalLM", "remote.Base"] + config_only = SimpleNamespace(auto_map={"AutoConfig": "remote.Config"}) + with pytest.raises(RuntimeError, match="auto_map"): + kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) + + +@pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") +def test_prepare_leaves_the_loader_alone_for_bf16_experts(): + """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves + transformers' conversion mapping unwrapped. (Needs ``compressed_tensors`` importable: the + decompressor imports it before it looks for the config.)""" + pytest.importorskip("compressed_tensors") + stock = mu.get_model_conversion_mapping + + kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace())) + + assert mu.get_model_conversion_mapping is stock + + +def test_the_kimi_k3_spec_routes_expert_parallelism_through_this_adapter(): + """kimi_k3 is remote code, so its whole EP setup is this adapter's ``prepare``.""" + spec = get_spec("kimi_k3") + assert spec.modeling_source == "remote_code" + assert spec.ep_spec.prepare is kimi_k3.prepare From aa5987ad95edb4e6f9bdc508a668676888c2ef9a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 09:50:34 +0000 Subject: [PATCH 02/18] Size Kimi-K3's fused experts from the MoE block's latent decision The experts read use_latent_moe off the config, but Kimi-K3's config has no such field: the model derives it from routed_expert_hidden_size. So every stock K3 config built the experts at the hidden width (7168) instead of the latent one (3584), and the expert-parallel load failed on the size mismatch. The block's own __init__ already decides latent MoE, and the patched forward follows that decision, so the fused experts now take their width from the block. The test stand-in derives the flag the same way instead of carrying a config field the real checkpoint does not have, which is how the tests missed this. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 17 ++++++++--------- tests/unit/torch/models/test_kimi_k3_ep.py | 17 ++++++++++------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 4ba29efcfdd..a2fb74879b0 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -172,15 +172,11 @@ def _build_fused_experts_class(kmod): _get_situ = kmod._get_situ_activation_params class KimiK3Experts(nn.Module): - def __init__(self, config): + def __init__(self, config, hidden_dim): super().__init__() self.num_experts = config.num_experts - # experts run in the LATENT space when use_latent_moe (routed_expert_hidden_size), else hidden. - self.hidden_dim = ( - config.routed_expert_hidden_size - if getattr(config, "use_latent_moe", False) - else config.hidden_size - ) + # The width the MoE block runs its experts at: the latent one under latent MoE, else hidden. + self.hidden_dim = hidden_dim self.intermediate_dim = config.moe_intermediate_size self.gate_up_proj = nn.Parameter( torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim) @@ -251,9 +247,12 @@ def _patch_kimi_moe_block(kmod) -> None: def _init(self, config): _orig_init(self, config) - # Replace the ModuleList of per-expert MLPs with the fused grouped-param module (EP-capable). + # Replace the ModuleList of per-expert MLPs with the fused grouped-param module (EP-capable), as + # wide as the space the block runs its experts in. The block decides latent MoE itself (Kimi-K3's + # config has no use_latent_moe field) and the forward below follows it, so read it off the block. del self.experts - self.experts = FusedExperts(config) + width = config.routed_expert_hidden_size if self.use_latent_moe else config.hidden_size + self.experts = FusedExperts(config, width) def _forward(self, hidden_states): identity = hidden_states diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 56b300aa7b2..b4e046fa2ae 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -99,7 +99,6 @@ def _config(**overrides): "hidden_act": "silu", "situ_beta": 1.5, "situ_linear_beta": 0.25, - "use_latent_moe": True, "latent_moe_use_norm": True, "num_shared_experts": 1, "_experts_implementation": "eager", @@ -165,15 +164,16 @@ class KimiSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() self.config = config - self.use_latent_moe = config.use_latent_moe + # 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 config.use_latent_moe else hidden + 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 config.use_latent_moe: + 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: @@ -285,7 +285,9 @@ def test_fused_experts_match_a_per_expert_reference(hidden_act): gate read from rows [0:I] and up from rows [I:2I] of ``gate_up_proj``. A ``situ`` model hands the full [gate; up] row to the remote ``SituAndMul``. Experts live in the latent width.""" config = _config(hidden_act=hidden_act) - experts = kimi_k3._build_fused_experts_class(_remote_module())(config) + experts = kimi_k3._build_fused_experts_class(_remote_module())( + config, config.routed_expert_hidden_size + ) n, inter, latent = 4, config.moe_intermediate_size, config.routed_expert_hidden_size assert experts.gate_up_proj.shape == (n, 2 * inter, latent) assert experts.down_proj.shape == (n, latent, inter) @@ -325,7 +327,9 @@ def probe(module, *args): monkeypatch.setitem(ALL_EXPERTS_FUNCTIONS, "modelopt_test_probe", probe) config = _config(hidden_act="situ") - experts = kimi_k3._build_fused_experts_class(_remote_module())(config) + experts = kimi_k3._build_fused_experts_class(_remote_module())( + config, config.routed_expert_hidden_size + ) _randomize(experts) gen = torch.Generator().manual_seed(1) tokens = torch.randn(len(ROUTES), config.routed_expert_hidden_size, generator=gen) @@ -345,7 +349,6 @@ def probe(module, *args): [ {}, { - "use_latent_moe": False, "latent_moe_use_norm": False, "routed_expert_hidden_size": None, "num_shared_experts": None, From 6f84fc191706a3e7d58ed1b87b699e03825a3abd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 10:16:57 +0000 Subject: [PATCH 03/18] Read KDA's beta-sigmoid flag off the call's keywords fla-core's chunk_kda and fused_recurrent_kda have no use_beta_sigmoid_in_kernel parameter: the flag lands in their **kwargs, which is why they ignore it. Binding a call to the kernel's signature therefore filed the flag under 'kwargs', the wrapper never saw it, and beta reached the recurrence raw. On Kimi-K3 the state overflowed and the first calibration forward produced NaN out of the first KDA layer. The flag is read off the call's keywords again, as before the binding; beta is still found by name or by position. The test kernels now take the flag through **kwargs, as fla-core's do, which is how the tests missed it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 36 +++++++++++----------- tests/unit/torch/models/test_kimi_k3_ep.py | 15 +++++---- 2 files changed, 27 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index a2fb74879b0..52191c9d8d9 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -342,29 +342,29 @@ def _patch_kda_beta_sigmoid(kmod) -> None: if orig is None or getattr(orig, "_kimi_k3_beta_wrapped", False): continue - def _squash(arguments): - if arguments.get("use_beta_sigmoid_in_kernel"): - beta = arguments.get("beta") - if isinstance(beta, torch.Tensor): - arguments["beta"] = torch.sigmoid(beta.float()).to(beta.dtype) - arguments["use_beta_sigmoid_in_kernel"] = False - def _make(orig_fn): - # Bind to the kernel's signature so beta and the flag are found however they are passed. + # Where beta sits when passed by position. The flag only ever arrives by name: fla-core + # takes it through **kwargs -- which is exactly why it ignores it -- so it is read off the + # call's keywords, never off the kernel's signature. try: - signature = inspect.signature(orig_fn) + params = list(inspect.signature(orig_fn).parameters) except (TypeError, ValueError): - signature = None + params = [] + beta_pos = params.index("beta") if "beta" in params else None def _wrapped(*args, **kwargs): - try: - bound = signature.bind(*args, **kwargs) # type: ignore[union-attr] - except (AttributeError, TypeError): # no signature, or a call it rejects - _squash(kwargs) - return orig_fn(*args, **kwargs) - bound.apply_defaults() - _squash(bound.arguments) - return orig_fn(*bound.args, **bound.kwargs) + if kwargs.get("use_beta_sigmoid_in_kernel"): + beta = kwargs.get("beta") + if isinstance(beta, torch.Tensor): + kwargs["beta"] = torch.sigmoid(beta.float()).to(beta.dtype) + kwargs["use_beta_sigmoid_in_kernel"] = False + elif beta_pos is not None and beta_pos < len(args): + beta = args[beta_pos] + if isinstance(beta, torch.Tensor): + squashed = torch.sigmoid(beta.float()).to(beta.dtype) + args = (*args[:beta_pos], squashed, *args[beta_pos + 1 :]) + kwargs["use_beta_sigmoid_in_kernel"] = False + return orig_fn(*args, **kwargs) setattr(_wrapped, "_kimi_k3_beta_wrapped", True) return _wrapped diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index b4e046fa2ae..a2c7960a7eb 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -207,12 +207,15 @@ def __init__(self, num_heads, head_dim): self.head_dim = head_dim self.A_log = nn.Parameter(torch.zeros(num_heads, dtype=torch.bfloat16)) # stale shape - def chunk_kda(q, k, v, g, beta=None, use_beta_sigmoid_in_kernel=False): - remote.calls.append(("chunk_kda", beta, use_beta_sigmoid_in_kernel)) + # fla-core 0.5's signatures: use_beta_sigmoid_in_kernel is no parameter of either kernel; it lands + # in **kwargs, which is why the kernels ignore it. + def chunk_kda(q, k, v, g, beta, scale=None, initial_state=None, **kwargs): + remote.calls.append(("chunk_kda", beta, kwargs.get("use_beta_sigmoid_in_kernel", False))) return "chunk_kda" - def fused_recurrent_kda(q, k, v, g, beta=None, use_beta_sigmoid_in_kernel=False): - remote.calls.append(("fused_recurrent_kda", beta, use_beta_sigmoid_in_kernel)) + def fused_recurrent_kda(q, k, v, g, beta, A_log=None, dt_bias=None, scale=None, **kwargs): # noqa: N803 + flag = kwargs.get("use_beta_sigmoid_in_kernel", False) + remote.calls.append(("fused_recurrent_kda", beta, flag)) return "fused_recurrent_kda" def create_causal_mask( @@ -425,12 +428,12 @@ def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(): def test_kda_beta_passed_positionally_is_squashed_too(): - """The fix must not depend on how the remote code spells the call.""" + """The fix must not depend on whether the remote code passes beta by position or by name.""" remote = _remote_module() kimi_k3._patch_kda_beta_sigmoid(remote) beta = torch.tensor([-8.0, 8.4], dtype=torch.bfloat16) - remote.chunk_kda(1, 2, 3, 4, beta, True) + remote.chunk_kda(1, 2, 3, 4, beta, use_beta_sigmoid_in_kernel=True) ((_, seen, flag),) = remote.calls assert flag is False From 7f039d81feb0396d0bd9b52bc926c8a79ec2ae9e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:25:57 +0000 Subject: [PATCH 04/18] Skip the tests below the transformers release the code requires The module skipped below 5.6, but the code under test requires 5.8 and raises below it, so a 5.6 or 5.7 run failed instead of skipping. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k3_ep.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index a2c7960a7eb..d7e8e82aa62 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -36,8 +36,8 @@ from packaging.version import Version -if Version(transformers.__version__) < Version("5.6"): - pytest.skip("Kimi-K3 expert parallelism needs transformers >= 5.6", allow_module_level=True) +if Version(transformers.__version__) < Version("5.8"): + pytest.skip("Kimi-K3 expert parallelism needs transformers >= 5.8", allow_module_level=True) import transformers.conversion_mapping as cm import transformers.core_model_loading as cml From 86deefe3e4f90826bf3a433d7782ad44c2d82e94 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 15:35:14 +0000 Subject: [PATCH 05/18] Leave Kimi-K3's latent-expert projections and residual gates unquantized by default vLLM's Kimi-K3 model builds routed_expert_down_proj, routed_expert_up_proj and the self_attention, mlp and output_attn residual gates with quant_config=None, so it refuses a checkpoint that quantizes them. The default formats quantized all 371 of them in a Kimi-K3 export. The [1, hidden] gates gain nothing from quantization either. Signed-off-by: Shengliang Xu --- .../units/default_disabled_quantizers.yaml | 13 +++++++ tests/unit/recipe/test_presets.py | 39 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 812b559ec78..041c2b8afab 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -50,6 +50,19 @@ 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 diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 16f5d31a2d6..a858f07b77f 100644 --- a/tests/unit/recipe/test_presets.py +++ b/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,43 @@ 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. + model = nn.Module() + model.layers = nn.ModuleList([_KimiK3Layer()]) + model.output_attn_res_proj = nn.Linear(32, 1, bias=False) + 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, + ): + 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 From 1b3f2c72de199c6bf150051cc1b94a9b9c97c276 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 16:02:33 +0000 Subject: [PATCH 06/18] Keep Kimi's mm_projector in BF16 by default, as the vision block intends The default unit means to keep every multimodal embedding projection in BF16, but it names only gemma's embed_vision and Llama-4's multi_modal_projector. Kimi-K2.5 and Kimi-K3 call theirs mm_projector, so every --recipe run, which keeps the whole VLM, quantized it. The published Kimi-K2.5 and Kimi-K3 NVFP4 checkpoints both keep it BF16. Signed-off-by: Shengliang Xu --- .../ptq/units/default_disabled_quantizers.yaml | 4 +++- tests/unit/recipe/test_presets.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 041c2b8afab..93412495ae4 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -68,7 +68,7 @@ # `*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, @@ -84,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 diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index a858f07b77f..333ad1f9609 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -155,6 +155,23 @@ def test_default_presets_leave_the_kimi_k3_projections_vllm_builds_unquantized(q assert layer.self_attn.o_proj.weight_quantizer.is_enabled +@pytest.mark.parametrize("qformat", ["fp8", "nvfp4"]) +def test_default_presets_keep_kimis_mm_projector_in_bf16(qformat): + # Kimi-K2.5 and Kimi-K3 name their multimodal projector ``mm_projector``; like the other + # projector spellings it stays BF16 unless a recipe re-enables it. + model = nn.Module() + model.mm_projector = nn.Sequential(nn.Linear(32, 32), nn.GELU(), nn.Linear(32, 32)) + model.language_model = nn.Linear(32, 32) + config = copy.deepcopy(presets.QUANT_CFG_CHOICES[qformat]) + config["algorithm"] = None # placement only, no calibration forward + mtq.quantize(model, config) + + for proj in (model.mm_projector[0], model.mm_projector[2]): + assert not proj.weight_quantizer.is_enabled + assert not proj.input_quantizer.is_enabled + assert model.language_model.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 From fd3b33978f988f5a9674d8dcee7f5a68e2c9d726 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:50:10 +0000 Subject: [PATCH 07/18] Patch each Kimi-K3 checkpoint's own remote module prepare patched the remote modeling only while a process-wide flag was unset, and found the module to patch as the first one in sys.modules that defines the MoE block, so a second checkpoint revision in the same process kept its per-expert layout, or had another revision's module patched in its place. Look the modules up in the package of the class the checkpoint's auto_map loads, and patch on every prepare: each patch is idempotent. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 45 ++++++------ tests/unit/torch/models/test_kimi_k3_ep.py | 81 +++++++++++++++++++--- 2 files changed, 94 insertions(+), 32 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 52191c9d8d9..93d669e4b8f 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -86,8 +86,6 @@ "layers.*.block_sparse_moe.experts": "moe_tp_experts", } -_PATCHED = False - def _install_output_recorder_shim() -> None: """Restore the generic-utils symbol the remote modeling imports. @@ -120,8 +118,12 @@ def __init__(self, *a, **k): _g.OutputRecorder = found -def _find_module_defining(class_name: str): - """Return the loaded module that DEFINES a class named ``class_name``, or None. +def _find_module_defining(class_name: str, package: str): + """Return the loaded module of ``package`` that DEFINES a class named ``class_name``, or None. + + ``package`` is where one checkpoint's remote code was imported (``transformers_modules....``); + another revision loaded in the same process has a package of its own, whose classes are not this + checkpoint's to patch. Identify the remote module by the CLASS it defines, NOT by its file name -- the plugin depends only on the class names it patches, never on the remote code's file layout, so this survives remote-file @@ -134,8 +136,8 @@ def _find_module_defining(class_name: str): ``__getattr__`` returns a namespace proxy for ANY name and raises on use (the proxy is not a ``type``, so it is rejected regardless). """ - for _name, mod in list(sys.modules.items()): - if mod is None: + for name, mod in list(sys.modules.items()): + if mod is None or name.rpartition(".")[0] != package: continue try: obj = getattr(mod, class_name, None) @@ -148,13 +150,13 @@ def _find_module_defining(class_name: str): return None -def _find_remote_module(): +def _find_remote_module(package: str): """Return the loaded remote backbone module. - the one that DEFINES ``KimiSparseMoeBlock`` (the Kimi MoE block this plugin patches) -- or None. Name-independent; - see ``_find_module_defining``. + the one in ``package`` that DEFINES ``KimiSparseMoeBlock`` (the Kimi MoE block this plugin patches) -- or + None. Name-independent; see ``_find_module_defining``. """ - return _find_module_defining("KimiSparseMoeBlock") + return _find_module_defining("KimiSparseMoeBlock", package) def _build_fused_experts_class(kmod): @@ -443,7 +445,7 @@ def _get_mask_sizes(self, query_length, layer_idx): ) -def _shim_tie_weights_kwargs() -> None: +def _shim_tie_weights_kwargs(package: str) -> None: """Absorb the ``recompute_mapping`` kwarg newer transformers passes to ``tie_weights``. Transformers 5.13's ``from_pretrained`` calls ``model.tie_weights(recompute_mapping=...)``, but @@ -465,7 +467,7 @@ def _tie(self, *args, **kwargs): setattr(_tie, "_kimi_k3_kwargs_shim", True) return _tie - mod = _find_module_defining("KimiK3ForConditionalGeneration") + mod = _find_module_defining("KimiK3ForConditionalGeneration", package) if mod is None: return cls = getattr(mod, "KimiK3ForConditionalGeneration", None) @@ -476,13 +478,13 @@ def _tie(self, *args, **kwargs): cls.tie_weights = _make_tie(tw) -def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> None: +def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> type: """Import the remote modeling now, so its classes can be patched before the model is built. Force the remote modeling into ``sys.modules`` (so we can monkey-patch its classes BEFORE ``from_pretrained`` instantiates the model) by resolving the model class the checkpoint's config declares in its ``auto_map`` -- no hardcoded remote file/module name. Runs the OutputRecorder + closure - shims first. + shims first. Returns that class, whose package holds this checkpoint's remote code. """ _install_output_recorder_shim() stage_full_remote_closure() @@ -495,7 +497,7 @@ def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bo "kimi_k3: config.auto_map has no AutoModel/AutoModelForCausalLM entry to force-load the remote " "modeling" ) - get_class_from_dynamic_module(class_ref, ckpt_path, trust_remote_code=trust_remote_code) + return get_class_from_dynamic_module(class_ref, ckpt_path, trust_remote_code=trust_remote_code) def _register_kimi_k3_conversions() -> None: @@ -576,7 +578,6 @@ def prepare( inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM (the memory bottleneck), but PTQ needs the bf16 weights. """ - global _PATCHED from modelopt.torch.utils.moe import require_transformers require_transformers("5.8", "kimi_k3 expert-parallel PTQ") @@ -586,9 +587,12 @@ def prepare( # Monkey-patch the remote modeling (needs ckpt_path to force-load it). If ckpt_path is missing we can # only set the config plan; the load would then fall back to the non-EP ModuleList layout -> callers # must thread ckpt_path through apply_distributed_support. - if ckpt_path is not None and not _PATCHED: - _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code) - kmod = _find_remote_module() + if ckpt_path is not None: + # This checkpoint's remote code, not whichever revision this process imported first. Every + # patch is idempotent, so a repeat prepare for the same checkpoint changes nothing. + package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ + package = package.rpartition(".")[0] + kmod = _find_remote_module(package) if kmod is None: raise RuntimeError( "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" @@ -598,8 +602,7 @@ def prepare( _patch_kda_beta_sigmoid(kmod) _shim_create_causal_mask(kmod) _patch_kimi_cache_get_mask_sizes(kmod) - _shim_tie_weights_kwargs() - _PATCHED = True + _shim_tie_weights_kwargs(package) if not getattr(text_cfg, "base_model_ep_plan", None): # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index d7e8e82aa62..6e04c85ed2a 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -493,12 +493,12 @@ def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): def test_tie_weights_shim_drops_the_new_kwargs_once(monkeypatch): """transformers >= 5.13 calls ``tie_weights(recompute_mapping=...)``; the remote VLM's override takes no kwargs. The shim drops them, delegates once per call, and wraps only once.""" - remote = _remote_module() + remote = _remote_module("transformers_modules.stub.modeling_kimi_k3") monkeypatch.setitem(sys.modules, remote.__name__, remote) vlm = remote.KimiK3ForConditionalGeneration - kimi_k3._shim_tie_weights_kwargs() + kimi_k3._shim_tie_weights_kwargs("transformers_modules.stub") shim = vlm.__dict__["tie_weights"] - kimi_k3._shim_tie_weights_kwargs() + kimi_k3._shim_tie_weights_kwargs("transformers_modules.stub") assert vlm.__dict__["tie_weights"] is shim model = vlm() @@ -515,12 +515,12 @@ class PreTrainedStub: def tie_weights(self, **kwargs): return kwargs - remote = types.ModuleType("kimi_k3_remote_inheriting_stub") + remote = types.ModuleType("transformers_modules.inheriting_stub.modeling_kimi_k3") vlm = type("KimiK3ForConditionalGeneration", (PreTrainedStub,), {"__module__": remote.__name__}) remote.KimiK3ForConditionalGeneration = vlm monkeypatch.setitem(sys.modules, remote.__name__, remote) - kimi_k3._shim_tie_weights_kwargs() + kimi_k3._shim_tie_weights_kwargs("transformers_modules.inheriting_stub") assert vlm().tie_weights(recompute_mapping=False) == {"recompute_mapping": False} @@ -635,8 +635,8 @@ def test_prepare_keeps_a_plan_the_config_already_has(monkeypatch): def test_prepare_with_a_checkpoint_patches_the_loaded_remote_module_once(monkeypatch): """Given a checkpoint, prepare imports the remote modeling through the class the config's ``auto_map`` names, then patches the module that *defines* the MoE block (found by class, not - by file name, and not a VLM wrapper that merely imports it) -- once per process -- and leaves - the remote-code staging shim installed.""" + by file name, and not a VLM wrapper that merely imports it) -- a repeat prepare patches nothing + twice -- and leaves the remote-code staging shim installed.""" remote = _remote_module() wrapper = types.ModuleType("kimi_k3_remote_stub_vlm") # imports the block, defines nothing wrapper.KimiSparseMoeBlock = remote.KimiSparseMoeBlock @@ -646,12 +646,12 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): loads.append((class_ref, ckpt_path, kwargs)) for module in (wrapper, remote): # what importing the remote code does monkeypatch.setitem(sys.modules, module.__name__, module) + return remote.KimiK3ForConditionalGeneration monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) staging = dmu.get_cached_module_file monkeypatch.setattr(dmu, "get_cached_module_file", staging) _restore(monkeypatch, hf_generic, "OutputRecorder") - monkeypatch.setattr(kimi_k3, "_PATCHED", False) monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) auto_map = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} @@ -659,7 +659,7 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) kimi_k3.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) - assert loads == [(auto_map["AutoModelForCausalLM"], "/ckpt", {"trust_remote_code": True})] + assert loads == [(auto_map["AutoModelForCausalLM"], "/ckpt", {"trust_remote_code": True})] * 2 assert dmu.get_cached_module_file is not staging assert hasattr(hf_generic, "OutputRecorder") # the remote code imports it from there # every incompatibility is patched on the module the checkpoint's code defined @@ -677,13 +677,72 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): def test_prepare_fails_loudly_when_no_loaded_module_defines_the_moe_block(monkeypatch): """Skipping the patch silently would load the per-expert layout that EP cannot shard.""" - monkeypatch.setattr(kimi_k3, "_PATCHED", False) - monkeypatch.setattr(kimi_k3, "_force_load_remote_modeling", lambda *args: None) + + class KimiK3ForConditionalGeneration: + __module__ = "transformers_modules.no_moe_block.modeling_kimi_k3" + + monkeypatch.setattr( + kimi_k3, "_force_load_remote_modeling", lambda *args: KimiK3ForConditionalGeneration + ) with pytest.raises(RuntimeError, match="KimiSparseMoeBlock"): kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace()), ckpt_path="/ckpt") +@pytest.mark.usefixtures("private_conversion_registry") +def test_each_checkpoint_revision_patches_its_own_remote_module(monkeypatch): + """Two revisions of the remote code import as two packages, each defining its own MoE block. Each + checkpoint's prepare patches its own revision's module, whichever was imported first, and a + repeat for the first checkpoint replaces nothing that is already patched.""" + revisions = { + f"/ckpt/{rev}": _remote_module(f"transformers_modules.{rev}.modeling_kimi") + for rev in ("rev_a", "rev_b") + } + + def import_remote_class(class_ref, ckpt_path, **kwargs): + remote = revisions[ckpt_path] + monkeypatch.setitem(sys.modules, remote.__name__, remote) + return remote.KimiK3ForConditionalGeneration + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + auto_map = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} + + def prepare(ckpt_path): + config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) + kimi_k3.prepare(config, ckpt_path=ckpt_path) + + def patched_surface(remote): + return ( + remote.KimiSparseMoeBlock.__init__, + remote.KimiSparseMoeBlock.forward, + remote.KimiDeltaAttention.__init__, + remote.chunk_kda, + remote.fused_recurrent_kda, + remote.create_causal_mask, + remote.KimiDynamicCache.get_mask_sizes, + remote.KimiK3ForConditionalGeneration.tie_weights, + ) + + rev_a, rev_b = revisions["/ckpt/rev_a"], revisions["/ckpt/rev_b"] + shipped_b = patched_surface(rev_b) + prepare("/ckpt/rev_a") + assert patched_surface(rev_b) == shipped_b + prepare("/ckpt/rev_b") + patched = patched_surface(rev_a), patched_surface(rev_b) + prepare("/ckpt/rev_a") + + assert (patched_surface(rev_a), patched_surface(rev_b)) == patched + assert all(new is not old for new, old in zip(patched_surface(rev_b), shipped_b)) + for remote in (rev_a, rev_b): + assert not isinstance(remote.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + assert ( + remote.KimiK3ForConditionalGeneration().tie_weights(recompute_mapping=False) == "tied" + ) + + def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares(monkeypatch): """``AutoModelForCausalLM`` is preferred, then ``AutoModel``; a config declaring neither is an error rather than a load that silently skips the patches.""" From 3b815d26706c14e813e77a854594fa769b0a01a7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:13:29 +0000 Subject: [PATCH 08/18] State the expert plan Kimi-K3's prepare attaches Kimi-K3's prepare patches the checkpoint's remote MoE block and attaches the plan for its fused experts. An EP degree of 1 skips prepare, so its spec supplied no plan there, and a checkpoint, whose config carries none, loaded whole on every rank with a replicated dense part instead of being refused. Its spec now states the plan prepare attaches, which finds the experts the checkpoint's own code builds, at the same paths. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/specs.py | 4 ++-- tests/unit/torch/models/test_kimi_k3_ep.py | 9 +++++---- tests/unit/torch/utils/test_resolve_distribution_plan.py | 5 ++++- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 6a9bfd90542..264dfdf74e0 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -20,7 +20,7 @@ """ from ..specs import EPSpec, ModelSpec, register -from .ep import prepare +from .ep import BASE_MODEL_EP_PLAN, prepare __all__: list[str] = [] @@ -28,6 +28,6 @@ ModelSpec( model_type="kimi_k3", modeling_source="remote_code", - ep_spec=EPSpec(prepare=prepare), + ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), ) ) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 6e04c85ed2a..49e0b1ab195 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -593,10 +593,10 @@ def style(name): def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the_decompressor( monkeypatch, ): - """The plan goes on the text config (its globs are written against the language model), the - experts forward is ``torch_a2a``, the kimi_linear fusion is registered, and the decompressor - sees the compressed-tensors config before it is dropped (so transformers builds no competing - quantizer).""" + """The plan goes on the text config (its globs are written against the language model) and is + the one the spec states for a run that skips prepare, the experts forward is ``torch_a2a``, the + kimi_linear fusion is registered, and the decompressor sees the compressed-tensors config before + it is dropped (so transformers builds no competing quantizer).""" seen = [] monkeypatch.setattr( _onfly_decompress, @@ -611,6 +611,7 @@ def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the assert text.base_model_ep_plan == kimi_k3.BASE_MODEL_EP_PLAN assert text.base_model_ep_plan is not kimi_k3.BASE_MODEL_EP_PLAN + assert text.base_model_ep_plan == get_spec("kimi_k3").ep_spec.prepared_ep_plan assert not hasattr(config, "base_model_ep_plan") assert text._experts_implementation == "torch_a2a" registered = cm.get_checkpoint_conversion_mapping("kimi_linear") diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 7b35d4af56b..da82d8ad679 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -616,7 +616,10 @@ def layer(): return model -@pytest.mark.parametrize(("model_type", "moe"), [("kimi_vl", "mlp"), ("kimi_k25", "mlp")]) +@pytest.mark.parametrize( + ("model_type", "moe"), + [("kimi_vl", "mlp"), ("kimi_k25", "mlp"), ("kimi_k3", "block_sparse_moe")], +) def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( run_as, serve, monkeypatch, tmp_path, model_type, moe ): From 694423bd9cc7667fe50859ff3e770f264952e839 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:07:07 +0000 Subject: [PATCH 09/18] Run the Kimi-K3 bf16-experts test without compressed-tensors Its skip predates the decompressor looking up the quantization config before it imports compressed_tensors, so a bf16 checkpoint no longer needs the package and the test passes without it. The skip kept it from running anywhere the package is missing, a plain dev-test install included. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k3_ep.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 49e0b1ab195..0aac1c1d744 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -769,9 +769,7 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") def test_prepare_leaves_the_loader_alone_for_bf16_experts(): """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves - transformers' conversion mapping unwrapped. (Needs ``compressed_tensors`` importable: the - decompressor imports it before it looks for the config.)""" - pytest.importorskip("compressed_tensors") + transformers' conversion mapping unwrapped, with or without ``compressed_tensors`` installed.""" stock = mu.get_model_conversion_mapping kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace())) From d6ac84795be78301561489632f05c42b06d42fba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:01:12 +0000 Subject: [PATCH 10/18] Trim the Kimi-K3 adapter tests: fold duplicates into covering tests - The per-expert reference test becomes the situ case of the patched-block test, which compares with the remote block itself. - The patch-twice test goes: the revisions test fails on a non-idempotent block patch. - The two beta tests become one, parametrized by how beta is passed, for both KDA entry points. - The A_log test goes: the prepare-with-checkpoint test checks the checkpoint's shape and now the dtype it keeps. - The tie_weights-shim-once test goes: the prepare test checks that the shim absorbs the kwargs and delegates exactly once, and the revisions test that it wraps once. - The mm_projector preset test joins the Kimi projections preset test. - The spec-routing test keeps only the prepare identity, which nothing else checks; test_specs_vs_transformers checks modeling_source. Signed-off-by: Shengliang Xu --- tests/unit/recipe/test_presets.py | 22 +--- tests/unit/torch/models/test_kimi_k3_ep.py | 132 +++++---------------- 2 files changed, 34 insertions(+), 120 deletions(-) diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 333ad1f9609..bab0c0a2477 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -135,9 +135,12 @@ def __init__(self, hidden=32, latent=16): @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) @@ -149,29 +152,14 @@ def test_default_presets_leave_the_kimi_k3_projections_vllm_builds_unquantized(q 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 -@pytest.mark.parametrize("qformat", ["fp8", "nvfp4"]) -def test_default_presets_keep_kimis_mm_projector_in_bf16(qformat): - # Kimi-K2.5 and Kimi-K3 name their multimodal projector ``mm_projector``; like the other - # projector spellings it stays BF16 unless a recipe re-enables it. - model = nn.Module() - model.mm_projector = nn.Sequential(nn.Linear(32, 32), nn.GELU(), nn.Linear(32, 32)) - model.language_model = nn.Linear(32, 32) - config = copy.deepcopy(presets.QUANT_CFG_CHOICES[qformat]) - config["algorithm"] = None # placement only, no calibration forward - mtq.quantize(model, config) - - for proj in (model.mm_projector[0], model.mm_projector[2]): - assert not proj.weight_quantizer.is_enabled - assert not proj.input_quantizer.is_enabled - assert model.language_model.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 diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 0aac1c1d744..91b847fc7cb 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -282,41 +282,6 @@ def _load(converters, checkpoint): # --------------------------------------------------------------------------- the fused experts -@pytest.mark.parametrize("hidden_act", ["silu", "situ"]) -def test_fused_experts_match_a_per_expert_reference(hidden_act): - """Per token: the routing-weighted sum over its top-k experts of down(act(gate x) * up x), with - gate read from rows [0:I] and up from rows [I:2I] of ``gate_up_proj``. A ``situ`` model hands - the full [gate; up] row to the remote ``SituAndMul``. Experts live in the latent width.""" - config = _config(hidden_act=hidden_act) - experts = kimi_k3._build_fused_experts_class(_remote_module())( - config, config.routed_expert_hidden_size - ) - n, inter, latent = 4, config.moe_intermediate_size, config.routed_expert_hidden_size - assert experts.gate_up_proj.shape == (n, 2 * inter, latent) - assert experts.down_proj.shape == (n, latent, inter) - _randomize(experts) - gen = torch.Generator().manual_seed(1) - tokens = torch.randn(len(ROUTES), latent, generator=gen) - route_weights = torch.rand(ROUTES.shape, generator=gen) - - def act(gate_up): - if hidden_act == "situ": - return _situ(gate_up, config.situ_beta, config.situ_linear_beta) - return F.silu(gate_up[:inter]) * gate_up[inter:] - - with torch.no_grad(): - expected = torch.stack( - [ - sum( - w * (experts.down_proj[e] @ act(experts.gate_up_proj[e] @ x)) - for e, w in zip(top_e, top_w) - ) - for x, top_e, top_w in zip(tokens, ROUTES, route_weights) - ] - ) - torch.testing.assert_close(experts(tokens, ROUTES, route_weights), expected) - - def test_fused_experts_dispatch_through_the_experts_interface(monkeypatch): """``config._experts_implementation`` picks the forward (``torch_a2a`` under EP), so the module must route through transformers' ``ExpertsInterface``; and a stock implementation reading the @@ -356,13 +321,15 @@ def probe(module, *args): "routed_expert_hidden_size": None, "num_shared_experts": None, }, + {"hidden_act": "situ"}, ], - ids=["latent-norm-shared", "plain"], + ids=["latent-norm-shared", "plain", "situ"], ) def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): """Swapping the per-expert ModuleList for the fused module keeps the block's output: the latent projections, the latent norm and the shared experts stay around the fused call, and w1 / w3 / - w2 are the gate rows, the up rows and the down of the fused params.""" + w2 are the gate rows, the up rows and the down of the fused params. A ``situ`` model hands the + full [gate; up] row to the remote ``SituAndMul``.""" remote = _remote_module() config = _config(**overrides) reference = remote.KimiSparseMoeBlock(config) @@ -386,24 +353,15 @@ def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): torch.testing.assert_close(block(hidden_states), expected) -def test_patching_the_block_twice_wraps_it_once(): - """A second ``prepare`` must not stack another fused-experts allocation onto ``__init__``.""" - remote = _remote_module() - kimi_k3._patch_kimi_moe_block(remote) - patched = (remote.KimiSparseMoeBlock.__init__, remote.KimiSparseMoeBlock.forward) - - kimi_k3._patch_kimi_moe_block(remote) - - assert (remote.KimiSparseMoeBlock.__init__, remote.KimiSparseMoeBlock.forward) == patched - - # ------------------------------------------------------------------- the linear-attention fixes -def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(): +@pytest.mark.parametrize("by", ["keyword", "position"]) +def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(by): """fla-core ignores ``use_beta_sigmoid_in_kernel``, so both KDA entry points apply - ``sigmoid(beta)`` themselves (in beta's dtype) and clear the flag; calls without the flag pass - beta through untouched. Each entry point is wrapped once.""" + ``sigmoid(beta)`` themselves (in beta's dtype) and clear the flag, whether the remote code + passes beta by name or by position; calls without the flag pass beta through untouched. Each + entry point is wrapped once.""" remote = _remote_module() kimi_k3._patch_kda_beta_sigmoid(remote) wrapped = (remote.chunk_kda, remote.fused_recurrent_kda) @@ -411,8 +369,13 @@ def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(): assert (remote.chunk_kda, remote.fused_recurrent_kda) == wrapped beta = torch.tensor([-8.0, 0.0, 8.4], dtype=torch.bfloat16) - assert remote.chunk_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=True) == "chunk_kda" - remote.fused_recurrent_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=True) + def call(kda, **flag): + return ( + kda(1, 2, 3, 4, beta=beta, **flag) if by == "keyword" else kda(1, 2, 3, 4, beta, **flag) + ) + + assert call(remote.chunk_kda, use_beta_sigmoid_in_kernel=True) == "chunk_kda" + call(remote.fused_recurrent_kda, use_beta_sigmoid_in_kernel=True) squashed = torch.sigmoid(beta.float()).to(torch.bfloat16) assert [name for name, *_ in remote.calls] == ["chunk_kda", "fused_recurrent_kda"] @@ -422,36 +385,11 @@ def test_kda_beta_is_squashed_in_torch_only_when_the_kernel_flag_asks_for_it(): assert torch.equal(seen, squashed) remote.calls.clear() - remote.chunk_kda(1, 2, 3, 4, beta=beta) - remote.fused_recurrent_kda(1, 2, 3, 4, beta=beta, use_beta_sigmoid_in_kernel=False) + call(remote.chunk_kda) + call(remote.fused_recurrent_kda, use_beta_sigmoid_in_kernel=False) assert [(seen is beta, flag) for _, seen, flag in remote.calls] == [(True, False)] * 2 -def test_kda_beta_passed_positionally_is_squashed_too(): - """The fix must not depend on whether the remote code passes beta by position or by name.""" - remote = _remote_module() - kimi_k3._patch_kda_beta_sigmoid(remote) - beta = torch.tensor([-8.0, 8.4], dtype=torch.bfloat16) - - remote.chunk_kda(1, 2, 3, 4, beta, use_beta_sigmoid_in_kernel=True) - - ((_, seen, flag),) = remote.calls - assert flag is False - assert torch.equal(seen, torch.sigmoid(beta.float()).to(torch.bfloat16)) - - -def test_kda_a_log_takes_the_checkpoints_head_dim_shape(): - """The stale port allocates ``A_log`` as [num_heads]; the checkpoint stores [head_dim], which - would fail to load (or be re-initialized). The reallocation keeps the dtype.""" - remote = _remote_module() - kimi_k3._patch_kda_a_log(remote) - - attention = remote.KimiDeltaAttention(num_heads=3, head_dim=4) - - assert attention.A_log.shape == (4,) - assert attention.A_log.dtype == torch.bfloat16 - - # ------------------------------------------------------------- transformers-version compatibility @@ -490,23 +428,6 @@ def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): assert remote.calls == [("get_seq_length", 2), ("get_seq_length", 1)] -def test_tie_weights_shim_drops_the_new_kwargs_once(monkeypatch): - """transformers >= 5.13 calls ``tie_weights(recompute_mapping=...)``; the remote VLM's override - takes no kwargs. The shim drops them, delegates once per call, and wraps only once.""" - remote = _remote_module("transformers_modules.stub.modeling_kimi_k3") - monkeypatch.setitem(sys.modules, remote.__name__, remote) - vlm = remote.KimiK3ForConditionalGeneration - kimi_k3._shim_tie_weights_kwargs("transformers_modules.stub") - shim = vlm.__dict__["tie_weights"] - kimi_k3._shim_tie_weights_kwargs("transformers_modules.stub") - assert vlm.__dict__["tie_weights"] is shim - - model = vlm() - - assert model.tie_weights(missing_keys=set(), recompute_mapping=False) == "tied" - assert remote.calls == [("tie_weights", model)] - - def test_tie_weights_shim_leaves_an_inherited_tie_weights_alone(monkeypatch): """Only an override defined on the VLM class has the old signature; an inherited ``tie_weights`` already takes the kwargs and must keep receiving them.""" @@ -665,7 +586,11 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): assert hasattr(hf_generic, "OutputRecorder") # the remote code imports it from there # every incompatibility is patched on the module the checkpoint's code defined assert not isinstance(remote.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) - assert remote.KimiDeltaAttention(num_heads=3, head_dim=4).A_log.shape == (4,) + a_log = remote.KimiDeltaAttention(num_heads=3, head_dim=4).A_log + assert (a_log.shape, a_log.dtype) == ( + (4,), + torch.bfloat16, + ) # the checkpoint's shape, same dtype remote.chunk_kda(1, 2, 3, 4, beta=torch.zeros(1), use_beta_sigmoid_in_kernel=True) _, beta, flag = remote.calls.pop() assert (beta.item(), flag) == (0.5, False) @@ -673,7 +598,9 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): config=None, input_embeds=None, attention_mask=None, past_key_values=None, cache_position=0 ) assert remote.KimiDynamicCache(seen_tokens=1).get_mask_sizes(2, 0) == (3, 0) - assert remote.KimiK3ForConditionalGeneration().tie_weights(recompute_mapping=False) == "tied" + vlm = remote.KimiK3ForConditionalGeneration() + assert vlm.tie_weights(recompute_mapping=False) == "tied" + assert [call for call in remote.calls if call[0] == "tie_weights"] == [("tie_weights", vlm)] def test_prepare_fails_loudly_when_no_loaded_module_defines_the_moe_block(monkeypatch): @@ -778,7 +705,6 @@ def test_prepare_leaves_the_loader_alone_for_bf16_experts(): def test_the_kimi_k3_spec_routes_expert_parallelism_through_this_adapter(): - """kimi_k3 is remote code, so its whole EP setup is this adapter's ``prepare``.""" - spec = get_spec("kimi_k3") - assert spec.modeling_source == "remote_code" - assert spec.ep_spec.prepare is kimi_k3.prepare + """kimi_k3 is remote code, so its whole EP setup is this adapter's ``prepare``: another Kimi + adapter's would load it without its own patches.""" + assert get_spec("kimi_k3").ep_spec.prepare is kimi_k3.prepare From 788b55e093f85f5764c6e12993d3d1f9c37ad5b3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:21:22 +0000 Subject: [PATCH 11/18] Refuse to import an untrusted checkpoint's code in the Kimi-K3 adapter Like the Kimi-K2.5 and Kimi-VL adapters, Kimi-K3's prepare loads the checkpoint's own modeling through get_class_from_dynamic_module, which runs it without checking trust_remote_code. Refuse unless the caller passed trust_remote_code=True, before anything is staged or patched. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 7 +++++++ tests/unit/torch/models/test_kimi_k3_ep.py | 19 ++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 93d669e4b8f..5934f4f3730 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -581,6 +581,13 @@ def prepare( from modelopt.torch.utils.moe import require_transformers require_transformers("5.8", "kimi_k3 expert-parallel PTQ") + if ckpt_path is not None and not trust_remote_code: + # Patching below imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + "kimi_k3: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) text_cfg = getattr(hf_config, "text_config", None) or hf_config diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 91b847fc7cb..0a78c97565e 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -614,6 +614,23 @@ class KimiK3ForConditionalGeneration: ) with pytest.raises(RuntimeError, match="KimiSparseMoeBlock"): + kimi_k3.prepare( + SimpleNamespace(text_config=SimpleNamespace()), + ckpt_path="/ckpt", + trust_remote_code=True, + ) + + +def test_prepare_refuses_an_untrusted_checkpoint(monkeypatch): + """transformers' get_class_from_dynamic_module runs the checkpoint's code without checking + trust_remote_code, so prepare must refuse before importing it.""" + monkeypatch.setattr( + kimi_k3, + "_force_load_remote_modeling", + lambda *args: pytest.fail("the checkpoint's code was imported"), + ) + + with pytest.raises(ValueError, match="trust_remote_code=True"): kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace()), ckpt_path="/ckpt") @@ -640,7 +657,7 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): def prepare(ckpt_path): config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) - kimi_k3.prepare(config, ckpt_path=ckpt_path) + kimi_k3.prepare(config, ckpt_path=ckpt_path, trust_remote_code=True) def patched_surface(remote): return ( From 254116073f1055e5586717cda9f67279634b115f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:14:56 +0000 Subject: [PATCH 12/18] Let a missing conversion API stop the Kimi-K3 expert-parallel load When transformers' weight-conversion API failed to import, _register_kimi_k3_conversions printed a warning and returned. The checkpoint's per-expert weights then loaded as unexpected and the fused experts as missing: newly initialized, at full size on every rank. prepare() already requires transformers 5.8, where the API exists, so the import now simply raises. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 14 ++++---------- tests/unit/torch/models/test_kimi_k3_ep.py | 9 +++++++++ 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 5934f4f3730..c8d5978e8a0 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -517,16 +517,10 @@ def _register_kimi_k3_conversions() -> None: w3=up (Concatenate gate-then-up on dim 1), w2=down. The on-the-fly decompressor then rewrites these ``.weight`` sources to the checkpoint's ``.weight_packed`` + scales (mxfp4 -> bf16). Idempotent. """ - try: - from transformers.conversion_mapping import register_checkpoint_conversion_mapping - from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter - except Exception as e: # pragma: no cover - version guard - print( - f"[kimi_k3] WARN: transformers conversion API unavailable ({type(e).__name__}: {e}); " - "per-expert fusion NOT registered -- experts will load as MISSING", - flush=True, - ) - return + # transformers is optional. Its conversion API is in every version prepare() accepts, and a + # load without it would leave the experts MISSING, so an ImportError here must stop the load. + from transformers.conversion_mapping import register_checkpoint_conversion_mapping + from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter def _mapping(): return [ diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 0a78c97565e..48a7d1b7ade 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -474,6 +474,15 @@ def test_registered_conversions_fuse_w1_w3_into_gate_up_and_stack_w2(): assert torch.equal(fused[prefix + "down_proj"], torch.stack(down)) +def test_the_registration_raises_without_the_conversion_api(monkeypatch): + """Unregistered, the checkpoint's experts would load as MISSING: newly initialized, at full + size on every rank. So a missing conversion API stops the load instead of warning.""" + monkeypatch.setitem(sys.modules, "transformers.conversion_mapping", None) + + with pytest.raises(ImportError): + kimi_k3._register_kimi_k3_conversions() + + @pytest.mark.usefixtures("private_conversion_registry") def test_the_plan_the_conversions_and_the_fused_params_name_the_same_tensors(): """Three separately written names must agree, or EP silently shards nothing: the fused module's From ea8cebe99faac503824f890b03880fc4fbe875ad Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:11:14 +0000 Subject: [PATCH 13/18] Decompress Kimi-K3 only with a quantization_config, and let its errors through prepare suppressed every ValueError from the decompressor, meaning to skip plain bf16 experts. But compressed_tensors raises pydantic's ValidationError, also a ValueError, for a quantization_config it cannot parse; prepare then dropped the config and the experts would load newly initialized. Install the decompressor only when the config carries a quantization_config, and let it raise. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 6 ++---- tests/unit/torch/models/test_kimi_k3_ep.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index c8d5978e8a0..e39d8383e48 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -67,7 +67,6 @@ transformers 5.13 + flash-linear-attention (FLA) 0.5.0, which the remote code's KDA layers need at runtime. """ -import contextlib import sys import torch @@ -618,9 +617,8 @@ def prepare( # the decompressor has converters to rewrite into compressed-aware (weight_packed -> bf16) ones. Drop # quantization_config afterwards so transformers builds no competing HF quantizer. _register_kimi_k3_conversions() - # ValueError == no compressed-tensors quantization_config, i.e. plain bf16 experts and nothing to - # decompress. - with contextlib.suppress(ValueError): + # Plain bf16 experts come without a quantization_config: nothing to decompress. + if any(getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg)): _onfly_decompress.install_onfly_decompress(hf_config) for cfg in (hf_config, text_cfg): if getattr(cfg, "quantization_config", None) is not None: diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 48a7d1b7ade..a29af7055b1 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -719,6 +719,21 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch): + """compressed_tensors raises pydantic's ValidationError, a ValueError, for a quantization_config + it cannot parse. Dropped after that, the config would load the experts newly initialized.""" + + def install(cfg): + raise ValueError("1 validation error for QuantizationConfig") + + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", install) + config = SimpleNamespace(text_config=SimpleNamespace(), quantization_config={"format": "new"}) + + with pytest.raises(ValueError, match="validation error"): + kimi_k3.prepare(config) + + @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") def test_prepare_leaves_the_loader_alone_for_bf16_experts(): """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves From deeffc04c09a22f79d86d0a20963b818e249564c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 21:11:30 +0000 Subject: [PATCH 14/18] Refuse a Kimi-K3 prepare without the checkpoint's path, and declare the release it was validated on prepare refuses, before it changes anything, a call with no checkpoint path: it patches the checkpoint's own MoE block, and without it the plan, the kimi_linear expert fusion and the decompressor still went in, naming fused experts the unpatched per-expert modeling never builds, whose own weights then loaded newly initialized. The trust_remote_code check covers every call. The spec declares transformers 5.13, the release the adapter's patches were developed and validated against, so check_distributed_support refuses an older one with its way around; prepare no longer checks a floor itself. The create_causal_mask shim passes every keyword through to a signature that takes **kwargs. The eager experts fallback keeps a class past the experts for the router's no-local-expert id, as transformers' reference loop does, the block forward reshapes as the remote one does, and the KDA patch's docstring keeps the reason, not the measurements. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 89 ++++++++++----------- modelopt/torch/models/kimi_k3/specs.py | 5 +- tests/unit/torch/models/test_kimi_k3_ep.py | 90 ++++++++++++++++++---- 3 files changed, 120 insertions(+), 64 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index e39d8383e48..227a3690ce8 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -63,8 +63,8 @@ Also shims ``OutputRecorder`` back into ``transformers.utils.generic`` (the remote code imports it there, but it moved to ``transformers.modeling_utils``) so the modeling imports at all. -Version floor: transformers >= 5.8, for the on-the-fly decompression. Developed and validated against -transformers 5.13 + flash-linear-attention (FLA) 0.5.0, which the remote code's KDA layers need at runtime. +Version floor: transformers >= 5.13, which it was developed and validated against (the spec declares it), +with flash-linear-attention (FLA) 0.5.0, which the remote code's KDA layers need at runtime. """ import sys @@ -209,7 +209,8 @@ def forward(self, hidden_states, top_k_index, top_k_weights): # DeepseekV3Experts/MiniMaxM3VLExperts reference loop, with situ gating. final = torch.zeros_like(hidden_states) with torch.no_grad(): - mask = nn.functional.one_hot(top_k_index, num_classes=self.num_experts).permute( + # One class past the experts: the id a router gives a token routed to no local expert. + mask = nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1).permute( 2, 1, 0 ) hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero() @@ -259,7 +260,7 @@ def _forward(self, hidden_states): identity = hidden_states orig_shape = hidden_states.shape topk_idx, topk_weight = self.gate(hidden_states) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) if self.use_latent_moe: hidden_states = self.routed_expert_down_proj(hidden_states) # Fused / a2a experts: returns the top-k-combined output in the latent space. @@ -268,7 +269,7 @@ def _forward(self, hidden_states): if self.latent_moe_use_norm: y = self.routed_expert_norm(y) y = self.routed_expert_up_proj(y) - y = y.view(*orig_shape) + y = y.reshape(orig_shape) if self.config.num_shared_experts is not None: y = y + self.shared_experts(identity) return y @@ -314,27 +315,15 @@ def _init(self, *args, **kwargs): def _patch_kda_beta_sigmoid(kmod) -> None: - """Fix the KDA calibration NaN. - - The remote ``KimiDeltaAttention.forward`` passes the RAW gate scalar ``beta = b_proj(hidden_states)`` to FLA's - ``chunk_kda`` / ``fused_recurrent_kda`` with ``use_beta_sigmoid_in_kernel=True``, expecting the kernel to apply - ``sigmoid`` (the delta-rule learning rate must be in ``(0,1)``). In fla-core 0.5.x that flag is a **silent no-op** - for the main recurrence -- the kernel uses the raw beta. Real trained weights produce ``beta`` up to ~8.4, so the - per-step delta-rule eigenvalue ``|1 - beta*||k||^2|`` exceeds 1 (``beta*||k||^2 ~ 7.7 >> 2``) and the recurrent - state grows exponentially (~1e11 over 128 tokens -> ``inf`` -> ``o_norm`` turns it into NaN -> ``detected nan values - in amax`` during calibration). Random-weight smokes never tripped it because their beta is small (``beta*||k||^2 < - 2``). Verified offline on the real layer-1 inputs: pre-applying ``sigmoid(beta)`` and passing - ``use_beta_sigmoid_in_kernel=False`` bounds the output (o~0.1, state~4.6, matching the pure-torch - ``naive_recurrent_kda`` reference), vs 8e9 unfixed. - - Fix: wrap the module-level ``chunk_kda`` / ``fused_recurrent_kda`` in the remote module; when invoked - with ``use_beta_sigmoid_in_kernel=True``, replace ``beta`` with ``beta.sigmoid()`` and set the flag - False so the transform happens (correctly) in torch. Idempotent. + """Apply the KDA gate's sigmoid in torch, and clear the kernel flag that asks for it. + + fla-core 0.5.x silently ignores ``use_beta_sigmoid_in_kernel``, so its ``chunk_kda`` / + ``fused_recurrent_kda`` use the raw gate, and on trained weights the delta-rule recurrence + (``|1 - beta*||k||^2| > 1``) diverges to inf, a NaN amax in calibration. Random weights keep + ``beta`` small enough to hide it. Idempotent. """ import inspect - import torch - if getattr(kmod, "_kimi_k3_beta_sigmoid_patched", False): return patched_any = False @@ -398,9 +387,15 @@ def _shim_create_causal_mask(kmod) -> None: if orig is None or getattr(orig, "_kimi_k3_ccm_shim", False): return try: - accepted = set(inspect.signature(orig).parameters) + params = inspect.signature(orig).parameters except (TypeError, ValueError): - accepted = None + params = {} + # None: every keyword passes, as for a signature that takes **kwargs. + accepted = ( + set(params) + if params and all(p.kind is not inspect.Parameter.VAR_KEYWORD for p in params.values()) + else None + ) @functools.wraps(orig) def _wrapped(*args, **kwargs): @@ -571,10 +566,12 @@ def prepare( inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM (the memory bottleneck), but PTQ needs the bf16 weights. """ - from modelopt.torch.utils.moe import require_transformers - - require_transformers("5.8", "kimi_k3 expert-parallel PTQ") - if ckpt_path is not None and not trust_remote_code: + if ckpt_path is None: + raise ValueError( + "kimi_k3: 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 below imports the checkpoint's own modeling code, and transformers' # get_class_from_dynamic_module runs it without checking trust_remote_code itself. raise ValueError( @@ -584,25 +581,21 @@ def prepare( text_cfg = getattr(hf_config, "text_config", None) or hf_config - # Monkey-patch the remote modeling (needs ckpt_path to force-load it). If ckpt_path is missing we can - # only set the config plan; the load would then fall back to the non-EP ModuleList layout -> callers - # must thread ckpt_path through apply_distributed_support. - if ckpt_path is not None: - # This checkpoint's remote code, not whichever revision this process imported first. Every - # patch is idempotent, so a repeat prepare for the same checkpoint changes nothing. - package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ - package = package.rpartition(".")[0] - kmod = _find_remote_module(package) - if kmod is None: - raise RuntimeError( - "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" - ) - _patch_kimi_moe_block(kmod) - _patch_kda_a_log(kmod) - _patch_kda_beta_sigmoid(kmod) - _shim_create_causal_mask(kmod) - _patch_kimi_cache_get_mask_sizes(kmod) - _shim_tie_weights_kwargs(package) + # This checkpoint's remote code, not whichever revision this process imported first. Every patch + # is idempotent, so a repeat prepare for the same checkpoint changes nothing. + package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ + package = package.rpartition(".")[0] + kmod = _find_remote_module(package) + if kmod is None: + raise RuntimeError( + "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" + ) + _patch_kimi_moe_block(kmod) + _patch_kda_a_log(kmod) + _patch_kda_beta_sigmoid(kmod) + _shim_create_causal_mask(kmod) + _patch_kimi_cache_get_mask_sizes(kmod) + _shim_tie_weights_kwargs(package) if not getattr(text_cfg, "base_model_ep_plan", None): # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 264dfdf74e0..caf76c9aba2 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -28,6 +28,9 @@ ModelSpec( model_type="kimi_k3", modeling_source="remote_code", - ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), + # The release prepare's patches to the checkpoint's code were validated against. + ep_spec=EPSpec( + prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, requires_transformers="5.13" + ), ) ) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index a29af7055b1..bfba062a740 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -49,6 +49,7 @@ from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.models.kimi_k3 import ep as kimi_k3 +from modelopt.torch.utils.moe import check_distributed_support ROUTES = torch.tensor([[0, 3], [1, 2], [3, 1], [2, 0], [1, 3]]) # top-2 experts of 5 tokens @@ -416,6 +417,17 @@ def test_create_causal_mask_shim_adapts_the_remote_call_once(): assert remote.calls == [("create_causal_mask", "config", "embeds", "mask_2d", "cache", "ids")] +def test_create_causal_mask_shim_keeps_what_a_kwargs_signature_takes(): + """A ``create_causal_mask`` that takes ``**kwargs`` accepts every keyword: none is dropped.""" + remote = types.ModuleType("kimi_k3_kwargs_stub") + remote.create_causal_mask = lambda config, inputs_embeds, **kwargs: kwargs + kimi_k3._shim_create_causal_mask(remote) + + kept = remote.create_causal_mask(config="config", input_embeds="embeds", cache_position=0) + + assert kept == {"cache_position": 0} + + def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): """transformers >= 5.13 passes the query length as an int; the remote cache expects a ``cache_position`` tensor. Either way kv_length is the query length plus the cached tokens.""" @@ -518,8 +530,30 @@ def style(name): # ------------------------------------------------------------------------------------- prepare +_AUTO_MAP = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} -@pytest.mark.usefixtures("private_conversion_registry") + +@pytest.fixture +def remote_checkpoint(monkeypatch): + """The checkpoint's remote code, a fresh stand-in that importing it registers; the shims + prepare installs on transformers are undone after the test.""" + remote = _remote_module("transformers_modules.stub.modeling_kimi_k3") + + def import_remote_class(class_ref, ckpt_path, **kwargs): + monkeypatch.setitem(sys.modules, remote.__name__, remote) + return remote.KimiK3ForConditionalGeneration + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + return remote + + +def _prepare(config): + return kimi_k3.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the_decompressor( monkeypatch, ): @@ -535,9 +569,9 @@ def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the ) quant = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} text = SimpleNamespace(quantization_config=quant.copy()) - config = SimpleNamespace(text_config=text, quantization_config=quant.copy()) + config = SimpleNamespace(text_config=text, quantization_config=quant.copy(), auto_map=_AUTO_MAP) - assert kimi_k3.prepare(config) is config + assert _prepare(config) is config assert text.base_model_ep_plan == kimi_k3.BASE_MODEL_EP_PLAN assert text.base_model_ep_plan is not kimi_k3.BASE_MODEL_EP_PLAN @@ -551,13 +585,13 @@ def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the assert text.quantization_config is None -@pytest.mark.usefixtures("private_conversion_registry") +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") def test_prepare_keeps_a_plan_the_config_already_has(monkeypatch): """A plan the checkpoint's config already carries describes its own layout and wins.""" monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) text = SimpleNamespace(base_model_ep_plan={"layers.*.custom": "grouped_gemm"}) - kimi_k3.prepare(SimpleNamespace(text_config=text)) + _prepare(SimpleNamespace(text_config=text, auto_map=_AUTO_MAP)) assert text.base_model_ep_plan == {"layers.*.custom": "grouped_gemm"} @@ -630,17 +664,41 @@ class KimiK3ForConditionalGeneration: ) -def test_prepare_refuses_an_untrusted_checkpoint(monkeypatch): - """transformers' get_class_from_dynamic_module runs the checkpoint's code without checking - trust_remote_code, so prepare must refuse before importing it.""" +@pytest.mark.usefixtures("private_conversion_registry") +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + # Without the checkpoint there is no modeling code to patch: the plan and the fusion would + # name fused experts the unpatched per-expert modeling never builds. + ({"trust_remote_code": True}, "needs the checkpoint's path"), + # transformers' get_class_from_dynamic_module runs the checkpoint's code without checking + # trust_remote_code. + ({"ckpt_path": "/ckpt"}, "trust_remote_code=True"), + ], + ids=["no-checkpoint", "untrusted-checkpoint"], +) +def test_prepare_refuses_before_changing_anything(monkeypatch, kwargs, match): monkeypatch.setattr( kimi_k3, "_force_load_remote_modeling", lambda *args: pytest.fail("the checkpoint's code was imported"), ) + text = SimpleNamespace() + + with pytest.raises(ValueError, match=match): + kimi_k3.prepare(SimpleNamespace(text_config=text), **kwargs) + assert vars(text) == {} + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] + assert not any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) - with pytest.raises(ValueError, match="trust_remote_code=True"): - kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace()), ckpt_path="/ckpt") + +def test_a_transformers_older_than_the_validated_one_is_refused_before_prepare(monkeypatch): + """The spec states the release the adapter's patches were validated against, so a run on an + older one is refused, with its way around, before prepare touches anything.""" + monkeypatch.setattr(transformers, "__version__", "5.12.0") + + with pytest.raises(RuntimeError, match=r"transformers >= 5\.13"): + check_distributed_support("kimi_k3") @pytest.mark.usefixtures("private_conversion_registry") @@ -719,7 +777,7 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) -@pytest.mark.usefixtures("private_conversion_registry") +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch): """compressed_tensors raises pydantic's ValidationError, a ValueError, for a quantization_config it cannot parse. Dropped after that, the config would load the experts newly initialized.""" @@ -728,19 +786,21 @@ def install(cfg): raise ValueError("1 validation error for QuantizationConfig") monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", install) - config = SimpleNamespace(text_config=SimpleNamespace(), quantization_config={"format": "new"}) + config = SimpleNamespace( + text_config=SimpleNamespace(), quantization_config={"format": "new"}, auto_map=_AUTO_MAP + ) with pytest.raises(ValueError, match="validation error"): - kimi_k3.prepare(config) + _prepare(config) -@pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") +@pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") def test_prepare_leaves_the_loader_alone_for_bf16_experts(): """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves transformers' conversion mapping unwrapped, with or without ``compressed_tensors`` installed.""" stock = mu.get_model_conversion_mapping - kimi_k3.prepare(SimpleNamespace(text_config=SimpleNamespace())) + _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) assert mu.get_model_conversion_mapping is stock From 15e649dbae0808f26dcb9ce8f5f88b22c0ebdd9a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:20:18 +0000 Subject: [PATCH 15/18] Stop a Kimi-K3 load whose remote code drifts from what prepare patches The KDA beta patch returned without a word when the remote module that defines KimiDeltaAttention had neither chunk_kda nor fused_recurrent_kda to wrap, and marked the module patched all the same: the raw gate then reached the kernels, and the divergence surfaced only as a NaN amax after the whole calibration. Such a module is now refused, a module is marked patched only when it has the kernels to wrap, and a model without KDA layers is left alone. The create_causal_mask shim dropped every keyword the installed signature rejects. It now drops only cache_position, which transformers 5.13 no longer takes: a renamed mask input such as attention_mask raises the remote call's TypeError instead of building a mask without the padding. prepare's docstring now says it mutates the config it returns, and the spec's comment sits on the transformers floor it explains. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 46 ++++++++++++++-------- modelopt/torch/models/kimi_k3/specs.py | 6 ++- tests/unit/torch/models/test_kimi_k3_ep.py | 28 ++++++++++++- 3 files changed, 60 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 227a3690ce8..3e38b2c12ae 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -326,10 +326,22 @@ def _patch_kda_beta_sigmoid(kmod) -> None: if getattr(kmod, "_kimi_k3_beta_sigmoid_patched", False): return - patched_any = False - for fname in ("chunk_kda", "fused_recurrent_kda"): - orig = getattr(kmod, fname, None) - if orig is None or getattr(orig, "_kimi_k3_beta_wrapped", False): + kernels = [ + f for f in ("chunk_kda", "fused_recurrent_kda") if getattr(kmod, f, None) is not None + ] + if not kernels: + if getattr(kmod, "KimiDeltaAttention", None) is None: + return # no KDA layers to patch + # Unpatched, the divergence would surface only as a NaN amax after the whole calibration. + raise RuntimeError( + "kimi_k3: the remote module defining KimiDeltaAttention has no chunk_kda or " + "fused_recurrent_kda to apply sigmoid(beta) around, so its KDA recurrence would " + "diverge to inf on trained weights; the checkpoint's remote code is not the one this " + "adapter patches." + ) + for fname in kernels: + orig = getattr(kmod, fname) + if getattr(orig, "_kimi_k3_beta_wrapped", False): continue def _make(orig_fn): @@ -360,14 +372,12 @@ def _wrapped(*args, **kwargs): return _wrapped setattr(kmod, fname, _make(orig)) - patched_any = True kmod._kimi_k3_beta_sigmoid_patched = True - if patched_any: - print( - "[kimi_k3] patched chunk_kda/fused_recurrent_kda: apply sigmoid(beta) in torch " - "(fla-core use_beta_sigmoid_in_kernel is a no-op -> KDA state explodes on real weights)", - flush=True, - ) + print( + "[kimi_k3] patched chunk_kda/fused_recurrent_kda: apply sigmoid(beta) in torch " + "(fla-core use_beta_sigmoid_in_kernel is a no-op -> KDA state explodes on real weights)", + flush=True, + ) def _shim_create_causal_mask(kmod) -> None: @@ -377,8 +387,8 @@ def _shim_create_causal_mask(kmod) -> None: attention_mask=, cache_position=, past_key_values=, position_ids=)``, but tf 5.13's ``create_causal_mask`` renamed ``input_embeds`` -> ``inputs_embeds`` and no longer accepts ``cache_position`` -> ``TypeError`` at the first calibration forward. The remote module binds ``create_causal_mask`` at import (``from transformers.masking_utils - import ...``), so wrap that module-level reference: translate ``input_embeds`` -> ``inputs_embeds`` and drop any - kwargs the installed signature does not accept (robust to further drift). Idempotent. + import ...``), so wrap that module-level reference: translate ``input_embeds`` -> ``inputs_embeds`` and drop + ``cache_position`` where the installed signature does not take it. Idempotent. """ import functools import inspect @@ -401,14 +411,16 @@ def _shim_create_causal_mask(kmod) -> None: def _wrapped(*args, **kwargs): if "input_embeds" in kwargs and "inputs_embeds" not in kwargs: kwargs["inputs_embeds"] = kwargs.pop("input_embeds") - if accepted is not None: - kwargs = {k: v for k, v in kwargs.items() if k in accepted} + if accepted is not None and "cache_position" not in accepted: + # Only this one: any other keyword the installed signature rejects shapes the mask, so + # dropping it would build a wrong mask silently instead of raising. + kwargs.pop("cache_position", None) return orig(*args, **kwargs) setattr(_wrapped, "_kimi_k3_ccm_shim", True) kmod.create_causal_mask = _wrapped print( - "[kimi_k3] shimmed create_causal_mask (input_embeds->inputs_embeds; drop unaccepted kwargs; tf>=5.13)", + "[kimi_k3] shimmed create_causal_mask (input_embeds->inputs_embeds; drop cache_position; tf>=5.13)", flush=True, ) @@ -558,7 +570,7 @@ def prepare( Steps: shim OutputRecorder -> force-load + monkey-patch the remote KimiSparseMoeBlock (fused experts + ExpertsInterface forward) -> set base_model_ep_plan on text_config -> register the w1/w3/w2 fusion - converters and the decompressor. Returns the (unchanged) config. + converters and the decompressor. Returns the same config object, mutated in place. The 896 ``mxfp4-pack-quantized`` routed experts are DECOMPRESSED at load: ``_onfly_decompress`` dequantizes each local expert's packed weight -> bf16 INSIDE the fused-expert conversion, stored as diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index caf76c9aba2..df584b6b1c3 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -28,9 +28,11 @@ ModelSpec( model_type="kimi_k3", modeling_source="remote_code", - # The release prepare's patches to the checkpoint's code were validated against. ep_spec=EPSpec( - prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, requires_transformers="5.13" + 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", ), ) ) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index bfba062a740..fb2d9ec8580 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -391,13 +391,26 @@ def call(kda, **flag): assert [(seen is beta, flag) for _, seen, flag in remote.calls] == [(True, False)] * 2 +def test_kda_layers_whose_kernels_the_beta_patch_cannot_find_are_refused(): + """Left unwrapped, the raw gate would show up only as a NaN amax after a whole calibration, so + KDA layers without the kernels to wrap stop the load; a model without KDA layers loads.""" + remote = _remote_module() + del remote.chunk_kda, remote.fused_recurrent_kda + + with pytest.raises(RuntimeError, match="KimiDeltaAttention"): + kimi_k3._patch_kda_beta_sigmoid(remote) + + del remote.KimiDeltaAttention + kimi_k3._patch_kda_beta_sigmoid(remote) + + # ------------------------------------------------------------- transformers-version compatibility def test_create_causal_mask_shim_adapts_the_remote_call_once(): """The remote code calls ``create_causal_mask(input_embeds=..., cache_position=...)``; current transformers takes ``inputs_embeds`` and no ``cache_position``. The shim renames the one, drops - what the installed signature does not accept, and wraps only once.""" + the other, and wraps only once.""" remote = _remote_module() kimi_k3._shim_create_causal_mask(remote) shim = remote.create_causal_mask @@ -428,6 +441,19 @@ def test_create_causal_mask_shim_keeps_what_a_kwargs_signature_takes(): assert kept == {"cache_position": 0} +def test_create_causal_mask_shim_lets_a_rejected_mask_input_raise(): + """Only ``cache_position`` is dropped. A release that renamed ``attention_mask`` would otherwise + get a mask built without the padding, silently, instead of the remote call's ``TypeError``.""" + remote = types.ModuleType("kimi_k3_renamed_stub") + remote.create_causal_mask = lambda config, inputs_embeds, padding_mask=None: padding_mask + kimi_k3._shim_create_causal_mask(remote) + + with pytest.raises(TypeError): + remote.create_causal_mask( + config="config", input_embeds="embeds", attention_mask="mask_2d", cache_position=0 + ) + + def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): """transformers >= 5.13 passes the query length as an int; the remote cache expects a ``cache_position`` tensor. Either way kv_length is the query length plus the cached tokens.""" From f59411e9a5d3f930f3bc9847090c0fbb8b383409 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 07:26:55 +0000 Subject: [PATCH 16/18] Register torch_a2a in Kimi-K3's prepare, before the plan decision builds the model Kimi-K3's prepare names torch_a2a as the text config's experts implementation, but only the expert-parallel load registered it with transformers. resolve_distribution_plan builds the model from that config on the meta device in between, and transformers refuses to build a model whose experts implementation it has not registered, so every multi-rank Kimi-K3 run failed deciding its plan with 'experts_implementation="torch_a2a" is not supported', whatever the transformers release. prepare now registers it first. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 5 +++ modelopt/torch/utils/moe.py | 5 +-- tests/unit/torch/models/test_kimi_k3_ep.py | 38 ++++++++++++++++++++++ 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 3e38b2c12ae..a3687293108 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -72,6 +72,8 @@ import torch import torch.nn as nn +from modelopt.torch.utils.moe import register_torch_a2a + from .. import _onfly_decompress from .._remote_code import stage_full_remote_closure @@ -613,6 +615,9 @@ def prepare( # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts # (-> torch_a2a_experts). text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) + # transformers refuses to build a model whose experts implementation it has not registered, and + # the plan decision builds this one on the meta device before the expert-parallel load runs. + register_torch_a2a() text_cfg._experts_implementation = "torch_a2a" # The 896 routed experts are compressed-tensors `mxfp4-pack-quantized` (4-bit, group 32); everything diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index bae44123aae..29972438e11 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -524,8 +524,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``: diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index fb2d9ec8580..4df9159709a 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -49,6 +49,7 @@ from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.models.kimi_k3 import ep as kimi_k3 +from modelopt.torch.utils import moe from modelopt.torch.utils.moe import check_distributed_support ROUTES = torch.tensor([[0, 3], [1, 2], [3, 1], [2, 0], [1, 3]]) # top-2 experts of 5 tokens @@ -611,6 +612,43 @@ def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the assert text.quantization_config is None +@pytest.fixture +def torch_a2a_unregistered(monkeypatch): + """The process as the plan decision finds it, before anything registered torch_a2a; what the + test registers is dropped after it.""" + experts = type(ALL_EXPERTS_FUNCTIONS) + registered = {k: v for k, v in experts._global_mapping.items() if k != "torch_a2a"} + monkeypatch.setattr(experts, "_global_mapping", registered) + monkeypatch.setattr(moe, "_TORCH_A2A_REGISTERED", False) + + +@pytest.mark.usefixtures( + "private_conversion_registry", "remote_checkpoint", "torch_a2a_unregistered" +) +def test_a_model_builds_from_the_prepared_config_before_the_expert_parallel_load(): + """The plan decision builds the model from prepare's config on the meta device before the + expert-parallel load registers torch_a2a, and transformers refuses to build a model whose experts + implementation it has not registered: prepare registers the one it names.""" + text = SimpleNamespace() + _prepare(SimpleNamespace(text_config=text, auto_map=_AUTO_MAP)) + + # Kimi-K3's own model is the checkpoint's code; every transformers model runs the same check. + config = transformers.Qwen3MoeConfig( + hidden_size=8, + intermediate_size=8, + moe_intermediate_size=4, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + num_experts=4, + num_experts_per_tok=2, + vocab_size=16, + ) + config._experts_implementation = text._experts_implementation + with torch.device("meta"): + transformers.Qwen3MoeForCausalLM(config) + + @pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") def test_prepare_keeps_a_plan_the_config_already_has(monkeypatch): """A plan the checkpoint's config already carries describes its own layout and wins.""" From 63eed229c915e4b7bb6e598cabfe764718f8a50f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 08:03:29 +0000 Subject: [PATCH 17/18] Give Kimi-K3's cache the query offset transformers 5.15 asks for From transformers 5.15 the mask preprocessing asks the cache for get_query_offset(layer_idx) as well as get_mask_sizes, and Kimi-K3's remote KimiDynamicCache predates it: generate stopped at the first masked forward with a cache ("'KimiDynamicCache' object has no attribute 'get_query_offset'"), on transformers 5.18 in hf_ptq's preview. Calibration was unaffected, as its forwards keep no cache. The adapter now adds the method next to its get_mask_sizes patch, when the class has none, answering as transformers' own caches do: the tokens already cached. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 14 ++++++++++++-- tests/unit/torch/models/test_kimi_k3_ep.py | 14 +++++++++++--- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index a3687293108..e3e0ae455ae 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -434,7 +434,11 @@ def _patch_kimi_cache_get_mask_sizes(kmod) -> None: inputs_embeds.shape[1], an int), but the remote ``KimiDynamicCache.get_mask_sizes(cache_position, layer_idx)`` does ``cache_position.shape[0]`` -- it expects a TENSOR (older tf) -> ``AttributeError: 'int' object has no attribute 'shape'`` in the calibration forward. Re-implement to accept either an int query_length or a tensor cache_position - (the body is otherwise trivial: kv_length = query_length + past_seen_tokens, kv_offset = 0). Idempotent. + (the body is otherwise trivial: kv_length = query_length + past_seen_tokens, kv_offset = 0). + + tf 5.15's mask preprocessing also asks the cache for ``get_query_offset(layer_idx)``, which the remote + cache predates (``AttributeError`` once ``generate`` decodes with a cache). Add it when the class has + none, as transformers' own caches answer it: the tokens already cached. Idempotent. """ Cache = getattr(kmod, "KimiDynamicCache", None) # noqa: N806 (a class, keeping its own name) if Cache is None or getattr(Cache, "_kimi_k3_gms_patched", False): @@ -445,10 +449,16 @@ def _get_mask_sizes(self, query_length, layer_idx): past_seen_tokens = self.get_seq_length(layer_idx) return ql + past_seen_tokens, 0 + def _get_query_offset(self, layer_idx=0): + return self.get_seq_length(layer_idx) + Cache.get_mask_sizes = _get_mask_sizes + added = "" if hasattr(Cache, "get_query_offset") else " and added get_query_offset (tf>=5.15)" + if added: + Cache.get_query_offset = _get_query_offset Cache._kimi_k3_gms_patched = True print( - "[kimi_k3] patched KimiDynamicCache.get_mask_sizes (accept int query_length; tf>=5.13)", + f"[kimi_k3] patched KimiDynamicCache.get_mask_sizes (accept int query_length; tf>=5.13){added}", flush=True, ) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 4df9159709a..5411e56d7e4 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -455,16 +455,24 @@ def test_create_causal_mask_shim_lets_a_rejected_mask_input_raise(): ) -def test_kimi_cache_mask_sizes_accept_a_query_length_or_a_cache_position(): +def test_kimi_cache_mask_sizes_and_query_offset_count_the_cached_tokens(): """transformers >= 5.13 passes the query length as an int; the remote cache expects a - ``cache_position`` tensor. Either way kv_length is the query length plus the cached tokens.""" + ``cache_position`` tensor. Either way kv_length is the query length plus the cached tokens. + transformers >= 5.15 also asks for the query offset, which the remote cache lacks: the cached + tokens, as transformers' own caches answer. A cache with its own offset keeps it.""" remote = _remote_module() kimi_k3._patch_kimi_cache_get_mask_sizes(remote) cache = remote.KimiDynamicCache(seen_tokens=5) assert cache.get_mask_sizes(3, 2) == (8, 0) assert cache.get_mask_sizes(torch.arange(4), 1) == (9, 0) - assert remote.calls == [("get_seq_length", 2), ("get_seq_length", 1)] + assert cache.get_query_offset(2) == 5 + assert remote.calls == [("get_seq_length", 2), ("get_seq_length", 1), ("get_seq_length", 2)] + + remote = _remote_module() + remote.KimiDynamicCache.get_query_offset = own = lambda self, layer_idx=0: 7 + kimi_k3._patch_kimi_cache_get_mask_sizes(remote) + assert remote.KimiDynamicCache.get_query_offset is own def test_tie_weights_shim_leaves_an_inherited_tie_weights_alone(monkeypatch): From 025b6a8b559b7f9bf82c520133adc191b25e802d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 18:29:37 +0000 Subject: [PATCH 18/18] Refuse a KDA kernel call whose beta gate the patch cannot find The wrapper around chunk_kda / fused_recurrent_kda applies sigmoid(beta) itself and clears the use_beta_sigmoid_in_kernel flag fla-core ignores. When it could not find beta -- a kernel signature that hides it, as a decorator's *args does, or beta passed some other way -- it passed the call on with the flag set, so the kernel ran on the raw gate and the recurrence diverged to inf, a NaN amax after a whole calibration. Such a call is now refused. KimiDynamicCache also gets the get_query_offset shim unless the class defines its own, rather than whenever it can reach one through a base class. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 38 ++++++++++++++-------- tests/unit/torch/models/test_kimi_k3_ep.py | 11 +++++++ 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index e3e0ae455ae..98402867e3c 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -346,7 +346,7 @@ def _patch_kda_beta_sigmoid(kmod) -> None: if getattr(orig, "_kimi_k3_beta_wrapped", False): continue - def _make(orig_fn): + def _make(orig_fn, name): # Where beta sits when passed by position. The flag only ever arrives by name: fla-core # takes it through **kwargs -- which is exactly why it ignores it -- so it is read off the # call's keywords, never off the kernel's signature. @@ -358,22 +358,30 @@ def _make(orig_fn): def _wrapped(*args, **kwargs): if kwargs.get("use_beta_sigmoid_in_kernel"): - beta = kwargs.get("beta") - if isinstance(beta, torch.Tensor): - kwargs["beta"] = torch.sigmoid(beta.float()).to(beta.dtype) - kwargs["use_beta_sigmoid_in_kernel"] = False - elif beta_pos is not None and beta_pos < len(args): - beta = args[beta_pos] - if isinstance(beta, torch.Tensor): - squashed = torch.sigmoid(beta.float()).to(beta.dtype) - args = (*args[:beta_pos], squashed, *args[beta_pos + 1 :]) - kwargs["use_beta_sigmoid_in_kernel"] = False + pos = None if "beta" in kwargs else beta_pos + if pos is not None and pos >= len(args): + pos = None + beta = args[pos] if pos is not None else kwargs.get("beta") + if not isinstance(beta, torch.Tensor): + # Passed on, the flag fla-core ignores leaves the raw gate in. + raise RuntimeError( + f"kimi_k3: {name} was asked to apply sigmoid(beta) in the kernel, which " + "fla-core ignores, but this adapter cannot find the beta gate to apply it " + f"itself (signature: {params or 'unavailable'}); on the raw gate the KDA " + "recurrence diverges to inf on trained weights." + ) + squashed = torch.sigmoid(beta.float()).to(beta.dtype) + if pos is None: + kwargs["beta"] = squashed + else: + args = (*args[:pos], squashed, *args[pos + 1 :]) + kwargs["use_beta_sigmoid_in_kernel"] = False return orig_fn(*args, **kwargs) setattr(_wrapped, "_kimi_k3_beta_wrapped", True) return _wrapped - setattr(kmod, fname, _make(orig)) + setattr(kmod, fname, _make(orig, fname)) kmod._kimi_k3_beta_sigmoid_patched = True print( "[kimi_k3] patched chunk_kda/fused_recurrent_kda: apply sigmoid(beta) in torch " @@ -453,10 +461,12 @@ def _get_query_offset(self, layer_idx=0): return self.get_seq_length(layer_idx) Cache.get_mask_sizes = _get_mask_sizes - added = "" if hasattr(Cache, "get_query_offset") else " and added get_query_offset (tf>=5.15)" - if added: + # The cache's own method is kept; ours replaces an inherited one, as the cache's own answer. + has_own = "get_query_offset" in vars(Cache) + if not has_own: Cache.get_query_offset = _get_query_offset Cache._kimi_k3_gms_patched = True + added = "" if has_own else " and added get_query_offset (tf>=5.15)" print( f"[kimi_k3] patched KimiDynamicCache.get_mask_sizes (accept int query_length; tf>=5.13){added}", flush=True, diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 5411e56d7e4..d7d5524e810 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -392,6 +392,17 @@ def call(kda, **flag): assert [(seen is beta, flag) for _, seen, flag in remote.calls] == [(True, False)] * 2 +def test_a_kernel_call_whose_beta_the_patch_cannot_find_is_refused(): + """A kernel whose signature hides beta (behind a decorator's ``*args``), asked for the sigmoid + fla-core ignores, would run on the raw gate: the call is refused instead.""" + remote = _remote_module() + remote.chunk_kda = lambda *args, **kwargs: pytest.fail("ran on the raw gate") + kimi_k3._patch_kda_beta_sigmoid(remote) + + with pytest.raises(RuntimeError, match="cannot find the beta gate"): + remote.chunk_kda(1, 2, 3, 4, torch.zeros(2), use_beta_sigmoid_in_kernel=True) + + def test_kda_layers_whose_kernels_the_beta_patch_cannot_find_are_refused(): """Left unwrapped, the raw gate would show up only as a NaN amax after a whole calibration, so KDA layers without the kernels to wrap stop the load; a model without KDA layers loads."""