code/models/experimental/gr00t/tt/megakernel/descriptors.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""``UnifiedKernelDescriptor`` / CB / semaphore builders of the DiT megakernel (IMPLEMENTATION_PLAN.md §5.1, §5.5, §5.6).

One ``ttnn.ProgramDescriptor`` over the whole ``(0,0)-(10,9)`` rectangle from one source ``dit_kernel.cpp``:

* **3 role flags as the only per-core compile-time args** (``is_reader`` / ``is_hub`` / ``is_compute``) -> at most 4
  kernel groups (readers, hub, compute, spare) and one binary per variant (JIT 0.4-1.8 s, mb3 F6);
* **structural constants as named CTAs shared by all cores** (``n_blocks``, ``m_logical``, ``l_*_pad``, ``a_pad``,
  ``h``, ``euler_fp32``, ``n_steps``, ``reader_mode``, ``{text,image,all}_mask_present``, CB ids/sizes, semaphore
  ids, ...);
* **per-core runtime args** ``(compute_idx, bank_id, vc, head_idx, group)`` (``0xFFFF`` = n/a) on every core and RISC;
* **common runtime args**: the io buffer addresses, the step index and the arena offsets table (<= 341 words);
* CBs per plan §5.5 (ids <= 15 < 64) and the semaphores of the §5.6 protocol.

Variants (plan §5.1): ``block`` (the first ``n_blocks`` DiT blocks -- unit tests), ``step`` (blocks + output head; the
decoder / Euler / action encoder run as ttnn ops between 4 launches -- K4 bring-up), ``steps4`` (the whole denoise,
``PersistentLoop<false>(4)``).

:func:`build_program` with ``dry_run=True`` builds the ttnn-free :class:`ProgramSpec` and returns
:class:`DescriptorStats`; without it, the spec is materialised into ``ttnn`` objects (lazy import) and the
``ttnn.ProgramDescriptor`` is returned.  ``ttnn`` is never imported at module level.

**K4 extension (WP-K4, plan §5.4 P1-P12 + "Per step", §5.5 with the K2/K3 CB requests):** :func:`build_program_spec_k4`
builds the spec of the *full* megakernel -- the K1 CBs plus the statistics CBs of the distributed LayerNorm / AdaLN
(``dist_layernorm.hpp``: ``cb_stat`` / ``cb_stat32`` / ``cb_const`` / the row-broadcast AdaLN tiles, and ``cb_h`` = the
own normalised tile the NCRISC gathers), the ``HeadSDPA`` CBs of ``head_sdpa.hpp`` (K3 §3.3), the tail-weight CB
``cb_tw`` (bf16 forever, plan D4) and the hub's Euler / velocity CBs (``euler_head.hpp``).  ``cb_x`` holds only the
core's tile row of the residual stream (``x_rep`` = HEIGHT_SHARDED ``[32, 1536]`` per compute core).  It is a
separate path on purpose: the default :func:`build_program_spec` keeps the WP-05 shape (ids <= 15, compute L1 <= 1 MB)
that ``tests/tt/test_mk_cpu.py`` pins; the K4 spec uses ids up to 31 (``NUM_CIRCULAR_BUFFERS``) and a compute-core L1
of ~1.3 MB (:class:`K4Budget`, checked against the 1,427 KiB usable minus a margin).  :func:`materialise` takes the
``unpack_to_dest_fp32_cbs`` list the fp32 ``copy_tile`` CBs need (K2 §7.2).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple

from models.experimental.gr00t.tt.megakernel.arena import (
    ArenaLayout,
    K_BLOCK_TILES,
    KEY_SETS,
    N_STEPS,
    TILE,
    TILE_BYTES,
    TILE_ORDERS,
    READER_MODES,
)
from models.experimental.gr00t.tt.megakernel.core_map import Core, CoreMap, Role

VARIANTS: Tuple[str, ...] = ("block", "step", "steps4")
ROLE_FLAGS: Tuple[Tuple[str, Role], ...] = (
    ("is_reader", Role.READER),
    ("is_hub", Role.HUB),
    ("is_compute", Role.COMPUTE),
)
KERNEL_SOURCE: str = "models/experimental/gr00t/tt/megakernel/kernels/dit_kernel.cpp"
MAX_COMMON_RT_ARGS: int = 341
MAX_CBS: int = 64
MAX_KERNEL_GROUPS: int = 4
NUM_CIRCULAR_BUFFERS: int = 32  # tt-metal CB ids 0..31 (the K4 spec uses all of them on the compute cores)
#: "compute-core L1 budget <= 1,000 KB" (WP-05 acceptance; plan §5.5: 850-980 KB of 1,427 KiB usable)
L1_COMPUTE_BUDGET_BYTES: int = 1_000_000
L1_USABLE_BYTES: int = 1427 * 1024
NA: int = 0xFFFF
BF16, FP32 = "bf16", "fp32"

#: CB ids (plan §5.5).  Reader-only / hub-only CBs reuse ids 14/15 on disjoint core ranges.
CB_IDS: Mapping[str, int] = {
    "cb_x": 0,  # residual [64,1536] bf16, tensor-backed (x_rep, HEIGHT_SHARDED replicated on the compute cores)
    "cb_a": 1,  # in0 arena A (same L1 address on all 110 cores; mcast target), tensor-backed (a_rep)
    "cb_w": 2,  # weight ring (K-blocks relayed by the readers)
    "cb_mm0": 3,  # matmul output ping
    "cb_mm1": 4,  # matmul output pong
    "cb_o": 5,  # attention output / gather source
    "cb_bias0": 6,
    "cb_bias1": 7,
    "cb_adaln": 8,  # 2 AdaLN rows (raw) + tilised staging
    "cb_mask": 9,  # self mask 8 KB + cross chunk ring 2 x 8 KB
    "cb_k": 10,
    "cb_v": 11,
    "cb_kv_ring": 12,  # cross K/V 64-key chunk ring 2 x 16 KB
    "cb_scores": 13,  # fp32 scores / probs scratch 2 x 16 KB
    "cb_relay": 14,  # readers only: landed K-blocks (NumBuffers 3 x 2)
    "cb_act": 14,  # hub only: action buffer [64, A_pad] fp32
    "cb_sa": 15,  # hub only: sa_embs staging [64, 1536] bf16
}
SEM_IDS: Mapping[str, int] = {
    "sem_w": 0,  # K-block landed in cb_w (reader BRISC -> compute NCRISC)
    "sem_credit": 1,  # cb_w slot credit (compute -> reader)
    "sem_kv": 2,  # K_h/V_h arrived (== 2)
    "sem_mcast_sender": 3,
    "sem_mcast_receiver": 4,
    "sem_gather_noc0": 5,
    "sem_gather_noc1": 6,
    "sem_ring_credit": 7,  # ff2 K-chunk ring
    "sem_sync": 8,  # sync_riscs_enter/exit around CB reconfig
    "sem_hub_done": 9,  # K4b: hub NCRISC -> hub BRISC, rounds whose NoC-1 multicast part is flushed (allgather.hpp)
    "sem_ready_odd": 10,  # K4b: hub receivers-ready count of the odd ff2 ring chunks (lookahead credits, dit_kernel P11)
    "sem_mcast_receiver1": 11,  # K4b: the receivers' round counter of hub 1's multicasts (hub 0: sem_mcast_receiver)
    "sem_land1": 12,  # K4b: hub 1 publishes its landing-ring base here (hub 0: sem_sync)
}
TILE_BYTES_BF16 = TILE_BYTES["bf16"]
TILE_BYTES_FP32 = TILE * TILE * 4


