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