# 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