Add files using upload-large-folder tool

README.md

@@ -5,25 +5,11 @@ tags:
 - tt-dit-server
 - tt-model-cache
 - tt-model-container
-- tenstorrent
-- ttnn
-- tt-metal
-- tt-nn
-- robotics
-- vla
-- gr00t
-- tt-model-catalog
-pipeline_tag: robotics
-license: other
-license_name: nvidia-license
-license_link: https://huggingface.co/nvidia/GR00T-N1.6-3B/blob/main/LICENSE
-base_model:
-- nvidia/GR00T-N1.6-3B
 ---
 
-# GR00T-N1.6-3B-p150
+# gr00t-n1.6-3b-p150
 
-NVIDIA Isaac GR00T N1.6 (3B vision-language-action policy) running on one Tenstorrent Blackhole p150a via tt-nn: one camera frame, proprioceptive state and a task instruction in, a 16-step chunk of GR1 arm, hand and waist actions out. Stage-1 release: TTNN ops replayed from Metal traces (DiT weights bfp8, everything else bf16); the persistent-megakernel denoise is in progress and not in this image.
+NVIDIA Isaac GR00T N1.6 (3B vision-language-action policy) running on one Tenstorrent Blackhole p150a via tt-nn: one camera frame, proprioceptive state and a task instruction in, a 16-step chunk of GR1 arm, hand and waist actions out. Stage-2 release: the whole DiT denoise (32 blocks × 4 flow-matching steps) runs as one persistent megakernel streaming bfp8 weights from a DRAM arena; vision, LLM and adapter are TTNN ops (bf16) replayed from Metal traces; `GR00T_DIT_BACKEND=ttnn` selects the Stage-1 denoise.
 Weights: [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) · Paper: [arXiv:2503.14734](https://arxiv.org/abs/2503.14734) · Project page: [gr00t-n1_6](https://research.nvidia.com/labs/gear/gr00t-n1_6/) · Upstream code: [NVIDIA/Isaac-GR00T](https://github.com/NVIDIA/Isaac-GR00T) (`n1.6.1-release`) · Port: `code/models/experimental/gr00t` and the server `code/gr00t_p150` in this repo
 
 Runs on **p150** (mesh `P150`).
@@ -37,8 +23,7 @@ tt-model pull  changh95/GR00T-N1.6-3B-p150 --with-weights
 tt-model serve changh95/GR00T-N1.6-3B-p150
 ```
 
-- Weights [`nvidia/GR00T-N1.6-3B`](https://huggingface.co/nvidia/GR00T-N1.6-3B) at `d0814e7ecb19` (6.57 GB) go to your HF cache; the image does not contain them.
-- Serves on port 20000 (or the next free port); ready when the log says `Application startup complete`.
+`pull --with-weights` downloads the Docker image and the [`nvidia/GR00T-N1.6-3B`](https://huggingface.co/nvidia/GR00T-N1.6-3B) weights at `d0814e7ecb19202e7c8468b46098b0b7ef3a6d61` (into your HF cache; they are not in the image). `serve` starts the model's own HTTP server on port 20000 (or the next free port, if that one is busy); the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs `Application startup complete`.
 
 ### Run with tt-cli
 
@@ -55,18 +40,18 @@ tt model stop changh95/GR00T-N1.6-3B-p150
 ```
 
 - `POST /predict`: `images` (`{"ego_view_bg_crop_pad_res256_freq20": <base64 PNG/JPEG>}`), `state` (`{group: [floats]}`: `left_arm` 7, `right_arm` 7, `left_hand` 6, `right_hand` 6, `waist` 3, raw joint values), `instruction` (≤ 24 tokens); optional `seed` (default 42) or `noise` (`[50, 128]`), `return_normalized`.
-- `GET /health`, `GET /info`, `GET /demo` (the shipped demo request and its golden actions).
+- `GET /health`, `GET /info` (`stage` names the denoise backend), `GET /demo` (the shipped demo request and its golden actions).
 
 ### Response
 
 ```json
-{"actions": {"left_arm":   [[-0.0289, 0.1453, 0.0224, -2.1383, -0.0165, -0.0260, -0.0414], "... 16 steps"],
-             "right_arm":  [[-0.0024, -0.1502, 0.0280, -2.1461, -0.0091, 0.0132, 0.0044], "..."],
-             "left_hand":  [[0.0097, -0.0177, 0.0041, 0.0078, 0.0539, 0.0281], "..."],
-             "right_hand": [[-0.2339, -0.2466, -0.2431, -0.2311, -0.2423, 0.3333], "..."],
-             "waist":      [[-0.0010, -0.0060, 0.0008], "..."]},
+{"actions": {"left_arm":   [[-0.0289, 0.1453, 0.0224, -2.1383, -0.0165, -0.0225, -0.0428], "... 16 steps"],
+             "right_arm":  [[-0.0024, -0.1512, 0.0269, -2.1461, -0.0091, 0.0124, 0.0044], "..."],
+             "left_hand":  [[0.0135, -0.0092, -0.0047, 0.0070, 0.0455, 0.0281], "..."],
+             "right_hand": [[-0.2253, -0.2421, -0.2345, -0.2311, -0.2338, 0.3385], "..."],
+             "waist":      [[0.0001, -0.0060, 0.0008], "..."]},
  "action_horizon": 16, "embodiment": "gr1", "seq_len": 116, "seed": 42,
- "timing_ms": {"decode": 0.8, "encode": 1.5, "device": 56.5, "decode_actions": 0.1, "total": 58.9}}
+ "timing_ms": {"decode": 0.9, "encode": 1.4, "device": 39.6, "decode_actions": 0.2, "total": 42.1}}
 ```
 
 - `actions` are physical, un-normalised joint targets, 16 steps × [7, 7, 6, 6, 3] dims, as `Gr00tPolicy.get_action` returns them for the GR1 data config (arm and hand groups are relative outputs composed onto the request's `state`, `waist` is absolute); `timing_ms.device` covers upload, the trace replays and the readback.
@@ -81,21 +66,23 @@ tt model stop changh95/GR00T-N1.6-3B-p150
 
 | Metric | Value |
 |---|---:|
-| Action chunk vs the fp32 reference (same observation, same initial noise) | PCC `left_arm` 0.99995 · `right_arm` 0.99991 · `right_hand` 0.99312 · `left_hand` / `waist` max\|d\| 0.059 / 0.017 (near-constant groups) · all 29 action dims 0.999209 |
-| Intermediate taps vs the fp32 reference (45 canonical taps, from pixels) | 52 of 56 gated rows pass; 4 marginal misses on the DiT state-token row that no action reads (`dit_out[k=1,2]` 0.999556 / 0.999496 vs gates 0.99961 / 0.99964, `action_decoder_out[k=1,2]` 0.999579 / 0.999556 vs 0.9997 / 0.99978), caused by bf16 error in the SigLIP2 tower |
-| Traced vs untraced, determinism | bit-identical; 5 trace replays bit-identical |
-| Inference on p150a (warm, batch 1, median of 50) | **59.4 ms end-to-end** = host encode 2.0 + upload 3.8 + traces 53.3 (vision 11.5 · LLM 9.5 · adapter 2.0 · denoise 30.4) + readback and decode 0.3 |
-| Served over HTTP by this image (warm, batch 1) | 56.8 ms device · 59.4 ms server-side incl. PNG decode |
-| Same policy on an RTX 5090 (same host, NVIDIA's unmodified `Gr00tPolicy`, eager PyTorch bf16 + flash-attn 2) | 63.1 ms end-to-end → p150a 1.06× faster; with `torch.compile` on the DiT the GPU reaches 33.1 ms → GPU 1.8× faster than this Stage-1 path |
+| Action chunk vs the fp32 reference (same observation, same initial noise) | PCC `left_arm` 0.99995 · `right_arm` 0.99990 · `right_hand` 0.99308 · `left_hand` / `waist` max\|d\| 0.059 / 0.013 (near-constant groups) · all 29 action dims 0.99923 (normalised valid slice 0.99922) |
+| Intermediate taps vs the fp32 reference (the 23 taps the megakernel path exposes, from pixels; DiT block internals are not observable inside the kernel) | 26 of 28 gated rows pass (29 of 31 rows incl. the determinism row and the two xfail info rows); 2 marginal misses on the DiT state-token row that no action reads (`action_decoder_out[k=1,2]` 0.999594 / 0.999598 vs gates 0.9997 / 0.99978), caused by bf16 error in the SigLIP2 tower; the Stage-1 card's `dit_out[k=1,2]` misses are now inside the kernel |
+| Megakernel vs the Stage-1 TTNN denoise, same device | `action_pred_normalized` PCC ≥ 0.999994 on every N1.6 reference sample |
+| Traced vs untraced, determinism | bit-identical (13 of 13 rows); 5 trace replays bit-identical; 2 CQs == 1 CQ |
+| Inference on p150a (warm, batch 1, median of 50) | **40.7 ms end-to-end** (p90 40.8) = host encode 1.4 + upload 1.4 + traces 37.8 (vision 10.8 · LLM 9.6 · adapter 2.2 · denoise 15.3 = 3.84 ms per step) + readback and decode 0.1 |
+| Served over HTTP by this image (warm, batch 1, 50 requests) | 39.6 ms device · 42.1 ms server-side incl. PNG decode (p90 42.5) |
+| Same policy on an RTX 5090 (same host, NVIDIA's unmodified `Gr00tPolicy`, eager PyTorch bf16 + flash-attn 2) | 63.1 ms end-to-end → p150a 1.55× faster; with `torch.compile` on the DiT the GPU reaches 33.1 ms → GPU 1.23× faster than this Stage-2 path |
 
 ### Caveats
 
-- Stage-1 path: TTNN ops in four Metal traces, DiT matmul weights bfp8, everything else bf16. The megakernel denoise (weight streaming measured at 464 GB/s bf16 / 414 GB/s bfp8 over the 32-block N1.6 weight set) is not in this image.
-- Fidelity margin: the served chunk matches the fp32 reference at PCC 0.99921 on the normalised valid slice (NVIDIA's own bf16 GPU run scores 0.99956), 2e-4 above the 0.999 boot self-check.
+- Stage-2 path: the `denoise` trace is one persistent `generic_op` on all 110 Tensix cores (8 readers stream the 32-block weight set once per step from a 0.98 GB bfp8 DRAM arena to 96 compute cores; distributed LayerNorm, per-head flash SDPA, Euler update and decoder inside the kernel). Against the ttnn denoise on the same tree (e2e 56.2 ms, denoise 30.4) it saves 27.6 % / 49.6 %, short of the planned 30 % e2e; `GR00T_DIT_BACKEND=ttnn` selects that path.
+- Fidelity margin: the served chunk matches the fp32 reference at PCC 0.99923 (physical actions; 0.99922 on the normalised valid slice, where NVIDIA's own bf16 GPU run scores 0.99956), 2.3e-4 above the 0.999 boot self-check.
+- Costs of the megakernel: 64 KiB less allocatable L1 per core (the kernel-config ring is 136,192 B, the compute-core binaries take ≈ 131.7 KB of it), the VL adapter's intermediates in DRAM (+0.2 ms adapter trace), and no DiT-internal taps (10 head taps + the backbone taps remain gated).
 - One embodiment and layout: GR1 (embodiment slot 20), one `ego_view_bg_crop_pad_res256_freq20` camera, instruction ≤ 24 tokens; anything else is refused with 400. Batch 1.
 - Noise is seeded (seed 42, as the deployed policy); pass `seed` or `noise` for another draw. Outputs are the base checkpoint's actions for the GR1 sim data config.
 - Not an OpenAI-compatible API; `GET /v1/models` is a stub.
-- First boot is cold (weight conversion + kernel JIT, 137 s); later boots take 13–21 s. Validated on tt-metal `668c2907575`, single p150a only. Details: [`SERVING.md`](SERVING.md), [`GPU_COMPARISON.md`](GPU_COMPARISON.md).
+- First boot is cold (weight conversion + kernel JIT incl. the megakernel, 58 s); later boots take 14 s. Validated on tt-metal `668c2907575`, single p150a only. Details: [`SERVING.md`](SERVING.md), [`GPU_COMPARISON.md`](GPU_COMPARISON.md).
 
 ### Licensing
 
@@ -110,5 +97,6 @@ The exact sources the image was built from — `code/` in this repo is byte-iden
 | component | built from |
 | --- | --- |
 | tt-metal | [`668c290757550588d0ce46b180c344a462a2aaf5`](https://github.com/tenstorrent/tt-metal/commit/668c290757550588d0ce46b180c344a462a2aaf5) |
-| `code/` digest | `7a456e1065da7159` (sha256, first 16 hex digits) |
-| built | 2026-09-14T16:25:52+00:00 by tt-model 0.1.0 |
+| `code/` digest | `8e6bd13a23654f91` (sha256, first 16 hex digits) |
+| built | 2026-09-17T20:40:28+00:00 by tt-model 0.1.0 |
+

code/models/experimental/gr00t/tests/tt/test_bfp8_emulation_cpu.py

@@ -0,0 +1,617 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""CPU emulation of ttnn's ``bfloat8_b`` / ``bfloat4_b`` weight rounding on the fp32 torch reference (WP-Q;
+IMPLEMENTATION_PLAN.md §1.2 "After WP-Q admission", §7.2; golden-tap-map-and-thresholds.md §6 "emulation method").
+
+No device, no ``ttnn`` import.  Two parts:
+
+* :func:`bfp_round` -- an exact torch re-implementation of tt-metal's host conversion
+  ``tt_metal/impl/data_format/blockfloat_common.cpp`` (``get_max_exp`` over the 16 consecutive elements of a tile
+  face row -> shared 8-bit exponent; ``convert_u32_to_bfp<Bfp8_b|Bfp4_b, truncate=false>``: hidden bit made
+  explicit, mantissa shifted right by the exponent difference, round-to-nearest-even to 7 (bfp8) / 3 (bfp4) bits,
+  saturated at ``2^bits - 1``, zero / denormal input -> +0).  In a ``[K, N]`` TILE tensor a face row is 16
+  consecutive elements along the last dim starting at column 0 (tile padding is zero and does not move the shared
+  exponent), so the emulation blocks the **last** dim; ``tests/tt/test_tt_dtype_policy.py`` checks it on the device
+  against ``ttnn.from_torch(x, bfloat8_b).to_torch()`` (1-mantissa-step tolerance, exactness recorded).
+* The class emulation: a :class:`RoundingCheckpoint` (a ``LazyCheckpoint`` whose ``get`` rounds the HF weights of
+  one dtype class exactly the way the device plan lays them out -- head padding, fused projections, meta-RoPE row
+  permutation of the LLM q/k, tile-interleaved gate/up: :func:`class_round_specs` derives the key set and the block
+  orientation from ``common.weights.enumerate_plan``) feeds ``reference.model.Gr00tReference`` with the canonical
+  golden inputs; every golden tap the reference produces is gated by ``tests.tt.harness.check_taps`` (multi-sample
+  gates, valid slices, per-dim action PCC) and the class -> min PCC / worst margin table is written through the
+  ``results`` fixture.  Activations, softmax, GELU, LayerNorm and accumulation stay fp32 (the device's bf16
+  activations and HiFi2 operand truncation are **not** emulated -- the device sweep of ``test_tt_dtype_policy.py``
+  is the authority; a class that fails here is strong evidence against admission).
+
+Run (CPU, ~1 min per (version, class))::
+
+    python -m pytest -q -rA -p no:cacheprovider models/experimental/gr00t/tests/tt/test_bfp8_emulation_cpu.py \\
+        --gr00t-version n15
+"""
+
+from __future__ import annotations
+
+import re
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from models.experimental.gr00t.common import weights as W
+from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
+from models.experimental.gr00t.common.configs import GR00T_VERSIONS, get_config
+from models.experimental.gr00t.common.golden import GoldenSet
+from models.experimental.gr00t.common.pcc import pcc
+from models.experimental.gr00t.tests.tt import harness
+from models.experimental.gr00t.tests.tt.conftest import skip_if_missing
+from models.experimental.gr00t.tt.policy import TTPolicy, plan_options
+
+torch.set_grad_enabled(False)
+
+BLOCK = 16  # elements sharing one exponent (a tile face row; blockfloat_common.cpp get_max_exp asserts size 16)
+MANTISSA_BITS: Dict[str, int] = {"bfp8_b": 7, "bfp4_b": 3, "bfp2_b": 1}
+#: Sweep classes -> plan categories (plan §1.2 rows 3-5 + the DiT row; "projector" = N1.6 mlp1 + N1.7 mergers).
+CLASS_CATEGORIES: Dict[str, Tuple[str, ...]] = {
+    "vision": ("vision",),
+    "projector": ("projector", "vision_merger"),
+    "llm": ("llm",),
+    "vl_self_attention": ("vl_self_attention",),
+    "dit": ("dit",),
+}
+#: Classes the plan expects to admit (§1.2 "After WP-Q admission"): the emulation must pass their gates.
+EXPECTED_ADMIT: Tuple[str, ...] = ("vision", "projector", "llm", "vl_self_attention", "dit")
+#: Emulation cases: name -> (classes, mantissa bits, expected to pass).  ``all_bfp8`` = every class at once (the
+#: plan's end state); ``dit_ff_bfp4`` = the plan §1.2 row-1 experiment (bfp4 on ff1/ff2 only; mb2 §B.8 FF-out
+#: 0.9986 and golden-tap-map §6 say "probably not").
+CASES: Dict[str, Tuple[Tuple[str, ...], int, bool]] = {
+    "vision": (("vision",), 7, True),
+    "projector": (("projector",), 7, True),
+    "llm": (("llm",), 7, True),
+    "vl_self_attention": (("vl_self_attention",), 7, True),
+    "dit": (("dit",), 7, True),
+    "all_bfp8": (EXPECTED_ADMIT, 7, True),
+    "dit_ff_bfp4": (("dit_ff",), 3, False),
+}
+_FF_KEY = re.compile(r"\.ff\.net\.(0\.proj|2)\.weight$")
+_PERHEAD_OUT = re.compile(r"^per-head pad (\d+)->(\d+) \(out\)")
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# 1. The rounding primitive (blockfloat_common.cpp)
+# --------------------------------------------------------------------------------------------------------------------
+def _blocked(x: torch.Tensor, dim: int) -> Tuple[torch.Tensor, Tuple[int, ...], int, int]:
+    """``x`` -> fp32 ``[-1, BLOCK]`` view along ``dim`` (zero padded to a multiple of BLOCK); returns the view, the
+    moved shape, the moved dim and the pad so :func:`_unblocked` can restore it."""
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"expected a torch.Tensor, got {type(x).__name__}")
+    if x.ndim == 0:
+        raise ValueError("bfp rounding needs at least one dim")
+    d = dim % x.ndim
+    xm = x.detach().to(torch.float32).movedim(d, -1)
+    n = xm.shape[-1]
+    pad = (-n) % BLOCK
+    if pad:
+        xm = F.pad(xm, (0, pad))
+    return xm.contiguous().reshape(-1, BLOCK), tuple(xm.shape), d, pad
+
+
+def _unblocked(blocks: torch.Tensor, moved_shape: Tuple[int, ...], d: int, pad: int) -> torch.Tensor:
+    out = blocks.reshape(moved_shape)
+    if pad:
+        out = out[..., : moved_shape[-1] - pad]
+    return out.movedim(-1, d).contiguous()
+
+
+def _decompose(blocks: torch.Tensor, mantissa_bits: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Per element of a ``[-1, BLOCK]`` fp32 view: the rounded ``mantissa_bits``-bit mantissa (hidden bit explicit),
+    the block's shared exponent (biased, as stored) and the sign -- ``convert_u32_to_bfp<_, false>`` step by step."""
+    if mantissa_bits not in MANTISSA_BITS.values():
+        raise ValueError(f"mantissa_bits must be one of {sorted(MANTISSA_BITS.values())}, got {mantissa_bits}")
+    bits = blocks.view(torch.int32)
+    exp = (bits >> 23) & 0xFF
+    sign = (bits >> 31) & 1
+    mant = bits & 0x7FFFFF
+    zero_or_denormal = exp == 0
+    shared = exp.amax(dim=-1, keepdim=True)  # get_max_exp over the 16 elements (is_exp_a = false for *_b)
+    mant_full = mant | (1 << 23)  # hidden one
+    shift = (shared - exp).clamp(min=0, max=31)
+    m = mant_full >> shift  # shift further down by the exponent difference
+    shift_bits = 24 - mantissa_bits  # MANTISSA_BFP_SHIFT
+    round_mask = (1 << shift_bits) - 1
+    tie = 1 << (shift_bits - 1)
+    round_value = m & round_mask
+    m = m >> shift_bits
+    guard = m & 1
+    round_up = (round_value > tie) | ((round_value == tie) & (guard == 1))
+    m = m + round_up.to(torch.int32)
+    m = torch.clamp(m, max=(1 << mantissa_bits) - 1)  # saturate (no carry into the exponent)
+    m = torch.where(zero_or_denormal, torch.zeros_like(m), m)
+    return m, shared.expand_as(m), sign
+
+
+def _compose(m: torch.Tensor, shared: torch.Tensor, sign: torch.Tensor, mantissa_bits: int) -> torch.Tensor:
+    """``convert_bfp_to_u32``: ``m x 2^(shared - 127 - (mantissa_bits - 1))`` with the sign (``+0`` for m == 0)."""
+    scale = torch.pow(2.0, (shared - 127 - (mantissa_bits - 1)).to(torch.float64))
+    val = (m.to(torch.float64) * scale).to(torch.float32)
+    neg = (sign == 1) & (m != 0)
+    return torch.where(neg, -val, val)
+
+
+def bfp_round(t: torch.Tensor, mantissa_bits: int = 7, *, dim: int = -1) -> torch.Tensor:
+    """Round ``t`` the way ``ttnn.from_torch(t, dtype=bfloat8_b | bfloat4_b, layout=TILE)`` stores it: blocks of 16
+    consecutive elements along ``dim`` (the last dim of the device tensor) share the maximum exponent, mantissas are
+    round-to-nearest-even to ``mantissa_bits`` (7 = bfp8_b, 3 = bfp4_b).  Returns fp32 in ``t``'s shape."""
+    blocks, moved_shape, d, pad = _blocked(t, dim)
+    m, shared, sign = _decompose(blocks, mantissa_bits)
+    return _unblocked(_compose(m, shared, sign, mantissa_bits), moved_shape, d, pad)
+
+
+def bfp8_round(t: torch.Tensor, *, dim: int = -1) -> torch.Tensor:
+    """ttnn ``bfloat8_b`` rounding (7-bit mantissa, 16-element shared exponent) along ``dim`` (default: last)."""
+    return bfp_round(t, 7, dim=dim)
+
+
+def bfp4_round(t: torch.Tensor, *, dim: int = -1) -> torch.Tensor:
+    """ttnn ``bfloat4_b`` rounding (3-bit mantissa, 16-element shared exponent) along ``dim`` (default: last)."""
+    return bfp_round(t, 3, dim=dim)
+
+
+def mantissa_step(t: torch.Tensor, mantissa_bits: int = 7, *, dim: int = -1) -> torch.Tensor:
+    """One mantissa step of every element's block: ``2^(shared_exp - 127 - (mantissa_bits - 1))`` (fp32, ``t``'s shape);
+    the tolerance of the device cross-check."""
+    blocks, moved_shape, d, pad = _blocked(t, dim)
+    _, shared, _ = _decompose(blocks, mantissa_bits)
+    step = torch.pow(2.0, (shared - 127 - (mantissa_bits - 1)).to(torch.float64)).to(torch.float32)
+    return _unblocked(step, moved_shape, d, pad)
+
+
+def round_linear_weight(
+    w_out_in: torch.Tensor,
+    mantissa_bits: int = 7,
+    *,
+    head_dim: Optional[int] = None,
+    row_perm: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Round an HF ``nn.Linear`` weight ``[out, in]`` exactly as its device plan tensor ``[in, out]`` (WL §4.1 "T")
+    is stored: the 16-element blocks run along **out** (dim 0).
+
+    * ``head_dim``: the plan zero-pads every head of a Q/K/V projection on the out axis (``pad_heads_out``: SigLIP
+      72 -> 96, DiT 48 -> 64) **before** transposing, so head ``h`` occupies its own padded column range and the
+      last block of a 72-wide head holds 8 real values + 8 zeros -- blocking ``[nh, head_dim, in]`` along dim 1 with
+      zero padding reproduces that exactly (zeros never move the shared exponent).
+    * ``row_perm``: the plan permutes the out rows of the LLM q/k projections into the Meta RoPE layout
+      (``reverse_permute``) before transposing; rounding on the permuted rows and un-permuting reproduces the device
+      blocks.
+    """
+    if w_out_in.ndim != 2:
+        raise ValueError(f"expected a 2-D [out, in] weight, got {tuple(w_out_in.shape)}")
+    w = w_out_in.to(torch.float32)
+    if row_perm is not None:
+        if head_dim is not None:
+            raise ValueError("row_perm and head_dim are exclusive")
+        if row_perm.shape != (w.shape[0],):
+            raise ValueError(f"row_perm {tuple(row_perm.shape)} does not index out={w.shape[0]}")
+        rounded_p = bfp_round(w[row_perm], mantissa_bits, dim=0)
+        out = torch.empty_like(w)
+        out[row_perm] = rounded_p
+        return out
+    if head_dim is not None:
+        n_out, n_in = w.shape
+        if n_out % head_dim:
+            raise ValueError(f"out={n_out} is not a multiple of head_dim={head_dim}")
+        per_head = bfp_round(w.reshape(n_out // head_dim, head_dim, n_in), mantissa_bits, dim=1)
+        return per_head.reshape(n_out, n_in)
+    return bfp_round(w, mantissa_bits, dim=0)
+
+
+def rope_row_permutation(n_out: int, n_heads: int) -> torch.Tensor:
+    """``perm`` with ``reverse_permute(w, n_heads) == w[perm]`` (rows of a ``[out, in]`` q/k projection)."""
+    idx = torch.arange(n_out, dtype=torch.float64).reshape(n_out, 1)
+    return W.reverse_permute(idx, n_heads).reshape(n_out).to(torch.long)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# 2. Class -> HF keys with their device block orientation
+# --------------------------------------------------------------------------------------------------------------------
+@dataclass(frozen=True)
+class RoundSpec:
+    """How one HF weight key is rounded: mantissa bits, per-head out blocking, meta-RoPE row permutation."""
+
+    key: str
+    plan_name: str
+    category: str
+    mantissa_bits: int
+    head_dim: Optional[int] = None
+    rope_heads: Optional[int] = None  # row_perm = rope_row_permutation(out, rope_heads)
+
+    def apply(self, w: torch.Tensor) -> torch.Tensor:
+        perm = None if self.rope_heads is None else rope_row_permutation(w.shape[0], self.rope_heads)
+        return round_linear_weight(w, self.mantissa_bits, head_dim=self.head_dim, row_perm=perm)
+
+
+def class_round_specs(version: str, classes: Tuple[str, ...], mantissa_bits: int = 7) -> Dict[str, RoundSpec]:
+    """``HF key -> RoundSpec`` for every bfp8 candidate of the plan categories behind ``classes`` (the canonical
+    embodiment's ``mixed`` plan under the default ``TTPolicy`` transform options, exactly what ``TTWeights`` uploads).
+
+    ``"dit_ff"`` is the FF-only subset of ``dit`` (``ff.net.0.proj`` / ``ff.net.2``) for the bfp4 experiment.
+    Raises on a class without candidates in this version only when the class is not a known no-op (N1.5 projector
+    ``fc1`` is ``big=False`` by plan §1.2; N1.6 has no VL-SA blocks).
+    """
+    if version not in GR00T_VERSIONS:
+        raise KeyError(f"unknown version {version!r}")
+    cfg = get_config(version)
+    cats: List[str] = []
+    ff_only = False
+    for c in classes:
+        if c == "dit_ff":
+            cats.append("dit")
+            ff_only = True
+        elif c in CLASS_CATEGORIES:
+            cats.extend(CLASS_CATEGORIES[c])
+        else:
+            raise KeyError(f"unknown class {c!r}; expected one of {sorted(CLASS_CATEGORIES)} or 'dit_ff'")
+    meta = W.enumerate_plan(version, None, "mixed", plan_options(TTPolicy(), version))
+    specs: Dict[str, RoundSpec] = {}
+    for name, entry in sorted(meta.items()):
+        if entry.category not in cats or not entry.bfp8_candidate:
+            continue
+        m = _PERHEAD_OUT.match(entry.transform)
+        head_dim = int(m.group(1)) if m else None
+        meta_rope = "meta rope rows" in entry.transform
+        for key in entry.source_keys:
+            if ff_only and not _FF_KEY.search(key):
+                continue
+            if key in specs:
+                raise KeyError(f"{key} feeds two plan tensors ({specs[key].plan_name}, {name})")
+            rope_heads: Optional[int] = None
+            if meta_rope:
+                if key.endswith("q_proj.weight"):
+                    rope_heads = int(cfg.llm.n_heads)
+                elif key.endswith("k_proj.weight"):
+                    rope_heads = int(cfg.llm.n_kv_heads)
+                elif not key.endswith("v_proj.weight"):
+                    raise KeyError(f"{name}: unexpected meta-rope source key {key}")
+            specs[key] = RoundSpec(key, name, entry.category, mantissa_bits, head_dim, rope_heads)
+    return specs
+
+
+class RoundingCheckpoint(LazyCheckpoint):
+    """``LazyCheckpoint`` whose ``get`` returns the weights of ``specs`` rounded (golden-tap-map §6 method, applied
+    through the reference's own weight reads so no module is re-implemented).  ``hits`` counts the reads per key so
+    the caller can assert that every candidate of the class was actually consumed."""
+
+    def __init__(self, version: str, specs: Mapping[str, RoundSpec]):
+        super().__init__(version)
+        self.specs: Dict[str, RoundSpec] = dict(specs)
+        self.hits: Dict[str, int] = {}
+        self.round_s = 0.0
+        missing = [k for k in self.specs if k not in self]
+        if missing:
+            raise KeyError(f"{version}: {len(missing)} spec keys are not in the checkpoint, e.g. {missing[:3]}")
+
+    def get(self, key: str, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        t = super().get(key, dtype)
+        spec = self.specs.get(key)
+        if spec is None:
+            return t
+        t0 = time.perf_counter()
+        r = spec.apply(t.to(torch.float32))
+        self.round_s += time.perf_counter() - t0
+        self.hits[key] = self.hits.get(key, 0) + 1
+        return r.to(t.dtype)
+
+    def slot(self, key: str, e: int, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        if key in self.specs:
+            raise NotImplementedError(f"{key}: slot reads of a rounded key are not part of the emulation")
+        return super().slot(key, e, dtype)
+
+    def rows(self, key: str, lo: int, hi: int, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        if key in self.specs:
+            raise NotImplementedError(f"{key}: row reads of a rounded key are not part of the emulation")
+        return super().rows(key, lo, hi, dtype)
+
+    def unread(self) -> List[str]:
+        return sorted(k for k in self.specs if k not in self.hits)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# 3. Reference forward -> golden-keyed candidates
+# --------------------------------------------------------------------------------------------------------------------
+def reference_candidates(
+    out: Any, gs: GoldenSet, ref: Any, obs: Any, exclude: Tuple[str, ...]
+) -> Dict[str, torch.Tensor]:
+    """Every golden tap instance of ``gs`` the reference produced (``reference.model.reference_tap_tensor``), keyed by
+    ``GoldenSet.threshold_key`` in the golden's shape, plus ``unnormalized_actions.<group>`` through ``ref.decode``;
+    taps in ``exclude`` (host inputs / reference-only) are skipped."""
+    from models.experimental.gr00t.reference.model import reconcile_shape, reference_tap_tensor
+
+    cand: Dict[str, torch.Tensor] = {}
+    for tap, kw in gs.available("fp32"):
+        if tap in exclude or tap == "unnormalized_actions":
+            continue
+        t = reference_tap_tensor(out, gs, tap, kw)
+        if t is None:
+            continue
+        cand[gs.threshold_key(tap, **kw)] = reconcile_shape(t.detach().to(torch.float32), gs.load(tap, **kw))
+    for group, t in ref.decode(out.action_pred_normalized, obs).items():
+        cand[f"unnormalized_actions.{group}"] = t.detach().to(torch.float32)
+    return cand
+
+
+def run_reference(
+    version: str, gs: GoldenSet, obs: Any, mi: Any, ck: Optional[LazyCheckpoint]
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, float]]:
+    """Build the reference on ``ck`` (None = unrounded fp32), run it on the canonical inputs with the golden noise
+    and return the golden-keyed candidates and the timings."""
+    from models.experimental.gr00t.reference.model import Gr00tReference
+    from models.experimental.gr00t.tt.model import HOST_OR_UNPRODUCED_TAPS
+
+    t0 = time.perf_counter()
+    ref = Gr00tReference(version, checkpoint=ck)
+    t1 = time.perf_counter()
+    out = ref.forward_from_inputs(mi, gs.load("initial_noise"), return_taps=True)
+    t2 = time.perf_counter()
+    cand = reference_candidates(out, gs, ref, obs, HOST_OR_UNPRODUCED_TAPS)
+    timing = {"build_s": t1 - t0, "forward_s": t2 - t1, **{k + "_s": v for k, v in out.timing.items()}}
+    return cand, timing
+
+
+_BASELINE: Dict[str, Tuple[Dict[str, torch.Tensor], harness.Report]] = {}
+
+
+def baseline(version: str, gs: GoldenSet, obs: Any, mi: Any) -> Tuple[Dict[str, torch.Tensor], harness.Report]:
+    """The unrounded fp32 reference's candidates and report (computed once per version, cached in the module)."""
+    if version not in _BASELINE:
+        cand, _ = run_reference(version, gs, obs, mi, None)
+        report = harness.check_taps(version, cand, gs, multi=True)
+        _BASELINE[version] = (cand, report)
+    return _BASELINE[version]
+
+
+def pcc_table(report: harness.Report) -> Dict[str, float]:
+    return {r["key"]: float(r["pcc"]) for r in report.rows if r.get("pcc") is not None}
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# 4. Unit tests of the primitive (no checkpoint needed)
+# --------------------------------------------------------------------------------------------------------------------
+def test_cpu_bfp_round_hand_values() -> None:
+    """Hand-checked blocks: the shared exponent is the block max; small values lose mantissa bits; RNE ties."""
+    # one block: max 1.0 (exp 127); 7-bit mantissa -> resolution 2^-6 = 0.015625 for every element of the block
+    x = torch.zeros(16)
+    x[0] = 1.0
+    x[1] = 0.5 + 0.015625 / 2  # 0.5078125: exactly between 0.5 and 0.515625 -> tie -> even (0.5 = mantissa 32)
+    x[2] = 0.5 + 0.015625 * 1.5  # -> rounds up to 0.53125 (mantissa 34, tie -> even)
+    x[3] = 0.001  # < one step -> 0
+    x[4] = -1.0
+    x[5] = 1.0 + 0.015625 * 0.75  # rounds up to 1.015625
+    r = bfp8_round(x)
+    assert r[0] == 1.0 and r[4] == -1.0
+    assert r[1] == 0.5, r[1]
+    assert r[2] == 0.53125, r[2]
+    assert r[3] == 0.0 and torch.signbit(r[3]).item() is False
+    assert r[5] == 1.015625, r[5]
+    # saturation: the max element with an all-ones mantissa that would round up stays at 127/64
+    y = torch.zeros(16)
+    y[0] = 1.0 + 63 / 64 + 0.9 / 64
+    assert bfp8_round(y)[0] == 1.0 + 63 / 64
+    # bfp4: 3-bit mantissa -> resolution 2^-2 of the block max
+    z = torch.zeros(16)
+    z[0], z[1], z[2] = 1.0, 0.6, 0.1
+    r4 = bfp4_round(z)
+    assert r4[0] == 1.0 and r4[1] == 0.5 and r4[2] == 0.0, r4[:3]
+    # 1 mantissa step
+    assert torch.equal(mantissa_step(x)[:16], torch.full((16,), 2.0**-6))
+    assert torch.equal(mantissa_step(z, 3)[:16], torch.full((16,), 2.0**-2))
+
+
+def test_cpu_bfp_round_properties() -> None:
+    """Idempotence, block independence (dim / padding), the golden-tap-map §6 self-test magnitude (N(0,3) rel-L2
+    6.4e-3 for bfp8, 0.10 for bfp4; bf16 1.7e-3) and exactness on values already on the bfp8 grid."""
+    g = torch.Generator().manual_seed(0)
+    x = torch.randn(64, 200, generator=g) * 3.0
+    r = bfp8_round(x)
+    assert torch.equal(bfp8_round(r), r), "not idempotent"
+    rel = float((r - x).norm() / x.norm())
+    assert 5e-3 < rel < 8e-3, rel  # §6: 6.4e-3
+    rel4 = float((bfp4_round(x) - x).norm() / x.norm())
+    assert 0.08 < rel4 < 0.13, rel4  # §6: 0.10
+    rel_bf16 = float((x.bfloat16().float() - x).norm() / x.norm())
+    assert 1.4e-3 < rel_bf16 < 2.0e-3, rel_bf16
+    # error never exceeds half a mantissa step of the block, except the saturated block maxima (mantissa 127
+    # cannot carry into the exponent: convert_u32_to_bfp clamps) which may sit one full step low
+    err = (r - x).abs()
+    step = mantissa_step(x)
+    assert bool((err <= step + 1e-12).all())
+    assert float((err > step / 2 + 1e-12).float().mean()) < 0.01
+    # blocks are along the chosen dim only: rounding the transpose along dim 0 equals rounding along dim 1
+    assert torch.equal(bfp8_round(x.t().contiguous(), dim=0), bfp8_round(x).t())
+    # ragged last block (200 = 12*16 + 8) equals explicit zero padding to 208
+    xp = F.pad(x, (0, 8))
+    assert torch.equal(bfp8_round(xp)[:, :200], r)
+    # a 4-D tensor blocks its last dim
+    x4 = torch.randn(1, 32, 100, 64, generator=g)
+    assert torch.equal(bfp8_round(x4), bfp8_round(x4.reshape(-1, 64)).reshape(x4.shape))
+    # values with <= 7 significant bits and a common exponent are exact (all elements in [1, 2) on the 1/64 grid)
+    y = (torch.randint(64, 128, (16, 16), generator=g).to(torch.float32)) / 64.0
+    assert torch.equal(bfp8_round(y), y)
+
+
+def test_cpu_round_linear_weight_layouts() -> None:
+    """The head-padded / row-permuted paths equal rounding of the plan's own transformed tensor along its last dim."""
+    g = torch.Generator().manual_seed(1)
+    nh, hd, hdp, n_in = 4, 72, 96, 40
+    w = torch.randn(nh * hd, n_in, generator=g)
+    dev = W.to_in_out(W.pad_heads_out(w, nh, hd, hdp))  # [in, nh*hdp]: what TTWeights uploads
+    dev_r = bfp8_round(dev)  # device blocks along the last dim
+    back = dev_r.t().reshape(nh, hdp, n_in)[:, :hd].reshape(nh * hd, n_in)
+    assert torch.equal(round_linear_weight(w, head_dim=hd), back)
+    assert not torch.equal(round_linear_weight(w), back), "72-wide heads straddle 16-blocks: plain blocking differs"
+    # meta-RoPE rows (LLM q/k): plan = to_in_out(reverse_permute(w)) -> blocks along out of the permuted rows
+    wq = torch.randn(8 * 128, 64, generator=g)
+    perm = rope_row_permutation(wq.shape[0], 8)
+    assert torch.equal(wq[perm], W.reverse_permute(wq, 8))
+    dev_q = bfp8_round(W.to_in_out(W.reverse_permute(wq, 8)))
+    back_q = torch.empty_like(wq)
+    back_q[perm] = dev_q.t()
+    assert torch.equal(round_linear_weight(wq, row_perm=perm), back_q)
+    # tile-interleaved gate/up: each 32-column tile belongs to one matrix -> plain per-matrix blocking is exact
+    wg, wu = torch.randn(96, 48, generator=g), torch.randn(96, 48, generator=g)
+    fused = bfp8_round(W.fuse_gate_up(wg, wu, "tile_interleaved"))
+    plain = bfp8_round(W.fuse_gate_up(round_linear_weight(wg), round_linear_weight(wu), "tile_interleaved"))
+    assert torch.equal(fused, plain)
+    # fused projections concatenate along out: per-matrix blocking is exact when out % 16 == 0
+    fq = bfp8_round(W.fuse_projections([wg, wu])[0])
+    assert torch.equal(fq, W.fuse_projections([round_linear_weight(wg), round_linear_weight(wu)])[0])
+
+
+@pytest.mark.parametrize("version", list(GR00T_VERSIONS))
+def test_cpu_class_round_specs_cover_candidates(version: str) -> None:
+    """Every bfp8 candidate of every class maps to HF keys with the right orientation; no key is claimed twice; the
+    known no-ops (N1.5 projector, N1.6 VL-SA) are empty; the meta-RoPE keys carry the head counts."""
+    cfg = get_config(version)
+    meta = W.enumerate_plan(version, None, "mixed", plan_options(TTPolicy(), version))
+    all_specs: Dict[str, RoundSpec] = {}
+    for cls in CLASS_CATEGORIES:
+        specs = class_round_specs(version, (cls,))
+        cands = [e for e in meta.values() if e.category in CLASS_CATEGORIES[cls] and e.bfp8_candidate]
+        assert {s.plan_name for s in specs.values()} == {e.name for e in cands}, cls
+        assert not (set(specs) & set(all_specs)), cls
+        all_specs.update(specs)
+        if cls == "projector" and version == "n15":
+            assert specs == {}, "N1.5 projector.fc1 is big=False (plan §1.2 row 2)"
+        if cls == "vl_self_attention" and version == "n16":
+            assert specs == {}, "N1.6 has no VL-SA blocks"
+    union = class_round_specs(version, EXPECTED_ADMIT)
+    assert set(union) == set(all_specs)
+    for spec in union.values():
+        if spec.rope_heads is not None:
+            assert spec.category == "llm" and spec.rope_heads in (cfg.llm.n_heads, cfg.llm.n_kv_heads)
+        if spec.head_dim is not None:
+            assert spec.category in ("vision", "dit")
+            assert spec.head_dim == (cfg.vision.head_dim if spec.category == "vision" else cfg.dit.head_dim)
+    if cfg.vision.family in ("siglip_v1", "siglip2_naflex"):
+        assert any(s.head_dim == 72 for s in union.values() if s.category == "vision")
+    ff = class_round_specs(version, ("dit_ff",), 3)
+    assert ff and all(_FF_KEY.search(k) and s.mantissa_bits == 3 for k, s in ff.items())
+    assert len(ff) == 2 * cfg.dit.n_blocks
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# 5. The emulation on the goldens (checkpoint + golden required; skipped otherwise)
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def canonical(version: str) -> Tuple[GoldenSet, Any, Any]:
+    skip_if_missing(version)
+    from models.experimental.gr00t.common.preprocessing import encode
+    from models.experimental.gr00t.reference.model import load_golden_sample
+
+    gs, obs = load_golden_sample(version)
+    return gs, obs, encode(version, obs)
+
+
+def test_cpu_emulation_baseline_is_fp32_exact(
+    version: str, canonical: Tuple[GoldenSet, Any, Any], results: Any
+) -> None:
+    """The unrounded reference passes every gate (foundation-f4: fp32-exact) -- the anchor of the class deltas."""
+    gs, obs, mi = canonical
+    cand, report = baseline(version, gs, obs, mi)
+    results.add_report(report)
+    results.set_meta(case="baseline_fp32", n_rows=len(report.rows), min_pcc=report.min_pcc())
+    print(report.summary(failures_only=True) if not report.ok else f"[{version}] baseline {len(report.rows)} rows ok")
+    report.assert_ok()
+
+
+@pytest.mark.parametrize("case", list(CASES))
+def test_cpu_bfp_emulation_class(version: str, case: str, canonical: Tuple[GoldenSet, Any, Any], results: Any) -> None:
+    """One dtype class (or the union / the bfp4 FF experiment) rounded on the fp32 reference; every golden tap gated
+    with the multi-sample gates; the table row (min PCC, worst margin, failures, deltas vs the fp32 baseline) goes to
+    the results JSON.  A class the plan expects to admit must pass; the bfp4 experiment is recorded and xfails."""
+    gs, obs, mi = canonical
+    classes, mbits, expect_pass = CASES[case]
+    specs = class_round_specs(version, classes, mbits)
+    if not specs:
+        results.set_meta(case=case, classes=list(classes), n_keys=0, note="no bfp8 candidates in this version")
+        pytest.skip(f"{version}: class {case} has no bfp8 candidates (plan §1.2)")
+    base_cand, base_report = baseline(version, gs, obs, mi)
+    ck = RoundingCheckpoint(version, specs)
+    cand, timing = run_reference(version, gs, obs, mi, ck)
+    unread = ck.unread()
+    assert not unread, f"{version}/{case}: {len(unread)} candidate keys were never read by the reference: {unread[:5]}"
+    assert set(cand) == set(base_cand), sorted(set(cand) ^ set(base_cand))
+    report = harness.check_taps(version, cand, gs, results=results, multi=True)
+    base_pcc = pcc_table(base_report)
+    this_pcc = pcc_table(report)
+    deltas = {k: this_pcc[k] - base_pcc[k] for k in this_pcc if k in base_pcc}
+    worst_key, worst_margin = report.worst_margin()
+    min_key, min_pcc = report.min_pcc()
+    # per-tap agreement with the *baseline* candidates (how much the class moved the tensors)
+    vs_base = {k: pcc(cand[k], base_cand[k]) for k in ("backbone_features", "action_pred_normalized") if k in cand}
+    results.set_meta(
+        case=case,
+        classes=list(classes),
+        mantissa_bits=mbits,
+        n_keys=len(specs),
+        plan_tensors=sorted({s.plan_name for s in specs.values()}),
+        n_rows=len(report.rows),
+        n_fail=len(report.failures()),
+        min_pcc={"key": min_key, "pcc": min_pcc},
+        worst_margin={"key": worst_key, "margin": worst_margin},
+        failing_keys=[r["key"] for r in report.failures()],
+        pcc_delta_vs_fp32=deltas,
+        pcc_vs_fp32_candidates=vs_base,
+        round_s=ck.round_s,
+    )
+    results.set_timing(**timing, round_s=ck.round_s)
+    print(
+        f"\n[{version}] {case}: {len(specs)} keys rounded ({mbits}-bit) in {ck.round_s:.1f} s; forward "
+        f"{timing['forward_s']:.1f} s; rows {len(report.rows)} fail {len(report.failures())}; "
+        f"min PCC {min_pcc:.6f} ({min_key}); worst margin {worst_margin:+.6f} ({worst_key})"
+    )
+    for key in (
+        "vit_block_last",
+        "projector_out",
+        "backbone_features",
+        "vlln_out",
+        "vl_self_attention_out",
+        "dit_block_last_out[k=3]",
+        "dit_out[k=3]",
+        "pred_velocity[k=3]",
+        "action_pred_normalized",
+        "action_pred_valid",
+    ):
+        if key in this_pcc:
+            print(f"    {key:28s} {this_pcc[key]:.6f} (delta {deltas[key]:+.2e}, gate {report.row(key)['gate_pcc']})")
+    if report.failures():
+        print(report.summary(failures_only=True))
+    if expect_pass:
+        report.assert_ok()
+    elif not report.ok:
+        pytest.xfail(
+            f"{version}/{case}: experiment fails {len(report.failures())} gate(s): "
+            + ", ".join(f"{r['key']} {r['pcc']:.6f} < {r['gate_pcc']}" for r in report.failures()[:6])
+        )
+
+
+__all__ = [
+    "BLOCK",
+    "CASES",
+    "CLASS_CATEGORIES",
+    "EXPECTED_ADMIT",
+    "MANTISSA_BITS",
+    "RoundSpec",
+    "RoundingCheckpoint",
+    "bfp4_round",
+    "bfp8_round",
+    "bfp_round",
+    "class_round_specs",
+    "mantissa_step",
+    "reference_candidates",
+    "rope_row_permutation",
+    "round_linear_weight",
+    "run_reference",
+]

code/models/experimental/gr00t/tests/tt/test_mk_e2e.py

@@ -0,0 +1,591 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K5 device test: ``Gr00tTT`` with ``TTPolicy(dit_backend="megakernel")`` against the Stage-1 ttnn backend on the
+same device, the same weights tier and the same golden sample(s) (IMPLEMENTATION_PLAN.md §5.7 K5 row, §7.4, §7.6;
+``tests/tt/results/mk_k4_summary.md`` §4.2.2 for the kernel-level oracle protocol this extends to the e2e level).
+
+Per version the module builds the two models **sequentially** (one device, opened with the megakernel's reduced
+``worker_l1_size`` so both backends run under the same configuration; the ttnn model is built with the DRAM adapter
+intermediates that configuration needs, ``adapter_mem="DRAM"``, values unchanged), runs every selected sample
+through ``run_untraced(return_taps=True)`` and the traced path (warm -> capture -> N replays) and releases the model;
+the tests then compare the recorded host results:
+
+* ``test_mk_e2e_backends_agree`` -- ``action_pred_normalized`` (untraced and traced) and every per-step
+  ``pred_velocity[k]`` of the megakernel vs the ttnn backend: PCC >= 0.9999 (the K4 oracle gate, plan §5.7; the
+  ``action_decoder_out[k]`` rows -- 64 rows including the state / future tokens nothing downstream reads -- are
+  recorded as information) and max |d| recorded; the megakernel's taps pass every ``gates_multi.json`` gate the ttnn
+  backend passes on that sample (``harness.check_taps``, the decoded ``unnormalized_actions.<group>`` included) and
+  it is at least as close to the fp32 golden on ``action_pred_normalized`` (K4's closeness rule, slack 2e-5).  On a
+  sample where the ttnn backend itself misses the ``action_pred_normalized`` golden gate (the Stage-1 divergent
+  multi samples, STATUS.md) the two bf16 pipelines are not each other's oracle: the oracle rows are informational
+  and the closeness rule decides.  :data:`MK_KNOWN_MARGINAL` lists the rows the megakernel misses while the ttnn
+  backend passes (``mk_k5_summary.md`` §4.1 / §7 item 4) with a PCC floor: such a run is an **xfail**, never a pass;
+* ``test_mk_e2e_trace_protocol`` -- megakernel traced == untraced bit-equal (raw device buffer and the kernel's
+  per-step ``velocity_out`` taps; K4 review NB 13), 10 traced replays bit-identical (plan D12: megakernel 10x),
+  ``backbone_features`` / ``K_0`` / ``V_0`` bit-equal traced vs untraced, the ttnn backend's 5 replays likewise;
+* ``test_mk_e2e_denoise_trace_ms`` -- informational: the split-pass trace times of both backends (median of a few
+  synchronised replays under alloc tracking -- **not** the latency source; ``benchmarks/bench_e2e.py --dit-backend``
+  is) and the ratio of the denoise trace.
+
+Run (one pytest process per version, through the device lock; alloc tracking is asserted, so set it in the shell)::
+
+    TT_METAL_TRACE_ALLOC_TRACKING=1 DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
+        python -m pytest -q -rA --timeout=3600 -p no:cacheprovider \\
+        models/experimental/gr00t/tests/tt/test_mk_e2e.py --gr00t-version n16 [--sample multi] \\
+        [GR00T_MK_ARENA_DTYPE=bf16|bfp8_b in the environment selects the megakernel arena; default: the policy table]
+"""
+
+from __future__ import annotations
+
+import os
+import statistics
+import time
+from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Set, Tuple
+
+import pytest
+import torch
+
+from models.experimental.gr00t.common.golden import GoldenSet
+from models.experimental.gr00t.common.pcc import max_abs, pcc
+from models.experimental.gr00t.tests.tt import harness
+from models.experimental.gr00t.tests.tt.conftest import (
+    DEFAULT_L1_SMALL_SIZE,
+    DEFAULT_TRACE_REGION_SIZE,
+    DEFAULT_TT_METAL_CACHE,
+    OPT_SAMPLE,
+    SAMPLE_SETS,
+    samples_for,
+    skip_if_missing,
+)
+from models.experimental.gr00t.tests.tt.test_tt_model_e2e import (
+    decoded_groups,
+    golden_keyed,
+    required_gated_keys,
+    split_info,
+)
+from models.experimental.gr00t.tt import model as M
+from models.experimental.gr00t.tt import tracer as T
+from models.experimental.gr00t.tt.layout import StaticShapeError
+from models.experimental.gr00t.tt.policy import TTPolicy
+
+torch.set_grad_enabled(False)
+
+BACKENDS: Tuple[str, ...] = ("ttnn", "megakernel")
+GATE_ORACLE = 0.9999  # plan §5.7 K4/K5: PCC vs the Stage-1 device path
+GOLDEN_SLACK = 2e-5  # mk_k4_summary.md §4.3: PCC(mk, golden) >= PCC(S1, golden) - slack
+#: rows the megakernel misses while the ttnn backend passes them (WP-K5 device facts, mk_k5_summary.md §4.1 / §7 item 4;
+#: both arenas, marginal gates): (version, sample) -> {key: (field, PCC floor)}.  ``field`` names the quantity of the
+#: ``harness.check_taps`` row that actually fails its gate -- ``"pcc"`` (the full-tensor PCC) or
+#: ``"per_dim.min_gated_pcc"`` (the worst gated dimension of a per-dimension gate) -- and the run is an **xfail**
+#: only while that quantity stays above the floor (WP-K5 review correction 2: the N1.7 row fails on dim 9 while its
+#: full-tensor PCC 0.99983 never failed, so a floor on ``pcc`` would let dim 9 sink to 0.99 unnoticed)
+MK_KNOWN_MARGINAL: Dict[Tuple[str, str], Dict[str, Tuple[str, float]]] = {
+    # 0.999872 (bf16) / 0.999875 (bfp8) vs gate 0.9999 (ttnn 0.999915); max|d| 0.034 / 0.029 within the 0.038 tolerance
+    ("n15", "t3_s300_seed2"): {"unnormalized_actions.right_arm": ("pcc", 0.9998)},
+    # per-dim gate 0.996 on dim 9: 0.99589 (bf16) / 0.99585 (bfp8) vs ttnn 0.99690; full-tensor PCC 0.99983 passes
+    ("n17", "t1_s120"): {"pred_velocity[k=2]": ("per_dim.min_gated_pcc", 0.995)},
+}
+
+
+def marginal_value(row: Mapping[str, Any], field: str) -> Optional[float]:
+    """The quantity of a ``check_taps`` row that :data:`MK_KNOWN_MARGINAL` puts a floor on: ``"pcc"`` or the dotted
+    ``"per_dim.min_gated_pcc"`` (``None`` when the row lacks it -- then the row is *not* a known marginal)."""
+    cur: Any = row
+    for part in field.split("."):
+        if not isinstance(cur, Mapping) or part not in cur:
+            return None
+        cur = cur[part]
+    return None if cur is None else float(cur)
+
+
+def known_marginal_regressions(
+    version: str, sample: str, regressions: Sequence[str], mk_rows: Mapping[str, Mapping[str, Any]]
+) -> Tuple[List[str], Dict[str, Dict[str, Any]]]:
+    """Split ``regressions`` (gates the ttnn backend passes and the megakernel fails) into the :data:`MK_KNOWN_MARGINAL`
+    rows that are above their floor (-> xfail) and the rest (-> fail); returns ``(known, facts)`` where ``facts``
+    records field / value / floor per known row."""
+    marginal = MK_KNOWN_MARGINAL.get((version, sample), {})
+    known: List[str] = []
+    facts: Dict[str, Dict[str, Any]] = {}
+    for key in regressions:
+        if key not in marginal:
+            continue
+        field, floor = marginal[key]
+        value = marginal_value(mk_rows.get(key, {}), field)
+        if value is not None and value >= floor:
+            known.append(key)
+            facts[key] = {"field": field, "value": value, "floor": floor}
+    return sorted(known), facts
+
+
+#: oracle-gated per-step taps (the action rows); action_decoder_out[k] (all 64 rows) is informational
+ORACLE_STEP_TAPS: Tuple[str, ...] = ("pred_velocity",)
+#: a sample whose ttnn-backend ``action_pred_normalized`` is below this PCC vs its golden is not a reference for either
+#: backend (the N1.7 extras: `droid_rtc` 0.849 for both); its rows are recorded and the test is skipped
+CLOSENESS_MIN_REFERENCE = 0.99
+N_REPLAYS: Dict[str, int] = {"ttnn": 5, "megakernel": 10}  # plan §7.4: 5 traced; megakernel 10x
+N_SPLIT_REPS = 5
+ACTION_KEYS: Tuple[str, ...] = ("action_pred_normalized", "action_pred_valid")
+#: the keys both backends produce (``tt.action_head.head_tap_names("megakernel", ...)``)
+COMMON_STEP_TAPS: Tuple[str, ...] = ("action_decoder_out", "pred_velocity")
+
+_RESULTS: Dict[str, Dict[str, Dict[str, Dict[str, Any]]]] = {}  # version -> backend -> sample -> record
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# CPU tier: the known-marginal resolver (WP-K5 review correction 2)
+# --------------------------------------------------------------------------------------------------------------------
+def test_cpu_known_marginal_floor_on_failing_quantity() -> None:
+    """The N1.7 floor sits on ``per_dim.min_gated_pcc`` (the quantity that fails), the N1.5 floor on ``pcc``; a row
+    below its floor, a row without the field, or an unlisted key never counts as a known marginal."""
+    n17_row = {
+        "key": "pred_velocity[k=2]",
+        "pcc": 0.999826,
+        "gate_pcc": 0.9997,
+        "ok": False,
+        "per_dim": {"gate_pcc": 0.996, "min_gated_pcc": 0.995846, "failed_dims": [9], "ok": False},
+    }
+    known, facts = known_marginal_regressions("n17", "t1_s120", ["pred_velocity[k=2]"], {n17_row["key"]: n17_row})
+    assert known == ["pred_velocity[k=2]"] and facts["pred_velocity[k=2]"]["field"] == "per_dim.min_gated_pcc"
+    assert facts["pred_velocity[k=2]"]["value"] == 0.995846 and facts["pred_velocity[k=2]"]["floor"] == 0.995
+    sunk = dict(n17_row, per_dim=dict(n17_row["per_dim"], min_gated_pcc=0.99))  # full-tensor pcc unchanged
+    assert known_marginal_regressions("n17", "t1_s120", ["pred_velocity[k=2]"], {sunk["key"]: sunk})[0] == []
+    no_field = {k: v for k, v in n17_row.items() if k != "per_dim"}
+    assert (
+        known_marginal_regressions("n17", "t1_s120", ["pred_velocity[k=2]"], {"pred_velocity[k=2]": no_field})[0] == []
+    )
+    n15_row = {"key": "unnormalized_actions.right_arm", "pcc": 0.999875, "gate_pcc": 0.9999, "ok": False}
+    known, facts = known_marginal_regressions(
+        "n15", "t3_s300_seed2", ["unnormalized_actions.right_arm", "action_pred_normalized"], {n15_row["key"]: n15_row}
+    )
+    assert known == ["unnormalized_actions.right_arm"] and facts[known[0]]["field"] == "pcc"
+    assert (
+        known_marginal_regressions("n15", "canonical", ["unnormalized_actions.right_arm"], {n15_row["key"]: n15_row})[0]
+        == []
+    )
+    for (v, smp), rows in MK_KNOWN_MARGINAL.items():
+        for key, (field, floor) in rows.items():
+            assert field in ("pcc", "per_dim.min_gated_pcc") and 0.99 <= floor < 1.0, (v, smp, key)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# fixtures
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def mk_policy(policy_name: str, trace_layout: str) -> TTPolicy:
+    """The megakernel policy of this run (dtype policy / trace layout from the options, arena from the environment)."""
+    pol = TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout, dit_backend="megakernel").with_env_overrides()
+    if pol.dit_backend != "megakernel":
+        raise pytest.UsageError("test_mk_e2e always builds the megakernel backend; unset GR00T_DIT_BACKEND")
+    return pol
+
+
+@pytest.fixture(scope="module")
+def mk_device(request: pytest.FixtureRequest, mk_policy: TTPolicy) -> Iterator[Any]:
+    """One p150a opened with the K4 ``worker_l1_size`` (both backends run under it; module docstring)."""
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+    device = M.open_model_device(
+        mk_policy, trace_region_size=DEFAULT_TRACE_REGION_SIZE, l1_small_size=DEFAULT_L1_SMALL_SIZE
+    )
+    facts = harness.device_facts(device)
+    facts["worker_l1_size"] = M.device_worker_l1_size(device)
+    request.config._gr00t_device_facts = facts  # type: ignore[attr-defined]
+    try:
+        yield device
+    finally:
+        _RESULTS.clear()
+        close_gr00t_device(device)
+
+
+def _selected_samples(request: pytest.FixtureRequest, version: str) -> Tuple[str, ...]:
+    opt = str(request.config.getoption(OPT_SAMPLE))
+    names = samples_for(version, opt) if opt in SAMPLE_SETS else (opt,)
+    return tuple(names)
+
+
+@pytest.fixture(scope="module")
+def backend_results(
+    request: pytest.FixtureRequest, version: str, mk_device: Any, mk_policy: TTPolicy
+) -> Dict[str, Dict[str, Dict[str, Any]]]:
+    """``backend -> sample -> record`` of every selected sample, both backends built once per version (module doc)."""
+    if version not in _RESULTS:
+        if not T.alloc_tracking_requested():
+            pytest.fail(
+                f"{T.ALLOC_TRACKING_ENV}=1 must be set in the shell before pytest starts (plan §4 / static-shape §10 "
+                "item 9): this test asserts traced == untraced"
+            )
+        samples = _selected_samples(request, version)
+        for s in samples:
+            skip_if_missing(version, s)
+        # the Stage-1 reference must be the ttnn backend explicitly: the TTPolicy default is the megakernel since
+        # 2026-09-18 (a defaulted policy here would compare the megakernel with itself)
+        ttnn_policy = TTPolicy(
+            dtype_policy=mk_policy.dtype_policy, trace_layout=mk_policy.trace_layout, dit_backend="ttnn"
+        )
+        _RESULTS[version] = {
+            "ttnn": _collect(version, samples, ttnn_policy, mk_device),
+            "megakernel": _collect(version, samples, mk_policy, mk_device),
+        }
+        if not (
+            _RESULTS[version]["ttnn"]["_model"]["recipe"]["dit_backend"] == "ttnn"
+            and _RESULTS[version]["megakernel"]["_model"]["recipe"]["dit_backend"] == "megakernel"
+        ):
+            raise AssertionError("the two collections must be the ttnn and the megakernel backend")
+    return _RESULTS[version]
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# collection (one model per backend, every selected sample, then release)
+# --------------------------------------------------------------------------------------------------------------------
+def _collect(version: str, samples: Tuple[str, ...], policy: TTPolicy, device: Any) -> Dict[str, Dict[str, Any]]:
+    from models.experimental.gr00t.reference.model import load_golden_sample
+
+    backend = policy.dit_backend
+    t0 = time.perf_counter()
+    # both models run on the megakernel's device configuration (reduced worker L1): the ttnn model therefore also
+    # needs the DRAM adapter intermediates (tt.model.ADAPTER_MEM_BY_BACKEND; values identical, placement only)
+    model = M.Gr00tTT.from_pretrained(version, policy=policy, device=device, adapter_mem="DRAM")
+    build_s = time.perf_counter() - t0
+    print(
+        f"\n[test_mk_e2e] {version}/{backend}: from_pretrained {build_s:.1f} s (weights {model.timing['weights_load_s']:.1f} s"
+        f"{'; arena ' + str(model.recipe.mk_arena_dtype) if backend == 'megakernel' else ''}; adapter_mem "
+        f"{model.adapter_mem}); head timing {({k: round(v, 2) for k, v in model.head.timing.items()})}"
+    )
+    out: Dict[str, Dict[str, Any]] = {}
+    try:
+        describe = model.describe()
+        for sample in samples:
+            gs, obs = load_golden_sample(version, sample)
+            try:
+                mi = model.encode(obs)  # raises for another embodiment (the N1.7 `g1_synthetic` extra)
+            except (ValueError, StaticShapeError) as exc:
+                out[sample] = {"sample": sample, "backend": backend, "skip": f"{type(exc).__name__}: {exc}"}
+                print(
+                    f"[test_mk_e2e] {version}/{backend}/{sample}: not this model's embodiment / layout -> skipped ({exc})"
+                )
+                continue
+            if mi.shape_key != model.dlayout.shape_key:
+                out[sample] = {"sample": sample, "backend": backend, "skip": f"shape key {mi.shape_key} != model's"}
+                continue
+            noise = gs.load("initial_noise")
+            rec: Dict[str, Any] = {"sample": sample, "backend": backend, "build_s": build_s, "n_text": int(mi.n_text)}
+            t0 = time.perf_counter()
+            pred, taps = model.run_untraced(mi, noise, return_taps=True, raw=True)
+            rec["untraced_s"] = time.perf_counter() - t0
+            rec["untraced_device_s"] = model.timing["last_untraced_s"]
+            rec["pred"] = pred
+            rec["raw"] = {k[len(M.RAW_PREFIX) :]: v for k, v in taps.items() if k.startswith(M.RAW_PREFIX)}
+            host_taps = {k: v for k, v in taps.items() if not k.startswith(M.RAW_PREFIX)}
+            expected = set(M.expected_tap_keys(model.cfg, model.dlayout, backend))
+            missing = sorted(expected - set(host_taps))
+            if missing:
+                raise AssertionError(f"{version}/{backend}: run_untraced lacks {missing}")
+            cand = golden_keyed(host_taps, gs)
+            cand.update(decoded_groups(model, pred, obs))
+            cand, info = split_info(version, cand)
+            rec["cand"] = cand
+            rec["required"] = required_gated_keys(version, gs, model.dlayout, backend) | {
+                k for k in cand if k.startswith(M.DECODE_TAP)
+            }
+            rec["head_taps"] = {k: v for k, v in host_taps.items() if k.split("[")[0] in COMMON_STEP_TAPS + ACTION_KEYS}
+            # traced path: capture on the first sample, replay every sample
+            if not model.captured:
+                t0 = time.perf_counter()
+                model.warm_and_capture(mi, noise)
+                rec["warm_and_capture_s"] = time.perf_counter() - t0
+            preds: List[torch.Tensor] = []
+            states: List[Dict[str, torch.Tensor]] = []
+            replay_s: List[float] = []
+            for _ in range(N_REPLAYS[backend]):
+                t0 = time.perf_counter()
+                preds.append(model.predict_normalized(mi, noise))
+                replay_s.append(time.perf_counter() - t0)
+                states.append(model.traced_state())
+            rec["traced_preds"], rec["traced_states"], rec["replay_s"] = preds, states, replay_s
+            if backend == "megakernel":
+                # the kernel's per-step decoder taps of the last traced replay vs the untraced run (K4 review NB 13);
+                # the untraced taps are golden-shaped ([1, M, A_max] per step), so compare on that slice
+                vel = model.head.velocity_taps_host()  # [n_steps, 64, A_pad]
+                m, a_max = int(model.dlayout.m_logical), int(model.cfg.encoders.max_action_dim)
+                rec["traced_velocity"] = vel[:, :m, :a_max].contiguous()
+                rec["untraced_velocity"] = torch.stack(
+                    [host_taps[f"action_decoder_out[k={k}]"][0] for k in range(vel.shape[0])]
+                )
+                if rec["traced_velocity"].shape != rec["untraced_velocity"].shape:
+                    raise AssertionError(
+                        f"velocity shapes {tuple(rec['traced_velocity'].shape)} vs {tuple(rec['untraced_velocity'].shape)}"
+                    )
+            # split pass: one synchronised replay per trace (informational under alloc tracking)
+            tr = model.tracer
+            per_stage: Dict[str, List[float]] = {s: [] for s in tr.stage_names}
+            for _ in range(N_SPLIT_REPS):
+                tr.upload(mi, noise)
+                tr.synchronize()
+                for st in tr.stage_names:
+                    t0 = time.perf_counter()
+                    tr.replay(st)
+                    tr.synchronize()
+                    per_stage[st].append((time.perf_counter() - t0) * 1e3)
+            rec["split_ms"] = {s: statistics.median(v) for s, v in per_stage.items()}
+            rec["split_raw_equal"] = bool(torch.equal(tr.read_action_pred(), states[0]["action_pred"]))
+            out[sample] = rec
+        out["_model"] = {
+            "describe": describe,
+            "timing": dict(model.timing),
+            "head_timing": dict(model.head.timing),
+            "recipe": model.recipe.describe(),
+            "executor": model.executor.describe(),
+        }
+    finally:
+        model.release()
+    return out
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# helpers
+# --------------------------------------------------------------------------------------------------------------------
+def _row(results: Any, tap: str, key: str, ok: bool, rule: str = "exact", **extra: Any) -> Dict[str, Any]:
+    r = {"tap": tap, "key": key, "rule": rule, "ok": bool(ok), **extra}
+    results.add_row(r)
+    return r
+
+
+def _pcc_row(
+    results: Any, key: str, got: torch.Tensor, ref: torch.Tensor, gate: Optional[float], **extra: Any
+) -> Dict[str, Any]:
+    """``gate=None`` records the row as information (``ok`` = finite only)."""
+    p = pcc(got.float(), ref.float())
+    finite = bool(torch.isfinite(got).all())
+    r = {
+        "tap": key.split("[")[0],
+        "key": f"{key}-mk-vs-ttnn",
+        "rule": "pcc" if gate is not None else "info",
+        "pcc": p,
+        "gate_pcc": gate,
+        "max_abs": max_abs(got.float(), ref.float()),
+        "finite": finite,
+        "ok": bool(finite and (gate is None or p >= gate)),
+        **extra,
+    }
+    results.add_row(r)
+    return r
+
+
+def _meta(results: Any, version: str, sample: str, res: Dict[str, Dict[str, Dict[str, Any]]]) -> None:
+    mk, s1 = res["megakernel"], res["ttnn"]
+    for rec in (mk.get(sample), s1.get(sample)):
+        if rec is None or "skip" in rec:
+            pytest.skip(f"{version}/{sample}: {rec['skip'] if rec else 'not collected'}")
+    results.set_meta(
+        sample=sample,
+        mk_recipe=mk["_model"]["recipe"],
+        ttnn_recipe=s1["_model"]["recipe"],
+        mk_head=mk["_model"]["describe"].get("head"),
+        mk_head_timing=mk["_model"]["head_timing"],
+        build_s={"ttnn": s1[sample]["build_s"], "megakernel": mk[sample]["build_s"]},
+        alloc_tracking_env=os.environ.get(T.ALLOC_TRACKING_ENV),
+        n_text=mk[sample]["n_text"],
+    )
+
+
+def _failures(
+    version: str, gs: GoldenSet, rec: Dict[str, Any], results: Optional[Any]
+) -> Tuple[harness.Report, Set[str]]:
+    report = harness.check_taps(version, rec["cand"], gs, results=results)
+    return report, {r["key"] for r in report.failures()}
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# tests
+# --------------------------------------------------------------------------------------------------------------------
+def test_mk_e2e_backends_agree(
+    version: str,
+    sample: str,
+    golden: Tuple[GoldenSet, Any, Any],
+    backend_results: Dict[str, Dict[str, Dict[str, Any]]],
+    results: harness.ResultsWriter,
+) -> None:
+    """Megakernel vs ttnn backend on the same sample: action_pred / per-step velocities >= 0.9999, golden gates."""
+    gs, _, _ = golden
+    mk, s1 = backend_results["megakernel"][sample], backend_results["ttnn"][sample]
+    _meta(results, version, sample, backend_results)
+    # golden gates first: the megakernel must pass every gate the ttnn backend passes (+ K4's closeness rule); a
+    # sample where the ttnn backend misses its own action_pred gate is a Stage-1 divergence -> no oracle there
+    report_mk, fail_mk = _failures(version, gs, mk, results)
+    report_s1, fail_s1 = _failures(version, gs, s1, None)
+    s1_is_oracle = "action_pred_normalized" not in fail_s1
+    oracle = GATE_ORACLE if s1_is_oracle else None
+    rows: Dict[str, Dict[str, Any]] = {}
+    rows["action_pred_normalized"] = _pcc_row(results, "action_pred_normalized", mk["pred"], s1["pred"], oracle)
+    rows["action_pred_normalized[traced]"] = _pcc_row(
+        results, "action_pred_normalized[traced]", mk["traced_preds"][0], s1["traced_preds"][0], oracle
+    )
+    for key, t in sorted(mk["head_taps"].items()):
+        if key in ACTION_KEYS:
+            continue
+        gate = oracle if key.split("[")[0] in ORACLE_STEP_TAPS else None
+        rows[key] = _pcc_row(results, key, t, s1["head_taps"][key], gate)
+    common = set(mk["cand"]) & set(s1["cand"])
+    regressions = sorted((fail_mk & common) - fail_s1)
+    mk_rows = {r["key"]: r for r in report_mk.rows}
+    known, known_facts = known_marginal_regressions(version, sample, regressions, mk_rows)
+    gold = gs.load("action_pred_normalized").float()
+    close_mk, close_s1 = pcc(mk["pred"], gold), pcc(s1["pred"], gold)
+    reference_ok = close_s1 >= CLOSENESS_MIN_REFERENCE
+    _row(
+        results,
+        "action_pred_normalized",
+        "mk-closeness-to-golden",
+        close_mk >= close_s1 - GOLDEN_SLACK,
+        rule="pcc",
+        pcc=close_mk,
+        s1_pcc_vs_golden=close_s1,
+        gate_pcc=close_s1 - GOLDEN_SLACK,
+        max_abs=max_abs(mk["pred"], gold),
+        s1_max_abs_vs_golden=max_abs(s1["pred"], gold),
+    )
+    _row(
+        results,
+        "gates",
+        "mk-passes-every-gate-ttnn-passes",
+        not regressions,
+        rule="set",
+        n_common=len(common),
+        n_required_mk=len(mk["required"]),
+        fail_mk=sorted(fail_mk),
+        fail_ttnn=sorted(fail_s1),
+        regressions=regressions,
+        known_marginal=known,
+        known_marginal_facts=known_facts,
+        fixed_by_mk=sorted((fail_s1 & common) - fail_mk),
+        ttnn_only_fail=sorted(fail_s1 - common),
+        ttnn_is_oracle=s1_is_oracle,
+    )
+    print(f"[{version}/{sample}] mk gates: {report_mk.summary()}")
+    print(f"[{version}/{sample}] ttnn gates: {report_s1.summary()}")
+    print(
+        f"[{version}/{sample}] mk vs ttnn: action_pred {rows['action_pred_normalized']['pcc']:.7f} "
+        f"(max|d| {rows['action_pred_normalized']['max_abs']:.3g}); mk/ttnn vs golden {close_mk:.6f}/{close_s1:.6f}; "
+        f"regressions {regressions}; fixed {sorted((fail_s1 & common) - fail_mk)}"
+    )
+    assert set(mk["cand"]) == mk["required"], (
+        sorted(mk["required"] - set(mk["cand"])),
+        sorted(set(mk["cand"]) - mk["required"]),
+    )
+    if not reference_ok:
+        pytest.skip(
+            f"{version}/{sample}: the ttnn backend is at PCC {close_s1:.4f} vs the golden (megakernel {close_mk:.4f}); "
+            f"neither pipeline is a reference for this sample (rows recorded)"
+        )
+    bad = {k: r["pcc"] for k, r in rows.items() if not r["ok"]}
+    assert not bad, f"megakernel vs ttnn below {GATE_ORACLE}: {bad}"
+    assert (
+        close_mk >= close_s1 - GOLDEN_SLACK
+    ), f"mk vs golden {close_mk:.6f} < ttnn vs golden {close_s1:.6f} - {GOLDEN_SLACK}"
+    unknown = [k for k in regressions if k not in known]
+    assert not unknown, f"gates the ttnn backend passes and the megakernel fails: {unknown} (known marginal: {known})"
+    if known:
+        pytest.xfail(
+            f"{version}/{sample}: only MK_KNOWN_MARGINAL rows regress (above their floors): "
+            + ", ".join(
+                f"{k} {known_facts[k]['field']} {known_facts[k]['value']:.6f} >= floor {known_facts[k]['floor']} "
+                f"(row pcc {mk_rows[k]['pcc']:.6f} vs gate {mk_rows[k].get('gate_pcc')})"
+                for k in known
+            )
+        )
+
+
+def test_mk_e2e_trace_protocol(
+    version: str,
+    sample: str,
+    golden: Tuple[GoldenSet, Any, Any],
+    backend_results: Dict[str, Dict[str, Dict[str, Any]]],
+    results: harness.ResultsWriter,
+) -> None:
+    """Traced == untraced (raw buffers), N replays bit-identical, for both backends (megakernel 10x)."""
+    _meta(results, version, sample, backend_results)
+    ok_all = True
+    for backend in BACKENDS:
+        rec = backend_results[backend][sample]
+        st0 = rec["traced_states"][0]
+        eq_pred = bool(
+            torch.equal(st0["action_pred"], rec["raw"]["action_pred"])
+            and torch.equal(rec["traced_preds"][0], rec["pred"])
+        )
+        eq_bb = bool(torch.equal(st0["backbone_features"], rec["raw"]["backbone_features"]))
+        eq_kv = bool(torch.equal(st0["K_0"], rec["raw"]["K_0"]) and torch.equal(st0["V_0"], rec["raw"]["V_0"]))
+        eq_replays = all(torch.equal(st0["action_pred"], s["action_pred"]) for s in rec["traced_states"][1:])
+        eq_host = all(torch.equal(p, rec["traced_preds"][0]) for p in rec["traced_preds"][1:])
+        _row(
+            results,
+            "action_pred",
+            f"{backend}-traced_equals_untraced",
+            eq_pred,
+            pcc=pcc(rec["traced_preds"][0], rec["pred"]),
+            max_abs=max_abs(rec["traced_preds"][0], rec["pred"]),
+        )
+        _row(results, "backbone_features", f"{backend}-traced_equals_untraced", eq_bb)
+        _row(results, "K_0/V_0", f"{backend}-traced_equals_untraced", eq_kv)
+        _row(
+            results,
+            "action_pred",
+            f"{backend}-{N_REPLAYS[backend]}_replays_bit_identical",
+            eq_replays and eq_host,
+            n_replays=N_REPLAYS[backend],
+            replay_ms=[round(s * 1e3, 3) for s in rec["replay_s"]],
+        )
+        _row(results, "action_pred", f"{backend}-split_replayer_equals_predict", rec["split_raw_equal"])
+        if "traced_velocity" in rec:
+            eq_vel = bool(torch.equal(rec["traced_velocity"], rec["untraced_velocity"]))
+            _row(results, "velocity_out", f"{backend}-traced_equals_untraced", eq_vel)
+            ok_all = ok_all and eq_vel
+        print(
+            f"[{version}/{sample}] {backend}: traced==untraced {eq_pred}, {N_REPLAYS[backend]} replays identical "
+            f"{eq_replays and eq_host}, replay+D2H ms {[round(s * 1e3, 2) for s in rec['replay_s']]}"
+        )
+        ok_all = ok_all and eq_pred and eq_bb and eq_kv and eq_replays and eq_host and rec["split_raw_equal"]
+    assert ok_all, "see the rows: a backend's traced path is not bit-equal to its untraced path / not deterministic"
+
+
+def test_mk_e2e_denoise_trace_ms(
+    version: str,
+    sample: str,
+    golden: Tuple[GoldenSet, Any, Any],
+    backend_results: Dict[str, Dict[str, Dict[str, Any]]],
+    results: harness.ResultsWriter,
+) -> None:
+    """Informational split-pass trace times of both backends under alloc tracking (bench_e2e is the latency source)."""
+    _meta(results, version, sample, backend_results)
+    mk, s1 = backend_results["megakernel"][sample], backend_results["ttnn"][sample]
+    stages = list(mk["split_ms"])
+    denoise = stages[-1]
+    ratio = mk["split_ms"][denoise] / s1["split_ms"][denoise]
+    for st in stages:
+        _row(
+            results,
+            "trace",
+            f"split_ms/{st}",
+            True,
+            rule="timing",
+            ttnn_ms=s1["split_ms"][st],
+            megakernel_ms=mk["split_ms"][st],
+            ratio_mk_over_ttnn=mk["split_ms"][st] / s1["split_ms"][st],
+            note="median of synchronised replays under TT_METAL_TRACE_ALLOC_TRACKING=1 (informational)",
+        )
+    results.set_timing(
+        mk_denoise_split_s=mk["split_ms"][denoise] / 1e3,
+        ttnn_denoise_split_s=s1["split_ms"][denoise] / 1e3,
+        mk_untraced_device_s=mk["untraced_device_s"],
+        ttnn_untraced_device_s=s1["untraced_device_s"],
+        mk_replay_median_s=sorted(mk["replay_s"])[len(mk["replay_s"]) // 2],
+        ttnn_replay_median_s=sorted(s1["replay_s"])[len(s1["replay_s"]) // 2],
+    )
+    results.set_meta(denoise_trace=denoise, denoise_ratio_mk_over_ttnn=ratio)
+    print(
+        f"[{version}/{sample}] split ms ttnn {({k: round(v, 2) for k, v in s1['split_ms'].items()})} | "
+        f"megakernel {({k: round(v, 2) for k, v in mk['split_ms'].items()})} | denoise ratio {ratio:.3f}"
+    )
+    assert ratio > 0 and all(v > 0 for v in mk["split_ms"].values())

code/models/experimental/gr00t/tests/tt/test_mk_head_sdpa.py

@@ -0,0 +1,754 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K3 tests of the megakernel's ``HeadSDPA`` (one head per core, flash-style over 64-key chunks) and the head
+exchange, run through the standalone harness kernel ``kernels/test/sdpa_harness_kernel.cpp`` (IMPLEMENTATION_PLAN.md
+§5.4 P3/P4 + cross-block paragraph, §5.7 K3 row, R13; mb2 §B.1 mask / scale semantics).
+
+* ``test_*_cpu`` -- the chunk arithmetic and DRAM tile-id maps the kernel uses, the test's own torch references
+  (masked == pad-key-excluded) and an fp32 emulation of the chunked running-softmax (half chunks included) -- no device;
+* ``test_head_sdpa`` -- 32 heads on the 32 head cores, ``M = 64`` query rows with the version's logical rows
+  (49 / 51 / 41), random Q/K/V with padded rows / keys, Sk in {64, 96, 320, 384, 448, 544}: **self** mode (Sk 64:
+  K_h / V_h delivered by the BRISC of cores 32+h / 64+h into core h, plan §5.4 P3) and **cross** mode (K^T / V / mask
+  64-key chunks fetched from the DRAM hoist tensors into the 2-deep rings, pre-transposed K^T hoist and natural K
+  with the on-TRISC transpose), with / without the ``-1e9`` pad-key mask, explicit ``48**-0.5`` scale and the folded
+  variant (Q pre-scaled in bf16 = ``arena.fold_attn_scale``).  Gate: PCC >= 0.9995 vs torch fp32 on the same bf16
+  inputs (valid rows; the masked case against the pad-key-excluded reference) **and** vs
+  ``ttnn.transformer.scaled_dot_product_attention``; no NaN anywhere (all 64 rows x 32 heads); second launch
+  bit-identical;
+* ``test_head_sdpa_timing`` -- per-head µs per Sk (mb3 difference method over in-kernel iterations), K/V/mask fetch
+  GB/s of the cross mode.
+
+Run (one pytest process per file, through the device lock)::
+
+    DEVICE_LOCK_TIMEOUT=14400 bin/with-device.sh python -m pytest -q -rA --timeout=1800 -p no:cacheprovider \\
+        models/experimental/gr00t/tests/tt/test_mk_head_sdpa.py --gr00t-version n16
+
+``ttnn`` is imported lazily (device tests only); the CPU tests import nothing from ``ttnn``.
+"""
+
+from __future__ import annotations
+
+import itertools
+import math
+import struct
+import time
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
+
+import pytest
+import torch
+
+from models.experimental.gr00t.common import pcc as pcc_mod
+from models.experimental.gr00t.tt.megakernel import arena as A
+from models.experimental.gr00t.tt.megakernel import dit_program as P
+from models.experimental.gr00t.tt.megakernel.core_map import N_HEADS, CoreMap
+
+torch.set_grad_enabled(False)
+
+KERNEL_SOURCE = "models/experimental/gr00t/tt/megakernel/kernels/test/sdpa_harness_kernel.cpp"
+GATE_PCC = 0.9995  # WP-K3 acceptance (plan §5.7 K3 row)
+SCALE = 48**-0.5  # cfg.dit.attention_scale (head_dim 48 padded to 64: the explicit scale is mandatory, mb2 §B.1)
+MASK_NEG = -1e9  # static-shape-plan §4.3: a large finite negative, never -inf
+TILE = 32
+M_ROWS = 64
+HEAD_DIM = 64
+#: logical query rows per version (LayoutConfig.dit_m_logical: 49 / 51 / 41; asserted against ArenaLayout on CPU)
+M_LOGICAL: Mapping[str, int] = {"n15": 49, "n16": 51, "n17": 41}
+SK_LIST: Tuple[int, ...] = (64, 96, 320, 384, 448, 544)
+#: valid keys per Sk (the mask excludes the rest): Sk 64 uses the version's logical rows (self mask), the cross values
+#: are real layouts (n16 image 81/96, n15 all 296/320 and 297/384, n17 alt image 528/544) or a synthetic pad (441/448)
+N_VALID: Mapping[int, int] = {96: 81, 320: 296, 384: 297, 448: 441, 544: 528}
+#: ttnn SDPA chunk pairs per Sk (q_chunk 64; (64, 64) verified for 64 / 96 / 320 / 512 in mb2 §B.1, (64, 128) for
+#: 384 / 448 in tt/layers.py SDPA_CHUNKS; (64, 320) / (64, 448) are the forbidden pairs)
+TTNN_CHUNKS: Mapping[int, Tuple[int, int]] = {
+    64: (64, 64),
+    96: (64, 64),
+    320: (64, 64),
+    384: (64, 128),
+    448: (64, 128),
+    544: (64, 64),
+}
+MODES: Tuple[str, ...] = ("self", "cross")
+K_LAYOUTS: Tuple[str, ...] = ("transposed", "natural")
+SEED = 1234
+CB_IDS: Mapping[str, int] = {
+    "cb_in": 0,
+    "cb_o": 1,
+    "cb_k": 2,
+    "cb_v": 3,
+    "cb_kt": 4,
+    "cb_mask": 5,
+    "cb_zero": 6,
+    "cb_ones": 7,
+    "cb_s": 8,
+    "cb_oacc": 9,
+    "cb_m0": 10,
+    "cb_m1": 11,
+    "cb_l0": 12,
+    "cb_l1": 13,
+    "cb_alpha": 14,
+}
+SEM_IDS: Mapping[str, int] = {"sem_kv": 0, "sem_credit": 1}
+TILE_BF16 = TILE * TILE * 2
+TILE_FP32 = TILE * TILE * 4
+KV_SLOTS = 2  # 2-deep K / V / mask chunk rings (plan §5.5)
+O_SLOTS = 2  # output slots of the tensor-backed cb_o (iteration i writes slot i % O_SLOTS)
+#: (name, bytes, page bytes, dtype) of the plain CBs of the harness kernel (head_sdpa.hpp CB contract)
+PLAIN_CBS: Tuple[Tuple[str, int, int, str], ...] = (
+    ("cb_k", KV_SLOTS * 4 * TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_v", KV_SLOTS * 4 * TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_kt", 4 * TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_mask", KV_SLOTS * 4 * TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_zero", TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_ones", TILE_BF16, TILE_BF16, "bf16"),
+    ("cb_s", 4 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_oacc", 4 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_m0", 2 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_m1", 2 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_l0", 2 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_l1", 2 * TILE_FP32, TILE_FP32, "fp32"),
+    ("cb_alpha", 2 * TILE_FP32, TILE_FP32, "fp32"),
+)
+#: common runtime args of the harness kernel (sdpa_harness_kernel.cpp enum RT), then compute_xy[96]
+RT_NAMES: Tuple[str, ...] = (
+    "mode",
+    "n_iters",
+    "sk_tiles",
+    "use_mask",
+    "scale_bits",
+    "k_natural",
+    "addr_k",
+    "addr_v",
+    "addr_mask",
+    "kv_slots",
+    "o_slots",
+    "debug_stage",
+    "debug_chunk",
+)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# CPU-side mirrors of the kernel contract
+# --------------------------------------------------------------------------------------------------------------------
+def n_valid_keys(sk: int, m_logical: int) -> int:
+    """Valid keys of the masked case: the version's logical rows for the self mask (Sk 64), :data:`N_VALID` else."""
+    if sk == M_ROWS:
+        return m_logical
+    if sk not in N_VALID:
+        raise KeyError(f"no valid-key count for Sk {sk}; expected one of {sorted(N_VALID)} or 64")
+    return N_VALID[sk]
+
+
+def chunk_plan(sk: int) -> List[int]:
+    """Key tiles per 64-key chunk (``HeadSDPAParams::chunk_kt``): ``[2, 2, ..., 1?]``; Sk must be a multiple of 32."""
+    if sk <= 0 or sk % TILE:
+        raise ValueError(f"Sk must be a positive multiple of {TILE}, got {sk}")
+    kt = sk // TILE
+    out = []
+    while kt > 0:
+        out.append(min(2, kt))
+        kt -= out[-1]
+    return out
+
+
+def fp32_bits(x: float) -> int:
+    """The fp32 bit pattern of ``x`` as the kernel's ``scale_bits`` runtime arg."""
+    return struct.unpack("<I", struct.pack("<f", float(x)))[0]
+
+
+def tile_id_kv(h: int, key_tile: int, d_tile: int, lt: int) -> int:
+    """Tile id of ``(h, key_tile, d_tile)`` in a TILE_LAYOUT ``[1, H, L_pad, 64]`` tensor (``HoistChunkFetcher``)."""
+    return (h * lt + key_tile) * 2 + d_tile
+
+
+def tile_id_kT(h: int, d_tile: int, key_tile: int, lt: int) -> int:
+    """Tile id of ``(h, d_tile, key_tile)`` in a TILE_LAYOUT ``[1, H, 64, L_pad]`` tensor (pre-transposed hoist)."""
+    return (h * 2 + d_tile) * lt + key_tile
+
+
+def tile_id_mask(r: int, key_tile: int, lt: int) -> int:
+    """Tile id of ``(r, key_tile)`` in a TILE_LAYOUT ``[1, 1, 64, L_pad]`` mask."""
+    return r * lt + key_tile
+
+
+def tiles_row_major(t: torch.Tensor) -> torch.Tensor:
+    """``[..., R, C]`` -> ``[n_tiles, 32, 32]`` in the TILE_LAYOUT page order (row-major over the flattened tile grid)."""
+    r, c = t.shape[-2], t.shape[-1]
+    if r % TILE or c % TILE:
+        raise ValueError(f"not tile aligned: {tuple(t.shape)}")
+    flat = t.reshape(-1, r, c)
+    tiles = flat.reshape(flat.shape[0], r // TILE, TILE, c // TILE, TILE).permute(0, 1, 3, 2, 4)
+    return tiles.reshape(-1, TILE, TILE)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# inputs and references
+# --------------------------------------------------------------------------------------------------------------------
+def make_inputs(sk: int, m_logical: int, seed: int) -> Dict[str, torch.Tensor]:
+    """Random bf16 ``q [32, 64, 64]``, ``k`` / ``v [32, Sk, 64]`` (pad rows / keys are random too, so a leak shows) and
+    the additive ``mask [64, Sk]`` (0 on the ``n_valid_keys`` valid keys, ``-1e9`` elsewhere)."""
+    g = torch.Generator().manual_seed(seed)
+    q = torch.randn(N_HEADS, M_ROWS, HEAD_DIM, generator=g).to(torch.bfloat16)
+    k = torch.randn(N_HEADS, sk, HEAD_DIM, generator=g).to(torch.bfloat16)
+    v = torch.randn(N_HEADS, sk, HEAD_DIM, generator=g).to(torch.bfloat16)
+    nv = n_valid_keys(sk, m_logical)
+    mask = torch.zeros(M_ROWS, sk, dtype=torch.bfloat16)
+    mask[:, nv:] = MASK_NEG
+    return {"q": q, "k": k, "v": v, "mask": mask}
+
+
+def fold_scale(q: torch.Tensor, scale: float) -> torch.Tensor:
+    """The folded variant: Q already carries the softmax scale (bf16-rounded like ``arena.fold_attn_scale`` applied to
+    Wq / bq)."""
+    return (q.float() * scale).to(torch.bfloat16)
+
+
+def torch_reference(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor], scale: float
+) -> torch.Tensor:
+    """fp32 ``softmax(scale * q k^T + mask) v`` on the bf16 inputs; ``[H, 64, 64]`` fp32."""
+    s = torch.matmul(q.float(), k.float().transpose(-1, -2)) * scale
+    if mask is not None:
+        s = s + mask.float()
+    p = torch.softmax(s, dim=-1)
+    return torch.matmul(p, v.float())
+
+
+def flash_emulation(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor], scale: float
+) -> torch.Tensor:
+    """fp32 emulation of the kernel's chunked loop (head_sdpa.hpp steps 2-6 + finalise) with the same chunk plan."""
+    qf, kf, vf = q.float(), k.float(), v.float()
+    sk = kf.shape[1]
+    o = torch.zeros(qf.shape[0], M_ROWS, HEAD_DIM)
+    m = torch.full((qf.shape[0], M_ROWS, 1), -float("inf"))
+    l = torch.zeros(qf.shape[0], M_ROWS, 1)
+    key = 0
+    for kt in chunk_plan(sk):
+        keys = slice(key, key + kt * TILE)
+        s = torch.matmul(qf, kf[:, keys].transpose(-1, -2)) * scale
+        if mask is not None:
+            s = s + mask[:, keys].float()
+        m_new = torch.maximum(m, s.amax(dim=-1, keepdim=True))
+        p = torch.exp(s - m_new)
+        alpha = torch.exp(m - m_new)
+        l = alpha * l + p.sum(dim=-1, keepdim=True)
+        o = alpha * o + torch.matmul(p, vf[:, keys])
+        m = m_new
+        key += kt * TILE
+    return o / l
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# CPU tests
+# --------------------------------------------------------------------------------------------------------------------
+def test_m_logical_matches_arena_cpu(version: str) -> None:
+    """The logical rows of the test (49 / 51 / 41) are the ArenaLayout's ``m_logical`` for the version."""
+    L = A.ArenaLayout.plan(version, tile_order="row_major")
+    assert L.m_logical == M_LOGICAL[version], (version, L.m_logical)
+    assert L.m_pad == M_ROWS and L.head_dim_pad == HEAD_DIM and L.n_heads == N_HEADS
+    assert math.isclose(L.cfg.dit.attention_scale, SCALE, rel_tol=1e-9)
+    for ks, lp in L.l_pad.items():
+        assert lp % TILE == 0 and (lp in SK_LIST), (ks, lp)  # every real key-set length is in the tested Sk list
+
+
+def test_chunk_plan_and_tile_ids_cpu() -> None:
+    """Chunk plan of every Sk (half chunk for 96 / 544) and the DRAM tile-id maps of ``HoistChunkFetcher``."""
+    assert chunk_plan(64) == [2] and chunk_plan(96) == [2, 1] and chunk_plan(320) == [2] * 5
+    assert chunk_plan(384) == [2] * 6 and chunk_plan(448) == [2] * 7 and chunk_plan(544) == [2] * 8 + [1]
+    with pytest.raises(ValueError):
+        chunk_plan(80)
+    h, sk = 5, 96
+    lt = sk // TILE
+    kv = torch.arange(N_HEADS * sk * HEAD_DIM, dtype=torch.float32).reshape(1, N_HEADS, sk, HEAD_DIM)
+    tiles = tiles_row_major(kv)
+    for key_t in range(lt):
+        for d in range(2):
+            ref = kv[0, h, key_t * TILE : (key_t + 1) * TILE, d * TILE : (d + 1) * TILE]
+            assert torch.equal(tiles[tile_id_kv(h, key_t, d, lt)], ref)
+    kT = kv.transpose(-1, -2).contiguous()  # [1, H, 64, Sk]
+    tiles_t = tiles_row_major(kT)
+    for d in range(2):
+        for key_t in range(lt):
+            ref = kT[0, h, d * TILE : (d + 1) * TILE, key_t * TILE : (key_t + 1) * TILE]
+            assert torch.equal(tiles_t[tile_id_kT(h, d, key_t, lt)], ref)
+    mask = torch.arange(M_ROWS * sk, dtype=torch.float32).reshape(1, 1, M_ROWS, sk)
+    tiles_m = tiles_row_major(mask)
+    for r in range(2):
+        for key_t in range(lt):
+            assert torch.equal(
+                tiles_m[tile_id_mask(r, key_t, lt)],
+                mask[0, 0, r * TILE : (r + 1) * TILE, key_t * TILE : (key_t + 1) * TILE],
+            )
+    assert fp32_bits(1.0) == 0x3F800000 and fp32_bits(SCALE) == 0x3E13CD3A
+
+
+@pytest.mark.parametrize("sk", SK_LIST)
+def test_references_cpu(sk: int) -> None:
+    """The masked torch reference equals the pad-key-excluded one; the chunked fp32 emulation reproduces the full
+    softmax for every Sk (incl. the half-chunk tails) with and without the mask."""
+    m_logical = M_LOGICAL["n16"]
+    x = make_inputs(sk, m_logical, SEED + sk)
+    nv = n_valid_keys(sk, m_logical)
+    masked = torch_reference(x["q"], x["k"], x["v"], x["mask"], SCALE)
+    excluded = torch_reference(x["q"], x["k"][:, :nv], x["v"][:, :nv], None, SCALE)
+    assert torch.allclose(masked, excluded, atol=1e-5, rtol=1e-5)
+    assert pcc_mod.pcc(masked, excluded) > 0.999999
+    for mask in (x["mask"], None):
+        ref = torch_reference(x["q"], x["k"], x["v"], mask, SCALE)
+        emu = flash_emulation(x["q"], x["k"], x["v"], mask, SCALE)
+        assert torch.isfinite(emu).all()
+        assert torch.allclose(emu, ref, atol=2e-5, rtol=1e-4), (sk, mask is None)
+    # the unmasked reference is a different function of the inputs (pad keys leak): the test can tell them apart
+    leak = torch_reference(x["q"], x["k"], x["v"], None, SCALE)
+    assert pcc_mod.pcc(leak, excluded) < 0.9995
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# the harness launcher (device)
+# --------------------------------------------------------------------------------------------------------------------
+class HeadSDPAHarness:
+    """Launches ``sdpa_harness_kernel.cpp`` over the 96 compute cores of the megakernel core map.
+
+    Tensors: ``qkv_in`` ``[96*64, 64]`` bf16 HEIGHT_SHARDED (shard of compute core ``c`` = Q_c for ``c < 32``, and in
+    self mode K_{c-32} / V_{c-64} for the sender cores), the DRAM hoists ``k_hoist`` (``[1,32,64,Sk]`` K^T or
+    ``[1,32,Sk,64]`` K), ``v_hoist`` ``[1,32,Sk,64]``, ``mask`` ``[1,1,64,Sk]``, and the output ``o_out``
+    ``[96*64, 64*O_SLOTS]`` bf16 HEIGHT_SHARDED (tensor-backed ``cb_o``).
+    """
+
+    def __init__(self, device: Any) -> None:
+        import ttnn  # lazy
+
+        self._ttnn = ttnn
+        self.device = device
+        self.core_map = CoreMap.from_device(device)
+        self._crs = self.core_map.to_core_range_sets()
+        self._shard_cores = self._device_shard_order()
+        self._xy = [self._pack_xy(c) for c in self.core_map.compute]
+        self._per_core = self.core_map.per_core_args()
+        self.launches = 0
+        self.compile_s: Optional[float] = None
+
+    # ------------------------------------------------------------------ helpers
+    def _device_shard_order(self) -> List[int]:
+        ttnn = self._ttnn
+        order = P.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 _pack_xy(self, core: Tuple[int, int]) -> int:
+        v = self.device.worker_core_from_logical_core(self._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)
+
+    def _sharded(self, t: torch.Tensor, cols: int) -> Any:
+        ttnn = self._ttnn
+        if t.shape != (self.core_map.n_compute * M_ROWS, cols):
+            raise ValueError(
+                f"sharded tensor must be [{self.core_map.n_compute * M_ROWS}, {cols}], got {tuple(t.shape)}"
+            )
+        spec = ttnn.ShardSpec(self._crs["compute"], [M_ROWS, cols], ttnn.ShardOrientation.ROW_MAJOR)
+        mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, spec)
+        return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=mem)
+
+    def _dram(self, t: torch.Tensor) -> Any:
+        ttnn = self._ttnn
+        dev = ttnn.from_torch(
+            t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+        )
+        page = int(dev.buffer_page_size())
+        if page != TILE_BF16:
+            raise RuntimeError(f"DRAM tile tensor page {page} B != {TILE_BF16}: the kernel reads 2 KB tile pages")
+        return dev
+
+    def head_output(self, out: torch.Tensor, h: int, slot: int) -> torch.Tensor:
+        """O_h ``[64, 64]`` from the ``o_out`` readback: the ``[64, 128]`` shard of head core ``h`` holds its 8 tile
+        pages row-major over a 2 x 4 tile grid, cb_o slot ``s`` = pages ``4s..4s+3`` = shard rows ``32s..32s+31``, and
+        page ``4s + p`` is the head's output tile ``(r, c) = (p // 2, p % 2)``."""
+        i = self._shard_cores.index(h)
+        shard = out[i * M_ROWS : (i + 1) * M_ROWS]
+        o = torch.empty(M_ROWS, HEAD_DIM, dtype=out.dtype)
+        for pg in range(4):
+            r, c = divmod(pg, 2)
+            o[r * TILE : (r + 1) * TILE, c * TILE : (c + 1) * TILE] = shard[
+                slot * TILE : (slot + 1) * TILE, pg * TILE : (pg + 1) * TILE
+            ]
+        return o
+
+    # ------------------------------------------------------------------ program
+    def _program(self, io: Mapping[str, Any], rt_values: Mapping[str, int]) -> Any:
+        ttnn = self._ttnn
+        from models.demos.deepseek_v3_b1.unified_kernel_descriptor import (
+            PerCoreRuntimeArgsDescriptor,
+            UnifiedKernelDescriptor,
+        )
+
+        crs = self._crs["compute"]
+        dt = {"bf16": ttnn.bfloat16, "fp32": ttnn.float32}
+        tile_desc = ttnn.TileDescriptor(ttnn.Tile([TILE, TILE]))
+        cbs = [
+            ttnn.cb_descriptor_from_sharded_tensor(CB_IDS["cb_in"], io["qkv_in"]),
+            ttnn.cb_descriptor_from_sharded_tensor(CB_IDS["cb_o"], io["o_out"]),
+        ]
+        for name, total, page, fmt in PLAIN_CBS:
+            f = ttnn.CBFormatDescriptor(buffer_index=CB_IDS[name], data_format=dt[fmt], page_size=page, tile=tile_desc)
+            cbs.append(ttnn.CBDescriptor(total_size=total, core_ranges=crs, format_descriptors=[f]))
+        sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=crs, initial_value=0) for i in SEM_IDS.values()]
+        named = [(k, v) for k, v in CB_IDS.items()] + [(k, v) for k, v in SEM_IDS.items()]
+        named += [("n_heads", N_HEADS), ("n_compute", self.core_map.n_compute)]
+        common = [int(rt_values[k]) for k in RT_NAMES] + list(self._xy)
+        for w in common:
+            if not 0 <= w < 2**32:
+                raise ValueError(f"runtime arg {w} does not fit uint32")
+        per_core = [
+            (ttnn.CoreCoord(x, y), list(self._per_core[(x, y)]))
+            for (x, y) in sorted(self.core_map.compute, key=lambda c: (c[1], c[0]))
+        ]
+        uk = UnifiedKernelDescriptor(
+            kernel_source=KERNEL_SOURCE,
+            core_ranges=crs,
+            ncrisc_named_compile_time_args=named,
+            brisc_named_compile_time_args=named,
+            trisc_named_compile_time_args=named,
+            ncrisc_common_runtime_args=common,
+            brisc_common_runtime_args=common,
+            trisc_common_runtime_args=common,
+            trisc_compute_config=ttnn.ComputeConfigDescriptor(
+                math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, math_approx_mode=False
+            ),
+            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) != 1:
+            raise RuntimeError(f"harness kernel compiled into {len(result.groups)} groups, expected 1")
+        return ttnn.ProgramDescriptor(kernels=result.kernels, semaphores=sems, cbs=cbs)
+
+    # ------------------------------------------------------------------ run
+    def run(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        *,
+        mode: str,
+        scale: Optional[float],
+        k_layout: str = "transposed",
+        n_iters: int = 1,
+        debug: Tuple[int, int] = (0, 0),
+    ) -> Tuple[torch.Tensor, float]:
+        """One launch; returns ``(O [32, 64, 64] bf16 of the last iteration, launch + synchronize wall seconds)``.
+
+        ``scale=None`` = folded (the kernel skips the SFPU scale); ``k_layout`` is the cross-mode hoist layout
+        (``"transposed"`` = K^T ``[1,32,64,Sk]``, ``"natural"`` = K ``[1,32,Sk,64]`` + on-TRISC transpose); self mode
+        always delivers natural K_h.  ``debug=(stage, chunk)`` (bring-up only) makes the TRISC dump the intermediate
+        of ``stage`` (head_sdpa.hpp ``debug_stage``) at chunk ``chunk`` into the output slot instead of O.
+        """
+        ttnn = self._ttnn
+        if mode not in MODES:
+            raise ValueError(f"mode must be in {MODES}, got {mode!r}")
+        if k_layout not in K_LAYOUTS:
+            raise ValueError(f"k_layout must be in {K_LAYOUTS}, got {k_layout!r}")
+        if q.shape != (N_HEADS, M_ROWS, HEAD_DIM) or q.dtype != torch.bfloat16:
+            raise ValueError(f"q must be bf16 [{N_HEADS}, 64, 64], got {tuple(q.shape)} {q.dtype}")
+        sk = int(k.shape[1])
+        if k.shape != (N_HEADS, sk, HEAD_DIM) or v.shape != (N_HEADS, sk, HEAD_DIM) or sk % TILE or sk < TILE:
+            raise ValueError(
+                f"k / v must be [{N_HEADS}, Sk (multiple of 32), 64], got {tuple(k.shape)} {tuple(v.shape)}"
+            )
+        if k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16:
+            raise ValueError("k / v must be bf16")
+        if mask is not None and (mask.shape != (M_ROWS, sk) or mask.dtype != torch.bfloat16):
+            raise ValueError(f"mask must be bf16 [64, {sk}] or None, got {tuple(mask.shape)}")
+        if mode == "self" and sk != M_ROWS:
+            raise ValueError(f"self mode exchanges one 64-key chunk; Sk {sk} needs cross mode")
+        if n_iters < 1:
+            raise ValueError("n_iters >= 1")
+        n_c = self.core_map.n_compute
+        k_natural = 1 if (mode == "self" or k_layout == "natural") else 0
+
+        qkv = torch.zeros(n_c * M_ROWS, HEAD_DIM, dtype=torch.bfloat16)
+        for i, c in enumerate(self._shard_cores):
+            rows = slice(i * M_ROWS, (i + 1) * M_ROWS)
+            if c < N_HEADS:
+                qkv[rows] = q[c]
+            elif mode == "self" and c < 2 * N_HEADS:
+                qkv[rows] = k[c - N_HEADS]
+            elif mode == "self":
+                qkv[rows] = v[c - 2 * N_HEADS]
+        io: Dict[str, Any] = {"qkv_in": self._sharded(qkv, HEAD_DIM)}
+        k_host = k.reshape(1, N_HEADS, sk, HEAD_DIM)
+        if not k_natural:
+            k_host = k_host.transpose(-1, -2).contiguous()
+        io["k_hoist"] = self._dram(k_host)
+        io["v_hoist"] = self._dram(v.reshape(1, N_HEADS, sk, HEAD_DIM))
+        io["mask"] = self._dram(
+            mask.reshape(1, 1, M_ROWS, sk) if mask is not None else torch.zeros(1, 1, TILE, TILE, dtype=torch.bfloat16)
+        )
+        io["o_out"] = self._sharded(
+            torch.zeros(n_c * M_ROWS, HEAD_DIM * O_SLOTS, dtype=torch.bfloat16), HEAD_DIM * O_SLOTS
+        )
+        rt = {
+            "mode": MODES.index(mode),
+            "n_iters": n_iters,
+            "sk_tiles": sk // TILE,
+            "use_mask": 0 if mask is None else 1,
+            "scale_bits": 0 if scale is None else fp32_bits(scale),
+            "k_natural": k_natural,
+            "addr_k": int(io["k_hoist"].buffer_address()),
+            "addr_v": int(io["v_hoist"].buffer_address()),
+            "addr_mask": int(io["mask"].buffer_address()),
+            "kv_slots": KV_SLOTS,
+            "o_slots": O_SLOTS,
+            "debug_stage": int(debug[0]),
+            "debug_chunk": int(debug[1]),
+        }
+        if scale is not None and rt["scale_bits"] == 0:
+            raise ValueError("an explicit scale of 0.0 is indistinguishable from the folded variant")
+        program = self._program(io, rt)
+        tensors = [io["qkv_in"], io["k_hoist"], io["v_hoist"], io["mask"], io["o_out"]]
+        t0 = time.perf_counter()
+        ttnn.generic_op(tensors, program)
+        ttnn.synchronize_device(self.device)
+        dt = time.perf_counter() - t0
+        if self.compile_s is None:
+            self.compile_s = dt
+        self.launches += 1
+        out = ttnn.to_torch(io["o_out"])
+        slot = (n_iters - 1) % O_SLOTS
+        o = torch.stack([self.head_output(out, h, slot) for h in range(N_HEADS)])
+        for t in io.values():
+            ttnn.deallocate(t)
+        return o, dt
+
+
+def ttnn_reference(
+    device: Any,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: Optional[torch.Tensor],
+    scale: float,
+    chunks: Tuple[int, int],
+) -> torch.Tensor:
+    """``ttnn.transformer.scaled_dot_product_attention`` (HiFi2, fp32 acc, exp_approx_mode False, explicit scale) on
+    ``[1, 32, 64, 64] x [1, 32, Sk, 64]`` with the additive ``[1, 1, 64, Sk]`` mask; returns ``[32, 64, 64]`` fp32."""
+    import ttnn
+
+    kw = dict(dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    tq = ttnn.from_torch(q.reshape(1, N_HEADS, M_ROWS, HEAD_DIM), **kw)
+    tk = ttnn.from_torch(k.reshape(1, N_HEADS, -1, HEAD_DIM), **kw)
+    tv = ttnn.from_torch(v.reshape(1, N_HEADS, -1, HEAD_DIM), **kw)
+    tm = ttnn.from_torch(mask.reshape(1, 1, M_ROWS, -1), **kw) if mask is not None else None
+    grid = device.compute_with_storage_grid_size()
+    pc = ttnn.SDPAProgramConfig(
+        compute_with_storage_grid_size=ttnn.CoreCoord(grid.x, grid.y),
+        q_chunk_size=chunks[0],
+        k_chunk_size=chunks[1],
+        exp_approx_mode=False,
+    )
+    ckc = ttnn.WormholeComputeKernelConfig(
+        math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=False, fp32_dest_acc_en=True, packer_l1_acc=False
+    )
+    out = ttnn.transformer.scaled_dot_product_attention(
+        tq, tk, tv, attn_mask=tm, is_causal=False, scale=float(scale), program_config=pc, compute_kernel_config=ckc
+    )
+    res = ttnn.to_torch(out).float().reshape(N_HEADS, M_ROWS, HEAD_DIM)
+    for t in (tq, tk, tv, tm, out):
+        if t is not None:
+            ttnn.deallocate(t)
+    return res
+
+
+class _Session:
+    """One harness per pytest session; ttnn references cached per (version, sk, mask, folded)."""
+
+    harness: Optional[HeadSDPAHarness] = None
+    refs: Dict[Tuple[Any, ...], torch.Tensor] = {}
+
+    @classmethod
+    def get(cls, device: Any) -> HeadSDPAHarness:
+        if cls.harness is None:
+            cls.harness = HeadSDPAHarness(device)
+        return cls.harness
+
+    @classmethod
+    def ttnn_ref(
+        cls,
+        device: Any,
+        key: Tuple[Any, ...],
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        scale: float,
+        chunks: Tuple[int, int],
+    ) -> torch.Tensor:
+        if key not in cls.refs:
+            cls.refs[key] = ttnn_reference(device, q, k, v, mask, scale, chunks)
+        return cls.refs[key]
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device tests
+# --------------------------------------------------------------------------------------------------------------------
+def _cases() -> List[Tuple[str, int, bool, bool, str]]:
+    """(mode, Sk, use_mask, folded, k_layout): self x Sk 64 x mask x scale, cross x every Sk x mask x scale x layout."""
+    cases: List[Tuple[str, int, bool, bool, str]] = []
+    for use_mask, folded in itertools.product((True, False), (False, True)):
+        cases.append(("self", M_ROWS, use_mask, folded, "natural"))
+    for sk, use_mask, folded, layout in itertools.product(SK_LIST, (True, False), (False, True), K_LAYOUTS):
+        cases.append(("cross", sk, use_mask, folded, layout))
+    return cases
+
+
+CASES = _cases()
+CASE_IDS = [
+    f"{m}-sk{sk}-{'mask' if um else 'nomask'}-{'folded' if f else 'explicit'}-{kl}" for m, sk, um, f, kl in CASES
+]
+
+
+@pytest.mark.parametrize("mode,sk,use_mask,folded,k_layout", CASES, ids=CASE_IDS)
+def test_head_sdpa(
+    gr00t_device: Any, results: Any, version: str, mode: str, sk: int, use_mask: bool, folded: bool, k_layout: str
+) -> None:
+    """32 heads through the harness kernel vs torch fp32 (valid rows; masked -> pad-key-excluded reference) and vs
+    ttnn SDPA: PCC >= 0.9995, no NaN over all 64 rows, second launch bit-identical."""
+    hz = _Session.get(gr00t_device)
+    m_logical = M_LOGICAL[version]
+    x = make_inputs(sk, m_logical, SEED + sk)
+    q = fold_scale(x["q"], SCALE) if folded else x["q"]
+    ref_scale = 1.0 if folded else SCALE
+    mask = x["mask"] if use_mask else None
+    nv = n_valid_keys(sk, m_logical)
+    # references: the masked case must equal the reference over the valid keys only (pad keys excluded)
+    if use_mask:
+        ref = torch_reference(q, x["k"][:, :nv], x["v"][:, :nv], None, ref_scale)
+    else:
+        ref = torch_reference(q, x["k"], x["v"], None, ref_scale)
+    ref_key = (version, sk, use_mask, folded)
+    ttnn_err: Optional[str] = None
+    try:
+        ref_ttnn: Optional[torch.Tensor] = _Session.ttnn_ref(
+            gr00t_device, ref_key, q, x["k"], x["v"], mask, ref_scale, TTNN_CHUNKS[sk]
+        )
+    except Exception as exc:  # noqa: BLE001 - recorded in the row, fails the gate
+        ref_ttnn, ttnn_err = None, f"{type(exc).__name__}: {exc}"[:400]
+
+    got, dt = hz.run(
+        x["q"] if not folded else q, x["k"], x["v"], mask, mode=mode, scale=None if folded else SCALE, k_layout=k_layout
+    )
+    got2, _ = hz.run(
+        x["q"] if not folded else q, x["k"], x["v"], mask, mode=mode, scale=None if folded else SCALE, k_layout=k_layout
+    )
+    gotf = got.float()
+    finite = bool(torch.isfinite(gotf).all())
+    valid = slice(0, m_logical)
+    m_torch = pcc_mod.metrics(gotf[:, valid], ref[:, valid])
+    pcc_all_rows = pcc_mod.pcc(gotf, ref) if finite else float("nan")
+    pcc_ttnn = pcc_mod.pcc(gotf[:, valid], ref_ttnn[:, valid]) if (finite and ref_ttnn is not None) else float("nan")
+    pcc_ttnn_vs_torch = pcc_mod.pcc(ref_ttnn[:, valid], ref[:, valid]) if ref_ttnn is not None else float("nan")
+    deterministic = bool(torch.equal(got, got2))
+    ok = bool(finite and m_torch["pcc"] >= GATE_PCC and pcc_ttnn >= GATE_PCC and deterministic)
+    row = {
+        "tap": "mk_k3_head_sdpa",
+        "key": f"{mode}-sk{sk}-{'mask' if use_mask else 'nomask'}-{'folded' if folded else 'explicit'}-{k_layout}",
+        "rule": "pcc",
+        "pcc": m_torch["pcc"],
+        "gate_pcc": GATE_PCC,
+        "pcc_ttnn": pcc_ttnn,
+        "pcc_ttnn_vs_torch": pcc_ttnn_vs_torch,
+        "pcc_all_rows": pcc_all_rows,
+        "max_abs": m_torch["max_abs"],
+        "rel_l2": m_torch["rel_l2"],
+        "finite": finite,
+        "deterministic": deterministic,
+        "mode": mode,
+        "sk": sk,
+        "n_chunks": len(chunk_plan(sk)),
+        "n_valid_keys": nv if use_mask else sk,
+        "m_logical": m_logical,
+        "use_mask": use_mask,
+        "folded": folded,
+        "k_layout": k_layout,
+        "ttnn_chunks": list(TTNN_CHUNKS[sk]),
+        "ttnn_error": ttnn_err,
+        "wall_s": dt,
+        "ok": ok,
+    }
+    results.add_row(row)
+    results.set_meta(compile_s=hz.compile_s, launches=hz.launches, core_map_hub=list(hz.core_map.hub))
+    assert finite, f"non-finite output ({(~torch.isfinite(gotf)).sum().item()} elements)"
+    assert deterministic, "second launch differs from the first"
+    assert (
+        m_torch["pcc"] >= GATE_PCC
+    ), f"PCC vs torch {m_torch['pcc']:.6f} < {GATE_PCC} (max|d| {m_torch['max_abs']:.4f})"
+    assert ttnn_err is None, f"ttnn SDPA reference failed: {ttnn_err}"
+    assert pcc_ttnn >= GATE_PCC, f"PCC vs ttnn SDPA {pcc_ttnn:.6f} < {GATE_PCC} (ttnn vs torch {pcc_ttnn_vs_torch:.6f})"
+
+
+def _timing_cases() -> List[Tuple[str, int, str]]:
+    cases = [("self", M_ROWS, "natural")]
+    cases += [("cross", sk, "transposed") for sk in SK_LIST]
+    cases += [("cross", sk, "natural") for sk in (M_ROWS, 544)]
+    return cases
+
+
+TIMING_CASES = _timing_cases()
+#: in-kernel iteration counts of the difference method: the per-head cost is a few us, so the window must be >= 1 ms
+#: (200 iterations) to sit well above the ~300 us generic_op launch overhead and its jitter (R21, mb3 F5)
+N_SMALL, N_LARGE, REPEATS = 20, 220, 3
+
+
+@pytest.mark.parametrize("mode,sk,k_layout", TIMING_CASES, ids=[f"{m}-sk{sk}-{kl}" for m, sk, kl in TIMING_CASES])
+def test_head_sdpa_timing(gr00t_device: Any, results: Any, version: str, mode: str, sk: int, k_layout: str) -> None:
+    """Per-head µs (all 32 heads run in parallel, so the per-iteration time of the launch is the per-head time) with
+    the mask and the explicit scale, mb3 difference method: ``(t(220) - t(20)) / 200`` in-kernel iterations (every
+    iteration re-delivers K/V/mask like a DiT block would), min of 3."""
+    hz = _Session.get(gr00t_device)
+    m_logical = M_LOGICAL[version]
+    x = make_inputs(sk, m_logical, SEED + sk)
+    ref = torch_reference(
+        x["q"], x["k"][:, : n_valid_keys(sk, m_logical)], x["v"][:, : n_valid_keys(sk, m_logical)], None, SCALE
+    )
+    last: Dict[str, torch.Tensor] = {}
+
+    def launch(n: int) -> float:
+        got, dt = hz.run(x["q"], x["k"], x["v"], x["mask"], mode=mode, scale=SCALE, k_layout=k_layout, n_iters=n)
+        last["got"] = got
+        return dt
+
+    kv_bytes = 2 * N_HEADS * sk * HEAD_DIM * 2  # K + V of all heads (cross: DRAM hoist; self: NoC head exchange)
+    mask_bytes = N_HEADS * M_ROWS * sk * 2  # every head fetches its own copy of the mask from DRAM
+    bytes_per_iter = kv_bytes + mask_bytes
+    rate = P.measure_rate(launch, N_SMALL, N_LARGE, bytes_per_iter, repeats=REPEATS)
+    pcc_last = pcc_mod.pcc(last["got"].float()[:, :m_logical], ref[:, :m_logical])
+    us = rate["per_iter_s"] * 1e6
+    row = {
+        "tap": "mk_k3_timing",
+        "key": f"{mode}-sk{sk}-{k_layout}",
+        "rule": "info",
+        "mode": mode,
+        "sk": sk,
+        "k_layout": k_layout,
+        "n_chunks": len(chunk_plan(sk)),
+        "us_per_head": us,
+        "us_per_chunk": us / len(chunk_plan(sk)),
+        "fetch_gbps": rate["gbps"],
+        "bytes_per_iter": bytes_per_iter,
+        "t_small_s": rate["t_small_s"],
+        "t_large_s": rate["t_large_s"],
+        "n_small": N_SMALL,
+        "n_large": N_LARGE,
+        "pcc_last_iter": pcc_last,
+        "ok": bool(math.isfinite(us) and us > 0 and pcc_last >= GATE_PCC),
+    }
+    results.add_row(row)
+    results.set_timing(**{f"us_{mode}_sk{sk}_{k_layout}": us / 1e6})
+    assert row["ok"], f"timing run: {us:.2f} us/iter, PCC of the last iteration {pcc_last:.6f}"

code/models/experimental/gr00t/tests/tt/test_mk_layernorm.py

@@ -0,0 +1,784 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K2 device tests of the megakernel's ``RowLayerNorm`` (plain LN / dynamic-affine AdaLN / residual-on-receive)
+through its own single-core harness kernel (IMPLEMENTATION_PLAN.md §5.7 K2; kernels/test/ln_harness_kernel.cpp).
+
+* ``test_helpers_cpu`` -- the host helpers (bf16 ulp metric, ``Single`` tile order pack/unpack, references);
+* ``test_ln_plain_golden`` -- ``LN_eps(x)`` on the golden ``sa_embs`` of steps 0 and 3 (bf16-rounded, zero-padded
+  to 64 rows), eps 1e-5 (blocks) and 1e-6 (output head): PCC >= 0.9999 and max|d| <= 2 bf16 ulps vs torch fp32;
+* ``test_adaln_golden`` -- ``LN(x) * (1 + s) + b`` with the in-kernel tilised ROW_MAJOR AdaLN rows, same gates, and
+  the **swapped (s, b) rows must fail** the PCC gate;
+* ``test_residual_adaln`` -- ``x <- x + g`` (g in the P12 ``Single`` gather order, in place on the resident x) then
+  AdaLN: ``x_after`` bit-exact vs ``bf16(x + g)``, ``h`` gated as above;
+* ``test_stress_rows`` -- synthetic rows with large |mean| / std, an outlier channel and a near-constant row;
+* ``test_ln_replicated_96`` -- the same LN on the 96 compute cores of the CoreMap (every core bit-identical);
+* ``test_ln_timing`` -- microseconds per LN (plain / AdaLN / residual + AdaLN) by the in-kernel iteration difference
+  method, plus the ``mk_ln*`` Tracy zones when ``GR00T_MK_PROFILE=1``;
+* ``test_ff_branch_vs_s1`` -- ``x + ff2(GELU(ff1(LN(x))))`` = harness LN -> K1 ``run_single_op`` ff1 / ff2 ->
+  harness residual, vs the Stage-1 ttnn path (``tt.layers``) on device: PCC >= 0.9999.
+
+Run (one pytest process per file, through the device lock)::
+
+    DEVICE_LOCK_TIMEOUT=14400 bin/with-device.sh python -m pytest -q -rA --timeout=1800 -p no:cacheprovider \\
+        models/experimental/gr00t/tests/tt/test_mk_layernorm.py --gr00t-version n16
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+import os
+import struct
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import pytest
+import torch
+
+from models.experimental.gr00t.common import pcc as pcc_mod
+from models.experimental.gr00t.common.configs import get_config
+from models.experimental.gr00t.common.golden import GoldenSet, golden_root
+from models.experimental.gr00t.tests.tt.conftest import have_golden
+
+if os.environ.get("GR00T_MK_PROFILE") == "1":  # before the device fixture imports ttnn
+    os.environ.setdefault("TT_METAL_DEVICE_PROFILER", "1")
+
+torch.set_grad_enabled(False)
+
+KERNEL = "models/experimental/gr00t/tt/megakernel/kernels/test/ln_harness_kernel.cpp"
+TT_METAL_HOME = Path(os.environ.get("TT_METAL_HOME", str(Path(__file__).resolve().parents[5])))
+PROFILE_CSV = TT_METAL_HOME / "generated" / "profiler" / ".logs" / "profile_log_device.csv"
+M_ROWS = 64
+WIDTH = 1536
+TILE = 32
+N_KT = WIDTH // TILE
+N_TILES = 2 * N_KT
+EPS_BLOCK = 1e-5
+EPS_HEAD = 1e-6
+GATE_PCC = 0.9999
+GATE_ULPS = 2.0
+SEED = 11
+STEPS = (0, 3)
+#: CB ids of the harness (kernel named CTAs ln_cb_*)
+CB: Dict[str, int] = {
+    "x": 0,
+    "g": 1,
+    "out": 2,
+    "const": 3,
+    "stat": 4,
+    "stat32": 5,
+    "y32": 6,
+    "adaln_raw": 7,
+    "adaln": 8,
+    "xc": 9,
+}
+G_ORDERS: Dict[str, int] = {"row_major": 0, "single": 2}  # row_layernorm.hpp LNGOrder
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# host helpers
+# --------------------------------------------------------------------------------------------------------------------
+def f32_bits(v: float) -> int:
+    """fp32 bit pattern of ``v`` as an unsigned 32-bit int (kernel scalar arguments)."""
+    return struct.unpack("<I", struct.pack("<f", float(v)))[0]
+
+
+def bf16_ulp(t: torch.Tensor) -> torch.Tensor:
+    """bf16 spacing at the magnitude of each element (``2^(floor(log2|t|) - 7)``; zero -> the smallest normal ulp)."""
+    a = t.detach().abs().double()
+    e = torch.floor(torch.log2(torch.clamp(a, min=2.0**-126)))
+    return torch.pow(2.0, e - 7)
+
+
+def ulp_metrics(got: torch.Tensor, ref: torch.Tensor) -> Dict[str, float]:
+    """``max|d|`` in bf16 ulps: elementwise at ``max(|ref_i|, rms(ref))`` (the acceptance metric: an element far
+    below the row's scale is judged at the ulp of the scale) and the plain per-element version for information."""
+    g, r = got.double().reshape(-1), ref.double().reshape(-1)
+    if g.numel() != r.numel():
+        raise ValueError(f"numel mismatch {g.numel()} vs {r.numel()}")
+    d = (g - r).abs()
+    rms = float(torch.sqrt(torch.mean(r * r)))
+    floor = torch.full_like(r, rms)
+    ulp_scale = bf16_ulp(torch.maximum(r.abs(), floor))
+    return {
+        "max_abs": float(d.max()),
+        "max_ulps": float((d / ulp_scale).max()),
+        "mean_ulps": float((d / ulp_scale).mean()),
+        "max_ulps_elementwise": float((d / bf16_ulp(r)).max()),
+        "rms_ref": rms,
+    }
+
+
+def bf16_round_ties_away(t: torch.Tensor) -> torch.Tensor:
+    """fp32 -> bf16 with ties rounded away from zero (the Tensix packer's fp32 -> bf16 conversion, measured in
+    test_residual_adaln: torch rounds ties to even), returned as fp32 values."""
+    b = t.float().contiguous().view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    upper = b >> 16
+    lower = b & 0xFFFF
+    upper = torch.where(lower >= 0x8000, upper + 1, upper)
+    return ((upper & 0xFFFF) << 16).to(torch.int32).view(torch.float32)
+
+
+def ref_layernorm(
+    x: torch.Tensor, eps: float, one_plus_scale: Optional[torch.Tensor] = None, shift: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """fp32 ``layer_norm(x) [* (1 + s) + b]`` of a bf16-representable ``x`` (reference.action_head.AdaLayerNorm),
+    rounded to bf16 (the device output dtype)."""
+    if x.shape != (M_ROWS, WIDTH):
+        raise ValueError(f"x must be [64, 1536], got {tuple(x.shape)}")
+    if (one_plus_scale is None) != (shift is None):
+        raise ValueError("one_plus_scale and shift must be given together")
+    y = torch.nn.functional.layer_norm(x.float(), (WIDTH,), eps=float(eps))
+    if one_plus_scale is not None:
+        y = y * one_plus_scale.float().reshape(1, WIDTH) + shift.float().reshape(1, WIDTH)
+    return y.to(torch.bfloat16).float()
+
+
+def to_single_order(g: torch.Tensor) -> torch.Tensor:
+    """[64, 1536] -> the [64, 1536] TILE-layout *shard image* whose page ``2k + r`` holds tile ``(r, k)`` (the P12
+    gather placement, dit_program.py ``Single`` order)."""
+    if g.shape != (M_ROWS, WIDTH):
+        raise ValueError(f"g must be [64, 1536], got {tuple(g.shape)}")
+    out = torch.empty_like(g)
+    for k in range(N_KT):
+        for r in range(2):
+            p = 2 * k + r
+            R, C = divmod(p, N_KT)
+            out[R * TILE : (R + 1) * TILE, C * TILE : (C + 1) * TILE] = g[
+                r * TILE : (r + 1) * TILE, k * TILE : (k + 1) * TILE
+            ]
+    return out
+
+
+def from_single_order(shard: torch.Tensor) -> torch.Tensor:
+    """Inverse of :func:`to_single_order`."""
+    if shard.shape != (M_ROWS, WIDTH):
+        raise ValueError(f"shard must be [64, 1536], got {tuple(shard.shape)}")
+    out = torch.empty_like(shard)
+    for k in range(N_KT):
+        for r in range(2):
+            p = 2 * k + r
+            R, C = divmod(p, N_KT)
+            out[r * TILE : (r + 1) * TILE, k * TILE : (k + 1) * TILE] = shard[
+                R * TILE : (R + 1) * TILE, C * TILE : (C + 1) * TILE
+            ]
+    return out
+
+
+def golden_sa_embs(version: str, step: int) -> torch.Tensor:
+    """``sa_embs[k=step]`` of the canonical golden sample, bf16-rounded and zero-padded to 64 rows (fp32 values)."""
+    gs = GoldenSet(version)
+    t = gs.load("sa_embs", "fp32", step=step)
+    if t.dim() != 3 or t.shape[0] != 1 or t.shape[2] != WIDTH or t.shape[1] > M_ROWS:
+        raise ValueError(f"{version}: sa_embs shape {tuple(t.shape)} unexpected")
+    x = torch.zeros(M_ROWS, WIDTH)
+    x[: t.shape[1]] = t[0]
+    return x.to(torch.bfloat16).float()
+
+
+def random_adaln(seed: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """``(1 + s, b)`` bf16 vectors of the magnitude of the checkpoint AdaLN outputs (s, b ~ N(0, 0.3^2))."""
+    gen = torch.Generator().manual_seed(seed)
+    s = 0.3 * torch.randn(WIDTH, generator=gen)
+    b = 0.3 * torch.randn(WIDTH, generator=gen)
+    return (1.0 + s).to(torch.bfloat16).float(), b.to(torch.bfloat16).float()
+
+
+def stress_rows() -> torch.Tensor:
+    """Synthetic [64, 1536]: N(0,1) rows plus a large-|mean| row, an outlier-channel row (mb2 §B.3 case), a
+    near-constant row, a large-scale row and a zero row -- bf16-rounded."""
+    gen = torch.Generator().manual_seed(3)
+    x = torch.randn(M_ROWS, WIDTH, generator=gen)
+    x[1] = 100.0 + torch.randn(WIDTH, generator=gen)  # |mean| / std = 100
+    x[2, 1279] = 1.53e4  # massive activation channel (mb2 §B.3 magnitudes, moved inside the 1536 width)
+    x[2, 1500] = -7.8e3
+    x[3] = 5.0 + 1e-2 * torch.randn(WIDTH, generator=gen)  # near-constant
+    x[4] = 300.0 * torch.randn(WIDTH, generator=gen)  # large scale (sa_embs absmax ~ 400-1000)
+    x[5] = 0.0  # zero row (a padded query row): LN -> 0
+    x[6] = 1e-3 * torch.randn(WIDTH, generator=gen)  # tiny scale
+    return x.to(torch.bfloat16).float()
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# the launcher
+# --------------------------------------------------------------------------------------------------------------------
+@dataclass
+class LNResult:
+    h: torch.Tensor  # [n_cores*64, 1536] bf16 -> fp32 (per-core LN outputs stacked)
+    x_after: torch.Tensor  # the resident x after the launch (== x + g with a residual)
+    wall_s: float
+
+
+class LNHarness:
+    """Launch ``ln_harness_kernel.cpp`` on ``cores`` (each core runs the same LN on its own replica of x).
+
+    Every call builds the program afresh (tensor-backed CB addresses change between allocations) and frees its
+    device tensors before returning; the kernel binary is JIT-cached per (CTAs) by tt-metal.
+    """
+
+    def __init__(self, device: Any, cores: Sequence[Tuple[int, int]]):
+        import ttnn  # lazy
+
+        self.ttnn = ttnn
+        self.device = device
+        self.cores = [(int(x), int(y)) for x, y in cores]
+        if not self.cores or len(set(self.cores)) != len(self.cores):
+            raise ValueError("cores must be a non-empty list of distinct logical cores")
+        self.crs = ttnn.CoreRangeSet(
+            [
+                ttnn.CoreRange(ttnn.CoreCoord(x, y), ttnn.CoreCoord(x, y))
+                for x, y in sorted(self.cores, key=lambda c: (c[1], c[0]))
+            ]
+        )
+        self.shard_cores = sorted(self.cores, key=lambda c: (c[1], c[0]))  # HEIGHT shard i <-> this core
+        self.launches = 0
+
+    # ------------------------------------------------------------------ tensors
+    def _sharded(self, t: torch.Tensor) -> Any:
+        ttnn = self.ttnn
+        spec = ttnn.ShardSpec(self.crs, [M_ROWS, WIDTH], ttnn.ShardOrientation.ROW_MAJOR)
+        mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, spec)
+        return ttnn.from_torch(
+            t.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.device, memory_config=mem
+        )
+
+    def _cb(self, name: str, total: int, page: int, fmt: Any) -> Any:
+        ttnn = self.ttnn
+        f = ttnn.CBFormatDescriptor(
+            buffer_index=CB[name], data_format=fmt, page_size=page, tile=ttnn.TileDescriptor(ttnn.Tile([TILE, TILE]))
+        )
+        return ttnn.CBDescriptor(total_size=total, core_ranges=self.crs, format_descriptors=[f])
+
+    # ------------------------------------------------------------------ launch
+    def run(
+        self,
+        x: torch.Tensor,
+        *,
+        eps: float,
+        one_plus_scale: Optional[torch.Tensor] = None,
+        shift: Optional[torch.Tensor] = None,
+        g: Optional[torch.Tensor] = None,
+        g_order: str = "single",
+        n_iters: int = 1,
+        swap_rows: bool = False,
+        xc_fp32: bool = True,
+    ) -> LNResult:
+        """LN of ``x`` [64, 1536] (bf16-representable) on every core; ``one_plus_scale``/``shift`` [1536] enable the
+        dynamic affine (fetched from a ROW_MAJOR DRAM table like the arena's ``adaln_rows``), ``g`` [64, 1536] the
+        in-place residual add (uploaded in ``g_order``), ``n_iters`` repeats the LN in-kernel (timing; with a residual
+        every iteration adds g again).  ``xc_fp32`` keeps the centred row scratch in fp32 (192 KB, <= 0.125 ulp of
+        centring error) instead of bf16 (96 KB, <= 0.5 ulp)."""
+        ttnn = self.ttnn
+        from models.demos.deepseek_v3_b1.unified_kernel_descriptor import (
+            UnifiedCompileTimeCoreDescriptor,
+            UnifiedKernelDescriptor,
+        )
+
+        if x.shape != (M_ROWS, WIDTH):
+            raise ValueError(f"x must be [64, 1536], got {tuple(x.shape)}")
+        if (one_plus_scale is None) != (shift is None):
+            raise ValueError("one_plus_scale and shift must be given together")
+        if g_order not in G_ORDERS:
+            raise KeyError(f"g_order must be one of {sorted(G_ORDERS)}")
+        if n_iters < 1:
+            raise ValueError("n_iters must be >= 1")
+        n = len(self.cores)
+        affine = one_plus_scale is not None
+        residual = g is not None
+        x_rep = self._sharded(x.repeat(n, 1))
+        g_img = torch.zeros(M_ROWS, WIDTH)
+        if residual:
+            if g.shape != (M_ROWS, WIDTH):
+                raise ValueError(f"g must be [64, 1536], got {tuple(g.shape)}")
+            g_img = to_single_order(g) if g_order == "single" else g.clone()
+        g_rep = self._sharded(g_img.repeat(n, 1))
+        out_rep = self._sharded(torch.zeros(n * M_ROWS, WIDTH))
+        # AdaLN table: 4 ROW_MAJOR rows [pad, 1+s, b, pad] (rows 1 / 2; swapped -> 2 / 1)
+        table = torch.zeros(4, WIDTH)
+        if affine:
+            table[1] = one_plus_scale.reshape(WIDTH)
+            table[2] = shift.reshape(WIDTH)
+        adaln_t = ttnn.from_torch(
+            table.to(torch.bfloat16),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=self.device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+        row_bytes = int(adaln_t.buffer_page_size())
+        if row_bytes != WIDTH * 2:
+            raise RuntimeError(f"adaln table page {row_bytes} B != one row ({WIDTH * 2} B): the kernel reads rows")
+        row_scale, row_shift = (2, 1) if swap_rows else (1, 2)
+
+        cbs = [
+            ttnn.cb_descriptor_from_sharded_tensor(CB["x"], x_rep),
+            ttnn.cb_descriptor_from_sharded_tensor(CB["g"], g_rep),
+            ttnn.cb_descriptor_from_sharded_tensor(CB["out"], out_rep),
+            self._cb("const", 3 * 2048, 2048, ttnn.bfloat16),
+            self._cb("stat", 2 * 4096, 4096, ttnn.float32),
+            self._cb("stat32", 2 * 4096, 4096, ttnn.float32),
+            self._cb("y32", 8 * 4096, 4096, ttnn.float32),
+            self._cb("adaln_raw", 4 * 2048, 2048, ttnn.bfloat16),
+            self._cb("adaln", 4 * 2048, 2048, ttnn.bfloat16),
+            (
+                self._cb("xc", N_KT * 4096, 4096, ttnn.float32)
+                if xc_fp32
+                else self._cb("xc", N_KT * 2048, 2048, ttnn.bfloat16)
+            ),
+        ]
+        named = [(f"ln_cb_{k}", v) for k, v in CB.items()] + [("ln_n_kt", N_KT)]
+        common = [
+            int(n_iters),
+            1 if affine else 0,
+            1 if residual else 0,
+            G_ORDERS[g_order],
+            f32_bits(eps),
+            f32_bits(1.0 / WIDTH),
+            int(adaln_t.buffer_address()),
+            row_scale,
+            row_shift,
+            row_bytes,
+        ]
+        cc = ttnn.ComputeConfigDescriptor(
+            math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False, fp32_dest_acc_en=True, dst_full_sync_en=False
+        )
+        modes = [ttnn.UnpackToDestMode.Default] * 64
+        modes[CB["stat32"]] = ttnn.UnpackToDestMode.UnpackToDestFp32  # copy_tile lands exact fp32 (row_layernorm.hpp)
+        cc.unpack_to_dest_mode = modes
+        uk = UnifiedKernelDescriptor(
+            kernel_source=KERNEL,
+            core_ranges=self.crs,
+            ncrisc_named_compile_time_args=named,
+            brisc_named_compile_time_args=named,
+            trisc_named_compile_time_args=named,
+            ncrisc_common_runtime_args=common,
+            brisc_common_runtime_args=common,
+            trisc_common_runtime_args=common,
+            trisc_compute_config=cc,
+            unified_compile_time_core_descriptors=[
+                UnifiedCompileTimeCoreDescriptor(
+                    named_compile_time_arg="is_active_core", core_range=self.crs, value=1, other_value=0
+                )
+            ],
+        )
+        program = ttnn.ProgramDescriptor(kernels=uk.get_kernel_descriptors().kernels, cbs=cbs)
+        t0 = time.perf_counter()
+        ttnn.generic_op([x_rep, g_rep, adaln_t, out_rep], program)
+        ttnn.synchronize_device(self.device)
+        wall = time.perf_counter() - t0
+        self.launches += 1
+        h = ttnn.to_torch(out_rep).float()
+        x_after = ttnn.to_torch(x_rep).float()
+        for t in (x_rep, g_rep, out_rep, adaln_t):
+            ttnn.deallocate(t)
+        return LNResult(h=h, x_after=x_after, wall_s=wall)
+
+    def per_core(self, t: torch.Tensor, i: int) -> torch.Tensor:
+        return t[i * M_ROWS : (i + 1) * M_ROWS]
+
+
+def measure_us(launch: Any, n_small: int = 2, n_large: int = 6, repeats: int = 3) -> Dict[str, float]:
+    """dit_program.measure_rate's difference method: per-iteration seconds = (t(n_large) - t(n_small)) / (n_large -
+    n_small), min over ``repeats`` (dispatch and readback cancel)."""
+    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 = (t_large - t_small) / (n_large - n_small)
+    return {"t_small_s": t_small, "t_large_s": t_large, "per_iter_us": per * 1e6}
+
+
+def read_ln_zones(path: Path = PROFILE_CSV) -> Dict[str, Dict[str, float]]:
+    """Mean duration (us) of the ``mk_ln*`` TRISC zones of the device profiler CSV over all launches (bench_mk_stream
+    read_zones, restricted to this kernel's zones)."""
+    if not path.is_file():
+        raise FileNotFoundError(f"{path} not found (GR00T_MK_PROFILE=1 sets TT_METAL_DEVICE_PROFILER)")
+    with open(path) as fh:
+        header = fh.readline()
+        freq_mhz = float(header.split("CHIP_FREQ[MHz]:")[1].split(",")[0])
+        reader = csv.reader(fh)
+        cols = [c.strip() for c in next(reader)]
+        idx = {c: i for i, c in enumerate(cols)}
+        starts: Dict[tuple, int] = {}
+        durations: Dict[str, List[int]] = defaultdict(list)
+        for row in reader:
+            if len(row) < len(cols) - 1:
+                continue
+            zone = row[idx["zone name"]].strip()
+            if not zone.startswith("mk_ln"):
+                continue
+            key = (
+                row[idx["run host ID"]].strip(),
+                row[idx["core_x"]],
+                row[idx["core_y"]],
+                row[idx["RISC processor type"]].strip(),
+                zone,
+            )
+            t = int(row[idx["time[cycles since reset]"]])
+            typ = row[idx["type"]].strip()
+            if typ == "ZONE_START":
+                starts[key] = t
+            elif typ == "ZONE_END" and key in starts:
+                durations[f"{key[3]}:{zone}"].append(t - starts.pop(key))
+    return {
+        name: {"count": len(ds), "mean_us": sum(ds) / len(ds) / freq_mhz, "max_us": max(ds) / freq_mhz}
+        for name, ds in durations.items()
+        if ds
+    }
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# CPU
+# --------------------------------------------------------------------------------------------------------------------
+def test_helpers_cpu() -> None:
+    g = torch.arange(M_ROWS * WIDTH, dtype=torch.float32).reshape(M_ROWS, WIDTH)
+    s = to_single_order(g)
+    assert not torch.equal(s, g) and torch.equal(from_single_order(s), g)
+    # page 2k + r holds tile (r, k): tile (1, 0) -> page 1 -> shard tile (0, 1)
+    assert torch.equal(s[:TILE, TILE : 2 * TILE], g[TILE:, :TILE])
+    ref = torch.tensor([1.0, 1.5, 3.0, 0.001, -2.0])
+    assert torch.allclose(
+        bf16_ulp(ref), torch.tensor([2.0**-7, 2.0**-7, 2.0**-6, 2.0**-17, 2.0**-6], dtype=torch.float64)
+    )
+    got = ref.clone()
+    got[0] += 2.0**-7  # 1 ulp at 1.0
+    m = ulp_metrics(got, ref)
+    assert abs(m["max_ulps_elementwise"] - 1.0) < 1e-6 and m["max_ulps"] <= 1.0 + 1e-6
+    assert f32_bits(1.0) == 0x3F800000 and f32_bits(1e-5) == 0x3727C5AC
+    x = stress_rows()
+    y = ref_layernorm(x, EPS_BLOCK)
+    assert torch.isfinite(y).all() and float(y[5].abs().max()) == 0.0
+    ops, sh = random_adaln(SEED)
+    y2 = ref_layernorm(x, EPS_BLOCK, ops, sh)
+    assert pcc_mod.pcc(y2, ref_layernorm(x, EPS_BLOCK, sh, ops)) < 0.95  # swapped rows are a different function
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device
+# --------------------------------------------------------------------------------------------------------------------
+def _harness(device: Any, n_cores: int = 1) -> LNHarness:
+    from models.experimental.gr00t.tt.megakernel.core_map import CoreMap
+
+    cm = CoreMap.from_device(device)
+    return LNHarness(device, list(cm.compute[:n_cores]))
+
+
+def _need_golden(version: str) -> None:
+    if not golden_root().is_dir() or not have_golden(version):
+        pytest.skip(f"{version}: canonical golden sample not present under {golden_root()}")
+
+
+def _row(
+    results: Any, key: str, got: torch.Tensor, ref: torch.Tensor, expect_pass: bool = True, **extra: Any
+) -> Dict[str, Any]:
+    m = pcc_mod.metrics(got.float(), ref.float())
+    u = ulp_metrics(got, ref)
+    ok = bool(m["finite"] and m["pcc"] >= GATE_PCC and u["max_ulps"] <= GATE_ULPS)
+    row = {
+        "tap": "mk_k2_ln",
+        "key": key,
+        "rule": "pcc+ulp",
+        "pcc": m["pcc"],
+        "gate_pcc": GATE_PCC,
+        "gate_ulps": GATE_ULPS,
+        "max_abs": u["max_abs"],
+        "max_ulps": u["max_ulps"],
+        "mean_ulps": u["mean_ulps"],
+        "max_ulps_elementwise": u["max_ulps_elementwise"],
+        "rms_ref": u["rms_ref"],
+        "finite": m["finite"],
+        "expect_pass": expect_pass,
+        "ok": ok if expect_pass else (not ok),
+    }
+    row.update(extra)
+    results.add_row(row)
+    return row
+
+
+@pytest.mark.parametrize("eps", (EPS_BLOCK, EPS_HEAD), ids=("eps1e-5", "eps1e-6"))
+@pytest.mark.parametrize("step", STEPS, ids=[f"k{s}" for s in STEPS])
+def test_ln_plain_golden(gr00t_device: Any, results: Any, version: str, step: int, eps: float) -> None:
+    """Plain LN of the golden sa_embs (P8 / output-head statistics) vs torch fp32: PCC >= 0.9999, <= 2 bf16 ulps."""
+    _need_golden(version)
+    x = golden_sa_embs(version, step)
+    hs = _harness(gr00t_device)
+    res = hs.run(x, eps=eps)
+    ref = ref_layernorm(x, eps)
+    row = _row(results, f"ln-{version}-k{step}-eps{eps:g}", res.h, ref, step=step, eps=eps, wall_s=res.wall_s)
+    assert torch.equal(res.x_after, x), "x must be untouched without a residual"
+    assert row["finite"], "non-finite LN output"
+    assert row["pcc"] >= GATE_PCC, f"PCC {row['pcc']:.6f} < {GATE_PCC}"
+    assert (
+        row["max_ulps"] <= GATE_ULPS
+    ), f"max |d| {row['max_ulps']:.2f} bf16 ulps > {GATE_ULPS} (abs {row['max_abs']:.4g})"
+
+
+@pytest.mark.parametrize("step", STEPS, ids=[f"k{s}" for s in STEPS])
+def test_adaln_golden(gr00t_device: Any, results: Any, version: str, step: int) -> None:
+    """AdaLN with in-kernel tilised ROW_MAJOR (1+s, b) rows on the golden sa_embs; swapped rows must fail."""
+    _need_golden(version)
+    x = golden_sa_embs(version, step)
+    ops, sh = random_adaln(SEED + step)
+    hs = _harness(gr00t_device)
+    res = hs.run(x, eps=EPS_BLOCK, one_plus_scale=ops, shift=sh)
+    ref = ref_layernorm(x, EPS_BLOCK, ops, sh)
+    row = _row(results, f"adaln-{version}-k{step}", res.h, ref, step=step, wall_s=res.wall_s)
+    swapped = hs.run(x, eps=EPS_BLOCK, one_plus_scale=ops, shift=sh, swap_rows=True)
+    row_sw = _row(results, f"adaln-swapped-{version}-k{step}", swapped.h, ref, expect_pass=False, step=step)
+    # the swapped launch computes LN * b + (1 + s): it must match *that* reference, i.e. the kernel reads the rows it
+    # is told to read
+    ref_sw = ref_layernorm(x, EPS_BLOCK, sh, ops)
+    row_sw2 = _row(results, f"adaln-swapped-selfcheck-{version}-k{step}", swapped.h, ref_sw, step=step)
+    assert row["finite"] and row["pcc"] >= GATE_PCC, f"AdaLN PCC {row['pcc']:.6f} < {GATE_PCC}"
+    assert row["max_ulps"] <= GATE_ULPS, f"AdaLN max |d| {row['max_ulps']:.2f} ulps > {GATE_ULPS}"
+    assert row_sw["pcc"] < GATE_PCC, f"swapped (scale, shift) still passes: PCC {row_sw['pcc']:.6f}"
+    assert row_sw2["ok"], "swapped rows do not match the swapped reference"
+
+
+@pytest.mark.parametrize("g_order", ("single", "row_major"))
+def test_residual_adaln(gr00t_device: Any, results: Any, version: str, g_order: str) -> None:
+    """x <- x + g in place (g in the P12 Single order or RowMajor), then AdaLN (P1 residual-on-receive).
+
+    The in-place sum is compared with torch's ``bf16(x + g)`` (round-to-nearest-even): the Tensix FPU add + fp32 ->
+    bf16 pack differs from RNE on a few percent of near-tie sums by exactly one bf16 ulp (measured: ties round away
+    from zero, plus a class of just-below-half sums that round up), so the gate is <= 1 ulp elementwise with the
+    fraction of differing elements recorded; the LN of the *device* sum is gated as every other LN."""
+    _need_golden(version)
+    x = golden_sa_embs(version, 0)
+    gen = torch.Generator().manual_seed(SEED)
+    g = (0.5 * x.std() * torch.randn(M_ROWS, WIDTH, generator=gen)).to(torch.bfloat16).float()
+    ops, sh = random_adaln(SEED)
+    hs = _harness(gr00t_device)
+    res = hs.run(x, eps=EPS_BLOCK, one_plus_scale=ops, shift=sh, g=g, g_order=g_order)
+    x_rne = (x + g).to(torch.bfloat16).float()
+    u = ulp_metrics(res.x_after, x_rne)
+    n_diff = int((res.x_after != x_rne).sum())
+    sum_ok = bool(u["max_ulps_elementwise"] <= 1.0 and n_diff <= 0.1 * x.numel())
+    results.add_row(
+        {
+            "tap": "mk_k2_ln",
+            "key": f"residual-x-{version}-{g_order}",
+            "rule": "ulp",
+            "ok": sum_ok,
+            "max_ulps_elementwise": u["max_ulps_elementwise"],
+            "n_diff_vs_rne": n_diff,
+            "frac_diff_vs_rne": n_diff / x.numel(),
+            "max_abs": u["max_abs"],
+        }
+    )
+    ref = ref_layernorm(res.x_after, EPS_BLOCK, ops, sh)  # LN of the sum the device actually formed
+    row = _row(results, f"residual-adaln-{version}-{g_order}", res.h, ref, wall_s=res.wall_s)
+    ref_rne = ref_layernorm(x_rne, EPS_BLOCK, ops, sh)
+    _row(results, f"residual-adaln-vs-rne-sum-{version}-{g_order}", res.h, ref_rne)
+    assert sum_ok, f"in-place residual: {n_diff} elements differ from bf16(x + g), max {u['max_ulps_elementwise']} ulps"
+    assert row["ok"], f"residual AdaLN: PCC {row['pcc']:.6f}, {row['max_ulps']:.2f} ulps"
+
+
+def test_stress_rows(gr00t_device: Any, results: Any) -> None:
+    """Synthetic rows: |mean|/std = 100, the mb2 §B.3 outlier channels, near-constant, x300 scale, zero, 1e-3 scale."""
+    x = stress_rows()
+    hs = _harness(gr00t_device)
+    for eps in (EPS_BLOCK, EPS_HEAD):
+        res = hs.run(x, eps=eps)
+        ref = ref_layernorm(x, eps)
+        row = _row(results, f"stress-eps{eps:g}", res.h, ref, eps=eps)
+        # per stress row detail
+        per_row = {int(r): ulp_metrics(res.h[r], ref[r])["max_ulps"] for r in range(7)}
+        results.set_meta(**{f"stress_eps{eps:g}_row_ulps": per_row})
+        assert row["finite"], "non-finite"
+        assert float(res.h[5].abs().max()) == 0.0, "the zero row must normalise to exactly 0"
+        assert row["pcc"] >= GATE_PCC, f"stress PCC {row['pcc']:.6f}"
+        assert row["max_ulps"] <= GATE_ULPS, f"stress max |d| {row['max_ulps']:.2f} ulps (rows {per_row})"
+
+
+def test_ln_replicated_96(gr00t_device: Any, results: Any, version: str) -> None:
+    """The same AdaLN on all 96 compute cores of the CoreMap: every core's result is bit-identical to core 0."""
+    _need_golden(version)
+    x = golden_sa_embs(version, 0)
+    ops, sh = random_adaln(SEED)
+    hs = _harness(gr00t_device, n_cores=96)
+    res = hs.run(x, eps=EPS_BLOCK, one_plus_scale=ops, shift=sh)
+    ref = ref_layernorm(x, EPS_BLOCK, ops, sh)
+    h0 = hs.per_core(res.h, 0)
+    identical = all(torch.equal(hs.per_core(res.h, i), h0) for i in range(96))
+    row = _row(results, f"adaln-96cores-{version}", h0, ref, identical=identical, wall_s=res.wall_s)
+    results.add_row({"tap": "mk_k2_ln", "key": f"adaln-96cores-identical-{version}", "rule": "exact", "ok": identical})
+    assert identical, "compute cores disagree on the replicated LN"
+    assert row["ok"]
+
+
+def test_ln_timing(gr00t_device: Any, results: Any, version: str) -> None:
+    """Microseconds per LN on one core (plain / AdaLN / residual + AdaLN): in-kernel iteration difference method;
+    Tracy zones when GR00T_MK_PROFILE=1."""
+    _need_golden(version)
+    x = golden_sa_embs(version, 0)
+    ops, sh = random_adaln(SEED)
+    gen = torch.Generator().manual_seed(SEED)
+    g = (0.01 * torch.randn(M_ROWS, WIDTH, generator=gen)).to(torch.bfloat16).float()
+    hs = _harness(gr00t_device)
+    if PROFILE_CSV.exists() and os.environ.get("TT_METAL_DEVICE_PROFILER") == "1":
+        PROFILE_CSV.unlink()
+    timings: Dict[str, Dict[str, float]] = {}
+    variants = {
+        "plain": dict(eps=EPS_BLOCK),
+        "adaln": dict(eps=EPS_BLOCK, one_plus_scale=ops, shift=sh),
+        "residual_adaln": dict(eps=EPS_BLOCK, one_plus_scale=ops, shift=sh, g=g),
+    }
+    for name, kw in variants.items():
+        t = measure_us(lambda n, kw=kw: hs.run(x, n_iters=n, **kw).wall_s)
+        timings[name] = t
+        results.add_row(
+            {
+                "tap": "mk_k2_ln",
+                "key": f"us-{name}-{version}",
+                "rule": "time",
+                "us": t["per_iter_us"],
+                "ok": bool(0 < t["per_iter_us"] < 1e4),
+                **t,
+            }
+        )
+    # the iterated residual is deterministic: x + 6 g in bf16 steps
+    res6 = hs.run(x, n_iters=6, **variants["residual_adaln"])
+    xi = x.clone()
+    for _ in range(6):
+        xi = (xi + g).to(torch.bfloat16).float()
+    # sanity of the in-place recurrence only (the single-step gate is test_residual_adaln): per step the device sum
+    # may differ from RNE by one ulp of |x| and cross a binade (g = 0.01 is far below the ulp of the largest |x|),
+    # so allow 2 ulps of |x| per step at the row scale; the unfloored per-element ulp is meaningless on the
+    # exactly-zero padded rows
+    u6 = ulp_metrics(res6.x_after, xi)
+    iter_ok = bool(u6["max_ulps"] <= 12.0 and pcc_mod.pcc(res6.x_after, xi) >= 0.99999)
+    results.add_row(
+        {
+            "tap": "mk_k2_ln",
+            "key": f"residual-6iters-{version}",
+            "rule": "ulp",
+            "ok": iter_ok,
+            "n_diff_vs_rne": int((res6.x_after != xi).sum()),
+            **{k: v for k, v in u6.items()},
+        }
+    )
+    zones: Dict[str, Any] = {}
+    if os.environ.get("TT_METAL_DEVICE_PROFILER") == "1":
+        hs.ttnn.ReadDeviceProfiler(gr00t_device)
+        try:
+            zones = read_ln_zones()
+        except FileNotFoundError as exc:
+            zones = {"error": str(exc)}
+    results.set_meta(timings=timings, zones=zones, profiler=os.environ.get("TT_METAL_DEVICE_PROFILER", "0"))
+    assert iter_ok, f"6 iterated residual adds drift {u6['max_ulps']} ulps (row scale) from the bf16 recurrence"
+    for name, t in timings.items():
+        assert 0 < t["per_iter_us"] < 1e4, f"{name}: implausible {t['per_iter_us']} us"
+
+
+def test_ff_branch_vs_s1(gr00t_device: Any, results: Any, version: str) -> None:
+    """FF branch x + ff2(GELU_tanh(ff1(LN(x)))): harness LN -> K1 BlockMatmul (run_single_op ff1, ff2) -> harness
+    residual, vs the Stage-1 ttnn path (tt.layers layer_norm / linear / residual_add) on device: PCC >= 0.9999."""
+    _need_golden(version)
+    from models.experimental.gr00t.tt import layers as L
+    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 import dit_program as P
+    from models.experimental.gr00t.tt.megakernel.core_map import CoreMap
+
+    ttnn = __import__("ttnn")
+    device = gr00t_device
+    x = golden_sa_embs(version, 0)
+    blk = 1  # the self block of the 2-block test arena
+    hs = _harness(device)
+    # 1. LN on the megakernel path
+    h2 = hs.run(x, eps=EPS_BLOCK).h
+    # 2. ff1 (+GELU_tanh) and ff2 with K1's BlockMatmul on a random 2-block arena
+    layout = A.ArenaLayout.plan(version, n_blocks=2, dtype_stream="bf16", reader_mode="direct", tile_order="row_major")
+    pack = layout.pack(P.random_arena_plan(layout, seed=SEED))
+    k_block_bytes = max(op.k_block_bytes for op in layout.stream_ops("stream"))
+    mk = P.DiTMegakernel(
+        get_config(version),
+        None,
+        None,
+        layout,
+        CoreMap.from_device(device),
+        device,
+        cb_budget=D.CBBudget(cb_w=3 * k_block_bytes),
+    )
+    try:
+        mk.upload(pack)
+        t0 = time.perf_counter()
+        f = mk.run_single_op(h2, blk, "ff1")
+        g = mk.run_single_op(f, blk, "ff2")
+        t_mm = time.perf_counter() - t0
+        w1, b1 = P.op_weight_and_bias(layout, pack, blk, "ff1")
+        w2, b2 = P.op_weight_and_bias(layout, pack, blk, "ff2")
+    finally:
+        mk.deallocate()
+    # 3. residual on the megakernel path (g in the P12 Single order)
+    res = hs.run(x, eps=EPS_BLOCK, g=g, g_order="single")
+    out_mk = res.x_after
+    # 4. Stage-1 ttnn path
+    xt = L.to_device(x.reshape(1, 1, M_ROWS, WIDTH), device, dtype="bf16", mem=L.L1)
+    w1t = L.to_device(w1.to(torch.bfloat16), device, dtype="bf16", mem=L.DRAM)
+    b1t = L.to_device(b1.to(torch.bfloat16).reshape(1, -1), device, dtype="bf16", mem=L.DRAM)
+    w2t = L.to_device(w2.to(torch.bfloat16), device, dtype="bf16", mem=L.DRAM)
+    b2t = L.to_device(b2.to(torch.bfloat16).reshape(1, -1), device, dtype="bf16", mem=L.DRAM)
+    h_s1 = L.layer_norm(xt, eps=EPS_BLOCK)
+    f_s1 = L.linear(h_s1, w1t, b1t, act="gelu_tanh")
+    g_s1 = L.linear(f_s1, w2t, b2t)
+    out_s1 = L.residual_add(xt, g_s1)
+    h_s1_t = L.to_torch(h_s1).reshape(M_ROWS, WIDTH)
+    g_s1_t = L.to_torch(g_s1).reshape(M_ROWS, WIDTH)
+    out_s1_t = L.to_torch(out_s1).reshape(M_ROWS, WIDTH)
+    for t in (xt, w1t, b1t, w2t, b2t, h_s1, f_s1, g_s1, out_s1):
+        ttnn.deallocate(t)
+    # 5. torch fp32 chain (bf16 rounding between ops, as both device paths)
+    bf = lambda t: t.to(torch.bfloat16).float()  # noqa: E731
+    h_ref = ref_layernorm(x, EPS_BLOCK)
+    f_ref = bf(torch.nn.functional.gelu(h_ref @ w1 + b1, approximate="tanh"))
+    g_ref = bf(f_ref @ w2 + b2)
+    out_ref = bf(x + g_ref)
+    rows = {
+        "ff-ln-vs-s1": (h2, h_s1_t),
+        "ff-g-vs-s1": (g, g_s1_t),
+        "ff-out-vs-s1": (out_mk, out_s1_t),
+        "ff-out-vs-torch": (out_mk, out_ref),
+        "ff-s1-out-vs-torch": (out_s1_t, out_ref),
+    }
+    got_rows = {}
+    # informational: the two paths add *different* g (their ff2 matmuls round differently), so the sums cannot be
+    # bit-identical; the row records how many elements differ and by how many bf16 ulps (residual rounding + g)
+    u_s1 = ulp_metrics(out_mk, out_s1_t)
+    n_diff_s1 = int((out_mk != out_s1_t).sum())
+    results.add_row(
+        {
+            "tap": "mk_k2_ff",
+            "key": f"ff-out-vs-s1-diff-{version}",
+            "rule": "info",
+            "ok": bool(u_s1["max_ulps"] <= 2.0),
+            "n_diff": n_diff_s1,
+            "frac_diff": n_diff_s1 / out_mk.numel(),
+            "max_ulps": u_s1["max_ulps"],
+            "g_n_diff": int((g != g_s1_t).sum()),
+            "note": "megakernel in-place residual of the megakernel g vs ttnn.add of the S1 g (g differs between paths)",
+        }
+    )
+    for key, (a, b) in rows.items():
+        m = pcc_mod.metrics(a.float(), b.float())
+        got_rows[key] = m["pcc"]
+        results.add_row(
+            {
+                "tap": "mk_k2_ff",
+                "key": f"{key}-{version}",
+                "rule": "pcc",
+                "pcc": m["pcc"],
+                "gate_pcc": GATE_PCC,
+                "max_abs": m["max_abs"],
+                "finite": m["finite"],
+                "ok": bool(m["finite"] and m["pcc"] >= GATE_PCC),
+            }
+        )
+    results.set_meta(matmul_wall_s=t_mm, ff_pcc=got_rows)
+    assert got_rows["ff-out-vs-s1"] >= GATE_PCC, f"FF branch vs S1: PCC {got_rows['ff-out-vs-s1']:.6f} < {GATE_PCC}"
+    assert got_rows["ff-ln-vs-s1"] >= GATE_PCC and got_rows["ff-out-vs-torch"] >= GATE_PCC

code/models/experimental/gr00t/tests/tt/test_tt_action_head.py

@@ -14,8 +14,13 @@ Two tiers (plan §7.1):
   (``sa_embs``, ``dit_block_last_out``, ``dit_out``, ``action_decoder_out``, ``pred_velocity``) gated by
   ``harness.check_taps``; ``weight_emulation='bfp8'`` emulates the ``mixed_dit`` bfp8 weights.  Select with ``-k cpu``.
 * device tests (``gr00t_device``; only through ``bin/with-device.sh``) -- the same 4-step denoise on the p150a with
-  ``TTWeights``, the K/V hoisted on device from the golden VL rows, buffer persistence and determinism, and the
-  megakernel backend raising until WP-K5.
+  ``TTWeights``, the K/V hoisted on device from the golden VL rows, buffer persistence and determinism.
+
+Every head built here is the **ttnn** backend (``TTPolicy(dit_backend="ttnn")``): this file tests the Stage-1 op
+sequence, and the session device has the firmware ``worker_l1_size`` the megakernel cannot run under.  The megakernel
+backend -- the ``TTPolicy`` default since 2026-09-18 -- is covered end to end by ``test_mk_e2e.py`` (both backends on
+one K5 device), ``test_mk_2cq.py`` (1-CQ vs 2-CQ bit-equality) and the whole-model files; the CPU contract here is
+only that its constructor refuses the stub's small configuration with ``ValueError`` instead of building an arena.
 
 Run (from the tt-metal root)::
 
@@ -201,7 +206,7 @@ def test_cpu_head_wiring_on_stub(version: str, golden, stub_ttnn: types.ModuleTy
         if l_pad is not None
     }
     kv = hoist_kv(cfg, dl, W, vl_sets_h, lambda t: FakeTensor(t.float(), stub.bfloat16))
-    policy = TTPolicy()
+    policy = TTPolicy(dit_backend="ttnn")  # the Stage-1 op path on the stub (module docstring)
     device = object()
 
     # own mask buffers: allocated from placeholders, refreshed in place by update_masks(slots)
@@ -290,9 +295,24 @@ def test_cpu_head_wiring_on_stub(version: str, golden, stub_ttnn: types.ModuleTy
     assert head.x_t.freed and head.action_pred.freed and head.masks == {}
     assert all(t.freed for t in own_masks), "release() must free the mask buffers the head allocated"
 
-    # contracts
-    with pytest.raises(NotImplementedError):
-        AH.ActionHeadTT(cfg, dl, W, TTPolicy(dit_backend="megakernel"), device, kv=kv).denoise(head2.x_t, state)
+    # contracts.  The megakernel backend builds its DRAM arena for the *real* version config at construction
+    # (``ArenaLayout.plan(cfg.version, ...)``: 16 / 32 blocks) and cross-checks it against the head's config, so the
+    # stub's 4-block ``small_cfg`` is refused with ValueError before any weight is touched (WP-K5 summary §7 item 6)
+    with pytest.raises(ValueError, match="megakernel arena"):
+        AH.ActionHeadTT(cfg, dl, W, TTPolicy(dit_backend="megakernel"), device, kv=kv)
+    with pytest.raises(ValueError, match="Stage-1"):
+        AH.ActionHeadTT(cfg, dl, W, TTPolicy(dit_backend="megakernel"), device, kv=kv, persist_sa_embs=True)
+    with pytest.raises(ValueError, match="megakernel-backend argument"):
+        AH.ActionHeadTT(cfg, dl, W, policy, device, kv=kv, mk_k4_budget=object())
+    mk_taps = [t for k in STEPS for t in (f"action_decoder_out[k={k}]", f"pred_velocity[k={k}]")]
+    assert AH.head_tap_names("megakernel", len(STEPS), cfg.dit.n_blocks) == (
+        *mk_taps,
+        "action_pred_normalized",
+        "action_pred_valid",
+    )
+    assert set(AH.head_tap_names("megakernel", len(STEPS), cfg.dit.n_blocks)) < set(
+        AH.head_tap_names("ttnn", len(STEPS), cfg.dit.n_blocks)
+    )
     h3 = AH.ActionHeadTT(cfg, dl, W, policy, device)
     with pytest.raises(ValueError):
         h3.denoise(h3.x_t, state)  # no K/V bound
@@ -444,7 +464,14 @@ def _device_setup(
     )
     t0 = time.perf_counter()
     dh = _DeviceHead(
-        gs, mi, cfg, dl, head_weights, TTPolicy(dtype_policy=head_weights.policy.dtype_policy), gr00t_device, **kw
+        gs,
+        mi,
+        cfg,
+        dl,
+        head_weights,
+        TTPolicy(dtype_policy=head_weights.policy.dtype_policy, dit_backend="ttnn"),
+        gr00t_device,
+        **kw,
     )
     import ttnn
 
@@ -556,40 +583,3 @@ def test_device_buffers_persistent_and_deterministic(
         assert det and changed
     finally:
         dh.release()
-
-
-def test_device_megakernel_backend_raises(
-    version: str, golden, gr00t_device: Any, head_weights: Any, results: ResultsWriter
-) -> None:
-    """``TTPolicy(dit_backend='megakernel')`` constructs but ``denoise`` raises ``NotImplementedError`` until WP-K5."""
-    gs, mi, cfg, dl = _golden_setup(version, golden)
-    from models.experimental.gr00t.tt import layout as LAY
-
-    masks = LAY.materialise(gr00t_device, dl, mi.slots)
-    kv = hoist_kv(
-        cfg,
-        dl,
-        head_weights,
-        gathered_vl_sets(gs, mi, dl),
-        lambda t: L.to_device(t, gr00t_device, dtype="bf16", mem="DRAM"),
-    )
-    head = AH.ActionHeadTT(cfg, dl, head_weights, TTPolicy(dit_backend="megakernel"), gr00t_device, masks=masks, kv=kv)
-    state = L.to_device(
-        gs.load("state_features").reshape(1, 1, 1, -1).to(torch.bfloat16), gr00t_device, dtype="bf16", mem="DRAM"
-    )
-    try:
-        assert head.backend == "megakernel" and head.megakernel is None
-        with pytest.raises(NotImplementedError):
-            head.denoise(head.x_t, state)
-        results.set_meta(backend="megakernel", raises="NotImplementedError")
-    finally:
-        import ttnn
-
-        head.release()
-        ttnn.deallocate(state)
-        for t in masks.values():
-            if t is not None:
-                ttnn.deallocate(t)
-        for k, v in kv.values():
-            ttnn.deallocate(k)
-            ttnn.deallocate(v)

code/models/experimental/gr00t/tests/tt/test_tt_backbone.py

@@ -72,10 +72,12 @@ REQUIRED_GATED = {
     "n16": ("llm_inputs_embeds", "projector_out", "backbone_features"),
     "n17": ("llm_inputs_embeds", "projector_out", "vlln_out"),
 }
-#: D0 device fact carried into D1 (test_tt_qwen3vl_vision.KNOWN_GATE_MISSES): the verbatim tt_dit N1.7 tower reaches
-#: PCC 0.99965 on ``vit_deepstack_mergers[0]`` against the frozen gate 0.9998 (bf16 block stack, HiFi2 linears); the
-#: miss is an **xfail** while the PCC stays above the documented floor, every other gated tap must pass, and the
-#: WP-12 acceptance keys (REQUIRED_GATED) are never in this table.  Plan owner: relax to 0.9995 or fork the tower (WP-P).
+#: D0 device fact carried into D1 (test_tt_qwen3vl_vision.KNOWN_GATE_MISSES): the **verbatim** tt_dit N1.7 tower
+#: reaches PCC 0.99965 on ``vit_deepstack_mergers[0]`` against the frozen gate 0.9998 (HiFi2 linears in the block
+#: stack).  Since 2026-09-17 the default tower is the fork with HiFi3 + fp32 acc on blocks 0-5 (``TTPolicy.vit_fork``,
+#: docs/plan/n17-fidelity.md) and passes the gate; the floor below only applies to ``GR00T_VIT_FORK=0`` (oracle) runs:
+#: the miss is an **xfail** while the PCC stays above the floor, every other gated tap must pass, and the WP-12
+#: acceptance keys (REQUIRED_GATED) are never in this table.
 KNOWN_GATE_MISSES: Dict[str, float] = {"vit_deepstack_mergers[0]": 0.9995}
 assert not any(k in KNOWN_GATE_MISSES for keys in REQUIRED_GATED.values() for k in keys)
 
@@ -727,6 +729,9 @@ def test_cpu_recipe_and_build_contract() -> None:
         r = B.BackboneRecipe.from_config(cfg, dl, pol)
         flags[v] = (r.tower, r.projector, r.llm_mask, r.rope_table, r.final_norm, r.deepstack, r.has_vlsa)
         assert r.describe()["per_call_inputs"] == list(r.per_call_inputs())
+        # the fork is the TTPolicy default and applies to the Qwen3-VL tower only (n17-fidelity.md)
+        assert r.vit_fork is (v == "n17") and r.describe()["vit_fork"] is r.vit_fork
+        assert B.BackboneRecipe.from_config(cfg, dl, TTPolicy(vit_fork=False)).vit_fork is False
         with pytest.raises(ValueError):
             B.BackboneRecipe.from_config(
                 cfg,

code/models/experimental/gr00t/tests/tt/test_tt_determinism.py

@@ -110,6 +110,42 @@ def test_cpu_all_equal_helper() -> None:
 # --------------------------------------------------------------------------------------------------------------------
 # device tier
 # --------------------------------------------------------------------------------------------------------------------
+# --------------------------------------------------------------------------------------------------------------------
+# DiT backend parametrisation (WP-K5): GR00T_DIT_BACKEND=ttnn|megakernel, GR00T_MK_ARENA_DTYPE=auto|bf16|bfp8_b
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def policy(policy_name: str, trace_layout: str) -> TTPolicy:
+    """conftest's ``TTPolicy(dtype_policy=--policy, trace_layout=--trace-layout)`` plus the environment overrides of
+    ``TTPolicy.with_env_overrides`` (the DiT backend and the megakernel arena dtype of this run)."""
+    return TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout).with_env_overrides()
+
+
+@pytest.fixture(scope="module")
+def gr00t_device(request: pytest.FixtureRequest, policy: TTPolicy) -> Iterator[Any]:
+    """The p150a opened per backend (``tt.model.open_model_device``: the megakernel needs the smaller K4
+    ``worker_l1_size``); same trace region / L1-small sizes as conftest's session fixture, one device per process."""
+    from models.experimental.gr00t.tests.tt.conftest import (
+        DEFAULT_L1_SMALL_SIZE,
+        DEFAULT_TRACE_REGION_SIZE,
+        DEFAULT_TT_METAL_CACHE,
+    )
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+    device = M.open_model_device(
+        policy, trace_region_size=DEFAULT_TRACE_REGION_SIZE, l1_small_size=DEFAULT_L1_SMALL_SIZE
+    )
+    facts = harness.device_facts(device)
+    facts["worker_l1_size"] = M.device_worker_l1_size(device)
+    facts["dit_backend"] = policy.dit_backend
+    facts["mk_arena_dtype"] = policy.mk_arena_dtype if policy.dit_backend == "megakernel" else None
+    request.config._gr00t_device_facts = facts  # type: ignore[attr-defined]
+    try:
+        yield device
+    finally:
+        close_gr00t_device(device)
+
+
 @pytest.fixture(scope="module")
 def model(request: pytest.FixtureRequest, version: str) -> Iterator[M.Gr00tTT]:
     request.getfixturevalue("golden")

code/models/experimental/gr00t/tests/tt/test_tt_dtype_policy.py

@@ -0,0 +1,904 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Device sweep of the weight dtype classes (WP-Q; IMPLEMENTATION_PLAN.md §1.2 "After WP-Q admission", §7.2).
+
+One pytest process per version, through the device lock, from the tt-metal root (the conftest option is
+``--gr00t-version``; ``TT_METAL_CACHE`` exported)::
+
+    DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
+        python -m pytest -q -rA --timeout=1800 --timeout-method=thread -p no:cacheprovider \\
+        models/experimental/gr00t/tests/tt/test_tt_dtype_policy.py --gr00t-version n15
+
+Device tier:
+
+* ``test_device_bfp_round_matches_ttnn[bfp8_b|bfp4_b]`` -- ``ttnn.from_torch(x, bfloat8_b | bfloat4_b, TILE)``
+  (host pack and a device round trip) against ``test_bfp8_emulation_cpu.bfp_round``: every element within one
+  mantissa step of its block, exactness fraction recorded (bf16-valued and fp32-valued inputs, ragged shapes).
+* ``test_device_dtype_sweep[<step>]`` -- the class-by-class admission sweep in :data:`SWEEP_STEPS` order
+  (``baseline`` = the shipped ``ADMITTED_CLASSES`` -> ``vision`` -> ``projector`` (N1.6 mlp1 / N1.7 mergers) ->
+  ``llm`` -> ``vl_self_attention`` -> the two experiments ``kv_bfp8`` (hoisted cross K/V rounded to bfp8 after the
+  hoist, plan §1.2 row 6) and ``dit_ff_bfp4`` (DiT ``ff1.W`` / ``ff2.W`` pre-rounded to bfp4, plan §1.2 row 1)).
+  Every step loads ``Gr00tTT`` with ``TTPolicy.with_admitted(<admitted so far + the class>, cache_tag="wpq")`` (its
+  own ``.tensorbin`` tier, removed at the end of the module), runs ``run_untraced(return_taps=True)`` on the canonical
+  golden and on every multi sample, gates everything with ``harness.check_taps`` (multi-sample gates, valid slices,
+  per-dim action PCC; the e2e module's ``KNOWN_GATE_MISSES`` floors apply exactly as in ``test_tt_model_e2e``),
+  measures the traced latency (``bench_e2e`` method: upload -> replays -> D2H, medians of 20 calls + a split pass)
+  and records the per-tap PCC deltas against the baseline step.  Admission rule of a weight class (never admit a
+  class whose worst tap is below its gate): canonical -- every gated row passes except the known misses, which must
+  stay above their floors; multi -- no row that passed at baseline fails, and ``action_pred_valid`` drops by at most
+  :data:`MULTI_ACTION_SLACK`; an admitted class is carried into the next step, a rejected one keeps bf16.  The sweep
+  **measures**; the shipped table ``tt/policy.py::ADMITTED_CLASSES`` is written by hand from
+  ``tests/tt/results/q_summary.md`` (a class must pass on all three versions, the table is version-independent).
+
+CPU tier (``-k cpu``): the sweep policies (``with_admitted`` / ``weights_tag`` / ``dtype_of`` histograms per step
+and version), the validation of the new ``TTPolicy`` fields, the step table.
+
+Workarounds in this file (owned files only; each one is marked ``WORKAROUND``): the K/V experiment rounds the
+adapter's persistent bf16 K/V buffers in place through ``ttnn.typecast`` (``tt/adapter.py`` refuses
+``kv_dtype="bfp8_b"``; a real bfp8 buffer needs the typecast after the head split there); the bfp4 experiment swaps
+the DiT blocks' ``w_ff1`` / ``w_ff2`` device tensors for bfp8 tensors holding bfp4-rounded values (bfp4 values are
+exactly representable in bfp8 with the same shared exponent, so the matmul numerics equal a bfp4 operand; the
+matmul table of ``tt/layers.py`` has no ``bfp4_b`` rows, so the byte / latency benefit is not measured).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import shutil
+import statistics
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Set, Tuple
+
+import pytest
+import torch
+
+from models.experimental.gr00t.common.configs import GR00T_VERSIONS
+from models.experimental.gr00t.common.golden import GoldenSet, list_multi_samples
+from models.experimental.gr00t.reference.model import load_golden_sample
+from models.experimental.gr00t.tests.tt import harness
+from models.experimental.gr00t.tests.tt import test_tt_model_e2e as E2E_T
+from models.experimental.gr00t.tests.tt.test_bfp8_emulation_cpu import bfp4_round, bfp_round, mantissa_step
+from models.experimental.gr00t.tt import model as M
+from models.experimental.gr00t.tt import weights as TTW
+from models.experimental.gr00t.tt.layout import StaticShapeError
+from models.experimental.gr00t.tt.policy import (
+    ADMITTED_CLASSES,
+    BFP8_CANDIDATE_CATEGORIES,
+    TTPolicy,
+    dtype_decisions,
+    dtype_histogram,
+)
+
+torch.set_grad_enabled(False)
+
+CACHE_TAG = "wpq"
+SWEEP_STEPS: Tuple[str, ...] = (
+    "baseline",
+    "vision",
+    "projector",
+    "llm",
+    "vl_self_attention",
+    "kv_bfp8",
+    "dit_ff_bfp4",
+)
+#: weight-class steps -> plan categories switched to bfp8_b on top of the admitted set
+WEIGHT_CLASS_STEPS: Dict[str, Tuple[str, ...]] = {
+    "vision": ("vision",),
+    "projector": ("projector", "vision_merger"),
+    "llm": ("llm",),
+    "vl_self_attention": ("vl_self_attention",),
+}
+EXPERIMENT_STEPS: Tuple[str, ...] = ("kv_bfp8", "dit_ff_bfp4")
+MULTI_ACTION_SLACK = 1e-3  # multi-sample admission: action_pred_valid PCC may drop by at most this vs the baseline
+LATENCY_WARMUP = 5
+LATENCY_REPS = 20
+MIN_FREE_DISK_GB = 8.0  # the sweep tier copies one version's tensors (<= 4.8 GB bf16); never fill the disk
+ACTION_KEY = "action_pred_valid"
+BFP_DTYPES: Dict[str, int] = {"bfp8_b": 7, "bfp4_b": 3}
+#: plan §1.2 "After WP-Q admission": the canonical taps each class's row names as its admission gate (the sweep's
+#: automatic decision is stricter -- every canonical tap + no new multi-sample failure; both verdicts are recorded).
+PLAN_GATE_TAPS: Dict[str, Dict[str, Tuple[str, ...]]] = {
+    "vision": {v: ("vit_block_last", "projector_out") for v in GR00T_VERSIONS},
+    "projector": {v: ("projector_out",) for v in GR00T_VERSIONS},
+    "llm": {"n15": ("backbone_features",), "n16": ("backbone_features",), "n17": ("vlln_out",)},
+    "vl_self_attention": {v: ("vl_self_attention_out",) for v in GR00T_VERSIONS},
+    "kv_bfp8": {
+        v: tuple(f"{t}[k={k}]" for t in ("dit_block_last_out", "pred_velocity") for k in range(4))
+        for v in GR00T_VERSIONS
+    },
+    "dit_ff_bfp4": {
+        v: tuple(f"{t}[k={k}]" for t in ("dit_block_last_out", "pred_velocity") for k in range(4))
+        for v in GR00T_VERSIONS
+    },
+}
+
+
+def plan_gate_verdict(step: str, version: str, pcc: Mapping[str, float], gates: Mapping[str, float]) -> Optional[bool]:
+    """``True`` when every plan §1.2 gate tap of ``step`` holds on the canonical golden (``None`` for the baseline)."""
+    taps = PLAN_GATE_TAPS.get(step, {}).get(version)
+    if taps is None:
+        return None
+    for t in taps:
+        if t not in pcc or t not in gates:
+            raise KeyError(f"{version}/{step}: plan gate tap {t!r} missing from the canonical report")
+    return all(pcc[t] >= gates[t] for t in taps)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# sweep policies (CPU)
+# --------------------------------------------------------------------------------------------------------------------
+def step_admitted(step: str, admitted_so_far: Mapping[str, str]) -> Dict[str, str]:
+    """The admitted-class table a step loads with: the classes admitted so far plus the step's class at bfp8_b
+    (experiments and the baseline load the table as is)."""
+    if step not in SWEEP_STEPS:
+        raise KeyError(f"unknown sweep step {step!r}; expected one of {SWEEP_STEPS}")
+    table = dict(admitted_so_far)
+    for cat in WEIGHT_CLASS_STEPS.get(step, ()):
+        table[cat] = "bfp8_b"
+    return table
+
+
+def category_dtype_counts(version: str, policy: TTPolicy) -> Dict[str, Dict[str, int]]:
+    """``{category: {device dtype: n tensors}}`` of the canonical embodiment's plan under ``policy`` (headers only)."""
+    entries = TTW.upload_plan(version, None, policy)
+    out: Dict[str, Dict[str, int]] = {}
+    for e in entries.values():
+        d = out.setdefault(e.category, {})
+        d[e.device_dtype] = d.get(e.device_dtype, 0) + 1
+    return {k: dict(sorted(v.items())) for k, v in sorted(out.items())}
+
+
+def category_bytes(version: str, policy: TTPolicy) -> Dict[str, float]:
+    entries = TTW.upload_plan(version, None, policy)
+    out: Dict[str, float] = {}
+    for e in entries.values():
+        out[e.category] = out.get(e.category, 0.0) + float(e.nbytes)
+    return out
+
+
+def test_cpu_policy_sweep_tables() -> None:
+    """``with_admitted`` / ``weights_tag`` / ``admitted`` / ``to_dict`` and the field validation of the WP-Q hooks."""
+    base = TTPolicy()
+    assert base.admitted() == ADMITTED_CLASSES and base.admitted() is not ADMITTED_CLASSES
+    assert base.weights_tag() == "mixed_dit" and base.to_dict()["admitted_classes"] is None
+    assert base.to_dict()["cache_tag"] is None
+    p = base.with_admitted({"dit": "bfp8_b", "llm": "bfp8_b"}, CACHE_TAG)
+    assert p.admitted_classes == (("dit", "bfp8_b"), ("llm", "bfp8_b"))
+    assert p.admitted() == {"dit": "bfp8_b", "llm": "bfp8_b"}
+    assert p.weights_tag() == f"mixed_dit+{CACHE_TAG}"
+    assert p.to_dict()["admitted_classes"] == {"dit": "bfp8_b", "llm": "bfp8_b"}
+    assert p.to_dict()["cache_tag"] == CACHE_TAG
+    assert p.trace_layout == base.trace_layout and p.dtype_policy == "mixed_dit"
+    # an explicit table is independent of the module table
+    assert base.with_admitted({}, None).admitted() == {}
+    for bad in (
+        dict(admitted_classes=[("dit", "bfp8_b")]),
+        dict(admitted_classes=(("encoders", "bfp8_b"),)),
+        dict(admitted_classes=(("dit", "fp32"),)),
+        dict(admitted_classes=(("dit", "bfp8_b"), ("dit", "bf16"))),
+        dict(admitted_classes=(("dit",),)),
+        dict(cache_tag=""),
+        dict(cache_tag="has space"),
+        dict(cache_tag="/etc"),
+    ):
+        with pytest.raises((TypeError, ValueError)):
+            TTPolicy(**bad)  # type: ignore[arg-type]
+    # the bf16 policy ignores the table; the mixed policy admits every candidate
+    for cat in BFP8_CANDIDATE_CATEGORIES:
+        p1 = TTPolicy(dtype_policy="bf16").with_admitted({cat: "bfp8_b"}, None)
+        assert set(dtype_histogram(dtype_decisions("n16", None, p1))) == {"bf16"}
+
+
+@pytest.mark.parametrize("version", list(GR00T_VERSIONS))
+def test_cpu_sweep_steps_grow_bfp8(version: str) -> None:
+    """Along the sweep the bfp8 tensor count grows monotonically and only inside the step's categories; the known
+    no-op steps (N1.5 projector, N1.6 VL-SA) change nothing; the sweep tier never shares the production tag."""
+    base = TTPolicy()
+    admitted = dict(base.admitted())
+    prev = category_dtype_counts(version, base.with_admitted(admitted, CACHE_TAG))
+    prev_n = sum(d.get("bfp8_b", 0) for d in prev.values())
+    for step in SWEEP_STEPS:
+        table = step_admitted(step, admitted)
+        pol = base.with_admitted(table, CACHE_TAG)
+        assert pol.weights_tag() == f"mixed_dit+{CACHE_TAG}" != base.weights_tag()
+        counts = category_dtype_counts(version, pol)
+        n = sum(d.get("bfp8_b", 0) for d in counts.values())
+        assert n >= prev_n, step
+        changed = {c for c in counts if counts[c] != prev.get(c)}
+        assert changed <= set(WEIGHT_CLASS_STEPS.get(step, ())), (step, changed)
+        if (version, step) in (("n15", "projector"), ("n16", "vl_self_attention")):
+            assert n == prev_n, f"{version}/{step} must be a no-op (plan §1.2)"
+        elif step in WEIGHT_CLASS_STEPS:
+            assert n > prev_n, f"{version}/{step} admitted no tensor"
+        for c in WEIGHT_CLASS_STEPS.get(step, ()):
+            gain = counts.get(c, {}).get("bfp8_b", 0) - prev.get(c, {}).get("bfp8_b", 0)
+            assert counts.get(c, {}).get("bf16", 0) == prev.get(c, {}).get("bf16", 0) - gain, (step, c)
+        admitted = table
+        prev, prev_n = counts, n
+    # every candidate category admitted == the host plan's own "mixed" dtype
+    full = base.with_admitted({c: "bfp8_b" for c in BFP8_CANDIDATE_CATEGORIES}, None)
+    assert dtype_decisions(version, None, full) == dtype_decisions(version, None, TTPolicy(dtype_policy="mixed"))
+
+
+def test_cpu_step_table() -> None:
+    assert SWEEP_STEPS[0] == "baseline" and set(WEIGHT_CLASS_STEPS) | set(EXPERIMENT_STEPS) == set(SWEEP_STEPS[1:])
+    for cats in WEIGHT_CLASS_STEPS.values():
+        assert set(cats) <= set(BFP8_CANDIDATE_CATEGORIES)
+    assert set(BFP8_CANDIDATE_CATEGORIES) - {"dit"} == {c for cats in WEIGHT_CLASS_STEPS.values() for c in cats}
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device: the rounding primitive vs ttnn
+# --------------------------------------------------------------------------------------------------------------------
+def _probe_tensors(seed: int = 0) -> Dict[str, torch.Tensor]:
+    g = torch.Generator().manual_seed(seed)
+    out: Dict[str, torch.Tensor] = {}
+    # per-row scale spread over 6 decades -> many distinct shared exponents, some blocks with huge exponent spread
+    x = torch.randn(64, 96, generator=g) * torch.pow(10.0, torch.rand(64, 1, generator=g) * 6 - 3)
+    x[3, :8] = 0.0
+    out["64x96_spread"] = x
+    out["200x72_ragged"] = torch.randn(200, 72, generator=g)
+    out["1x32x100x64"] = torch.randn(1, 32, 100, 64, generator=g) * 4.0
+    w = torch.randn(1152, 4608, generator=g) * 0.02
+    w[::7, ::5] *= 40.0  # outlier weights inside blocks
+    out["1152x4608_weightlike"] = w
+    return out
+
+
+@pytest.mark.parametrize("dtype_name", list(BFP_DTYPES))
+def test_device_bfp_round_matches_ttnn(gr00t_device: Any, results: harness.ResultsWriter, dtype_name: str) -> None:
+    """``ttnn.from_torch(x, <bfp>, TILE)`` (host pack, and the same through the device) vs :func:`bfp_round`."""
+    import ttnn
+
+    mbits = BFP_DTYPES[dtype_name]
+    tt_dt = {"bfp8_b": ttnn.bfloat8_b, "bfp4_b": ttnn.bfloat4_b}[dtype_name]
+    worst_steps = 0.0
+    for name, x in _probe_tensors().items():
+        for kind, xin in (("bf16", x.to(torch.bfloat16).to(torch.float32)), ("fp32", x)):
+            emu = bfp_round(xin, mbits)
+            step = mantissa_step(xin, mbits)
+            host = ttnn.to_torch(ttnn.from_torch(xin, dtype=tt_dt, layout=ttnn.TILE_LAYOUT)).to(torch.float32)
+            dev_t = ttnn.from_torch(xin, dtype=tt_dt, layout=ttnn.TILE_LAYOUT, device=gr00t_device)
+            dev = ttnn.to_torch(dev_t).to(torch.float32)
+            ttnn.deallocate(dev_t)
+            assert host.shape == emu.shape == dev.shape, (name, host.shape, emu.shape)
+            assert torch.equal(host, dev), f"{name}/{kind}: host pack != device round trip"
+            diff_steps = (host - emu).abs() / step
+            max_steps = float(diff_steps.max())
+            exact = float((host == emu).to(torch.float32).mean())
+            worst_steps = max(worst_steps, max_steps)
+            ok = max_steps <= 1.0 + 1e-6
+            results.add_row(
+                {
+                    "tap": "bfp_round",
+                    "key": f"{dtype_name}/{name}/{kind}",
+                    "rule": "abs",
+                    "ok": ok,
+                    "shape": list(xin.shape),
+                    "exact_fraction": exact,
+                    "max_mantissa_steps": max_steps,
+                    "n_mismatch": int((host != emu).sum()),
+                    "max_abs": float((host - emu).abs().max()),
+                    "rel_l2_vs_input": float((host - xin).norm() / xin.norm()),
+                }
+            )
+            print(f"[{dtype_name}] {name:22s} {kind}: exact {exact:.6f}, max |d| {max_steps:.3f} mantissa steps")
+            assert ok, f"{dtype_name}/{name}/{kind}: emulation off by {max_steps} mantissa steps"
+    results.set_meta(dtype=dtype_name, mantissa_bits=mbits, worst_mantissa_steps=worst_steps)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device: the class sweep
+# --------------------------------------------------------------------------------------------------------------------
+@dataclass
+class StepResult:
+    step: str
+    admitted_loaded: Dict[str, str]
+    pcc: Dict[str, float] = field(default_factory=dict)  # canonical: key -> pcc
+    gates: Dict[str, float] = field(default_factory=dict)
+    fails: List[str] = field(default_factory=list)  # canonical failing keys (before known-miss handling)
+    known: List[str] = field(default_factory=list)
+    unknown: List[str] = field(default_factory=list)
+    multi_pcc: Dict[str, Dict[str, float]] = field(default_factory=dict)  # sample -> key -> pcc
+    multi_fails: Dict[str, List[str]] = field(default_factory=dict)
+    multi_errors: Dict[str, str] = field(default_factory=dict)
+    latency_ms: Dict[str, Any] = field(default_factory=dict)
+    timing: Dict[str, float] = field(default_factory=dict)
+    device_bytes: float = 0.0
+    dtype_hist: Dict[str, int] = field(default_factory=dict)
+    canonical_ok: bool = False
+    multi_ok: bool = False
+    admitted: bool = False
+    note: str = ""
+
+
+@dataclass
+class SweepState:
+    version: str
+    admitted: Dict[str, str]
+    steps: Dict[str, StepResult] = field(default_factory=dict)
+    tensorbin_dirs: Set[Path] = field(default_factory=set)
+
+
+_STATE: Dict[str, SweepState] = {}
+
+
+def state_for(version: str, base_policy: TTPolicy) -> SweepState:
+    st = _STATE.get(version)
+    if st is None:
+        st = SweepState(version=version, admitted=dict(base_policy.admitted()))
+        _STATE[version] = st
+    return st
+
+
+def _stats_ms(samples: Sequence[float]) -> Dict[str, float]:
+    s = sorted(samples)
+    return {
+        "median": statistics.median(s),
+        "p90": s[min(len(s) - 1, int(round(0.9 * (len(s) - 1))))],
+        "min": s[0],
+        "n": len(s),
+    }
+
+
+def traced_latency(model: M.Gr00tTT, mi: Any, noise: torch.Tensor, warmup: int, reps: int) -> Dict[str, Any]:
+    """``benchmarks/bench_e2e.py`` method on an already captured model: per call ``upload (+sync) -> every trace
+    replay -> one sync -> D2H``; medians / p90 of the main pass, then a split pass with a sync after every trace."""
+    tracer = model.tracer
+    stages = list(tracer.stage_names)
+
+    def one(split: bool) -> Dict[str, float]:
+        t: Dict[str, float] = {}
+        t0 = time.perf_counter()
+        tracer.upload(mi, noise)
+        tracer.synchronize()
+        t1 = time.perf_counter()
+        for st in stages:
+            ts = time.perf_counter()
+            tracer.replay(st)
+            if split:
+                tracer.synchronize()
+                t[f"trace/{st}"] = (time.perf_counter() - ts) * 1e3
+        tracer.synchronize()
+        t2 = time.perf_counter()
+        tracer.read_action_pred()
+        t3 = time.perf_counter()
+        t.update(upload=(t1 - t0) * 1e3, traces=(t2 - t1) * 1e3, d2h=(t3 - t2) * 1e3, device=(t3 - t0) * 1e3)
+        return t
+
+    for _ in range(warmup):
+        one(False)
+    main = [one(False) for _ in range(reps)]
+    split_pass = [one(True) for _ in range(reps)]
+    out: Dict[str, Any] = {k: _stats_ms([c[k] for c in main]) for k in ("device", "traces", "upload", "d2h")}
+    for st in stages:
+        out[f"trace/{st}"] = _stats_ms([c[f"trace/{st}"] for c in split_pass])
+    out["device_split_pass"] = _stats_ms([c["device"] for c in split_pass])
+    return out
+
+
+def install_kv_rounding(model: M.Gr00tTT, mantissa_bits: int = 7) -> Any:
+    """WORKAROUND (tt/adapter.py not owned; it raises on ``kv_dtype="bfp8_b"``): wrap the executor's adapter stage so
+    that, right after the hoist, every persistent ``K_i`` / ``V_i`` buffer is rounded in place to bfp8 values
+    (``typecast -> bfloat8_b -> bfloat16 -> copy``).  The buffers stay bf16 for the DiT consumers, so the numerics
+    equal a bfp8 K/V store (plan §1.2 row 6) while the byte traffic does not.  Returns the restore function."""
+    import ttnn
+
+    ex = model.executor
+    idx = [i for i, st in enumerate(ex._stages) if st.name == M.STAGE_ADAPTER]
+    if len(idx) != 1:
+        raise RuntimeError(f"expected one adapter stage, found {len(idx)}")
+    orig = ex._stages[idx[0]]
+    tt_dt = {7: ttnn.bfloat8_b, 3: ttnn.bfloat4_b}[mantissa_bits]
+
+    def wrapped(inputs: Mapping[str, Any], state: Mapping[str, Any]) -> Dict[str, Any]:
+        out = orig.fn(inputs, state)
+        for buffers in (model.adapter.K, model.adapter.V):
+            for _, buf in sorted(buffers.items()):
+                low = ttnn.typecast(buf, tt_dt)
+                back = ttnn.typecast(low, ttnn.bfloat16)
+                ttnn.copy(back, buf)
+                ttnn.deallocate(low)
+                ttnn.deallocate(back)
+        return out
+
+    ex._stages[idx[0]] = dataclasses.replace(orig, fn=wrapped)
+
+    def restore() -> None:
+        ex._stages[idx[0]] = orig
+
+    return restore
+
+
+def swap_dit_ff_bfp4(model: M.Gr00tTT, device: Any) -> Tuple[Any, Dict[str, Any]]:
+    """WORKAROUND (tt/weights.py / tt/layers.py not owned: no ``bfp4_b`` plan dtype or matmul rows): replace every DiT
+    block's ``w_ff1`` / ``w_ff2`` by a bfp8_b device tensor whose values were rounded to bfp4_b by ttnn (bfp4 values
+    are exactly representable in bfp8 under the same shared exponent).  Returns ``(restore, facts)``; ``facts``
+    records the agreement of ttnn's bfp4 pack with ``bfp4_round`` on the real weights."""
+    import ttnn
+
+    originals: List[Tuple[Any, str, Any]] = []
+    swapped: List[Any] = []
+    n_exact = 0
+    n_total = 0
+    max_steps = 0.0
+    rel_l2: List[float] = []
+    for blk in model.head.step.blocks:
+        for attr in ("w_ff1", "w_ff2"):
+            t = getattr(blk, attr)
+            host = ttnn.to_torch(t).to(torch.float32)  # the bfp8 weight values
+            h4 = ttnn.to_torch(ttnn.from_torch(host, dtype=ttnn.bfloat4_b, layout=ttnn.TILE_LAYOUT)).to(torch.float32)
+            emu = bfp4_round(host)
+            n_exact += int((h4 == emu).sum())
+            n_total += h4.numel()
+            max_steps = max(max_steps, float(((h4 - emu).abs() / mantissa_step(host, 3)).max()))
+            rel_l2.append(float((h4 - host).norm() / host.norm()))
+            dev = ttnn.from_torch(
+                h4, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+            )
+            originals.append((blk, attr, t))
+            swapped.append(dev)
+            setattr(blk, attr, dev)
+
+    def restore() -> None:
+        for blk, attr, t in originals:
+            setattr(blk, attr, t)
+        for dev in swapped:
+            ttnn.deallocate(dev)
+
+    facts = {
+        "n_tensors": len(swapped),
+        "bfp4_exact_fraction_vs_emulation": n_exact / max(1, n_total),
+        "bfp4_max_mantissa_steps_vs_emulation": max_steps,
+        "bfp4_rel_l2_vs_bfp8_weights_mean": sum(rel_l2) / len(rel_l2),
+        "bfp4_rel_l2_vs_bfp8_weights_max": max(rel_l2),
+    }
+    return restore, facts
+
+
+def gated_candidates(
+    model: M.Gr00tTT, gs: GoldenSet, obs: Any, taps: Mapping[str, torch.Tensor], pred: torch.Tensor
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+    """Exactly the e2e test's candidate set: golden-keyed device taps + decoded groups, N1.7 raw LLM taps split off
+    as informational rows (``test_tt_model_e2e.INFO_TAPS``)."""
+    cand = E2E_T.golden_keyed(dict(taps), gs)
+    cand.update(E2E_T.decoded_groups(model, pred, obs))
+    cand, info = E2E_T.split_info(model.version, cand)
+    return cand, info
+
+
+def check_free_disk(path: Path) -> float:
+    free_gb = shutil.disk_usage(str(path)).free / 1e9
+    if free_gb < MIN_FREE_DISK_GB:
+        pytest.fail(f"only {free_gb:.1f} GB free under {path}; the sweep tier needs {MIN_FREE_DISK_GB} GB")
+    return free_gb
+
+
+@pytest.fixture(scope="module")
+def sweep_cleanup(version: str) -> Iterator[None]:
+    """Remove the ``+wpq`` tensorbin tiers this module created (one version's tensors, up to 4.8 GB)."""
+    yield
+    st = _STATE.get(version)
+    if st is None:
+        return
+    for d in sorted(st.tensorbin_dirs):
+        if d.is_dir() and d.name.endswith(f"+{CACHE_TAG}"):
+            size = sum(p.stat().st_size for p in d.rglob("*") if p.is_file())
+            shutil.rmtree(d, ignore_errors=True)
+            print(f"\n[test_tt_dtype_policy] removed sweep tier {d} ({size / 1e9:.2f} GB)")
+
+
+def _known_miss_split(version: str, gs: GoldenSet, report: harness.Report, cand: Mapping[str, torch.Tensor]):
+    known, unknown = E2E_T.known_miss_rows(version, gs, report, cand)
+    return known, unknown
+
+
+@pytest.mark.parametrize("step", list(SWEEP_STEPS))
+def test_device_dtype_sweep(
+    version: str,
+    step: str,
+    golden: Tuple[Any, Any, Any],
+    gr00t_device: Any,
+    policy: TTPolicy,
+    results: harness.ResultsWriter,
+    sweep_cleanup: None,
+) -> None:
+    """One sweep step (module docstring): load with the step's table, gate canonical + multi, time the traces, decide."""
+    import ttnn
+
+    if policy.dtype_policy != "mixed_dit":
+        pytest.fail(f"the sweep needs --policy mixed_dit (got {policy.dtype_policy!r})")
+    st = state_for(version, policy)
+    if step != "baseline" and "baseline" not in st.steps:
+        pytest.fail(f"{version}: the baseline step did not run; the sweep steps depend on it")
+    table = step_admitted(step, st.admitted)
+    pol = policy.with_admitted(table, CACHE_TAG)
+    counts = category_dtype_counts(version, pol)
+    new_cats = WEIGHT_CLASS_STEPS.get(step, ())
+    n_new = sum(counts.get(c, {}).get("bfp8_b", 0) for c in new_cats)
+    res = StepResult(step=step, admitted_loaded=dict(table))
+    results.set_meta(
+        step=step,
+        admitted_loaded=dict(table),
+        admitted_before=dict(st.admitted),
+        policy=pol.to_dict(),
+        category_dtype_counts=counts,
+        category_mib={k: v / 2**20 for k, v in category_bytes(version, pol).items()},
+    )
+    if step in WEIGHT_CLASS_STEPS and n_new == 0:
+        res.note = f"no bfp8 candidates in {new_cats} for {version} (plan §1.2)"
+        st.steps[step] = res
+        results.set_meta(note=res.note)
+        pytest.skip(f"{version}/{step}: {res.note}")
+
+    gs, obs, mi_auto = golden
+    free_gb = check_free_disk(Path(TTW.W.default_cache_root()))
+    results.set_meta(free_disk_gb_before=free_gb)
+
+    # ---- load
+    t0 = time.perf_counter()
+    model = M.Gr00tTT.from_pretrained(version, policy=pol, device=gr00t_device)
+    res.timing["from_pretrained_s"] = time.perf_counter() - t0
+    st.tensorbin_dirs.add(Path(model.W.tensorbin_dir))
+    res.timing.update({k: v for k, v in model.timing.items() if k.endswith("_s")})
+    res.device_bytes = float(model.W.device_bytes)
+    res.dtype_hist = model.W.dtype_histogram()
+    results.set_meta(
+        weights=dict(
+            n_tensors=len(model.W),
+            device_bytes=res.device_bytes,
+            dtype_histogram=res.dtype_hist,
+            load_stats=dict(getattr(model.W, "load_stats", {})),
+            tensorbin_dir=str(model.W.tensorbin_dir),
+        )
+    )
+    print(
+        f"\n[{version}/{step}] loaded {len(model.W)} tensors {res.device_bytes / 1e6:.0f} MB {res.dtype_hist} in "
+        f"{res.timing['from_pretrained_s']:.1f} s (weights {model.timing['weights_load_s']:.1f} s) table {table}"
+    )
+    restore = None
+    try:
+        if step == "kv_bfp8":
+            restore = install_kv_rounding(model, 7)
+        elif step == "dit_ff_bfp4":
+            t1 = time.perf_counter()
+            restore, facts = swap_dit_ff_bfp4(model, gr00t_device)
+            res.timing["ff_bfp4_swap_s"] = time.perf_counter() - t1
+            results.set_meta(ff_bfp4=facts)
+            print(f"[{version}/{step}] FF weights -> bfp4 values: {facts}")
+
+        # ---- canonical golden, untraced taps
+        mi = model.encode(obs)
+        assert mi.shape_key == mi_auto.shape_key, (mi.shape_key, mi_auto.shape_key)
+        noise = gs.load("initial_noise")
+        dl = model.dlayout
+        t2 = time.perf_counter()
+        pred, taps = model.run_untraced(mi, noise, return_taps=True)
+        res.timing["untraced_first_s"] = time.perf_counter() - t2
+        res.timing["untraced_device_s"] = float(model.timing["last_untraced_s"])
+        assert bool(torch.isfinite(pred).all()), "non-finite action_pred"
+        cand, info = gated_candidates(model, gs, obs, taps, pred)
+        required = E2E_T.required_gated_keys(version, gs, dl) | {k for k in cand if k.startswith(M.DECODE_TAP)}
+        assert set(cand) == required, (sorted(required - set(cand)), sorted(set(cand) - required))
+        report = harness.check_taps(version, cand, gs, results=results)
+        for r_ in E2E_T.info_rows(gs, info):
+            results.add_row(r_)
+            res.pcc[r_["key"] + "[info]"] = float(r_["pcc"])
+        res.pcc.update({r["key"]: float(r["pcc"]) for r in report.rows if r.get("pcc") is not None})
+        res.gates.update({r["key"]: float(r["gate_pcc"]) for r in report.rows if r.get("gate_pcc") is not None})
+        res.fails = [r["key"] for r in report.failures()]
+        known, unknown = _known_miss_split(version, gs, report, cand)
+        for r_ in known:
+            results.add_row(r_)
+        res.known = [r_["key"] for r_ in known]
+        res.unknown = [r_["key"] for r_ in unknown]
+        res.canonical_ok = not unknown
+        print(report.summary(failures_only=True) if report.failures() else f"[{version}/{step}] canonical: all pass")
+
+        # ---- multi samples
+        for sample in list_multi_samples(version):
+            try:
+                gs_s, obs_s = load_golden_sample(version, sample)
+                mi_s = model.encode(obs_s)
+                pred_s, taps_s = model.run_untraced(mi_s, gs_s.load("initial_noise"), return_taps=True)
+                cand_s, _ = gated_candidates(model, gs_s, obs_s, taps_s, pred_s)
+                rep_s = harness.check_taps(version, cand_s, gs_s, results=results)
+            except (StaticShapeError, FileNotFoundError, KeyError) as exc:  # recorded, never silent
+                res.multi_errors[sample] = f"{type(exc).__name__}: {exc}"
+                results.add_row(
+                    E2E_T.row("multi", f"{sample}/error", False, rule="info", error=res.multi_errors[sample])
+                )
+                print(f"[{version}/{step}] multi {sample}: {res.multi_errors[sample]}")
+                continue
+            res.multi_pcc[sample] = {r["key"]: float(r["pcc"]) for r in rep_s.rows if r.get("pcc") is not None}
+            res.multi_fails[sample] = [r["key"] for r in rep_s.failures()]
+            print(
+                f"[{version}/{step}] multi {sample}: rows {len(rep_s.rows)} fail {len(rep_s.failures())} "
+                f"{ACTION_KEY} {res.multi_pcc[sample].get(ACTION_KEY, float('nan')):.6f}"
+            )
+
+        # ---- traced latency (bench_e2e method) -- the K/V hook is removed first (its typecasts are not production)
+        if step == "kv_bfp8" and restore is not None:
+            restore()
+            restore = None
+        t3 = time.perf_counter()
+        model.warm_and_capture(mi, noise)
+        res.timing["warm_and_capture_s"] = time.perf_counter() - t3
+        res.latency_ms = traced_latency(model, mi, noise, LATENCY_WARMUP, LATENCY_REPS)
+        print(
+            f"[{version}/{step}] traced: traces (device only) ms median {res.latency_ms['traces']['median']:.1f}, "
+            f"upload {res.latency_ms['upload']['median']:.1f}, device (upload..D2H) {res.latency_ms['device']['median']:.1f} "
+            f"(p90 {res.latency_ms['device']['p90']:.1f}); split "
+            + " ".join(
+                f"{k.split('/')[1]}={v['median']:.1f}" for k, v in res.latency_ms.items() if k.startswith("trace/")
+            )
+        )
+    finally:
+        if restore is not None:
+            restore()
+        model.release()
+        ttnn.synchronize_device(gr00t_device)
+
+    # ---- decision vs the baseline
+    base = st.steps.get("baseline")
+    deltas: Dict[str, float] = {}
+    multi_new_fails: Dict[str, List[str]] = {}
+    multi_action_drop: Dict[str, float] = {}
+    if base is not None and step != "baseline":
+        deltas = {k: res.pcc[k] - base.pcc[k] for k in res.pcc if k in base.pcc}
+        for sample, fails in res.multi_fails.items():
+            base_fails = set(base.multi_fails.get(sample, ()))
+            multi_new_fails[sample] = sorted(set(fails) - base_fails)
+            a, b = res.multi_pcc[sample].get(ACTION_KEY), base.multi_pcc.get(sample, {}).get(ACTION_KEY)
+            if a is not None and b is not None:
+                multi_action_drop[sample] = b - a
+        res.multi_ok = (
+            not res.multi_errors
+            and all(not v for v in multi_new_fails.values())
+            and all(d <= MULTI_ACTION_SLACK for d in multi_action_drop.values())
+            and set(res.multi_fails) == set(base.multi_fails)
+        )
+    else:
+        res.multi_ok = not res.multi_errors
+    res.admitted = res.canonical_ok and res.multi_ok
+    # a plan tap missing from the canonical report is a reporting defect: recorded first (the device work of the step
+    # is done), re-raised after the JSON meta is complete (never silent, never a lost step).
+    plan_gate_error: Optional[str] = None
+    try:
+        plan_gate_ok = plan_gate_verdict(step, version, res.pcc, res.gates)
+    except KeyError as exc:
+        plan_gate_ok, plan_gate_error = None, str(exc)
+    lat_delta = None
+    traces_delta = None
+    if base is not None and step != "baseline" and base.latency_ms and res.latency_ms:
+        lat_delta = res.latency_ms["device"]["median"] - base.latency_ms["device"]["median"]
+        traces_delta = res.latency_ms["traces"]["median"] - base.latency_ms["traces"]["median"]
+    st.steps[step] = res
+    if step in WEIGHT_CLASS_STEPS and res.admitted:
+        st.admitted = dict(table)
+    results.set_meta(
+        canonical_ok=res.canonical_ok,
+        canonical_fails=res.fails,
+        known_misses=res.known,
+        unknown_misses=res.unknown,
+        multi_ok=res.multi_ok,
+        multi_fails=res.multi_fails,
+        multi_new_fails=multi_new_fails,
+        multi_action_pred_valid_drop=multi_action_drop,
+        multi_errors=res.multi_errors,
+        admitted=res.admitted,
+        plan_gate_ok=plan_gate_ok,
+        plan_gate_error=plan_gate_error,
+        plan_gate_taps=list(PLAN_GATE_TAPS.get(step, {}).get(version, ())),
+        latency_delta_note=(
+            "n/a by construction: experiment (kv_bfp8 hook removed before capture; dit_ff_bfp4 runs bfp8 tensors)"
+            if step in EXPERIMENT_STEPS
+            else "traces = device-only replay time (bench_e2e method); 'device' includes the host upload/D2H"
+        ),
+        admitted_after=dict(st.admitted),
+        pcc_canonical=res.pcc,
+        gates_canonical=res.gates,
+        pcc_delta_vs_baseline=deltas,
+        latency_ms=res.latency_ms,
+        latency_device_median_delta_ms=lat_delta,
+        latency_traces_median_delta_ms=traces_delta,
+        free_disk_gb_after=shutil.disk_usage(str(TTW.W.default_cache_root())).free / 1e9,
+    )
+    results.set_timing(**res.timing)
+    results.add_row(
+        E2E_T.row(
+            "sweep",
+            f"step/{step}",
+            True,
+            rule="info",
+            admitted=res.admitted,
+            canonical_ok=res.canonical_ok,
+            multi_ok=res.multi_ok,
+            device_ms=res.latency_ms.get("device", {}).get("median"),
+            traces_ms=res.latency_ms.get("traces", {}).get("median"),
+            latency_delta_ms=lat_delta,
+            traces_delta_ms=traces_delta,
+        )
+    )
+    verdict = ("ADMITTED" if res.admitted else "REJECTED (bf16 kept)") + f" [plan §1.2 gate taps hold: {plan_gate_ok}]"
+    if step in EXPERIMENT_STEPS:
+        verdict = "gates hold" if res.admitted else "gates fail"
+    elif step == "baseline":
+        verdict = "baseline ok" if res.canonical_ok else "baseline FAILS"
+    print(
+        f"[{version}/{step}] {verdict}: canonical_ok={res.canonical_ok} (unknown misses {res.unknown}, known "
+        f"{res.known}) multi_ok={res.multi_ok} new multi fails {multi_new_fails} action drop {multi_action_drop}; "
+        f"device ms {res.latency_ms.get('device', {}).get('median')} (delta {lat_delta}); admitted now {st.admitted}"
+    )
+    if step == "baseline" and not res.canonical_ok:
+        raise AssertionError(
+            f"{version}: the baseline (shipped policy) misses gates beyond KNOWN_GATE_MISSES: {res.unknown}"
+        )
+    if step == "baseline" and res.known:
+        results.set_meta(baseline_known_misses=res.known)
+    if plan_gate_error is not None:
+        raise AssertionError(
+            f"{version}/{step}: plan §1.2 gate tap missing from the canonical report: {plan_gate_error}"
+        )
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# summary helper (CPU): markdown over the sweep JSONs of a results directory
+# --------------------------------------------------------------------------------------------------------------------
+KEY_TAPS: Tuple[str, ...] = (
+    "vit_block_last",
+    "vit_post_ln",
+    "vit_deepstack_mergers[0]",
+    "projector_out",
+    "llm_layer_last",
+    "backbone_features",
+    "vlln_out",
+    "vl_self_attention_out",
+    "dit_encoder_hidden_states",
+    "dit_block_last_out[k=3]",
+    "dit_out[k=3]",
+    "pred_velocity[k=3]",
+    "action_pred_normalized",
+    "action_pred_valid",
+)
+
+
+def sweep_docs(results_dir: Path) -> Dict[str, Dict[str, Dict[str, Any]]]:
+    """``{version: {step: latest results JSON}}`` of ``test_device_dtype_sweep`` under ``results_dir``."""
+    out: Dict[str, Dict[str, Dict[str, Any]]] = {}
+    for f in sorted(Path(results_dir).glob("test_device_dtype_sweep_*.json")):
+        d = harness.ResultsWriter.load(f)
+        step = d.get("meta", {}).get("step")
+        if step is None:
+            continue
+        cur = out.setdefault(d["version"], {}).get(step)
+        if cur is None or d["created"] > cur["created"]:
+            out[d["version"]][step] = d
+    return out
+
+
+def summarise_sweep(results_dir: Optional[Path] = None) -> str:
+    docs = sweep_docs(Path(results_dir) if results_dir is not None else harness.RESULTS_DIR)
+    lines: List[str] = []
+    for version in GR00T_VERSIONS:
+        steps = docs.get(version)
+        if not steps:
+            continue
+        lines.append(f"### {version}")
+        lines.append("")
+        lines.append(
+            "| step | outcome | strict verdict (all canonical taps + no new multi fails) | plan §1.2 gate taps hold | "
+            "canonical unknown misses | known misses | multi new fails | traces ms, device only (delta vs baseline) | "
+            "upload ms (host) |"
+        )
+        lines.append("|---|---|---|---|---|---|---|---|---|")
+        base_lat = (steps.get("baseline", {}).get("meta", {}).get("latency_ms") or {}).get("traces", {}).get("median")
+        for step in SWEEP_STEPS:
+            d = steps.get(step)
+            if d is None:
+                continue
+            m = d["meta"]
+            lat = (m.get("latency_ms") or {}).get("traces", {}).get("median")
+            upload = (m.get("latency_ms") or {}).get("upload", {}).get("median")
+            delta = (lat - base_lat) if (step != "baseline" and lat is not None and base_lat is not None) else None
+            if step in EXPERIMENT_STEPS:
+                delta = None  # n/a by construction: the experiments run bf16 K/V / bfp8 FF tensors in the traces
+            canon = {r["key"]: r for r in d["rows"] if r.get("sample") == "canonical" and r.get("gate_pcc") is not None}
+            try:
+                pg = plan_gate_verdict(
+                    step, version, {k: r["pcc"] for k, r in canon.items()}, {k: r["gate_pcc"] for k, r in canon.items()}
+                )
+            except KeyError:
+                pg = m.get("plan_gate_ok")
+            lines.append(
+                "| {s} | {o} | {a} | {pg} | {u} | {k} | {nf} | {lat} | {up} |".format(
+                    s=step,
+                    o=d.get("outcome"),
+                    a=(
+                        "-"
+                        if step == "baseline" or d.get("outcome") == "skipped"
+                        else ("admitted" if m.get("admitted") else "rejected")
+                    ),
+                    pg="-" if pg is None else pg,
+                    u=", ".join(m.get("unknown_misses") or []) or "-",
+                    k=", ".join(m.get("known_misses") or []) or "-",
+                    nf="; ".join(f"{k}: {v}" for k, v in (m.get("multi_new_fails") or {}).items() if v) or "-",
+                    lat=(f"{lat:.1f}" if isinstance(lat, (int, float)) else "-")
+                    + (
+                        f" ({delta:+.1f})"
+                        if isinstance(delta, (int, float))
+                        else (" (n/a: experiment)" if step in EXPERIMENT_STEPS and lat is not None else "")
+                    ),
+                    up=f"{upload:.1f}" if isinstance(upload, (int, float)) else "-",
+                )
+            )
+        present = [s for s in SWEEP_STEPS if s in steps]
+        lines.append("")
+        lines.append("Canonical golden, PCC vs the fp32 golden (gate in parentheses; `-` = step skipped / no row):")
+        lines.append("")
+        lines.append("| tap (gate) | " + " | ".join(present) + " |")
+        lines.append("|---|" + "---|" * len(present))
+        for key in KEY_TAPS:
+            cells = []
+            gate = None
+            for s in present:
+                m = steps[s]["meta"]
+                p = (m.get("pcc_canonical") or {}).get(key)
+                gate = (m.get("gates_canonical") or {}).get(key, gate)
+                cells.append(f"{p:.6f}" if isinstance(p, (int, float)) else "-")
+            if any(c != "-" for c in cells):
+                g = f" ({gate:g})" if isinstance(gate, (int, float)) else ""
+                lines.append(f"| {key}{g} | " + " | ".join(cells) + " |")
+        # multi samples: failing rows / action_pred_valid PCC per sample and step (rows gated by gates_multi.json)
+        samples = sorted({smp for s in present for smp in (steps[s]["meta"].get("multi_fails") or {})})
+        if samples:
+            lines.append("")
+            lines.append(
+                "Multi-sample goldens: failing gated rows / `action_pred_valid` PCC per sample "
+                "(`+n` = rows failing that passed at the baseline step):"
+            )
+            lines.append("")
+            lines.append("| sample | " + " | ".join(present) + " |")
+            lines.append("|---|" + "---|" * len(present))
+            for smp in samples:
+                cells = []
+                for s in present:
+                    m = steps[s]["meta"]
+                    fails = (m.get("multi_fails") or {}).get(smp)
+                    err = (m.get("multi_errors") or {}).get(smp)
+                    if err:
+                        cells.append(f"error: {err.split(':')[0]}")
+                        continue
+                    if fails is None:
+                        cells.append("-")
+                        continue
+                    new = (m.get("multi_new_fails") or {}).get(smp) or []
+                    apv = None
+                    for r in steps[s]["rows"]:
+                        if r.get("sample") == smp and r.get("key") == ACTION_KEY and r.get("pcc") is not None:
+                            apv = float(r["pcc"])
+                    cell = f"{len(fails)}"
+                    if new:
+                        cell += f" (+{len(new)})"
+                    cell += f" / {apv:.6f}" if apv is not None else " / -"
+                    cells.append(cell)
+                lines.append(f"| {smp} | " + " | ".join(cells) + " |")
+        lines.append("")
+    return "\n".join(lines)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    import argparse
+
+    ap = argparse.ArgumentParser(description="Summarise the WP-Q sweep JSONs as markdown.")
+    ap.add_argument("results_dir", nargs="?", default=str(harness.RESULTS_DIR))
+    args = ap.parse_args(argv)
+    print(summarise_sweep(Path(args.results_dir)))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/models/experimental/gr00t/tests/tt/test_tt_model_e2e.py

@@ -7,7 +7,12 @@ untraced), §7.5 step 1 (the untraced tap chain is the bisection walk).
 
 Device tier (WP-D2; one pytest process per version, through the device lock, from the tt-metal root; the conftest
 option is ``--gr00t-version``; ``TT_METAL_TRACE_ALLOC_TRACKING=1`` must be set **in the shell** -- ttnn reads it once
-at import and the tt-metal root ``conftest.py`` imports ttnn during collection)::
+at import and the tt-metal root ``conftest.py`` imports ttnn during collection).  The DiT backend is parametrised
+through the environment (WP-K5): the module-level ``policy`` / ``gr00t_device`` fixtures build ``Gr00tTT`` with
+``TTPolicy(...).with_env_overrides()`` on a device opened through ``tt.model.open_model_device`` -- the **megakernel**
+(the ``TTPolicy`` default since 2026-09-18, reduced ``worker_l1_size``, bfp8_b arena) unless ``GR00T_DIT_BACKEND=ttnn``
+selects the Stage-1 path (+ ``GR00T_MK_ARENA_DTYPE=bf16|bfp8_b`` for the arena); the expected tap set follows the
+backend (``tt.model.expected_tap_keys``)::
 
     TT_METAL_TRACE_ALLOC_TRACKING=1 DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
         python -m pytest -q -rA --timeout=1800 -p no:cacheprovider \\
@@ -133,10 +138,10 @@ def info_rows(gs: GoldenSet, info: Mapping[str, torch.Tensor]) -> List[Dict[str,
     return out
 
 
-def required_gated_keys(version: str, gs: GoldenSet, dl: DeviceLayout) -> Set[str]:
+def required_gated_keys(version: str, gs: GoldenSet, dl: DeviceLayout, dit_backend: str = "ttnn") -> Set[str]:
     """Every expected untraced tap key that has a golden in this sample, minus the informational ones."""
     cfg = get_config(version)
-    expected = set(M.expected_tap_keys(cfg, dl))
+    expected = set(M.expected_tap_keys(cfg, dl, dit_backend))
     return {k for k in expected if k in available_keys(gs)} - set(INFO_TAPS[version])
 
 
@@ -468,7 +473,9 @@ def build_emulated_model(version: str, gs: GoldenSet, mi: Any, device: Any, trac
 
     cfg = get_config(version)
     dl = DeviceLayout.for_inputs(cfg, mi)
-    pol = TTPolicy(dtype_policy="bf16", trace_layout=trace_layout)
+    # the emulation runs the Stage-1 ttnn op sequence on the stub (the megakernel -- the TTPolicy default since
+    # 2026-09-18 -- is a real generic_op the stub cannot execute)
+    pol = TTPolicy(dtype_policy="bf16", trace_layout=trace_layout, dit_backend="ttnn")
     key = f"{version}/{pol.rope_layout}/{pol.gate_up_layout}"
     if key not in _FULL_FAKE_WEIGHTS:
         t0 = time.perf_counter()
@@ -647,6 +654,43 @@ def test_cpu_emulated_e2e_vs_golden(
     report.assert_ok()
 
 
+# --------------------------------------------------------------------------------------------------------------------
+# DiT backend parametrisation (WP-K5): GR00T_DIT_BACKEND=ttnn|megakernel, GR00T_MK_ARENA_DTYPE=auto|bf16|bfp8_b
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def policy(policy_name: str, trace_layout: str) -> TTPolicy:
+    """conftest's ``TTPolicy(dtype_policy=--policy, trace_layout=--trace-layout)`` plus the environment overrides of
+    ``TTPolicy.with_env_overrides`` (the DiT backend and the megakernel arena dtype of this run)."""
+    return TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout).with_env_overrides()
+
+
+@pytest.fixture(scope="module")
+def gr00t_device(request: pytest.FixtureRequest, policy: TTPolicy) -> Iterator[Any]:
+    """The p150a opened per backend (``tt.model.open_model_device``: the megakernel needs the smaller K4
+    ``worker_l1_size``); same trace region / L1-small sizes as conftest's session fixture, one device per process."""
+    from models.experimental.gr00t.tests.tt.conftest import (
+        DEFAULT_L1_SMALL_SIZE,
+        DEFAULT_TRACE_REGION_SIZE,
+        DEFAULT_TT_METAL_CACHE,
+    )
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+    device = M.open_model_device(
+        policy, trace_region_size=DEFAULT_TRACE_REGION_SIZE, l1_small_size=DEFAULT_L1_SMALL_SIZE
+    )
+    facts = harness.device_facts(device)
+    facts["worker_l1_size"] = M.device_worker_l1_size(device)
+    facts["dit_backend"] = policy.dit_backend
+    # the policy field ("auto" = the per-version table); ``meta.recipe.mk_arena_dtype`` carries the resolved value
+    facts["mk_arena_dtype"] = policy.mk_arena_dtype if policy.dit_backend == "megakernel" else None
+    request.config._gr00t_device_facts = facts  # type: ignore[attr-defined]
+    try:
+        yield device
+    finally:
+        close_gr00t_device(device)
+
+
 # --------------------------------------------------------------------------------------------------------------------
 # device fixtures
 # --------------------------------------------------------------------------------------------------------------------
@@ -662,7 +706,8 @@ def model(request: pytest.FixtureRequest, version: str) -> Iterator[M.Gr00tTT]:
     print(
         f"\n[test_tt_model_e2e] {version}: from_pretrained {m.timing['from_pretrained_s']:.1f} s "
         f"(weights {m.timing['weights_load_s']:.1f} s, {len(m.W)} tensors, {m.W.device_bytes / 1e6:.0f} MB, "
-        f"{m.W.load_stats['path']}); trace layout {policy.trace_layout}"
+        f"{m.W.load_stats['path']}); trace layout {policy.trace_layout}; dit_backend {policy.dit_backend} "
+        f"(arena {m.recipe.mk_arena_dtype})"
     )
     yield m
     m.release()
@@ -700,7 +745,7 @@ def test_e2e_untraced_taps_vs_golden(
     results.set_timing(untraced_first_s=time.perf_counter() - t0, untraced_device_s=model.timing["last_untraced_s"])
     assert tuple(pred.shape) == tuple(dl.noise_shape) and pred.dtype == torch.float32
     assert bool(torch.isfinite(pred).all()), "non-finite action_pred"
-    expected = set(M.expected_tap_keys(model.cfg, dl))
+    expected = set(M.expected_tap_keys(model.cfg, dl, model.policy.dit_backend))
     produced = {k for k in taps if not k.startswith(M.RAW_PREFIX)}
     assert expected <= produced, sorted(expected - produced)
     assert torch.equal(taps["action_pred_normalized"], pred)
@@ -708,7 +753,9 @@ def test_e2e_untraced_taps_vs_golden(
     cand = golden_keyed(taps, gs)
     cand.update(decoded_groups(model, pred, obs))
     cand, info = split_info(version, cand)
-    required = required_gated_keys(version, gs, dl) | {k for k in cand if k.startswith(M.DECODE_TAP)}
+    required = required_gated_keys(version, gs, dl, model.policy.dit_backend) | {
+        k for k in cand if k.startswith(M.DECODE_TAP)
+    }
     assert set(cand) == required, (sorted(required - set(cand)), sorted(set(cand) - required))
     report = harness.check_taps(version, cand, gs, results=results)
     print(report.summary())

code/models/experimental/gr00t/tests/tt/test_tt_perf.py

@@ -17,9 +17,23 @@ Device tests (one pytest process, through the lock)::
 Baseline = the last **committed** ``benchmarks/results/<bench>_<version>_*.json`` (``git ls-files``); because this
 package is currently untracked in ``tt-metal`` the lookup falls back to the newest *measured* JSON present (dry-run
 documents never count), and ``bench_stage_ops`` also accepts a ``_all_`` document (its stages carry ``versions``).
-``GR00T_PERF_BASELINE`` overrides: ``committed`` (default) | ``latest`` | ``none`` | ``/path/to/baseline.json``.
-Every device test **skips** (never fails) when no baseline exists, and writes its own JSON through the benchmark
-writer plus the comparison rows through the ``results`` fixture (``tests/tt/results/``).
+For ``bench_e2e`` the baseline must have been measured with the **same DiT backend** as the run (``meta.dit_backend``;
+documents from before WP-K5 carry none and count as ``ttnn``), so the megakernel default -- ``TTPolicy.dit_backend``
+since 2026-09-18 -- compares with the megakernel-default bench (``bench_e2e_<v>_20260918-*`` of the megakernel-default
+pass, ``tests/tt/results/mk_default_summary.md``) and ``GR00T_DIT_BACKEND=ttnn`` with the Stage-1 bench.
+``GR00T_PERF_BASELINE`` overrides: ``committed`` (default) | ``latest`` | ``none`` | ``/path/to/baseline.json`` (its
+backend must match).  Every device test **skips** (never fails) when no baseline exists, and writes its own JSON
+through the benchmark writer plus the comparison rows through the ``results`` fixture (``tests/tt/results/``).
+``ms`` stages whose baseline is below :data:`PERF_MIN_BASELINE_MS` (0.5 ms: the per-tensor ``upload/<input>`` copies,
+``d2h``) are recorded but not compared -- one 20 us host-timer jitter is +20 % there (megakernel-default pass).
+
+Devices (one open per test, closed after it): the whole-model bench runs on ``tt.model.open_model_device(policy,
+num_command_queues=2)`` -- the production configuration (the megakernel's reduced ``worker_l1_size``, CQ-1 input writes)
+-- for the policy of the run (``TTPolicy(dtype_policy=--policy, trace_layout=--trace-layout).with_env_overrides()``);
+``bench_stage_ops`` and ``bench_denoise`` run on the firmware-default device their baselines were measured on (the
+standalone N1.7 Qwen-ViT ``ff1`` op does not fit its static CBs under the 64 KiB cut).  ``bench_denoise`` times the
+Stage-1 (ttnn) denoise trace on a Stage-1 model, so ``test_denoise_within_10pct`` runs only under
+``GR00T_DIT_BACKEND=ttnn`` (skipped otherwise).
 
 CPU tests (no ttnn, run in every CI pass) pin the comparison rule, the baseline lookup and the three ``--dry-run``
 paths; the ``version`` fixture comes from ``tests/tt/conftest.py`` (``--gr00t-version``).
@@ -30,27 +44,74 @@ from __future__ import annotations
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, List, Mapping, Optional, Sequence
+from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence
 
 import pytest
 
 from models.experimental.gr00t.benchmarks import bench_denoise as BD
 from models.experimental.gr00t.benchmarks import bench_e2e as E2E
+from models.experimental.gr00t.benchmarks import bench_load as BL
 from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
+from models.experimental.gr00t.tt.policy import TTPolicy
 
 TOL = 0.10  # plan §7.1: +10 %
+#: stages whose baseline is below this many ms are recorded but not compared: the +10 % rule is meaningless on the
+#: per-tensor ``upload/<input>`` copies and ``d2h`` (20-150 us, where one 20 us host-timer jitter is +20 %); the
+#: stage totals (``upload``, ``device``, ``traces``, per trace, host encode / decode, ``e2e``) are all above it
+PERF_MIN_BASELINE_MS = 0.5
 BASELINE_ENV = "GR00T_PERF_BASELINE"
 BASELINE_MODES = ("committed", "latest", "none")
 RULE = "perf+10%"
+#: ``meta.dit_backend`` of perf documents written before WP-K5 (the Stage-1 path was the only one)
+LEGACY_BACKEND = "ttnn"
+
+
+def doc_backend(doc: Mapping[str, Any]) -> str:
+    """The DiT backend a ``bench_e2e`` document was measured with (``meta.dit_backend``; :data:`LEGACY_BACKEND` when the
+    document predates the field)."""
+    meta = doc.get("meta") or {}
+    b = meta.get("dit_backend")
+    if b is None:
+        b = (meta.get("policy") or {}).get("dit_backend")
+    return str(b) if b else LEGACY_BACKEND
 
 
 # --------------------------------------------------------------------------------------------------------------------
 # Baseline lookup
 # --------------------------------------------------------------------------------------------------------------------
+def _newest_matching(
+    bench: str, version: str, results_dir: Path, prefer_committed: bool, exclude: Sequence[Path], backend: str
+) -> Optional[Path]:
+    """Newest measured ``<bench>_<version>_*.json`` whose :func:`doc_backend` is ``backend`` (committed first)."""
+    ex = {Path(p).resolve() for p in exclude}
+
+    def ok(p: Path) -> bool:
+        if p.resolve() in ex:
+            return False
+        try:
+            doc = SO.load_perf_json(p)
+        except (OSError, ValueError, KeyError):
+            return False
+        return not doc["dry_run"] and doc_backend(doc) == backend
+
+    if prefer_committed:
+        committed = [p for p in SO.committed_perf_json(bench, version, results_dir) if ok(p)]
+        if committed:
+            return committed[-1]
+    present = [p for p in SO.list_perf_json(bench, version, results_dir) if ok(p)]
+    return present[-1] if present else None
+
+
 def baseline_for(
-    bench: str, version: str, results_dir: Path = SO.RESULTS_DIR, exclude: Sequence[Path] = ()
+    bench: str,
+    version: str,
+    results_dir: Path = SO.RESULTS_DIR,
+    exclude: Sequence[Path] = (),
+    backend: Optional[str] = None,
 ) -> Optional[Path]:
-    """The baseline JSON for ``(bench, version)`` or ``None`` (module docstring rules + ``GR00T_PERF_BASELINE``)."""
+    """The baseline JSON for ``(bench, version)`` or ``None`` (module docstring rules + ``GR00T_PERF_BASELINE``).
+    ``backend`` (``bench_e2e`` only) restricts the lookup to documents measured with that DiT backend and refuses an
+    explicit ``GR00T_PERF_BASELINE`` file of another backend."""
     mode = os.environ.get(BASELINE_ENV, "committed").strip()
     if mode == "none":
         return None
@@ -61,8 +122,12 @@ def baseline_for(
         doc = SO.load_perf_json(p)
         if doc["bench"] != bench:
             raise ValueError(f"{p}: bench {doc['bench']!r} != {bench!r}")
+        if backend is not None and doc_backend(doc) != backend:
+            raise ValueError(f"{p}: measured with dit_backend {doc_backend(doc)!r}, this run is {backend!r}")
         return p
     prefer = mode == "committed"
+    if backend is not None:
+        return _newest_matching(bench, version, results_dir, prefer, exclude, backend)
     p = SO.latest_perf_json(bench, version, results_dir, prefer_committed=prefer, exclude=exclude)
     if p is None and bench == SO.BENCH_NAME and version != "all":
         p = SO.latest_perf_json(bench, "all", results_dir, prefer_committed=prefer, exclude=exclude)
@@ -83,6 +148,20 @@ def stages_for_version(doc: Mapping[str, Any], version: str) -> Dict[str, Dict[s
     return out
 
 
+def apply_floor(cmp: Sequence[Mapping[str, Any]], floor_ms: float = PERF_MIN_BASELINE_MS) -> List[Dict[str, Any]]:
+    """Mark the compared rows whose ``ms`` baseline is below ``floor_ms`` as not compared (recorded as information,
+    :data:`PERF_MIN_BASELINE_MS`); ``us`` rows (``bench_stage_ops``) are never floored."""
+    out: List[Dict[str, Any]] = []
+    for c in cmp:
+        c = dict(c)
+        if c.get("compared") and c.get("metric") == "ms" and c.get("baseline") is not None and c["baseline"] < floor_ms:
+            c["compared"] = False
+            c["ok"] = None
+            c["note"] = f"baseline {c['baseline']:.3f} ms below the {floor_ms} ms comparison floor (host-timer jitter)"
+        out.append(c)
+    return out
+
+
 def comparison_rows(bench: str, version: str, cmp: Sequence[Mapping[str, Any]], baseline: Path) -> List[Dict[str, Any]]:
     """Rows in the ``ResultsWriter.add_row`` contract (``tap``, ``key``, ``rule``, ``ok`` + the perf facts)."""
     rows = []
@@ -149,6 +228,18 @@ def test_compare_stages_rule():
         SO.compare_stages({}, {}, tol=-0.1)
 
 
+def test_apply_floor_rule():
+    base = {"upload/patches": _st(0.093, "ms"), "d2h": _st(0.061, "ms"), "upload": _st(1.295, "ms"), "op": _st(100.0)}
+    cur = {"upload/patches": _st(0.113, "ms"), "d2h": _st(0.067, "ms"), "upload": _st(1.5, "ms"), "op": _st(125.0)}
+    rows = {r["stage"]: r for r in apply_floor(SO.compare_stages(cur, base, TOL))}
+    assert rows["upload/patches"]["compared"] is False and rows["upload/patches"]["ok"] is None  # 93 us: floored
+    assert rows["d2h"]["compared"] is False and "floor" in rows["d2h"]["note"]
+    assert rows["upload"]["compared"] is True and rows["upload"]["ok"] is False  # +15.8 % on a 1.3 ms stage fails
+    assert rows["op"]["compared"] is True and rows["op"]["ok"] is False  # us rows are never floored
+    with pytest.raises(pytest.skip.Exception):
+        _assert_within([rows["d2h"]], "x")  # nothing left to compare -> skip, never a silent pass
+
+
 def test_perf_payload_validation_and_roundtrip(tmp_path: Path):
     with pytest.raises(KeyError):
         SO.perf_payload("x", "n16", dry_run=True, args={}, method={}, stages={"s": {"value": 1.0}}, rows=[])
@@ -224,6 +315,42 @@ def test_baseline_for_env_modes(tmp_path: Path, monkeypatch: pytest.MonkeyPatch)
     assert set(stages_for_version(SO.load_perf_json(pa), "all")) == {"siglip_352/bf16", "llm_384/bf16"}
 
 
+def test_baseline_for_backend_aware(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    """``bench_e2e`` baselines are matched on ``meta.dit_backend`` (legacy documents count as ttnn)."""
+
+    def write(stamp: str, backend: Optional[str], value: float = 1.0) -> Path:
+        meta = {} if backend is None else {"dit_backend": backend}
+        payload = SO.perf_payload(
+            E2E.BENCH_NAME,
+            "n16",
+            dry_run=False,
+            args={},
+            method={},
+            stages={"e2e": _st(value, "ms")},
+            rows=[],
+            meta=meta,
+        )
+        return SO.write_perf_json(payload, tmp_path, path=tmp_path / f"{E2E.BENCH_NAME}_n16_{stamp}.json")
+
+    monkeypatch.setenv(BASELINE_ENV, "latest")
+    legacy = write("20260101-000000", None)
+    ttnn_doc = write("20260102-000000", "ttnn")
+    mk_doc = write("20260103-000000", "megakernel")
+    assert doc_backend(SO.load_perf_json(legacy)) == "ttnn"
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path) == mk_doc  # backend-agnostic: the newest
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="megakernel") == mk_doc
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="ttnn") == ttnn_doc
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="ttnn", exclude=[ttnn_doc]) == legacy
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="megakernel", exclude=[mk_doc]) is None
+    monkeypatch.setenv(BASELINE_ENV, str(mk_doc))
+    assert baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="megakernel") == mk_doc
+    with pytest.raises(ValueError):
+        baseline_for(E2E.BENCH_NAME, "n16", tmp_path, backend="ttnn")
+    # the policy default is what the device tier compares against unless GR00T_DIT_BACKEND says otherwise
+    assert TTPolicy().with_env_overrides({}).dit_backend == "megakernel"
+    assert TTPolicy().with_env_overrides({"GR00T_DIT_BACKEND": "ttnn"}).dit_backend == "ttnn"
+
+
 def test_stage_ops_dry_run_all_versions():
     """Every stage x dtype resolves against the tables without ttnn (the WP-06 acceptance shape list)."""
     payload = SO.dry_run(version="all", verbose=False)
@@ -363,31 +490,112 @@ def test_replayer_contract_is_checked():
 
 
 def test_existing_results_documents_load():
+    """Every perf-schema document under ``benchmarks/results`` loads and names a known bench; documents of the other
+    schemas that share the directory (the K1 / K4 megakernel benches, ``gr00t-mk-k1/1`` / ``gr00t-mk-k4/1``) are
+    counted, not parsed as perf documents."""
     files = sorted(SO.RESULTS_DIR.glob("*.json"))
     if not files:
         pytest.skip("no benchmark results present yet")
+    n_perf, other_schemas = 0, set()
     for f in files:
+        with open(f) as fh:
+            schema = json.load(fh).get("schema")
+        if schema != SO.PERF_SCHEMA:
+            other_schemas.add(str(schema))
+            continue
         doc = SO.load_perf_json(f)
-        assert doc["bench"] in (SO.BENCH_NAME, E2E.BENCH_NAME, BD.BENCH_NAME), f
+        assert doc["bench"] in (SO.BENCH_NAME, E2E.BENCH_NAME, BD.BENCH_NAME, BL.BENCH_NAME), f
+        n_perf += 1
+    assert n_perf > 0, files
+    # the K1 / K4 megakernel benches (bench_mk_*.py) share the directory; anything new must be listed on purpose
+    assert other_schemas <= {"gr00t-mk-k1/1", "gr00t-mk-k4/1"}, other_schemas
 
 
 # --------------------------------------------------------------------------------------------------------------------
 # Device tests: re-measure and compare with the baseline (skipped without a baseline)
 # --------------------------------------------------------------------------------------------------------------------
-def _load_baseline(bench: str, version: str) -> Any:
-    p = baseline_for(bench, version)
+@pytest.fixture(scope="module")
+def perf_policy(policy_name: str, trace_layout: str) -> TTPolicy:
+    """The policy of this perf run: the option values plus ``GR00T_DIT_BACKEND`` / ``GR00T_MK_ARENA_DTYPE`` (the
+    code default is the megakernel backend)."""
+    return TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout).with_env_overrides()
+
+
+def _open(policy: TTPolicy, num_command_queues: int) -> Any:
+    from models.experimental.gr00t.tests.tt.conftest import (
+        DEFAULT_L1_SMALL_SIZE,
+        DEFAULT_TRACE_REGION_SIZE,
+        DEFAULT_TT_METAL_CACHE,
+    )
+    from models.experimental.gr00t.tt import model as M
+
+    os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+    return M.open_model_device(
+        policy,
+        trace_region_size=DEFAULT_TRACE_REGION_SIZE,
+        l1_small_size=DEFAULT_L1_SMALL_SIZE,
+        num_command_queues=num_command_queues,
+    )
+
+
+@pytest.fixture
+def perf_device(perf_policy: TTPolicy) -> Iterator[Any]:
+    """The production device configuration for ``perf_policy`` (the whole-model bench): ``tt.model.open_model_device``
+    (the megakernel's ``worker_l1_size`` when that is the backend) with **two** command queues, as ``bench_e2e
+    --num-cqs 2`` and the policy server open it.  Function-scoped (one device open per test, closed after it) and
+    requested lazily after the baseline check, so a baseline-less run never opens the device and the two device
+    configurations of this module never coexist."""
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    device = _open(perf_policy, 2)
+    try:
+        yield device
+    finally:
+        close_gr00t_device(device)
+
+
+@pytest.fixture
+def default_device() -> Iterator[Any]:
+    """The firmware-default device (``worker_l1_size`` 1,461,248 B, one CQ) for the Stage-1 microbenches:
+    ``bench_stage_ops`` places its standalone backbone ops with their own L1 budgets (the N1.7 Qwen-ViT ``ff1`` matmul
+    clashes with its static CBs under the megakernel's 64 KiB cut, ``program.cpp:2149``) and ``bench_denoise`` times
+    the Stage-1 model; their baselines were measured on this configuration.  Function-scoped like :func:`perf_device`.
+    """
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    device = _open(TTPolicy(dit_backend="ttnn"), 1)
+    try:
+        yield device
+    finally:
+        close_gr00t_device(device)
+
+
+def _load_baseline(bench: str, version: str, backend: Optional[str] = None) -> Any:
+    p = baseline_for(bench, version, backend=backend)
     if p is None:
-        pytest.skip(f"{bench}: no baseline JSON for {version} in {SO.RESULTS_DIR} (set {BASELINE_ENV} to override)")
+        what = f" measured with dit_backend {backend!r}" if backend else ""
+        pytest.skip(
+            f"{bench}: no baseline JSON for {version}{what} in {SO.RESULTS_DIR} (set {BASELINE_ENV} to override)"
+        )
     return p, SO.load_perf_json(p)
 
 
-def _device(request: pytest.FixtureRequest, results: Any) -> Any:
-    """The session device, requested **after** the baseline check so that a baseline-less run skips without opening
-    the device (a session-scoped fixture in the signature would be instantiated first)."""
+def _device(request: pytest.FixtureRequest, results: Any, fixture: str = "perf_device") -> Any:
+    """The device of ``fixture`` (``perf_device`` = the run policy's configuration, ``default_device`` = the firmware
+    default), requested **after** the baseline check so that a baseline-less run skips without opening the device (a
+    fixture in the signature would be instantiated first)."""
     from models.experimental.gr00t.tests.tt.harness import device_facts
-
-    device = request.getfixturevalue("gr00t_device")
-    results.set_device_facts(device_facts(device))
+    from models.experimental.gr00t.tt import model as M
+
+    device = request.getfixturevalue(fixture)
+    policy: TTPolicy = request.getfixturevalue("perf_policy")
+    facts = device_facts(device)
+    facts["worker_l1_size"] = M.device_worker_l1_size(device)
+    facts["device_fixture"] = fixture
+    facts["dit_backend"] = policy.dit_backend if fixture == "perf_device" else "ttnn"
+    facts["mk_arena_dtype"] = policy.mk_arena_dtype if fixture == "perf_device" else None
+    facts["num_command_queues"] = 2 if fixture == "perf_device" else 1
+    results.set_device_facts(facts)
     return device
 
 
@@ -402,43 +610,79 @@ def _record(
 def test_stage_ops_within_10pct(request: pytest.FixtureRequest, version: str, results: Any):
     """Plan §7.1 perf tier for the backbone blocks: every canonical stage of ``version`` within +10 % of the baseline."""
     baseline_path, baseline = _load_baseline(SO.BENCH_NAME, version)
-    device = _device(request, results)
+    device = _device(request, results, "default_device")
     with results.time("bench_stage_ops"):
         payload = SO.run(device, version=version, verbose=False)
     current = SO.write_perf_json(payload)
     assert not payload["errors"], f"bench_stage_ops reported op failures: {payload['errors']}"
     cur = stages_for_version(payload, version)
     base = stages_for_version(baseline, version)
-    cmp = SO.compare_stages(cur, base, TOL)
+    cmp = apply_floor(SO.compare_stages(cur, base, TOL))
     _record(results, SO.BENCH_NAME, version, cmp, baseline_path, current)
     _assert_within(cmp, SO.BENCH_NAME)
 
 
 def test_denoise_within_10pct(
-    request: pytest.FixtureRequest, version: str, results: Any, policy_name: str, trace_layout: str
+    request: pytest.FixtureRequest,
+    version: str,
+    results: Any,
+    policy_name: str,
+    trace_layout: str,
+    perf_policy: TTPolicy,
 ):
     if E2E.import_gr00t_tt() is None:
         pytest.skip("tt.model.Gr00tTT not importable (WP-14 pending)")
+    if perf_policy.dit_backend != "ttnn":
+        pytest.skip(
+            "bench_denoise times the Stage-1 (ttnn) denoise trace on a Stage-1 model; run with GR00T_DIT_BACKEND=ttnn "
+            f"(this run's backend: {perf_policy.dit_backend}, whose device configuration cannot host the L1 adapter)"
+        )
     baseline_path, baseline = _load_baseline(BD.BENCH_NAME, version)
-    device = _device(request, results)
+    device = _device(request, results, "default_device")
     with results.time("bench_denoise"):
         payload = BD.run(device, version=version, policy=policy_name, trace_layout=trace_layout, verbose=False)
     current = SO.write_perf_json(payload)
-    cmp = SO.compare_stages(stages_for_version(payload, version), stages_for_version(baseline, version), TOL)
+    cmp = apply_floor(
+        SO.compare_stages(stages_for_version(payload, version), stages_for_version(baseline, version), TOL)
+    )
     _record(results, BD.BENCH_NAME, version, cmp, baseline_path, current)
     _assert_within(cmp, BD.BENCH_NAME)
 
 
 def test_e2e_within_10pct(
-    request: pytest.FixtureRequest, version: str, results: Any, policy_name: str, trace_layout: str
+    request: pytest.FixtureRequest,
+    version: str,
+    results: Any,
+    policy_name: str,
+    trace_layout: str,
+    perf_policy: TTPolicy,
 ):
+    """Plan §7.1 perf tier for the whole traced model: every ``bench_e2e`` stage within +10 % of the newest baseline
+    measured with the same DiT backend (module docstring)."""
     if E2E.import_gr00t_tt() is None:
         pytest.skip("tt.model.Gr00tTT not importable (WP-14 pending)")
-    baseline_path, baseline = _load_baseline(E2E.BENCH_NAME, version)
+    baseline_path, baseline = _load_baseline(E2E.BENCH_NAME, version, backend=perf_policy.dit_backend)
     device = _device(request, results)
     with results.time("bench_e2e"):
-        payload = E2E.run(device, version=version, policy=policy_name, trace_layout=trace_layout, verbose=False)
+        payload = E2E.run(
+            device,
+            version=version,
+            policy=policy_name,
+            trace_layout=trace_layout,
+            verbose=False,
+            dit_backend=perf_policy.dit_backend,
+            mk_arena_dtype=perf_policy.mk_arena_dtype,
+        )
     current = SO.write_perf_json(payload)
-    cmp = SO.compare_stages(stages_for_version(payload, version), stages_for_version(baseline, version), TOL)
+    results.set_meta(
+        dit_backend=perf_policy.dit_backend,
+        baseline_backend=doc_backend(baseline),
+        mk_arena_dtype=payload["meta"].get("mk_arena_dtype"),
+        num_command_queues=payload["meta"].get("num_command_queues"),
+    )
+    assert doc_backend(baseline) == perf_policy.dit_backend, (doc_backend(baseline), perf_policy.dit_backend)
+    cmp = apply_floor(
+        SO.compare_stages(stages_for_version(payload, version), stages_for_version(baseline, version), TOL)
+    )
     _record(results, E2E.BENCH_NAME, version, cmp, baseline_path, current)
     _assert_within(cmp, E2E.BENCH_NAME)

code/models/experimental/gr00t/tests/tt/test_tt_stale_state.py

@@ -426,6 +426,42 @@ def test_cpu_emulated_n_text_toggle(
 # --------------------------------------------------------------------------------------------------------------------
 # device tier
 # --------------------------------------------------------------------------------------------------------------------
+# --------------------------------------------------------------------------------------------------------------------
+# DiT backend parametrisation (WP-K5): GR00T_DIT_BACKEND=ttnn|megakernel, GR00T_MK_ARENA_DTYPE=auto|bf16|bfp8_b
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.fixture(scope="module")
+def policy(policy_name: str, trace_layout: str) -> TTPolicy:
+    """conftest's ``TTPolicy(dtype_policy=--policy, trace_layout=--trace-layout)`` plus the environment overrides of
+    ``TTPolicy.with_env_overrides`` (the DiT backend and the megakernel arena dtype of this run)."""
+    return TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout).with_env_overrides()
+
+
+@pytest.fixture(scope="module")
+def gr00t_device(request: pytest.FixtureRequest, policy: TTPolicy) -> Iterator[Any]:
+    """The p150a opened per backend (``tt.model.open_model_device``: the megakernel needs the smaller K4
+    ``worker_l1_size``); same trace region / L1-small sizes as conftest's session fixture, one device per process."""
+    from models.experimental.gr00t.tests.tt.conftest import (
+        DEFAULT_L1_SMALL_SIZE,
+        DEFAULT_TRACE_REGION_SIZE,
+        DEFAULT_TT_METAL_CACHE,
+    )
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+    device = M.open_model_device(
+        policy, trace_region_size=DEFAULT_TRACE_REGION_SIZE, l1_small_size=DEFAULT_L1_SMALL_SIZE
+    )
+    facts = harness.device_facts(device)
+    facts["worker_l1_size"] = M.device_worker_l1_size(device)
+    facts["dit_backend"] = policy.dit_backend
+    facts["mk_arena_dtype"] = policy.mk_arena_dtype if policy.dit_backend == "megakernel" else None
+    request.config._gr00t_device_facts = facts  # type: ignore[attr-defined]
+    try:
+        yield device
+    finally:
+        close_gr00t_device(device)
+
+
 @pytest.fixture(scope="module")
 def pair(version: str) -> List[Tuple[str, GoldenSet, Any]]:
     """``[(label, GoldenSet, Observation), ...]`` of the pair (skipped when a sample, its ``obs/`` or the snapshot is

image/blobs/sha256/0be2c39e4a01f307b061822a3a6d0b098e495a20c8ca0870953f0a73a05e4e60

@@ -0,0 +1 @@
+{"architecture":"amd64","config":{"User":"tt","Env":["PATH=/opt/tt-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin","VENV=/opt/tt-venv","VIRTUAL_ENV=/opt/tt-venv","TT_METAL_RUNTIME_ROOT=/opt/tt-metal","TT_METAL_HOME=/opt/tt-metal","PYTHONPATH=/opt/tt-metal","LD_LIBRARY_PATH=/opt/tt-metal/build/lib:/opt/openmpi-v5.0.7-ulfm/lib","EXTRA_MODELS_DIR=","TT_VLLM_BUILTIN_MODELS=","TT_MODEL_KIND=tt-dit-server","HF_HOME=/hf","TT_METAL_CACHE=/cache","HOME=/home/tt","USER=tt","LOGNAME=tt"],"Entrypoint":["/usr/local/bin/entrypoint.sh"],"Cmd":["/usr/local/bin/serve-default.sh"],"WorkingDir":"/home/tt/work","Labels":{"org.opencontainers.image.revision":"668c290757550588d0ce46b180c344a462a2aaf5","org.opencontainers.image.version":"22.04","org.tenstorrent.tt-model":"gr00t-n1.6-3b-p150","org.tenstorrent.tt-model.arch":"blackhole","org.tenstorrent.tt-model.kind":"tt-dit-server","org.tenstorrent.tt-model.plugin":"","org.tenstorrent.tt-model.profiles":"default","org.tenstorrent.tt-model.repo":"changh95/GR00T-N1.6-3B-p150","org.tenstorrent.tt-model.tt-metal":"v0.79.0-dev20260914-1-g668c290757","org.tenstorrent.tt-model.weights":"nvidia/GR00T-N1.6-3B"},"ArgsEscaped":true},"created":"2026-09-18T05:45:28.071165066+09:00","history":[{"created":"2026-09-03T12:14:51.139348462Z","created_by":"/bin/sh -c #(nop)  ARG RELEASE","empty_layer":true},{"created":"2026-09-03T12:14:51.170427469Z","created_by":"/bin/sh -c #(nop)  ARG LAUNCHPAD_BUILD_ARCH","empty_layer":true},{"created":"2026-09-03T12:14:51.200307479Z","created_by":"/bin/sh -c #(nop)  LABEL org.opencontainers.image.version=22.04","empty_layer":true},{"created":"2026-09-03T12:14:53.528689977Z","created_by":"/bin/sh -c #(nop) ADD file:81c01921c5f642ac2fcbfae682e489e8e64b347467d9fa1587707e310e64d790 in / "},{"created":"2026-09-03T12:14:54.045714841Z","created_by":"/bin/sh -c #(nop)  CMD [\"/bin/bash\"]","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG OMPI_DIR=/opt/openmpi-v5.0.7-ulfm","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG EXTRA_MODELS_DIR=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG TT_MODEL_KIND","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_NAME","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_REPO","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_WEIGHTS","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_ARCH","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_PROFILES","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_TT_METAL_SHA","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_TT_METAL_DESCRIBE","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"ARG MODEL_PLUGIN_SHA","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-15T00:17:47.388335211+09:00","created_by":"RUN |11 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=gr00t-n1.6-3b-p150 MODEL_REPO=changh95/GR00T-N1.6-3B-p150 MODEL_WEIGHTS=nvidia/GR00T-N1.6-3B MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=668c290757550588d0ce46b180c344a462a2aaf5 MODEL_TT_METAL_DESCRIBE=v0.79.0-dev20260914-1-g668c290757 MODEL_PLUGIN_SHA= /bin/sh -c apt-get update \u0026\u0026 apt-get install -y --no-install-recommends     libhwloc15 libnuma1 libatomic1 libudev1 libcap2 zlib1g     libmpc3 libmpfr6 libgmp10 libzstd1     libevent-core-2.1-7 libevent-pthreads-2.1-7     libgl1 libsndfile1 ca-certificates     \u0026\u0026 apt-get clean \u0026\u0026 rm -rf /var/lib/apt/lists/* # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-15T00:17:47.741172096+09:00","created_by":"RUN |11 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=gr00t-n1.6-3b-p150 MODEL_REPO=changh95/GR00T-N1.6-3B-p150 MODEL_WEIGHTS=nvidia/GR00T-N1.6-3B MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=668c290757550588d0ce46b180c344a462a2aaf5 MODEL_TT_METAL_DESCRIBE=v0.79.0-dev20260914-1-g668c290757 MODEL_PLUGIN_SHA= /bin/sh -c existing=\"$(getent passwd 1000 | cut -d: -f1)\"     \u0026\u0026 if [ -n \"$existing\" ]; then          userdel -r \"$existing\" 2\u003e/dev/null || userdel \"$existing\";        fi     \u0026\u0026 useradd --uid 1000 --create-home --home-dir /home/tt --shell /bin/bash tt     \u0026\u0026 mkdir -p /home/tt/work/logs /cache /opt/tt-metal     \u0026\u0026 chown -R tt:tt /home/tt /cache /opt/tt-metal     \u0026\u0026 chmod 1777 /home/tt /home/tt/work /home/tt/work/logs /cache # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-15T00:26:50.424500566+09:00","created_by":"COPY /opt/openmpi-v5.0.7-ulfm /opt/openmpi-v5.0.7-ulfm # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-15T00:26:56.847494353+09:00","created_by":"COPY /opt/tenstorrent /opt/tenstorrent # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:31:20.970763383+09:00","created_by":"COPY /usr/local/share/uv /usr/local/share/uv # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:43:55.192165457+09:00","created_by":"COPY /opt/tt-venv /opt/tt-venv # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:43:55.320045222+09:00","created_by":"COPY /opt/vllm /opt/vllm # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:00.176020413+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/runtime /opt/tt-metal/runtime # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:02.947444345+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/build_Release /opt/tt-metal/build_Release # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:05.896621192+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/build /opt/tt-metal/build # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:11.794830059+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/tt_metal /opt/tt-metal/tt_metal # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:16.009790261+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/ttnn /opt/tt-metal/ttnn # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:17.508565672+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/tools /opt/tt-metal/tools # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:18.531013694+09:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/setup.py /opt/tt-metal/pyproject.toml /opt/tt-metal/ # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:20.210963815+09:00","created_by":"COPY --chown=tt:tt code/ /opt/tt-metal/ # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:21.69976708+09:00","created_by":"COPY entrypoint.sh /usr/local/bin/entrypoint.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"COPY --chmod=0755 serve-default.sh /usr/local/bin/serve-default.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV VENV=/opt/tt-venv","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV VIRTUAL_ENV=/opt/tt-venv","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV PATH=/opt/tt-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV TT_METAL_RUNTIME_ROOT=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV TT_METAL_HOME=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV PYTHONPATH=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV LD_LIBRARY_PATH=/opt/tt-metal/build/lib:/opt/openmpi-v5.0.7-ulfm/lib","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV EXTRA_MODELS_DIR=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ARG TT_VLLM_BUILTIN_MODELS=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV TT_VLLM_BUILTIN_MODELS=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV TT_MODEL_KIND=tt-dit-server","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV HF_HOME=/hf","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV TT_METAL_CACHE=/cache","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV HOME=/home/tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV USER=tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"ENV LOGNAME=tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:22.388246477+09:00","created_by":"USER tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:44:23.450079066+09:00","created_by":"WORKDIR /home/tt/work","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:44:24.890053981+09:00","created_by":"COPY verify.sh /ctx/verify.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:45:27.833267805+09:00","created_by":"RUN |12 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=gr00t-n1.6-3b-p150 MODEL_REPO=changh95/GR00T-N1.6-3B-p150 MODEL_WEIGHTS=nvidia/GR00T-N1.6-3B MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=668c290757550588d0ce46b180c344a462a2aaf5 MODEL_TT_METAL_DESCRIBE=v0.79.0-dev20260914-1-g668c290757 MODEL_PLUGIN_SHA= TT_VLLM_BUILTIN_MODELS= /bin/sh -c bash /ctx/verify.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:45:27.833267805+09:00","created_by":"USER root","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:45:28.071165066+09:00","created_by":"RUN |12 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=gr00t-n1.6-3b-p150 MODEL_REPO=changh95/GR00T-N1.6-3B-p150 MODEL_WEIGHTS=nvidia/GR00T-N1.6-3B MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=668c290757550588d0ce46b180c344a462a2aaf5 MODEL_TT_METAL_DESCRIBE=v0.79.0-dev20260914-1-g668c290757 MODEL_PLUGIN_SHA= TT_VLLM_BUILTIN_MODELS= /bin/sh -c chmod -R a+rwX /home/tt # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-09-18T05:45:28.071165066+09:00","created_by":"USER tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:45:28.071165066+09:00","created_by":"LABEL org.tenstorrent.tt-model=gr00t-n1.6-3b-p150 org.tenstorrent.tt-model.repo=changh95/GR00T-N1.6-3B-p150 org.tenstorrent.tt-model.weights=nvidia/GR00T-N1.6-3B org.tenstorrent.tt-model.arch=blackhole org.tenstorrent.tt-model.kind=tt-dit-server org.tenstorrent.tt-model.profiles=default org.opencontainers.image.revision=668c290757550588d0ce46b180c344a462a2aaf5 org.tenstorrent.tt-model.tt-metal=v0.79.0-dev20260914-1-g668c290757 org.tenstorrent.tt-model.plugin=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:45:28.071165066+09:00","created_by":"ENTRYPOINT [\"/usr/local/bin/entrypoint.sh\"]","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-09-18T05:45:28.071165066+09:00","created_by":"CMD [\"/usr/local/bin/serve-default.sh\"]","comment":"buildkit.dockerfile.v0","empty_layer":true}],"os":"linux","rootfs":{"type":"layers","diff_ids":["sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34","sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f","sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102","sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0","sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89","sha256:ea7f9e8dc4b32aa84e39e01475e6e4ad24f6d4c175388103a9e93d1b021b3b7d","sha256:dddd09bc1156c04e4431e15e7adc290e18581a04d7545030b8115be07d8ba8aa","sha256:569f156189a8b1b0abf25d42c095d4342c66bad7f359f8f44d174a5b7d06d6f1","sha256:27b911cd5c2d8d210063fc08ff3f4326c701ddc06965cdf21c47f64fb5fa7c9f","sha256:f64479d7869b23160b3944a14abfb4b9a40ae5e70fd1b58eba3319d3253f8853","sha256:7cfb2f74d1f012e7c735a943ac188c0b72761cf2c12f815a9768b127b8350f23","sha256:6c625c305958b1775ab7779936b8bb3be1493dce685f9b26f68abebd74c46d98","sha256:15f1b153eb697ee10f2a6b85ef43a386725544ab4e38327b19429b7f5bcf149c","sha256:5c18d58ef8b510cafbfd183ad353acfd6ec015e633ac202ffb0559fc462c048d","sha256:f0ef8bb8e669c94f302f8c3bb22f260910987963e3b64dd82d715f00df9ce9e5","sha256:660e84fb235e1a8e1281edec1404998d8a70dade16785315b5857ca9b7e48e83","sha256:a5a17edd1988763853752ee15d5bd4d6d2584a3393aec695e8d8b3862758a279","sha256:6f2c01a2e64a5954458cbbe49c32378dcf04517278ee030ae5c00424324f3529","sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef","sha256:91a77e841802c172ee3d5b6b162272bf0732cd8a54a7600d7250c5b788f2bbd4","sha256:f07d2440aabdc8a01341c7dbf2a12c605d15a367a972f3c3035401bedeb433c5","sha256:909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418"]}}

image/blobs/sha256/0c65a290603e6797836c24eaf4dbad448eb486ddd37d5b22ec615fd70cb133f2

@@ -0,0 +1 @@
+{"id":"b541132a377e8633a4983b59756ce33da5a2151a7b644a20f8e9fe4dc2719ddc","parent":"847be5b832733bb8d6e8ea3d80c5a09566773f3d6677ddce762e7e11bbf3f13e","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/1c590b200182d3f28659786c6b802d5c9fb63bef9a1aa09e5efdde2591a12c20

@@ -0,0 +1 @@
+{"id":"7813508801c5c352cf188ed0de962c1d8096aadaab2def10d03e4bd2139876b6","parent":"d76bfba1989c444294dd385520690ddd3375bf472931f08eb356bf17de161d5b","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/252a9530fcab61e4872b019330842abd14e288837bb494eabf7d961c2c623112

@@ -0,0 +1 @@
+{"id":"f699578393b121282585192635bfea1ff2a5107ed5590bf51a89ec10f8efa818","parent":"18d9869bf524b4834b53475ddb58956c35b519fcf3f76047114982881b3abfb9","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/3c4361cee2aeb8ae93f4ef84a1a82ff11ad20abeea0aeb0c97e9bf3be6859591

@@ -0,0 +1 @@
+{"id":"3d53fe9bdacbd53d9801d22baf7ddd7efdacaab8a8b84a631bdbdf597ddebdf5","parent":"7c9732243332e5b463ca4352fdb3cf27f3b357bddbeec5ae377cc9a3303a08b9","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/3f7cf805baca40c6075bef88421935d03e2e28cbdb906610cd51b7b761523c12

@@ -0,0 +1 @@
+{"id":"b72df2a436cbacbf641552f5620b58733416506634f7459ef7eac667ddbc993d","parent":"3d53fe9bdacbd53d9801d22baf7ddd7efdacaab8a8b84a631bdbdf597ddebdf5","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/5e8401978135bdfadc3dd329a54851b7fa7a06c24bb959b536e75966fea6b278

@@ -0,0 +1 @@
+{"schemaVersion":2,"mediaType":"application/vnd.oci.image.manifest.v1+json","config":{"mediaType":"application/vnd.oci.image.config.v1+json","digest":"sha256:0be2c39e4a01f307b061822a3a6d0b098e495a20c8ca0870953f0a73a05e4e60","size":14445},"layers":[{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34","size":80641024},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f","size":204086784},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102","size":352256},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0","size":33464832},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89","size":454126592},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:ea7f9e8dc4b32aa84e39e01475e6e4ad24f6d4c175388103a9e93d1b021b3b7d","size":187652608},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:dddd09bc1156c04e4431e15e7adc290e18581a04d7545030b8115be07d8ba8aa","size":1237841408},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:569f156189a8b1b0abf25d42c095d4342c66bad7f359f8f44d174a5b7d06d6f1","size":2048},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:27b911cd5c2d8d210063fc08ff3f4326c701ddc06965cdf21c47f64fb5fa7c9f","size":454626304},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:f64479d7869b23160b3944a14abfb4b9a40ae5e70fd1b58eba3319d3253f8853","size":136048640},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:7cfb2f74d1f012e7c735a943ac188c0b72761cf2c12f815a9768b127b8350f23","size":136048640},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:6c625c305958b1775ab7779936b8bb3be1493dce685f9b26f68abebd74c46d98","size":195663360},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:15f1b153eb697ee10f2a6b85ef43a386725544ab4e38327b19429b7f5bcf149c","size":69021184},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:5c18d58ef8b510cafbfd183ad353acfd6ec015e633ac202ffb0559fc462c048d","size":5885952},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:f0ef8bb8e669c94f302f8c3bb22f260910987963e3b64dd82d715f00df9ce9e5","size":23040},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:660e84fb235e1a8e1281edec1404998d8a70dade16785315b5857ca9b7e48e83","size":15580672},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:a5a17edd1988763853752ee15d5bd4d6d2584a3393aec695e8d8b3862758a279","size":4608},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:6f2c01a2e64a5954458cbbe49c32378dcf04517278ee030ae5c00424324f3529","size":5120},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef","size":1024},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:91a77e841802c172ee3d5b6b162272bf0732cd8a54a7600d7250c5b788f2bbd4","size":8704},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:f07d2440aabdc8a01341c7dbf2a12c605d15a367a972f3c3035401bedeb433c5","size":1763840},{"mediaType":"application/vnd.oci.image.layer.v1.tar","digest":"sha256:909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418","size":39936}]}

image/blobs/sha256/7f8ab8adfddcbb515e20c539f42c507b2bd9f9ad8e27e87719e5487da6125f89

@@ -0,0 +1 @@
+{"id":"0a2097b734d0b43c3ff587ac77bd60b79fab2bade954c1b98eb6e92294f5d8d3","parent":"7813508801c5c352cf188ed0de962c1d8096aadaab2def10d03e4bd2139876b6","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/8dbbd756d63b80f87f9c009aeb284a337a9457d21e97317066c14cdd8f21349f

@@ -0,0 +1 @@
+{"id":"18d9869bf524b4834b53475ddb58956c35b519fcf3f76047114982881b3abfb9","parent":"227145cc5d36b8622ca2eb65f3ed53c2b53ca29eec7b9ee823f530d34a63e74e","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/8eedd596896beb613d02b359b843aadeb06cb5ec8effbdbdf697632d1f238328

@@ -0,0 +1 @@
+{"id":"7c9732243332e5b463ca4352fdb3cf27f3b357bddbeec5ae377cc9a3303a08b9","parent":"a6e66eeb79c678923678ed7399fe67597aa0f66125b3477d930ce351ad58ce3c","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/97a4a5f03aa47006f66965fb4a4f06010d695fc04d7683459bfdd729381c399d

@@ -0,0 +1 @@
+{"id":"da6a3c97d1c3758896789dd2b1486d422669701f2cf7f7625feda5c2f5daca7a","parent":"b72df2a436cbacbf641552f5620b58733416506634f7459ef7eac667ddbc993d","created":"2026-09-18T05:45:28.071165066+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"config":{"Hostname":"","Domainname":"","User":"tt","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":["PATH=/opt/tt-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin","VENV=/opt/tt-venv","VIRTUAL_ENV=/opt/tt-venv","TT_METAL_RUNTIME_ROOT=/opt/tt-metal","TT_METAL_HOME=/opt/tt-metal","PYTHONPATH=/opt/tt-metal","LD_LIBRARY_PATH=/opt/tt-metal/build/lib:/opt/openmpi-v5.0.7-ulfm/lib","EXTRA_MODELS_DIR=","TT_VLLM_BUILTIN_MODELS=","TT_MODEL_KIND=tt-dit-server","HF_HOME=/hf","TT_METAL_CACHE=/cache","HOME=/home/tt","USER=tt","LOGNAME=tt"],"Cmd":["/usr/local/bin/serve-default.sh"],"ArgsEscaped":true,"Image":"","Volumes":null,"WorkingDir":"/home/tt/work","Entrypoint":["/usr/local/bin/entrypoint.sh"],"OnBuild":null,"Labels":{"org.opencontainers.image.revision":"668c290757550588d0ce46b180c344a462a2aaf5","org.opencontainers.image.version":"22.04","org.tenstorrent.tt-model":"gr00t-n1.6-3b-p150","org.tenstorrent.tt-model.arch":"blackhole","org.tenstorrent.tt-model.kind":"tt-dit-server","org.tenstorrent.tt-model.plugin":"","org.tenstorrent.tt-model.profiles":"default","org.tenstorrent.tt-model.repo":"changh95/GR00T-N1.6-3B-p150","org.tenstorrent.tt-model.tt-metal":"v0.79.0-dev20260914-1-g668c290757","org.tenstorrent.tt-model.weights":"nvidia/GR00T-N1.6-3B"}},"architecture":"amd64","os":"linux"}

image/blobs/sha256/be72b093968ed7fd244f0b1dff9dc05fd87d0031245c3a710ad2a3c0e55a772f

@@ -0,0 +1 @@
+{"id":"c456cdd8530267fc6889f8c1cf87762230e39652fcbad134f88ebc619ab0fc65","parent":"112dc954cd6c04a85cecc65f8188681e5493149aa0b4b7a62698ccda4edff3cc","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/c601ad8dea9cf1cb5c826fd622d446767719a4bdcc005207d9883b0bae3ec2f2

@@ -0,0 +1 @@
+{"id":"c89f024f2263f097b6bd9a82c8a69267c64e82f8cf00e97b0ff7b2970b6e877b","parent":"b541132a377e8633a4983b59756ce33da5a2151a7b644a20f8e9fe4dc2719ddc","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/e219f70f6edabc2003d4bee846c0f04bec4f41d668285011fda98a2ae63bdec7

@@ -0,0 +1 @@
+{"id":"847be5b832733bb8d6e8ea3d80c5a09566773f3d6677ddce762e7e11bbf3f13e","parent":"71d75d164a539e9fa8f69d3f211c58ff68c1d4cd7433c93704bee9980ec4c369","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/e3c4bbf0e1a172c093407d72bb7f93bfe42b5da41e10fbb57164701db74b06af

@@ -0,0 +1 @@
+{"id":"d76bfba1989c444294dd385520690ddd3375bf472931f08eb356bf17de161d5b","parent":"f699578393b121282585192635bfea1ff2a5107ed5590bf51a89ec10f8efa818","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/blobs/sha256/fb55b0f1dfab1bc8d7a83da0cfb21ce42baa0f8f3b85eaffab714c5fc4119cdc

@@ -0,0 +1 @@
+{"id":"227145cc5d36b8622ca2eb65f3ed53c2b53ca29eec7b9ee823f530d34a63e74e","parent":"c89f024f2263f097b6bd9a82c8a69267c64e82f8cf00e97b0ff7b2970b6e877b","created":"1970-01-01T09:00:00+09:00","container_config":{"Hostname":"","Domainname":"","User":"","AttachStdin":false,"AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false,"StdinOnce":false,"Env":null,"Cmd":null,"Image":"","Volumes":null,"WorkingDir":"","Entrypoint":null,"OnBuild":null,"Labels":null},"os":"linux"}

image/index.json

@@ -1 +1 @@
-{"schemaVersion":2,"mediaType":"application/vnd.oci.image.index.v1+json","manifests":[{"mediaType":"application/vnd.oci.image.manifest.v1+json","digest":"sha256:fbd897255c55c5be409bc1041428ff0b3c9ee37b59522a9962bd18c3884662c0","size":3616,"annotations":{"io.containerd.image.name":"docker.io/tt-model/gr00t-n1.6-3b-p150:feb77048a2ed","org.opencontainers.image.ref.name":"feb77048a2ed"}}]}
+{"schemaVersion":2,"mediaType":"application/vnd.oci.image.index.v1+json","manifests":[{"mediaType":"application/vnd.oci.image.manifest.v1+json","digest":"sha256:5e8401978135bdfadc3dd329a54851b7fa7a06c24bb959b536e75966fea6b278","size":3616,"annotations":{"io.containerd.image.name":"docker.io/tt-model/gr00t-n1.6-3b-p150:0be2c39e4a01","org.opencontainers.image.ref.name":"0be2c39e4a01"}}]}

image/manifest.json

@@ -1 +1 @@
-[{"Config":"blobs/sha256/feb77048a2ed1772037207da1e3fbf571555db9c9554509bad74076e7c94fa55","RepoTags":["tt-model/gr00t-n1.6-3b-p150:feb77048a2ed"],"Layers":["blobs/sha256/ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34","blobs/sha256/2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f","blobs/sha256/0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102","blobs/sha256/12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0","blobs/sha256/20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89","blobs/sha256/456928125bdb6641fde495b5089401d6eefdf8a6b8c9c52f14b0836c29f931d7","blobs/sha256/6ef3d5458860fc96e862cfc6df467f80a4db334ebe6bae7da8f09fda2d1cd758","blobs/sha256/197405c374ca76fc01e68a8f55f41a379d40389f10e9cfe029cca8e00a5abd0f","blobs/sha256/58c2b8faf6877c94ecd1468ab1ad99d4c4c73fb376613d8517b659fc09309496","blobs/sha256/a08d44c169e46f95b649984d7adb7cbf9171a6c3343c1ceafdf3cc23512b008c","blobs/sha256/ea07a201af443d5a600b22c02de0009e907acab32e5b033ed8df86d0a0113eb5","blobs/sha256/256e0cf0f7d85bbf0d3362be27dee0534059bc45011db9b5aad683550e991905","blobs/sha256/425f5a58dfa6749ff395b01ab459d4cc946046a8ef0d87280808494adea78ea0","blobs/sha256/da48df271e2ffd8c89608ed1e938fd4f5c1d2599b134f3218b4751f3fdc7acda","blobs/sha256/3c99da1ba1a01d71278c50c0d1a0571372f3a5026536fbe7e3428615b0274bd7","blobs/sha256/7f6c689c3e7be5fee57c085083155ae0721468b07eb39d97dbcb2f3b010404c2","blobs/sha256/658567390c2a80da8263482cbb57632be7bd0386b8a9c71d8f02b1657b2fdf82","blobs/sha256/77ef632adb98b5a998318d8a04d1d5dd6a73010530c50bfe6c18c47a92032ca2","blobs/sha256/5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef","blobs/sha256/71bfc9a481add3ed85649644d61cdbdc195c0cb06ddf2a79f0704eb71ac4e93d","blobs/sha256/f996a3012184d72d82d3fce69523ef7e4f47d317339d5286c93ea6a463d63466","blobs/sha256/aa10b8db9ce55f80f832cc35d23d3776ad72b6dc6dd95bc6b8dd9f1a9cde331b"],"LayerSources":{"sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":352256,"digest":"sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102"},"sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":33464832,"digest":"sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0"},"sha256:197405c374ca76fc01e68a8f55f41a379d40389f10e9cfe029cca8e00a5abd0f":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":2048,"digest":"sha256:197405c374ca76fc01e68a8f55f41a379d40389f10e9cfe029cca8e00a5abd0f"},"sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":454126592,"digest":"sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89"},"sha256:256e0cf0f7d85bbf0d3362be27dee0534059bc45011db9b5aad683550e991905":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":195663360,"digest":"sha256:256e0cf0f7d85bbf0d3362be27dee0534059bc45011db9b5aad683550e991905"},"sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":204086784,"digest":"sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f"},"sha256:3c99da1ba1a01d71278c50c0d1a0571372f3a5026536fbe7e3428615b0274bd7":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":23040,"digest":"sha256:3c99da1ba1a01d71278c50c0d1a0571372f3a5026536fbe7e3428615b0274bd7"},"sha256:425f5a58dfa6749ff395b01ab459d4cc946046a8ef0d87280808494adea78ea0":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":69021184,"digest":"sha256:425f5a58dfa6749ff395b01ab459d4cc946046a8ef0d87280808494adea78ea0"},"sha256:456928125bdb6641fde495b5089401d6eefdf8a6b8c9c52f14b0836c29f931d7":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":187652608,"digest":"sha256:456928125bdb6641fde495b5089401d6eefdf8a6b8c9c52f14b0836c29f931d7"},"sha256:58c2b8faf6877c94ecd1468ab1ad99d4c4c73fb376613d8517b659fc09309496":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":454626304,"digest":"sha256:58c2b8faf6877c94ecd1468ab1ad99d4c4c73fb376613d8517b659fc09309496"},"sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1024,"digest":"sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef"},"sha256:658567390c2a80da8263482cbb57632be7bd0386b8a9c71d8f02b1657b2fdf82":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":4608,"digest":"sha256:658567390c2a80da8263482cbb57632be7bd0386b8a9c71d8f02b1657b2fdf82"},"sha256:6ef3d5458860fc96e862cfc6df467f80a4db334ebe6bae7da8f09fda2d1cd758":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1237653504,"digest":"sha256:6ef3d5458860fc96e862cfc6df467f80a4db334ebe6bae7da8f09fda2d1cd758"},"sha256:71bfc9a481add3ed85649644d61cdbdc195c0cb06ddf2a79f0704eb71ac4e93d":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":6656,"digest":"sha256:71bfc9a481add3ed85649644d61cdbdc195c0cb06ddf2a79f0704eb71ac4e93d"},"sha256:77ef632adb98b5a998318d8a04d1d5dd6a73010530c50bfe6c18c47a92032ca2":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":5120,"digest":"sha256:77ef632adb98b5a998318d8a04d1d5dd6a73010530c50bfe6c18c47a92032ca2"},"sha256:7f6c689c3e7be5fee57c085083155ae0721468b07eb39d97dbcb2f3b010404c2":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":10250752,"digest":"sha256:7f6c689c3e7be5fee57c085083155ae0721468b07eb39d97dbcb2f3b010404c2"},"sha256:a08d44c169e46f95b649984d7adb7cbf9171a6c3343c1ceafdf3cc23512b008c":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":136048640,"digest":"sha256:a08d44c169e46f95b649984d7adb7cbf9171a6c3343c1ceafdf3cc23512b008c"},"sha256:aa10b8db9ce55f80f832cc35d23d3776ad72b6dc6dd95bc6b8dd9f1a9cde331b":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":39936,"digest":"sha256:aa10b8db9ce55f80f832cc35d23d3776ad72b6dc6dd95bc6b8dd9f1a9cde331b"},"sha256:da48df271e2ffd8c89608ed1e938fd4f5c1d2599b134f3218b4751f3fdc7acda":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":5885952,"digest":"sha256:da48df271e2ffd8c89608ed1e938fd4f5c1d2599b134f3218b4751f3fdc7acda"},"sha256:ea07a201af443d5a600b22c02de0009e907acab32e5b033ed8df86d0a0113eb5":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":136048640,"digest":"sha256:ea07a201af443d5a600b22c02de0009e907acab32e5b033ed8df86d0a0113eb5"},"sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":80641024,"digest":"sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34"},"sha256:f996a3012184d72d82d3fce69523ef7e4f47d317339d5286c93ea6a463d63466":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1316864,"digest":"sha256:f996a3012184d72d82d3fce69523ef7e4f47d317339d5286c93ea6a463d63466"}}}]
+[{"Config":"blobs/sha256/0be2c39e4a01f307b061822a3a6d0b098e495a20c8ca0870953f0a73a05e4e60","RepoTags":["tt-model/gr00t-n1.6-3b-p150:0be2c39e4a01"],"Layers":["blobs/sha256/ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34","blobs/sha256/2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f","blobs/sha256/0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102","blobs/sha256/12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0","blobs/sha256/20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89","blobs/sha256/ea7f9e8dc4b32aa84e39e01475e6e4ad24f6d4c175388103a9e93d1b021b3b7d","blobs/sha256/dddd09bc1156c04e4431e15e7adc290e18581a04d7545030b8115be07d8ba8aa","blobs/sha256/569f156189a8b1b0abf25d42c095d4342c66bad7f359f8f44d174a5b7d06d6f1","blobs/sha256/27b911cd5c2d8d210063fc08ff3f4326c701ddc06965cdf21c47f64fb5fa7c9f","blobs/sha256/f64479d7869b23160b3944a14abfb4b9a40ae5e70fd1b58eba3319d3253f8853","blobs/sha256/7cfb2f74d1f012e7c735a943ac188c0b72761cf2c12f815a9768b127b8350f23","blobs/sha256/6c625c305958b1775ab7779936b8bb3be1493dce685f9b26f68abebd74c46d98","blobs/sha256/15f1b153eb697ee10f2a6b85ef43a386725544ab4e38327b19429b7f5bcf149c","blobs/sha256/5c18d58ef8b510cafbfd183ad353acfd6ec015e633ac202ffb0559fc462c048d","blobs/sha256/f0ef8bb8e669c94f302f8c3bb22f260910987963e3b64dd82d715f00df9ce9e5","blobs/sha256/660e84fb235e1a8e1281edec1404998d8a70dade16785315b5857ca9b7e48e83","blobs/sha256/a5a17edd1988763853752ee15d5bd4d6d2584a3393aec695e8d8b3862758a279","blobs/sha256/6f2c01a2e64a5954458cbbe49c32378dcf04517278ee030ae5c00424324f3529","blobs/sha256/5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef","blobs/sha256/91a77e841802c172ee3d5b6b162272bf0732cd8a54a7600d7250c5b788f2bbd4","blobs/sha256/f07d2440aabdc8a01341c7dbf2a12c605d15a367a972f3c3035401bedeb433c5","blobs/sha256/909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418"],"LayerSources":{"sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":352256,"digest":"sha256:0c2acf106608fb8af04e412572f17c510576f1086dfcd90a4789ba4fe6ab3102"},"sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":33464832,"digest":"sha256:12ce46944e0430cb86189d84b92ea4ba5110c40e2b965a27fd5dabea785158d0"},"sha256:15f1b153eb697ee10f2a6b85ef43a386725544ab4e38327b19429b7f5bcf149c":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":69021184,"digest":"sha256:15f1b153eb697ee10f2a6b85ef43a386725544ab4e38327b19429b7f5bcf149c"},"sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":454126592,"digest":"sha256:20cad74cf2e0718048ce7cad5032e4837d231af0ead9812353883170b27a8c89"},"sha256:27b911cd5c2d8d210063fc08ff3f4326c701ddc06965cdf21c47f64fb5fa7c9f":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":454626304,"digest":"sha256:27b911cd5c2d8d210063fc08ff3f4326c701ddc06965cdf21c47f64fb5fa7c9f"},"sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":204086784,"digest":"sha256:2c86a37deecae8c6c4d64575b68d161cde8651d057ce91ee6ad9df65ce22535f"},"sha256:569f156189a8b1b0abf25d42c095d4342c66bad7f359f8f44d174a5b7d06d6f1":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":2048,"digest":"sha256:569f156189a8b1b0abf25d42c095d4342c66bad7f359f8f44d174a5b7d06d6f1"},"sha256:5c18d58ef8b510cafbfd183ad353acfd6ec015e633ac202ffb0559fc462c048d":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":5885952,"digest":"sha256:5c18d58ef8b510cafbfd183ad353acfd6ec015e633ac202ffb0559fc462c048d"},"sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1024,"digest":"sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef"},"sha256:660e84fb235e1a8e1281edec1404998d8a70dade16785315b5857ca9b7e48e83":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":15580672,"digest":"sha256:660e84fb235e1a8e1281edec1404998d8a70dade16785315b5857ca9b7e48e83"},"sha256:6c625c305958b1775ab7779936b8bb3be1493dce685f9b26f68abebd74c46d98":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":195663360,"digest":"sha256:6c625c305958b1775ab7779936b8bb3be1493dce685f9b26f68abebd74c46d98"},"sha256:6f2c01a2e64a5954458cbbe49c32378dcf04517278ee030ae5c00424324f3529":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":5120,"digest":"sha256:6f2c01a2e64a5954458cbbe49c32378dcf04517278ee030ae5c00424324f3529"},"sha256:7cfb2f74d1f012e7c735a943ac188c0b72761cf2c12f815a9768b127b8350f23":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":136048640,"digest":"sha256:7cfb2f74d1f012e7c735a943ac188c0b72761cf2c12f815a9768b127b8350f23"},"sha256:909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":39936,"digest":"sha256:909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418"},"sha256:91a77e841802c172ee3d5b6b162272bf0732cd8a54a7600d7250c5b788f2bbd4":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":8704,"digest":"sha256:91a77e841802c172ee3d5b6b162272bf0732cd8a54a7600d7250c5b788f2bbd4"},"sha256:a5a17edd1988763853752ee15d5bd4d6d2584a3393aec695e8d8b3862758a279":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":4608,"digest":"sha256:a5a17edd1988763853752ee15d5bd4d6d2584a3393aec695e8d8b3862758a279"},"sha256:dddd09bc1156c04e4431e15e7adc290e18581a04d7545030b8115be07d8ba8aa":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1237841408,"digest":"sha256:dddd09bc1156c04e4431e15e7adc290e18581a04d7545030b8115be07d8ba8aa"},"sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":80641024,"digest":"sha256:ea16cace89338c84eb6bcb91a7efdfcae6838fff359efe951858227436486c34"},"sha256:ea7f9e8dc4b32aa84e39e01475e6e4ad24f6d4c175388103a9e93d1b021b3b7d":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":187652608,"digest":"sha256:ea7f9e8dc4b32aa84e39e01475e6e4ad24f6d4c175388103a9e93d1b021b3b7d"},"sha256:f07d2440aabdc8a01341c7dbf2a12c605d15a367a972f3c3035401bedeb433c5":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":1763840,"digest":"sha256:f07d2440aabdc8a01341c7dbf2a12c605d15a367a972f3c3035401bedeb433c5"},"sha256:f0ef8bb8e669c94f302f8c3bb22f260910987963e3b64dd82d715f00df9ce9e5":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":23040,"digest":"sha256:f0ef8bb8e669c94f302f8c3bb22f260910987963e3b64dd82d715f00df9ce9e5"},"sha256:f64479d7869b23160b3944a14abfb4b9a40ae5e70fd1b58eba3319d3253f8853":{"mediaType":"application/vnd.oci.image.layer.v1.tar","size":136048640,"digest":"sha256:f64479d7869b23160b3944a14abfb4b9a40ae5e70fd1b58eba3319d3253f8853"}}}]

image/repositories

@@ -1 +1 @@
-{"tt-model/gr00t-n1.6-3b-p150":{"feb77048a2ed":"aa10b8db9ce55f80f832cc35d23d3776ad72b6dc6dd95bc6b8dd9f1a9cde331b"}}
+{"tt-model/gr00t-n1.6-3b-p150":{"0be2c39e4a01":"909e145eca13ac01c3f52611518262d779e8afdbf1d6f9330502727bcf54a418"}}

requirements.lock

@@ -12,7 +12,7 @@ distro==1.9.0
 elastic-transport==9.4.2
 elasticsearch==9.5.1
 fastapi==0.141.1
-filelock==3.32.6
+filelock==3.32.7
 fonttools==4.65.0
 fsspec==2026.7.0
 graphviz==0.21
@@ -20,9 +20,9 @@ h11==0.16.0
 hf-xet==1.6.0
 httpcore==1.0.9
 httpx==0.28.1
-huggingface_hub==1.31.0
+huggingface_hub==1.32.0
 identify==2.6.19
-idna==3.19
+idna==3.20
 Jinja2==3.1.6
 kiwisolver==1.5.1
 linkify-it-py==2.2.0
@@ -41,7 +41,7 @@ opencv-python-headless==4.8.1.78
 packaging==26.3
 pandas==3.0.5
 pillow==12.3.0
-platformdirs==4.11.8
+platformdirs==4.11.9
 pre_commit==4.6.2
 psutil==7.2.2
 pydantic==2.13.5
@@ -77,7 +77,7 @@ ttnn==0.65.2.dev10011
 typer==0.27.2
 typing-inspection==0.4.4
 typing_extensions==4.16.0
-urllib3==2.7.0
+urllib3==2.8.0
 uvicorn==0.53.0
-virtualenv==21.7.9
+virtualenv==21.7.11
 wheel==0.48.0

tt_kernel_manifest.json

@@ -6,7 +6,7 @@
   "device_count": 1,
   "producer": {
     "tt_kernel_version": "0.1.0",
-    "created_at": "2026-09-14T16:29:57.941499+00:00",
+    "created_at": "2026-09-17T20:45:45.250616+00:00",
     "hostname": "deepgadget"
   },
   "weights": {
@@ -35,8 +35,8 @@
     "image": {
       "registry": "hf",
       "repository": "gr00t-n1.6-3b-p150",
-      "tag": "tt-model/gr00t-n1.6-3b-p150:feb77048a2ed",
-      "digest": "sha256:feb77048a2ed1772037207da1e3fbf571555db9c9554509bad74076e7c94fa55"
+      "tag": "tt-model/gr00t-n1.6-3b-p150:0be2c39e4a01",
+      "digest": "sha256:0be2c39e4a01f307b061822a3a6d0b098e495a20c8ca0870953f0a73a05e4e60"
     },
     "kind": "tt-dit-server",
     "runtime": {
@@ -72,6 +72,9 @@
         "GR00T_LAYOUT": "gr1",
         "GR00T_POLICY": "mixed_dit",
         "GR00T_TRACE_LAYOUT": "per_stage",
+        "GR00T_DIT_BACKEND": "megakernel",
+        "GR00T_MK_ARENA_DTYPE": "auto",
+        "GR00T_NUM_CQS": "2",
         "GR00T_TOKENIZER_DIR": "/opt/tt-metal/gr00t_p150/assets/tokenizer/n16",
         "GR00T_PROJECT_ROOT": "/nonexistent",
         "GR00T_GOLDEN_ROOT": "/nonexistent",
@@ -103,6 +106,12 @@
       "import gr00t_p150.server.app as a; assert a.app",
       "import sys, gr00t_p150.server.app; assert 'ttnn' not in sys.modules and not any(m.startswith('models.experimental.gr00t') for m in sys.modules), 'the server module must import without side effects'",
       "import models.experimental.gr00t.tt.model as m; assert m.Gr00tTT",
+      "from models.experimental.gr00t.tt.policy import TTPolicy, DEFAULT_DIT_BACKEND; p = TTPolicy(); assert p.dit_backend == 'megakernel' == DEFAULT_DIT_BACKEND and p.mk_arena_dtype == 'auto' and p.mk_arena_dtype_for('n16') == 'bfp8_b', (p.dit_backend, p.mk_arena_dtype)",
+      "import models.experimental.gr00t.tt.megakernel.dit_program as d; from models.experimental.gr00t.tt.megakernel import arena, core_map, descriptors; assert d.DiTMegakernel and descriptors.KERNEL_SOURCE.endswith('tt/megakernel/kernels/dit_kernel.cpp')",
+      "import os; from pathlib import Path; from models.experimental.gr00t.tt.megakernel.descriptors import KERNEL_SOURCE; k = Path(os.environ.get('TT_METAL_HOME', '/opt/tt-metal')) / KERNEL_SOURCE; assert k.is_file() and len(list((k.parent / 'ops').glob('*.hpp'))) == 7, str(k)",
+      "import os, re; from pathlib import Path; from models.experimental.gr00t.tt.megakernel.descriptors import KERNEL_SOURCE; home = Path(os.environ.get('TT_METAL_HOME', '/opt/tt-metal')); k = home / KERNEL_SOURCE; incs = {m for f in [k, *(k.parent / 'ops').glob('*.hpp')] for m in re.findall(r'#include \"([^\"]*unified_kernels/[^\"]+)\"', f.read_text())}; assert incs, 'no unified_kernels includes found'; missing = [i for i in incs if not (k.parent / i).resolve().is_file() and not (home / i).is_file()]; assert not missing, missing",
+      "from models.experimental.gr00t.tt.model import open_model_device, device_worker_l1_size, worker_l1_size_for, MK_L1_CUT_BYTES_DEFAULT; from models.experimental.gr00t.tt.policy import TTPolicy; assert MK_L1_CUT_BYTES_DEFAULT == 65536 and worker_l1_size_for(TTPolicy()) is not None and worker_l1_size_for(TTPolicy(dit_backend='ttnn')) is None",
+      "from gr00t_p150.server.app import STAGE_SERVED, DEFAULT_NUM_CQS, DIT_BACKENDS, MK_ARENA_DTYPES; assert STAGE_SERVED['megakernel'] == 'stage2-megakernel' and STAGE_SERVED['ttnn'] == 'stage1-ttnn-traces' and DEFAULT_NUM_CQS == 2 and set(DIT_BACKENDS) == {'ttnn', 'megakernel'} and 'bfp8_b' in MK_ARENA_DTYPES",
       "from models.experimental.gr00t.common import configs, preprocessing, normalization, prompts, checkpoint, weights; assert configs.HF_SNAPSHOT_SHAS['n16'] == 'd0814e7ecb19202e7c8468b46098b0b7ef3a6d61' and configs.HF_REPOS['n16'] == 'nvidia/GR00T-N1.6-3B'",
       "import models.tt_dit.utils.matmul as mm; assert mm.get_matmul_config",
       "import models.demos.deepseek_v3_b1.unified_kernel_descriptor as u; assert u.UnifiedKernelDescriptor",
@@ -114,14 +123,14 @@
       "from gr00t_p150.server.app import load_demo_request; import numpy as np; r = load_demo_request('n16'); assert set(r['images']) == {'ego_view_bg_crop_pad_res256_freq20'} and np.asarray(r['noise']).squeeze().shape == (50, 128)",
       "from pathlib import Path; assert Path('/opt/tt-metal/models/experimental/gr00t/common/golden_tap_map.json').is_file() and Path('/opt/tt-metal/models/experimental/gr00t/tests/tt/gates/gates_multi.json').is_file()",
       "from pathlib import Path; assert Path('/opt/tt-metal/gr00t_p150/server/smoke_test.py').is_file() and Path('/opt/tt-metal/scripts/bench_http.py').is_file() and Path('/opt/tt-metal/scripts/download_weights.sh').is_file()",
-      "import ttnn; assert all(hasattr(ttnn, a) for a in ('begin_trace_capture', 'end_trace_capture', 'execute_trace', 'release_trace', 'copy_host_to_device_tensor', 'MinimalMatmulConfig'))",
+      "import ttnn; assert all(hasattr(ttnn, a) for a in ('begin_trace_capture', 'end_trace_capture', 'execute_trace', 'release_trace', 'copy_host_to_device_tensor', 'MinimalMatmulConfig', 'generic_op', 'ProgramDescriptor'))",
       "from gr00t_p150.server.app import parse_mesh_shape as p; assert p('1x1') == p('(1, 1)') == p('1,1') == (1, 1)"
     ],
     "built": {
-      "image": "tt-model/gr00t-n1.6-3b-p150:feb77048a2ed",
+      "image": "tt-model/gr00t-n1.6-3b-p150:0be2c39e4a01",
       "repo": "changh95/GR00T-N1.6-3B-p150",
       "tt_model_version": "0.1.0",
-      "created_at": "2026-09-14T16:25:52+00:00",
+      "created_at": "2026-09-17T20:40:28+00:00",
       "tt_metal": {
         "sha": "668c290757550588d0ce46b180c344a462a2aaf5",
         "describe": "v0.79.0-dev20260914-1-g668c290757",
@@ -132,8 +141,8 @@
         "branch": "HEAD",
         "pushed": true
       },
-      "code_sha256": "7a456e1065da7159a7c2bb83c70e115542e43c5a655f5fe8c14e310f446da0d5",
-      "image_digest": "sha256:feb77048a2ed1772037207da1e3fbf571555db9c9554509bad74076e7c94fa55"
+      "code_sha256": "8e6bd13a23654f91a229745d001c12b7972f9215a95c75fe05258181f7c324b8",
+      "image_digest": "sha256:0be2c39e4a01f307b061822a3a6d0b098e495a20c8ca0870953f0a73a05e4e60"
     }
   }
 }