@dataclass(frozen=True)
class CBBudget:
    """CB sizes in bytes (plan §5.5).  ``cb_w`` defaults to ``cb_w_slots[dtype]`` K-blocks of the largest streamed
    slice (15 x 17,408 B = 255 KiB bfp8 / 9 x 34,816 B = 306 KiB bf16: >= 1 op slice of lookahead, 104 / 197 KB per
    core for ``Wqkv``); an explicit ``cb_w`` must be a multiple of the K-block bytes."""

    cb_x: int = 64 * 1536 * 2  # 196,608
    cb_a: int = 64 * 2048 * 2  # 262,144 = max(h 192 KB, attn 256 KB, g 192 KB, ff2 ring 2 x 131 KB)
    cb_w: Optional[int] = None
    cb_w_slots: Mapping[str, int] = field(default_factory=lambda: {"bfp8_b": 15, "bf16": 9})
    cb_mm: int = 4 * TILE_BYTES_BF16  # [64,64] bf16 (x2)
    cb_o: int = 4 * TILE_BYTES_BF16
    cb_bias: int = 2 * TILE_BYTES_BF16  # ct <= 2 tiles (x2)
    cb_adaln: int = 16 * 1024
    cb_mask: int = 4 * TILE_BYTES_BF16 + 2 * 4 * TILE_BYTES_BF16  # self [64,64] + 2 cross chunks [64,64]
    cb_k: int = 4 * TILE_BYTES_BF16
    cb_v: int = 4 * TILE_BYTES_BF16
    cb_kv_ring: int = 2 * 8 * TILE_BYTES_BF16  # 2 x [64,128] bf16 (K^T and V of a 64-key chunk)
    cb_scores: int = 2 * 4 * TILE_BYTES_FP32  # 2 x [64,64] fp32
    cb_relay_slots: int = 6  # NumBuffers 3 (trid ring) x 2 (relay lookahead)
    cb_sa: int = 64 * 1536 * 2
    l1_budget_compute: int = L1_COMPUTE_BUDGET_BYTES

    def cb_w_bytes(self, dtype_stream: str, k_block_bytes: int) -> int:
        if self.cb_w is not None:
            if self.cb_w <= 0 or self.cb_w % k_block_bytes:
                raise ValueError(f"cb_w {self.cb_w} is not a positive multiple of the K-block ({k_block_bytes} B)")
            return int(self.cb_w)
        try:
            return int(self.cb_w_slots[dtype_stream]) * k_block_bytes
        except KeyError:
            raise KeyError(f"no default cb_w slot count for stream dtype {dtype_stream!r}") from None


@dataclass(frozen=True)
class CBSpec:
    """One ``CBDescriptor`` to build: ``role_set`` names a :meth:`CoreMap.to_core_range_sets` key."""

    name: str
    cb_id: int
    role_set: str
    total_bytes: int
    page_bytes: int
    data_format: str
    tensor_backed: Optional[str] = None  # io name of the backing sharded tensor


@dataclass(frozen=True)
class SemSpec:
    name: str
    sem_id: int
    initial_value: int = 0


@dataclass(frozen=True)
class DescriptorStats:
    """What the dry run reports (WP-05 acceptance: <= 64 CBs, <= 341 common RT args, <= 4 kernel groups, compute
    L1 <= 1,000 KB)."""

    n_cbs: int
    n_common_rt_args: int
    n_kernel_groups: int
    l1_bytes_compute: int
    l1_bytes_reader: int
    l1_bytes_hub: int
    n_named_ctas: int
    n_semaphores: int
    n_io_tensors: int
    n_per_core_rt_args: int

    def check(self, budget: CBBudget) -> None:
        """Raise if any acceptance limit is exceeded."""
        if self.n_cbs > MAX_CBS:
            raise ValueError(f"{self.n_cbs} CBs > {MAX_CBS}")
        if self.n_common_rt_args > MAX_COMMON_RT_ARGS:
            raise ValueError(f"{self.n_common_rt_args} common RT args > {MAX_COMMON_RT_ARGS}")
        if self.n_kernel_groups > MAX_KERNEL_GROUPS:
            raise ValueError(f"{self.n_kernel_groups} kernel groups > {MAX_KERNEL_GROUPS}")
        if self.l1_bytes_compute > budget.l1_budget_compute:
            raise ValueError(f"compute L1 {self.l1_bytes_compute} B > {budget.l1_budget_compute}")
        for role, b in (("reader", self.l1_bytes_reader), ("hub", self.l1_bytes_hub)):
            if b > L1_USABLE_BYTES:
                raise ValueError(f"{role} L1 {b} B > usable {L1_USABLE_BYTES}")


@dataclass(frozen=True)
class IoSlot:
    """One entry of the io tensor list: ``required`` tensors must be present (non-``None``) when materialising."""

    name: str
    required: bool
    description: str


def mask_presence(layout: ArenaLayout) -> Dict[str, bool]:
    """Which cross masks the kernel of ``layout`` reads: ``{"text", "image", "all"}`` -> present (static-shape-plan.md
    §6.2; mirrors ``DeviceLayout.mask_presence`` in ``tt/layout.py``).

    * ``text`` / ``all``: per-call masks, present whenever a block of the layout uses that key set;
    * ``image``: the constant pad-key mask of the ``idx % 4 == 2`` blocks, present when the layout uses the ``image``
      key set **and** ``L_img_real != L_img_pad`` (``cross_image_keys = (real, padded)``): n16 gr1 ``(81, 96)`` has 15
      zero-pad key columns that must be masked, n17 droid ``(448, 448)`` has none and
      ``DeviceLayout.cross_image_mask()`` returns ``None`` there.
    """
    used = set(layout.key_sets)
    keys = layout.layout_cfg.cross_image_keys
    if "image" in used and keys is None:
        raise ValueError(f"layout {layout.layout_cfg.name!r} of {layout.version} uses the image key set without keys")
    return {
        "text": "text" in used,
        "image": "image" in used and int(keys[0]) != int(keys[1]),
        "all": "all" in used,
    }


