"""Depth decoder: FastDepth._compute_fused (serve/fast_depth.py:143-203) for one seat, fp16 in production. Weights: the release's bf16 values -> fp32 (engine.py:161) -> the depth dtype (fast_depth.py:44). Per frame: prefix from the backbone state and cb0, 31 levels in weight groups (GROUP_OF_LEVEL), incremental K/V inside a shared group, Triton-kernel arithmetic for add+LayerNorm and the top-k Gumbel sampler, then the backbone feedback vector (audio_in.frames). See NOTES.md §5. """ import math import mlx.core as mx import numpy as np from . import quant as Q from . import weights as W GROUP_OF_LEVEL = [None] + [i for i in range(8)] + [8] * 8 + [9] * 8 + [10] * 7 # depth_v2.py:45 LAYERS, D, H = 4, 1024, 16 # depth_v2.py:47 HEAD = D // H N_LEVELS = 32 LN_EPS = 1e-5 # nn.LayerNorm default def gelu(x, dt): """torch F.gelu (erf) on a half/bf16 tensor: x*0.5*(1+erf(x*sqrt(1/2))) in fp32, one rounding (ActivationGeluKernel.cu:37-38).""" xf = x.astype(mx.float32) return (xf * 0.5 * (1.0 + mx.erf(xf * math.sqrt(0.5)))).astype(dt) def layer_norm(x, w, b, dt): """_add_ln's LayerNorm (fast_depth_kernels.py:27-32): fp32 statistics, biased variance, one rounding.""" xf = x.astype(mx.float32) mean = mx.mean(xf, -1, keepdims=True) d = xf - mean var = mx.mean(d * d, -1, keepdims=True) return (d / mx.sqrt(var + LN_EPS) * w + b).astype(dt) class Depth: def __init__(self, meta, audio_in, dtype=mx.float16, temp=0.7, top_k=64, path=W.DEFAULT_DIR): """audio_in: (table [65536, 2048], norm weight [2048]) in the backbone dtype (shared, engine dtype).""" self.dt = dtype self.top_k = int(top_k) self.inv_temp = float(np.float32(1.0 / temp)) # topk_sample passes 1.0/temp as an fp32 argument (:77) w = W.load(["depth.safetensors"], path) recipe = Q.load_recipe(path) def c(name): # bf16 -> fp32 -> depth dtype; a quantized weight was made a = Q.take(w, name, recipe) # from the depth-dtype copy (export/quantize.py) return a if isinstance(a, Q.QW) else a.astype(mx.float32).astype(dtype) self.in_proj = c("in_proj.weight") lvl = c("level_emb.weight") code_lo, code_up = (x.dense(dtype) if isinstance(x, Q.QW) else x for x in (c("code_lo.weight"), c("code_up.weight"))) # fast_depth.py:49-51: embed[k][code] = code_up(code_lo[code + k*2048]) + level_emb[k+1] table = (code_lo @ code_up.T).reshape(N_LEVELS - 1, 2048, D) self.embed = table + lvl[1:N_LEVELS, None, :] self.level0 = lvl[0] self.groups = [] for g in range(11): blocks = [] for li in range(LAYERS): p = f"groups.{g}.{li}." blocks.append({"n1w": w[p + "n1.weight"].astype(mx.float32), "n1b": w[p + "n1.bias"].astype(mx.float32), "n2w": w[p + "n2.weight"].astype(mx.float32), "n2b": w[p + "n2.bias"].astype(mx.float32), "qkv": c(p + "qkv.weight"), "o": c(p + "o.weight"), "f1": c(p + "f1.weight"), "f2": c(p + "f2.weight")}) self.groups.append(blocks) # LayerNorm weights are read as fp32 by the kernel from the depth-dtype copy: round them to the dtype first for blocks in self.groups: for b in blocks: for k in ("n1w", "n1b", "n2w", "n2b"): b[k] = b[k].astype(dtype).astype(mx.float32) self.norm_w = c("norm.weight").astype(mx.float32) self.norm_b = c("norm.bias").astype(mx.float32) self.heads = [c(f"heads.{i}.weight") for i in range(N_LEVELS - 1)] self.cross_kv = [c(f"cross.{i}.kv.weight") for i in range(LAYERS)] self.cross_o = [c(f"cross.{i}.o.weight") for i in range(LAYERS)] self.null_mem = w["null_mem.weight" if "null_mem.weight" in w else "null_mem"].astype(mx.float32) self.table, self.fb_w = audio_in self.fb_offsets = mx.arange(N_LEVELS) * 2048 self.fb_inv = float(np.float32(1.0) / np.float32(math.sqrt(N_LEVELS))) # CUDA: x * (1/b) (NOTES §5) self.cross = None mx.eval(Q.arrays([self.in_proj, self.embed, self.level0, self.groups, self.norm_w, self.norm_b, self.heads, self.cross_kv, self.cross_o, self.null_mem, self.table, self.fb_w])) def set_speaker(self, mem=None): """FastDepth.set_speaker (fast_depth.py:73-82) with one reference, ref_mask True: cross[layer] = o(v(mem)). mem None = the null token (named and plain voices: depth_v2.py:253-254).""" mem = (self.null_mem if mem is None else mem).reshape(1, D).astype(self.dt) self.cross = [] for li in range(LAYERS): kv = Q.mm(mem, self.cross_kv[li]) v = kv[:, D:] self.cross.append(Q.mm(v, self.cross_o[li])[0]) mx.eval(self.cross) def frame(self, h, cb0, uniform, force=None, logits_out=None): """h [hidden] (backbone dtype), cb0 scalar int array, uniform [31, 2048] fp32 (Philox, clamped) -> (codes [32] int32, feedback [hidden] in the backbone dtype). Gates only: `force` [32] int makes each level's prefix input the given code instead of the sampled one (teacher forcing inside the frame); `logits_out` (a list) receives every level's head output.""" dt = self.dt x0 = h.astype(dt)[None] # bf16 -> depth dtype (fast_depth.py:254) prefix = [None] * N_LEVELS prefix[0] = Q.mm(x0, self.in_proj)[0] + self.level0 prefix[1] = self.embed[0][cb0] kv = [[None] * N_LEVELS for _ in range(LAYERS)] # per layer: (k, v) [H, HEAD] per position codes = [cb0] previous = None for level in range(1, N_LEVELS): group = GROUP_OF_LEVEL[level] n = level + 1 incremental = group == previous reused_later = level + 1 < N_LEVELS and GROUP_OF_LEVEL[level + 1] == group x = prefix[level][None] if incremental else mx.stack(prefix[:n]) blocks = self.groups[group] y = layer_norm(x, blocks[0]["n1w"], blocks[0]["n1b"], dt) for li, b in enumerate(blocks): T = x.shape[0] qkv = Q.mm(y, b["qkv"]).reshape(T, 3, H, HEAD) q = qkv[:, 0].transpose(1, 0, 2)[None] # [1, H, T, 64] if incremental: kv[li][level] = (qkv[0, 1], qkv[0, 2]) k = mx.stack([kv[li][j][0] for j in range(n)], 1)[None] v = mx.stack([kv[li][j][1] for j in range(n)], 1)[None] a = mx.fast.scaled_dot_product_attention(q, k, v, scale=HEAD ** -0.5) else: if reused_later: for j in range(n): kv[li][j] = (qkv[j, 1], qkv[j, 2]) k = qkv[:, 1].transpose(1, 0, 2)[None] v = qkv[:, 2].transpose(1, 0, 2)[None] if li == LAYERS - 1: # only the final position is read after the last layer (fast_depth.py:116-119) a = mx.fast.scaled_dot_product_attention(q[:, :, -1:], k, v, scale=HEAD ** -0.5) x = x[-1:] else: a = mx.fast.scaled_dot_product_attention(q, k, v, scale=HEAD ** -0.5, mask="causal") T = x.shape[0] x = x + Q.mm(a[0].transpose(1, 0, 2).reshape(T, D), b["o"]) # add_ln: x' = dt(x + a) y2 = layer_norm(x, b["n2w"], b["n2b"], dt) f = Q.mm(gelu(Q.mm(y2, b["f1"]), dt), b["f2"]) x = x + f # dt(x + f) x = x + self.cross[li] # dt(x' + c) if li + 1 < LAYERS: nb = blocks[li + 1] y = layer_norm(x, nb["n1w"], nb["n1b"], dt) else: y = layer_norm(x, self.norm_w, self.norm_b, dt) logits = Q.mm(y[-1:], self.heads[level - 1])[0] # [2048], depth dtype if logits_out is not None: logits_out.append(logits) xs = logits.astype(mx.float32) * self.inv_temp # _topk_sample :53 kth = mx.topk(xs, self.top_k).min() # the K-th largest (:54-55) xm = mx.where(xs >= kth, xs, -mx.inf) score = xm - mx.log(-mx.log(uniform[level - 1])) code = mx.argmax(score).astype(mx.int32) codes.append(code) if n < N_LEVELS: prefix[n] = self.embed[level][code if force is None else int(force[level])] previous = group codes = mx.stack(codes) # feedback == audio_in.frames(codes, 32) (fast_depth.py:200-202, audio_io.py:27-35) e = Q.rows(self.table, codes + self.fb_offsets) # [32, hidden] s = e.astype(mx.float32).sum(0).astype(self.fb_w.dtype) # sum: fp32 accumulation, one rounding s = (s.astype(mx.float32) * self.fb_inv).astype(self.fb_w.dtype) # / sqrt(32) as CUDA: x * (1/b) sf = s.astype(mx.float32) fb = (sf * mx.rsqrt(mx.mean(sf * sf) + np.float32(2.0 ** -23)) * self.fb_w.astype(mx.float32)) return codes, fb.astype(self.fb_w.dtype)