# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import json import math from dataclasses import asdict, dataclass from pathlib import Path import torch import ttnn from models.common.lightweightmodule import LightweightModule TILE = 32 BLOCK = 32 CONTEXT = 2 * BLOCK MASK_NEG = -1e9 CONV0_CHUNK = 1024 CONV0_SINGLE_ROWS = 2048 AUDIO_PREFIX = "audio_tower." EMBEDDER_PREFIX = "embed_audio." DTYPES = {"bfloat16": ttnn.bfloat16, "float32": ttnn.float32, "bfloat8_b": ttnn.bfloat8_b, "bfloat4_b": ttnn.bfloat4_b} FIDELITY = { "HiFi4": ttnn.MathFidelity.HiFi4, "HiFi3": ttnn.MathFidelity.HiFi3, "HiFi2": ttnn.MathFidelity.HiFi2, "LoFi": ttnn.MathFidelity.LoFi, } @dataclass(frozen=True) class AudioConfig: hidden_size: int num_hidden_layers: int num_attention_heads: int head_dim: int conv_channels: tuple conv_kernel_size: int chunk_size: int window: int logit_cap: float residual_weight: float rms_norm_eps: float output_proj_dims: int text_hidden_size: int feature_size: int @classmethod def from_hf(cls, hf_config): if isinstance(hf_config, (str, Path)): hf_config = json.loads(Path(hf_config).read_text()) a = hf_config["audio_config"] if a["attention_context_right"] != 0: raise ValueError("audio attention with right context is not supported") window = a["attention_context_left"] - 1 if window - 1 > BLOCK: raise ValueError(f"audio window {window} exceeds block {BLOCK}") return cls( hidden_size=a["hidden_size"], num_hidden_layers=a["num_hidden_layers"], num_attention_heads=a["num_attention_heads"], head_dim=a["hidden_size"] // a["num_attention_heads"], conv_channels=tuple(a["subsampling_conv_channels"]), conv_kernel_size=a["conv_kernel_size"], chunk_size=a["attention_chunk_size"], window=window, logit_cap=float(a["attention_logit_cap"]), residual_weight=float(a["residual_weight"]), rms_norm_eps=float(a["rms_norm_eps"]), output_proj_dims=a["output_proj_dims"], text_hidden_size=hf_config["text_config"]["hidden_size"], feature_size=128, ) WEIGHT_GROUPS = ("conv0", "qkv", "post", "ffn_up", "ffn_down", "lconv_start", "lconv_end", "proj") @dataclass(frozen=True) class AudioPolicy: residual: str = "float32" attention: str = "bfloat16" weights: tuple = tuple((g, "bfloat16") for g in WEIGHT_GROUPS) fidelities: tuple = tuple((g, "HiFi4") for g in WEIGHT_GROUPS) attention_fidelity: str = "HiFi4" conv_out: str = "bfloat16" small_m: int = 256 l1_rows: int = 256 small_m_glu_slices: bool = True matmul_table: str = "r1" def dtype(self, group): return DTYPES[dict(self.weights)[group]] def fidelity(self, group): return FIDELITY[dict(self.fidelities)[group]] def describe(self): return asdict(self) def padded_tokens(n): return BLOCK * max(1, math.ceil(n / BLOCK)) def subsample_len(n): return (n + 1) // 2 @dataclass class AudioHostInputs: features: torch.Tensor mask1: torch.Tensor attn_bias: torch.Tensor gather_idx: torch.Tensor batch: int tokens_padded: int valid_tokens: tuple out_padded: int @dataclass class AudioDeviceInputs: features: ttnn.Tensor mask1: ttnn.Tensor attn_bias: ttnn.Tensor gather_idx: ttnn.Tensor host: AudioHostInputs def attention_bias(valid_tokens, tokens_padded, window): nb = tokens_padded // BLOCK q = torch.arange(tokens_padded).view(nb, BLOCK, 1) k = (torch.arange(nb).view(nb, 1, 1) - 1) * BLOCK + torch.arange(CONTEXT).view(1, 1, CONTEXT) dist = q - k band = (dist >= 0) & (dist < window) & (k >= 0) out = torch.full((len(valid_tokens), 1, tokens_padded, CONTEXT), MASK_NEG) for b, n in enumerate(valid_tokens): ok = (band & (k < n)).reshape(tokens_padded, CONTEXT) out[b, 0][ok] = 0.0 return out def prepare_host_inputs(config, input_features, input_features_mask, tokens_padded=None, out_padded=None): b, t, f = input_features.shape if f != config.feature_size: raise ValueError(f"feature size {f} != {config.feature_size}") mask = input_features_mask.bool() frames = mask.sum(dim=-1).tolist() for i, n in enumerate(frames): if not bool(mask[i, :n].all()): raise ValueError("padding frames must follow valid frames") valid = tuple(subsample_len(subsample_len(n)) for n in frames) tp = tokens_padded if tokens_padded is not None else padded_tokens(subsample_len(subsample_len(t))) if max(valid) > tp or subsample_len(subsample_len(t)) > tp: raise ValueError(f"{t} frames exceed tokens_padded {tp}") t1 = 2 * tp feats = torch.zeros(b, 2 * t1, f) feats[:, :t] = input_features.float() * mask[..., None].float() mask1 = torch.zeros(b, t1) for i, n in enumerate(frames): mask1[i, : subsample_len(n)] = 1.0 total = sum(valid) op = out_padded if out_padded is not None else padded_tokens(total) if total > op: raise ValueError(f"{total} audio tokens exceed out_padded {op}") gather = torch.full((1, op), b * tp, dtype=torch.int32) gather[0, :total] = torch.cat([i * tp + torch.arange(n, dtype=torch.int32) for i, n in enumerate(valid)]) return AudioHostInputs( features=feats.reshape(b, t1, 2 * f), mask1=mask1.reshape(1, 1, b * t1, 1), attn_bias=attention_bias(valid, tp, config.window), gather_idx=gather, batch=b, tokens_padded=tp, valid_tokens=valid, out_padded=op, ) INPUT_SPEC = { "features": (ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT), "mask1": (ttnn.bfloat16, ttnn.TILE_LAYOUT), "attn_bias": (ttnn.bfloat16, ttnn.TILE_LAYOUT), "gather_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), } def host_inputs_to_device(host, mesh_device): tensors = { name: ttnn.from_torch( getattr(host, name).contiguous(), dtype=dt, layout=layout, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) for name, (dt, layout) in INPUT_SPEC.items() } return AudioDeviceInputs(host=host, **tensors) def copy_host_inputs(host, device_inputs): old = device_inputs.host for f in ("batch", "tokens_padded", "out_padded"): if getattr(old, f) != getattr(host, f): raise ValueError(f"{f} differs from the captured bucket: {getattr(host, f)} vs {getattr(old, f)}") for name, (dt, layout) in INPUT_SPEC.items(): src = ttnn.from_torch(getattr(host, name).contiguous(), dtype=dt, layout=layout) ttnn.copy_host_to_device_tensor(src, getattr(device_inputs, name)) device_inputs.host = host def _compute_config(mesh_device, fidelity=ttnn.MathFidelity.HiFi4, fp32_acc=True, l1_acc=True): return ttnn.init_device_compute_kernel_config( mesh_device.arch(), math_fidelity=fidelity, math_approx_mode=False, fp32_dest_acc_en=fp32_acc, packer_l1_acc=l1_acc, ) def _put(t, mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT): return ttnn.from_torch( t.contiguous(), dtype=dtype, layout=layout, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG ) def _norm_weight(w, mesh_device): return _put(w.float().reshape(1, 1, w.numel() // TILE, TILE), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT) def _tile_weight(w, mesh_device): return _put(w.float().reshape(1, 1, 1, -1), mesh_device) def _bmm_config(grid, in0_block_w, per_core_n, subblock_w): return ttnn.MatmulMultiCoreReuseProgramConfig( compute_with_storage_grid_size=grid, in0_block_w=in0_block_w, out_subblock_h=1, out_subblock_w=subblock_w, per_core_M=1, per_core_N=per_core_n, ) def _small_m_config(grid, m_rows, k, n, activation): mt, kt, nt = m_rows // TILE, k // TILE, n // TILE cores = grid.x * grid.y pcn = -(-nt // cores) while nt % pcn: pcn += 1 cols = nt // pcn gx = min(grid.x, cols) gy = -(-cols // gx) kw = next(v for v in (8, 4, 2, 1) if kt % v == 0) sw = next(v for v in (4, 2, 1) if pcn % v == 0 and v * mt <= 4) if mt <= 4 else 1 return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig( compute_with_storage_grid_size=(gx, gy), in0_block_w=kw, out_subblock_h=1, out_subblock_w=sw, per_core_M=mt, per_core_N=pcn, fuse_batch=True, fused_activation=activation, mcast_in0=True, ) MATMUL_BLOCKS = { "r1": { "ffn_up": (None, (4, 8, 16, 1, 4)), "ffn_down": ((16, 8, 4, 1, 4), (4, 8, 16, 1, 4)), "qkv": ((8, 8, 8, 2, 2), (4, 4, 24, 2, 2)), "post": ((8, 8, 8, 2, 2), (4, 8, 16, 1, 4)), "lconv_end": ((8, 4, 8, 2, 2), (4, 8, 16, 1, 4)), "conv0": ((16, 4, 8, 1, 4), (16, 4, 8, 1, 4)), "in_proj": ((8, 8, 8, 2, 2), (4, 8, 16, 1, 4)), "embed": ((4, 8, 16, 1, 4), (4, 8, 16, 1, 4)), }, "r2": { "ffn_up": ((8, 8, 4, 2, 2), (4, 8, 16, 1, 4)), "ffn_down": ((4, 8, 16, 1, 4), (4, 16, 8, 1, 4)), "qkv": ((8, 8, 4, 2, 2), (4, 8, 16, 1, 4)), "post": ((8, 4, 8, 2, 2), (4, 8, 16, 1, 4)), "lconv_end": ((8, 4, 8, 2, 2), (8, 16, 4, 2, 2)), "lconv_start": ((8, 4, 8, 2, 2), (8, 8, 8, 2, 2)), "out_proj": ((8, 4, 8, 2, 2), (4, 4, 24, 2, 2)), "conv0": ((16, 4, 8, 1, 4), (16, 4, 8, 1, 4)), "in_proj": ((8, 8, 4, 2, 2), (4, 8, 16, 1, 4)), "embed": ((4, 8, 16, 1, 4), (4, 8, 16, 1, 4)), }, } FP32_BLOCK = (8, 4, 8, 2, 2) LARGE_M = 4096 def _minimal_config(grid, role, rows, fp32_io, table="r1"): blocks = MATMUL_BLOCKS[table].get(role) blk = blocks[0 if rows >= LARGE_M else 1] if blocks else None if blk is None and fp32_io: blk = FP32_BLOCK if blk is None: return None mb, kb, nb, sh, sw = blk return ttnn.MinimalMatmulConfig( M_block_size=mb, K_block_size=kb, N_block_size=nb, subblock_h=sh, subblock_w=sw, compute_with_storage_grid_size=grid, ) def _matmul( x, w, grid, compute_cfg, small_m, dtype=ttnn.bfloat16, bias=None, bounds=None, activation=None, role=None, mem=ttnn.DRAM_MEMORY_CONFIG, table="r1", ): rows = x.padded_shape[-2] if rows <= small_m: y = ttnn.linear( x, w, bias=bias, program_config=_small_m_config(grid, rows, w.padded_shape[-2], w.padded_shape[-1], activation), compute_kernel_config=compute_cfg, dtype=dtype, memory_config=mem, ) if bounds is None: return y z = ttnn.clamp(y, min=bounds[0], max=bounds[1]) ttnn.deallocate(y) return z cfg = _minimal_config(grid, role, rows, x.dtype == ttnn.float32 and dtype == ttnn.float32, table) return ttnn.experimental.minimal_matmul( x, w, bias_tensor=bias, fused_activation=_clamp(bounds) if bounds is not None else activation, config=cfg, compute_kernel_config=compute_cfg, dtype=dtype, memory_config=mem, ) def _clamp(bounds): return ttnn.UnaryWithParam(ttnn.UnaryOpType.CLAMP_TSS, float(bounds[0]), float(bounds[1])) def relative_tables(config, rel_weight): half = config.hidden_size // 2 inc = math.log(10000.0) / max(half - 1, 1) inv = torch.exp(torch.arange(half, dtype=torch.float32) * -inc) d = torch.arange(config.window, dtype=torch.float32)[:, None] * inv[None] pe = torch.cat([torch.sin(d), torch.cos(d)], dim=-1) return (pe @ rel_weight.float().T).view(config.window, config.num_attention_heads, config.head_dim) def skew_tables(window): e = torch.zeros(1, 1, BLOCK, BLOCK * window) w2 = torch.zeros(BLOCK * window, CONTEXT) for r in range(BLOCK): for d in range(window): e[0, 0, r, r * window + d] = 1.0 w2[r * window + d, r + BLOCK - d] = 1.0 return e, w2 def conv0_matrix(w, freq): cout = w.shape[0] fo = freq // 2 m = torch.zeros(3, freq, fo, cout) for i in range(3): for j in range(3): for f in range(fo): src = 2 * f - 1 + j if 0 <= src < freq: m[i, src, f] = w[:, 0, i, j] return m.reshape(3 * freq, fo * cout) class AudioEncoderLayer(LightweightModule): def __init__(self, config, layer_idx, mesh_device, weights, bounds, policy): super().__init__() self.config = config self.layer_idx = layer_idx self.mesh_device = mesh_device self.w = weights self.bounds = bounds self.policy = policy self.eps = config.rms_norm_eps self.res_dtype = DTYPES[policy.residual] self.attn_dtype = DTYPES[policy.attention] self.norm_cfg = _compute_config(mesh_device) self.attn_cfg = _compute_config(mesh_device, FIDELITY[policy.attention_fidelity]) self.mm_cfg = {g: _compute_config(mesh_device, policy.fidelity(g)) for g in WEIGHT_GROUPS} self._w1_cache = {} self._dw_weights = {} self._mem = ttnn.DRAM_MEMORY_CONFIG self.conv_cfg = _compute_config(mesh_device, l1_acc=False) grid = mesh_device.compute_with_storage_grid_size() self.grid = grid self.num_cores = grid.x * grid.y hd_t = config.head_dim // TILE self.bmm = { "rel": _bmm_config(grid, hd_t, BLOCK * config.window // TILE, 4), "qk": _bmm_config(grid, hd_t, CONTEXT // TILE, 2), "pv": _bmm_config(grid, CONTEXT // TILE, hd_t, 4), } @classmethod def from_state_dict(cls, state_dict, *, hf_config, layer_idx, mesh_device, prefix=None, policy=None): config = hf_config if isinstance(hf_config, AudioConfig) else AudioConfig.from_hf(hf_config) policy = policy or AudioPolicy() p = prefix if prefix is not None else f"{AUDIO_PREFIX}layers.{layer_idx}." def g(name): return state_dict[p + name].float() def lin(name, group): return _put(g(name + ".linear.weight").T, mesh_device, dtype=policy.dtype(group)) def bound(name): return tuple(float(g(f"{name}.{k}")) for k in ("input_min", "input_max", "output_min", "output_max")) qkv_bounds = {bound(f"self_attn.{n}_proj") for n in "qkv"} if len(qkv_bounds) != 1: raise ValueError(f"layer {layer_idx}: q/k/v clip bounds differ") att = "self_attn." hd, cap = config.head_dim, config.logit_cap q_scale = (hd**-0.5) / math.log(2) * torch.nn.functional.softplus(g(att + "per_dim_scale")) k_scale = math.log(1 + math.e) / math.log(2) rel = relative_tables(config, g(att + "relative_k_proj.weight")) * (q_scale / cap) w1 = ( rel.permute(1, 2, 0)[:, :, None, :] .expand(-1, -1, BLOCK, -1) .reshape(config.num_attention_heads, hd, BLOCK * config.window) ) qkv = torch.cat([g(att + f"{n}_proj.linear.weight") for n in "qkv"], dim=0) rw = config.residual_weight dw = g("lconv1d.depthwise_conv1d.weight").reshape(config.hidden_size, config.conv_kernel_size) weights = { "ffn1_pre": _norm_weight(g("feed_forward1.pre_layer_norm.weight"), mesh_device), "ffn1_post": _tile_weight(g("feed_forward1.post_layer_norm.weight") * rw, mesh_device), "ffn1_w1": lin("feed_forward1.ffw_layer_1", "ffn_up"), "ffn1_w2": lin("feed_forward1.ffw_layer_2", "ffn_down"), "ffn2_pre": _norm_weight(g("feed_forward2.pre_layer_norm.weight"), mesh_device), "ffn2_post": _tile_weight(g("feed_forward2.post_layer_norm.weight") * rw, mesh_device), "ffn2_w1": lin("feed_forward2.ffw_layer_1", "ffn_up"), "ffn2_w2": lin("feed_forward2.ffw_layer_2", "ffn_down"), "norm_pre_attn": _norm_weight(g("norm_pre_attn.weight"), mesh_device), "norm_post_attn": _tile_weight(g("norm_post_attn.weight"), mesh_device), "norm_out": _norm_weight(g("norm_out.weight"), mesh_device), "qkv": _put(qkv.T, mesh_device, dtype=policy.dtype("qkv")), "post": lin(att + "post", "post"), "k_scale": _put((q_scale * k_scale / cap).reshape(1, 1, 1, hd), mesh_device, dtype=ttnn.float32), "w1_host": w1, "lconv_pre": _norm_weight(g("lconv1d.pre_layer_norm.weight"), mesh_device), "lconv_norm": _norm_weight(g("lconv1d.conv_norm.weight"), mesh_device), "lconv_start": lin("lconv1d.linear_start", "lconv_start"), "lconv_end": lin("lconv1d.linear_end", "lconv_end"), "dw": ttnn.from_torch(dw.reshape(config.hidden_size, 1, config.conv_kernel_size), dtype=ttnn.bfloat16), } bounds = { "qkv": qkv_bounds.pop(), "post": bound(att + "post"), "ffn1_w1": bound("feed_forward1.ffw_layer_1"), "ffn1_w2": bound("feed_forward1.ffw_layer_2"), "ffn2_w1": bound("feed_forward2.ffw_layer_1"), "ffn2_w2": bound("feed_forward2.ffw_layer_2"), "lconv_start": bound("lconv1d.linear_start"), "lconv_end": bound("lconv1d.linear_end"), } return cls(config, layer_idx, mesh_device, weights, bounds, policy) def _norm_in(self, x, weight, bounds): y = ttnn.rms_norm(x, epsilon=self.eps, weight=weight, compute_kernel_config=self.norm_cfg) z = ttnn.clamp(y, min=bounds[0], max=bounds[1]) ttnn.deallocate(y) return z def _norm_add(self, x, weight, residual): return ttnn.experimental.dit_fused_distributed_rmsnorm( x, None, self.mesh_device, [], epsilon=self.eps, num_heads_per_device=1, per_head_norm=False, weight=weight, bias=residual, memory_config=self._mem, compute_kernel_config=self.norm_cfg, ) def _mm(self, x, wname, group, bounds, dtype=ttnn.bfloat16): return _matmul( x, self.w[wname], self.grid, self.mm_cfg[group], self.policy.small_m, dtype=dtype, bounds=bounds[2:], role=group, mem=self._mem, table=self.policy.matmul_table, ) def feed_forward(self, x, name): b1, b2 = self.bounds[f"{name}_w1"], self.bounds[f"{name}_w2"] h = self._norm_in(x, self.w[f"{name}_pre"], b1) u = self._mm(h, f"{name}_w1", "ffn_up", b1) ttnn.deallocate(h) a = ttnn.unary_chain(u, [ttnn.UnaryWithParam(ttnn.UnaryOpType.SILU), _clamp(b2)]) ttnn.deallocate(u) d = self._mm(a, f"{name}_w2", "ffn_down", b2, dtype=self.res_dtype) ttnn.deallocate(a) out = self._norm_add(d, self.w[f"{name}_post"], x) ttnn.deallocate(d) return out def prepare(self, batch): if batch not in self._w1_cache: w = self.w["w1_host"][None].expand(batch, -1, -1, -1) self._w1_cache[batch] = _put(w, self.mesh_device, dtype=self.attn_dtype) def _context(self, t, bh, nb, zeros): cur = ttnn.reshape(t, (bh, nb, BLOCK, t.shape[-1])) if nb == 1: return ttnn.concat([zeros, cur], dim=2) head = ttnn.slice(cur, (0, 0, 0, 0), (bh, nb - 1, BLOCK, t.shape[-1])) prev = ttnn.concat([zeros, head], dim=1) ttnn.deallocate(head) ctx = ttnn.concat([prev, cur], dim=2) ttnn.deallocate(prev) return ctx def attention(self, x, inputs, tables): cfg = self.config b, tp = inputs.host.batch, inputs.host.tokens_padded nh, hd, nb = cfg.num_attention_heads, cfg.head_dim, tp // BLOCK bh = b * nh ad = self.attn_dtype qkv = self._mm(x, "qkv", "qkv", self.bounds["qkv"], dtype=ad) qkv = ttnn.reshape(qkv, (b, 1, tp, 3 * cfg.hidden_size)) q, k, v = ttnn.experimental.nlp_create_qkv_heads( qkv, num_heads=nh, num_kv_heads=nh, transpose_k_heads=False, memory_config=self._mem ) ttnn.deallocate(qkv) ks = ttnn.multiply(k, self.w["k_scale"], dtype=ad) ttnn.deallocate(k) rel_pc = self.bmm["rel"] if bh * nb <= self.num_cores else None rel_x = ttnn.matmul( q, self._w1_cache[b], dtype=ad, compute_kernel_config=self.attn_cfg, program_config=rel_pc, memory_config=self._mem, ) rel_e = ttnn.multiply(ttnn.reshape(rel_x, (bh * nb, 1, BLOCK, BLOCK * cfg.window)), tables["e"], dtype=ad) ttnn.deallocate(rel_x) rel = ttnn.matmul( ttnn.reshape(rel_e, (1, 1, bh * tp, BLOCK * cfg.window)), tables["w2"], dtype=ad, compute_kernel_config=self.attn_cfg, ) ttnn.deallocate(rel_e) zeros = tables["zeros"][bh] qb = ttnn.reshape(q, (bh, nb, BLOCK, hd)) kc = self._context(ks, bh, nb, zeros) vc = self._context(v, bh, nb, zeros) ttnn.deallocate(ks) ttnn.deallocate(v) scores = ttnn.matmul( qb, kc, transpose_b=True, dtype=ad, compute_kernel_config=self.attn_cfg, program_config=self.bmm["qk"], memory_config=self._mem, ) ttnn.deallocate(q) ttnn.deallocate(kc) capped = ttnn.add( scores, ttnn.reshape(rel, (bh, nb, BLOCK, CONTEXT)), activations=[ ttnn.UnaryWithParam(ttnn.UnaryOpType.TANH), ttnn.UnaryWithParam(ttnn.UnaryOpType.MUL_UNARY_SFPU, cfg.logit_cap), ], dtype=ad, ) ttnn.deallocate(scores) ttnn.deallocate(rel) masked = ttnn.add(ttnn.reshape(capped, (b, nh, tp, CONTEXT)), inputs.attn_bias, dtype=ad) ttnn.deallocate(capped) probs = ttnn.softmax(masked, dim=-1, numeric_stable=True, compute_kernel_config=self.attn_cfg) ttnn.deallocate(masked) ctx = ttnn.matmul( ttnn.reshape(probs, (bh, nb, BLOCK, CONTEXT)), vc, dtype=ttnn.bfloat16, compute_kernel_config=self.attn_cfg, program_config=self.bmm["pv"], memory_config=self._mem, ) ttnn.deallocate(probs) ttnn.deallocate(vc) merged = ttnn.experimental.nlp_concat_heads(ttnn.reshape(ctx, (b, nh, tp, hd)), memory_config=self._mem) ttnn.deallocate(ctx) merged = ttnn.reshape(merged, (1, 1, b * tp, cfg.hidden_size)) pb = self.bounds["post"] mc = ttnn.clamp(merged, min=pb[0], max=pb[1]) ttnn.deallocate(merged) out = self._mm(mc, "post", "post", pb, dtype=self.res_dtype) ttnn.deallocate(mc) return out def _depthwise(self, u, batch, tokens): cfg = self.config d, k = cfg.hidden_size, cfg.conv_kernel_size urm = ttnn.to_layout(ttnn.reshape(u, (batch, tokens, d)), ttnn.ROW_MAJOR_LAYOUT) key = (batch, tokens) out, _, (wd, _) = ttnn.conv1d( input_tensor=urm, weight_tensor=self._dw_weights.get(key, self.w["dw"]), in_channels=d, out_channels=d, device=self.mesh_device, kernel_size=k, stride=1, padding=(k - 1, 0), batch_size=batch, input_length=tokens, groups=d, dtype=DTYPES[self.policy.conv_out], conv_config=ttnn.Conv1dConfig( weights_dtype=ttnn.bfloat16, shard_layout=ttnn.TensorMemoryLayout.HEIGHT_SHARDED ), compute_config=self.conv_cfg, return_output_dim=True, return_weights_and_bias=True, ) ttnn.deallocate(urm) self._dw_weights[key] = wd if out.is_sharded(): inter = ttnn.sharded_to_interleaved(out, self._mem) ttnn.deallocate(out) out = inter return ttnn.reshape(out, (1, 1, batch * tokens, d)) def light_conv(self, x, inputs, tables): cfg = self.config b, tp, d = inputs.host.batch, inputs.host.tokens_padded, cfg.hidden_size k = cfg.conv_kernel_size bs, be = self.bounds["lconv_start"], self.bounds["lconv_end"] h = self._norm_in(x, self.w["lconv_pre"], bs) rows = h.padded_shape[-2] if rows <= self.policy.small_m and self.policy.small_m_glu_slices: s = _matmul( h, self.w["lconv_start"], self.grid, self.mm_cfg["lconv_start"], self.policy.small_m, bounds=bs[2:], mem=self._mem, ) a = ttnn.slice(s, (0, 0, 0, 0), (1, 1, rows, d), memory_config=self._mem) gt = ttnn.slice(s, (0, 0, 0, d), (1, 1, rows, 2 * d), memory_config=self._mem) ttnn.deallocate(s) else: a, gt = ttnn.experimental.minimal_matmul_split( h, self.w["lconv_start"], chunks=2, dim=-1, fused_activation=_clamp(bs[2:]), config=_minimal_config(self.grid, "lconv_start", rows, False, self.policy.matmul_table), compute_kernel_config=self.mm_cfg["lconv_start"], dtype=ttnn.bfloat16, memory_config=self._mem, ) ttnn.deallocate(h) u = ttnn.multiply(a, gt, input_tensor_b_activations=[ttnn.UnaryWithParam(ttnn.UnaryOpType.SIGMOID)]) ttnn.deallocate(a) ttnn.deallocate(gt) conv = self._depthwise(u, b, tp) ttnn.deallocate(u) n = ttnn.rms_norm(conv, epsilon=self.eps, weight=self.w["lconv_norm"], compute_kernel_config=self.norm_cfg) ttnn.deallocate(conv) act = ttnn.unary_chain(n, [ttnn.UnaryWithParam(ttnn.UnaryOpType.SILU), _clamp(be)]) ttnn.deallocate(n) e = self._mm(act, "lconv_end", "lconv_end", be, dtype=self.res_dtype) ttnn.deallocate(act) out = ttnn.add(x, e, dtype=self.res_dtype) ttnn.deallocate(e) return out def forward(self, x, inputs, tables): self._mem = tables["mem"] h1 = self.feed_forward(x, "ffn1") n = self._norm_in(h1, self.w["norm_pre_attn"], self.bounds["qkv"]) a = self.attention(n, inputs, tables) ttnn.deallocate(n) h2 = self._norm_add(a, self.w["norm_post_attn"], h1) ttnn.deallocate(a) ttnn.deallocate(h1) h3 = self.light_conv(h2, inputs, tables) ttnn.deallocate(h2) h4 = self.feed_forward(h3, "ffn2") ttnn.deallocate(h3) out = ttnn.rms_norm(h4, epsilon=self.eps, weight=self.w["norm_out"], compute_kernel_config=self.norm_cfg) ttnn.deallocate(h4) return out class AudioEncoder(LightweightModule): def __init__(self, config, mesh_device, weights, layers, policy): super().__init__() self.config = config self.mesh_device = mesh_device self.w = weights self.layers = layers self.policy = policy self.eps = config.rms_norm_eps self.res_dtype = DTYPES[policy.residual] self.compute_cfg = _compute_config(mesh_device) self.grid = mesh_device.compute_with_storage_grid_size() self.proj_cfg = _compute_config(mesh_device, policy.fidelity("proj")) self.conv_cfg = _compute_config(mesh_device, l1_acc=False) self.conv0_cfg = _compute_config(mesh_device, policy.fidelity("conv0")) e, w2 = skew_tables(config.window) ad = DTYPES[policy.attention] self.tables = { "e": _put(e, mesh_device, dtype=ad), "w2": _put(w2, mesh_device, dtype=ad), "zeros": {}, "feat_zeros": {}, "conv_ready": set(), "mem": ttnn.DRAM_MEMORY_CONFIG, } self.zero_rows = ttnn.zeros( (1, 1, TILE, config.text_hidden_size), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) self.conv_weights = {} @classmethod def from_state_dict(cls, state_dict, *, hf_config, mesh_device, policy=None, num_layers=None): config = AudioConfig.from_hf(hf_config) policy = policy or AudioPolicy() def g(name): return state_dict[name].float() sp = AUDIO_PREFIX + "subsample_conv_projection." weights = { "conv0": _put( conv0_matrix(g(sp + "layer0.conv.weight"), config.feature_size), mesh_device, dtype=policy.dtype("conv0"), ), "conv0_chunks": [ _put(chunk, mesh_device, dtype=policy.dtype("conv0")) for chunk in conv0_matrix(g(sp + "layer0.conv.weight"), config.feature_size).split(CONV0_CHUNK, dim=1) ], "ln0": _norm_weight(g(sp + "layer0.norm.weight"), mesh_device), "conv1": ttnn.from_torch(g(sp + "layer1.conv.weight"), dtype=ttnn.bfloat16), "ln1": _norm_weight(g(sp + "layer1.norm.weight"), mesh_device), "in_proj": _put(g(sp + "input_proj_linear.weight").T, mesh_device, dtype=policy.dtype("proj")), "out_proj": _put(g(AUDIO_PREFIX + "output_proj.weight").T, mesh_device, dtype=policy.dtype("proj")), "out_bias": _put(g(AUDIO_PREFIX + "output_proj.bias").reshape(1, -1), mesh_device), "embed": _put( g(EMBEDDER_PREFIX + "embedding_projection.weight").T, mesh_device, dtype=policy.dtype("proj") ), } n = config.num_hidden_layers if num_layers is None else num_layers layers = [ AudioEncoderLayer.from_state_dict( state_dict, hf_config=config, layer_idx=i, mesh_device=mesh_device, policy=policy ) for i in range(n) ] return cls(config, mesh_device, weights, layers, policy) def prepare(self, batch): cfg = self.config bh = batch * cfg.num_attention_heads ad = DTYPES[self.policy.attention] if bh not in self.tables["zeros"]: self.tables["zeros"][bh] = _put(torch.zeros(bh, 1, BLOCK, cfg.head_dim), self.mesh_device, dtype=ad) if batch not in self.tables["feat_zeros"]: self.tables["feat_zeros"][batch] = _put( torch.zeros(batch, 1, cfg.feature_size), self.mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT ) for layer in self.layers: layer.prepare(batch) self.tables["conv_ready"].add(batch) def _conv1(self, x, batch, height, width): c0, c1 = self.config.conv_channels key = (batch, height) out, (oh, ow), (wd, _) = ttnn.conv2d( input_tensor=x, weight_tensor=self.conv_weights.get(key, self.w["conv1"]), device=self.mesh_device, in_channels=c0, out_channels=c1, batch_size=batch, input_height=height, input_width=width, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), compute_config=self.conv_cfg, conv_config=ttnn.Conv2dConfig(weights_dtype=ttnn.bfloat16), dtype=ttnn.bfloat16, return_output_dim=True, return_weights_and_bias=True, ) self.conv_weights[key] = wd return out, oh, ow def subsample(self, inputs): cfg = self.config b, tp = inputs.host.batch, inputs.host.tokens_padded c0, c1 = cfg.conv_channels f = cfg.feature_size t1 = 2 * tp fo = f // 2 xp = inputs.features odd = ttnn.slice(xp, (0, 0, f), (b, t1 - 1, 2 * f)) prev = ttnn.concat([self.tables["feat_zeros"][b], odd], dim=1) ttnn.deallocate(odd) x3 = ttnn.concat([prev, xp], dim=2) ttnn.deallocate(prev) x3t = ttnn.to_layout(ttnn.reshape(x3, (1, 1, b * t1, 3 * f)), ttnn.TILE_LAYOUT) ttnn.deallocate(x3) x3m = ttnn.multiply(x3t, inputs.mask1) ttnn.deallocate(x3t) parts = [] weights = [self.w["conv0"]] if b * t1 <= CONV0_SINGLE_ROWS else self.w["conv0_chunks"] for w in weights: y = ttnn.experimental.minimal_matmul( x3m, w, config=_minimal_config(self.grid, "conv0", b * t1, False, self.policy.matmul_table), compute_kernel_config=self.conv0_cfg, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) parts.append(ttnn.to_layout(y, ttnn.ROW_MAJOR_LAYOUT)) ttnn.deallocate(y) ttnn.deallocate(x3m) if len(parts) == 1: y0rm = parts[0] else: y0rm = ttnn.concat(parts, dim=3) for t in parts: ttnn.deallocate(t) y0c = ttnn.to_layout(ttnn.reshape(y0rm, (1, 1, b * t1 * fo, c0)), ttnn.TILE_LAYOUT) ttnn.deallocate(y0rm) n0 = ttnn.layer_norm(y0c, epsilon=self.eps, weight=self.w["ln0"], compute_kernel_config=self.compute_cfg) ttnn.deallocate(y0c) r0 = ttnn.relu(n0) ttnn.deallocate(n0) y1, h1, w1 = self._conv1(r0, b, t1, fo) ttnn.deallocate(r0) y1 = ttnn.to_layout(ttnn.reshape(y1, (1, 1, b * h1 * w1, c1)), ttnn.TILE_LAYOUT) n1 = ttnn.layer_norm(y1, epsilon=self.eps, weight=self.w["ln1"], compute_kernel_config=self.compute_cfg) ttnn.deallocate(y1) r1 = ttnn.relu(n1) ttnn.deallocate(n1) flat = ttnn.reshape(r1, (1, 1, b * h1, w1 * c1)) ttnn.deallocate(r1) h = _matmul( flat, self.w["in_proj"], self.grid, self.proj_cfg, self.policy.small_m, dtype=self.res_dtype, role="in_proj", table=self.policy.matmul_table, ) ttnn.deallocate(flat) return h def project(self, h, inputs): o = _matmul( h, self.w["out_proj"], self.grid, self.proj_cfg, self.policy.small_m, bias=self.w["out_bias"], role="out_proj", table=self.policy.matmul_table, ) n = ttnn.rms_norm(o, epsilon=self.eps, compute_kernel_config=self.compute_cfg) e = _matmul( n, self.w["embed"], self.grid, self.proj_cfg, self.policy.small_m, role="embed", table=self.policy.matmul_table, ) ttnn.deallocate(n) table = ttnn.concat([e, self.zero_rows], dim=2) ttnn.deallocate(e) table_rm = ttnn.to_layout(ttnn.reshape(table, (table.shape[-2], table.shape[-1])), ttnn.ROW_MAJOR_LAYOUT) ttnn.deallocate(table) tokens = ttnn.embedding( inputs.gather_idx, table_rm, layout=ttnn.TILE_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG ) ttnn.deallocate(table_rm) tokens = ttnn.reshape(tokens, (1, 1, tokens.shape[-2], tokens.shape[-1])) return tokens, o def allocate_output(self, inputs): return ttnn.zeros( (1, 1, inputs.host.out_padded, self.config.text_hidden_size), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) def forward(self, inputs, return_intermediates=False, output=None): if inputs.host.batch not in self.tables["conv_ready"]: raise ValueError(f"call prepare({inputs.host.batch}) before forward") rows = inputs.host.batch * inputs.host.tokens_padded self.tables["mem"] = ttnn.L1_MEMORY_CONFIG if rows <= self.policy.l1_rows else ttnn.DRAM_MEMORY_CONFIG h = self.subsample(inputs) inter = {"subsample": ttnn.clone(h, memory_config=ttnn.DRAM_MEMORY_CONFIG)} if return_intermediates else None for i, layer in enumerate(self.layers): nxt = layer(h, inputs, self.tables) if return_intermediates: inter[f"layer{i:02d}"] = ttnn.clone(nxt, memory_config=ttnn.DRAM_MEMORY_CONFIG) ttnn.deallocate(h) h = nxt tokens, tower = self.project(h, inputs) ttnn.deallocate(h) if output is not None: ttnn.copy(tokens, output) ttnn.deallocate(tokens) tokens = output if not return_intermediates: ttnn.deallocate(tower) return {"tokens": tokens} inter["tower"] = tower inter["tokens"] = tokens return inter