Add files using upload-large-folder tool

code/models/experimental/gr00t/benchmarks/bench_denoise.py

@@ -225,7 +225,11 @@ def table_estimate_us(cfg: Any, matmul_rows: Sequence[Mapping[str, Any]]) -> Dic
 
 def analyse(version: str, policy: str, layout_name: Optional[str], embodiment: Optional[str]) -> Dict[str, Any]:
     """CPU analysis shared by the dry run and the device run."""
-    setup = E2E.resolve_setup(version, policy=policy, layout_name=layout_name, embodiment=embodiment)
+    # this benchmark measures the Stage-1 (ttnn) trace B -- its streamed bytes and its replay time; the megakernel
+    # backend (the TTPolicy default since 2026-09-18) has no separate denoise op sequence to count
+    setup = E2E.resolve_setup(
+        version, policy=policy, layout_name=layout_name, embodiment=embodiment, dit_backend="ttnn"
+    )
     cfg, layout = setup["cfg"], setup["layout"]
     n_steps = cfg.sampler.num_inference_timesteps
     entries = trace_b_entries(version, setup["embodiment_id"], setup["policy"])
@@ -382,7 +386,12 @@ def run(
         raise ValueError("warmup must be >= 0 and reps > 0")
     an = analyse(version, policy, layout_name, embodiment)
     setup = E2E.resolve_setup(
-        version, policy=policy, trace_layout=trace_layout, layout_name=layout_name, embodiment=embodiment
+        version,
+        policy=policy,
+        trace_layout=trace_layout,
+        layout_name=layout_name,
+        embodiment=embodiment,
+        dit_backend="ttnn",  # the Stage-1 denoise trace (see analyse); the device is the default-L1 one
     )
     ok, reason = E2E.golden_available(version, sample)
     if not ok:

code/models/experimental/gr00t/benchmarks/bench_e2e.py

@@ -20,7 +20,12 @@ copy ms) is reported as ``upload/<input>`` rows and ``upload/total`` (docs/plan/
 
 ``--num-cqs`` (default 2) opens the device with that many command queues; ``--cq1-uploads auto|on|off`` and
 ``--embed-gather device|host`` are forwarded to ``Gr00tTT``; ``--profile-encode`` adds the per-step host-encode
-profile (``meta.encode_profile``; also in ``--dry-run``).
+profile (``meta.encode_profile``; also in ``--dry-run``).  ``--dit-backend ttnn|megakernel`` (WP-K5) selects the
+denoise backend (``TTPolicy.dit_backend``; default = the policy default, the **megakernel** since 2026-09-18; the
+device is opened through ``tt.model.open_model_device`` so the megakernel gets its ``worker_l1_size``) and
+``--mk-arena-dtype auto|bf16|bfp8_b`` the megakernel's weight arena; ``meta.dit_backend`` / ``meta.mk_arena_dtype`` /
+``meta.megakernel`` / ``meta.weights`` (tensors and bytes ``TTWeights`` uploaded, what the megakernel skipped) record
+what ran.  ``--load-ttnn-dit-weights`` keeps the Stage-1 DiT upload under the megakernel (startup A/B).
 
 The model comes from ``tt.model.Gr00tTT`` (WP-14).  Until that module exists the import is guarded: ``--dry-run``
 validates ``--version`` / ``--policy`` / ``--trace-layout`` / ``--layout`` / ``--embodiment`` against ``common.configs``
@@ -61,6 +66,7 @@ from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
 from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
 from models.experimental.gr00t.benchmarks import mb1_common as MB
 from models.experimental.gr00t.tt import device as D
+from models.experimental.gr00t.tt.policy import DEFAULT_DIT_BACKEND, DIT_BACKENDS, MK_ARENA_DTYPES  # torch-free
 
 BENCH_NAME = "bench_e2e"
 VERSIONS = SO.VERSIONS
@@ -88,6 +94,8 @@ def resolve_setup(
     trace_layout: str = "per_stage",
     layout_name: Optional[str] = None,
     embodiment: Optional[str] = None,
+    dit_backend: str = DEFAULT_DIT_BACKEND,
+    mk_arena_dtype: str = "auto",
 ) -> Dict[str, Any]:
     """Validate the CLI choices against ``common.configs`` / ``tt.policy`` and return the resolved facts
     (``cfg``, ``layout``, ``policy`` object, ``embodiment_tag`` / ``embodiment_id``, expected trace names)."""
@@ -102,7 +110,9 @@ def resolve_setup(
     layout = cfg.canonical_layout if layout_name is None else cfg.layout(layout_name)
     tag = embodiment if embodiment is not None else layout.embodiment_tag
     emb_id = cfg.embodiment_id(tag)  # raises on an unknown tag
-    pol = TTPolicy(dtype_policy=policy, trace_layout=trace_layout)  # raises on unknown values
+    pol = TTPolicy(
+        dtype_policy=policy, trace_layout=trace_layout, dit_backend=dit_backend, mk_arena_dtype=mk_arena_dtype
+    )  # raises on unknown values
     return {
         "cfg": cfg,
         "layout": layout,
@@ -121,6 +131,8 @@ def resolve_setup(
             "dit_m_pad": layout.dit_m_pad,
             "noise_shape": list(layout.noise_shape),
             "policy": pol.to_dict(),
+            "dit_backend": pol.dit_backend,
+            "mk_arena_dtype": pol.mk_arena_dtype_for(version) if pol.dit_backend == "megakernel" else None,
             "expected_stage_names": list(STAGE_NAMES[trace_layout]),
         },
     }
@@ -190,6 +202,7 @@ def load_model(
     *,
     embed_gather: str = "device",
     cq1_uploads: Optional[bool] = None,
+    load_ttnn_dit_weights: Optional[bool] = None,
 ) -> Any:
     Gr00tTT = import_gr00t_tt()
     if Gr00tTT is None:
@@ -207,6 +220,8 @@ def load_model(
         embed_gather=embed_gather,
         cq1_uploads=cq1_uploads,
     )
+    if load_ttnn_dit_weights is not None:
+        kw["load_ttnn_dit_weights"] = bool(load_ttnn_dit_weights)
     if cache_root is not None:
         kw["cache_root"] = cache_root
     return Gr00tTT.from_pretrained(setup["cfg"].version, **kw)
@@ -235,6 +250,15 @@ def resolve_stage_replayer(model: Any, expected_stages: Sequence[str]) -> Any:
     return tracer
 
 
+def _worker_l1_size(device: Any) -> Optional[int]:
+    """``tt.model.device_worker_l1_size`` (the allocatable L1 the device was opened with; recorded in ``meta``)."""
+    try:
+        from models.experimental.gr00t.tt.model import device_worker_l1_size
+    except ImportError:
+        return None
+    return int(device_worker_l1_size(device))
+
+
 def denoise_stage_name(stage_names: Sequence[str]) -> str:
     """The trace that holds the 4-step denoise: the last one of either layout (``denoise`` / ``B``)."""
     if not stage_names:
@@ -406,6 +430,9 @@ def run(
     embed_gather: str = "device",
     cq1_uploads: Optional[bool] = None,
     profile_encode: bool = False,
+    dit_backend: str = DEFAULT_DIT_BACKEND,
+    mk_arena_dtype: str = "auto",
+    load_ttnn_dit_weights: Optional[bool] = None,
 ) -> Dict[str, Any]:
     """Measure the traced end-to-end latency and its split on an open device; returns the perf payload (not written).
     Main pass = production path (``tracer.call``), split pass = serialised writes + per-trace sync (module docstring).
@@ -418,7 +445,13 @@ def run(
     if warmup < 0 or reps <= 0:
         raise ValueError("warmup must be >= 0 and reps > 0")
     setup = resolve_setup(
-        version, policy=policy, trace_layout=trace_layout, layout_name=layout_name, embodiment=embodiment
+        version,
+        policy=policy,
+        trace_layout=trace_layout,
+        layout_name=layout_name,
+        embodiment=embodiment,
+        dit_backend=dit_backend,
+        mk_arena_dtype=mk_arena_dtype,
     )
     ok, reason = golden_available(version, sample)
     if not ok:
@@ -427,7 +460,14 @@ def run(
     obs, mi = inputs["obs"], inputs["mi"]
     noise = make_noise(setup["layout"], noise_seed)
     t_load = time.perf_counter()
-    model = load_model(setup, device, cache_root, embed_gather=embed_gather, cq1_uploads=cq1_uploads)
+    model = load_model(
+        setup,
+        device,
+        cache_root,
+        embed_gather=embed_gather,
+        cq1_uploads=cq1_uploads,
+        load_ttnn_dit_weights=load_ttnn_dit_weights,
+    )
     load_s = time.perf_counter() - t_load
     tracer = resolve_stage_replayer(model, setup["stage_names"])
     missing = [m for m in PRODUCTION_MEMBERS if not hasattr(tracer, m)]
@@ -510,6 +550,19 @@ def run(
         per_call_inputs=[s.describe() for s in executor.plan().inputs],
         per_call_bytes=ex_desc.get("plan", {}).get("per_call_bytes"),
         vit_fork=getattr(getattr(model.backbone, "tower", None), "fork", None),
+        dit_backend=model.policy.dit_backend,
+        mk_arena_dtype=model.recipe.mk_arena_dtype,
+        megakernel=model.head.describe().get("megakernel"),
+        worker_l1_size=_worker_l1_size(device),
+        head_timing={k: v for k, v in model.head.timing.items()},
+        model_timing={k: v for k, v in model.timing.items() if k.endswith("_s")},
+        weights={
+            "n_tensors": len(model.W),
+            "device_bytes": float(model.W.device_bytes),
+            "load_path": model.W.load_stats.get("path"),
+            "load_s": model.timing.get("weights_load_s"),
+            "skipped": dict(getattr(model, "weights_skipped", {})),
+        },
     )
     if profile_encode:
         meta["encode_profile"] = encode_profile(version, inputs, setup["layout"])
@@ -530,6 +583,9 @@ def run(
             "noise_seed": noise_seed,
             "embed_gather": embed_gather,
             "cq1_uploads": cq1_uploads,
+            "dit_backend": dit_backend,
+            "mk_arena_dtype": mk_arena_dtype,
+            "load_ttnn_dit_weights": load_ttnn_dit_weights,
         },
         method={
             "main_pass": f"{warmup} warm-up calls, then {reps} calls: encode -> tracer.call (TraceExecutor.replay: "
@@ -586,10 +642,18 @@ def dry_run(
     host_reps: int = DEFAULT_HOST_REPS,
     verbose: bool = True,
     profile_encode: bool = False,
+    dit_backend: str = DEFAULT_DIT_BACKEND,
+    mk_arena_dtype: str = "auto",
 ) -> Dict[str, Any]:
     """Validate the setup and time the host side (encode / decode) when the golden data is present; no device."""
     setup = resolve_setup(
-        version, policy=policy, trace_layout=trace_layout, layout_name=layout_name, embodiment=embodiment
+        version,
+        policy=policy,
+        trace_layout=trace_layout,
+        layout_name=layout_name,
+        embodiment=embodiment,
+        dit_backend=dit_backend,
+        mk_arena_dtype=mk_arena_dtype,
     )
     stages: Dict[str, Dict[str, Any]] = {}
     meta = dict(setup["meta"])
@@ -621,6 +685,8 @@ def dry_run(
             "embodiment": setup["embodiment_tag"],
             "layout": setup["layout"].name,
             "host_reps": host_reps,
+            "dit_backend": dit_backend,
+            "mk_arena_dtype": mk_arena_dtype,
         },
         method={"host": f"{host_reps} host calls after 1 warm-up (encode / decode), median / p90 ms; no device"},
         stages=stages,
@@ -655,13 +721,15 @@ def print_summary(payload: Mapping[str, Any]) -> None:
     )
     m = payload["meta"]
     print(
-        "   layout {l} | embodiment {e} | S_pad {s} | shape_key {sk} | policy {p} | stages {st} | embed_gather {g} | "
-        "cq1_uploads {c} ({n} CQ){skip}".format(
+        "   layout {l} | embodiment {e} | S_pad {s} | shape_key {sk} | policy {p} | dit_backend {b} (arena {a}) | "
+        "stages {st} | embed_gather {g} | cq1_uploads {c} ({n} CQ){skip}".format(
             l=m.get("layout_name"),
             e=m.get("embodiment_tag"),
             s=m.get("s_pad_llm"),
             sk=m.get("shape_key", "-"),
             p=(m.get("policy") or {}).get("dtype_policy"),
+            b=m.get("dit_backend", (m.get("policy") or {}).get("dit_backend")),
+            a=m.get("mk_arena_dtype", "-"),
             st=m.get("stage_names", m.get("expected_stage_names")),
             g=m.get("embed_gather", "-"),
             c=m.get("cq1_uploads", "-"),
@@ -680,6 +748,24 @@ def build_parser() -> argparse.ArgumentParser:
     )
     ap.add_argument("--version", required=True, choices=VERSIONS)
     ap.add_argument("--policy", default="mixed_dit", choices=POLICIES, help="TTPolicy.dtype_policy")
+    ap.add_argument(
+        "--dit-backend",
+        default=DEFAULT_DIT_BACKEND,
+        choices=DIT_BACKENDS,
+        help=f"TTPolicy.dit_backend: Stage-1 ttnn ops | megakernel (default: the policy default, {DEFAULT_DIT_BACKEND})",
+    )
+    ap.add_argument(
+        "--load-ttnn-dit-weights",
+        action="store_true",
+        help="keep the Stage-1 DiT / precompute / encoder upload under the megakernel backend (startup A/B; "
+        "default: the megakernel skips them, tt.model.megakernel_weight_names)",
+    )
+    ap.add_argument(
+        "--mk-arena-dtype",
+        default="auto",
+        choices=MK_ARENA_DTYPES,
+        help="TTPolicy.mk_arena_dtype: the megakernel's weight arena (auto = the per-version default)",
+    )
     ap.add_argument(
         "--trace-layout", default="per_stage", choices=TRACE_LAYOUTS, help="TTPolicy.trace_layout (plan §4)"
     )
@@ -733,13 +819,24 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             layout_name=args.layout,
             host_reps=args.host_reps,
             profile_encode=args.profile_encode,
+            dit_backend=args.dit_backend,
+            mk_arena_dtype=args.mk_arena_dtype,
         )
         SO.write_perf_json(payload, Path(args.out))
         return 0
     if import_gr00t_tt() is None:
         raise ImportError("tt.model.Gr00tTT is not importable (WP-14 pending) -- run with --dry-run")
     SO.ensure_tt_metal_cache()
-    device = D.open_gr00t_device(
+    from models.experimental.gr00t.tt import model as M
+    from models.experimental.gr00t.tt.policy import TTPolicy
+
+    device = M.open_model_device(
+        TTPolicy(
+            dtype_policy=args.policy,
+            trace_layout=args.trace_layout,
+            dit_backend=args.dit_backend,
+            mk_arena_dtype=args.mk_arena_dtype,
+        ),
         trace_region_size=args.trace_region_mb << 20,
         l1_small_size=D.DEFAULT_L1_SMALL_SIZE,
         num_command_queues=args.num_cqs,
@@ -761,6 +858,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             embed_gather=args.embed_gather,
             cq1_uploads=resolve_cq1(args.cq1_uploads),
             profile_encode=args.profile_encode,
+            dit_backend=args.dit_backend,
+            mk_arena_dtype=args.mk_arena_dtype,
+            load_ttnn_dit_weights=True if args.load_ttnn_dit_weights else None,
         )
     finally:
         D.close_gr00t_device(device)
@@ -772,6 +872,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
 __all__ = [
     "BENCH_NAME",
     "CQ1_MODES",
+    "DIT_BACKENDS",
+    "MK_ARENA_DTYPES",
     "EMBED_GATHERS",
     "POLICIES",
     "PRODUCTION_MEMBERS",

code/models/experimental/gr00t/benchmarks/bench_load.py

@@ -0,0 +1,236 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Model start-up cost of ``Gr00tTT.from_pretrained`` per DiT backend (megakernel-default pass; WP-K5 review NB 7 /
+§8 pre-condition 5): wall time of the weight upload and the module builds, the tensors and bytes ``TTWeights`` puts
+on the device, and what the megakernel backend skips (``tt.model.megakernel_weight_names``).
+
+Per invocation the device is opened once through ``tt.model.open_model_device(policy)`` and the model is built
+``--rounds`` times for each configuration -- ``skip`` (the default: the megakernel packs the DiT / precompute / encoder
+tensors into its arena and ``TTWeights`` leaves them out) and ``keep`` (``load_ttnn_dit_weights=True``: the full
+Stage-1 upload next to the arena, what WP-K5 measured) -- alternating ``keep, skip, keep, skip, ...`` so both see the
+same OS page-cache state of the ``.tensorbin`` tier; every model is released before the next build.  With
+``--dit-backend ttnn`` only the ``keep`` configuration exists (the ttnn head reads those weights).
+
+Output ``benchmarks/results/bench_load_<version>_<ts>.json`` (schema ``gr00t-tt-perf/1``): stages
+``from_pretrained/<cfg>`` (median ms over the rounds), ``weights_load/<cfg>``, ``head_build/<cfg>`` (the megakernel
+arena plan + pack + upload sit inside it); ``meta.configs`` holds the per-round numbers, the tensor counts, the bytes
+on the device and the skipped set, ``meta.device_mb`` the MB of plan tensors per configuration and ``meta.delta`` the
+keep - skip savings.
+
+Run::
+
+    DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
+        python -m models.experimental.gr00t.benchmarks.bench_load --version n16 [--rounds 2] [--dit-backend ttnn]
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence
+
+from models.experimental.gr00t.benchmarks import bench_e2e as E2E
+from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
+from models.experimental.gr00t.tt import device as D
+from models.experimental.gr00t.tt.policy import DEFAULT_DIT_BACKEND, DIT_BACKENDS, MK_ARENA_DTYPES
+
+BENCH_NAME = "bench_load"
+CONFIGS: Dict[str, Optional[bool]] = {"keep": True, "skip": None}  # -> from_pretrained(load_ttnn_dit_weights=)
+
+
+def configs_for(dit_backend: str) -> Dict[str, Optional[bool]]:
+    """The configurations that exist for a backend: both for the megakernel, ``keep`` only for ttnn."""
+    if dit_backend == "megakernel":
+        return dict(CONFIGS)
+    return {"keep": None}
+
+
+def one_build(
+    setup: Dict[str, Any], device: Any, cache_root: Optional[Path], load_ttnn: Optional[bool]
+) -> Dict[str, Any]:
+    """Build the model once, record the timings / weight facts, release it."""
+    t0 = time.perf_counter()
+    model = E2E.load_model(setup, device, cache_root, load_ttnn_dit_weights=load_ttnn)
+    total = time.perf_counter() - t0
+    try:
+        timing = {k: float(v) for k, v in model.timing.items() if k.endswith("_s")}
+        head = {k: float(v) for k, v in model.head.timing.items()}
+        rec = {
+            "from_pretrained_s": total,
+            "weights_load_s": timing.get("weights_load_s"),
+            "backbone_build_s": timing.get("backbone_build_s"),
+            "adapter_build_s": timing.get("adapter_build_s"),
+            "embed_gather_build_s": timing.get("embed_gather_build_s"),
+            "executor_and_head_build_s": timing.get("executor_and_head_build_s"),
+            "head_timing": head,
+            "n_tensors": len(model.W),
+            "device_bytes": float(model.W.device_bytes),
+            "load_path": model.W.load_stats.get("path"),
+            "n_warm": model.W.load_stats.get("n_warm"),
+            "n_cold": model.W.load_stats.get("n_cold"),
+            "skipped": dict(model.weights_skipped),
+            "dit_backend": model.policy.dit_backend,
+            "mk_arena_dtype": model.recipe.mk_arena_dtype,
+            "adapter_mem": model.adapter_mem,
+        }
+    finally:
+        model.release()
+    import ttnn
+
+    ttnn.synchronize_device(device)
+    return rec
+
+
+def run(
+    device: Any,
+    *,
+    version: str,
+    policy: str = "mixed_dit",
+    trace_layout: str = "per_stage",
+    dit_backend: str = DEFAULT_DIT_BACKEND,
+    mk_arena_dtype: str = "auto",
+    rounds: int = 2,
+    cache_root: Optional[Path] = None,
+    verbose: bool = True,
+) -> Dict[str, Any]:
+    if rounds <= 0:
+        raise ValueError("rounds must be positive")
+    setup = E2E.resolve_setup(
+        version, policy=policy, trace_layout=trace_layout, dit_backend=dit_backend, mk_arena_dtype=mk_arena_dtype
+    )
+    cfgs = configs_for(dit_backend)
+    per_cfg: Dict[str, List[Dict[str, Any]]] = {name: [] for name in cfgs}
+    for r in range(rounds):
+        for name, load_ttnn in cfgs.items():
+            rec = one_build(setup, device, cache_root, load_ttnn)
+            rec["round"] = r
+            per_cfg[name].append(rec)
+            if verbose:
+                print(
+                    f"[bench_load] {version}/{dit_backend}/{name} round {r}: from_pretrained {rec['from_pretrained_s']:.2f} s "
+                    f"(weights {rec['weights_load_s']:.2f} s, {rec['n_tensors']} tensors, {rec['device_bytes'] / 1e6:.0f} MB, "
+                    f"{rec['load_path']}; head {rec['executor_and_head_build_s']:.2f} s"
+                    f"{'; skipped ' + str(rec['skipped'].get('n_tensors')) + ' tensors / ' + format(rec['skipped'].get('device_bytes', 0) / 1e6, '.0f') + ' MB' if rec['skipped'] else ''})"
+                )
+    stages: Dict[str, Dict[str, Any]] = {}
+    for name, recs in per_cfg.items():
+
+        def med(key: str) -> float:
+            vals = [float(x[key]) for x in recs if x.get(key) is not None]
+            return statistics.median(vals) if vals else float("nan")
+
+        stages[f"from_pretrained/{name}"] = {
+            "metric": "ms",
+            "value": med("from_pretrained_s") * 1e3,
+            "min": min(float(x["from_pretrained_s"]) for x in recs) * 1e3,
+            "per_round": [round(float(x["from_pretrained_s"]) * 1e3, 1) for x in recs],
+        }
+        stages[f"weights_load/{name}"] = {"metric": "ms", "value": med("weights_load_s") * 1e3}
+        stages[f"head_build/{name}"] = {
+            "metric": "ms",
+            "value": med("executor_and_head_build_s") * 1e3,
+            "note": "executor + ActionHeadTT build (megakernel: arena plan + pack + upload + io inside)",
+        }
+    facts = D.device_facts(device)
+    facts["worker_l1_size"] = E2E._worker_l1_size(device)
+    meta = dict(
+        setup["meta"],
+        configs={name: recs for name, recs in per_cfg.items()},
+        device_mb={
+            name: {"mb": recs[0]["device_bytes"] / 1e6, "n_tensors": recs[0]["n_tensors"]}
+            for name, recs in per_cfg.items()
+        },
+        rounds=rounds,
+        worker_l1_size=facts["worker_l1_size"],
+    )
+    if "skip" in per_cfg and "keep" in per_cfg:
+        k, s = per_cfg["keep"][0], per_cfg["skip"][0]
+        meta["delta"] = {
+            "device_mb_saved": (k["device_bytes"] - s["device_bytes"]) / 1e6,
+            "n_tensors_saved": k["n_tensors"] - s["n_tensors"],
+            "from_pretrained_ms_saved_median": stages["from_pretrained/keep"]["value"]
+            - stages["from_pretrained/skip"]["value"],
+            "weights_load_ms_saved_median": stages["weights_load/keep"]["value"] - stages["weights_load/skip"]["value"],
+        }
+    payload = SO.perf_payload(
+        BENCH_NAME,
+        version,
+        dry_run=False,
+        args={
+            "version": version,
+            "policy": policy,
+            "trace_layout": trace_layout,
+            "dit_backend": dit_backend,
+            "mk_arena_dtype": mk_arena_dtype,
+            "rounds": rounds,
+        },
+        method={
+            "builds": f"{rounds} round(s) of alternating keep / skip Gr00tTT.from_pretrained on one open device, "
+            "each model released before the next build; medians over the rounds",
+            "clock": "host wall clock (time.perf_counter)",
+        },
+        stages=stages,
+        rows=[],
+        device=facts,
+        meta=meta,
+    )
+    return payload
+
+
+def build_parser() -> argparse.ArgumentParser:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    ap.add_argument("--version", required=True, choices=E2E.VERSIONS)
+    ap.add_argument("--policy", default="mixed_dit", choices=E2E.POLICIES)
+    ap.add_argument("--trace-layout", default="per_stage", choices=E2E.TRACE_LAYOUTS)
+    ap.add_argument("--dit-backend", default=DEFAULT_DIT_BACKEND, choices=DIT_BACKENDS)
+    ap.add_argument("--mk-arena-dtype", default="auto", choices=MK_ARENA_DTYPES)
+    ap.add_argument("--rounds", type=int, default=2)
+    ap.add_argument("--cache-root", default=None)
+    ap.add_argument("--out", default=str(SO.RESULTS_DIR))
+    ap.add_argument("--trace-region-mb", type=int, default=64)
+    ap.add_argument("--num-cqs", type=int, default=2, choices=(1, 2))
+    return ap
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    args = build_parser().parse_args(argv)
+    if E2E.import_gr00t_tt() is None:
+        raise ImportError("tt.model.Gr00tTT is not importable")
+    SO.ensure_tt_metal_cache()
+    from models.experimental.gr00t.tt import model as M
+    from models.experimental.gr00t.tt.policy import TTPolicy
+
+    device = M.open_model_device(
+        TTPolicy(
+            dtype_policy=args.policy,
+            trace_layout=args.trace_layout,
+            dit_backend=args.dit_backend,
+            mk_arena_dtype=args.mk_arena_dtype,
+        ),
+        trace_region_size=args.trace_region_mb << 20,
+        l1_small_size=D.DEFAULT_L1_SMALL_SIZE,
+        num_command_queues=args.num_cqs,
+    )
+    try:
+        payload = run(
+            device,
+            version=args.version,
+            policy=args.policy,
+            trace_layout=args.trace_layout,
+            dit_backend=args.dit_backend,
+            mk_arena_dtype=args.mk_arena_dtype,
+            rounds=args.rounds,
+            cache_root=Path(args.cache_root) if args.cache_root else None,
+        )
+    finally:
+        D.close_gr00t_device(device)
+    path = SO.write_perf_json(payload, Path(args.out))
+    print(f"[bench_load] wrote {path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/models/experimental/gr00t/benchmarks/bench_mk_step.py

@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K4 timing / zone bench of the DiT megakernel (IMPLEMENTATION_PLAN.md §5.7 gate G2, §5.4 timing chain).
+
+Modes::
+
+    # G2: the steps4 program captured in a trace, N replays -> ms per step vs 1.3 x arena bytes / measured GB/s
+    bench_mk_step.py --version n16 --dtype bf16 --replays 20
+    # per-phase DeviceZoneScopedN zones of the block variant over a few blocks (TT_METAL_DEVICE_PROFILER=1 is set
+    # before ttnn is imported; the profiler cannot be combined with the watcher on Blackhole)
+    bench_mk_step.py --version n16 --dtype bf16 --profile --blocks 4
+
+Every run writes ``benchmarks/results/bench_mk_step_<version>_<stamp>.json`` (schema ``gr00t-mk-k4/1``).  The
+arena is packed from the real plan tensors (``dit_program.real_arena_plan``), the K/V are hoisted on device from the
+canonical golden sample and the masks are the call's (``tests/tt/test_mk_dit_block.K4Session``).  Run through
+``bin/with-device.sh`` from the tt-metal root; ``ttnn`` is imported lazily in :func:`main`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import os
+import statistics
+import sys
+import time
+from collections import defaultdict
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import torch
+
+TT_METAL_HOME = Path(os.environ.get("TT_METAL_HOME", str(Path(__file__).resolve().parents[4])))
+PROFILE_CSV = TT_METAL_HOME / "generated" / "profiler" / ".logs" / "profile_log_device.csv"
+RESULTS_DIR = Path(__file__).resolve().parent / "results"
+SCHEMA = "gr00t-mk-k4/1"
+MEASURED_GBPS: Dict[str, float] = {"bf16": 464.0, "bfp8_b": 414.0}  # mk_k1_summary.md §3 (direct mode)
+G2_FACTOR = 1.3
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--version", default="n16", choices=("n15", "n16", "n17"))
+    ap.add_argument("--dtype", default="bf16", choices=("bf16", "bfp8_b"), help="arena stream dtype")
+    ap.add_argument("--replays", type=int, default=20, help="traced steps4 replays for the G2 timing")
+    ap.add_argument("--profile", action="store_true", help="collect the per-phase zones of a short block launch")
+    ap.add_argument(
+        "--profile-variant",
+        default="block",
+        choices=("block", "step"),
+        help="profiled launch: `block` = blocks [0, --blocks) of the block variant; `step` = one whole denoising step "
+        "(all blocks + output head; set GR00T_MK_ZONE_ROUNDS=0 so the hub BRISC's 250 profiler markers hold every block)",
+    )
+    ap.add_argument("--blocks", type=int, default=4, help="blocks of the profiled block launch")
+    ap.add_argument(
+        "--first-block",
+        type=int,
+        default=0,
+        help="first block of the profiled block launch (the launch's first block streams its weights cold: profile "
+        "from block 1 or 2 to see the steady-state prefetch)",
+    )
+    ap.add_argument("--step", type=int, default=0, help="denoising step of the profiled block launch")
+    ap.add_argument(
+        "--cb-w-slots", type=int, default=None, help="weight-ring depth in K-block slots (K4Budget default)"
+    )
+    ap.add_argument(
+        "--k5-device",
+        action="store_true",
+        help="open the device as the K5 model does (tt.model.worker_l1_size_for(TTPolicy(dit_backend='megakernel')), the "
+        "64 KiB worker-L1 cut) instead of the K4 test device (96 KiB): the compute-core binaries must fit K5's "
+        "kernel-config ring buffer (mk_k4b_summary.md section 5) -- a program too large fails at the first launch",
+    )
+    ap.add_argument("--results-dir", default=str(RESULTS_DIR))
+    return ap.parse_args(argv)
+
+
+def read_zones(path: Path = PROFILE_CSV) -> Dict[str, Any]:
+    """``{"freq_mhz", "zones": {"<RISC>:<zone>": {count, median_us, mean_us, min_us, max_us, cores}}}`` over the
+    ``mk_*`` zones of every launch in the profiler CSV (nested zones handled with a per-(run, core, RISC) stack)."""
+    if not path.is_file():
+        raise FileNotFoundError(f"{path} not found (set TT_METAL_DEVICE_PROFILER=1 before importing ttnn)")
+    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)}
+        stacks: Dict[tuple, List[Tuple[str, int]]] = defaultdict(list)
+        durations: Dict[str, List[int]] = defaultdict(list)
+        cores: Dict[str, set] = defaultdict(set)
+        for row in reader:
+            if len(row) < len(cols) - 1:
+                continue
+            zone = row[idx["zone name"]].strip()
+            if not zone.startswith("mk_"):
+                continue
+            core = (
+                row[idx["run host ID"]].strip(),
+                row[idx["core_x"]],
+                row[idx["core_y"]],
+                row[idx["RISC processor type"]].strip(),
+            )
+            t = int(row[idx["time[cycles since reset]"]])
+            typ = row[idx["type"]].strip()
+            if typ == "ZONE_START":
+                stacks[core].append((zone, t))
+            elif typ == "ZONE_END":
+                st = stacks[core]
+                for j in range(len(st) - 1, -1, -1):
+                    if st[j][0] == zone:
+                        name = f"{core[3]}:{zone}"
+                        durations[name].append(t - st[j][1])
+                        cores[name].add(core[1:3])
+                        del st[j]
+                        break
+    zones = {}
+    for name, ds in durations.items():
+        us = [d / freq_mhz for d in ds]
+        zones[name] = {
+            "count": len(us),
+            "cores": len(cores[name]),
+            "median_us": statistics.median(us),
+            "mean_us": sum(us) / len(us),
+            "min_us": min(us),
+            "max_us": max(us),
+        }
+    return {"freq_mhz": freq_mhz, "zones": zones}
+
+
+def g2_record(per_step_ms: float, bytes_per_step: int, dtype: str) -> Dict[str, Any]:
+    floor_ms = bytes_per_step / MEASURED_GBPS[dtype] / 1e6
+    return {
+        "per_step_ms": per_step_ms,
+        "bytes_per_step": bytes_per_step,
+        "measured_gbps": MEASURED_GBPS[dtype],
+        "floor_ms_per_step": floor_ms,
+        "gate_ms_per_step": G2_FACTOR * floor_ms,
+        "ratio_to_floor": per_step_ms / floor_ms,
+        "effective_gbps": bytes_per_step / per_step_ms / 1e6,
+        "pass": per_step_ms <= G2_FACTOR * floor_ms,
+    }
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    args = parse_args(argv)
+    if args.profile:
+        os.environ["TT_METAL_DEVICE_PROFILER"] = "1"
+        # the device CSV of the last run and the profiler's accumulated zone-source-location logs: stale entries of
+        # earlier builds of the same zone names at other line numbers can collide in its 16-bit hash (profiler.cpp
+        # `populateZoneSrcLocations` throws at device close) -- they are rebuilt from this run's compile logs
+        for stale in (PROFILE_CSV, *PROFILE_CSV.parent.glob("*zone_src_locations.log")):
+            if stale.exists():
+                stale.unlink()
+    os.environ.setdefault("TT_METAL_CACHE", str(Path.home() / ".cache" / "tt-metal-cache-gr00t"))
+    torch.set_grad_enabled(False)
+    import ttnn  # noqa: F401  (after the profiler env var)
+
+    from models.experimental.gr00t.common.preprocessing import encode
+    from models.experimental.gr00t.reference.model import load_golden_sample
+    from models.experimental.gr00t.tests.tt import test_mk_dit_block as TB
+    from models.experimental.gr00t.tests.tt.harness import device_facts
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+    from models.experimental.gr00t.tt.megakernel import descriptors as D
+    from models.experimental.gr00t.tt.megakernel import dit_program as P
+
+    TB.STREAM_DTYPE = args.dtype  # the session packs the arena in this dtype and picks the matching oracle policy
+    if args.cb_w_slots is not None:
+        os.environ["GR00T_MK_CB_W_SLOTS"] = str(args.cb_w_slots)
+    gs, obs = load_golden_sample(args.version, "canonical")
+    mi = encode(args.version, obs)
+    out: Dict[str, Any] = {
+        "schema": SCHEMA,
+        "created": datetime.now().astimezone().isoformat(timespec="seconds"),
+        "version": args.version,
+        "dtype": args.dtype,
+        "cb_w_slots": args.cb_w_slots,
+        "profile": bool(args.profile),
+        "profile_variant": args.profile_variant if args.profile else None,
+        "env": {k: v for k, v in os.environ.items() if k.startswith("GR00T_MK_")},
+        "argv": list(argv) if argv is not None else sys.argv[1:],
+    }
+    worker_l1 = None
+    if args.k5_device:
+        from models.experimental.gr00t.tt.model import worker_l1_size_for
+        from models.experimental.gr00t.tt.policy import TTPolicy
+
+        worker_l1 = int(worker_l1_size_for(TTPolicy(dit_backend="megakernel")))
+    out["k5_device"] = bool(args.k5_device)
+    out["worker_l1_size"] = worker_l1 if worker_l1 is not None else P.K4_WORKER_L1_SIZE
+    device = P.open_k4_device(worker_l1_size=out["worker_l1_size"])
+    try:
+        out["device"] = device_facts(device)
+        s = TB.K4Session(args.version, (gs, obs, mi), device, variant="block")
+        out["arena"] = {
+            "n_blocks": s.arena.n_blocks,
+            "bytes_per_step": s.mk.arena_bytes_per_step(),
+            "stream_bytes_per_step": s.arena.bytes_per_bank("stream") * 8,
+            "small_bytes_per_step": s.arena.bytes_per_bank("small") * 8,
+            "weights_load_s": s.weights_s,
+            "arena_pack_upload_s": s.arena_s,
+        }
+        if args.profile:
+            # one short launch: the zones of every RISC of every core land in the profiler CSV
+            tap = "sa_embs" if args.first_block == 0 else f"dit_block{args.first_block - 1}_out"
+            x = s.golden_rows(tap, args.step)
+            xr = s.mk.make_x_rows(x)
+            if args.profile_variant == "step":
+                s.mk.set_variant("step")
+                t0 = time.perf_counter()
+                s.mk.run_step(xr, args.step)
+                wall = time.perf_counter() - t0
+                n_blk = s.arena.n_blocks
+            else:
+                t0 = time.perf_counter()
+                s.mk.run_block(xr, args.first_block, matmul_only=False, n_blk=args.blocks, step=args.step)
+                wall = time.perf_counter() - t0
+                n_blk = args.blocks
+            s.mk._ttnn.deallocate(xr)
+            out["profile_launch"] = {
+                "variant": args.profile_variant,
+                "blocks": n_blk,
+                "first_block": args.first_block if args.profile_variant == "block" else 0,
+                "step": args.step,
+                "wall_s": wall,
+                "zone_rounds": P.K4_ZONE_ROUNDS,
+            }
+        else:
+            s.mk.set_variant("steps4")
+            noise = gs.load("initial_noise")
+            state = gs.load("state_features").reshape(-1)
+            t0 = time.perf_counter()
+            ref = s.mk.run_steps4(noise, state)
+            out["steps4_untraced_wall_s"] = time.perf_counter() - t0
+            s.mk.set_steps4_inputs(noise, state)
+            tid = s.mk.capture_steps4()
+            try:
+                times = [s.mk.replay(tid) for _ in range(args.replays)]
+                traced = s.mk.read_steps4_outputs()
+            finally:
+                s.mk.release_trace(tid)
+            ms = sorted(t * 1e3 for t in times)
+            out["steps4_trace"] = {
+                "replays": args.replays,
+                "ms_min": ms[0],
+                "ms_median": statistics.median(ms),
+                "ms_max": ms[-1],
+                "traced_eq_untraced": bool(torch.equal(traced["action_pred"], ref["action_pred"])),
+            }
+            out["g2"] = g2_record(statistics.median(ms) / P.STEPS4_N_STEPS, s.mk.arena_bytes_per_step(), args.dtype)
+            out["g2_min"] = g2_record(ms[0] / P.STEPS4_N_STEPS, s.mk.arena_bytes_per_step(), args.dtype)
+        out["stats"] = s.mk.stats().__dict__
+        out["k4b"] = {  # the K4b schedule knobs of this run (dit_program / descriptors)
+            "mcast_split": P.K4_MCAST_SPLIT,
+            "zone_rounds": P.K4_ZONE_ROUNDS,
+            "land_pages": D.K4_LAND_PAGES,
+            "hub": list(s.cm.hub),
+            "hub2": list(s.cm.hub2),
+            "split_blocks": [b for b in range(s.arena.n_blocks) if P.is_split_block(s.arena, b)],
+            "cb_w_slots": s.mk.k4.cb_w_slots[args.dtype],
+        }
+        s.release()
+    finally:
+        close_gr00t_device(device)
+    if args.profile:
+        zones = read_zones()
+        out["zones"] = zones
+        print(f"zones ({zones['freq_mhz']} MHz), {args.blocks} blocks:")
+        for name, z in sorted(zones["zones"].items()):
+            print(
+                f"  {name:28s} n={z['count']:5d} cores={z['cores']:3d} median {z['median_us']:8.2f} us  min {z['min_us']:8.2f}  max {z['max_us']:8.2f}"
+            )
+    else:
+        g = out["g2"]
+        print(
+            f"{args.version} {args.dtype}: steps4 trace median {out['steps4_trace']['ms_median']:.3f} ms "
+            f"(min {out['steps4_trace']['ms_min']:.3f}) -> {g['per_step_ms']:.3f} ms/step; floor {g['floor_ms_per_step']:.3f}, "
+            f"gate {g['gate_ms_per_step']:.3f}, ratio {g['ratio_to_floor']:.2f}, {g['effective_gbps']:.0f} GB/s -> G2 {'PASS' if g['pass'] else 'FAIL'}"
+        )
+    Path(args.results_dir).mkdir(parents=True, exist_ok=True)
+    path = Path(args.results_dir) / f"bench_mk_step_{args.version}_{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
+    with open(path, "w") as fh:
+        json.dump(out, fh, indent=1, default=str)
+    print("wrote", path)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/models/experimental/gr00t/benchmarks/sweep_backbone.py

@@ -1394,7 +1394,8 @@ def sweep_cq(
         row: Dict[str, Any] = {"cq1_uploads": flag}
         model = None
         try:
-            model = Gr00tTT.from_pretrained(version, policy=TTPolicy(), device=device)
+            # the Stage-1 model on the default-L1 device (the megakernel default needs open_model_device)
+            model = Gr00tTT.from_pretrained(version, policy=TTPolicy(dit_backend="ttnn"), device=device)
             model.executor.set_cq1_uploads(flag)
             model.warm_and_capture(mi, noise)
             times = []

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

@@ -0,0 +1,384 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Command-queue bit-equality of the traced ``Gr00tTT`` under the megakernel DiT backend (megakernel-default pass;
+``docs/plan/reviews/WP-K5-review-round1.md`` §2 "Command-queue paths", NB 3 and §8 pre-condition 3).
+
+The production path opens the device with **two** command queues (``Gr00tTT.from_pretrained`` default
+``num_command_queues=2``, the policy server's ``GR00T_NUM_CQS`` default, ``bench_e2e --num-cqs 2``): the per-call input
+writes of call ``i + 1`` go to CQ 1 while the traces of call ``i`` run on CQ 0, fenced by ``record_event`` /
+``wait_for_event`` (``tt.tracer`` ``cq1_uploads``).  Every WP-K5 device test ran on one CQ and only the bench exercised
+two, checking finiteness and shape.  This module runs the same version / sample on a device opened **twice in one
+process** through ``tt.model.open_model_device(policy, num_command_queues=...)`` -- first with 2 CQs (CQ-1 writes on),
+then with 1 CQ -- and asserts, per selected sample:
+
+* on each device: traced == untraced bit for bit on the raw ``action_pred`` buffer, the kernel's per-step
+  ``velocity_out`` taps, ``backbone_features`` and ``K_0`` / ``V_0``; ``N_REPLAYS`` traced replays bit-identical;
+* across the devices: the untraced raw tensors, the traced raw tensors, the host ``action_pred_normalized`` and the
+  per-step velocities of the 2-CQ run equal the 1-CQ run's bit for bit (``torch.equal``; PCC 1.0, max |d| 0 recorded);
+* the facts: ``cq1_uploads`` on and ``num_command_queues == 2`` on the first device, off / 1 on the second, the
+  policy (backend, arena) and the ``worker_l1_size`` the device was opened with.
+
+The backend follows the ``TTPolicy`` default (the megakernel since 2026-09-18) and ``GR00T_DIT_BACKEND`` /
+``GR00T_MK_ARENA_DTYPE`` like the other whole-model files, so the same test doubles as the ttnn backend's 2-CQ check
+when asked (``test_tt_uploads.py`` covers the ttnn path's CQ-1 writes vs the host gather independently).
+
+Run (one pytest process per version, through the device lock; alloc tracking is asserted because traced == untraced
+is)::
+
+    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_2cq.py --gr00t-version n16 [--sample multi]
+
+CPU tier (``-k cpu``): the comparison helper.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Dict, Iterator, List, Mapping, Sequence, 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.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)
+
+CQ_CONFIGS: Tuple[int, ...] = (2, 1)  # the production path first, then the reference
+N_REPLAYS = 5
+RAW_KEYS: Tuple[str, ...] = ("action_pred", "backbone_features", "K_0", "V_0")
+_RESULTS: Dict[str, Dict[int, Dict[str, Any]]] = {}  # version -> num_cqs -> {sample: record, "_facts": {...}}
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# helpers (CPU-testable)
+# --------------------------------------------------------------------------------------------------------------------
+def equal_tensors(a: Mapping[str, torch.Tensor], b: Mapping[str, torch.Tensor], keys: Sequence[str]) -> Dict[str, Any]:
+    """``{key: {"ok", "pcc", "max_abs", "shape"}}`` -- ``ok`` is ``torch.equal`` (same shape, dtype and bits); the
+    PCC / max |d| are recorded for the diagnosis when it is not."""
+    out: Dict[str, Any] = {}
+    for k in keys:
+        x, y = a[k], b[k]
+        same_shape = tuple(x.shape) == tuple(y.shape)
+        ok = bool(same_shape and x.dtype == y.dtype and torch.equal(x, y))
+        out[k] = {
+            "ok": ok,
+            "pcc": pcc(x.float(), y.float()) if same_shape else None,
+            "max_abs": max_abs(x.float(), y.float()) if same_shape else None,
+            "shape": list(x.shape),
+            "dtype": str(x.dtype),
+        }
+    return out
+
+
+def raw_from_untraced(taps: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    out = {k[len(M.RAW_PREFIX) :]: v for k, v in taps.items() if k.startswith(M.RAW_PREFIX)}
+    missing = [k for k in RAW_KEYS if k not in out]
+    if missing:
+        raise KeyError(f"run_untraced(raw=True) lacks {missing}")
+    return out
+
+
+def test_cpu_equal_tensors_helper() -> None:
+    a = {k: torch.arange(6, dtype=torch.bfloat16).reshape(2, 3) for k in RAW_KEYS}
+    res = equal_tensors(a, {k: t.clone() for k, t in a.items()}, RAW_KEYS)
+    assert all(v["ok"] and v["pcc"] == 1.0 and v["max_abs"] == 0.0 for v in res.values())
+    b = dict(a)
+    b["K_0"] = a["K_0"].clone()
+    b["K_0"][0, 0] += 1
+    res = equal_tensors(a, b, RAW_KEYS)
+    assert not res["K_0"]["ok"] and res["K_0"]["max_abs"] == 1.0 and res["action_pred"]["ok"]
+    c = dict(a)
+    c["V_0"] = a["V_0"].float()  # a dtype change is not bit-equality
+    assert not equal_tensors(a, c, RAW_KEYS)["V_0"]["ok"]
+    d = dict(a)
+    d["action_pred"] = a["action_pred"].reshape(3, 2)
+    assert not equal_tensors(a, d, RAW_KEYS)["action_pred"]["ok"]
+    with pytest.raises(KeyError):
+        raw_from_untraced({M.RAW_PREFIX + "action_pred": a["action_pred"]})
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# fixtures
+# --------------------------------------------------------------------------------------------------------------------
+@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
+    (``GR00T_DIT_BACKEND`` / ``GR00T_MK_ARENA_DTYPE``); the code default is the megakernel."""
+    return TTPolicy(dtype_policy=policy_name, trace_layout=trace_layout).with_env_overrides()
+
+
+def _selected_samples(request: pytest.FixtureRequest, version: str) -> Tuple[str, ...]:
+    opt = str(request.config.getoption(OPT_SAMPLE))
+    return tuple(samples_for(version, opt) if opt in SAMPLE_SETS else (opt,))
+
+
+@pytest.fixture(scope="module")
+def cq_results(request: pytest.FixtureRequest, version: str, policy: TTPolicy) -> Dict[int, Dict[str, Any]]:
+    """``num_cqs -> {sample: record, "_facts": {...}}`` of every selected sample: one device open per CQ
+    configuration (sequentially, module docstring), the model built once per configuration."""
+    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)
+        os.environ.setdefault("TT_METAL_CACHE", str(DEFAULT_TT_METAL_CACHE))
+        _RESULTS[version] = {n: _collect(version, samples, policy, n) for n in CQ_CONFIGS}
+        request.config._gr00t_device_facts = _RESULTS[version][CQ_CONFIGS[0]]["_facts"]["device"]  # type: ignore[attr-defined]
+    return _RESULTS[version]
+
+
+def _collect(version: str, samples: Tuple[str, ...], policy: TTPolicy, num_cqs: int) -> Dict[str, Any]:
+    """Open the device with ``num_cqs`` command queues, build the model, run every sample untraced + traced, release."""
+    from models.experimental.gr00t.reference.model import load_golden_sample
+    from models.experimental.gr00t.tt.device import close_gr00t_device
+
+    device = M.open_model_device(
+        policy,
+        trace_region_size=DEFAULT_TRACE_REGION_SIZE,
+        l1_small_size=DEFAULT_L1_SMALL_SIZE,
+        num_command_queues=num_cqs,
+    )
+    out: Dict[str, Any] = {}
+    try:
+        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_for(version) if policy.dit_backend == "megakernel" else None
+        facts["num_command_queues_requested"] = num_cqs
+        t0 = time.perf_counter()
+        model = M.Gr00tTT.from_pretrained(version, policy=policy, device=device)
+        build_s = time.perf_counter() - t0
+        try:
+            n_cqs = int(model.executor.backend.num_command_queues(device))
+            ex_desc = model.executor.describe()
+            print(
+                f"\n[test_mk_2cq] {version}/{num_cqs}cq: from_pretrained {build_s:.1f} s (weights "
+                f"{model.timing['weights_load_s']:.1f} s, {len(model.W)} tensors, {model.W.device_bytes / 1e6:.0f} MB; "
+                f"backend {policy.dit_backend}, arena {model.recipe.mk_arena_dtype}); device reports {n_cqs} CQ(s), "
+                f"cq1_uploads {ex_desc.get('cq1_uploads')}"
+            )
+            out["_facts"] = {
+                "device": facts,
+                "num_command_queues": n_cqs,
+                "cq1_uploads": bool(ex_desc.get("cq1_uploads")),
+                "build_s": build_s,
+                "weights": {"n_tensors": len(model.W), "device_bytes": float(model.W.device_bytes)},
+                "recipe": model.recipe.describe(),
+                "head_timing": dict(model.head.timing),
+            }
+            for sample in samples:
+                gs, obs = load_golden_sample(version, sample)
+                try:
+                    mi = model.encode(obs)
+                except (ValueError, StaticShapeError) as exc:  # another embodiment / layout (the N1.7 extras)
+                    out[sample] = {"skip": f"{type(exc).__name__}: {exc}"}
+                    continue
+                if mi.shape_key != model.dlayout.shape_key:
+                    out[sample] = {"skip": f"shape key {mi.shape_key} != model's"}
+                    continue
+                noise = gs.load("initial_noise")
+                rec: Dict[str, Any] = {"n_text": int(mi.n_text)}
+                pred, taps = model.run_untraced(mi, noise, return_taps=True, raw=True)
+                rec["untraced_pred"] = pred
+                rec["untraced_raw"] = raw_from_untraced(taps)
+                rec["untraced_velocity"] = _velocity_from_taps(model, taps)
+                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):
+                    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
+                rec["traced_velocity"] = _velocity_traced(model)
+                rec["golden_pred"] = gs.load("action_pred_normalized").float()
+                out[sample] = rec
+        finally:
+            model.release()
+    finally:
+        close_gr00t_device(device)
+    return out
+
+
+def _velocity_from_taps(model: M.Gr00tTT, taps: Mapping[str, torch.Tensor]) -> torch.Tensor:
+    """The untraced per-step decoder output ``[n_steps, M, A_max]`` from the head's golden-shaped taps (the megakernel's
+    ``velocity_out`` host copies; the ttnn backend's decoder taps under the same key)."""
+    keys = [f"action_decoder_out[k={k}]" for k in range(int(model.recipe.n_steps))]
+    missing = [k for k in keys if k not in taps]
+    if missing:
+        raise KeyError(f"run_untraced lacks {missing}")
+    return torch.stack([taps[k][0] for k in keys]).contiguous()
+
+
+def _velocity_traced(model: M.Gr00tTT) -> torch.Tensor:
+    """The per-step decoder output of the last traced replay on the same ``[n_steps, M, A_max]`` slice: the kernel's
+    ``velocity_out`` tap tensor (megakernel) or ``None`` for the ttnn backend (its decoder outputs are not persistent).
+    """
+    if model.policy.dit_backend != "megakernel":
+        return torch.empty(0)
+    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)
+    return vel[:, :m, :a_max].contiguous()
+
+
+def _row(results: Any, tap: str, key: str, ok: bool, **extra: Any) -> None:
+    results.add_row({"tap": tap, "key": key, "rule": "exact", "ok": bool(ok), **extra})
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device tier
+# --------------------------------------------------------------------------------------------------------------------
+def test_mk_2cq_bit_equality(
+    version: str,
+    sample: str,
+    golden: Tuple[GoldenSet, Any, Any],
+    policy: TTPolicy,
+    cq_results: Dict[int, Dict[str, Any]],
+    results: harness.ResultsWriter,
+) -> None:
+    """2-CQ (CQ-1 input writes) vs 1-CQ: traced == untraced on each device, replays bit-identical, and the two
+    devices' results bit-equal (module docstring)."""
+    for n in CQ_CONFIGS:
+        rec = cq_results[n].get(sample)
+        if rec is None or "skip" in rec:
+            pytest.skip(f"{version}/{sample}/{n}cq: {rec['skip'] if rec else 'not collected'}")
+    facts = {n: cq_results[n]["_facts"] for n in CQ_CONFIGS}
+    results.set_meta(
+        sample=sample,
+        policy=policy.to_dict(),
+        alloc_tracking_env=os.environ.get(T.ALLOC_TRACKING_ENV),
+        n_replays=N_REPLAYS,
+        cq_facts={str(n): {k: v for k, v in f.items() if k != "device"} for n, f in facts.items()},
+        device_2cq=facts[2]["device"],
+        device_1cq=facts[1]["device"],
+        n_text=cq_results[2][sample]["n_text"],
+    )
+    ok_all = True
+    # the CQ facts themselves: the production configuration really ran with two queues and the CQ-1 writes on
+    fact_ok = (
+        facts[2]["num_command_queues"] == 2
+        and facts[2]["cq1_uploads"] is True
+        and facts[1]["num_command_queues"] == 1
+        and facts[1]["cq1_uploads"] is False
+    )
+    _row(results, "device", "cq_facts", fact_ok, **{f"{n}cq": facts[n]["cq1_uploads"] for n in CQ_CONFIGS})
+    ok_all = ok_all and fact_ok
+    for n in CQ_CONFIGS:
+        rec = cq_results[n][sample]
+        st0 = rec["traced_states"][0]
+        eq_raw = equal_tensors(st0, rec["untraced_raw"], RAW_KEYS)
+        for k, r in eq_raw.items():
+            _row(results, k, f"{n}cq-traced_equals_untraced", r["ok"], pcc=r["pcc"], max_abs=r["max_abs"])
+        eq_host = bool(torch.equal(rec["traced_preds"][0], rec["untraced_pred"]))
+        _row(results, "action_pred_normalized", f"{n}cq-traced_equals_untraced", eq_host)
+        eq_replays = all(
+            all(equal_tensors(st0, s, RAW_KEYS)[k]["ok"] for k in RAW_KEYS) for s in rec["traced_states"][1:]
+        ) and all(torch.equal(p, rec["traced_preds"][0]) for p in rec["traced_preds"][1:])
+        _row(
+            results,
+            "action_pred",
+            f"{n}cq-{N_REPLAYS}_replays_bit_identical",
+            eq_replays,
+            replay_ms=[round(s * 1e3, 3) for s in rec["replay_s"]],
+        )
+        ok_all = ok_all and all(r["ok"] for r in eq_raw.values()) and eq_host and eq_replays
+        if rec["traced_velocity"].numel():
+            eq_vel = bool(torch.equal(rec["traced_velocity"], rec["untraced_velocity"]))
+            _row(
+                results,
+                "velocity_out",
+                f"{n}cq-traced_equals_untraced",
+                eq_vel,
+                shape=list(rec["traced_velocity"].shape),
+            )
+            ok_all = ok_all and eq_vel
+        print(
+            f"[{version}/{sample}] {n}cq: traced==untraced {all(r['ok'] for r in eq_raw.values()) and eq_host}, "
+            f"{N_REPLAYS} replays identical {eq_replays}, replay+D2H ms {[round(s * 1e3, 2) for s in rec['replay_s']]}"
+        )
+    # across the two devices: untraced, traced, host prediction and the per-step velocities
+    two, one = cq_results[2][sample], cq_results[1][sample]
+    cross_untraced = equal_tensors(two["untraced_raw"], one["untraced_raw"], RAW_KEYS)
+    cross_traced = equal_tensors(two["traced_states"][0], one["traced_states"][0], RAW_KEYS)
+    for k in RAW_KEYS:
+        _row(
+            results,
+            k,
+            "2cq_equals_1cq-untraced",
+            cross_untraced[k]["ok"],
+            **{x: cross_untraced[k][x] for x in ("pcc", "max_abs")},
+        )
+        _row(
+            results,
+            k,
+            "2cq_equals_1cq-traced",
+            cross_traced[k]["ok"],
+            **{x: cross_traced[k][x] for x in ("pcc", "max_abs")},
+        )
+    eq_pred = bool(torch.equal(two["traced_preds"][0], one["traced_preds"][0]))
+    _row(
+        results,
+        "action_pred_normalized",
+        "2cq_equals_1cq-host_pred",
+        eq_pred,
+        pcc=pcc(two["traced_preds"][0], one["traced_preds"][0]),
+        max_abs=max_abs(two["traced_preds"][0], one["traced_preds"][0]),
+    )
+    eq_vel_cross = True
+    if two["traced_velocity"].numel() and one["traced_velocity"].numel():
+        eq_vel_cross = bool(torch.equal(two["traced_velocity"], one["traced_velocity"]))
+        _row(results, "velocity_out", "2cq_equals_1cq-traced", eq_vel_cross)
+    ok_cross = all(r["ok"] for r in cross_untraced.values()) and all(r["ok"] for r in cross_traced.values())
+    ok_all = ok_all and ok_cross and eq_pred and eq_vel_cross
+    # sanity, not the subject: the prediction is a valid answer for this sample (the e2e files gate it)
+    p_gold = pcc(two["traced_preds"][0], two["golden_pred"])
+    results.add_row(
+        {
+            "tap": "action_pred_normalized",
+            "key": "2cq-vs-golden",
+            "rule": "info",
+            "ok": bool(p_gold > 0.99),
+            "pcc": p_gold,
+        }
+    )
+    results.set_timing(
+        replay_median_2cq_s=sorted(two["replay_s"])[len(two["replay_s"]) // 2],
+        replay_median_1cq_s=sorted(one["replay_s"])[len(one["replay_s"]) // 2],
+        build_2cq_s=facts[2]["build_s"],
+        build_1cq_s=facts[1]["build_s"],
+    )
+    print(
+        f"[{version}/{sample}] 2cq == 1cq: untraced {all(r['ok'] for r in cross_untraced.values())}, traced "
+        f"{all(r['ok'] for r in cross_traced.values())}, host pred {eq_pred}, velocity {eq_vel_cross}; "
+        f"pcc vs golden {p_gold:.6f}"
+    )
+    assert fact_ok, {n: facts[n] for n in CQ_CONFIGS}
+    assert ok_all, "see the rows: a CQ configuration is not bit-equal to its untraced run or to the other one"
+    assert p_gold > 0.99

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

@@ -174,7 +174,8 @@ def test_core_map_vc_and_per_core_args(core_map: CM.CoreMap):
     for c, core in enumerate(core_map.compute):
         g = CM.group_of(c)
         assert args[core] == (c, g, core_map.vc[g], c % 32, g)
-    assert args[core_map.hub] == (0xFFFF,) * 5
+    assert args[core_map.hub] == (0xFFFF,) * 4 + (0,)  # K4b: the group word of a hub is its hub index
+    assert args[core_map.hub2] == (0xFFFF,) * 4 + (1,)
 
 
 def test_core_map_other_worker_lists_and_errors():
@@ -612,7 +613,10 @@ def test_descriptor_spec_contents(core_map):
         assert ctas[name] == cb_id
     for name, sem_id in D.SEM_IDS.items():
         assert ctas[name] == sem_id
-    assert len(spec.semaphores) == 9 and sorted(s.sem_id for s in spec.semaphores) == list(range(9))
+    # 9 K1 semaphores + the 4 K4b ones (hub done / odd ready counter / hub-1 flag / hub-1 landing base)
+    n_sems = len(D.SEM_IDS)
+    assert n_sems == 13 and len(spec.semaphores) == n_sems
+    assert sorted(s.sem_id for s in spec.semaphores) == list(range(n_sems))
     # common RT args: addresses (0 in the dry run), step index, offsets table
     io = spec.io_order()
     assert io[:3] == ["x_rep", "a_rep", "arena_stream"] and io[-3:] == ["x_t", "state_features", "action_pred"]

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

@@ -0,0 +1,398 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K4 device tests of the megakernel's full DiT blocks (``block`` variant; IMPLEMENTATION_PLAN.md §5.4 P1-P12,
+§5.7 K4 acceptance rows ``dit_block0_out`` / ``dit_block1_out``).
+
+* ``test_kernel_cta_names_cpu`` -- every ``get_named_compile_time_arg_val`` name of ``dit_kernel.cpp`` / ``ops/*.hpp``
+  is emitted by the K4 program spec (or by ``materialise``), so a missing CTA fails on the CPU, not as a JIT error;
+* ``test_k4_spec_cpu`` -- the K4 spec of every version / stream dtype keeps the descriptor limits, the tail geometry
+  mirror and the tile-row ``x_rep`` shard bookkeeping;
+* ``test_block_vs_oracle`` -- one self block (block 1 from the golden ``dit_block0_out[k]``) and one cross block
+  (block 0 from the golden ``sa_embs[k]``) with the real bf16 weights, device-hoisted K/V and the call's masks, steps 0
+  and 3: PCC >= 0.9999 against the Stage-1 ``DiTSelfBlockTT`` / ``DiTCrossBlockTT`` on device (all 64 rows), the
+  golden ``dit_block{0,1}_out[k]`` gates of ``gates_multi.json`` (rows < M), every core's resident tile row
+  bit-identical, and a second launch bit-identical;
+* ``test_blocks_chain_vs_golden`` -- blocks [0, 2) in one launch from ``sa_embs[k]`` vs ``dit_block1_out[k]`` and
+  vs the chained oracle.
+
+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_dit_block.py --gr00t-version n16
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterator, 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 DIT_QUERY_ROWS_PAD, GR00TConfig, get_config
+from models.experimental.gr00t.common.golden import GoldenSet
+from models.experimental.gr00t.tests.tt.harness import check_taps
+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, default_core_map
+
+torch.set_grad_enabled(False)
+
+KERNEL_DIR = Path(P.__file__).resolve().parent / "kernels"
+GATE_ORACLE = 0.9999  # plan §5.7 K4: PCC vs the S1 device oracle
+STEPS_TESTED: Tuple[int, ...] = (0, 3)
+SELF_BLK, CROSS_BLK = 1, 0
+#: arena stream dtype of the session (bring-up order: bf16 first, bfp8 after G2); ``GR00T_MK_STREAM_DTYPE=bfp8_b``
+#: runs the same tests on the bfp8 arena with the ``mixed_dit`` (bfp8 DiT weights) Stage-1 oracle
+STREAM_DTYPE = os.environ.get("GR00T_MK_STREAM_DTYPE", "bf16")
+if STREAM_DTYPE not in ("bf16", "bfp8_b"):
+    raise ValueError(f"GR00T_MK_STREAM_DTYPE must be bf16 or bfp8_b, got {STREAM_DTYPE!r}")
+HEAD_CATEGORIES: Tuple[str, ...] = ("dit", "dit_precompute", "encoders")
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# host helpers
+# --------------------------------------------------------------------------------------------------------------------
+def pad_rows(t: torch.Tensor) -> torch.Tensor:
+    """Golden ``[1, M, C]`` fp32 -> ``[64, C]`` bf16-representable fp32 with exactly-zero pad rows (plan D5)."""
+    if t.ndim != 3 or t.shape[0] != 1 or t.shape[1] > DIT_QUERY_ROWS_PAD:
+        raise ValueError(f"expected [1, M <= 64, C], got {tuple(t.shape)}")
+    out = torch.zeros(DIT_QUERY_ROWS_PAD, t.shape[2], dtype=torch.float32)
+    out[: t.shape[1]] = t[0].to(torch.bfloat16).float()
+    return out
+
+
+def kernel_cta_names(kernel_dir: Path = KERNEL_DIR) -> List[str]:
+    """Every distinct ``get_named_compile_time_arg_val("<name>")`` of the megakernel sources (kernel + ops headers)."""
+    pat = re.compile(r'get_named_compile_time_arg_val\("([A-Za-z0-9_]+)"\)')
+    names: List[str] = []
+    for f in sorted(list(kernel_dir.glob("*.cpp")) + list((kernel_dir / "ops").glob("*.hpp"))):
+        for m in pat.finditer(f.read_text()):
+            if m.group(1) not in names:
+                names.append(m.group(1))
+    return names
+
+
+def oracle_policy() -> Any:
+    from models.experimental.gr00t.tt.policy import TTPolicy
+
+    return TTPolicy(dtype_policy="bf16" if STREAM_DTYPE == "bf16" else "mixed_dit")
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# CPU
+# --------------------------------------------------------------------------------------------------------------------
+def test_kernel_cta_names_cpu() -> None:
+    """The K4 spec (+ ``fill_missing_ctas`` + ``materialise``'s NoC names) covers every CTA the sources read."""
+    cm = default_core_map()
+    L = A.ArenaLayout.plan("n16", dtype_stream="bf16", reader_mode="direct", tile_order="row_major")
+    spec = D.build_program_spec_k4(cm, L, "steps4", {}, dry_run=True)
+    spec.named_ctas += [(f"mk_io_{n}", 0) for n in P.MK_IO_CTAS] + [
+        ("mk_k1", 0),
+        ("mk_ring_depth", 4),
+        ("mk_relay_window", 6),
+        ("mk_rt_base", 0),
+        ("mk_io_kv_base", 0),
+    ]
+    P.fill_missing_ctas(spec)
+    emitted = {n for n, _ in spec.named_ctas} | {flag for flag, _ in spec.role_flags}
+    emitted |= set(P.MATERIALISE_CTAS)
+    names = kernel_cta_names()
+    missing = [n for n in names if n not in emitted]
+    assert len(names) > 60, names
+    assert not missing, f"CTAs read by the kernel sources but never emitted: {missing}"
+    # and the union list itself only names things the sources read or the K1 spec emits
+    k1 = D.build_program_spec(cm, L, "block", {}, dry_run=True)
+    k1_emitted = {n for n, _ in k1.named_ctas}
+    stale = [n for n in P.KERNEL_CTA_UNION if n not in names and n not in k1_emitted]
+    assert not stale, f"KERNEL_CTA_UNION entries no source reads: {stale}"
+
+
+@pytest.mark.parametrize("dtype", ("bf16", "bfp8_b"))
+def test_k4_spec_cpu(version: str, dtype: str) -> None:
+    """Descriptor limits, tail geometry mirror and the x_rep tile-row bookkeeping of the K4 spec."""
+    cm = default_core_map()
+    L = A.ArenaLayout.plan(version, dtype_stream=dtype, reader_mode="direct", tile_order="row_major")
+    assert P.check_kernel_geometry(L) == 4 * L.n_blocks
+    assert P.check_tail_geometry(L) == len(P.TAIL_OP_NAMES)
+    for variant in D.VARIANTS:
+        spec = D.build_program_spec_k4(cm, L, variant, {}, dry_run=True)
+        st = spec.stats()
+        assert st.n_common_rt_args + len(P.MK4_RT) + cm.n_compute <= D.MAX_COMMON_RT_ARGS, st
+        assert st.n_kernel_groups <= D.MAX_KERNEL_GROUPS and st.n_cbs <= D.MAX_CBS
+        assert st.l1_bytes_compute <= D.K4Budget().l1_budget_compute
+        cb = {c.name: c for c in spec.cbs}
+        assert cb["cb_x"].total_bytes == 48 * 2048 and cb["cb_x"].tensor_backed == "x_rep"
+        assert cb["cb_a"].total_bytes == 128 * 2048
+        assert cb["cb_stat32"].cb_id == D.K4_HUB_CB_IDS["cb_act"]
+    # the tile row of every compute core (dist_layernorm.hpp: r = c // 48, t = c % 48)
+    order = P.shard_order(cm)
+    assert sorted(order) == list(range(cm.n_compute))
+    assert all((c // 48) in (0, 1) for c in order) and sum(1 for c in order if c // 48 == 0) == 48
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device session (one arena + weights + hoisted K/V per version)
+# --------------------------------------------------------------------------------------------------------------------
+class K4Session:
+    """Real bf16 weights on device (``TTWeights``, the Stage-1 oracle's), the same plan tensors packed into the
+    megakernel arena, the call's masks and the device-hoisted K/V of the golden sample, one ``DiTMegakernel``."""
+
+    def __init__(self, version: str, golden: Tuple[GoldenSet, Any, Any], device: Any, variant: str = "block") -> None:
+        import ttnn
+        from models.experimental.gr00t.tests.tt.test_tt_dit_ops import check_timesteps, gathered_vl_sets, hoist_kv
+        from models.experimental.gr00t.tt import layers as TL
+        from models.experimental.gr00t.tt import layout as LAY
+        from models.experimental.gr00t.tt import weights as TW
+        from models.experimental.gr00t.tt.dit_ops import dit_masks
+
+        gs, obs, mi = golden
+        self.gs, self.mi = gs, mi
+        self.cfg: GR00TConfig = get_config(version)
+        self.dl = LAY.DeviceLayout.for_inputs(self.cfg, mi)
+        self.dl.assert_call(mi)
+        check_timesteps(gs, self.cfg)
+        if self.dl.layout_name != self.cfg.canonical_layout.name:
+            raise RuntimeError(f"{version}: golden sample layout {self.dl.layout_name} is not the canonical layout")
+        self.device = device
+        self.policy = oracle_policy()
+        t0 = time.perf_counter()
+        entries = TW.upload_plan(version, None, self.policy)
+        names = sorted(n for n, e in entries.items() if e.category in HEAD_CATEGORIES)
+        self.W = TW.TTWeights.load(version, None, self.policy, device, names=names)
+        self.weights_s = time.perf_counter() - t0
+        # masks (materialise dictionary) and hoisted K/V from the golden VL rows (the WP-10 hoist recipe on device)
+        self.masks = LAY.materialise(device, self.dl, mi.slots)
+        self.kv = hoist_kv(
+            self.cfg, self.dl, self.W, gathered_vl_sets(gs, mi, self.dl), lambda t: TL.to_device(t, device, mem="DRAM")
+        )
+        # the megakernel arena from the same host plan tensors
+        t0 = time.perf_counter()
+        self.arena = A.ArenaLayout.plan(
+            version, dtype_stream=STREAM_DTYPE, reader_mode="direct", tile_order="row_major"
+        )
+        plan = P.real_arena_plan(version, policy=self.policy)
+        self.pack = self.arena.pack(plan)
+        self.cm = CoreMap.from_device(device)
+        # GR00T_MK_CB_W_SLOTS overrides the weight-ring depth (K-block slots) for the G2 sweeps (K4Budget default 13 bf16 / 15 bfp8)
+        slots = os.environ.get("GR00T_MK_CB_W_SLOTS")
+        k4_budget = D.K4Budget(cb_w_slots={STREAM_DTYPE: int(slots)}) if slots else None
+        self.mk = P.DiTMegakernel(
+            self.cfg, self.dl, None, self.arena, self.cm, device, variant=variant, k1_extras=False, k4_budget=k4_budget
+        )
+        self.mk.upload(self.pack)
+        self.mk.bind_kv(self.kv)
+        self.mk.bind_masks(dit_masks(self.masks))
+        ttnn.synchronize_device(device)
+        self.arena_s = time.perf_counter() - t0
+        self.dit_masks = dit_masks(self.masks)
+        self._blocks: Dict[int, Any] = {}
+
+    def oracle_block(self, i: int) -> Any:
+        from models.experimental.gr00t.tt.dit_ops import DiTCrossBlockTT, DiTSelfBlockTT
+
+        if i not in self._blocks:
+            if self.arena.is_cross(i):
+                self._blocks[i] = DiTCrossBlockTT(self.cfg, i, self.W, self.dl, self.masks, self.arena.key_sets[i])
+            else:
+                self._blocks[i] = DiTSelfBlockTT(self.cfg, i, self.W, self.dl, self.masks)
+        return self._blocks[i]
+
+    def oracle_blocks(self, x: torch.Tensor, blocks: Sequence[int], step: int) -> torch.Tensor:
+        """Stage-1 device path: blocks ``blocks`` in order from the host ``x [64, 1536]`` -> host fp32 ``[64, 1536]``."""
+        import ttnn
+        from models.experimental.gr00t.tt import layers as TL
+
+        cur = TL.to_device(x.reshape(1, 1, DIT_QUERY_ROWS_PAD, -1).to(torch.bfloat16), self.device, mem="L1")
+        for i in blocks:
+            blk = self.oracle_block(i)
+            nxt = blk.forward(cur, step, self.kv[i] if self.arena.is_cross(i) else None)
+            ttnn.deallocate(cur)
+            cur = nxt
+        out = TL.to_torch(cur).reshape(DIT_QUERY_ROWS_PAD, -1).float()
+        ttnn.deallocate(cur)
+        return out
+
+    def golden_rows(self, tap: str, step: int) -> torch.Tensor:
+        return pad_rows(self.gs.load(tap, step=step))
+
+    def release(self) -> None:
+        import ttnn
+
+        self.mk.deallocate()
+        for t in self.masks.values():
+            if t is not None:
+                ttnn.deallocate(t)
+        for k, v in self.kv.values():
+            ttnn.deallocate(k)
+            ttnn.deallocate(v)
+        self.W.release()
+
+
+_SESSIONS: Dict[str, K4Session] = {}
+
+
+@pytest.fixture(scope="module")
+def k4_device() -> Iterator[Any]:
+    """The p150a opened with the smaller ``worker_l1_size`` the K4 binaries need (``dit_program.open_k4_device``: the
+    compute-core BRISC + NCRISC + 3 TRISC binaries are ~91 KB, the default kernel-config ring buffer 70,656 B).  The
+    session-scoped device fixture of conftest must not be open in the same process (one device at a time)."""
+    import os
+
+    from models.experimental.gr00t.tests.tt.conftest import 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 = P.open_k4_device()
+    try:
+        yield device
+    finally:
+        for v, s in list(_SESSIONS.items()):
+            s.release()
+            del _SESSIONS[v]
+        close_gr00t_device(device)
+
+
+@pytest.fixture(scope="module")
+def k4(version: str, golden: Tuple[GoldenSet, Any, Any], k4_device: Any) -> Iterator[K4Session]:
+    for v, s in list(_SESSIONS.items()):
+        if v != version:
+            s.release()
+            del _SESSIONS[v]
+    if version not in _SESSIONS:
+        _SESSIONS[version] = K4Session(version, golden, k4_device)
+    yield _SESSIONS[version]
+
+
+def _row(results: Any, key: str, got: torch.Tensor, ref: torch.Tensor, gate: float, **extra: Any) -> Dict[str, Any]:
+    m = pcc_mod.metrics(got.float(), ref.float())
+    row = {
+        "tap": "mk_k4_block",
+        "key": key,
+        "rule": "pcc",
+        "pcc": m["pcc"],
+        "gate_pcc": gate,
+        "max_abs": m["max_abs"],
+        "rel_l2": m["rel_l2"],
+        "finite": m["finite"],
+        "ok": bool(m["finite"] and m["pcc"] >= gate),
+    }
+    row.update(extra)
+    results.add_row(row)
+    return row
+
+
+def _session_meta(results: Any, s: K4Session) -> None:
+    from models.experimental.gr00t.tests.tt.harness import device_facts
+
+    results.set_device_facts(device_facts(s.device))
+    results.set_meta(
+        worker_l1_size=P.K4_WORKER_L1_SIZE,
+        layout=s.dl.layout_name,
+        embodiment_id=s.dl.embodiment_id,
+        stream_dtype=STREAM_DTYPE,
+        weights_load_s=s.weights_s,
+        arena_pack_upload_s=s.arena_s,
+        key_sets=list(s.arena.key_sets[:4]),
+        stats=s.mk.stats().__dict__ if s.mk.last_spec is not None else None,
+        compile_s={str(k): v for k, v in s.mk.compile_seconds.items()},
+    )
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# device
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.mark.parametrize("step", STEPS_TESTED)
+@pytest.mark.parametrize("kind", ("self", "cross"))
+def test_block_vs_oracle(results: Any, version: str, k4: K4Session, kind: str, step: int) -> None:
+    """One full block (P1-P12) from a golden input vs the S1 block on device (all 64 rows, PCC >= 0.9999) and vs the
+    golden block output (rows < M, ``gates_multi.json``); every core's tile row bit-identical; deterministic."""
+    s = k4
+    blk = SELF_BLK if kind == "self" else CROSS_BLK
+    assert s.arena.is_cross(blk) == (kind == "cross")
+    x = s.golden_rows("sa_embs" if blk == 0 else f"dit_block{blk - 1}_out", step)
+    x_rows = s.mk.make_x_rows(x)
+    try:
+        t0 = time.perf_counter()
+        got = s.mk.run_block(x_rows, blk, matmul_only=False, step=step)
+        dt = time.perf_counter() - t0
+        all_rows = s.mk.read_x_rows_all(x_rows)
+        # second launch from the same input: bit-identical
+        x_rows2 = s.mk.make_x_rows(x)
+        try:
+            got2 = s.mk.run_block(x_rows2, blk, matmul_only=False, step=step)
+        finally:
+            s.mk._ttnn.deallocate(x_rows2)
+    finally:
+        s.mk._ttnn.deallocate(x_rows)
+    ref = s.oracle_blocks(x, [blk], step)
+    m = s.dl.m_logical
+    row = _row(
+        results,
+        f"{kind}-blk{blk}-k{step}-vs-oracle",
+        got,
+        ref,
+        GATE_ORACLE,
+        blk=blk,
+        step=step,
+        kind=kind,
+        pcc_valid_rows=pcc_mod.pcc(got[:m], ref[:m]),
+        wall_s=dt,
+    )
+    rows_ok = all(torch.equal(all_rows[c], all_rows[48 * (c // 48)]) for c in range(s.cm.n_compute))
+    results.add_row(
+        {
+            "tap": "mk_k4_block",
+            "key": f"{kind}-blk{blk}-k{step}-consistency",
+            "rule": "exact",
+            "rows_bit_identical": rows_ok,
+            "deterministic": bool(torch.equal(got, got2)),
+            "ok": bool(rows_ok and torch.equal(got, got2)),
+        }
+    )
+    taps = {s.gs.threshold_key(f"dit_block{blk}_out", step=step): got[:m].reshape(1, m, -1)}
+    report = check_taps(version, taps, s.gs, results=results)
+    _session_meta(results, s)
+    assert row["finite"], "non-finite block output"
+    assert rows_ok, "the compute cores of one tile row hold different residual rows"
+    assert torch.equal(got, got2), "two launches of the same block differ"
+    assert row["pcc"] >= GATE_ORACLE, f"{kind} block {blk} k={step}: PCC vs S1 {row['pcc']:.6f} < {GATE_ORACLE}"
+    assert report.ok, report.summary(failures_only=True)
+
+
+@pytest.mark.parametrize("step", STEPS_TESTED)
+def test_blocks_chain_vs_golden(results: Any, version: str, k4: K4Session, step: int) -> None:
+    """Blocks [0, 2) in one launch from ``sa_embs[k]``: the pending-residual hand-over between blocks (P12 -> P1) vs
+    the golden ``dit_block1_out[k]`` and the chained S1 oracle."""
+    s = k4
+    x = s.golden_rows("sa_embs", step)
+    x_rows = s.mk.make_x_rows(x)
+    try:
+        t0 = time.perf_counter()
+        got = s.mk.run_block(x_rows, 0, matmul_only=False, n_blk=2, step=step)
+        dt = time.perf_counter() - t0
+    finally:
+        s.mk._ttnn.deallocate(x_rows)
+    ref = s.oracle_blocks(x, [0, 1], step)
+    m = s.dl.m_logical
+    row = _row(
+        results,
+        f"chain-blk0-1-k{step}-vs-oracle",
+        got,
+        ref,
+        GATE_ORACLE,
+        step=step,
+        pcc_valid_rows=pcc_mod.pcc(got[:m], ref[:m]),
+        wall_s=dt,
+    )
+    taps = {s.gs.threshold_key("dit_block1_out", step=step): got[:m].reshape(1, m, -1)}
+    report = check_taps(version, taps, s.gs, results=results)
+    _session_meta(results, s)
+    assert row["finite"] and row["pcc"] >= GATE_ORACLE, f"chain k={step}: PCC vs S1 {row['pcc']:.6f}"
+    assert report.ok, report.summary(failures_only=True)

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

@@ -0,0 +1,335 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""WP-K4 device tests of the megakernel's ``step`` and ``steps4`` variants (IMPLEMENTATION_PLAN.md §5.4 "Per step",
+§5.7 K4 acceptance rows ``dit_block_last_out`` / ``dit_out`` / ``pred_velocity`` / ``action_pred_*``, gate G2).
+
+* ``test_step_vs_golden[k]`` -- all ``n_blocks`` blocks + output AdaLN + ``proj_out_2`` of denoising step ``k`` from
+  the golden ``sa_embs[k]``: ``dit_block_last_out[k]`` / ``dit_out[k]`` against the golden gates and against the
+  Stage-1 ``DiTStepTT`` on device (``dit_block_last_out`` PCC >= 0.9999; ``dit_out`` >= 0.9998 **and** not further from
+  the fp32 golden than Stage 1 is, see ``GATE_ORACLE_HEAD`` / ``GOLDEN_SLACK``); a second launch bit-identical;
+* ``test_steps4_vs_golden`` -- the whole denoise in one launch (action encoder -> sa_embs -> blocks -> head ->
+  decoder -> Euler, 4 steps) from the golden ``initial_noise`` / ``state_features``: ``pred_velocity[k]``,
+  ``action_decoder_out[k]``, ``action_pred_normalized``, ``action_pred_valid`` gates + PCC vs the Stage-1
+  ``ActionHeadTT.denoise`` on device;
+* ``test_steps4_determinism_and_trace`` -- 10 untraced launches bit-identical, the traced replay bit-identical to the
+  untraced result (``action_pred``, and ``velocity`` when the taps are allocated), and the G2 timing: 20 traced replays
+  (min / median ms), per-step ms vs ``1.3 x bytes / GB/s``.  G2 is **recorded, not asserted**: the ``G2-step-time`` row
+  carries ``ok = pass`` (a miss shows as ``ok: False`` in the results JSON while the test stays green -- plan section
+  5.7 fallback, ratified 2026-09-18 in docs/plan/reviews/WP-K4-review-round1.md).
+
+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_dit_step.py --gr00t-version n16
+"""
+
+from __future__ import annotations
+
+import statistics
+import time
+from typing import Any, Dict, Iterator, List, Tuple
+
+import pytest
+import torch
+
+from models.experimental.gr00t.common import pcc as pcc_mod
+from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
+from models.experimental.gr00t.common.configs import DIT_QUERY_ROWS_PAD
+from models.experimental.gr00t.common.golden import GoldenSet
+from models.experimental.gr00t.tests.tt.harness import check_taps, device_facts
+from models.experimental.gr00t.tests.tt.test_mk_dit_block import K4Session, STREAM_DTYPE, pad_rows
+from models.experimental.gr00t.tt import encoders as E
+from models.experimental.gr00t.tt.megakernel import dit_program as P
+
+torch.set_grad_enabled(False)
+
+GATE_ORACLE = 0.9999  # plan §5.7 K4: PCC vs the S1 device oracle (blocks, residual stream, action taps)
+#: the output head after 32 blocks: the two bf16 device pipelines (different op / rounding order per block) drift to
+#: 1 - PCC ~ 1e-4 on dit_out (measured 0.99990-0.99997 on N1.6, mk_k4_summary.md §4); gated at 0.9998 against the
+#: oracle AND "not further from the fp32 golden than Stage 1 is" (GOLDEN_SLACK).  Post-hoc change from the plan's
+#: 0.9999: N1.6 k=3 measured 0.9998980 / 0.9998996 under 0.9999 (test_step_vs_golden_3_n16_20260917-233300 / -233501,
+#: failed); ratified by the plan owner 2026-09-18 (WP-K4-review-round1, "Plan-owner decisions" 1; summary §9).  The
+#: golden gates (gates_multi.json) are untouched.
+GATE_ORACLE_HEAD = 0.9998
+GOLDEN_SLACK = 2e-5
+STEPS: Tuple[int, ...] = (0, 1, 2, 3)
+N_DETERMINISM = 10  # plan D12: megakernel 10x bit-identical
+N_REPLAYS = 20  # G2 timing (work_packages.json WP-K4)
+#: measured direct-mode arena rates (mk_k1_summary.md §3: 464 GB/s bf16, 414 GB/s bfp8) and the G2 factor (plan §5.7)
+MEASURED_GBPS: Dict[str, float] = {"bf16": 464.0, "bfp8_b": 414.0}
+G2_FACTOR = 1.3
+
+_SESSIONS: Dict[str, K4Session] = {}
+
+
+@pytest.fixture(scope="module")
+def k4_device() -> Iterator[Any]:
+    """The p150a opened with the smaller ``worker_l1_size`` the K4 binaries need (see test_mk_dit_block)."""
+    import os
+
+    from models.experimental.gr00t.tests.tt.conftest import 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 = P.open_k4_device()
+    try:
+        yield device
+    finally:
+        for v, s in list(_SESSIONS.items()):
+            s.release()
+            del _SESSIONS[v]
+        close_gr00t_device(device)
+
+
+@pytest.fixture(scope="module")
+def k4(version: str, golden: Tuple[GoldenSet, Any, Any], k4_device: Any) -> Iterator[K4Session]:
+    for v, s in list(_SESSIONS.items()):
+        if v != version:
+            s.release()
+            del _SESSIONS[v]
+    if version not in _SESSIONS:
+        _SESSIONS[version] = K4Session(version, golden, k4_device, variant="step")
+    yield _SESSIONS[version]
+
+
+def _row(results: Any, key: str, got: torch.Tensor, ref: torch.Tensor, gate: float, **extra: Any) -> Dict[str, Any]:
+    m = pcc_mod.metrics(got.float(), ref.float())
+    row = {
+        "tap": "mk_k4_step",
+        "key": key,
+        "rule": "pcc",
+        "pcc": m["pcc"],
+        "gate_pcc": gate,
+        "max_abs": m["max_abs"],
+        "rel_l2": m["rel_l2"],
+        "finite": m["finite"],
+        "ok": bool(m["finite"] and m["pcc"] >= gate),
+    }
+    row.update(extra)
+    results.add_row(row)
+    return row
+
+
+def _meta(results: Any, s: K4Session) -> None:
+    results.set_device_facts(device_facts(s.device))
+    results.set_meta(
+        layout=s.dl.layout_name,
+        embodiment_id=s.dl.embodiment_id,
+        stream_dtype=STREAM_DTYPE,
+        n_blocks=s.arena.n_blocks,
+        key_sets=sorted(set(s.arena.key_sets)),
+        worker_l1_size=P.K4_WORKER_L1_SIZE,
+        stats=s.mk.stats().__dict__ if s.mk.last_spec is not None else None,
+        compile_s={str(k): v for k, v in s.mk.compile_seconds.items()},
+    )
+
+
+def _b3_bias(s: K4Session) -> torch.Tensor:
+    """``action_encoder.W3.b[e]`` (folded into ``b3_pos`` on device; the golden ``action_encoder_out`` includes it)."""
+    ck = LazyCheckpoint(s.cfg.version)
+    try:
+        return ck.slot(
+            s.cfg.checkpoint.action_head_prefix + "action_encoder.W3.b", int(s.dl.embodiment_id), torch.float32
+        )
+    finally:
+        ck.close()
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# step variant
+# --------------------------------------------------------------------------------------------------------------------
+@pytest.mark.parametrize("step", STEPS)
+def test_step_vs_golden(results: Any, version: str, k4: K4Session, step: int) -> None:
+    """One denoising step (all blocks + head) from the golden ``sa_embs[k]`` vs the golden taps and the S1 step."""
+    import ttnn
+    from models.experimental.gr00t.tt import layers as TL
+    from models.experimental.gr00t.tt.dit_ops import DiTStepTT
+
+    s = k4
+    m = s.dl.m_logical
+    x = s.golden_rows("sa_embs", step)
+    # S1 oracle first (blocks + head, every block output kept as a tap), with the megakernel's L1 residents released:
+    # under the K4 worker_l1_size the oracle's layer_norm CBs do not fit next to a_rep on 32-block versions
+    s.mk.release_l1_residents()
+    oracle = DiTStepTT(s.cfg, s.W, s.dl, s.masks, s.kv, collect_taps=True)
+    xd = TL.to_device(x.reshape(1, 1, DIT_QUERY_ROWS_PAD, -1).to(torch.bfloat16), s.device, mem="L1")
+    y_dev = oracle.forward(xd, step)
+    ref_y = TL.to_torch(y_dev).float().reshape(DIT_QUERY_ROWS_PAD, -1)
+    ref_x = TL.to_torch(oracle.debug_taps["dit_block_last_out"]).float().reshape(DIT_QUERY_ROWS_PAD, -1)
+    for t in list(oracle.debug_taps.values()) + [xd]:
+        ttnn.deallocate(t)
+    s.mk.set_variant("step")
+    t0 = time.perf_counter()
+    y, x_last = s.mk.run_step(x, step)
+    dt = time.perf_counter() - t0
+    y2, x_last2 = s.mk.run_step(x, step)
+    row_x = _row(results, f"dit_block_last_out-k{step}-vs-oracle", x_last, ref_x, GATE_ORACLE, step=step, wall_s=dt)
+    row_y = _row(results, f"dit_out-k{step}-vs-oracle", y, ref_y, GATE_ORACLE_HEAD, step=step)
+    det = bool(torch.equal(y, y2) and torch.equal(x_last, x_last2))
+    results.add_row({"tap": "mk_k4_step", "key": f"step-k{step}-deterministic", "rule": "exact", "ok": det})
+    taps = {
+        s.gs.threshold_key("dit_block_last_out", step=step): x_last[:m].reshape(1, m, -1),
+        s.gs.threshold_key("dit_out", step=step): y[:m].reshape(1, m, -1),
+    }
+    report = check_taps(version, taps, s.gs, results=results)
+    # the megakernel must be at least as close to the fp32 golden as the Stage-1 path is (both bf16 pipelines)
+    gold_x = s.golden_rows("dit_block_last_out", step)
+    gold_y = pad_rows(s.gs.load("dit_out", step=step))
+    closeness = {
+        "dit_block_last_out": (pcc_mod.pcc(x_last[:m], gold_x[:m]), pcc_mod.pcc(ref_x[:m], gold_x[:m])),
+        "dit_out": (pcc_mod.pcc(y[:m], gold_y[:m]), pcc_mod.pcc(ref_y[:m], gold_y[:m])),
+    }
+    for tap, (mk_pcc, s1_pcc) in closeness.items():
+        results.add_row(
+            {
+                "tap": "mk_k4_step",
+                "key": f"{tap}-k{step}-mk-vs-s1-closeness-to-golden",
+                "rule": "pcc",
+                "pcc": mk_pcc,
+                "s1_pcc_vs_golden": s1_pcc,
+                "gate_pcc": s1_pcc - GOLDEN_SLACK,
+                "ok": bool(mk_pcc >= s1_pcc - GOLDEN_SLACK),
+            }
+        )
+    _meta(results, s)
+    assert row_x["finite"] and row_y["finite"], "non-finite step output"
+    assert det, "two launches of the same step differ"
+    assert report.ok, report.summary(failures_only=True)
+    assert row_x["pcc"] >= GATE_ORACLE, f"dit_block_last_out k={step}: PCC vs S1 {row_x['pcc']:.6f} < {GATE_ORACLE}"
+    assert row_y["pcc"] >= GATE_ORACLE_HEAD, f"dit_out k={step}: PCC vs S1 {row_y['pcc']:.6f} < {GATE_ORACLE_HEAD}"
+    for tap, (mk_pcc, s1_pcc) in closeness.items():
+        assert mk_pcc >= s1_pcc - GOLDEN_SLACK, f"{tap} k={step}: mk vs golden {mk_pcc:.6f} < S1 vs golden {s1_pcc:.6f}"
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# steps4 variant
+# --------------------------------------------------------------------------------------------------------------------
+def _steps4_inputs(s: K4Session) -> Tuple[torch.Tensor, torch.Tensor]:
+    noise = s.gs.load("initial_noise")
+    state = s.gs.load("state_features").reshape(-1)
+    return noise, state
+
+
+def _steps4_taps(s: K4Session, out: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Megakernel outputs -> golden-keyed host tensors (``taps_to_golden`` shapes)."""
+    lo, hi = E.action_rows(s.cfg, s.dl)
+    m = s.dl.m_logical
+    a_max = int(s.cfg.encoders.max_action_dim)
+    pred = E.xt_to_action_pred(out["action_pred"], s.cfg, s.dl)
+    taps = {"action_pred_normalized": pred, "action_pred_valid": pred}
+    v = out["velocity"]  # [4, 64, A_pad]
+    for k in STEPS:
+        taps[s.gs.threshold_key("pred_velocity", step=k)] = v[k, lo:hi, :a_max].reshape(1, hi - lo, a_max)
+        taps[s.gs.threshold_key("action_decoder_out", step=k)] = v[k, :m, :a_max].reshape(1, m, a_max)
+    return taps
+
+
+def test_steps4_vs_golden(results: Any, version: str, k4: K4Session) -> None:
+    """The 4-step denoise in one launch vs the golden action taps and the S1 ``ActionHeadTT`` on device."""
+    import ttnn
+    from models.experimental.gr00t.tt import action_head as AH
+    from models.experimental.gr00t.tt import layers as TL
+    from models.experimental.gr00t.tt.policy import TTPolicy
+
+    s = k4
+    noise, state = _steps4_inputs(s)
+    # S1 oracle first, with the megakernel's L1 residents released (a_rep 256 KB on all cores, the steps4 rows): the
+    # same inputs through ActionHeadTT (ttnn ops, same weights / K/V / masks, every intermediate kept as a tap).  Under
+    # the K4 worker_l1_size its layer_norm CBs (~980 KB on the op's 2 cores) do not fit next to a_rep.
+    s.mk.release_l1_residents()
+    head = AH.ActionHeadTT(
+        s.cfg,
+        s.dl,
+        s.W,
+        TTPolicy(dtype_policy=s.policy.dtype_policy),
+        s.device,
+        masks=s.masks,
+        kv=s.kv,
+        collect_taps=True,
+    )
+    try:
+        head.upload_noise(noise)
+        state_dev = TL.to_device(state.reshape(1, 1, 1, -1).to(torch.bfloat16), s.device, mem="DRAM")
+        head.denoise(head.x_t, state_dev)
+        ref = head.golden_taps(b3=_b3_bias(s), block_taps=False)
+        ttnn.deallocate(state_dev)
+        head.free_taps()
+    finally:
+        head.release()
+    s.mk.set_variant("steps4")
+    t0 = time.perf_counter()
+    out = s.mk.run_steps4(noise, state)
+    dt = time.perf_counter() - t0
+    taps = _steps4_taps(s, out)
+    rows = {}
+    for key in ("action_pred_normalized",) + tuple(s.gs.threshold_key("pred_velocity", step=k) for k in STEPS):
+        rows[key] = _row(results, f"{key}-vs-oracle", taps[key], ref[key], GATE_ORACLE, wall_s=dt)
+    report = check_taps(version, taps, s.gs, results=results)
+    _meta(results, s)
+    results.set_meta(steps4_wall_s=dt, oracle_min_pcc=min(r["pcc"] for r in rows.values()))
+    assert all(r["finite"] for r in rows.values()), "non-finite steps4 output"
+    assert report.ok, report.summary(failures_only=True)
+    bad = {k: r["pcc"] for k, r in rows.items() if r["pcc"] < GATE_ORACLE}
+    assert not bad, f"steps4 vs S1 below {GATE_ORACLE}: {bad}"
+
+
+def test_steps4_determinism_and_trace(results: Any, version: str, k4: K4Session) -> None:
+    """10 untraced launches bit-identical; the trace replay bit-identical to the untraced result; G2 timing."""
+    s = k4
+    s.mk.set_variant("steps4")
+    noise, state = _steps4_inputs(s)
+    outs: List[Dict[str, torch.Tensor]] = [s.mk.run_steps4(noise, state) for _ in range(N_DETERMINISM)]
+    identical = all(
+        torch.equal(o["action_pred"], outs[0]["action_pred"]) and torch.equal(o["velocity"], outs[0]["velocity"])
+        for o in outs[1:]
+    )
+    results.add_row({"tap": "mk_k4_step", "key": "steps4-10x-bit-identical", "rule": "exact", "ok": identical})
+    # trace: capture once, replay N_REPLAYS times (the inputs are the persistent tensors, refreshed before capture)
+    s.mk.set_steps4_inputs(noise, state)
+    tid = s.mk.capture_steps4()
+    try:
+        times = [s.mk.replay(tid) for _ in range(N_REPLAYS)]
+        traced = s.mk.read_steps4_outputs()
+    finally:
+        s.mk.release_trace(tid)
+    traced_ok = bool(torch.equal(traced["action_pred"], outs[0]["action_pred"]))
+    cmp_velocity = "velocity" in traced and "velocity" in outs[0]
+    if cmp_velocity:
+        traced_ok = traced_ok and bool(torch.equal(traced["velocity"], outs[0]["velocity"]))
+    results.add_row(
+        {
+            "tap": "mk_k4_step",
+            "key": "steps4-traced-eq-untraced",
+            "rule": "exact",
+            "ok": traced_ok,
+            "velocity_compared": cmp_velocity,
+        }
+    )
+    ms = sorted(t * 1e3 for t in times)
+    per_step_ms = statistics.median(ms) / P.STEPS4_N_STEPS
+    bytes_step = s.mk.arena_bytes_per_step()
+    floor_ms = bytes_step / MEASURED_GBPS[STREAM_DTYPE] / 1e6
+    gate_ms = G2_FACTOR * floor_ms
+    g2 = {
+        "replays": N_REPLAYS,
+        "trace_ms_min": ms[0],
+        "trace_ms_median": statistics.median(ms),
+        "trace_ms_max": ms[-1],
+        "per_step_ms_median": per_step_ms,
+        "per_step_ms_min": ms[0] / P.STEPS4_N_STEPS,
+        "bytes_per_step": bytes_step,
+        "measured_gbps": MEASURED_GBPS[STREAM_DTYPE],
+        "floor_ms_per_step": floor_ms,
+        "gate_ms_per_step": gate_ms,
+        "ratio_to_floor": per_step_ms / floor_ms,
+        "pass": per_step_ms <= gate_ms,
+    }
+    # recorded, not asserted (plan §5.7 fallback ratified 2026-09-18): the row's ok follows the gate so a miss is
+    # visible as ok: False / n_fail in the JSON while the determinism assertions below decide the test outcome
+    results.add_row({"tap": "mk_k4_step", "key": "G2-step-time", "rule": "timing", "ok": bool(g2["pass"]), **g2})
+    results.set_timing(steps4_trace_median_s=statistics.median(times), steps4_trace_min_s=min(times))
+    _meta(results, s)
+    results.set_meta(g2=g2)
+    assert identical, "10 steps4 launches are not bit-identical"
+    assert traced_ok, "the traced steps4 replay differs from the untraced launch"

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

@@ -24,17 +24,45 @@ and refreshed per call with :meth:`ActionHeadTT.update_masks`) and the hoisted K
 (``HoistOut.kv``).  Weights and AdaLN vector tables live in ``TTWeights`` (DRAM).
 
 ``backend`` (``TTPolicy.dit_backend``): ``"ttnn"`` runs the Stage-1 op sequence above; ``"megakernel"`` (Stage 2,
-one ``generic_op`` for the whole denoise, plan §5) raises ``NotImplementedError`` from :meth:`denoise` until WP-K5
-wires ``DiTMegakernel`` -- construction succeeds so that ``Gr00tTT`` can be assembled under either policy.
-
-Taps (``collect_taps=True``) are device tensors under the golden keys (``sa_embs[k=0]``, ``dit_block{i}_out[k=0]``,
-``dit_block_last_out[k=0]``, ``dit_out[k=0]``, ``action_decoder_out[k=0]``, ``actions_out[k=0]``,
-``action_pred_normalized``; plus ``action_encoder_pre_pos[k=0]``); :func:`taps_to_golden` slices them to the goldens'
-shapes for ``tests.tt.harness.check_taps``.  ``ttnn`` is imported lazily (constructor and forwards).
+plan §5, WP-K5) runs the whole denoise as **one** persistent ``ttnn.generic_op`` -- the ``steps4`` variant of
+``tt.megakernel.dit_program.DiTMegakernel`` (action encoder -> ``sa_embs`` -> blocks -> head -> decoder -> Euler,
+four steps in one launch; ``tests/tt/results/mk_k4_summary.md`` §3, §8).  Under that backend the constructor packs the
+DiT / precompute / encoder host plan tensors into the DRAM weight arena (``TTPolicy.mk_arena_dtype_for(version)``:
+bf16 or bfp8_b stream), uploads it once, binds the hoisted K/V handles and the mask buffers (the same persistent
+objects trace A / the executor write), allocates the kernel's io tensors once (rank-4 shapes so that ``ttnn.copy`` can
+feed them from the model's buffers) and builds the cached program descriptor.  Per call (inside trace B)::
+
+    ttnn.copy(x_t -> mk x_t fp32 [1,1,64,A_pad])      the version dtype -> the kernel's fp32 Euler state (TILE typecast)
+    ttnn.to_layout(state_features, ROW_MAJOR) -> copy  bf16 [1,1,1,1536] row the hub reads
+    ttnn.generic_op(tensors, program)                  encoder / 16-32 blocks / head / decoder / Euler x 4
+    ttnn.copy(mk action_pred fp32 -> action_pred)      back to the version dtype: the same buffer contract as "ttnn"
+
+so ``Gr00tTT``'s denoise stage, decode and the executor are unchanged.  The kernel's two L1 tensors -- ``a_rep``
+(the multicast target, 256 KB on all 110 cores) and the ``x_rep`` tile row (96 KB per compute core) -- are **not**
+kept resident: next to them the Stage-1 backbone programs (matmuls with ~1 MB of static CBs) do not fit under the K4
+``worker_l1_size``, so :meth:`ActionHeadTT._denoise_megakernel` allocates them (``ttnn.allocate_tensor_on_device``,
+no host write) right before the generic_op, rebuilds the program descriptor on their addresses (the addresses are
+common runtime args, ``descriptors._address_of``) and frees them after -- exactly like any traced ttnn op's
+intermediates, so trace B replays them at the captured addresses.  The device must be opened with a ``worker_l1_size``
+reduced by ``tt.model.MK_L1_CUT_BYTES_DEFAULT`` (**64 KiB**: the compute-core K4 binaries -- 108,912 B on N1.5 -- need a
+kernel-config ring buffer larger than the p150a default of 70,656 B, ``mk_k4_summary.md`` §7.2 fact 2;
+``tt.model.open_model_device`` does it, and ``Gr00tTT`` moves the VL adapter's intermediates to DRAM because no cut
+that fits the binaries leaves the N1.5 / N1.7 adapter programs room for their L1 activations; the K4 unit tests'
+own 96 KiB cut, ``dit_program.K4_WORKER_L1_SIZE``, is not what the model uses).  The megakernel is the
+``TTPolicy.dit_backend`` **default** since 2026-09-18 (``tt.policy.DEFAULT_DIT_BACKEND``).
+
+Taps (``collect_taps=True``): the ttnn backend keeps device tensors under the golden keys (``sa_embs[k=0]``,
+``dit_block{i}_out[k=0]``, ``dit_block_last_out[k=0]``, ``dit_out[k=0]``, ``action_decoder_out[k=0]``,
+``actions_out[k=0]``, ``action_pred_normalized``; plus ``action_encoder_pre_pos[k=0]``); the megakernel exposes the
+per-step decoder output (``action_decoder_out[k]`` from its ``velocity_out`` tap tensor, host copies) and
+``action_pred_normalized`` -- the residual-stream intermediates live inside the kernel (:func:`head_tap_names`).
+:func:`taps_to_golden` slices them to the goldens' shapes for ``tests.tt.harness.check_taps``.  ``ttnn`` is imported
+lazily (constructor and forwards).
 """
 
 from __future__ import annotations
 
+import time
 from dataclasses import dataclass
 from typing import Any, Dict, List, Mapping, Optional, Tuple
 
@@ -59,6 +87,40 @@ STEP_TAPS: Tuple[str, ...] = (
     "actions_out",
 )
 FINAL_TAPS: Tuple[str, ...] = ("action_pred_normalized",)
+#: per-step taps the megakernel backend produces (``velocity_out`` = the decoder output of every step; module docstring)
+MK_STEP_TAPS: Tuple[str, ...] = ("action_decoder_out",)
+#: golden-keyed taps derived by :func:`taps_to_golden` from a per-step tap (``<tap>[k=i]`` -> ``<derived>[k=i]``)
+DERIVED_STEP_TAPS: Dict[str, Tuple[str, ...]] = {
+    "action_encoder_pre_pos": ("action_encoder_out",),
+    "action_decoder_out": ("pred_velocity",),
+}
+#: the kernel's persistent io tensors :class:`ActionHeadTT` re-allocates with rank-4 shapes (module docstring)
+MK_IO_RANK4: Tuple[str, ...] = ("x_t", "state_features", "action_pred")
+
+
+def head_tap_names(backend: str, n_steps: int, n_blocks: int) -> Tuple[str, ...]:
+    """The golden-keyed tap names ``ActionHeadTT.golden_taps`` produces for ``backend`` in step order
+    (``tt.model.expected_tap_keys`` appends them; ``GoldenSet.threshold_key`` form).
+
+    * ``"ttnn"``: ``action_encoder_out[k]``, ``sa_embs[k]``, every ``dit_block{i}_out[k]``, ``dit_block_last_out[k]``,
+      ``dit_out[k]``, ``action_decoder_out[k]``, ``pred_velocity[k]``, ``actions_out[k]`` per step, then
+      ``action_pred_normalized`` / ``action_pred_valid``;
+    * ``"megakernel"``: ``action_decoder_out[k]`` / ``pred_velocity[k]`` per step, then ``action_pred_normalized`` /
+      ``action_pred_valid`` (the residual stream never leaves the kernel; ``mk_k4_summary.md`` §8 "Not needed").
+    """
+    if backend not in DIT_BACKENDS:
+        raise ValueError(f"unknown dit_backend {backend!r}; expected one of {DIT_BACKENDS}")
+    keys: List[str] = []
+    for k in range(int(n_steps)):
+        if backend == "ttnn":
+            keys += [f"action_encoder_out[k={k}]", f"sa_embs[k={k}]"]
+            keys += [f"dit_block{i}_out[k={k}]" for i in range(int(n_blocks))]
+            keys += [f"dit_block_last_out[k={k}]", f"dit_out[k={k}]"]
+        keys += [f"action_decoder_out[k={k}]", f"pred_velocity[k={k}]"]
+        if backend == "ttnn":
+            keys.append(f"actions_out[k={k}]")
+    keys += ["action_pred_normalized", "action_pred_valid"]
+    return tuple(keys)
 
 
 def _shape(t: Any) -> Tuple[int, ...]:
@@ -236,15 +298,17 @@ class ActionHeadTT:
     Args:
         cfg, dlayout: version config and the baked static layout of the call shape.
         W: ``tt.weights.TTWeights`` (categories ``dit``, ``dit_precompute``, ``encoders`` resident; ``constants``).
-        policy: ``TTPolicy`` -- ``dit_backend`` selects the Stage-1 op path or the megakernel (WP-K5).
-        device: the open mesh device (buffers are allocated here).
-        encoders: a prebuilt :class:`HeadEncoders` (default: built from ``W``).
+        policy: ``TTPolicy`` -- ``dit_backend`` selects the Stage-1 op path or the megakernel (WP-K5);
+            ``mk_arena_dtype_for(version)`` the megakernel's weight arena.
+        device: the open mesh device (buffers are allocated here; the megakernel needs ``tt.model.open_model_device``).
+        encoders: a prebuilt :class:`HeadEncoders` (default: built from ``W``; ttnn backend only).
         masks: the ``tt.layout.materialise`` dictionary shared with the model (default: own buffers from
             :func:`allocate_mask_buffers`, refreshed with :meth:`update_masks`).
         kv: block index -> persistent ``(K_i, V_i)`` (``VLAdapterTT`` ``HoistOut.kv``); or :meth:`bind_kv` later.
-        weights: name -> tensor overrides passed to the DiT modules (negative controls, tests).
+        weights: name -> tensor overrides passed to the DiT modules (negative controls, tests; ttnn backend only).
         persist_sa_embs: keep a persistent DRAM ``sa_embs`` buffer updated every step (one extra copy per step).
         collect_taps: keep the golden-named device taps in ``taps`` (never inside a trace capture).
+        mk_k4_budget: a ``descriptors.K4Budget`` for the megakernel's CB / L1 budget (default: the K4 defaults).
     """
 
     def __init__(
@@ -267,6 +331,7 @@ class ActionHeadTT:
         fuse_swish: bool = False,
         persist_sa_embs: bool = False,
         collect_taps: bool = False,
+        mk_k4_budget: Any = None,
     ) -> None:
         if not isinstance(cfg, GR00TConfig) or not isinstance(dlayout, DeviceLayout):
             raise TypeError("ActionHeadTT(cfg: GR00TConfig, dlayout: DeviceLayout, W, policy: TTPolicy, device)")
@@ -290,6 +355,7 @@ class ActionHeadTT:
         self.width = self.plan.width
         self.rows = E.action_rows(cfg, dlayout)
         self.buckets = tuple(int(b) for b in cfg.sampler.timestep_buckets)
+        self.timing: Dict[str, float] = {}
 
         # masks: shared dictionary or own buffers (materialise-compatible, refreshed in place)
         if masks is None:
@@ -309,18 +375,34 @@ class ActionHeadTT:
             self.owns_masks = False
         self.dit_masks = dit_masks(self.masks)
 
-        # modules
-        self.encoders = (
-            encoders if encoders is not None else HeadEncoders.build(cfg, W, dlayout, mem=mem, fuse_swish=fuse_swish)
-        )
-        self._check_encoders(self.encoders)
-        self.step = DiTStepTT(
-            cfg, W, dlayout, self.masks, kv, weights=weights, mem=mem, ckc=ckc, head_ckc=head_ckc, sdpa_ckc=sdpa_ckc
-        )
-        self.step.set_collect_taps(self.collect_taps)
-        self.encoders.set_collect_taps(self.collect_taps)
+        # modules (ttnn backend) / the megakernel launcher (megakernel backend)
+        self.encoders: Optional[HeadEncoders] = None
+        self.step: Optional[DiTStepTT] = None
         self._tap_pool: List[Any] = []  # encoder intermediates kept while taps are on (freed by free_taps)
-        self.megakernel: Any = None  # WP-K5 builds DiTMegakernel here for backend == "megakernel"
+        self.megakernel: Any = None  # tt.megakernel.dit_program.DiTMegakernel (megakernel backend)
+        self.arena: Any = None  # tt.megakernel.arena.ArenaLayout (megakernel backend)
+        self.mk_arena_dtype: Optional[str] = None
+        self._kv: Optional[Dict[int, Tuple[Any, Any]]] = None
+        if self.backend == "ttnn":
+            if mk_k4_budget is not None:
+                raise ValueError("mk_k4_budget is a megakernel-backend argument")
+            self.encoders = (
+                encoders
+                if encoders is not None
+                else HeadEncoders.build(cfg, W, dlayout, mem=mem, fuse_swish=fuse_swish)
+            )
+            self._check_encoders(self.encoders)
+            self.step = DiTStepTT(
+                cfg, W, dlayout, self.masks, kv, weights=weights, mem=mem, ckc=ckc, head_ckc=head_ckc, sdpa_ckc=sdpa_ckc
+            )
+            self.step.set_collect_taps(self.collect_taps)
+            self.encoders.set_collect_taps(self.collect_taps)
+        else:
+            if encoders is not None or weights is not None or persist_sa_embs:
+                raise ValueError(
+                    "encoders= / weights= / persist_sa_embs are Stage-1 (ttnn backend) arguments: the megakernel "
+                    "assembles sa_embs and runs the encoders / decoder / Euler in-kernel"
+                )
 
         # persistent buffers
         zeros_x = torch.zeros(1, 1, DIT_QUERY_ROWS_PAD, self.a_pad, dtype=E.TORCH_DTYPES[self.xt_dtype])
@@ -331,6 +413,8 @@ class ActionHeadTT:
             zeros_sa = torch.zeros(1, 1, DIT_QUERY_ROWS_PAD, self.width, dtype=torch.bfloat16)
             self.sa_embs = L.to_device(zeros_sa, device, dtype="bf16", mem=buffer_mem)
         self.n_calls = 0
+        if self.backend == "megakernel":
+            self._build_megakernel(kv, mk_k4_budget)
 
     # ------------------------------------------------------------------ checks
     def _check_encoders(self, enc: HeadEncoders) -> None:
@@ -351,14 +435,183 @@ class ActionHeadTT:
                 f"action encoder has {enc.action_encoder.n_steps} step biases, sampler has {self.n_steps} steps"
             )
 
+    # ------------------------------------------------------------------ megakernel build (WP-K5; mk_k4_summary.md §8)
+    def _build_megakernel(self, kv: Optional[Mapping[int, Any]], k4_budget: Any) -> None:
+        """Arena (packed from the same bf16 host plan ``TTWeights`` uploads) + upload, K/V and mask binding, the
+        persistent io tensors and the cached ``steps4`` program (built once K/V are bound)."""
+        import ttnn
+
+        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 CoreMap
+
+        cfg, dl = self.cfg, self.dlayout
+        self.mk_arena_dtype = self.policy.mk_arena_dtype_for(cfg.version)
+        t0 = time.perf_counter()
+        self.arena = A.ArenaLayout.plan(
+            cfg.version,
+            dtype_stream=self.mk_arena_dtype,
+            reader_mode="direct",  # the K4 steps4 program streams in direct mode (mk_k4_summary.md §8)
+            tile_order="row_major",
+            layout_name=dl.layout_name,
+        )
+        for name, got, want in (
+            ("m_pad", int(self.arena.m_pad), DIT_QUERY_ROWS_PAD),
+            ("a_pad", int(self.arena.a_pad), self.a_pad),
+            ("width", int(self.arena.width), self.width),
+            ("n_blocks", int(self.arena.n_blocks), int(cfg.dit.n_blocks)),
+            ("action_rows", tuple(self.arena.action_rows), tuple(self.rows)),
+        ):
+            if got != want:
+                raise ValueError(f"megakernel arena {name} {got} != head {want}")
+        plan = P.real_arena_plan(cfg.version, int(dl.embodiment_id), self.policy)
+        pack = self.arena.pack(plan)
+        t1 = time.perf_counter()
+        mk = P.DiTMegakernel(
+            cfg,
+            dl,
+            None,
+            self.arena,
+            CoreMap.from_device(self.device),
+            self.device,
+            variant="steps4",
+            k1_extras=False,
+            k4_budget=k4_budget,
+        )
+        mk.upload(pack)
+        mk.bind_masks(self.dit_masks)
+        t2 = time.perf_counter()
+        # the kernel's persistent io: rank-4 shapes (same DRAM pages: fp32 tile 4096 B, one bf16 row 3072 B) so that
+        # ttnn.copy (identical logical shape + layout, dtype conversion on TILE) feeds them from the model's buffers
+        io = mk.allocate_steps4_io(taps=True)
+        shapes = {
+            "x_t": ((1, 1, DIT_QUERY_ROWS_PAD, self.a_pad), ttnn.float32, ttnn.TILE_LAYOUT, torch.float32, 4096),
+            "state_features": (
+                (1, 1, 1, self.width),
+                ttnn.bfloat16,
+                ttnn.ROW_MAJOR_LAYOUT,
+                torch.bfloat16,
+                self.width * 2,
+            ),
+            "action_pred": (
+                (1, 1, DIT_QUERY_ROWS_PAD, self.a_pad),
+                ttnn.float32,
+                ttnn.TILE_LAYOUT,
+                torch.float32,
+                4096,
+            ),
+        }
+        for name in MK_IO_RANK4:
+            shape, dt, lay, tdt, page = shapes[name]
+            ttnn.deallocate(io[name])
+            t = ttnn.from_torch(
+                torch.zeros(*shape, dtype=tdt),
+                dtype=dt,
+                layout=lay,
+                device=self.device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            )
+            if int(t.buffer_page_size()) != page:
+                raise RuntimeError(f"megakernel io {name}: DRAM page {int(t.buffer_page_size())} B != {page}")
+            io[name] = t
+        if "velocity_out" not in io:
+            raise RuntimeError("allocate_steps4_io(taps=True) gave no velocity_out tensor")
+        # the upload-time a_rep is freed: the L1 residents live only inside the denoise (module docstring)
+        ttnn.deallocate(mk.io["a_rep"])
+        mk.io["a_rep"] = None
+        self.megakernel = mk
+        if kv is not None:
+            self.bind_kv(kv)
+        ttnn.synchronize_device(self.device)
+        self.timing.update(mk_arena_plan_pack_s=t1 - t0, mk_upload_s=t2 - t1, mk_io_s=time.perf_counter() - t2)
+
+    def _alloc_residents(self) -> Tuple[Any, Any]:
+        """Allocate the kernel's L1 tensors (``a_rep`` ``[110*64, 2048]`` HEIGHT_SHARDED on all cores = ``cb_a``; the
+        ``x_rep`` tile rows ``[96*32, width]`` on the compute cores = ``cb_x``) without a host write and hand them to
+        the launcher (``io["a_rep"]`` / ``k4_io["x_rep"]``; ``dit_program.DiTMegakernel._alloc_a_rep`` /
+        ``make_x_rows`` shapes).  Their contents are never read before the kernel writes them (the multicasts and the
+        ``LoadX`` of ``sa_embs`` fill them)."""
+        import ttnn
+
+        from models.experimental.gr00t.tt.megakernel import dit_program as P
+
+        mk = self.megakernel
+        crs = mk.core_map.to_core_range_sets()
+        a_cols = int(mk.budget.cb_a) // (P.M_ROWS * 2)
+
+        def alloc(rows: int, cols: int, shard_rows: int, grid: Any) -> Any:
+            mem = ttnn.MemoryConfig(
+                ttnn.TensorMemoryLayout.HEIGHT_SHARDED,
+                ttnn.BufferType.L1,
+                ttnn.ShardSpec(grid, [shard_rows, cols], ttnn.ShardOrientation.ROW_MAJOR),
+            )
+            return ttnn.allocate_tensor_on_device(
+                ttnn.Shape([rows, cols]), ttnn.bfloat16, ttnn.TILE_LAYOUT, self.device, mem
+            )
+
+        a_rep = alloc(mk.core_map.n_cores * P.M_ROWS, a_cols, P.M_ROWS, crs["all"])
+        x_rep = alloc(mk.core_map.n_compute * P.TILE, self.width, P.TILE, crs["compute"])
+        mk.io["a_rep"] = a_rep
+        mk.k4_io["x_rep"] = x_rep
+        mk.k4_program = None  # the descriptor carries the addresses: rebuild on these tensors
+        mk.k4_tensors = []
+        return a_rep, x_rep
+
+    def _free_residents(self) -> None:
+        import ttnn
+
+        mk = self.megakernel
+        for holder, name in ((mk.io, "a_rep"), (mk.k4_io, "x_rep")):
+            t = holder.get(name)
+            if t is not None:
+                ttnn.deallocate(t)
+                holder[name] = None
+        mk.k4_program = None
+        mk.k4_tensors = []
+
+    @property
+    def mk_io(self) -> Dict[str, Any]:
+        """The megakernel's persistent io tensors (``x_t`` fp32, ``state_features`` bf16 row, ``action_pred`` fp32,
+        ``velocity_out`` bf16 ``[4*64, A_pad]``, ``x_rep``)."""
+        if self.megakernel is None:
+            raise RuntimeError("no megakernel: this head runs the ttnn backend")
+        return self.megakernel.k4_io
+
+    def mk_program(self) -> Tuple[Any, List[Any]]:
+        """The ``(ProgramDescriptor, io tensors)`` of the ``steps4`` generic_op on the currently allocated L1
+        residents (:meth:`_alloc_residents` first; what trace B captures)."""
+        if self.megakernel is None:
+            raise RuntimeError("no megakernel: this head runs the ttnn backend")
+        if self._kv is None:
+            raise ValueError("no hoisted K/V bound (VLAdapterTT HoistOut.kv): pass kv= or call bind_kv first")
+        mk = self.megakernel
+        if mk.io.get("a_rep") is None or mk.k4_io.get("x_rep") is None:
+            raise RuntimeError("the megakernel's L1 residents are not allocated (_alloc_residents)")
+        t0 = time.perf_counter()
+        out = mk.steps4_program()
+        self.timing["mk_program_build_s"] = time.perf_counter() - t0
+        return out
+
     # ------------------------------------------------------------------ per-call state
     @property
     def kv(self) -> Optional[Dict[int, Tuple[Any, Any]]]:
-        return self.step.kv
+        return self.step.kv if self.step is not None else self._kv
 
     def bind_kv(self, kv: Mapping[int, Any]) -> Dict[int, Tuple[Any, Any]]:
-        """Bind the persistent hoisted K/V (``VLAdapterTT`` ``HoistOut.kv``) once; the objects are what trace B captures."""
-        return self.step.bind_kv(kv)
+        """Bind the persistent hoisted K/V (``VLAdapterTT`` ``HoistOut.kv``) once; the objects are what trace B captures.
+        Under the megakernel the program descriptor is (re)built on the bound handles."""
+        if self.step is not None:
+            return self.step.bind_kv(kv)
+        bound = {int(i): (k, v) for i, (k, v) in kv.items()}
+        self.megakernel.bind_kv(bound)
+        self._kv = bound
+        # validate the whole descriptor once (limits, CTAs, runtime args) on transient residents
+        self._alloc_residents()
+        try:
+            self.mk_program()
+        finally:
+            self._free_residents()
+        return bound
 
     def update_masks(self, slots: SlotLayout) -> Dict[str, Any]:
         """Write this call's per-call masks (``cross_text`` / ``cross_all`` / ``llm`` / ``vl``) into the persistent
@@ -378,8 +631,10 @@ class ActionHeadTT:
     def set_collect_taps(self, flag: bool) -> None:
         """Toggle tap collection on the head and its DiT step (off before a trace capture)."""
         self.collect_taps = bool(flag)
-        self.step.set_collect_taps(flag)
-        self.encoders.set_collect_taps(flag)
+        if self.step is not None:
+            self.step.set_collect_taps(flag)
+        if self.encoders is not None:
+            self.encoders.set_collect_taps(flag)
         if not flag:
             self.taps = {}
             self._tap_pool = []
@@ -403,16 +658,7 @@ class ActionHeadTT:
             for m in modules:
                 self._tap_pool.extend(m.debug_taps.values())
 
-    def denoise(self, x_t: Any, state_features: Any, kv: Optional[Mapping[int, Any]] = None) -> Any:
-        """Run the 4 unrolled steps from ``x_t`` (``[1,1,64,A_pad]`` in the version dtype; the persistent buffer or a
-        tensor that is copied into it) with ``state_features [1,1,1,1536]`` bf16 (``VLAdapterTT`` ``HoistOut``) and
-        the hoisted ``kv`` (default: the bound handles) -> the persistent ``action_pred`` buffer."""
-        import ttnn
-
-        if self.backend != "ttnn":
-            raise NotImplementedError(
-                f"ActionHeadTT backend {self.backend!r}: the megakernel denoise (DiTMegakernel.run) is wired by WP-K5"
-            )
+    def _check_inputs(self, x_t: Any, state_features: Any) -> None:
         want = (1, 1, DIT_QUERY_ROWS_PAD, self.a_pad)
         if _shape(x_t) != want:
             raise ValueError(f"x_t must be {want}, got {_shape(x_t)}")
@@ -420,15 +666,31 @@ class ActionHeadTT:
             raise TypeError(f"x_t must be {self.xt_dtype} (cfg.sampler.euler_dtype), got {L.dtype_name(x_t.dtype)}")
         if _shape(state_features) != (1, 1, 1, self.width):
             raise ValueError(f"state_features must be (1, 1, 1, {self.width}), got {_shape(state_features)}")
-        if kv is not None and kv is not self.step.kv:
-            self.step.bind_kv(kv)
-        if self.step.kv is None:
+
+    def denoise(self, x_t: Any, state_features: Any, kv: Optional[Mapping[int, Any]] = None) -> Any:
+        """Run the 4 unrolled steps from ``x_t`` (``[1,1,64,A_pad]`` in the version dtype; the persistent buffer or a
+        tensor that is copied into it) with ``state_features [1,1,1,1536]`` bf16 (``VLAdapterTT`` ``HoistOut``) and
+        the hoisted ``kv`` (default: the bound handles) -> the persistent ``action_pred`` buffer.  Both backends return
+        the same buffer (module docstring)."""
+        self._check_inputs(x_t, state_features)
+        if self.backend == "megakernel":
+            return self._denoise_megakernel(x_t, state_features, kv)
+        return self._denoise_ttnn(x_t, state_features, kv)
+
+    def _denoise_ttnn(self, x_t: Any, state_features: Any, kv: Optional[Mapping[int, Any]]) -> Any:
+        import ttnn
+
+        step, enc = self.step, self.encoders
+        if step is None or enc is None:
+            raise RuntimeError("ttnn backend without its modules")
+        if kv is not None and kv is not step.kv:
+            step.bind_kv(kv)
+        if step.kv is None:
             raise ValueError("no hoisted K/V bound (VLAdapterTT HoistOut.kv): pass kv= or call bind_kv first")
         if x_t is not self.x_t:
             ttnn.copy(x_t, self.x_t)
         self.taps = {}
         self._tap_pool = []
-        enc = self.encoders
         x = self.x_t
         for k in range(self.n_steps):
             a = enc.action_encoder.forward(x, k)
@@ -446,8 +708,8 @@ class ActionHeadTT:
                     ttnn.deallocate(sa)
                 sa = self.sa_embs
             self._tap(f"sa_embs[k={k}]", sa)
-            y = self.step.forward(sa, k)
-            for name, t in self.step.debug_taps.items():
+            y = step.forward(sa, k)
+            for name, t in step.debug_taps.items():
                 self._tap(f"{name}[k={k}]", t)
             if sa is not self.sa_embs:
                 self._free(sa)
@@ -468,6 +730,47 @@ class ActionHeadTT:
         self.n_calls += 1
         return self.action_pred
 
+    def _denoise_megakernel(self, x_t: Any, state_features: Any, kv: Optional[Mapping[int, Any]]) -> Any:
+        """The module docstring's four-op sequence around the ``steps4`` generic_op (all on persistent tensors).
+
+        The program descriptor is rebuilt on every *untraced* call (``mk_program`` -> ``steps4_program()`` on the
+        fresh L1 resident addresses, ~1.4 ms host per ``head_timing.mk_program_build_s`` of the K5 benches); inside
+        a trace replay nothing runs on the host, so the cost never reaches the production path (WP-K5 review NB 8).
+        """
+        import ttnn
+
+        if kv is not None and kv is not self._kv:
+            bound = {int(i): (k, v) for i, (k, v) in kv.items()}
+            if self._kv is None or any(
+                bound[i][0] is not self._kv[i][0] or bound[i][1] is not self._kv[i][1] for i in bound
+            ):
+                self.bind_kv(bound)
+        io = self.mk_io
+        self.taps = {}
+        # 1. the version-dtype x_t -> the kernel's fp32 Euler state (a TILE copy; bf16 -> fp32 exact)
+        ttnn.copy(x_t, io["x_t"])
+        # 2. state_features [1,1,1,1536] bf16 TILE -> the one ROW_MAJOR row the hub reads
+        sf_rm = ttnn.to_layout(state_features, ttnn.ROW_MAJOR_LAYOUT)
+        ttnn.copy(sf_rm, io["state_features"])
+        ttnn.deallocate(sf_rm)
+        # 3. the denoise on freshly allocated L1 residents (module docstring)
+        self._alloc_residents()
+        try:
+            program, tensors = self.mk_program()
+            ttnn.generic_op(tensors, program)
+        finally:
+            self._free_residents()
+        # 4. fp32 action_pred -> the version dtype buffer (N1.6 / N1.7 values are bf16-rounded per step in-kernel)
+        ttnn.copy(io["action_pred"], self.action_pred)
+        if self.collect_taps:
+            vel = L.to_torch(io["velocity_out"])  # [4*64, A_pad] fp32 host (bf16 values)
+            for k in range(self.n_steps):
+                rows = vel[k * DIT_QUERY_ROWS_PAD : (k + 1) * DIT_QUERY_ROWS_PAD]
+                self._tap(f"action_decoder_out[k={k}]", rows.reshape(1, 1, DIT_QUERY_ROWS_PAD, self.a_pad).contiguous())
+            self._tap("action_pred_normalized", self.action_pred)
+        self.n_calls += 1
+        return self.action_pred
+
     __call__ = denoise
 
     # ------------------------------------------------------------------ host side
@@ -479,26 +782,44 @@ class ActionHeadTT:
         """:func:`taps_to_golden` of the last ``denoise`` call's taps."""
         return taps_to_golden(self.taps, self.cfg, self.dlayout, b3=b3, block_taps=block_taps)
 
+    def velocity_taps_host(self) -> torch.Tensor:
+        """Megakernel backend: the last launch's per-step decoder outputs ``[n_steps, 64, A_pad]`` fp32 (bf16 values;
+        the kernel's ``velocity_out`` tensor, written on every launch, traced or not)."""
+        vel = L.to_torch(self.mk_io["velocity_out"])
+        return vel.reshape(self.n_steps, DIT_QUERY_ROWS_PAD, self.a_pad)
+
     def free_taps(self) -> None:
-        """Deallocate every collected tap that is not a persistent buffer (each object once)."""
+        """Deallocate every collected device tap that is not a persistent buffer (each object once; host copies skipped)."""
         import ttnn
 
         keep = {id(self.x_t), id(self.action_pred), id(self.sa_embs)}
         seen: set = set()
-        for t in list(self.taps.values()) + list(self.step.debug_taps.values()) + self._tap_pool:
-            if t is None or id(t) in seen or id(t) in keep:
+        step_taps = list(self.step.debug_taps.values()) if self.step is not None else []
+        for t in list(self.taps.values()) + step_taps + self._tap_pool:
+            if t is None or isinstance(t, torch.Tensor) or id(t) in seen or id(t) in keep:
                 continue
             seen.add(id(t))
             ttnn.deallocate(t)
         self.taps = {}
         self._tap_pool = []
-        self.step.debug_taps = {}
-        self.encoders.set_collect_taps(self.collect_taps)  # clears the modules' debug_taps dicts
+        if self.step is not None:
+            self.step.debug_taps = {}
+        if self.encoders is not None:
+            self.encoders.set_collect_taps(self.collect_taps)  # clears the modules' debug_taps dicts
 
     def release(self) -> None:
-        """Deallocate the persistent buffers this head owns (masks only when it allocated them; never the K/V)."""
+        """Deallocate the persistent buffers this head owns (masks only when it allocated them; never the K/V) and,
+        under the megakernel backend, the arena, tables, io tensors and L1 residents of the launcher."""
         import ttnn
 
+        if self.megakernel is not None:
+            mk = self.megakernel
+            # the bound K/V and mask handles belong to the adapter / executor: unhook them before the launcher frees its io
+            for name in list(mk.io):
+                if name.startswith("kv_") or name.endswith("_mask"):
+                    mk.io[name] = None
+            mk.deallocate()
+            self.megakernel = None
         for t in (self.x_t, self.action_pred, self.sa_embs):
             if t is not None:
                 ttnn.deallocate(t)
@@ -509,7 +830,7 @@ class ActionHeadTT:
         self.masks = {}
 
     def describe(self) -> Dict[str, Any]:
-        d = self.step.describe()
+        d: Dict[str, Any] = self.step.describe() if self.step is not None else self.plan.describe()
         d.update(
             {
                 "backend": self.backend,
@@ -518,19 +839,39 @@ class ActionHeadTT:
                 "buckets": list(self.buckets),
                 "persist_sa_embs": self.sa_embs is not None,
                 "owns_masks": self.owns_masks,
-                "fuse_swish": self.encoders.action_encoder.fuse_swish,
+                "fuse_swish": self.encoders.action_encoder.fuse_swish if self.encoders is not None else None,
                 "buffer_mem": getattr(self.buffer_mem, "kind", str(self.buffer_mem)),
+                "timing": dict(self.timing),
             }
         )
+        if self.megakernel is not None:
+            mk = self.megakernel
+            st = mk.stats() if mk.last_spec is not None else None
+            d["megakernel"] = {
+                "arena_dtype": self.mk_arena_dtype,
+                "reader_mode": self.arena.reader_mode,
+                "n_blocks": int(self.arena.n_blocks),
+                "key_sets": sorted(set(self.arena.key_sets)),
+                "arena_bytes_per_step": int(mk.arena_bytes_per_step()),
+                "cb_w_slots": dict(mk.k4.cb_w_slots),
+                "stats": st.__dict__ if st is not None else None,
+                "compile_s": {str(k): v for k, v in mk.compile_seconds.items()},
+                "io": {k: (list(int(x) for x in v.shape) if v is not None else None) for k, v in mk.k4_io.items()},
+                "launches": int(mk.launches),
+            }
         return d
 
 
 __all__ = [
+    "DERIVED_STEP_TAPS",
     "FINAL_TAPS",
+    "MK_IO_RANK4",
+    "MK_STEP_TAPS",
     "STEP_TAPS",
     "ActionHeadTT",
     "HeadEncoders",
     "allocate_mask_buffers",
+    "head_tap_names",
     "placeholder_masks",
     "taps_to_golden",
 ]

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

@@ -141,20 +141,28 @@ def open_gr00t_device(
     l1_small_size: int = DEFAULT_L1_SMALL_SIZE,
     num_command_queues: int = 1,
     device_id: int = 0,
+    worker_l1_size: Optional[int] = None,
 ) -> Any:
     """Open device ``device_id`` with the program cache enabled and assert the p150a facts this plan relies on.
 
-    ``ttnn.open_device(device_id, l1_small_size, trace_region_size[, num_command_queues])``
+    ``ttnn.open_device(device_id, l1_small_size, trace_region_size[, num_command_queues][, worker_l1_size])``
     (``ttnn-nanobind/device.cpp:71-80``, mb2 §5) returns the 1x1 ``MeshDevice``; the compute grid must be 11x10 and
     the DRAM grid 8x1 (mb1 Table 5a) -- every core map and every program config in ``tt.layers`` is sized for that
-    card, so anything else raises instead of silently running a wrong configuration.  Callers hold the device lock
-    (``bin/with-device.sh``).
+    card, so anything else raises instead of silently running a wrong configuration.  ``worker_l1_size`` (bytes of
+    allocatable L1 per worker core; ``None`` = the firmware default, 1,461,248 B on the p150a) is what
+    ``tt.model.open_model_device`` reduces for the megakernel DiT backend (the kernel-config ring buffer grows by the
+    same amount, ``mk_k4_summary.md`` §7.2 fact 2); the value actually in force is checked there.  Callers hold the
+    device lock (``bin/with-device.sh``).
     """
     import ttnn
 
     kwargs: Dict[str, Any] = dict(device_id=device_id, l1_small_size=l1_small_size, trace_region_size=trace_region_size)
     if num_command_queues != 1:
-        kwargs["num_command_queues"] = num_command_queues
+        kwargs["num_command_queues"] = int(num_command_queues)
+    if worker_l1_size is not None:
+        if int(worker_l1_size) <= 0:
+            raise ValueError(f"worker_l1_size must be a positive number of bytes, got {worker_l1_size}")
+        kwargs["worker_l1_size"] = int(worker_l1_size)
     device = ttnn.open_device(**kwargs)
     try:
         device.enable_program_cache()

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

@@ -27,7 +27,19 @@ layout groups them into ``A`` / ``B``)::
                                               side effect: hoisted K_i / V_i in the adapter's persistent buffers
     denoise  (B)  ActionHeadTT.denoise        inputs: x_t, cross_text_mask | cross_all_mask; constants self_mask,
                                               cross_image_mask (when L_img_pad != L_img_real)  -> action_pred (a clone
-                                              of the head's persistent buffer, so the executor owns the trace output)
+                                              of the head's persistent buffer, so the executor owns the trace output).
+                                              ``policy.dit_backend``: the Stage-1 ttnn ops or the ``steps4``
+                                              megakernel ``generic_op`` (WP-K5) -- same inputs, same output buffer.
+
+The megakernel backend -- the ``TTPolicy`` **default** since 2026-09-18 (``tt.policy.DEFAULT_DIT_BACKEND``; the ttnn
+Stage-1 path stays selectable through ``TTPolicy(dit_backend="ttnn")`` / ``GR00T_DIT_BACKEND=ttnn``) -- needs the
+device opened with a smaller ``worker_l1_size`` (the kernel-config ring buffer of ``mk_k4_summary.md`` §7.2 fact 2;
+:data:`MK_L1_CUT_BYTES_DEFAULT`): :func:`open_model_device` does that from the policy and
+``from_pretrained(device=None)`` uses it; a device that was opened another way is refused at build
+(:func:`device_worker_l1_size`).  Under the megakernel the ttnn DiT weights are not uploaded
+(:func:`megakernel_weight_names`: the arena replaces the ``dit`` / ``dit_precompute`` / ``encoders`` tensors except
+the adapter's ``dit.block.{i}.Wkv`` / ``bkv``, ``enc.state.*`` and ``enc.action.b3_pos`` that ``TTWeights.constants``
+reads); ``from_pretrained(load_ttnn_dit_weights=True)`` keeps the full Stage-1 upload.
 
 Every per-call tensor -- the masks included -- is an executor buffer (static-shape §9 / §10: written unconditionally
 before every replay, address-checked, presence fixed per layout), so ``ActionHeadTT`` is built on the executor's
@@ -69,7 +81,7 @@ import torch
 
 from models.experimental.gr00t.common.configs import DIT_QUERY_ROWS_PAD, GR00TConfig, LayoutConfig, get_config
 from models.experimental.gr00t.tt import layers as L
-from models.experimental.gr00t.tt.action_head import ActionHeadTT
+from models.experimental.gr00t.tt.action_head import ActionHeadTT, head_tap_names
 from models.experimental.gr00t.tt.adapter import STATE_ROWS, HoistOut, VLAdapterTT, slots_to_rows, state_rows_host
 from models.experimental.gr00t.tt.backbone import (
     IN_COS,
@@ -86,7 +98,15 @@ from models.experimental.gr00t.tt.backbone import (
     BackboneRecipe,
     BackboneTT,
 )
-from models.experimental.gr00t.tt.device import DRAM, close_gr00t_device, open_gr00t_device, resolve_mem
+from models.experimental.gr00t.tt.device import (
+    DEFAULT_L1_SMALL_SIZE,
+    DEFAULT_TRACE_REGION_SIZE,
+    DRAM,
+    L1,
+    close_gr00t_device,
+    open_gr00t_device,
+    resolve_mem,
+)
 from models.experimental.gr00t.tt.encoders import noise_to_xt, xt_dtype, xt_to_action_pred
 from models.experimental.gr00t.tt.layout import MASK_NAMES, DeviceLayout, StaticShapeError
 from models.experimental.gr00t.tt.policy import TTPolicy
@@ -203,6 +223,7 @@ class ModelRecipe:
     trace_layout: str
     dit_backend: str
     dtype_policy: str
+    mk_arena_dtype: Optional[str] = None  # the megakernel's weight arena (resolved per version); None for "ttnn"
 
     @classmethod
     def from_config(cls, cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy) -> "ModelRecipe":
@@ -239,6 +260,7 @@ class ModelRecipe:
             trace_layout=policy.trace_layout,
             dit_backend=policy.dit_backend,
             dtype_policy=policy.dtype_policy,
+            mk_arena_dtype=policy.mk_arena_dtype_for(cfg.version) if policy.dit_backend == "megakernel" else None,
         )
 
     @property
@@ -308,10 +330,14 @@ def model_input_specs(
     return specs
 
 
-def expected_tap_keys(cfg: GR00TConfig, dlayout: DeviceLayout) -> Tuple[str, ...]:
+def expected_tap_keys(cfg: GR00TConfig, dlayout: DeviceLayout, dit_backend: str = "ttnn") -> Tuple[str, ...]:
     """The golden-named keys ``run_untraced(return_taps=True)`` produces for this version (plan §7.2 matrix; keys in
-    ``GoldenSet.threshold_key`` form).  Per-block DiT taps cover every block; the goldens hold blocks 0 / 1 / last."""
-    r = ModelRecipe.from_config(cfg, dlayout, TTPolicy())
+    ``GoldenSet.threshold_key`` form).  Per-block DiT taps cover every block; the goldens hold blocks 0 / 1 / last.
+    ``dit_backend`` selects the head's tap set (``tt.action_head.head_tap_names``): the megakernel exposes only the
+    per-step decoder output and the final prediction (WP-K5).  The parameter default stays ``"ttnn"`` on purpose:
+    it is the **full** Stage-1 set (the superset the golden tap map is checked against on CPU); model-level callers
+    pass ``model.policy.dit_backend``."""
+    r = ModelRecipe.from_config(cfg, dlayout, TTPolicy(dit_backend=dit_backend))
     keys: List[str] = ["vit_embeddings", "vit_block_first", "vit_block_last"]
     if r.tower == "siglip":
         keys.append("vit_post_ln")
@@ -325,22 +351,161 @@ def expected_tap_keys(cfg: GR00TConfig, dlayout: DeviceLayout) -> Tuple[str, ...
     if r.has_vlsa:
         keys.append("vl_self_attention_out")
     keys += ["dit_encoder_hidden_states", "state_features"]
-    for k in range(r.n_steps):
-        keys += [f"action_encoder_out[k={k}]", f"sa_embs[k={k}]"]
-        keys += [f"dit_block{i}_out[k={k}]" for i in range(r.n_blocks)]
-        keys += [
-            f"dit_block_last_out[k={k}]",
-            f"dit_out[k={k}]",
-            f"action_decoder_out[k={k}]",
-            f"pred_velocity[k={k}]",
-            f"actions_out[k={k}]",
-        ]
-    keys += ["action_pred_normalized", "action_pred_valid"]
+    keys += list(head_tap_names(r.dit_backend, r.n_steps, r.n_blocks))
     if len(set(keys)) != len(keys):
         raise AssertionError("duplicate tap keys")
     return tuple(keys)
 
 
+# --------------------------------------------------------------------------------------------------------------------
+# device opening per policy (WP-K5: the megakernel's worker_l1_size)
+# --------------------------------------------------------------------------------------------------------------------
+
+
+#: Bytes taken off every core's allocatable L1 (``worker_l1_size``) for the megakernel backend: the kernel-config
+#: ring buffer grows by the same amount (``mk_k4_summary.md`` §7.2 fact 2).  Measured compute-core K4 program sizes
+#: (``program.cpp:3224``): N1.6 91,168 B, N1.5 108,912 B (its future-token / mask paths) against the p150a default
+#: ring of 70,656 B, so >= 38,256 B are needed; 64 KiB leaves 27 KB of headroom (K4's own tests use 96 KiB,
+#: ``dit_program.K4_WORKER_L1_SIZE``).  Any cut >= 32 KiB breaks the N1.5 / N1.7 **adapter** stage with L1
+#: intermediates: its VL-SA programs (S_pad 544 / 480) allocate static CBs up to 1,405,952 B and the stage's L1
+#: activations extend down to <= 1,339,392 B (``program.cpp:2149``; WP-K5 bring-up sweeps at 96 / 48 / 40 KiB), i.e.
+#: the two requirements leave no common value -- hence :data:`ADAPTER_MEM_BY_BACKEND` moves the adapter's
+#: intermediates to DRAM under the megakernel backend.  ``GR00T_MK_L1_CUT_KIB`` overrides the value (sweeps).
+MK_L1_CUT_BYTES_DEFAULT: int = 64 * 1024
+#: Memory of the ``VLAdapterTT`` intermediates per DiT backend (``VLAdapterTT(mem=)``; the persistent K/V stay in
+#: DRAM either way).  ``"DRAM"`` under the megakernel frees the ~230 KB of per-core L1 the VL-SA stage needs next to
+#: its static CBs (comment above); the values are identical (same ops, same dtypes), only the placement changes.
+ADAPTER_MEM_BY_BACKEND: Dict[str, str] = {"ttnn": "L1", "megakernel": "DRAM"}
+MK_L1_CUT_ENV = "GR00T_MK_L1_CUT_KIB"
+
+
+def mk_l1_cut_bytes() -> int:
+    """:data:`MK_L1_CUT_BYTES_DEFAULT` or the :data:`MK_L1_CUT_ENV` override (KiB, positive integer)."""
+    import os
+
+    raw = os.environ.get(MK_L1_CUT_ENV)
+    if not raw:
+        return MK_L1_CUT_BYTES_DEFAULT
+    kib = int(raw)
+    if kib <= 0:
+        raise ValueError(f"{MK_L1_CUT_ENV}={raw!r} must be a positive number of KiB")
+    return kib * 1024
+
+
+def worker_l1_size_for(policy: TTPolicy) -> Optional[int]:
+    """The ``worker_l1_size`` the device must be opened with for ``policy``: ``descriptors.L1_USABLE_BYTES`` (the
+    p150a default, 1,461,248 B) minus :func:`mk_l1_cut_bytes` for the megakernel backend (the K4 binaries need the
+    larger kernel-config ring buffer that the cut buys, module constant above), ``None`` (the ttnn default)
+    otherwise."""
+    if not isinstance(policy, TTPolicy):
+        raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}")
+    if policy.dit_backend != "megakernel":
+        return None
+    from models.experimental.gr00t.tt.megakernel.descriptors import L1_USABLE_BYTES
+
+    return int(L1_USABLE_BYTES) - mk_l1_cut_bytes()
+
+
+def device_worker_l1_size(device: Any) -> int:
+    """The user-allocatable L1 per worker core the device was opened with: ``ttnn.reports.get_device_info(device)``
+    ``.cb_limit`` (= the ``worker_l1_size`` of ``open_device``; the ``.worker_l1_size`` field there is the physical
+    1.5 MiB).  Measured on the p150a: 1,461,248 B by default, 1,362,944 B under ``K4_WORKER_L1_SIZE``."""
+    import ttnn
+
+    return int(ttnn._ttnn.reports.get_device_info(device).cb_limit)
+
+
+def open_model_device(
+    policy: TTPolicy,
+    *,
+    trace_region_size: int = DEFAULT_TRACE_REGION_SIZE,
+    l1_small_size: int = DEFAULT_L1_SMALL_SIZE,
+    num_command_queues: int = 1,
+    device_id: int = 0,
+) -> Any:
+    """``tt.device.open_gr00t_device`` (program cache, p150a grid asserts) with the ``worker_l1_size`` the policy's
+    DiT backend needs: the firmware default for the ttnn backend, :func:`worker_l1_size_for` (the 64 KiB cut) for the
+    megakernel -- checked against :func:`device_worker_l1_size` after the open.  ``num_command_queues=2`` is the
+    production path (``Gr00tTT.from_pretrained`` default; CQ-1 input writes overlap the traces).  Callers hold the
+    device lock (``bin/with-device.sh``)."""
+    wl1 = worker_l1_size_for(policy)
+    device = open_gr00t_device(
+        trace_region_size=trace_region_size,
+        l1_small_size=l1_small_size,
+        num_command_queues=num_command_queues,
+        device_id=device_id,
+        worker_l1_size=wl1,
+    )
+    if wl1 is not None:
+        got = device_worker_l1_size(device)
+        if got != wl1:
+            close_gr00t_device(device)
+            raise RuntimeError(f"device opened with worker_l1_size {got}, requested {wl1} (megakernel backend)")
+    return device
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# ttnn weights the megakernel backend does not need (startup cost, WP-K5 review NB 7)
+# --------------------------------------------------------------------------------------------------------------------
+
+
+#: Plan categories the Stage-1 head reads from ``TTWeights`` and the megakernel arena replaces
+#: (``tt.megakernel.dit_program.real_arena_plan`` packs the same host plan tensors into the DRAM arena).
+MK_ARENA_CATEGORIES: Tuple[str, ...] = ("dit", "dit_precompute", "encoders")
+
+
+def megakernel_needs_ttnn_tensor(name: str, category: str) -> bool:
+    """Whether plan tensor ``name`` of ``category`` must stay in ``TTWeights`` under the megakernel backend: every
+    tensor outside :data:`MK_ARENA_CATEGORIES`, plus the ones the *other* modules read from those categories --
+    ``dit.block.{i}.Wkv`` / ``bkv`` (``VLAdapterTT`` hoists the cross K/V with them), ``enc.state.*`` (the adapter's
+    state encoder) and ``enc.action.b3_pos`` (``TTWeights.constants`` builds ``b3_pos_rows`` from the resident
+    tensor)."""
+    if category not in MK_ARENA_CATEGORIES:
+        return True
+    if name.startswith("dit.block.") and (name.endswith(".Wkv") or name.endswith(".bkv")):
+        return True
+    if name.startswith("enc.state."):
+        return True
+    return name == "enc.action.b3_pos"
+
+
+def megakernel_weight_names(
+    version: str, embodiment_id: int, policy: TTPolicy, layout_name: Optional[str] = None
+) -> Tuple[Tuple[str, ...], Dict[str, Any]]:
+    """``(names to upload, facts)`` for ``TTWeights.load(names=...)`` under the megakernel backend: the full plan of
+    ``(version, embodiment, policy)`` minus the tensors :func:`megakernel_needs_ttnn_tensor` rejects.  ``facts``
+    records what is skipped (``n_tensors``, ``device_bytes``, per category) for ``Gr00tTT.describe``."""
+    from models.experimental.gr00t.tt.policy import plan_options
+    from models.experimental.gr00t.tt.weights import upload_plan
+
+    entries = upload_plan(version, int(embodiment_id), policy, plan_options(policy, version, layout_name))
+    keep: List[str] = []
+    skipped: Dict[str, Dict[str, float]] = {}
+    n_skip, bytes_skip = 0, 0.0
+    for name, e in entries.items():
+        if megakernel_needs_ttnn_tensor(name, e.category):
+            keep.append(name)
+            continue
+        c = skipped.setdefault(e.category, {"n_tensors": 0, "device_bytes": 0.0})
+        c["n_tensors"] += 1
+        c["device_bytes"] += float(e.nbytes)
+        n_skip += 1
+        bytes_skip += float(e.nbytes)
+    if not keep:
+        raise AssertionError("megakernel_weight_names selected no plan tensor")
+    facts = {
+        "n_tensors": n_skip,
+        "device_bytes": bytes_skip,
+        "categories": skipped,
+        "n_kept": len(keep),
+        "n_full": len(entries),
+        "kept_from_arena_categories": sorted(
+            n for n in keep if entries[n].category in MK_ARENA_CATEGORIES and not n.startswith("dit.block.")
+        ),
+    }
+    return tuple(keep), facts
+
+
 # --------------------------------------------------------------------------------------------------------------------
 # bench / serve replayer contract (benchmarks/bench_e2e.py REPLAYER_MEMBERS)
 # --------------------------------------------------------------------------------------------------------------------
@@ -438,6 +603,7 @@ class Gr00tTT:
         backbone: Optional[BackboneTT] = None,
         embed_gather: str = "host",
         cq1_uploads: Optional[bool] = None,
+        adapter_mem: Optional[str] = None,
     ) -> None:
         """Assemble the modules on an open device (``from_pretrained`` is the public entry point).
 
@@ -447,6 +613,8 @@ class Gr00tTT:
         ``"host"`` (the WP-12 host gather; the constructor default, so stand-in weights keep working -- the production
         factory :meth:`from_pretrained` defaults to :data:`DEFAULT_EMBED_GATHER`; module docstring).
         ``cq1_uploads``: per-call input writes on CQ 1 (``None`` = on iff the device has two command queues).
+        ``adapter_mem``: ``"L1"`` | ``"DRAM"`` for the adapter's intermediates (``None`` = :data:`ADAPTER_MEM_BY_BACKEND`
+        of the policy's DiT backend; module docstring).
         """
         if not isinstance(cfg, GR00TConfig) or not isinstance(dlayout, DeviceLayout):
             raise TypeError("Gr00tTT(cfg: GR00TConfig, dlayout: DeviceLayout, policy: TTPolicy, W, device)")
@@ -467,6 +635,17 @@ class Gr00tTT:
                 f"weights are for {getattr(W, 'version', '?')}/slot {getattr(W, 'embodiment_id', '?')}, "
                 f"model is {cfg.version}/slot {e}"
             )
+        wl1 = worker_l1_size_for(policy)
+        if wl1 is not None:
+            got = device_worker_l1_size(device)
+            if got > wl1:
+                raise RuntimeError(
+                    f"dit_backend='megakernel' (the TTPolicy default) needs the device opened with worker_l1_size <= "
+                    f"{wl1}; this device was opened with {got} (the K4 kernel binaries do not fit its kernel-config "
+                    "ring buffer, mk_k4_summary.md section 7.2).  Open it with tt.model.open_model_device(policy, ...) "
+                    "or let Gr00tTT.from_pretrained(device=None) open it; for the Stage-1 path on a default device "
+                    "use TTPolicy(dit_backend='ttnn') / GR00T_DIT_BACKEND=ttnn"
+                )
         self.cfg, self.dlayout, self.policy, self.W, self.device = cfg, dlayout, policy, W, device
         self.version: str = cfg.version
         self.layout: LayoutConfig = dlayout.layout
@@ -479,6 +658,7 @@ class Gr00tTT:
         self.embed_gather = embed_gather
         self.specs: Dict[str, InputSpec] = model_input_specs(cfg, dlayout, policy, embed_gather)
         self.timing: Dict[str, float] = {}
+        self.weights_skipped: Dict[str, Any] = {}  # set by from_pretrained under the megakernel (module docstring)
         self.n_predict = 0
         self._b3: Optional[torch.Tensor] = None
         self._hoist: Optional[HoistOut] = None
@@ -497,7 +677,19 @@ class Gr00tTT:
                 raise ValueError("the injected backbone must be built with collect_taps=False and hold no taps")
             self.backbone = backbone
         t1 = time.perf_counter()
-        self.adapter = VLAdapterTT(cfg, dlayout, W, policy, device, gather_mode=gather_mode, collect_taps=False)
+        self.adapter_mem: str = ADAPTER_MEM_BY_BACKEND[policy.dit_backend] if adapter_mem is None else str(adapter_mem)
+        if self.adapter_mem not in ("L1", "DRAM"):
+            raise ValueError(f"adapter_mem must be 'L1' or 'DRAM', got {adapter_mem!r}")
+        self.adapter = VLAdapterTT(
+            cfg,
+            dlayout,
+            W,
+            policy,
+            device,
+            gather_mode=gather_mode,
+            collect_taps=False,
+            mem=DRAM if self.adapter_mem == "DRAM" else L1,
+        )
         t2 = time.perf_counter()
         if embed_gather == "device":
             if not hasattr(W, "_embed_table_tensor"):
@@ -621,6 +813,8 @@ class Gr00tTT:
         embed_gather: str = DEFAULT_EMBED_GATHER,
         cq1_uploads: Optional[bool] = None,
         num_command_queues: int = 2,
+        adapter_mem: Optional[str] = None,
+        load_ttnn_dit_weights: Optional[bool] = None,
     ) -> "Gr00tTT":
         """Load the weights and assemble the model (plan §2.2 row ``tt/model.py``).
 
@@ -628,14 +822,22 @@ class Gr00tTT:
             version: ``"n15" | "n16" | "n17"``.
             embodiment: embodiment tag (default: the layout's ``embodiment_tag``); must be the slot the layout is baked for.
             layout_name: one of ``cfg.layouts`` (default ``cfg.canonical_layout_name``, D2).
-            policy: ``TTPolicy`` (default ``TTPolicy()``: mixed_dit, ttnn backend, meta RoPE, per_stage traces).
-            device: an open mesh device; ``None`` opens one with ``tt.device.open_gr00t_device(num_command_queues=
-                num_command_queues)`` (closed by ``release``).
+            policy: ``TTPolicy`` (default ``TTPolicy()``: mixed_dit, **megakernel** backend with the bfp8_b arena, meta
+                RoPE, per_stage traces; ``TTPolicy().with_env_overrides()`` honours ``GR00T_DIT_BACKEND`` /
+                ``GR00T_MK_ARENA_DTYPE`` -- the factory itself never reads the environment).
+            device: an open mesh device; ``None`` opens one with :func:`open_model_device` (``num_command_queues``;
+                the megakernel backend's ``worker_l1_size``), closed by ``release``.  A device opened another way is
+                refused for the megakernel backend (constructor).
             cache_root: ``HostPlanCache`` root for ``TTWeights.load`` (default: the project cache).
             warm_inputs: a ``ModelInputs`` (or ``Observation``) to warm and capture on immediately; otherwise the
                 first ``predict`` / ``predict_normalized`` / ``tracer.upload`` call warms and captures on its inputs.
             noise: the noise for that warm pass (default: the version's seeded noise).
-            embed_gather, cq1_uploads: see the constructor (module docstring).
+            embed_gather, cq1_uploads, adapter_mem: see the constructor (module docstring; ``adapter_mem="DRAM"`` is
+                what a ttnn-backend model needs on a device opened for the megakernel backend).
+            load_ttnn_dit_weights: ``None`` (default) uploads the Stage-1 DiT / precompute / encoder tensors only for
+                the ttnn backend (the megakernel packs them into its arena instead, :func:`megakernel_weight_names`;
+                ``describe()["weights"]["skipped"]`` and :attr:`weights_skipped` record what was left out);
+                ``True`` forces the full upload under either backend, ``False`` is an error for the ttnn backend.
         """
         from models.experimental.gr00t.common.checkpoint import LazyCheckpoint
         from models.experimental.gr00t.tt.weights import TTWeights
@@ -653,11 +855,17 @@ class Gr00tTT:
                 f"embodiment {tag!r} is slot {e} -- pick a layout of that embodiment"
             )
         dlayout = DeviceLayout.from_layout(cfg, layout)
+        skip_facts: Dict[str, Any] = {}
+        names: Optional[Tuple[str, ...]] = None
+        if load_ttnn_dit_weights is False and pol.dit_backend != "megakernel":
+            raise ValueError("load_ttnn_dit_weights=False needs dit_backend='megakernel' (the ttnn head reads them)")
+        if pol.dit_backend == "megakernel" and not load_ttnn_dit_weights:
+            names, skip_facts = megakernel_weight_names(version, e, pol, layout.name)
         owns = device is None
-        dev = open_gr00t_device(num_command_queues=int(num_command_queues)) if owns else device
+        dev = open_model_device(pol, num_command_queues=int(num_command_queues)) if owns else device
         try:
             t0 = time.perf_counter()
-            W = TTWeights.load(version, e, pol, dev, cache_root=cache_root, layout_name=layout.name)
+            W = TTWeights.load(version, e, pol, dev, cache_root=cache_root, names=names, layout_name=layout.name)
             load_s = time.perf_counter() - t0
             ck = LazyCheckpoint(version)
             model = cls(
@@ -671,12 +879,14 @@ class Gr00tTT:
                 owns_device=owns,
                 embed_gather=embed_gather,
                 cq1_uploads=cq1_uploads,
+                adapter_mem=adapter_mem,
             )
         except Exception:
             if owns:
                 close_gr00t_device(dev)
             raise
         model.timing["weights_load_s"] = load_s
+        model.weights_skipped = dict(skip_facts)
         if warm_inputs is not None:
             mi = warm_inputs if not _is_observation(warm_inputs) else model.encode(warm_inputs)
             model.warm_and_capture(mi, noise)
@@ -1005,7 +1215,7 @@ class Gr00tTT:
             else:
                 taps[name] = slots_to_rows(L.to_torch(t), mi.slots)
         taps.update(self.head.golden_taps(b3=self.b3_bias(), block_taps=True))
-        missing = [k for k in expected_tap_keys(cfg, dl) if k not in taps]
+        missing = [k for k in expected_tap_keys(cfg, dl, self.policy.dit_backend) if k not in taps]
         if missing:
             raise RuntimeError(f"run_untraced produced no tap for {missing}")
         return taps
@@ -1077,6 +1287,8 @@ class Gr00tTT:
             "policy": self.policy.to_dict(),
             "recipe": self.recipe.describe(),
             "embed_gather": self.embed_gather,
+            "adapter_mem": self.adapter_mem,
+            "worker_l1_size": _worker_l1_size_or_none(self.device),
             "gather": None if self.gather is None else self.gather.describe(),
             "input_specs": {n: s.describe() for n, s in self.specs.items()},
             "executor": self.executor.describe(),
@@ -1087,6 +1299,7 @@ class Gr00tTT:
                 "device_bytes": float(self.W.device_bytes),
                 "load_stats": dict(getattr(self.W, "load_stats", {})),
                 "dtype_histogram": self.W.dtype_histogram() if hasattr(self.W, "dtype_histogram") else None,
+                "skipped": dict(self.weights_skipped),
             },
         }
         for name in ("backbone", "adapter", "head"):
@@ -1101,6 +1314,14 @@ class Gr00tTT:
 # --------------------------------------------------------------------------------------------------------------------
 
 
+def _worker_l1_size_or_none(device: Any) -> Optional[int]:
+    """:func:`device_worker_l1_size` for a real mesh device; ``None`` for the CPU tests' stand-in devices."""
+    try:
+        return device_worker_l1_size(device)
+    except (ImportError, TypeError, AttributeError, RuntimeError):
+        return None
+
+
 def _is_observation(x: Any) -> bool:
     return hasattr(x, "images") and hasattr(x, "instruction") and hasattr(x, "embodiment_tag")
 
@@ -1139,6 +1360,7 @@ __all__ = [
     "IN_VL_MASK",
     "IN_XT",
     "MASK_BUFFER_NAMES",
+    "MK_ARENA_CATEGORIES",
     "NOISE_SEED",
     "N_STEPS",
     "OUT_ACTION_PRED",
@@ -1150,8 +1372,17 @@ __all__ = [
     "Gr00tTT",
     "ModelRecipe",
     "TraceReplayer",
+    "ADAPTER_MEM_BY_BACKEND",
+    "MK_L1_CUT_BYTES_DEFAULT",
+    "MK_L1_CUT_ENV",
+    "device_worker_l1_size",
     "expected_tap_keys",
+    "megakernel_needs_ttnn_tensor",
+    "megakernel_weight_names",
+    "mk_l1_cut_bytes",
     "model_input_specs",
+    "open_model_device",
     "text_ids_spec",
     "unproduced_taps",
+    "worker_l1_size_for",
 ]

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

@@ -42,6 +42,29 @@ from typing import Any, Dict, Mapping, Optional, Tuple
 
 DTYPE_POLICIES: Tuple[str, ...] = ("bf16", "mixed_dit", "mixed")
 DIT_BACKENDS: Tuple[str, ...] = ("ttnn", "megakernel")
+#: ``TTPolicy.dit_backend`` default.  **``"megakernel"`` since 2026-09-18** (plan-owner ratification in
+#: ``docs/plan/reviews/WP-K5-review-round1.md`` "Plan-owner decisions" item 1, overriding plan §5.7's "G3 fails ->
+#: Stage 1 stays production"): the whole DiT denoise runs as one ``steps4`` ``generic_op`` inside trace B
+#: (``tests/tt/results/mk_k5_summary.md``: gate outcomes equal to the ttnn backend's on 21 / 23 (version, sample)
+#: pairs and every canonical sample, mk vs ttnn ``action_pred`` >= 0.999993, e2e -13 / -23 / -17 %).  The ttnn
+#: backend stays fully selectable (``TTPolicy(dit_backend="ttnn")`` or ``GR00T_DIT_BACKEND=ttnn``).
+DEFAULT_DIT_BACKEND: str = "megakernel"
+#: ``TTPolicy.mk_arena_dtype``: the DRAM weight arena of the megakernel backend (plan §5.3 ``dtype_stream``);
+#: ``"auto"`` resolves per version through :data:`MK_ARENA_DTYPE_DEFAULT` (:meth:`TTPolicy.mk_arena_dtype_for`).
+MK_ARENA_DTYPES: Tuple[str, ...] = ("auto", "bf16", "bfp8_b")
+#: Per-version arena dtype behind ``mk_arena_dtype="auto"``: the fastest arena whose e2e gate outcomes equal the
+#: bf16 arena's on the canonical and every multi-sample golden (WP-K5, ``tests/tt/results/mk_k5_summary.md`` §3 / §4.1:
+#: bf16 and bfp8_b have identical pass / fail sets on all 23 (version, sample) pairs and agree to <= 8e-4 PCC on
+#: ``action_pred_normalized`` (<= 1e-6 on N1.6); bfp8_b denoise trace 10.41 / 17.16 / 18.80 ms vs bf16 11.36 / 18.92 /
+#: 20.66 ms, §5).  Stage 1 already runs the DiT projections in bfp8 (``ADMITTED_CLASSES["dit"]``), so this is the same
+#: weight dtype.
+MK_ARENA_DTYPE_DEFAULT: Dict[str, str] = {"n15": "bfp8_b", "n16": "bfp8_b", "n17": "bfp8_b"}
+#: Environment overrides of :meth:`TTPolicy.with_env_overrides` (WP-K5; the device test files, the published policy
+#: server and ``Gr00tTT`` callers that want an A/B against the Stage-1 path use them):
+#: ``GR00T_DIT_BACKEND=ttnn|megakernel`` and ``GR00T_MK_ARENA_DTYPE=auto|bf16|bfp8_b``.  Unset = the code defaults
+#: (:data:`DEFAULT_DIT_BACKEND`, ``"auto"``).
+DIT_BACKEND_ENV = "GR00T_DIT_BACKEND"
+MK_ARENA_DTYPE_ENV = "GR00T_MK_ARENA_DTYPE"
 READER_MODES: Tuple[str, ...] = ("relay", "direct")
 ROPE_LAYOUTS: Tuple[str, ...] = ("meta", "hf")
 TRACE_LAYOUTS: Tuple[str, ...] = ("per_stage", "two")
@@ -109,8 +132,19 @@ class TTPolicy:
     """Run-time policy of the device port (plan §2.2 row ``tt/policy.py``).
 
     * ``dtype_policy``: weight dtype class table, see the module docstring.
-    * ``dit_backend``: ``"ttnn"`` (Stage 1 trace B of ttnn ops) | ``"megakernel"`` (Stage 2 ``generic_op``; D1/D10).
-    * ``reader_mode``: megakernel weight-stream topology, ``"relay"`` ring (default, plan §5.3) | ``"direct"``.
+    * ``dit_backend``: ``"megakernel"`` (**default** since 2026-09-18, :data:`DEFAULT_DIT_BACKEND`: Stage 2, the whole
+      denoise as one ``steps4`` ``generic_op`` inside trace B, ``tt.action_head.ActionHeadTT`` +
+      ``tt.megakernel.dit_program``; D1/D10, WP-K5) | ``"ttnn"`` (Stage 1: trace B of ttnn ops, the previous default;
+      selectable for A/B).  The megakernel needs the device opened with a ``worker_l1_size`` reduced by
+      ``tt.model.MK_L1_CUT_BYTES_DEFAULT`` (``tt.model.open_model_device``; ``Gr00tTT.from_pretrained(device=None)``
+      does it) and the VL adapter's intermediates in DRAM; ``Gr00tTT`` refuses a device opened otherwise.
+    * ``mk_arena_dtype``: the megakernel's DRAM weight arena, ``"bf16"`` | ``"bfp8_b"`` | ``"auto"`` (default: the
+      per-version table :data:`MK_ARENA_DTYPE_DEFAULT`, :meth:`mk_arena_dtype_for`).  Independent of ``dtype_policy``
+      (which governs the ttnn weights ``TTWeights`` uploads); not part of the weights cache key (the arena is packed
+      from the bf16 host plan at model build).  Ignored by the ttnn backend.
+    * ``reader_mode``: K1 megakernel weight-stream topology, ``"relay"`` ring (default, plan §5.3) | ``"direct"``;
+      the K4 / K5 ``steps4`` program always streams in direct mode (``mk_k4_summary.md`` §8), so this knob only
+      reaches the K1 bench / unit-test launches.
     * ``rope_layout``: ``"meta"`` (``rotary_embedding_llama`` on the permuted layout, D8 default) | ``"hf"``
       (``rotary_embedding_hf`` on the HF layout).  Byte-changing: part of ``TransformOptions``.
     * ``trace_layout``: ``"per_stage"`` (vision / llm / adapter / denoise traces, bring-up) | ``"two"`` (D1 end state).
@@ -134,7 +168,8 @@ class TTPolicy:
     """
 
     dtype_policy: str = "mixed_dit"
-    dit_backend: str = "ttnn"
+    dit_backend: str = DEFAULT_DIT_BACKEND
+    mk_arena_dtype: str = "auto"
     reader_mode: str = "relay"
     rope_layout: str = "meta"
     trace_layout: str = "per_stage"
@@ -149,6 +184,7 @@ class TTPolicy:
     def __post_init__(self) -> None:
         _check("dtype_policy", self.dtype_policy, DTYPE_POLICIES)
         _check("dit_backend", self.dit_backend, DIT_BACKENDS)
+        _check("mk_arena_dtype", self.mk_arena_dtype, MK_ARENA_DTYPES)
         _check("reader_mode", self.reader_mode, READER_MODES)
         _check("rope_layout", self.rope_layout, ROPE_LAYOUTS)
         _check("trace_layout", self.trace_layout, TRACE_LAYOUTS)
@@ -184,6 +220,31 @@ class TTPolicy:
             raise ValueError(f"unknown vision family {vision_family!r}; expected one of {VISION_FAMILIES}")
         return bool(self.vit_fork) and vision_family == VIT_FORK_FAMILY
 
+    def mk_arena_dtype_for(self, version: str) -> str:
+        """The megakernel arena dtype of ``version`` (``"bf16"`` | ``"bfp8_b"``): the explicit field, or the
+        per-version default :data:`MK_ARENA_DTYPE_DEFAULT` behind ``"auto"``.  Unknown versions raise."""
+        if self.mk_arena_dtype != "auto":
+            return self.mk_arena_dtype
+        try:
+            return MK_ARENA_DTYPE_DEFAULT[version]
+        except KeyError:
+            raise KeyError(f"no default megakernel arena dtype for version {version!r}") from None
+
+    def with_env_overrides(self, env: Optional[Mapping[str, str]] = None) -> "TTPolicy":
+        """Copy with ``dit_backend`` / ``mk_arena_dtype`` taken from :data:`DIT_BACKEND_ENV` /
+        :data:`MK_ARENA_DTYPE_ENV` when set (the device test files and the published policy server parametrise the
+        DiT backend this way, WP-K5: ``GR00T_DIT_BACKEND=ttnn pytest ...`` runs the Stage-1 path, unset = the
+        megakernel default).  Unset variables change nothing; an unknown value raises through ``__post_init__``."""
+        import os
+
+        src = os.environ if env is None else env
+        changes: Dict[str, str] = {}
+        if src.get(DIT_BACKEND_ENV):
+            changes["dit_backend"] = str(src[DIT_BACKEND_ENV])
+        if src.get(MK_ARENA_DTYPE_ENV):
+            changes["mk_arena_dtype"] = str(src[MK_ARENA_DTYPE_ENV])
+        return replace(self, **changes) if changes else self
+
     def admitted(self) -> Dict[str, str]:
         """The category -> dtype table in force for this policy: ``admitted_classes`` when given, else the module
         table :data:`ADMITTED_CLASSES` (read at call time so a table change is seen by every later ``dtype_of``)."""
@@ -211,6 +272,7 @@ class TTPolicy:
         return {
             "dtype_policy": self.dtype_policy,
             "dit_backend": self.dit_backend,
+            "mk_arena_dtype": self.mk_arena_dtype,
             "reader_mode": self.reader_mode,
             "rope_layout": self.rope_layout,
             "trace_layout": self.trace_layout,
@@ -358,6 +420,11 @@ def dtype_histogram(decisions: Mapping[str, str]) -> Dict[str, int]:
 __all__ = [
     "DTYPE_POLICIES",
     "DIT_BACKENDS",
+    "DEFAULT_DIT_BACKEND",
+    "MK_ARENA_DTYPES",
+    "MK_ARENA_DTYPE_DEFAULT",
+    "DIT_BACKEND_ENV",
+    "MK_ARENA_DTYPE_ENV",
     "READER_MODES",
     "ROPE_LAYOUTS",
     "TRACE_LAYOUTS",