File size: 9,718 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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)
|