Blame for code/tests/test_fused_host.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Host-only tests of the ``TT_FUSED=1`` path (torch only -- NO device, NO ttnn tensors).

What is proven here (BRIEF deliverable 3):

* every exact reformulation of ``hamer.fused`` against the reference math
  (``torch.equal`` / fp64 max|d|, or <= 1 bf16 ULP with the reason);
* the torch emulation of the whole fused ViT block and of the whole fused MANO
  head (padded heads, ``[Q|K|V]`` head-split layout, V-only self-attention,
  zero-Q fake K/V weight, constant ``x0``, fused regressor) equals the reference
  modules in fp64;
* knob plumbing: ``TT_FUSED=0`` -> legacy path selected, unset/1 -> fused (default), sub-knobs parsed;
* the fused DEVICE code (``vit._forward_fused`` / ``mano_head._forward_fused``)
  walked with a mock ``ttnn`` that enforces the validate rules read from this
  tree (rank/S%32/width rules of the head ops, the fused matmul's bias and
  ternary shape + dtype + buffer-type rules, SDPA without a mask, fp32
  ROW_MAJOR tilize input) and detects use-after-deallocate; launch counts
  match the device-validated graph (10 per ViT block, 16 per head layer, 420 per forward);
* ``model.py``'s per-request plumbing on the same mock (device injected, never
  opened): ``_host_input`` dtype/layout per knob, ``_upload_tokens`` persistent
  buffer + content cache + ``TT_FUSED_HASH=0`` bypass + shape-change re-allocation,
  and the batch-1 guard of the fused path raising even in non-strict mode.

Run (host python of the tt-metal tree; the real ttnn is imported by ``hamer.model``
but never asked for a tensor)::

    cd models/hamer-p150 && PYTHONPATH=code TT_METAL_HOME=<tree> HAMER_USE_TT=0 \
        <tree>/python_env/bin/python -m pytest code/tests/test_fused_host.py -q
