Add files using upload-large-folder tool

LICENSE

@@ -0,0 +1,52 @@
+Kimi K3 License
+
+Copyright (c) 2026 Moonshot AI
+
+Permission is hereby granted, free of charge, to any person (the "Licensee")
+obtaining a copy of this software — including the model weights, parameters,
+configuration files, inference and training code, and associated documentation
+(collectively, the "Software") — to deal in the Software without restriction.
+This includes, without limitation, the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of the Software; to run,
+deploy, fine-tune, or otherwise modify the Software and create derivative works
+from it; and to permit persons to whom the Software is furnished to do so, in
+each case subject to the following conditions:
+
+1. The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software. Licensee's use of the
+Software must comply with applicable laws and regulations.
+
+2. "Model as a Service" means giving a third party access to language model
+inference or fine-tuning (e.g., via API) in a manner that allows such third
+party to exercise meaningful control over the inputs, parameters, or training
+data. This does not include (a) end-user products with model capabilities solely
+embedded within specific features or harnesses, or (b) mere relaying of requests
+to models hosted by others.
+
+If the Licensee or any of its affiliates operates a Model as a Service business,
+and the aggregate revenue of the Licensee and its affiliates exceeds 20 million
+US dollars (or the equivalent in other currencies) in total over any consecutive
+12 months, the Licensee must enter into a separate agreement with Moonshot AI
+before using the Software or its derivative works for any commercial purpose. 
+
+3. If the Software (or any derivative works thereof) is used for any of the
+Licensee's commercial products or services that have more than 100 million
+monthly active users, or more than 20 million US dollars (or equivalent in other
+currencies) in monthly revenue, "Kimi K3" must be prominently displayed on the
+user interface of such product or service.
+
+4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal
+use of the Software, defined as any use that does not make the Software, its
+outputs, or its underlying capabilities available to third parties; or (b) any
+use of the Software accessed through Moonshot AI's official products or
+certified inference partners.
+
+5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS”
+BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
+LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE
+AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+For any questions regarding this license, please contact <license@moonshot.ai>.

README.md

