"""Qwen3 backbone with the official engine's arithmetic (serve/fast_backbone.py, fuse_rope=True path). One seat: a static KV cache of `capacity` positions. Per layer (NOTES.md ยง3): h = h + pend (one rounding) y = RMSNorm(h) (mx.fast.rms_norm: Qwen3's two roundings) qkv = y @ [Wq;Wk;Wv]^T q, k = per-head RMSNorm (two roundings), RoPE in fp32 on the bf16 cos/sin table, attention over keys 0..pos (fp32 softmax, one rounding) h = h + o_proj(attn) y2 = RMSNorm(h) pend = down(silu(g) * u), silu in fp32 rounded once final: h = h + pend, out = RMSNorm(h, norm). `dtype` is bf16 in production, fp32 for the strict gates. """ import mlx.core as mx from . import quant as Q from . import weights as W class Backbone: def __init__(self, meta, dtype=mx.bfloat16, path=W.DEFAULT_DIR, capacity=None): cfg = meta["backbone"] self.dt = dtype self.layers_n = cfg["num_hidden_layers"] self.heads = cfg["num_attention_heads"] self.kv_heads = cfg["num_key_value_heads"] self.head_dim = cfg["head_dim"] self.hidden = cfg["hidden_size"] self.eps = float(cfg["rms_norm_eps"]) self.scale = self.head_dim ** -0.5 # fast_backbone.py:125 self.capacity = int(capacity or meta["capacity"]) w = W.load(W.backbone_files(meta), path) recipe = Q.load_recipe(path) # quantized weights (quant.py); None = the bf16 release cast = (lambda a: a) if dtype == mx.bfloat16 else (lambda a: Q.cast(a, dtype)) take = lambda n: cast(Q.take(w, n, recipe)) self.embed = take("embed_tokens.weight") self.norm = take("norm.weight") self.cb0_w = take("cb0_head.weight") # the table StaticBackbone holds (fast_backbone.py:151-153) for this model dtype: in bf16 it is built from the # bf16-rounded inv_freq buffer (the engine's model.to(bf16) casts it) and stored in bf16; the RoPE maths # reads it as fp32 tag = "bf16" if dtype == mx.bfloat16 else "fp32" if dtype not in (mx.bfloat16, mx.float32): raise ValueError("the backbone runs in bf16 (production) or fp32 (strict)") self.cos = w.pop(f"rope_{tag}.cos")[: self.capacity].astype(mx.float32) self.sin = w.pop(f"rope_{tag}.sin")[: self.capacity].astype(mx.float32) for t in ("fp32", "bf16"): w.pop(f"rope_{t}.cos", None) w.pop(f"rope_{t}.sin", None) self.layers = [] for i in range(self.layers_n): p = f"layers.{i}." g = lambda n: take(p + n) L = {"in_norm": g("input_layernorm.weight"), "post_norm": g("post_attention_layernorm.weight"), "q_norm": g("self_attn.q_norm.weight"), "k_norm": g("self_attn.k_norm.weight"), "o": g("self_attn.o_proj.weight"), "down": g("mlp.down_proj.weight")} # fuse_proj (fast_backbone.py:131-141): one GEMM for q|k|v and one for gate|up. A quantized export stores # them fused (rows concatenated before quantizing, the same groups as quantizing each part) if any(k.startswith(p + "self_attn.qkv_proj.weight") for k in w): L["qkv"], L["gu"] = g("self_attn.qkv_proj.weight"), g("mlp.gate_up_proj.weight") else: L["qkv"] = mx.concatenate([g("self_attn.q_proj.weight"), g("self_attn.k_proj.weight"), g("self_attn.v_proj.weight")], 0) L["gu"] = mx.concatenate([g("mlp.gate_proj.weight"), g("mlp.up_proj.weight")], 0) self.layers.append(L) if w: raise RuntimeError(f"unused backbone tensors: {sorted(w)[:8]}") mx.eval(Q.arrays([self.embed, self.norm, self.cb0_w, self.cos, self.sin, self.layers])) self.reset() def reset(self): shape = (1, self.kv_heads, self.capacity, self.head_dim) self.k = [mx.zeros(shape, self.dt) for _ in range(self.layers_n)] self.v = [mx.zeros(shape, self.dt) for _ in range(self.layers_n)] self.length = 0 # ------------------------------------------------------------------ pieces def embed_tokens(self, ids): return Q.rows(self.embed, mx.array(ids, mx.int32)).astype(self.dt) def _rope(self, x, pos): """x [T, heads, 128] normed (model dtype); fp32 maths on the bf16 cos/sin, one rounding (fast_backbone.py:79-84).""" half = self.head_dim // 2 c = self.cos[pos][:, None, :] s = self.sin[pos][:, None, :] xf = x.astype(mx.float32) y1, y2 = xf[..., :half], xf[..., half:] o1 = y1 * c[..., :half] - y2 * s[..., :half] o2 = y2 * c[..., half:] + y1 * s[..., half:] return mx.concatenate([o1, o2], -1).astype(self.dt) def forward(self, x): """x [T, hidden] (model dtype) at positions length..length+T-1 -> final-norm hidden states [T, hidden].""" T = x.shape[0] p0 = self.length if p0 + T > self.capacity: raise ValueError(f"context {p0 + T} exceeds capacity {self.capacity}") pos = mx.arange(p0, p0 + T) H, KV, D = self.heads, self.kv_heads, self.head_dim h = x pend = None for i, L in enumerate(self.layers): if pend is not None: h = h + pend # _add_rms: one rounding (:99) y = mx.fast.rms_norm(h, L["in_norm"], self.eps) qkv = Q.mm(y, L["qkv"]) q = mx.fast.rms_norm(qkv[:, : H * D].reshape(T, H, D), L["q_norm"], self.eps) k = mx.fast.rms_norm(qkv[:, H * D:(H + KV) * D].reshape(T, KV, D), L["k_norm"], self.eps) v = qkv[:, (H + KV) * D:].reshape(T, KV, D) q = self._rope(q, pos) k = self._rope(k, pos) self.k[i][:, :, p0:p0 + T, :] = k.transpose(1, 0, 2)[None] self.v[i][:, :, p0:p0 + T, :] = v.transpose(1, 0, 2)[None] K = self.k[i][:, :, : p0 + T] V = self.v[i][:, :, : p0 + T] a = mx.fast.scaled_dot_product_attention(q.transpose(1, 0, 2)[None], K, V, scale=self.scale, mask="causal" if T > 1 else None) h = h + Q.mm(a[0].transpose(1, 0, 2).reshape(T, H * D), L["o"]) y2 = mx.fast.rms_norm(h, L["post_norm"], self.eps) gu = Q.mm(y2, L["gu"]) n = gu.shape[-1] // 2 g = gu[:, :n].astype(mx.float32) act = (g / (1.0 + mx.exp(-g))).astype(self.dt) # torch silu: fp32, one rounding pend = Q.mm(act * gu[:, n:], L["down"]) h = h + pend self.length = p0 + T return mx.fast.rms_norm(h, self.norm, self.eps) def cb0_logits(self, h): """cb0_head (2048 -> 2049, no bias), model dtype, then fp32 (engine.py:657).""" return Q.mm(h, self.cb0_w).astype(mx.float32)