code/alpamayo_tt/expert.py

"""Action expert (flow-matching trajectory decoder) of Alpamayo2-Super on one Tenstorrent chip, batch 1.

Per Euler step: action_in_proj (host torch, tiny) -> 64-layer Qwen3-VL text transformer (device, traced) whose
layer i attends non-causally to [VLM KV cache of layer i ; its own 64 tokens] -> final RMSNorm (device) ->
action_out_proj (host). See EXPERT_NOTES.md for layouts, masking and measurements.

KV cache interface (provided by the text model): ``kv_cache[i] = (K, V)``, each ``[1, n_kv, Smax, head_dim]``
TILE DRAM (bf16 or bfp8) in HF head order, rows ``[0, S)`` valid. The expert writes its own 64 roped keys /
values into rows ``[Smax-64, Smax)`` of every layer's cache (``ttnn.fill_cache`` with a tile-aligned offset) and
runs one attention op over the whole cache with an additive mask that hides rows ``[S, Smax-64)``. Hence the
traced step does not depend on S; only the mask / cos / sin buffers (rewritten from host per trajectory) do.

Attention op (``attn_mode``):
  * "decode" (default): flash-decode kernel ``scaled_dot_product_attention_decode`` with the cache viewed as 8
    users x 1 kv head ``[8, 1, Smax, 128]`` and the queries viewed as ``[1, 8, 128, 128]`` (user g = kv head g,
    its 128 "heads" = the 2 q heads of the group x 64 tokens). K/V are split across cores (~16 per kv head), so the
    cache is read once, in parallel: ~170 us for Smax = 8192 vs ~520 us for the standard kernel.
  * "standard": ``scaled_dot_product_attention`` (non-causal, GQA); every q chunk re-reads the whole K/V.
"""
from __future__ import annotations

import time
from dataclasses import dataclass

import torch
import ttnn

from .action_proj import ActionInProj, ActionOutProj, expert_positions, rope_cos_sin
from .config import ModelConfig
from .weights import DTYPES, EXPERT, Checkpoint, TTWeights, fused_qkv, linear_weight

MASK_MIN = -1e9  # additive mask value for hidden cache rows (bf16/bfp4 representable, exp() underflows to 0)


def _norm_weight(w: torch.Tensor) -> torch.Tensor:
    """HF RMSNorm weight [dim] -> ttnn.rms_norm gamma layout [1, 1, 1, dim] (TILE; width must equal input width)."""
    return w.reshape(1, 1, 1, w.numel())


@dataclass
class LayerWeights:
    ln1: ttnn.Tensor
    ln2: ttnn.Tensor
    qn: ttnn.Tensor
    kn: ttnn.Tensor
    wqkv: ttnn.Tensor
    wo: ttnn.Tensor
    w1: ttnn.Tensor
    w3: ttnn.Tensor
    w2: ttnn.Tensor


