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