Add files using upload-large-folder tool

23c283a verified about 19 hours ago

6.98 kB

"""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)

	"""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)