code/models/experimental/gr00t/tt/device.py

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

"""Device open / close, compute-kernel-config presets and memory presets (IMPLEMENTATION_PLAN.md §2.2 ``tt/device.py``).

``ttnn`` is imported lazily: the presets below are CPU-side *descriptors* (:class:`CKCPreset`, :class:`MemPreset`)
whose :meth:`build` creates the real ``ttnn`` object on first use, so ``tt.layers`` can use them as default
arguments without importing ``ttnn`` at module import time (plan §2.1 rule).

Presets (plan §1.2 dtype / fidelity policy, mb2 §B.3 / §B.5 / §B.8):

* :data:`CKC_HIFI2_FP32ACC` -- projections (bfp8 or bf16 weights): HiFi2 + ``fp32_dest_acc_en`` is mandatory
  (mb2 §B.8: without fp32 acc the FF-out PCC drops to 0.9991).
* :data:`CKC_HIFI4_FP32ACC` -- bf16-forever tensors (``proj_out_2``, decoders, encoders).
* :data:`CKC_HIFI4_FP32ACC_NOAPPROX` -- every ``layer_norm`` / ``rms_norm`` call: ``rms_norm``'s default is
  approx-on / fp32acc-off and loses max|d| 0.21 on the massive-activation row (mb2 §B.5).
* :data:`CKC_LOFI_FP32ACC` -- gated N1.7 ViT FF lever only (plan §1.2; LoFi is forbidden on attention projections).
* :data:`CKC_SDPA` -- SDPA: HiFi2, fp32 acc off (mb2 §B.8 (v): the SDPA precision knobs change nothing);
  ``exp_approx_mode=False`` lives in the ``SDPAProgramConfig`` (``tt.layers.sdpa``).
* :data:`CKC_SDPA_HIFI4` -- LLM SDPA (plan §1.2 "LLM SDPA HiFi4").
"""

from __future__ import annotations

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

# Device facts of the p150a this plan is written for (mb1 §5, Table 5a; plan §2.1 ``DEVICE_COMPUTE_GRID``).
EXPECTED_COMPUTE_GRID: Tuple[int, int] = (11, 10)
EXPECTED_DRAM_GRID: Tuple[int, int] = (8, 1)
DEFAULT_TRACE_REGION_SIZE: int = 64 << 20  # plan §3.5: <= 16 MiB per trace B + trace A + a second layout
DEFAULT_L1_SMALL_SIZE: int = 32768  # mb1/mb2 convention (mb2 §1)

FIDELITIES: Tuple[str, ...] = ("LoFi", "HiFi2", "HiFi3", "HiFi4")


@dataclass(frozen=True)
class CKCPreset:
    """CPU-side description of a ``ttnn.WormholeComputeKernelConfig`` (alias of the Blackhole config,
    ``ttnn-nanobind/operations/core.cpp:58-84``); :meth:`build` creates the ``ttnn`` object lazily."""

    name: str
    fidelity: str
    math_approx_mode: bool
    fp32_dest_acc_en: bool
    packer_l1_acc: bool = False

    def __post_init__(self) -> None:
        if self.fidelity not in FIDELITIES:
            raise ValueError(f"unknown math fidelity {self.fidelity!r}; expected one of {FIDELITIES}")

    def build(self) -> Any:
        import ttnn

        fid = getattr(ttnn.MathFidelity, self.fidelity)
        return ttnn.WormholeComputeKernelConfig(
            math_fidelity=fid,
            math_approx_mode=self.math_approx_mode,
            fp32_dest_acc_en=self.fp32_dest_acc_en,
            packer_l1_acc=self.packer_l1_acc,
        )

    def as_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "fidelity": self.fidelity,
            "math_approx_mode": self.math_approx_mode,
            "fp32_dest_acc_en": self.fp32_dest_acc_en,
            "packer_l1_acc": self.packer_l1_acc,
        }