"""
from __future__ import annotations

import math
import os
import sys
import types
from pathlib import Path

# The wrapper must never try to open a chip from a host test.
os.environ.setdefault("HAMER_USE_TT", "0")
os.environ.pop("TT_FUSED", None)

import pytest  # noqa: E402
import numpy as np
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

from hamer import fused as fz  # noqa: E402
from hamer import mano_head as mh  # noqa: E402
from hamer import vit  # noqa: E402
from hamer import tt_kernels as tk  # noqa: E402
from hamer.fused import FusedOptions  # noqa: E402
from hamer.reference import Hamer, HamerConfig, sample_input  # noqa: E402

torch.manual_seed(0)
torch.set_grad_enabled(False)

# Small-depth reference with the REAL widths (1280 / 16 heads x 80 / 5120; head 1024 / 8 x 64):
# every shape, padding and layout below is the production one, only fewer blocks.
DEPTH, HEAD_DEPTH = 2, 2


@pytest.fixture(scope="module")
def ref() -> Hamer:
    torch.manual_seed(0)
    m = Hamer(HamerConfig(depth=DEPTH, head_depth=HEAD_DEPTH))
    # Random LN affine / pos so that nothing is trivially identity or zero.
    for n, p in m.named_parameters():
        if n.endswith(("norm1.weight", "norm2.weight", "last_norm.weight", "norm_sa.weight", "norm_ca.weight",
                       "norm_ff.weight")):
            p.copy_(1.0 + 0.1 * torch.randn_like(p))
        elif n.endswith(("norm1.bias", "norm2.bias", "last_norm.bias", "norm_sa.bias", "norm_ca.bias", "norm_ff.bias")):
            p.copy_(0.1 * torch.randn_like(p))
    m.backbone.pos_embed.copy_(0.02 * torch.randn_like(m.backbone.pos_embed))
    # the reference init sets the MANO mean buffers to zero; give them values so host_finalize is exercised
    m.head.init_hand_pose.copy_(0.1 * torch.randn_like(m.head.init_hand_pose))
    m.head.init_betas.copy_(0.1 * torch.randn_like(m.head.init_betas))
    m.head.init_cam.copy_(torch.tensor([[5.0, 0.01, -0.02]]))
    return m.eval()


@pytest.fixture(scope="module")
def image() -> torch.Tensor:
    return sample_input(batch=1, seed=3)


def _maxd(a: torch.Tensor, b: torch.Tensor) -> float:
    return float((a.double() - b.double()).abs().max())


# =========================================================================== #
# 1. Knob plumbing
# =========================================================================== #
def test_knob_default_on_and_parsing():
    # default ON since the p150a pass; TT_FUSED=0 is the legacy path
    assert FusedOptions.from_env({}).enabled is True
    assert FusedOptions.from_env({}) == FusedOptions.from_env({"TT_FUSED": "1"})
    assert FusedOptions.from_env({"TT_FUSED": "0"}) == FusedOptions()
    assert FusedOptions.from_env({"TT_FUSED": "0"}).enabled is False
    # a stray sub-knob never affects the legacy path (not even parsed)
    assert FusedOptions.from_env({"TT_FUSED": "0", "TT_FUSED_MM": "garbage"}) == FusedOptions()
    on = FusedOptions.from_env({"TT_FUSED": "1"})
    assert on.enabled and on.matmul == "tuned" and on.fc1 == "bfp8" and on.fc2 == "default"
    # tuned (chip-25 pass): qkv / fc2 / proj bfp8 with fp32 accumulation (proj bf16 without it);
    # minimal keeps the bf16 fusable set
    assert on.wdtype("fc2") == "bfp8" and on.wdtype("qkv_w") == "bfp8" and on.wdtype("proj_w") == "bfp8"
    noacc = FusedOptions.from_env({"TT_FUSED": "1", "TT_FUSED_FP32ACC": "0"})
    assert noacc.fp32_acc is False and noacc.wdtype("proj_w") == "bf16" and noacc.wdtype("fc2") == "bfp8"
    mini = FusedOptions.from_env({"TT_FUSED": "1", "TT_FUSED_MM": "minimal"})
    assert mini.wdtype("fc2") == mini.wdtype("qkv_w") == mini.wdtype("proj_w") == "bf16"
    # defaults: fp32 patch weight (p150a pass), 32x192 SDPA chunking (chip-25 pass)
    assert on.patch == "device" and on.device_im2col and on.sdpa == "32x192" and on.fidelity == "HiFi2" and on.head_concat == "nlp"
    assert on.head == "hg" and on.heads == "custom" and on.gelu == "kernel" and on.vit_shard is True
    assert on.l1 is True and on.head_l1 is False and on.fp32_acc is True and on.hash_cache is True
    assert on.device_patch_embed is True
    custom = FusedOptions.from_env({
        "TT_FUSED": "true", "TT_FUSED_MM": "linear", "TT_FUSED_FC1": "bf16", "TT_FUSED_FC2": "bfp8",
        "TT_FUSED_L1": "0", "TT_FUSED_PATCH": "host", "TT_FUSED_SDPA": "default", "TT_FUSED_FIDELITY": "hifi4",
        "TT_FUSED_FP32ACC": "1", "TT_FUSED_HEAD_CONCAT": "permute", "TT_FUSED_HASH": "off",
    })
    assert custom == FusedOptions(enabled=True, matmul="linear", fc1="bf16", fc2="bfp8", l1=False, patch="host",
                                  sdpa="default", fidelity="HiFi4", fp32_acc=True, head_concat="permute",
                                  hash_cache=False)
    assert custom.device_patch_embed is False
    with pytest.raises(ValueError):
        FusedOptions.from_env({"TT_FUSED": "1", "TT_FUSED_MM": "fast"})
    with pytest.raises(ValueError):
        FusedOptions.from_env({"TT_FUSED": "1", "TT_FUSED_L1": "maybe"})
    assert "TT_FUSED=0" in FusedOptions().describe() and "matmul=tuned" in on.describe()


def test_wrapper_reads_knob_at_build_not_import(ref, monkeypatch):
    """``TtHamer`` (HAMER_USE_TT=0 -> no device) picks the legacy path with ``TT_FUSED=0``,
    the fused one without the env (default) or with ``TT_FUSED=1``, and the explicit
    ``fused=`` argument wins over the env."""
    from hamer.model import TtHamer, build_from_reference

    monkeypatch.setenv("TT_FUSED", "0")
    m = TtHamer(ref)
    assert m.fused is False and m.fused_opts == FusedOptions() and m.device is None and m._vit_params is None
    monkeypatch.delenv("TT_FUSED", raising=False)
    assert TtHamer(ref).fused is True  # default ON
    monkeypatch.setenv("TT_FUSED", "1")
    monkeypatch.setenv("TT_FUSED_PATCH", "host")
    m2 = build_from_reference(ref)
    assert m2.fused is True and m2.fused_opts.patch == "host" and m2.fused_opts.device_patch_embed is False
    m3 = build_from_reference(ref, fused=FusedOptions())
    assert m3.fused is False  # explicit legacy object wins over the env
    monkeypatch.setenv("TT_FUSED", "0")
    assert TtHamer(ref).fused is False
    # legacy public surfaces are intact
    assert vit.forward.__defaults__[-1] is False and mh.forward_device.__defaults__[-1] is False
    assert vit.build_parameters_from_reference.__defaults__ == (None,)
    assert mh.build_parameters_from_reference.__defaults__ == (None,)


# =========================================================================== #
# 2. Patch embedding: exact im2col + constant table
# =========================================================================== #
def test_im2col_ordering_is_c_kh_kw(image):
    """Column j of the im2col = padded pixel (c, gh*16 + kh, gw*16 + kw) with j = c*256 + kh*16 + kw,
    the same (c, kh, kw) row order ``patch_embed_weight`` flattens the conv kernel to."""
    cols = fz.host_im2col(image)  # [1, 1, 192, 768]
    assert tuple(cols.shape) == (1, 1, fz.N_TOKENS, fz.PATCH_DIM) and cols.dtype == torch.float32
    padded = F.pad(image, (fz.PATCH_PAD,) * 4)  # [1, 3, 260, 196]; the stride-16 conv covers rows/cols < 256/192
    padded = padded[:, :, : fz.GRID_H * fz.PATCH, : fz.GRID_W * fz.PATCH]
    manual = padded.reshape(1, 3, fz.GRID_H, fz.PATCH, fz.GRID_W, fz.PATCH)          # c gh kh gw kw
    manual = manual.permute(0, 2, 4, 1, 3, 5).reshape(1, 1, fz.N_TOKENS, fz.PATCH_DIM)  # (gh gw) (c kh kw)
    assert torch.equal(cols, manual)


def test_im2col_matmul_equals_conv2d(ref, image):
    conv = ref.backbone.patch_embed.proj
    w = fz.patch_embed_weight(conv)                       # [768, 1280]
    assert torch.equal(w, vit._patch_embed_weight(conv))  # same flattening as the legacy helper
    tokens_ref, _ = ref.backbone.patch_embed(image)       # [1, 192, 1280] fp32 conv
    tokens = (fz.host_im2col(image)[0] @ w + conv.bias)   # [1, 192, 1280]
    d = _maxd(tokens, tokens_ref)
    assert d <= 1e-5, d                                   # fp32 accumulation order only
    # exact in fp64
    d64 = _maxd(fz.host_im2col(image).double()[0] @ w.double() + conv.bias.double(),
                F.conv2d(image.double(), conv.weight.double(), conv.bias.double(), stride=16, padding=2)
                .flatten(2).transpose(1, 2))
    assert d64 <= 1e-9, d64


def test_pos_plus_bias_constant(ref, image):
    """The fused patch op computes ``pos_plus_bias + im2col @ W``; that must equal the
    reference ``patch_embed(x) + pos[:, 1:] + pos[:, :1]`` (what the legacy path had after
    its device pos add)."""
    w, ppb = fz.patch_constants(ref.backbone)
    assert tuple(ppb.shape) == (1, 1, fz.N_TOKENS, fz.EMBED)
    pos = ref.backbone.pos_embed
    assert torch.equal(fz.pos_combined(ref.backbone), (pos[0, 1:] + pos[0, :1]))
    fused_tokens = fz.host_im2col(image)[0] @ w + ppb[0]
    d = _maxd(fused_tokens, fz.patch_tokens_reference(ref.backbone, image)[0])
    assert d <= 1e-5, d


# =========================================================================== #
# 3. ViT block: padding + head-split layout + fused graph == reference (fp64)
# =========================================================================== #
def _emulate_create_qkv_heads(qkv: torch.Tensor, num_heads: int):
    """Torch model of ``nlp_create_qkv_heads(transpose_k_heads=False)``: the reader kernel
    consumes each tile-row as [Q tiles | K tiles | V tiles] (q_num_tiles = H*dh/32 first),
    per head contiguous -> ``view(B, S, 3, H, dh)``."""
    B, one, S, W = qkv.shape
    assert one == 1 and S % 32 == 0 and W % (3 * num_heads) == 0
    dh = W // (3 * num_heads)
    t = qkv.reshape(B, S, 3, num_heads, dh).permute(2, 0, 3, 1, 4)
    return t[0], t[1], t[2]  # [B, H, S, dh]


def _emulate_concat_heads(x: torch.Tensor):
    B, H, S, dh = x.shape
    return x.permute(0, 2, 1, 3).reshape(B, 1, S, H * dh)


def test_qkv_padding_layout_matches_head_split(ref):
    blk = ref.backbone.blocks[0]
    H = blk.attn.num_heads
    w = blk.attn.qkv.weight.t()                                   # [1280, 3840]
    wp = vit._pad_qkv_weight(w, H)                                # [1280, 4608]
    assert tuple(wp.shape) == (fz.EMBED, 3 * H * fz.VIT_HEAD_DIM_PAD)
    v = wp.view(fz.EMBED, 3, H, fz.VIT_HEAD_DIM_PAD)
    assert torch.equal(v[..., : fz.VIT_HEAD_DIM], w.view(fz.EMBED, 3, H, fz.VIT_HEAD_DIM))
    assert torch.equal(v[..., fz.VIT_HEAD_DIM:], torch.zeros_like(v[..., fz.VIT_HEAD_DIM:]))
    bp = vit._pad_qkv_bias(blk.attn.qkv.bias, H).view(3, H, fz.VIT_HEAD_DIM_PAD)
    assert torch.equal(bp[..., : fz.VIT_HEAD_DIM], blk.attn.qkv.bias.view(3, H, fz.VIT_HEAD_DIM))
    assert bool((bp[..., fz.VIT_HEAD_DIM:] == 0).all())
    pw = vit._pad_proj_weight(blk.attn.proj.weight.t(), H).view(H, fz.VIT_HEAD_DIM_PAD, fz.EMBED)
    assert torch.equal(pw[:, : fz.VIT_HEAD_DIM], blk.attn.proj.weight.t().view(H, fz.VIT_HEAD_DIM, fz.EMBED))
    assert bool((pw[:, fz.VIT_HEAD_DIM:] == 0).all())
    # the padded columns of q/k/v produced through the fused-qkv layout are exactly zero
    x = torch.randn(1, 1, fz.N_TOKENS, fz.EMBED)
    q, k, vv = _emulate_create_qkv_heads(x @ wp + vit._pad_qkv_bias(blk.attn.qkv.bias, H), H)
    for t in (q, k, vv):
        assert tuple(t.shape) == (1, H, fz.N_TOKENS, fz.VIT_HEAD_DIM_PAD)
        assert bool((t[..., fz.VIT_HEAD_DIM:] == 0).all())
    # ... and equal the reference split of the unpadded qkv in the first 80 columns
    qr, kr, vr = (x[0] @ w + blk.attn.qkv.bias).reshape(1, fz.N_TOKENS, 3, H, fz.VIT_HEAD_DIM).permute(2, 0, 3, 1, 4)
    assert torch.equal(q[..., :80], qr) and torch.equal(k[..., :80], kr) and torch.equal(vv[..., :80], vr)


def _fused_block_torch(x4: torch.Tensor, blk, H: int) -> torch.Tensor:
    """fp64 torch model of ``vit._block_fused`` (4-D [1,1,192,1280] in/out)."""
    eps = blk.norm1.eps
    wqkv = vit._pad_qkv_weight(blk.attn.qkv.weight.t(), H).double()
    bqkv = vit._pad_qkv_bias(blk.attn.qkv.bias, H).double()
    wproj = vit._pad_proj_weight(blk.attn.proj.weight.t(), H).double()
    h = F.layer_norm(x4, (fz.EMBED,), blk.norm1.weight.double(), blk.norm1.bias.double(), eps)
    q, k, v = _emulate_create_qkv_heads(h @ wqkv + bqkv, H)                 # [1,16,192,96]
    scores = (q @ k.transpose(-1, -2)) * (fz.VIT_HEAD_DIM ** -0.5)          # scale = REAL head dim
    ctx = torch.softmax(scores, dim=-1) @ v                                # [1,16,192,96]
    ctx = _emulate_concat_heads(ctx)                                        # [1,1,192,1536]
    x4 = x4 + (ctx @ wproj + blk.attn.proj.bias.double())                   # dit(proj,+b,+res)
    h = F.layer_norm(x4, (fz.EMBED,), blk.norm2.weight.double(), blk.norm2.bias.double(), eps)
    h = F.gelu(h @ blk.mlp.fc1.weight.t().double() + blk.mlp.fc1.bias.double())
    return x4 + (h @ blk.mlp.fc2.weight.t().double() + blk.mlp.fc2.bias.double())  # dit(fc2,+b,+res)


def test_fused_vit_block_equals_reference_fp64(ref):
    blk = ref.backbone.blocks[0].double()
    x = torch.randn(1, fz.N_TOKENS, fz.EMBED, dtype=torch.float64)
    y_ref = blk(x)                                                          # reference Block.forward
    y_fused = _fused_block_torch(x.reshape(1, 1, fz.N_TOKENS, fz.EMBED), blk, blk.attn.num_heads)
    d = _maxd(y_fused.reshape(1, fz.N_TOKENS, fz.EMBED), y_ref)
    assert d <= 1e-9, d
    ref.backbone.blocks[0].float()


# =========================================================================== #
# 4. MANO head exact re-layouts + fused graph == reference (fp64)
# =========================================================================== #
def test_head_x0_constant(ref):
    td = ref.head.transformer
    x0 = fz.head_x0(td)
    assert tuple(x0.shape) == (1, 1, 1, fz.HEAD_DIM)
    token = torch.zeros(1, 1, 1)
    ref_x0 = td.to_token_embedding(token) + td.pos_embedding[:, :1]         # reference prologue
    assert torch.equal(x0.reshape(-1), ref_x0.reshape(-1))                  # 0 @ W is exactly 0
    # The legacy DEVICE chain rounds to bf16 three times (0@W -> +b -> +pos); the constant is
    # rounded once from the fp32 sum.  Where b and pos nearly cancel the legacy value carries
    # the rounding errors of both operands, so the bound is 2 bf16 ULP of the LARGER OPERAND
    # (0.5 ulp(b) + 0.5 ulp(pos) + 0.5 ulp(sum) + 0.5 ulp(sum)), not 1 ULP of the sum -- and
    # the constant is never farther from the fp32 value than the legacy chain (max and mean).
    bf = torch.bfloat16
    b, pos = td.to_token_embedding.bias.detach(), td.pos_embedding.detach()[0, 0]
    legacy = ((torch.zeros(fz.HEAD_DIM, dtype=bf) + b.to(bf)).float().to(bf) + pos.to(bf)).float()
    new = x0.reshape(-1).to(bf).float()
    ulp_operand = fz.bf16_ulp(torch.maximum(b.abs(), pos.abs()))
    assert bool(((legacy - new).abs() <= 2.0 * ulp_operand).all())
    exact = ref_x0.reshape(-1)
    assert _maxd(new, exact) <= _maxd(legacy, exact) + 1e-12
    assert float((new - exact).abs().mean()) <= float((legacy - exact).abs().mean())


def test_head_self_attention_is_v_for_one_token(ref):
    """N=1 => softmax over one key == 1.0 EXACTLY (x/x), so Attn(q,k,v) == v and the
    self-attention output is ``to_out(LN(x) @ Wv)`` with the V column slice of to_qkv."""
    assert torch.equal(torch.softmax(torch.randn(8, 1, 1) * 50, dim=-1), torch.ones(8, 1, 1))
    blk = ref.head.transformer.layers[0]
    wv = fz.head_sa_v_weight(blk)
    assert tuple(wv.shape) == (fz.HEAD_DIM, fz.HEAD_INNER)
    assert torch.equal(wv, blk.sa.to_qkv.weight.t()[:, 2 * fz.HEAD_INNER:])
    h = torch.randn(1, 1, fz.HEAD_DIM)
    d = _maxd(blk.sa.to_out(h @ wv), blk.sa(h))
    assert d <= 1e-5, d
    blk.double()
    d64 = _maxd(blk.sa.to_out(h.double() @ wv.double()), blk.sa(h.double()))
    assert d64 <= 1e-12, d64
    blk.float()


def test_head_fake_kv_weight_layout(ref):
    """``[0 | Wk | Wv]`` through the ``[Q|K|V]`` head split gives zero Q and the reference K/V heads."""
    blk = ref.head.transformer.layers[0]
    wf = fz.head_ca_kv_fake_weight(blk)
    assert tuple(wf.shape) == (fz.HEAD_CONTEXT, 3 * fz.HEAD_INNER)
    assert bool((wf[:, : fz.HEAD_INNER] == 0).all())
    assert torch.equal(wf[:, fz.HEAD_INNER:], blk.ca.to_kv.weight.t())
    assert wf.shape[1] % (3 * fz.HEAD_HEADS) == 0 and wf.shape[1] // (3 * fz.HEAD_HEADS) == fz.HEAD_DIM_HEAD
    ctx = torch.randn(1, 1, fz.N_TOKENS, fz.HEAD_CONTEXT, dtype=torch.float64)
    q0, k, v = _emulate_create_qkv_heads(ctx @ wf.double(), fz.HEAD_HEADS)
    assert bool((q0 == 0).all())
    kr, vr = blk.ca.to_kv.double()(ctx[0]).chunk(2, dim=-1)                # reference: chunk(2) then heads
    kr = kr.reshape(1, fz.N_TOKENS, fz.HEAD_HEADS, fz.HEAD_DIM_HEAD).permute(0, 2, 1, 3)
    vr = vr.reshape(1, fz.N_TOKENS, fz.HEAD_HEADS, fz.HEAD_DIM_HEAD).permute(0, 2, 1, 3)
    assert _maxd(k, kr) <= 1e-12 and _maxd(v, vr) <= 1e-12
    blk.ca.to_kv.float()
    # concat_heads on the one-token context == the reference "b h n d -> b n (h d)"
    ca = torch.randn(1, fz.HEAD_HEADS, 1, fz.HEAD_DIM_HEAD)
    assert torch.equal(_emulate_concat_heads(ca).reshape(1, 1, -1), ca.permute(0, 2, 1, 3).reshape(1, 1, -1))


def _fused_head_torch(context4: torch.Tensor, head) -> torch.Tensor:
    """fp64 torch model of ``mano_head._forward_fused``: returns dec [1,1,1,109]."""
    td = head.transformer
    x = fz.head_x0(td).double()
    eps = td.layers[0].norm_sa.eps
    for blk in td.layers:
        h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_sa.weight.double(), blk.norm_sa.bias.double(), eps)
        v = h @ fz.head_sa_v_weight(blk).double()
        x = x + (v @ blk.sa.to_out.weight.t().double() + blk.sa.to_out.bias.double())
        h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_ca.weight.double(), blk.norm_ca.bias.double(), eps)
        q = h @ blk.ca.to_q.weight.t().double()                                        # [1,1,1,512]
        q_h = q.reshape(1, 1, fz.HEAD_HEADS, fz.HEAD_DIM_HEAD).permute(0, 2, 1, 3)   # [1,8,1,64]
        _, k_h, v_h = _emulate_create_qkv_heads(context4 @ fz.head_ca_kv_fake_weight(blk).double(), fz.HEAD_HEADS)
        att = torch.softmax((q_h @ k_h.transpose(-1, -2)) * (fz.HEAD_DIM_HEAD ** -0.5), dim=-1) @ v_h
        ca = _emulate_concat_heads(att)                                                # [1,1,1,512]
        x = x + (ca @ blk.ca.to_out.weight.t().double() + blk.ca.to_out.bias.double())
        h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_ff.weight.double(), blk.norm_ff.bias.double(), eps)
        h = F.gelu(h @ blk.ff.net[0].weight.t().double() + blk.ff.net[0].bias.double())
        x = x + (h @ blk.ff.net[2].weight.t().double() + blk.ff.net[2].bias.double())
    dec_w, dec_b = fz.head_dec_weight_bias(head)
    return x @ dec_w.double() + dec_b.double()


def test_fused_head_equals_reference_fp64(ref):
    head = ref.head.double()
    feats = torch.randn(1, fz.EMBED, fz.GRID_H, fz.GRID_W, dtype=torch.float64)   # reference input [B,C,h,w]
    context = feats.flatten(2).transpose(1, 2)                                      # [1,192,1280] (b (h w) c)
    tok = head.transformer(torch.zeros(1, 1, 1, dtype=torch.float64), context=context).squeeze(1)
    dec_ref = torch.cat([head.decpose(tok), head.decshape(tok), head.deccam(tok)], dim=-1)  # [1,109]
    dec = _fused_head_torch(context.reshape(1, 1, fz.N_TOKENS, fz.EMBED), head)
    assert tuple(dec.shape) == (1, 1, 1, fz.DEC_OUT)
    d = _maxd(dec.reshape(1, -1), dec_ref)
    assert d <= 1e-9, d
    # host_finalize on that dec reproduces the reference head output (init params + rot6d)
    out = mh.host_finalize(dec.reshape(1, -1).float(), (96, 10, 3), head.init_hand_pose.float(),
                           head.init_betas.float(), head.init_cam.float())
    ref_out = head(feats)
    assert tuple(out.shape) == (1, 157)
    assert _maxd(out, ref_out) <= 1e-4
    # numpy host_finalize == the original torch one
    dd = torch.randn(3, 109)
    args = ((96, 10, 3), head.init_hand_pose.float(), head.init_betas.float(), head.init_cam.float())
    assert _maxd(mh.host_finalize(dd, *args), mh.host_finalize_torch(dd, *args)) <= 2e-6
    ref.head.float()


def test_ds_head_relayouts_equal_reference_fp64(ref):
    """fp64 model of ``mano_head.forward_ds`` (TT_FUSED_HEAD=ds): host-folded layer-0 SA / query,
    ONE batched K/V projection split into (K_0, V_0, ..., K_5, V_5), the token replicated over a
    tile's 32 rows, the regressor N zero-padded to the DRAM banks -> equals the reference head."""
    head = ref.head.double()
    td = head.transformer
    feats = torch.randn(1, fz.EMBED, fz.GRID_H, fz.GRID_W, dtype=torch.float64)
    context = feats.flatten(2).transpose(1, 2)                                      # [1,192,1280]
    tok = td(torch.zeros(1, 1, 1, dtype=torch.float64), context=context).squeeze(1)
    dec_ref = torch.cat([head.decpose(tok), head.decshape(tok), head.deccam(tok)], dim=-1)
    x1, q0 = fz.head_layer0_constants(td)
    assert tuple(x1.shape) == (1, 1, 1, fz.HEAD_DIM) and tuple(q0.shape) == (1, fz.HEAD_HEADS, 1, fz.HEAD_DIM_HEAD)
    kv = context.reshape(1, 1, fz.N_TOKENS, fz.EMBED) @ fz.head_kv_all_weight(td)   # [1,1,192,6144]
    depth = len(td.layers)
    kvs = kv.reshape(1, fz.N_TOKENS, 2 * depth, fz.HEAD_HEADS, fz.HEAD_DIM_HEAD).permute(2, 0, 3, 1, 4)
    rep = lambda t: t.expand(*t.shape[:-2], 32, t.shape[-1])                       # noqa: E731
    x = rep(x1)
    eps = td.layers[0].norm_sa.eps
    for i, blk in enumerate(td.layers):
        if i > 0:
            h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_sa.weight, blk.norm_sa.bias, eps)
            x = x + blk.sa.to_out(h @ fz.head_sa_v_weight(blk))
            h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_ca.weight, blk.norm_ca.bias, eps)
            q_h = (h @ blk.ca.to_q.weight.t()).reshape(1, 32, fz.HEAD_HEADS, fz.HEAD_DIM_HEAD).permute(0, 2, 1, 3)
        else:
            q_h = rep(q0)
        att = torch.softmax((q_h @ kvs[2 * i].transpose(-1, -2)) * fz.HEAD_DIM_HEAD ** -0.5, dim=-1) @ kvs[2 * i + 1]
        x = x + blk.ca.to_out(att.permute(0, 2, 1, 3).reshape(1, 1, 32, fz.HEAD_INNER))
        h = F.layer_norm(x, (fz.HEAD_DIM,), blk.norm_ff.weight, blk.norm_ff.bias, eps)
        x = x + blk.ff(h)
    dec_w, dec_b = fz.head_dec_weight_bias(head)
    dec = x @ F.pad(dec_w, (0, 256 - fz.DEC_OUT)) + F.pad(dec_b, (0, 256 - fz.DEC_OUT))
    assert float(dec[..., fz.DEC_OUT:].abs().max()) == 0.0
    for r in (0, 17, 31):  # every replicated row carries the token
        assert _maxd(dec[0, 0, r:r + 1, : fz.DEC_OUT], dec_ref) <= 1e-9
    ref.head.float()


