smallthinker.py

# Copyright © 2026 Apple Inc.

# SmallThinker-4BA0.6B MLX-LM port.
#
# Architecture notes (checkpoint: PowerInfer/Tiiny SmallThinker-4BA0.6B-Instruct):
#   - ReLU-gated MoE experts: down(up(x) * relu(gate(x)))  -- NOT SwiGLU.
#   - Router normalization: top-k select, then (default) sigmoid(selected) / sum,
#     or softmax(selected) if moe_primary_router_apply_softmax is set.
#   - UNUSUAL router input: the router sees the block's ORIGINAL pre-attention input x,
#     while the experts see post_attention_layernorm(x + attn).
#   - GQA: 12 query heads, 2 KV heads, explicit head_dim 128 (NOT derived from hidden).
#
# Phase-1 checkpoint-compatible simplification: this exact checkpoint has
# rope_layout all-ones (RoPE on every layer) and sliding_window_layout all-zeros
# (full causal attention on every layer), so we apply RoPE and full causal attention
# uniformly and use a standard per-layer KVCache. layer_idx is plumbed through so that
# arbitrary rope_layout / sliding_window_layout (hybrid caches) can be added later.

from dataclasses import dataclass, field
from typing import Any, List, Optional

import mlx.core as mx
import mlx.nn as nn

from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
from .cache import KVCache
from .switch_layers import SwitchGLU


@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str = "smallthinker"
    hidden_size: int = 1536
    num_hidden_layers: int = 32
    num_attention_heads: int = 12
    num_key_value_heads: int = 2
    head_dim: int = 128
    vocab_size: int = 151936
    rms_norm_eps: float = 1e-6
    rope_theta: float = 1_500_000.0
    max_position_embeddings: int = 32768
    # MoE
    moe_num_primary_experts: int = 32
    moe_num_active_primary_experts: int = 4
    moe_ffn_hidden_size: int = 768
    moe_primary_router_apply_softmax: bool = False
    norm_topk_prob: bool = True
    # Layouts (per-layer). All-ones rope_layout / all-zero sliding_window_layout for
    # this checkpoint. Kept for validation + future hybrid support.
    rope_layout: List[int] = field(default_factory=lambda: [1] * 32)
    sliding_window_layout: List[int] = field(default_factory=lambda: [0] * 32)
    sliding_window_size: int = 4096
    tie_word_embeddings: bool = True

    def __post_init__(self):
        n = self.num_hidden_layers
        if len(self.rope_layout) != n:
            raise ValueError(
                f"rope_layout length {len(self.rope_layout)} != num_hidden_layers {n}"
            )
        if len(self.sliding_window_layout) != n:
            raise ValueError(
                f"sliding_window_layout length {len(self.sliding_window_layout)} "
                f"!= num_hidden_layers {n}"
            )
        for v in self.rope_layout:
            if v not in (0, 1):
                raise ValueError(f"rope_layout values must be in {{0,1}}, got {v}")
        for v in self.sliding_window_layout:
            if v not in (0, 1):
                raise ValueError(
                    f"sliding_window_layout values must be in {{0,1}}, got {v}"
                )
        if self.num_attention_heads % self.num_key_value_heads != 0:
            raise ValueError(
                f"num_attention_heads {self.num_attention_heads} not divisible by "
                f"num_key_value_heads {self.num_key_value_heads}"
            )
        if self.moe_num_active_primary_experts > self.moe_num_primary_experts:
            raise ValueError(
                f"moe_num_active_primary_experts "
                f"{self.moe_num_active_primary_experts} > moe_num_primary_experts "
                f"{self.moe_num_primary_experts}"
            )


class ReLUGLU(nn.Module):
    """Activation for SmallThinker MoE experts.

    SwitchGLU calls activation(x_up, x_gate); we return up * relu(gate), matching the
    reference expert: down(up(x) * relu(gate(x))).
    """

    def __call__(self, x_up: mx.array, x_gate: mx.array) -> mx.array:
        return x_up * nn.relu(x_gate)


class SmallThinkerAttention(nn.Module):
    def __init__(self, args: ModelArgs, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        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  # explicit, do NOT derive from hidden_size
        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)

        # Phase 1: this checkpoint has rope on every layer (rope_layout all-ones).
        self.use_rope = bool(args.rope_layout[layer_idx])
        self.rope = nn.RoPE(self.head_dim, traditional=False, base=args.rope_theta)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[Any] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        B, L, _ = x.shape

        queries = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
        keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
        values = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)

        if self.use_rope:
            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 SmallThinkerMoeBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.hidden_dim = args.hidden_size
        self.ffn_dim = args.moe_ffn_hidden_size
        self.num_experts = args.moe_num_primary_experts
        self.top_k = args.moe_num_active_primary_experts
        self.apply_softmax = args.moe_primary_router_apply_softmax

        self.primary_router = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
        self.switch_mlp = SwitchGLU(
            self.hidden_dim,
            self.ffn_dim,
            self.num_experts,
            activation=ReLUGLU(),
            bias=False,
        )

    def __call__(self, router_input: mx.array, expert_input: mx.array) -> mx.array:
        # Router sees the block's ORIGINAL pre-attention input (router_input), NOT the
        # expert input. This is the SmallThinker-specific wiring.
        gates = self.primary_router(router_input)

        k = self.top_k
        inds = mx.stop_gradient(
            mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k]
        )
        scores = mx.take_along_axis(gates, inds, axis=-1)

        if self.apply_softmax:
            scores = mx.softmax(scores, axis=-1, precise=True)
        else:
            scores = mx.sigmoid(scores)
            scores = scores / scores.sum(axis=-1, keepdims=True)

        y = self.switch_mlp(expert_input, inds)
        y = (y * scores[..., None]).sum(axis=-2)
        return y