def io_spec(variant: str, layout: ArenaLayout) -> List[IoSlot]:
    """The io tensors of ``variant`` in the order their addresses appear in the common runtime args (plan §5.1).

    The three cross masks are required exactly when :func:`mask_presence` says the kernel reads them; a required tensor
    that is missing makes :func:`build_program_spec` raise before any address is emitted (a silent zero address for the
    image mask would let the pad keys of the ``idx % 4 == 2`` blocks leak into the softmax).
    """
    if variant not in VARIANTS:
        raise ValueError(f"variant must be in {VARIANTS}, got {variant!r}")
    masks = mask_presence(layout)
    slots = [
        IoSlot("x_rep", True, "residual [64,1536] bf16 HEIGHT_SHARDED replicated on the compute cores (cb_x)"),
        IoSlot("a_rep", True, "in0 arena [64,2048] bf16 HEIGHT_SHARDED on all 110 cores (cb_a, mcast target)"),
        IoSlot("arena_stream", True, "block projections, WIDTH_SHARDED over the 8 DRAM banks"),
        IoSlot("arena_small", variant != "block", "tail ops (proj_out_2, decoder, encoder), WIDTH_SHARDED"),
        IoSlot("adaln_rows", True, "AdaLN table [4*(2*n_blocks+2), 1536] bf16 ROW_MAJOR"),
        IoSlot("bias_rows", True, "bias table [4*n_blocks+4, 6144] bf16 ROW_MAJOR"),
        IoSlot("b2_step", variant == "steps4", "[4,1536] bf16 action-encoder step biases"),
        IoSlot("b3_pos_rows", variant == "steps4", "[64,1536] bf16 (b3 + position rows at M-H..M)"),
        IoSlot("future_tokens", variant == "steps4" and layout.n_future_tokens > 0, "[F,1536] bf16 (n15)"),
        IoSlot("self_mask", True, "[1,1,64,64] bf16 additive self mask"),
        IoSlot("text_mask", masks["text"], "[1,1,64,L_txt_pad] bf16 per-call text mask"),
        IoSlot(
            "image_mask",
            masks["image"],
            "[1,1,64,L_img_pad] bf16 constant image pad-key mask (absent when L_img_real == L_img_pad, e.g. n17 droid)",
        ),
        IoSlot("all_mask", masks["all"], "[1,1,64,L_all_pad] bf16 per-call all-token mask (n15)"),
    ]
    for i in range(layout.n_blocks):
        if layout.is_cross(i):
            slots.append(IoSlot(f"kv_k_{i}", True, f"hoisted K of cross block {i} [1,32,L_pad,64] bf16"))
            slots.append(IoSlot(f"kv_v_{i}", True, f"hoisted V of cross block {i} [1,32,L_pad,64] bf16"))
    if variant == "step":
        slots.append(IoSlot("y_out", True, "output head result [64,1024] bf16"))
    if variant == "steps4":
        slots += [
            IoSlot("x_t", True, "Euler state [64,A_pad] (fp32 for n15, bf16 otherwise); noise placed at rows M-H..M"),
            IoSlot("state_features", True, "[1,1536] bf16 (row 0 of sa_embs)"),
            IoSlot("action_pred", True, "[64,A_pad] output"),
        ]
    return slots


@dataclass
class ProgramSpec:
    """ttnn-free description of the program: everything :func:`materialise` needs plus the statistics."""

    variant: str
    kernel_source: str
    core_map: CoreMap
    layout: ArenaLayout
    budget: CBBudget
    named_ctas: List[Tuple[str, int]]
    role_flags: List[Tuple[str, Any]]  # (flag CTA, Role | core-set name)
    per_core_rt_args: Dict[Core, List[int]]
    common_rt_names: List[str]
    common_rt_args: List[int]
    cbs: List[CBSpec]
    semaphores: List[SemSpec]
    io: List[IoSlot]
    l1_bytes: Dict[str, int]

    def kernel_groups(self) -> List[Tuple[int, ...]]:
        """Distinct role-flag combinations present on the grid (one kernel group each).  A flag maps to a
        :class:`Role` or (K4b) to a core set name: ``"hubs"`` = the hub and :attr:`CoreMap.hub2`."""
        cm = self.core_map
        hubs = {cm.hub, cm.hub2} if cm.spare else {cm.hub}

        def has(core: Core, target: Any) -> bool:
            if target == "hubs":
                return core in hubs
            return cm.role_of(core) is target

        combos = []
        for core in cm.all_cores():
            key = tuple(1 if has(core, r) else 0 for _, r in self.role_flags)
            if key not in combos:
                combos.append(key)
        return combos

    def stats(self) -> DescriptorStats:
        return DescriptorStats(
            n_cbs=len(self.cbs),
            n_common_rt_args=len(self.common_rt_args),
            n_kernel_groups=len(self.kernel_groups()),
            l1_bytes_compute=self.l1_bytes["compute"],
            l1_bytes_reader=self.l1_bytes["reader"],
            l1_bytes_hub=self.l1_bytes["hub"],
            n_named_ctas=len(self.named_ctas),
            n_semaphores=len(self.semaphores),
            n_io_tensors=len(self.io),
            n_per_core_rt_args=len(next(iter(self.per_core_rt_args.values()))),
        )

    def io_order(self) -> List[str]:
        """Names in the order ``dit_program.py`` must pass the tensors to ``ttnn.generic_op``."""
        return [s.name for s in self.io]

    def manifest(self) -> Dict[str, Any]:
        st = self.stats()
        return {
            "variant": self.variant,
            "kernel_source": self.kernel_source,
            "stats": st.__dict__,
            "named_ctas": dict(self.named_ctas),
            "common_rt_names": self.common_rt_names,
            "cbs": [cb.__dict__ for cb in self.cbs],
            "semaphores": [s.__dict__ for s in self.semaphores],
            "io": [s.name for s in self.io],
            "l1_bytes": dict(self.l1_bytes),
        }


def _cb_specs(layout: ArenaLayout, budget: CBBudget, core_map: CoreMap) -> Tuple[List[CBSpec], Dict[str, int]]:
    """CBs of plan §5.5 and the per-role L1 totals."""
    k_block_bytes = max(op.k_block_bytes for op in layout.stream_ops("stream"))
    cb_w = budget.cb_w_bytes(layout.dtype_stream, k_block_bytes)
    stream_tile = TILE_BYTES[layout.dtype_stream]
    cb_act = layout.m_pad * layout.a_pad * 4
    compute_cbs = [
        CBSpec("cb_x", CB_IDS["cb_x"], "compute", budget.cb_x, TILE_BYTES_BF16, BF16, tensor_backed="x_rep"),
        CBSpec("cb_a", CB_IDS["cb_a"], "all", budget.cb_a, TILE_BYTES_BF16, BF16, tensor_backed="a_rep"),
        CBSpec("cb_w", CB_IDS["cb_w"], "compute", cb_w, stream_tile, layout.dtype_stream),
        CBSpec("cb_mm0", CB_IDS["cb_mm0"], "compute", budget.cb_mm, TILE_BYTES_BF16, BF16),
        CBSpec("cb_mm1", CB_IDS["cb_mm1"], "compute", budget.cb_mm, TILE_BYTES_BF16, BF16),
        CBSpec("cb_o", CB_IDS["cb_o"], "compute", budget.cb_o, TILE_BYTES_BF16, BF16),
        CBSpec("cb_bias0", CB_IDS["cb_bias0"], "compute", budget.cb_bias, TILE_BYTES_BF16, BF16),
        CBSpec("cb_bias1", CB_IDS["cb_bias1"], "compute", budget.cb_bias, TILE_BYTES_BF16, BF16),
        CBSpec("cb_adaln", CB_IDS["cb_adaln"], "compute", budget.cb_adaln, TILE_BYTES_BF16, BF16),
        CBSpec("cb_mask", CB_IDS["cb_mask"], "compute", budget.cb_mask, TILE_BYTES_BF16, BF16),
        CBSpec("cb_k", CB_IDS["cb_k"], "compute", budget.cb_k, TILE_BYTES_BF16, BF16),
        CBSpec("cb_v", CB_IDS["cb_v"], "compute", budget.cb_v, TILE_BYTES_BF16, BF16),
        CBSpec("cb_kv_ring", CB_IDS["cb_kv_ring"], "compute", budget.cb_kv_ring, TILE_BYTES_BF16, BF16),
        CBSpec("cb_scores", CB_IDS["cb_scores"], "compute", budget.cb_scores, TILE_BYTES_FP32, FP32),
    ]
    reader_cbs = [
        CBSpec(
            "cb_relay",
            CB_IDS["cb_relay"],
            "readers",
            budget.cb_relay_slots * k_block_bytes,
            stream_tile,
            layout.dtype_stream,
        )
    ]
    hub_cbs = [
        CBSpec("cb_act", CB_IDS["cb_act"], "hub", cb_act, TILE_BYTES_FP32, FP32),
        CBSpec("cb_sa", CB_IDS["cb_sa"], "hub", budget.cb_sa, TILE_BYTES_BF16, BF16),
    ]
    cbs = compute_cbs + reader_cbs + hub_cbs
    for cb in cbs:
        if cb.total_bytes <= 0 or cb.total_bytes % cb.page_bytes:
            raise ValueError(
                f"{cb.name}: total {cb.total_bytes} is not a positive multiple of the page {cb.page_bytes}"
            )
    ids_by_set: Dict[str, List[int]] = {}
    for cb in cbs:
        for role_set in (("compute", "readers", "hub", "spare") if cb.role_set == "all" else (cb.role_set,)):
            ids = ids_by_set.setdefault(role_set, [])
            if cb.cb_id in ids:
                raise ValueError(f"CB id {cb.cb_id} used twice on {role_set}")
            ids.append(cb.cb_id)
    on = lambda s: sum(cb.total_bytes for cb in cbs if cb.role_set in ("all", s))  # noqa: E731
    l1 = {"compute": on("compute"), "reader": on("readers"), "hub": on("hub"), "spare": on("spare")}
    if core_map.n_compute and l1["compute"] > L1_USABLE_BYTES:
        raise ValueError(f"compute L1 {l1['compute']} B exceeds usable {L1_USABLE_BYTES}")
    return cbs, l1


