code/alpamayo_tt/text.py

"""Qwen3-VL-32B text model (Alpamayo2-Super VLM backbone) for one Blackhole chip, batch 1.

Prefill: the whole prompt in one pass (L padded to a multiple of 128), row-parallel matmuls via
ttnn.experimental.minimal_matmul on the 11x10 grid, causal SDPA, KV cache filled with ttnn.fill_cache.
Decode: M=32 padded single token, DRAM-sharded matmuls (weights width-sharded over the 8 DRAM banks, 2 workers per
bank) fed by width-sharded L1 activations, sharded RMSNorm, flash-decode SDPA; the whole step is captured as a
device trace with persistent input buffers (embedding row, cos, sin, position) and a row-major logits row output.
Layouts / configs / measurements: TEXT_NOTES.md.
"""
from __future__ import annotations

import math
import time

import torch
import ttnn

from .config import ModelConfig
from .matmul_configs import (compute_config, decode_1d_config, fidelity_for, minimal_config, sdpa_decode_config,
                             sdpa_prefill_config)
from .rope import mrope_cos_sin, rope_delta as _rope_delta, text_cos_sin
from .weights import VLM_TXT, Checkpoint, TTWeights, fused_qkv, linear_weight

TILE = 32
PREFILL_GRID = (11, 10)
EMBED_KEY = VLM_TXT + "embed_tokens.weight"
NORM_KEY = VLM_TXT + "norm.weight"
LM_HEAD_KEY = "vlm.lm_head.weight"


def _ceil(a: int, b: int) -> int:
    return (a + b - 1) // b * b


def _core_range(n: int = 1):
    return ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(n - 1, 0))})


DRAM_BANKS = 8
DRAM_GRID = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(DRAM_BANKS - 1, 0))})