class SmallThinkerDecoderLayer(nn.Module):
    def __init__(self, args: ModelArgs, layer_idx: int):
        super().__init__()
        self.self_attn = SmallThinkerAttention(args, layer_idx)
        self.block_sparse_moe = SmallThinkerMoeBlock(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[Any] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        router_input = x
        h = x + self.self_attn(self.input_layernorm(x), mask, cache)
        out = h + self.block_sparse_moe(
            router_input, self.post_attention_layernorm(h)
        )
        return out


class SmallThinkerModel(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
        self.layers = [
            SmallThinkerDecoderLayer(args, i) for i in range(args.num_hidden_layers)
        ]
        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)

    def __call__(
        self,
        inputs: mx.array,
        cache=None,
        input_embeddings: Optional[mx.array] = None,
    ) -> mx.array:
        if input_embeddings is not None:
            h = input_embeddings
        else:
            h = self.embed_tokens(inputs)

        if cache is None:
            cache = [None] * len(self.layers)

        mask = create_attention_mask(h, cache[0])

        for layer, c in zip(self.layers, cache):
            h = layer(h, mask, c)

        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 = SmallThinkerModel(args)
        if not args.tie_word_embeddings:
            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)

    def __call__(
        self,
        inputs: mx.array,
        cache=None,
        input_embeddings: Optional[mx.array] = None,
    ) -> mx.array:
        out = self.model(inputs, cache, input_embeddings)
        if self.args.tie_word_embeddings:
            return self.model.embed_tokens.as_linear(out)
        return self.lm_head(out)

    def sanitize(self, weights):
        # 1) Tied head handling. The checkpoint ships a separate lm_head.weight even
        #    though tie_word_embeddings=True. Assert it equals embed_tokens before
        #    dropping; if it differs, KEEP it (and flip to an untied head) and report.
        if "lm_head.weight" in weights:
            lm = weights["lm_head.weight"]
            emb = weights.get("model.embed_tokens.weight")
            if self.args.tie_word_embeddings:
                tied = emb is not None and lm.shape == emb.shape and mx.array_equal(lm, emb)
                if tied:
                    weights.pop("lm_head.weight", None)
                else:
                    print(
                        "[smallthinker.sanitize] WARNING: tie_word_embeddings=True but "
                        "lm_head.weight != embed_tokens.weight; keeping separate head."
                    )
                    self.args.tie_word_embeddings = False
                    if not hasattr(self, "lm_head"):
                        self.lm_head = nn.Linear(
                            self.args.hidden_size, self.args.vocab_size, bias=False
                        )

        # 2) Pack per-expert HF tensors into stacked SwitchGLU weights.
        #    HF:   model.layers.{l}.block_sparse_moe.experts.{e}.{up,gate,down}.{suffix}
        #    MLX:  model.layers.{l}.block_sparse_moe.switch_mlp.{up_proj,gate_proj,down_proj}.{suffix}
        # Idempotent: if already packed (or unpacked experts absent), leave as-is.
        prefix0 = "model.layers.0.block_sparse_moe.experts.0.up.weight"
        if prefix0 not in weights:
            return weights

        name_map = [("up", "up_proj"), ("gate", "gate_proj"), ("down", "down_proj")]
        for l in range(self.args.num_hidden_layers):
            base = f"model.layers.{l}.block_sparse_moe"
            for hf_name, mlx_name in name_map:
                for suffix in ("weight", "scales", "biases"):
                    first = f"{base}.experts.0.{hf_name}.{suffix}"
                    if first not in weights:
                        continue
                    to_join = [
                        weights.pop(
                            f"{base}.experts.{e}.{hf_name}.{suffix}"
                        )
                        for e in range(self.args.moe_num_primary_experts)
                    ]
                    weights[f"{base}.switch_mlp.{mlx_name}.{suffix}"] = mx.stack(
                        to_join
                    )
        return weights

    def make_cache(self):
        return [KVCache() for _ in range(self.args.num_hidden_layers)]

    @property
    def layers(self):
        return self.model.layers