def _named_ctas(
    layout: ArenaLayout,
    core_map: CoreMap,
    variant: str,
    budget: CBBudget,
    cbs: Sequence[CBSpec],
    n_io: int,
    cb_ids: Optional[Mapping[str, int]] = None,
) -> List[Tuple[str, int]]:
    """Structural constants shared by all cores (plan §5.1), incl. the :func:`mask_presence` flags.  ``cb_ids`` (default
    :data:`CB_IDS`) is the name -> id map emitted as named CTAs (the K4 spec passes :data:`K4_CB_IDS`)."""
    cfg = layout.cfg
    masks = mask_presence(layout)
    (x0, y0), (x1, y1) = core_map.mcast_rect
    cb_bytes = {cb.name: cb.total_bytes for cb in cbs}
    ctas: List[Tuple[str, int]] = [
        ("variant", VARIANTS.index(variant)),
        ("n_steps", N_STEPS if variant == "steps4" else 1),
        ("n_blocks", layout.n_blocks),
        ("m_logical", layout.m_logical),
        ("m_pad", layout.m_pad),
        ("l_txt_pad", layout.l_pad.get("text", 0)),
        ("l_img_pad", layout.l_pad.get("image", 0)),
        ("l_all_pad", layout.l_pad.get("all", 0)),
        ("a_pad", layout.a_pad),
        ("h", layout.h),
        ("action_row0", layout.action_rows[0]),
        ("n_future_tokens", layout.n_future_tokens),
        ("euler_fp32", 1 if cfg.sampler.euler_dtype == "fp32_actions_bf16_velocity" else 0),
        ("reader_mode", READER_MODES.index(layout.reader_mode)),
        ("tile_order", TILE_ORDERS.index(layout.tile_order)),
        ("fold_attn_scale", 1 if layout.fold_attn_scale else 0),
        ("text_mask_present", 1 if masks["text"] else 0),
        ("image_mask_present", 1 if masks["image"] else 0),
        ("all_mask_present", 1 if masks["all"] else 0),
        ("width", layout.width),
        ("n_heads", layout.n_heads),
        ("head_dim_pad", layout.head_dim_pad),
        ("ff_dim", layout.ff_dim),
        ("out_dim", layout.out_dim),
        ("hidden", layout.hidden),
        ("stream_tile_bytes", TILE_BYTES[layout.dtype_stream]),
        ("small_tile_bytes", TILE_BYTES[layout.dtype_small]),
        ("k_block_tiles", K_BLOCK_TILES),
        ("n_compute", core_map.n_compute),
        ("n_readers", core_map.n_readers),
        ("n_spare", len(core_map.spare)),
        ("grid_x", core_map.grid[0]),
        ("grid_y", core_map.grid[1]),
        ("hub_x", core_map.hub[0]),
        ("hub_y", core_map.hub[1]),
        ("mcast_x0", x0),
        ("mcast_y0", y0),
        ("mcast_x1", x1),
        ("mcast_y1", y1),
        ("cb_relay_slots", budget.cb_relay_slots),
        ("n_io_tensors", n_io),
        ("n_key_sets", len(KEY_SETS)),
    ]
    ids = CB_IDS if cb_ids is None else cb_ids
    ctas += [(name, cb_id) for name, cb_id in ids.items()]
    ctas += [(f"{name}_bytes", cb_bytes[name]) for name in cb_bytes]
    ctas += [(name, sem_id) for name, sem_id in SEM_IDS.items()]
    seen = set()
    for name, value in ctas:
        if name in seen:
            raise ValueError(f"duplicate named CTA {name}")
        seen.add(name)
        if not 0 <= int(value) < 2**32:
            raise ValueError(f"named CTA {name}={value} does not fit uint32")
    return ctas


def _address_of(io: Mapping[str, Any], slot: IoSlot, dry_run: bool) -> int:
    t = io.get(slot.name)
    if t is None:
        if slot.required and not dry_run:
            raise KeyError(f"io tensor {slot.name!r} is required for this variant/version but missing or None")
        return 0
    if dry_run:
        return 0
    return int(t.buffer_address())


def build_program_spec(
    core_map: CoreMap,
    arena_layout: ArenaLayout,
    variant: str,
    io: Mapping[str, Any],
    cb_budget: Optional[CBBudget] = None,
    *,
    dry_run: bool = False,
    step_index: int = 0,
    kernel_source: str = KERNEL_SOURCE,
) -> ProgramSpec:
    """The ttnn-free :class:`ProgramSpec` of ``variant`` (see module docstring); ``dry_run`` skips address reads."""
    if variant not in VARIANTS:
        raise ValueError(f"variant must be in {VARIANTS}, got {variant!r}")
    if not 0 <= step_index < N_STEPS:
        raise ValueError(f"step_index {step_index} out of range")
    budget = cb_budget if cb_budget is not None else CBBudget()
    slots = io_spec(variant, arena_layout)
    unknown = set(io) - {s.name for s in slots}
    if unknown:
        raise KeyError(f"io has tensors unknown to variant {variant!r}: {sorted(unknown)}")
    cbs, l1 = _cb_specs(arena_layout, budget, core_map)
    ctas = _named_ctas(arena_layout, core_map, variant, budget, cbs, len(slots))
    names = [f"addr_{s.name}" for s in slots] + ["step_index", "offsets_table_len"]
    args = [_address_of(io, s, dry_run) for s in slots] + [step_index]
    offsets = arena_layout.offsets_table()
    args.append(len(offsets))
    args += offsets
    names += [f"offsets[{i}]" for i in range(len(offsets))]
    ctas.append(("offsets_table_base", len(slots) + 2))
    ctas.append(("common_rt_args_len", len(args)))
    if len(args) > MAX_COMMON_RT_ARGS:
        raise ValueError(f"{len(args)} common runtime args exceed {MAX_COMMON_RT_ARGS}")
    per_core = {core: list(a) for core, a in core_map.per_core_args().items()}
    sems = [SemSpec(name, sem_id) for name, sem_id in SEM_IDS.items()]
    return ProgramSpec(
        variant=variant,
        kernel_source=kernel_source,
        core_map=core_map,
        layout=arena_layout,
        budget=budget,
        named_ctas=ctas,
        role_flags=list(ROLE_FLAGS),
        per_core_rt_args=per_core,
        common_rt_names=names,
        common_rt_args=args,
        cbs=cbs,
        semaphores=sems,
        io=slots,
        l1_bytes=l1,
    )