CKC_HIFI2_FP32ACC = CKCPreset("hifi2_fp32acc", "HiFi2", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_HIFI4_FP32ACC = CKCPreset("hifi4_fp32acc", "HiFi4", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_HIFI4_FP32ACC_NOAPPROX = CKCPreset(
    "hifi4_fp32acc_noapprox", "HiFi4", math_approx_mode=False, fp32_dest_acc_en=True, packer_l1_acc=True
)
CKC_LOFI_FP32ACC = CKCPreset("lofi_fp32acc", "LoFi", math_approx_mode=False, fp32_dest_acc_en=True)
CKC_SDPA = CKCPreset("sdpa_hifi2", "HiFi2", math_approx_mode=False, fp32_dest_acc_en=False)
CKC_SDPA_HIFI4 = CKCPreset("sdpa_hifi4", "HiFi4", math_approx_mode=False, fp32_dest_acc_en=False)

CKC_PRESETS: Dict[str, CKCPreset] = {
    p.name: p
    for p in (
        CKC_HIFI2_FP32ACC,
        CKC_HIFI4_FP32ACC,
        CKC_HIFI4_FP32ACC_NOAPPROX,
        CKC_LOFI_FP32ACC,
        CKC_SDPA,
        CKC_SDPA_HIFI4,
    )
}


def resolve_ckc(ckc: Any) -> Any:
    """``CKCPreset`` / preset name / ready ``ttnn`` config -> ``ttnn`` compute-kernel config (``None`` stays ``None``)."""
    if ckc is None:
        return None
    if isinstance(ckc, CKCPreset):
        return ckc.build()
    if isinstance(ckc, str):
        if ckc not in CKC_PRESETS:
            raise KeyError(f"unknown compute-kernel preset {ckc!r}; known: {sorted(CKC_PRESETS)}")
        return CKC_PRESETS[ckc].build()
    return ckc  # an already built ttnn.WormholeComputeKernelConfig


@dataclass(frozen=True)
class MemPreset:
    """``"L1"`` or ``"DRAM"`` interleaved memory config, built lazily (plan §3 memory legend L / D)."""

    kind: str

    def __post_init__(self) -> None:
        if self.kind not in ("L1", "DRAM"):
            raise ValueError(f"memory preset must be 'L1' or 'DRAM', got {self.kind!r}")

    def build(self) -> Any:
        import ttnn

        return ttnn.L1_MEMORY_CONFIG if self.kind == "L1" else ttnn.DRAM_MEMORY_CONFIG


L1 = MemPreset("L1")
DRAM = MemPreset("DRAM")


def resolve_mem(mem: Any) -> Any:
    """``MemPreset`` / ``"L1"`` / ``"DRAM"`` / ready ``ttnn.MemoryConfig`` -> ``ttnn.MemoryConfig`` (``None`` stays)."""
    if mem is None:
        return None
    if isinstance(mem, MemPreset):
        return mem.build()
    if isinstance(mem, str):
        return MemPreset(mem).build()
    return mem


def open_gr00t_device(
    trace_region_size: int = DEFAULT_TRACE_REGION_SIZE,
    l1_small_size: int = DEFAULT_L1_SMALL_SIZE,
    num_command_queues: int = 1,
    device_id: int = 0,
    worker_l1_size: Optional[int] = None,
) -> Any:
    """Open device ``device_id`` with the program cache enabled and assert the p150a facts this plan relies on.

    ``ttnn.open_device(device_id, l1_small_size, trace_region_size[, num_command_queues][, worker_l1_size])``
    (``ttnn-nanobind/device.cpp:71-80``, mb2 §5) returns the 1x1 ``MeshDevice``; the compute grid must be 11x10 and
    the DRAM grid 8x1 (mb1 Table 5a) -- every core map and every program config in ``tt.layers`` is sized for that
    card, so anything else raises instead of silently running a wrong configuration.  ``worker_l1_size`` (bytes of
    allocatable L1 per worker core; ``None`` = the firmware default, 1,461,248 B on the p150a) is what
    ``tt.model.open_model_device`` reduces for the megakernel DiT backend (the kernel-config ring buffer grows by the
    same amount, ``mk_k4_summary.md`` §7.2 fact 2); the value actually in force is checked there.  Callers hold the
    device lock (``bin/with-device.sh``).
    """
    import ttnn

    kwargs: Dict[str, Any] = dict(device_id=device_id, l1_small_size=l1_small_size, trace_region_size=trace_region_size)
    if num_command_queues != 1:
        kwargs["num_command_queues"] = int(num_command_queues)
    if worker_l1_size is not None:
        if int(worker_l1_size) <= 0:
            raise ValueError(f"worker_l1_size must be a positive number of bytes, got {worker_l1_size}")
        kwargs["worker_l1_size"] = int(worker_l1_size)
    device = ttnn.open_device(**kwargs)
    try:
        device.enable_program_cache()
        grid = device.compute_with_storage_grid_size()
        dram = device.dram_grid_size()
        got_grid = (int(grid.x), int(grid.y))
        got_dram = (int(dram.x), int(dram.y))
        if got_grid != EXPECTED_COMPUTE_GRID:
            raise RuntimeError(
                f"compute_with_storage_grid_size() = {got_grid}, this plan requires {EXPECTED_COMPUTE_GRID}"
            )
        if got_dram != EXPECTED_DRAM_GRID:
            raise RuntimeError(f"dram_grid_size() = {got_dram}, this plan requires {EXPECTED_DRAM_GRID}")
    except Exception:
        ttnn.close_device(device)
        raise
    return device


def close_gr00t_device(device: Any) -> None:
    """Synchronise and close (``ttnn.close_device``); errors during the final synchronise are re-raised after closing."""
    import ttnn

    sync_error: Optional[BaseException] = None
    try:
        ttnn.synchronize_device(device)
    except Exception as exc:  # noqa: BLE001 - we still want to close the device
        sync_error = exc
    ttnn.close_device(device)
    if sync_error is not None:
        raise sync_error


def dram_worker_cores(device: Any, noc: Any = None) -> List[Any]:
    """Runtime query of the bank-adjacent worker cores (mb1 §5: NOC_0 -> (0,9),(0,0),(0,7),(0,3),(6,9),(6,1),(6,6),(6,4)
    on this card; never use a pinned list -- one harvested Tensix column sits at NOC0 x in 1..7)."""
    import ttnn

    if noc is None:
        noc = ttnn.NOC.NOC_0
    cores = device.get_optimal_dram_bank_to_logical_worker_assignment(noc)
    return list(cores)


def device_facts(device: Any) -> Dict[str, Any]:
    """Small JSON-able summary for the results files (plan §7.1 perf tier: "device facts")."""
    import ttnn

    grid = device.compute_with_storage_grid_size()
    dram = device.dram_grid_size()
    facts: Dict[str, Any] = {
        "compute_with_storage_grid_size": [int(grid.x), int(grid.y)],
        "dram_grid_size": [int(dram.x), int(dram.y)],
        "arch": str(device.arch()),
    }
    for name in ("NOC_0", "NOC_1"):
        cores = dram_worker_cores(device, getattr(ttnn.NOC, name))
        facts[f"dram_workers_{name}"] = [[int(c.x), int(c.y)] for c in cores]
    return facts


def core_grid(xy: Sequence[int] = EXPECTED_COMPUTE_GRID) -> Any:
    """``ttnn.CoreCoord(x, y)`` of a ``(cols, rows)`` pair (the program configs take the grid as a ``CoreCoord``)."""
    import ttnn

    if len(xy) != 2:
        raise ValueError(f"grid must be (cols, rows), got {tuple(xy)}")
    return ttnn.CoreCoord(int(xy[0]), int(xy[1]))


__all__ = [
    "CKCPreset",
    "CKC_HIFI2_FP32ACC",
    "CKC_HIFI4_FP32ACC",
    "CKC_HIFI4_FP32ACC_NOAPPROX",
    "CKC_LOFI_FP32ACC",
    "CKC_SDPA",
    "CKC_SDPA_HIFI4",
    "CKC_PRESETS",
    "DEFAULT_L1_SMALL_SIZE",
    "DEFAULT_TRACE_REGION_SIZE",
    "DRAM",
    "EXPECTED_COMPUTE_GRID",
    "EXPECTED_DRAM_GRID",
    "L1",
    "MemPreset",
    "close_gr00t_device",
    "core_grid",
    "device_facts",
    "dram_worker_cores",
    "open_gr00t_device",
    "resolve_ckc",
    "resolve_mem",
]