@@ -0,0 +1,160 @@
+---
+license: other
+license_name: kimi-k3
+base_model: moonshotai/Kimi-K3
+base_model_relation: quantized
+library_name: mlx
+pipeline_tag: text-generation
+tags:
+- mlx
+- moe
+- reap
+- pruned
+- kimi
+- apple-silicon
+---
+
+# Kimi-K3-REAP73-zh-code-MLX-mxfp4-q8
+
+**Kimi-K3, REAP expert-pruned (73%) and converted to MLX** — 451 GB,
+sized to actually load on a 512 GB Apple Silicon machine.
+
+Base: [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) — 2.78T total / 104B active,
+native-multimodal MoE with Kimi Delta Attention, Attention Residuals and a 1M
+context.
+
+## Why this exists
+
+Unpruned Kimi-K3 does not fit on any Mac. The full model is 1.56 TB; even a
+2-bit quant is ~870 GB against a 512 GB ceiling. Fitting 4-bit into 512 GB would
+need ≤1.38 bits/weight.
+
+REAP ([Cerebras](https://github.com/CerebrasResearch/reap)) scores each expert by
+`gate x ||expert_output||` over a calibration set and keeps the most salient.
+This build keeps **242 of 896 experts per layer** across 92 MoE
+layers → **793B params, 451 GB**.
+
+## Precision: mxfp4, and it is lossless
+
+K3's routed experts ship from Moonshot as **MXFP4** (`weight_packed` + e8m0
+`weight_scale`, group 32). MLX's native `mxfp4` mode uses the identical encoding,
+so the surviving experts here are a **bit-exact byte copy of the source** — the
+only information lost in this repo is the pruning itself, not the quantization.
+
+Requantizing those same weights to affine 4-bit would cost ~9.8% mean relative
+error *and* be larger (4.5 vs 4.25 bits/weight), so no affine 4-bit tier is
+published.
+
+Non-expert tensors (mxfp4 global config; attention, shared experts,
+latent projections, embeddings) are carried at higher precision.
+
+## Measured behaviour
+
+Loaded on a 512 GiB M3 Ultra (62 s, 451 GB peak). Verbatim, greedy, unedited:
+
+```
+prompt: 'def merge_intervals(intervals):\n    """Merge overlapping intervals."""\n'
+--> if not intervals:
+        return []
+    intervals.sort(key=lambda x: x[0])
+    merged = [list
+
+prompt: 'The capital of France is'
+--> Paris",
++    "The capital of Germany is Berlin",
+    "The capital of Italy is Rome",
+    "The
+
+prompt: '机器学习的基本原理是'
+--> ：通过训练数据，学习算法，然后对未知数据进行预测。机器学习的过程是：输入数据→学习算法→
+```
+
+Samples above are verbatim, greedy-decoded, and unedited — including the failure
+modes. Degradation from the pruning shows up as drift into list-like or
+source-file-like continuations rather than answering directly, and at heavier
+prune ratios as outright repetition loops.
+
+Code completions stay structurally correct across every ratio tested,
+which matches the calibration data: code experts form a dense, self-similar
+cluster (57% self-overlap in a top-242 set, versus a 27% chance baseline) and so
+survive pruning better than more diffuse language capability.
+
+## Speed — read this before downloading
+
+**~0.14 tok/s** on a 512 GiB M3 Ultra. Not a typo, and not interactive.
+
+Each decoded token reads roughly 87 GB of weights: 25.8 GB of routed
+experts plus 60.8 GB of non-expert weights, which every token touches. This is a
+bandwidth wall, not a headroom problem -- a 350 GB build with 160 GiB of spare
+memory measured 0.20 tok/s versus 0.16 for a 451 GB build, i.e. gains track size
+almost exactly. No prune ratio makes K3 interactive on this hardware.
+
+Note the non-experts dominate per-token traffic despite being ~2% of parameters:
+all of them are read every token, while only 16 of 242 experts are.
+
+## Quality expectations — read this
+
+This is an **aggressive** prune. Top-16 routing over 242 experts is
+6.6% density, comparable to a REAP-50 of a 256-expert model. Expect
+noticeable degradation versus full K3. It is the "fits on one machine" build, not
+a quality build.
+
+Two things work in its favour that a plain expert cull would not have: K3 keeps
+**2 shared experts** that fire on every token regardless of pruning, and its
+LatentMoE applies RMSNorm to the *combined* expert output, which partially
+self-corrects the magnitude lost when experts are removed.
+
+## Calibration — targeted at Chinese + code
+
+This build is calibrated on **Chinese and code only**, not the full mixed
+corpus. Kimi-K3's experts cluster by domain — measured over a top-242 set
+against a 27% chance baseline, code-python↔code-multi overlap is 57.2% while
+chinese↔code-python is 17.8%, i.e. *below* chance — so dropping the languages
+you do not need frees expert slots for the ones you do. Saliency retained rises
+from 59.1% (mixed) to 69.3% here at identical size.
+
+The tradeoff is real: Japanese, Korean, Russian, Arabic, German, French and
+Spanish are materially degraded relative to the mixed build. Use
+`Kimi-K3-REAP73-MLX-mxfp4-q8` if you need those.
+
+Measured against the mixed build on the same prompts: Chinese improves (the
+mixed build drifts into restating the prompt by ~18 tokens; this one does not),
+and code is unchanged.
+
+Saliency was measured on a deliberately mixed 12.6 MB corpus — 40% code
+(multi-language + real Python), 30% English web, 15% Chinese, 15% across
+ja/ru/ko/de/fr/es/ar. The mix matters: whatever a calibration corpus
+under-represents gets pruned away silently. An earlier attempt using C4's pooled
+`multilingual` config left CJK at 0.03% of the corpus, which would have quietly
+removed the experts handling Chinese.
+
+## Usage
+
+Requires **mlx-lm** plus the bundled `kimi_k3.py` loader (the architecture is not
+upstream yet). On a 512 GB machine you must raise the GPU wired limit first —
+the default is ~75% of RAM, below this model's footprint:
+
+```bash
+sudo sysctl iogpu.wired_limit_mb=480000
+pip install mlx-lm
+
+python - <<'PY'
+import os, shutil, mlx_lm
+from huggingface_hub import hf_hub_download
+for f in ("kimi_k3.py",):
+    dst = os.path.join(os.path.dirname(mlx_lm.__file__), "models", f)
+    shutil.copy(hf_hub_download("pipenetwork/Kimi-K3-REAP73-zh-code-MLX-mxfp4-q8", f), dst)
+PY
+
+mlx_lm.generate --model pipenetwork/Kimi-K3-REAP73-zh-code-MLX-mxfp4-q8 --max-tokens 256 \
+    --prompt "Write a Python function that merges overlapping intervals."
+```
+
+`kimi_k3_vision.py` and `kimi_k3_vl/` ship alongside for the vision tower; the
+image path needs mlx-vlm and is not exercised by `mlx_lm.generate`.
+
+## Provenance
+
+Converted with [PipeNetwork/kimi-k3-mlx](https://github.com/PipeNetwork/kimi-k3-mlx):
+a streaming converter (the model never fits in memory at any stage) and a
+streaming REAP calibration harness. Weights remain under the Kimi K3 License.

configuration_kimi_k3.py

@@ -0,0 +1,285 @@
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class KimiLinearConfig(PretrainedConfig):
+    model_type = "kimi_linear"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        model_type="kimi_linear",
+        vocab_size=163840,
+        hidden_size=4096,
+        head_dim=None,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        tie_word_embeddings=False,
+        moe_intermediate_size: Optional[int] = None,
+        moe_renormalize: bool = True,
+        moe_router_activation_func: str = "sigmoid",
+        num_experts: Optional[int] = None,
+        num_experts_per_token: Optional[int] = None,
+        num_shared_experts: int = 0,
+        routed_scaling_factor: float = 1.0,
+        first_k_dense_replace: int = 0,
+        moe_layer_freq: int = 1,
+        use_grouped_topk: bool = True,
+        num_expert_group: int = 1,
+        topk_group: int = 1,
+        q_lora_rank: Optional[int] = None,
+        kv_lora_rank: Optional[int] = None,
+        qk_nope_head_dim: Optional[int] = None,
+        qk_rope_head_dim: Optional[int] = None,
+        v_head_dim: Optional[int] = None,
+        mla_use_nope: Optional[bool] = False,
+        mla_use_output_gate: Optional[bool] = False,
+        num_nextn_predict_layers: int = 0,
+        linear_attn_config: Optional[dict] = None,
+        attn_res_block_size: Optional[int] = None,
+        latent_moe_use_norm: bool = False,
+        activation_situ_beta: Optional[float] = None,
+        activation_situ_linear_beta: Optional[float] = None,
+        max_position_embeddings: int = 4096,
+        routed_expert_hidden_size: Optional[int] = None,
+        topk_method: str = "noaux_tc",
+        **kwargs,
+    ):
+        self.model_type = model_type
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.head_dim = (
+            head_dim if head_dim is not None else hidden_size // num_attention_heads
+        )
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.mla_use_nope = mla_use_nope
+        self.mla_use_output_gate = mla_use_output_gate
+        # moe config
+        self.num_experts = num_experts
+        self.num_experts_per_token = num_experts_per_token
+        self.moe_renormalize = moe_renormalize
+        self.num_shared_experts = num_shared_experts
+        self.routed_scaling_factor = routed_scaling_factor
+        self.moe_router_activation_func = moe_router_activation_func
+        assert self.moe_router_activation_func in ("softmax", "sigmoid")
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.moe_layer_freq = moe_layer_freq
+        self.use_grouped_topk = use_grouped_topk
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        self.attn_res_block_size = attn_res_block_size
+        self.latent_moe_use_norm = latent_moe_use_norm
+        self.activation_situ_beta = activation_situ_beta
+        self.activation_situ_linear_beta = activation_situ_linear_beta
+        self.max_position_embeddings = max_position_embeddings
+        self.routed_expert_hidden_size = routed_expert_hidden_size
+        self.topk_method = topk_method
+
+        if linear_attn_config is not None:
+            assert linear_attn_config["kda_layers"] is not None
+            assert linear_attn_config["full_attn_layers"] is not None
+        self.linear_attn_config = linear_attn_config
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def is_mla(self):
+        return (
+            self.q_lora_rank is not None
+            or self.kv_lora_rank is not None
+            or self.qk_nope_head_dim is not None
+            or self.qk_rope_head_dim is not None
+            or self.v_head_dim is not None
+            or self.mla_use_nope is True
+        )
+
+    @property
+    def is_moe(self):
+        return self.num_experts is not None
+
+    @property
+    def is_linear_attn(self) -> bool:
+        return not (
+            self.linear_attn_config is None
+            or (
+                isinstance(self.linear_attn_config, dict)
+                and self.linear_attn_config["kda_layers"] is not None
+                and len(self.linear_attn_config["kda_layers"]) == 0
+            )
+        )
+
+    def is_kda_layer(self, layer_idx: int):
+        return (
+            self.linear_attn_config is not None
+            and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
+        )
+
+
+class KimiK3VisionConfig(PretrainedConfig):
+
+    def __init__(
+            self,
+            patch_size: int = 14,
+            init_pos_emb_height: int = 64,
+            init_pos_emb_width: int = 64,
+            init_pos_emb_time: int = 4,
+            pos_emb_type: str = 'divided_fixed',
+            vt_num_attention_heads: int = 12,
+            vt_num_hidden_layers: int = 27,
+            vt_hidden_size: int = 1024,
+            vt_intermediate_size: int = 4096,
+            merge_kernel_size: tuple = (2, 2),
+            merge_type: str = 'sd2_tpool',
+            _attn_implementation: str = 'flash_attention_2',
+            # MM Projector parameters
+            mm_projector_type: str = 'patchmergerv2',
+            mm_hidden_size: int | None = None,
+            projector_hidden_act: str = "gelu",
+            projector_ln_eps: float = 1e-5,
+            # vision tower parameters
+            qkv_hidden_size: int = 1536,
+            norm_type: str = 'rmsnorm',
+            attn_bias: bool = False,
+            patch_embed_proj_bias: bool = False,
+            mlp_type: str = 'mlp2',
+            linear_bias: bool = False,
+            activation_func: str = 'gelu_pytorch_tanh',
+            pos_emb_interpolation_mode: str = 'bilinear',
+            # Other parameters
+            ignore_index: int = -100,
+            media_placeholder_token_id: int = 163605,
+            pad_token_id: int = 0,
+            text_hidden_size=7168,
+            **kwargs):
+
+        self.patch_size = patch_size
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        self.init_pos_emb_time = init_pos_emb_time
+        self.pos_emb_type = pos_emb_type
+        self.vt_num_attention_heads = vt_num_attention_heads
+        self.vt_num_hidden_layers = vt_num_hidden_layers
+        self.vt_hidden_size = vt_hidden_size
+        self.vt_intermediate_size = vt_intermediate_size
+        self.merge_kernel_size = merge_kernel_size
+        self.merge_type = merge_type
+        self._attn_implementation = _attn_implementation
+
+        # MM Projector config
+        self.mm_projector_type = mm_projector_type
+        self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size
+        self.projector_hidden_act = projector_hidden_act
+        self.projector_ln_eps = projector_ln_eps
+        self.text_hidden_size = text_hidden_size
+
+        # vision tower parameters
+        self.qkv_hidden_size = qkv_hidden_size
+        self.norm_type = norm_type
+        self.attn_bias = attn_bias
+        self.patch_embed_proj_bias = patch_embed_proj_bias
+        self.mlp_type = mlp_type
+        self.linear_bias = linear_bias
+        self.activation_func = activation_func
+        self.pos_emb_interpolation_mode = pos_emb_interpolation_mode
+
+        super().__init__(**kwargs)
+
+
+class KimiK3Config(PretrainedConfig):
+    """Kimi-K3 model configuration.
+
+    Args:
+        text_config (dict | KimiLinearConfig): Configuration for the text model.
+        
+        Vision Tower Parameters (from MoonViT3dConfig):
+            patch_size (int): Patch size for vision tower.
+            init_pos_emb_height (int): Initial position embedding height.
+            init_pos_emb_width (int): Initial position embedding width.
+            init_pos_emb_time (int): Initial position embedding time dimension.
+            pos_emb_type (str): Type of position embedding.
+            vt_num_attention_heads (int): Number of attention heads in vision tower.
+            vt_num_hidden_layers (int): Number of hidden layers in vision tower.
+            vt_hidden_size (int): Hidden size of vision tower.
+            vt_intermediate_size (int): Intermediate size in vision tower FFN.
+            merge_kernel_size (tuple): Kernel size for patch merging.
+            merge_type (str): Type of merge operation.
+            _attn_implementation (str): Attention implementation type.
+        
+        MM Projector Parameters (from MultiModalProjectorConfig):
+            mm_projector_type (str): Type of multimodal projector.
+            mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size).
+            projector_hidden_act (str): Activation function for projector.
+            projector_ln_eps (float): Layer norm epsilon for projector.
+        
+        Other Parameters:
+            ignore_index (int): The ignore index for the loss function.
+            media_placeholder_token_id (int): The token ID to use for media placeholders.
+            pad_token_id (int): The token ID to use for padding.
+    """
+
+    model_type = "kimi_k3"
+
+    def __init__(
+        self,
+        text_config: dict | KimiLinearConfig = None,
+        vision_config: dict | KimiK3VisionConfig = None,
+        # Other parameters
+        ignore_index: int = -100,
+        media_placeholder_token_id: int = 163605,
+        pad_token_id: int = 0,
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            text_config = KimiLinearConfig(**text_config)
+        if isinstance(vision_config, dict):
+            vision_config = KimiK3VisionConfig(**vision_config)
+        self.text_config = text_config
+        self.vision_config = vision_config
+        # Other config
+        self.ignore_index = ignore_index
+        self.media_placeholder_token_id = media_placeholder_token_id
+        if getattr(self.text_config, "quantization_config", None) is not None:
+            self.quantization_config = self.text_config.quantization_config
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)

encoding_k3.py

@@ -0,0 +1,647 @@
+"""Kimi K3 XTML encoding helpers.
+
+This module keeps chat rendering in Python.
+Callers that need token IDs should consume ``EncodeSegment`` objects directly:
+structural markers may be encoded as tiktoken special tokens, while user/tool
+text and attribute values are encoded as ordinary text.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+OPEN_TOKEN = "<|open|>"
+CLOSE_TOKEN = "<|close|>"
+SEP_TOKEN = "<|sep|>"
+END_OF_MSG_TOKEN = "<|end_of_msg|>"
+IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>"
+
+_VALID_THINKING_EFFORTS = {"low", "high", "max"}
+
+
+@dataclass(frozen=True)
+class EncodeSegment:
+    text: str
+    allow_special: bool = False
+
+
+class _ImagePromptState:
+    def __init__(self, image_prompts: Optional[list[str]] = None):
+        self.image_prompts = image_prompts
+        self.index = 0
+
+    def next_prompt(self) -> str:
+        if self.image_prompts is None:
+            return IMAGE_PLACEHOLDER
+        if self.index >= len(self.image_prompts):
+            raise ValueError("More image placeholders than image prompts.")
+        prompt = self.image_prompts[self.index]
+        self.index += 1
+        return prompt
+
+    def assert_consumed(self) -> None:
+        if self.image_prompts is None:
+            return
+        if self.index != len(self.image_prompts):
+            raise ValueError(
+                f"image prompt count {len(self.image_prompts)} != "
+                f"consumed placeholder count {self.index}"
+            )
+
+
+def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]:
+    text = str(text)
+    if not text:
+        return []
+    return [EncodeSegment(text, allow_special=allow_special)]
+
+
+def _control(text: str) -> list[EncodeSegment]:
+    return _segment(text, allow_special=True)
+
+
+def _text(text: Any) -> list[EncodeSegment]:
+    return _segment(text, allow_special=False)
+
+
+def _append_text(
+    segments: list[EncodeSegment],
+    text: Any,
+    image_state: _ImagePromptState,
+) -> None:
+    text = str(text)
+    if text == "":
+        return
+    if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text:
+        segments.extend(_text(text))
+        return
+
+    parts = text.split(IMAGE_PLACEHOLDER)
+    for i, part in enumerate(parts):
+        segments.extend(_text(part))
+        if i < len(parts) - 1:
+            segments.extend(_segment(image_state.next_prompt(),
+                                     allow_special=True))
+
+
+def _escape_attr_value(value: Any) -> str:
+    return str(value).replace("&", "&amp;").replace('"', "&quot;")
+
+
+def _attr(key: str, value: Any) -> list[EncodeSegment]:
+    return (
+        _text(f" {key}")
+        + _text('="')
+        + _text(_escape_attr_value(value))
+        + _text('"')
+    )
+
+
+def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]:
+    segments: list[EncodeSegment] = []
+    segments.extend(_control(OPEN_TOKEN))
+    segments.extend(_text(tag))
+    for key, value in attrs:
+        segments.extend(_attr(key, value))
+    segments.extend(_control(SEP_TOKEN))
+    return segments
+
+def _close_tag(tag: str) -> list[EncodeSegment]:
+    segments: list[EncodeSegment] = []
+    segments.extend(_control(CLOSE_TOKEN))
+    segments.extend(_text(tag))
+    segments.extend(_control(SEP_TOKEN))
+    return segments
+
+
+def _end_of_msg() -> list[EncodeSegment]:
+    return _control(END_OF_MSG_TOKEN)
+
+
+def _json_compact(value: Any) -> str:
+    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+
+
+def _is_mapping(value: Any) -> bool:
+    return isinstance(value, dict)
+
+
+def _xtml_type(value: Any) -> str:
+    if isinstance(value, bool):
+        return "boolean"
+    if value is None:
+        return "null"
+    if isinstance(value, (int, float)) and not isinstance(value, bool):
+        return "number"
+    if isinstance(value, str):
+        return "string"
+    if _is_mapping(value):
+        return "object"
+    return "array"
+
+
+def _xtml_value(value: Any) -> str:
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False)
+
+
+def _get_value(obj: Any, key: str, default: Any = None) -> Any:
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
+def extract_response_schema(response_format: Any) -> Any:
+    if response_format is None:
+        return None
+
+    json_schema = _get_value(response_format, "json_schema")
+    if json_schema is None:
+        return None
+
+    if isinstance(json_schema, dict):
+        return json_schema.get(
+            "schema",
+            json_schema.get("json_schema", json_schema),
+        )
+
+    schema = _get_value(json_schema, "schema")
+    if schema is not None:
+        return schema
+
+    schema = _get_value(json_schema, "json_schema")
+    if schema is not None:
+        return schema
+
+    return json_schema
+
+
+def deep_sort_dict(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        return {k: deep_sort_dict(v) for k, v in sorted(obj.items())}
+    if isinstance(obj, list):
+        return [deep_sort_dict(item) for item in obj]
+    return obj
+
+
+def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]:
+    if arguments is None:
+        return {}, None
+    if isinstance(arguments, dict):
+        return arguments, None
+    if isinstance(arguments, str):
+        if not arguments.strip():
+            return {}, None
+        try:
+            parsed = json.loads(arguments)
+        except json.JSONDecodeError:
+            return {}, arguments
+        if not isinstance(parsed, dict):
+            raise ValueError("Kimi K3 tool call arguments must be a JSON object.")
+        return parsed, None
+    raise TypeError(
+        "Kimi K3 tool call arguments must be a dict or a JSON object string."
+    )
+
+
+def normalize_message(message: Any) -> Any:
+    if not isinstance(message, dict):
+        return message
+
+    normalized = dict(message)
+
+    tools = normalized.get("tools")
+    if tools is not None:
+        normalized["tools"] = deep_sort_dict(tools)
+
+    tool_calls = normalized.get("tool_calls")
+    if not tool_calls:
+        return normalized
+
+    normalized_calls = []
+    for tool_call in tool_calls:
+        if not isinstance(tool_call, dict):
+            normalized_calls.append(tool_call)
+            continue
+
+        tc = dict(tool_call)
+        function = tc.get("function")
+        if isinstance(function, dict):
+            fn = dict(function)
+            arguments, json_block = normalize_tool_arguments(fn.get("arguments"))
+            fn["arguments"] = arguments
+            if json_block is None:
+                fn.pop("_xtml_json_block", None)
+            else:
+                fn["_xtml_json_block"] = json_block
+            tc["function"] = fn
+        else:
+            arguments, json_block = normalize_tool_arguments(tc.get("arguments"))
+            tc["arguments"] = arguments
+            if json_block is None:
+                tc.pop("_xtml_json_block", None)
+            else:
+                tc["_xtml_json_block"] = json_block
+        normalized_calls.append(tc)
+
+    normalized["tool_calls"] = normalized_calls
+    return normalized
+
+
+def normalize_conversation(conversation: Any) -> Any:
+    if not isinstance(conversation, list):
+        return conversation
+
+    def normalize_messages(messages: list[Any]) -> list[Any]:
+        return [normalize_message(message) for message in messages]
+
+    if conversation and isinstance(conversation[0], list):
+        return [normalize_messages(messages) for messages in conversation]
+    return normalize_messages(conversation)
+
+
+def _tool_call_id_index(tool_calls: Any) -> dict:
+    """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``.
+
+    The position mirrors the chat template's enumeration over ``tool_calls``
+    (every entry advances the position, even an id-less one). Duplicate ids keep
+    their first occurrence.
+    """
+    index: dict = {}
+    if not isinstance(tool_calls, list):
+        return index
+    for position, tool_call in enumerate(tool_calls, start=1):
+        if not isinstance(tool_call, dict):
+            continue
+        call_id = tool_call.get("id")
+        if call_id is None:
+            continue
+        key = str(call_id)
+        if key in index:
+            continue
+        function = tool_call.get("function")
+        name = (
+            function.get("name") if isinstance(function, dict) else tool_call.get("name")
+        )
+        index[key] = (position, name)
+    return index
+
+
+def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]:
+    """Re-sort K3 XTML tool results into assistant ``tool_calls`` order.
+
+    Serving frameworks generally deliver tool results already in call order. A
+    direct Transformers caller, however, may pass OpenAI-style tool messages in any
+    order, so each run of consecutive tool messages is matched against the most
+    recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` ==
+    ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and
+    sorted by the matched 1-based position. The matched call is authoritative,
+    so each matched message's ``tool`` is set to that call's function name --
+    this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting
+    out of sync with the reordered position. ``index`` is still derived from the
+    rendered position by the chat template. A run that cannot be fully matched is
+    left untouched. Re-running is idempotent.
+
+    This function is side-effect free: matched tool messages are shallow-copied
+    before their ``tool``/``name`` is rewritten, and every other message is
+    appended to the output as-is. The input list and its message objects are
+    never mutated.
+    """
+    if not isinstance(messages, list):
+        return messages
+
+    output: list[Any] = []
+    current_index: dict = {}
+    i = 0
+    n = len(messages)
+
+    while i < n:
+        message = messages[i]
+
+        if isinstance(message, dict) and message.get("role") == "assistant":
+            tool_calls = message.get("tool_calls")
+            current_index = _tool_call_id_index(tool_calls) if tool_calls else {}
+            output.append(message)
+            i += 1
+            continue
+
+        if not isinstance(message, dict) or message.get("role") != "tool":
+            output.append(message)
+            i += 1
+            continue
+
+        run: list[tuple] = []  # (position, original_offset, message, name)
+        unresolved = False
+        offset = 0
+        while (
+            i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool"
+        ):
+            tool_message = messages[i]
+            call_id = tool_message.get("tool_call_id", tool_message.get("id"))
+            matched = current_index.get(str(call_id)) if call_id is not None else None
+            if matched is None:
+                unresolved = True
+                run.append((None, offset, tool_message, None))
+            else:
+                position, name = matched
+                run.append((position, offset, tool_message, name))
+            offset += 1
+            i += 1
+
+        if unresolved:
+            output.extend(item[2] for item in run)
+        else:
+            run.sort(key=lambda item: (item[0], item[1]))
+            for _, _, tool_message, name in run:
+                if name is None:
+                    output.append(tool_message)
+                    continue
+                # The id-matched call is authoritative: align tool (and any
+                # explicit name) so the rendered XTML tool attribute cannot
+                # disagree with the reordered position. Copy first so the
+                # caller's message object is never mutated.
+                resolved = dict(tool_message)
+                resolved["tool"] = name
+                if "name" in resolved:
+                    resolved["name"] = name
+                output.append(resolved)
+
+    return output
+
+
+def is_batched_conversation(conversation: Any) -> bool:
+    return (
+        isinstance(conversation, list)
+        and bool(conversation)
+        and isinstance(conversation[0], list)
+    )
+
+
+def _render_content_segments(
+    content: Any,
+    image_state: _ImagePromptState,
+) -> list[EncodeSegment]:
+    segments: list[EncodeSegment] = []
+    if isinstance(content, str):
+        _append_text(segments, content, image_state)
+    elif content is not None:
+        for part in content:
+            if part["type"] in ["image", "image_url"]:
+                segments.extend(
+                    _segment(image_state.next_prompt(), allow_special=True))
+            else:
+                _append_text(segments, part["text"], image_state)
+    return segments
+
+
+def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]:
+    segments: list[EncodeSegment] = []
+    segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)]))
+    segments.extend(_text(body.strip()))
+    segments.extend(_close_tag("message"))
+    segments.extend(_end_of_msg())
+    return segments
+
+
+def _render_assistant_segments(
+    message: dict[str, Any],
+    image_state: _ImagePromptState,
+    thinking: bool = True,
+) -> list[EncodeSegment]:
+    segments: list[EncodeSegment] = []
+    # The <think> channel is structural: in thinking mode every assistant
+    # message carries the open/close tags even when there is no reasoning
+    # content to fill in. In non-thinking mode the channel is dropped
+    # entirely.
+    if thinking:
+        reasoning_content = message.get("reasoning_content") or message.get(
+            "reasoning"
+        )
+        segments.extend(_open_tag("think"))
+        if reasoning_content is not None and str(reasoning_content).strip():
+            _append_text(segments, reasoning_content, image_state)
+        segments.extend(_close_tag("think"))
+
+    segments.extend(_open_tag("response"))
+    segments.extend(_render_content_segments(message.get("content"), image_state))
+    segments.extend(_close_tag("response"))
+
+    tool_calls = message.get("tool_calls")
+    if tool_calls:
+        segments.extend(_open_tag("tools"))
+        for index, tool_call in enumerate(tool_calls, start=1):
+            fn = tool_call.get("function", tool_call)
+            segments.extend(
+                _open_tag("call", [("tool", fn["name"]), ("index", index)])
+            )
+            args = fn.get("arguments", {})
+            json_block = fn.get("_xtml_json_block")
+            if json_block is not None:
+                segments.extend(_open_tag("json", [("type", "object")]))
+                _append_text(segments, json_block, image_state)
+                segments.extend(_close_tag("json"))
+            elif _is_mapping(args):
+                for key, value in args.items():
+                    segments.extend(
+                        _open_tag(
+                            "argument",
+                            [("key", key), ("type", _xtml_type(value))],
+                        )
+                    )
+                    _append_text(segments, _xtml_value(value), image_state)
+                    segments.extend(_close_tag("argument"))
+            segments.extend(_close_tag("call"))
+        segments.extend(_close_tag("tools"))
+
+    return segments
+
+
+def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]:
+    if dynamic:
+        body = (
+            "## New Tools Available\n"
+            "The system dynamically extends the toolset via lazy-loading.\n"
+            "You have access to all existing and extended tools.\n"
+            "Here are the specs for the extended tools.\n\n"
+            "```json\n"
+            f"{_json_compact(tools)}\n"
+            "```"
+        )
+    else:
+        body = (
+            "# Tools\n"
+            "Here are the available tools, described in JSONSchema.\n\n"
+            "```json\n"
+            f"{_json_compact(tools)}\n"
+            "```"
+        )
+    segments: list[EncodeSegment] = []
+    segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")]))
+    segments.extend(_text(body))
+    segments.extend(_close_tag("message"))
+    segments.extend(_end_of_msg())
+    return segments
+
+
+def build_chat_segments(
+    messages: list[Any],
+    tools: Optional[list[dict]] = None,
+    *,
+    add_generation_prompt: bool = True,
+    thinking: bool = True,
+    image_prompts: Optional[list[str]] = None,
+    **kwargs: Any,
+) -> list[EncodeSegment]:
+    # Re-sort tool results by tool_call_id at the lowest layer so every caller
+    # (processor or direct tokenizer) gets correctly ordered XTML. The helper is
+    # side-effect free, so the caller's message objects are left untouched.
+    messages = normalize_xtml_tool_result_messages(messages)
+    messages = normalize_conversation(messages)
+    tools = deep_sort_dict(tools)
+
+    kwargs = dict(kwargs)
+    response_format = kwargs.get("response_format")
+    if "response_schema" not in kwargs:
+        response_schema = extract_response_schema(response_format)
+        if response_schema is not None:
+            kwargs["response_schema"] = response_schema
+    if kwargs.get("response_schema") is not None:
+        kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"])
+
+    image_state = _ImagePromptState(image_prompts)
+    segments: list[EncodeSegment] = []
+
+    tool_calls = None
+    tool_index = 0
+
+    if tools:
+        segments.extend(_render_tool_declare(tools))
+
+    thinking_effort = kwargs.get("thinking_effort")
+    if thinking and thinking_effort is not None:
+        assert thinking_effort in _VALID_THINKING_EFFORTS, (
+            f"Unsupported thinking_effort={thinking_effort!r}; "
+            f"supported values are {sorted(_VALID_THINKING_EFFORTS)}."
+        )
+    if thinking and thinking_effort in _VALID_THINKING_EFFORTS:
+        segments.extend(
+            _internal_system_message(
+                "thinking-effort",
+                "`thinking_effort` guides on how much to think in your "
+                "thinking channel (not including the response channel), "
+                "supported values include `low`, `medium`, `high`, and `max`.\n"
+                f"Now the system is invoked with `thinking_effort={thinking_effort}`.",
+            )
+        )
+
+    for message_index, message in enumerate(messages):
+        if not isinstance(message, dict):
+            continue
+
+        role = message["role"]
+        if role == "user":
+            attrs = [("role", "user")]
+            if message.get("name"):
+                attrs.append(("name", message["name"]))
+            segments.extend(_open_tag("message", attrs))
+            segments.extend(_render_content_segments(message.get("content"), image_state))
+            segments.extend(_close_tag("message"))
+            segments.extend(_end_of_msg())
+        elif role == "system" and message.get("tools"):
+            segments.extend(_render_tool_declare(message["tools"], dynamic=True))
+        elif role == "system":
+            attrs = [("role", "system")]
+            if message.get("name"):
+                attrs.append(("name", message["name"]))
+            segments.extend(_open_tag("message", attrs))
+            segments.extend(_render_content_segments(message.get("content"), image_state))
+            segments.extend(_close_tag("message"))
+            segments.extend(_end_of_msg())
+        elif role == "tool":
+            tool_index += 1
+            tool_name = message.get("tool", message.get("name"))
+            if (
+                tool_name is None
+                and tool_calls is not None
+                and tool_index <= len(tool_calls)
+            ):
+                tc = tool_calls[tool_index - 1]
+                fn = tc.get("function", tc)
+                tool_name = fn["name"]
+            if tool_name is None:
+                raise ValueError(
+                    "Kimi K3 tool messages need a resolvable tool name: "
+                    "carry `tool`/`name`, or match a preceding assistant "
+                    "tool_call by order."
+                )
+            segments.extend(
+                _open_tag(
+                    "message",
+                    [("role", "tool"), ("tool", tool_name), ("index", tool_index)],
+                )
+            )
+            segments.extend(_render_content_segments(message.get("content"), image_state))
+            segments.extend(_close_tag("message"))
+            segments.extend(_end_of_msg())
+        elif role == "assistant":
+            tool_calls = message.get("tool_calls")
+            tool_index = 0
+            attrs = [("role", "assistant")]
+            if message.get("name"):
+                attrs.append(("name", message["name"]))
+            segments.extend(_open_tag("message", attrs))
+            segments.extend(_render_assistant_segments(message, image_state, thinking))
+            segments.extend(_close_tag("message"))
+            segments.extend(_end_of_msg())
+
+    tool_choice = kwargs.get("tool_choice")
+    if tool_choice == "required":
+        segments.extend(
+            _internal_system_message(
+                "tool-choice",
+                "The system is invoked with `tool_choice=required`.\n"
+                "You MUST call tools in the next message.",
+            )
+        )
+    elif tool_choice == "none":
+        segments.extend(
+            _internal_system_message(
+                "tool-choice",
+                "The system is invoked with `tool_choice=none`.\n"
+                "You MUST NOT call any tools in the next message.",
+            )
+        )
+
+    rf = kwargs.get("response_format")
+    rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf
+    if rf_type == "json_object":
+        segments.extend(
+            _internal_system_message(
+                "response-format",
+                "The system is invoked with `response_format=json_object`.\n"
+                "Your response must be raw JSON data without markdown code "
+                "blocks (```json) or any additional formatting.",
+            )
+        )
+    elif rf_type == "json_schema":
+        schema = _json_compact(kwargs.get("response_schema"))
+        segments.extend(
+            _internal_system_message(
+                "response-format",
+                "The system is invoked with `response_format=json_schema`.\n"
+                "Your response must be raw JSON data without markdown code "
+                "blocks (```json) or any additional formatting.\n"
+                "The JSON data must match the following schema:\n"
+                f"```json\n{schema}\n```",
+            )
+        )
+
+    if add_generation_prompt:
+        segments.extend(_open_tag("message", [("role", "assistant")]))
+        segments.extend(_open_tag("think" if thinking else "response"))
+
+    image_state.assert_consumed()
+    return segments

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "max_length": 1048576,
+  "eos_token_id": 163586
+}