def materialise(spec: ProgramSpec, io: Mapping[str, Any], unpack_to_dest_fp32_cbs: Sequence[int] = ()) -> Any:
    """Turn a :class:`ProgramSpec` into a ``ttnn.ProgramDescriptor`` (lazy ttnn; needs the device tensors of ``io``).

    ``unpack_to_dest_fp32_cbs``: CB ids whose ``copy_tile`` must land exact fp32 in DST
    (``ComputeConfigDescriptor.unpack_to_dest_mode[cb] = UnpackToDestFp32``; row_layernorm.hpp ``cb_stat32``,
    euler_head.hpp ``cb_act``; K2 §7.2).  The mode list is program-wide (indexed by CB id on every core), so the K4
    spec gives those CBs an id no other-format CB uses on another role (:data:`K4_CB_IDS` / :data:`K4_HUB_CB_IDS`).
    """
    import ttnn  # lazy
    from models.demos.deepseek_v3_b1.unified_kernel_descriptor import (
        PerCoreRuntimeArgsDescriptor,
        UnifiedCompileTimeCoreDescriptor,
        UnifiedKernelDescriptor,
    )

    cm = spec.core_map
    crs = cm.to_core_range_sets()
    dt = {BF16: ttnn.bfloat16, FP32: ttnn.float32, "bfp8_b": ttnn.bfloat8_b}
    tile_desc = ttnn.TileDescriptor(ttnn.Tile([TILE, TILE]))
    cbs = []
    for cb in spec.cbs:
        if cb.tensor_backed is not None:
            t = io.get(cb.tensor_backed)
            if t is None:
                raise KeyError(f"CB {cb.name} is backed by io tensor {cb.tensor_backed!r}, which is missing")
            cbs.append(ttnn.cb_descriptor_from_sharded_tensor(cb.cb_id, t))
            continue
        fmt = ttnn.CBFormatDescriptor(
            buffer_index=cb.cb_id, data_format=dt[cb.data_format], page_size=cb.page_bytes, tile=tile_desc
        )
        cbs.append(ttnn.CBDescriptor(total_size=cb.total_bytes, core_ranges=crs[cb.role_set], format_descriptors=[fmt]))
    sems = [
        ttnn.SemaphoreDescriptor(id=s.sem_id, core_ranges=crs["all"], initial_value=s.initial_value)
        for s in spec.semaphores
    ]

    device = io["x_rep"].device()
    (x0, y0), (x1, y1) = cm.mcast_rect
    noc_start = device.worker_core_from_logical_core(ttnn.CoreCoord(x0, y0))
    noc_end = device.worker_core_from_logical_core(ttnn.CoreCoord(x1, y1))
    hub_noc = device.worker_core_from_logical_core(ttnn.CoreCoord(*cm.hub))
    hub1_noc = device.worker_core_from_logical_core(ttnn.CoreCoord(*(cm.hub2 if cm.spare else cm.hub)))
    named = list(spec.named_ctas) + [
        ("mcast_noc_x0", noc_start.x),
        ("mcast_noc_y0", noc_start.y),
        ("mcast_noc_x1", noc_end.x),
        ("mcast_noc_y1", noc_end.y),
        ("hub_noc_x", hub_noc.x),
        ("hub_noc_y", hub_noc.y),
        ("hub1_noc_x", hub1_noc.x),  # K4b: the second hub (ff2 chunks 2..5); == the hub in the K1 program
        ("hub1_noc_y", hub1_noc.y),
    ]
    per_core = [(ttnn.CoreCoord(x, y), args) for (x, y), args in spec.per_core_rt_args.items()]
    role_sets: Dict[Any, str] = {Role.READER: "readers", Role.HUB: "hub", Role.COMPUTE: "compute", "hubs": "hubs"}
    cc = ttnn.ComputeConfigDescriptor(
        math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, math_approx_mode=False
    )
    if unpack_to_dest_fp32_cbs:
        modes = [ttnn.UnpackToDestMode.Default] * MAX_CBS
        for cb_id in unpack_to_dest_fp32_cbs:
            if not 0 <= int(cb_id) < NUM_CIRCULAR_BUFFERS:
                raise ValueError(f"unpack_to_dest_fp32 CB id {cb_id} out of range")
            modes[int(cb_id)] = ttnn.UnpackToDestMode.UnpackToDestFp32
        cc.unpack_to_dest_mode = modes
    uk = UnifiedKernelDescriptor(
        kernel_source=spec.kernel_source,
        core_ranges=crs["all"],
        ncrisc_named_compile_time_args=named,
        brisc_named_compile_time_args=named,
        trisc_named_compile_time_args=named,
        ncrisc_common_runtime_args=list(spec.common_rt_args),
        brisc_common_runtime_args=list(spec.common_rt_args),
        trisc_common_runtime_args=list(spec.common_rt_args),
        trisc_compute_config=cc,
        unified_compile_time_core_descriptors=[
            UnifiedCompileTimeCoreDescriptor(
                named_compile_time_arg=flag, core_range=crs[role_sets[role]], value=1, other_value=0
            )
            for flag, role in spec.role_flags
        ],
        per_core_runtime_args_descriptor=PerCoreRuntimeArgsDescriptor(
            ncrisc_args=per_core, brisc_args=per_core, trisc_args=per_core
        ),
    )
    result = uk.get_kernel_descriptors()
    if len(result.groups) > MAX_KERNEL_GROUPS:
        raise RuntimeError(f"{len(result.groups)} kernel groups > {MAX_KERNEL_GROUPS}")
    return ttnn.ProgramDescriptor(kernels=result.kernels, semaphores=sems, cbs=cbs)


