Upload folder using huggingface_hub

README.md

@@ -0,0 +1,38 @@
+---
+license: apache-2.0
+base_model: tencent/Hy3
+library_name: mlx
+pipeline_tag: text-generation
+language: en
+tags:
+  - mlx
+  - apple-silicon
+  - hy3
+  - mixture-of-experts
+  - mtp
+  - speculative-decoding
+  - mtplx
+---
+
+# Hy3-Demolition-MLX lite-v1-mtp
+
+The **MTP-equipped** variant of [lite-v1](https://huggingface.co/philipjohnbasile/hy3-demolition-mlx-lite-v1): its fused trunk with the Hy3 **NextN (Multi-Token-Prediction) sidecar** grafted back on (`num_nextn_predict_layers=1`, num_experts=192), for self-speculative decoding on MTPLX.
+
+## Status — honest
+
+This artifact is **MTP-equipped but not yet end-to-end runnable**. `mtplx inspect` (release v2.0.1) recognizes it — *"HY V3 MTP markers recognized, but MTPLX does not yet have a native MLX runtime backend for this family"* — i.e. it is **`recognized-backend-pending`**. The native `hy_v3` runtime backend is in flight as [MTPLX PR #142](https://github.com/youssofal/MTPLX/pull/142), gated on `hy_v3` reaching MTPLX's mlx-lm pin. Until that lands, use the AR daily-driver ([lite-v1](https://huggingface.co/philipjohnbasile/hy3-demolition-mlx-lite-v1)) instead.
+
+## How it was built
+
+The MTP head consumes the trunk's final hidden state (hidden_size 4096, unchanged by pruning) and runs its own MoE on the global `num_experts`. So the base checkpoint's `mtp.*` sidecar grafts directly onto the fused AR trunk — no re-heal, no re-prune of the trunk.
+
+
+Graft script + `mtplx inspect` receipts: https://github.com/PhilipJohnBasile/hy3-demolition-mlx (`scripts/38_mtp_sidecar_graft.py`, `eval/receipts/mtplx_inspect_*.json`).
+
+## Limitations
+
+- **Does not run on stock mlx_lm as MTP.** The fast MTP path needs the MTPLX backend; mlx-lm's own per-token self-speculative loop is ~4.7× *slower* than AR (measured), which is why the MTPLX batched-verify backend is the target.
+- Everything from the base lite-v1 card applies (quantized MoE, English/agent focus, no tool execution).
+- End-to-end MTP behavior is **unverified** until the backend loads it; recognition is structural (`mtplx inspect`), not a live run.
+
+Base recipe + receipts: https://github.com/PhilipJohnBasile/hy3-demolition-mlx

chat_template.jinja

@@ -0,0 +1,222 @@
+{#- ----------‑‑‑ special token variables ‑‑‑---------- -#}
+{%- set HYTK = ':opensource' %}
+{%- set eos_token = '<｜hy_eos{}｜>'.format(HYTK) %}
+{%- set bos_token = '<｜hy_begin_of_sentence{}｜>'.format(HYTK) %}
+{%- set pad_token = '<｜hy_pad{}｜>'.format(HYTK) %}
+{%- set user_token = '<｜hy_User{}｜>'.format(HYTK) %}
+{%- set assistant_token = '<｜hy_Assistant{}｜>'.format(HYTK) %}
+{%- set think_begin_token = '<think{}>'.format(HYTK) %}
+{%- set think_end_token = '</think{}>'.format(HYTK) %}
+{%- set toolcalls_begin_token = '<tool_calls{}>'.format(HYTK) %}
+{%- set toolcalls_end_token = '</tool_calls{}>'.format(HYTK) %}
+{%- set toolcall_begin_token = '<tool_call{}>'.format(HYTK) %}
+{%- set toolcall_end_token = '</tool_call{}>'.format(HYTK) %}
+{%- set toolsep_token = '<tool_sep{}>'.format(HYTK) %}
+{%- set argkey_begin_token = '<arg_key{}>'.format(HYTK) %}
+{%- set argkey_end_token = '</arg_key{}>'.format(HYTK) %}
+{%- set argvalue_begin_token = '<arg_value{}>'.format(HYTK) %}
+{%- set argvalue_end_token = '</arg_value{}>'.format(HYTK) %}
+{%- set toolresponses_begin_token = '<tool_responses{}>'.format(HYTK) %}
+{%- set toolresponses_end_token = '</tool_responses{}>'.format(HYTK) %}
+{%- set toolresponse_begin_token = '<tool_response{}>'.format(HYTK) %}
+{%- set toolresponse_end_token = '</tool_response{}>'.format(HYTK) %}
+{%- set reasoning_mode_token = '<｜reasoning_mode{}｜>'.format(HYTK) %}
+
+{#- ----------‑‑‑ hyperparameters variables ‑‑‑---------- -#}
+{%- if not add_generation_prompt is defined %}
+    {%- set add_generation_prompt = false %}
+{%- endif %}
+{%- if not preserved_thinking is defined %}
+    {%- if not tools %}
+        {%- set preserved_thinking = false %}
+    {%- else %}
+        {%- set preserved_thinking = true %}
+    {%- endif %}
+{%- endif %}
+{%- if not is_training is defined %}
+    {%- set is_training = false %}
+{%- endif %}
+
+{%- if not reasoning_effort is defined %}
+    {%- set reasoning_effort = 'high' %}
+{%- elif reasoning_effort not in ['high', 'low', 'no_think'] %}
+    {%- if reasoning_effort is none %}
+        {{- raise_exception('reasoning_effort error : None, should be no_think/low/high') }}
+    {%- else %}
+        {{- raise_exception('reasoning_effort error : ' + reasoning_effort + ', should be no_think/low/high') }}
+    {%- endif %}
+{%- endif %}
+
+{%- if fallback_strategy is defined and fallback_strategy == 'reasoning_toolcall_retry' %}
+    {%- set reasoning_effort = 'high' %}
+    {%- set add_generation_prompt = false %}
+{%- endif %}
+{%- if not raw_last_assistant is defined %}
+    {%- set raw_last_assistant = false %}
+{%- endif %}
+
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- elif content is none -%}
+        {{- '' }}
+    {%- else -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- set ns = namespace(last_user_index=-1) %}
+{%- set sp_ns = namespace(system_prompt='', is_first_sp=true) %}
+{%- for message in messages %}
+    {%- if message['role'] == 'system' %}
+        {%- set sp_ns.system_prompt = sp_ns.system_prompt + visible_text(message['content']) %}
+    {%- endif %}
+    {%- if message['role'] == 'user' %}
+        {%- set ns.last_user_index = loop.index0 %}
+    {%- endif %}
+{%- endfor %}
+{%- if reasoning_effort is defined and reasoning_effort is string and reasoning_effort != '' and not tools %}
+    {%- set sp_ns.system_prompt = sp_ns.system_prompt + reasoning_mode_token + 'reasoning_effort:' + reasoning_effort %}
+{%- endif %}
+{{- bos_token }}
+{{- sp_ns.system_prompt }}
+{%- if tools %}
+    {%- if sp_ns.system_prompt != '' %}
+        {{- '\n\n# Tools\n\nYou may call one or more functions to assist with the user query.' }}
+    {%- else %}
+        {{- '# Tools\n\nYou may call one or more functions to assist with the user query.' }}
+    {%- endif %}
+    {{- '\n\nYou are provided with function signatures within <tools></tools> XML tags:' }}
+    {{- '\n<tools>\n' }}
+    {%- for tool in tools %}
+        {%- if loop.index0 > 0 %}
+            {{- '\n' }}
+        {%- endif %}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- '\n</tools>\n\n' }}
+    {{- 'For function call returns, you should first print ' + toolcalls_begin_token + '\n' }}
+    {{- 'For each function call, you should return object like:\n' }}
+    {{- toolcall_begin_token + '{function-name}' + toolsep_token + '\n' }}
+    {{- argkey_begin_token + '{arg-key-1}' + argkey_end_token + '\n' }}
+    {{- argvalue_begin_token + '{arg-value-1}' + argvalue_end_token + '\n' }}
+    {{- argkey_begin_token + '{arg-key-2}' + argkey_end_token + '\n' }}
+    {{- argvalue_begin_token + '{arg-value-2}' + argvalue_end_token + '\n' }}
+    {{- '...\n' }}
+    {{- toolcall_end_token + '\n' }}
+    {%- if reasoning_effort is defined and reasoning_effort is string and reasoning_effort != '' %}
+        {{- 'At the end of function call returns, you should print ' + toolcalls_end_token + reasoning_mode_token + 'reasoning_effort:' + reasoning_effort }}
+    {%- else %}
+        {{- 'At the end of function call returns, you should print ' + toolcalls_end_token }}
+    {%- endif %}
+{%- endif %}
+
+{%- set prev_ns = namespace(is_tool=false, is_tool_first=true) %}
+{%- set last_ns = namespace(last_is_assistant=false) %}
+{%- for message in messages %}
+    {%- if message['role'] == 'user' %}
+        {%- if prev_ns.is_tool %}
+            {{- toolresponses_end_token }}
+        {%- endif %}
+        {{- user_token + visible_text(message['content']) }}
+        {%- set prev_ns.is_tool = false %}
+    {%- endif %}
+    {%- if message['role'] == 'assistant' %}
+        {%- if is_training %}
+            {%- if 'reasoning_content' in message and message['reasoning_content'] is string %}
+                {%- set rc = message['reasoning_content'] %}
+            {%- elif 'reasoning' in message and message['reasoning'] is string %}
+                {%- set rc = message['reasoning'] %}
+            {%- else %}
+                {%- set rc = none %}
+            {%- endif %}
+            {%- if rc is not none %}
+                {%- set content = think_begin_token + rc + think_end_token + visible_text(message['content']) %}
+            {%- else %}
+                {%- set content = think_begin_token + think_end_token + visible_text(message['content']) %}
+            {%- endif %}
+        {%- else %}
+            {%- if ((preserved_thinking is defined and preserved_thinking) or loop.index0 > ns.last_user_index) %}
+                {%- if 'reasoning_content' in message and message['reasoning_content'] is string %}
+                    {%- set rc = message['reasoning_content'] %}
+                {%- elif 'reasoning' in message and message['reasoning'] is string %}
+                    {%- set rc = message['reasoning'] %}
+                {%- else %}
+                    {%- set rc = none %}
+                {%- endif %}
+                {%- if rc is not none %}
+                    {%- set content = think_begin_token + rc + think_end_token + visible_text(message['content']) %}
+                {%- else %}
+                    {%- set content = think_begin_token + think_end_token + visible_text(message['content']) %}
+                {%- endif %}
+            {%- else %}
+                {%- set content = think_begin_token + think_end_token + visible_text(message['content']) %}
+            {%- endif %}
+        {%- endif %}
+        {%- if prev_ns.is_tool %}
+            {{- toolresponses_end_token }}
+        {%- endif %}
+        {{- assistant_token }}
+        {%- if message['tool_calls'] is defined and message['tool_calls'] %}
+            {%- set prev_ns.is_tool_first = true %}
+            {{- content }}
+            {{- toolcalls_begin_token + '\n' }}
+            {%- for tool in message['tool_calls'] %}
+                {%- set arguments = tool['function']['arguments'] %}
+                {{- toolcall_begin_token + tool['function']['name'] + toolsep_token + '\n' }}
+                {%- for key, value in arguments.items() %}
+                    {{- argkey_begin_token + key + argkey_end_token + '\n' }}
+                    {%- if value is not string %}
+                        {%- set value = value | tojson(ensure_ascii=False) %}
+                    {%- endif %}
+                    {{- argvalue_begin_token + value + argvalue_end_token + '\n' }}
+                {%- endfor %}
+                {{- toolcall_end_token + '\n' }}
+            {%- endfor %}
+            {{- toolcalls_end_token + eos_token }}
+        {%- else %}
+            {%- if loop.last and raw_last_assistant %}
+                {{- visible_text(message['content']) }}
+            {%- elif not loop.last or is_training %}
+                {{- content + eos_token }}
+            {%- else %}
+                {{- content }}
+            {%- endif %}
+        {%- endif %}
+        {%- set prev_ns.is_tool = false %}
+    {%- endif %}
+    {%- if message['role'] == 'tool' %}
+        {%- set prev_ns.is_tool = true %}
+        {%- if prev_ns.is_tool_first %}
+            {{- toolresponses_begin_token + '\n' }}
+            {%- set prev_ns.is_tool_first = false %}
+        {%- endif %}
+        {{- toolresponse_begin_token + '\n' + visible_text(message['content']) + '\n' + toolresponse_end_token + '\n' }}
+    {%- endif %}
+    {%- if loop.last and message['role'] == 'assistant' %}
+        {%- set last_ns.last_is_assistant = true %}
+    {%- endif %}
+
+{%- endfor %}
+{%- if prev_ns.is_tool %}
+    {{- toolresponses_end_token }}
+{%- endif %}
+{%- if add_generation_prompt %}
+    {%- if not last_ns.last_is_assistant %}
+        {%- if reasoning_effort is defined and reasoning_effort in ['low', 'high'] %}
+            {{- assistant_token + think_begin_token }}
+        {%- elif reasoning_effort is defined and reasoning_effort == 'no_think' %}
+            {{- assistant_token + think_begin_token + think_end_token }}
+        {%- else %}
+            {{- assistant_token }}
+        {%- endif %}
+    {%- endif %}
+{%- endif %}

hy_v3.py

@@ -0,0 +1,467 @@
+# Copyright © 2026 Apple Inc.
+#
+# Tencent Hunyuan 3 (hy_v3). Base model support follows the community work in
+# ml-explore/mlx-lm#1211 (kernelpool); this file additionally *keeps and uses*
+# the Multi-Token-Prediction (MTP) layer for self-speculative decoding instead
+# of stripping it.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache
+from .pipeline import PipelineMixin
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    num_experts: int
+    num_experts_per_tok: int
+    num_shared_experts: int
+    expert_hidden_dim: int
+    first_k_dense_replace: int
+    rms_norm_eps: float
+    rope_parameters: Dict[str, Any]
+    router_scaling_factor: float = 1.0
+    qk_norm: bool = True
+    route_norm: bool = True
+    moe_router_use_sigmoid: bool = True
+    moe_router_enable_expert_bias: bool = True
+    tie_word_embeddings: bool = False
+    num_nextn_predict_layers: int = 0
+    max_position_embeddings: int = 262144
+    enable_moe_fp32_combine: bool = False
+    enable_lm_head_fp32: bool = False
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
+
+        self.use_qk_norm = args.qk_norm
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+            self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            dims=self.head_dim,
+            base=args.rope_parameters["rope_theta"],
+            traditional=False,
+            scaling_config=args.rope_parameters,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim)
+        keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+        values = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+
+        if self.use_qk_norm:
+            queries = self.q_norm(queries)
+            keys = self.k_norm(keys)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.transpose(0, 2, 1, 3)
+
+        offset = cache.offset if cache is not None else 0
+        queries = self.rope(queries, offset=offset)
+        keys = self.rope(keys, offset=offset)
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+@mx.compile
+def expert_select(
+    gates,
+    expert_bias,
+    top_k,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + expert_bias
+
+    inds = mx.argpartition(scores, kth=-top_k, axis=-1)[..., -top_k:]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20)
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.route_norm
+        self.routed_scaling_factor = args.router_scaling_factor
+        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.expert_bias = mx.zeros((args.num_experts,))
+
+    def __call__(self, x):
+        return expert_select(
+            self.gate(x),
+            self.expert_bias,
+            self.top_k,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts_per_tok = args.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.expert_hidden_dim,
+            args.num_experts,
+        )
+        self.router = MoEGate(args)
+        if args.num_shared_experts > 0:
+            self.shared_mlp = MLP(
+                args.hidden_size,
+                args.expert_hidden_dim * args.num_shared_experts,
+            )
+        else:
+            self.shared_mlp = None
+
+        self.fp32_combine = args.enable_moe_fp32_combine
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        inds, scores = self.router(x)
+        if not self.fp32_combine:
+            scores = scores.astype(x.dtype)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        if self.shared_mlp is not None:
+            y = y + self.shared_mlp(x)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y.astype(x.dtype)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(args)
+        if layer_idx < args.first_k_dense_replace:
+            self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        else:
+            self.mlp = MoE(args)
+        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
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class MTPBlock(nn.Module):
+    """Hy3 Multi-Token-Prediction block (the layer after the main stack).
+
+    Projects concat[norm(next-token embedding), norm(hidden state)] through
+    ``eh_proj`` and one full decoder layer to produce the hidden state for the
+    speculatively-drafted next token.
+    """
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.enorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.hnorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.eh_proj = nn.Linear(args.hidden_size * 2, args.hidden_size, bias=False)
+        self.layer = DecoderLayer(args, layer_idx=args.num_hidden_layers)
+        self.final_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        h_N: mx.array,
+        e_N1: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        # Order matters: [normed embedding, normed hidden state].
+        x = mx.concatenate([self.enorm(e_N1), self.hnorm(h_N)], axis=-1)
+        y = self.layer(self.eh_proj(x), mask, cache)
+        return self.final_layernorm(y)
+
+
+class HYV3Model(PipelineMixin, nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.vocab_size = args.vocab_size
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(args, idx) for idx in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+        return_hidden_states: bool = False,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
+        if cache is None:
+            cache = [None] * len(self.pipeline_layers)
+        mask = create_attention_mask(h, cache[0])
+
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for layer, c in zip(self.pipeline_layers, cache):
+            h = layer(h, mask, cache=c)
+
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        out = self.norm(h)
+        if return_hidden_states:
+            return out, h
+        return out
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = HYV3Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+        self.num_nextn_predict_layers = getattr(args, "num_nextn_predict_layers", 0)
+        if self.num_nextn_predict_layers > 0:
+            self.mtp = MTPBlock(args)
+
+    def _logits(self, out):
+        if self.args.enable_lm_head_fp32:
+            out = out.astype(mx.float32)
+        if self.args.tie_word_embeddings:
+            return self.model.embed_tokens.as_linear(out)
+        return self.lm_head(out)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        return_hidden_states: bool = False,
+    ):
+        if return_hidden_states:
+            out, h = self.model(inputs, cache, return_hidden_states=True)
+            return self._logits(out), h
+        out = self.model(inputs, cache)
+        return self._logits(out)
+
+    def predict_next_tokens(self, h_N: mx.array, token_ids: mx.array, cache=None):
+        """Run the MTP head to draft the next token from a hidden state."""
+        if not hasattr(self, "mtp"):
+            raise ValueError("MTP is not enabled or its weights are not loaded.")
+        e_N1 = self.model.embed_tokens(token_ids)
+        mask = create_attention_mask(e_N1, cache)
+        h_mtp = self.mtp(h_N, e_N1, mask, cache)
+        return self._logits(h_mtp)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [KVCache() for _ in self.layers]
+
+    def sanitize(self, weights):
+        n_layers = self.args.num_hidden_layers
+        n_mtp = self.args.num_nextn_predict_layers
+
+        # Keep the MTP layer (the base model drops it). If the checkpoint stores
+        # it under model.layers.{n_layers}.*, remap it onto the mtp.* submodule;
+        # if it is already stored under mtp.*, leave it as-is.
+        if n_mtp > 0:
+            mtp_src = f"model.layers.{n_layers}."
+            for k in list(weights.keys()):
+                if k.startswith(mtp_src):
+                    rest = k[len(mtp_src):]
+                    if any(
+                        t in rest
+                        for t in ("enorm", "hnorm", "eh_proj", "final_layernorm")
+                    ):
+                        weights["mtp." + rest] = weights.pop(k)
+                    else:
+                        weights["mtp.layer." + rest] = weights.pop(k)
+
+        def fix_moe(prefix):
+            bias_key = f"{prefix}.mlp.expert_bias"
+            if bias_key in weights:
+                weights[f"{prefix}.mlp.router.expert_bias"] = weights.pop(bias_key)
+            for m in ("gate_proj", "down_proj", "up_proj"):
+                for k in ("weight", "scales", "biases"):
+                    per_expert = f"{prefix}.mlp.experts.0.{m}.{k}"
+                    stacked = f"{prefix}.mlp.experts.{m}.{k}"
+                    if per_expert in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+                    elif stacked in weights:
+                        # Already stacked (MLX-converted checkpoint): just rename.
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = weights.pop(
+                            stacked
+                        )
+
+        for l in range(n_layers):
+            fix_moe(f"model.layers.{l}")
+        if n_mtp > 0:
+            fix_moe("mtp.layer")
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return weights
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads = max(1, layer.self_attn.n_kv_heads // N)
+
+            if isinstance(layer.mlp, MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.mlp.sharding_group = group
+                if layer.mlp.shared_mlp is not None:
+                    shard_inplace(
+                        layer.mlp.shared_mlp.gate_proj, "all-to-sharded", group=group
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_mlp.down_proj, "sharded-to-all", group=group
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_mlp.up_proj, "all-to-sharded", group=group
+                    )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
+    @property
+    def layers(self):
+        return self.model.pipeline_layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.router.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "expert_bias" not in k
+
+        return predicate

model-00001-of-00022.safetensors

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

model-00002-of-00022.safetensors

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

model-00003-of-00022.safetensors

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

model-00004-of-00022.safetensors

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

model-00005-of-00022.safetensors

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

model-00006-of-00022.safetensors

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

model-00007-of-00022.safetensors

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

model-00008-of-00022.safetensors

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

model-00009-of-00022.safetensors

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

model-00010-of-00022.safetensors

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

model-00011-of-00022.safetensors

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

model-00012-of-00022.safetensors

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

model-00013-of-00022.safetensors

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

model-00014-of-00022.safetensors

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

model-00015-of-00022.safetensors

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

model-00016-of-00022.safetensors

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

model-00017-of-00022.safetensors

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

model-00018-of-00022.safetensors

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

model-00019-of-00022.safetensors

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

model-00020-of-00022.safetensors

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

model-00021-of-00022.safetensors

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

model-00022-of-00022.safetensors

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

model-mtp-sidecar.safetensors

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

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<｜hy_begin_of_sentence:opensource｜>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<｜hy_eos:opensource｜>",
+  "fix_mistral_regex": true,
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<｜hy_pad:opensource｜>",
+  "token_suffix": ":opensource",
+  "tokenizer_class": "TokenizersBackend",
+  "tool_parser_type": "hy_v3_opensource"
+}