File size: 6,980 Bytes
23c283a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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)
|