# ====================================================================================================================
# K4: the full megakernel spec (WP-K4; plan §5.4 P1-P12 + per-step tail, §5.5 + K2 §7.2 / K3 §3.3 CB requests)
# ====================================================================================================================
#: Compute-core CB ids of the K4 kernel (``dit_kernel.cpp`` / ``ops/*.hpp`` CB structs).  Ids 0-3, 6, 7, 13 keep the
#: K1 meaning; 4, 9-12 are re-purposed (K1 never used ``cb_mm1`` / ``cb_kv_ring``); 14 / 15 are free on the compute
#: cores (K1 used them on readers / hub only); 16-31 are new.  ``cb_stat32`` (19) is the only compute-core CB read
#: through ``UnpackToDestFp32``; the hub's ``cb_act`` shares id 19 so the program-wide mode list stays consistent.
K4_CB_IDS: Mapping[str, int] = {
    "cb_x": 0,  # this core's tile row of the residual [32,1536] bf16 (48 tiles), tensor-backed x_rep, updated in place
    "cb_a": 1,  # in0 arena A (mcast target), tensor-backed a_rep on all 110 cores
    "cb_w": 2,  # weight ring (direct-mode PageRing pages of 8 tiles)
    "cb_mm0": 3,  # matmul output (4-page transactions; the NCRISC gathers / pops)
    "cb_q": 4,  # Q_h [64,64] of the head cores (TRISC-only: P2 packs, HeadSDPA pops)
    "cb_o": 5,  # HeadSDPA output O_h [64,64] (TRISC packs, NCRISC gathers / pops)
    "cb_bias0": 6,  # row-broadcast bias tiles (BRISC builds, TRISC pops)
    "cb_bias1": 7,  # BRISC raw scratch (no CB ops)
    "cb_adaln_raw": 8,  # the 2 row-broadcast AdaLN tiles (1+s, b) of the own column tile (BRISC builds, TRISC pops)
    "cb_mask": 9,  # additive mask chunks, 2 x 4 tiles (BRISC pushes, TRISC pops)
    "cb_k": 10,  # K / K^T chunks, 2 x 4 tiles (NCRISC pushes, TRISC pops)
    "cb_v": 11,  # V chunks, 2 x 4 tiles
    "cb_kt": 12,  # K^T transpose scratch (TRISC-only)
    "cb_scores": 13,  # fp32 matmul partials (K1 cb_acc) 2 x 4 pages
    "cb_tw": 14,  # tail weights: one fixed 48-tile bf16 transaction per tail op (NCRISC pushes, TRISC pops)
    "cb_tconst": 15,  # tail constants: b3_pos tiles, action-row mask, state/future rows (BRISC once)
    "cb_h": 16,  # the own normalised LN tile h_(r,t) (TRISC packs, NCRISC gathers / pops)
    "cb_part": 17,  # K4b: the two remote (m, l, O) attention parts of a split cross block, 2 x 8 fp32 tiles (NCRISC pushes, TRISC pops)
    "cb_stat": 18,  # LN fp32 GI staging + the residual fence page
    "cb_stat32": 19,  # LN -mean (UnpackToDestFp32)
    "cb_const": 21,  # LN ONES / IDENT / ZERO tiles
    "cb_zero": 23,  # SDPA zero tile (also the tail epilogue's acc_to_dest operand)
    "cb_ones": 24,  # SDPA reduce scaler
    "cb_s": 25,  # SDPA scores / probabilities fp32
    "cb_oacc": 26,  # SDPA running O fp32
    "cb_m0": 27,
    "cb_m1": 28,
    "cb_l0": 29,
    "cb_l1": 30,
    "cb_alpha": 31,
}
#: Hub-only CB ids (disjoint core range: may reuse compute ids; ``cb_act`` = 19 shares the UnpackToDestFp32 mode).
K4_HUB_CB_IDS: Mapping[str, int] = {
    "cb_act": 19,  # Euler action state x_t [64,A_pad] fp32 (NCRISC reads x_t once; TRISC updates in place)
    "cb_vel": 16,  # decoder output v per step, bf16, n_steps x 2*nt tiles (TRISC packs; BRISC writes the taps)
    "cb_xin": 17,  # bf16 copy of x_t for the action encoder (TRISC packs, BRISC multicasts / pops)
    "cb_hconst": 18,  # dt row tiles (2*nt), dec b2 bias tiles (nt), zero tile (BRISC once)
    "cb_hacc": 20,  # fp32 fence / scratch page
    "cb_tok_in": 21,  # NCRISC -> TRISC token: "the dec_W2 partials landed in cb_a"
    "cb_tok_out": 22,  # TRISC -> BRISC token: "partials consumed, cb_a may be popped"
    "cb_xbf": 23,  # bf16 round-trip scratch of the bf16 Euler mode (TRISC-only)
    "cb_land": 24,  # K4b: the gathers' landing ring (K4_LAND_PAGES x 2 KB; the multicast source), NCRISC pushes, BRISC pops
}
#: K4b: pages of the hub's gather landing ring ``cb_land`` -- 2 x the 128-page transaction, i.e. four 64-page ff2 chunk
#: slots, so every chunk gather can land while earlier chunks are still being multicast (the K4 hub landed everything
#: in its 128-page ``a_rep``: chunk kc + 2 could only be sent after chunk kc's multicast, and its gather then ran
#: concurrently with the next multicast at half speed, mk_k4b_summary.md).  Must be a multiple of 128.
K4_LAND_PAGES: int = 256
#: fixed transaction of the tail weight CB (the largest tail slice per core: proj_out_2 / enc_W2a / enc_W3 = 48 tiles)
K4_TW_TILES: int = 48
#: tail constant tiles per compute core: b3_pos rows (2), action-row mask (2), state + future rows (2)
K4_TCONST_TILES: int = 6
K4_SDPA_SLOTS: int = 2  # K / V / mask chunk rings (plan §5.5 "2 x")
#: K4b split cross attention: cross blocks with at least this many 64-key chunks share every head's keys over the three
#: cores h / 32+h / 64+h (head_sdpa.hpp run_partial + merge_part); N1.5 `all` (6 chunks) and N1.7 `image` (7) split,
#: N1.6 `image` (2) and every `text` set (1) do not.  Mirrored by dit_kernel.cpp `mk::split_blk`.
K4_SPLIT_MIN_CHUNKS: int = 3
K4_SPLIT_WAYS: int = 3
#: fp32 pages of `cb_part`: the two remote parts (m 2, l 2, O 4 tiles each) landing on the head core
K4_PART_PAGES: int = 2 * 8


@dataclass(frozen=True)
class K4Budget:
    """L1 budget of the K4 spec (bytes / counts).  Defaults: ~1.15 MB (bf16 arena) / ~0.98 MB (bfp8) of the 1,427 KiB
    usable on the compute cores with a weight-ring depth of 13 bf16 / 15 bfp8 K-block slots of lookahead (K1 used 9
    bf16; 13 buys 6 % of step time on the bf16 arena, mk_k4_summary.md §5.2, and still fits under ``l1_budget_compute``),
    the distributed LN of ``dist_layernorm.hpp`` (no centred-row scratch, no fp32 affine ring: the K2 ``cb_xc`` /
    ``cb_y32`` are gone) and the per-core tile row of x (96 KB instead of the replicated 192 KB); ``l1_budget_compute``
    is the acceptance limit checked by :meth:`DescriptorStats.check`."""

    cb_w_slots: Mapping[str, int] = field(default_factory=lambda: {"bfp8_b": 15, "bf16": 13})
    cb_w: Optional[int] = None
    cb_mm: int = 8 * TILE_BYTES_BF16  # 2 slots of 4 pages
    cb_scores_pages: int = 8  # 2 slots of 4 fp32 partial pages
    cb_adaln_slots: int = 2  # (1+s, b) row-broadcast tile pairs the BRISC may run ahead
    cb_h_tiles: int = 4  # own LN tiles in flight (TRISC -> NCRISC)
    l1_budget_compute: int = 1_340_000
    #: below L1_USABLE_BYTES: dit_program.K4_KERNEL_CONFIG_EXTRA_BYTES (96 KiB of the worker L1 go to the kernel-config
    #: ring buffer) + the L1 allocator's headroom for the sharded buffers
    l1_margin: int = 96 * 1024 + 24_000

    def cb_w_bytes(self, dtype_stream: str, k_block_bytes: int) -> int:
        if self.cb_w is not None:
            if self.cb_w <= 0 or self.cb_w % k_block_bytes:
                raise ValueError(f"cb_w {self.cb_w} is not a positive multiple of the K-block ({k_block_bytes} B)")
            return int(self.cb_w)
        try:
            return int(self.cb_w_slots[dtype_stream]) * k_block_bytes
        except KeyError:
            raise KeyError(f"no K4 cb_w slot count for stream dtype {dtype_stream!r}") from None


