# 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", ]