class TTExpert:
    """Op-graph implementation. ``layer_fn`` (one layer: (i, x, k_cache, v_cache) -> x) is the unit a fused kernel
    can replace; everything else (buffers, context, trace, Euler loop) stays."""

    def __init__(self, device, cfg: ModelConfig, ckpt: Checkpoint, tw: TTWeights, dtype: str = "bfp8",
                 layers: range | None = None, attn_mode: str = "decode", mask_dtype: str | None = None,
                 sdpa_k_chunk: int | None = None, sdpa_q_chunk: int = 32, fold_heads: bool = False,
                 use_trace: bool = True, in_proj_dtype: torch.dtype = torch.float32, tuned_matmul: bool = True):
        assert attn_mode in ("decode", "standard")
        self.device = device
        self.cfg = cfg
        self.tc = cfg.expert.text
        self.n_tok = cfg.expert.n_waypoints
        self.layers = list(layers) if layers is not None else list(range(self.tc.n_layers))
        self.dtype = dtype
        self.attn_mode = attn_mode
        self.mask_dtype = mask_dtype or ("bfp4" if attn_mode == "decode" else "bf16")
        self.use_trace = use_trace
        self.scale = self.tc.head_dim ** -0.5
        tc = self.tc
        assert tc.n_heads % tc.n_kv_heads == 0
        self.rep = tc.n_heads // tc.n_kv_heads
        # decode mode always groups the q heads of a kv head onto rows; standard mode optionally does (fold_heads)
        self.fold_heads = True if attn_mode == "decode" else fold_heads
        self.sq = self.n_tok * self.rep if self.fold_heads else self.n_tok  # mask rows / query rows per kv head
        grid = device.compute_with_storage_grid_size()
        self.grid = grid
        if attn_mode == "decode":
            self.sdpa_cfg = ttnn.SDPAProgramConfig(
                compute_with_storage_grid_size=grid, q_chunk_size=32, k_chunk_size=sdpa_k_chunk or 128,
                exp_approx_mode=False, max_cores_per_head_batch=16)
        else:
            self.sdpa_cfg = ttnn.SDPAProgramConfig(
                compute_with_storage_grid_size=grid, q_chunk_size=sdpa_q_chunk, k_chunk_size=sdpa_k_chunk or 512,
                exp_approx_mode=False)
        self.ck_mm = ttnn.WormholeComputeKernelConfig(
            math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=False, fp32_dest_acc_en=True, packer_l1_acc=True)
        self.ck_sdpa = ttnn.WormholeComputeKernelConfig(
            math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=False, fp32_dest_acc_en=False, packer_l1_acc=False)
        # rms_norm: the op default config is ~9 us faster per call than HiFi4+fp32 acc on [64, 1536] (same PCC)
        self.ck_norm = None
        self.mem = ttnn.DRAM_MEMORY_CONFIG
        # matmul program configs (measured in EXPERT_NOTES.md; None -> ttnn default heuristics)
        H, Q, I = tc.hidden, tc.q_dim, tc.inter
        if tuned_matmul:
            self.pc_qkv = self._mm1d(H, tc.qkv_dim, in0_block_w=8)
            self.pc_o = self._mm1d(Q, H, in0_block_w=8)
            self.pc_w1 = self._mm1d(H, I, in0_block_w=8, activation=ttnn.UnaryOpType.SILU)
            self.pc_w3 = self._mm1d(H, I, in0_block_w=8)
            self.pc_w2 = self._mm1d(I, H, in0_block_w=48)
        else:
            self.pc_qkv = self.pc_o = self.pc_w1 = self.pc_w3 = self.pc_w2 = None

        # ---- weights -----------------------------------------------------------------------------
        t0 = time.time()
        self.w: dict[int, LayerWeights] = {}
        for i in self.layers:
            self.w[i] = self._load_layer(i, ckpt, tw)
        self.final_norm = tw.to_device("expert.norm_g", _norm_weight(ckpt.get(EXPERT + "norm.weight")), "bf16")
        self.weight_load_s = time.time() - t0
        self.in_proj = ActionInProj(cfg.expert, ckpt, dtype=in_proj_dtype)
        self.out_proj = ActionOutProj(ckpt)

        # ---- persistent device buffers (inputs of the traced graph) ------------------------------
        self.x_dev = self._dev_zeros((1, 1, self.n_tok, H))
        self.cos_dev = self._dev_zeros((1, 1, self.n_tok, tc.head_dim))
        self.sin_dev = self._dev_zeros((1, 1, self.n_tok, tc.head_dim))
        # The mask must exist before ANY trace is captured in this process (a trace's intermediates are placed in
        # memory that is free at capture time; a buffer allocated later can alias them and get clobbered on replay).
        self.mask_dev: ttnn.Tensor = self._dev_zeros((1, 1, self.sq, cfg.max_seq_len), self.mask_dtype)
        self.ctx: tuple | None = None  # (S, rope_delta, Smax) currently loaded into cos/sin/mask
        self.tail: int | None = None  # row offset of own K/V inside the cache
        self._trace: dict | None = None
        self._views: dict[tuple, tuple] = {}  # cache buffer addresses -> [8, 1, Smax, dh] views (decode mode)
        self.layer_fn = self.layer  # swap point for a fused implementation

    # ------------------------------------------------------------------------------------------ setup
    def _mm1d(self, K: int, N: int, in0_block_w: int, activation=None):
        """1D multicast matmul config for x [64, K] @ W [K, N]: in0 broadcast, N split over all cores."""
        nc = self.grid.x * self.grid.y
        Kt, Nt = K // 32, N // 32
        assert Kt % in0_block_w == 0, (K, in0_block_w)
        per_core_N = -(-Nt // nc)
        return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
            compute_with_storage_grid_size=(self.grid.x, self.grid.y), in0_block_w=in0_block_w,
            out_subblock_h=1, out_subblock_w=per_core_N if per_core_N <= 4 else 1,
            out_block_h=self.n_tok // 32, out_block_w=per_core_N, per_core_M=self.n_tok // 32, per_core_N=per_core_N,
            fuse_batch=True, fused_activation=ttnn.UnaryWithParam(activation) if activation is not None else None,
            mcast_in0=True)

    def _linear(self, x, w, pc, activation: str | None = None):
        if pc is not None:
            return ttnn.linear(x, w, dtype=ttnn.bfloat16, program_config=pc, compute_kernel_config=self.ck_mm,
                               memory_config=self.mem)
        return ttnn.linear(x, w, dtype=ttnn.bfloat16, activation=activation, compute_kernel_config=self.ck_mm,
                           memory_config=self.mem)

    def _dev_zeros(self, shape, dtype="bf16"):
        return ttnn.from_torch(torch.zeros(*shape), dtype=DTYPES[dtype], layout=ttnn.TILE_LAYOUT,
                               device=self.device, memory_config=self.mem)

    def _load_layer(self, i: int, ckpt: Checkpoint, tw: TTWeights) -> LayerWeights:
        w = ckpt.text_layer(i, prefix=EXPERT)
        p = f"expert.L{i}."
        d = self.dtype
        return LayerWeights(
            ln1=tw.to_device(p + "ln1_g", _norm_weight(w["input_layernorm.weight"]), "bf16"),
            ln2=tw.to_device(p + "ln2_g", _norm_weight(w["post_attention_layernorm.weight"]), "bf16"),
            qn=tw.to_device(p + "qn_g", _norm_weight(w["self_attn.q_norm.weight"]), "bf16"),
            kn=tw.to_device(p + "kn_g", _norm_weight(w["self_attn.k_norm.weight"]), "bf16"),
            wqkv=tw.to_device(p + "wqkv", fused_qkv(w["self_attn.q_proj.weight"], w["self_attn.k_proj.weight"],
                                                    w["self_attn.v_proj.weight"]), d),
            wo=tw.to_device(p + "wo", linear_weight(w["self_attn.o_proj.weight"]), d),
            w1=tw.to_device(p + "w1", linear_weight(w["mlp.gate_proj.weight"]), d),
            w3=tw.to_device(p + "w3", linear_weight(w["mlp.up_proj.weight"]), d),
            w2=tw.to_device(p + "w2", linear_weight(w["mlp.down_proj.weight"]), d),
        )

    def _upload(self, dev: ttnn.Tensor, t: torch.Tensor):
        host = ttnn.from_torch(t.reshape(tuple(dev.shape)).to(torch.float32), dtype=dev.dtype, layout=ttnn.TILE_LAYOUT)
        ttnn.copy_host_to_device_tensor(host, dev)

    def build_mask(self, S: int, Smax: int) -> torch.Tensor:
        """Additive mask [1, 1, sq, Smax]: 0 on valid cache rows [0, S) and on own rows [Smax-64, Smax)."""
        m = torch.full((1, 1, self.sq, Smax), MASK_MIN, dtype=torch.float32)
        m[..., :S] = 0.0
        m[..., Smax - self.n_tok:] = 0.0
        return m

    def set_context(self, S: int, rope_delta: int, Smax: int, force: bool = False):
        """Load cos/sin (positions i + rope_delta + S) and the cache mask for a trajectory (~1 ms).

        Skips the upload when (S, rope_delta, Smax) is unchanged unless ``force``: other traces (the text model's
        decode) may run between two trajectories, so ``sample()`` always forces a refresh at its start."""
        assert S <= Smax - self.n_tok, f"cache too small: S={S}, Smax={Smax} (last {self.n_tok} rows are expert scratch)"
        key = (int(S), int(rope_delta), int(Smax))
        if self.ctx == key and not force:
            return
        cos, sin = rope_cos_sin(expert_positions(S, rope_delta, self.n_tok), self.tc.head_dim, self.tc.rope_theta)
        self._upload(self.cos_dev, cos)
        self._upload(self.sin_dev, sin)
        if self.mask_dev.shape[-1] != Smax:  # cache longer/shorter than cfg.max_seq_len: reallocate (unsafe if traces exist)
            print(f"[TTExpert] warning: reallocating mask for Smax={Smax} (was {self.mask_dev.shape[-1]}); "
                  "do this before any trace capture")
            self.release_trace()
            self.mask_dev.deallocate()
            self.mask_dev = self._dev_zeros((1, 1, self.sq, Smax), self.mask_dtype)
        self._upload(self.mask_dev, self.build_mask(S, Smax))
        self.tail = Smax - self.n_tok
        self.ctx = key

    def _kv_views(self, k_cache: ttnn.Tensor, v_cache: ttnn.Tensor):
        """[1, n_kv, Smax, dh] -> zero-copy views [n_kv, 1, Smax, dh] for the decode kernel (cached per buffer)."""
        key = (id(k_cache), id(v_cache), k_cache.buffer_address(), v_cache.buffer_address())
        v = self._views.get(key)
        if v is None:
            shp = (self.tc.n_kv_heads, 1, k_cache.shape[2], self.tc.head_dim)
            v = (ttnn.view(k_cache, shp), ttnn.view(v_cache, shp))
            self._views[key] = v
        return v

    # ------------------------------------------------------------------------------------------ layer
    def attention_core(self, q: ttnn.Tensor, k_cache: ttnn.Tensor, v_cache: ttnn.Tensor) -> ttnn.Tensor:
        """q [1, n_heads, 64, dh] (roped) -> attention output [1, n_heads, 64, dh] over the full cache + mask."""
        tc = self.tc
        if self.attn_mode == "decode":
            k8, v8 = self._kv_views(k_cache, v_cache)
            qd = ttnn.view(q, (1, tc.n_kv_heads, self.sq, tc.head_dim))  # [1, users=8, heads=128, dh]
            a = ttnn.transformer.scaled_dot_product_attention_decode(
                qd, k8, v8, is_causal=False, attn_mask=self.mask_dev, scale=self.scale,
                program_config=self.sdpa_cfg, compute_kernel_config=self.ck_sdpa, memory_config=self.mem)
            return ttnn.view(a, (1, tc.n_heads, self.n_tok, tc.head_dim))
        if self.fold_heads:
            q = ttnn.view(q, (1, tc.n_kv_heads, self.sq, tc.head_dim))
        a = ttnn.transformer.scaled_dot_product_attention(
            q, k_cache, v_cache, attn_mask=self.mask_dev, is_causal=False, scale=self.scale,
            program_config=self.sdpa_cfg, compute_kernel_config=self.ck_sdpa, memory_config=self.mem)
        if self.fold_heads:
            a = ttnn.view(a, (1, tc.n_heads, self.n_tok, tc.head_dim))
        return a

    def layer(self, i: int, x: ttnn.Tensor, k_cache: ttnn.Tensor, v_cache: ttnn.Tensor) -> ttnn.Tensor:
        """One decoder layer. x [1, 1, 64, hidden] bf16 -> new tensor of the same shape (x is not freed)."""
        L = self.w[i]
        tc = self.tc
        eps = tc.rms_eps
        # --- attention
        h = ttnn.rms_norm(x, epsilon=eps, weight=L.ln1, compute_kernel_config=self.ck_norm)
        qkv = self._linear(h, L.wqkv, self.pc_qkv)
        h.deallocate()
        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=self.mem)
        qkv.deallocate()
        qn = ttnn.rms_norm(q, epsilon=eps, weight=L.qn, compute_kernel_config=self.ck_norm)
        q.deallocate()
        kn = ttnn.rms_norm(k, epsilon=eps, weight=L.kn, compute_kernel_config=self.ck_norm)
        k.deallocate()
        q = ttnn.experimental.rotary_embedding_hf(qn, self.cos_dev, self.sin_dev, is_decode_mode=False)
        qn.deallocate()
        k = ttnn.experimental.rotary_embedding_hf(kn, self.cos_dev, self.sin_dev, is_decode_mode=False)
        kn.deallocate()
        if k.dtype != k_cache.dtype:
            k = ttnn.typecast(k, k_cache.dtype)
            v = ttnn.typecast(v, v_cache.dtype)
        ttnn.fill_cache(k_cache, k, 0, update_idx=self.tail)
        ttnn.fill_cache(v_cache, v, 0, update_idx=self.tail)
        k.deallocate()
        v.deallocate()
        a = self.attention_core(q, k_cache, v_cache)
        q.deallocate()
        o = ttnn.experimental.nlp_concat_heads(a, memory_config=self.mem)  # [1, 1, 64, n_heads*head_dim]
        a.deallocate()
        o2 = self._linear(o, L.wo, self.pc_o)
        o.deallocate()
        x1 = ttnn.add(x, o2, memory_config=self.mem)
        o2.deallocate()
        # --- MLP
        h2 = ttnn.rms_norm(x1, epsilon=eps, weight=L.ln2, compute_kernel_config=self.ck_norm)
        g = self._linear(h2, L.w1, self.pc_w1, activation="silu")
        u = self._linear(h2, L.w3, self.pc_w3)
        h2.deallocate()
        gu = ttnn.multiply(g, u, memory_config=self.mem)
        g.deallocate()
        u.deallocate()
        m = self._linear(gu, L.w2, self.pc_w2)
        gu.deallocate()
        x2 = ttnn.add(x1, m, memory_config=self.mem)
        x1.deallocate()
        m.deallocate()
        return x2

    def _forward_dev(self, kv_cache, layers, final_norm: bool) -> ttnn.Tensor:
        x = self.x_dev
        for i in layers:
            k_cache, v_cache = kv_cache[i]
            x_new = self.layer_fn(i, x, k_cache, v_cache)
            if x is not self.x_dev:
                x.deallocate()
            x = x_new
        if final_norm:
            out = ttnn.rms_norm(x, epsilon=self.tc.rms_eps, weight=self.final_norm, compute_kernel_config=self.ck_norm)
            if x is not self.x_dev:
                x.deallocate()
            return out
        return x

    # ------------------------------------------------------------------------------------------ trace
    def _trace_key(self, kv_cache):
        return tuple((kv_cache[i][0].buffer_address(), kv_cache[i][1].buffer_address()) for i in self.layers) + (
            self.mask_dev.buffer_address(),)

    def _capture(self, kv_cache):
        out = self._forward_dev(kv_cache, self.layers, True)  # warm-up compiles all kernels
        ttnn.synchronize_device(self.device)
        out.deallocate()
        tid = ttnn.begin_trace_capture(self.device, cq_id=0)
        out = self._forward_dev(kv_cache, self.layers, True)
        ttnn.end_trace_capture(self.device, tid, cq_id=0)
        ttnn.synchronize_device(self.device)
        self._trace = {"tid": tid, "out": out, "key": self._trace_key(kv_cache)}

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

    def forward_hidden(self, emb: torch.Tensor, kv_cache, S: int, rope_delta: int) -> torch.Tensor:
        """emb [1, 64, hidden] (action_in_proj output) -> final-normed hidden [1, 64, hidden] fp32."""
        Smax = kv_cache[self.layers[0]][0].shape[2]
        self.set_context(S, rope_delta, Smax)
        self._upload(self.x_dev, emb)
        if self.use_trace:
            if self._trace is None or self._trace["key"] != self._trace_key(kv_cache):
                self.release_trace()
                self._capture(kv_cache)
            ttnn.execute_trace(self.device, self._trace["tid"], cq_id=0, blocking=False)
            res = ttnn.to_torch(self._trace["out"]).float()
        else:
            out = self._forward_dev(kv_cache, self.layers, True)
            res = ttnn.to_torch(out).float()
            out.deallocate()
        return res.reshape(1, self.n_tok, self.tc.hidden)

    # ------------------------------------------------------------------------------------------ public API
    def step(self, x_t: torch.Tensor, t: float, kv_cache, S: int, rope_delta: int) -> torch.Tensor:
        """One expert forward: x_t [1, 64, 2], t scalar -> velocity v [1, 64, 2] fp32."""
        emb = self.in_proj(x_t, t)
        hidden = self.forward_hidden(emb, kv_cache, S, rope_delta)
        return self.out_proj(hidden)

    def sample(self, x0: torch.Tensor, kv_cache, S: int, rope_delta: int, n_steps: int = 10,
               record: list | None = None) -> torch.Tensor:
        """Euler flow matching: t = linspace(0, 1, n_steps+1)[:-1], x <- x + dt * v(x, t). Returns the action."""
        ts = torch.linspace(0.0, 1.0, n_steps + 1)
        x = x0.clone().float()
        self.set_context(S, rope_delta, kv_cache[self.layers[0]][0].shape[2], force=True)  # see set_context
        for i in range(n_steps):
            v = self.step(x, float(ts[i]), kv_cache, S, rope_delta)
            if record is not None:
                record.append({"t": float(ts[i]), "x": x.clone(), "v": v.clone()})
            x = x + (ts[i + 1] - ts[i]) * v
        return x

    # ------------------------------------------------------------------------------------------ testing helpers
    def run_layers(self, x: torch.Tensor, kv_cache, S: int, rope_delta: int, layers=None,
                   final_norm: bool = False) -> torch.Tensor:
        """Eager: x [1, 64, hidden] -> output of the given layers (default: all loaded), optional final norm."""
        layers = list(layers) if layers is not None else self.layers
        Smax = kv_cache[layers[0]][0].shape[2]
        self.set_context(S, rope_delta, Smax)
        self._upload(self.x_dev, x)
        out = self._forward_dev(kv_cache, layers, final_norm)
        res = ttnn.to_torch(out).float().reshape(1, self.n_tok, self.tc.hidden)
        if out is not self.x_dev:
            out.deallocate()
        return res

    def bench_step(self, kv_cache, S: int, rope_delta: int, iters: int = 10, warmup: int = 2) -> dict:
        """Wall time of the device part of one Euler step (traced, warm). Returns ms statistics."""
        emb = torch.randn(1, self.n_tok, self.tc.hidden)
        for _ in range(warmup):
            self.forward_hidden(emb, kv_cache, S, rope_delta)
        ts = []
        for _ in range(iters):
            t0 = time.perf_counter()
            self.forward_hidden(emb, kv_cache, S, rope_delta)
            ts.append((time.perf_counter() - t0) * 1e3)
        ts.sort()
        return {"ms_median": ts[len(ts) // 2], "ms_min": ts[0], "ms_max": ts[-1], "n_layers": len(self.layers)}

    def _trace_us(self, fn, reps: int, protected: set[int]) -> float:
        """Device time (us) per call of fn(), measured as a trace of `reps` back-to-back calls."""
        dev = self.device
        r = fn()
        ttnn.synchronize_device(dev)
        tid = ttnn.begin_trace_capture(dev, cq_id=0)
        tmp = [fn() for _ in range(reps)]
        ttnn.end_trace_capture(dev, tid, cq_id=0)
        ttnn.synchronize_device(dev)
        ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
        t0 = time.perf_counter()
        for _ in range(3):
            ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
        us = (time.perf_counter() - t0) / 3 / reps * 1e6
        ttnn.release_trace(dev, tid)
        for t in tmp:
            for tt in (t if isinstance(t, (tuple, list)) else (t,)):
                if isinstance(tt, ttnn.Tensor) and tt.buffer_address() not in protected:
                    tt.deallocate()
        return us, r

    def bench_layer_traced(self, i: int, kv_cache, S: int, rope_delta: int, reps: int = 8) -> float:
        """Device time of one full layer (us): trace of `reps` back-to-back executions of layer i."""
        k_cache, v_cache = kv_cache[i]
        self.set_context(S, rope_delta, k_cache.shape[2])
        us, r = self._trace_us(lambda: self.layer_fn(i, self.x_dev, k_cache, v_cache), reps, set())
        r.deallocate()
        return us

    def layer_breakdown(self, i: int, kv_cache, S: int, rope_delta: int, reps: int = 10) -> dict[str, float]:
        """Per-op device time (us) of layer i: each op is captured `reps` times in its own trace (no host overhead)."""
        k_cache, v_cache = kv_cache[i]
        self.set_context(S, rope_delta, k_cache.shape[2])
        L = self.w[i]
        tc = self.tc
        eps = tc.rms_eps
        x = self.x_dev
        prot = {k_cache.buffer_address(), v_cache.buffer_address()}
        res: dict[str, float] = {}

        def timed(name, fn):
            us, r = self._trace_us(fn, reps, prot)
            res[name] = us
            return r

        h = timed("rms_norm", lambda: ttnn.rms_norm(x, epsilon=eps, weight=L.ln1, compute_kernel_config=self.ck_norm))
        qkv = timed("wqkv", lambda: self._linear(h, L.wqkv, self.pc_qkv))
        q, k, v = timed("split_heads", lambda: ttnn.experimental.nlp_create_qkv_heads(
            qkv, num_heads=tc.n_heads, num_kv_heads=tc.n_kv_heads, transpose_k_heads=False))
        qn = timed("q_norm", lambda: ttnn.rms_norm(q, epsilon=eps, weight=L.qn, compute_kernel_config=self.ck_norm))
        kn = timed("k_norm", lambda: ttnn.rms_norm(k, epsilon=eps, weight=L.kn, compute_kernel_config=self.ck_norm))
        qr = timed("rope_q", lambda: ttnn.experimental.rotary_embedding_hf(qn, self.cos_dev, self.sin_dev, is_decode_mode=False))
        kr = timed("rope_k", lambda: ttnn.experimental.rotary_embedding_hf(kn, self.cos_dev, self.sin_dev, is_decode_mode=False))
        timed("fill_k", lambda: ttnn.fill_cache(k_cache, kr, 0, update_idx=self.tail))
        timed("fill_v", lambda: ttnn.fill_cache(v_cache, v, 0, update_idx=self.tail))
        a = timed("attention", lambda: self.attention_core(qr, k_cache, v_cache))
        o = timed("concat_heads", lambda: ttnn.experimental.nlp_concat_heads(a))
        o2 = timed("wo", lambda: self._linear(o, L.wo, self.pc_o))
        x1 = timed("add", lambda: ttnn.add(x, o2))
        h2 = timed("rms_norm2", lambda: ttnn.rms_norm(x1, epsilon=eps, weight=L.ln2, compute_kernel_config=self.ck_norm))
        g = timed("w1_silu", lambda: self._linear(h2, L.w1, self.pc_w1, activation="silu"))
        u = timed("w3", lambda: self._linear(h2, L.w3, self.pc_w3))
        gu = timed("mul", lambda: ttnn.multiply(g, u))
        m = timed("w2", lambda: self._linear(gu, L.w2, self.pc_w2))
        timed("add2", lambda: ttnn.add(x1, m))
        res["sum_of_ops"] = sum(res.values())
        return res