def k4_n_tiles_a(layout: ArenaLayout) -> int:
    """Column tiles of the action buffer ``[64, A_pad]`` (``nt`` of dec_W2 / enc_W1; 1 / 4 / 5 for N1.5 / N1.6 / N1.7)."""
    if layout.a_pad % TILE:
        raise ValueError(f"A_pad {layout.a_pad} is not a tile multiple")
    return layout.a_pad // TILE


def k4_cb_specs(layout: ArenaLayout, k4: K4Budget, core_map: CoreMap) -> Tuple[List[CBSpec], Dict[str, int]]:
    """CBs of the K4 kernel per role (compute / hub / spare; readers idle in direct mode) and the per-role L1 totals."""
    ids = K4_CB_IDS
    hid = K4_HUB_CB_IDS
    k_block_bytes = max(op.k_block_bytes for op in layout.stream_ops("stream"))
    cb_w = k4.cb_w_bytes(layout.dtype_stream, k_block_bytes)
    stream_tile = TILE_BYTES[layout.dtype_stream]
    if layout.dtype_small != "bf16":
        raise ValueError("K4 streams the tail ops as bf16 tiles (plan D4)")
    nt = k4_n_tiles_a(layout)
    n_steps = N_STEPS
    row_tiles = 2 * (layout.width // TILE)  # 96 tiles of the [64,1536] stream
    b = TILE_BYTES_BF16
    f = TILE_BYTES_FP32
    compute = [
        CBSpec("cb_x", ids["cb_x"], "compute", (row_tiles // 2) * b, b, BF16, tensor_backed="x_rep"),
        CBSpec("cb_a", ids["cb_a"], "all", 64 * 2048 * 2, b, BF16, tensor_backed="a_rep"),
        CBSpec("cb_w", ids["cb_w"], "compute", cb_w, stream_tile, layout.dtype_stream),
        CBSpec("cb_mm0", ids["cb_mm0"], "compute", k4.cb_mm, b, BF16),
        CBSpec("cb_q", ids["cb_q"], "compute", 4 * b, b, BF16),
        CBSpec("cb_o", ids["cb_o"], "compute", 4 * b, b, BF16),
        CBSpec("cb_bias0", ids["cb_bias0"], "compute", 4 * b, b, BF16),
        CBSpec("cb_bias1", ids["cb_bias1"], "compute", 4 * b, b, BF16),
        CBSpec("cb_adaln_raw", ids["cb_adaln_raw"], "compute", k4.cb_adaln_slots * 2 * b, b, BF16),
        CBSpec("cb_mask", ids["cb_mask"], "compute", K4_SDPA_SLOTS * 4 * b, b, BF16),
        CBSpec("cb_k", ids["cb_k"], "compute", K4_SDPA_SLOTS * 4 * b, b, BF16),
        CBSpec("cb_v", ids["cb_v"], "compute", K4_SDPA_SLOTS * 4 * b, b, BF16),
        CBSpec("cb_kt", ids["cb_kt"], "compute", 4 * b, b, BF16),
        CBSpec("cb_scores", ids["cb_scores"], "compute", k4.cb_scores_pages * f, f, FP32),
        CBSpec("cb_tw", ids["cb_tw"], "compute", K4_TW_TILES * b, b, BF16),
        CBSpec("cb_tconst", ids["cb_tconst"], "compute", K4_TCONST_TILES * b, b, BF16),
        CBSpec("cb_h", ids["cb_h"], "compute", k4.cb_h_tiles * b, b, BF16),
        CBSpec("cb_part", ids["cb_part"], "compute", K4_PART_PAGES * f, f, FP32),
        CBSpec("cb_stat", ids["cb_stat"], "compute", 2 * f, f, FP32),
        CBSpec("cb_stat32", ids["cb_stat32"], "compute", f, f, FP32),
        CBSpec("cb_const", ids["cb_const"], "compute", 3 * b, b, BF16),
        CBSpec("cb_zero", ids["cb_zero"], "compute", b, b, BF16),
        CBSpec("cb_ones", ids["cb_ones"], "compute", b, b, BF16),
        CBSpec("cb_s", ids["cb_s"], "compute", 4 * f, f, FP32),
        CBSpec("cb_oacc", ids["cb_oacc"], "compute", 4 * f, f, FP32),
        CBSpec("cb_m0", ids["cb_m0"], "compute", 2 * f, f, FP32),
        CBSpec("cb_m1", ids["cb_m1"], "compute", 2 * f, f, FP32),
        CBSpec("cb_l0", ids["cb_l0"], "compute", 2 * f, f, FP32),
        CBSpec("cb_l1", ids["cb_l1"], "compute", 2 * f, f, FP32),
        CBSpec("cb_alpha", ids["cb_alpha"], "compute", 2 * f, f, FP32),
    ]
    hub = [
        CBSpec("cb_act", hid["cb_act"], "hubs", 2 * nt * f, f, FP32),
        CBSpec("cb_vel", hid["cb_vel"], "hubs", n_steps * 2 * nt * b, b, BF16),
        CBSpec("cb_xin", hid["cb_xin"], "hubs", 2 * nt * b, b, BF16),
        CBSpec("cb_hconst", hid["cb_hconst"], "hubs", (3 * nt + 1) * b, b, BF16),
        CBSpec("cb_hacc", hid["cb_hacc"], "hubs", f, f, FP32),
        CBSpec("cb_tok_in", hid["cb_tok_in"], "hubs", 2 * b, b, BF16),
        CBSpec("cb_tok_out", hid["cb_tok_out"], "hubs", 2 * b, b, BF16),
        CBSpec("cb_xbf", hid["cb_xbf"], "hubs", 2 * nt * b, b, BF16),
        CBSpec("cb_land", hid["cb_land"], "hubs", K4_LAND_PAGES * b, b, BF16),
    ]
    if K4_LAND_PAGES % 128:
        raise ValueError("K4_LAND_PAGES must be a multiple of the 128-page hub transaction")
    cbs = compute + hub
    for cb in cbs:
        if cb.total_bytes <= 0 or cb.total_bytes % cb.page_bytes:
            raise ValueError(
                f"{cb.name}: total {cb.total_bytes} is not a positive multiple of the page {cb.page_bytes}"
            )
        if not 0 <= cb.cb_id < NUM_CIRCULAR_BUFFERS:
            raise ValueError(f"{cb.name}: CB id {cb.cb_id} out of range")
    ids_by_set: Dict[str, List[int]] = {}
    for cb in cbs:
        for role_set in (("compute", "readers", "hubs", "spare") if cb.role_set == "all" else (cb.role_set,)):
            lst = ids_by_set.setdefault(role_set, [])
            if cb.cb_id in lst:
                raise ValueError(f"CB id {cb.cb_id} used twice on {role_set}")
            lst.append(cb.cb_id)
    on = lambda s: sum(cb.total_bytes for cb in cbs if cb.role_set in ("all", s))  # noqa: E731
    l1 = {"compute": on("compute"), "reader": on("readers"), "hub": on("hubs"), "spare": on("spare")}
    limit = L1_USABLE_BYTES - k4.l1_margin
    for role, used in l1.items():
        if used > limit:
            raise ValueError(f"K4 {role} L1 {used} B exceeds usable {L1_USABLE_BYTES} - margin {k4.l1_margin}")
    if l1["compute"] > k4.l1_budget_compute:
        raise ValueError(f"K4 compute L1 {l1['compute']} B > budget {k4.l1_budget_compute}")
    return cbs, l1


def k4_io_extra(variant: str, layout: ArenaLayout) -> List[IoSlot]:
    """Io tensors the K4 kernel needs beyond :func:`io_spec`: the per-step velocity taps of ``steps4`` (optional;
    ``pred_velocity[k]`` acceptance) -- the ``x_t`` / ``action_pred`` slots of :func:`io_spec` carry fp32 ``[64,A_pad]``
    DRAM tensors in K4 (the hub keeps the Euler state in fp32 for every version, rounding to bf16 per step when
    ``euler_fp32 == 0``; ``euler_head.hpp``)."""
    if variant not in VARIANTS:
        raise ValueError(f"variant must be in {VARIANTS}, got {variant!r}")
    if variant != "steps4":
        return []
    return [IoSlot("velocity_out", False, "[4*64, A_pad] bf16 TILE: decoder output of every step (taps; optional)")]


def build_program_spec_k4(
    core_map: CoreMap,
    arena_layout: ArenaLayout,
    variant: str,
    io: Mapping[str, Any],
    k4: Optional[K4Budget] = None,
    *,
    dry_run: bool = False,
    step_index: int = 0,
    kernel_source: str = KERNEL_SOURCE,
) -> ProgramSpec:
    """The K4 :class:`ProgramSpec` (module docstring): K1 io order + :func:`k4_io_extra`, :func:`k4_cb_specs`,
    named CTAs with :data:`K4_CB_IDS` (+ ``k4_*`` hub ids / sizes), the same semaphores and common RT layout as
    :func:`build_program_spec` (``dit_program.py`` appends its own RT args after the offsets table)."""
    if variant not in VARIANTS:
        raise ValueError(f"variant must be in {VARIANTS}, got {variant!r}")
    if not 0 <= step_index < N_STEPS:
        raise ValueError(f"step_index {step_index} out of range")
    if arena_layout.reader_mode != "direct":
        raise ValueError("the K4 kernel streams weights in direct mode only (K1 GO decision, mk_k1_summary.md §0)")
    k4b = k4 if k4 is not None else K4Budget()
    slots = io_spec(variant, arena_layout) + k4_io_extra(variant, arena_layout)
    unknown = set(io) - {s.name for s in slots}
    if unknown:
        raise KeyError(f"io has tensors unknown to variant {variant!r}: {sorted(unknown)}")
    cbs, l1 = k4_cb_specs(arena_layout, k4b, core_map)
    budget = CBBudget(l1_budget_compute=k4b.l1_budget_compute)
    ctas = _named_ctas(arena_layout, core_map, variant, budget, cbs, len(slots), cb_ids=K4_CB_IDS)
    ctas += [(f"k4_{name}", cb_id) for name, cb_id in K4_HUB_CB_IDS.items()]
    ctas += [
        ("k4_tw_tiles", K4_TW_TILES),
        ("k4_tconst_tiles", K4_TCONST_TILES),
        ("k4_sdpa_slots", K4_SDPA_SLOTS),
        ("k4_nt_a", k4_n_tiles_a(arena_layout)),
        ("k4_velocity_out_present", 1 if io.get("velocity_out") is not None else 0),
    ]
    seen = set()
    for name, _ in ctas:
        if name in seen:
            raise ValueError(f"duplicate named CTA {name}")
        seen.add(name)
    names = [f"addr_{s.name}" for s in slots] + ["step_index", "offsets_table_len"]
    args = [_address_of(io, s, dry_run) for s in slots] + [step_index]
    offsets = arena_layout.offsets_table()
    args.append(len(offsets))
    args += offsets
    names += [f"offsets[{i}]" for i in range(len(offsets))]
    ctas.append(("offsets_table_base", len(slots) + 2))
    ctas.append(("common_rt_args_len", len(args)))
    if len(args) > MAX_COMMON_RT_ARGS:
        raise ValueError(f"{len(args)} common runtime args exceed {MAX_COMMON_RT_ARGS}")
    per_core = {core: list(a) for core, a in core_map.per_core_args().items()}
    sems = [SemSpec(name, sem_id) for name, sem_id in SEM_IDS.items()]
    # K4b: `is_hub` on the hub AND CoreMap.hub2 (the ff2-chunk hub); the per-core `group` word tells them apart
    role_flags = [(flag, "hubs" if role is Role.HUB else role) for flag, role in ROLE_FLAGS]
    return ProgramSpec(
        variant=variant,
        kernel_source=kernel_source,
        core_map=core_map,
        layout=arena_layout,
        budget=budget,
        named_ctas=ctas,
        role_flags=role_flags,
        per_core_rt_args=per_core,
        common_rt_names=names,
        common_rt_args=args,
        cbs=cbs,
        semaphores=sems,
        io=slots,
        l1_bytes=l1,
    )


def k4_unpack_fp32_cbs() -> Tuple[int, ...]:
    """CB ids read through ``UnpackToDestFp32`` in the K4 program (compute ``cb_stat32`` == hub ``cb_act``)."""
    if K4_CB_IDS["cb_stat32"] != K4_HUB_CB_IDS["cb_act"]:
        raise AssertionError("cb_stat32 and cb_act must share the UnpackToDestFp32 id (program-wide mode list)")
    return (K4_CB_IDS["cb_stat32"],)


def build_program(
    core_map: CoreMap,
    arena_layout: ArenaLayout,
    variant: str,
    io: Mapping[str, Any],
    cb_budget: Optional[CBBudget] = None,
    *,
    dry_run: bool = False,
    step_index: int = 0,
    kernel_source: str = KERNEL_SOURCE,
) -> Any:
    """Build the megakernel program (plan §5.1).

    Returns :class:`DescriptorStats` when ``dry_run`` (no ttnn import, ``io`` values may be ``None``), otherwise the
    ``ttnn.ProgramDescriptor`` whose io tensor order is :meth:`ProgramSpec.io_order` (see :func:`build_program_spec`).
    """
    spec = build_program_spec(
        core_map,
        arena_layout,
        variant,
        io,
        cb_budget,
        dry_run=dry_run,
        step_index=step_index,
        kernel_source=kernel_source,
    )
    if dry_run:
        return spec.stats()
    return materialise(spec, io)