Add files using upload-large-folder tool

23c283a verified about 19 hours ago

9.72 kB

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

	"""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: x0.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)