kimi_k3.py

@@ -0,0 +1,954 @@
+# MLX port of moonshotai/Kimi-K3 (2.78T-total / 104B-active multimodal MoE).
+#
+# Text tower only — the vision tower rides along in the checkpoint and is handled
+# separately (see scripts/convert.py, which passes it through untouched).
+#
+# Derived from mlx-lm's `kimi_linear.py` (Kimi-Linear-48B-A3B), which already
+# covers KDA + gated MLA + DeepSeek-style sparse MoE. K3 adds five things that
+# do not exist anywhere in MLX; each is marked [K3] below:
+#
+#   [K3-1] SiTU-GLU activation           beta=4.0, linear_beta=25.0
+#   [K3-2] Attention Residuals (AttnRes) rank-1 [1,H] score projections, block 12
+#   [K3-3] Stable LatentMoE              routed experts live in a 3584-d latent
+#   [K3-4] q-LoRA + output-gated MLA     q_a/q_b rank 1536, sigmoid g_proj
+#   [K3-5] per-channel KDA decay         A_log is [head_dim], not [num_heads]
+#
+# NOTE on [K3-5]: the reference `modeling_kimi_linear.py` allocates
+# `A_log = Parameter(log(empty(num_heads)))` — i.e. [96] — but every checkpoint
+# shard ships A_log as [128] == head_dim, while b_proj is [96, H] and dt_bias is
+# [12288] == 96*128. fla's kernel broadcasts A_log against g of shape
+# (B,T,96,128), so a length-128 vector can only align with the trailing head-dim
+# axis. The decay rate is therefore per-channel and shared across heads. The
+# reference init code is simply stale w.r.t. the released weights; the shapes are
+# authoritative. tests/test_kimi_k3.py asserts this against the real shard.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import ArraysCache, KVCache
+from .gated_delta import gated_delta_kernel, gated_delta_ops
+from .mla import MultiLinear
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    intermediate_size: int
+    rms_norm_eps: float
+    linear_attn_config: Dict[str, Any]
+    num_experts: int
+    moe_intermediate_size: int
+    kv_lora_rank: int
+    q_lora_rank: Optional[int] = None
+    tie_word_embeddings: bool = False
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None
+    v_head_dim: Optional[int] = None
+    mla_use_nope: bool = True
+    mla_use_output_gate: bool = False
+    num_experts_per_token: int = 1
+    num_shared_experts: int = 0
+    moe_router_activation_func: str = "sigmoid"
+    moe_renormalize: bool = True
+    routed_scaling_factor: float = 1.0
+    first_k_dense_replace: int = 0
+    moe_layer_freq: int = 1
+    num_expert_group: int = 1
+    topk_group: int = 1
+    max_position_embeddings: int = 1048576
+    # [K3-1]
+    hidden_act: str = "situ"
+    activation_situ_beta: float = 1.0
+    activation_situ_linear_beta: Optional[float] = None
+    # [K3-2]
+    attn_res_block_size: Optional[int] = None
+    # [K3-3]
+    routed_expert_hidden_size: Optional[int] = None
+    latent_moe_use_norm: bool = False
+    # REAP-pruned checkpoints keep a different number of experts per layer
+    # (0 for dense layers). None => every MoE layer has `num_experts`.
+    expert_counts: Optional[List[int]] = None
+    # Graded REAP: per layer, how many of that layer's experts sit in the
+    # high-precision bank. 0 (or None) => a single uniform-width bank.
+    expert_bank_split: Optional[List[int]] = None
+    # multimodal wrapper fields (ignored by the text tower)
+    vision_config: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, params: Dict[str, Any]) -> "ModelArgs":
+        # K3's config.json nests the language model under `text_config`.
+        if "text_config" in params:
+            merged = dict(params["text_config"])
+            merged["model_type"] = params.get("model_type", "kimi_k3")
+            merged["vision_config"] = params.get("vision_config", {})
+            for k in ("bos_token_id", "eos_token_id", "pad_token_id"):
+                merged.setdefault(k, params.get(k))
+            params = merged
+        return super().from_dict(params)
+
+
+# ---------------------------------------------------------------- [K3-1] SiTU
+
+
+class SituGLU(nn.Module):
+    """beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta).
+
+    Matches SituAndMul in the reference, including the float32 upcast — the
+    tanh saturation points (beta=4, linear_beta=25) are far enough out that
+    bf16 rounding visibly changes the result.
+    """
+
+    def __init__(self, beta: float = 1.0, linear_beta: Optional[float] = None):
+        super().__init__()
+        self.beta = beta
+        self.linear_beta = linear_beta
+
+    def __call__(self, x_up: mx.array, x_gate: mx.array) -> mx.array:
+        dtype = x_up.dtype
+        gate = x_gate.astype(mx.float32)
+        up = x_up.astype(mx.float32)
+        out = self.beta * mx.tanh(gate / self.beta) * mx.sigmoid(gate)
+        if self.linear_beta is not None:
+            up = self.linear_beta * mx.tanh(up / self.linear_beta)
+        return (out * up).astype(dtype)
+
+
+def _situ_from_args(args: ModelArgs) -> SituGLU:
+    return SituGLU(args.activation_situ_beta or 1.0, args.activation_situ_linear_beta)
+
+
+class KimiMLP(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        hidden_size: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+    ):
+        super().__init__()
+        h = hidden_size or args.hidden_size
+        i = intermediate_size or args.intermediate_size
+        self.gate_proj = nn.Linear(h, i, bias=False)
+        self.up_proj = nn.Linear(h, i, bias=False)
+        self.down_proj = nn.Linear(i, h, bias=False)
+        self.act = _situ_from_args(args)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(self.act(self.up_proj(x), self.gate_proj(x)))
+
+
+# ------------------------------------------------------------- [K3-2] AttnRes
+
+
+def _apply_attn_res(
+    prefix_sum: mx.array,
+    block_residual: mx.array,
+    proj_weight: mx.array,
+    norm_weight: mx.array,
+    eps: float,
+) -> mx.array:
+    """Softmax mixture over [block residuals..., prefix_sum].
+
+    prefix_sum:     (N, H)
+    block_residual: (N, K, H)  ->  returns (N, H)
+
+    Each candidate is RMS-normed, scored against a single learned direction
+    (norm.weight * proj.weight, both length H), and the candidates are combined
+    by a softmax over K+1. All in float32, as in the reference.
+    """
+    v = mx.concatenate([block_residual, mx.expand_dims(prefix_sum, 1)], axis=1)
+    vf = v.astype(mx.float32)
+    variance = mx.mean(vf * vf, axis=-1, keepdims=True)
+    k = vf * mx.rsqrt(variance + eps)
+    score_weight = norm_weight.astype(mx.float32) * proj_weight.reshape(-1).astype(
+        mx.float32
+    )
+    scores = mx.sum(k * score_weight, axis=-1)
+    probs = mx.expand_dims(mx.softmax(scores, axis=-1), 1)
+    return (probs @ vf).squeeze(1).astype(v.dtype)
+
+
+class AttnResProj(nn.Module):
+    """Holds the rank-1 [1, H] score projection and its paired RMSNorm gain.
+
+    Kept as plain arrays rather than nn.Linear/nn.RMSNorm so the quantizer skips
+    them — they are 7168 floats each and feed a softmax, so quantizing them is
+    all downside.
+    """
+
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.proj_weight = mx.zeros((1, hidden_size))
+        self.norm_weight = mx.ones((hidden_size,))
+
+    def __call__(
+        self, prefix_sum: mx.array, block_residual: mx.array, eps: float
+    ) -> mx.array:
+        return _apply_attn_res(
+            prefix_sum, block_residual, self.proj_weight, self.norm_weight, eps
+        )
+
+
+# ------------------------------------------------------ [K3-4] q-LoRA gated MLA
+
+
+class KimiMLAAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_heads = args.num_attention_heads
+        self.qk_nope_head_dim = args.qk_nope_head_dim or 128
+        self.qk_rope_head_dim = args.qk_rope_head_dim or 0
+        self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.v_head_dim = args.v_head_dim or 128
+        self.kv_lora_rank = args.kv_lora_rank
+        self.q_lora_rank = args.q_lora_rank
+        self.scale = self.q_head_dim**-0.5
+
+        h = args.hidden_size
+        if self.q_lora_rank is not None:
+            self.q_a_proj = nn.Linear(h, self.q_lora_rank, bias=False)
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=args.rms_norm_eps)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+            )
+        else:
+            self.q_proj = nn.Linear(h, self.num_heads * self.q_head_dim, bias=False)
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            h, self.kv_lora_rank + self.qk_rope_head_dim, bias=False
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=args.rms_norm_eps)
+
+        # kv_b_proj is absorbed into these two at load time (see sanitize), so the
+        # KV cache holds the 512-d latent instead of 96 full heads. At 1M context
+        # that is the difference between ~110 GB and ~1.2 GB of cache per layer.
+        self.embed_q = MultiLinear(self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads)
+        self.unembed_out = MultiLinear(self.kv_lora_rank, self.v_head_dim, self.num_heads)
+
+        self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, h, bias=False)
+
+        self.use_output_gate = args.mla_use_output_gate
+        if self.use_output_gate:
+            self.g_proj = nn.Linear(h, self.num_heads * self.v_head_dim, bias=False)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        if self.q_lora_rank is not None:
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+        else:
+            q = self.q_proj(x)
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv_latent = mx.expand_dims(self.kv_a_layernorm(compressed_kv), axis=1)
+
+        if cache is not None:
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
+
+        # NoPE: q_pe/k_pe carry no rotary embedding at all (the reference asserts
+        # use_nope and leaves rotary_emb=None). They are a plain extra 64 dims,
+        # folded in as an additive score term alongside the causal mask.
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask, pe_scores, mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype)
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        out = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
+        )
+        if L == 1:
+            out = self.unembed_out(out)
+
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        if self.use_output_gate:
+            out = out * mx.sigmoid(self.g_proj(x))
+        return self.o_proj(out)
+
+
+# ------------------------------------------------------------- [K3-5] KDA
+
+
+class ShortConv1d(nn.Module):
+    def __init__(self, channels: int, kernel_size: int):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.conv = nn.Conv1d(
+            in_channels=channels,
+            out_channels=channels,
+            kernel_size=kernel_size,
+            groups=channels,
+            padding=0,
+            bias=False,
+        )
+
+    def __call__(self, x, state, mask=None, lengths=None):
+        if mask is not None and state is not None:
+            x = x * mx.expand_dims(mask, -1) if mask.ndim == 2 else x
+        conv_input = mx.concatenate([state, x], axis=1)
+        new_state = conv_input[:, -(self.kernel_size - 1) :, :]
+        return nn.silu(self.conv(conv_input)), new_state
+
+
+class KimiDeltaAttention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        cfg = args.linear_attn_config
+        self.layer_idx = layer_idx
+        self.num_heads = cfg["num_heads"]
+        self.head_dim = cfg["head_dim"]
+        self.conv_kernel = cfg.get("short_conv_kernel_size", 4)
+        self.gate_lower_bound = cfg.get("gate_lower_bound", None)
+        self.use_full_rank_gate = cfg.get("use_full_rank_gate", False)
+
+        self.projection_dim = self.num_heads * self.head_dim
+        h = args.hidden_size
+        self.scale = float(self.head_dim) ** -0.5
+
+        self.q_proj = nn.Linear(h, self.projection_dim, bias=False)
+        self.k_proj = nn.Linear(h, self.projection_dim, bias=False)
+        self.v_proj = nn.Linear(h, self.projection_dim, bias=False)
+
+        self.q_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+        self.k_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+        self.v_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+
+        self.f_a_proj = nn.Linear(h, self.head_dim, bias=False)
+        self.f_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False)
+        self.b_proj = nn.Linear(h, self.num_heads, bias=False)
+
+        # [K3] full-rank output gate (use_full_rank_gate=True) replaces the
+        # low-rank g_a/g_b pair used by Kimi-Linear-48B.
+        if self.use_full_rank_gate:
+            self.g_proj = nn.Linear(h, self.projection_dim, bias=False)
+        else:
+            self.g_a_proj = nn.Linear(h, self.head_dim, bias=False)
+            self.g_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False)
+
+        # [K3-5] per-channel, shared across heads -> broadcast as (1, head_dim).
+        self.A_log = mx.zeros((self.head_dim,))
+        self.dt_bias = mx.zeros((self.projection_dim,))
+
+        self.o_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.o_proj = nn.Linear(self.projection_dim, h, bias=False)
+
+    def _compute_g(self, a: mx.array) -> mx.array:
+        # log-decay, then the safe-gate clamp (fla's `lower_bound`), then exp.
+        log_g = -mx.exp(self.A_log.reshape(1, self.head_dim).astype(mx.float32)) * nn.softplus(
+            a + self.dt_bias.reshape(self.num_heads, self.head_dim)
+        )
+        if self.gate_lower_bound is not None:
+            log_g = mx.maximum(log_g, self.gate_lower_bound)
+        return mx.exp(log_g)
+
+    def __call__(self, x, mask=None, cache=None):
+        B, T, _ = x.shape
+        dtype = x.dtype
+
+        if cache is not None:
+            q_state, k_state, v_state, ssm_state = cache
+            lengths = cache.lengths
+        else:
+            q_state = k_state = v_state = ssm_state = None
+            lengths = None
+
+        if q_state is None:
+            s = mx.zeros((B, self.conv_kernel - 1, self.projection_dim), dtype=dtype)
+            q_state = k_state = v_state = s
+
+        q_conv, q_state = self.q_conv(self.q_proj(x), q_state, mask, lengths)
+        k_conv, k_state = self.k_conv(self.k_proj(x), k_state, mask, lengths)
+        v_conv, v_state = self.v_conv(self.v_proj(x), v_state, mask, lengths)
+
+        if cache is not None:
+            cache[0], cache[1], cache[2] = q_state, k_state, v_state
+
+        q = q_conv.reshape(B, T, self.num_heads, self.head_dim)
+        k = k_conv.reshape(B, T, self.num_heads, self.head_dim)
+        v = v_conv.reshape(B, T, self.num_heads, self.head_dim)
+
+        # use_qk_l2norm_in_kernel=True in the reference
+        q = (self.scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+        k = self.scale * mx.fast.rms_norm(k, None, 1e-6)
+
+        a_logits = self.f_b_proj(self.f_a_proj(x)).reshape(
+            B, T, self.num_heads, self.head_dim
+        )
+        beta = mx.sigmoid(self.b_proj(x).reshape(B, T, self.num_heads))
+        g = self._compute_g(a_logits)
+
+        if ssm_state is None:
+            ssm_state = mx.zeros(
+                (B, self.num_heads, self.head_dim, self.head_dim), dtype=mx.float32
+            )
+        use_kernel = mx.default_device() == mx.gpu and mx.metal.is_available()
+        fn = gated_delta_kernel if use_kernel else gated_delta_ops
+        out, ssm_state = fn(q, k, v, g, beta, ssm_state, mask)
+
+        if cache is not None:
+            cache[3] = ssm_state
+            cache.advance(T)
+
+        if self.use_full_rank_gate:
+            gate = self.g_proj(x)
+        else:
+            gate = self.g_b_proj(self.g_a_proj(x))
+        gate = gate.reshape(B, T, self.num_heads, self.head_dim)
+        out = (
+            self.o_norm(out.reshape(B, T, self.num_heads, self.head_dim))
+            * mx.sigmoid(gate)
+        ).reshape(B, T, -1)
+        return self.o_proj(out)
+
+
+# --------------------------------------------------------- [K3-3] LatentMoE
+
+
+def _group_expert_select(
+    scores: mx.array,
+    bias: mx.array,
+    top_k: int,
+    num_expert_group: int,
+    topk_group: int,
+    routed_scaling_factor: float,
+    renormalize: bool,
+    activation: str,
+):
+    if activation == "sigmoid":
+        scores = mx.sigmoid(scores.astype(mx.float32))
+    else:
+        scores = mx.softmax(scores.astype(mx.float32), axis=-1)
+
+    orig = scores
+    scores = scores + bias
+    if num_expert_group > 1 and num_expert_group > topk_group:
+        k = scores.shape[-1] // num_expert_group
+        scores = mx.unflatten(scores, -1, (num_expert_group, k))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        gate = mx.argpartition(-group_scores, kth=topk_group - 1, axis=-2)
+        gate = mx.put_along_axis(
+            mx.ones_like(group_scores),
+            gate[..., :topk_group, :],
+            mx.zeros((1,), scores.dtype),
+            axis=-2,
+        )
+        scores = mx.flatten(scores - gate * mx.finfo(scores.dtype).max, -2, -1)
+
+    inds = mx.stop_gradient(mx.argpartition(-scores, kth=top_k - 1, axis=-1)[..., :top_k])
+    weights = mx.take_along_axis(orig, inds, axis=-1)
+    if top_k > 1 and renormalize:
+        weights = weights / (mx.sum(weights, axis=-1, keepdims=True) + 1e-20)
+    weights = weights * routed_scaling_factor
+    return inds, weights.astype(orig.dtype)
+
+
+class TwoBankSwitchGLU(nn.Module):
+    """SwitchGLU whose experts are split into two banks at different bit widths.
+
+    Experts [0, n_hi) live in `bank_hi`, [n_hi, n_hi+n_lo) in `bank_lo`. This
+    exists because MLX pins one bit width per expert tensor: QuantizedSwitchLinear
+    stores scalar bits/group_size/mode and gather_qmm takes them as scalars. Two
+    modules means two widths, which is what graded REAP needs -- most-salient
+    experts at mxfp4 (free on K3, it is the source encoding), the rest lower.
+
+    The naive way to do this costs 2x expert compute: run both banks over every
+    routed pair and select. Instead we exploit the sort SwitchGLU already
+    performs. After sorting (token, slot) pairs by expert index, a bank owning a
+    contiguous index range owns a contiguous *slice* of the sorted pairs, so
+    each bank does exactly its own share of the work and no more.
+
+    Cost: the split point is data-dependent, so it must be read back to the
+    host to slice on. Measured at K3 REAP dims (3584->3072, 240 experts,
+    top-16), that sync is essentially the *entire* overhead -- the second
+    dispatch is nearly free:
+
+        prefill (256 tok)   1.05x vs single-bank
+        decode  (1 tok)     1.16x
+
+    Decode routes only top_k pairs, too few to amortise two kernel launches
+    plus a device sync, so below PARTITION_ON_HOST pairs we pull the indices
+    down in one transfer and partition with numpy instead. That alone took
+    decode from 1.44x to 1.16x.
+    """
+
+    # Below this many (token, slot) pairs, partition on the host.
+    PARTITION_ON_HOST = 1024
+
+    def __init__(self, input_dims, hidden_dims, n_hi, n_lo, activation):
+        super().__init__()
+        if n_hi <= 0 or n_lo <= 0:
+            raise ValueError(f"both banks must be non-empty, got {n_hi}/{n_lo}")
+        self.n_hi, self.n_lo = n_hi, n_lo
+        self.bank_hi = SwitchGLU(input_dims, hidden_dims, n_hi, activation=activation)
+        self.bank_lo = SwitchGLU(input_dims, hidden_dims, n_lo, activation=activation)
+
+    @property
+    def num_experts(self):
+        return self.n_hi + self.n_lo
+
+    @staticmethod
+    def _run(bank, xp, idx):
+        """One bank over an already-sorted slice of pairs."""
+        up = bank.up_proj(xp, idx, sorted_indices=True)
+        gate = bank.gate_proj(xp, idx, sorted_indices=True)
+        return bank.down_proj(bank.activation(up, gate), idx, sorted_indices=True)
+
+    def __call__(self, x: mx.array, indices: mx.array) -> mx.array:
+        lead = indices.shape[:-1]
+        K = indices.shape[-1]
+        flat = indices.reshape(-1)
+
+        total = flat.size
+        if total <= self.PARTITION_ON_HOST:
+            # one host transfer beats argsort-kernel + sum-kernel + sync
+            host = np.array(flat.tolist(), dtype=np.int32)
+            order = mx.array(np.argsort(host, kind="stable"))
+            split = int((host < self.n_hi).sum())
+        else:
+            order = mx.argsort(flat)
+            split = int(mx.sum(flat < self.n_hi))
+
+        inv_order = mx.argsort(order)
+        sorted_idx = flat[order]
+        # one row per (token, slot) pair, in expert order
+        pairs = mx.expand_dims(x, (-2, -3)).flatten(0, -3)[order // K]
+
+        if split == 0:
+            out = self._run(self.bank_lo, pairs, sorted_idx - self.n_hi)
+        elif split == total:
+            out = self._run(self.bank_hi, pairs, sorted_idx)
+        else:
+            out = mx.concatenate(
+                [
+                    self._run(self.bank_hi, pairs[:split], sorted_idx[:split]),
+                    self._run(self.bank_lo, pairs[split:], sorted_idx[split:] - self.n_hi),
+                ],
+                axis=0,
+            )
+
+        return mx.unflatten(out[inv_order], 0, (*lead, K)).squeeze(-2)
+
+
+class KimiSparseMoE(nn.Module):
+    """LatentMoE: 896 experts run in a 3584-d latent, not the 7168-d residual.
+
+    hidden -> down_proj -> [top-16 of 896 experts] -> norm -> up_proj -> hidden
+    plus 2 shared experts applied to the *undown-projected* input.
+    """
+
+    def __init__(
+        self,
+        args: ModelArgs,
+        num_experts: Optional[int] = None,
+        bank_split: int = 0,
+    ):
+        super().__init__()
+        self.args = args
+        h = args.hidden_size
+        self.num_experts = num_experts if num_experts is not None else args.num_experts
+        self.bank_split = bank_split
+        self.use_latent = args.routed_expert_hidden_size is not None
+        self.moe_hidden = args.routed_expert_hidden_size if self.use_latent else h
+        self.use_norm = args.latent_moe_use_norm
+
+        self.gate = nn.Linear(h, self.num_experts, bias=False)
+        self.e_score_correction_bias = mx.zeros((self.num_experts,), dtype=mx.float32)
+        if bank_split:
+            self.switch_mlp = TwoBankSwitchGLU(
+                self.moe_hidden,
+                args.moe_intermediate_size,
+                bank_split,
+                self.num_experts - bank_split,
+                activation=_situ_from_args(args),
+            )
+        else:
+            self.switch_mlp = SwitchGLU(
+                self.moe_hidden,
+                args.moe_intermediate_size,
+                self.num_experts,
+                activation=_situ_from_args(args),
+            )
+
+        if self.use_latent:
+            self.routed_expert_down_proj = nn.Linear(h, self.moe_hidden, bias=False)
+            self.routed_expert_up_proj = nn.Linear(self.moe_hidden, h, bias=False)
+            if self.use_norm:
+                self.routed_expert_norm = nn.RMSNorm(
+                    self.moe_hidden, eps=args.rms_norm_eps
+                )
+
+        if args.num_shared_experts:
+            self.shared_experts = KimiMLP(
+                args, intermediate_size=args.moe_intermediate_size * args.num_shared_experts
+            )
+        else:
+            self.shared_experts = None
+
+        # Optional REAP saliency tap (scripts/reap_calibrate.py). Left as None in
+        # normal use and costs nothing. It exists so the calibration harness reads
+        # gates and per-expert outputs from *this* forward pass rather than a
+        # reimplementation that could drift from it. MLX keeps non-array
+        # attributes out of parameters()/children(), so this never touches
+        # quantization, saving or loading.
+        self.expert_stats_hook = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        identity = x
+        inds, weights = _group_expert_select(
+            self.gate(x),
+            self.e_score_correction_bias,
+            min(self.args.num_experts_per_token, self.num_experts),
+            self.args.num_expert_group,
+            self.args.topk_group,
+            self.args.routed_scaling_factor,
+            self.args.moe_renormalize,
+            self.args.moe_router_activation_func,
+        )
+        y = self.routed_expert_down_proj(x) if self.use_latent else x
+        y = self.switch_mlp(y, inds)
+        if self.expert_stats_hook is not None:
+            # y here is the per-expert output BEFORE gate weighting, shape
+            # (..., top_k, moe_hidden) -- exactly what REAP saliency needs.
+            self.expert_stats_hook(inds, weights, y)
+        y = (y * weights[..., None]).sum(axis=-2)
+        if self.use_latent:
+            if self.use_norm:
+                y = self.routed_expert_norm(y)
+            y = self.routed_expert_up_proj(y)
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(identity)
+        return y
+
+
+# ------------------------------------------------- checkpoint key remapping
+#
+# Single source of truth, shared by Model.sanitize (whole-checkpoint, used when
+# mlx-lm loads a converted repo) and scripts/convert.py (per-layer, used when
+# streaming the 1.56 TB source). Keeping one implementation is the point: the
+# streaming converter never instantiates a Model, so a second copy of these
+# rules would be free to drift out of sync and silently mis-place weights.
+
+
+def is_kda_layer(args: ModelArgs, idx: int) -> bool:
+    # config lists kda_layers/full_attn_layers 1-BASED: config n -> layer n-1.
+    return (idx + 1) in args.linear_attn_config["kda_layers"]
+
+
+def is_moe_layer(args: ModelArgs, idx: int) -> bool:
+    return (
+        args.num_experts > 0
+        and idx >= args.first_k_dense_replace
+        and idx % args.moe_layer_freq == 0
+    )
+
+
+def experts_in_layer(args: ModelArgs, idx: int) -> int:
+    """Expert count for one layer — varies after REAP pruning."""
+    if args.expert_counts:
+        return int(args.expert_counts[idx])
+    return args.num_experts
+
+
+def bank_split_in_layer(args: ModelArgs, idx: int) -> int:
+    """How many of a layer's experts sit in the high-precision bank (0 = one bank)."""
+    if not args.expert_bank_split:
+        return 0
+    return int(args.expert_bank_split[idx])
+
+
+def remap_checkpoint(
+    weights: Dict[str, mx.array],
+    args: ModelArgs,
+    layer_indices: Optional[List[int]] = None,
+    stack_experts: bool = True,
+) -> Dict[str, mx.array]:
+    """Map Kimi-K3 checkpoint keys onto this module tree.
+
+    `layer_indices` restricts the layer loop (the converter passes the single
+    layer it is holding). `stack_experts=False` leaves per-expert tensors alone
+    so the converter can quantize each expert before stacking, which keeps peak
+    memory at ~6 GB per layer instead of ~59 GB.
+    """
+    out: Dict[str, mx.array] = {}
+    for k, v in weights.items():
+        # K3 wraps the text tower as `language_model.*`; the vision tower and
+        # projector are handled outside mlx-lm.
+        if k.startswith("vision_tower.") or k.startswith("mm_projector."):
+            continue
+        out[k[len("language_model.") :] if k.startswith("language_model.") else k] = v
+    weights = out
+
+    if args.tie_word_embeddings:
+        weights.pop("lm_head.weight", None)
+
+    if layer_indices is None:
+        layer_indices = range(args.num_hidden_layers)
+
+    for i in layer_indices:
+        p = f"model.layers.{i}"
+
+        # [K3-2] AttnRes: fold `*_res_proj` / `*_res_norm` into the AttnResProj
+        # containers so the quantizer leaves them alone.
+        for src in ("self_attention_res", "mlp_res"):
+            for suffix, field_name in (("proj", "proj_weight"), ("norm", "norm_weight")):
+                key = f"{p}.{src}_{suffix}.weight"
+                if key in weights:
+                    weights[f"{p}.{src}.{field_name}"] = weights.pop(key)
+
+        if is_moe_layer(args, i):
+            src = f"{p}.block_sparse_moe"
+            if stack_experts:
+                for w, dst in (("w1", "gate_proj"), ("w3", "up_proj"), ("w2", "down_proj")):
+                    if f"{src}.experts.0.{w}.weight" in weights:
+                        weights[f"{src}.switch_mlp.{dst}.weight"] = mx.stack(
+                            [
+                                weights.pop(f"{src}.experts.{e}.{w}.weight")
+                                for e in range(experts_in_layer(args, i))
+                            ]
+                        )
+            bias_key = f"{src}.gate.e_score_correction_bias"
+            if bias_key in weights:
+                weights[f"{src}.e_score_correction_bias"] = weights.pop(bias_key)
+
+        attn = f"{p}.self_attn"
+        if is_kda_layer(args, i):
+            for src, dst in (("q_conv1d", "q_conv"), ("k_conv1d", "k_conv"), ("v_conv1d", "v_conv")):
+                key = f"{attn}.{src}.weight"
+                if key in weights:
+                    w = weights.pop(key)
+                    if w.ndim == 3:
+                        w = w.moveaxis(2, 1)
+                    weights[f"{attn}.{dst}.conv.weight"] = w
+            dt = f"{attn}.dt_bias"
+            if dt in weights and weights[dt].ndim > 1:
+                weights[dt] = weights[dt].reshape(-1)
+        else:
+            # [K3-4] absorb kv_b_proj -> embed_q / unembed_out so the KV cache
+            # stores the 512-d latent rather than 96 materialised heads.
+            kv_b = f"{attn}.kv_b_proj.weight"
+            if kv_b in weights:
+                qk_nope = args.qk_nope_head_dim or 128
+                v_head = args.v_head_dim or 128
+                v = weights.pop(kv_b).reshape(args.num_attention_heads, qk_nope + v_head, -1)
+                weights[f"{attn}.embed_q.weight"] = mx.contiguous(
+                    v[:, :qk_nope, :].swapaxes(-1, -2)
+                )
+                weights[f"{attn}.unembed_out.weight"] = mx.contiguous(v[:, qk_nope:, :])
+
+    for suffix, field_name in (("proj", "proj_weight"), ("norm", "norm_weight")):
+        key = f"model.output_attn_res_{suffix}.weight"
+        if key in weights:
+            weights[f"model.output_attn_res.{field_name}"] = weights.pop(key)
+
+    return weights
+
+
+# ------------------------------------------------------------------- layers
+
+
+class KimiDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.args = args
+        self.layer_idx = layer_idx
+        # config lists full_attn_layers/kda_layers 1-BASED: config n -> layer n-1.
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.is_linear = (layer_idx + 1) in kda_layers
+
+        if self.is_linear:
+            self.self_attn = KimiDeltaAttention(args, layer_idx)
+        else:
+            self.self_attn = KimiMLAAttention(args)
+
+        is_moe = (
+            args.num_experts > 0
+            and layer_idx >= args.first_k_dense_replace
+            and layer_idx % args.moe_layer_freq == 0
+        )
+        if is_moe:
+            self.block_sparse_moe = KimiSparseMoE(
+                args,
+                experts_in_layer(args, layer_idx),
+                bank_split_in_layer(args, layer_idx),
+            )
+        else:
+            self.mlp = KimiMLP(args)
+        self.is_moe = is_moe
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+        self.use_attn_res = args.attn_res_block_size is not None
+        if self.use_attn_res:
+            self.attn_res_block_size = args.attn_res_block_size
+            self.self_attention_res = AttnResProj(args.hidden_size)
+            self.mlp_res = AttnResProj(args.hidden_size)
+
+    def _ffn(self, x: mx.array) -> mx.array:
+        return self.block_sparse_moe(x) if self.is_moe else self.mlp(x)
+
+    def __call__(self, x, mask=None, cache=None, block_residual=None):
+        if not self.use_attn_res:
+            h = x + self.self_attn(self.input_layernorm(x), mask, cache)
+            return h + self._ffn(self.post_attention_layernorm(h)), block_residual
+
+        B, T, H = x.shape
+        eps = self.args.rms_norm_eps
+        prefix_sum = x
+
+        if block_residual is not None and block_residual.shape[1] > 0:
+            x = self.self_attention_res(
+                prefix_sum.reshape(-1, H), block_residual, eps
+            ).reshape(B, T, H)
+
+        pushed = self.layer_idx % self.attn_res_block_size == 0
+        if pushed:
+            block_residual = mx.concatenate(
+                [block_residual, mx.expand_dims(prefix_sum.reshape(-1, H), 1)], axis=1
+            )
+
+        y = self.self_attn(self.input_layernorm(x), mask, cache)
+        prefix_sum = y if pushed else prefix_sum + y
+
+        x = self.mlp_res(prefix_sum.reshape(-1, H), block_residual, eps).reshape(B, T, H)
+        z = self._ffn(self.post_attention_layernorm(x))
+        return prefix_sum + z, block_residual
+
+
+class KimiK3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [KimiDecoderLayer(args, i) for i in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.use_attn_res = args.attn_res_block_size is not None
+        if self.use_attn_res:
+            self.output_attn_res = AttnResProj(args.hidden_size)
+
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.ssm_idx = kda_layers[0] - 1
+        self.attn_idx = next(
+            i for i in range(len(self.layers)) if (i + 1) not in kda_layers
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+        inputs_embeds: Optional[mx.array] = None,
+    ) -> mx.array:
+        # inputs_embeds lets the multimodal wrapper splice image features in
+        # place of the <|media_pad|> placeholder before the stack runs.
+        h = self.embed_tokens(inputs) if inputs_embeds is None else inputs_embeds
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
+        attn_mask = create_attention_mask(h, cache[self.attn_idx], return_array=True)
+
+        B, T, H = h.shape
+        block_residual = mx.zeros((B * T, 0, H), dtype=h.dtype) if self.use_attn_res else None
+
+        for layer, layer_cache in zip(self.layers, cache):
+            mask = ssm_mask if layer.is_linear else attn_mask
+            h, block_residual = layer(h, mask, layer_cache, block_residual)
+
+        if self.use_attn_res:
+            h = self.output_attn_res(
+                h.reshape(-1, H), block_residual, self.args.rms_norm_eps
+            ).reshape(B, T, H)
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = KimiK3Model(args)
+        self.lm_head = (
+            None
+            if args.tie_word_embeddings
+            else nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+        inputs_embeds: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, inputs_embeds=inputs_embeds)
+        if self.lm_head is None:
+            return self.model.embed_tokens.as_linear(out)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            ArraysCache(size=4) if layer.is_linear else KVCache()
+            for layer in self.layers
+        ]
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        return remap_checkpoint(weights, self.args)
+
+    @property
+    def cast_predicate(self):
+        def predicate(path: str):
+            if "e_score_correction_bias" in path:
+                return False
+            if path.endswith("A_log") or path.endswith("dt_bias"):
+                return False
+            return True
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            # router gate drives top-16-of-896 selection; 8 bits is cheap insurance
+            if path.endswith("block_sparse_moe.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

kimi_k3_processor.py

@@ -0,0 +1,187 @@
+"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface.
+
+Chat rendering (including XTML tool-result ordering) is handled by the
+tokenizer's Python encoder; this processor adds multimodal media preprocessing.
+"""
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+from transformers.utils import logging
+
+from .media_utils import ensure_media_type
+
+logger = logging.get_logger(__name__)
+
+# ── KimiK3Processor ───────────────────────────────────────────────────
+
+
+class KimiK3Processor(ProcessorMixin):
+    r"""
+    Constructs a KimiK3 processor which wraps a KimiK3 image processor
+    and a tokenizer into a single processor.
+
+    [`KimiK3Processor`] offers all the functionalities of
+    [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`].
+
+    Args:
+        image_processor ([`KimiK3VisionProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`TikTokenTokenizer`], *optional*):
+            The tokenizer is a required input.
+        chat_template (`str`, *optional*): Kept for ProcessorMixin
+            compatibility. Kimi K3 chat encoding is implemented in Python by
+            the tokenizer.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = ["chat_template"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        **kwargs,
+    ):
+        super().__init__(image_processor,
+                         tokenizer,
+                         chat_template=chat_template)
+        self.media_processor = image_processor
+        self.image_placeholder = "<|kimi_image_placeholder|>"
+
+    # ── Media preprocessing ────────────────────────────────────────────
+
+    def update_raw_text(self, text: str, image_prompts: list[str]) -> str:
+        # Replace image placeholders
+        image_count = text.count(self.image_placeholder)
+        if image_count > 0:
+            assert image_count == len(image_prompts), (
+                f"image placeholder count {image_count} != "
+                f"image_prompts count {len(image_prompts)}")
+            text_parts = text.split(self.image_placeholder)
+            assert len(text_parts) == len(image_prompts) + 1
+            text = "".join([
+                text_parts[i] + image_prompts[i]
+                for i in range(len(image_prompts))
+            ])
+            text += text_parts[-1]
+
+        return text
+
+    def preprocess_medias(self,
+                          medias: list[dict]) -> tuple[list[dict], list[str]]:
+        """Process media items and generate corresponding prompts.
+
+        Returns:
+            A tuple of (updated_medias, image_prompts).
+        """
+        updated_medias = []
+        image_prompts = []
+        for media in medias:
+            if media['type'] == 'image':
+                updated_medias.append(media)
+                img = ensure_media_type(
+                    media,
+                    transparent_bg_config=self.media_processor.
+                    _transparent_bg_config,
+                    transparent_bg_fill_stage=self.media_processor.
+                    _transparent_bg_fill_stage,
+                )['image']
+                w, h = img.size
+                image_prompts.append(
+                    self.media_processor.make_image_prompt(w, h))
+            else:
+                raise ValueError(f"unsupported media type: {media['type']}")
+        return updated_medias, image_prompts
+
+    # ── Main entry points ──────────────────────────────────────────────
+
+    def __call__(self,
+                 messages: list[dict] = None,
+                 medias: list[dict] = None,
+                 text: str = None,
+                 return_tensors: str = "pt",
+                 **kwargs) -> BatchFeature:
+        """
+        Process multimodal inputs for Kimi-K3 model.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content' fields.
+                     If provided, medias and text will be extracted automatically.
+            medias: Pre-extracted list of media dicts.
+            text: Pre-formatted text string.
+            return_tensors: Format of returned tensors. Default: 'pt'.
+            **kwargs: Additional arguments passed to apply_chat_template.
+
+        Returns:
+            BatchFeature with fields: input_ids, attention_mask,
+            pixel_values, grid_thws.
+        """
+        if messages is None and (medias is None or text is None):
+            raise ValueError(
+                "Provide either 'messages' or both 'medias' and 'text'")
+
+        if medias is not None and text is not None:
+            updated_medias, image_prompts = (self.preprocess_medias(medias))
+            preprocessed = self.media_processor.preprocess(
+                updated_medias, return_tensors=return_tensors)
+            text = self.update_raw_text(text, image_prompts)
+            text_inputs = self.tokenizer(text, return_tensors=return_tensors)
+            return BatchFeature(data={**text_inputs, **preprocessed.data})
+
+        if medias is None:
+            medias = self._extract_medias_from_messages(messages)
+        updated_medias, image_prompts = (self.preprocess_medias(medias))
+        preprocessed = self.media_processor.preprocess(
+            updated_medias, return_tensors=return_tensors)
+
+        if text is None:
+            text_inputs = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                return_tensors=return_tensors,
+                return_dict=True,
+                image_prompts=image_prompts,
+                **kwargs)
+            return BatchFeature(data={**text_inputs, **preprocessed.data})
+
+        text = self.update_raw_text(text, image_prompts)
+        text_inputs = self.tokenizer(text, return_tensors=return_tensors)
+        return BatchFeature(data={**text_inputs, **preprocessed.data})
+
+    @staticmethod
+    def _extract_medias_from_messages(messages: list[dict]) -> list[dict]:
+        """Extract media items from messages in a single pass."""
+        medias = []
+        for msg in messages:
+            if msg['role'] != 'user' or not msg.get('content'):
+                continue
+
+            for content_part in msg['content']:
+                if not isinstance(content_part, dict):
+                    continue
+
+                content_type = content_part.get('type')
+                if content_type in ['image_url', 'image']:
+                    image_data = content_part.get(content_type)
+                    assert image_data is not None, f"image data is missing for content part: {content_part}"
+                    medias.append({
+                        'type': 'image',
+                        'image': image_data,
+                    })
+        return medias
+
+    def apply_chat_template(self, messages, **kwargs):
+        return self.tokenizer.apply_chat_template(messages, **kwargs)
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        return ['input_ids', 'attention_mask', 'pixel_values', 'grid_thws']

kimi_k3_vision.py

@@ -0,0 +1,404 @@
+# MLX port of Kimi-K3's vision tower — a 3D (video-capable) MoonViT variant plus
+# the patchmergerv2 projector.
+#
+# 0.49 B params, 0.02% of the model, and the only reason K3 is multimodal.
+#
+# Closest existing MLX code is mlx-vlm's `kimi_vl/vision.py` (MoonViT for
+# Kimi-VL). This shares its complex64 2D-rope approach but differs in six ways,
+# each marked [K3-V*] below:
+#
+#   [K3-V1] 3D grids (t,h,w): a 1D sin-cos *time* embedding is added on top of
+#           the learnable 2D grid, and rope frequencies repeat per frame.
+#   [K3-V2] bilinear position-embedding interpolation (Kimi-VL used bicubic);
+#           mlx-vlm ships bicubic/nearest kernels only, so it is written here.
+#   [K3-V3] qkv_hidden_size 1536 != hidden 1024 — attention runs wider than the
+#           residual stream, so wo is [1024, 1536] and is NOT square.
+#   [K3-V4] RMSNorm, not LayerNorm.
+#   [K3-V5] no biases anywhere (attn_bias / linear_bias / patch_embed_proj_bias
+#           are all false).
+#   [K3-V6] sd2_tpool merge: mean-pool over frames, then 2x2 spatial merge, then
+#           a 2-layer MLP projector with *exact* (erf) GELU and a post-RMSNorm.
+#
+# The ViT blocks use tanh-approximate GELU (`PytorchGELUTanh` in the reference)
+# while the projector uses exact `nn.GELU()`. They are different functions and
+# the reference deliberately uses both; do not unify them.
+
+import math
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+# RMSNorm eps for the ViT blocks. The reference constructs `nn.RMSNorm(dim)`
+# with no eps, so PyTorch falls back to finfo(dtype).eps — 1.19e-7 in fp32 but
+# 7.8e-3 in bf16, which are materially different normalizers. We use the fp32
+# value, matching an upcast implementation; `projector_ln_eps` (1e-5) is given
+# explicitly by the config and is used only for the projector's post_norm.
+VIT_NORM_EPS = 1.1920929e-07
+
+
+@dataclass
+class VisionConfig:
+    vt_hidden_size: int = 1024
+    vt_intermediate_size: int = 4096
+    vt_num_hidden_layers: int = 27
+    vt_num_attention_heads: int = 12
+    qkv_hidden_size: int = 1536
+    patch_size: int = 14
+    num_channels: int = 3
+    init_pos_emb_height: int = 64
+    init_pos_emb_width: int = 64
+    init_pos_emb_time: int = 4
+    pos_emb_type: str = "divided_fixed"
+    pos_emb_interpolation_mode: str = "bilinear"
+    merge_kernel_size: Tuple[int, int] = (2, 2)
+    merge_type: str = "sd2_tpool"
+    norm_type: str = "rmsnorm"
+    mlp_type: str = "mlp2"
+    mm_projector_type: str = "patchmergerv2"
+    mm_hidden_size: int = 1024
+    text_hidden_size: int = 7168
+    projector_hidden_act: str = "gelu"
+    projector_ln_eps: float = 1e-5
+    attn_bias: bool = False
+    linear_bias: bool = False
+    patch_embed_proj_bias: bool = False
+    rope_theta_base: float = 10000.0
+    rope_max_size: int = 512
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "VisionConfig":
+        known = {f for f in cls.__dataclass_fields__}
+        out = {k: v for k, v in d.items() if k in known}
+        if "merge_kernel_size" in out:
+            out["merge_kernel_size"] = tuple(out["merge_kernel_size"])
+        return cls(**out)
+
+    @property
+    def head_dim(self) -> int:
+        return self.qkv_hidden_size // self.vt_num_attention_heads
+
+
+# ------------------------------------------------------- [K3-V2] bilinear
+
+
+def bilinear_interpolate(x: mx.array, out_h: int, out_w: int) -> mx.array:
+    """Resize (H, W, C) -> (out_h, out_w, C), matching torch's bilinear
+    `F.interpolate(..., align_corners=False)`.
+
+    torch maps destination index d to source coordinate
+    `max(0, scale * (d + 0.5) - 0.5)` and lerps between the two neighbours,
+    clamping the upper index at the edge. Getting the half-pixel offset wrong
+    shifts every position embedding by half a patch, which is the kind of error
+    that degrades quality without ever looking obviously broken.
+    """
+    in_h, in_w = x.shape[0], x.shape[1]
+    if (in_h, in_w) == (out_h, out_w):
+        return x
+
+    def axis_index(in_n: int, out_n: int):
+        scale = in_n / out_n
+        src = mx.maximum(
+            (mx.arange(out_n, dtype=mx.float32) + 0.5) * scale - 0.5,
+            mx.array(0.0, mx.float32),
+        )
+        i0 = mx.floor(src)
+        w1 = src - i0
+        i0 = i0.astype(mx.int32)
+        i1 = mx.minimum(i0 + 1, mx.array(in_n - 1, mx.int32))
+        return i0, i1, w1
+
+    h0, h1, wh = axis_index(in_h, out_h)
+    w0, w1i, ww = axis_index(in_w, out_w)
+
+    xf = x.astype(mx.float32)
+    top = mx.take(xf, h0, axis=0)      # (out_h, in_w, C)
+    bot = mx.take(xf, h1, axis=0)
+    wh = wh.reshape(out_h, 1, 1)
+    rows = top * (1.0 - wh) + bot * wh
+
+    left = mx.take(rows, w0, axis=1)   # (out_h, out_w, C)
+    right = mx.take(rows, w1i, axis=1)
+    ww = ww.reshape(1, out_w, 1)
+    return left * (1.0 - ww) + right * ww
+
+
+def sincos_1d(dim: int, length: int) -> mx.array:
+    """(length, dim) sin-cos table; matches get_1d_sincos_pos_embed."""
+    omega = mx.arange(dim // 2, dtype=mx.float32) / (dim / 2.0)
+    omega = 1.0 / (10000.0**omega)
+    out = mx.outer(mx.arange(length, dtype=mx.float32), omega)
+    return mx.concatenate([mx.sin(out), mx.cos(out)], axis=1)
+
+
+# ------------------------------------------------ [K3-V1] 3D position embed
+
+
+class Learnable2DInterpPosEmbDivided(nn.Module):
+    """Learnable 2D grid, bilinearly resized per image, plus a sin-cos time term."""
+
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.cfg = cfg
+        # reference does nn.init.normal_(weight); a zero init would make an
+        # unloaded tower emit exactly zero, which hides a failed weight load
+        self.weight = mx.random.normal(
+            (cfg.init_pos_emb_height, cfg.init_pos_emb_width, cfg.vt_hidden_size)
+        )
+        self._time = sincos_1d(cfg.vt_hidden_size, cfg.init_pos_emb_time)
+
+    def __call__(self, x: mx.array, grid_thws: List[Tuple[int, int, int]]) -> mx.array:
+        embs = []
+        for t, h, w in grid_thws:
+            if (h, w) == self.weight.shape[:2]:
+                emb2d = self.weight.reshape(-1, self.weight.shape[-1])
+            else:
+                emb2d = bilinear_interpolate(self.weight, h, w).reshape(
+                    -1, self.weight.shape[-1]
+                )
+            if t == 1:
+                emb = emb2d
+            else:
+                emb = mx.expand_dims(emb2d, 0) + mx.expand_dims(self._time[:t], 1)
+                emb = emb.reshape(-1, emb.shape[-1])
+            embs.append(emb)
+        return x + mx.concatenate(embs, axis=0).astype(x.dtype)
+
+
+class MoonVision3dPatchEmbed(nn.Module):
+    """Patch projection. The reference is a Conv2d whose kernel equals the patch
+    size applied to pre-extracted (L, 3, 14, 14) patches, which is exactly a
+    linear map over the flattened patch — done as a matmul here to sidestep
+    MLX's NHWC conv weight layout entirely."""
+
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.cfg = cfg
+        fan_in = cfg.num_channels * cfg.patch_size * cfg.patch_size
+        # torch Conv2d default: uniform(-1/sqrt(fan_in), 1/sqrt(fan_in))
+        bound = 1.0 / math.sqrt(fan_in)
+        self.proj_weight = mx.random.uniform(
+            low=-bound, high=bound, shape=(cfg.vt_hidden_size, fan_in)
+        )
+        self.pos_emb = Learnable2DInterpPosEmbDivided(cfg)
+
+    def __call__(self, patches: mx.array, grid_thws) -> mx.array:
+        x = patches.reshape(patches.shape[0], -1).astype(self.proj_weight.dtype)
+        x = x @ self.proj_weight.T
+        return self.pos_emb(x, grid_thws)
+
+
+# --------------------------------------------------------- [K3-V1] 2D rope
+
+
+class Rope2DPosEmbRepeated(nn.Module):
+    """freqs_cis[h, w] interleaves x (width) and y (height) frequencies:
+    entry 2i is cis(w * theta^(-4i/dim)), entry 2i+1 is cis(h * ...).
+    Repeated t times so every frame shares the same spatial rope."""
+
+    def __init__(self, dim: int, max_h: int, max_w: int, theta_base: float = 10000.0):
+        super().__init__()
+        assert dim % 4 == 0, "rope dim must be divisible by 4"
+        self.dim, self.max_h, self.max_w, self.theta_base = dim, max_h, max_w, theta_base
+        self._cache: Optional[mx.array] = None
+
+    def _full(self) -> mx.array:
+        if self._cache is None:
+            n = self.max_h * self.max_w
+            flat = mx.arange(0, n, dtype=mx.float32)
+            x_pos = flat % self.max_w
+            y_pos = mx.floor(flat / self.max_w)
+            dim_range = mx.arange(0, self.dim, 4)[: self.dim // 4].astype(mx.float32)
+            freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))
+            xf, yf = mx.outer(x_pos, freqs), mx.outer(y_pos, freqs)
+            x_cis = mx.cos(xf) + 1j * mx.sin(xf)
+            y_cis = mx.cos(yf) + 1j * mx.sin(yf)
+            self._cache = mx.stack([x_cis, y_cis], axis=-1).reshape(
+                self.max_h, self.max_w, -1
+            )
+        return self._cache
+
+    def __call__(self, grid_thws: List[Tuple[int, int, int]]) -> mx.array:
+        full = self._full()
+        parts = []
+        for t, h, w in grid_thws:
+            assert 1 <= h <= self.max_h and 1 <= w <= self.max_w, (h, w)
+            base = full[:h, :w].reshape(-1, self.dim // 2)
+            parts.append(mx.tile(base, (t, 1)) if t > 1 else base)
+        return mx.concatenate(parts, axis=0)
+
+
+def apply_rope(q: mx.array, k: mx.array, freqs_cis: mx.array):
+    """q, k: (L, heads, head_dim); freqs_cis: (L, head_dim/2) complex64."""
+    f = mx.expand_dims(freqs_cis, axis=-2)
+
+    def rot(t):
+        pair = t.astype(mx.float32).reshape(*t.shape[:-1], -1, 2)
+        c = (pair[..., 0] + 1j * pair[..., 1]) * f
+        out = mx.stack([mx.real(c), mx.imag(c)], axis=-1)
+        return out.reshape(*t.shape).astype(t.dtype)
+
+    return rot(q), rot(k)
+
+
+# ----------------------------------------------------------- encoder
+
+
+class MLP2(nn.Module):
+    def __init__(self, dims: List[int], bias: bool = False):
+        super().__init__()
+        self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
+        self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # PytorchGELUTanh -- the tanh approximation, not exact erf GELU.
+        return self.fc1(nn.gelu_approx(self.fc0(x)))
+
+
+class MoonViTEncoderLayer(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.cfg = cfg
+        h, qkv = cfg.vt_hidden_size, cfg.qkv_hidden_size
+        self.num_heads = cfg.vt_num_attention_heads
+        self.head_dim = cfg.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.norm0 = nn.RMSNorm(h, eps=VIT_NORM_EPS)
+        self.norm1 = nn.RMSNorm(h, eps=VIT_NORM_EPS)
+        self.wqkv = nn.Linear(h, qkv * 3, bias=cfg.attn_bias)
+        self.wo = nn.Linear(qkv, h, bias=cfg.attn_bias)  # [K3-V3] not square
+        self.mlp = MLP2([h, cfg.vt_intermediate_size, h], bias=cfg.linear_bias)
+
+    def _attend(self, x, freqs_cis, spans):
+        L = x.shape[0]
+        qkv = self.wqkv(x).reshape(L, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
+        q, k = apply_rope(q, k, freqs_cis)
+
+        # Images are packed into one sequence but must not attend across each
+        # other. Looping the (few) images costs sum(L_i^2) instead of the
+        # (sum L_i)^2 an explicit block-diagonal mask would.
+        outs = []
+        for s, e in spans:
+            qi = q[s:e].transpose(1, 0, 2)[None]
+            ki = k[s:e].transpose(1, 0, 2)[None]
+            vi = v[s:e].transpose(1, 0, 2)[None]
+            o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=self.scale)
+            outs.append(o[0].transpose(1, 0, 2).reshape(e - s, -1))
+        return self.wo(mx.concatenate(outs, axis=0))
+
+    def __call__(self, x, freqs_cis, spans):
+        x = x + self._attend(self.norm0(x), freqs_cis, spans)
+        return x + self.mlp(self.norm1(x))
+
+
+class MoonViT3dEncoder(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.rope_2d = Rope2DPosEmbRepeated(
+            cfg.head_dim, cfg.rope_max_size, cfg.rope_max_size, cfg.rope_theta_base
+        )
+        self.blocks = [MoonViTEncoderLayer(cfg) for _ in range(cfg.vt_num_hidden_layers)]
+        self.final_layernorm = nn.RMSNorm(cfg.vt_hidden_size, eps=VIT_NORM_EPS)
+
+    def __call__(self, x: mx.array, grid_thws) -> mx.array:
+        freqs_cis = self.rope_2d(grid_thws)
+        spans, pos = [], 0
+        for t, h, w in grid_thws:
+            spans.append((pos, pos + t * h * w))
+            pos += t * h * w
+        for blk in self.blocks:
+            x = blk(x, freqs_cis, spans)
+        return self.final_layernorm(x)
+
+
+# ---------------------------------------------------- [K3-V6] merge + project
+
+
+def tpool_patch_merger(x: mx.array, grid_thws, merge_kernel_size=(2, 2)):
+    """-> list of (h/2 * w/2, 4, C): mean over frames, then 2x2 spatial merge."""
+    d = x.shape[-1]
+    kh, kw = merge_kernel_size
+    out, pos = [], 0
+    for t, h, w in grid_thws:
+        seq = x[pos : pos + t * h * w]
+        pos += t * h * w
+        nh, nw = h // kh, w // kw
+        s = seq.reshape(t, nh, kh, nw, kw, d).transpose(0, 1, 3, 2, 4, 5)
+        s = mx.mean(s, axis=0)                      # temporal pooling
+        out.append(s.reshape(nh * nw, kh * kw, d))
+    return out
+
+
+class PatchMergerMLPV2(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        inner = cfg.mm_hidden_size * cfg.merge_kernel_size[0] * cfg.merge_kernel_size[1]
+        self.inner = inner
+        # nn.Sequential(Linear, GELU, Linear) in the reference -> proj.0 / proj.2
+        self.proj_0 = nn.Linear(inner, inner, bias=False)
+        self.proj_2 = nn.Linear(inner, cfg.text_hidden_size, bias=False)
+        self.post_norm = nn.RMSNorm(cfg.text_hidden_size, eps=cfg.projector_ln_eps)
+
+    def __call__(self, feats: List[mx.array]) -> List[mx.array]:
+        lengths = [f.shape[0] for f in feats]
+        x = mx.concatenate([f.reshape(f.shape[0], -1) for f in feats], axis=0)
+        # exact erf GELU here -- deliberately different from the ViT's tanh GELU
+        x = self.post_norm(self.proj_2(nn.gelu(self.proj_0(x))))
+        out, pos = [], 0
+        for n in lengths:
+            out.append(x[pos : pos + n])
+            pos += n
+        return out
+
+
+class VisionModel(nn.Module):
+    """vision_tower + mm_projector: pixels -> text-space embeddings."""
+
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.cfg = cfg
+        if cfg.pos_emb_type != "divided_fixed":
+            raise NotImplementedError(cfg.pos_emb_type)
+        if cfg.merge_type != "sd2_tpool":
+            raise NotImplementedError(cfg.merge_type)
+        if cfg.mm_projector_type != "patchmergerv2":
+            raise NotImplementedError(cfg.mm_projector_type)
+        if cfg.norm_type != "rmsnorm":
+            raise NotImplementedError(cfg.norm_type)
+        self.patch_embed = MoonVision3dPatchEmbed(cfg)
+        self.encoder = MoonViT3dEncoder(cfg)
+        self.mm_projector = PatchMergerMLPV2(cfg)
+
+    def __call__(self, pixel_values: mx.array, grid_thws) -> List[mx.array]:
+        grid_thws = [tuple(int(v) for v in g) for g in grid_thws]
+        x = self.patch_embed(pixel_values, grid_thws)
+        x = self.encoder(x, grid_thws)
+        merged = tpool_patch_merger(x, grid_thws, self.cfg.merge_kernel_size)
+        return self.mm_projector(merged)
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        out: Dict[str, mx.array] = {}
+        for k, v in weights.items():
+            if k.startswith("vision_tower."):
+                k = k[len("vision_tower.") :]
+                if k == "patch_embed.proj.weight":
+                    # Conv2d (out, in, kH, kW) -> flat linear (out, in*kH*kW).
+                    # torch sums over (c, i, j) in that order, so a plain
+                    # reshape is the correct flattening.
+                    out["patch_embed.proj_weight"] = v.reshape(v.shape[0], -1)
+                elif k == "patch_embed.pos_emb.weight":
+                    out["patch_embed.pos_emb.weight"] = v
+                else:
+                    out[k] = v
+            elif k.startswith("mm_projector."):
+                k = k[len("mm_projector.") :]
+                if k == "proj.0.weight":
+                    out["mm_projector.proj_0.weight"] = v
+                elif k == "proj.2.weight":
+                    out["mm_projector.proj_2.weight"] = v
+                else:
+                    out["mm_projector." + k] = v
+        return out

kimi_k3_vision_processing.py

@@ -0,0 +1,179 @@
+"""Image processor class for Kimi-K3.
+"""
+
+import json
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers.image_processing_utils import (BaseImageProcessor,
+                                                 BatchFeature)
+from transformers.utils import TensorType
+
+from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor,
+                          ensure_media_type, image_to_np, navit_patchify,
+                          navit_resize_image, normalize)
+
+
+class KimiK3VisionProcessor(BaseImageProcessor):
+    model_type = "kimi_k3"
+
+    def __init__(
+        self,
+        media_proc_cfg: dict,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.media_proc_cfg = media_proc_cfg
+
+    @property
+    def _transparent_bg_config(self) -> Optional[TransparentBgConfig]:
+        cfg = self.media_proc_cfg.get("transparent_bg_config")
+        if cfg is None:
+            return None
+        if isinstance(cfg, TransparentBgConfig):
+            return cfg
+        return TransparentBgConfig(**cfg)
+
+    @property
+    def _transparent_bg_fill_stage(self) -> str:
+        return self.media_proc_cfg.get("transparent_bg_fill_stage",
+                                       "before_resize")
+
+    def media_tokens_calculator(self, media: MediaInput):
+        media = ensure_media_type(
+            media,
+            transparent_bg_config=self._transparent_bg_config,
+            transparent_bg_fill_stage=self._transparent_bg_fill_stage,
+        )
+        ret = self.get_resize_config(media)
+        return ret['num_tokens']
+
+    @classmethod
+    def make_image_prompt(cls, width: int, height: int) -> str:
+        """Build the K3 image placeholder with resolution info."""
+        return (f"<|media_begin|>image {width}x{height}"
+                f"<|media_content|><|media_pad|><|media_end|>")
+
+    def get_resize_config(self, media_input: MediaInput) -> dict:
+        if media_input['type'] == 'image':
+            w, h = media_input['image'].size
+            ret = navit_resize_image(
+                w, h, self.media_proc_cfg['patch_size'],
+                self.media_proc_cfg['merge_kernel_size'],
+                self.media_proc_cfg['in_patch_limit'],
+                self.media_proc_cfg['patch_limit_on_one_side'],
+                self.media_proc_cfg['fixed_output_tokens'])
+            return ret
+        else:
+            raise ValueError("Unsupported type: {}".format(
+                media_input['type']))
+
+    def resize_image(self, image: Image.Image, new_width: int, new_height: int,
+                     pad_width: int, pad_height: int) -> np.ndarray:
+        image_np = image_to_np(
+            image,
+            (new_width, new_height),
+            "resize",
+            transparent_bg_config=self._transparent_bg_config,
+            transparent_bg_fill_stage=self._transparent_bg_fill_stage,
+        )
+        image_np = np.pad(
+            image_np,
+            ((0, pad_height), (0, pad_width), (0, 0)),
+            mode="constant",
+            constant_values=0,
+        )
+        return image_np
+
+    def preprocess(
+        self,
+        medias: list[MediaInput],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ) -> BatchFeature:
+        """
+        Preprocess a atom vision input (images) into model-ready tensors.
+
+        Args:
+            medias: List of MediaInput.
+            return_tensors: Desired output format ('pt', 'np', 'tf', or None).
+
+        Returns:
+            BatchFeature containing 'pixel_values' and 'grid_thws' tensors.
+        """
+        if not isinstance(medias, list):
+            medias = [medias]
+        if medias:
+            pixel_values = []
+            for item in medias:
+                item = ensure_media_type(
+                    item,
+                    transparent_bg_config=self._transparent_bg_config,
+                    transparent_bg_fill_stage=self._transparent_bg_fill_stage,
+                )
+                resize_config = self.get_resize_config(item)
+                new_width, new_height, pad_width, pad_height = resize_config[
+                    'new_width'], resize_config['new_height'], resize_config[
+                        'pad_width'], resize_config['pad_height']
+                if item['type'] == 'image':
+                    image = item['image']
+                    image_np = self.resize_image(image, new_width, new_height,
+                                                 pad_width, pad_height)
+                    pixel_values.append(np.expand_dims(image_np, axis=0))
+                else:
+                    raise ValueError("Unsupported type: {}".format(
+                        item['type']))
+            normalized_pixel_values = []
+            image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std'])
+            image_mean = np.array(self.media_proc_cfg['image_mean'])
+            for pixels in pixel_values:
+                pixels = normalize(pixels, image_mean, image_std_inv)
+                pixels_and_thw = navit_patchify(
+                    pixels,
+                    self.media_proc_cfg['patch_size'],
+                )
+                normalized_pixel_values.append(pixels_and_thw)
+
+            pixel_values = torch.cat([
+                _to_tensor(pixel_value['pixel_values'])
+                for pixel_value in normalized_pixel_values
+            ])
+            grid_thws = torch.cat([
+                _to_tensor(pixel_value['grid_thw'],
+                           dtype=torch.int64).unsqueeze(0)
+                for pixel_value in normalized_pixel_values
+            ])
+
+            data = {
+                'pixel_values': pixel_values,
+                'grid_thws': grid_thws,
+            }
+
+        else:
+            data = {}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def __repr__(self):
+        return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})"
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = super().to_dict()
+        output["media_proc_cfg"] = self.media_proc_cfg
+        if "media_processor" in output:
+            del output["media_processor"]
+        return output
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any], **kwargs):
+        config = config_dict.copy()
+        media_proc_cfg = config.pop("media_proc_cfg", {})
+        return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)
+
+    def to_json_string(self):
+        dictionary = self.to_dict()
+        for key, value in dictionary.items():
+            if hasattr(value, 'tolist'):
+                dictionary[key] = value.tolist()
+        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"

