# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """Host side of the DiT megakernel launch -- K1 rung (IMPLEMENTATION_PLAN.md §5.1, §5.3, §5.7 K1; WP-K1) and the K4 step schedule (plan §5.4 P1-P12 + "Per step", §5.7 K4; WP-K4). :class:`DiTMegakernel` owns the device tensors of one arena (``arena_stream`` WIDTH_SHARDED over the 8 banks, the ``bias_rows`` / ``adaln_rows`` tables, the mask / hoist placeholders of the ``block`` variant, the ``a_rep`` in0 arena on all 110 cores) plus the K1 working tensors (``w_ring`` = the tensor-backed ``cb_w`` weight ring, ``mm_out`` = the tensor-backed ``cb_mm0`` [64,64] output per compute core, ``sync`` = 4 KB of zeroed L1 per core for the relay credits) and launches ``kernels/dit_kernel.cpp`` through ``descriptors.build_program_spec`` + ``materialise`` as one ``ttnn.generic_op``. Three launch modes (``kernels/dit_kernel.cpp`` header): :meth:`run_block` = the matmul-only block chain ``[Wqkv|Wq] -> gather Q -> mcast -> Wo -> gather -> mcast -> ff1+GELU -> 6 K-chunk gathers/mcasts -> ff2 -> gather -> mcast`` (plan §5.4 P2/P5/P6/P7/P9/P10-12 without LN / attention); :meth:`run_single_op` = one projection with ``in0 = x_rep`` (the PCC acceptance shapes); :meth:`run_stream` = the four projections of every block with no gathers (the pure weight-stream rate, G1). **Adapter over WP-05's descriptors (K1 additions, listed in the WP report):** ``cb_w`` and ``cb_mm0`` are made tensor-backed (the reader cores must know the compute cores' ``cb_w`` address; the host reads ``mm_out`` back), the named CTAs ``mk_*`` and the common runtime args :data:`MK_RT` (+ reader / compute NoC coordinate tables) are appended to the :class:`descriptors.ProgramSpec` before ``materialise``; the extra tensors are appended to the io list (``mm_out`` last = the ``generic_op`` output). Nothing in ``descriptors.py`` / ``arena.py`` is modified. The kernel derives every op's slice geometry structurally (``kernel_geometry``); :func:`check_kernel_geometry` asserts that mirror against :class:`arena.ArenaLayout` for every op of a layout (CPU test). **K4 (WP-K4):** the same :class:`DiTMegakernel` launches the full kernel through :func:`descriptors.build_program_spec_k4` (direct reader only): :meth:`run_block(x, blk, step=k, matmul_only=False)` (one full block, ``block`` variant), :meth:`run_step(sa_embs, k)` (all blocks + output head -> ``dit_out``, ``step`` variant), :meth:`run_steps4(noise, state_features)` (the 4-step denoise incl. decoder / Euler / action encoder on device, ``steps4`` variant) with :meth:`capture_steps4` / :meth:`replay` for the G2 timing. The hoisted K/V and the per-call masks are bound with :meth:`bind_kv` / :meth:`bind_masks` (device tensors of ``tt.adapter`` / ``tt.layout.materialise``); the real weights come from :func:`real_arena_plan` (the ``common.weights`` plan the Stage-1 oracle uploads). The K1 kernel modes stay behind the ``mk_k1`` CTA; both specs emit the union :data:`KERNEL_CTA_UNION` of named CTAs (a missing name is a JIT failure). The tail geometry of ``mk4::tail_geom`` is mirrored by :func:`tail_kernel_geometry` / :func:`check_tail_geometry`. ``ttnn`` is imported lazily; everything up to :meth:`DiTMegakernel.__init__` runs on the CPU. """ from __future__ import annotations import math import os import time from dataclasses import dataclass, replace from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import torch from models.experimental.gr00t.common.weights import enumerate_plan from models.experimental.gr00t.tt.megakernel import arena as A from models.experimental.gr00t.tt.megakernel import descriptors as D from models.experimental.gr00t.tt.megakernel.core_map import N_BANKS, CoreMap, Role, group_of TILE = A.TILE PAGE_TILES = 8 # weight page = one NoC burst (weight_stream.hpp) M_ROWS = 64 RT_TILES = 2 CHUNK_COLS = 1024 # ff2 K-chunk [64, 1024] (plan §5.4 P10/11) N_CHUNKS = 6 MODES: Mapping[str, int] = {"chain": 0, "single": 1, "stream": 2} OP_NAMES: Tuple[str, ...] = ("qkv", "Wo", "ff1", "ff2") CORE_SETS: Tuple[str, ...] = ("all96", "first32", "even48") # kernel CoreSet enum order #: K1 common runtime args appended after WP-05's (dit_kernel.cpp ``enum RT``); then reader_xy[8], compute_xy[96] MK_RT: Tuple[str, ...] = ( "mode", "single_op", "blk_first", "n_blk", "n_iters", "x_tiles", "addr_sync", "addr_w_ring", "addr_mm_out", "bias_row_bytes", "pop_mm", ) DEFAULT_RING_DEPTH = 4 # DRAM pages in flight per PageRing DEFAULT_RELAY_WINDOW = 6 # relay writes in flight per reader BRISC IN0_ORDERS: Tuple[str, ...] = ("RowMajor", "Pairs", "Single", "RowMajorWrap") # -------------------------------------------------------------------------------------------------------------------- # kernel geometry mirror (dit_kernel.cpp mk::op_geom & friends) # -------------------------------------------------------------------------------------------------------------------- @dataclass(frozen=True) class KernelOpGeom: """What ``mk::op_geom(blk, j)`` computes on device for op ``j`` of block ``blk``.""" blk: int j: int kt: int nt: int ct: int coreset: str per_bank: int pages_core: int off_bank: int bias_row: int gelu: bool @property def n_cores(self) -> int: return self.nt // self.ct @property def page_bytes(self) -> int: raise AttributeError("page bytes depend on the stream dtype; use ArenaLayout") def is_split_block(layout: A.ArenaLayout, blk: int) -> bool: """K4b: a cross block whose key set has >= :data:`descriptors.K4_SPLIT_MIN_CHUNKS` 64-key chunks shares every head over cores h / 32+h / 64+h (``dit_kernel.cpp`` ``mk::split_blk``: ``ceil(l_pad / 64) >= 3``).""" if not layout.is_cross(blk): return False return -(-layout.l_pad[layout.key_sets[blk]] // 64) >= D.K4_SPLIT_MIN_CHUNKS def _op_shape(layout: A.ArenaLayout, is_cross: bool, j: int) -> Tuple[int, int, str, int]: """``(K, N, core set, ct)`` of op ``j`` exactly as the kernel's switch (structural constants).""" w, inner, ff = layout.width, layout.inner, layout.ff_dim if j == 0: return (w, inner, "first32", 2) if is_cross else (w, 3 * inner, "all96", 2) if j == 1: return inner, w, "even48", 1 if j == 2: return w, ff, "all96", 2 if j == 3: return ff, w, "even48", 1 raise IndexError(f"op index {j} out of range 0..3") def kernel_geometry(layout: A.ArenaLayout, blk: int, j: int) -> KernelOpGeom: """Mirror of ``mk::op_geom``: byte offsets from the block's first-op word plus structural op sizes.""" if not 0 <= blk < layout.n_blocks: raise IndexError(f"block {blk} out of range for {layout.n_blocks} blocks") is_cross = layout.is_cross(blk) tile_bytes = A.TILE_BYTES[layout.dtype_stream] words = layout.offsets_table() base = len(A.OFFSETS_HEADER) + A.WORDS_PER_BLOCK * blk off = words[base + 0] bias_row_base = words[base + 1] for i in range(j): k, n, _, ct = _op_shape(layout, is_cross, i) per_bank_i = (n // TILE // ct) // N_BANKS off += per_bank_i * (k // TILE) * ct * tile_bytes k, n, cs, ct = _op_shape(layout, is_cross, j) kt, nt = k // TILE, n // TILE return KernelOpGeom( blk=blk, j=j, kt=kt, nt=nt, ct=ct, coreset=cs, per_bank=(nt // ct) // N_BANKS, pages_core=kt * ct // PAGE_TILES, off_bank=off, bias_row=bias_row_base + j, gelu=(j == 2), ) def participates(coreset: str, c: int) -> bool: if coreset == "all96": return True if coreset == "first32": return c < 32 if coreset == "even48": return c % 2 == 0 raise KeyError(coreset) def slot_of(coreset: str, c: int) -> int: """Position of compute core ``c`` in the op's core list (N tiles ``[slot*ct, (slot+1)*ct)``).""" if not participates(coreset, c): raise KeyError(f"core {c} is not in {coreset}") return c // 2 if coreset == "even48" else c def rank_of(coreset: str, c: int) -> int: """Position of ``c`` among the op's cores of its own bank group (``OpGeometry.cores_in_bank`` order).""" if not participates(coreset, c): raise KeyError(f"core {c} is not in {coreset}") if coreset == "first32": return c & 3 if coreset == "even48": return 2 * (c >> 5) + ((c & 3) >> 1) return 4 * (c >> 5) + (c & 3) def member_core(coreset: str, bank: int, r: int) -> int: if coreset == "first32": return 4 * bank + r if coreset == "even48": return 32 * (r >> 1) + 4 * bank + 2 * (r & 1) return 32 * (r >> 2) + 4 * bank + (r & 3) def rank12(c: int) -> int: """Stable index of a compute core inside its bank group (credit word, relay bookkeeping).""" return 4 * (c >> 5) + (c & 3) def check_kernel_geometry(layout: A.ArenaLayout) -> int: """Assert the kernel's structural geometry against :class:`arena.ArenaLayout` for every streamed op; returns the number of ops checked. Raises ``AssertionError`` with the offending op on any mismatch.""" if layout.tile_order != "row_major": raise ValueError("the kernel consumes matmul_block in1 order: plan the arena with tile_order='row_major'") # dit_kernel.cpp hard-codes k_offsets_header = 11 and k_words_per_block = 4 (mk::block_word); a WP-05 header change # must fail here on the CPU, not as garbage offsets on the device (WP-K1 review item 8) if len(A.OFFSETS_HEADER) != 11 or A.WORDS_PER_BLOCK != 4: raise AssertionError( f"arena offsets table header {len(A.OFFSETS_HEADER)} words / {A.WORDS_PER_BLOCK} words per block differ " "from dit_kernel.cpp (k_offsets_header 11, k_words_per_block 4): update the kernel constants" ) n = 0 for blk in range(layout.n_blocks): for j, name in enumerate(OP_NAMES): op = layout.op(blk, name) g = kernel_geometry(layout, blk, j) assert g.kt == op.k_tiles and g.nt == op.n_tiles and g.ct == op.n_tiles_per_core, (blk, name, g, op) assert g.coreset == op.core_set, (blk, name, g.coreset, op.core_set) assert g.off_bank == op.byte_offset, (blk, name, g.off_bank, op.byte_offset) assert g.per_bank == len(op.cores_in_bank(0)), (blk, name) assert g.pages_core * PAGE_TILES == op.tiles_per_core, (blk, name) assert g.bias_row == layout.bias_row(blk, name), (blk, name) assert op.k_block_tiles == PAGE_TILES, (blk, name, op.k_block_tiles) for bank in range(N_BANKS): slots = op.cores_in_bank(bank) for r, p in enumerate(slots): c = op.cores[p] assert group_of(c) == bank assert rank_of(g.coreset, c) == r, (blk, name, c, r) assert member_core(g.coreset, bank, r) == c, (blk, name, c, r) assert slot_of(g.coreset, c) == p, (blk, name, c, p) # stream order of the pages: relay = unit-major, direct = core-major (arena bank_stream_slots) units = op.bank_stream_slots(bank) if layout.reader_mode == "relay": assert units == [(kb, p) for kb in range(op.n_k_blocks) for p in slots] else: assert units == [(kb, p) for p in slots for kb in range(op.n_k_blocks)] n += 1 return n def random_arena_plan(layout: A.ArenaLayout, seed: int = 0, scale: float = 0.02) -> Dict[str, torch.Tensor]: """``{plan name: bf16 tensor}`` with distinct random values for every tensor the layout packs (no checkpoint). Weights are ``N(0, scale^2)`` (scale 0.02 keeps the four chained projections O(1)); biases / vectors ``N(0, 1)`` scaled by 0.1. Deterministic in ``seed``. """ entries = enumerate_plan(layout.version) g = torch.Generator().manual_seed(seed) plan: Dict[str, torch.Tensor] = {} for name in layout.required_plan_names(): shape = entries[name].shape t = torch.randn(*shape, generator=g, dtype=torch.float32) t = t * (scale if len(shape) == 2 and min(shape) > 1 else 0.1) plan[name] = t.to(torch.bfloat16) return plan # -------------------------------------------------------------------------------------------------------------------- # torch references # -------------------------------------------------------------------------------------------------------------------- def _bf16(t: torch.Tensor) -> torch.Tensor: return t.to(torch.bfloat16).to(torch.float32) def gelu_tanh(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.gelu(x, approximate="tanh") def op_weight_and_bias( layout: A.ArenaLayout, pack: A.ArenaPack, blk: int, name: str, arena_stream: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: """``(W [K,N] fp32, b [N] fp32)`` of op ``name`` of block ``blk`` from the packed arena (bit-exact unshuffle) and the bias table -- i.e. exactly the bytes the kernel streams (attention scale already folded).""" stream = pack.arena_stream if arena_stream is None else arena_stream w = layout.unpack(stream, blk, name).to(torch.float32) b = pack.tables["bias_rows"][layout.bias_row(blk, name), : w.shape[1]].to(torch.float32) return w, b def reference_single_op( layout: A.ArenaLayout, pack: A.ArenaPack, blk: int, name: str, x: torch.Tensor, arena_stream=None ) -> torch.Tensor: """bf16-rounded ``gelu?(x @ W + b)`` in fp32 maths for one op.""" w, b = op_weight_and_bias(layout, pack, blk, name, arena_stream) y = _bf16(x) @ w + b if name == "ff1": y = gelu_tanh(y) return _bf16(y) def reference_block( layout: A.ArenaLayout, pack: A.ArenaPack, blk: int, x: torch.Tensor, arena_stream=None ) -> Dict[str, torch.Tensor]: """The matmul-only block chain in fp32 with bf16 rounding between ops (what the device does): ``q = (x @ Wqkv + b)[:, :inner]`` (or ``x @ Wq + bq``), ``o = q @ Wo + bo``, ``f = gelu(o @ W1 + b1)``, ``g = f @ W2 + b2``. Returns ``{"q", "o", "f", "g"}``.""" q_full = reference_single_op(layout, pack, blk, "qkv", x, arena_stream) q = q_full[:, : layout.inner] o = reference_single_op(layout, pack, blk, "Wo", q, arena_stream) f = reference_single_op(layout, pack, blk, "ff1", o, arena_stream) g = reference_single_op(layout, pack, blk, "ff2", f, arena_stream) return {"q": q, "o": o, "f": f, "g": g} # -------------------------------------------------------------------------------------------------------------------- # tile-order helpers (host readback) # -------------------------------------------------------------------------------------------------------------------- def unpack_single_order(shard: torch.Tensor, n_tiles: int) -> torch.Tensor: """``Single`` order (48 gathered [64,32] columns): tile ``(r, n)`` sits at page ``2n + r`` of the [64, 2048] TILE-layout shard (pages row-major over its 2 x 64 tiles). Returns ``[64, 32*n_tiles]``.""" if shard.shape != (M_ROWS, 2048): raise ValueError(f"expected a [64, 2048] shard, got {tuple(shard.shape)}") out = torch.empty(M_ROWS, TILE * n_tiles, dtype=shard.dtype) tiles_per_row = shard.shape[1] // TILE for n in range(n_tiles): for r in range(RT_TILES): p = 2 * n + r R, C = divmod(p, tiles_per_row) out[r * TILE : (r + 1) * TILE, n * TILE : (n + 1) * TILE] = shard[ R * TILE : (R + 1) * TILE, C * TILE : (C + 1) * TILE ] return out def unpack_pairs_order(shard: torch.Tensor, n_tiles: int) -> torch.Tensor: """``Pairs`` order (gathered [64,64] blocks): tile ``(r, n)`` at page ``4*(n//2) + 2r + n%2``.""" if shard.shape[0] != M_ROWS: raise ValueError(f"expected 64 rows, got {tuple(shard.shape)}") out = torch.empty(M_ROWS, TILE * n_tiles, dtype=shard.dtype) tiles_per_row = shard.shape[1] // TILE for n in range(n_tiles): for r in range(RT_TILES): p = 4 * (n // 2) + 2 * r + (n % 2) R, C = divmod(p, tiles_per_row) out[r * TILE : (r + 1) * TILE, n * TILE : (n + 1) * TILE] = shard[ R * TILE : (R + 1) * TILE, C * TILE : (C + 1) * TILE ] return out def shard_order(core_map: CoreMap) -> List[int]: """Compute indices in the order ttnn assigns HEIGHT/WIDTH shards to the ``compute`` CoreRangeSet: the set is built from single-core ranges sorted by ``(y, x)`` (``CoreMap.to_core_range_sets``), so shard ``i`` belongs to the ``i``-th compute core in row-major logical order -- **not** compute index ``i``.""" return [core_map.compute_idx_of(core) for core in sorted(core_map.compute, key=lambda c: (c[1], c[0]))] def mm_out_to_matrix(mm_out: torch.Tensor, geom: KernelOpGeom, shard_cores: Sequence[int]) -> torch.Tensor: """``mm_out`` readback ``[n_compute*64, 64]`` -> the op's ``[64, N]`` output; ``shard_cores[i]`` is the compute index that owns shard ``i`` (:func:`shard_order`). ``ct == 2``: the [64,64] shard block is the slice; ``ct == 1``: DST tiles 0/1 (rows 0-31 / 32-63) were packed into shard tiles (0,0)/(0,1).""" n_compute = len(shard_cores) if mm_out.shape != (n_compute * M_ROWS, M_ROWS): raise ValueError(f"mm_out shape {tuple(mm_out.shape)} != {(n_compute * M_ROWS, M_ROWS)}") out = torch.zeros(M_ROWS, geom.nt * TILE, dtype=mm_out.dtype) for i, c in enumerate(shard_cores): if not participates(geom.coreset, c): continue p = slot_of(geom.coreset, c) block = mm_out[i * M_ROWS : (i + 1) * M_ROWS] cols = slice(p * geom.ct * TILE, (p + 1) * geom.ct * TILE) if geom.ct == 2: out[:, cols] = block else: out[:TILE, cols] = block[:TILE, :TILE] out[TILE:, cols] = block[:TILE, TILE : 2 * TILE] return out # -------------------------------------------------------------------------------------------------------------------- # K4: tail geometry mirror (dit_kernel.cpp mk4::tail_geom) and the CTA / RT contracts of the K4 kernel # -------------------------------------------------------------------------------------------------------------------- #: gr00t_ops::TailOp order == arena.TAIL_OPS names TAIL_OP_NAMES: Tuple[str, ...] = ("proj_out_2", "dec_W1", "dec_W2", "enc_W1", "enc_W2a", "enc_W3") #: K4 common runtime args after WP-05's io / step / offsets words (dit_kernel.cpp ``enum RT4``), then compute_xy[96] MK4_RT: Tuple[str, ...] = ( "blk_first", "n_blk", "bias_row_bytes", "adaln_row_bytes", "state_row_bytes", "b2_row_bytes", "future_row_bytes", "debug", ) #: bring-up knobs of the ``block`` variant (common RT arg ``debug``; dit_kernel.cpp ``enum Debug``) K4_DEBUG_OFF = 0 K4_DEBUG_SKIP_G_RESIDUAL = 1 # the host reads x + attn_out (the ff_out residual is popped, not added) K4_DEBUG_DUMP_H1 = 2 # every core's own P1 AdaLN tile replaces its x tile (P8 residual skipped): read_x_rows_all -> h K4_DEBUG_DUMP_H2 = 3 # every core's own P8 LN tile replaces its x tile: read_x_rows_all -> h2 K4_DEBUG_DUMP_F1 = 4 # every core's ff1 output tiles of its tile row (hidden columns 64c..64c+64) -> x tiles 0, 1 K4_DEBUG_DUMP_G = 5 # the even cores' ff2 output tiles (rows 0..31 / 32..63 of column c/2) -> x tiles 0, 1 K4_DEBUG_DUMP_A_H1 = 6 # every core's *received* h multicast (its tile row) -> x: read_x_rows_all K4_DEBUG_DUMP_A_H2 = 7 # every core's *received* h2 multicast (its tile row) -> x K4_DEBUG_F1_NO_GELU = 8 # ff1 without GELU; its pre-activation tiles of the own row -> x tiles 0, 1 K4_DEBUG_DUMP_Q = 9 # head cores: the 4 Q tiles of P2 -> x tiles 0..3 (assemble_dumped_q) K4_DEBUG_DUMP_O = 10 # head cores: the 4 attention output tiles of P4 -> x tiles 0..3 (assemble_dumped_q) #: io tensors whose slot index the kernel needs as a named CTA ``mk_io_`` (0xFFFF when the variant lacks it) MK_IO_CTAS: Tuple[str, ...] = ( "x_rep", "a_rep", "arena_stream", "arena_small", "adaln_rows", "bias_rows", "b2_step", "b3_pos_rows", "future_tokens", "self_mask", "text_mask", "image_mask", "all_mask", "x_t", "state_features", "action_pred", "velocity_out", ) #: every ``get_named_compile_time_arg_val`` name of dit_kernel.cpp / ops/*.hpp that neither :func:`descriptors.build_program_spec` #: (K1) nor :func:`descriptors.build_program_spec_k4` emits by itself for *both* specs: a missing name is a JIT failure #: (compile_time_args.h ``__builtin_unreachable``), so the launcher fills the union with 0 (``fill_missing_ctas``). #: ``tests/tt/test_mk_dit_block.py::test_kernel_cta_names_cpu`` greps the sources against this list. KERNEL_CTA_UNION: Tuple[str, ...] = ( ("mk_k1", "mk_ring_depth", "mk_relay_window", "mk_rt_base") + tuple(f"mk_io_{n}" for n in MK_IO_CTAS) + ("mk_io_kv_base",) + tuple(D.CB_IDS) + tuple(D.K4_CB_IDS) + tuple(f"k4_{n}" for n in D.K4_HUB_CB_IDS) + ("k4_tw_tiles", "k4_sdpa_slots", "k4_nt_a", "k4_velocity_out_present", "k4_zone_rounds", "k4_mcast_split") + ("cb_land_bytes", "k4_debug_knobs") # K4-only: hub landing ring (descriptors.K4_LAND_PAGES), bring-up knobs ) #: named CTAs :func:`descriptors.materialise` emits itself (device NoC coordinates) -- never in the union above (a 0 #: filler would collide with the real value: named_kernel_args.cpp `inserted || it->second == value`) MATERIALISE_CTAS: Tuple[str, ...] = ( "mcast_noc_x0", "mcast_noc_y0", "mcast_noc_x1", "mcast_noc_y1", "hub_noc_x", "hub_noc_y", "hub1_noc_x", "hub1_noc_y", ) #: K4b: the hub NCRISC's share of every hub multicast in quarters (``dit_kernel.cpp`` ``k_mcast_split_q``; 0 = the K4 #: single-NoC multicast from the hub BRISC, 2 = half / half over the two NoCs). ``GR00T_MK_MCAST_SPLIT`` overrides. K4_MCAST_SPLIT: int = int(os.environ.get("GR00T_MK_MCAST_SPLIT", "2")) #: K4b profiling knob (``dit_kernel.cpp`` ``k_zone_rounds``): 1 = a ``DeviceZoneScopedN`` per hub round (``mk_hR_*``) #: inside the hub BRISC's block / tail zones -- 11 extra zones per block, so whole-step profiles (16-32 blocks) must run #: with 0 (the profiler's L1 buffer holds 250 markers per RISC). Free without ``TT_METAL_DEVICE_PROFILER``. K4_ZONE_ROUNDS: int = int(os.environ.get("GR00T_MK_ZONE_ROUNDS", "1")) STEPS4_N_STEPS = A.N_STEPS #: The compute-core K4 binaries (BRISC + NCRISC + the 3 TRISC threads: LN + SDPA + matmuls + tail) exceed the default #: kernel-config ring buffer of the p150a (70,656 B; ``program.cpp`` ``finalize_kernel_bins`` = max over the kernel #: groups of one core). The ring buffer is ``l1_unreserved_base - KERNEL_CONFIG`` with ``l1_unreserved_base = L1 end - #: worker_l1_size`` (``metal_context.cpp``; ``dev_mem_map.h``: "configurable via worker_l1_size"), so :func:`open_k4_device` #: opens the device with ``worker_l1_size`` reduced by this many bytes -- the K4 L1 budget (``descriptors.K4Budget``) #: leaves that much headroom below ``L1_USABLE_BYTES``. K4_KERNEL_CONFIG_EXTRA_BYTES: int = 96 * 1024 K4_WORKER_L1_SIZE: int = D.L1_USABLE_BYTES - K4_KERNEL_CONFIG_EXTRA_BYTES def open_k4_device( trace_region_size: int = 64 << 20, l1_small_size: int = 32768, device_id: int = 0, worker_l1_size: int = K4_WORKER_L1_SIZE, ) -> Any: """``tt.device.open_gr00t_device`` with the smaller ``worker_l1_size`` the K4 kernel binaries need (module constant); same p150a asserts (11x10 compute grid, 8x1 DRAM grid) and program cache. Callers hold the device lock.""" import ttnn # lazy device = ttnn.open_device( device_id=device_id, l1_small_size=l1_small_size, trace_region_size=trace_region_size, worker_l1_size=int(worker_l1_size), ) try: device.enable_program_cache() grid = device.compute_with_storage_grid_size() dram = device.dram_grid_size() if (int(grid.x), int(grid.y)) != (11, 10): raise RuntimeError(f"compute_with_storage_grid_size() = {(grid.x, grid.y)}, the megakernel needs (11, 10)") if (int(dram.x), int(dram.y)) != (8, 1): raise RuntimeError(f"dram_grid_size() = {(dram.x, dram.y)}, the megakernel needs (8, 1)") except Exception: ttnn.close_device(device) raise return device @dataclass(frozen=True) class TailKernelGeom: """What ``mk4::tail_geom(op)`` computes on device for tail op ``op`` (K-split ``dec_W2``: per-lead K range).""" op: str off_bank: int kt: int # K tiles this core multiplies nt_w: int # N tiles of the slice (in1 column stride) tiles: int # tiles of the per-core slice coreset: str # all96 | first32 | even48 | leads8 bias_row: Optional[int] def tail_kernel_geometry(layout: A.ArenaLayout, op: str) -> TailKernelGeom: """Mirror of ``mk4::tail_geom``: byte offsets from the offsets-table header, structural sizes.""" if op not in TAIL_OP_NAMES: raise KeyError(f"tail op must be one of {TAIL_OP_NAMES}, got {op!r}") words = layout.offsets_table() hdr = {name: words[i] for i, name in enumerate(A.OFFSETS_HEADER)} nt = D.k4_n_tiles_a(layout) n_kt = layout.width // TILE hid = layout.hidden // TILE base = hdr["tail_bias_row_base"] table = { "proj_out_2": (hdr["off_proj_out_2"], n_kt, 1, n_kt, "first32", base + 0), "dec_W1": (hdr["off_dec_W1"], hid, 1, hid, "first32", base + 1), "dec_W2": (hdr["off_dec_W2"], hid // N_BANKS, nt, (hid // N_BANKS) * nt, "leads8", None), "enc_W1": (hdr["off_enc_W1"], nt, 1, nt, "even48", base + 3), "enc_W2a": (hdr["off_enc_W2a"], n_kt, 1, n_kt, "even48", None), "enc_W3": (hdr["off_enc_W3"], n_kt, 1, n_kt, "even48", None), } off, kt, nt_w, tiles, cs, bias = table[op] return TailKernelGeom(op, off, kt, nt_w, tiles, cs, bias) def participates4(coreset: str, c: int) -> bool: if coreset == "leads8": return c % 4 == 0 and c < 32 return participates(coreset, c) def rank4(coreset: str, c: int) -> int: return 0 if coreset == "leads8" else rank_of(coreset, c) def slot4(coreset: str, c: int) -> int: return c // 4 if coreset == "leads8" else slot_of(coreset, c) def check_tail_geometry(layout: A.ArenaLayout) -> int: """Assert ``mk4::tail_geom`` against the arena's small-stream ops (offsets, per-core tiles, K ranges, ranks).""" n = 0 for op_name in TAIL_OP_NAMES: op = layout.op(None, op_name) g = tail_kernel_geometry(layout, op_name) assert g.off_bank == op.byte_offset, (op_name, g.off_bank, op.byte_offset) assert g.tiles == op.tiles_per_core, (op_name, g.tiles, op.tiles_per_core) assert g.kt == op.k_tiles_per_core and g.nt_w == op.n_tiles_per_core, (op_name, g, op) assert g.coreset == op.core_set, (op_name, g.coreset, op.core_set) assert g.tiles <= D.K4_TW_TILES, (op_name, g.tiles) if g.bias_row is not None: assert g.bias_row == layout.bias_row(None, op_name), (op_name, g.bias_row) for bank in range(N_BANKS): slots = op.cores_in_bank(bank) for r, p in enumerate(slots): c = op.cores[p] assert group_of(c) == bank and participates4(g.coreset, c) assert rank4(g.coreset, c) == r and slot4(g.coreset, c) == p, (op_name, c, r, p) # the K4 kernel streams the tail in direct mode (core-major); a relay arena (K1 modes only) is K-block-major if layout.reader_mode == "direct": assert op.bank_stream_slots(bank) == [(kb, p) for p in slots for kb in range(op.n_k_blocks)] else: assert op.bank_stream_slots(bank) == [(kb, p) for kb in range(op.n_k_blocks) for p in slots] # the slice is contiguous in the bank stream: row-major (k, n) tiles, one K-block per lead for dec_W2 for p in slots: tiles = op.tiles_of(0, p) if op.split == "k": k0 = p * op.k_tiles_per_core assert tiles == [(k, nn) for k in range(k0, k0 + op.k_tiles_per_core) for nn in range(op.n_tiles)] else: assert tiles == [(k, p) for k in range(op.k_block_tiles)] n += 1 return n def fill_missing_ctas(spec: D.ProgramSpec, names: Sequence[str] = KERNEL_CTA_UNION, value: int = 0) -> int: """Append ``(name, value)`` for every union name the spec lacks; returns how many were added.""" have = {n for n, _ in spec.named_ctas} added = [(n, value) for n in names if n not in have] spec.named_ctas += added return len(added) def real_arena_plan(version: str, embodiment_id: Optional[int] = None, policy: Any = None) -> Dict[str, torch.Tensor]: """The host plan tensors (bf16) of the ``dit`` / ``dit_precompute`` / ``encoders`` categories -- the bytes ``TTWeights`` uploads for the Stage-1 oracle (``common.weights.build_device_plan`` under ``TTPolicy.plan_options``) -- as the ``ArenaLayout.pack`` input. ``policy`` defaults to the bf16 policy (the megakernel arena is bf16 first, plan §5.7 K4 bring-up order).""" from models.experimental.gr00t.common.configs import get_config from models.experimental.gr00t.common.weights import build_device_plan from models.experimental.gr00t.tt.policy import TTPolicy, plan_options cfg = get_config(version) emb = cfg.canonical_embodiment_id if embodiment_id is None else int(embodiment_id) pol = TTPolicy(dtype_policy="bf16") if policy is None else policy plan = build_device_plan( version, emb, "bf16", options=plan_options(pol, version), categories=("dit", "dit_precompute", "encoders") ) return {name: pt.tensor for name, pt in plan.items()} # -------------------------------------------------------------------------------------------------------------------- # the launcher # -------------------------------------------------------------------------------------------------------------------- class DiTMegakernel: """Launcher of the DiT megakernel (module docstring): the K1 matmul-only modes and the K4 step schedule. Args: cfg: ``GR00TConfig`` of the version (``arena.cfg`` is used for the checks). dlayout: ``tt.layout.DeviceLayout`` of the call layout or ``None`` (K1 / block unit tests). policy: ``tt.policy.TTPolicy`` -- ``reader_mode`` must agree with the arena's (``None`` = no check). arena: the planned :class:`arena.ArenaLayout` (``tile_order="row_major"``; K4 paths need ``reader_mode="direct"``). core_map: :class:`core_map.CoreMap` built from the runtime DRAM-worker query. device: the open mesh device. ring_depth / relay_window: DRAM pages in flight per ring / relay writes in flight per reader (K1 relay only). variant: ``"block"`` (K1 modes + K4 ``run_block``), ``"step"`` (``run_step``), ``"steps4"`` (``run_steps4``). cb_budget: the K1 :class:`descriptors.CBBudget` (K1 launches only). k4_budget: the :class:`descriptors.K4Budget` of the K4 launches. """ def __init__( self, cfg: Any, dlayout: Optional[Any], policy: Any, arena: A.ArenaLayout, core_map: CoreMap, device: Any, ring_depth: int = DEFAULT_RING_DEPTH, relay_window: int = DEFAULT_RELAY_WINDOW, variant: str = "block", cb_budget: Optional[D.CBBudget] = None, k4_budget: Optional[D.K4Budget] = None, k1_extras: bool = True, ) -> None: import ttnn # lazy self._ttnn = ttnn if getattr(cfg, "version", arena.version) != arena.version: raise ValueError(f"arena is {arena.version}, cfg is {cfg.version}") if policy is not None and getattr(policy, "reader_mode", arena.reader_mode) != arena.reader_mode: raise ValueError(f"policy.reader_mode={policy.reader_mode!r} != arena.reader_mode={arena.reader_mode!r}") if variant not in D.VARIANTS: raise ValueError(f"variant must be one of {D.VARIANTS}, got {variant!r}") if not 1 <= ring_depth <= 15 or not 1 <= relay_window <= 15: raise ValueError("ring_depth / relay_window must be in 1..15 (NoC transaction ids)") check_kernel_geometry(arena) check_tail_geometry(arena) budget = cb_budget if cb_budget is not None else D.CBBudget() k4 = k4_budget if k4_budget is not None else D.K4Budget() k_block_bytes_max = max(op.k_block_bytes for op in arena.stream_ops("stream")) page_bytes = PAGE_TILES * A.TILE_BYTES[arena.dtype_stream] # direct-mode ring progress condition (weight_stream.hpp PageRing::try_issue needs `ring_depth` free pages to # issue; the TRISC needs up to ct = 2 pages per K-block): ring_depth + 2 <= cb_w pages, else a full ring can # never refill -> the two RISCs stall (WP-K1 review item 4); checked for both budgets for name, cb_w_bytes in ( ("cb_budget", budget.cb_w_bytes(arena.dtype_stream, k_block_bytes_max)), ("k4_budget", k4.cb_w_bytes(arena.dtype_stream, k_block_bytes_max)), ): ring_pages = cb_w_bytes // page_bytes if ring_depth + 2 > ring_pages: raise ValueError( f"{name}: ring_depth {ring_depth} + 2 > cb_w pages {ring_pages}: PageRing could never issue into a " "ring that holds one un-consumed K-block (weight_stream.hpp)" ) if arena.reader_mode == "relay": relay_pages = budget.cb_relay_slots * k_block_bytes_max // page_bytes if relay_window + ring_depth > relay_pages + 1: raise ValueError( f"relay_window {relay_window} + ring_depth {ring_depth} > relay pages {relay_pages} + 1: the reader " "NCRISC/BRISC pair would deadlock (weight_stream.hpp Relay/PageRing)" ) self.cfg = cfg self.dlayout = dlayout self.policy = policy self.arena = arena self.core_map = core_map self.device = device self.variant = variant self.ring_depth = int(ring_depth) self.relay_window = int(relay_window) self.budget = budget self.k4 = k4 # K1 launches allocate `sync` / `w_ring` / `mm_out` right after `a_rep` (the WP-K1 order): allocating them # lazily under a first, short-lived x_rep leaves a hole above them that a later 768 KB x_rep cannot use and the # buffer then clashes with the static CB region (test_mk_layernorm::test_ff_branch_vs_s1). K4 launchers pass # False: they never build the K1 program and must not carry the 116 KB next to the K4 CBs. self.k1_extras = bool(k1_extras) self.tile_bytes = A.TILE_BYTES[arena.dtype_stream] self.page_bytes = page_bytes self.io: Dict[str, Any] = {} self.extra: Dict[str, Any] = {} self.k4_io: Dict[str, Any] = {} # persistent steps4 tensors (x_t, state_features, action_pred, velocity_out) self.pack: Optional[A.ArenaPack] = None self.last_spec: Optional[D.ProgramSpec] = None self.launches = 0 self.compile_seconds: Dict[Tuple[int, ...], float] = {} self.k4_program: Optional[Any] = None # the cached steps4 ProgramDescriptor (trace capture) self.k4_tensors: List[Any] = [] self._crs = core_map.to_core_range_sets() self._xy = self._noc_coordinates() self._shard_cores = self._device_shard_order() self.nt_a = D.k4_n_tiles_a(arena) # ------------------------------------------------------------------ device tensors def _device_shard_order(self) -> List[int]: """:func:`shard_order`, cross-checked against ``ttnn.corerange_to_cores`` of the compute CoreRangeSet.""" ttnn = self._ttnn order = shard_order(self.core_map) cores = ttnn.corerange_to_cores(self._crs["compute"], None, True) got = [self.core_map.compute_idx_of((int(c.x), int(c.y))) for c in cores] if got != order: raise RuntimeError("ttnn shard order of the compute CoreRangeSet differs from CoreMap row-major order") return order def _noc_coordinates(self) -> Dict[str, List[int]]: ttnn = self._ttnn def pack_xy(core: Tuple[int, int]) -> int: v = self.device.worker_core_from_logical_core(ttnn.CoreCoord(core[0], core[1])) if not (0 <= v.x < 256 and 0 <= v.y < 256): raise ValueError(f"NoC coordinate {v} does not fit a byte") return int(v.x) | (int(v.y) << 8) return { "readers": [pack_xy(c) for c in self.core_map.readers], "compute": [pack_xy(c) for c in self.core_map.compute], } def _sharded( self, t: torch.Tensor, role_set: str, shard: Tuple[int, int], dtype: Any, layout_kind: str = "HEIGHT", ) -> Any: ttnn = self._ttnn grid = self._crs[role_set] spec = ttnn.ShardSpec(grid, list(shard), ttnn.ShardOrientation.ROW_MAJOR) kind = ( ttnn.TensorMemoryLayout.HEIGHT_SHARDED if layout_kind == "HEIGHT" else ttnn.TensorMemoryLayout.WIDTH_SHARDED ) mem = ttnn.MemoryConfig(kind, ttnn.BufferType.L1, spec) return ttnn.from_torch(t, dtype=dtype, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=mem) def _dram(self, t: torch.Tensor, dtype: Any, layout: str = "TILE") -> Any: ttnn = self._ttnn lay = ttnn.TILE_LAYOUT if layout == "TILE" else ttnn.ROW_MAJOR_LAYOUT return ttnn.from_torch(t, dtype=dtype, layout=lay, device=self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG) def make_x_rep(self, x: torch.Tensor) -> Any: """``[64, K]`` (bf16) replicated HEIGHT_SHARDED on the 96 compute cores (the K1 ``cb_x``).""" if x.dim() != 2 or x.shape[0] != M_ROWS or x.shape[1] % TILE: raise ValueError(f"x must be [64, 32k], got {tuple(x.shape)}") rep = x.to(torch.bfloat16).repeat(self.core_map.n_compute, 1) return self._sharded(rep, "compute", (M_ROWS, x.shape[1]), self._ttnn.bfloat16) def row_of_core(self, c: int) -> int: """Tile row of the residual stream compute core ``c`` keeps (``dist_layernorm.hpp``: ``c // 48``).""" return int(c) // (self.arena.width // TILE) def make_x_rows(self, x: torch.Tensor) -> Any: """``x [64, 1536]`` (bf16) as the K4 ``x_rep``: HEIGHT_SHARDED ``[32, 1536]`` per compute core, shard = the core's tile row (cores 0..47 rows 0..31, cores 48..95 rows 32..63; shards in the device's ``(y, x)`` core order, :func:`shard_order`).""" if tuple(x.shape) != (M_ROWS, self.arena.width): raise ValueError(f"x must be [64, {self.arena.width}], got {tuple(x.shape)}") xb = x.to(torch.bfloat16) rows = torch.cat([xb[TILE * self.row_of_core(c) : TILE * (self.row_of_core(c) + 1)] for c in self._shard_cores]) return self._sharded(rows, "compute", (TILE, self.arena.width), self._ttnn.bfloat16) def read_x_rows(self, x_rows: Any, cores: Tuple[int, int] = (0, 48)) -> torch.Tensor: """The residual stream ``[64, 1536]`` (fp32) from the two compute cores ``cores`` = (a row-0 core, a row-1 core).""" t = self._ttnn.to_torch(x_rows) parts = [] for r, c in enumerate(cores): if self.row_of_core(c) != r: raise ValueError(f"core {c} keeps tile row {self.row_of_core(c)}, not {r}") i = self._shard_cores.index(int(c)) parts.append(t[i * TILE : (i + 1) * TILE].float()) return torch.cat(parts) def assemble_dumped_tiles(self, rows_all: torch.Tensor) -> torch.Tensor: """``read_x_rows_all`` output of a ``K4_DEBUG_DUMP_H*`` launch -> the assembled ``[64, 1536]`` LN output (core c's tile ``(c // 48, c % 48)``).""" n_kt = self.arena.width // TILE out = torch.empty(M_ROWS, self.arena.width) for c in range(self.core_map.n_compute): r, t = divmod(c, n_kt) out[TILE * r : TILE * (r + 1), TILE * t : TILE * (t + 1)] = rows_all[c][:, TILE * t : TILE * (t + 1)] return out def assemble_dumped_ff1(self, rows_all: torch.Tensor) -> torch.Tensor: """``K4_DEBUG_DUMP_F1``: core c's x tiles 0, 1 = ff1 output tiles (r, 2c), (r, 2c + 1) -> ``[64, 6144]``.""" n_kt = self.arena.width // TILE out = torch.empty(M_ROWS, self.arena.ff_dim) for c in range(self.core_map.n_compute): r = c // n_kt out[TILE * r : TILE * (r + 1), 64 * c : 64 * (c + 1)] = rows_all[c][:, : 2 * TILE] return out def assemble_dumped_g(self, rows_all: torch.Tensor) -> torch.Tensor: """``K4_DEBUG_DUMP_G``: even core c's x tiles 0, 1 = ff2 output tiles (0, c/2), (1, c/2) -> ``[64, 1536]``.""" out = torch.empty(M_ROWS, self.arena.width) for c in range(0, self.core_map.n_compute, 2): t = c // 2 out[:TILE, TILE * t : TILE * (t + 1)] = rows_all[c][:, :TILE] out[TILE:, TILE * t : TILE * (t + 1)] = rows_all[c][:, TILE : 2 * TILE] return out def assemble_dumped_q(self, rows_all: torch.Tensor) -> torch.Tensor: """``K4_DEBUG_DUMP_Q``: head core h's x tiles 0..3 = Q_h [64,64] row-major tiles -> ``[64, 2048]``.""" out = torch.empty(M_ROWS, self.arena.inner) for h in range(self.arena.n_heads): out[:TILE, 64 * h : 64 * h + TILE] = rows_all[h][:, :TILE] out[:TILE, 64 * h + TILE : 64 * (h + 1)] = rows_all[h][:, TILE : 2 * TILE] out[TILE:, 64 * h : 64 * h + TILE] = rows_all[h][:, 2 * TILE : 3 * TILE] out[TILE:, 64 * h + TILE : 64 * (h + 1)] = rows_all[h][:, 3 * TILE : 4 * TILE] return out def read_x_rows_all(self, x_rows: Any) -> torch.Tensor: """``[96, 32, 1536]`` fp32: every compute core's resident tile row (index = compute index) for consistency checks.""" t = self._ttnn.to_torch(x_rows).float() out = torch.empty(self.core_map.n_compute, TILE, self.arena.width) for i, c in enumerate(self._shard_cores): out[c] = t[i * TILE : (i + 1) * TILE] return out def upload(self, pack: A.ArenaPack) -> Dict[str, Any]: """Upload the packed arena (both streams), the tables, the constant self mask, the cross-mask / K-V placeholders (zeros; replaced by :meth:`bind_masks` / :meth:`bind_kv`) and the ``a_rep`` arena; returns the io dict (also kept in ``self.io``). The K1 working tensors are allocated lazily by the K1 launches.""" ttnn = self._ttnn L = self.arena cm = self.core_map if pack.arena_stream.shape[1] != L.tiles_per_bank["stream"]: raise ValueError("pack does not belong to this arena layout") if pack.arena_small.shape[1] != L.tiles_per_bank["small"]: raise ValueError("pack's small stream does not belong to this arena layout") self.pack = pack io: Dict[str, Any] = {} io["arena_stream"] = A.upload_arena(pack.arena_stream, L.dtype_stream, self.device) io["arena_small"] = A.upload_arena(pack.arena_small, L.dtype_small, self.device) io["adaln_rows"] = self._dram(pack.tables["adaln_rows"], ttnn.bfloat16, "ROW_MAJOR") io["bias_rows"] = self._dram(pack.tables["bias_rows"], ttnn.bfloat16, "ROW_MAJOR") io["b2_step"] = self._dram(pack.tables["b2_step"], ttnn.bfloat16, "ROW_MAJOR") b3 = pack.tables["b3_pos_rows"] if tuple(b3.shape) != (L.m_pad, L.width): raise ValueError(f"b3_pos_rows table {tuple(b3.shape)} != {(L.m_pad, L.width)}") io["b3_pos_rows"] = self._dram(b3, ttnn.bfloat16, "TILE") if pack.tables.get("future_tokens") is not None: io["future_tokens"] = self._dram(pack.tables["future_tokens"], ttnn.bfloat16, "ROW_MAJOR") io["self_mask"] = self._dram(pack.tables["self_mask"], ttnn.bfloat16, "TILE") masks = D.mask_presence(L) for ks in ("text", "image", "all"): if masks[ks]: lp = L.l_pad[ks] io[f"{ks}_mask"] = self._dram(torch.zeros(1, 1, L.m_pad, lp, dtype=torch.bfloat16), ttnn.bfloat16) for i in range(L.n_blocks): if L.is_cross(i): lp = L.l_pad[L.key_sets[i]] for kind in ("k", "v"): io[f"kv_{kind}_{i}"] = self._dram( torch.zeros(1, L.n_heads, lp, L.head_dim_pad, dtype=torch.bfloat16), ttnn.bfloat16 ) self.io = io self._alloc_a_rep() self._bound = {"kv": set(), "masks": set()} # K1 ring bookkeeping (the tensors themselves are allocated on the first K1 launch) k_block_bytes = max(op.k_block_bytes for op in L.stream_ops("stream")) cb_w_bytes = self.budget.cb_w_bytes(L.dtype_stream, k_block_bytes) if cb_w_bytes % self.page_bytes: raise ValueError(f"cb_w {cb_w_bytes} B is not a whole number of {self.page_bytes} B pages") self.ring_tiles = cb_w_bytes // self.tile_bytes self.ring_pages = self.ring_tiles // PAGE_TILES if self.k1_extras: self._ensure_k1_extras() return io def _alloc_a_rep(self) -> Any: """The in0 arena ``a_rep`` (``[64, 2048]`` bf16 HEIGHT_SHARDED on all 110 cores = ``cb_a``, the multicast target): allocated by :meth:`upload` and re-created lazily after :meth:`release_l1_residents`.""" ttnn = self._ttnn cm = self.core_map a_cols = self.budget.cb_a // (M_ROWS * 2) self.io["a_rep"] = self._sharded( torch.zeros(cm.n_cores * M_ROWS, a_cols, dtype=torch.bfloat16), "all", (M_ROWS, a_cols), ttnn.bfloat16 ) return self.io["a_rep"] def release_l1_residents(self) -> None: """Free every L1 tensor the launcher keeps between launches (``a_rep`` on all 110 cores, the K1 extras, the steps4 residents) so that Stage-1 ttnn ops can use the whole L1 in the same session; the next launch re-creates them (a cached steps4 program is dropped).""" self.release_steps4_io() if self.extra: self._free_k1_extras() if self.io.get("a_rep") is not None: self._ttnn.deallocate(self.io["a_rep"]) self.io["a_rep"] = None def _row_page_bytes(self, name: str, row_elems: int) -> int: """DRAM page of a ROW_MAJOR table as allocated (one row per page); the kernel's ``InterleavedAddrGen`` must use the allocator's page size, not an assumption.""" page = int(self.io[name].buffer_page_size()) if name in self.io else int(self.k4_io[name].buffer_page_size()) if page != row_elems * 2: raise ValueError(f"{name} page {page} B != one row ({row_elems * 2} B): the kernel reads rows") return page def _bias_page_bytes(self) -> int: return self._row_page_bytes("bias_rows", A.BIAS_TABLE_WIDTH) def _check_tile_page(self, t: Any, name: str, page: int) -> None: got = int(t.buffer_page_size()) if got != page: raise ValueError(f"{name}: DRAM page {got} B != {page} (the kernel reads whole tiles)") # ------------------------------------------------------------------ K4 inputs def set_variant(self, variant: str) -> None: """Switch the K4 program variant (``block`` / ``step`` / ``steps4``) of this launcher; the uploaded arena, the bound K/V and masks stay, the cached ``steps4`` program is dropped.""" if variant not in D.VARIANTS: raise ValueError(f"variant must be one of {D.VARIANTS}, got {variant!r}") if variant != self.variant: self.variant = variant self.k4_program = None self.k4_tensors = [] def bind_kv(self, kv: Mapping[int, Tuple[Any, Any]]) -> None: """Bind the hoisted ``(K_i, V_i)`` device tensors (``[1, 32, L_pad, 64]`` bf16 TILE DRAM, natural K) of every cross block of the arena (``tt.adapter`` / ``hoist_kv``); the zero placeholders are deallocated.""" L = self.arena for i in range(L.n_blocks): if not L.is_cross(i): continue if i not in kv: raise KeyError(f"kv lacks cross block {i}") k, v = kv[i] lp = L.l_pad[L.key_sets[i]] for nm, t in (("K", k), ("V", v)): shape = tuple(int(d) for d in t.shape) if shape != (1, L.n_heads, lp, L.head_dim_pad): raise ValueError( f"block {i} hoisted {nm} must be {(1, L.n_heads, lp, L.head_dim_pad)}, got {shape}" ) if t.dtype != self._ttnn.bfloat16: raise TypeError(f"block {i} hoisted {nm} must be bf16") self._check_tile_page(t, f"kv_{nm}_{i}", 2048) for kind, t in (("k", k), ("v", v)): old = self.io.get(f"kv_{kind}_{i}") if old is not None and old is not t and i not in self._bound["kv"]: self._ttnn.deallocate(old) self.io[f"kv_{kind}_{i}"] = t self._bound["kv"].add(i) extra = set(kv) - {i for i in range(L.n_blocks) if L.is_cross(i)} if extra: raise KeyError(f"kv has entries for non-cross blocks {sorted(extra)}") def bind_masks(self, masks: Mapping[str, Any]) -> None: """Bind the additive masks (``tt.dit_ops.dit_masks`` canonical form ``{"self","text","image","all"}`` of ``[1,1,64,L_pad]`` bf16 TILE DRAM tensors): the ones :func:`descriptors.mask_presence` says the kernel reads. The constant self mask of the arena pack stays unless a ``self`` entry is given.""" L = self.arena presence = D.mask_presence(L) for ks in ("text", "image", "all"): if not presence[ks]: continue t = masks.get(ks) if t is None: raise KeyError(f"mask {ks!r} is read by the kernel of this layout but missing / None") lp = L.l_pad[ks] shape = tuple(int(d) for d in t.shape) if shape != (1, 1, L.m_pad, lp): raise ValueError(f"mask {ks!r} must be {(1, 1, L.m_pad, lp)}, got {shape}") self._check_tile_page(t, f"{ks}_mask", 2048) old = self.io.get(f"{ks}_mask") if old is not None and old is not t and ks not in self._bound["masks"]: self._ttnn.deallocate(old) self.io[f"{ks}_mask"] = t self._bound["masks"].add(ks) if masks.get("self") is not None: t = masks["self"] shape = tuple(int(d) for d in t.shape) if shape != (1, 1, L.m_pad, L.m_pad): raise ValueError(f"self mask must be {(1, 1, L.m_pad, L.m_pad)}, got {shape}") self._check_tile_page(t, "self_mask", 2048) old = self.io.get("self_mask") if old is not None and old is not t and "self" not in self._bound["masks"]: self._ttnn.deallocate(old) self.io["self_mask"] = t self._bound["masks"].add("self") def allocate_steps4_io(self, taps: bool = True) -> Dict[str, Any]: """Persistent tensors of the ``steps4`` variant: ``x_t`` fp32 ``[64, A_pad]`` (the noise at the action rows), ``state_features`` bf16 ``[1, 1536]`` ROW_MAJOR, ``action_pred`` fp32 ``[64, A_pad]``, ``velocity_out`` bf16 ``[4*64, A_pad]`` (per-step decoder outputs; ``taps=False`` skips it). Allocated once; refreshed in place.""" ttnn = self._ttnn L = self.arena if self.k4_io: return self.k4_io k4_io: Dict[str, Any] = { "x_t": self._dram(torch.zeros(L.m_pad, L.a_pad, dtype=torch.float32), ttnn.float32, "TILE"), "state_features": self._dram(torch.zeros(1, L.width, dtype=torch.bfloat16), ttnn.bfloat16, "ROW_MAJOR"), "action_pred": self._dram(torch.zeros(L.m_pad, L.a_pad, dtype=torch.float32), ttnn.float32, "TILE"), } if taps: k4_io["velocity_out"] = self._dram( torch.zeros(STEPS4_N_STEPS * L.m_pad, L.a_pad, dtype=torch.bfloat16), ttnn.bfloat16, "TILE" ) self._check_tile_page(k4_io["x_t"], "x_t", 4096) self._check_tile_page(k4_io["action_pred"], "action_pred", 4096) if taps: self._check_tile_page(k4_io["velocity_out"], "velocity_out", 2048) self.k4_io = k4_io self._row_page_bytes("state_features", L.width) return k4_io def release_steps4_io(self) -> None: """Deallocate the persistent ``steps4`` tensors and drop the cached program (they are re-created lazily); frees the 96 KB zero ``x_rep`` rows so that Stage-1 ttnn ops (whose LayerNorm CBs need ~980 KB on their cores) fit next to ``a_rep`` under the K4 ``worker_l1_size``.""" for t in self.k4_io.values(): if t is not None: self._ttnn.deallocate(t) self.k4_io = {} self.k4_program = None self.k4_tensors = [] def set_steps4_inputs(self, noise: torch.Tensor, state_features: torch.Tensor) -> None: """Write the call inputs into the persistent tensors: ``noise`` ``[1, H, A_max]`` (placed at rows M-H..M, columns 0..A_max, fp32) and ``state_features`` ``[1536]`` / ``[1, 1536]`` (bf16).""" ttnn = self._ttnn L = self.arena if not self.k4_io: raise RuntimeError("call allocate_steps4_io() first") a_max = int(self.cfg.encoders.max_action_dim) if tuple(noise.shape) != (1, L.h, a_max): raise ValueError(f"noise must be [1, {L.h}, {a_max}], got {tuple(noise.shape)}") x_t = torch.zeros(L.m_pad, L.a_pad, dtype=torch.float32) r0, r1 = L.action_rows x_t[r0:r1, :a_max] = noise[0].to(torch.float32) sf = state_features.reshape(-1) if sf.numel() != L.width: raise ValueError(f"state_features must have {L.width} elements, got {sf.numel()}") ttnn.copy_host_to_device_tensor( ttnn.from_torch(x_t, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT), self.k4_io["x_t"] ) ttnn.copy_host_to_device_tensor( ttnn.from_torch( sf.reshape(1, L.width).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT ), self.k4_io["state_features"], ) # ------------------------------------------------------------------ K1 working tensors (lazy) def _ensure_k1_extras(self) -> None: if self.extra: return ttnn = self._ttnn cm = self.core_map tt_stream = {"bf16": ttnn.bfloat16, "bfp8_b": ttnn.bfloat8_b}[self.arena.dtype_stream] self.extra = { "sync": ttnn.from_torch( torch.zeros(cm.n_cores * TILE, TILE, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=ttnn.MemoryConfig( ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(self._crs["all"], [TILE, TILE], ttnn.ShardOrientation.ROW_MAJOR), ), ), "w_ring": self._sharded( torch.zeros(TILE, self.ring_tiles * TILE * cm.n_compute, dtype=torch.bfloat16), "compute", (TILE, self.ring_tiles * TILE), tt_stream, layout_kind="WIDTH", ), "mm_out": self._sharded( torch.zeros(cm.n_compute * M_ROWS, M_ROWS, dtype=torch.bfloat16), "compute", (M_ROWS, M_ROWS), ttnn.bfloat16, ), } def _free_k1_extras(self) -> None: for t in self.extra.values(): self._ttnn.deallocate(t) self.extra = {} def zero_sync(self) -> None: """Re-zero the relay credit words (a new run's counters start at 0 on both sides).""" ttnn = self._ttnn cm = self.core_map self._ensure_k1_extras() zeros = ttnn.from_torch( torch.zeros(cm.n_cores * TILE, TILE, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT ) ttnn.copy_host_to_device_tensor(zeros, self.extra["sync"]) # ------------------------------------------------------------------ K1 program def _spec(self, x_rep: Any, controls: Mapping[str, int]) -> D.ProgramSpec: if not self.io: raise RuntimeError("call upload(pack) first") if self.io.get("a_rep") is None: self._alloc_a_rep() self._ensure_k1_extras() io = dict(self.io) io["x_rep"] = x_rep spec = D.build_program_spec(self.core_map, self.arena, "block", io, self.budget, dry_run=False) # --- K1 adapter over WP-05's spec (documented in the module docstring) cbs = [] for cb in spec.cbs: if cb.name == "cb_w": cb = replace(cb, tensor_backed="w_ring") elif cb.name == "cb_mm0": cb = replace(cb, tensor_backed="mm_out") cbs.append(cb) spec.cbs = cbs io_order = spec.io_order() rt_base = len(spec.common_rt_args) spec.named_ctas += [ ("mk_k1", 1), ("mk_ring_depth", self.ring_depth), ("mk_relay_window", self.relay_window), ("mk_rt_base", rt_base), ("mk_io_x_rep", io_order.index("x_rep")), ("mk_io_a_rep", io_order.index("a_rep")), ("mk_io_arena_stream", io_order.index("arena_stream")), ("mk_io_bias_rows", io_order.index("bias_rows")), ] fill_missing_ctas(spec) values = { "mode": int(controls["mode"]), "single_op": int(controls.get("single_op", 0)), "blk_first": int(controls["blk_first"]), "n_blk": int(controls["n_blk"]), "n_iters": int(controls["n_iters"]), "x_tiles": int(controls["x_tiles"]), "addr_sync": int(self.extra["sync"].buffer_address()), "addr_w_ring": int(self.extra["w_ring"].buffer_address()), "addr_mm_out": int(self.extra["mm_out"].buffer_address()), "bias_row_bytes": self._bias_page_bytes(), "pop_mm": int(controls.get("pop_mm", 0)), } args = [values[k] for k in MK_RT] + list(self._xy["readers"]) + list(self._xy["compute"]) names = list(MK_RT) + [f"reader_xy[{i}]" for i in range(len(self._xy["readers"]))] names += [f"compute_xy[{i}]" for i in range(len(self._xy["compute"]))] spec.common_rt_args = list(spec.common_rt_args) + args spec.common_rt_names = list(spec.common_rt_names) + names if len(spec.common_rt_args) > D.MAX_COMMON_RT_ARGS: raise ValueError(f"{len(spec.common_rt_args)} common runtime args exceed {D.MAX_COMMON_RT_ARGS}") for w in args: if not 0 <= w < 2**32: raise ValueError(f"runtime arg {w} does not fit uint32") self.last_spec = spec return spec def _launch(self, x_rep: Any, controls: Mapping[str, int]) -> float: """Build + run one K1 generic_op; returns the wall time (s) of launch + synchronize.""" ttnn = self._ttnn spec = self._spec(x_rep, controls) io = dict(self.io) io["x_rep"] = x_rep io.update(self.extra) program = D.materialise(spec, io) tensors = [io[name] for name in spec.io_order() if io.get(name) is not None] tensors += [self.extra["sync"], self.extra["w_ring"], self.extra["mm_out"]] t0 = time.perf_counter() ttnn.generic_op(tensors, program) ttnn.synchronize_device(self.device) dt = time.perf_counter() - t0 key = (int(controls["mode"]), int(controls.get("single_op", 0))) if key not in self.compile_seconds: self.compile_seconds[key] = dt self.launches += 1 return dt def stats(self) -> D.DescriptorStats: if self.last_spec is None: raise RuntimeError("no program built yet") return self.last_spec.stats() def kernel_groups(self) -> int: if self.last_spec is None: raise RuntimeError("no program built yet") return len(self.last_spec.kernel_groups()) # ------------------------------------------------------------------ K4 program def _k4_variant_io(self, x_rep: Optional[Any]) -> Dict[str, Any]: # the step variant's y_out placeholder is not an io tensor of the other variants io = {k: v for k, v in self.io.items() if k != "y_out" or self.variant == "step"} L = self.arena if self.variant == "steps4": if not self.k4_io: raise RuntimeError("steps4: call allocate_steps4_io() first") io.update(self.k4_io) if x_rep is None: # cb_x is loaded from the multicast sa_embs (LoadX); zero tile rows are the tensor backing x_rep = self.make_x_rows(torch.zeros(M_ROWS, L.width, dtype=torch.bfloat16)) self.k4_io["x_rep"] = x_rep io["x_rep"] = x_rep else: if x_rep is None: raise ValueError("block / step launches need the x_rep tensor") io["x_rep"] = x_rep if self.variant == "step": # y_out is required by io_spec("step") (plan §5.1); the K4 hub keeps y in its a_rep shard and this slot # is only an address the kernel never reads -> a 4 KB placeholder if "y_out" not in io: io["y_out"] = self._dram(torch.zeros(TILE, TILE, dtype=torch.bfloat16), self._ttnn.bfloat16) self.io["y_out"] = io["y_out"] return io def _spec_k4(self, io: Mapping[str, Any], controls: Mapping[str, int]) -> D.ProgramSpec: if not self.io: raise RuntimeError("call upload(pack) first") L = self.arena missing_kv = [i for i in range(L.n_blocks) if L.is_cross(i) and i not in self._bound["kv"]] if missing_kv and not controls.get("allow_placeholder_kv", 0): raise RuntimeError(f"hoisted K/V not bound for cross blocks {missing_kv} (bind_kv)") presence = D.mask_presence(L) missing_masks = [ks for ks in ("text", "all") if presence[ks] and ks not in self._bound["masks"]] if missing_masks and not controls.get("allow_placeholder_masks", 0): raise RuntimeError(f"per-call masks not bound: {missing_masks} (bind_masks)") spec = D.build_program_spec_k4( self.core_map, L, self.variant, io, self.k4, dry_run=False, step_index=int(controls["step"]) ) io_order = spec.io_order() def slot(name: str) -> int: return io_order.index(name) if name in io_order else D.NA cross = [i for i in range(L.n_blocks) if L.is_cross(i)] kv_base = slot(f"kv_k_{cross[0]}") if cross else D.NA rt_base = len(spec.common_rt_args) spec.named_ctas += [ ("mk_k1", 0), ("mk_ring_depth", self.ring_depth), ("mk_relay_window", self.relay_window), ("mk_rt_base", rt_base), ("mk_io_kv_base", kv_base), ("k4_zone_rounds", K4_ZONE_ROUNDS), ("k4_mcast_split", K4_MCAST_SPLIT), ("k4_debug_knobs", 1 if int(controls.get("debug", K4_DEBUG_OFF)) != K4_DEBUG_OFF else 0), ] spec.named_ctas += [(f"mk_io_{name}", slot(name)) for name in MK_IO_CTAS] fill_missing_ctas(spec) row = lambda name, n: ( self._row_page_bytes(name, n) if (name in io and io[name] is not None) else 0 ) # noqa: E731 values = { "blk_first": int(controls["blk_first"]), "n_blk": int(controls["n_blk"]), "bias_row_bytes": self._bias_page_bytes(), "adaln_row_bytes": self._row_page_bytes("adaln_rows", L.width), "state_row_bytes": row("state_features", L.width), "b2_row_bytes": self._row_page_bytes("b2_step", L.width), "future_row_bytes": row("future_tokens", L.width), "debug": int(controls.get("debug", K4_DEBUG_OFF)), } args = [values[k] for k in MK4_RT] + list(self._xy["compute"]) names = list(MK4_RT) + [f"compute_xy[{i}]" for i in range(len(self._xy["compute"]))] spec.common_rt_args = list(spec.common_rt_args) + args spec.common_rt_names = list(spec.common_rt_names) + names if len(spec.common_rt_args) > D.MAX_COMMON_RT_ARGS: raise ValueError(f"{len(spec.common_rt_args)} common runtime args exceed {D.MAX_COMMON_RT_ARGS}") for w in args: if not 0 <= w < 2**32: raise ValueError(f"runtime arg {w} does not fit uint32") self.last_spec = spec return spec def build_k4(self, x_rep: Optional[Any], controls: Mapping[str, int]) -> Tuple[Any, List[Any]]: """Build the K4 ``ttnn.ProgramDescriptor`` + its io tensor list (no launch).""" if self.arena.reader_mode != "direct": raise ValueError("K4 launches need a direct-mode arena (ArenaLayout.plan(reader_mode='direct'))") if self.extra: self._free_k1_extras() # the K1 w_ring / mm_out would not fit next to the K4 CBs if self.io.get("a_rep") is None: self._alloc_a_rep() io = self._k4_variant_io(x_rep) spec = self._spec_k4(io, controls) program = D.materialise(spec, io, unpack_to_dest_fp32_cbs=D.k4_unpack_fp32_cbs()) tensors = [io[name] for name in spec.io_order() if io.get(name) is not None] return program, tensors def launch_k4(self, x_rep: Optional[Any], controls: Mapping[str, int], sync: bool = True) -> float: """Build + enqueue one K4 generic_op; returns the wall seconds (launch + synchronize when ``sync``).""" ttnn = self._ttnn program, tensors = self.build_k4(x_rep, controls) t0 = time.perf_counter() ttnn.generic_op(tensors, program) if sync: ttnn.synchronize_device(self.device) dt = time.perf_counter() - t0 key = (10 + D.VARIANTS.index(self.variant), int(controls["step"])) if key not in self.compile_seconds: self.compile_seconds[key] = dt self.launches += 1 return dt # ------------------------------------------------------------------ readbacks def read_x_rep(self, x_rep: Any, compute_idx: int = 0) -> torch.Tensor: """The resident ``x`` ``[64, 1536]`` (bf16 -> fp32) of compute core ``compute_idx`` after a launch.""" t = self._ttnn.to_torch(x_rep) i = self._shard_cores.index(compute_idx) return t[i * M_ROWS : (i + 1) * M_ROWS].float() def read_hub_a(self, n_tiles: int) -> torch.Tensor: """``[64, 32 * n_tiles]`` from the hub's ``a_rep`` shard (the last full gather, Single order).""" a = self._ttnn.to_torch(self.io["a_rep"]) lid = self.core_map.linear_id(self.core_map.hub) shard = a[lid * M_ROWS : (lid + 1) * M_ROWS] return unpack_single_order(shard, n_tiles).float() def read_block_output(self, core: Optional[Tuple[int, int]] = None) -> torch.Tensor: """``g`` ``[64, 1536]`` from the ``a_rep`` shard of ``core`` (default: the hub) -- the K1 chain output.""" a = self._ttnn.to_torch(self.io["a_rep"]) lid = self.core_map.linear_id(core if core is not None else self.core_map.hub) shard = a[lid * M_ROWS : (lid + 1) * M_ROWS] return unpack_single_order(shard, self.arena.width // TILE) # ------------------------------------------------------------------ K1 runs (unchanged API) def run_single_op(self, x: torch.Tensor, blk: int, op: str, n_iters: int = 1) -> torch.Tensor: """One projection ``op`` of block ``blk`` with ``in0 = x`` ``[64, K]``: returns the bf16 ``[64, N]`` result read from ``mm_out`` (``n_iters > 1`` repeats the op in-kernel for timing; the result is the last iteration's).""" if op not in OP_NAMES: raise KeyError(f"op must be one of {OP_NAMES}") j = OP_NAMES.index(op) g = kernel_geometry(self.arena, blk, j) if x.shape != (M_ROWS, g.kt * TILE): raise ValueError(f"x must be [64, {g.kt * TILE}] for {op}, got {tuple(x.shape)}") x_rep = self.make_x_rep(x) self._launch( x_rep, { "mode": MODES["single"], "single_op": j, "blk_first": blk, "n_blk": 1, "n_iters": n_iters, "x_tiles": RT_TILES * g.kt, "pop_mm": 1 if n_iters > 1 else 0, }, ) out = self._ttnn.to_torch(self.extra["mm_out"]) self._ttnn.deallocate(x_rep) return mm_out_to_matrix(out, g, self._shard_cores) def run_block( self, x_rep: Any, blk: int, matmul_only: bool = True, n_blk: int = 1, n_iters: int = 1, step: int = 0, debug: int = K4_DEBUG_OFF, ) -> torch.Tensor: """``matmul_only=True`` (K1): the matmul-only chain of blocks ``[blk, blk + n_blk)`` (``n_iters`` times) from the device tensor ``x_rep``; returns the last block's ``g`` ``[64, 1536]`` from the hub's ``a_rep`` shard. ``matmul_only=False`` (K4, plan §5.4 P1-P12): the full blocks ``[blk, blk + n_blk)`` of denoising step ``step`` from ``x_rep`` (a host ``[64, 1536]`` tensor or the device tile rows of :meth:`make_x_rows`) with the bound K/V and masks; returns the block output (``x`` after the final residual) as fp32 ``[64, 1536]`` read from compute cores 0 (rows 0..31) and 48 (rows 32..63). """ if not 0 <= blk < self.arena.n_blocks or blk + n_blk > self.arena.n_blocks: raise IndexError(f"blocks [{blk}, {blk + n_blk}) out of range for {self.arena.n_blocks}") if matmul_only: self._launch( x_rep, { "mode": MODES["chain"], "blk_first": blk, "n_blk": n_blk, "n_iters": n_iters, "x_tiles": RT_TILES * (self.arena.width // TILE), }, ) return self.read_block_output() if self.variant != "block": raise ValueError(f"run_block(matmul_only=False) needs variant='block', this launcher is {self.variant!r}") if n_iters != 1: raise ValueError("the K4 block variant runs once per launch") if not 0 <= step < STEPS4_N_STEPS: raise ValueError(f"step {step} out of range") own = isinstance(x_rep, torch.Tensor) dev = self.make_x_rows(x_rep) if own else x_rep try: self.launch_k4(dev, {"blk_first": blk, "n_blk": n_blk, "step": step, "debug": int(debug)}) return self.read_x_rows(dev) finally: if own: self._ttnn.deallocate(dev) def run_step(self, sa_embs: Any, step: int) -> Tuple[torch.Tensor, torch.Tensor]: """The ``step`` variant (plan §5.1): all blocks of denoising step ``step`` + output LN + ``proj_out_2`` from ``sa_embs`` (host ``[64, 1536]`` or the device tile rows of :meth:`make_x_rows`). Returns ``(dit_out [64, 1024], dit_block_last_out [64, 1536])`` fp32 (``dit_out`` from the hub's gathered ``y``, the residual stream from compute cores 0 / 48).""" if self.variant != "step": raise ValueError(f"run_step needs variant='step', this launcher is {self.variant!r}") if not 0 <= step < STEPS4_N_STEPS: raise ValueError(f"step {step} out of range") own = isinstance(sa_embs, torch.Tensor) dev = self.make_x_rows(sa_embs) if own else sa_embs try: self.launch_k4(dev, {"blk_first": 0, "n_blk": self.arena.n_blocks, "step": step}) y = self.read_hub_a(self.arena.out_dim // TILE) x_last = self.read_x_rows(dev) return y, x_last finally: if own: self._ttnn.deallocate(dev) def steps4_program(self) -> Tuple[Any, List[Any]]: """The (cached) ``steps4`` program over the persistent io tensors (the same object is what a trace captures).""" if self.variant != "steps4": raise ValueError(f"steps4 launches need variant='steps4', this launcher is {self.variant!r}") if self.k4_program is None: self.k4_program, self.k4_tensors = self.build_k4( self.k4_io.get("x_rep"), {"blk_first": 0, "n_blk": self.arena.n_blocks, "step": 0} ) return self.k4_program, self.k4_tensors def run_steps4( self, noise: torch.Tensor, state_features: torch.Tensor, sync: bool = True ) -> Dict[str, torch.Tensor]: """The whole denoise (4 steps) from ``noise`` ``[1, H, A_max]`` and ``state_features`` ``[1536]``: returns ``{"action_pred": fp32 [1, 1, 64, A_pad] (the x_t buffer layout of tt.encoders), "velocity": fp32 [4, 64, A_pad] (decoder output per step, bf16 values; absent without velocity taps)}``.""" ttnn = self._ttnn self.allocate_steps4_io() self.set_steps4_inputs(noise, state_features) program, tensors = self.steps4_program() t0 = time.perf_counter() ttnn.generic_op(tensors, program) ttnn.synchronize_device(self.device) self.last_wall_s = time.perf_counter() - t0 self.launches += 1 return self.read_steps4_outputs() def read_steps4_outputs(self) -> Dict[str, torch.Tensor]: ttnn = self._ttnn L = self.arena out = {"action_pred": ttnn.to_torch(self.k4_io["action_pred"]).float().reshape(1, 1, L.m_pad, L.a_pad)} if "velocity_out" in self.k4_io: v = ttnn.to_torch(self.k4_io["velocity_out"]).float() out["velocity"] = v.reshape(STEPS4_N_STEPS, L.m_pad, L.a_pad) return out def capture_steps4(self) -> Any: """Capture the ``steps4`` generic_op into a trace (the inputs are the persistent tensors); returns the trace id object (``ttnn.MeshTraceId``) that :meth:`replay` / :meth:`release_trace` take.""" ttnn = self._ttnn program, tensors = self.steps4_program() ttnn.synchronize_device(self.device) tid = ttnn.begin_trace_capture(self.device, cq_id=0) try: ttnn.generic_op(tensors, program) finally: ttnn.end_trace_capture(self.device, tid, cq_id=0) ttnn.synchronize_device(self.device) return tid def replay(self, tid: Any) -> float: """Execute a captured trace and wait; returns the wall seconds.""" ttnn = self._ttnn t0 = time.perf_counter() ttnn.execute_trace(self.device, tid, cq_id=0, blocking=False) ttnn.synchronize_device(self.device) return time.perf_counter() - t0 def release_trace(self, tid: Any) -> None: self._ttnn.release_trace(self.device, tid) def run_stream(self, x_rep: Any, blk: int, n_blk: int, n_iters: int) -> float: """Stream-only launch (no gathers / mcasts); returns launch + sync wall seconds.""" return self._launch( x_rep, { "mode": MODES["stream"], "blk_first": blk, "n_blk": n_blk, "n_iters": n_iters, "x_tiles": RT_TILES * (self.arena.width // TILE), "pop_mm": 1, }, ) def run_chain_timed(self, x_rep: Any, blk: int, n_blk: int, n_iters: int) -> float: return self._launch( x_rep, { "mode": MODES["chain"], "blk_first": blk, "n_blk": n_blk, "n_iters": n_iters, "x_tiles": RT_TILES * (self.arena.width // TILE), }, ) def stream_bytes(self, blk: int, n_blk: int) -> int: """Weight bytes streamed from DRAM (all 8 banks) by one iteration over blocks ``[blk, blk + n_blk)``.""" return sum(op.bytes_per_bank for b in range(blk, blk + n_blk) for op in self.arena.ops_of_block(b)) * N_BANKS def arena_bytes_per_step(self) -> int: """DRAM bytes one denoising step streams: every block op of every block + the tail ops (both arenas, all banks).""" return (self.arena.bytes_per_bank("stream") + self.arena.bytes_per_bank("small")) * N_BANKS def deallocate(self) -> None: ttnn = self._ttnn seen: set = set() for t in list(self.io.values()) + list(self.extra.values()) + list(self.k4_io.values()): if t is not None and id(t) not in seen: seen.add(id(t)) ttnn.deallocate(t) self.io = {} self.extra = {} self.k4_io = {} self.k4_program = None self.k4_tensors = [] def measure_rate(launch: Any, n_small: int, n_large: int, bytes_per_iter: int, repeats: int = 3) -> Dict[str, float]: """The mb3 timing method (device-stack-smoke.md §4.1): ``launch(n)`` runs ``n`` in-kernel iterations; the per- iteration time is ``(t(n_large) - t(n_small)) / (n_large - n_small)`` (min over ``repeats``), so launch and dispatch overhead cancel. Returns seconds and GB/s.""" if n_large <= n_small or n_small < 1: raise ValueError("need 1 <= n_small < n_large") launch(n_small) launch(n_large) t_small = min(launch(n_small) for _ in range(repeats)) t_large = min(launch(n_large) for _ in range(repeats)) per_iter = (t_large - t_small) / (n_large - n_small) gbps = bytes_per_iter / per_iter / 1e9 if per_iter > 0 else math.inf return { "t_small_s": t_small, "t_large_s": t_large, "per_iter_s": per_iter, "bytes_per_iter": float(bytes_per_iter), "gbps": gbps, } __all__ = [ "CORE_SETS", "KERNEL_CTA_UNION", "MK4_RT", "MK_IO_CTAS", "TAIL_OP_NAMES", "TailKernelGeom", "check_tail_geometry", "fill_missing_ctas", "participates4", "rank4", "real_arena_plan", "slot4", "tail_kernel_geometry", "DiTMegakernel", "K4_KERNEL_CONFIG_EXTRA_BYTES", "K4_WORKER_L1_SIZE", "open_k4_device", "IN0_ORDERS", "KernelOpGeom", "MK_RT", "MODES", "OP_NAMES", "PAGE_TILES", "check_kernel_geometry", "kernel_geometry", "measure_rate", "member_core", "mm_out_to_matrix", "op_weight_and_bias", "participates", "random_arena_plan", "rank12", "rank_of", "reference_block", "shard_order", "reference_single_op", "slot_of", "unpack_pairs_order", "unpack_single_order", ]