| |
| """Device and host helpers shared with the ACE-Step port. |
| |
| Extracted verbatim from `tt-acestep` (Apache-2.0), where the same code opens the card, reports |
| DRAM, moves tensors, and prepares checkpoint tensors for ttnn: the matmul layout, the |
| adjacent-pair head permutation ttnn's rotary kernel expects, and the tiled rotation matrix. |
| |
| The two ports also share `worker.py`, `remote.py`, `wire.py` and `oobleck_vae.py`; this module |
| exists so that this repository stands alone. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| import ttnn |
|
|
| MEM = ttnn.DRAM_MEMORY_CONFIG |
|
|
|
|
| def open_device(l1_small_size: int = 98304): |
| return ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(1, 1), |
| dispatch_core_config=ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.WORKER), |
| l1_small_size=l1_small_size) |
|
|
|
|
| def close_device(dev): |
| ttnn.close_mesh_device(dev) |
|
|
|
|
| def dram_stats(dev) -> dict: |
| try: |
| s = ttnn.get_memory_view(dev, ttnn.BufferType.DRAM) |
| return {"allocated_bytes": s.total_bytes_allocated_per_bank * s.num_banks, |
| "free_bytes": s.total_bytes_free_per_bank * s.num_banks} |
| except Exception: |
| return {} |
|
|
|
|
| @dataclass |
| class Precision: |
| weight_dtype: ttnn.DataType = ttnn.bfloat16 |
| mm_fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi2 |
| sdpa_fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi4 |
| sdpa_fp32_acc: bool = True |
|
|
|
|
| def _dev(dev, t: torch.Tensor, dtype=ttnn.bfloat16): |
| return ttnn.from_torch(t.contiguous(), dtype=dtype, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=MEM) |
|
|
|
|
| def _row(dev, v: torch.Tensor): |
| """[D] -> [1, 1, 1, D] bf16 row (norm gamma, shift, gate, bias).""" |
| return _dev(dev, v.to(torch.bfloat16).reshape(1, 1, 1, -1)) |
|
|
|
|
| TILE = 32 |
| SEQ_MULTIPLE = 128 |
|
|
|
|
| NEG = -1e9 |
|
|
|
|
| def linear_to_mm(w: torch.Tensor) -> torch.Tensor: |
| """nn.Linear weight [out, in] -> matmul weight [in, out].""" |
| return w.t().contiguous() |
|
|
|
|
| def interleave_pairs_permutation(head_dim: int) -> torch.Tensor: |
| """new[j] = old[p[j]]: rotate-half pairs (i, i + D/2) -> adjacent pairs (2i, 2i + 1). |
| |
| Applying it to the rows of q_proj/k_proj (per head) and to q_norm/k_norm leaves every q.k product unchanged and |
| lets tt's rotary_embedding_llama (adjacent pairs) replace upstream's rotate_half RoPE.""" |
| half = head_dim // 2 |
| p = torch.empty(head_dim, dtype=torch.long) |
| p[0::2] = torch.arange(half) |
| p[1::2] = torch.arange(half) + half |
| return p |
|
|
|
|
| def permute_heads_rows(w_out_in: torch.Tensor, n_heads: int, head_dim: int, perm: torch.Tensor) -> torch.Tensor: |
| out, inp = w_out_in.shape |
| assert out == n_heads * head_dim |
| return w_out_in.view(n_heads, head_dim, inp)[:, perm, :].reshape(out, inp).contiguous() |
|
|
|
|
| def rot_transformation_mat(tile: int = TILE) -> torch.Tensor: |
| """[1, 1, 32, 32] T with (x @ T)[2k] = -x[2k+1], (x @ T)[2k+1] = x[2k] (adjacent-pair rotation).""" |
| m = torch.zeros(1, 1, tile, tile) |
| m[..., torch.arange(0, tile, 2), torch.arange(1, tile, 2)] = 1.0 |
| m[..., torch.arange(1, tile, 2), torch.arange(0, tile, 2)] = -1.0 |
| return m |
|
|
|
|
| def rms(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * w |
|
|
|
|
| def round_up(n: int, m: int = SEQ_MULTIPLE) -> int: |
| return (n + m - 1) // m * m |
|
|
|
|
| def swiglu_interleave(gate_out_in: torch.Tensor, up_out_in: torch.Tensor, tile: int = TILE) -> torch.Tensor: |
| """[K, 2N] weight for minimal_matmul(fuse_swiglu=True): column tile 2p = gate tile p, 2p+1 = up tile p.""" |
| g, u = linear_to_mm(gate_out_in), linear_to_mm(up_out_in) |
| K, N = g.shape |
| assert N % tile == 0, N |
| return torch.stack([g.view(K, N // tile, tile), u.view(K, N // tile, tile)], dim=2).reshape(K, 2 * N).contiguous() |
|
|
|
|
| def apply_rope_adjacent(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| rot = torch.stack([-x[..., 1::2], x[..., 0::2]], dim=-1).flatten(-2) |
| return x * cos + rot * sin |
|
|