kimi_k3_vl/__init__.py

@@ -0,0 +1,14 @@
+from .config import ModelConfig, TextConfig, VisionConfig
+from .kimi_k3_vl import Model, merge_image_features
+from .language import LanguageModel
+from .vision import VisionModel
+
+__all__ = [
+    "ModelConfig",
+    "TextConfig",
+    "VisionConfig",
+    "Model",
+    "LanguageModel",
+    "VisionModel",
+    "merge_image_features",
+]

kimi_k3_vl/config.py

@@ -0,0 +1,85 @@
+"""Config plumbing for the Kimi-K3 mlx-vlm wrapper."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class BaseCfg:
+    @classmethod
+    def from_dict(cls, params: Dict[str, Any]):
+        known = {f for f in cls.__dataclass_fields__}
+        return cls(**{k: v for k, v in params.items() if k in known})
+
+
+@dataclass
+class TextConfig(BaseCfg):
+    """Thin holder: the real parsing lives in kimi_k3.ModelArgs, which already
+    understands K3's nested `text_config`. Kept so mlx-vlm's config machinery
+    has something with the shape it expects."""
+
+    model_type: str = "kimi_linear"
+    hidden_size: int = 7168
+    num_hidden_layers: int = 93
+    vocab_size: int = 163840
+    raw: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, params: Dict[str, Any]):
+        return cls(
+            model_type=params.get("model_type", "kimi_linear"),
+            hidden_size=params.get("hidden_size", 7168),
+            num_hidden_layers=params.get("num_hidden_layers", 93),
+            vocab_size=params.get("vocab_size", 163840),
+            raw=dict(params),
+        )
+
+
+@dataclass
+class VisionConfig(BaseCfg):
+    model_type: str = "moonvit3d"
+    raw: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, params: Dict[str, Any]):
+        return cls(model_type=params.get("model_type", "moonvit3d"), raw=dict(params))
+
+
+@dataclass
+class ModelConfig(BaseCfg):
+    text_config: TextConfig
+    vision_config: VisionConfig
+    model_type: str = "kimi_k3"
+    # <|media_pad|>. Exactly ONE of these per image in the prompt; it expands to
+    # that image's full token count. See kimi_k3_vl.merge_image_features.
+    media_placeholder_token_id: int = 163605
+    image_token_index: Optional[int] = None
+    ignore_index: int = -100
+    pad_token_id: int = 163839
+    bos_token_id: int = 163584
+    eos_token_id: Optional[Any] = 163586
+    vocab_size: int = 163840
+    quantization: Optional[Dict[str, Any]] = None
+    raw: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, params: Dict[str, Any]):
+        # TextConfig.raw carries the *whole* config: kimi_k3.ModelArgs.from_dict
+        # already knows how to unwrap K3's nested `text_config`, so reusing that
+        # path avoids a second, divergent copy of the same unwrapping rules.
+        text = TextConfig.from_dict(params.get("text_config", {}))
+        text.raw = dict(params)
+        return cls(
+            text_config=text,
+            vision_config=VisionConfig.from_dict(params.get("vision_config", {})),
+            model_type=params.get("model_type", "kimi_k3"),
+            media_placeholder_token_id=params.get("media_placeholder_token_id", 163605),
+            image_token_index=params.get("image_token_index"),
+            ignore_index=params.get("ignore_index", -100),
+            pad_token_id=params.get("pad_token_id", 163839),
+            bos_token_id=params.get("bos_token_id", 163584),
+            eos_token_id=params.get("eos_token_id", 163586),
+            vocab_size=params.get("text_config", {}).get("vocab_size", 163840),
+            quantization=params.get("quantization"),
+            raw=dict(params),
+        )

kimi_k3_vl/kimi_k3_vl.py

@@ -0,0 +1,191 @@
+"""Kimi-K3 multimodal glue for mlx-vlm: vision tower -> prompt -> text tower.
+
+The only genuinely interesting part is `merge_image_features`.
+
+K3's processor rewrites each `<|kimi_image_placeholder|>` in the raw text into
+
+    <|media_begin|>image {W}x{H}<|media_content|><|media_pad|><|media_end|>
+
+so the tokenized prompt carries **exactly one** `<|media_pad|>` (163605) per
+image, which must then EXPAND into that image's full token count. That differs
+from the more common LLaVA-style arrangement (and from mlx-vlm's existing
+Kimi-VL glue), where the processor has already emitted one placeholder per image
+token and merging is a same-length scatter.
+
+Getting this wrong is quiet rather than loud: a same-length scatter would write
+one image token and silently drop the rest, leaving a model that still generates
+fluent text while being effectively blind to most of the picture.
+"""
+
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import ModelConfig
+from .language import LanguageModel
+from .vision import VisionModel
+
+try:
+    from mlx_vlm.models.base import InputEmbeddingsFeatures
+except ImportError:  # allow importing the glue without mlx-vlm present
+    from dataclasses import dataclass
+
+    @dataclass
+    class InputEmbeddingsFeatures:  # type: ignore[no-redef]
+        inputs_embeds: mx.array
+        attention_mask: Optional[mx.array] = None
+        position_ids: Optional[mx.array] = None
+
+
+def merge_image_features(
+    inputs_embeds: mx.array,
+    input_ids: mx.array,
+    image_features: List[mx.array],
+    placeholder_id: int,
+    pad_token_id: Optional[int] = None,
+):
+    """Expand each placeholder token into its image's feature block.
+
+    inputs_embeds : (B, S, D)   token embeddings
+    input_ids     : (B, S)
+    image_features: list of (n_i, D), one per image, in prompt order
+    returns       : (B, S', D) merged embeddings and a (B, S') attention mask,
+                    where S' = S - K + sum(n_i) for K placeholders
+
+    Rows are built by concatenating the spans between placeholders, which is
+    O(images) concatenations rather than O(sequence length) scatter indices, and
+    is far easier to check than the reference's index arithmetic. With B > 1 the
+    rows are LEFT-padded to the longest merged length, matching the reference's
+    `left_padding` branch (it is what generation with a KV cache expects).
+    """
+    B, S = input_ids.shape
+    ids = input_ids.tolist()
+
+    n_ph_total = sum(1 for row in ids for t in row if t == placeholder_id)
+    if n_ph_total != len(image_features):
+        raise ValueError(
+            f"{n_ph_total} <|media_pad|> placeholder(s) in the prompt but "
+            f"{len(image_features)} image feature block(s). K3 uses exactly one "
+            f"placeholder per image."
+        )
+
+    rows, masks, taken = [], [], 0
+    for b in range(B):
+        parts, mask_parts, prev = [], [], 0
+        for j, tok in enumerate(ids[b]):
+            if tok != placeholder_id:
+                continue
+            if j > prev:
+                parts.append(inputs_embeds[b, prev:j])
+                mask_parts.append(mx.ones((j - prev,), dtype=mx.int32))
+            feat = image_features[taken].astype(inputs_embeds.dtype)
+            taken += 1
+            parts.append(feat)
+            mask_parts.append(mx.ones((feat.shape[0],), dtype=mx.int32))
+            prev = j + 1
+        if prev < S:
+            parts.append(inputs_embeds[b, prev:])
+            mask_parts.append(mx.ones((S - prev,), dtype=mx.int32))
+        rows.append(mx.concatenate(parts, axis=0))
+        masks.append(mx.concatenate(mask_parts, axis=0))
+
+    if B == 1:
+        merged = rows[0][None]
+        mask = masks[0][None]
+    else:
+        longest = max(r.shape[0] for r in rows)
+        padded, padded_masks = [], []
+        for r, m in zip(rows, masks):
+            gap = longest - r.shape[0]
+            if gap:
+                r = mx.concatenate([mx.zeros((gap, r.shape[-1]), r.dtype), r], axis=0)
+                m = mx.concatenate([mx.zeros((gap,), mx.int32), m], axis=0)
+            padded.append(r)
+            padded_masks.append(m)
+        merged = mx.stack(padded)
+        mask = mx.stack(padded_masks)
+
+    # zero out embeddings at pad positions, as the reference does
+    if pad_token_id is not None:
+        merged = mx.where(mask[..., None] == 0, mx.zeros_like(merged), merged)
+    return merged, mask
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        self.vision_tower = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config)
+
+    # ------------------------------------------------------------------ api
+
+    @property
+    def layers(self):
+        return self.language_model.model.model.layers
+
+    def make_cache(self):
+        return self.language_model.model.make_cache()
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ) -> InputEmbeddingsFeatures:
+        embed = self.language_model.model.model.embed_tokens
+        inputs_embeds = embed(input_ids)
+        if pixel_values is None:
+            return InputEmbeddingsFeatures(inputs_embeds=inputs_embeds)
+
+        grids = (
+            kwargs.pop("grid_thws", None)
+            or kwargs.pop("image_grid_thw", None)
+            or kwargs.pop("image_grid_hws", None)
+        )
+        if grids is None:
+            raise ValueError("grid_thws is required alongside pixel_values")
+        if isinstance(grids, mx.array):
+            grids = grids.tolist()
+        grids = [tuple(int(v) for v in g) for g in grids]
+
+        feats = kwargs.get("cached_image_features")
+        if feats is None:
+            feats = self.vision_tower(pixel_values, grids)
+
+        placeholder = kwargs.pop("image_token_id", None) or self.config.media_placeholder_token_id
+        merged, mask = merge_image_features(
+            inputs_embeds, input_ids, feats, placeholder, self.config.pad_token_id
+        )
+        return InputEmbeddingsFeatures(inputs_embeds=merged, attention_mask=mask)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        cache: Optional[List[Any]] = None,
+        **kwargs,
+    ) -> mx.array:
+        feats = self.get_input_embeddings(input_ids, pixel_values, **kwargs)
+        embeds = feats.inputs_embeds
+        # After merging, the token axis no longer lines up with input_ids, so the
+        # text tower must be driven by the embeddings' own length.
+        ids = mx.zeros(embeds.shape[:2], dtype=mx.int32)
+        return self.language_model(ids, cache=cache, inputs_embeds=embeds)
+
+    def sanitize(self, weights):
+        """Route source keys to the two towers.
+
+        `vision_tower.*` / `mm_projector.*` are consumed by the vision tower;
+        `language_model.*` by the text tower. Both delegate to the same sanitize
+        functions the standalone models use, so there is one implementation of
+        the key mapping rather than a wrapper-specific copy that could drift.
+        """
+        vis = {k: v for k, v in weights.items()
+               if k.startswith(("vision_tower.", "mm_projector."))}
+        txt = {k: v for k, v in weights.items() if not k.startswith(("vision_tower.", "mm_projector."))}
+        out = {f"vision_tower.{k}": v for k, v in self.vision_tower.sanitize(vis).items()}
+        out.update({f"language_model.{k}": v for k, v in self.language_model.sanitize(txt).items()})
+        return out

kimi_k3_vl/language.py

@@ -0,0 +1,46 @@
+"""Text-tower adapter: wraps kimi_k3.Model for mlx-vlm."""
+
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import TextConfig
+
+try:
+    from mlx_lm.models import kimi_k3 as _k3
+except ImportError as e:  # pragma: no cover
+    raise ImportError(
+        "kimi_k3 is not registered with mlx-lm. It uses mlx-lm-relative imports "
+        "(.base, .cache, .switch_layers) so it must live in mlx_lm/models/. Run "
+        "`scripts/install_model.sh` once, or copy kimi_k3.py there."
+    ) from e
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        args = _k3.ModelArgs.from_dict(config.raw)
+        self.args = args
+        self.model = _k3.Model(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+        inputs_embeds: Optional[mx.array] = None,
+        **kwargs,
+    ) -> mx.array:
+        return self.model(inputs, cache=cache, inputs_embeds=inputs_embeds)
+
+    @property
+    def layers(self):
+        return self.model.model.layers
+
+    def make_cache(self):
+        return self.model.make_cache()
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        # kimi_k3.Model.sanitize already strips `language_model.` and drops the
+        # vision keys; prefix the result to sit under this module.
+        return {f"model.{k}": v for k, v in self.model.sanitize(weights).items()}

kimi_k3_vl/vision.py

@@ -0,0 +1,52 @@
+"""Vision-tower adapter: wraps kimi_k3_vision.VisionModel for mlx-vlm."""
+
+import importlib.util
+import os
+import sys
+from typing import Dict, List, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import VisionConfig as WrapperVisionConfig
+
+
+def _load_vision_module():
+    """kimi_k3_vision.py ships alongside the weights; import it from wherever
+    this package lives so a downloaded repo works without installation."""
+    if "kimi_k3_vision" in sys.modules:
+        return sys.modules["kimi_k3_vision"]
+    try:
+        import kimi_k3_vision  # noqa: F401
+
+        return sys.modules["kimi_k3_vision"]
+    except ImportError:
+        pass
+    for d in (os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+              os.path.dirname(os.path.abspath(__file__))):
+        p = os.path.join(d, "kimi_k3_vision.py")
+        if os.path.exists(p):
+            spec = importlib.util.spec_from_file_location("kimi_k3_vision", p)
+            mod = importlib.util.module_from_spec(spec)
+            sys.modules["kimi_k3_vision"] = mod
+            spec.loader.exec_module(mod)
+            return mod
+    raise ImportError("kimi_k3_vision.py not found next to kimi_k3_vl/")
+
+
+_v = _load_vision_module()
+
+
+class VisionModel(nn.Module):
+    def __init__(self, config: WrapperVisionConfig):
+        super().__init__()
+        self.cfg = _v.VisionConfig.from_dict(config.raw)
+        self.tower = _v.VisionModel(self.cfg)
+
+    def __call__(self, pixel_values: mx.array, grid_thws) -> List[mx.array]:
+        """pixel_values: (L, 3, patch, patch) pre-extracted patches.
+        Returns one (n_i, text_hidden) array per image."""
+        return self.tower(pixel_values, grid_thws)
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        return {f"tower.{k}": v for k, v in self.tower.sanitize(weights).items()}

media_utils.py

@@ -0,0 +1,376 @@
+import base64
+import functools
+import io
+import math
+from dataclasses import dataclass
+from typing import Literal, TypedDict
+
+import numpy as np
+from PIL import Image
+
+
+class ImageInput(TypedDict):
+    type: Literal['image']
+    image: Image.Image
+
+
+MediaInput = ImageInput
+
+
+@dataclass
+class TransparentBgConfig:
+    """The config of the transparent background."""
+
+    pattern: Literal["white", "black", "gray", "chessboard"] = "black"
+    """The pattern of the transparent background."""
+
+    chessboard_square_size: int = 16
+    """The size of the squares in the chessboard background."""
+
+    chessboard_square_on_top_left: bool = True
+    """Whether to start the chessboard with a white square on the top left."""
+
+    chessboard_white_value: int = 255
+    """The value of the white pixels in the background."""
+
+    chessboard_gray_value: int = 200
+    """The value of the gray pixels in the background."""
+
+
+@functools.lru_cache(maxsize=256)
+def _create_chessboard_background(
+    height: int,
+    width: int,
+    square_size: int,
+    square_on_top_left: bool,
+    white_value: int,
+    gray_value: int,
+) -> np.ndarray:
+    """Create a chessboard background."""
+    bg = np.ones((height, width, 3), dtype=np.uint8) * white_value
+    for y in range(0, height, square_size):
+        for x in range(0, width, square_size):
+            if (y // square_size + x // square_size) % 2 == (
+                    1 if square_on_top_left else 0):
+                bg[y:y + square_size, x:x + square_size] = gray_value
+    return bg
+
+
+def fill_transparent_bg_with(
+    image: Image.Image,
+    transparent_bg_config: TransparentBgConfig | None = None,
+) -> Image.Image:
+    """Composite a (possibly) transparent image onto a configured background.
+
+    When ``transparent_bg_config`` is ``None``, the image is simply converted
+    to RGB (preserving the historical behavior). Otherwise the alpha channel
+    is alpha-composited over a background generated according to the config.
+    """
+    if transparent_bg_config is None:
+        return image.convert("RGB")
+
+    if image.mode == "RGB":
+        return image
+
+    has_alpha = "A" in image.getbands() or "transparency" in image.info
+    if not has_alpha:
+        return image.convert("RGB")
+
+    img = np.array(image.convert("RGBA"))
+    height, width = img.shape[:2]
+    bg_pattern = transparent_bg_config.pattern
+    if bg_pattern == "white":
+        bg = np.full((height, width, 3), 255, dtype=np.uint8)
+    elif bg_pattern == "black":
+        bg = np.zeros((height, width, 3), dtype=np.uint8)
+    elif bg_pattern == "gray":
+        bg = np.full((height, width, 3), 128, dtype=np.uint8)
+    elif bg_pattern == "chessboard":
+        bg = _create_chessboard_background(
+            height,
+            width,
+            transparent_bg_config.chessboard_square_size,
+            transparent_bg_config.chessboard_square_on_top_left,
+            transparent_bg_config.chessboard_white_value,
+            transparent_bg_config.chessboard_gray_value,
+        )
+    else:
+        raise ValueError(f"Invalid background pattern: {bg_pattern}")
+
+    alpha = img[:, :, 3]
+    img_rgb = img[:, :, :3]
+    alpha_normalized = alpha.astype(np.float32) / 255.0
+    alpha_3d = np.stack([alpha_normalized] * 3, axis=2)
+    result = alpha_3d * img_rgb + (1 - alpha_3d) * bg
+    result = result.astype(np.uint8)
+    return Image.fromarray(result)
+
+
+def navit_resize_image(
+    width: int,
+    height: int,
+    patch_size: int,
+    merge_kernel_size: int,
+    in_patch_limit: int,
+    patch_limit_on_one_side: int,
+    fixed_output_tokens: int | None,
+):
+    # Apply the patch limits.
+    s1 = math.sqrt(
+        in_patch_limit /
+        (max(1.0, width // patch_size) * max(1.0, height // patch_size)))
+    s2 = patch_limit_on_one_side * patch_size / width
+    s3 = patch_limit_on_one_side * patch_size / height
+    scale = min(1.0, s1, s2, s3)
+    new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale))
+    new_w = min(new_w, patch_limit_on_one_side * patch_size)
+    new_h = min(new_h, patch_limit_on_one_side * patch_size)
+
+    # Calculate the padding to make the height and width divisible by the merge kernel size and patch size.
+    factor = merge_kernel_size * patch_size
+
+    pad_height = (factor - new_h % factor) % factor
+    pad_width = (factor - new_w % factor) % factor
+
+    if fixed_output_tokens is not None:
+        num_tokens = fixed_output_tokens
+    else:
+        # Calculate new dimensions after padding and patching
+        token_height = (new_h + pad_height) // factor
+        token_width = (new_w + pad_width) // factor
+
+        assert token_height * merge_kernel_size <= patch_limit_on_one_side, (
+            f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}"
+        )
+        assert token_width * merge_kernel_size <= patch_limit_on_one_side, (
+            f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}"
+        )
+
+        num_tokens = token_height * token_width
+    return {
+        "num_tokens": num_tokens,
+        "new_width": new_w,
+        "new_height": new_h,
+        "pad_width": pad_width,
+        "pad_height": pad_height,
+        "sampled_nframes": 1,
+    }
+
+
+def _to_pil(
+    data: str | bytes | Image.Image,
+    transparent_bg_config: TransparentBgConfig | None = None,
+    to_rgb: bool = True,
+) -> Image.Image:
+    """Load an image and (optionally) composite its transparent background.
+
+    Args:
+        data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes.
+        transparent_bg_config: The config used to fill the transparent
+            background. ``None`` keeps the historical behavior of converting
+            to RGB without compositing.
+        to_rgb: If ``False`` the image is returned as-is (the
+            ``transparent_bg_config`` is ignored). The caller is then
+            expected to call :func:`fill_transparent_bg_with` later — e.g.
+            after a resize.
+    """
+    if isinstance(data, Image.Image):
+        image = data
+    elif isinstance(data, str):
+        if data.startswith("data:"):
+            raw_base64 = data.split(",")[1]
+            image = Image.open(io.BytesIO(base64.b64decode(raw_base64)))
+        else:
+            image = Image.open(data)
+    elif isinstance(data, bytes):
+        image = Image.open(io.BytesIO(data))
+    else:
+        raise ValueError(f"Unsupported data type: {type(data)}")
+
+    if not to_rgb:
+        return image
+
+    return fill_transparent_bg_with(image, transparent_bg_config)
+
+
+def ensure_media_type(
+    media: MediaInput,
+    transparent_bg_config: TransparentBgConfig | None = None,
+    transparent_bg_fill_stage: Literal["before_resize",
+                                       "after_resize"] = "before_resize",
+) -> MediaInput:
+    if media['type'] == 'image':
+        media['image'] = _to_pil(
+            media['image'],
+            transparent_bg_config=transparent_bg_config,
+            to_rgb=transparent_bg_fill_stage == "before_resize",
+        )
+        return media
+    else:
+        raise ValueError(f"Unsupported media type: {media['type']}")
+
+
+def image_to_np(
+    image: Image.Image,
+    resize_to: tuple[int, int] | None = None,
+    mode: str = "resize",
+    raise_error_for_ill_resize: bool = True,
+    transparent_bg_config: TransparentBgConfig | None = None,
+    transparent_bg_fill_stage: Literal["before_resize",
+                                       "after_resize"] = "before_resize",
+) -> np.ndarray:
+    """Convert an image to a numpy array.
+
+    Args:
+        content: The image to convert.
+        resize_to: The size to resize the image to.
+        mode: The mode to resize the image to.
+        raise_error_for_ill_resize: Whether to raise an error for ill-sized resize.
+        transparent_bg_config: The config of the transparent background. Only
+            used when ``transparent_bg_fill_stage == "after_resize"`` (the
+            caller is responsible for filling before resize otherwise).
+        transparent_bg_fill_stage: When to composite the transparent
+            background — before or after the resize step.
+
+    Returns:
+        A numpy array.
+    """
+    assert isinstance(image, Image.Image), "image must be a PIL Image"
+    if resize_to is not None:
+        if mode == "resize":
+            image = image.resize(resize_to, resample=Image.Resampling.BICUBIC)
+            if transparent_bg_fill_stage == "after_resize":
+                image = fill_transparent_bg_with(image, transparent_bg_config)
+
+        elif mode == "rescale_and_pad_to_center":
+            scale = min(resize_to[0] / image.width,
+                        resize_to[1] / image.height, 1.0)
+            new_width = round(image.width * scale)
+            new_height = round(image.height * scale)
+            if new_width == 0 or new_height == 0:
+                if raise_error_for_ill_resize:
+                    raise ValueError(
+                        f"Invalid resize to: {resize_to}, from image size: {image.size}"
+                    )
+                else:
+                    return np.zeros((resize_to[1], resize_to[0], 3),
+                                    dtype=np.uint8)
+
+            image = image.resize((new_width, new_height),
+                                 resample=Image.Resampling.BICUBIC)
+            if transparent_bg_fill_stage == "after_resize":
+                image = fill_transparent_bg_with(image, transparent_bg_config)
+            padding_left = (resize_to[0] - new_width) // 2
+            padding_right = resize_to[0] - new_width - padding_left
+            padding_top = (resize_to[1] - new_height) // 2
+            padding_bottom = resize_to[1] - new_height - padding_top
+            image = np.asarray(image)
+            image = np.pad(
+                image,
+                ((padding_top, padding_bottom), (padding_left, padding_right),
+                 (0, 0)),
+                mode="constant",
+                constant_values=0,
+            )
+            assert image.shape == (resize_to[1], resize_to[0], 3)
+
+        elif mode == "rescale_and_pad_to_rightbottom":
+            scale = min(resize_to[0] / image.width,
+                        resize_to[1] / image.height, 1.0)
+            new_width = round(image.width * scale)
+            new_height = round(image.height * scale)
+            if new_width == 0 or new_height == 0:
+                if raise_error_for_ill_resize:
+                    raise ValueError(
+                        f"Invalid resize to: {resize_to}, from image size: {image.size}"
+                    )
+                else:
+                    return np.zeros((resize_to[1], resize_to[0], 3),
+                                    dtype=np.uint8)
+
+            image = image.resize((new_width, new_height),
+                                 resample=Image.Resampling.BICUBIC)
+            if transparent_bg_fill_stage == "after_resize":
+                image = fill_transparent_bg_with(image, transparent_bg_config)
+            padding_right = resize_to[0] - new_width
+            padding_bottom = resize_to[1] - new_height
+            image = np.asarray(image)
+            image = np.pad(
+                image,
+                ((0, padding_bottom), (0, padding_right), (0, 0)),
+                mode="constant",
+                constant_values=0,
+            )
+            assert image.shape == (resize_to[1], resize_to[0], 3)
+
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+
+    if isinstance(image, Image.Image):
+        return np.asarray(image)
+    else:
+        return image
+
+
+def navit_patchify(pixel_values: np.ndarray,
+                   patch_size: int) -> dict[str, np.ndarray]:
+    """Reshape the pixel values to a navit shape.
+
+    Args:
+        pixel_values: np.ndarray, shape (t, h, w, c)
+        patch_size: int
+
+    Returns:
+        dict[str, np.ndarray]
+        - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size)
+        - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size)
+    """
+    T, H, W, C = pixel_values.shape
+    assert C == 3, "pixel_values must have 3 channels"
+
+    patches = pixel_values.reshape(T, H // patch_size, patch_size,
+                                   W // patch_size, patch_size, C)
+    # (T, H//patch_size, W//patch_size, C, patch_size, patch_size)
+    patches = patches.transpose(0, 1, 3, 5, 2, 4)
+    patches = patches.reshape(-1, C, patch_size, patch_size)
+    grid_thw = np.array([T, H // patch_size, W // patch_size])
+    return {"pixel_values": patches, "grid_thw": grid_thw}
+
+
+def normalize(x: np.ndarray,
+              mean,
+              std_inv,
+              pixels_dtype: np.dtype = np.float32) -> np.ndarray:
+    """Normalize the image.
+
+    Args:
+        x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255].
+        mean: The mean of the image.
+        std_inv: The inverse of the std of the image.
+        pixels_dtype: The dtype of the image.
+    Returns:
+        The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype.
+    """
+    x = (x / 255.0).astype(pixels_dtype)
+    x -= mean
+    x *= std_inv
+    return x
+
+
+def _to_tensor(data, **kwargs):
+    import torch
+
+    if isinstance(data, np.ndarray):
+        return torch.from_numpy(data).to(**kwargs)
+    elif isinstance(data, torch.Tensor):
+        return data.to(**kwargs)
+    elif isinstance(data, list):
+        return [_to_tensor(item, **kwargs) for item in data]
+    elif isinstance(data, tuple):
+        return tuple(_to_tensor(item, **kwargs) for item in data)
+    elif isinstance(data, dict):
+        return {k: _to_tensor(v, **kwargs) for k, v in data.items()}
+    elif data is None:
+        return None
+    else:
+        raise ValueError(f"Unsupported data type: {type(data)}")

model-00003-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee2a6f6b3b8818de3fa93332eb9d34557be603e7f7d2bb4630da4f4016b91c8a
+size 5436028813

model-00006-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6448dfa1de5cd56e634b6c5acdca9ccc1e7559b73f96609738565e7074626336
+size 6105541561

model-00009-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0fe6e59ab5b98b11e7f9c592e92827ae1e74802aaa99a165b45e5c0d37632e5
+size 6105541618

model-00011-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcd286a358ee62cc9a4bae5584adb6ed295f0f07c1429388cd0f204563cf984b
+size 5379479293

model-00012-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14fd63b64a66c8bc9749b5f9e1a9b1bd044d0c1531d37b15aff966ad68242f96
+size 6309626568

model-00014-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:415f3fb772cf0cfe8a03b3d2edb3b073e69ad045286aa694d4b4929dabada170
+size 6330481117

model-00017-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b97fa6ceeb9b0b9a03a973942ba88c1ab2d452a4bfa74feb4498728d6509351c
+size 6330481113

model-00018-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f391b109cd133843386191e999944c534bd463e3c8ae70292c916e186b9789db
+size 6330481117

model-00020-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd6c6bb53d90d0702e46999d4bd81cba9b960ca410a1f3c924c356b759f6a8f1
+size 6253076934

model-00025-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36ac781716d59e87c5062adbee45106affb467ce107fb7d32725d9b0915da767
+size 6105541618

model-00029-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b6553530aed30872566fe901550585d0dffefe725017b21d7170680d4ee1b58
+size 6330481117

model-00031-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cbd1012010bafcb00271c19414156c9d74668c677490aab480ee9c0c8de9555
+size 5411251298

model-00032-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1c8bcd6a40a00656f30f78f0efeae50916ede029c40eaafa83599c87d98ca65
+size 6277854537

model-00034-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a0efc2be3d2a3eb00ced9ff009e75f607b8ba2a81469d4a66eb4c670de91c71
+size 6330481117

model-00037-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b85a77ab16156cea9916598e385c350511169777d2f96503566ae29de5ba3f6
+size 6330481117

model-00038-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88bca405e82263f293ec8589927a9947cd6853a818f46c00fefc76d959f9231e
+size 6105541618

model-00041-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dcd9787752bbab263c8c3ad9858417080f839353eda5c0ad28434ecb59a4e04
+size 6105541618

model-00042-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:450bb4e324453add694f6240e70a7392cadb651f58993be14e301ffba1298362
+size 6330481113

model-00044-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44d6e2abe1472823323f64b3121f7caec9a4cf60defd2f45575362dd53e118fe
+size 6309626462

model-00047-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:074f2e094d6529f2e299d8fa2a3da9a908c8bf1a1333d56e30db1bea874c6698
+size 5411251308

model-00048-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80889374efb3add94a2abb8f370cc4c1eaca8724b1800210ea5a5490883f2ca1
+size 6277854585

model-00053-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1eb7ee303b48d504ec1603a28732775692a7ba63b8ee8894c9f19ceca7d9e1e
+size 6330481117

model-00055-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f0341a8bfa50479923e1f3daa907e9a53d6933f37ccd5eba6e7856942b497f4
+size 5379479249

model-00056-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2312406654989559f52b41069b9c8ffb6cd77f3c9efef98be07a07992f90147e
+size 6534566067

model-00061-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28b201cb736c90c8e31231bf5a0e4955a6a6412befeeb586262787566f8436f1
+size 6330481117

model-00064-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7271ea6542fcfb833be60c5562066aef5bc0c42d98a1633c7cb28b0d573964b1
+size 6277854585

model-00068-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29ea631bae5696c2d12a6629c0a272178ae0dcd5317d258615f6950f6b1c95a1
+size 6253076934

model-00070-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d82aaa0ec50044b5f246da82f776ae4fd8ce763f19cf3059e485bdc852a31104
+size 6105541614

model-00073-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb3b5df114d34f45e4edb9fb3dbd39c96c02f90fbdb223def58c72ef1279d89d
+size 6105541618

model-00075-of-00075.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da4abec65dd10c84a43b425a45e9c867557eaedab7735991e0f20b05df72e359
+size 2225850591

preprocessor_config.json

@@ -0,0 +1,38 @@
+{
+  "media_proc_cfg": {
+    "in_patch_limit": 65536,
+    "patch_size": 14,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "merge_kernel_size": 2,
+    "fixed_output_tokens": null,
+    "patch_limit_on_one_side": 512,
+    "in_patch_limit_each_frame": 16384,
+    "in_patch_limit_video": 655360,
+    "sample_fps": 8.0,
+    "max_num_frames_each_video": null,
+    "temporal_merge_kernel_size": 4,
+    "timestamp_mode": "hh:mm:ss.fff",
+    "transparent_bg_config": {
+      "pattern": "chessboard",
+      "chessboard_square_size": 8,
+      "chessboard_square_on_top_left": true,
+      "chessboard_white_value": 255,
+      "chessboard_gray_value": 180
+    },
+    "transparent_bg_fill_stage": "after_resize",
+    "config_type": "media_proc.processors.moonvit.MoonViTMediaProcessorConfig"
+  },
+  "auto_map": {
+    "AutoProcessor": "kimi_k3_processor.KimiK3Processor",
+    "AutoImageProcessor": "kimi_k3_vision_processing.KimiK3VisionProcessor"
+  }
+}

tokenization_kimi.py

@@ -0,0 +1,408 @@
+import os
+from logging import getLogger
+from pathlib import Path
+from shutil import copyfile
+from typing import Dict, Iterator, List, Optional, Tuple, Union, cast
+
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
+from tokenizers import AddedToken
+from transformers.convert_slow_tokenizer import bytes_to_unicode
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+try:
+    from .encoding_k3 import build_chat_segments, is_batched_conversation
+except ImportError:  # pragma: no cover - supports direct file execution/import.
+    from encoding_k3 import build_chat_segments, is_batched_conversation
+
+logger = getLogger(__name__)
+VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"}
+
+
+class TikTokenTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            The path to the Tiktoken model file.
+        bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`):
+            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+        eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`):
+            The end of sequence token.
+        unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead. The second to last item in special_tokens.
+        pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        additional_special_tokens (list of `str`, *optional*):
+            A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be
+            skipped when decoding if `skip_special_tokens` is set to `True`.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    special_tokens: Dict[str, int]
+
+    num_reserved_special_tokens = 256
+
+    pat_str = "|".join([
+        r"""[\p{Han}]+""",
+        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+        r"""\p{N}{1,3}""",
+        r""" ?[^\s\p{L}\p{N}]+[\r\n]*""",
+        r"""\s*[\r\n]+""",
+        r"""\s+(?!\S)""",
+        r"""\s+""",
+    ])
+
+    def __init__(
+        self,
+        vocab_file,
+        bos_token: Union[str, AddedToken] = "[BOS]",
+        eos_token: Union[str, AddedToken] = "[EOS]",
+        unk_token: Union[str, AddedToken, None] = None,
+        pad_token: Union[str, AddedToken, None] = None,
+        additional_special_tokens: List[str] = None,
+        added_tokens_decoder: Optional[dict] = None,
+        **kwargs,
+    ):
+        assert os.path.isfile(vocab_file), vocab_file
+
+        if additional_special_tokens is None:
+            additional_special_tokens = [
+                "<|im_end|>",
+                "<|im_user|>",
+                "<|im_assistant|>",
+                "<|start_header_id|>",
+                "<|end_header_id|>",
+                "[EOT]",
+                "<|im_system|>",
+                "<|im_middle|>",
+            ]
+
+        if added_tokens_decoder:
+            special_tokens_mapping = {
+                i: added_tokens_decoder[i].content
+                for i in added_tokens_decoder
+            }
+        else:
+            special_tokens_mapping = {}
+
+        self.vocab_file = vocab_file
+        mergeable_ranks = load_tiktoken_bpe(vocab_file)
+        num_base_tokens = len(mergeable_ranks)
+        self.special_tokens = {
+            special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i
+            for i in range(num_base_tokens, num_base_tokens +
+                           self.num_reserved_special_tokens)
+        }
+
+        self.model = tiktoken.Encoding(
+            name=Path(vocab_file).name,
+            pat_str=self.pat_str,
+            mergeable_ranks=mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        logger.info(f"Reloaded tiktoken model from {vocab_file}")
+
+        self.n_words: int = self.model.n_vocab
+        # BOS / EOS token IDs
+        self.bos_id: int = self.special_tokens[str(bos_token)]
+        self.eos_id: int = self.special_tokens[str(eos_token)]
+        logger.info(
+            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
+        )
+
+        self.pad_id: int = self.special_tokens[str(pad_token)]
+        self.unk_id: int = self.special_tokens[str(unk_token)]
+
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+
+        self.decoder = {}
+        for i in range(self.n_words):
+            # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
+            decoding = ''.join([
+                self.byte_encoder[ord(char)] for char in
+                self.model.decode_single_token_bytes(i).decode('latin-1')
+            ])
+            self.decoder[i] = decoding
+
+        self.encoder = {}
+        for i in range(self.n_words):
+            if i in self.decoder:
+                self.encoder[self.decoder[i]] = i
+
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            additional_special_tokens=additional_special_tokens,
+            added_tokens_decoder=added_tokens_decoder,
+            **kwargs,
+        )
+        self.all_special_ids_set = set(self.all_special_ids)
+
+    def _encode_text_piece(self, text: str,
+                           allow_special_tokens: bool = True) -> List[int]:
+        # The tiktoken tokenizer can handle <=400k chars without
+        # pyo3_runtime.PanicException.
+        TIKTOKEN_MAX_ENCODE_CHARS = 400_000
+
+        # https://github.com/openai/tiktoken/issues/195
+        # Here we iterate over subsequences and split if we exceed the limit
+        # of max consecutive non-whitespace or whitespace characters.
+        MAX_NO_WHITESPACES_CHARS = 25_000
+
+        t: List[int] = []
+        for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS):
+            for substr in self._split_whitespaces_or_nonwhitespaces(
+                text[i:i + TIKTOKEN_MAX_ENCODE_CHARS],
+                MAX_NO_WHITESPACES_CHARS,
+            ):
+                if allow_special_tokens:
+                    t.extend(
+                        # structural markers: encode <|...|> as their special token IDs
+                        self.model.encode(
+                            substr,
+                            allowed_special="all",
+                        ))
+                else:
+                    t.extend(
+                        # user/tool text: encode any <|...|> as ordinary BPE tokens (never as control tokens)
+                        self.model.encode(
+                            substr,
+                            disallowed_special=(),
+                        ))
+
+        return t
+
+    def encode(self,
+               text: str,
+               allow_special_tokens: bool = True,
+               **kwargs) -> List[int]:
+        """
+        Encodes a string into a list of token IDs.
+
+        Args:
+            text (str): The input string to be encoded.
+
+        Returns:
+            list[int]: A list of token IDs.
+        """
+        # If there are other args, we should call super().encode because there are a lot of code
+        # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id.
+        # NOTE: our encode method is not compatible with the super().encode method,
+        #   e.g. split_special_tokens' default is True in our encode method.
+        if len(kwargs) > 0:
+            logger.warning(f"Calling super().encode with {kwargs}")
+            return super().encode(text, **kwargs)
+
+        assert type(text) is str
+        return self._encode_text_piece(text,
+                                       allow_special_tokens=allow_special_tokens)
+
+    def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str:
+        """
+        Decodes a list of token IDs into a string.
+
+        Args:
+            token_ids (List[int]): The list of token IDs to be decoded.
+
+        Returns:
+            str: The decoded string.
+        """
+        # If there are other args, we should call super().decode because there are a lot of code
+        # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token.
+        if len(kwargs) > 0:
+            return super().decode(token_ids, **kwargs)
+
+        if type(token_ids) is int:
+            token_ids = [token_ids]
+
+        return self.model.decode(cast(List[int], token_ids))
+
+    @staticmethod
+    def _split_whitespaces_or_nonwhitespaces(
+            s: str, max_consecutive_slice_len: int) -> Iterator[str]:
+        """
+        Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`
+        consecutive whitespaces or consecutive non-whitespaces.
+        """
+        current_slice_len = 0
+        current_slice_is_space = s[0].isspace() if len(s) > 0 else False
+        slice_start = 0
+
+        for i in range(len(s)):
+            is_now_space = s[i].isspace()
+
+            if current_slice_is_space ^ is_now_space:
+                current_slice_len = 1
+                current_slice_is_space = is_now_space
+            else:
+                current_slice_len += 1
+                if current_slice_len > max_consecutive_slice_len:
+                    yield s[slice_start:i]
+                    slice_start = i
+                    current_slice_len = 1
+        yield s[slice_start:]
+
+    def _encode_chat_segments(self, segments) -> List[int]:
+        token_ids: List[int] = []
+        for segment in segments:
+            token_ids.extend(
+                self._encode_text_piece(
+                    segment.text,
+                    allow_special_tokens=segment.allow_special,
+                ))
+        return token_ids
+
+    @staticmethod
+    def _truncate(ids: List[int],
+                  truncation: bool = False,
+                  max_length: Optional[int] = None) -> List[int]:
+        if truncation and max_length is not None:
+            return ids[:max_length]
+        return ids
+
+    def _format_chat_token_output(self,
+                                  encoded_inputs: List[List[int]],
+                                  *,
+                                  is_batched: bool,
+                                  padding=False,
+                                  truncation: bool = False,
+                                  max_length: Optional[int] = None,
+                                  return_tensors=None,
+                                  return_dict: bool = False):
+        encoded_inputs = [
+            self._truncate(ids, truncation=truncation, max_length=max_length)
+            for ids in encoded_inputs
+        ]
+
+        needs_batch_encoding = (
+            is_batched or padding or return_tensors is not None or return_dict)
+        if not needs_batch_encoding:
+            return encoded_inputs[0]
+
+        features = [{
+            "input_ids": ids,
+            "attention_mask": [1] * len(ids)
+        } for ids in encoded_inputs]
+        batch = self.pad(features,
+                         padding=padding,
+                         max_length=max_length if padding else None,
+                         return_attention_mask=True,
+                         return_tensors=return_tensors)
+
+        if return_dict:
+            return batch
+        if is_batched:
+            return batch["input_ids"]
+        return batch["input_ids"][0] if return_tensors is None else batch[
+            "input_ids"]
+
+    """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """
+
+    @property
+    def vocab_size(self) -> int:
+        return self.n_words
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self.encoder
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        return [self.decoder[t] for t in self.encode(text)]
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.encoder.get(token, self.unk_id)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.decoder.get(index)
+
+    @staticmethod
+    def clean_up_tokenization(out_string: str) -> str:
+        return out_string
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        text = ''.join(tokens)
+        text = bytearray([self.byte_decoder[c]
+                          for c in text]).decode('utf-8', 'replace')
+        return text
+
+    def save_vocabulary(self,
+                        save_directory: str,
+                        filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            raise ValueError(
+                f"vocabulary path ({save_directory}) should be a directory")
+        out_vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") +
+            VOCAB_FILES_NAMES["vocab_file"])
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(
+                out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+
+        return (out_vocab_file, )
+
+    def apply_chat_template(self,
+                            conversation,
+                            tools: Optional[list[dict]] = None,
+                            tokenize: bool = False,
+                            add_generation_prompt: bool = True,
+                            thinking: bool = True,
+                            padding=False,
+                            truncation: bool = False,
+                            max_length: Optional[int] = None,
+                            return_tensors=None,
+                            return_dict: bool = False,
+                            **kwargs):
+        # Tokenizer-level rendering reorders tool result messages to match
+        # assistant tool_calls, normalizes per-call arguments and response
+        # schema, then encodes the resulting XTML structure segment-by-segment.
+        is_batched = is_batched_conversation(conversation)
+        conversations = conversation if is_batched else [conversation]
+        image_prompts = kwargs.pop("image_prompts", None)
+        if is_batched and image_prompts is not None:
+            raise ValueError("image_prompts is only supported for one chat.")
+        
+        # by default set thinking effort to max
+        kwargs.setdefault("thinking_effort", "max")
+
+        segment_batches = [
+            build_chat_segments(
+                messages,
+                tools=tools,
+                add_generation_prompt=add_generation_prompt,
+                thinking=thinking,
+                image_prompts=image_prompts,
+                **kwargs,
+            ) for messages in conversations
+        ]
+
+        if not tokenize:
+            rendered = ["".join(segment.text for segment in segments)
+                        for segments in segment_batches]
+            return rendered if is_batched else rendered[0]
+
+        encoded_inputs = [
+            self._encode_chat_segments(segments) for segments in segment_batches
+        ]
+        return self._format_chat_token_output(
+            encoded_inputs,
+            is_batched=is_batched,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_dict=return_dict,
+        )

tokenizer_config.json

@@ -0,0 +1,157 @@
+{
+  "added_tokens_decoder": {
+    "163584": {
+      "content": "[BOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163585": {
+      "content": "[EOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163586": {
+      "content": "<|end_of_msg|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163587": {
+      "content": "<|open|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "163588": {
+      "content": "<|close|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "163589": {
+      "content": "<|sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "163590": {
+      "content": "[start_header_id]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163591": {
+      "content": "[end_header_id]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163593": {
+      "content": "[EOT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163602": {
+      "content": "<|media_begin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163603": {
+      "content": "<|media_content|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163604": {
+      "content": "<|media_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163605": {
+      "content": "<|media_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163649": {
+      "content": "<osagent_mode>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163838": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "163839": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|end_of_msg|>",
+    "[start_header_id]",
+    "[end_header_id]",
+    "[EOT]",
+    "<|media_begin|>",
+    "<|media_content|>",
+    "<|media_end|>",
+    "<|media_pad|>",
+    "<osagent_mode>"
+  ],
+  "bos_token": "[BOS]",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "[EOS]",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TikTokenTokenizer",
+  "unk_token": "[UNK]",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_kimi.TikTokenTokenizer",
+      null
+    ]
+  }
+}