# =========================================================================== #
# 5. The fused DEVICE code, walked with a mock ttnn (shapes, validate rules, launches)
# =========================================================================== #
TILE = 32


def _pad32(n: int) -> int:
    return -(-n // TILE) * TILE


class FakeTensor:
    def __init__(self, shape, dtype, layout, mem, alive_set):
        self.shape = tuple(int(s) for s in shape)
        self.dtype = dtype
        self.layout = layout
        self.mem = mem
        self._alive = alive_set
        self._id = object()
        alive_set.add(self._id)

    # --- what the real code calls on tensors
    def memory_config(self):
        return self.mem

    def padded(self):
        if self.layout != "TILE" or len(self.shape) < 2:
            return self.shape
        return self.shape[:-2] + (_pad32(self.shape[-2]), _pad32(self.shape[-1]))

    def check(self):
        assert self._id in self._alive, "use of a deallocated tensor"
        return self

    def device(self):  # legacy vit.block asks the hidden tensor for its device
        return FakeDevice()

    # legacy path operators
    def __matmul__(self, other):
        return FAKE.matmul(self, other)

    def __add__(self, other):
        return FAKE.add(self, other)

    def __getitem__(self, i):
        FAKE.launches += 1
        self.check()
        return FAKE._new(self.shape[1:], self.dtype, self.layout, self.mem)


class Mem:
    def __init__(self, buffer_type, sharded=False):
        self.buffer_type = buffer_type
        self.sharded = sharded

    def __repr__(self):
        return f"Mem({self.buffer_type})"


class FakeTTNN(types.SimpleNamespace):
    """Just enough of ttnn for ``vit`` / ``mano_head`` (build + forward), enforcing the
    validate rules of this tree that a host cannot otherwise check.  Every op call counts
    as one device launch (``deallocate`` / build-time uploads do not)."""

    def __init__(self):
        super().__init__()
        self.launches = 0
        self.alive = set()
        self.calls = []
        self.uploads, self.copies, self.released = [], [], []
        self.bfloat16, self.bfloat8_b, self.float32, self.uint32 = "bf16", "bfp8", "fp32", "u32"
        self.TILE_LAYOUT, self.ROW_MAJOR_LAYOUT = "TILE", "RM"
        self.BufferType = types.SimpleNamespace(L1="L1", DRAM="DRAM")
        self.L1_MEMORY_CONFIG, self.DRAM_MEMORY_CONFIG = Mem("L1"), Mem("DRAM")
        self.MathFidelity = types.SimpleNamespace(LoFi="LoFi", HiFi2="HiFi2", HiFi4="HiFi4")
        self.experimental = types.SimpleNamespace(
            minimal_matmul=self.minimal_matmul,
            dit_minimal_matmul_addcmul_fused=self.dit_minimal_matmul_addcmul_fused,
            nlp_create_qkv_heads=self.nlp_create_qkv_heads,
            nlp_concat_heads=self.nlp_concat_heads,
        )
        self.transformer = types.SimpleNamespace(scaled_dot_product_attention=self.sdpa)

    # ---- construction (not launches)
    def _new(self, shape, dtype, layout, mem):
        return FakeTensor(shape, dtype, layout, mem, self.alive)

    def from_torch(self, t, device=None, layout="RM", dtype="bf16", memory_config=None, mesh_mapper=None):
        assert (dtype == "fp32") == (t.dtype == torch.float32), "host dtype must match the ttnn dtype"
        out = self._new(t.shape, dtype, layout, memory_config or Mem("DRAM"))
        out.on_device = device is not None
        self.uploads.append(out)
        return out

    # ---- model.py plumbing (persistent input buffer, trace release)
    def ReplicateTensorToMesh(self, device):
        return ("replicate", device)

    def copy_host_to_device_tensor(self, host_t, dev_t):
        host_t.check(); dev_t.check()
        assert host_t.on_device is False and dev_t.on_device is True
        # the persistent buffer's spec must match what the host builds (§5 triage row)
        assert (host_t.shape, host_t.dtype, host_t.layout) == (dev_t.shape, dev_t.dtype, dev_t.layout), (
            (host_t.shape, host_t.dtype, host_t.layout), (dev_t.shape, dev_t.dtype, dev_t.layout))
        self.copies.append((host_t, dev_t))

    def release_trace(self, device, trace_id):
        self.released.append(trace_id)

    def to_memory_config(self, t, mem):
        t.check()
        out = self._new(t.shape, t.dtype, t.layout, mem)
        return self._launch("reshard", out) if mem.sharded else out

    def deallocate(self, t):
        assert t._id in self.alive, "double deallocate"
        self.alive.discard(t._id)

    def CoreGrid(self, x, y):
        return (x, y)

    def CoreCoord(self, x, y):
        return types.SimpleNamespace(x=x, y=y)

    ShardStrategy = types.SimpleNamespace(BLOCK="BLOCK", WIDTH="WIDTH", HEIGHT="HEIGHT")
    ShardOrientation = types.SimpleNamespace(ROW_MAJOR="ROW_MAJOR", COL_MAJOR="COL_MAJOR")

    TensorMemoryLayout = types.SimpleNamespace(WIDTH_SHARDED="WIDTH_SHARDED", BLOCK_SHARDED="BLOCK_SHARDED")

    def num_cores_to_corerangeset(self, n, grid, row_wise=True):
        assert n <= grid.x * grid.y
        return ("cores", n)

    def ShardSpec(self, grid, shape, orientation):
        return types.SimpleNamespace(grid=grid, shape=tuple(shape), orientation=orientation)

    def MemoryConfig(self, layout, buffer_type, shard_spec=None):
        assert layout == "WIDTH_SHARDED" and buffer_type == "L1"
        n = shard_spec.grid[1]
        assert shard_spec.shape[1] % TILE == 0 and shard_spec.shape[0] % TILE == 0
        m = Mem("L1", sharded=True)
        m.width_shards = (n, shard_spec.shape)
        return m

    def create_sharded_memory_config(self, shape, *, core_grid, strategy, orientation):
        gx, gy = core_grid
        assert gx <= 12 and gy <= 10, "12x10 grid policy"
        assert strategy == "BLOCK" and shape[0] % (gy * TILE) == 0 and shape[1] % (gx * TILE) == 0
        return Mem("L1", sharded=True)

    def LayerNormShardedMultiCoreProgramConfig(self, *, compute_with_storage_grid_size, subblock_w, block_h,
                                               block_w, inplace):
        assert block_w % subblock_w == 0 and not inplace
        return types.SimpleNamespace(grid=tuple(compute_with_storage_grid_size), bh=block_h, bw=block_w)

    def sharded_to_interleaved(self, t, mem, output_dtype=None):
        t.check()
        assert t.mem.sharded and not mem.sharded
        return self._launch("s2i", self._new(t.shape, output_dtype or t.dtype, t.layout, mem))

    @staticmethod
    def _check_subblock(pcm, pcn, sh, sw):
        assert pcm % sh == 0 and pcn % sw == 0 and sh * sw <= 8, (pcm, pcn, sh, sw)
        assert sw == pcn or sh == 1, "out_subblock_w must equal per_core_N unless out_subblock_h == 1"

    def MatmulMultiCoreReuseMultiCast1DProgramConfig(self, *, compute_with_storage_grid_size, in0_block_w,
                                                     out_subblock_h, out_subblock_w, per_core_M, per_core_N,
                                                     fuse_batch, fused_activation, mcast_in0):
        self._check_subblock(per_core_M, per_core_N, out_subblock_h, out_subblock_w)
        g = compute_with_storage_grid_size
        assert g.x <= 12 and g.y <= 10, "12x10 grid policy"
        return types.SimpleNamespace(kind="1d", grid=g, ibw=in0_block_w, pcm=per_core_M, pcn=per_core_N)

    def MatmulMultiCoreReuseMultiCastProgramConfig(self, *, compute_with_storage_grid_size, in0_block_w,
                                                   out_subblock_h, out_subblock_w, per_core_M, per_core_N,
                                                   transpose_mcast, fused_activation):
        self._check_subblock(per_core_M, per_core_N, out_subblock_h, out_subblock_w)
        g = compute_with_storage_grid_size
        assert g.x <= 12 and g.y <= 10, "12x10 grid policy"
        return types.SimpleNamespace(kind="2d", grid=g, ibw=in0_block_w, pcm=per_core_M, pcn=per_core_N)

    @staticmethod
    def _check_pc(pc, M, K, N):
        Mt, Kt, Nt = _pad32(M) // TILE, _pad32(K) // TILE, _pad32(N) // TILE
        assert Kt % pc.ibw == 0, ("in0_block_w must divide K tiles", Kt, pc.ibw)
        if pc.kind == "1d":
            assert pc.pcm == Mt and -(-Nt // pc.pcn) <= pc.grid.x * pc.grid.y, (Mt, Nt, pc)
        else:
            assert pc.pcm * pc.grid.y >= Mt and pc.pcn * pc.grid.x >= Nt, (Mt, Nt, pc)

    def init_device_compute_kernel_config(self, arch, *, math_fidelity, fp32_dest_acc_en, packer_l1_acc):
        return types.SimpleNamespace(fidelity=math_fidelity, fp32=fp32_dest_acc_en)

    def MinimalMatmulConfig(self, *, M_block_size, K_block_size, N_block_size, subblock_h, subblock_w,
                            compute_with_storage_grid_size):
        assert M_block_size % subblock_h == 0 and N_block_size % subblock_w == 0
        assert subblock_h * subblock_w <= 4, "DST cap with fp32 accumulation is 4 tiles"
        return types.SimpleNamespace(M=M_block_size, K=K_block_size, N=N_block_size)

    def SDPAProgramConfig(self, *, compute_with_storage_grid_size, q_chunk_size, k_chunk_size, exp_approx_mode):
        assert q_chunk_size % TILE == 0 and k_chunk_size % TILE == 0
        return types.SimpleNamespace(q=q_chunk_size, k=k_chunk_size)

    # ---- launches
    def _launch(self, name, out):
        self.launches += 1
        self.calls.append(name)
        return out

    def layer_norm(self, x, *, weight, bias, epsilon=None, memory_config=None, program_config=None):
        x.check(); weight.check(); bias.check()
        if program_config is not None:  # sharded LN: block-sharded in/out, RM [1,1,C/32,32] gamma/beta
            assert x.mem.sharded and memory_config is not None and memory_config.sharded
            assert weight.layout == bias.layout == "RM" and weight.shape[-1] == 32
            assert weight.shape[-2] * 32 == x.shape[-1] == bias.shape[-2] * 32
            gx, gy = program_config.grid
            assert x.shape[-1] // TILE % gx == 0 and _pad32(x.shape[-2]) // TILE % gy == 0
            return self._launch("layer_norm", self._new(x.shape, x.dtype, x.layout, memory_config))
        assert weight.shape[-1] == x.shape[-1] == bias.shape[-1]
        assert epsilon is None or 0 < epsilon <= 1e-5  # fused path passes the reference eps; legacy none
        return self._launch("layer_norm", self._new(x.shape, x.dtype, x.layout, memory_config or x.mem))

    def _mm_checks(self, a, w, bias, ternary=None):
        a.check(); w.check()
        assert a.layout == "TILE" and w.layout == "TILE", "minimal_matmul requires TILE"
        assert a.dtype in ("bf16", "bfp8", "fp32") and w.dtype in ("bf16", "bfp8", "fp32")
        assert all(d == 1 for d in w.shape[:-2]), "weight leading dims must be 1"
        M, K, N = a.shape[-2], a.shape[-1], w.shape[-1]
        assert w.shape[-2] == K, (K, w.shape)
        if bias is not None:
            bias.check()
            assert bias.layout == "TILE" and bias.dtype in ("bf16", "bfp8")
            assert all(d == 1 for d in bias.shape[:-1]) and bias.shape[-1] == N, bias.shape
        if ternary is not None:
            res, ones = ternary
            res.check(); ones.check()
            assert res.layout == "TILE" and ones.layout == "TILE"
            assert res.shape[-2:] == (M, N), (res.shape, M, N)
            assert ones.shape[-2:] == (1, N) and all(d == 1 for d in ones.shape[:-1]), ones.shape
            assert res.dtype == w.dtype, f"dit fused: residual dtype {res.dtype} must equal weight dtype {w.dtype}"
            # Code invariant, NOT a validate rule of this tree: the minimal_matmul program factory
            # appends separate TensorAccessorArgs for the residual and the scale row, so mixed
            # placement would be legal.  The pairing is a precaution carried over from the rf-detr
            # BRIEF (DEVICE_VALIDATION.md §7); this asserts that ``_ones`` keeps doing what it says.
            assert res.mem.buffer_type == ones.mem.buffer_type, "_ones(): scale row placement should follow the residual (precaution)"
        return M, K, N

    def minimal_matmul(self, a, w, bias_tensor=None, *, config, memory_config=None, compute_kernel_config):
        _, _, N = self._mm_checks(a, w, bias_tensor)
        return self._launch("minimal_matmul", self._new(a.shape[:-1] + (N,), a.dtype, "TILE", memory_config or a.mem))

    def dit_minimal_matmul_addcmul_fused(self, a, w, scalar, res, ones, *, bias_tensor=None, config,
                                         memory_config=None, dtype=None, compute_kernel_config):
        assert scalar == 1.0
        _, _, N = self._mm_checks(a, w, bias_tensor, ternary=(res, ones))
        return self._launch("dit_fused", self._new(a.shape[:-1] + (N,), dtype or a.dtype, "TILE", memory_config or a.mem))

    def linear(self, a, w, *, bias=None, activation=None, memory_config=None, core_grid=None,
               program_config=None, compute_kernel_config=None, dtype=None):
        a.check(); w.check()
        assert a.layout == "TILE" and w.layout == "TILE" and a.shape[-1] == w.shape[-2]
        if program_config is not None:
            self._check_pc(program_config, a.shape[-2], a.shape[-1], w.shape[-1])
        if bias is not None:
            bias.check()
            assert len(bias.shape) >= 2 and all(d == 1 for d in bias.shape[:-1]), "row-vector bias [1,..,1,N]"
            assert _pad32(bias.shape[-1]) == _pad32(w.shape[-1])
        assert activation in (None, "gelu")
        if getattr(a.mem, "width_shards", None):  # width-sharded in0: 1D mcast-in0, one in0 block per shard
            n, (sh, sw) = a.mem.width_shards
            assert program_config is not None and program_config.kind == "1d" and program_config.ibw * TILE == sw
            assert sh == a.shape[-2] and n * sw == a.shape[-1]
        odt = {None: a.dtype, self.bfloat16: "bf16", self.float32: "fp32"}.get(dtype, dtype)
        return self._launch("linear", self._new(a.shape[:-1] + (w.shape[-1],), odt, "TILE", memory_config or a.mem))

    def matmul(self, a, b, core_grid=None):
        a.check(); b.check()
        assert a.shape[-1] == b.shape[-2]
        return self._launch("matmul", self._new(a.shape[:-1] + (b.shape[-1],), a.dtype, "TILE", a.mem))

    def add(self, a, b, memory_config=None):
        a.check(); b.check()
        shape = torch.broadcast_shapes(a.shape, b.shape)
        return self._launch("add", self._new(shape, a.dtype, a.layout, memory_config or a.mem))

    def gelu(self, x, fast_and_approximate_mode=False, memory_config=None):
        x.check()
        return self._launch("gelu", self._new(x.shape, x.dtype, x.layout, memory_config or x.mem))

    def reshape(self, x, shape, memory_config=None):
        x.check()
        assert math.prod(x.shape) == math.prod(shape), (x.shape, shape)
        return self._launch("reshape", self._new(shape, x.dtype, x.layout, memory_config or x.mem))

    def permute(self, x, dims):
        x.check()
        assert sorted(dims) == list(range(len(x.shape)))
        return self._launch("permute", self._new(tuple(x.shape[d] for d in dims), x.dtype, x.layout, x.mem))

    def tilize_with_zero_padding(self, x, *, memory_config=None, use_multicore=True):
        x.check()
        assert x.layout == "RM", "tilize input must be ROW_MAJOR"
        assert x.dtype in ("bf16", "fp32", "u32")
        return self._launch("tilize", self._new(x.shape, x.dtype, "TILE", memory_config or x.mem))

    def nlp_create_qkv_heads(self, x, *, num_heads, transpose_k_heads=True, memory_config=None):
        x.check()
        assert len(x.shape) == 4 and x.shape[1] == 1, x.shape
        assert x.layout == "TILE" and x.dtype in ("bf16", "bfp8", "fp32")
        assert x.padded()[2] % TILE == 0, "S must be tile aligned"
        assert x.padded()[3] % (3 * num_heads) == 0
        dh = x.padded()[3] // (3 * num_heads)
        B, S = x.shape[0], x.shape[2]
        mem = memory_config or x.mem
        q = self._new((B, num_heads, S, dh), x.dtype, "TILE", mem)
        k = self._new((B, num_heads, dh, S) if transpose_k_heads else (B, num_heads, S, dh), x.dtype, "TILE", mem)
        v = self._new((B, num_heads, S, dh), x.dtype, "TILE", mem)
        return self._launch("create_qkv_heads", (q, k, v))

    # hamer.tt_kernels generic_op tile permutes (exact; interleaved TILE in/out, batch 1)
    def head_split(self, x, num_heads, head_dim, n_out=3, memory_config=None, grid=None):
        x.check()
        assert len(x.shape) == 4 and x.shape[:2] == (1, 1) and x.layout == "TILE" and not x.mem.sharded
        assert x.shape[2] % TILE == 0 and head_dim % TILE == 0 and x.shape[3] == n_out * num_heads * head_dim
        mem = memory_config or x.mem
        outs = [self._new((1, num_heads, x.shape[2], head_dim), x.dtype, "TILE", mem) for _ in range(n_out)]
        return self._launch("head_split", outs)

    def head_concat(self, x, memory_config=None, grid=None):
        x.check()
        B, H, S, dh = x.shape
        assert B == 1 and S % TILE == 0 and dh % TILE == 0 and x.layout == "TILE"
        return self._launch("head_concat", self._new((1, 1, S, H * dh), x.dtype, "TILE", memory_config or x.mem))

    def addln_constants(self, gamma, beta, device):
        W = gamma.numel()
        up = lambda t: self.from_torch(t.reshape(1, 1, 1, W).to(torch.bfloat16).expand(1, 1, 32, W), device=device,
                                       layout="TILE", dtype="bf16")
        return {"gamma": up(gamma), "beta": up(beta), "W": W}

    def add_layernorm(self, x, y, consts, eps, out_memory_config=None, out_dtype=None, prefetch=None):
        x.check()
        assert x.mem.sharded and x.layout == "TILE" and x.dtype == "bf16" and x.shape[-1] == consts["W"]
        if y is not None:
            y.check()
            assert y.mem.sharded and y.shape == x.shape and y.dtype == x.dtype
        s = self._new(x.shape, x.dtype, "TILE", x.mem) if y is not None else None
        h = self._new(x.shape, out_dtype or "bf16", "TILE", out_memory_config or self.L1_MEMORY_CONFIG)
        self._launch("add_layernorm", h)
        return s, h

    def allocate_tensor_on_device(self, shape, dtype, layout, device, mem):
        return self._new(tuple(shape), dtype, layout, mem)

    def cmm_linear(self, a, w, b, out_memory_config, kb=8, slots=3, unicast=False, fidelity=None, defines=(),
                   layout="out", rotate=False, gelu=False, prefetch=None, a_heads=None, v1=False, ss=3, gelu_from=None,
                   helper_prefetch=None, addln=None, out_dtype=None, gelu_pf=None, idle_pf=None, idle_pf_remote=False):
        assert layout in ("out", "k2")  # round 11: k2 = the 2-way K-split fc2 (same output / epilogue contract)
        for pf_ in (gelu_pf, idle_pf, (addln or {}).get("prefetch2")):  # round 11: the next qkv's XmmPrefetch
            assert pf_ is None or isinstance(pf_, FakeXmmPrefetch)
        if gelu_from is not None:  # round 8: GELU fused into fc2 (helper cores); `a` is the freshly allocated buffer
            gelu_from.check()
            assert gelu_from.dtype == "bf16" and gelu_from.layout == "TILE" and a.shape == gelu_from.shape
        if addln is not None:  # round 8: the add+LN epilogue -> (new residual s, LN output h); x is consumed by the caller
            addln["x"].check()
            assert addln["x"].mem.sharded and addln["x"].shape[-1] == addln["consts"]["W"]
            a.check(); b.check()
            s_ = self._new(addln["x"].shape, "bf16", "TILE", addln["x"].mem)
            h_ = self._new(addln["x"].shape, addln.get("out_dtype") or "bf16", "TILE",
                           addln.get("out_mem") or Mem("L1"))
            self._launch("cmm", h_)  # (the cmm program with its add+LN epilogue)
            return s_, h_
        # hamer.tt_kernels.cmm_linear (TT_FUSED_CMM, default since round 5): packed weight, sharded output
        a.check(); b.check()
        assert isinstance(w, FakeCmmWeight) and w.kb == kb and a.layout == "TILE" and a.dtype in ("bf16", "bfp8")
        if a_heads is not None:
            H, D = a_heads
            assert len(a.shape) == 4 and a.shape[1] == H and a.shape[3] == D and H * D == w.K
            M = a.shape[2]
        else:
            assert a.shape[-1] == w.K
            M = a.shape[-2]
        assert out_memory_config.sharded and b.shape[-1] == w.N
        if prefetch is not None:
            assert prefetch.w is w
        return self._launch("cmm", self._new((1, 1, M, w.N), "bf16", "TILE", out_memory_config))

    def xmm_linear(self, a, w, b, rows, out_dtype=None, qkv=None, slots=3, fidelity=None, defines=(), out=None, la=1,
                   transposed=False, prefetch=None):
        # round 11: hamer.tt_kernels.xmm_linear (qkv): width-sharded in0, packed weight, head-split q / k / v outputs
        a.check(); b.check()
        assert isinstance(w, FakeCmmWeight) and a.layout == "TILE" and a.dtype in ("bf16", "bfp8")  # any placement (accessor)
        assert a.shape[-1] == w.K and b.shape[-1] == w.N and w.GX == 4 * rows and w.NT == 4
        assert prefetch is None or (isinstance(prefetch, FakeXmmPrefetch) and prefetch.w is w and transposed)
        H, Dh = qkv
        M = a.shape[-2]
        assert w.N == 3 * H * Dh
        outs = [self._new((1, H, M, Dh), out_dtype or "bf16", "TILE", Mem("L1")) for _ in range(3)]
        self._launch("xmm", outs[0])
        return outs

    def gelu_kernel(self, x, memory_config=None, grid=None, defines=(), fp32_dest=True, prefetch=None, out_dtype=None):
        # hamer.tt_kernels.gelu (custom erf-GELU generic_op, the default since round 5): TILE bf16 in, bf16 / bfp8 out
        x.check()
        assert x.layout == "TILE" and x.dtype == "bf16"
        return self._launch("gelu", self._new(x.shape, out_dtype or x.dtype, x.layout, memory_config or x.mem))

    def im2col_tile(self, img, memory_config=None, grid=None):
        img.check()
        assert img.shape == (1, 3, 256, 192) and img.dtype == "fp32" and img.layout == "RM"
        assert img.mem.buffer_type == "L1", "16 B-aligned NOC windows: the crop buffer must be in L1"
        return self._launch("im2col", self._new((1, 1, 192, 768), "fp32", "TILE", memory_config or Mem("L1")))

    def nlp_concat_heads(self, x, *, memory_config=None):
        x.check()
        assert len(x.shape) == 4 and x.layout == "TILE" and x.dtype in ("bf16", "bfp8", "fp32")
        B, H, S, dh = x.shape
        return self._launch("concat_heads", self._new((B, 1, S, H * dh), x.dtype, "TILE", memory_config or x.mem))

    def sdpa(self, q, k, v, *, attn_mask=None, is_causal, scale, program_config=None, memory_config=None):
        assert attn_mask is None, "never pass an explicit padding mask (wrong on Blackhole)"
        assert is_causal is False and scale is not None
        for t in (q, k, v):
            t.check()
            assert len(t.shape) == 4 and t.layout == "TILE" and t.dtype in ("bf16", "bfp8")
            assert t.shape[3] % TILE == 0, "no padding allowed on head_dim"
        assert k.shape == v.shape and q.shape[0] == k.shape[0] and q.shape[1] % k.shape[1] == 0 and q.shape[3] == k.shape[3]
        return self._launch("sdpa", self._new(q.shape, q.dtype, "TILE", memory_config or q.mem))


FAKE = FakeTTNN()


class FakeCmmWeight:
    def __init__(self, w_t, dev, GX, NT, kb, dtype, nbanks=0):
        self.K, self.N = w_t.shape
        assert self.N == GX * NT * 32 and (self.K // 32) % kb == 0 and ((self.K // 32) // kb) % 2 == 0
        self.kb, self.NT, self.GX, self.dtype = kb, NT, GX, dtype


class FakeCmmPrefetch:
    def __init__(self, dev, GX, GY, w, pairs, single=False):
        self.w, self.pairs = w, pairs

    def with_weight(self, w):
        o = FakeCmmPrefetch(None, 0, 0, w, self.pairs)
        return o

    def for_k2(self):
        return FakeCmmPrefetch(None, 0, 0, self.w, self.pairs)  # round 11: the k2 fc2's two senders per column

    def with_range(self, lo, hi):
        # round 8: the prefetch blocks split between the add+LN before fc1 and the GELU
        assert 0 <= lo <= hi <= 2 * self.pairs
        return FakeCmmPrefetch(None, 0, 0, self.w, self.pairs)


class FakeXmmPrefetch:
    def __init__(self, dev, w, rows, nblk):
        self.w, self.rows, self.nblk = w, rows, nblk

    def with_weight(self, w):
        return FakeXmmPrefetch(None, w, self.rows, self.nblk)

    def with_range(self, lo, hi):  # round 12: the qkv prefetch split between the proj and fc2 programs
        assert 0 <= lo <= hi <= self.nblk
        o = FakeXmmPrefetch(None, self.w, self.rows, self.nblk)
        o.rng = (lo, hi)
        return o


class FakeDevice:
    def compute_with_storage_grid_size(self):
        return types.SimpleNamespace(x=12, y=10)

    def arch(self):
        return "BLACKHOLE"


@pytest.fixture()
def fake(monkeypatch):
    global FAKE
    FAKE = FakeTTNN()
    monkeypatch.setattr(vit, "ttnn", FAKE)
    monkeypatch.setattr(mh, "ttnn", FAKE)
    monkeypatch.setattr(tk, "head_split", FAKE.head_split)
    monkeypatch.setattr(tk, "head_concat", FAKE.head_concat)
    monkeypatch.setattr(tk, "im2col_tile", FAKE.im2col_tile)
    monkeypatch.setattr(tk, "add_layernorm", FAKE.add_layernorm)
    monkeypatch.setattr(tk, "gelu", FAKE.gelu_kernel)
    monkeypatch.setattr(tk, "cmm_linear", FAKE.cmm_linear)
    monkeypatch.setattr(tk, "CmmWeight", FakeCmmWeight)
    monkeypatch.setattr(tk, "CmmPrefetch", FakeCmmPrefetch)
    monkeypatch.setattr(tk, "XmmPrefetch", FakeXmmPrefetch)
    monkeypatch.setattr(tk, "xmm_linear", FAKE.xmm_linear)
    monkeypatch.setattr(tk, "addln_constants", FAKE.addln_constants)
    return FAKE


@pytest.mark.parametrize("env", [dict(e, TT_FUSED_HEAD="minimal") for e in [
    {"TT_FUSED": "1"},                                                          # defaults: tuned, L1, device patch
    {"TT_FUSED": "1", "TT_FUSED_ADDLN": "0"},                                   # library add + LN + s2i
    {"TT_FUSED": "1", "TT_FUSED_IN0WS": "0"},                                   # addln -> interleaved qkv/fc1 input
    {"TT_FUSED": "1", "TT_FUSED_MM": "minimal", "TT_FUSED_SDPA": "default"},    # the p150a graph (420 launches)
    {"TT_FUSED": "1", "TT_FUSED_L1": "0"},                                      # DRAM placement
    {"TT_FUSED": "1", "TT_FUSED_FC2": "bfp8"},                                  # fc2 unfusable -> linear + add
    {"TT_FUSED": "1", "TT_FUSED_PATCH": "fp32w"},                               # fp32 patch weight / pos / ones
    {"TT_FUSED": "1", "TT_FUSED_PATCH": "host", "TT_FUSED_SDPA": "default"},    # legacy token upload
    {"TT_FUSED": "1", "TT_FUSED_MM": "linear", "TT_FUSED_HEAD_CONCAT": "permute", "TT_FUSED_FC1": "bf16"},
]])
def test_fused_device_graph_walk(ref, image, fake, env):
    """Build the fused parameters and run the fused ViT + head on the mock ttnn: every op's
    documented constraints hold, no tensor is used after deallocation, the input buffer is
    never freed, and the launch count is the device-validated one (2 ingest + 10/block + 1 LN;
    16/layer + 1 dec = 420 at depth 32/6)."""
    opts = FusedOptions.from_env(env)
    dev = FakeDevice()
    vp = vit.build_parameters_from_reference(ref, dev, fused=opts)
    hp = mh.build_parameters_from_reference(ref, dev, fused=opts)
    assert "fused" in vp and "fused" in hp and len(vp["blocks"]) == DEPTH and len(hp["layers"]) == HEAD_DEPTH
    fake.launches = 0
    fake.calls.clear()
    if opts.device_im2col:
        x_in = fake.from_torch(image.float(), device=dev, layout=fake.ROW_MAJOR_LAYOUT, dtype=fake.float32,
                               memory_config=fake.L1_MEMORY_CONFIG)
        ingest_launches = 2  # device im2col + patch matmul
    elif opts.device_patch_embed:
        host = fz.host_im2col(image)
        x_in = fake.from_torch(host, device=dev, layout=fake.ROW_MAJOR_LAYOUT, dtype=fake.float32)
        ingest_launches = 2
    else:
        host = ref.backbone.patch_embed(image)[0].to(torch.bfloat16)
        x_in = fake.from_torch(host, device=dev, layout=fake.TILE_LAYOUT, dtype=fake.bfloat16)
        ingest_launches = 2  # view reshape + pos add
    feat = vit.forward(x_in, vp, depth=DEPTH, fused=True)
    assert feat.shape == (1, 1, fz.N_TOKENS, fz.EMBED) and feat.dtype == "bf16" and feat.layout == "TILE"
    assert feat.mem.buffer_type == ("L1" if opts.l1 else "DRAM")
    vit_launches = fake.launches
    if opts.matmul == "tuned":
        per_block_extra = 2  # proj / fc2: linear(+bias) + residual add
        if opts.vit_shard and opts.l1 and opts.addln and opts.heads == "custom":
            per_block_extra = 0   # add + LN + s2i are ONE fused op (2 per block, counted in the 10)
            if vp["fused"].get("patch_lin") is None:
                ingest_launches += 1  # interleaved -> block-sharded residual stream
            # (round 8 patch_lin: the patch linear writes the sharded stream itself; pos is block 0's add+LN x)
            if vp["fused"].get("cmm"):
                per_block_extra = -1  # cmm proj reads the head-split SDPA output: no head_concat
                if vp["fused"].get("gelu_help"):
                    per_block_extra = -2  # round 8: the GELU runs inside the fc2 cmm program
                if vp["fused"].get("cmm_addln") and vp["fused"].get("patch_lin") is not None:
                    per_block_extra = -4  # round 8: both add+LNs are cmm epilogues (block 0's LN1 replaces the last LN)
                    if vp["blocks"][0]["attn"].get("qkv_xmm") is not None:
                        per_block_extra = -5  # round 11: qkv (xmm) writes the head-split q / k / v (no head_split)
                    if vp["fused"]["patch_lin"].get("cmm") is not None:
                        ingest_launches -= 1  # round 15: patch matmul + block 0's add+LN are ONE cmm program
        elif opts.vit_shard and opts.l1:
            per_block_extra += 2  # sharded LN -> interleaved copy before qkv / fc1
            ingest_launches += 1  # interleaved -> block-sharded residual stream
            ingest_launches += 1  # last LN sharded -> interleaved (counted with the ingest for brevity)
    else:
        per_block_extra = (1 if opts.wdtype("fc2") == "bfp8" or opts.matmul == "linear" else 0) + (1 if opts.matmul == "linear" else 0)
    assert vit_launches == ingest_launches + DEPTH * (10 + per_block_extra) + 1, fake.calls
    dec = mh.forward_device(feat, hp, dev, depth=HEAD_DEPTH, fused=True)
    head_launches = fake.launches - vit_launches
    per_layer_extra = 3 if opts.matmul == "linear" else 0  # 3 residual fusions -> linear + add each
    per_layer_extra += 1 if opts.head_concat == "permute" else 0  # permute + reshape instead of concat_heads
    # 16 per layer since the device pass: fc1 (+bias) and the exact gelu are separate launches
    # (ttnn.linear's fused gelu is not placement-neutral on the p150a, see mano_head._decoder_block_fused)
    assert head_launches == HEAD_DEPTH * (16 + per_layer_extra) + 1, fake.calls[-head_launches:]
    assert dec.shape == (1, 1, 1, fz.DEC_OUT)
    assert x_in._id in fake.alive, "the persistent input buffer must never be deallocated"
    assert hp["x0"]._id in fake.alive, "the constant x0 must never be deallocated"
    if opts.matmul == "minimal" and opts.wdtype("fc2") == "bf16":
        # the full-depth count (evaluation §5.1 said 414 with the fused-gelu linear; +6 since)
        full = 2 + 32 * 10 + 1 + 6 * 16 + 1
        assert full == 420
    if opts.sdpa == "32x192":
        assert vp["fused"]["sdpa_pc"].q == 32 and vp["fused"]["sdpa_pc"].k == 192
    else:
        assert vp["fused"]["sdpa_pc"] is None
    if opts.device_patch_embed:
        pt = vp["fused"]["patch"]
        want = "bf16" if opts.patch == "bf16w" else "fp32"
        assert pt["w"].dtype == pt["pos"].dtype == want and pt["pos"].shape == (1, 1, fz.N_TOKENS, fz.EMBED)
        assert pt["compute"].fidelity == "HiFi4" and pt["compute"].fp32 is True
    if vp["fused"].get("cmm"):
        assert vp["blocks"][0]["mlp"]["fc2_w"] is None and vp["blocks"][0]["mlp"]["fc2_cmm"].dtype == "bfp8"
        assert "cmm" in fake.calls and "concat_heads" not in fake.calls[:vit_launches]
    else:
        assert vp["blocks"][0]["mlp"]["fc2_w"].dtype == opts.wdtype("fc2")
    if vp["blocks"][0]["attn"].get("qkv_xmm") is not None:  # round 11: qkv as xmm_linear (packed bfp8 weight only)
        assert vp["blocks"][0]["attn"]["qkv_w"] is None and vp["blocks"][0]["attn"]["qkv_xmm"].dtype == "bfp8"
        assert "xmm" in fake.calls and "head_split" not in fake.calls[:vit_launches]
    else:
        assert vp["blocks"][0]["attn"]["qkv_w"].dtype == opts.wdtype("qkv_w")
    if vp["fused"].get("cmm"):
        assert vp["blocks"][0]["attn"]["proj_w"] is None and vp["blocks"][0]["attn"]["proj_cmm"].dtype == "bfp8"
    else:
        assert vp["blocks"][0]["attn"]["proj_w"].dtype == opts.wdtype("proj_w")
    assert vp["blocks"][0]["mlp"]["fc1_w"].dtype == ("bfp8" if opts.fc1 == "bfp8" else "bf16")
    assert vp["blocks"][0]["attn"]["qkv_b"].shape == (1, 1, 1, 3 * 16 * fz.VIT_HEAD_DIM_PAD)
    assert hp["layers"][0]["ca_kv_fake_w"].shape == (fz.HEAD_CONTEXT, 3 * fz.HEAD_INNER)
    assert hp["layers"][0]["sa_v_w"].shape == (fz.HEAD_DIM, fz.HEAD_INNER)
    assert hp["dec_b"].shape == (1, 1, 1, fz.DEC_OUT)


def test_legacy_graph_untouched_on_mock(ref, image, fake):
    """The legacy functions still run their 21-launch block / 32-launch layer graphs (and take
    no fused parameters)."""
    dev = FakeDevice()
    vp = vit.build_parameters_from_reference(ref, dev)
    hp = mh.build_parameters_from_reference(ref, dev)
    assert "fused" not in vp and "fused" not in hp and "sa_qkv_w" in hp["layers"][0]
    fake.launches = 0
    x_in = fake.from_torch(ref.backbone.patch_embed(image)[0].to(torch.bfloat16), device=dev,
                           layout=fake.TILE_LAYOUT, dtype=fake.bfloat16)
    monkey_grid = vit._MLP_GRID
    vit._MLP_GRID = None
    try:
        feat = vit.forward(x_in, vp, depth=DEPTH)
    finally:
        vit._MLP_GRID = monkey_grid
    assert feat.shape == (1, fz.N_TOKENS, fz.EMBED)
    assert fake.launches == 1 + DEPTH * 21 + 1
    n0 = fake.launches
    token = fake.from_torch(torch.zeros(1, 1, 1, dtype=torch.bfloat16), device=dev, layout="TILE", dtype="bf16")
    dec = mh.forward_device(feat, hp, dev, depth=HEAD_DEPTH, cached_token=(token,))
    assert fake.launches - n0 == 3 + HEAD_DEPTH * 32 + 2
    assert dec.shape == (1, 1, fz.DEC_OUT)


def test_head_cached_token_none_passthrough_when_fused(ref, image, fake):
    """``TtHamer._forward_device`` passes ``cached_token=(self._head_token,)`` on both paths; on
    the fused one ``_head_token`` is ``None`` (the head starts from the constant ``x0``) and
    ``forward_device`` must never look at it."""
    dev = FakeDevice()
    opts = FusedOptions.from_env({"TT_FUSED": "1", "TT_FUSED_HEAD": "minimal"})
    vp = vit.build_parameters_from_reference(ref, dev, fused=opts)
    hp = mh.build_parameters_from_reference(ref, dev, fused=opts)
    x_in = fake.from_torch(image.float(), device=dev, layout=fake.ROW_MAJOR_LAYOUT, dtype=fake.float32,
                           memory_config=fake.L1_MEMORY_CONFIG)
    feat = vit.forward(x_in, vp, depth=DEPTH, fused=True)
    dec = mh.forward_device(feat, hp, dev, depth=HEAD_DEPTH, cached_token=(None,), fused=True)
    assert dec.shape == (1, 1, 1, fz.DEC_OUT)


# =========================================================================== #
# 6. model.py plumbing on the mock (no device is ever opened: HAMER_USE_TT=0, device injected)
# =========================================================================== #
def _wrapper_on_fake(ref, fake, monkeypatch, env):
    """``TtHamer`` built without a chip (``HAMER_USE_TT=0``), then given the FakeDevice so the
    per-request plumbing (``_host_input`` / ``_upload_tokens`` / ``_forward_device`` guard) runs
    against the mock ``ttnn`` that ``model.py`` imports lazily."""
    from hamer.model import TtHamer

    monkeypatch.setitem(sys.modules, "ttnn", fake)
    for k in list(os.environ):
        if k.startswith("TT_FUSED"):
            monkeypatch.delenv(k, raising=False)
    for k, v in env.items():
        monkeypatch.setenv(k, v)
    m = TtHamer(ref, strict=False, use_trace=True)
    assert m.device is None and m._vit_params is None  # HAMER_USE_TT=0: nothing opened
    m.device = FakeDevice()
    m._is_mesh = False
    return m


@pytest.mark.parametrize("env, want", [
    ({"TT_FUSED": "1"}, ("fp32", "RM", (1, fz.IN_CHANS, fz.IMG_H, fz.IMG_W))),             # device im2col: raw crop
    ({"TT_FUSED": "1", "TT_FUSED_PATCH": "fp32w"}, ("fp32", "RM", (1, 1, fz.N_TOKENS, fz.PATCH_DIM))),
    ({"TT_FUSED": "1", "TT_FUSED_PATCH": "host"}, ("bf16", "TILE", (1, fz.N_TOKENS, fz.EMBED))),  # legacy CPU conv
    ({"TT_FUSED": "0"}, ("bf16", "TILE", (1, fz.N_TOKENS, fz.EMBED))),                        # legacy path
])
def test_host_input_selects_dtype_layout_per_knob(ref, image, fake, monkeypatch, env, want):
    m = _wrapper_on_fake(ref, fake, monkeypatch, env)
    host, dtype, layout = m._host_input(image)
    assert (dtype, layout, tuple(host.shape)) == want
    assert host.is_contiguous()
    if dtype == "fp32" and m.fused_opts.device_im2col:
        assert host.dtype == torch.float32 and torch.equal(host, image.float())
    elif dtype == "fp32":
        assert host.dtype == torch.float32 and torch.equal(host, fz.host_im2col(image))
    else:
        assert host.dtype == torch.bfloat16
        assert torch.equal(host, ref.backbone.patch_embed(image)[0].to(torch.bfloat16))


@pytest.mark.parametrize("env", [{"TT_FUSED": "1"}, {"TT_FUSED": "1", "TT_FUSED_HASH": "0"}, {"TT_FUSED": "0"}])
def test_upload_tokens_persistent_buffer_and_hash_cache(ref, image, fake, monkeypatch, env):
    """First upload allocates the persistent buffer; a repeated frame is served from the
    content cache (no upload) unless ``TT_FUSED_HASH=0``; a new frame of the same shape is
    copied INTO the same buffer (identity preserved -> the trace stays valid)."""
    m = _wrapper_on_fake(ref, fake, monkeypatch, env)
    use_cache = not (m.fused and not m.fused_opts.hash_cache)
    other = sample_input(batch=1, seed=11)

    buf = m._upload_tokens(image)
    assert buf.on_device and len(fake.uploads) == 1 and fake.copies == []
    assert (m._tokens_key is not None) == use_cache
    again = m._upload_tokens(image)                       # repeated frame
    assert again is buf
    assert len(fake.copies) == (0 if use_cache else 1)
    fresh = m._upload_tokens(other)                       # new content, same shape
    assert fresh is buf, "same shape -> copy into the persistent buffer, never a new allocation"
    assert len(fake.copies) == (1 if use_cache else 2)
    assert len(fake.uploads) == 1 + len(fake.copies)      # copies go through a host-only from_torch
    assert all(not h.on_device for h, _ in fake.copies)
    back = m._upload_tokens(image)
    assert back is buf and len(fake.copies) == (2 if use_cache else 3)
    assert fake.released == []


def test_upload_tokens_shape_change_reallocates_and_voids_trace(ref, image, fake, monkeypatch):
    """Legacy path (batching is legacy-only): a different batch size (re)allocates the buffer
    and releases the trace bound to the old one."""
    m = _wrapper_on_fake(ref, fake, monkeypatch, {"TT_FUSED": "0"})
    buf1 = m._upload_tokens(image)
    m._trace = (7, None)
    buf2 = m._upload_tokens(sample_input(batch=2, seed=5))
    assert buf2 is not buf1 and buf2.shape == (2, fz.N_TOKENS, fz.EMBED)
    assert m._trace is None and fake.released == [7]
    assert fake.copies == [] and len(fake.uploads) == 2


def test_fused_batch_guard_is_loud_even_when_not_strict(ref, fake, monkeypatch):
    """Review fix: the fused graph is batch-1 only (minimal_matmul folds B into M, validate
    compares the logical [-2] only, the head hard-codes the (1,1,H,dh) q reshape).  The guard
    must raise -- not become the non-strict CPU fallback that pins the harness accuracy to 100."""
    m = _wrapper_on_fake(ref, fake, monkeypatch, {"TT_FUSED": "1"})
    m._vit_params = {}  # pretend the upload succeeded so _forward_device reaches the guard
    m.strict = False
    two = sample_input(batch=2, seed=1)
    with pytest.raises(ValueError, match="batch 1"):
        m(two)
    assert fake.uploads == [] and fake.launches == 0, "the guard fires before any device work"
    # the legacy wrapper still accepts batch > 1 at the same point (falls through to the try)
    legacy = _wrapper_on_fake(ref, fake, monkeypatch, {"TT_FUSED": "0"})
    legacy._vit_params = None
    assert legacy._forward_device(two) is None  # no params -> documented non-strict None
    # and the direct fused ViT entry rejects it too
    opts = FusedOptions.from_env({"TT_FUSED": "1"})
    vp = vit.build_parameters_from_reference(ref, FakeDevice(), fused=opts)
    x2 = fake.from_torch(fz.host_im2col(two), device=FakeDevice(), layout=fake.ROW_MAJOR_LAYOUT, dtype=fake.float32)
    with pytest.raises(ValueError, match="batch 1"):
        vit.forward(x2, vp, depth=DEPTH, fused=True)


if __name__ == "__main__":  # plain-script fallback when pytest is unavailable
    sys.exit(pytest.main([__file__, "-q"]))


def test_u8_input_folding_is_exact(ref):
    """TT_FUSED_INPUT=u8: im2col(u - o) @ (a * W) + (pos_plus_bias + im2col(B) @ W) == im2col(x) @ W + pos_plus_bias
    for the normalized x of an 8-bit crop u (fp64 algebra; the padding stays zero in normalized space)."""
    from hamer.fused import (centred_pixels, host_im2col, patch_constants, patch_constants_u8,
                             pixels_from_normalized)
    from hamer.preprocess import normalize_raw
    g = torch.Generator().manual_seed(3)
    u = torch.randint(0, 256, (1, 3, 256, 192), generator=g, dtype=torch.uint8)
    x = normalize_raw(u)
    assert torch.equal(pixels_from_normalized(x), u.float())
    w, ppb = patch_constants(ref.backbone)
    w8, ppb8 = patch_constants_u8(ref.backbone)
    a = host_im2col(x).double() @ w.double() + ppb.double()
    v = centred_pixels(u)
    assert torch.equal(v.float() + torch.tensor([124., 116., 104.]).view(1, 3, 1, 1), u.float())
    b = host_im2col(v.float()).double() @ w8.double() + ppb8.double()
    assert float((a - b).abs().max()) < 2e-4 * float(a.abs().max())
    with pytest.raises(ValueError):
        pixels_from_normalized(torch.randn(1, 3, 256, 192))


def test_preprocess_box_resize_bit_identical():
    """Round 6: preprocess_pil resizes only the centre 192 columns (PIL resize box) -- bit-identical to the
    original resize-to-256x256 + flip + [32:224] slice, for random boxes (inside / crossing the frame), both hands."""
    import random
    import numpy as np
    from PIL import Image
    from hamer.preprocess import crop_bbox, preprocess_pil, W_SLICE
    media = Path(__file__).resolve().parents[2] / "media"
    rng = random.Random(0)
    ims = [Image.open(media / f"sample{i}.jpg").convert("RGB") for i in range(1, 5)]
    for _ in range(60):
        im = rng.choice(ims)
        x1, y1 = rng.uniform(-60, im.width - 10), rng.uniform(-60, im.height - 10)
        bb = (x1, y1, x1 + rng.uniform(8, 300), y1 + rng.uniform(8, 300))
        right = rng.random() < 0.5
        c256, meta = crop_bbox(im, bb, 2.0)
        if not right:
            c256 = c256.transpose(Image.FLIP_LEFT_RIGHT)
        want = np.asarray(c256)[:, W_SLICE:256 - W_SLICE].transpose(2, 0, 1)[None]
        got, meta2 = preprocess_pil(im, bb, is_right=right, raw=True)
        assert meta2 == meta and np.array_equal(got.numpy(), want)


def test_host_finalizer_tile_gather_bit_identical():
    """Round 13: ``HostFinalizer`` (row 0 gathered from the raw fp32 TILE bytes of a [1, 1, 32, 128] output, the
    init vector held by the object) == ``host_finalize`` on the untilized row, bit for bit."""
    rng = np.random.default_rng(13)
    split = (96, 10, 3)
    ip, ib, ic = (torch.from_numpy(rng.standard_normal((1, n)).astype(np.float32)) for n in split)
    init = mh.finalize_init_vector(ip, ib, ic)
    fin = mh.HostFinalizer(split, init, mh.HostFinalizer.tile_row0_index(109))
    for k in range(300):
        rows = (rng.standard_normal((32, 128)) * rng.choice([1e-3, 1.0, 40.0])).astype(np.float32)
        # tilize: 4 tiles of 32x32 along the width, each 4 faces of 16x16 (row-major faces)
        flat = rows.reshape(2, 16, 4, 2, 16).transpose(2, 0, 3, 1, 4).reshape(-1)
        a = fin(flat)
        b = mh.host_finalize(torch.from_numpy(rows[:1, :109].copy()), split, ip, ib, ic)
        c = mh.host_finalize(torch.from_numpy(rows[:1, :109].copy()), split, ip, ib, ic, init=init)
        assert torch.equal(a, b) and torch.equal(b, c)


def test_banded_resize_bit_identical():
    """Round 13: the threaded row-band resize (preprocess.RESIZE_BANDS) == the one-call resize, bit for bit, on random
    images / boxes (out-of-frame squares, odd rescale factors, left hands); the contiguous=False view has the same values."""
    from PIL import Image
    from hamer import preprocess as pp
    rng = np.random.default_rng(7)
    for k in range(120):
        W, H = int(rng.integers(64, 1200)), int(rng.integers(64, 1200))
        img = Image.fromarray(rng.integers(0, 256, (H, W, 3), dtype=np.uint8))
        x1, y1 = rng.uniform(0, W - 8), rng.uniform(0, H - 8)
        bbox = (x1, y1, x1 + rng.uniform(4, W - x1), y1 + rng.uniform(4, H - y1))
        right = bool(rng.integers(0, 2))
        rf = float(rng.choice([2.0, rng.uniform(0.6, 5.0)]))
        crop, _ = pp._square_crop(img, bbox, rf)
        one = pp.resize_crop_hw(crop, right, bands=1)
        for nb in (2, 3, 4):
            assert np.array_equal(one, pp.resize_crop_hw(crop, right, bands=nb)), (W, H, bbox, rf, right, nb)
        a = pp.preprocess_pil(img, bbox, rescale_factor=rf, is_right=right, raw=True)[0]
        b = pp.preprocess_pil(img, bbox, rescale_factor=rf, is_right=right, raw=True, contiguous=False)[0]
        assert a.is_contiguous() and torch.equal(a, b)


def test_native_crop_resize_bit_identical(monkeypatch):
    """Round 14 (P1): hamer/native/resize.c (crop + grey pad + PIL BILINEAR box resize + mirror in one call) == the PIL
    path, bit for bit: random frames, boxes inside / crossing / beyond the frame, odd rescale factors, both hands, plus
    the media crops; the bbox meta is identical too.  Skipped where no C compiler is available (the server then uses PIL)."""
    from PIL import Image
    from hamer import native
    from hamer import preprocess as pp
    if native.load() is None:
        pytest.skip(f"native preprocess unavailable: {native.STATUS}")
    rng = np.random.default_rng(14)
    cases = []
    media = Path(__file__).resolve().parents[2] / "media"
    import json
    for name, ents in json.loads((media / "bboxes.json").read_text()).items():
        im = Image.open(media / name).convert("RGB")
        cases += [(im, tuple(e["bbox"]), 2.0, bool(e.get("is_right", 1))) for e in ents]
    for k in range(300):
        W, H = int(rng.integers(16, 1100)), int(rng.integers(16, 1100))
        a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8)
        if k % 2:
            a = np.cumsum(a, axis=1, dtype=np.uint8)
        w, h = rng.uniform(1.5, 1.3 * W), rng.uniform(1.5, 1.3 * H)
        x1, y1 = rng.uniform(-0.7 * w, W - 0.3 * w), rng.uniform(-0.7 * h, H - 0.3 * h)
        rf = float(rng.choice([2.0, 1.0, 0.75, rng.uniform(0.51, 5.0)]))
        cases.append((Image.fromarray(a), (x1, y1, x1 + w, y1 + h), rf, bool(k % 3)))
    n = 0
    for im, bb, rf, right in cases:
        if pp.hamer_bbox_size(bb, rf) < 2:
            continue
        got = native.crop_resize(np.asarray(im), bb, rf, right)
        crop, meta = pp._square_crop(im, bb, rf)
        want = pp.resize_crop_hw(crop, right, bands=1)
        if got is None:  # degenerate box: the PIL path handles it
            continue
        n += 1
        assert got[1] == meta and np.array_equal(got[0], want), (im.size, bb, rf, right)
        # round 15: the zero-copy path (the frame's RGBX buffer read in place through Pillow's Arrow export) == the same bytes
        zc = native.crop_resize_image(im, bb, rf, right)
        assert zc is not None and zc[1] == meta and np.array_equal(zc[0], want), ("zero-copy", im.size, bb, rf, right)
    assert n > 250
    # preprocess_pil end to end (raw u8 CHW, raw=False normalized fp32), native vs HAMER_PRE_NATIVE=0
    for im, bb, rf, right in cases[:8] + cases[40:48]:
        if pp.hamer_bbox_size(bb, rf) < 2:
            continue
        monkeypatch.delenv("HAMER_PRE_NATIVE", raising=False)
        a = [pp.preprocess_pil(im, bb, rescale_factor=rf, is_right=right, raw=r) for r in (True, False)]
        monkeypatch.setenv("HAMER_PRE_NATIVE", "0")
        b = [pp.preprocess_pil(im, bb, rescale_factor=rf, is_right=right, raw=r) for r in (True, False)]
        assert a[0][0].is_contiguous() and tuple(a[0][0].shape) == (1, 3, 256, 192)
        for x, y in zip(a, b):
            assert torch.equal(x[0], y[0]) and x[1] == y[1]


def test_tiled_input_im2col_bit_identical():
    """Round 14 (TT_FUSED_TILE_IN): the host-built TILE-order im2col of the centred crop (index table and the native C
    writer) == ttnn's host tilize of ``host_im2col(centred_pixels(u))``, for contiguous CHW crops and strided HWC views."""
    import ttnn
    from hamer import native
    from hamer.fused import PIXEL_OFFSET, centred_pixels, host_im2col, im2col_tile_index
    from hamer.model import TtHamer
    idx = im2col_tile_index()
    io = {"tile_idx": idx, "v_ext": torch.zeros(3 * 256 * 192 + 1, dtype=torch.bfloat16),
          "in_t": torch.empty(192 * 768, dtype=torch.bfloat16), "im2col": native.im2col_tile_fn(),
          "offset": torch.tensor(PIXEL_OFFSET, dtype=torch.bfloat16).view(1, 3, 1, 1)}
    g = torch.Generator().manual_seed(3)
    for k in range(6):
        hwc = torch.randint(0, 256, (256, 192, 3), generator=g, dtype=torch.uint8)
        u = hwc.permute(2, 0, 1).unsqueeze(0)
        if k % 2:
            u = u.contiguous()
        m = host_im2col(centred_pixels(u).float()).to(torch.bfloat16)
        h = ttnn.from_torch(m, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16)
        want = torch.utils.dlpack.from_dlpack(h.host_buffer().get_shard(ttnn.MeshCoordinate(0, 0))).view(torch.bfloat16).view(-1).clone()
        TtHamer._write_input(None, io, u, force_torch=True)
        assert torch.equal(io["in_t"], want)
        if io["im2col"] is not None:
            io["in_t"].zero_()
            TtHamer._write_input(None, io, u)
            assert torch.equal(io["in_t"], want)