def dram_sharded_weight_mem(K: int, N: int):
    """WIDTH_SHARDED DRAM memory config: [K, N/8] per bank (N must be a multiple of 256)."""
    assert N % (TILE * DRAM_BANKS) == 0, N
    return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM,
                             ttnn.ShardSpec(DRAM_GRID, (K, N // DRAM_BANKS), ttnn.ShardOrientation.ROW_MAJOR))


def width_sharded(K: int, grid: ttnn.CoreGrid):
    """L1 WIDTH_SHARDED activation config for a [32, K] row split over grid.num_cores cores."""
    n = grid.x * grid.y
    assert K % (n * TILE) == 0, (K, n)
    return ttnn.create_sharded_memory_config(shape=(TILE, K // n), core_grid=grid, strategy=ttnn.ShardStrategy.WIDTH,
                                             orientation=ttnn.ShardOrientation.ROW_MAJOR,
                                             use_height_and_width_as_shard_shape=True)


def ds_matmul_config(in0_block_w: int, per_core_N: int, act=None, workers: int = 2):
    return ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(in0_block_w=in0_block_w, per_core_M=1,
                                                                     per_core_N=per_core_N, fused_activation=act,
                                                                     num_workers_per_dram_bank=workers)


def _height_sharded(rows: int, cols: int, cores=None):
    cores = cores if cores is not None else _core_range(1)
    return ttnn.create_sharded_memory_config(shape=(rows, cols), core_grid=cores, strategy=ttnn.ShardStrategy.HEIGHT,
                                             orientation=ttnn.ShardOrientation.ROW_MAJOR,
                                             use_height_and_width_as_shard_shape=True)


class TTText:
    """See module docstring. `layers` restricts which decoder layers are loaded (tests); kv_cache has one entry
    per model layer (None for layers that are not loaded)."""

    def __init__(self, device, cfg: ModelConfig, ckpt: Checkpoint, tw: TTWeights, layers: range | None = None,
                 mlp_dtype: str | None = None, attn_dtype: str | None = None, load_lm_head: bool = True,
                 max_seq_len: int | None = None, verbose: bool = True):
        self.device = device
        self.cfg = cfg
        self.tc = cfg.text
        self.ckpt = ckpt
        self.tw = tw
        self.layers = list(layers) if layers is not None else list(range(self.tc.n_layers))
        self.mlp_dtype = mlp_dtype or cfg.precision.vlm_mlp
        self.attn_dtype = attn_dtype or cfg.precision.vlm_attn
        self.kv_dtype = {"bf16": ttnn.bfloat16, "bfp8": ttnn.bfloat8_b}[cfg.precision.kv_cache]
        self.max_seq_len = max_seq_len or cfg.max_seq_len
        self.verbose = verbose
        tc = self.tc
        self.scale = 1.0 / math.sqrt(tc.head_dim)

        # compute configs
        self.ck_attn = compute_config(fidelity_for(self.attn_dtype))
        self.ck_mlp = compute_config(fidelity_for(self.mlp_dtype))
        self.ck_head = compute_config(fidelity_for(cfg.precision.vlm_lm_head))
        self.ck_norm = compute_config("hifi4")
        self.ck_sdpa_prefill = compute_config("hifi4", fp32_acc=False, l1_acc=False)  # fp32 acc overflows L1 at 4.6K
        self.ck_sdpa = compute_config("hifi4", fp32_acc=True, l1_acc=False)
        self.ck_exact = compute_config("hifi2")  # placement matmul (0/1 bfp8 matrix x bf16): exact, measured
        self.mm_cfg = minimal_config(PREFILL_GRID, 8)

        # decode: DRAM-sharded matmuls (2 workers per bank) fed by width-sharded activations; measured in
        # scripts/probe_text_ops3/4.py. Grids: hidden (5120 = 160 tiles) over 8x10 = 80 cores (2 tiles each),
        # attention output (8192 = 256 tiles) over 8x4 = 32 cores (8 tiles each), inter (25600) over 80 cores (10).
        g80 = ttnn.CoreGrid(x=8, y=10)
        self.res_mem = width_sharded(tc.hidden, g80)  # residual stream + norm outputs
        self.attn_in_mem = width_sharded(tc.q_dim, ttnn.CoreGrid(x=8, y=4))  # wo input (32 cores)
        n80 = g80.x * g80.y
        self.pc_qkv = ds_matmul_config(tc.hidden // TILE // n80, tc.qkv_dim // TILE // n80)  # bw2, pcN4: 0.140 ms
        self.pc_wo = ds_matmul_config(tc.q_dim // TILE // 32, tc.hidden // TILE // 32)  # bw8, pcN5: 0.124 ms
        self.pc_w1 = ds_matmul_config(8, tc.inter // TILE // n80, act=ttnn.UnaryOpType.SILU)  # bw8, pcN10: 0.19 (+0.05 silu)
        self.pc_w3 = ds_matmul_config(8, tc.inter // TILE // n80)
        self.pc_w2 = ds_matmul_config(5, tc.hidden // TILE // n80)  # bw5, pcN2: 0.175 ms
        self.pc_head = decode_1d_config(tc.hidden, tc.vocab, 11, 10)  # interleaved lm_head, 2.3 ms
        self.sdpa_dec_pc = sdpa_decode_config((8, 8), 128)
        bw = tc.hidden // n80 // TILE  # 2 tiles per core
        self.norm_pc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=[g80.x, g80.y],
                                                                   subblock_w=bw, block_h=1, block_w=bw, inplace=False)
        # decode sharded layouts (batch 1 -> core (0,0))
        self.heads_mem = _height_sharded(_ceil(tc.n_heads, TILE), tc.head_dim, ttnn.CoreGrid(y=4, x=8))
        self.cos_mem = _height_sharded(TILE, tc.head_dim)
        self.attn_out_mem = _height_sharded(_ceil(tc.n_heads, TILE), tc.head_dim)

        # weights
        t0 = time.time()
        self.layer_w = {}
        for i in self.layers:
            self.layer_w[i] = self._load_layer(i)
            self._log(f"layer {i} weights loaded ({time.time() - t0:.0f}s, {tw.stats})")
        self.norm_w = self._norm_weight("text.norm", ckpt.get(NORM_KEY))
        self.norm_w_tile = self._norm_weight_tile("text.norm", ckpt.get(NORM_KEY))
        self.lm_head = None
        if load_lm_head:
            self.lm_head = tw.to_device("text.lm_head", linear_weight(ckpt.get(LM_HEAD_KEY)), cfg.precision.vlm_lm_head)
        self._embed = None
        self._log(f"all weights on device in {time.time() - t0:.0f}s ({tw.stats})")

        # KV cache: [1, n_kv_heads, max_seq_len, head_dim] per layer, HF head order, tile layout
        self.kv_cache: list = [None] * tc.n_layers
        for i in self.layers:
            self.kv_cache[i] = tuple(
                ttnn.zeros((1, tc.n_kv_heads, self.max_seq_len, tc.head_dim), dtype=self.kv_dtype,
                           layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)
                for _ in range(2))
        self.cur_len = 0
        self.rope_delta = 0
        self._dec = None  # decode trace state
        self.timings: dict[str, float] = {}
        self.debug_capture: dict | None = None  # when a dict: prefill stores host copies of norm.in / norm.out
        # Seam for the fused (megakernel) decode layer: decode() calls self.layer_fn(i, x, cos_s, sin_s, pos) per layer
        # with x = residual stream [1,1,32,hidden] bf16 width-sharded over 8x10 cores (self.res_mem), cos_s/sin_s
        # [1,1,1(32),head_dim] height-sharded on core (0,0), pos [1] int32 DRAM; it must return x in the same layout
        # and write K/V row `pos` of self.kv_cache[i]. Layer weights are self.layer_w[i] (dict: wqkv, wo, w1, w3, w2
        # as DRAM width-sharded TILE tensors; ln1, ln2, qn, kn as [1,1,dim/32,32] ROW_MAJOR bf16). Override per layer
        # via self.layer_fns[i] before the first decode() call (the trace binds whatever is installed then).
        self.layer_fns: dict[int, object] = {}

    # ------------------------------------------------------------------ weights
    def _log(self, *a):
        if self.verbose:
            print(f"[TTText {time.strftime('%H:%M:%S')}]", *a, flush=True)

    def _norm_weight(self, name: str, w: torch.Tensor):
        """Row-major gamma [1,1,dim/32,32] (sharded decode norm + the 128-wide q/k norms)."""
        return self.tw.to_device(name, w.reshape(1, 1, -1, TILE).to(torch.bfloat16), "bf16", layout=ttnn.ROW_MAJOR_LAYOUT)

    def _norm_weight_tile(self, name: str, w: torch.Tensor):
        """TILE gamma [1,1,1,dim] for the interleaved prefill norms: with a row-major gamma ttnn.rms_norm skips its
        L1 budget check and allocates ~1.45 MB of circular buffers per core (clashes with L1-resident tensors of
        other components); with a TILE gamma it takes the budgeted large-tensor path (~0.15 ms slower per norm)."""
        return self.tw.to_device(name + ".tile", w.reshape(1, 1, 1, -1).to(torch.bfloat16), "bf16")

    def _load_layer(self, i: int) -> dict:
        w = self.ckpt.text_layer(i)
        p = f"text.layers.{i}."
        tw = self.tw
        tc = self.tc

        def ds(name, t, dtype):  # DRAM width-sharded over the 8 banks (decode reads contiguous shards; prefill
            # minimal_matmul consumes the same tensors, measured identical numerics and speed)
            return tw.to_device(p + name + ".ds", t, dtype, mem_config=dram_sharded_weight_mem(t.shape[0], t.shape[1]))

        out = {
            "wqkv": ds("wqkv", fused_qkv(w["self_attn.q_proj.weight"], w["self_attn.k_proj.weight"],
                                         w["self_attn.v_proj.weight"]), self.attn_dtype),
            "wo": ds("wo", linear_weight(w["self_attn.o_proj.weight"]), self.attn_dtype),
            "w1": ds("w1", linear_weight(w["mlp.gate_proj.weight"]), self.mlp_dtype),
            "w3": ds("w3", linear_weight(w["mlp.up_proj.weight"]), self.mlp_dtype),
            "w2": ds("w2", linear_weight(w["mlp.down_proj.weight"]), self.mlp_dtype),
            "ln1": self._norm_weight(p + "ln1", w["input_layernorm.weight"]),
            "ln2": self._norm_weight(p + "ln2", w["post_attention_layernorm.weight"]),
            "ln1_tile": self._norm_weight_tile(p + "ln1", w["input_layernorm.weight"]),
            "ln2_tile": self._norm_weight_tile(p + "ln2", w["post_attention_layernorm.weight"]),
            "qn": self._norm_weight(p + "qn", w["self_attn.q_norm.weight"]),
            "kn": self._norm_weight(p + "kn", w["self_attn.k_norm.weight"]),
        }
        return out

    @property
    def embed_table(self) -> torch.Tensor:
        if self._embed is None:
            self._embed = self.ckpt.get(EMBED_KEY).to(torch.bfloat16)
        return self._embed

    def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Host embedding lookup -> [1, L, hidden] bf16."""
        return self.embed_table[input_ids.reshape(-1).long()].unsqueeze(0)

    # ------------------------------------------------------------------ rope
    def rope_cos_sin(self, position_ids3: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Interleaved mrope cos/sin [L, head_dim] float32 for position_ids [3, 1, L] or [3, L]."""
        return mrope_cos_sin(position_ids3, self.tc.head_dim, self.tc.rope_theta, self.tc.mrope_section)

    # ------------------------------------------------------------------ prefill
    def _mm(self, x, w, ck, act=None):
        return ttnn.experimental.minimal_matmul(x, w, compute_kernel_config=ck, config=self.mm_cfg, dtype=ttnn.bfloat16,
                                                memory_config=ttnn.DRAM_MEMORY_CONFIG)

    def _rms(self, x, w, mem=ttnn.DRAM_MEMORY_CONFIG):
        return ttnn.rms_norm(x, epsilon=self.tc.rms_eps, weight=w, compute_kernel_config=self.ck_norm, memory_config=mem)

    def _placement(self, positions: torch.Tensor, Lp: int) -> ttnn.Tensor:
        """0/1 matrix P [1,1,Lp,Nm] (bfp8, exact) with P[pos_j, j] = 1 so that P @ D scatters D rows to positions."""
        n = positions.numel()
        P = torch.zeros(Lp, _ceil(n, TILE), dtype=torch.bfloat16)
        P[positions.long(), torch.arange(n)] = 1.0
        return ttnn.from_torch(P.view(1, 1, Lp, -1), dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=self.device,
                               memory_config=ttnn.DRAM_MEMORY_CONFIG)

    def _to_dev(self, t, pad_rows: int | None = None):
        if isinstance(t, ttnn.Tensor):
            return t
        t = t.to(torch.bfloat16).reshape(1, 1, t.shape[-2], t.shape[-1])
        if pad_rows is not None and t.shape[2] != pad_rows:
            t = torch.nn.functional.pad(t, (0, 0, 0, pad_rows - t.shape[2]))
        return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.device,
                               memory_config=ttnn.DRAM_MEMORY_CONFIG)

    def _scatter_add(self, h, P, feats):
        """h += P @ feats (feats [1,1,Nm,hidden] bf16 on device)."""
        d = self._to_dev(feats)
        add = self._mm(P, d, self.ck_exact)
        out = ttnn.add(h, add, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(add)
        ttnn.deallocate(h)
        if d is not feats:
            ttnn.deallocate(d)
        return out

    def _fill_cache(self, cache, x, Lp: int):
        """cache[0, :, :Lp] = x ([1, n_kv, Lp, head_dim]). ttnn.fill_cache mis-addresses interleaved multi-head inputs
        (measured: heads 1..7 corrupted), so the input is height-sharded over 8x8 cores first (tt_transformers recipe)."""
        if x.dtype != cache.dtype:
            x = ttnn.typecast(x, cache.dtype)
        shard = ttnn.create_sharded_memory_config((self.tc.n_kv_heads * Lp // 64, self.tc.head_dim), ttnn.CoreGrid(y=8, x=8),
                                                  ttnn.ShardStrategy.HEIGHT, ttnn.ShardOrientation.ROW_MAJOR,
                                                  use_height_and_width_as_shard_shape=True)
        xs = ttnn.interleaved_to_sharded(x, shard)
        ttnn.fill_cache(cache, xs, 0)
        ttnn.deallocate(xs)

    def _layer_prefill(self, i: int, h, cos, sin, Lp: int, sdpa_pc):
        w = self.layer_w[i]
        tc = self.tc
        n1 = self._rms(h, w["ln1_tile"])
        qkv = self._mm(n1, w["wqkv"], self.ck_attn)
        ttnn.deallocate(n1)
        q, k, v = ttnn.experimental.nlp_create_qkv_heads(qkv, num_heads=tc.n_heads, num_kv_heads=tc.n_kv_heads,
                                                         transpose_k_heads=False, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(qkv)
        qn = self._rms(q, w["qn"])
        ttnn.deallocate(q)
        kn = self._rms(k, w["kn"])
        ttnn.deallocate(k)
        q = ttnn.experimental.rotary_embedding_hf(qn, cos, sin, is_decode_mode=False)
        ttnn.deallocate(qn)
        k = ttnn.experimental.rotary_embedding_hf(kn, cos, sin, is_decode_mode=False)
        ttnn.deallocate(kn)
        K, V = self.kv_cache[i]
        self._fill_cache(K, k, Lp)
        self._fill_cache(V, v, Lp)
        a = ttnn.transformer.scaled_dot_product_attention(q, k, v, is_causal=True, scale=self.scale, program_config=sdpa_pc,
                                                          compute_kernel_config=self.ck_sdpa_prefill)
        ttnn.deallocate(q)
        ttnn.deallocate(k)
        ttnn.deallocate(v)
        ac = ttnn.experimental.nlp_concat_heads(a, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(a)
        o = self._mm(ac, w["wo"], self.ck_attn)
        ttnn.deallocate(ac)
        h2 = ttnn.add(h, o, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(o)
        ttnn.deallocate(h)
        n2 = self._rms(h2, w["ln2_tile"])
        g = self._mm(n2, w["w1"], self.ck_mlp)
        u = self._mm(n2, w["w3"], self.ck_mlp)
        ttnn.deallocate(n2)
        m = ttnn.mul(g, u, input_tensor_a_activations=[ttnn.UnaryOpType.SILU], memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(g)
        ttnn.deallocate(u)
        d = self._mm(m, w["w2"], self.ck_mlp)
        ttnn.deallocate(m)
        h3 = ttnn.add(h2, d, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(d)
        ttnn.deallocate(h2)
        return h3

    def _prefill_setup(self, L: int, position_ids: torch.Tensor):
        Lp = _ceil(L, 128)
        cos, sin = self.rope_cos_sin(position_ids)
        delta = _rope_delta(position_ids, L)
        if Lp > L:  # padded rows get the continuation positions (never attended by real rows: causal)
            c2, s2 = text_cos_sin(torch.arange(L, Lp) + delta, self.tc.head_dim, self.tc.rope_theta)
            cos, sin = torch.cat([cos, c2]), torch.cat([sin, s2])
        cos_t = self._to_dev(cos.view(1, 1, Lp, -1))
        sin_t = self._to_dev(sin.view(1, 1, Lp, -1))
        chunk = 256 if Lp % 256 == 0 else 128
        sdpa_pc = sdpa_prefill_config(PREFILL_GRID, chunk, chunk)
        return Lp, cos_t, sin_t, delta, sdpa_pc

    def prefill_hidden(self, embeds: torch.Tensor | ttnn.Tensor, position_ids: torch.Tensor, image_embeds=None,
                       deepstack=None, image_positions: torch.Tensor | None = None, keep_on_device: bool = False):
        """Run the loaded layers on inputs_embeds [1, L, hidden] (image rows must be zero if image_embeds is given).
        Returns the hidden state after the last loaded layer (before the final norm): torch [1, L, hidden] bf16, or
        the device tensor [1,1,Lp,hidden] when keep_on_device. Fills the KV cache rows [0, L)."""
        L = position_ids.shape[-1]
        t0 = time.perf_counter()
        Lp, cos_t, sin_t, delta, sdpa_pc = self._prefill_setup(L, position_ids)
        h = self._to_dev(embeds, pad_rows=Lp)
        P = None
        if image_positions is not None and (image_embeds is not None or deepstack):
            P = self._placement(image_positions, Lp)
        if image_embeds is not None:
            h = self._scatter_add(h, P, image_embeds)
        self.timings["prefill_setup"] = time.perf_counter() - t0
        t1 = time.perf_counter()
        for i in self.layers:
            h = self._layer_prefill(i, h, cos_t, sin_t, Lp, sdpa_pc)
            if deepstack is not None and i < len(deepstack):
                h = self._scatter_add(h, P, deepstack[i])
        ttnn.deallocate(cos_t)
        ttnn.deallocate(sin_t)
        if P is not None:
            ttnn.deallocate(P)
        self.cur_len = L
        self.rope_delta = delta
        if keep_on_device:
            ttnn.synchronize_device(self.device)
            self.timings["prefill_layers"] = time.perf_counter() - t1
            return h
        out = ttnn.to_torch(h)[0, 0, :L].unsqueeze(0)
        self.timings["prefill_layers"] = time.perf_counter() - t1
        ttnn.deallocate(h)
        return out

    def prefill(self, input_ids: torch.Tensor, position_ids: torch.Tensor, image_embeds=None, deepstack=None,
                image_positions: torch.Tensor | None = None, embeds: torch.Tensor | None = None) -> torch.Tensor:
        """Full prefill: logits [1, V] float32 at the last position; KV cache rows [0, L) filled.
        `embeds` overrides the host embedding lookup (golden inputs_embeds with the image rows already scattered)."""
        t0 = time.perf_counter()
        L = int(input_ids.shape[-1])
        if embeds is None:
            embeds = self.embed_tokens(input_ids).clone()
            if image_embeds is not None:
                embeds[0, image_positions.long()] = 0
        h = self.prefill_hidden(embeds, position_ids, image_embeds, deepstack, image_positions, keep_on_device=True)
        t1 = time.perf_counter()
        hn = self._rms(h, self.norm_w_tile)
        if self.debug_capture is not None:
            self.debug_capture["norm.in"] = ttnn.to_torch(h)[0]
            self.debug_capture["norm.out"] = ttnn.to_torch(hn)[0]
        ttnn.deallocate(h)
        r0 = (L - 1) // TILE * TILE
        last = ttnn.slice(hn, [0, 0, r0, 0], [1, 1, r0 + TILE, self.tc.hidden])
        ttnn.deallocate(hn)
        last = ttnn.to_memory_config(last, ttnn.L1_MEMORY_CONFIG)
        logits = ttnn.linear(last, self.lm_head, compute_kernel_config=self.ck_head, program_config=self.pc_head,
                             memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn.bfloat16)
        ttnn.deallocate(last)
        r = (L - 1) % TILE
        row = ttnn.untilize_with_unpadding(logits, [0, 0, r, self.tc.vocab - 1], memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(logits)
        out = ttnn.to_torch(row)[0, 0, r].float().unsqueeze(0)
        ttnn.deallocate(row)
        self.timings["prefill_head"] = time.perf_counter() - t1
        self.timings["prefill_total"] = time.perf_counter() - t0
        return out

    # ------------------------------------------------------------------ decode
    def _norm_sharded(self, x, w):
        mem = x.memory_config()
        xi = ttnn.to_memory_config(x, ttnn.L1_MEMORY_CONFIG)
        ttnn.deallocate(x)
        xn = self._rms(xi, w, mem=ttnn.L1_MEMORY_CONFIG)  # [64, 128]: small CBs
        ttnn.deallocate(xi)
        out = ttnn.to_memory_config(xn, mem)
        ttnn.deallocate(xn)
        return out

    def _res_norm(self, x, w):
        """Sharded RMSNorm of the residual stream (both width-sharded over 80 cores)."""
        return ttnn.rms_norm(x, epsilon=self.tc.rms_eps, weight=w, program_config=self.norm_pc,
                             compute_kernel_config=self.ck_norm, memory_config=self.res_mem)

    def _ds_linear(self, x, w, pc, ck):
        return ttnn.linear(x, w, compute_kernel_config=ck, program_config=pc, dtype=ttnn.bfloat16,
                           memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG)

    def _layer_decode(self, i: int, x, cos_s, sin_s, pos):
        """x: residual [1,1,32,hidden] width-sharded (res_mem); returns the same layout."""
        w = self.layer_w[i]
        tc = self.tc
        L1 = ttnn.L1_MEMORY_CONFIG
        n1 = self._res_norm(x, w["ln1"])
        qkv = self._ds_linear(n1, w["wqkv"], self.pc_qkv, self.ck_attn)  # [32, 10240] sharded 80c
        ttnn.deallocate(n1)
        qkv_i = ttnn.to_memory_config(qkv, L1)
        ttnn.deallocate(qkv)
        qkv_i = ttnn.reshape(qkv_i, (1, 1, 1, tc.qkv_dim), (1, 1, TILE, tc.qkv_dim))
        q, k, v = ttnn.experimental.nlp_create_qkv_heads_decode(qkv_i, num_heads=tc.n_heads, num_kv_heads=tc.n_kv_heads,
                                                                memory_config=self.heads_mem)
        ttnn.deallocate(qkv_i)
        q = self._norm_sharded(q, w["qn"])
        k = self._norm_sharded(k, w["kn"])
        qr = ttnn.experimental.rotary_embedding_hf(q, cos_s, sin_s, is_decode_mode=True)
        ttnn.deallocate(q)
        kr = ttnn.experimental.rotary_embedding_hf(k, cos_s, sin_s, is_decode_mode=True)
        ttnn.deallocate(k)
        K, V = self.kv_cache[i]
        ttnn.experimental.paged_update_cache(K, kr, update_idxs_tensor=pos)
        ttnn.experimental.paged_update_cache(V, v, update_idxs_tensor=pos)
        ttnn.deallocate(kr)
        ttnn.deallocate(v)
        a = ttnn.transformer.scaled_dot_product_attention_decode(qr, K, V, cur_pos_tensor=pos, scale=self.scale,
                                                                 program_config=self.sdpa_dec_pc,
                                                                 compute_kernel_config=self.ck_sdpa,
                                                                 memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(qr)
        a_sh = ttnn.to_memory_config(a, self.attn_out_mem)
        ttnn.deallocate(a)
        ac = ttnn.experimental.nlp_concat_heads_decode(a_sh, num_heads=tc.n_heads)  # [32, 8192] sharded by head
        ttnn.deallocate(a_sh)
        ac32 = ttnn.to_memory_config(ac, self.attn_in_mem)
        ttnn.deallocate(ac)
        o = self._ds_linear(ac32, w["wo"], self.pc_wo, self.ck_attn)  # [32, 5120] sharded 32c
        ttnn.deallocate(ac32)
        o80 = ttnn.to_memory_config(o, self.res_mem)
        ttnn.deallocate(o)
        x2 = ttnn.add(x, o80, memory_config=self.res_mem)
        ttnn.deallocate(o80)
        ttnn.deallocate(x)
        n2 = self._res_norm(x2, w["ln2"])
        g = self._ds_linear(n2, w["w1"], self.pc_w1, self.ck_mlp)  # silu fused, [32, 25600] sharded 80c
        u = self._ds_linear(n2, w["w3"], self.pc_w3, self.ck_mlp)
        ttnn.deallocate(n2)
        m = ttnn.mul(g, u, memory_config=g.memory_config())
        ttnn.deallocate(g)
        ttnn.deallocate(u)
        d = self._ds_linear(m, w["w2"], self.pc_w2, self.ck_mlp)  # [32, 5120] sharded 80c
        ttnn.deallocate(m)
        x3 = ttnn.add(x2, d, memory_config=self.res_mem)
        ttnn.deallocate(d)
        ttnn.deallocate(x2)
        return x3

    def _decode_graph(self, x_in, cos_in, sin_in, pos, with_head: bool = True):
        """x_in [1,1,32,hidden] bf16 (row 0 = token embedding), cos/sin [1,1,1,head_dim] interleaved, pos [1] int32."""
        x = ttnn.to_memory_config(x_in, self.res_mem)
        cos_s = ttnn.to_memory_config(cos_in, self.cos_mem)
        sin_s = ttnn.to_memory_config(sin_in, self.cos_mem)
        for i in self.layers:
            x = self.layer_fns.get(i, self._layer_decode)(i, x, cos_s, sin_s, pos)
        ttnn.deallocate(cos_s)
        ttnn.deallocate(sin_s)
        if not with_head:
            xi = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG)
            ttnn.deallocate(x)
            return xi
        xn = self._res_norm(x, self.norm_w)
        ttnn.deallocate(x)
        xi = ttnn.to_memory_config(xn, ttnn.L1_MEMORY_CONFIG)  # 1D mcast matmul wants interleaved in0
        ttnn.deallocate(xn)
        logits = ttnn.linear(xi, self.lm_head, compute_kernel_config=self.ck_head, program_config=self.pc_head,
                             memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn.bfloat16)
        ttnn.deallocate(xi)
        # row 0 only, row-major [1,1,1,V]: 0.03 ms + 0.06 ms read vs 2.6 ms for the padded tile
        row = ttnn.untilize_with_unpadding(logits, [0, 0, 0, self.tc.vocab - 1], memory_config=ttnn.DRAM_MEMORY_CONFIG)
        ttnn.deallocate(logits)
        return row

    def _decode_inputs_host(self, embed_row: torch.Tensor, pos: int):
        x = torch.zeros(1, 1, TILE, self.tc.hidden, dtype=torch.bfloat16)
        x[0, 0, 0] = embed_row.reshape(-1).to(torch.bfloat16)
        cos, sin = text_cos_sin(torch.tensor([pos + self.rope_delta]), self.tc.head_dim, self.tc.rope_theta)
        return (ttnn.from_torch(x, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT),
                ttnn.from_torch(cos.view(1, 1, 1, -1).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT),
                ttnn.from_torch(sin.view(1, 1, 1, -1).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT),
                ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32))

    def _decode_alloc(self):
        dev, tc = self.device, self.tc
        D = ttnn.DRAM_MEMORY_CONFIG
        return {
            "x": ttnn.allocate_tensor_on_device(ttnn.Shape((1, 1, TILE, tc.hidden)), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, D),
            "cos": ttnn.allocate_tensor_on_device(ttnn.Shape((1, 1, 1, tc.head_dim)), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, D),
            "sin": ttnn.allocate_tensor_on_device(ttnn.Shape((1, 1, 1, tc.head_dim)), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, D),
            "pos": ttnn.allocate_tensor_on_device(ttnn.Shape((1,)), ttnn.int32, ttnn.ROW_MAJOR_LAYOUT, dev, D),
        }

    def _decode_upload(self, bufs, embed_row: torch.Tensor, pos: int):
        hx, hc, hs, hp = self._decode_inputs_host(embed_row, pos)
        ttnn.copy_host_to_device_tensor(hx, bufs["x"])
        ttnn.copy_host_to_device_tensor(hc, bufs["cos"])
        ttnn.copy_host_to_device_tensor(hs, bufs["sin"])
        ttnn.copy_host_to_device_tensor(hp, bufs["pos"])

    def decode_hidden(self, embed_row: torch.Tensor, pos: int | None = None) -> torch.Tensor:
        """Untraced single decode step through the loaded layers; returns hidden [1, 1, hidden] (tests)."""
        pos = self.cur_len if pos is None else pos
        bufs = self._decode_alloc()
        self._decode_upload(bufs, embed_row, pos)
        x = self._decode_graph(bufs["x"], bufs["cos"], bufs["sin"], bufs["pos"], with_head=False)
        out = ttnn.to_torch(x)[0, :, :1].float()
        ttnn.deallocate(x)
        for t in bufs.values():
            ttnn.deallocate(t)
        self.cur_len = pos + 1
        return out

    def decode(self, token_id: int) -> torch.Tensor:
        """Traced decode step for `token_id` at position cur_len -> logits [1, V] float32; cur_len += 1."""
        t0 = time.perf_counter()
        dev = self.device
        row = self.embed_table[int(token_id)]
        if self._dec is None:
            bufs = self._decode_alloc()
            self._decode_upload(bufs, row, self.cur_len)
            # compile run (also produces valid logits for this step)
            out = self._decode_graph(bufs["x"], bufs["cos"], bufs["sin"], bufs["pos"])
            logits = ttnn.to_torch(out).reshape(1, -1).float()
            ttnn.deallocate(out)
            # capture: same inputs, the KV write for this position is idempotent
            tid = ttnn.begin_trace_capture(dev, cq_id=0)
            out = self._decode_graph(bufs["x"], bufs["cos"], bufs["sin"], bufs["pos"])
            ttnn.end_trace_capture(dev, tid, cq_id=0)
            ttnn.synchronize_device(dev)
            self._dec = {"bufs": bufs, "tid": tid, "out": out}
            self.cur_len += 1
            self.timings["decode_compile"] = time.perf_counter() - t0
            return logits
        d = self._dec
        self._decode_upload(d["bufs"], row, self.cur_len)
        t1 = time.perf_counter()
        ttnn.execute_trace(dev, d["tid"], cq_id=0, blocking=False)
        logits = ttnn.to_torch(d["out"]).reshape(1, -1).float()
        self.timings["decode_trace_exec"] = time.perf_counter() - t1
        self.timings["decode_step"] = time.perf_counter() - t0
        self.cur_len += 1
        return logits

    def release(self):
        if self._dec is not None:
            ttnn.release_trace(self.device, self._dec["tid"])
            self._dec = None