Add files using upload-large-folder tool

README.md

@@ -0,0 +1,75 @@
+---
+tags:
+- blackhole
+- tt-model-cache
+- tt-model-container
+- vllm-plugin
+---
+
+# qwen3.8-flash-next-p300x2
+
+Qwen3.8-Flash-Next served as an OpenAI-compatible endpoint on two P300
+boards (four Blackhole devices). The optimized runtime uses TP4+EP4 with
+all 512 routed experts resident on-device; only the PLE n-gram table and
+its sparse row assembly remain host-backed.
+
+Runs on **p300x2** (mesh `(4, 1)`) — 262,144-token context, up to 2 concurrent sequences.
+
+Packaged and published with [tt-model-manager](https://github.com/tenstorrent/tt-model-manager) 0.1.0 (manifest schema 5.1).
+
+## Quickstart
+
+```bash
+tt-model pull  tt-hous/qwen3.8-flash-next-p300x2
+tt-model serve tt-hous/qwen3.8-flash-next-p300x2
+```
+
+`pull` downloads the Docker image and the [`Qwen/Qwen3.8-Flash-Next`](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) weights at `f5d08274bafd880402bd16f5e3e6c514136ec06c` (into your HF cache; they are not in the image). `serve` starts an OpenAI-compatible server on port 8000; the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs `Application startup complete`.
+
+### Serve on two P300 boards (four chips)
+
+```bash
+tt-model serve tt-hous/qwen3.8-flash-next-p300x2 --follow
+```
+
+The package pins the exact Hugging Face checkpoint revision, vLLM 0.24.0,
+Transformers 5.16.0, and the Qwen3.8-capable TT plugin revision. It exposes
+two virtual request slots over a physical batch-1 decode trace.
+
+## Endpoint performance
+
+Measured with `vllm bench serve` against this packaged endpoint on four
+devices. Requests use exact random token lengths, greedy temperature 0,
+EOS ignored, and an unmeasured warmup probe. Decode tok/s/user is derived
+from mean TPOT; aggregate output throughput includes TTFT.
+
+| ISL | OSL | Max concurrency | Requests | Median TTFT | Mean TPOT | Decode tok/s/user | Aggregate output tok/s |
+| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 128 | 128 | 1 | 3/3 | 504 ms | 181.61 ms | 5.506 | 5.229 |
+| 1,024 | 128 | 1 | 3/3 | 3.818 s | 181.70 ms | 5.504 | 4.625 |
+| 4,096 | 128 | 1 | 3/3 | 15.052 s | 181.81 ms | 5.500 | 3.289 |
+| 128 | 128 | 2 | 4/4 | 1.955 s | 368.36 ms | 2.715 | 5.253 |
+
+The two virtual request slots share a physical-B1 decode trace. Concurrency
+2 is state-safe but does not improve aggregate throughput, so concurrency 1
+is the headline per-user profile. The canonical direct model harness is
+faster at 99.940 ms TPOT / 10.006 tok/s/user; the endpoint figures include
+vLLM scheduling, virtual-slot state, PLE service, and plugin handoff.
+
+Runtime telemetry reported all 48 expert layers resident with
+16,986,931,200 expert bytes per device, zero expert host-store bytes, zero
+expert H2D, and zero route D2H/stall time. Only PLE n-gram row lookup and
+selected-row DMA remain host-backed.
+
+## Provenance
+
+The exact sources the image was built from — `code/` in this repo is byte-identical to the model code inside the image:
+
+| component | built from |
+| --- | --- |
+| tt-metal | a local checkout — commit not published *(dirty tree — the image includes uncommitted changes)* |
+| vLLM | [`v0.24.0`](https://github.com/vllm-project/vllm/releases/tag/v0.24.0) |
+| vllm-tt-plugin | a local checkout — commit not published |
+| `code/` digest | `bc605f7e1f6f2342` (sha256, first 16 hex digits) |
+| built | 2026-09-03T04:02:56+00:00 by tt-model 0.1.0 |
+

code/models/common/README.md

@@ -0,0 +1,319 @@
+# TTNN Validation & Testing Utilities
+
+Helpers for validating TTNN computations against reference implementations and
+for moving tensors between TTNN and PyTorch. The public API is implemented
+across `models.common.validation_tools`, `models.common.metrics`,
+`models.common.auto_compose`, and `models.common.distribute_as`, and is
+exercised in:
+
+- `models/common/tests/test_validation_tools.py`
+- `models/common/tests/test_metrics.py`
+- `models/common/tests/test_auto_compose.py`
+- `models/common/tests/test_distribute_as.py`
+- `models/common/tests/host/test_metrics_pytorch_only.py`
+
+The examples in these tests are the most up‑to‑date reference for usage.
+
+## Quick Start – host reference (`compare_to_torch`)
+
+Use `compare_to_torch` when your reference implementation is a PyTorch function.
+Inputs and outputs are automatically converted between TTNN and PyTorch.
+
+```python
+import torch
+import ttnn
+from models.common.validation_tools import compare_to_torch, Metric, get_validation_registry
+
+
+@compare_to_torch(
+    reference_fn=torch.matmul,
+    metric_tolerances={
+        Metric.MAX_ABS_ERROR: 1e-1,
+        Metric.PCC: 0.99,
+    },
+)
+def ttnn_matmul(a, b):
+    # a, b are TTNN tensors (possibly sharded)
+    return ttnn.matmul(a, b)
+
+
+def run_example(device: ttnn.MeshDevice):
+    m, n, k = 16, 24, 12
+    a = torch.randn(1, m, k, dtype=torch.bfloat16)
+    b = torch.randn(1, k, n, dtype=torch.bfloat16)
+
+    a_tt = ttnn.from_torch(a.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    _ = ttnn_matmul(a_tt, b_tt)
+
+    registry = get_validation_registry()
+    registry.print_report()
+```
+
+Note:
+- When the signature of the decorated function is different from the reference function, `input_to_torch` and `output_to_torch` can be used to map the inputs and outputs between the decorated function and the reference function.
+- See `models/common/tests/test_validation_tools.py::test_validation_matmul` for a real test using
+this pattern.
+
+## Quick Start – TTNN reference (`compare_to_ttnn`)
+
+Use `compare_to_ttnn` when both your implementation and reference are TTNN‑based
+and you want metrics computed directly on device.
+
+```python
+import torch
+import ttnn
+from models.common.validation_tools import compare_to_ttnn
+
+
+def torch_rms_norm(x, weight, eps=1e-6):
+    var = x.pow(2).mean(-1, keepdim=True)
+    return weight * x * torch.rsqrt(var + eps)
+
+
+class DeviceValidatedRMSNorm:
+    def __init__(self, weight: torch.Tensor, eps: float, device: ttnn.MeshDevice):
+        self.eps = eps
+        self.device = device
+        self.weight_torch = weight
+        self.weight = ttnn.from_torch(
+            weight.unsqueeze(0).unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+        )
+
+    def _reference_impl(self, x):
+        x_torch = ttnn.to_torch(x).squeeze(0)
+        y_torch = torch_rms_norm(x_torch, self.weight_torch, self.eps)
+        return ttnn.from_torch(
+            y_torch.unsqueeze(0), device=self.device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+        )
+
+    @compare_to_ttnn(reference_fn=lambda self, x: self._reference_impl(x))
+    def __call__(self, x):
+        x_sq = ttnn.mul(x, x)
+        mean_x_sq = ttnn.mean(x_sq, dim=-1, keepdim=True)
+        rms = ttnn.sqrt(ttnn.add(mean_x_sq, self.eps))
+        x_norm = ttnn.mul(x, ttnn.reciprocal(rms))
+        return ttnn.mul(x_norm, self.weight)
+```
+Note:
+- When the signature of the decorated function is different from the reference function, `input_to_ttnn` and `output_to_ttnn` can be used to map the inputs and outputs between the decorated function and the reference function.
+- This mirrors the pattern used by `DeviceValidatedRMSNorm` in
+`models/common/tests/test_validation_tools.py`.
+
+## Features
+
+- **Decorator‑based validation** – Wrap TTNN functions or methods and compare
+  them against PyTorch or TTNN references.
+- **Host and device modes** – `compare_to_torch` (PyTorch reference) and
+  `compare_to_ttnn` (TTNN reference).
+- **TTNN‑native metrics** – When both outputs are TTNN tensors, metrics are
+  computed on device with minimal host transfer.
+- **Configurable tolerances** – Per‑metric tolerances via the `Metric` enum,
+  string keys, or `MetricSpec`.
+- **Custom metrics** – Inject your own metric functions.
+- **Global registry** – Collects all validation runs for reporting.
+- **Easy disabling** – Turn validation on/off globally without changing call
+  sites.
+
+## Core Components
+
+### Validation decorators
+
+All decorators live in `models.common.validation_tools`:
+
+- `compare_to_torch(reference_fn, *, input_to_torch=None, output_to_torch=None, metric_tolerances=None, enabled=True, raise_exceptions=False, return_reference_output=False)`
+  - Use when `reference_fn` is a PyTorch implementation.
+  - By default, all TTNN tensors in the arguments/outputs are converted to
+    PyTorch via `to_torch_auto_compose`.
+  - Optional `input_to_torch(*args, **kwargs)` lets you override how inputs
+    are mapped to the reference.
+  - Optional `output_to_torch(output)` converts the implementation output
+    before metrics are computed.
+
+- `compare_to_ttnn(reference_fn, *, input_to_ttnn=None, output_to_ttnn=None, metric_tolerances=None, enabled=True, raise_exceptions=False, return_reference_output=False)`
+  - Use when `reference_fn` consumes and returns TTNN tensors.
+  - Optional `input_to_ttnn(*args, **kwargs)` lets you override how inputs
+    are mapped to the reference.
+  - Optional `output_to_ttnn(output)` converts the implementation output
+    before metrics are computed.
+  - If both implementation and reference return TTNN tensors, metrics run
+    entirely on device.
+
+In both cases, decorating a function records a `ValidationResult` in the global
+`ValidationRegistry` every time the function is called (unless disabled).
+
+### Metrics
+
+Metric utilities are implemented in `models.common.metrics`:
+
+- `compute_max_abs_error(impl, ref)` – max absolute error.
+- `compute_mean_abs_error(impl, ref)` – mean absolute error.
+- `compute_pcc(impl, ref)` – Pearson correlation coefficient; uses TTNN
+  operations when possible and falls back to host.
+- `comp_allclose(impl, ref, rtol=..., atol=...)` – allclose check plus a
+  detailed delta string.
+- `DEFAULT_METRICS` – dict with built‑in metrics (`"max_abs_error"`,
+  `"mean_abs_error"`, `"pcc"`).
+
+Metrics support both TTNN and PyTorch tensors.
+
+### Registry and control functions
+
+From `models.common.validation_tools`:
+
+- `get_validation_registry() -> ValidationRegistry`
+  - Holds all `ValidationResult` objects.
+  - Provides `get_summary()` and `print_report(verbose: bool = False)`.
+
+- `enable_validation(enabled: bool = True)`
+  - Globally enable/disable validation; when disabled, decorators become
+    transparent wrappers.
+
+- `clear_validation_results()`
+  - Clear all accumulated validation results.
+
+`ValidationResult` includes:
+
+- `function_name`
+- `passed` (bool)
+- `metrics` – map of metric name → per‑metric result (value, passed, error)
+- `execution_time_impl`, `execution_time_ref`
+- `timestamp`
+- `logs` – optional debug strings
+
+### Auto‑compose helper
+
+`to_torch_auto_compose` lives in `models.common.auto_compose`.
+
+It converts an arbitrary TTNN tensor (including sharded/replicated multi‑device
+tensors) to a single PyTorch tensor by automatically choosing the appropriate
+mesh composer.
+
+It is heavily used in:
+
+- `test_auto_compose.py`
+- `test_distribute_as.py`
+- all `compare_to_torch`‑based examples.
+
+## Usage Patterns
+
+High‑level patterns illustrated in the tests:
+
+1. **Host reference with explicit input mapping**
+   - See `HostValidatedRMSNorm` in `models/common/tests/test_validation_tools.py`.
+   - Uses `compare_to_torch` with `input_to_torch` to map TTNN inputs and
+     TTNN weights to a pure‑PyTorch reference function.
+
+2. **TTNN reference (on‑device metrics)**
+   - See `DeviceValidatedRMSNorm` in `models/common/tests/test_validation_tools.py`.
+   - Uses `compare_to_ttnn` where both implementation and reference return
+     TTNN tensors; metrics run on device.
+
+3. **Simple library calls**
+   - See `ttnn_matmul` and `ttnn_matmul_reverse` in `models/common/tests/test_validation_tools.py`.
+   - `compare_to_torch(reference_fn=torch.matmul, ...)` with optional
+     `input_to_torch` remapping.
+
+4. **Checkpoint / `from_torch` validation**
+   - See `from_torch_checkpoint` in `models/common/tests/test_validation_tools.py`.
+   - Validates a direct `ttnn.from_torch(...)` call using `compare_to_torch`
+     and `output_to_torch`.
+
+5. **Custom metric via `MetricSpec`**
+   - See `ttnn_matmul_metric_spec` in `models/common/tests/test_validation_tools.py`
+     and `MetricSpec` usage in `models/common/tests/host/test_metrics_pytorch_only.py`.
+   - Use `MetricSpec(tolerance=..., higher_is_better=..., compute_fn=...)`
+     in `metric_tolerances`.
+
+6. **Non‑decorator usage**
+   - `test_validation_non_decorator_class_vs_class_torch` demonstrates calling
+     `compare_to_torch` in a more manual, non‑decorator style between two
+     callable classes.
+
+## Default Metrics and Tolerances
+
+When `metric_tolerances` is omitted, the framework uses sensible defaults:
+
+- `Metric.MAX_ABS_ERROR` with tolerance `1e-2`
+- `Metric.PCC` with tolerance `0.99`
+
+If you pass a `metric_tolerances` dict, keys can be:
+
+- `Metric` enum members (recommended), e.g. `Metric.MAX_ABS_ERROR`
+- strings (`"max_abs_error"`, `"mean_abs_error"`, `"pcc"`)
+- arbitrary names when used with `MetricSpec`
+
+Values can be:
+
+- a float tolerance (uses the built‑in metric)
+- a `MetricSpec` instance to define a custom metric and tolerance
+
+Example:
+
+```python
+from models.common.validation_tools import Metric, MetricSpec
+from models.common.metrics import compute_pcc
+
+
+@compare_to_torch(
+    reference_fn=torch.matmul,
+    metric_tolerances={
+        Metric.PCC: MetricSpec(tolerance=0.99, higher_is_better=True, compute_fn=compute_pcc),
+        Metric.MAX_ABS_ERROR: 1.5e-1,
+    },
+)
+def ttnn_matmul_metric_spec(a, b):
+    return ttnn.matmul(a, b)
+```
+
+## Testing
+
+The local test suite in `models/common/tests` shows end‑to‑end usage:
+
+- `test_validation_tools.py`
+  - Core decorator usage, registry behaviour, error handling, custom metrics.
+- `test_metrics.py`
+  - Numerical correctness of device and host metric functions.
+- `host/test_metrics_pytorch_only.py`
+  - Pure‑PyTorch metric tests.
+- `test_auto_compose.py`
+  - Auto‑composition of sharded/replicated TTNN tensors into PyTorch.
+- `test_distribute_as.py`
+  - Distribution helpers (`from_torch_dist_as`) that mirror an existing TTNN
+    tensor’s topology.
+
+Example commands (run from the repo root, with TTNN available):
+
+```bash
+python -m pytest models/common/tests/test_validation_tools.py -v
+python -m pytest models/common/tests/test_metrics.py -v
+python -m pytest models/common/tests/host/test_metrics_pytorch_only.py -v
+```
+
+## API Reference (public surface)
+
+All symbols below are imported from `models.common.validation_tools` and `models.common.metrics`:
+
+- Decorators:
+  - `compare_to_torch`
+  - `compare_to_ttnn`
+- Registry and control:
+  - `ValidationResult`
+  - `ValidationRegistry`
+  - `get_validation_registry`
+  - `enable_validation`
+  - `clear_validation_results`
+- Metrics:
+  - `Metric` (enum: `MAX_ABS_ERROR`, `MEAN_ABS_ERROR`, `PCC`)
+  - `MetricSpec`
+  - `compute_max_abs_error`
+  - `compute_mean_abs_error`
+  - `compute_pcc`
+  - `comp_allclose`
+  - `DEFAULT_METRICS`
+- Auto‑compose:
+  - `to_torch_auto_compose`
+
+For concrete, runnable examples of each API, see the tests listed at the top
+of this document.

code/models/common/auto_compose.py

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Automatic composition of multi-device sharded tensors using TensorTopology.
+
+This module provides utilities to infer the correct MeshToTensor composer from a
+sharded ttnn.Tensor's topology metadata and use it to compose shards on host.
+"""
+
+from typing import Optional
+
+import torch
+from loguru import logger
+
+import ttnn
+
+# ======================================================================================
+# Public API
+# ======================================================================================
+
+
+def to_torch_auto_compose(tensor: ttnn.Tensor, device: Optional[ttnn.MeshDevice] = None) -> torch.Tensor:
+    """
+    Convert a (possibly multi-device) TTNN tensor to torch, automatically
+    composing shards based on the tensor's topology.
+
+    Args:
+        tensor: The distributed tensor to convert
+        device: Optional MeshDevice to use when the tensor lives on host
+
+    Returns:
+        PyTorch tensor with shards composed
+    """
+    composer = _infer_mesh_composer_from_topology(tensor, device=device)
+    try:
+        return ttnn.to_torch(tensor, mesh_composer=composer)
+    except Exception as e:
+        logger.error(f"Failed to convert tensor to torch with mesh_composer: {e}")
+        raise
+
+
+def extract_tensor_topology_info(
+    tensor: ttnn.Tensor,
+) -> tuple[list[object], list[int]]:
+    """
+    Extract placements and distribution shape from a tensor's topology.
+
+    Returns:
+        (placements, dist_shape)
+    """
+    topology = tensor.tensor_topology()
+    placements = topology.placements()
+    dist_shape = list(topology.distribution_shape())
+    return placements, dist_shape
+
+
+def get_device_from_tensor(tensor: ttnn.Tensor) -> Optional[ttnn.MeshDevice]:
+    """Get device from tensor or fallback to provided mesh_device."""
+    device = tensor.device()
+    # tensor.device() returns None if the tensor is on the host (ttnn/core/tensor/tensor.cpp --> Tensor::device())
+    if device is None:
+        logger.debug("tensor.device() returns None, tensor is on the host")
+    else:
+        logger.debug(f"tensor.device() returns {device}")
+
+    return device
+
+
+# ======================================================================================
+# Private Implementation
+# ======================================================================================
+
+
+def _infer_mesh_composer_from_topology(
+    tensor: ttnn.Tensor, *, device: Optional[ttnn.MeshDevice] = None
+) -> Optional[ttnn.CppMeshToTensor]:
+    """
+    Return a MeshToTensor composer inferred from the tensor's TensorTopology,
+    or None if no composition is needed (fully replicated, single-device).
+
+    Note: For ND meshes with replicated dimensions, the composer will concatenate
+    all replicas, resulting in duplicated data. Callers may want to slice the
+    result if only one copy is desired.
+
+    Args:
+        tensor: The distributed tensor to infer composer for
+
+    Returns:
+        MeshToTensor composer or None if no composition needed
+    """
+    placements, dist_shape = extract_tensor_topology_info(tensor)
+
+    # No distribution or trivial 1-device case
+    if len(dist_shape) == 0 or (len(dist_shape) == 1 and dist_shape[0] == 1):
+        return None
+
+    tensor_device = get_device_from_tensor(tensor)
+    mesh_device = tensor_device or device
+    if mesh_device is None:
+        # As a last resort, try default device for backward-compatibility
+        mesh_device = ttnn.GetDefaultDevice()
+        if mesh_device is None:
+            raise RuntimeError(
+                "Tensor is on host and no mesh_device provided. "
+                "Pass device=... to to_torch_auto_compose or set a default via ttnn.SetDefaultDevice(...)."
+            )
+
+    # Must match length (should be guaranteed by C++ TT_FATAL in ttnn/core/distributed/distributed_tensor.cpp)
+    assert len(dist_shape) == len(placements)
+
+    if len(dist_shape) == 1 and mesh_device.shape.dims() == 1:
+        return _compose_1d_sharded(mesh_device, placements, dist_shape)
+    else:
+        # N >= 2 dimensions
+        return _compose_nd_sharded(mesh_device, placements, dist_shape)
+
+
+def _compose_1d_sharded(
+    device: ttnn.MeshDevice,
+    placements: list[object],
+    dist_shape: list[int],
+) -> Optional[ttnn.CppMeshToTensor]:
+    """Handle 1D case - returns None if fully replicated."""
+    p = placements[0]
+    if isinstance(p, ttnn.PlacementShard):
+        # Use ND composer with shape override to match the tensor's distribution
+        composer_cfg = ttnn.MeshComposerConfig(dims=[p.dim], mesh_shape_override=ttnn.MeshShape(dist_shape))
+        return ttnn.create_mesh_composer(device, composer_cfg)
+    # Fully replicated - no composition needed
+    return None
+
+
+def _compose_nd_sharded(
+    device: ttnn.MeshDevice,
+    placements: list[object],
+    dist_shape: list[int],
+) -> ttnn.CppMeshToTensor:
+    """
+    Handle ND (N>=2) case.
+
+    For replicated mesh dims, we use dim 0 as convention (the composed result
+    will include all replicas concatenated, which is typically not desired but
+    is how the C++ API works).
+    """
+    dims = []
+    shape_override = []
+    for i, p in enumerate(placements):
+        if isinstance(p, ttnn.PlacementShard):
+            dims.append(p.dim)
+            shape_override.append(dist_shape[i])
+        else:
+            assert isinstance(p, ttnn.PlacementReplicate)
+            # [INFO] steal from TensorDistribution2x4Test test case in test_distributed_tensor.cpp
+            # Replicated: use dim 0 as convention
+            dims.append(0)
+            # Replicated: use shape 1 to skip concatenation
+            shape_override.append(1)
+
+    composer_cfg = ttnn.MeshComposerConfig(dims=dims, mesh_shape_override=ttnn.MeshShape(shape_override))
+    return ttnn.create_mesh_composer(device, composer_cfg)

code/models/common/device_utils.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Device topology naming helpers shared by TTTv2 modules."""
+
+from __future__ import annotations
+
+import ttnn
+
+
+def is_blackhole() -> bool:
+    return "blackhole" in ttnn.get_arch_name()
+
+
+def get_device_name(mesh_device: ttnn.MeshDevice, num_devices: int | None = None) -> str:
+    """Return the product/topology name for a TT mesh device.
+
+    By default, the full mesh device count is used. CCL callers can pass a
+    host-local device count when they need link-count tuning for the current
+    process rather than for the full mesh.
+    """
+    num_devices = mesh_device.get_num_devices() if num_devices is None else num_devices
+    dram_grid_size = mesh_device.dram_grid_size()
+
+    if ttnn.device.is_blackhole(mesh_device):
+        device_names = {
+            1: "P100" if dram_grid_size and dram_grid_size.x == 7 else "P150",
+            2: "P300",
+            4: "P150x4",
+            8: "P150x8",
+            32: "BHGLX",
+        }
+    elif ttnn.device.is_wormhole_b0(mesh_device):
+        device_names = {
+            1: "N150",
+            2: "N300",
+            4: "N150x4",
+            8: "T3K",
+            32: "TG",
+        }
+    else:
+        raise ValueError(f"Unsupported architecture: {ttnn.get_arch_name()}")
+
+    if num_devices in device_names:
+        return device_names[num_devices]
+    raise ValueError(f"Unsupported number of devices: {num_devices} for {ttnn.get_arch_name()}")

code/models/common/distribute_as.py

@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Distribute a torch.Tensor over a mesh using the same topology as a reference TTNN tensor.
+
+This mirrors the composition logic in `auto_compose.py` but in reverse: we infer a
+TensorToMesh mapper from the reference tensor's TensorTopology and use it to distribute
+the torch tensor accordingly.
+"""
+
+from typing import Optional
+
+import torch
+
+import ttnn
+
+from .auto_compose import extract_tensor_topology_info, get_device_from_tensor
+
+# ======================================================================================
+# Public API
+# ======================================================================================
+
+
+def from_torch_dist_as(
+    from_tensor_pt: torch.Tensor, as_tensor_tt: ttnn.Tensor, device: Optional[ttnn.MeshDevice] = None
+) -> ttnn.Tensor:
+    """
+    Distribute a torch.Tensor over a mesh using the same topology as an existing TTNN tensor.
+
+    Args:
+        from_tensor_pt: Source PyTorch tensor on host.
+        as_tensor_tt: Reference TTNN tensor whose topology (placements + distribution shape) will be mirrored.
+        device: Optional mesh device. If omitted, inferred from ``as_tensor_tt`` when possible.
+
+    Returns:
+        A TTNN tensor distributed according to ``as_tensor_tt``'s topology and memory configuration
+        (e.g. height-sharded decode heads).
+    """
+    mapper, device = _infer_mesh_mapper_from_topology(as_tensor_tt, device=device)
+
+    # Usage Patterns: unlike ttnn.to_torch, `device` is required here!
+    # Pattern 1: Using mesh_mapper without device (tensor stays in host memory) Programming_Mesh_of_Devices_with_TT-NN.md:370-375
+    # Then transfer to device separately: Programming_Mesh_of_Devices_with_TT-NN.md:404-405
+    # Pattern 2: Using both mesh_mapper and device together (direct to device) llms.md:1204-1218
+    mem_cfg = as_tensor_tt.memory_config()
+    return ttnn.from_torch(
+        from_tensor_pt,
+        dtype=getattr(as_tensor_tt, "dtype", None),
+        layout=getattr(as_tensor_tt, "layout", None),
+        device=device,
+        memory_config=mem_cfg,
+        mesh_mapper=mapper,
+    )
+
+
+# ======================================================================================
+# Private Implementation
+# ======================================================================================
+
+
+def _infer_mesh_mapper_from_topology(
+    tensor: ttnn.Tensor, *, device: Optional[ttnn.MeshDevice] = None
+) -> Optional[ttnn.CppTensorToMesh]:
+    """
+    Return a TensorToMesh mapper inferred from the tensor's TensorTopology,
+    or (None, mesh_device) if no distribution is needed (fully replicated, single-device).
+    """
+    placements, dist_shape = extract_tensor_topology_info(tensor)
+
+    tensor_device = get_device_from_tensor(tensor)
+    mesh_device = tensor_device or device
+    if mesh_device is None:
+        mesh_device = ttnn.GetDefaultDevice()
+        if mesh_device is None:
+            raise RuntimeError(
+                "Tensor is on host and no mesh_device provided. " "Set a default via ttnn.SetDefaultDevice(...)."
+            )
+
+    assert len(dist_shape) == len(placements)
+
+    if len(dist_shape) == 1 and mesh_device.shape.dims() == 1:
+        return _map_1d(mesh_device, placements, dist_shape), mesh_device
+    else:
+        return _map_nd(mesh_device, placements, dist_shape), mesh_device
+
+
+def _map_1d(
+    device: ttnn.MeshDevice,
+    placements: list[object],
+    dist_shape: list[int],
+) -> Optional[ttnn.CppTensorToMesh]:
+    """
+    Build a 1D TensorToMesh mapper. Returns None if fully trivial (handled earlier).
+    """
+    p = placements[0]
+    if isinstance(p, ttnn.PlacementShard):
+        mapper_cfg = ttnn.MeshMapperConfig(
+            placements=[ttnn.PlacementShard(p.dim)],
+            mesh_shape_override=ttnn.MeshShape(dist_shape),
+        )
+        return ttnn.create_mesh_mapper(device, mapper_cfg)
+    else:
+        # Replicate across the 1D mesh extent
+        mapper_cfg = ttnn.MeshMapperConfig(
+            placements=[ttnn.PlacementReplicate()],
+            mesh_shape_override=ttnn.MeshShape(dist_shape),
+        )
+        return ttnn.create_mesh_mapper(device, mapper_cfg)
+
+
+def _map_nd(
+    device: ttnn.MeshDevice,
+    placements: list[object],
+    dist_shape: list[int],
+) -> ttnn.CppTensorToMesh:
+    """
+    Build an ND TensorToMesh mapper that mirrors the tensor's placements and distribution shape.
+    """
+    mapper_placements = []
+    for p in placements:
+        if isinstance(p, ttnn.PlacementShard):
+            mapper_placements.append(ttnn.PlacementShard(p.dim))
+        else:
+            assert isinstance(p, ttnn.PlacementReplicate)
+            mapper_placements.append(ttnn.PlacementReplicate())
+
+    mapper_cfg = ttnn.MeshMapperConfig(placements=mapper_placements, mesh_shape_override=ttnn.MeshShape(dist_shape))
+    return ttnn.create_mesh_mapper(device, mapper_cfg)

code/models/common/generation_utils.py

@@ -0,0 +1,279 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from loguru import logger
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.logits_process import (  # ForceTokensLogitsProcessor,
+    EncoderNoRepeatNGramLogitsProcessor,
+    EncoderRepetitionPenaltyLogitsProcessor,
+    ExponentialDecayLengthPenalty,
+    ForcedBOSTokenLogitsProcessor,
+    ForcedEOSTokenLogitsProcessor,
+    InfNanRemoveLogitsProcessor,
+    LogitNormalization,
+    LogitsProcessorList,
+    MinLengthLogitsProcessor,
+    MinNewTokensLengthLogitsProcessor,
+    NoBadWordsLogitsProcessor,
+    NoRepeatNGramLogitsProcessor,
+    PrefixConstrainedLogitsProcessor,
+    RepetitionPenaltyLogitsProcessor,
+    SuppressTokensAtBeginLogitsProcessor,
+    SuppressTokensLogitsProcessor,
+)
+
+# HammingDiversityLogitsProcessor (diverse beam search) was removed in
+# transformers 5.x with no replacement. Import it optionally so this module
+# still loads; it's only used when diversity_penalty > 0, which TT generation
+# paths don't exercise.
+try:
+    from transformers.generation.logits_process import HammingDiversityLogitsProcessor
+except ImportError:  # transformers >= 5.x
+    HammingDiversityLogitsProcessor = None
+
+
+def _merge_criteria_processor_list(
+    default_list,  # Union[LogitsProcessorList, StoppingCriteriaList],
+    custom_list,  # Union[LogitsProcessorList, StoppingCriteriaList],
+):  # -> Union[LogitsProcessorList, StoppingCriteriaList]:
+    if len(custom_list) == 0:
+        return default_list
+
+    for default in default_list:
+        for custom in custom_list:
+            if type(custom) is type(default):
+                object_type = "stopping criteria" if isinstance(custom, StoppingCriteria) else "logits processor"
+                raise ValueError(
+                    f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to"
+                    f" `generate`, but it has already been created with the values {default}. {default} has been"
+                    " created by passing the corresponding arguments to generate or by the model's config default"
+                    f" values. If you just want to change the default values of {object_type} consider passing"
+                    f" them as arguments to `generate` instead of using a custom {object_type}."
+                )
+    default_list.extend(custom_list)
+    return default_list
+
+
+def _get_logits_processor(
+    generation_config: GenerationConfig,
+    input_ids_seq_length: int,
+    encoder_input_ids,  # torch.LongTensor
+    prefix_allowed_tokens_fn,  # Callable[[int, torch.Tensor], List[int]],
+    logits_processor,  # Optional[LogitsProcessorList]
+):  # -> LogitsProcessorList:
+    """
+    This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsProcessor`]
+    instances used to modify the scores of the language model head.
+    """
+    # instantiate processors list
+    processors = LogitsProcessorList()
+
+    # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files
+    # all samplers can be found in `generation_utils_samplers.py`
+    if generation_config.diversity_penalty is not None and generation_config.diversity_penalty > 0.0:
+        if HammingDiversityLogitsProcessor is None:
+            raise NotImplementedError(
+                "diversity_penalty > 0 (diverse beam search) requires HammingDiversityLogitsProcessor, "
+                "which was removed in transformers 5.x."
+            )
+        processors.append(
+            HammingDiversityLogitsProcessor(
+                diversity_penalty=generation_config.diversity_penalty,
+                num_beams=generation_config.num_beams,
+                num_beam_groups=generation_config.num_beam_groups,
+            )
+        )
+    if generation_config.encoder_repetition_penalty is not None and generation_config.encoder_repetition_penalty != 1.0:
+        processors.append(
+            EncoderRepetitionPenaltyLogitsProcessor(
+                penalty=generation_config.encoder_repetition_penalty,
+                encoder_input_ids=encoder_input_ids,
+            )
+        )
+    if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0:
+        processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty))
+    if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0:
+        processors.append(NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size))
+    if (
+        generation_config.encoder_no_repeat_ngram_size is not None
+        and generation_config.encoder_no_repeat_ngram_size > 0
+    ):
+        if len(encoder_input_ids.shape) == 2:
+            processors.append(
+                EncoderNoRepeatNGramLogitsProcessor(generation_config.encoder_no_repeat_ngram_size, encoder_input_ids)
+            )
+        else:
+            raise ValueError("It's impossible to use `encoder_no_repeat_ngram_size` with decoder-only architecture")
+    if generation_config.bad_words_ids is not None:
+        processors.append(NoBadWordsLogitsProcessor(generation_config.bad_words_ids, generation_config.eos_token_id))
+    if (
+        generation_config.min_length is not None
+        and generation_config.eos_token_id is not None
+        and generation_config.min_length > 0
+    ):
+        processors.append(MinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id))
+    if (
+        generation_config.min_new_tokens is not None
+        and generation_config.eos_token_id is not None
+        and generation_config.min_new_tokens > 0
+    ):
+        processors.append(
+            MinNewTokensLengthLogitsProcessor(
+                input_ids_seq_length,
+                generation_config.min_new_tokens,
+                generation_config.eos_token_id,
+            )
+        )
+    if prefix_allowed_tokens_fn is not None:
+        processors.append(
+            PrefixConstrainedLogitsProcessor(
+                prefix_allowed_tokens_fn,
+                generation_config.num_beams // generation_config.num_beam_groups,
+            )
+        )
+    if generation_config.forced_bos_token_id is not None:
+        processors.append(ForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id))
+    if generation_config.forced_eos_token_id is not None:
+        processors.append(
+            ForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id)
+        )
+    if generation_config.remove_invalid_values is True:
+        processors.append(InfNanRemoveLogitsProcessor())
+    if generation_config.exponential_decay_length_penalty is not None:
+        processors.append(
+            ExponentialDecayLengthPenalty(
+                generation_config.exponential_decay_length_penalty,
+                generation_config.eos_token_id,
+                input_ids_seq_length,
+            )
+        )
+    if generation_config.suppress_tokens is not None:
+        processors.append(SuppressTokensLogitsProcessor(generation_config.suppress_tokens))
+    if generation_config.begin_suppress_tokens is not None:
+        begin_index = input_ids_seq_length
+        begin_index = (
+            begin_index
+            if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None)
+            else begin_index + 1
+        )
+        processors.append(SuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index))
+    processors = _merge_criteria_processor_list(processors, logits_processor)
+    # `LogitNormalization` should always be the last logit processor, when present
+    if generation_config.renormalize_logits is True:
+        processors.append(LogitNormalization())
+    return processors
+
+
+def get_logits_processor(input_ids, config):
+    generation_config = GenerationConfig.from_model_config(config)
+    input_ids_seq_length = input_ids.shape[-1]
+
+    logits_processor = _get_logits_processor(
+        generation_config=generation_config,
+        input_ids_seq_length=input_ids_seq_length,
+        encoder_input_ids=input_ids,
+        prefix_allowed_tokens_fn=None,
+        logits_processor=LogitsProcessorList(),
+    )
+
+    return logits_processor
+
+
+def pad_input_32(tensor, value):
+    len = tensor.shape[1]
+
+    if len % 32 == 0:
+        return tensor
+
+    padded_len = ((len // 32) + 1) * 32
+
+    pad_tensor = (value * torch.ones(tensor.shape[0], padded_len - len)).to(torch.long)
+    tensor = torch.cat([tensor, pad_tensor], dim=1)
+
+    return tensor
+
+
+def run_generate(
+    input_sentance,
+    tokenizer,
+    tt_model_constructor,
+    device,
+    run_tt_model=True,
+    log=True,
+    comp_pcc=None,
+):
+    tt_model, hf_reference_model = tt_model_constructor(device)
+
+    # Prepare input
+    tokenized = tokenizer(input_sentance, return_tensors="pt")  # Batch size 1
+
+    input_ids = pad_input_32(tokenized.input_ids, hf_reference_model.generation_config.pad_token_id)
+    attention_mask = pad_input_32(tokenized.attention_mask, 0)
+
+    if log:
+        logger.debug(f"input_ids {input_ids.shape} {input_ids}")
+        logger.debug(f"attention_mask {attention_mask.shape} {attention_mask}")
+
+    logits_processor = get_logits_processor(input_ids, hf_reference_model.config)
+
+    decoder_start_values = hf_reference_model.generation_config.pad_token_id * torch.ones(1, 32).to(torch.long)
+    decoder_input_ids = hf_reference_model.generation_config.pad_token_id * torch.ones(1, 64).to(torch.long)
+
+    if log:
+        logger.debug(f"decoder_input_ids {decoder_input_ids}")
+
+    encoder_outputs = None
+    use_cache = False
+
+    for i in range(64):
+        # PyTorch forward pass
+        pt_out = hf_reference_model(
+            input_ids=input_ids,
+            decoder_input_ids=decoder_input_ids,
+            attention_mask=attention_mask,
+        )
+
+        if run_tt_model:
+            tt_out = tt_model(
+                input_ids=input_ids,
+                decoder_input_ids=decoder_input_ids,
+                attention_mask=attention_mask,
+                encoder_outputs=encoder_outputs,
+                return_dict=True,
+                use_cache=use_cache,
+            )
+            encoder_outputs = tt_out.encoder_outputs
+            next_token_logits = tt_out.logits
+
+            if comp_pcc is not None:
+                does_pass, pcc_message = comp_pcc(pt_out.logits, tt_out.logits, 0.98)
+
+                if log:
+                    logger.info(pcc_message)
+        else:
+            next_token_logits = pt_out.logits
+
+        # pre-process distribution
+        next_tokens_scores = logits_processor(input_ids, next_token_logits)
+
+        # argmax
+        next_tokens = torch.argmax(next_tokens_scores, dim=-1)
+
+        if log:
+            logger.debug(f"next_tokens {next_tokens}")
+
+        if next_tokens[0][i] == hf_reference_model.generation_config.eos_token_id:
+            break
+
+        # We need to expand decoder_input_ids
+        if (i + 1) % 32 == 0:
+            decoder_input_ids = torch.cat([decoder_input_ids, decoder_start_values], dim=1)
+
+        decoder_input_ids[0][i + 1] = next_tokens[0][i]
+
+        if log:
+            logger.debug(f"decoder_input_ids {decoder_input_ids[0]}")
+
+    return tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)

code/models/common/helper_funcs.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import ttnn
+
+
+def Linear(
+    in_features: int,
+    out_features: int,
+    weight: ttnn.Tensor,
+    bias: Optional[ttnn.Tensor] = None,
+    output_mem_config=ttnn.DRAM_MEMORY_CONFIG,
+):
+    """
+    Returns a function that performs a Linear operation with optional bias.
+
+    ``weight`` must be tt_tensor.
+    """
+    assert weight.padded_shape == [
+        1,
+        1,
+        out_features,
+        in_features,
+    ], "weight does not have the expected shape"
+
+    if bias is not None:
+        assert bias.padded_shape[-1] == out_features, "bias does not have the expected shape"
+
+    weight = weight
+    bias = bias
+    weight_T = ttnn.transpose(weight, -2, -1)
+
+    def linear_(activation):
+        nonlocal bias
+        assert activation.padded_shape[-1] == in_features, "activation tensor do not have the expected shape"
+        if bias is not None and bias.get_layout() != ttnn.TILE_LAYOUT:
+            bias = ttnn.to_layout(bias, ttnn.TILE_LAYOUT)
+        return ttnn.linear(activation, weight_T, bias=bias, memory_config=output_mem_config)
+
+    return linear_

code/models/common/lightweightmodule.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+class LightweightModule:
+    """Torch modules add a surprising amount of host overhead for attribute
+    access and method calls. This class is a lightweight alternative that
+    just wraps a forward function for now."""
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)

code/models/common/llama_models.py

@@ -0,0 +1,212 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Union
+
+import torch
+from PIL import Image
+from pydantic import BaseModel, validator
+
+# ``AutoModelForImageTextToText`` (the replacement for ``AutoModelForVision2Seq``,
+# which was removed in transformers 5.x) is only consumed by the
+# ``GeneratorChat``/``GeneratorText`` constructors below — defer the imports
+# so loading this module doesn't break every downstream import chain
+# (e.g. ``tt_transformers.tt.generator``, used by every TT vLLM bridge)
+# under transformers >= 5.
+
+
+class Role(Enum):
+    system = "system"
+    user = "user"
+    assistant = "assistant"
+    ipython = "ipython"
+
+
+class StopReason(Enum):
+    end_of_turn = "end_of_turn"
+    end_of_message = "end_of_message"
+    out_of_tokens = "out_of_tokens"
+
+
+@dataclass
+class TokenResult:
+    token: int
+    text: str
+    logprobs: Optional[List[float]] = None
+
+
+@dataclass
+class CompletionMessage:
+    content: str
+    role: Role = Role.assistant.value
+
+
+class BuiltinTool(Enum):
+    brave_search = "brave_search"
+    wolfram_alpha = "wolfram_alpha"
+    photogen = "photogen"
+    code_interpreter = "code_interpreter"
+
+
+Primitive = Union[str, int, float, bool, None]
+RecursiveType = Union[Primitive, List[Primitive], Dict[str, Primitive]]
+
+
+class ToolCall(BaseModel):
+    call_id: str
+    tool_name: Union[BuiltinTool, str]
+    arguments: Dict[str, RecursiveType]
+
+    @validator("tool_name", pre=True)
+    @classmethod
+    def validate_field(cls, v):
+        if isinstance(v, str):
+            try:
+                return BuiltinTool(v)
+            except ValueError:
+                return v
+        return v
+
+
+class ChatPrediction:
+    generation: CompletionMessage
+    decoded_tokens: Optional[List[str]] = None
+    logprobs: Optional[List[List[float]]] = None
+
+
+class CompletionPrediction:
+    generation: str
+    decoded_tokens: Optional[List[str]] = None
+    logprobs: Optional[List[List[float]]] = None
+
+
+def sample_top_p(probs, p):
+    """
+    Perform top-p (nucleus) sampling on a probability distribution.
+
+    Args:
+        probs (torch.Tensor): Probability distribution tensor.
+        p (float): Probability threshold for top-p sampling.
+
+    Returns:
+        torch.Tensor: Sampled token indices.
+
+    Note:
+        Top-p sampling selects the smallest set of tokens whose cumulative probability mass
+        exceeds the threshold p. The distribution is renormalized based on the selected tokens.
+    From: https://github.com/meta-llama/llama-models/blob/v0.1.5/models/llama3/reference_impl/generation.py#L450-L472
+    """
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort[mask] = 0.0
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    next_token = torch.multinomial(probs_sort, num_samples=1)
+    next_token = torch.gather(probs_idx, -1, next_token)
+    return next_token
+
+
+def extract_images_from_messages(messages):
+    images = []
+    for message in messages:
+        if "content" in message:
+            contents = message["content"]
+            for content in contents:
+                if (content["type"] == "image") and ("image" in content):
+                    images.append(content["image"])
+    return images
+
+
+def create_vision_mask(
+    tokens: List[int],
+    vision_token: int,
+) -> List[List[int]]:
+    """From: https://github.com/meta-llama/llama-models/blob/v0.1.5/models/llama3/api/chat_format.py#L253-L276"""
+    vision_token_locations = [i for i, token in enumerate(tokens) if token == vision_token]
+    if len(vision_token_locations) == 0:
+        return []
+
+    if len(vision_token_locations) == 1:
+        # only one image present, unmask until end of sequence
+        return [[vision_token_locations[0], -1]]
+    vision_masks = [[loc1, loc2] for loc1, loc2 in zip(vision_token_locations[:-1], vision_token_locations[1:])]
+    # last image will attend to all subsequent text
+    vision_masks.append([vision_token_locations[-1], len(tokens)])
+
+    # if there are two or more consecutive vision tokens,
+    # they should all attend to all subsequent
+    # text present
+    last_mask_end = vision_masks[-1][1]
+    for vision_mask in vision_masks[::-1]:
+        if vision_mask[0] == vision_mask[1] - 1:
+            vision_mask[1] = last_mask_end
+        last_mask_end = vision_mask[1]
+    return vision_masks
+
+
+def encode_content(content, images, image_token):
+    if isinstance(content, Image):
+        images.append(content)
+        assert image_token is not None
+        return image_token
+    if isinstance(content, str):
+        return content
+    if isinstance(content, (list, tuple)):
+        return "\n".join(encode_content(item, images) for item in content)
+    if isinstance(content, dict):
+        content_type = content.get("type")
+        if content_type == "text":
+            return content["text"]
+        if content_type == "image":
+            # TBD: support url
+            images.append(content["image"])
+            assert image_token is not None
+            return image_token
+    raise ValueError(f"Unknown content format: {content}")
+
+
+class GeneratorChat:
+    def __init__(self, model_name, max_batch_size=1):
+        from transformers import pipeline
+
+        self.pipe = pipeline("image-text-to-text", model=model_name, batch_size=max_batch_size)
+
+    def chat_completion(
+        self,
+        messages,
+        temperature=0.6,
+        top_p: float = 0.9,
+        max_gen_len=None,
+    ):
+        generation_output = self.pipe(
+            text=messages, temperature=temperature, top_p=top_p, max_new_tokens=max_gen_len, return_full_text=False
+        )
+
+        if len(generation_output) == 1:
+            return CompletionMessage(content=generation_output[0]["generated_text"])
+        return [CompletionMessage(content=output[0]["generated_text"]) for output in generation_output]
+
+
+class GeneratorText:
+    def __init__(self, model_name):
+        from transformers import AutoModelForImageTextToText, AutoProcessor
+
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model = AutoModelForImageTextToText.from_pretrained(model_name)
+
+    def text_completion(
+        self,
+        content: Union[str, Image.Image, Dict, List[Dict]],
+        temperature: float = 0.6,
+        top_p: float = 0.9,
+        max_gen_len=None,
+    ):
+        images = []
+        text = encode_content(content, images, self.processor.image_token)
+        model_input = self.processor(text=text, images=images or None, return_tensors="pt", add_special_tokens=False)
+        tokens = self.model.generate(**model_input, temperature=temperature, top_p=top_p, max_new_tokens=max_gen_len)[0]
+        tokens = tokens[model_input["input_ids"].shape[-1] :]
+        return self.processor.decode(tokens, skip_special_tokens=True)

code/models/common/metrics.py

@@ -0,0 +1,488 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+TTNN Metric Functions
+
+Metric functions for comparing TTNN implementations against reference implementations.
+All metrics support both PyTorch and TTNN tensors, with TTNN-native computation
+that stays on device until the final scalar result.
+
+Key Features:
+- Automatic detection of tensor type (PyTorch vs TTNN)
+- TTNN-native computation using device operations
+- Minimal host transfers (only final scalar)
+- Graceful fallback to PyTorch for non-tensor inputs
+"""
+
+import numpy as np
+import torch
+from loguru import logger
+
+import ttnn
+
+from .auto_compose import to_torch_auto_compose
+
+# ======================================================================================
+# Public API
+# ======================================================================================
+
+
+def compute_max_abs_error(impl, ref):
+    """
+    Compute maximum absolute error between two tensors.
+
+    Supports both PyTorch and TTNN tensors. For TTNN tensors, computation
+    stays on device until the final scalar value.
+
+    Args:
+        impl: Implementation output (PyTorch or TTNN tensor)
+        ref: Reference output (PyTorch or TTNN tensor)
+
+    Returns:
+        float: Maximum absolute difference between tensors
+
+    Examples:
+        >>> a = torch.tensor([1.0, 2.0, 3.0])
+        >>> b = torch.tensor([1.1, 2.0, 2.9])
+        >>> _compute_max_abs_error(a, b)
+        0.10000002384185791
+    """
+    try:
+        if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref):
+            # TTNN path - stay on device
+            diff = ttnn.subtract(impl, ref)
+            abs_diff = _ttnn_op_layout_invariant(diff, ttnn.abs)
+            return _ttnn_max_scalar_all_dtype(abs_diff)
+        elif _is_ttnn_tensor(impl):
+            return (to_torch_auto_compose(impl) - ref).abs().max().item()
+        elif _is_ttnn_tensor(ref):
+            return (impl - to_torch_auto_compose(ref)).abs().max().item()
+        else:
+            # PyTorch path - both must be torch tensors
+            return (impl - ref).abs().max().item()
+    except Exception as e:
+        return float("inf")
+
+
+def compute_mean_abs_error(impl, ref):
+    """
+    Compute mean absolute error between two tensors.
+
+    Supports both PyTorch and TTNN tensors. For TTNN tensors, computation
+    stays on device until the final scalar value.
+
+    Args:
+        impl: Implementation output (PyTorch or TTNN tensor)
+        ref: Reference output (PyTorch or TTNN tensor)
+
+    Returns:
+        float: Mean absolute difference between tensors
+
+    Examples:
+        >>> a = torch.tensor([1.0, 2.0, 3.0])
+        >>> b = torch.tensor([1.1, 2.0, 2.9])
+        >>> _compute_mean_abs_error(a, b)
+        0.06666667014360428
+    """
+    try:
+        if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref):
+            # TTNN path - stay on device
+            diff = ttnn.subtract(impl, ref)
+            abs_diff = _ttnn_op_layout_invariant(diff, ttnn.abs)
+            return _ttnn_mean_scalar_all_dtype(abs_diff)
+        elif _is_ttnn_tensor(impl):
+            return (to_torch_auto_compose(impl) - ref).abs().mean().item()
+        elif _is_ttnn_tensor(ref):
+            return (impl - to_torch_auto_compose(ref)).abs().mean().item()
+        else:
+            # PyTorch path - both must be torch tensors
+            return (impl - ref).abs().mean().item()
+    except Exception as e:
+        return float("inf")
+
+
+def compute_pcc(impl, ref):
+    """
+    Compute Pearson Correlation Coefficient (PCC) between two tensors.
+
+    PCC measures the linear correlation between two tensors. Returns 1.0 for
+    perfect positive correlation, 0.0 for no correlation, and -1.0 for perfect
+    negative correlation. This is a common metric used in tt-metal for validating
+    tensor computations.
+
+    For TTNN tensors: Uses on-device computation (100-1000× faster). If edge cases
+    produce non-finite values (NaN/inf), raises ValueError with instructions to use
+    Pattern 2 (PyTorch reference with output_map for robust CPU-based PCC).
+
+    For PyTorch tensors: Uses robust numpy implementation (from tt-metal
+    comparison_funcs.py) that handles complex, NaN, inf, and constant tensors.
+
+    Args:
+        impl: Implementation output (PyTorch or TTNN tensor)
+        ref: Reference output (PyTorch or TTNN tensor)
+
+    Returns:
+        float: PCC value in range [-1.0, 1.0], or 0.0 on error
+
+    Raises:
+        ValueError: If TTNN-native computation produces non-finite values
+
+    Examples:
+        >>> a = torch.tensor([1.0, 2.0, 3.0])
+        >>> b = torch.tensor([1.0, 2.0, 3.0])
+        >>> _compute_pcc(a, b)
+        1.0
+
+        >>> a = torch.tensor([1.0, 2.0, 3.0])
+        >>> b = torch.tensor([3.0, 2.0, 1.0])
+        >>> _compute_pcc(a, b)
+        -1.0
+    """
+    try:
+        # TTNN fast path - compute on device
+        if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref):
+            return compute_pcc_device(impl, ref)
+        elif _is_ttnn_tensor(impl):
+            return compute_pcc_host(to_torch_auto_compose(impl), ref)
+        elif _is_ttnn_tensor(ref):
+            return compute_pcc_host(impl, to_torch_auto_compose(ref))
+        else:
+            return compute_pcc_host(impl, ref)
+    except Exception:
+        return 0.0
+
+
+# Default metrics dictionary for easy import
+DEFAULT_METRICS = {
+    "max_abs_error": compute_max_abs_error,
+    "mean_abs_error": compute_mean_abs_error,
+    "pcc": compute_pcc,
+}
+
+
+# Allclose comparison with detailed delta string
+def comp_allclose(impl, ref, rtol=1e-05, atol=1e-08):
+    """
+    Compare two tensors using an allclose criterion and return (passing, details).
+
+    Provides both a TTNN-native on-device implementation and a PyTorch fallback.
+    Mirrors semantics of torch.allclose(..., equal_nan=True) and reports the
+    maximum absolute and relative deltas similar to comparison_funcs.py.
+
+    Args:
+        impl: Implementation output (PyTorch or TTNN tensor)
+        ref: Reference output (PyTorch or TTNN tensor)
+        rtol (float): Relative tolerance
+        atol (float): Absolute tolerance
+
+    Returns:
+        tuple[bool, str]: (passing, "Max ATOL Delta: x, Max RTOL Delta: y[,...]")
+    """
+    try:
+        # TTNN-native path: compute deltas and allclose on device, then transfer final scalars
+        if _is_ttnn_tensor(impl) and _is_ttnn_tensor(ref):
+            # Compute deltas (device)
+            diff = _ttnn_op_layout_invariant(ttnn.subtract(impl, ref), ttnn.abs)
+            cal_atol = _ttnn_max_scalar_all_dtype(diff)
+            # For rtol delta, divide by abs(ref) (may produce inf for zeros; acceptable for reporting)
+            divided_by_ref = ttnn.divide(diff, _ttnn_op_layout_invariant(ref, ttnn.abs))
+            cal_rtol = _ttnn_max_scalar_all_dtype(divided_by_ref)
+
+            # equal_nan=True semantics and finite/infinite handling
+            isnan_impl = _ttnn_op_layout_invariant(impl, ttnn.isnan)
+            isnan_ref = _ttnn_op_layout_invariant(ref, ttnn.isnan)
+            both_nan = ttnn.logical_and(isnan_impl, isnan_ref)
+
+            isinf_impl = _ttnn_op_layout_invariant(impl, ttnn.isinf)
+            isinf_ref = _ttnn_op_layout_invariant(ref, ttnn.isinf)
+            impl_sign = _ttnn_op_layout_invariant(impl, ttnn.sign)
+            ref_sign = _ttnn_op_layout_invariant(ref, ttnn.sign)
+            same_sign_inf = ttnn.eq(impl_sign, ref_sign)
+            both_inf_same_sign = ttnn.logical_and(ttnn.logical_and(isinf_impl, isinf_ref), same_sign_inf)
+
+            # Finite elements where numeric closeness applies
+            any_nan = ttnn.logical_or(isnan_impl, isnan_ref)
+            any_inf = ttnn.logical_or(isinf_impl, isinf_ref)
+            finite_both = _ttnn_op_layout_invariant(ttnn.logical_or(any_nan, any_inf), ttnn.logical_not)
+
+            # |impl - ref| <= atol + rtol * |ref|
+            bound = ttnn.add(ttnn.mul(_ttnn_op_layout_invariant(ref, ttnn.abs), rtol, dtype=ttnn.bfloat16), atol)
+            close_numeric = ttnn.le(diff, bound)
+            finite_and_close = ttnn.logical_and(finite_both, close_numeric)
+
+            ok_mask = ttnn.logical_or(ttnn.logical_or(both_nan, both_inf_same_sign), finite_and_close)
+            fail_mask = _ttnn_op_layout_invariant(ok_mask, ttnn.logical_not)
+
+            # Reduce to scalar: any failure -> 1.0 else 0.0
+            fail_indicator = _ttnn_op_layout_invariant(fail_mask, ttnn.where, true_value=1.0, false_value=0.0)
+            any_fail = _ttnn_max_scalar_all_dtype(fail_indicator)
+            passing = any_fail == 0.0
+
+            output_str = f"Max ATOL Delta: {cal_atol}, Max RTOL Delta: {cal_rtol}"
+            if not passing:
+                output_str += ", Allclose check failed"
+            return passing, output_str
+
+        # Fallback: compute with PyTorch (handles mixed inputs by converting TTNN -> torch)
+        impl_torch = to_torch_auto_compose(impl) if _is_ttnn_tensor(impl) else impl
+        ref_torch = to_torch_auto_compose(ref) if _is_ttnn_tensor(ref) else ref
+
+        if torch.is_tensor(impl_torch) and torch.is_tensor(ref_torch):
+            # Match dtype for fair comparison
+            if impl_torch.dtype != ref_torch.dtype:
+                ref_torch = ref_torch.to(impl_torch.dtype)
+
+            atol_delta = torch.max(torch.abs(impl_torch - ref_torch)).item()
+            # May produce inf where ref == 0; this mirrors comparison_funcs.py behavior
+            rtol_delta = torch.max(torch.abs(impl_torch - ref_torch) / torch.abs(ref_torch)).item()
+            passing = torch.allclose(impl_torch, ref_torch, rtol, atol, True)
+            output_str = f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}"
+            if not passing:
+                output_str += ", Allclose check failed"
+            return passing, output_str
+
+        # Unsupported types
+        return False, "Unsupported input types for comp_allclose"
+    except Exception as e:
+        return False, f"Error computing comp_allclose: {e}"
+
+
+def compute_pcc_device(impl, ref):
+    """Fast on-device PCC for TTNN tensors only."""
+    try:
+        # Early edge-case handling to mirror CPU semantics
+        # - All NaNs → 1.0; mixed NaNs → 0.0
+        # - One tensor all zero and the other not → 0.0
+        # - Both constant → 1.0 if equal, else 0.0
+
+        # Any nonzero check (all-zero detection)
+        impl_abs_max = _ttnn_max_scalar_all_dtype(_ttnn_op_layout_invariant(impl, ttnn.abs))
+        ref_abs_max = _ttnn_max_scalar_all_dtype(_ttnn_op_layout_invariant(ref, ttnn.abs))
+
+        impl_has_any = impl_abs_max != 0.0
+        ref_has_any = ref_abs_max != 0.0
+        if impl_has_any != ref_has_any:
+            return 0.0
+
+        # Min/Max scalars for constant and NaN detection
+        impl_min = _ttnn_min_scalar_all_dtype(impl)
+        impl_max = _ttnn_max_scalar_all_dtype(impl)
+        ref_min = _ttnn_min_scalar_all_dtype(ref)
+        ref_max = _ttnn_max_scalar_all_dtype(ref)
+
+        impl_min_finite = np.isfinite(impl_min)
+        impl_max_finite = np.isfinite(impl_max)
+        ref_min_finite = np.isfinite(ref_min)
+        ref_max_finite = np.isfinite(ref_max)
+
+        impl_all_nan = (not impl_min_finite) and (not impl_max_finite)
+        ref_all_nan = (not ref_min_finite) and (not ref_max_finite)
+        if impl_all_nan and ref_all_nan:
+            return 1.0
+        if impl_all_nan != ref_all_nan:
+            return 0.0
+
+        # Constant tensors
+        if impl_min_finite and impl_max_finite and ref_min_finite and ref_max_finite:
+            if impl_min == impl_max and ref_min == ref_max:
+                return (
+                    1.0
+                    if torch.isclose(
+                        torch.tensor(impl_max, dtype=torch.float32), torch.tensor(ref_max, dtype=torch.float32)
+                    )
+                    else 0.0
+                )
+
+        # Standard PCC formula on device
+        mean_impl = _ttnn_mean_scalar_all_dtype(impl)
+        mean_ref = _ttnn_mean_scalar_all_dtype(ref)
+
+        impl_centered = ttnn.subtract(impl, mean_impl)
+        ref_centered = ttnn.subtract(ref, mean_ref)
+
+        # todo)) ttnn.sum() does local reduction only; need CCL reduction for global sum when adding support for multiple-devices
+        # [INFO] we cast to float32 to avoid overflow when impl and ref are in bfloat8_b or bfloat4_b
+        numerator = ttnn.sum(ttnn.mul(impl_centered, ref_centered, dtype=ttnn.float32))
+        impl_sq_sum = ttnn.sum(ttnn.mul(impl_centered, impl_centered, dtype=ttnn.float32))
+        ref_sq_sum = ttnn.sum(ttnn.mul(ref_centered, ref_centered, dtype=ttnn.float32))
+        denominator = ttnn.sqrt(ttnn.mul(impl_sq_sum, ref_sq_sum, dtype=ttnn.float32))
+
+        # Safe divide
+        denom_scalar = denominator.item()
+        if denom_scalar == 0.0 or not np.isfinite(denom_scalar):
+            return 0.0
+
+        pcc = numerator.item() / denom_scalar
+        if not np.isfinite(pcc):
+            return 0.0
+        return pcc
+    except Exception as e:
+        # todo)) maybe return a string for logging: f"Error computing PCC on device: {e}, impl: {impl}, ref: {ref}"
+        return 0.0
+
+
+# code stolen from tests/tt_eager/python_api_testing/sweep_tests/comparison_funcs.py
+# and models/common/utility_functions.py
+def compute_pcc_host(impl, ref):
+    """Robust CPU PCC for PyTorch tensors only."""
+    try:
+        calculated = impl
+        golden = ref
+        if golden.dtype != calculated.dtype:
+            calculated = calculated.type(golden.dtype)
+
+        # Handle complex tensors
+        if golden.is_complex() and calculated.is_complex():
+            golden = torch.view_as_real(golden.clone())
+            calculated = torch.view_as_real(calculated.clone())
+
+        # Convert to float if needed
+        if not (golden.is_floating_point() or calculated.is_floating_point()):
+            golden = golden.to(torch.float)
+            calculated = calculated.to(torch.float)
+
+        # Both tensors are nan
+        if torch.all(torch.isnan(golden)) and torch.all(torch.isnan(calculated)):
+            return 1.0
+
+        # One tensor is all nan, the other is not
+        if torch.all(torch.isnan(golden)) or torch.all(torch.isnan(calculated)):
+            return 0.0
+
+        # One tensor is all zero, the other is not — also a zero-variance case.
+        if torch.any(golden.bool()) != torch.any(calculated.bool()):
+            logger.warning("One tensor is all zero. PCC undefined; falling back to allclose.")
+            return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04))
+
+        # Mask all infs and nans
+        golden = golden.clone()
+        golden[
+            torch.logical_or(
+                torch.isnan(golden),
+                torch.logical_or(torch.isinf(golden), torch.isneginf(golden)),
+            )
+        ] = 0
+        calculated = calculated.clone()
+        calculated[
+            torch.logical_or(
+                torch.isnan(calculated),
+                torch.logical_or(torch.isinf(calculated), torch.isneginf(calculated)),
+            )
+        ] = 0
+
+        if torch.equal(golden, calculated):
+            return 1.0
+
+        if golden.dtype == torch.bfloat16:
+            golden = golden.type(torch.float32)
+            calculated = calculated.type(torch.float32)
+
+        # Single element or constant tensor: PCC is undefined.
+        if golden.numel() == 1:
+            return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04))
+
+        if torch.max(golden) == torch.min(golden) or torch.max(calculated) == torch.min(calculated):
+            logger.warning("One or both tensors are constant (zero std dev). PCC undefined; falling back to allclose.")
+            return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04))
+
+        # Compute PCC using numpy's corrcoef
+        cal_pcc = np.ma.corrcoef(
+            np.ma.masked_invalid(torch.squeeze(golden).detach().numpy()).flatten(),
+            np.ma.masked_invalid(torch.squeeze(calculated).detach().numpy()).flatten(),
+        )
+        # Read off-diagonal directly to avoid diagonal contamination.
+        cal_pcc = cal_pcc[0, 1]
+
+        if isinstance(cal_pcc, np.ma.core.MaskedConstant) or np.isnan(float(cal_pcc)):
+            logger.warning("PCC returned NaN/masked. Falling back to allclose.")
+            return float(torch.allclose(golden, calculated, rtol=1e-05, atol=1e-04))
+
+        return float(cal_pcc)
+    except Exception:
+        return 0.0
+
+
+# ======================================================================================
+# Private Implementation
+# ======================================================================================
+
+
+def _is_ttnn_tensor(x):
+    """Safely detect TTNN tensors even if ttnn.Tensor is not defined in this environment."""
+    return isinstance(x, ttnn.Tensor)
+
+
+def _ttnn_op_layout_invariant(x, op_func, **kwargs):
+    """
+    Generic helper for TTNN operations that require TILE layout for non-sharded tensors.
+
+    Args:
+        x: TTNN tensor
+        op_func: TTNN operation function to call (e.g., ttnn.sign, ttnn.abs, ttnn.typecast)
+        **kwargs: Additional keyword arguments to pass to op_func (e.g., dtype=ttnn.bfloat16 for typecast)
+
+    Returns:
+        Result of op_func applied to x, with layout preserved.
+    """
+    # [ttnn contract] all TTNN operations require tensors to be in TILE layout when working with non-sharded tensors:
+    # ttnn.sign(), ttnn.abs(), ttnn.isinf(), ttnn.isnan(), ttnn.typecast(), ttnn.where(), ttnn.logical_not()
+    layout = x.get_layout()
+    if layout == ttnn.TILE_LAYOUT or x.is_sharded():  # sharded tensors can use either layout
+        return op_func(x, **kwargs)
+    else:
+        return ttnn.to_layout(op_func(ttnn.to_layout(x, ttnn.TILE_LAYOUT), **kwargs), layout)
+
+
+def _ttnn_max_scalar_all_dtype(x):
+    x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16)
+
+    # [ttnn contract] ttnn.max() internally calls a FillPad operation, which only supports the following dtypes:
+    # BFLOAT16
+    # FLOAT32
+    # UINT16, UINT32, INT32
+    # UINT8
+    # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details
+    # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes
+
+    # [ttnn contract] When called without a dim parameter, ttnn.max() returns the maximum value across the entire tensor as a scalar.
+    max_val_tensor = ttnn.max(x_bf16)
+
+    # [ttnn contract] When called without a dim parameter, ttnn.max() returns the maximum value across the entire tensor as a scalar.
+    # The method supports multiple data types:
+    # FLOAT32 → Python float
+    # BFLOAT16 → Python float (cast from bfloat16)
+    # BFLOAT8_B and BFLOAT4_B → Python float
+    return max_val_tensor.item()
+
+
+def _ttnn_mean_scalar_all_dtype(x):
+    x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16)
+
+    # [ttnn contract] ttnn.mean() internally calls a FillPad operation, which only supports the following dtypes:
+    # BFLOAT16
+    # FLOAT32
+    # UINT16, UINT32, INT32
+    # UINT8
+    # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details
+    # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes
+
+    # [ttnn contract] When called without a dim parameter, ttnn.mean() returns the mean value across the entire tensor as a scalar.
+    mean_val_tensor = ttnn.mean(x_bf16)
+    return mean_val_tensor.item()
+
+
+def _ttnn_min_scalar_all_dtype(x):
+    x_bf16 = _ttnn_op_layout_invariant(x, ttnn.typecast, dtype=ttnn.bfloat16)
+
+    # [ttnn contract] ttnn.mean() internally calls a FillPad operation, which only supports the following dtypes:
+    # BFLOAT16
+    # FLOAT32
+    # UINT16, UINT32, INT32
+    # UINT8
+    # see ttnn/cpp/ttnn/operations/data_movement/fill_pad/device/fill_pad_program_factory.hpp for more details
+    # [INFO] we cast to bfloat16 to avoid the limitation of the FillPad operation when dealing with e.g., bfloat8_b and bfloat4_b dtypes
+
+    # [ttnn contract] When called without a dim parameter, ttnn.mean() returns the mean value across the entire tensor as a scalar.
+    min_val_tensor = ttnn.min(x_bf16)
+    return min_val_tensor.item()

code/models/common/model_capabilities.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Maximum number of tokens that can simultaneously occupy the KV cache
+# across all concurrent users. This fallback applies to model/device
+# configurations not covered by a model-specific override.
+# Derived from the default branch of the per-model KV-cache rules in
+# the TT vLLM worker (tenstorrent/vllm#315).
+# See also: https://github.com/tenstorrent/vllm/issues/315
+FALLBACK_MAX_TOKENS_ALL_USERS = 131_072
+
+
+class ModelCapabilitiesMixin:
+    """Defines interface for hardware- or model-specific configurations.
+
+    NOTE: The default values here and per-model overrides will eventually be
+    unified with the corresponding vLLM scheduler configuration so that both
+    paths derive from the same source of truth.
+    """
+
+    @classmethod
+    def get_max_tokens_all_users(cls, **kwargs) -> int:
+        """Returns the fallback all-user KV-cache token capacity.
+
+        Used when no model- or device-specific override applies.
+        """
+        return FALLBACK_MAX_TOKENS_ALL_USERS

code/models/common/rmsnorm.py

@@ -0,0 +1,258 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import ttnn
+from models.common.lightweightmodule import LightweightModule
+from models.common.utility_functions import copy_to_buffer
+from models.tt_transformers.tt.common import Mode
+
+TILE = 32
+SHARD_HEIGHT = TILE  # Current ttnn.rms_norm implementation requires shard height to be a single tile
+
+
+class RMSNorm(LightweightModule):
+    """
+    RMSNorm supporting replication over a MeshDevice and sharding within devices.
+
+    This class implements a Root Mean Square Normalization (RMSNorm) that can be
+    distributed across multiple devices and cores. If the `device` parameter is a
+    MeshDevice, the weights and computations are replicated across all devices in
+    the mesh. Expects an interleaved input tensor, can optionally output a sharded tensor.
+
+    Args:
+        device: The device or MeshDevice on which to perform the computations.
+        state_dict: The state dictionary containing the model parameters.
+        dim: Input dimension (e.g. model hidden dimension size).
+        layer_num: The layer number to determine the weight key in the state dictionary.
+        weight_key: The key for retrieving the weight from the state dictionary.
+        weight_cache_path: Optional path for caching the tilized weights.
+        weight_memory_config: Configuration for the weight memory, default is DRAM_MEMORY_CONFIG.
+        weight_dtype: The data type for the tensors, bfp8_b hits >0.999 PCC in the models we tested.
+        model_config: Optional configuration dictionary for the model.
+        eps (float): Small value to avoid division by zero in normalization, default is 1e-05.
+
+    If model_config is provided, it must specify SHARDED_NORM_INPUT_MEMCFG, SHARDED_NORM_PRGM_CFG
+    and SHARDED_NORM_OUTPUT_MEMCFG. If not provided, default configurations will be generated.
+    """
+
+    def __init__(
+        self,
+        device,
+        dim,
+        state_dict,
+        weight_key,
+        layer_num=None,
+        state_dict_prefix=None,
+        weight_cache_path=None,
+        weight_memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        weight_dtype=ttnn.bfloat16,
+        is_distributed=None,
+        eps: float = 1e-05,
+        add_unit_offset=False,
+        sharded_program_config=None,
+        sharded_output_config=None,
+        output_mem_config=None,
+        ccl_topology=ttnn.Topology.Ring,
+        tt_ccl=None,
+        fp32_dest_acc_en=True,
+    ):
+        super().__init__()
+        self.device = device
+        self.eps = eps
+        self.is_distributed = is_distributed
+        self.ccl_topology = ccl_topology
+        self.tt_ccl = tt_ccl
+        self.add_unit_offset = add_unit_offset
+
+        if state_dict_prefix:
+            weight_name = f"{state_dict_prefix}{weight_key}.weight"
+        else:
+            if layer_num is None:
+                weight_name = f"{weight_key}.weight"
+            else:
+                weight_name = f"layers.{layer_num}.{weight_key}.weight"
+
+        torch_weight = (
+            state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT])
+        )
+
+        # Add offset before caching
+        if add_unit_offset:
+            torch_weight = torch_weight + 1.0
+
+        # Compatibility with models that don't use mesh devices (e.g. single-chip Mistral-7b)
+        is_mesh_device = device.__class__.__name__ == "MeshDevice"
+
+        self.weight = ttnn.as_tensor(
+            torch_weight,
+            device=device,
+            dtype=weight_dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            memory_config=weight_memory_config,
+            cache_file_name=None if weight_cache_path is None else weight_cache_path / weight_name,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(device) if is_mesh_device else None,
+        )
+
+        if self.is_distributed:
+            self.weight_distributed = ttnn.as_tensor(
+                torch_weight,
+                device=device,
+                dtype=weight_dtype,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                memory_config=weight_memory_config,
+                cache_file_name=(
+                    None if weight_cache_path is None else weight_cache_path / (weight_name + "_distributed")
+                ),
+                mesh_mapper=(
+                    ttnn.ShardTensor2dMesh(device, dims=(None, 2), mesh_shape=list(device.shape))
+                    if is_mesh_device
+                    else None
+                ),
+            )
+
+        self.sharded_output_config = sharded_output_config
+        self.sharded_program_config = sharded_program_config
+        self.output_mem_config = output_mem_config
+
+        self.compute_kernel_config_hifi2 = ttnn.WormholeComputeKernelConfig(
+            math_fidelity=ttnn.MathFidelity.HiFi2,
+            math_approx_mode=False,
+            fp32_dest_acc_en=fp32_dest_acc_en,
+            packer_l1_acc=True,
+        )
+
+    def update(self, *, weight: ttnn.Tensor) -> None:
+        """In-place replace the RMSNorm gamma via ``ttnn.copy``.
+
+        HF-format input: ``weight`` is HF ``...norm.weight``, shape
+        ``(1, 1, 1, dim)``, bf16, TILE, DRAM-interleaved, replicated.
+
+        ``copy_to_buffer`` reshapes to the storage shape
+        ``(1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT)`` and TILE -> ROW_MAJOR to
+        match ``self.weight``. ``add_unit_offset`` is not supported (see
+        assert): the caller must ship a gamma that already includes the +1.
+
+        When ``self.weight_distributed`` (the column-sharded mirror) exists it's
+        kept in sync on device: project ``self.weight`` into the sharded layout
+        via ``ttnn.mesh_partition`` (the inverse of the constructor's
+        ``ShardTensor2dMesh(dims=(None, 2))``, hence ``dim=2, cluster_axis=1``)
+        and ``ttnn.copy`` into it. Both buffers keep their address, so captured
+        traces and the prefetcher's recorded addresses stay valid.
+        """
+        assert not self.add_unit_offset, "RMSNorm.update does not support add_unit_offset=True"
+        copy_to_buffer(weight, self.weight, self.weight.dtype)
+
+        if getattr(self, "weight_distributed", None) is not None:
+            partitioned = ttnn.mesh_partition(
+                self.weight,
+                memory_config=self.weight_distributed.memory_config(),
+                dim=2,
+                cluster_axis=1,
+            )
+            copy_to_buffer(partitioned, self.weight_distributed, self.weight_distributed.dtype)
+
+    def forward(
+        self,
+        x: ttnn.Tensor,
+        mode: Mode | str,
+        in_sharded=False,
+        out_sharded=False,
+        norm_config=None,
+    ) -> ttnn.Tensor:
+        if isinstance(mode, str):
+            try:
+                mode = Mode(mode)
+            except ValueError:
+                raise ValueError(f"Invalid mode: {mode}")
+        elif not isinstance(mode, Mode):
+            raise ValueError(f"Invalid mode: {mode}")
+
+        sharded_program_config = norm_config.get("sharded_program_config") if norm_config else None
+        sharded_output_config = norm_config.get("sharded_output_config") if norm_config else None
+        output_mem_config = norm_config.get("output_mem_config") if norm_config else None
+        # Optional L1 placement for the distributed 3-op outputs (pre/gather/post); None -> DRAM default.
+        distributed_out_mc = norm_config.get("distributed_output_mem_config") if norm_config else None
+
+        # If input is sharded do sharded RMSNorm and optionally return sharded output
+        program_config = sharded_program_config if in_sharded else None
+        memory_config = sharded_output_config if out_sharded else None
+        distributed = self.is_distributed and self.is_distributed(mode)
+        weight = self.weight_distributed if distributed else self.weight
+
+        if in_sharded:
+            assert not distributed, "Distributed RMSNorm does not support sharded inputs"
+        else:
+            assert not out_sharded, "Non-sharded version of RMSNorm cannot output a sharded tensor"
+
+        if distributed:
+            x = self._distributed_rmsnorm(
+                x,
+                epsilon=self.eps,
+                weight=weight,
+                compute_kernel_config=self.compute_kernel_config_hifi2,
+                output_memory_config=distributed_out_mc,
+            )
+        else:
+            x = ttnn.rms_norm(
+                x,
+                epsilon=self.eps,
+                weight=weight,
+                program_config=program_config,
+                memory_config=memory_config,
+                compute_kernel_config=self.compute_kernel_config_hifi2,
+            )
+
+        if in_sharded and not out_sharded:
+            return ttnn.sharded_to_interleaved(x)
+        else:
+            if output_mem_config is not None:
+                x = ttnn.to_memory_config(x, output_mem_config)
+            return x
+
+    def _distributed_rmsnorm(
+        self,
+        inp,
+        epsilon=None,
+        weight=None,
+        program_config=None,
+        memory_config=None,
+        compute_kernel_config=None,
+        output_memory_config=None,
+    ):
+        assert program_config is None, "Distributed RMSNorm does not support sharded inputs"
+        assert memory_config is None, "Distributed RMSNorm does not support sharded outputs"
+        assert self.tt_ccl is not None, "Distributed RMSNorm requires tt_ccl"
+
+        # Interleaved output placement for the 3 ops; default DRAM (matches the prior hardcoded behavior).
+        mc = output_memory_config if output_memory_config is not None else ttnn.DRAM_MEMORY_CONFIG
+
+        # Run distributed rmsnorm part 1
+        tt_stats = ttnn.rms_norm_pre_all_gather(
+            inp, compute_kernel_config=compute_kernel_config, dtype=ttnn.bfloat16, memory_config=mc
+        )
+        # AllGather stats
+        tt_stats = ttnn.experimental.all_gather_async(
+            tt_stats,
+            persistent_output_buffer=None,
+            dim=3,
+            multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(),
+            num_links=1,
+            topology=self.ccl_topology,
+            memory_config=mc,
+            barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(),
+            chunks_per_sync=10,
+            num_workers_per_link=2,
+            num_buffers_per_channel=2,
+        )
+        # Run distributed rmsnorm part 2
+        tt_out = ttnn.rms_norm_post_all_gather(
+            inp,
+            tt_stats,
+            epsilon=epsilon,
+            weight=weight,
+            compute_kernel_config=compute_kernel_config,
+            memory_config=mc,
+        )
+        tt_stats.deallocate(True)
+
+        return tt_out

code/models/common/tensor_utils.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tensor utility functions for TTTv2 modules.
+"""
+
+import json
+import math
+import re
+
+import torch
+
+import ttnn
+
+# Standard tile size - hardware constant
+TILE_SIZE = ttnn.TILE_SIZE  # 32
+
+
+def nearest_multiple(value: int, multiple: int) -> int:
+    return math.ceil(value / multiple) * multiple
+
+
+def nearest_32(value: int) -> int:
+    return nearest_multiple(value, TILE_SIZE)
+
+
+def num_to_core_range_set(num_cores: int):
+    assert num_cores < 8 or num_cores % 8 == 0
+    num_x = min(num_cores, 8)
+    num_y = num_cores // num_x
+    assert num_x * num_y == num_cores
+    return ttnn.CoreRangeSet(
+        {
+            ttnn.CoreRange(
+                ttnn.CoreCoord(0, 0),
+                ttnn.CoreCoord(num_x - 1, num_y - 1),
+            )
+        }
+    )
+
+
+def get_out_subblock_w(per_core_n: int, out_subblock_h: int = 1) -> int:
+    out_subblock_w = 4
+    while out_subblock_w > 1:
+        if out_subblock_w * out_subblock_h <= 4 and per_core_n % out_subblock_w == 0:
+            break
+        out_subblock_w -= 1
+    return out_subblock_w
+
+
+def align_shape_to_tile(shape, tile_size: int = TILE_SIZE):
+    """Round up the last two dimensions of *shape* to multiples of *tile_size*.
+
+    This is the recommended replacement for the deprecated ``ttnn.pad_to_tile_shape``.
+
+    Args:
+        shape: An iterable of dimension sizes (list, tuple, or ttnn.Shape).
+        tile_size: Tile dimension to align to (default 32).
+
+    Returns:
+        List[int]: A new shape with the last two dims tile-aligned.
+    """
+    import math
+
+    result = list(shape)
+    if len(result) >= 1:
+        result[-1] = math.ceil(result[-1] / tile_size) * tile_size
+    if len(result) >= 2:
+        result[-2] = math.ceil(result[-2] / tile_size) * tile_size
+    return result
+
+
+def get_rot_transformation_mat(dhead: int = TILE_SIZE) -> torch.Tensor:
+    """
+    Create rotation transformation matrix for RoPE.
+
+    Constructs a permutation matrix that pairs adjacent dimensions with
+    signs (+1, -1) for the RoPE rotation:
+        [0, 1] → +1 at (0,1), -1 at (1,0)
+        [2, 3] → +1 at (2,3), -1 at (3,2)
+        ...
+
+    Used by ttnn.experimental.rotary_embedding_llama.
+
+    Args:
+        dhead: Matrix dimension. Must equal TILE_SIZE. Use TILE_SIZE for decode.
+
+    Returns:
+        torch.Tensor of shape [1, 1, dhead, dhead].
+    """
+    rot_emb_matrix = torch.zeros(1, 1, dhead, dhead)
+    rot_emb_matrix[..., torch.arange(0, dhead, 2), torch.arange(1, dhead, 2)] = 1
+    rot_emb_matrix[..., torch.arange(1, dhead, 2), torch.arange(0, dhead, 2)] = -1
+    return rot_emb_matrix
+
+
+def zeros_like_kv_cache(batch_size: int, n_kv_heads: int, max_seq_len: int, head_dim: int) -> torch.Tensor:
+    """Create zeros tensor for standard KV cache."""
+    return torch.zeros((batch_size, n_kv_heads, max_seq_len, head_dim))
+
+
+def zeros_like_paged_cache(paged_config, n_kv_heads: int, head_dim: int) -> torch.Tensor:
+    """Create zeros tensor for paged KV cache."""
+    return torch.zeros((paged_config.max_num_blocks, n_kv_heads, paged_config.block_size, head_dim))
+
+
+# todo)) add a on-device pad_dim_to_size function?
+def pad_dim_to_size(x: "torch.Tensor", dim: int, size: int) -> "torch.Tensor":
+    """Pads the specified dimension of the input tensor with zeros."""
+    if dim < 0:
+        dim = x.dim() + dim
+    current_size = x.size(dim)
+    pad_size = size - current_size
+
+    if pad_size < 0:
+        raise ValueError(f"Target size {size} is smaller than current size {current_size} on dim {dim}")
+
+    if pad_size == 0:
+        return x
+
+    pad = [0] * (2 * x.dim())
+    pad_index = 2 * (x.dim() - dim - 1)
+    pad[pad_index + 1] = pad_size
+
+    return torch.nn.functional.pad(x, pad, mode="constant", value=0)
+
+
+def pad_to_shape(x: "torch.Tensor", target_shape: tuple[int, ...], pad_value: float = 0.0) -> "torch.Tensor":
+    """Pad tensor to target_shape in a single F.pad call (more efficient than per-dim padding)."""
+    if x.shape == target_shape:
+        return x
+
+    # F.pad expects: (left_last, right_last, left_second_last, right_second_last, ...)
+    pad = []
+    for orig, target in zip(reversed(x.shape), reversed(target_shape)):
+        if target < orig:
+            raise ValueError(f"Target size {target} is smaller than current size {orig}")
+        pad.extend([0, target - orig])
+
+    return torch.nn.functional.pad(x, pad, mode="constant", value=pad_value)
+
+
+def get_padded_hidden_dim(hidden_dim: int, num_devices: int, tile_size: int = 32) -> int:
+    """
+    Compute padded hidden_dim to satisfy ttnn.from_torch's tile alignment constraint.
+
+    ttnn.from_torch requires physical shard shapes to be tile-aligned. When sharding
+    a tensor across devices, each shard_dim = hidden_dim / num_devices must be
+    divisible by tile_size.
+
+    We pad the global tensor first, then shard evenly so only the last shard has padding.
+    """
+    shard_dim = hidden_dim // num_devices
+    padded_shard = ((shard_dim + tile_size - 1) // tile_size) * tile_size
+    return padded_shard * num_devices
+
+
+def parse_shard_dims_from_mesh_mapper_config(mesh_mapper_config: ttnn.MeshMapperConfig) -> list[int]:
+    """
+    Parse shard dimensions from MeshMapperConfig's repr.
+
+    MeshMapperConfig doesn't expose .placements directly, but repr shows them:
+        'MeshMapperConfig(placements: [PlacementShard(-1)], mesh_shape_override=MeshShape([8]))'
+
+    This parses out the shard dimensions (e.g., [-1]) from PlacementShard entries.
+    Returns empty list if no PlacementShard found (e.g., replicated).
+
+    Note: This is a workaround until TTNN exposes .placements directly.
+    """
+    config_repr = repr(mesh_mapper_config)
+    matches = re.findall(r"PlacementShard\((-?\d+)\)", config_repr)
+    return [int(d) for d in matches]
+
+
+def memory_config_to_dict(memory_config: ttnn.MemoryConfig):
+    # Convert to plain types for deterministic serialization.
+    return {
+        "memory_layout": str(memory_config.memory_layout),
+        "buffer_type": str(memory_config.buffer_type),
+        "shard_spec": str(memory_config.shard_spec),
+        "is_sharded": bool(memory_config.is_sharded()),
+        "interleaved": bool(memory_config.interleaved),
+        "hash": int(memory_config.__hash__()),
+    }
+
+
+def compute_kernel_config_to_str(compute_kernel_config: ttnn.WormholeComputeKernelConfig):
+    # Backward compat shim; prefer compute_kernel_config_to_dict + serialize_config.
+    cfg = compute_kernel_config_to_dict(compute_kernel_config)
+    return serialize_config(cfg)
+
+
+def compute_kernel_config_to_dict(compute_kernel_config: ttnn.WormholeComputeKernelConfig):
+    return {
+        "math_fidelity": str(compute_kernel_config.math_fidelity),
+        "math_approx_mode": str(compute_kernel_config.math_approx_mode),
+        "fp32_dest_acc_en": bool(compute_kernel_config.fp32_dest_acc_en),
+        "packer_l1_acc": bool(compute_kernel_config.packer_l1_acc),
+        "dst_full_sync_en": bool(compute_kernel_config.dst_full_sync_en),
+        "throttle_level": str(compute_kernel_config.throttle_level),
+    }
+
+
+def program_config_to_str(program_config: ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig):
+    # Backward compat shim; prefer program_config_to_dict + serialize_config.
+    cfg = program_config_to_dict(program_config)
+    return serialize_config(cfg)
+
+
+def program_config_to_dict(program_config):
+    if hasattr(program_config, "to_json"):
+        d = json.loads(program_config.to_json())
+        d["type"] = type(program_config).__name__
+        return d
+    else:
+        return {"type": type(program_config).__name__, "repr": repr(program_config)}
+
+
+def serialize_config(cfg_dict: dict, fmt: str = "json") -> str:
+    if fmt == "json":
+        return json.dumps(cfg_dict, sort_keys=True)
+    if fmt == "yaml":
+        try:
+            import yaml
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise RuntimeError("PyYAML is required for yaml serialization") from exc
+        return yaml.safe_dump(cfg_dict, sort_keys=True)
+    raise ValueError(f"Unsupported format: {fmt}")

code/models/common/utility_functions.py

@@ -0,0 +1,1266 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import os
+import struct
+import time
+from typing import Union
+
+import numpy as np
+import pytest
+import torch
+from loguru import logger
+from ttnn.device import Arch
+from typing_extensions import deprecated
+
+import ttnn
+
+
+def get_mesh_device():
+    """Fixture to provide mesh device configuration."""
+    mesh_device = os.environ.get("MESH_DEVICE", "N150")
+    mesh_config = {
+        "N150": (1, 1),
+        "N300": (2, 1),
+        "T3K": (8, 1),
+        "TG": (8, 4),
+    }.get(mesh_device, (ttnn.get_num_devices(), 1))
+    return mesh_config
+
+
+### Math operations ###
+def _nearest_32(x):
+    return math.ceil(x / 32) * 32
+
+
+def nearest_32(
+    x,
+):  # needs refctoring; to match alias called in some scripts (e.g. test_padding_test in unit tests)
+    return _nearest_32(x)
+
+
+def _nearest_y(x, y):
+    return math.ceil(x / y) * y
+
+
+def nearest_y(x, y):
+    return _nearest_y(x, y)
+
+
+def divup(a, b):
+    return (a + b - 1) // b
+
+
+def roundup(a, b):
+    result = divup(a, b) * b
+    return result
+
+
+def roundup32(a):
+    return roundup(a, 32)
+
+
+def float_to_bits(x):
+    s = struct.pack(">f", x)
+    return struct.unpack(">l", s)[0]
+
+
+def torch_random(shape, low, high, dtype):
+    if dtype in [torch.int64, torch.int32, torch.int16, torch.int8]:
+        return torch.randint(low, high, shape, dtype=dtype)
+    return torch.zeros(shape, dtype=dtype).uniform_(low, high)
+
+
+def torch_random_with_zeros(shape, low, high, dtype, zero_fraction=0.1):
+    total_elements = torch.prod(torch.tensor(shape)).item()
+    num_zeros = int(total_elements * zero_fraction)
+    num_random = total_elements - num_zeros
+
+    # Generate random values between low and high
+    random_values = torch.empty(num_random).uniform_(low, high)
+    zeros = torch.zeros(num_zeros)
+
+    # Combine zeros and random values
+    combined = torch.cat([zeros, random_values])
+
+    # Shuffle the tensor
+    shuffled = combined[torch.randperm(combined.size(0))]
+
+    # Reshape to the desired shape
+    result_tensor = shuffled.view(shape)
+    result_tensor.to(dtype)
+    return result_tensor
+
+
+### Profiling ###
+class Profiler:
+    def __init__(self):
+        self.start_times = dict()
+        self.times = dict()
+        self.disabled = False
+
+    def clear(self):
+        self.start_times = dict()
+        self.times = dict()
+        self.disabled = False
+
+    def enable(self):
+        self.disabled = False
+
+    def disable(self):
+        self.disabled = True
+
+    def start(self, key, force_enable=False):
+        if self.disabled and not force_enable:
+            return
+
+        self.start_times[key] = time.time()
+
+    def end(self, key, PERF_CNT=1, force_enable=False):
+        if self.disabled and not force_enable:
+            return
+
+        if key not in self.start_times:
+            return
+
+        diff = time.time() - self.start_times[key]
+
+        if key not in self.times:
+            self.times[key] = []
+
+        self.times[key].append(diff / PERF_CNT)
+
+    def get(self, key):
+        if key not in self.times:
+            return 0
+
+        return sum(self.times[key]) / len(self.times[key])
+
+    def print(self, units="s"):
+        for key in self.times:
+            average = self.get(key)
+            if units == "s":
+                pass
+            elif units == "ms":
+                average *= 1000
+            elif units == "us":
+                average *= 1000000
+            elif units == "ns":
+                average *= 1000000000
+            else:
+                raise ValueError(f"Invalid units: {units}")
+            print(f"{key}: {average:.3f}{units}")
+
+
+profiler = Profiler()
+
+
+### Turn flags on/off ###
+def enable_memory_reports():
+    """
+    Enables generating reports of memory allocation statistics in .reports/tt_metal dir
+    """
+    return ttnn.device.EnableMemoryReports()
+
+
+def disable_memory_reports():
+    """
+    Disables generating reports of memory allocation statistics
+    """
+    return ttnn.device.DisableMemoryReports()
+
+
+### Tensor conversion ###
+def torch2tt_tensor(
+    py_tensor: torch.Tensor,
+    tt_device,
+    tt_layout=ttnn.TILE_LAYOUT,
+    tt_memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED),
+    tt_dtype=ttnn.bfloat16,
+):
+    size = list(py_tensor.size())
+
+    while len(size) < 4:
+        size.insert(0, 1)
+
+    tt_tensor = ttnn.Tensor(py_tensor.reshape(size), tt_dtype)
+    tt_tensor = tt_tensor.to(tt_layout)
+
+    if tt_device is not None:
+        tt_tensor = tt_tensor.to(tt_device, tt_memory_config)
+    else:
+        tt_tensor = tt_tensor.cpu()
+
+    return tt_tensor
+
+
+def tt_tensors_to_torch_tensors(
+    tt_tensors_device: ttnn.Tensor, mesh_device: Union[ttnn.MeshDevice, ttnn.Device], concat_dim: int = 0
+):
+    # Convert tensors to interleaved
+    if tt_tensors_device.is_sharded():
+        tt_tensors_device = ttnn.sharded_to_interleaved(tt_tensors_device)
+
+    # Convert tensors to RM layout
+    if tt_tensors_device.layout == ttnn.TILE_LAYOUT:
+        # Convert to bfloat16 to ensure untilize works
+        if tt_tensors_device.dtype != ttnn.bfloat16:
+            tt_tensors_device = ttnn.clone(
+                tt_tensors_device, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG
+            )
+        # Untilize using singlecore since multicore version runs out of l1 memory (Issue #9022)
+        tt_tensors_device = ttnn.untilize(tt_tensors_device, use_multicore=False)
+
+    return torch.cat([t.to_torch() for t in ttnn.get_device_tensors(tt_tensors_device.cpu())], dim=concat_dim)
+
+
+def tt2torch_tensor(tt_tensor):
+    tt_output = tt_tensor.cpu()
+    if tt_output.get_layout() != ttnn.ROW_MAJOR_LAYOUT:
+        tt_output = tt_output.to(ttnn.ROW_MAJOR_LAYOUT)
+    return tt_output.to_torch()
+
+
+def tt_to_torch_tensor(tt_tensor):
+    tt_output = tt_tensor.cpu().to(ttnn.ROW_MAJOR_LAYOUT)
+    return tt_output.to_torch()
+
+
+def torch_to_tt_tensor_rm(py_tensor, device, shape=None, put_on_device=True):
+    if shape is None:
+        shape = list(py_tensor.size())
+        while len(shape) < 4:
+            shape.insert(0, 1)
+
+    tt_tensor = ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16)
+    if put_on_device:
+        tt_tensor = tt_tensor.to(device)
+    return tt_tensor
+
+
+def torch_to_tt_tensor(py_tensor, device):
+    shape = list(py_tensor.size())
+    while len(shape) < 4:
+        shape.insert(0, 1)
+
+    tt_tensor = (
+        ttnn.Tensor(py_tensor.reshape(shape), ttnn.bfloat16)
+        .to(
+            ttnn.TILE_LAYOUT
+        )  # change memory layout of TT Tensor to TILE (as operation that will use it expects TILE layout)
+        .to(device)  # move TT Tensor from host to TT accelerator device (device is of type ttnn.device.Device)
+    )
+
+    return tt_tensor
+
+
+def unpad_from_zero(x, desired_shape):
+    if x.padded_shape[-1] == desired_shape[-1] and x.padded_shape[-2] == desired_shape[-2]:
+        x = tt2torch_tensor(x)
+    else:
+        x = x.cpu()
+        if x.get_layout() != ttnn.ROW_MAJOR_LAYOUT:
+            x = x.to(ttnn.ROW_MAJOR_LAYOUT)
+        x = x.unpad(
+            (0, 0, 0, 0),
+            (
+                desired_shape[0],
+                desired_shape[1],
+                desired_shape[2],
+                desired_shape[3],
+            ),
+        )
+
+        x = x.to_torch()
+    return x
+
+
+def pad_activation(x):
+    """
+    This function pads an activation with 0s as a pre-preprocessing step to tilization.
+
+    In the 2d case, it pads a vector to the right with 0s, and in the 2+d case,
+    it pads the bottom and right corners of the last two dimensions.
+
+    :param x: Input PyTorch Tensor
+    :type x: class:`torch.Tensor`
+
+    WARNING: This function should eventually be retired in favour of padding on device
+    """
+    nearest_32 = _nearest_32
+
+    assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor"
+    assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported"
+    if len(x.shape) == 1:  # (num_features,)
+        padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0]))
+        padded_tensor[:, 0, 0, : x.shape[0]] = x
+    elif len(x.shape) == 2:  # (batch, num features)
+        padded_tensor = torch.zeros(x.shape[0], 1, 32, nearest_32(x.shape[1]))
+        padded_tensor[:, 0, 0, : x.shape[1]] = x
+    elif len(x.shape) == 3:  # (batch, num features y, num features x)
+        padded_tensor = torch.zeros(x.shape[0], 1, nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
+        padded_tensor[..., 0, : x.shape[-2], : x.shape[-1]] = x
+    else:  # (batch, num channels, num features y, num features x)
+        padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
+        padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x
+    return padded_tensor
+
+
+def pad_weight(x):
+    """
+    This function pads a weight/bias with 0s as a pre-preprocessing step to tilization.
+
+    tt_tensor = ttnn.Tensor(
+        py_tensor.reshape(shape), ttnn.bfloat16
+    In the 2d case, it pads a vector to the right with 0s, and in the 2+d case,
+    it pads the bottom and right corners of the last two dimensions.
+
+    :param x: Input PyTorch Tensor
+    :type x: class:`torch.Tensor`
+
+    WARNING: This function should eventually be retired in favour of padding on device
+    """
+    nearest_32 = _nearest_32
+
+    assert isinstance(x, torch.Tensor), "Input to this function must be an instance of torch.Tensor"
+    assert len(x.shape) >= 1 and len(x.shape) <= 4, "Only tensors with dimension 1-4 supported"
+
+    if len(x.shape) == 1:  # (num_features,)
+        padded_tensor = torch.zeros(1, 1, 32, nearest_32(x.shape[0]))
+        padded_tensor[:, 0, 0, : x.shape[0]] = x
+    elif len(x.shape) == 2:  # (r_features, c_features)
+        padded_tensor = torch.zeros(1, 1, nearest_32(x.shape[0]), nearest_32(x.shape[1]))
+        padded_tensor[:, 0, : x.shape[0], : x.shape[1]] = x
+    else:
+        padded_tensor = torch.zeros(*x.shape[:-2], nearest_32(x.shape[-2]), nearest_32(x.shape[-1]))
+        padded_tensor[..., : x.shape[-2], : x.shape[-1]] = x
+
+    return padded_tensor
+
+
+def convert_weights_2d_matrix(weights, w_shape):
+    """
+    :param weights: Input PyTorch Tensor
+    :type weights: class:`torch.Tensor`
+    """
+    ret_shape = [1, 1, w_shape[0], w_shape[1] * w_shape[2] * w_shape[3]]
+    if isinstance(weights, torch.Tensor):
+        ret = torch.zeros(np.prod(ret_shape))
+    else:
+        ret = np.zeros(np.prod(ret_shape))
+    idx = 0
+    for k in range(w_shape[0]):
+        for r in range(w_shape[2]):
+            for s in range(w_shape[3]):
+                for c in range(w_shape[1]):
+                    ret[idx] = weights[k][c][r][s]
+                    idx += 1
+    assert idx == np.prod(ret_shape)
+    return ret.reshape(ret_shape).transpose(2, 3)
+
+
+def convert_act_2d_matrix(activation, kernel_y, kernel_x, stride_y, stride_x, pad_y, pad_x):
+    """
+    :param activation: Input PyTorch Tensor
+    :type activation: class:`torch.Tensor`
+    """
+    N = activation.shape[0]
+    C = activation.shape[1]
+    H = activation.shape[2]
+    W = activation.shape[3]
+
+    OH = (int)((H - kernel_y + 2 * pad_y) // stride_y) + 1
+    OW = ((W - kernel_x + 2 * pad_x) // stride_x) + 1
+    nrows = OH * OW
+    ncols = C * kernel_x * kernel_y
+    ret_shape = [1, N, nrows, ncols]
+    if isinstance(activation, torch.Tensor):
+        ret = torch.zeros(np.prod(ret_shape))
+    else:
+        ret = np.zeros(np.prod(ret_shape))
+    idx = 0
+    for n in range(N):
+        for h in range(-1 * pad_y, H + pad_y - kernel_y + 1, stride_y):
+            for w in range(-1 * pad_x, W + pad_x - kernel_x + 1, stride_x):
+                for r in range(kernel_y):
+                    for s in range(kernel_x):
+                        for c in range(C):
+                            h_offs = h + r
+                            w_offs = w + s
+                            pad = h_offs < 0 or h_offs >= H or w_offs < 0 or w_offs >= W
+                            ret[idx] = 0 if pad else activation[n][c][h_offs][w_offs]
+                            idx += 1
+    assert idx == np.prod(ret_shape)
+    return ret.reshape(ret_shape)
+
+
+### Tilizing / Untilizing ###
+@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
+def tilize(x):
+    return x
+
+
+@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
+def tilize_to_list(x):
+    """
+    Returns a flattened list of the tensor
+    """
+    return tilize(x).reshape(-1).tolist()
+
+
+@deprecated("PyTorch data is handled automatically in tensor infra. This function does nothing now:")
+def untilize(x):
+    return x
+
+
+### Measuring accuracy and other metrics ###
+def is_close(a, b, rtol=1e-2, atol=1e-2, max_mag=2.0, max_mag_fraction=0.02):
+    """
+    A variant of np.isclose with logging.
+    """
+    absdiff = (a - b).abs()
+    reldiff1 = (a.abs() / b.abs()) - 1.0
+    reldiff2 = (a.abs() + 1.0) / (b.abs() + 1.0) - 1.0  # in case b.abs() is 0
+    reldiff_or = torch.logical_or(reldiff1.abs() < rtol, reldiff2.abs() < rtol)
+    max_mag_ok = absdiff < max_mag * max_mag_fraction
+
+    or_abs_rel = torch.logical_or(absdiff < atol, reldiff_or)
+    or_abs_rel = torch.logical_or(or_abs_rel, max_mag_ok)
+    debug_index = or_abs_rel.to(torch.int32).argmin().item()
+
+    if not or_abs_rel.reshape(-1)[debug_index]:
+        logger.info(f"isclose mismatch at index={debug_index}")
+        logger.info(a.reshape(-1)[debug_index])
+        logger.info(b.reshape(-1)[debug_index])
+        logger.info(f"reldiff1={reldiff1.reshape(-1)[debug_index]}")
+        logger.info(f"reldiff2={reldiff2.reshape(-1)[debug_index]}")
+        logger.info(f"absdiff={absdiff.reshape(-1)[debug_index]}")
+
+        HT = a.shape[-2] // 32
+        WT = a.shape[-1] // 32
+        hwt = debug_index // 1024
+        wt = hwt % WT
+        ht = hwt // WT
+        h = (debug_index % 1024) // 32
+        w = (debug_index % 1024) % 32
+
+        logger.info(f"****    at {debug_index} --- HTWT={ht} {wt} HW={h} {w}")
+
+    return torch.all(or_abs_rel)
+
+
+def _comp_nonfinite(golden, calculated):
+    """
+    Returns True if tensors contain the same non-finite values (nan, inf, -inf) at the same positions. Also returns True if all elements are finite.
+    Returns False if non-finite values differ between both tensors.
+    """
+
+    # torch.equal(['nan'], ['nan']] => False
+    # For this reason, we check for nan and inf separately
+    if torch.not_equal(torch.isnan(golden), torch.isnan(calculated)).any():
+        return False
+
+    golden_inf_mask = torch.isinf(golden)
+    calculated_inf_mask = torch.isinf(calculated)
+
+    if torch.not_equal(golden_inf_mask, calculated_inf_mask).any():
+        return False
+
+    golden_inf = golden[golden_inf_mask]
+    calculated_inf = calculated[calculated_inf_mask]
+    return torch.equal(golden_inf, calculated_inf)
+
+
+def comp_allclose(golden, calculated, rtol=1e-05, atol=1e-08):
+    if golden.dtype != calculated.dtype:
+        calculated = calculated.type(golden.dtype)
+
+    atol_delta = torch.max(torch.abs(golden - calculated)).item()
+    rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item()
+    return (
+        torch.allclose(golden, calculated, rtol, atol, True),
+        f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}",
+    )
+
+
+def comp_pcc(golden, calculated, pcc=0.99, rtol=1e-05, atol=1e-04):
+    golden = torch.Tensor(golden)
+    calculated = torch.Tensor(calculated)
+
+    if golden.dtype != calculated.dtype:
+        calculated = calculated.type(golden.dtype)
+
+    if torch.all(torch.isnan(golden)) and torch.all(torch.isnan(calculated)):
+        logger.warning("Both tensors are 'nan'")
+        return True, 1.0
+
+    if torch.all(torch.isnan(golden)) or torch.all(torch.isnan(calculated)):
+        logger.error("One tensor is all nan, the other is not.")
+        return False, 0.0
+
+    # Test if either is completely zero — but a zero tensor is also a constant tensor,
+    # so fall back to allclose instead of a hard 0.0: zero-vs-small-constant may be
+    # within the caller's tolerances.
+    if torch.any(golden.bool()) != torch.any(calculated.bool()):
+        logger.warning("One tensor is all zero. PCC undefined; falling back to allclose.")
+        result = torch.allclose(golden, calculated, rtol=rtol, atol=atol)
+        return result, float(result)
+
+    golden = torch.squeeze(golden).flatten()
+    calculated = torch.squeeze(calculated).flatten()
+
+    # For now, mask all infs and nans (to zero) so that we check the rest... TODO
+    # Skip this for integer types which don't have NaN/Inf values.
+    if golden.dtype.is_floating_point:
+        # FP8 doesn't support isfinite/nan_to_num and bfloat16 products lose precision,
+        # so correlate these in float32.
+        if golden.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.bfloat16):
+            golden = golden.to(torch.float32)
+            calculated = calculated.to(torch.float32)
+
+        # Zero out NaN/Inf, preserving the historical PCC values. nan_to_num allocates a
+        # full-size copy of each tensor, so only do it when invalid values are actually
+        # present; on the common all-finite path the tensors stay as views and no copy is
+        # made (this short-circuit is what keeps peak memory near 1x of one input).
+        if not bool((torch.isfinite(golden) & torch.isfinite(calculated)).all()):
+            golden = torch.nan_to_num(golden, nan=0.0, posinf=0.0, neginf=0.0)
+            calculated = torch.nan_to_num(calculated, nan=0.0, posinf=0.0, neginf=0.0)
+
+    if torch.equal(golden, calculated):
+        return True, 1.0
+
+    # Integer tensors must be correlated in floating point (centering/products would
+    # otherwise truncate/overflow). float32 keeps the working set small.
+    if not golden.dtype.is_floating_point:
+        golden = golden.to(torch.float32)
+        calculated = calculated.to(torch.float32)
+
+    # Pearson r with float64 *accumulation* (dtype= on the reductions) over the float32
+    # data: no float64 copy of either tensor is materialized, so peak memory stays near
+    # 1x of one input on large tensors while matching a full-float64 correlation to
+    # |Δ|<1e-9 across the high-PCC (>=0.999) range.
+    n = golden.numel()
+    g_centered = golden - (golden.sum(dtype=torch.float64) / n).to(golden.dtype)
+    c_centered = calculated - (calculated.sum(dtype=torch.float64) / n).to(calculated.dtype)
+    cov = (g_centered * c_centered).sum(dtype=torch.float64)
+    g_sq_sum = g_centered.pow(2).sum(dtype=torch.float64)
+    c_sq_sum = c_centered.pow(2).sum(dtype=torch.float64)
+    denom = torch.sqrt(g_sq_sum * c_sq_sum)
+    # pow/sum stay in float32 before the reduction; large-magnitude tensors (e.g. ldexp)
+    # can overflow to inf here even though float64 accumulation would be finite.
+    if not math.isfinite(denom.item()) or not math.isfinite(cov.item()):
+        g_centered64 = g_centered.to(torch.float64)
+        c_centered64 = c_centered.to(torch.float64)
+        cov = (g_centered64 * c_centered64).sum()
+        denom = torch.sqrt(g_centered64.pow(2).sum() * c_centered64.pow(2).sum())
+    cal_pcc = (cov / denom).item()
+
+    # Zero variance -> denom == 0 -> cal_pcc is nan: PCC is undefined for constant tensors.
+    # Fall back to allclose rather than returning a misleading 1.0.
+    if math.isnan(cal_pcc):
+        logger.warning("PCC is NaN (zero variance / constant tensor). Falling back to allclose check.")
+        result = torch.allclose(golden, calculated, rtol=rtol, atol=atol)
+        return result, float(result)
+
+    return cal_pcc >= pcc, cal_pcc
+
+
+def ulp(x: Union[ttnn.Tensor, torch.Tensor]) -> Union[ttnn.Tensor, torch.Tensor]:
+    "Return Unit of Least Precision for each element of a given tensor"
+
+    received_ttnn_input = False
+    if isinstance(x, ttnn.Tensor):
+        x = ttnn.to_torch(x)
+        received_ttnn_input = True
+
+    # Notes:
+    # - This should be identical to the definition of ULP by Goldberg
+    #   "What every computer scientist should know about floating-point arithmetic"
+    #   https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html
+    # - We use torch.abs(x) to ensure symmetry ULP(-x) == ULP(x)
+    # - For x powers of 2, x + ULP(x) is not closest number but second closest (previous number is 2x closer)
+    #   However, this avoids rounding-to-nearest-tie-to-even issues on addition (i.e. x + ULP(x) != x)
+    abs_x = torch.abs(x)
+    next = torch.nextafter(
+        abs_x, torch.tensor(math.inf, dtype=x.dtype)
+    )  # 1 ULP ~ Difference between two consecutive floating point numbers
+    ulp_value = next - abs_x
+
+    # Special case: if abs_x == torch.finfo(x.dtype).max, then next == math.inf, which leads to ULP(x) == inf rather than finite number
+    # We fix this problem by manually calculating ULP at max value, and masking tensor when input == max
+    dtype_max = torch.finfo(x.dtype).max
+    max_epsilon = dtype_max - torch.nextafter(
+        torch.tensor(dtype_max, dtype=x.dtype), torch.tensor(-math.inf, dtype=x.dtype)
+    )
+    ulp_value = torch.where(abs_x == dtype_max, max_epsilon, ulp_value)
+
+    if received_ttnn_input:  # Ensures that type(input) == type(output)
+        ulp_value = ttnn.from_torch(ulp_value)
+
+    return ulp_value
+
+
+def comp_ulp(golden, calculated, ulp_threshold, allow_nonfinite=False):
+    """
+    Compute absolute error between two tensors in Units of Least Precision (ULP)
+    """
+
+    # If both tensors are empty, then we can return True
+    if torch.numel(golden) == 0 and torch.numel(calculated) == 0:
+        return True, "Both tensors are empty"
+
+    if not allow_nonfinite and not torch.all(torch.isfinite(calculated)):
+        return False, "Calculated tensor contains non-finite values"
+
+    if not _comp_nonfinite(golden, calculated):
+        return False, "Tensors are not finite at the same positions"
+    # nonfinite elements can interfere with ULP error calculation
+    # To avoid this, replace nan, +inf, -inf with 0
+    # (we have already checked that both tensors have the same nonfinite elements)
+    mask_finite = ~torch.isfinite(golden)
+    golden = golden.clone()
+    calculated = calculated.clone()
+    golden[mask_finite] = 0
+    calculated[mask_finite] = 0
+
+    # ULP is measured according to the golden tensor
+    # In most cases, data type of golden tensor should be the same as calculated tensor.
+    # However, in some cases, we may want to measure < 1 ULP differences, which requires golden tensor
+    # to have higher precision than calculated tensor.
+    # If we passed golden tensor to ulp() as is, we would get ULP of higher precision.
+    # e.g. ulp of float32 rather bfloat16 calculation, which would give us a wrong value.
+    ulp_value = ulp(golden.type(calculated.dtype))
+
+    if golden.dtype != calculated.dtype:  # Note: assumes that golden has higher precision than calculated tensor
+        calculated = calculated.type(golden.dtype)
+        ulp_value = ulp_value.type(golden.dtype)  # Convert ULP to higher precision (for sub-1 ULP measurements)
+
+    ulp_tensor = torch.abs(calculated - golden) / ulp_value
+    ulp_delta = torch.max(ulp_tensor)
+    within_threshold = ulp_delta <= ulp_threshold
+    message = f"Max ULP Delta: {ulp_delta}"
+    if not within_threshold:
+        ulp_index = torch.argmax(ulp_tensor)
+        ulp_index_tuple = tuple(int(idx) for idx in torch.unravel_index(ulp_index, golden.shape))
+        message += (
+            f" @ {list(ulp_index_tuple)} = "
+            f"|calculated {calculated[ulp_index_tuple]} - golden {golden[ulp_index_tuple]}| "
+            f"/ ULP(golden) {ulp_value[ulp_index_tuple]}"
+        )
+    return (within_threshold, message)
+
+
+def calculate_detailed_ulp_stats(expected, actual):
+    """
+    Calculate detailed ULP statistics for analysis.
+
+    Returns:
+        dict: Dictionary with ULP statistics including max, mean, std, and percentiles
+    """
+    if isinstance(actual, ttnn.Tensor):
+        actual = ttnn.to_torch(actual)
+    if isinstance(expected, ttnn.Tensor):
+        expected = ttnn.to_torch(expected)
+
+    # Convert to bfloat16 if not already
+    expected = expected.to(torch.bfloat16)
+    actual = actual.to(torch.bfloat16)
+
+    # Handle special cases
+    if torch.allclose(expected, actual, rtol=0, atol=0, equal_nan=True):
+        return {
+            "max_ulp": 0.0,
+            "mean_ulp": 0.0,
+            "median_ulp": 0.0,
+            "std_ulp": 0.0,
+            "p95_ulp": 0.0,
+            "p99_ulp": 0.0,
+            "perfect_matches": 1.0,
+        }
+
+    # Convert bfloat16 to uint16 representation for bit manipulation
+    expected_bits = expected.view(torch.int16).to(torch.int32)
+    actual_bits = actual.view(torch.int16).to(torch.int32)
+
+    # Handle sign differences
+    expected_sign = expected_bits < 0
+    actual_sign = actual_bits < 0
+    same_sign = expected_sign == actual_sign
+
+    # Calculate ULP differences
+    expected_abs_bits = torch.where(expected_sign, -expected_bits, expected_bits)
+    actual_abs_bits = torch.where(actual_sign, -actual_bits, actual_bits)
+
+    ulp_diff = torch.where(same_sign, torch.abs(expected_bits - actual_bits), expected_abs_bits + actual_abs_bits)
+
+    # Handle non-finite values
+    expected_finite = torch.isfinite(expected)
+    actual_finite = torch.isfinite(actual)
+    both_finite = expected_finite & actual_finite
+
+    ulp_diff = torch.where(both_finite, ulp_diff, torch.tensor(float("inf")))
+
+    # Handle same non-finite values
+    both_nan = torch.isnan(expected) & torch.isnan(actual)
+    both_posinf = torch.isposinf(expected) & torch.isposinf(actual)
+    both_neginf = torch.isneginf(expected) & torch.isneginf(actual)
+    same_nonfinite = both_nan | both_posinf | both_neginf
+
+    ulp_diff = torch.where(same_nonfinite, torch.tensor(0.0), ulp_diff)
+
+    # Calculate statistics only on finite ULP differences
+    finite_ulp = ulp_diff[torch.isfinite(ulp_diff)]
+
+    if len(finite_ulp) == 0:
+        return {
+            "max_ulp": float("inf"),
+            "mean_ulp": float("inf"),
+            "median_ulp": float("inf"),
+            "std_ulp": float("inf"),
+            "p95_ulp": float("inf"),
+            "p99_ulp": float("inf"),
+            "perfect_matches": 0.0,
+        }
+
+    finite_ulp_float = finite_ulp.float()
+    perfect_matches = (finite_ulp == 0).float().mean().item()
+
+    return {
+        "max_ulp": torch.max(finite_ulp).item(),
+        "mean_ulp": torch.mean(finite_ulp_float).item(),
+        "median_ulp": torch.median(finite_ulp_float).item(),
+        "std_ulp": torch.std(finite_ulp_float).item(),
+        "p95_ulp": torch.quantile(finite_ulp_float, 0.95).item(),
+        "p99_ulp": torch.quantile(finite_ulp_float, 0.99).item(),
+        "perfect_matches": perfect_matches,
+    }
+
+
+def comp_allclose_and_pcc(golden, calculated, rtol=1e-05, atol=1e-08, pcc=0.99):
+    # 0-volume tensors are special because they don't have elements, so we can't compute PCC, etc.
+    # If one of the tensors is a 0-volume tensor, simply call torch.equal to check if they are equal
+    # (i.e. that both are 0-volume tensors and they have equal shapes).
+    if golden.numel() == 0 or calculated.numel() == 0:
+        return torch.equal(golden, calculated), f"{golden} != {calculated}"
+
+    if golden.dtype != calculated.dtype:
+        calculated = calculated.type(golden.dtype)
+
+    passing = True
+    output = ""
+    passing_allclose, output_allclose = comp_allclose(golden, calculated, rtol, atol)
+    passing &= passing_allclose
+    output += output_allclose
+    if torch.numel(golden) != 1:
+        passing_pcc, output_pcc = comp_pcc(golden, calculated, pcc, rtol=rtol, atol=atol)
+        passing &= passing_pcc
+        output += f", pcc={output_pcc}"
+
+    return passing, output
+
+
+def comp_equal(golden, calculated):
+    if golden.dtype != calculated.dtype:
+        calculated = calculated.type(golden.dtype)
+
+    # If either tensor is zero-volume, broadcasting can still yield an empty delta and
+    # crash torch.max(); defer entirely to torch.equal (False on shape mismatch).
+    if golden.numel() == 0 or calculated.numel() == 0:
+        return torch.equal(golden, calculated), f"{golden} != {calculated}"
+
+    atol_delta = torch.max(torch.abs(golden - calculated)).item()
+    rtol_delta = torch.max(torch.abs(golden - calculated) / torch.abs(calculated)).item()
+    return (
+        torch.equal(golden, calculated),
+        f"Max ATOL Delta: {atol_delta}, Max RTOL Delta: {rtol_delta}",
+    )
+
+
+def get_oom_of_float(float_lst):
+    """
+    Given a list of floats, returns a list of the order or magnitudes
+    of the floats. Useful when you want to make sure that even if your
+    tt outputs don't match pytorch all that well, they are at least
+    on the same order of magnitude
+    """
+    ooms = []
+    for el in float_lst:
+        str_el = str(el)
+        if "e" in str_el:
+            oom = int(str_el.split("e")[1])
+        elif str_el[:2] == "0.":
+            str_el = str_el.split(".")[1]
+
+            oom = -1
+            for e in str_el:
+                if e != "0":
+                    break
+                oom -= 1
+        else:
+            oom = len(str_el.split(".")[0])
+
+        ooms.append(oom)
+
+    return ooms
+
+
+def print_diff_argmax(a, b, annotation=""):
+    """
+    Prints out the value of both tensors at a point where the absolute difference is the largest.
+    """
+    absdiff = (a - b).abs()
+    argmax = absdiff.argmax().item()
+    diff = absdiff.reshape(-1)[argmax]
+    rela = a.abs() / (torch.max(a.abs(), b.abs()))
+    relb = b.abs() / (torch.max(a.abs(), b.abs()))
+    HT = a.shape[-2] // 32
+    WT = a.shape[-1] // 32
+    hwt = argmax // 1024
+    wt = hwt % WT
+    ht = hwt // WT
+    h = (argmax % 1024) // 32
+    w = (argmax % 1024) % 32
+    print(
+        "Abs diff=",
+        diff,
+        " at ",
+        argmax,
+        " --- ",
+        annotation,
+        "HTWT=",
+        ht,
+        wt,
+        "HW=",
+        h,
+        w,
+    )
+    print("  (a=", a.reshape(-1)[argmax].item(), ")")
+    print("  (b=", b.reshape(-1)[argmax].item(), ")")
+    print("  Rel a=", rela.reshape(-1)[argmax], " at ", argmax)
+    print("  Rel b=", relb.reshape(-1)[argmax], " at ", argmax)
+    return diff.item()
+
+
+def print_diff_tt_pyt(a, b, annotation=""):
+    # first convert a pytorch tensor argument b to tt
+    padded_b = pad_weight(b)
+    pyt_a = tt2torch(a)  # untilizes also
+    return print_diff_argmax(pyt_a, padded_b, annotation)
+
+
+def ttP(x, count=4, offset=0, stride=1):
+    if type(x) == torch.Tensor:
+        t1 = x.reshape(-1)
+    else:
+        tt_out = x.cpu()
+        torch_out = untilize(tt_out.to_torch())
+        t1 = torch_out.reshape(-1)
+    print("Tensor vals: (", end="")
+    for j in range(offset, offset + count * stride, stride):
+        print(t1[j].item(), " ", end="")
+    print(")")
+
+
+### Conv related helpers ###
+def read_conv_act_into_mm_act_block(
+    conv_act,
+    act_address_map_index,
+    address_map,
+    address_map_this_block_size,
+    act_block_h,
+    act_block_w,
+):
+    mm_act_block_shape = [1, 1, act_block_h * 32, act_block_w * 32]
+    mm_act_block_size = act_block_h * act_block_w * 1024
+    mm_act_block = torch.zeros(mm_act_block_size, dtype=torch.bfloat16).float()
+    for i in range(0, address_map_this_block_size, 4):
+        src_address = address_map[act_address_map_index]
+        dst_address = address_map[act_address_map_index + 1]
+        read_size = address_map[act_address_map_index + 2]
+        pad = address_map[act_address_map_index + 3]
+        for s in range(read_size):
+            assert dst_address + s < mm_act_block_size
+            if pad:
+                mm_act_block[dst_address + s] = 0
+            else:
+                assert src_address + s < len(conv_act)
+                mm_act_block[dst_address + s] = conv_act[src_address + s]
+        act_address_map_index += 4
+    return (mm_act_block.reshape(mm_act_block_shape), act_address_map_index)
+
+
+def read_conv_weight_into_mm_weight_block(
+    conv_weight,
+    weight_address_map_index,
+    weight_address_map,
+    weight_address_map_this_block_size,
+    weight_block_h,
+    weight_block_w,
+):
+    mm_weight_block_shape = [1, 1, weight_block_h * 32, weight_block_w * 32]
+    mm_weight_block_size = weight_block_h * weight_block_w * 1024
+    mm_weight_block = torch.zeros(mm_weight_block_size, dtype=torch.bfloat16).float()
+    for i in range(0, weight_address_map_this_block_size, 4):
+        src_address = weight_address_map[weight_address_map_index]
+        dst_address = weight_address_map[weight_address_map_index + 1]
+        read_size = weight_address_map[weight_address_map_index + 2]
+        pad = weight_address_map[weight_address_map_index + 3]
+        for s in range(read_size):
+            assert dst_address + s < mm_weight_block_size
+            if pad:
+                mm_weight_block[dst_address + s] = 0
+            else:
+                assert src_address + s < len(conv_weight)
+                mm_weight_block[dst_address + s] = conv_weight[src_address + s]
+        weight_address_map_index += 4
+    return (mm_weight_block.reshape(mm_weight_block_shape), weight_address_map_index)
+
+
+def blocked_mm_with_conv_act(
+    conv_act,
+    mm_weight,
+    act_address_map,
+    weight_address_map,
+    num_blocks_act_h,
+    num_blocks_act_w,
+    num_blocks_weight_w,
+    act_block_h,
+    act_block_w,
+    weight_block_w,
+):
+    # act refers to conv activation tensor
+    # weight refers to conv weight tensor
+    mm_output_shape = [
+        1,
+        1,
+        num_blocks_act_h * act_block_h * 32,
+        num_blocks_weight_w * weight_block_w * 32,
+    ]
+    ret = torch.zeros(mm_output_shape, dtype=torch.bfloat16).float()
+    mm_output_block_shape = [1, 1, act_block_h * 32, weight_block_w * 32]
+    act_address_map_index = 0
+    weight_address_map_index = 0
+    weight_block_h = act_block_w
+    num_groups = act_address_map[act_address_map_index]
+    assert num_groups == num_blocks_act_h * num_blocks_act_w * num_blocks_weight_w
+    weight_num_groups = act_address_map[weight_address_map_index]
+    assert weight_num_groups == num_groups
+    act_address_map_index += 1
+    weight_address_map_index += 1
+    for block_act_h in range(num_blocks_act_h):
+        # Reset weight (weight) to the starting tile in this column
+        for block_weight_w in range(num_blocks_weight_w):
+            output_block = torch.zeros(mm_output_block_shape, dtype=torch.bfloat16).float()
+            for block_act_w in range(num_blocks_act_w):
+                address_map_this_block_size = act_address_map[act_address_map_index]
+                act_address_map_index += 1
+                weight_address_map_this_block_size = weight_address_map[weight_address_map_index]
+                weight_address_map_index += 1
+                (mm_act_block, act_address_map_index) = read_conv_act_into_mm_act_block(
+                    conv_act,
+                    act_address_map_index,
+                    act_address_map,
+                    address_map_this_block_size,
+                    act_block_h,
+                    act_block_w,
+                )
+                (
+                    mm_weight_block,
+                    weight_address_map_index,
+                ) = read_conv_weight_into_mm_weight_block(
+                    mm_weight,
+                    weight_address_map_index,
+                    weight_address_map,
+                    weight_address_map_this_block_size,
+                    weight_block_h,
+                    weight_block_w,
+                )
+                # Untilize weight block (this CPU reference does matmul on untilized blocks)
+                mm_weight_block = untilize(mm_weight_block)
+                for out_h_block in range(act_block_h * 32):
+                    for out_w_block in range(weight_block_w * 32):
+                        output_block[0][0][out_h_block][out_w_block] += torch.dot(
+                            mm_act_block[0, 0, out_h_block, :].reshape(-1),
+                            mm_weight_block[0, 0, :, out_w_block].reshape(-1),
+                        )
+            start_oh = block_act_h * act_block_h * 32
+            start_ow = block_weight_w * weight_block_w * 32
+            end_oh = start_oh + (act_block_h * 32)
+            end_ow = start_ow + (weight_block_w * 32)
+            ret[0, 0, start_oh:end_oh, start_ow:end_ow] = output_block
+
+    return ret
+
+
+def is_conv_supported_on_device(conv_params):
+    K, C, R, S, U, V, P_H, P_W, dilation, groups = [conv_params[i] for i in range(10)]
+
+    if K % 32 != 0 or dilation != 1 or groups != 1:
+        logger.warning("DOES NOT HAVE SUPPORT FOR Conv with following parameters -")
+        logger.warning(
+            "K="
+            + str(K)
+            + " C="
+            + str(C)
+            + " R="
+            + str(R)
+            + " S="
+            + str(S)
+            + " U="
+            + str(U)
+            + " V="
+            + str(V)
+            + " PH="
+            + str(P_H)
+            + " PW="
+            + str(P_W)
+            + " dilation="
+            + str(dilation)
+            + " groups="
+            + str(groups)
+        )
+        return False
+
+    return True
+
+
+def is_x2_harvested(device):
+    grid = device.compute_with_storage_grid_size()
+    return device.arch() == Arch.WORMHOLE_B0 and (grid.x, grid.y) == (8, 7)
+
+
+def is_single_chip():
+    return ttnn.GetNumAvailableDevices() == 1
+
+
+def is_quasar():
+    ARCH_NAME = ttnn.get_arch_name()
+    return "quasar" in ARCH_NAME
+
+
+def is_blackhole():
+    ARCH_NAME = ttnn.get_arch_name()
+    return "blackhole" in ARCH_NAME
+
+
+def is_wormhole_b0():
+    ARCH_NAME = ttnn.get_arch_name()
+    return "wormhole_b0" in ARCH_NAME
+
+
+def is_watcher_enabled():
+    watcher = os.environ.get("TT_METAL_WATCHER")
+    lightweight_asserts = os.environ.get("TT_METAL_LIGHTWEIGHT_KERNEL_ASSERTS")
+    return (watcher is not None and watcher != "") or lightweight_asserts == "1"
+
+
+def is_llk_assert_enabled():
+    llk_assert = os.environ.get("TT_METAL_LLK_ASSERTS")
+    return llk_assert == "1"
+
+
+def is_n300():
+    return os.environ.get("MESH_DEVICE", "N150") == "N300"
+
+
+def is_slow_dispatch():
+    return os.environ.get("TT_METAL_SLOW_DISPATCH_MODE") == "1"
+
+
+def ti_skip(condition, reason="Invalid test parameters"):
+    return pytest.mark.skipif(condition, reason="Skipping unsupported case: " + reason)
+
+
+def skip_for_blackhole(reason_str="not a blackhole test"):
+    return ti_skip(is_blackhole(), reason=reason_str)
+
+
+def skip_for_wormhole_b0(reason_str="not a wormhole test"):
+    return ti_skip(is_wormhole_b0(), reason=reason_str)
+
+
+def skip_with_watcher(reason_str="Test is not passing with watcher enabled"):
+    return ti_skip(is_watcher_enabled(), reason=reason_str)
+
+
+def skip_with_llk_assert(reason_str="Test is not passing with LLK asserts enabled"):
+    return ti_skip(is_llk_assert_enabled(), reason=reason_str)
+
+
+def run_for_blackhole(reason_str="only runs for Blackhole"):
+    return ti_skip(not is_blackhole(), reason=reason_str)
+
+
+def run_for_wormhole_b0(reason_str="only runs for Wormhole B0"):
+    return ti_skip(not is_wormhole_b0(), reason=reason_str)
+
+
+def run_for_wormhole_b0_or_blackhole(reason_str="only runs for Wormhole B0 or Blackhole"):
+    return ti_skip(not (is_wormhole_b0() or is_blackhole()), reason=reason_str)
+
+
+def run_for_n_dev(n, reason_str="Test is not meant for this number of devices"):
+    return ti_skip(ttnn.get_num_devices() != n, reason=reason_str)
+
+
+def skip_for_n_dev(n, reason_str="Test is not meant for this number of devices"):
+    return ti_skip(ttnn.get_num_devices() == n, reason=reason_str)
+
+
+def skip_for_n_or_less_dev(n, reason_str="Test is not meant for this number of devices"):
+    return ti_skip(ttnn.get_num_devices() <= n, reason=reason_str)
+
+
+def skip_for_slow_dispatch(reason_str="not working for slow dispatch"):
+    return ti_skip(is_slow_dispatch(), reason=reason_str)
+
+
+def ttl_complex_2_torch_complex(tt_tensor):
+    torch_tensor = tt2torch_tensor(tt_tensor)
+
+    # extract real and imag parts of the complex tensor
+    real = torch_tensor[:, :, :, : torch_tensor.shape[-1] // 2].to(torch.bfloat16).to(torch.float)
+    imag = torch_tensor[:, :, :, torch_tensor.shape[-1] // 2 :].to(torch.bfloat16).to(torch.float)
+
+    # create torch complex tensor
+    result = torch.complex(real, imag)
+    return result
+
+
+def pad_and_fold_conv_filters_for_unity_stride(filter_pyt_nchw_tensor, stride_h, stride_w, align_c=4):
+    assert stride_h == stride_w
+    assert filter_pyt_nchw_tensor.shape[2] == filter_pyt_nchw_tensor.shape[3]
+    assert isinstance(align_c, int) and align_c > 0
+    # Fold activation for unity stride
+    # Pad channel size to align_c. This keeps L1 read addresses aligned; extra channels become
+    # zero-valued weights that contribute nothing to the convolution. align_c=4 is the WH/BH default
+    # (16B alignment for bf16 gives C a multiple of 4 with a tiled conv reader). Quasar's row-major
+    # fold needs align_c=8 (bf16 row-major shard width must be a multiple of 8) so the first conv
+    # folds to groups*8 input channels and consumes the aligned output without per-group padding strip.
+    C = _nearest_y(filter_pyt_nchw_tensor.shape[1], align_c)
+    # Pad filter to nearest stride
+    Padded_filter_height = _nearest_y(filter_pyt_nchw_tensor.shape[2], stride_h)
+    Padded_filter_width = _nearest_y(filter_pyt_nchw_tensor.shape[3], stride_w)
+    filter_pyt_padded = torch.nn.functional.pad(
+        filter_pyt_nchw_tensor,
+        (
+            0,
+            Padded_filter_width - filter_pyt_nchw_tensor.shape[3],
+            0,
+            Padded_filter_height - filter_pyt_nchw_tensor.shape[2],
+            0,
+            C - filter_pyt_nchw_tensor.shape[1],
+        ),
+    )
+    # Fold filter for unity stride.
+    filter_pyt_padded_folded = torch.zeros(
+        [
+            filter_pyt_padded.shape[0],
+            C * stride_h * stride_w,
+            (int)(filter_pyt_padded.shape[2] / stride_h),
+            (int)(filter_pyt_padded.shape[3] / stride_w),
+        ]
+    )
+    for h in range(0, filter_pyt_padded.shape[2], stride_h):
+        for w in range(0, filter_pyt_padded.shape[3], stride_w):
+            folded_h = (int)(h / stride_h)
+            folded_w = (int)(w / stride_w)
+            for i in range(4):
+                start_c = i * C
+                filter_pyt_padded_folded[:, start_c : start_c + C, folded_h, folded_w] = filter_pyt_padded[
+                    :, :, h + (int)(i / stride_w), w + (int)(i % stride_w)
+                ]
+    return filter_pyt_padded_folded
+
+
+# produces a tensor where each element in a page is the page number
+# this tensor is easy to debug and visualize
+def get_debug_tensor(num_pages_width, num_pages_height, dtype, page_width=32, page_height=32):
+    torch_tensor = None
+    for row_idx in range(0, int(num_pages_height)):
+        tile_row = None
+        for col_idx in range(0, int(num_pages_width)):
+            tile_idx = col_idx + num_pages_width * row_idx
+            tile = torch.full((1, 1, page_width, page_height), tile_idx + 1, dtype=dtype)
+            if tile_row == None:
+                tile_row = tile
+            else:
+                tile_row = torch.cat((tile_row, tile), 3)
+        if torch_tensor == None:
+            torch_tensor = tile_row
+        else:
+            torch_tensor = torch.cat((torch_tensor, tile_row), 2)
+
+    return torch_tensor
+
+
+# ── transformers 5.x Cache API compatibility ────────────────────────────────
+# transformers 5.x removed the legacy Cache API: DynamicCache no longer exposes
+# from_legacy_cache / to_legacy_cache / key_cache / value_cache (per-layer KV now
+# lives at cache.layers[i].keys/.values). These helpers work on both 4.x and 5.x.
+def hf_cache_layer_kv(cache, layer_idx):
+    """Return (key, value) tensors for a layer of a transformers Cache.
+
+    Handles the legacy tuple-of-tuples past_key_values, transformers <5 Cache
+    (key_cache/value_cache), and transformers >=5 Cache (layers[i].keys/.values).
+    """
+    if isinstance(cache, (tuple, list)):  # legacy tuple-of-tuples past_key_values
+        return cache[layer_idx][0], cache[layer_idx][1]
+    if hasattr(cache, "key_cache"):  # transformers < 5.x Cache
+        return cache.key_cache[layer_idx], cache.value_cache[layer_idx]
+    layer = cache.layers[layer_idx]  # transformers >= 5.x Cache
+    return layer.keys, layer.values
+
+
+def hf_cache_to_legacy(cache):
+    """Export a transformers Cache to the legacy tuple-of-(key, value) format."""
+    if hasattr(cache, "to_legacy_cache"):  # transformers < 5.x
+        return cache.to_legacy_cache()
+    return tuple((layer.keys, layer.values) for layer in cache.layers)  # transformers >= 5.x
+
+
+def hf_dynamic_cache_from_legacy(layer_kvs):
+    """Build a transformers DynamicCache from per-layer (key, value) tuples."""
+    from transformers import DynamicCache
+
+    layer_kvs = tuple(layer_kvs)
+    if hasattr(DynamicCache, "from_legacy_cache"):  # transformers < 5.x
+        return DynamicCache.from_legacy_cache(layer_kvs)
+    return DynamicCache(layer_kvs)  # transformers >= 5.x
+
+
+def hf_cache_num_layers(cache):
+    """Number of populated layers in a transformers Cache (version-tolerant)."""
+    return len(cache.key_cache) if hasattr(cache, "key_cache") else len(cache.layers)
+
+
+def hf_empty_encoder_decoder_cache():
+    """Create an empty transformers EncoderDecoderCache (version-tolerant)."""
+    from transformers import DynamicCache, EncoderDecoderCache
+
+    if hasattr(EncoderDecoderCache, "from_legacy_cache"):  # transformers < 5.x
+        return EncoderDecoderCache.from_legacy_cache(None)
+    return EncoderDecoderCache(DynamicCache(), DynamicCache())  # transformers >= 5.x
+
+
+def copy_to_buffer(src: "ttnn.Tensor", dst: "ttnn.Tensor", target_dtype) -> None:
+    """Convert ``src`` to ``dst``'s layout/dtype/shape/memcfg and write it into
+    ``dst``. ``dst``'s device buffer is preserved (no reallocation) so any
+    captured trace and the DRAM prefetcher's recorded buffer addresses remain
+    valid. The final ``ttnn.to_memory_config`` with ``output_tensor=dst`` both
+    reshards to ``dst``'s memory config and copies into ``dst``'s buffer.
+    """
+    converted = src
+
+    if converted.layout != dst.layout:
+        converted = ttnn.to_layout(converted, layout=dst.layout)
+
+    if converted.dtype != target_dtype:
+        converted = ttnn.typecast(converted, dtype=target_dtype)
+
+    if tuple(converted.shape) != tuple(dst.shape):
+        converted = ttnn.reshape(converted, list(dst.shape))
+
+    ttnn.to_memory_config(converted, dst.memory_config(), output_tensor=dst)

code/models/common/utils.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+# Backward compatibility: filter_none moved to models.common.sampling._utils
+from models.common.sampling._utils import filter_none  # noqa: F401
+
+# Backward compatibility: LogProbsCalculator moved to models.common.sampling.tt_log_probs
+from models.common.sampling.tt_log_probs import LogProbsCalculator  # noqa: F401
+
+
+def top_k_top_p_filtering(
+    logits: Tensor,
+    top_k: int = 0,
+    top_p: float = 1.0,
+    filter_value: float = -float("Inf"),
+    min_tokens_to_keep: int = 1,
+) -> Tensor:
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+    return logits

code/models/common/validation_tools.py

@@ -0,0 +1,719 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+TTNN Validation Framework
+
+A decorator-based validation system for comparing TTNN implementations against
+reference implementations (in PyTorch). Supports automatic input/output
+mapping, metric computation, and result collection.
+
+Key Features:
+- Automatic comparison of TTNN vs reference implementations
+- TTNN-native metric computation (stays on device until final scalar)
+- Flexible input/output mapping
+- Built-in metrics: max_abs_error, mean_abs_error, cosine_similarity
+- Performance tracking
+- Result registry for batch reporting
+"""
+
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+
+import ttnn
+
+from .auto_compose import to_torch_auto_compose
+from .distribute_as import from_torch_dist_as
+from .metrics import DEFAULT_METRICS
+
+# ============================================================================
+# Public API
+# ============================================================================
+
+# Module exports are defined at the package level in __init__.py
+
+
+def get_validation_registry() -> "ValidationRegistry":
+    """Get the global validation registry"""
+    return _validation_registry
+
+
+def enable_validation(enabled: bool = True):
+    """Enable or disable validation globally"""
+    _validation_registry.enabled = enabled
+
+
+def clear_validation_results():
+    """Clear all validation results"""
+    _validation_registry.results.clear()
+
+
+def compare_to_ttnn(
+    reference_fn: Callable,
+    *,
+    input_to_ttnn: Optional[Callable] = None,
+    output_to_ttnn: Optional[Callable] = None,
+    metric_tolerances: Optional[Dict[Any, Any]] = None,
+    enabled: bool = True,
+    raise_exceptions: bool = False,
+    return_reference_output: bool = False,
+):
+    """
+    Convenience wrapper for TTNN-on-device comparison. Provides useful visual cue to users that the reference function is a TTNN-native function.
+
+    Args:
+        reference_fn: Reference function to compare against
+        input_to_ttnn: Maps decorated function inputs to reference function inputs
+        output_to_ttnn: Maps decorated function outputs to reference function outputs
+        metric_tolerances: Dictionary specifying tolerances and optionally custom metrics.
+        enabled: Whether validation is enabled (can disable globally via registry)
+        raise_exceptions: When True, re-raise any exceptions encountered during
+            reference execution, output mapping, or metric computation instead
+            of logging them into validation results.
+
+    Examples:
+        @compare_to_ttnn(
+            reference_fn=lambda self, x: ttnn.matmul(x, self.weight),
+            input_to_ttnn=lambda self, x: (self, x),
+        )
+        def __call__(self, x):
+            return torch.matmul(x, self.torch_weight)
+            # alternatively, the decorated function can return a TTNN tensor: return ttnn.from_torch(x) @ self.weight
+
+    NOTES:
+        - The reference function is expected to accepts TTNN tensors and returns a TTNN tensor
+        - The decorated function inputs/outputs TTNN tensors, Torch tensors, or mixed TTNN and Torch tensors
+        - When decorated function returns torch tensors:
+          - the reference function's inputs will be constructed through either input_to_ttnn or from_torch(decorated function inputs, device=ttnn.GetDefaultDevice())
+          - the metric on output tensor will be computed on the host
+        - Experimental support for on-device metric computation is provided and used when both the decorated function and the reference function return TTNN tensors
+    """
+
+    # Default converters: recursively convert any TTNN tensors to torch, auto-compose shards.
+    # Non-tensor objects are passed through unchanged.
+
+    def _to_ttnn_auto(x: Any) -> Any:
+        if torch.is_tensor(x):
+            # Use auto-compose; relies on tensor.device() or a globally-set default device
+            assert (
+                ttnn.GetDefaultDevice() is not None
+            ), "Default device is not set. It is required by compare_to_ttnn. Please set it via ttnn.SetDefaultDevice(...)."
+            return ttnn.from_torch(x, device=ttnn.GetDefaultDevice())
+        return x
+
+    def _default_input_map(*args, **kwargs):
+        ref_args = _map_structure(args, _to_ttnn_auto)
+        ref_kwargs = _map_structure(kwargs, _to_ttnn_auto)
+        return ref_args, ref_kwargs
+
+    map_fn_to_match_sig = lambda tt_tensor, filler: to_torch_auto_compose(tt_tensor)
+
+    return __validate_against(
+        reference_fn=reference_fn,
+        input_map=input_to_ttnn or _default_input_map,
+        output_map=output_to_ttnn,
+        metric_tolerances=metric_tolerances,
+        enabled=enabled,
+        raise_exceptions=raise_exceptions,
+        reference_output_map_fn=map_fn_to_match_sig if return_reference_output else None,
+    )
+
+
+def compare_to_torch(
+    reference_fn: Callable,
+    *,
+    input_to_torch: Optional[Callable] = None,
+    output_to_torch: Optional[Callable] = None,
+    metric_tolerances: Optional[Dict[Any, Any]] = None,
+    enabled: bool = True,
+    raise_exceptions: bool = False,
+    return_reference_output: Optional[Callable[..., bool] | bool] = False,
+):
+    """
+    Convenience wrapper for host/CPU comparison using torch.
+
+    # Args:
+    #     reference_fn: Reference function to compare against
+    #     input_to_torch: Maps decorated function inputs to reference function inputs
+    #     output_to_torch: Maps decorated function outputs to reference function outputs
+    #     metric_tolerances: Dictionary specifying tolerances and optionally custom metrics.
+    #     enabled: Whether validation is enabled (can disable globally via registry)
+    #     raise_exceptions: When True, re-raise any exceptions encountered during
+    #         reference execution, output mapping, or metric computation instead
+    #         of logging them into validation results.
+    #
+    # Notes:
+    # - compare_to_torch is used when the reference function is a PyTorch function
+    # - the reference function takes as inputs to_torch_auto_compose(decorated function inputs) and compares the outputs with to_torch_auto_compose(decorated function outputs)
+    # - the decorated function inputs/outputs TTNN tensors, Torch tensors, or mixed TTNN and Torch tensors
+    """
+
+    # Default converters: recursively convert any TTNN tensors to torch, auto-compose shards.
+    # Non-tensor objects are passed through unchanged.
+
+    def _to_torch_auto(x: Any) -> Any:
+        if isinstance(x, ttnn.Tensor):
+            # Use auto-compose; relies on tensor.device() or a globally-set default device
+            return to_torch_auto_compose(x)
+        return x
+
+    def _default_input_map(*args, **kwargs):
+        ref_args = _map_structure(args, _to_torch_auto)
+        ref_kwargs = _map_structure(kwargs, _to_torch_auto)
+        return ref_args, ref_kwargs
+
+    def _default_output_map(output):
+        return _map_structure(output, _to_torch_auto)
+
+    return __validate_against(
+        reference_fn=reference_fn,
+        input_map=input_to_torch or _default_input_map,
+        output_map=output_to_torch or _default_output_map,
+        metric_tolerances=metric_tolerances,
+        enabled=enabled,
+        raise_exceptions=raise_exceptions,
+        reference_output_map_fn=from_torch_dist_as if return_reference_output else None,
+    )
+
+
+# ============================================================================
+# Data Structures
+# ============================================================================
+
+
+@dataclass
+class MetricResult:
+    """Per-metric validation outcome"""
+
+    value: float = float("inf")
+    passed: bool = False
+    error: str = ""
+
+
+@dataclass
+class ValidationResult:
+    """Results from a single validation run"""
+
+    function_name: str
+    passed: bool
+    # Map of metric name to its result (value/pass/fail/error)
+    metrics: Dict[Any, MetricResult] = field(default_factory=dict)
+    execution_time_impl: float = 0.0
+    execution_time_ref: float = 0.0
+    timestamp: float = field(default_factory=time.time)
+    logs: List[str] = field(default_factory=list)
+
+
+class ValidationRegistry:
+    """Global registry for validation results"""
+
+    def __init__(self):
+        self.results: List[ValidationResult] = []
+        self.enabled = True
+
+    def add_result(self, result: ValidationResult):
+        self.results.append(result)
+
+    def get_summary(self) -> Dict[str, Any]:
+        """Get summary statistics of all validations"""
+        if not self.results:
+            return {"total": 0, "passed": 0, "failed": 0}
+
+        passed = sum(1 for r in self.results if r.passed)
+        failed = len(self.results) - passed
+
+        return {
+            "total": len(self.results),
+            "passed": passed,
+            "failed": failed,
+            "pass_rate": passed / len(self.results) if self.results else 0.0,
+            "avg_speedup": (
+                sum(r.execution_time_ref / r.execution_time_impl for r in self.results if r.execution_time_impl > 0)
+                / len(self.results)
+                if self.results
+                else 0.0
+            ),
+        }
+
+    def print_report(self, verbose: bool = False):
+        """Print detailed validation report"""
+        summary = self.get_summary()
+        print("\n" + "=" * 80)
+        print("VALIDATION REPORT")
+        print("=" * 80)
+        print()
+
+        for result in self.results:
+            status = "✓ PASS" if result.passed else "✗ FAIL"
+            print(f"{status} - {result.function_name}")
+            print(
+                f"  Execution time: impl={result.execution_time_impl*1000:.2f}ms, ref={result.execution_time_ref*1000:.2f}ms"
+            )
+
+            if result.metrics:
+                print(f"  Metrics:")
+                for metric_name, mres in result.metrics.items():
+                    # Use enum value for readability if metric is an Enum
+                    name_str = metric_name.value if hasattr(metric_name, "value") else str(metric_name)
+                    if mres.value is not None:
+                        try:
+                            val_str = f"{mres.value:.6f}"
+                        except Exception:
+                            val_str = str(mres.value)
+                    else:
+                        val_str = "-"
+                    status = "PASS" if mres.passed else "FAIL"
+                    print(f"    {name_str}: {val_str} — {status}")
+                    if mres.error:
+                        print(f"      error: {mres.error}")
+
+            # Print any collected logs for this validation
+            if result.logs and verbose:
+                print("  Logs:")
+                for entry in result.logs:
+                    try:
+                        msg = str(entry)
+                    except Exception:
+                        msg = "<unprintable log entry>"
+                    print(f"    {msg}")
+
+            # All errors are reported via per-metric entries
+            print()
+
+        print("-" * 36 + "Summary:" + "-" * 36)
+        print(f"Total validations: {summary['total']}")
+        print(f"Passed: {summary['passed']} ({summary['pass_rate']*100:.1f}%)")
+        print(f"Failed: {summary['failed']}")
+        print(f"Average speedup: {summary['avg_speedup']:.2f}x")
+        print()
+        print("=" * 80 + "\n")
+
+
+# Global validation registry
+_validation_registry = ValidationRegistry()
+
+
+# ============================================================================
+# Validation Decorator
+# ============================================================================
+
+
+class Metric(str, Enum):
+    """Enumeration of supported metric names, values match current string keys."""
+
+    MAX_ABS_ERROR = "max_abs_error"
+    MEAN_ABS_ERROR = "mean_abs_error"
+    PCC = "pcc"
+
+
+@dataclass
+class MetricSpec:
+    """Metric specification: name, tolerance, direction, and compute function."""
+
+    tolerance: float
+    higher_is_better: bool
+    compute_fn: Callable[[Any, Any], float]
+    name: str = field(default="")
+
+
+# Registry of built-in metrics with defaults. Tolerances here are sensible
+# defaults; callers can override per-validation via `tolerances`.
+METRIC_SPECS: Dict[Metric, MetricSpec] = {
+    Metric.MAX_ABS_ERROR: MetricSpec(
+        name=Metric.MAX_ABS_ERROR.value,
+        tolerance=0.0,
+        higher_is_better=False,
+        compute_fn=DEFAULT_METRICS[Metric.MAX_ABS_ERROR.value],
+    ),
+    Metric.MEAN_ABS_ERROR: MetricSpec(
+        name=Metric.MEAN_ABS_ERROR.value,
+        tolerance=0.0,
+        higher_is_better=False,
+        compute_fn=DEFAULT_METRICS[Metric.MEAN_ABS_ERROR.value],
+    ),
+    Metric.PCC: MetricSpec(
+        name=Metric.PCC.value,
+        tolerance=0.0,
+        higher_is_better=True,
+        compute_fn=DEFAULT_METRICS[Metric.PCC.value],
+    ),
+}
+
+# Convenience groupings for quick checks
+HIGHER_IS_BETTER_METRICS = {m.value for m, spec in METRIC_SPECS.items() if spec.higher_is_better}
+LOWER_IS_BETTER_METRICS = {m.value for m, spec in METRIC_SPECS.items() if not spec.higher_is_better}
+
+
+# Helper: prefer Metric enum as dict key when possible
+def _metric_key(key: Any) -> Any:
+    try:
+        return Metric(key)
+    except Exception:
+        return key
+
+
+# Helper: Build active metrics map (name -> compute fn). Accept Metric enum keys for tolerances.
+def _normalize_key(k: Any) -> str:
+    try:
+        # Enum or similar objects with .value as canonical string
+        return k.value if hasattr(k, "value") else str(k)
+    except Exception:
+        return str(k)
+
+
+# Helper: Prepare metrics, tolerances, and directionality
+def _prepare_metric_config(metric_tolerances_input):
+    metrics_map = {name: fn for name, fn in DEFAULT_METRICS.items()}
+    hib = set(HIGHER_IS_BETTER_METRICS)
+    logs_local: List[str] = []
+    tol_map: Dict[str, float] = {}
+
+    if not isinstance(metric_tolerances_input, dict):
+        logs_local.append(f"metric_tolerances_input must be a dict, got {type(metric_tolerances_input)}")
+        metric_tolerances_input = dict()
+
+    if not metric_tolerances_input:
+        logs_local.append("no metric tolerances provided")
+        metric_tolerances_input = dict()
+
+    for raw_key, spec in metric_tolerances_input.items():
+        name = _normalize_key(raw_key)
+        if isinstance(spec, MetricSpec):
+            tol_map[name] = float(spec.tolerance)
+            metrics_map[name] = spec.compute_fn
+            spec.name = name if spec.name == "" else spec.name
+            if spec.higher_is_better:
+                hib.add(name)
+            else:
+                hib.discard(name)
+            continue
+        try:
+            tol_map[name] = float(spec)
+        except Exception:
+            logs_local.append(f"unrecognized tolerance: {raw_key}: {spec}")
+
+    return metrics_map, hib, tol_map, logs_local
+
+
+# todo)) also allow raise an exception from the a failed metric!
+
+# todo)) add support for multiple outputs from the reference function and the decorated function!
+# e.g., return logits, past_key_values, etc.
+
+# todo)) make sure the dtypes are taken care of in the validate_against decorator!
+# e.g., if the decorated function is of dtype bfp4, what is the dtype of the to_torch_auto_compose output?
+
+# todo)) add file line number to the validation results!
+
+# todo)) add function to export the validation results to a csv file!
+
+# todo)) enhance report to use file line number as index to summarize the validation results
+# e.g., ✗ FAIL - __main__.Attention.__call__ (line 100) -> 100 failed validations
+
+# todo)) remove compile time from speed up calculation -- e.g., 9118.15ms should be removed in the example below:
+# ================================================================================
+# VALIDATION REPORT
+# ================================================================================
+# Total validations: 1400
+# Passed: 1400 (100.0%)
+# Failed: 0
+# Average speedup: 0.97x
+
+# ✓ PASS - __main__.TransformerBlock.__call__
+#   Execution time: impl=9118.15ms, ref=14.84ms
+#   Metrics:
+#     pcc: 0.999743 — PASS
+
+# ✓ PASS - __main__.TransformerBlock.__call__
+#   Execution time: impl=3.05ms, ref=12.73ms
+#   Metrics:
+#     pcc: 0.999913 — PASS
+
+# ✓ PASS - __main__.TransformerBlock.__call__
+#   Execution time: impl=3.31ms, ref=12.48ms
+#   Metrics:
+#     pcc: 0.999962 — PASS
+
+# ✓ PASS - __main__.TransformerBlock.__call__
+#   Execution time: impl=3.11ms, ref=12.89ms
+#   Metrics:
+#     pcc: 1.000000 — PASS
+
+# ✓ PASS - __main__.TransformerBlock.__call__
+#   Execution time: impl=3.16ms, ref=12.97ms
+#   Metrics:
+#     pcc: 0.999998 — PASS
+
+
+# todo)) stretch goals:
+# - generate unit test automatically from the failed validations
+def __validate_against(
+    reference_fn: Callable,
+    *,
+    input_map: Optional[Callable] = None,
+    output_map: Optional[Callable] = None,
+    metric_tolerances: Optional[Dict[Any, Any]] = None,
+    enabled: bool = True,
+    raise_exceptions: bool = False,
+    reference_output_map_fn: Optional[Callable] = None,
+):
+    """
+    Decorator to validate a function against a reference implementation.
+
+    Args:
+        reference_fn: Reference function to compare against
+        input_map: Maps decorated function inputs to reference function inputs
+                   Signature: (args, kwargs) -> (ref_args, ref_kwargs)
+                   If None, inputs are passed as-is
+        output_map: Converts impl output to match ref output's type
+                    Signature: (output) -> comparable_output
+                    Applied ONLY to impl_output to convert it to ref_output's type
+                    Common use: lambda x: ttnn.to_torch(x).squeeze() to convert ttnn → torch
+                    If None, outputs are used as-is (both must already be same type)
+        metric_tolerances: Dictionary specifying tolerances and optionally custom metrics.
+            Accepts the following per metric key (str or Metric):
+              - float: tolerance only (uses built-in compute + direction)
+              - MetricSpec instance
+            Validation fails if any metric exceeds its tolerance
+        enabled: Whether validation is enabled (can disable globally via registry)
+        raise_exceptions: When True, re-raise any exceptions encountered during
+            reference execution, output mapping, or metric computation instead
+            of logging them into validation results.
+
+    Examples:
+        # Pattern 1: TTNN-native metrics (recommended, 100-1000× faster!)
+        # Both impl and ref return ttnn.Tensor, no output_map needed
+        def _reference_impl(self, x):
+            x_torch = ttnn.to_torch(x).squeeze(0)
+            result_torch = torch.matmul(x_torch, self.weight_torch)
+            # Convert back to TTNN for on-device metrics!
+            return ttnn.from_torch(result_torch.unsqueeze(0), device=self.device, ...)
+
+        @validate_against(
+            reference_fn=lambda self, x: self._reference_impl(x),
+            tolerances={'max_abs_error': 1e-3}
+        )
+        def __call__(self, x):
+            return ttnn.matmul(x, self.weight)
+
+        # Pattern 2: PyTorch metrics (when reference returns torch.Tensor)
+        # Use output_map to convert impl output (ttnn.Tensor) to match ref (torch.Tensor)
+        @validate_against(
+            reference_fn=torch.nn.functional.rms_norm,
+            input_map=lambda args, kwargs: (
+                (ttnn.to_torch(args[1]).squeeze(),),
+                {'eps': args[0].eps}
+            ),
+            output_map=lambda x: ttnn.to_torch(x).squeeze(),  # Convert impl: ttnn → torch
+            tolerances={'max_abs_error': 1e-3}
+        )
+        def __call__(self, x):
+            return ttnn.rms_norm(x, self.weight, self.eps)  # Returns ttnn.Tensor
+    """
+
+    if metric_tolerances is None:
+        metric_tolerances = {
+            Metric.MAX_ABS_ERROR: 1e-2,
+            Metric.PCC: 0.99,
+        }
+
+    metrics_to_use, higher_is_better_effective, tolerances_map, pre_logs = _prepare_metric_config(metric_tolerances)
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Check if validation is enabled
+            if not enabled or not _validation_registry.enabled:
+                return func(*args, **kwargs)
+
+            # Execute implementation
+            start_time = time.perf_counter()
+            impl_output = func(*args, **kwargs)
+            impl_time = time.perf_counter() - start_time
+            logs: List[str] = pre_logs.copy()
+
+            # Map inputs for reference function: prefer input_map, else pass-through
+            if input_map:
+                _nm = getattr(input_map, "__name__", None) or type(input_map).__name__
+                logs.append(f"input_map={_nm}")
+                try:
+                    mapped = input_map(*args, **kwargs)
+                except Exception as e:
+                    # If input mapping fails, log error, record result, and return impl output
+                    logs.append(f"input_mapping_error={str(e)}")
+                    result = ValidationResult(
+                        function_name=f"{func.__module__}.{func.__qualname__}",
+                        passed=False,
+                        metrics={
+                            "input_mapping": MetricResult(
+                                value=None, passed=False, error=f"Input mapping failed: {str(e)}"
+                            )
+                        },
+                        execution_time_impl=impl_time,
+                        execution_time_ref=0.0,
+                        logs=logs,
+                    )
+                    _validation_registry.add_result(result)
+                    # Re-raise exception if raise_exceptions is True
+                    if raise_exceptions:
+                        raise
+                    return impl_output
+                # Normalize mapper output:
+                # - If (ref_args, ref_kwargs) with kwargs as dict, use directly
+                # - Otherwise, treat return as positional args and use empty kwargs
+                if isinstance(mapped, tuple) and len(mapped) == 2 and isinstance(mapped[1], dict):
+                    ref_args, ref_kwargs = mapped
+                else:
+                    ref_args = mapped if isinstance(mapped, (list, tuple)) else (mapped,)
+                    ref_kwargs = {}
+            else:
+                logs.append("input_map=pass-through")
+                ref_args, ref_kwargs = args, kwargs
+
+            # Execute reference
+            try:
+                start_time = time.perf_counter()
+                ref_output = reference_fn(*ref_args, **ref_kwargs)
+                ref_time = time.perf_counter() - start_time
+            except Exception as e:
+                # If reference fails, just return impl output and log error via metrics
+                logs.append(f"reference_execution_error={str(e)}")
+                # Record elapsed time until failure
+                ref_time = time.perf_counter() - start_time
+                result = ValidationResult(
+                    function_name=f"{func.__module__}.{func.__qualname__}",
+                    passed=False,
+                    metrics={
+                        "reference_execution": MetricResult(
+                            value=None, passed=False, error=f"Reference execution failed: {str(e)}"
+                        )
+                    },
+                    execution_time_impl=impl_time,
+                    execution_time_ref=ref_time,
+                    logs=logs,
+                )
+                _validation_registry.add_result(result)
+                # Re-raise exception if raise_exceptions is True
+                if raise_exceptions:
+                    raise
+                return impl_output
+
+            # Map outputs for comparison
+            # Note: output_map only applies to impl_output to convert it to match ref_output's type
+            try:
+                _nm = getattr(output_map, "__name__", None) or type(output_map).__name__
+                logs.append(f"output_map={_nm}")
+                impl_comparable = output_map(impl_output) if output_map else impl_output
+                ref_comparable = ref_output  # Reference output is always used as-is
+            except Exception as e:
+                logs.append(f"output_mapping_error={str(e)}")
+                result = ValidationResult(
+                    function_name=f"{func.__module__}.{func.__qualname__}",
+                    passed=False,
+                    metrics={
+                        "output_mapping": MetricResult(
+                            value=None, passed=False, error=f"Output mapping failed: {str(e)}"
+                        )
+                    },
+                    execution_time_impl=impl_time,
+                    execution_time_ref=ref_time,
+                    logs=logs,
+                )
+                _validation_registry.add_result(result)
+                # Re-raise exception if raise_exceptions is True
+                if raise_exceptions:
+                    raise
+                return impl_output
+
+            # Compute metrics
+            computed_metrics: Dict[Any, MetricResult] = {}
+            passed = True
+
+            for metric_name, threshold in tolerances_map.items():
+                try:
+                    metric_fn = metrics_to_use.get(metric_name)
+
+                    # Store results keyed by enum when available
+                    metric_key = _metric_key(metric_name)
+                    # If metric function isn't known, record an error
+                    if metric_fn is None:
+                        computed_metrics[metric_key] = MetricResult(
+                            value=None, passed=False, error=f"Unknown metric: {metric_name}"
+                        )
+                        passed = False
+                        continue
+
+                    value = metric_fn(impl_comparable, ref_comparable)
+
+                    # Determine direction using registry when available
+                    if metric_name in higher_is_better_effective:
+                        ok = value >= threshold
+                        err = None
+                        if not ok:
+                            passed = False
+                            err = f"{metric_name}={value:.6e} below threshold {threshold:.6e}"
+                        computed_metrics[metric_key] = MetricResult(value=value, passed=ok, error=err)
+                    else:
+                        ok = value <= threshold
+                        err = None
+                        if not ok:
+                            passed = False
+                            err = f"{metric_name}={value:.6e} exceeds tolerance {threshold:.6e}"
+                        computed_metrics[metric_key] = MetricResult(value=value, passed=ok, error=err)
+                except Exception as e:
+                    msg = f"Metric {metric_name} failed: {str(e)}"
+                    computed_metrics[metric_key] = MetricResult(value=None, passed=False, error=msg)
+                    passed = False
+                    if raise_exceptions:
+                        raise
+
+            # Optionally return the (aligned) reference output instead of impl output
+            backup_impl_output = impl_output
+            try:
+                if reference_output_map_fn:
+                    impl_output = reference_output_map_fn(ref_output, impl_output)
+                    logs.append(f"reference_output_mapping_fn={reference_output_map_fn.__name__}")
+            except Exception as e:
+                # If alignment fails, fall back to impl output
+                impl_output = backup_impl_output
+                # Re-raise exception if raise_exceptions is True after logging the error
+                logs.append(f"reference_output_mapping_error={str(e)}")
+                if raise_exceptions:
+                    raise
+
+            # Record results
+            pass_count = sum(1 for v in computed_metrics.values() if v.passed)
+            fail_count = sum(1 for v in computed_metrics.values() if not v.passed)
+            logs.append(f"metrics={pass_count}_pass,{fail_count}_fail")
+            result = ValidationResult(
+                function_name=f"{func.__module__}.{func.__qualname__}",
+                passed=passed,
+                metrics=computed_metrics,
+                execution_time_impl=impl_time,
+                execution_time_ref=ref_time,
+                logs=logs,
+            )
+            _validation_registry.add_result(result)
+
+            return impl_output
+
+        return wrapper
+
+    return decorator
+
+
+def _map_structure(obj: Any, fn: Callable[[Any], Any]) -> Any:
+    """
+    Map a structure of objects to a new structure using a function.
+    """
+    if isinstance(obj, (list, tuple)):
+        mapped = [_map_structure(x, fn) for x in obj]
+        return type(obj)(mapped)
+    if isinstance(obj, dict):
+        return {k: _map_structure(v, fn) for k, v in obj.items()}
+    return fn(obj)

code/models/common/weight_cache.py

@@ -0,0 +1,426 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Shared, ModelArgs-agnostic warm ttnn weight-cache helpers (generalizes PR #50550 / #48531
+to forked model loaders — issue #45400 follow-up).
+
+On a warm cache, a model can be built from its on-disk ``.tensorbin`` files without the
+expensive host-side HF ``from_pretrained`` load (the load that OOMs/hangs during prefill,
+#48509): ``ttnn.as_tensor(torch_weight, cache_file_name=...)`` loads the cached tensor and
+ignores ``torch_weight`` on a cache hit (see ttnn/operations/core.py). So most weights only
+need a dataless placeholder (``torch.empty`` of the right shape/dtype) to satisfy the modules'
+host-side reshape ops before ``as_tensor``.
+
+Some forks (e.g. gemma4) additionally consume a *small* set of weights on the host — token
+embeddings used via ``F.embedding``, per-layer scalars read via ``.item()``, etc. Those must be
+real. ``mark_weight_cache_complete`` persists exactly those tensors to a sidecar at cold-build
+time (write-access run), and ``build_cached_state_dict`` serves them real on later warm runs
+while placeholdering the rest — a HYBRID state_dict. The host subset is a tiny fraction of the
+weight bytes, so the full from_pretrained (and its OOM) is still avoided.
+"""
+
+import collections.abc
+import hashlib
+import json
+import os
+from pathlib import Path
+
+import torch
+from loguru import logger
+
+WEIGHT_CACHE_MARKER = ".weights_complete"
+HOST_WEIGHTS_SIDECAR = ".host_weights.pt"
+# Bump when the set/naming/layout of cached weights, or this marker schema, changes such that an
+# existing cache would not satisfy a new build. A marker written by an older format is rejected,
+# so the run cold-loads and regenerates rather than building from an incompatible cache.
+#   v2: model/n_layers/mesh_shape validation + a {key: [shape, dtype]} manifest.
+#   v3: canonical mesh_shape encoding shared with ModelArgs (the two writers previously encoded it
+#       differently and each rejected the other's marker), a `components` field so a text-only seed
+#       cannot certify a cache for a build that also needs the vision tower, and `cache_files` --
+#       the recursive list of .tensorbin files the completed build actually produced, verified
+#       per-file on read. That last one is load-bearing: ttnn.as_tensor PERSISTS whatever tensor it
+#       is handed on a cache miss, so a marker that outlives some of its tensorbins would otherwise
+#       dump placeholders to disk as real cache entries -- silent, permanent corruption. Verifying
+#       the recorded file set turns every such case back into a plain cold load. Also `build_variant`
+#       -- the build options (prefetcher, precision) that change an as_tensor cache FILENAME, matched
+#       exactly, because a different variant needs different files rather than fewer.
+WEIGHT_CACHE_FORMAT_VERSION = 3
+
+DEFAULT_FORCE_ENV = "TT_TRANSFORMERS_FORCE_MODEL_LOAD"
+
+
+def _variant_digest(build_variant):
+    """Stable short digest of a build_variant dict ("none" for None)."""
+    if build_variant is None:
+        return "none"
+    return hashlib.sha1(json.dumps(build_variant, sort_keys=True, default=str).encode()).hexdigest()[:12]
+
+
+def _variant_unverifiable(build_variant):
+    return bool(build_variant) and bool(build_variant.get("unverifiable"))
+
+
+def marker_path(cache_path, build_variant=None):
+    """The marker file for one (cache dir, build variant).
+
+    The variant digest is part of the FILENAME, not just a field compared inside one shared
+    marker. A cache dir legitimately serves several build variants (the Llama CI job runs
+    eval-32 with and without the DRAM prefetcher against the same instruct cache), and a single
+    marker matched exactly would make each variant's seed evict the other's on every run -- both
+    then cold-load forever with nothing going red. One marker per variant lets them coexist.
+    (#45400 review, finding B3)"""
+    return Path(cache_path) / f"{WEIGHT_CACHE_MARKER}.{_variant_digest(build_variant)}"
+
+
+def _dtype_from_str(s):
+    return getattr(torch, s.rsplit(".", 1)[-1])
+
+
+def normalize_mesh_shape(mesh_shape):
+    """Canonical marker encoding for a mesh shape.
+
+    ``ttnn.MeshShape`` stringifies as ``MeshShape([1, 8])`` while callers that pass a plain tuple
+    stringify as ``(1, 8)``. Both writers must agree or each rejects the other's marker and the
+    model cold-loads forever (gemma3 inherits ModelArgs but its demos call this module). Normalize
+    everything to a plain tuple-of-ints string."""
+    try:
+        return str(tuple(int(d) for d in mesh_shape))
+    except TypeError:
+        return str(mesh_shape)
+
+
+def _normalize_components(components):
+    """Canonical component list. ``None`` means "the whole model as this loader builds it" and is
+    encoded as a single implicit component so old-style callers stay self-consistent."""
+    if components is None:
+        return ["all"]
+    if isinstance(components, str):
+        return [components]
+    return sorted(str(c) for c in components)
+
+
+def list_cache_files(cache_path):
+    """Every ``.tensorbin`` under ``cache_path``, recursively, as sorted relative POSIX paths.
+
+    Recursive because forked loaders nest per-layer weights in subdirectories (qwen36
+    ``layers.{n}/``, gemma4 ``layer_{i}/``); a top-level ``glob`` would call a cache complete when
+    only the root-level ``output.weight`` survived an interrupted seed."""
+    cache_path = Path(cache_path)
+    return sorted(p.relative_to(cache_path).as_posix() for p in cache_path.rglob("*.tensorbin"))
+
+
+# One-entry cache so the completeness gate's validation load is reused by
+# build_cached_state_dict instead of torch.load-ing the same multi-GB file twice per warm run
+# (gemma-4-31b's embedding alone is ~2.8 GB). Keyed on (path, mtime, size) so a republished
+# sidecar is never served stale; the builder consumes the entry so the tensors are not pinned
+# past the build. (#45400 review, finding R1)
+_SIDECAR_CACHE = {}
+
+
+def load_host_sidecar(cache_path, *, consume=False):
+    """Load the host-weights sidecar, or None if absent/unreadable.
+
+    ``consume=True`` drops the memoized entry after returning it (the caller takes ownership)."""
+    sidecar = Path(cache_path) / HOST_WEIGHTS_SIDECAR
+    if not sidecar.is_file():
+        return None
+    try:
+        st = sidecar.stat()
+        key = (str(sidecar), st.st_mtime_ns, st.st_size)
+        host = _SIDECAR_CACHE.get(key)
+        if host is None:
+            host = torch.load(sidecar, map_location="cpu", weights_only=True)
+            _SIDECAR_CACHE.clear()
+            _SIDECAR_CACHE[key] = host
+        if consume:
+            _SIDECAR_CACHE.pop(key, None)
+        return host
+    except Exception:
+        return None
+
+
+def weight_cache_is_complete(
+    cache_path,
+    *,
+    model_name,
+    n_layers,
+    mesh_shape,
+    components=None,
+    build_variant=None,
+    force_env=DEFAULT_FORCE_ENV,
+):
+    """True when the on-disk ttnn weight cache at ``cache_path`` was fully built by a previous run
+    for this exact build, and every tensorbin that build produced is still present.
+
+    ``components`` names the model parts this build will construct (e.g. ``"text"`` vs
+    ``"text+vision"``); a marker written by a narrower build does not satisfy a wider one, because
+    the wider build needs tensorbins the narrower one never wrote. ``force_env=...=1`` forces a
+    cold load."""
+    if force_env and os.getenv(force_env) == "1":
+        return False
+    # A variant we could not compute is a variant we cannot verify: accepting it could hand a
+    # placeholder to a build whose cache-filename set we did not check, and as_tensor would
+    # persist that placeholder to disk. Fail closed, loudly. (#45400 review, finding R3)
+    if _variant_unverifiable(build_variant):
+        logger.warning(
+            f"Warm-cache check for {cache_path}: build_variant could not be computed "
+            f"({build_variant.get('error', 'unknown error')}); forcing a cold load."
+        )
+        return False
+    cache_path = Path(cache_path)
+    marker = marker_path(cache_path, build_variant)
+    if not marker.is_file():
+        return False
+    try:
+        meta = json.loads(marker.read_text())
+    except (ValueError, OSError):
+        return False
+    if meta.get("format_version") != WEIGHT_CACHE_FORMAT_VERSION:
+        return False
+    if meta.get("model_name") != model_name or meta.get("n_layers") != n_layers:
+        return False
+    if meta.get("mesh_shape") != normalize_mesh_shape(mesh_shape):
+        return False
+    # The recorded build must cover every component this build needs. Superset is fine (a
+    # text+vision seed wrote the text tensorbins too, so it satisfies a text-only build); a subset
+    # is not (a text-only seed never wrote the vision tower's tensorbins, and accepting it would
+    # make as_tensor dump placeholders for them).
+    if not set(_normalize_components(components)).issubset(set(meta.get("components") or [])):
+        return False
+    # Build options that change an as_tensor cache FILENAME (prefetcher, precision) must match
+    # exactly. A superset rule is wrong here: a different variant does not need fewer files, it
+    # needs DIFFERENT ones, and any it is missing would be regenerated from the placeholder.
+    if meta.get("build_variant") != build_variant:
+        return False
+    if not meta.get("weights"):
+        return False
+    # Every tensorbin the completed build produced must still be on disk. Any missing file would
+    # otherwise be regenerated by as_tensor FROM THE PLACEHOLDER we are about to hand it, writing
+    # garbage into the cache permanently. Missing file => cold load, which rebuilds it correctly.
+    recorded = meta.get("cache_files")
+    if not recorded:
+        return False
+    present = set(list_cache_files(cache_path))
+    if not all(f in present for f in recorded):
+        return False
+    # If host weights were captured, the sidecar must be present AND loadable. A torn/corrupt
+    # sidecar (interrupted or racing seed) must fall back to a cold load -- the way a torn marker
+    # already does via the except above -- rather than pass this gate and then crash torch.load on
+    # every subsequent run, bricking the cache dir. Checked LAST so the load it performs is
+    # memoized only when the gate is about to pass, for build_cached_state_dict to consume.
+    # (#45400 review)
+    if meta.get("host_weights") and load_host_sidecar(cache_path) is None:
+        return False
+    return True
+
+
+def mark_weight_cache_complete(
+    cache_path,
+    state_dict,
+    *,
+    model_name,
+    n_layers,
+    mesh_shape,
+    components=None,
+    build_variant=None,
+    is_moe=False,
+    is_host_weight=None,
+):
+    """Record that the ttnn weight cache at ``cache_path`` is fully built.
+
+    Writes a ``.weights_complete`` marker holding a ``{key: [shape, dtype]}`` manifest of every
+    weight plus the recursive list of ``.tensorbin`` files this build produced (verified per-file
+    on read). If ``is_host_weight(key)`` is provided, the (real) tensors it matches are also saved
+    to a ``.host_weights.pt`` sidecar so a later warm run can serve them for real (hybrid).
+
+    Call this only AFTER the model has been constructed, so the tensorbins exist to be recorded."""
+    if _variant_unverifiable(build_variant):
+        # Never certify a cache under an identity we could not compute -- a later run computing
+        # the same error string would otherwise warm-match it. (#45400 review, finding R3)
+        logger.warning(
+            f"Not marking weight cache complete at {cache_path}: build_variant could not be "
+            f"computed ({build_variant.get('error', 'unknown error')})."
+        )
+        return
+    cache_path = Path(cache_path)
+    marker = marker_path(cache_path, build_variant)
+    weights = {}
+    host = {}
+    for k, v in state_dict.items():
+        shape = getattr(v, "shape", None)
+        dt = getattr(v, "dtype", None)
+        if shape is None or dt is None:
+            continue  # skip non-tensor entries
+        weights[k] = [list(shape), str(dt)]
+        if is_host_weight is not None and is_host_weight(k):
+            host[k] = v
+    try:
+        cache_path.mkdir(parents=True, exist_ok=True)
+        cache_files = list_cache_files(cache_path)
+        if not cache_files:
+            logger.warning(f"Not marking weight cache complete: no .tensorbin files under {cache_path}")
+            return
+        # Write both the sidecar and the marker atomically (temp file + os.replace, atomic on
+        # POSIX). Two jobs can seed the same (model, dtype, mesh) dir on one host concurrently, and
+        # an interrupted write must never leave a torn file that a later run picks up: a half-written
+        # sidecar would otherwise pass the is_file() gate and crash torch.load on every subsequent
+        # run. The temp name is pid-unique so two concurrent seeders cannot write the SAME temp
+        # inode -- with a fixed name, B could publish the file while A was still writing into it.
+        # Sidecar first, then marker, so the completeness gate only appears once its sidecar is
+        # fully in place. (#45400 review)
+        uniq = os.getpid()
+        if host:
+            sidecar = cache_path / HOST_WEIGHTS_SIDECAR
+            sidecar_tmp = sidecar.with_suffix(sidecar.suffix + f".tmp.{uniq}")
+            torch.save(host, sidecar_tmp)
+            os.replace(sidecar_tmp, sidecar)
+        marker_body = json.dumps(
+            {
+                "format_version": WEIGHT_CACHE_FORMAT_VERSION,
+                "model_name": model_name,
+                "n_layers": n_layers,
+                "mesh_shape": normalize_mesh_shape(mesh_shape),
+                "components": _normalize_components(components),
+                "build_variant": build_variant,
+                "cache_files": cache_files,
+                "is_moe": bool(is_moe),
+                "host_weights": sorted(host.keys()),
+                "weights": weights,
+            }
+        )
+        marker_tmp = marker.with_suffix(marker.suffix + f".tmp.{uniq}")
+        marker_tmp.write_text(marker_body)
+        os.replace(marker_tmp, marker)
+        logger.info(f"Marked ttnn weight cache complete: {marker} ({len(weights)} weights, {len(host)} host-loaded)")
+    except Exception as e:
+        # Deliberately broad: this function only RECORDS completion -- failing to record must
+        # never kill a build that already succeeded. The concrete case: on a read-only
+        # /mnt/MLPerf, torch.save of the host sidecar raises RuntimeError from torch's C++
+        # serializer (inline_container.cc "Read-only file system"), not OSError, and the narrow
+        # except crashed every read-only cold run of the sidecar models (gemma4/gemma3) right
+        # after a successful build. (#45400 review, finding R5; seen on Gemma-4-E2B bh_p150,
+        # run 32511945147)
+        logger.warning(f"Could not write weight-cache completion marker {marker}: {e}")
+
+
+class CachedStateDict(collections.abc.MutableMapping):
+    """A stand-in ``state_dict`` for warm-cache builds.
+
+    Serves the real tensor for keys captured in the host-weights sidecar; for every other key it
+    returns a fresh dataless ``torch.empty`` of the manifest shape/dtype (which ``ttnn.as_tensor``
+    discards on the guaranteed cache hit). Mutable (some loaders ``setdefault`` missing KV-shared
+    weights) and truthy (some loaders gate real-weight loading on ``if state_dict:``)."""
+
+    # Explicit marker that this is a warm-cache stand-in, NOT real weights. Callers that must tell
+    # "warm-cache placeholder" apart from "real weights" MUST branch on this attribute, never on
+    # truthiness: this mapping is truthy (non-zero __len__) but tt_transformers' _PlaceholderStateDict
+    # is falsy (__bool__ -> False), so a truthiness test silently means opposite things for the two.
+    # If tt_transformers is ever collapsed onto this class (a listed follow-up), the attribute keeps
+    # `if is_placeholder(...)` reload sites (e.g. test_model_prefill) correct. (#45400 review)
+    is_placeholder = True
+
+    def __init__(self, manifest, host):
+        self._manifest = manifest  # key -> (shape, dtype_str)
+        self._host = dict(host or {})  # key -> real torch.Tensor
+        self._overrides = {}  # keys set by the caller at build time
+        self._deleted = set()
+
+    def __getitem__(self, key):
+        if key in self._deleted:
+            raise KeyError(key)
+        if key in self._overrides:
+            return self._overrides[key]
+        if key in self._host:
+            return self._host[key]
+        spec = self._manifest.get(key)
+        if spec is None:
+            raise KeyError(key)
+        shape, dt = spec
+        return torch.empty(tuple(shape), dtype=_dtype_from_str(dt))
+
+    def __setitem__(self, key, value):
+        self._deleted.discard(key)
+        self._overrides[key] = value
+
+    def __delitem__(self, key):
+        if key not in self:
+            raise KeyError(key)
+        self._overrides.pop(key, None)
+        if key in self._host or key in self._manifest:
+            self._deleted.add(key)
+
+    def __iter__(self):
+        seen = set()
+        for k in list(self._overrides) + list(self._host) + list(self._manifest):
+            if k in self._deleted or k in seen:
+                continue
+            seen.add(k)
+            yield k
+
+    def __len__(self):
+        return sum(1 for _ in self)
+
+    # Mapping's default __contains__/get/items route through __getitem__, which allocates a
+    # full-size torch.empty for EVERY key touched -- including multi-GB ones like lm_head.weight.
+    # substate() (models/tt_dit/utils/substate.py) iterates .items() and filters by prefix, so a
+    # 62-layer gemma4 build would allocate the entire model once per layer just to discard it.
+    # Answer membership from the key sets, and make items() lazy so only matching keys materialize.
+    def __contains__(self, key):
+        if key in self._deleted:
+            return False
+        return key in self._overrides or key in self._host or key in self._manifest
+
+    def keys(self):
+        return list(self)
+
+    def items(self):
+        for k in self:
+            yield k, self[k]
+
+    def get(self, key, default=None):
+        if key not in self:
+            return default
+        return self[key]
+
+
+def build_cached_state_dict(cache_path, host=None, args=None, build_variant=None):
+    """Build the warm-cache stand-in ``state_dict`` from the marker manifest + host sidecar.
+
+    ``host`` may be a sidecar dict already loaded by ``weight_cache_is_complete``'s validation, to
+    avoid a second multi-GB ``torch.load`` of the same file on every warm run (gemma-4-31b's
+    embedding alone is ~2.8 GB).
+
+    ``args`` (a ModelArgs-like) has ``is_mixture_of_experts`` restored from the marker. That flag is
+    normally set as a side effect of ``load_state_dict`` (by sniffing for ``.experts.`` keys), which
+    the warm path skips -- so without this a MoE checkpoint would build a dense decoder and die on a
+    missing ``feed_forward.w1.weight``. (#45400 review)"""
+    cache_path = Path(cache_path)
+    meta = json.loads(marker_path(cache_path, build_variant).read_text())
+    manifest = meta["weights"]
+    if args is not None and hasattr(args, "__dict__"):
+        args.is_mixture_of_experts = bool(meta.get("is_moe", False))
+        # fuse_qkv / fuse_mlp are normally sniffed from the checkpoint keys inside load_state_dict,
+        # which the warm path skips -- leaving them at their __init__ defaults and silently changing
+        # how the decoder is built. The manifest holds the same key set, so derive them identically.
+        keys = manifest.keys()
+        args.fuse_qkv = any("qkv" in k for k in keys)
+        args.fuse_mlp = any("gate_up" in k for k in keys)
+        if args.is_mixture_of_experts:
+            args.moe = True
+            expert_indices = [int(k[-11]) + 1 for k in keys if "block_sparse_moe.experts" in k]
+            if expert_indices:
+                args.num_experts = max(expert_indices)
+            elif hasattr(args, "num_local_experts"):
+                args.num_experts = args.num_local_experts
+    if host is None and meta.get("host_weights"):
+        # consume=True: reuse the load the completeness gate just performed and release the
+        # memoized entry, so the sidecar is read from NAS once per warm run, not twice. (R1)
+        host = load_host_sidecar(cache_path, consume=True)
+    host = host or {}
+    logger.info(
+        f"Warm ttnn weight cache: built state_dict for {len(manifest)} weights "
+        f"({len(host)} real host weights, no full HF load)."
+    )
+    return CachedStateDict(manifest, host)

code/models/tt_transformers/Host_Mem_Profile.md

@@ -0,0 +1,88 @@
+# Model memory on host for performance and accuracy
+
+[Host memory usage results](../sample_data/host_mem_profiling/) at peak memory (in MB) are collected from host_mem_profiler.py by running the target demo under memory_profiler at
+100 ms sampling intervals. You can generate these results by running the script with --mesh-device, --hf-model, and
+--test arguments in [demo/simple_text_demo.py](demo/simple_text_demo.py); it produces a timestamped PNG plot under profiling_results/ and prints peak and baseline RSS to the
+console.
+
+Note that all measurements include child processes spawned by the TT device runtime, giving a realistic view of total
+host-side footprint. Peak memory is defined as the maximum RSS observed across the full test lifetime — including
+model weight loading, KV-cache allocation, and decode iterations. Baseline is the RSS at process start before any
+model work begins.
+
+Host memory (RSS) was profiled across three configurations — [Performance](#performance), [Accuracy](#accuracy), and [Multimodal](#multimodal-models) — using up to 200
+decode iterations on T3K (8-chip) and N300 (2-chip) devices.
+
+All [numbers](../sample_data/host_mem_profiling/) below were captured with a maximum generation of 200 tokens (200 decode iterations) to match the
+conditions used for centralized targets in [models/model_targets.yaml](../model_targets.yaml).
+
+## Performance
+
+This configuration uses bfp4 MLP and bfp8 attention weights for all models except:
+* Qwen-2.5-7B, which uses bfp8 MLP and bfp16 attention weights in all decoder layers
+* Llama-3.1-8B which uses bfp8 MLP in only the 32nd decoder layer and bfp4 MLP elsewhere
+
+| Model             | Device      | batch-1 max mem (MB)  | batch-32 max mem (MB) | Speed (t/s/u) | TTFT (ms) |
+|-------------------|-------------|-----------------------|-----------------------|---------------|-----------|
+| Llama-3.2-1B      | T3K         |         3532          |         2712          | 119.8         | 32        |
+| Llama-3.2-3B      | T3K         |         7846          |         5388          | 68.5          | 52        |
+| Llama-3.1-8B      | T3K         |         13683         |         6922          | 64.3          | 53        |
+| Llama-3.2-11B     | T3K         |         13885         |         7254          | 62.7          | 47        |
+| Llama-3.1-70B     | T3K         |         35820         |         35187         | 16.6          | 164       |
+| Llama-3.3-70B     | T3K         |         101640        |         36176         | 16.6          | 164       |
+| Llama-3.2-90B     | T3K         |         34427         |         34842         | 6             | 5535      |
+| Qwen2.5-7B        | N300        |         16597         |         17754         | 24.6          | 92        |
+| Qwen2.5-72B       | T3K         |         147106        |         145432        | 15.2          | 225       |
+| Qwen2.5-32B       | T3K         |         53434         |         15594         | 22.4          | 190       |
+| Qwen3-32B         | T3K         |         56033         |         20422         | 22.9          | 123       |
+| QwQ-32B           | T3K         |         69900         |         67186         | 20.7          | 105       |
+
+## Accuracy
+
+This configuration uses bfp8 MLP and BF16 attention weights (70B+ models use bfp8 attention and bfp4 MLP).
+Llama 3 models test as insensitive to attention precision and so we use bfp8 attention and kv-cache for them even in accuracy mode.
+
+| Model             | Device      | batch-1 max mem (MB)  | batch-32 max mem (MB) | Speed (t/s/u) | TTFT (ms) |
+|-------------------|-------------|----------------------|------------------------|---------------|-----------|
+| Llama-3.2-1B      | T3K         |        4654          |          2707          | 120.5         | 28        |
+| Llama-3.2-3B      | T3K         |        10374         |          5270          | 67.9          | 69        |
+| Llama-3.1-8B      | T3K         |        20406         |          7704          | 60.8          | 81        |
+| Llama-3.2-11B     | T3K         |        20265         |          7051          | 61.4          | 53        |
+| Llama-3.1-70B     | T3K         |        154122        |          37370         | 16.5          | 168       |
+| Llama-3.3-70B     | T3K         |        100112        |          35026         | 16.5          | 168       |
+| Llama-3.2-90B     | T3K         |        155139        |          34997         | 6             | 5600      |
+| Qwen2.5-7B        | N300        |        20189         |          17956         | 24.6          | 92        |
+| Qwen2.5-72B       | T3K         |        146204        |          147032        | 15.1          | 216       |
+| Qwen2.5-32B       | T3K         |        74162         |          14811         | 19.7          | 183       |
+| Qwen3-32B         | T3K         |        76288         |          20785         | 19.6          | 119       |
+| QwQ-32B           | T3K         |        72627         |          67278         | 18.3          | 120       |
+
+
+## Multimodal Models
+The results are collected using this script [demo/simple_vision_demo.py](demo/simple_vision_demo.py)); with host memory profiler and evaluated on both language and vision branches.
+The memory is reported in MB.
+
+| Model             | Device      | batch-1 max mem (MB)  | batch-32 max mem (MB) | Speed (t/s/u) | TTFT (ms) |
+|-------------------|-------------|-----------------------|-----------------------|---------------|-----------|
+| Llama-3.2-11B     | T3K         |         14309         |         8610          | 61.4          | 53        |
+| Llama-3.2-90B     | T3K         |         75451         |         -             | 6             | 5600      |
+| Qwen2.5-VL-3B     | N300        |         9319          |         9068          | 24.6          | 92        |
+| Qwen2.5-VL-7B     | N300        |         18745         |         10946         | 24.6          | 92        |
+| Qwen2.5-VL-72B    | T3K         |         160437        |         147925        | 15.1          | 216       |
+| Qwen2.5-VL-32B    | T3K         |         71960         |         68464         | 19.7          | 183       |
+| Qwen3-VL-32B      | T3K         |         76297         |         21514         | 19.6          | 119       |
+
+
+[Performance](#performance) vs. [Accuracy](#accuracy) modes trade memory for precision: accuracy mode uses BF16 attention weights versus bfp8/bfp4
+in performance mode, and typically raises peak memory. The increase is clearest at batch-1, where most models grow noticeably Llama-3.1-8B rises from 13.7 GB to 20.4 GB and Qwen2.5-32B from 53 GB to 74 GB — while 70B+ models show a smaller relative increase since attention weight precision is already reduced in both modes. At batch-32 the picture is mixed: KV-cache and activations dominate, and several models (Llama-3.2-1B/3B/11B, Llama-3.3-70B, Qwen2.5-32B) report slightly lower peaks in accuracy mode rather than higher.
+
+Memory scales roughly with model size, ranging from ~3–4 GB for 1B models up to ~147–155 GB for 72–90B models. The
+largest models (Qwen2.5-72B, Llama-3.2-90B) show little difference between batch-1 and batch-32, suggesting KV-cache
+and weight loading dominate over activation memory. Mid-size models (32B class) exhibit a larger batch-1 to batch-32
+gap, likely due to KV-cache growth with sequence length at batch-1.
+
+Throughput and TTFT follow expected trends: smaller models deliver higher tokens/s/user (up to 120 t/s/u for 1B) and
+low TTFT (~30 ms), while 90B models drop to 6 t/s/u with TTFT exceeding 5.5 seconds.
+
+[Multimodal](#multimodal-models) models carry a modest overhead versus their text-only counterparts — Llama-3.2-11B adds ~400 MB at batch-1
+though Qwen2.5-VL-72B reaches the highest overall footprint at ~160 GB.

code/models/tt_transformers/PERF.md

@@ -0,0 +1,247 @@
+Note: DEPRECATED. PLEASE CHECK THE YAML FILE [models/model_targets.yaml](models/model_targets.yaml)
+
+# Model performance and accuracy
+
+Performance and token accuracy using teacher forcing is collected from [demo/simple_text_demo.py](demo/simple_text_demo.py) with the `ci-token-matching` test case. You can generate this table by running these tests with the `lt` tool (tell it to run `table` or `pareto`) and pressing `m` whilst in the results section to export to markdown.
+
+Note that token accuracy parses the below to determine expected values +- 0.5. In May 2025 we switched the default to measuring the accuracy by prefilling 512 tokens and generating another 511, rather than generating 128 tokens in earlier versions. This caused overall accuracy values to drop slightly.
+
+Also note that all the performance metrics below were taken for a maximum generation of 200 tokens, i.e., 200 decode iterations.
+
+## Performance
+
+This configuration uses bfp4 MLP and bfp8 attention weights for all models except:
+* Qwen-2.5-7B, which uses bfp8 MLP and bfp16 attention weights in all decoder layers
+* Llama-3.1-8B which uses bfp8 MLP in only the 32nd decoder layer and bfp4 MLP elsewhere
+
+| Model             | Device      | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|-------------------|-------------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B      | N150        | 79        | 97        | 87.8          | 26        |
+| Llama-3.2-1B      | N300        | 79        | 97        | 105.9         | 22        |
+| Llama-3.2-1B      | T3K         | 80        | 97        | 119.8         | 32        |
+| Llama-3.2-1B      | TG          | 77        | 96        | 51.0          |           |
+| Llama-3.2-3B      | N150        | 89        | 98        | 54.0          | 55        |
+| Llama-3.2-3B      | N300        | 89        | 98        | 68.0          | 39        |
+| Llama-3.2-3B      | T3K         | 91        | 99        | 68.5          | 52        |
+| Llama-3.2-3B      | TG          | 87        | 97        | 33.5          |           |
+| Llama-3.1-8B      | N150        | 90        | 97        | 28.3          | 104       |
+| Llama-3.1-8B      | N300        | 90        | 97        | 44.2          | 67        |
+| Llama-3.1-8B      | P100        | 90        | 98        | 29.5          | 84        |
+| Llama-3.1-8B      | P150        | 90        | 98        | 33.6          | 76        |
+| Llama-3.1-8B      | T3K         | 90        | 98        | 64.3          | 53        |
+| Llama-3.1-8B      | T3K  (DP=4) |           |           | 39.6          | 58        |
+| Llama-3.1-8B      | T3K  (DP=8) |           |           | 24.9          | 86        |
+| Llama-3.1-8B      | TG          | 88        | 97        | 29.5          |           |
+| Llama-3.2-11B     | N150        | 90        | 98        | 55.5          | 58        |
+| Llama-3.2-11B     | N300        | 90        | 98        | 44.1          | 67        |
+| Llama-3.2-11B     | T3K         | 90        | 98        | 62.7          | 47        |
+| Llama-3.2-11B     | TG          | 87        | 97        | 29.5          |           |
+| Llama-3.1-70B     | T3K         | 96        | 100       | 16.6          | 164       |
+| Llama-3.3-70B     | T3K         | 96        | 100       | 16.6          | 164       |
+| Llama-3.1-70B     | TG          | 95        | 100       | 12.7          |           |
+| Llama-3.3-70B     | TG          | 95        | 100       | 12.7          |           |
+| Llama-3.1-70B     | TG   (DP=4) |           |           | 14.8          | 189       |
+| Llama-3.2-90B     | T3K         | 96        | 100       | 6             | 5535      |
+| Qwen2.5-7B        | N300        | 84        | 96        | 24.6          | 92        |
+| Qwen2.5-72B       | T3K         | 99        | 100       | 15.2          | 225       |
+| Qwen2.5-32B       | T3K         | 98        | 99        | 22.4          | 190       |
+| Qwen2.5-72B       | P150x8      | 99        | 100       |               |           |
+| Qwen2.5-32B       | P150x8      | 97        | 99        |               |           |
+| Qwen2.5-Coder-32B | T3K         | 96        | 99        | 22.4          | 190       |
+| Qwen3-32B         | T3K         | 89        | 97        | 22.9          | 123       |
+| QwQ-32B           | T3K         | 96        | 100       | 20.7          | 105       |
+| Phi3.5-mini       | N150        |           |           | 43.2          | 98        |
+| Phi3.5-mini       | N300        |           |           | 57.8          | 62        |
+| Phi3.5-mini       | T3K         |           |           | 48.8          | 51        |
+| Mistral-7B        | N150        | 95        | 99        | 29.75         | 100.24    |
+| Mistral-7B        | N300        | 95        | 100       | 47.01         | 65.95     |
+| Mistral-7B        | T3K         | 95        | 100       | 67.82         | 53.93     |
+| Mistral-Small-3.1-24B | T3K      | 95        | 99        |               |           |
+| Phi-3-mini-128k-instruct | N150        | 89        | 99        | 45.0          | 73.32     |
+| Phi-3-mini-128k-instruct | N300        | 89        | 99        | 60.87         | 114.94    |
+| Phi-4 | N300 | 97 | 100 | 37.34 | 123.33 |
+| Mixtral-8x7B-v0.1 | T3K         | 98        | 100       | 67.82         | 53.93     |
+| Ministral-8B      | N300        | 93        | 98        | 22.15         | 79.3      |
+
+## Accuracy
+
+This configuration uses bfp8 MLP and BF16 attention weights (70B+ models use bfp8 attention and bfp4 MLP).
+Llama 3 models test as insensitive to attention precision and so we use bfp8 attention and kv-cache for them even in accuracy mode.
+
+| Model             | Device      | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|-------------------|-------------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B      | N150        | 87        | 99        | 84.7          | 29        |
+| Llama-3.2-1B      | N300        | 87        | 98        | 102.8         | 21        |
+| Llama-3.2-1B      | T3K         | 88        | 99        | 120.5         | 28        |
+| Llama-3.2-1B      | TG          | 85        | 98        | 48.4          |           |
+| Llama-3.2-3B      | N150        | 96        | 100       | 47.6          | 63        |
+| Llama-3.2-3B      | N300        | 96        | 100       | 63.5          | 41        |
+| Llama-3.2-3B      | T3K         | 96        | 100       | 67.9          | 69        |
+| Llama-3.2-3B      | TG          | 92        | 99        | 33.6          |           |
+| Llama-3.1-8B      | N150        | 96        | 100       | 25.2          | 138       |
+| Llama-3.1-8B      | N300        | 96        | 100       | 38.8          | 79        |
+| Llama-3.1-8B      | T3K         | 97        | 100       | 60.8          | 81        |
+| Llama-3.1-8B      | TG          | 95        | 100       | 29.5          |           |
+| Llama-3.2-11B     | N150        | 95        | 100       | 56.7          | 62        |
+| Llama-3.2-11B     | N300        | 95        | 100       | 38.3          | 78        |
+| Llama-3.2-11B     | T3K         | 96        | 100       | 61.4          | 53        |
+| Llama-3.2-11B     | TG          | 94        | 100       | 29.5          |           |
+| Llama-3.1-70B     | T3K         | 96        | 100       | 16.5          | 168       |
+| Llama-3.1-70B     | TG          | 95        | 100       | 12.7          |           |
+| Llama-3.2-90B     | T3K         | 96        | 100       | 6             | 5600      |
+| Qwen2.5-7B        | N300        | 84        | 96        | 24.6          | 92        |
+| Qwen2.5-72B       | T3K         | 99        | 100       | 15.1          | 216       |
+| Qwen2.5-32B       | T3K         | 95        | 99        | 19.7          | 183       |
+| Qwen2.5-72B       | P150x8      | 99        | 100       |               |           |
+| Qwen2.5-32B       | P150x8      | 97        | 99        |               |           |
+| Qwen2.5-Coder-32B | T3K         | 95        | 99        | 19.7          | 183       |
+| Qwen3-32B         | T3K         | 95        | 100       | 19.6          | 119       |
+| QwQ-32B           | T3K         | 99        | 100       | 18.3          | 120       |
+| Phi3.5-mini       | N150        |           |           | 38.8          | 92        |
+| Phi3.5-mini       | N300        |           |           | 53.9          | 63        |
+| Phi3.5-mini       | T3K         |           |           | 48.6          | 53        |
+| Mistral-7B        | N150        | 96        | 100       | 29.75         | 100.24    |
+| Mistral-7B        | N300        | 97        | 100       | 47.01         | 65.95     |
+| Mistral-7B        | T3K         | 98        | 100       | 67.82         | 53.93     |
+| Mistral-Small-3.1-24B | T3K      | 95        | 99        |               |           |
+| Phi-3-mini-128k-instruct | N150        | 94        | 99        | 40.41         | 82.58     |
+| Phi-3-mini-128k-instruct | N300        | 94        | 99        | 57.0          | 115.36    |
+| Phi-4 | N300 | 99 | 100 | 20.48 | 146.32 |
+| Mixtral-8x7B-v0.1 | T3K         | 95        | 100       | 67.82         | 53.93     |
+| Ministral-8B      | N300        | 97        | 100       | 19.95         | 93.22     |
+
+##  Long-context (64K Tokens)
+
+This configuration uses bfp4 MLP FF1+FF3 for all models. **Batch_size=1 and prefill_length is 64k tokens.**
+
+| Model          | Device | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|---------------|-----------|
+| Llama-3.2-1B   | N150   | 53.0          | 20066     |
+| Llama-3.2-1B   | N300   | 65.2          | 10949     |
+| Llama-3.2-1B   | T3K    | 73.7          | 5271      |
+| Llama-3.2-1B   | TG     |               |           |
+| Llama-3.2-3B   | N150   | 25.3          | 46743     |
+| Llama-3.2-3B   | N300   | 34.8          | 22921     |
+| Llama-3.2-3B   | T3K    | 41.0          | 10677     |
+| Llama-3.2-3B   | TG     |               |           |
+| Llama-3.1-8B   | N150   | 16.9          | 64385     |
+| Llama-3.1-8B   | N300   | 26.1          | 36229     |
+| Llama-3.1-8B   | T3K    | 38.1          | 16165     |
+| Llama-3.1-8B   | TG     |               |           |
+| Llama-3.2-11B  | N300   | 26.1          | 36247     |
+| Llama-3.2-11B  | T3K    | 38.4          | 16167     |
+| Llama-3.2-11B  | TG     |               |           |
+| Llama-3.1-70B  | T3K    | 11.9          | 74363     |
+| Llama-3.1-70B  | TG     |               |           |
+| Qwen2.5-7B     | N300   |               |           |
+| Qwen2.5-72B    | T3K    |               |           |
+
+##  Long-context (32K Tokens)
+
+This configuration uses bfp4 MLP FF1+FF3 for all models. **Batch_size=1 and prefill_length is 32k tokens.**
+
+| Model                     | Device | Speed (t/s/u) | TTFT (ms) |
+|---------------------------|--------|---------------|-----------|
+| Phi-3-mini-128k-instruct  | N300   | 26.1          | 10072     |
+
+## Short-Context, Batch-32
+
+This configuration uses bfp4 MLP FF1+FF3 for all models. **Batch_size=32 and prefill_length is 128 tokens.**
+
+| Model          | Device | Speed (t/s/u) | avg TTFT (ms) |
+|----------------|--------|---------------|---------------|
+| Llama-3.2-1B   | N150   | 54.7          | 38            |
+| Llama-3.2-1B   | N300   | 64.2          | 34            |
+| Llama-3.2-1B   | T3K    | 69.9          | 42            |
+| Llama-3.2-1B   | TG     |               |               |
+| Llama-3.2-3B   | N150   | 36.5          | 69            |
+| Llama-3.2-3B   | N300   | 45.8          | 51            |
+| Llama-3.2-3B   | T3K    | 47.8          | 63            |
+| Llama-3.2-3B   | TG     |               |               |
+| Llama-3.1-8B   | N150   | 22.3          | 119           |
+| Llama-3.1-8B   | N300   | 33.5          | 80            |
+| Llama-3.1-8B   | T3K    | 45.6          | 64            |
+| Llama-3.1-8B   | TG     |               |               |
+| Llama-3.2-11B  | N300   | 33.4          | 79            |
+| Llama-3.2-11B  | T3K    | 45.1          | 64            |
+| Llama-3.2-11B  | TG     |               |               |
+| Llama-3.1-70B  | T3K    | 14.8          | 192           |
+| Llama-3.1-70B  | TG     |               |               |
+| Qwen2.5-7B     | N300   |               |               |
+| Qwen2.5-72B    | T3K    |               |               |
+| Phi-3-mini-128k-instruct  | 150    | 25.66         | 68.58         |
+| Phi-3-mini-128k-instruct  | N300   | 39.4          | 85.99         |
+
+# Llama 3 model precision and math fidelity
+
+## precision_cfg = {ff1_3: bfp4, ff2: bfp4, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: lofi, li_ff2: lofi, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 85        | 98        | 100.3         | 69        |
+
+## precision_cfg = {ff1_3: bfp4, ff2: bfp8, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: lofi, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 88        | 98        | 100.3         | 55        |
+
+## precision_cfg = {ff1_3: bfp4, ff2: bf16, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: lofi, li_ff2: hifi4, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 87        | 98        | 96.8          | 51        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bfp4, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: lofi, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 87        | 98        | 98.5          | 50        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bfp8, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 91        | 98        | 99.0          | 60        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bf16, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: hifi4, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 89        | 99        | 95.2          | 49        |
+
+## precision_cfg = {ff1_3: bf16, ff2: bfp4, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi4, li_ff2: lofi, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 89        | 98        | 95.2          | 53        |
+
+## precision_cfg = {ff1_3: bf16, ff2: bfp8, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi4, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 91        | 98        | 94.4          | 57        |
+
+## precision_cfg = {ff1_3: bf16, ff2: bf16, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi4, li_ff2: hifi4, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 90        | 98        | 91.2          | 60        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bfp8, wqkv: bfp8, wo: bfp4, kv_cache: bfp8, activation: bf16}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: lofi, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: lofi, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 88        | 98        | 98.2          | 45        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bfp8, wqkv: bfp8, wo: bfp4, kv_cache: bfp8, activation: bfp8}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: lofi, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: lofi, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 90        | 98        | 101.0         | 57        |
+
+## precision_cfg = {ff1_3: bfp8, ff2: bfp8, wqkv: bfp8, wo: bfp4, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: hifi2, li_qkv_decode: hifi2, li_o_decode: lofi, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: lofi, sdpa_prefill: hifi4}
+
+| Model          | Device | Top-1 (%) | Top-5 (%) | Speed (t/s/u) | TTFT (ms) |
+|----------------|--------|-----------|-----------|---------------|-----------|
+| Llama-3.2-1B   | N300   | 90        | 98        | 99.4          | 73        |

code/models/tt_transformers/README.md

@@ -0,0 +1,388 @@
+# TT-Transformers
+
+This code can run large language models such as the Llama3 family, Qwen2.5, Mistral, DeepSeek-R1-Distill variants and similar. Tensor-parallelism automatically distributes workloads across all available chips.
+
+The current version is verified to work with the following models:
+| Model                                                                                            | Hardware                    | <org/model>                                      |
+|--------------------------------------------------------------------------------------------------|-----------------------------|-------------------------------------------------|
+| [DeepSeek R1 Distill Llama 70B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B)| LoudBox / QuietBox / Galaxy | ```deepseek-ai/DeepSeek-R1-Distill-Llama-70B``` |
+| [Llama 3.1 8B](https://huggingface.co/meta-llama/Llama-3.1-8B)                                   | n150 / p100 / p150          | ```meta-llama/Llama-3.1-8B```                   |
+| [Llama 3.1 70B](https://huggingface.co/meta-llama/Llama-3.1-70B)                                 | LoudBox / QuietBox / Galaxy | ```meta-llama/Llama-3.1-70B```                  |
+| [Llama 3.2 1B](https://huggingface.co/meta-llama/Llama-3.2-1B)                                   | n150                        | ```meta-llama/Llama-3.2-1B```                   |
+| [Llama 3.2 3B](https://huggingface.co/meta-llama/Llama-3.2-3B)                                   | n150                        | ```meta-llama/Llama-3.2-3B```                   |
+| [Llama 3.2 11B Vision](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision)                   | n300                        | ```meta-llama/Llama-3.2-11B-Vision```           |
+| [Llama 3.2 90B Vision](https://huggingface.co/meta-llama/Llama-3.2-90B-Vision)                   | LoudBox / QuietBox          | ```meta-llama/Llama-3.2-90B-Vision```           |
+| [Mistral 7B Instruct v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3)            | n150                        | ```mistralai/Mistral-7B-Instruct-v0.3```        |
+| [Mistral Small 3.1 24B Instruct](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503) | T3K                   | ```mistralai/Mistral-Small-3.1-24B-Instruct-2503``` |
+| [Mixtral 8x7B Instruct v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1)        | LoudBox / QuietBox          | ```mistralai/Mixtral-8x7B-Instruct-v0.1```        |
+| [Qwen 2.5 7B](https://huggingface.co/Qwen/Qwen2.5-7B)                                            | n300                        | ```Qwen/Qwen2.5-7B```                           |
+| [Qwen 2.5 Coder 32B](https://huggingface.co/Qwen/Qwen2.5-Coder-32B)                              | LoudBox / QuietBox          | ```Qwen/Qwen2.5-Coder-32B```                    |
+| [Qwen 2.5 72B](https://huggingface.co/Qwen/Qwen2.5-72B)                                          | LoudBox / QuietBox          | ```Qwen/Qwen2.5-72B```                          |
+| [Qwen 3 32B](https://huggingface.co/Qwen/Qwen3-32B)                                              | LoudBox / QuietBox          | ```Qwen/Qwen3-32B```                            |
+| [Phi-3-mini-128k-instruct](https://huggingface.co/microsoft/Phi-3-mini-128k-instruct)            | n150 / n300                 | ```microsoft/Phi-3-mini-128k-instruct```        |
+
+<details>
+<summary>  Explore an extended list of compatible models (experimental) </summary>
+
+| Model                                                                                                                     | Hardware | <org/model>.                                                 |
+|---------------------------------------------------------------------------------------------------------------------------|----------|-------------------------------------------------------------|
+| [Bellatrix-Tiny-0.5B](https://huggingface.co/prithivMLmods/Bellatrix-Tiny-0.5B)                                           | N150     | `prithivMLmods/Bellatrix-Tiny-0.5B`                         |
+| [Bellatrix-Tiny-1B](https://huggingface.co/prithivMLmods/Bellatrix-Tiny-1B)                                               | N150     | `prithivMLmods/Bellatrix-Tiny-1B`                           |
+| [Bellatrix-Tiny-1B-R1](https://huggingface.co/prithivMLmods/Bellatrix-Tiny-1B-R1)                                         | N150     | `prithivMLmods/Bellatrix-Tiny-1B-R1`                        |
+| [deepseek-coder-6.7b-instruct](https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct)                           | N300     | `deepseek-ai/deepseek-coder-6.7b-instruct`                  |
+| [deepseek-llm-7b-base](https://huggingface.co/deepseek-ai/deepseek-llm-7b-base)                                           | N300     | `deepseek-ai/deepseek-llm-7b-base`                          |
+| [deepseek-llm-7b-chat](https://huggingface.co/deepseek-ai/deepseek-llm-7b-chat)                                           | N300     | `deepseek-ai/deepseek-llm-7b-chat`                          |
+| [DeepSeek-R1-Distill-Llama-8B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B)                           | N300     | `deepseek-ai/DeepSeek-R1-Distill-Llama-8B`                  |
+| [DeepSeek-R1-Distill-Llama-8B](https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B)                               | N300     | `unsloth/DeepSeek-R1-Distill-Llama-8B`                      |
+| [DeepSeek-R1-Distill-Llama-8B-abliterated](https://huggingface.co/huihui-ai/DeepSeek-R1-Distill-Llama-8B-abliterated)     | N300     | `huihui-ai/DeepSeek-R1-Distill-Llama-8B-abliterated`        |
+| [DeepSeek-R1-Distill-Qwen-1.5B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B)                         | N150     | `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B`                 |
+| [DeepSeek-R1-Distill-Qwen-32B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B)                           | T3K      | `deepseek-ai/DeepSeek-R1-Distill-Qwen-32B`                  |
+| [DeepSeek-R1-Distill-Qwen-32B-Japanese](https://huggingface.co/cyberagent/DeepSeek-R1-Distill-Qwen-32B-Japanese)          | T3K      | `cyberagent/DeepSeek-R1-Distill-Qwen-32B-Japanese`          |
+| [DeepSeek-R1-Distill-Qwen-32B-abliterated](https://huggingface.co/huihui-ai/DeepSeek-R1-Distill-Qwen-32B-abliterated)     | T3K      | `huihui-ai/DeepSeek-R1-Distill-Qwen-32B-abliterated`        |
+| [DeepSeek-R1-Distill-Qwen-7B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B)                             | N300     | `deepseek-ai/DeepSeek-R1-Distill-Qwen-7B`                   |
+| [DeepSeek-R1-Distill-Qwen-7B-abliterated-v2](https://huggingface.co/huihui-ai/DeepSeek-R1-Distill-Qwen-7B-abliterated-v2) | N300     | `huihui-ai/DeepSeek-R1-Distill-Qwen-7B-abliterated-v2`      |
+| [DeepSeek-R1-Distill-Qwen-7B-Japanese](https://huggingface.co/lightblue/DeepSeek-R1-Distill-Qwen-7B-Japanese)             | N300     | `lightblue/DeepSeek-R1-Distill-Qwen-7B-Japanese`            |
+| [Dobby-Mini-Leashed-Llama-3.1-8B](https://huggingface.co/SentientAGI/Dobby-Mini-Leashed-Llama-3.1-8B)                     | N300     | `SentientAGI/Dobby-Mini-Leashed-Llama-3.1-8B`               |
+| [Dobby-Mini-Unhinged-Llama-3.1-8B](https://huggingface.co/SentientAGI/Dobby-Mini-Unhinged-Llama-3.1-8B)                   | N300     | `SentientAGI/Dobby-Mini-Unhinged-Llama-3.1-8B`              |
+| [Dolphin3.0-Mistral-24B](https://huggingface.co/cognitivecomputations/Dolphin3.0-Mistral-24B)                             | T3K      | `cognitivecomputations/Dolphin3.0-Mistral-24B`              |
+| [Dolphin3.0-R1-Mistral-24B](https://huggingface.co/cognitivecomputations/Dolphin3.0-R1-Mistral-24B)                       | T3K      | `cognitivecomputations/Dolphin3.0-R1-Mistral-24B`           |
+| [EuroLLM-9B-Instruct](https://huggingface.co/utter-project/EuroLLM-9B-Instruct)                                           | N300     | `utter-project/EuroLLM-9B-Instruct`                         |
+| [GPT4chan-24B](https://huggingface.co/v2ray/GPT4chan-24B)                                                                 | T3K      | `v2ray/GPT4chan-24B`                                        |
+| [Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)                                                     | N300     | `meta-llama/Llama-2-7b-hf`                                  |
+| [Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)                                                          | N300     | `meta-llama/Llama-2-7b-hf`                                  |
+| [Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)                                          | N300     | `meta-llama/Llama-3.1-8B-Instruct`                          |
+| [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)                                          | N150     | `meta-llama/Llama-3.2-1B-Instruct`                          |
+| [Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)                                          | N150     | `meta-llama/Llama-3.2-3B-Instruct`                          |
+| [Llama-3.2-6B-AlgoCode](https://huggingface.co/prithivMLmods/Llama-3.2-6B-AlgoCode)                                       | N150     | `prithivMLmods/Llama-3.2-6B-AlgoCode`                       |
+| [Llama-Krikri-8B-Base](https://huggingface.co/ilsp/Llama-Krikri-8B-Base)                                                  | N150     | `ilsp/Llama-Krikri-8B-Base`                                 |
+| [Llama-Krikri-8B-Instruct](https://huggingface.co/ilsp/Llama-Krikri-8B-Instruct)                                          | N300     | `ilsp/Llama-Krikri-8B-Instruct`                             |
+| [LwQ-10B-Instruct](https://huggingface.co/prithivMLmods/LwQ-10B-Instruct)                                                 | N300     | `prithivMLmods/LwQ-10B-Instruct`                            |
+| [Magnum-v4-Cydonia-vXXX-22B](https://huggingface.co/Kaoeiri/Magnum-v4-Cydonia-vXXX-22B)                                   | T3K      | `Kaoeiri/Magnum-v4-Cydonia-vXXX-22B`                        |
+| [Megatron-Opus-7B-Exp](https://huggingface.co/prithivMLmods/Megatron-Opus-7B-Exp)                                         | N300     | `prithivMLmods/Megatron-Opus-7B-Exp`                        |
+| [Meta-Llama-3-8B](https://huggingface.co/meta-llama/Meta-Llama-3-8B)                                                      | N300     | `meta-llama/Meta-Llama-3-8B`                                |
+| [Meta-Llama-3-8B-Instruct](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct)                                    | N300     | `meta-llama/Meta-Llama-3-8B-Instruct`                       |
+| [Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)                                     | N150     | `mistralai/Mistral-7B-Instruct-v0.1`                        |
+| [Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2)                                     | N150     | `mistralai/Mistral-7B-Instruct-v0.2`                        |
+| [Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)                                                       | N150     | `mistralai/Mistral-7B-v0.1`                                 |
+| [Mistral-Small-24B-Base-2501](https://huggingface.co/mistralai/Mistral-Small-24B-Base-2501)                               | T3K      | `mistralai/Mistral-Small-24B-Base-2501`                     |
+| [Mistral-Small-24B-Instruct-2501](https://huggingface.co/mistralai/Mistral-Small-24B-Instruct-2501)                       | T3K      | `mistralai/Mistral-Small-24B-Instruct-2501`                 |
+| [Mistral-Small-3.1-24B-Instruct-2503](https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503)             | T3K      | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`             |
+| [NuminaMath-7B-TIR](https://huggingface.co/AI-MO/NuminaMath-7B-TIR)                                                       | N300     | `AI-MO/NuminaMath-7B-TIR`                                   |
+| [OpenThinker-7B](https://huggingface.co/open-thoughts/OpenThinker-7B)                                                     | N300     | `open-thoughts/OpenThinker-7B`                              |
+| [Primal-Mini-3B-Exp](https://huggingface.co/prithivMLmods/Primal-Mini-3B-Exp)                                             | N150     | `prithivMLmods/Primal-Mini-3B-Exp`                          |
+| [Qwen2.5-0.5B](https://huggingface.co/Qwen/Qwen2.5-0.5B)                                                                  | N150     | `Qwen/Qwen2.5-0.5B`                                         |
+| [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct)                                                | N150     | `Qwen/Qwen2.5-0.5B-Instruct`                                |
+| [Qwen2.5-1.5B](https://huggingface.co/Qwen/Qwen2.5-1.5B)                                                                  | N150     | `Qwen/Qwen2.5-1.5B`                                         |
+| [Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct)                                                | N150     | `Qwen/Qwen2.5-1.5B-Instruct`                                |
+| [Qwen2.5-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-3B-Instruct)                                                    | N150     | `Qwen/Qwen2.5-3B-Instruct`                                  |
+| [Qwen2.5-32B](https://huggingface.co/Qwen/Qwen2.5-32B)                                                                    | T3K      | `Qwen/Qwen2.5-32B`                                          |
+| [Qwen2.5-Coder-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct)                                        | N300     | `Qwen/Qwen2.5-Coder-7B-Instruct`                            |
+| [Qwen2.5-Math-1.5B](https://huggingface.co/Qwen/Qwen2.5-Math-1.5B)                                                        | N150     | `Qwen/Qwen2.5-Math-1.5B`                                    |
+| [Qwen2.5-Math-7B](https://huggingface.co/Qwen/Qwen2.5-Math-7B)                                                            | N300     | `Qwen/Qwen2.5-Math-7B`                                      |
+| [Reasoning-Distilled-ta-7B](https://huggingface.co/prithivMLmods/Reasoning-Distilled-ta-7B)                               | N300     | `prithivMLmods/Reasoning-Distilled-ta-7B`                   |
+| [s1-32B](https://huggingface.co/simplescaling/s1-32B)                                                                     | T3K      | `simplescaling/s1-32B`                                      |
+| [Selene-1-Mini-Llama-3.1-8B](https://huggingface.co/AtlaAI/Selene-1-Mini-Llama-3.1-8B)                                    | N300     | `AtlaAI/Selene-1-Mini-Llama-3.1-8B`                         |
+| [Sky-T1-32B-Preview](https://huggingface.co/NovaSky-AI/Sky-T1-32B-Preview)                                                | T3K      | `NovaSky-AI/Sky-T1-32B-Preview`                             |
+| [SmallThinker-3B-Preview](https://huggingface.co/PowerInfer/SmallThinker-3B-Preview)                                      | N150     | `PowerInfer/SmallThinker-3B-Preview`                        |
+| [SmolLM2-1.7B-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct)                                       | N150     | `HuggingFaceTB/SmolLM2-1.7B-Instruct`                       |
+| [Sqweeks-7B-Instruct](https://huggingface.co/prithivMLmods/Sqweeks-7B-Instruct)                                           | N300     | `prithivMLmods/Sqweeks-7B-Instruct`                         |
+| [TinyLlama-1.1B-Chat-v1.0](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0)                                     | N150     | `TinyLlama/TinyLlama-1.1B-Chat-v1.0`                        |
+| [TinySwallow-1.5B](https://huggingface.co/SakanaAI/TinySwallow-1.5B)                                                      | N150     | `SakanaAI/TinySwallow-1.5B`                                 |
+| [TinySwallow-1.5B-Instruct](https://huggingface.co/SakanaAI/TinySwallow-1.5B-Instruct)                                    | N150     | `SakanaAI/TinySwallow-1.5B-Instruct`                        |
+| [Triangulum-v2-10B](https://huggingface.co/prithivMLmods/Triangulum-v2-10B)                                               | N300     | `prithivMLmods/Triangulum-v2-10B`                           |
+| [WebMind-7B-v0.1](https://huggingface.co/prithivMLmods/WebMind-7B-v0.1)                                                   | N300     | `prithivMLmods/WebMind-7B-v0.1`                             |
+| [WizardLM-7B-Uncensored](https://huggingface.co/cognitivecomputations/WizardLM-7B-Uncensored)                             | N300     | `cognitivecomputations/WizardLM-7B-Uncensored`              |
+| [Xwen-7B-Chat](https://huggingface.co/xwen-team/Xwen-7B-Chat)                                                             | N300     | `xwen-team/Xwen-7B-Chat`                                    |
+| [YuE-s1-7B-anneal-en-cot](https://huggingface.co/m-a-p/YuE-s1-7B-anneal-en-cot)                                           | N300     | `m-a-p/YuE-s1-7B-anneal-en-cot`                             |
+| [YuE-s2-1B-general](https://huggingface.co/m-a-p/YuE-s2-1B-general)                                                       | N150     | `m-a-p/YuE-s2-1B-general`                                   |
+
+</details>
+
+## Prerequisites
+
+Install [TT-Metalium and TTNN](../../INSTALLING.md).
+
+
+## How to Run
+
+To configure the weights and run a demo, choose [Automatic Download](#automatic-download) for a quick setup or [Manual Download](#manual-download) for more control over the files.
+
+---
+
+### Automatic Setup
+
+#### Login to HuggingFace using your token
+
+Use `huggingface-cli login` or set the token with the command `export HF_TOKEN=<token>`
+- To obtain a HuggingFace token visit: https://huggingface.co/docs/hub/security-tokens
+
+#### Choose your Model
+
+Set the `HF_MODEL` environment variable to the HuggingFace org/name of the model to be run.
+
+This will automatically download the weights into your HuggingFace cache directory and run the model directly.
+
+- Check the models chart on the top of the page and substitute the `<org/model>` on the following command:
+```
+export HF_MODEL=<org/name>
+```
+
+#### You are all set!
+
+- Skip to section: [Run the Demo](#run-the-demo)
+
+---
+
+### Manual Setup
+
+If you wish, you can manually download the weights [from HuggingFace](#download-weights-from-huggingface) as described by the following section:
+
+#### Download Weights from HuggingFace
+
+Navigate to [HuggingFace](https://huggingface.co) and download the weights of the model you wish to run (check [models chart](#tt-transformers) for compatibility and direct links to HuggingFace models).
+
+- Ensure your model directory has the following structure:
+```
+/path_to/org/model/
+    config.json
+    generation_config.json
+    model-00001-of-00062.safetensors
+    ...
+```
+
+- Set `HF_MODEL` to the directory of the downloaded weights:
+```
+export HF_MODEL=<path_to_downloaded_directory>
+```
+
+#### You are all set!
+
+- Skip to section: [Run the Demo](#run-the-demo)
+
+---
+
+### Run the Demo
+
+- Make sure `HF_MODEL` has been set:
+
+`export HF_MODEL=<hf_model_name or hf_downloaded_directory>`
+
+- Run the Demo
+
+These are some example commands on how to run the demo. For more information about the parameters, check [parameters description.](#parameters-description)
+
+```
+# Batch-1
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1"
+
+# Batch-32
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-32"
+
+# Long context with custom parameters
+pytest models/tt_transformers/demo/simple_text_demo.py -k "long-context" --max_seq_len=16384
+
+# Long-context
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and long"
+```
+
+The above examples are run in `ModelOptimizations.performance` mode. You can override this by setting the `optimizations` or the `decoder_config_file` argument in the demo. To use instead the accuracy mode you can call the above tests with `-k "accuracy and ..."` instead of performance.
+
+NOTE: trace region sizes are declared in [`models/model_trace_region_sizes.yaml`](../model_trace_region_sizes.yaml) and resolved at device-open time via [`get_supported_trace_region_size`](demo/trace_region_config.py) (which delegates to [`resolve_trace_region_size`](../demos/utils/trace_region_sizes.py)). A `(model, SKU)` pair without a YAML entry is not an error: resolution logs an info message and falls back to `TRACE_REGION_SIZE_DYNAMIC` (`0`, dynamic allocation). Add an explicit entry when a model needs a fixed reserved trace region.
+
+## Details
+
+### Extra compatibility settings for non-Llama models
+
+If you are bringing up a new model that is similar to these but is not listed above, you may also need to set additional environment variables:
+- `MAX_PREFILL_CHUNK_SIZE` - this determines how many thousands of tokens are prefilled in one go. For optimal performance pick 128. Depending on the model dimensions and hardware you're running on, there may not be enough L1 to prefill 128K tokens at once, in which case you can reduce this in powers of 2 down to 4.
+- `PAD_MLP_CORES` - models with a hidden_dim that is not a nice power of 2 may not have a valid layout or may run with lower performance. You can set this to a multiple of 8 between 8 and 64; `16` and `32` commonly work well if this is required.
+
+You should also watch out for:
+- RoPE encoding style. `llama3`, `yarn` and of course `none` are supported. HuggingFace models encode the complex numbers in RoPE as r1, r2, ..., i1, i2, ... whereas Meta models encode the complex numbers as r1, i1, r2, i2, ... - TTT uses a Meta-style implementation of the RoPE op and when loading a HuggingFace it will reshuffle the weights of the pre-rope attention weights to interleave their outputs in this style (see `reverse_permute` in [load_checkpoints.py](tt/load_checkpoints.py)). When _using_ TTT this happens invisibly and should not affect you, but when bringing up a new model or modifying TTT being aware of this will make some otherwise confusing things clearer.
+- Our [accuracy test](demo/simple_text_demo.py -k ci-token-matching) will require you to [generate some reference logits](tests/generate_reference_hf.py) and perhaps update the test to use them.
+- We parallelise attention over the number of heads. If this number is e.g. 14 then you will not be able to run it on more than 2 chips (because 14/2=7, a prime number). We do not support head-padding or similar mitigations at this time but a PR would be cool.
+
+Huggingface models specify their architecture in the `config.json` file. The following architectures are known to work:
+
+- LlamaForCausalLM
+- Qwen2ForCausalLM
+- Qwen3ForCausalLM
+- MistralForCausalLM
+- Mistral3ForConditionalGeneration
+- Phi3ForCausalLM
+
+At the time of writing this covers the majority of popular HuggingFace text-generation models. If you find another architecture that works or extend TT-Transformers to support one we would love to accept a PR!
+
+---
+
+### Environment Variables Description
+
+- `HF_MODEL` is the HuggingFace org/name of the model you want to run or the path to the downloaded Huggingface weights.
+- `TT_CACHE_PATH` is optional. It sets the path for ttnn's weight cache files. See below for more details.
+- `MESH_DEVICE` is optional. It allows you to use fewer devices than are available. See below for more details.
+
+On the first execution of each model, TTNN will create weight cache files for that model, to speed up future runs. These cache files only need to be created once for each model and device. These files are stored in one of three places:
+
+1. `TT_CACHE_PATH` if you have set it.
+2. `HF_MODEL/device_name` if a path to downloaded weights was specified using `HF_MODEL`.
+3. `model_cache/HF_MODEL/device_name` if a HuggingFace model name was specified using `HF_MODEL`.
+
+The device name used is:
+
+- `N150` for N150
+- `N300` for N300
+- `T3K` for LoudBox / QuietBox
+- `TG` for Galaxy
+
+By default tensor parallelism is used to run the model over all available chips. You can instead run on a smaller mesh either for testing or for performance reasons (for very small models the communication overhead of tensor parallelism may be larger than the performance gained). To use a smaller mesh, set `MESH_DEVICE` to one of the supported devices: `N150`, `N300`, `T3K` or `TG`.
+
+Example: `export MESH_DEVICE=N150`, will enable running one a single chip of a multi-chip system.
+
+---
+
+### Parameters Description
+
+The `simple_text_demo.py` script includes the following main modes of operation and is parametrized to support other configurations.
+
+- `batch-1`: Runs a small prompt (128 tokens) for a single user
+- `batch-32`: Runs a small prompt (128 tokens) for a a batch of 32 users
+- `long-context`: Runs a large prompt (64k tokens) for a single user
+- `reasoning-1`: Runs a reasoning prompt for a single user (generates up to 15k tokens)
+
+If you want to provide your own demo configuration, please take a look at the pytest parametrize calls in `models/tt_transformers/demo/simple_text_demo.py`. For convenience we list all the supported params below:
+
+- `input_prompts (string)`: input json file with prompts to process. See `models/tt_transformers/demo/*.json` for a list of input files
+- `instruct (bool)`: Whether to use Llama instruct weights or general weights
+- `repeat_batches (int)`: Number of consecutive batches of users to run (default: 1)
+- `max_seq_len (int)`: Maximum context length supported by the model (refer to the table above)
+- `batch_size (int)`: Number of users in a batch (Supports 1/2/4/8/16/32 batches)
+- `max_generated_tokens (int)`: Maximum number of tokens to generate for each user (Note that the users will stop generation before this limit if they reach a eos token)
+- `paged_attention (bool)`: Whether to use paged attention or default attention (vLLM support (WIP) requires paged attention)
+- `page_params (dict)`: Page parameters for paged attention - [`block_size`, `max_num_blocks`]. For smaller context lengths use `block_size=32` and `max_num_blocks=1024`, for larger context use block_size=64 and max_num_blocks=2048
+- `sampling_params (dict)`: Sampling parameters for decoding -[`temperature`, `top_p`]. If temperature is set to 0, argmax (greedy decode) is used.
+- `stop_at_eos (bool)`: Flag to stop decoding when the model generates an EoS token
+- `optimizations (ModelOptimizations)`: Optimization level to use for the model [`accuracy`, `performance`]. Applied uniformly across all decoders unless an override config exists in `models/tt_transformers/model_params/<model-name>`
+- `decoder_config_file (DecodersPrecision)`: Fine-grained optimization control that allows specifying a configuration file to set different settings for each decoder.
+
+Please note that using `argmax` with `batch_size > 1` or using `top-p` sampling with any batch size, these ops will be run on host. This is because those ops are not yet fully supported on device. A decrease in performance is expected when these configurations are enabled.
+
+---
+
+### Optimization overrides
+Some models require a unique set of optimizations defined in `models/tt_transformers/model_params/<model-name>`. To override the default optimizations, you can define files named `models/tt_transformers/tt/model_config/PERFORMANCE_DECODER_CONFIG_FILENAME` and `models/tt_transformers/tt/model_config/ACCURACY_DECODER_CONFIG_FILENAME` in the appropriate `models/tt_transformers/model_params/<model-name>` directory to override the `ModelOptimizations.performance` and `ModelOptimizations.accuracy` optimizations respectively. For example, to override the default "performance" optimizations for Llama3.1-8B-Instruct, a file named `performance_decoder_config.json` has been created in the `models/tt_transformers/model_params/Llama3.1-8B-Instruct` directory. The content to write in override files is described in [the custom optimizations section](#custom-optimizations). Optimizations are applied with the following prioritization:
+1. from override config (if it exists)
+2. from the `optimizations` argument
+
+---
+
+### Custom input arguments
+To facilitate testing different configurations, `simple_text_demo.py` supports argument overrides. The full list of overrides is included in `models/tt_transformers/demo/conftest.py`.
+
+An example usage where the `batch-1` test is modified to run with 16 users and keep generating tokens until 1024 are generated:
+
+```
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --batch_size 16 --max_generated_tokens 1024 --stop_at_eos 0
+```
+
+---
+
+### Custom optimizations
+To apply the same settings across all decoders, the `optimizations` argument can be used. `optimizations` offers a wide range of configurations for precision and math fidelity. The user can override the configurations of the data types of the weight tensors and activation tensors and the math fidelity of the kernels that works on those tensors, using the `--optimizations` argument on the command line. For example:
+
+```
+pytest models/tt_transformers/demo/simple_text_demo.py -k "accuracy and batch-1" --optimizations 'precision_cfg = {ff1_3: bfp4, ff2: bfp4, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: lofi, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}'
+```
+
+Please refer to [model_config.py](models/tt_transformers/tt/model_config.py) for the full list of supported key-value pairs in the `--optimizations` argument. Centralized performance and accuracy targets are defined in [models/model_targets.yaml](../model_targets.yaml). The `lt` utility can still export markdown snapshots for local Pareto analysis.
+
+To apply non-uniform settings across the decoders, the user can provide a JSON file using the `decoder_config_file` argument to specify the configuration for each decoder. For example
+
+```
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --decoder_config_file 'models/tt_transformers/demo/config_16_decoders.json'
+```
+
+When a component is not specified (e.g., FF2 is missing for decoder 2 in `models/tt_transformers/demo/config_16_decoders.json`), the baseline configuration is used for that component.
+
+---
+
+### Expected performance and accuracy
+
+See [models/model_targets.yaml](../model_targets.yaml) for expected performance and accuracy targets across supported configurations.
+Accuracy of the network architectures is measured by exact token matching using teacher forcing method. During inference the previous token is replaced by the ground truth token while the network generates the next token. This allows to avoid accumulating errors when comparisons on a finer level (tokens) assessed in comparison to other known metrics that compare quality and context of the answer. Token accuracy can be reported by passing the argument shown below:
+
+```
+pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --token_accuracy True
+```
+
+---
+
+### Implementation notes
+
+**Chunked prefill (text-only)**: All of the compatible model/device combinations support a max prefill context-length of 128k, with the exception of Llama3.1-8B and Llama3.2-11B on N150 which have a max of 32k (due to a lack of memory). To support these large max context-lengths, chunked prefill is performed with different max chunk sizes as shown in the table below.
+
+Max Prefill Chunk Sizes (text-only):
+|              |      N150     |      N300     |      T3K       |      TG     |
+|--------------|---------------|---------------|----------------|-------------|
+| Llama3.2-1B  | 128k tokens   | 128k tokens   | 128k tokens    | 128k tokens |
+| Llama3.2-3B  | 8k tokens     | 128k tokens   | 128k tokens    | 128k tokens |
+| Qwen2.5-7B | 4k tokens | 32k tokens | 128k tokens | 128k tokens |
+| Llama3.1-8B  | 4k tokens     | 64k tokens    | 128k tokens    | 128k tokens |
+| Llama3.2-11B | 4k tokens     | 64k tokens    | 128k tokens    | 128k tokens |
+| Mistral-Small-3.1-24B | 8k tokens | 128k tokens | 128k tokens | 128k tokens |
+| Llama3.1-70B | Not supported | Not supported | 32k tokens     | 128k tokens |
+| Llama3.2-90B | Not supported | Not supported | 32k tokens     | Not supported |
+| DeepSeek-R1-Distill-Llama3.3-70B | Not supported | Not supported | 32k tokens | 128k tokens |
+
+
+- These max chunk sizes are specific to max context length 128k and are configured via `MAX_PREFILL_CHUNK_SIZES_DIV1024` in [model_config.py](https://github.com/tenstorrent/tt-metal/blob/main/models/demos/llama3/tt/model_config.py). If the max context length is set to a smaller value using the `max_seq_len` flag (see [Run the demo](#run-the-demo)), these chunk sizes can possibly be increased due to using a smaller KV cache.
+
+**Chunked prefill (Llama3.2-11B multimodal)**: Llama3.2-11B multimodal is currently only supported on N300 and T3000. On N300, a max prefill context length of 8k is supported, while T3000 supports a max context length of 128k.
+
+**Chunked prefill (Mistral-Small-3.1-24B multimodal)**: Mistral-Small-3.1-24B-Instruct-2503 (Pixtral vision) is currently supported on T3000. On T3000, a max prefill context length of 128k is supported.
+
+---
+
+### Memory Optimization
+
+#### HuggingFace Model Caching Control
+
+To help manage memory usage, you can control whether the HuggingFace model is cached in memory using the `cache_hf` parameter via command line or code:
+
+```python
+# Default: disables caching to conserve memory usage
+model_args = ModelArgs(
+    mesh_device,
+    cache_hf=False,  # Default: Reduces memory usage by not keeping HF model in memory
+    max_batch_size=1,
+    max_seq_len=2048
+)
+
+# Optional: enables caching for faster repeated access
+model_args = ModelArgs(
+    mesh_device,
+    cache_hf=True,  # Cache HF model for better performance running reference tests
+    max_batch_size=4,
+    max_seq_len=4096
+)
+```
+
+**When to disable caching (`cache_hf=False`):**
+- Running on systems with limited memory (< 256GB)
+- Loading large models (70B+ parameters)
+- Using the model for single inference runs
+- When you don't need reference model comparisons
+
+**When to keep caching enabled (`cache_hf=True`, default):**
+- Sufficient memory available
+- Comparisons with torch model is needed
+- Minimizing test duration is prioritized over memory usage
+- Running reference model tests
+
+The `cache_hf` parameter affects:
+- `load_state_dict()` method: Controls whether HF model is cached after loading
+- `reference_transformer()` method: Controls whether to reuse cached model or load fresh
+
+**Memory Impact:**
+- Disabling caching saves approximately the full model size in memory
+- For a 70B model, this can save ~140GB+ of memory usage
+- Increased test duration as model needs to be reloaded for reference operations

code/models/tt_transformers/conftest.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+
+import ttnn
+
+
+@pytest.fixture
+def device_params(request, galaxy_type):
+    # Get param dict passed in from test parametrize (or default to empty dict).
+    # Any TRACE_MODEL_KEY_PARAM is left in place; the mesh_device fixture resolves it
+    # to trace_region_size using the logical submesh SKU.
+    params = getattr(request, "param", {}).copy()
+
+    mesh_device = {"N150": (1, 1), "N300": (1, 2), "N150x4": (1, 4), "T3K": (1, 8), "TG": (8, 4), "P150x8": (1, 8)}.get(
+        os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+    )
+    is_single_device = (mesh_device == (1, 1)) if isinstance(mesh_device, tuple) else (mesh_device == 1)
+
+    if "fabric_config" in params:
+        if is_single_device:
+            params["fabric_config"] = None
+        elif params["fabric_config"] == True:
+            cluster_type = ttnn.cluster.get_cluster_type()
+            if cluster_type == ttnn.cluster.ClusterType.BLACKHOLE_GALAXY:
+                # The 8x4 decode path uses Ring collectives along both mesh axes.
+                params["fabric_config"] = ttnn.FabricConfig.FABRIC_2D_TORUS_XY
+            else:
+                params["fabric_config"] = (
+                    ttnn.FabricConfig.FABRIC_1D_RING if galaxy_type == "6U" else ttnn.FabricConfig.FABRIC_1D
+                )
+
+    return params

code/models/tt_transformers/host_mem_profiler.py

@@ -0,0 +1,188 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+host_mem_profiler.py — Host-side memory profiler for model tests
+===============================================================
+
+What this script does
+---------------------
+Wraps a pytest run with `memory_profiler` to record the host process
+memory (RSS) at 100 ms intervals while the model test executes.
+After the run it produces a PNG plot showing memory over time with a
+peak-memory reference line.  The script exits with pytest's own return
+code, so CI pipelines see failures correctly and no plot is saved for a
+failed run.
+
+Measured memory includes all child processes spawned by the test
+(e.g. the TT device runtime), giving a realistic view of total host
+footprint during inference.
+
+Dependencies
+------------
+`memory-profiler` and `matplotlib` are auto-installed at startup if
+absent.  Everything else must already be present in the active
+environment (loguru, pytest, and the tt-metal Python stack).
+
+    pip install memory-profiler matplotlib   # manual one-time install
+
+Arguments
+---------
+  --mesh-device   Value forwarded as the MESH_DEVICE environment variable
+                  to the test process (e.g. N150, N300, T3K, TG).
+                  Default: N150
+  --hf-model      Hugging Face model ID forwarded as HF_MODEL.
+                  Default: meta-llama/Llama-3.2-1B-Instruct
+  --test          Pytest target — a file path, directory, or node id.
+                  Default: models/tt_transformers/demo/simple_text_demo.py
+  -k              Optional pytest -k filter expression to select a subset
+                  of tests within the target.
+
+Output
+------
+A timestamped directory is created under profiling_results/:
+
+    profiling_results/<name>_<YYYYMMDD_HHMMSS>/
+        memory_profile.png   ← time-series plot (peak annotated in red)
+
+The script exits with pytest's return code; no plot is saved on failure.
+
+Usage examples
+--------------
+# Minimal — profile the default Llama-3.2-1B demo on a single N150 card:
+    python host_mem_profiler.py
+
+# Llama-3.1-8B on a T3K (8-chip) host, full test file:
+     python host_mem_profiler.py \\
+        --mesh-device T3K \\
+         --hf-model    meta-llama/Llama-3.1-8B-Instruct \\
+         --test        models/tt_transformers/demo/simple_text_demo.py
+
+# Same model but only the decode token accuracy test case:
+     python host_mem_profiler.py \\
+         --mesh-device T3K \\
+        --hf-model    meta-llama/Llama-3.1-8B-Instruct \\
+        --test        models/tt_transformers/demo/simple_text_demo.py \\
+        -k            performance-ci-token-mathcing
+"""
+
+import argparse
+import importlib
+import os
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+
+from loguru import logger
+
+_ALLOWED_PACKAGES = {"memory-profiler", "matplotlib"}
+
+
+def _require(package, import_name=None):
+    if package not in _ALLOWED_PACKAGES:
+        raise ValueError(f"Package '{package}' is not in the allowed list: {_ALLOWED_PACKAGES}")
+    import_name = import_name or package
+    try:
+        importlib.import_module(import_name)
+    except ImportError:
+        print(f"Installing missing dependency: {package}")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+        importlib.invalidate_caches()  # flush finder cache so the new package is visible
+
+
+_require("memory-profiler", "memory_profiler")
+_require("matplotlib")
+
+import matplotlib.pyplot as plt
+from memory_profiler import memory_usage
+
+PYTEST_EXTRA_ARGS = ["-v", "-s"]
+OUTPUT_ROOT = Path("profiling_results")
+
+
+# ── Runner ────────────────────────────────────────────────────────────────────
+def make_runner(model_cfg, k_filter=None):
+    """Return a zero-arg callable that runs a model's test suite."""
+
+    def run():
+        env = os.environ.copy()
+        env.update(model_cfg["env"])
+        cmd = [sys.executable, "-m", "pytest", model_cfg["test"], *PYTEST_EXTRA_ARGS]
+        if k_filter:
+            cmd += ["-k", k_filter]
+        result = subprocess.run(cmd, capture_output=False, env=env)
+        return result.returncode
+
+    return run
+
+
+# ── Plot ──────────────────────────────────────────────────────────────────────
+def plot_results(name, elapsed, mib, out_dir):
+    fig, ax = plt.subplots(figsize=(14, 6))
+    ax.plot(elapsed, mib, "+-k", linewidth=1.0, markersize=4, label=name)
+    ax.axhline(max(mib), color="red", linestyle="--", linewidth=0.8, label=f"Peak: {max(mib):.1f} MiB")
+
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel("Memory (MiB)")
+    ax.set_title("Memory Usage Over Time")
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(os.path.join(out_dir, "memory_profile.png"), dpi=150)
+    plt.close(fig)
+    logger.info(f"Saved at {os.path.join(out_dir, 'memory_profile.png')}")
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Profile host-side memory usage of a model test")
+    parser.add_argument("--mesh-device", default="N150", help="MESH_DEVICE value (e.g. N150, T3K)")
+    parser.add_argument("--hf-model", default="meta-llama/Llama-3.2-1B-Instruct", help="Hugging Face model ID")
+    parser.add_argument("--test", default="models/tt_transformers/demo/simple_text_demo.py", help="Pytest target path")
+    parser.add_argument("-k", default=None, help="pytest -k filter expression")
+    args = parser.parse_args()
+
+    name = args.hf_model.split("/")[-1] if "/" in args.hf_model else args.hf_model
+    if args.k:
+        name = f"{name}_{args.k}"
+
+    model = {
+        "name": name,
+        "env": {
+            "MESH_DEVICE": args.mesh_device,
+            "HF_MODEL": args.hf_model,
+        },
+        "test": args.test,
+    }
+
+    logger.info(f"\n{'='*60}\nProfiling: {model['name']}\n{'='*60}")
+
+    mem_ts, returncode = memory_usage(
+        (make_runner(model, args.k), [], {}),
+        interval=0.1,
+        retval=True,
+        timestamps=True,
+        include_children=True,
+    )
+    if returncode:
+        logger.error(f"pytest exited with code {returncode} — skipping plot")
+        sys.exit(returncode)
+
+    mib = [m for m, _ in mem_ts]
+    ts = [t for _, t in mem_ts]
+    t0 = ts[0]
+    elapsed = [t - t0 for t in ts]
+
+    logger.info(f"  Peak:     {max(mib):.1f} MiB")
+    logger.info(f"  Baseline: {min(mib):.1f} MiB")
+
+    safe_name = name.replace("/", "_").replace(" ", "_")
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_root = OUTPUT_ROOT.resolve()
+    out_dir = (output_root / f"{safe_name}_{timestamp}").resolve()
+    if not out_dir.is_relative_to(output_root):
+        raise ValueError(f"Refusing to write outside output root: {out_dir}")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    plot_results(model["name"], elapsed, mib, out_dir)

code/models/tt_transformers/scripts/op_perf_results.py

@@ -0,0 +1,190 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import csv
+from argparse import ArgumentParser
+from collections import defaultdict
+
+
+def main():
+    parser = ArgumentParser(
+        "Parse an op perf results CSV and show performance data using the min allgather time and max other time over devices, optionally only for a specific signpost region."
+    )
+    parser.add_argument("csv", help="Input CSV file")
+    parser.add_argument("--all", help="Show all times for each device", action="store_true")
+    parser.add_argument("--signpost", help="Only include data after this signpost and before any others")
+    parser.add_argument("--skip-last", help="Do not include timings from the last N ops", type=int, default=0)
+    parser.add_argument("--skip-first", help="Do not include timings from the first N ops", type=int, default=0)
+    parser.add_argument("--prefill", help="Prefill mode: will compute tok/s", action="store_true")
+    parser.add_argument("--seqlen", help="Sequence length used for prefill statistics.", type=int, default=0)
+    parser.add_argument(
+        "--estimate-full-model",
+        help="Estimate the full model performance by multiplying by N and adding back in the skipped ops",
+        type=int,
+        default=0,
+    )
+    parser.add_argument("--write-ops-to-csv", help="Write the summarized ops to a CSV file", type=str, default=None)
+    args = parser.parse_args()
+
+    header, rows = read_rows(args.csv)
+    blocks, signposts_seen = make_blocks(header, rows, args.signpost)
+
+    if args.signpost and not args.signpost in signposts_seen:
+        print(f'Error: signpost "{args.signpost}" was not found in this file')
+        print(f"Valid signposts are: {signposts_seen}")
+        return
+
+    print(f'{"Op":20} {"Time (us)"}')
+
+    if args.skip_first:
+        print(f"The following ops from the start of the run are not included in summary statistics:")
+        for block in blocks[: args.skip_first] if args.skip_first else blocks:
+            print(block.long_str() if args.all else block.short_str())
+        print(f"Ops included in the summary statistics:")
+        skipped_ops = blocks[: args.skip_first]
+        blocks = blocks[args.skip_first :]
+    else:
+        skipped_ops = []
+
+    for block in blocks[: -args.skip_last] if args.skip_last else blocks:
+        print(block.long_str() if args.all else block.short_str())
+
+    if args.skip_last:
+        print(f"The following ops from the end of the run are not included in summary statistics below:")
+        for block in blocks[-args.skip_last :]:
+            print(block.long_str() if args.all else block.short_str())
+        skipped_ops += blocks[-args.skip_last :]
+        blocks = blocks[: -args.skip_last]
+
+    total_time_ns = sum(block.time() for block in blocks)
+    total_time_s = total_time_ns / 1e9
+    tokens_per_s = 1 / total_time_s
+    if args.prefill:
+        sequences_per_s = tokens_per_s
+        tokens_per_s *= args.seqlen
+        print(f"Tokens/s: {tokens_per_s:.2f} ({total_time_s*1000*1000:.1f} us latency, {sequences_per_s:.2f} seq/s)")
+    else:
+        print(f"Tokens/s/user: {tokens_per_s:.2f} ({total_time_s*1000*1000:.1f} us latency)")
+
+    if args.estimate_full_model:
+        total_time_ns *= args.estimate_full_model
+        total_time_ns += sum(block.time() for block in skipped_ops)
+        total_time_s = total_time_ns / 1e9
+        tokens_per_s = 1 / total_time_s
+        if args.prefill:
+            sequences_per_s = tokens_per_s
+            tokens_per_s *= args.seqlen
+            print(
+                f"Estimated full model ({args.estimate_full_model} * above + skipped ops) tokens/s: {tokens_per_s:.2f} ({total_time_s*1000*1000:.1f} us latency, {sequences_per_s:.2f} seq/s)"
+            )
+        else:
+            print(
+                f"Estimated full model ({args.estimate_full_model} * above + skipped ops) tokens/s/user: {tokens_per_s:.2f} ({total_time_s*1000*1000:.1f} us latency)"
+            )
+
+    if signposts_seen and not args.signpost:
+        print(f"Warning - this file contains the following signposts that were not used for this analysis:")
+        for s in signposts_seen:
+            print(f'   "{s}"')
+        print("Rerun with --signpost to show only the performance for a specific signpost region")
+
+    if args.write_ops_to_csv:
+        write_blocks_to_csv(blocks, args.write_ops_to_csv)
+
+    return tokens_per_s
+
+
+def read_rows(csv_file):
+    with open(csv_file, "r") as f:
+        reader = csv.reader(f)
+        header = next(reader)
+        rows = list(reader)
+    return header, rows
+
+
+class Block:
+    def __init__(self, op_name, times):
+        self.op_name = op_name
+        self.times = times
+
+    def time(self):
+        return min(self.times) if "AllGather" in self.op_name or "ReduceScatter" in self.op_name else max(self.times)
+
+    def short_str(self):
+        short_name = self.op_name.split("::")[-1].split(")")[0]
+        time_range = max(self.times) - min(self.times)
+        return f"{short_name:20} {self.time()/1000:-6.0f} ± {time_range/1000:-5.0f}"
+
+    def long_str(self):
+        short_name = self.op_name.split("::")[-1].split(")")[0]
+        return f"{short_name:20} {self.time()/1000:-6.0f} <-" + " | ".join(f"{t/1000:-5.0f}" for t in self.times)
+
+    def __repr__(self):
+        return f"Block({self.op_name}, {self.times})"
+
+
+def make_blocks(header, rows, signpost):
+    """Perf dumps have one row per device in order, repeated for each op
+    This returns a list of blocks, where each block has an op name
+    and a list of times for each device.
+    """
+
+    # group rows by device then merge them together
+    block_by_device = defaultdict(list)
+    stop_on_signpost = False
+    signposts_seen = []
+
+    OP_CODE = header.index("OP CODE")
+    OP_TYPE = header.index("OP TYPE")
+    DEVICE_ID = header.index("DEVICE ID")
+    FW_DURATION = header.index("DEVICE FW DURATION [ns]")
+
+    block_op_name = None
+    for row in rows:
+        op_name = row[OP_CODE]
+        op_type = row[OP_TYPE]
+
+        if op_type == "signpost":
+            signposts_seen.append(op_name)
+            if stop_on_signpost:
+                break
+            elif op_name == signpost:
+                # clear any previous data and stop on the next signpost
+                stop_on_signpost = True
+                block_by_device = defaultdict(list)
+        elif op_type == "tt_dnn_device":
+            device_id = int(row[DEVICE_ID])
+            time = int(row[FW_DURATION])
+            block_by_device[device_id].append(Block(op_name, [time]))
+
+    # merge each device block into a single block with all the device times,
+    # checking that the op name matches
+    # blocks_by_device is a dict of device_id -> Block
+    # we want to get a list of Block (with all device times)
+
+    device_ids = list(sorted(block_by_device.keys()))
+    merged_blocks = block_by_device[device_ids[0]]
+
+    for device_id in device_ids[1:]:
+        assert len(block_by_device[device_id]) == len(
+            merged_blocks
+        ), f"Device {device_id} has {len(block_by_device[device_id])} ops, expected {len(merged_blocks)} from previous devices"
+        for row, b in enumerate(block_by_device[device_id]):
+            assert (
+                b.op_name == merged_blocks[row].op_name
+            ), f"Op name mismatch at row {row}: device {device_id} has {b.op_name} != {merged_blocks[row].op_name}"
+            merged_blocks[row].times += b.times
+
+    return merged_blocks, signposts_seen
+
+
+def write_blocks_to_csv(blocks, csv_file):
+    with open(csv_file, "w") as f:
+        writer = csv.writer(f)
+        writer.writerow(["Op", "Time (us)"])
+        for block in blocks:
+            writer.writerow([block.op_name, block.time()])
+
+
+if __name__ == "__main__":
+    main()

code/models/tt_transformers/scripts/repack_weights_70b.py

@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Llama2-70B weights are saved as 8 sharded checkpoints. Loading weights for a
+single layer is slow since we load all 80 layers into memory to construct the
+model. This script repacks the weights into checkpoints chunked by layers to
+speed up development.
+"""
+import argparse
+import math
+import shutil
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+
+def layer_num(key):
+    if "layers" in key:
+        return int(key.split("layers.")[1].split(".")[0])
+    return 0
+
+
+def chunk_key(key, chunk_size):
+    """
+    Return the chunk number that a key should go into
+    """
+    chunk_id = layer_num(key) // chunk_size
+    print(f"Key: {key} -> chunk_id: {chunk_id}")
+    return chunk_id
+
+
+def repack(in_dir, out_dir, chunk_size, num_layers, hidden_size):
+    """
+    Repack llama2-70b weights into checkpoints chunked by layers.
+    Non-layer weights are saved in the first checkpoint.
+    """
+    num_chunks = math.ceil(num_layers / chunk_size)
+    print(f"Repacking {num_layers} layers into {num_chunks} chunks of size {chunk_size}")
+    checkpoints = sorted(Path(in_dir).glob("*.pth"))
+    merged_checkpoints = defaultdict(list)
+    assert len(checkpoints) > 0, f"no checkpoint files found in {in_dir}"
+    print(f"Loading {len(checkpoints)} checkpoint files")
+    for ckpt in tqdm(checkpoints):
+        print(f"Checkpoint file: {ckpt}")
+        loaded_ckpt = torch.load(ckpt, map_location="cpu")
+        for key, value in loaded_ckpt.items():
+            merged_checkpoints[key].append(value)
+
+    # concat checkpoint values
+    chunks = [dict() for _ in range(num_chunks)]
+    for key, value in merged_checkpoints.items():
+        if len(value) == 1 or "norm" in key:
+            val = value[0]
+        else:
+            if (key == "tok_embeddings.weight" or key == "output.weight") and value[0].shape[1] == hidden_size:
+                # Concatenate along dimension 0 for llama3 token embeddings weight and lm head
+                val = torch.cat(value, dim=0)
+            else:
+                # cat_dim is index of the smallest dimension in value[0].shape
+                cat_dim = torch.argmin(torch.tensor(value[0].shape))
+                val = torch.cat(value, dim=cat_dim)
+
+        chunk_id = chunk_key(key, chunk_size)
+        chunks[chunk_id][key] = val
+
+    # save chunks and copy params.json if needed
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    params_file = Path(in_dir) / "params.json"
+    if params_file.exists() and not (out_dir / "params.json").exists():
+        shutil.copy(params_file, out_dir)
+        print(f"Copied params.json to {out_dir}")
+    for i, chunk in enumerate(chunks):
+        # each chunk file name should tell which layers are in it
+        start_layer = i * chunk_size
+        end_layer = (i + 1) * chunk_size - 1
+        end_layer = min(end_layer, num_layers - 1)
+        out_file = out_dir / f"layers_{start_layer}-{end_layer}.pth"
+        torch.save(chunk, out_file)
+        print(f"Saved {out_file}")
+
+
+if __name__ == "__main__":
+    # Take in command line arguments
+    parser = argparse.ArgumentParser(description="Repack llama2-70b weights")
+    parser.add_argument("in_dir", type=str, help="input directory")
+    parser.add_argument("out_dir", type=str, help="output directory")
+    parser.add_argument("chunk_size", type=int, default=10, help="number of layers per chunk")
+    parser.add_argument("-n", "--num_layers", type=int, default=80, help="total number of layers")
+    parser.add_argument("-hs", "--hidden_size", type=int, default=8192, help="hidden size of the model")
+    args = parser.parse_args()
+    repack(args.in_dir, args.out_dir, args.chunk_size, args.num_layers, args.hidden_size)

code/models/tt_transformers/scripts/repack_weights_90b.py

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Llama-3.2-90B weights are saved as 8 sharded checkpoints. Loading weights for a
+single layer is slow since we load all layers into memory to construct the
+model. This script repacks the weights into checkpoints chunked by layers to
+speed up development.
+"""
+import argparse
+import asyncio
+import json
+import math
+import shutil
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+from models.tt_transformers.tt.load_checkpoints import is_param_replicated_across_shards
+
+
+def layer_num(key):
+    if "layers" in key:
+        return int(key.split("layers.")[1].split(".")[0])
+    return -1
+
+
+def chunk_key(key, chunk_size):
+    """
+    Return the chunk number that a key should go into
+    """
+    layer_id = layer_num(key)
+    assert layer_id >= 0, f"Unexpected key {key}"
+    chunk_id = layer_id // chunk_size
+    print(f"Key: {key} -> chunk_id: {chunk_id}")
+    return chunk_id
+
+
+def get_unified_tensor(key, value, hidden_size):
+    res = None
+    if len(value) == 1 or is_param_replicated_across_shards(key):
+        res = value[0]
+    else:
+        if key.endswith("tok_embeddings.weight") or key.endswith("output.weight"):
+            assert value[0].shape[1] == hidden_size
+            res = torch.cat(value, dim=0)
+        else:
+            cat_dim = torch.argmin(torch.tensor(value[0].shape))
+            res = torch.cat(value, dim=cat_dim)
+
+    assert res is not None, f"Failed to unify tensor for key {key}"
+    return res
+
+
+def copy_file_if_no_exist(src_path: Path, dst_path: Path, file_name: str) -> None:
+    src_file = src_path / file_name
+    if src_file.exists() and not (dst_path / file_name).exists():
+        shutil.copy(src_file, dst_path)
+        print(f"Copied {file_name} to {dst_path}")
+
+
+async def torch_save_async(chunk, file_full_path):
+    loop = asyncio.get_running_loop()
+    await loop.run_in_executor(None, torch.save, chunk, file_full_path)
+
+
+async def repack(in_dir, out_dir, chunk_size, stop_after: int = None):
+    """
+    Repack llama3.2-90b weights into checkpoints chunked by layers.
+    Non-layer weights are saved in the first checkpoint.
+
+    Args:
+        in_dir: input directory containing llama3.2-90b weights from Meta
+        out_dir: output directory to save the chunked checkpoints
+        chunk_size: number of layers per chunk
+        stop_at: stop repacking at this many chunks
+    """
+    assert stop_after is None or stop_after > 0, f"Invalid stop_at value: {stop_after}"
+
+    # load model params
+    params_file = Path(in_dir) / "params.json"
+    assert params_file.exists(), f"params.json not found in {in_dir}"
+    with open(params_file, "r") as f:
+        params = json.load(f)
+    num_layers = params["n_layers"]
+    hidden_size = params["dim"]
+
+    # chunk the vision_model and the first FIVE decoder layers into the first checkpoint
+    # the rest of the decoder layers are chunked based on chunk_size
+
+    # first load the Meta checkpoints
+    checkpoints = sorted(Path(in_dir).glob("*.pth"))
+    merged_checkpoints = defaultdict(list)
+    assert len(checkpoints) > 0, f"no checkpoint files found in {in_dir}"
+    print(f"Loading {len(checkpoints)} checkpoint files:")
+    for ckpt in tqdm(checkpoints, leave=True):
+        tqdm.write(f"Checkpoint file: {ckpt}")
+        loaded_ckpt = torch.load(ckpt, map_location="cpu")
+        for key, value in loaded_ckpt.items():
+            merged_checkpoints[key].append(value)
+
+    # next we iterate over the merged checkpoints and get all the vision model tensors,
+    # the first decoder layer tensors, and all the non-layer tensors
+    num_decoder_layers_in_first_chunk = 1
+    chunk = {}
+    for key in list(merged_checkpoints.keys()):
+        if (
+            key.startswith("vision_model")
+            or layer_num(key) in range(num_decoder_layers_in_first_chunk)
+            or "layers." not in key
+        ):
+            chunk[key] = get_unified_tensor(key, merged_checkpoints[key], hidden_size)
+            del merged_checkpoints[key]
+
+    save_tasks = []
+    # save the first chunk
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    copy_file_if_no_exist(Path(in_dir), out_dir, "params.json")
+    copy_file_if_no_exist(Path(in_dir), out_dir, "tokenizer.model")
+    out_file = out_dir / f"vision-model-and-layers_{0}-{num_decoder_layers_in_first_chunk - 1}.pth"
+    save_tasks.append(asyncio.create_task(torch_save_async(chunk, out_file)))
+    print(f"Saved the following layers in {out_file}:")
+    for key in chunk.keys():
+        print("\t" + key)
+    del chunk
+
+    if stop_after is not None and stop_after == 1:
+        await wait_with_progress(save_tasks, desc="Writing chunked checkpoints to files")
+        return  # early return to stop at the first chunk
+
+    # save the rest of the merged checkpoints into chunks
+    num_chunks = math.ceil((num_layers - num_decoder_layers_in_first_chunk) / chunk_size)
+    # set stop_after to num_chunks if it is None, which means repacking all layers
+    stop_after = num_chunks if stop_after is None else stop_after - 1  # [INFO] -1 because already saved the 1st chunk
+
+    chunks = [list() for _ in range(num_chunks)]
+    for key in merged_checkpoints.keys():
+        assert key.startswith("text_model"), f"Unexpected key: {key}"
+        layer_id = layer_num(key)
+        assert layer_id != -1, f"Unexpected key: {key}"
+        chunk_id = (layer_id - num_decoder_layers_in_first_chunk) // chunk_size  # the first few layers is already saved
+        chunks[chunk_id].append(key)
+
+    print(f"Repacking {num_layers} layers into {num_chunks} chunks of size {chunk_size}")
+    for chunk_id in tqdm(range(num_chunks)):
+        if chunk_id >= stop_after:
+            break
+
+        chunk = {}
+        for key in chunks[chunk_id]:
+            chunk[key] = get_unified_tensor(key, merged_checkpoints[key], hidden_size)
+            del merged_checkpoints[key]
+
+        # save the chunk
+        start_layer = chunk_id * chunk_size + num_decoder_layers_in_first_chunk
+        end_layer = (chunk_id + 1) * chunk_size + num_decoder_layers_in_first_chunk - 1
+        end_layer = min(end_layer, num_layers - 1)
+        out_file = out_dir / f"layers_{start_layer}-{end_layer}.pth"
+        save_tasks.append(asyncio.create_task(torch_save_async(chunk, out_file)))
+        print(f"Saving the following layers in {out_file}:")
+        for key in chunk.keys():
+            print("\t" + key)
+        del chunk
+
+    await wait_with_progress(save_tasks, desc="Writing chunked checkpoints to files")
+
+
+async def wait_with_progress(tasks, desc):
+    """Wait for tasks to finish, updating a progress bar as each completes."""
+    total = len(tasks)
+    with tqdm(total=total, desc=desc, leave=True) as pbar:
+        pending = set(tasks)
+        while pending:
+            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+            pbar.update(len(done))
+
+
+if __name__ == "__main__":
+    # Take in command line arguments
+    parser = argparse.ArgumentParser(description="Repack llama3.2-90b weights")
+    parser.add_argument("in_dir", type=str, help="input directory")
+    parser.add_argument("out_dir", type=str, help="output directory")
+    parser.add_argument("chunk_size", type=int, default=10, help="number of layers per chunk")
+    parser.add_argument(
+        "--stop_after", type=int, default=None, help="stop repacking after this many chunks are saved (default to all)"
+    )
+    args = parser.parse_args()
+
+    asyncio.run(repack(args.in_dir, args.out_dir, args.chunk_size, args.stop_after))

code/models/tt_transformers/tests/conftest.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import gc
+
+import pytest
+
+from models.tt_transformers.tt.model_config import parse_optimizations
+
+# transformers 5.x removed Cache.get_usable_length, but some trust_remote_code reference models
+# still call it -- e.g. microsoft/Phi-3-mini-128k-instruct's modeling_phi3.py does
+# `kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)`, which raises
+# AttributeError under transformers 5.10.2. For an unbounded cache (DynamicCache) the old method
+# simply returned get_seq_length(layer_idx), so restore it as that alias to keep those reference
+# models working. Scoped to DynamicCache only -- bounded caches had different (max-length) logic.
+try:
+    from transformers.cache_utils import DynamicCache
+
+    if not hasattr(DynamicCache, "get_usable_length"):
+
+        def _get_usable_length(self, new_seq_length=0, layer_idx=0):
+            return self.get_seq_length(layer_idx)
+
+        DynamicCache.get_usable_length = _get_usable_length
+except Exception:  # defensive: transformers cache internals may move
+    pass
+
+
+@pytest.fixture(autouse=True)
+def ensure_gc():
+    gc.collect()
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--optimizations",
+        action="store",
+        default=None,
+        type=parse_optimizations,
+        help="Precision and fidelity configuration diffs over default (i.e., accuracy)",
+    )
+
+    parser.addoption(
+        "--decoder_config_file",
+        action="store",
+        default=None,
+        type=str,
+        help="Provide a JSON file defining per-decoder precision and fidelity settings",
+    )
+    parser.addoption(
+        "--use_hf_rope",
+        action="store_true",
+        default=False,
+        help="Whether to use HF-style rope, if not passed, the default mllama will be used",
+    )

code/models/tt_transformers/tests/generate_reference_outputs.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import argparse
+import bz2
+import os
+
+import torch
+from loguru import logger
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+from models.tt_transformers.tt.model_config import ModelArgs
+
+
+def generate_reference_outputs(total_length, output_file, hf_model_name=None):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    if hf_model_name:
+        # HuggingFace path
+        tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
+        config = AutoConfig.from_pretrained(hf_model_name)
+        # Qwen only: add rope scaling to the config
+        # https://huggingface.co/Qwen/Qwen2.5-7B-Instruct#processing-long-texts
+        if "Qwen" in hf_model_name:
+            config.rope_scaling = {"factor": 4.0, "original_max_position_embeddings": 32768, "type": "yarn"}
+        model = AutoModelForCausalLM.from_pretrained(
+            hf_model_name, config=config, torch_dtype=torch.float32 if device == "cpu" else None, device_map="auto"
+        )
+        model.eval()
+        model_args = ModelArgs(mesh_device=None, cache_hf=True)
+
+    else:
+        # Original path - load reference model
+        model_args = ModelArgs(mesh_device=None, cache_hf=True)
+        model_args.max_seq_len = total_length
+        tokenizer = model_args.tokenizer
+        assert tokenizer is not None, "Tokenizer must be provided for non-dummy weights"
+
+    reference_model = model_args.reference_transformer(load_checkpoint=True, wrap=False)
+    reference_model.to(device)  # Move model to device
+    reference_model.eval()  # Set to evaluation mode
+    embd = reference_model.model.embed_tokens
+    embd.to(device)  # Move embedding to device
+
+    # Load the book text and encode tokens
+    current_file_path = os.path.abspath(__file__)
+    current_file_dir = os.path.dirname(current_file_path)
+    prompt_file = os.path.join(current_file_dir, "tale-of-two-cities.txt.bz2")
+
+    with bz2.open(prompt_file, "rt", encoding="utf-8") as f:
+        text = f.read()
+
+    # Encode text to tokens
+    encoded_tokens = model_args.encode_prompt(text, instruct=False)
+    encoded_tokens_tensor = torch.tensor(encoded_tokens, device=device).unsqueeze(0)  # Move to device
+
+    print(f"{'Progress':<15}{'Correct':<8}{'Actual':<15}{'Top 5 Predictions':<75}")
+    print("-" * 113)
+
+    # Initialize lists to store results
+    all_top1_correct = []
+    all_top5_correct = []
+    all_top5_tokens = []
+    segment_accuracies = []
+    chunk_size = 1024
+
+    with torch.no_grad():
+        for chunk_start in range(0, total_length - 1, chunk_size):
+            chunk_end = min(chunk_start + chunk_size, total_length)
+            # Get input and target chunks, ensuring they have matching lengths
+            chunk_tokens = encoded_tokens_tensor[:, chunk_start:chunk_end]
+            chunk_next_tokens = encoded_tokens[chunk_start + 1 : chunk_end + 1]
+            actual_chunk_size = min(len(chunk_tokens[0]), len(chunk_next_tokens))
+
+            # Trim input chunk if needed
+            chunk_tokens = chunk_tokens[:, :actual_chunk_size]
+
+            # Process chunk based on model type
+            chunk_tokens = chunk_tokens.to(device)
+            if hf_model_name:
+                outputs = model(chunk_tokens)
+                ref_output = outputs.logits
+            else:
+                pt_decode_input = embd(chunk_tokens).view(1, actual_chunk_size, -1)
+                ref_output = reference_model(pt_decode_input, start_pos=chunk_start)
+
+            # Compute top-5 predictions
+            probs = torch.softmax(ref_output, dim=-1)
+            _, chunk_top5_tokens = torch.topk(probs, k=5, dim=-1)  # Shape: [1, chunk_size, 5]
+            chunk_top5_tokens = chunk_top5_tokens.squeeze(0)  # Shape: [chunk_size, 5]
+
+            # Get next tokens tensor, ensuring same length as predictions
+            chunk_next_tokens_tensor = torch.tensor(chunk_next_tokens[:actual_chunk_size], device=device)
+
+            # Calculate correctness
+            chunk_top1_correct = chunk_top5_tokens[:, 0] == chunk_next_tokens_tensor
+            chunk_top5_correct = torch.any(chunk_top5_tokens == chunk_next_tokens_tensor.unsqueeze(1), dim=1)
+
+            # Store results
+            all_top1_correct.extend(chunk_top1_correct.tolist())
+            all_top5_correct.extend(chunk_top5_correct.tolist())
+            all_top5_tokens.append(chunk_top5_tokens)
+
+            # Print predictions for this chunk
+            for i in range(len(chunk_next_tokens)):
+                global_pos = chunk_start + i
+                next_token = chunk_next_tokens[i]
+
+                sanitize = lambda x: x.replace("\n", "").replace("\r", "").replace("\x0c", "")
+                actual_token = sanitize(tokenizer.decode([next_token]))
+                top5_tokens = [sanitize(tokenizer.decode([t.item()])) for t in chunk_top5_tokens[i]]
+                correct = "x" if chunk_top1_correct[i] else ("-" if chunk_top5_correct[i] else " ")
+                top5_str = " ".join(f"{t:<14}" for t in top5_tokens)
+
+                progress_str = f"{global_pos+1}/{total_length-1}"
+                print(f"{progress_str:<15}{correct:<8}{actual_token:<15}{top5_str}")
+
+                # Calculate and store segment accuracies every 100 tokens
+                if (global_pos + 1) % 100 == 0 or global_pos == total_length - 2:
+                    start_idx = (global_pos // 100) * 100
+                    end_idx = min(start_idx + 100, len(all_top1_correct))
+                    segment_top1_acc = sum(all_top1_correct[start_idx:end_idx]) / (end_idx - start_idx) * 100
+                    segment_top5_acc = sum(all_top5_correct[start_idx:end_idx]) / (end_idx - start_idx) * 100
+                    if len(segment_accuracies) <= global_pos // 100:
+                        segment_accuracies.append((segment_top1_acc, segment_top5_acc))
+
+        # Concatenate all top5 tokens into a single tensor
+        all_top5_tokens = torch.cat(all_top5_tokens, dim=0)  # Shape: [total_tokens, 5]
+
+    # Move tensors back to CPU before saving
+    data = {
+        "top5_tokens": all_top5_tokens.cpu(),
+        "reference_tokens": encoded_tokens_tensor[:, :total_length].clone().cpu(),
+    }
+
+    torch.save(data, output_file)
+    logger.info(f"Saved reference outputs to {output_file}")
+
+    # Print all segment accuracy summaries as a table
+    print("\nSegment Accuracy Summaries:")
+    print(f"{'Tokens':<15}{'Top-1 Accuracy':<20}{'Top-5 Accuracy':<20}")
+    print("-" * 55)
+    for i, (top1_acc, top5_acc) in enumerate(segment_accuracies):
+        start_token = i * 100 + 1
+        end_token = min((i + 1) * 100, total_length)
+        print(f"{f'{start_token}-{end_token}':<15}{f'{top1_acc:.2f}%':<20}{f'{top5_acc:.2f}%':<20}")
+
+    # Calculate overall accuracy
+    overall_top1_acc = sum(acc[0] for acc in segment_accuracies) / len(segment_accuracies)
+    overall_top5_acc = sum(acc[1] for acc in segment_accuracies) / len(segment_accuracies)
+    print("-" * 55)
+    print(f"{'Overall':<15}{f'{overall_top1_acc:.2f}%':<20}{f'{overall_top5_acc:.2f}%':<20}")
+
+
+# New main function with argparse
+def main():
+    parser = argparse.ArgumentParser(description="Generate reference outputs for LLaMA accuracy testing.")
+    parser.add_argument("--total_length", type=int, default=1024, help="Total length of tokens to process")
+    parser.add_argument(
+        "--output_file", type=str, default="reference_outputs.pt", help="Output file path for reference data"
+    )
+    parser.add_argument(
+        "--model", type=str, help="Optional: HuggingFace model name (e.g., 'meta-llama/Llama-3.1-8B-Instruct')"
+    )
+    args = parser.parse_args()
+
+    generate_reference_outputs(total_length=args.total_length, output_file=args.output_file, hf_model_name=args.model)
+
+
+if __name__ == "__main__":
+    main()

code/models/tt_transformers/tests/generate_reference_outputs.sh

@@ -0,0 +1,82 @@
+#!/bin/bash
+
+# Parse command line arguments
+TOTAL_LENGTH=1024  # Default value
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --total-length)
+            TOTAL_LENGTH="$2"
+            shift 2
+            ;;
+        --help|-h)
+            echo "Usage: $0 [OPTIONS]"
+            echo
+            echo "Generate reference outputs for Llama models"
+            echo
+            echo "Options:"
+            echo "  --total-length N    Set the total sequence length (default: 1024)"
+            echo "  --help, -h          Show this help message"
+            exit 0
+            ;;
+        *)
+            echo "Unknown option: $1"
+            echo "Use --help to see available options"
+            exit 1
+            ;;
+    esac
+done
+
+# Define model directories from environment variables with fallbacks
+HF_MODELS=(
+    "${LLAMA_32_1B_DIR:-meta-llama/Llama-3.2-1B-Instruct}"
+    "${LLAMA_32_3B_DIR:-meta-llama/Llama-3.2-3B-Instruct}"
+    "${LLAMA_31_8B_DIR:-meta-llama/Llama-3.1-8B-Instruct}"
+    "${LLAMA_32_11B_DIR:-meta-llama/Llama-3.2-11B-Vision-Instruct}"
+    "${LLAMA_33_70B_DIR:-meta-llama/Llama-3.3-70B-Instruct}"
+    "${LLAMA_32_90B_DIR:-meta-llama/Llama-3.2-90B-Vision-Instruct}"
+    "${QWEN_25_7B_DIR:-Qwen/Qwen2.5-7B-Instruct}"
+    "${QWEN_25_72B_DIR:-Qwen/Qwen2.5-72B-Instruct}"
+    "${QWEN_25_32B_DIR:-Qwen/Qwen2.5-32B-Instruct}"
+    "${MIXTRAL_8X7B_DIR:-mistralai/Mixtral-8x7B-Instruct-v0.1}"
+    "${QWEN_25_CODER_32B_DIR:-Qwen/Qwen2.5-Coder-32B-Instruct}"
+)
+
+# Create reference_outputs directory if it doesn't exist
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT_DIR="${SCRIPT_DIR}/reference_outputs"
+mkdir -p "$OUTPUT_DIR"
+
+# Function to get model name from directory path
+get_model_name() {
+    local dir_name=$(basename "$1")
+    # If the path ends in /repacked, use the parent directory name instead
+    if [ "$dir_name" = "repacked" ]; then
+        dir_name=$(basename "$(dirname "$1")")
+    fi
+    echo "$dir_name"
+}
+
+# Loop through each LLAMA directory
+for DIR in "${HF_MODELS[@]}"; do
+    # TBD: do check using HF_HOME
+    # if [ ! -d "$DIR" ]; then
+    #     echo "Warning: Directory $DIR does not exist, skipping..."
+    #     continue
+    # fi
+
+    # Get model size for output filename
+    MODEL_NAME=$(get_model_name "$DIR")
+    OUTPUT_FILE="${OUTPUT_DIR}/${MODEL_NAME}_full.refpt"
+
+    echo "Generating reference outputs for ${MODEL_SIZE} model..."
+    echo "Using weights from: ${DIR}"
+    echo "Output will be saved to: ${OUTPUT_FILE}"
+
+    # Set HF_MODEL environment variable and run the Python script
+    HF_MODEL="$DIR" python3 "${SCRIPT_DIR}/generate_reference_outputs.py" \
+        --total_length "$TOTAL_LENGTH" \
+        --output_file "$OUTPUT_FILE" \
+        --model "$DIR"
+done
+
+echo "All reference outputs have been generated!"

code/models/tt_transformers/tests/test_attention_prefill.py

@@ -0,0 +1,277 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+from models.tt_transformers.tests.test_utils import get_ref_model_dype
+from models.tt_transformers.tt.attention import Attention
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.common import Mode, PagedAttentionConfig, get_rot_transformation_mat, precompute_freqs
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.prefetcher import Prefetcher
+from models.tt_transformers.tt.rope import get_rot_mats, get_rot_mats_hf
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+# Model and attention prefill tests should run both with and without paged attention to debug any issues that may occur with default attention
+@pytest.mark.parametrize(
+    "paged_attention",
+    (
+        True,
+        False,
+    ),
+    ids=(
+        "paged_attention",
+        "default_attention",
+    ),
+)
+@pytest.mark.parametrize(
+    "page_params",
+    [{"page_block_size": 32, "page_max_num_blocks": 1024}],
+)
+@pytest.mark.parametrize(
+    "max_seq_len",
+    (
+        256,  # 4096,
+        # 1024 * 32,
+        # 1024 * 64,
+    ),
+)
+@pytest.mark.parametrize(
+    "use_prefetcher",
+    ([False]),
+)
+@pytest.mark.parametrize("use_hf_rope", (True, False), ids=("hf_rope", "mllama_rope"))
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_attention_inference(
+    max_seq_len,
+    paged_attention,
+    page_params,
+    mesh_device,
+    use_hf_rope,
+    reset_seeds,
+    ensure_gc,
+    use_prefetcher,
+):
+    dtype = ttnn.bfloat8_b
+    pcc = 0.99
+    batch_size = 1  # For prefill we only support batch_size = 1
+
+    # In prefill mode, we do not use prefetcher but we test the prefetcher interface for completeness and
+    num_tensors = 0
+    prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None
+    if use_prefetcher:
+        prefetcher.init(mode=Mode.PREFILL)
+
+    model_args = ModelArgs(
+        mesh_device, max_batch_size=batch_size, max_seq_len=max_seq_len, cache_hf=True, use_hf_rope=use_hf_rope
+    )
+    model_args.n_layers = 1
+    state_dict = model_args.load_state_dict()
+
+    # Ref model needs partial state dict, but our models use full state dict keys as cached weight names
+    first_layer_prefix = model_args.get_state_dict_prefix("Attention", 0) + "."
+    partial_state_dict = {
+        k[len(first_layer_prefix) :]: v for k, v in state_dict.items() if (k.startswith(first_layer_prefix))
+    }
+    reference_model = model_args.reference_attention(load_checkpoint=True)
+
+    rot_mats_fn = get_rot_mats_hf if model_args.use_hf_rope else get_rot_mats
+
+    # pre-compute the rotational embedding matrix and send to device
+    rot_mats = rot_mats_fn(
+        head_dim=model_args.head_dim,
+        device=mesh_device,
+        seq_len=max_seq_len,
+        theta=model_args.rope_theta,
+        rope_scaling=model_args.rope_scaling,
+    )
+
+    transformation_mats = {}
+    if not model_args.use_hf_rope:
+        transformation_mat_torch = get_rot_transformation_mat(model_args.head_dim)
+        transformation_mats_prefill = ttnn.as_tensor(
+            transformation_mat_torch,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+        transformation_mats = {"prefill": transformation_mats_prefill}
+
+    generation_start_pos = 0
+    generation_length = 3
+    all_tests_pass = True
+
+    # Setup page table
+    page_table_tt = None
+    paged_attention_config = None
+
+    if paged_attention:
+        paged_attention_config = PagedAttentionConfig(
+            block_size=page_params["page_block_size"],
+            max_num_blocks=page_params["page_max_num_blocks"],
+        )
+        # Implied shuffling of blocks
+        permutation = torch.randperm(paged_attention_config.max_num_blocks)
+        # Page table which maps virtual blocks to physical
+        reverse_permutation = torch.argsort(permutation)
+        page_table = reverse_permutation.reshape(
+            model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size
+        )
+        page_table_tt = ttnn.from_torch(
+            page_table,
+            device=mesh_device,
+            dtype=ttnn.int32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+
+    tt_ccl = TT_CCL(mesh_device)
+    tt_model = Attention(
+        mesh_device,
+        tt_ccl,
+        model_args,
+        state_dict,
+        weight_cache_path=model_args.weight_cache_path(dtype),
+        layer_num=0,
+        dtype=dtype,
+        transformation_mats=transformation_mats,
+        configuration=model_args,
+        paged_attention_config=paged_attention_config,
+        prefetcher=prefetcher,
+    )
+
+    pt_attention_input = (
+        torch.rand(
+            batch_size, max_seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
+        )
+        * 2
+    ) - 1
+    tt_attention_input = pt_attention_input.clone()
+    attention_input = model_args.prepare_residual_tensor_prefill(
+        tt_attention_input,
+        force_replicated=False if model_args.is_galaxy else True,
+    )
+
+    tt_out = tt_model(
+        attention_input,
+        current_pos=None,
+        rot_mats=rot_mats,
+        user_id=0,
+        mode=Mode.PREFILL,
+        page_table=page_table_tt,
+    )
+    tt_out = ttnn.to_torch(
+        tt_out, mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape)
+    )
+    tt_output_torch = tt_out[:, 0:1, :, : model_args.dim].view(batch_size, max_seq_len, -1)  # [ batch, seq, hidden_dim]
+    positions = torch.LongTensor(range(max_seq_len))
+
+    cos, sin = precompute_freqs(
+        model_args.head_dim,
+        model_args.max_seq_len * 2,
+        model_args.rope_theta,
+        model_args.rope_scaling.factor if model_args.rope_scaling else None,
+        model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None,
+        model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3",
+    )
+    freqs_cis_i = torch.complex(cos, sin)[positions]
+
+    attn_mask = torch.full((max_seq_len, max_seq_len), torch.finfo(torch.float32).min)
+    attn_mask_torch = torch.triu(attn_mask, diagonal=1)
+    reference_output = reference_model(pt_attention_input, positions[0], freqs_cis_i, mask=attn_mask_torch)
+
+    passing, pcc_message = comp_pcc(reference_output, tt_output_torch, pcc)
+
+    logger.info(comp_allclose(reference_output, tt_output_torch))
+    logger.info(f"PCC: {pcc_message}")
+    if passing:
+        logger.info(f"Attention Passed!")
+    else:
+        logger.warning(f"Attention Failed!")
+        all_tests_pass = False
+
+    check_kv_cache = True  # May want to disable: Issue #10648
+    if check_kv_cache:
+        # PyTorch output --------------------------------------------------------------------
+        pytorch_layer_present = [
+            reference_model.cache_k.clone().permute(0, 2, 1, 3),  # [batch_size, n_kv_heads, seq, head_dim]
+            reference_model.cache_v.clone().permute(0, 2, 1, 3),  # [batch_size, n_kv_heads, seq, head_dim]
+        ]
+        # TT hardware execution -------------------------------------------------------------
+        if paged_attention:
+            tt_layer_present = [
+                (
+                    ttnn.to_torch(
+                        cache,
+                        mesh_composer=ttnn.ConcatMesh2dToTensor(
+                            mesh_device,
+                            dims=(1, 3) if model_args.is_galaxy else (0, 1),
+                            mesh_shape=model_args.cluster_shape,
+                        ),
+                    )[reverse_permutation][:, : model_args.n_kv_heads, :, : model_args.head_dim]
+                    .reshape(
+                        model_args.max_batch_size,
+                        paged_attention_config.max_num_blocks // model_args.max_batch_size,
+                        model_args.n_kv_heads,
+                        paged_attention_config.block_size,
+                        model_args.head_dim,
+                    )
+                    .transpose(1, 2)
+                    .reshape(model_args.max_batch_size, model_args.n_kv_heads, -1, model_args.head_dim)[
+                        :batch_size, ...
+                    ]
+                )
+                for cache in tt_model.layer_past
+            ]
+        else:
+            tt_layer_present = [
+                ttnn.to_torch(
+                    cache,
+                    mesh_composer=ttnn.ConcatMesh2dToTensor(
+                        mesh_device,
+                        dims=(1, 0) if model_args.is_galaxy else (0, 1),
+                        mesh_shape=model_args.cluster_shape,
+                    ),
+                )[:batch_size, :, :, :]
+                for cache in tt_model.layer_past
+            ]
+
+        for i, (cache_pt, cache_tt) in enumerate(zip(pytorch_layer_present, tt_layer_present)):
+            cache_length_to_check = min(model_args.max_seq_len, generation_start_pos + generation_length + 1)
+            cache_pt = cache_pt[:, :, generation_start_pos:cache_length_to_check, :]
+            cache_tt = cache_tt[:, :, generation_start_pos:cache_length_to_check, :]
+            does_pass, output_pcc = comp_pcc(cache_pt, cache_tt, pcc)
+            if i == 0:
+                logger.info(f"K cache output: {output_pcc}")
+            else:
+                logger.info(f"V cache output: {output_pcc}")
+
+            if does_pass:
+                logger.info(f"KV Cache Passed!")
+            else:
+                logger.warning(f"KV Cache Failed! PCC value is lower than {pcc}")
+                all_tests_pass = False
+
+    if all_tests_pass:
+        logger.info("Attention output Passed!")
+    else:
+        logger.warning("Attention output Failed!")
+        assert all_tests_pass, f"PCC value is lower than {pcc} for some of the outputs. Check Warnings!"

code/models/tt_transformers/tests/test_batched_prefill_slots.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Batched prefill lays its device rows out by physical slot, not by prefill position.
+
+``empty_slots[i]`` is the device slot that owns request ``i``'s per-slot state, and
+every slot-indexed buffer the batched path builds (``prefill_ids``,
+``padded_last_token_idx``, the padded page table) is bounded by ``padded_batch``. vLLM
+hands out the slot a request already owns, so a batch of N requests can land on slots
+above N and the device batch has to span them. The arrays handed back to the caller
+stay in prefill order, so the readback reads by slot and writes by position. Pure host
+index bookkeeping, no device execution.
+"""
+
+import torch
+
+from models.common.sampling import SamplingParams, slice_sampling_params
+from models.common.sampling.tt_log_probs import LogProbsResult
+from models.tt_transformers.tt.generator import batched_prefill_padded_batch, gather_batched_prefill_samples
+
+
+def test_dense_slots_keep_todays_batch_shape():
+    """The common case is unchanged, so existing shapes and traces are reused."""
+    assert batched_prefill_padded_batch(7, list(range(7)), 32) == 8
+    assert batched_prefill_padded_batch(2, [0, 1], 32) == 2
+    assert batched_prefill_padded_batch(32, list(range(32)), 32) == 32
+
+
+def test_batch_spans_the_highest_slot_in_use():
+    """A live off-batch request holding a low slot pushes a prefill onto a high one."""
+    # THE BUG: seven requests whose slots reach 7. The count-based rule returned 8,
+    # which is fine here, but a request on slot 20 got a 1-row batch.
+    assert batched_prefill_padded_batch(7, [0, 1, 2, 3, 4, 5, 7], 32) == 8
+    assert batched_prefill_padded_batch(1, [20], 32) == 32
+    assert batched_prefill_padded_batch(3, [3, 4, 5], 32) == 8
+
+
+def test_no_slots_means_the_request_count_is_the_span():
+    """Callers that omit the slots get ``range(N)``, so N bounds the rows."""
+    assert batched_prefill_padded_batch(4, None, 32) == 4
+    assert batched_prefill_padded_batch(4, [], 32) == 4
+
+
+class _SlotTaggedLogProbs(LogProbsResult):
+    """Stands in for a device top-k result: reports which slot it was read from."""
+
+    def __init__(self):
+        super().__init__(topk_logprobs=None, topk_indices=None, topk_logprobs_host=None, topk_indices_host=None)
+
+    def extract_user(self, user_batch_idx: int):
+        return f"slot{int(user_batch_idx)}"
+
+
+def test_samples_come_back_in_prefill_order_not_slot_order():
+    """Read the device row by slot, write the caller's row by position.
+
+    Row i of the device batch holds slot i's sample, so a request on slot 5 has to
+    end up at output row 0 if it prefilled first.
+    """
+    slots = [5, 0, 3]
+    # Device rows: index == slot, so slot 5 sampled token 105, slot 0 token 100, ...
+    tokens_host = torch.tensor([100, 101, 102, 103, 104, 105, 106, 107])
+    plain_log_probs_host = torch.tensor([-0.0, -0.1, -0.2, -0.3, -0.4, -0.5, -0.6, -0.7])
+    output_tokens = torch.zeros(len(slots), 1, dtype=torch.int64)
+    output_log_probs = [None] * len(slots)
+
+    gather_batched_prefill_samples(slots, tokens_host, None, plain_log_probs_host, output_tokens, output_log_probs)
+
+    assert [int(t) for t in output_tokens.reshape(-1)] == [105, 100, 103]
+    assert [round(float(lp), 1) for lp in output_log_probs] == [-0.5, -0.0, -0.3]
+
+
+def test_a_slot_at_the_request_count_does_not_overflow_the_output():
+    """THE CRASH: three requests reaching slot 7 wrote past a 3-row output."""
+    slots = [0, 1, 7]
+    tokens_host = torch.tensor([200, 201, 202, 203, 204, 205, 206, 207])
+    output_tokens = torch.zeros(len(slots), 1, dtype=torch.int64)
+    output_log_probs = [None] * len(slots)
+
+    gather_batched_prefill_samples(slots, tokens_host, None, None, output_tokens, output_log_probs)
+
+    assert [int(t) for t in output_tokens.reshape(-1)] == [200, 201, 207]
+    assert output_log_probs == [None, None, None]
+
+
+def test_topk_logprobs_are_extracted_from_the_slot_row():
+    slots = [4, 1]
+    tokens_host = torch.arange(8)
+    output_tokens = torch.zeros(len(slots), 1, dtype=torch.int64)
+    output_log_probs = [None] * len(slots)
+
+    gather_batched_prefill_samples(slots, tokens_host, _SlotTaggedLogProbs(), None, output_tokens, output_log_probs)
+
+    assert output_log_probs == ["slot4", "slot1"]
+
+
+def test_slice_sampling_params_gives_each_chunk_its_own_requests():
+    """A chunked prefill must not hand every chunk the first N requests' params."""
+    params = SamplingParams(
+        temperature=[0.1, 0.2, 0.3, 0.4], top_k=[1, 2, 3, 4], top_p=[0.5, 0.6, 0.7, 0.8], seed=[11, 12, 13, 14]
+    )
+
+    second = slice_sampling_params(params, 2, 4)
+
+    assert second.temperature == [0.3, 0.4]
+    assert second.top_k == [3, 4]
+    assert second.top_p == [0.7, 0.8]
+    assert second.seed == [13, 14]
+    assert params.temperature == [0.1, 0.2, 0.3, 0.4]
+    assert slice_sampling_params(None, 0, 2) is None
+
+
+def test_a_span_no_bucket_covers_reports_the_span():
+    """The caller's ``> max_batch_size`` guard has to fire and pick sequential prefill.
+
+    Reporting ``max_batch_size`` instead would leave batched prefill enabled and
+    scatter into a row the buffers do not have.
+    """
+    assert batched_prefill_padded_batch(2, [40], 32) == 41
+    assert batched_prefill_padded_batch(2, [40], 32) > 32
+    # A wider model still covers the slot, so batching stays on as it did before.
+    assert batched_prefill_padded_batch(2, [40], 64) == 64

code/models/tt_transformers/tests/test_ci_dispatch.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+from loguru import logger
+
+from models.tt_transformers.tt.common import get_hf_tt_cache_path
+
+
+# This test will run all the nightly fast dispatch tests for all supported TTT models in CI [N150 / N300 only]
+@pytest.mark.parametrize(
+    "model_weights",
+    [
+        "meta-llama/Llama-3.2-1B-Instruct",
+        "meta-llama/Llama-3.2-3B-Instruct",
+        "meta-llama/Llama-3.1-8B-Instruct",
+        "meta-llama/Llama-3.2-11B-Vision-Instruct",
+        "mistralai/Mistral-7B-Instruct-v0.3",
+    ],
+    ids=[
+        "ttt-llama3.2-1B",
+        "ttt-llama3.2-3B",
+        "ttt-llama3.1-8B",
+        "ttt-llama3.2-11B",
+        "ttt-mistral-7B-v0.3",
+    ],
+)
+def test_ci_dispatch(model_weights):
+    logger.info(f"Running fast dispatch tests for {model_weights}")
+
+    os.environ["HF_MODEL"] = model_weights
+    os.environ["TT_CACHE_PATH"] = get_hf_tt_cache_path(model_weights)
+
+    # Pass the exit code of pytest to proper keep track of failures during runtime
+    exit_code = pytest.main(
+        [
+            "models/tt_transformers/tests/test_embedding.py",
+            "models/tt_transformers/tests/test_rms_norm.py",
+            "models/tt_transformers/tests/test_mlp.py",
+            "models/tt_transformers/tests/test_attention.py",
+            "models/tt_transformers/tests/test_attention_prefill.py",
+            "models/tt_transformers/tests/test_decoder.py",
+            "models/tt_transformers/tests/test_decoder_prefill.py",
+        ]
+        + ["-x"]  # Fail if one of the tests fails
+        + (["--timeout", "600"] if "mistral" in model_weights.lower() else [])
+    )
+    if exit_code == pytest.ExitCode.TESTS_FAILED:
+        pytest.fail(
+            f"One or more CI dispatch tests failed for {model_weights}. Please check the log above for more info",
+            pytrace=False,
+        )

code/models/tt_transformers/tests/test_decoder.py

@@ -0,0 +1,281 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+from models.tt_transformers.tests.test_utils import get_ref_model_dype
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.common import Mode, PagedAttentionConfig, precompute_freqs
+from models.tt_transformers.tt.decoder import TransformerBlock
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.prefetcher import Prefetcher
+from models.tt_transformers.tt.rope import HfRotarySetup, RotarySetup
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "use_prefetcher",
+    ([False]),
+)
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "paged_attention",
+    (
+        True,
+        # False
+    ),
+    ids=(
+        "paged_attention",
+        # "default_attention"
+    ),
+)
+@pytest.mark.parametrize(
+    "page_params",
+    [{"page_block_size": 32, "page_max_num_blocks": 1024}],
+)
+@pytest.mark.parametrize(
+    "batch_size",
+    (1, 32),
+)
+@pytest.mark.parametrize(
+    "max_seq_len",
+    (256,),  # For decode-only unit test, there's no need to run with large sequence lengths
+)
+@pytest.mark.parametrize(
+    "generation_length",
+    (10,),  # For decode-only unit test, there's no need to run with large sequence lengths
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_decoder_inference(
+    max_seq_len,
+    batch_size,
+    paged_attention,
+    page_params,
+    mesh_device,
+    reset_seeds,
+    ensure_gc,
+    generation_length,
+    use_prefetcher,
+):
+    dtype = ttnn.bfloat8_b
+
+    mode = Mode.DECODE
+    num_tensors = 5 if use_prefetcher else 0
+    prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None
+
+    if use_prefetcher:
+        prefetcher.init(mode=mode)
+
+    model_args = ModelArgs(
+        mesh_device,
+        max_batch_size=batch_size,
+        max_seq_len=max_seq_len,
+        cache_hf=True,
+        prefetcher=prefetcher,
+        use_hf_rope=False,
+    )
+    model_args.n_layers = 1
+
+    state_dict = model_args.load_state_dict()
+    reference_model = model_args.reference_decoder(load_checkpoint=True)
+
+    generation_start_pos = 0
+    all_tests_pass = True
+
+    # Setup RoPE transformation matrices
+    DefaultRopeSetup = HfRotarySetup if model_args.use_hf_rope else RotarySetup
+    rope_setup = DefaultRopeSetup(
+        mesh_device,
+        model_args.max_batch_size,
+        model_args.head_dim,
+        model_args.max_seq_len,
+        model_args.rope_theta,
+        model_args.rope_scaling,
+        model_args.use_qk_fused,
+        prefetcher=prefetcher,
+    )
+
+    if model_args.rope_theta_local is not None:
+        rope_setup_local = RotarySetup(
+            mesh_device,
+            model_args.max_batch_size,
+            model_args.head_dim,
+            model_args.max_seq_len,
+            model_args.rope_theta_local,
+            None,
+            # Match model.py: when fused QK is on, the local rope must also double
+            # its cos/sin batch so one tensor serves both Q and K (Gemma-2 sliding layers).
+            use_qk_fused=model_args.use_qk_fused,
+        )
+    else:
+        rope_setup_local = None
+
+    transformation_mats = rope_setup.get_both_trans_mats()
+
+    # Prepare page table for paged attention
+    page_table_tt = None
+    paged_attention_config = None
+
+    if paged_attention:
+        paged_attention_config = PagedAttentionConfig(
+            block_size=page_params["page_block_size"],
+            max_num_blocks=page_params["page_max_num_blocks"],
+        )
+        # Implied shuffling of blocks
+        permutation = torch.randperm(paged_attention_config.max_num_blocks)
+        # Page table which maps virtual blocks to physical
+        reverse_permutation = torch.argsort(permutation)
+        page_table = reverse_permutation.reshape(
+            model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size
+        )
+        page_table_tt = ttnn.from_torch(
+            page_table,
+            device=mesh_device,
+            dtype=ttnn.int32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(
+                mesh_device,
+                dims=(None, -2) if (model_args.is_galaxy and batch_size > 1) else (None, None),
+                mesh_shape=model_args.cluster_shape,
+            ),
+        )
+
+    # Initialize TT model
+    tt_ccl = TT_CCL(mesh_device)
+    tt_model = TransformerBlock(
+        args=model_args,
+        mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
+        dtype=dtype,
+        state_dict=state_dict,
+        layer_num=0,
+        weight_cache_path=model_args.weight_cache_path(dtype),
+        transformation_mats=transformation_mats,
+        paged_attention_config=paged_attention_config,
+        prefetcher=prefetcher,
+    )
+    if use_prefetcher:
+        tt_model.prefetcher.prefetch()
+
+    seqlen = 1
+
+    # Precompute freqs_cis for reference model
+    cos, sin = precompute_freqs(
+        model_args.head_dim,
+        model_args.max_seq_len * 2,
+        model_args.rope_theta,
+        model_args.rope_scaling.factor if model_args.rope_scaling else None,
+        model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None,
+        model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3",
+    )
+    freqs_cis = torch.complex(cos, sin)
+
+    # Initial positions
+    current_pos = torch.tensor([generation_start_pos for _ in range(batch_size)])
+    current_pos_tensor = ttnn.from_torch(
+        current_pos,
+        device=mesh_device,
+        dtype=ttnn.int32,
+        mesh_mapper=ttnn.ShardTensor2dMesh(
+            mesh_device,
+            dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None),
+            mesh_shape=model_args.cluster_shape,
+        ),
+    )
+    for i in range(generation_length):
+        logger.info(f"[Decoder] Generating token {i}")
+
+        if prefetcher is not None:
+            prefetcher.run()
+
+        # input = torch.randn(1, 32, 4096)
+        pt_decode_input = (
+            torch.rand(
+                batch_size, seqlen, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
+            )
+            * 2
+        ) - 1
+        tt_decode_input = pt_decode_input.clone()
+
+        decode_input = model_args.prepare_residual_tensor_decode(
+            tt_decode_input,
+            model_args.get_residual_mem_config(mode, prefetcher),
+        )
+
+        # Get cos/sin matrices for the current position of each user
+        rot_mats = rope_setup.get_rot_mats(current_pos)
+        rot_mats_local = None if rope_setup_local is None else rope_setup_local.get_rot_mats(current_pos)
+
+        # Run TT model
+        tt_out = tt_model(
+            decode_input,
+            current_pos_tensor,
+            rot_mats_global=rot_mats,
+            rot_mats_local=rot_mats_local,
+            mode=mode,
+            page_table=page_table_tt,
+        )
+
+        tt_out = ttnn.to_torch(
+            tt_out,
+            mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),
+        )
+
+        tt_output_torch = tt_out[:, 0:1, : model_args.max_batch_size, : model_args.dim].view(-1, 1, model_args.dim)
+
+        # In this test all users have the same position
+        freqs_cis_i = freqs_cis[current_pos[0], :].unsqueeze(0)
+
+        # Reference model
+        ref_output = reference_model(pt_decode_input, current_pos[0], freqs_cis_i, mask=None)
+        if ref_output.dim() == 2:
+            ref_output = ref_output.unsqueeze(1)
+
+        # For some model variants the HF decoder returns output only for the first batch item.
+        # Since all users share the same position in this test, compare the first ref_output.shape[0]
+        # items from TT output to ref_output.
+        batch_cmp = ref_output.shape[0]
+        tt_output_cmp = tt_output_torch[:batch_cmp]
+        passing, pcc_message = comp_pcc(ref_output, tt_output_cmp)
+
+        logger.info(comp_allclose(ref_output, tt_output_cmp))
+        logger.info(f"PCC: {pcc_message}")
+
+        if passing:
+            logger.info("Decoder Block Passed!")
+        else:
+            logger.warning("Decoder Block Failed!")
+            all_tests_pass = False
+
+        # Increment position
+        current_pos = torch.tensor([generation_start_pos + i + 1 for _ in range(batch_size)])
+        current_pos_tensor = ttnn.from_torch(
+            current_pos,
+            device=mesh_device,
+            dtype=ttnn.int32,
+            mesh_mapper=ttnn.ShardTensor2dMesh(
+                mesh_device,
+                dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None),
+                mesh_shape=model_args.cluster_shape,
+            ),
+        )
+
+    if all_tests_pass:
+        logger.info(f"All {generation_length} decode iterations Passed!")
+    else:
+        logger.warning("One or more iterations of decode Failed!")
+        assert all_tests_pass, f"PCC value is lower than {0.99} for some of the outputs. Check Warnings!"

code/models/tt_transformers/tests/test_device_perf.py

@@ -0,0 +1,362 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+
+import pandas as pd
+import pytest
+from loguru import logger
+
+from models.perf.benchmarking_utils import BenchmarkData, BenchmarkProfiler
+from models.perf.device_perf_utils import run_device_perf
+from models.tt_transformers.tests.test_utils import (
+    merge_device_rows,
+    print_dict,
+    process_measurements,
+    split_compile_and_trace,
+    verify_value_within_margin,
+)
+from tools.tracy.common import PROFILER_DEFAULT_OP_SUPPORT_COUNT
+from tools.tracy.process_model_log import get_latest_ops_log_filename
+
+
+# This pytest flag is necessary to ensure that we do NOT open the device in the main process for device perf tests that run
+# the test inside a subprocess since UMD does not allow multiple subprocesses opening the device at the same time.
+@pytest.mark.no_reset_default_device
+@pytest.mark.timeout(600)
+@pytest.mark.parametrize("export_measurements", [True, False])
+@pytest.mark.parametrize("batch_size", [1, 32])
+@pytest.mark.parametrize("data_parallel", [1, 2, 4, 8])
+@pytest.mark.parametrize("num_layers", [2, 10])
+@pytest.mark.parametrize("num_runs", [2])
+@pytest.mark.parametrize("max_seq_len", [1024, 32768, 131072])
+@pytest.mark.parametrize("max_generated_tokens", [2])
+@pytest.mark.parametrize("model_name", ["llama3_70b", "llama3_8b", "qwen3_32b"])  # Add more models here as needed
+@pytest.mark.parametrize("mode", ["prefill", "decode"])
+def test_device_perf_one_iter(
+    num_layers,
+    model_name,
+    batch_size,
+    data_parallel,
+    max_seq_len,
+    mode,
+    num_runs,
+    max_generated_tokens,
+    export_measurements,
+):
+    cmd = f"pytest models/tt_transformers/demo/simple_text_demo.py -k 'device-perf and performance' --num_layers {num_layers} --data_parallel {data_parallel} --max_seq_len {max_seq_len} --max_generated_tokens {max_generated_tokens} --paged_attention 1  --batch_size {batch_size} --mode {mode} --use_prefetcher True"
+    cols = ["DEVICE FW", "DEVICE KERNEL", "DEVICE BRISC KERNEL"]
+    device_analysis_types = ["device_kernel_duration", "device_kernel_first_to_last_start"]
+    subdir = f"ttt-device-perf-{mode}"
+    profiler = BenchmarkProfiler()
+    benchmark_data = BenchmarkData()
+    profiler.start("run")
+    profiler.start("decoder-perf-op-metrics")
+
+    # Load perf targets
+    perf_targets = {}
+    try:
+        with open(f"models/tt_transformers/tests/perf_targets/device_perf_{mode}.json", "r") as f:
+            perf_targets = json.load(f)
+    except FileNotFoundError:
+        logger.error(
+            f"Perf targets file not found, device perf test will proceed without performance target comparison"
+        )
+
+    # Large sequence lengths dispatch many more programs than the default profiler buffer
+    # (1333 slots) can hold, causing buffer overflow and missing device perf data.
+    # Scale op_support_count proportionally with max_seq_len (empirically, seq=131072 needs ~10000).
+    op_support_count = max(
+        int(PROFILER_DEFAULT_OP_SUPPORT_COUNT * 1.333),  # preserve previous minimum
+        int(max_seq_len / 10),
+    )
+
+    _ = run_device_perf(
+        cmd,
+        subdir,
+        num_iterations=1,
+        cols=cols,
+        batch_size=batch_size,
+        device_analysis_types=device_analysis_types,
+        op_support_count=op_support_count,
+    )
+
+    profiler.end("decoder-perf-op-metrics")
+    profiler.end("run")
+
+    # Parse the latest ops CSV and aggregate per-op metrics
+    filename = get_latest_ops_log_filename(subdir)
+    df = pd.read_csv(filename)
+    df = df[df["OP TYPE"].isin(["tt_dnn_device"])]
+    df = merge_device_rows(df)
+
+    # Split compile and trace
+    (
+        df_model_compilation,
+        df_model_trace,
+        df_first_layer_compilation,
+        df_first_layer_trace,
+        df_mid_layers_compilation,
+        df_mid_layers_trace,
+        df_model_tail_compilation,
+        df_model_tail_trace,
+    ) = split_compile_and_trace(
+        df,
+        mode=mode,
+        num_runs=num_runs,
+        num_layers=num_layers,
+    )
+
+    (
+        kernel_agg_first_layer_compile,
+        dispatch_agg_first_layer_compile,
+        firstlast_agg_first_layer_compile,
+    ) = process_measurements(df_first_layer_compilation, 1)
+    (
+        kernel_agg_first_layer_trace,
+        dispatch_agg_first_layer_trace,
+        firstlast_agg_first_layer_trace,
+    ) = process_measurements(df_first_layer_trace, 1)
+
+    if num_layers > 1:
+        (
+            kernel_agg_mid_layers_compile,
+            dispatch_agg_mid_layers_compile,
+            firstlast_agg_mid_layers_compile,
+        ) = process_measurements(
+            df_mid_layers_compilation, num_layers - 1
+        )  # we dont count the first layer
+
+        (
+            kernel_agg_mid_layers_trace,
+            dispatch_agg_mid_layers_trace,
+            firstlast_agg_mid_layers_trace,
+        ) = process_measurements(df_mid_layers_trace, num_layers - 1)
+
+    if df_model_tail_compilation is not None:
+        (
+            kernel_agg_model_tail_compile,
+            dispatch_agg_model_tail_compile,
+            firstlast_agg_model_tail_compile,
+        ) = process_measurements(df_model_tail_compilation, 1)
+        (
+            kernel_agg_model_tail_trace,
+            dispatch_agg_model_tail_trace,
+            firstlast_agg_model_tail_trace,
+        ) = process_measurements(df_model_tail_trace, 1)
+
+    # Print measurements
+    print_dict(kernel_agg_first_layer_compile, "KERNEL AVERAGE DURATION FOR FIRST LAYER COMPILE")
+    print_dict(kernel_agg_first_layer_trace, "KERNEL AVERAGE DURATION FOR FIRST LAYER TRACE")
+
+    if num_layers > 1:
+        print_dict(kernel_agg_mid_layers_compile, "KERNEL AVERAGE DURATION FOR MID LAYERS COMPILE")
+        print_dict(kernel_agg_mid_layers_trace, "KERNEL AVERAGE DURATION FOR MID LAYERS TRACE")
+        print_dict(dispatch_agg_mid_layers_trace, "DISPATCH AVERAGE DURATION FOR MID LAYERS TRACE")
+        print_dict(firstlast_agg_mid_layers_trace, "FIRST TO LAST AVERAGE START TIME FOR MID LAYERS TRACE")
+
+    if df_model_tail_compilation is not None:
+        print_dict(kernel_agg_model_tail_compile, "KERNEL AVERAGE DURATION FOR MODEL TAIL COMPILE")
+        print_dict(kernel_agg_model_tail_trace, "KERNEL AVERAGE DURATION FOR MODEL TAIL TRACE")
+        print_dict(dispatch_agg_model_tail_trace, "DISPATCH AVERAGE DURATION FOR MODEL TAIL TRACE")
+        print_dict(firstlast_agg_model_tail_trace, "FIRST TO LAST AVERAGE START TIME FOR MODEL TAIL TRACE")
+
+    # Prefer trace for collectives, compile for others
+    def is_collective(op_code: str) -> bool:
+        return any(x in op_code for x in ("AllGather", "ReduceScatter", "AllReduce", "Matmul_RS"))
+
+    # Prepare export structure and default margins
+    perf_measurements_export: dict[str, dict[str, dict[str, float]]] = {}
+
+    # Export metrics for an op group (first layer, mid layers, model tail)
+    def export_group(
+        group_name: str,
+        kernel_agg_compile: dict,
+        kernel_agg_trace: dict,
+        dispatch_agg_trace: dict,
+        firstlast_agg_trace: dict | None,
+    ):
+        all_passing = True
+        perf_measurements_export[group_name] = {}
+        op_codes = set(list(kernel_agg_compile["avg"].keys()) + list(kernel_agg_trace["avg"].keys()))
+        for op_code in op_codes:
+            # kernel avg
+            k_avg_trace = kernel_agg_trace["avg"].get(op_code)
+            k_min_trace = kernel_agg_trace["min"].get(op_code)
+            k_max_trace = kernel_agg_trace["max"].get(op_code)
+            k_avg_comp = kernel_agg_compile["avg"].get(op_code)
+            k_min_comp = kernel_agg_compile["min"].get(op_code)
+            k_max_comp = kernel_agg_compile["max"].get(op_code)
+
+            if is_collective(op_code):
+                k_avg = k_avg_trace if k_avg_trace is not None else k_avg_comp
+                k_min = k_min_trace if k_min_trace is not None else k_min_comp
+                k_max = k_max_trace if k_max_trace is not None else k_max_comp
+            else:
+                k_avg = k_avg_comp if k_avg_comp is not None else k_avg_trace
+                k_min = k_min_comp if k_min_comp is not None else k_min_trace
+                k_max = k_max_comp if k_max_comp is not None else k_max_trace
+
+            if k_avg is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-avg", float(k_avg)
+                )
+            if k_min is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-min", float(k_min)
+                )
+            if k_max is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-max", float(k_max)
+                )
+
+            # Initialize export entry for this op
+            export_entry: dict[str, float] = {}
+            if k_avg is not None:
+                export_entry["kernel_duration"] = float(k_avg)
+
+            # Check that perf_targets, group_name, and op_code exist and keys exist
+            if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]:
+                passing = verify_value_within_margin(
+                    k_avg,
+                    perf_targets[group_name][op_code]["kernel_duration"],
+                    perf_targets[group_name][op_code]["kernel_duration_relative_margin"],
+                    op_code,
+                    "kernel",
+                )
+                all_passing = all_passing and passing
+            else:
+                logger.warning(f"Warning: {op_code}-{group_name}-kernel not found in perf_targets")
+            # dispatch from trace only
+            d_avg = dispatch_agg_trace["avg"].get(op_code)
+            d_min = dispatch_agg_trace["min"].get(op_code)
+            d_max = dispatch_agg_trace["max"].get(op_code)
+            if d_avg is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-avg", float(d_avg)
+                )
+            if d_min is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-min", float(d_min)
+                )
+            if d_max is not None:
+                benchmark_data.add_measurement(
+                    profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-max", float(d_max)
+                )
+
+            if d_avg is not None:
+                export_entry["op_to_op"] = float(d_avg)
+
+            if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]:
+                passing = verify_value_within_margin(
+                    d_avg,
+                    perf_targets[group_name][op_code]["op_to_op"],
+                    perf_targets[group_name][op_code]["op_to_op_duration_relative_margin"],
+                    op_code,
+                    "op_to_op",
+                )
+                all_passing = all_passing and passing
+            else:
+                logger.warning(f"Warning: {op_code}-{group_name}-op_to_op not found in perf_targets")
+
+            # first_to_last from trace only (if provided)
+            if firstlast_agg_trace is not None:
+                fl_avg = firstlast_agg_trace["avg"].get(op_code)
+                fl_min = firstlast_agg_trace["min"].get(op_code)
+                fl_max = firstlast_agg_trace["max"].get(op_code)
+                if fl_avg is not None:
+                    benchmark_data.add_measurement(
+                        profiler,
+                        0,
+                        "decoder-perf-op-metrics",
+                        f"{op_code}-{group_name}-first_to_last-avg",
+                        float(fl_avg),
+                    )
+                if fl_min is not None:
+                    benchmark_data.add_measurement(
+                        profiler,
+                        0,
+                        "decoder-perf-op-metrics",
+                        f"{op_code}-{group_name}-first_to_last-min",
+                        float(fl_min),
+                    )
+                if fl_max is not None:
+                    benchmark_data.add_measurement(
+                        profiler,
+                        0,
+                        "decoder-perf-op-metrics",
+                        f"{op_code}-{group_name}-first_to_last-max",
+                        float(fl_max),
+                    )
+
+                if fl_avg is not None:
+                    export_entry["first_to_last_start"] = float(fl_avg)
+
+                if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]:
+                    passing = verify_value_within_margin(
+                        fl_avg,
+                        perf_targets[group_name][op_code]["first_to_last_start"],
+                        perf_targets[group_name][op_code]["first_to_last_start_relative_margin"],
+                        op_code,
+                        "first_to_last_start",
+                    )
+                    all_passing = all_passing and passing
+                else:
+                    logger.warning(f"Warning: {op_code}-{group_name}-first_to_last not found in perf_targets")
+            # Save the entry for this op if we collected any metrics
+            if export_entry:
+                perf_measurements_export[group_name][op_code] = export_entry
+        return all_passing
+
+    # Export per-op metrics for each group
+    all_passing = True
+    all_passing = all_passing and export_group(
+        group_name=f"decoder-first",
+        kernel_agg_compile=kernel_agg_first_layer_compile,
+        kernel_agg_trace=kernel_agg_first_layer_trace,
+        dispatch_agg_trace=dispatch_agg_first_layer_trace,
+        firstlast_agg_trace=firstlast_agg_first_layer_trace,
+    )
+
+    if num_layers > 1:
+        all_passing = all_passing and export_group(
+            group_name=f"decoder-mid",
+            kernel_agg_compile=kernel_agg_mid_layers_compile,
+            kernel_agg_trace=kernel_agg_mid_layers_trace,
+            dispatch_agg_trace=dispatch_agg_mid_layers_trace,
+            firstlast_agg_trace=firstlast_agg_mid_layers_trace,
+        )
+    if df_model_tail_compilation is not None:
+        all_passing = all_passing and export_group(
+            group_name=f"model-tail",
+            kernel_agg_compile=kernel_agg_model_tail_compile,
+            kernel_agg_trace=kernel_agg_model_tail_trace,
+            dispatch_agg_trace=dispatch_agg_model_tail_trace,
+            firstlast_agg_trace=None,  # align with decoder tail export (no first_to_last)
+        )
+
+    # Write out current measurements in the perf-measurements JSON schema
+    if export_measurements:
+        try:
+            output_dir = os.path.join(os.path.dirname(__file__), "perf_measurements")
+            os.makedirs(output_dir, exist_ok=True)
+            output_path = os.path.join(
+                output_dir,
+                f"device_perf_measurements_{model_name}_{mode}_bs{batch_size}_dp{data_parallel}_layers{num_layers}_seq{max_seq_len}.json",
+            )
+            with open(output_path, "w", encoding="utf-8") as f:
+                json.dump(perf_measurements_export, f, indent=4, sort_keys=True)
+            logger.info(f"Wrote perf measurements export to {output_path}")
+        except Exception as e:
+            logger.error(f"Failed to write perf measurements export: {e}")
+
+    # Save partial run
+    benchmark_data.save_partial_run_json(
+        profiler,
+        run_type="ttnn_decoder_unit",
+        ml_model_name=f"{model_name}-{mode}-{data_parallel}dp-{num_layers}layers-{max_seq_len}seq",
+    )
+
+    # No strict assertions on perf; test succeeds if profiling and export ran
+    assert True

code/models/tt_transformers/tests/test_hybrid_attention_for_causal_lm.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for ``HybridAttentionForCausalLM``.
+
+The class is the vLLM wrapper base for hybrid attention models (Gemma3,
+Gemma4, GPT-OSS, ...). The bulk of its responsibility is the
+``get_kv_cache_spec`` classmethod that translates ``layer_types`` from
+HF config into per-layer KVCacheSpecs that upstream's hybrid kv cache
+manager groups by attention type. This test pins that translation
+across the typical patterns we'll see on real models.
+"""
+
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+# Stub ttnn so importing generator_vllm doesn't blow up on the local
+# tt-metal C++ extension. We don't exercise any real ttnn behaviour here.
+sys.modules.setdefault("ttnn", MagicMock(name="ttnn-test-mock"))
+sys.modules.setdefault("ttnn._ttnn", MagicMock(name="ttnn._ttnn-test-mock"))
+
+
+def _make_vllm_config(layer_types, sliding_window=1024, num_kv_heads=8, head_size=128):
+    text_config = SimpleNamespace(layer_types=layer_types, sliding_window=sliding_window)
+    hf_config = SimpleNamespace(text_config=text_config)
+    cfg = MagicMock()
+    cfg.model_config.hf_config = hf_config
+    cfg.model_config.dtype = torch.bfloat16
+    cfg.model_config.get_num_kv_heads.return_value = num_kv_heads
+    cfg.model_config.get_head_size.return_value = head_size
+    cfg.cache_config.cache_dtype = "auto"
+    cfg.cache_config.block_size = 64
+    return cfg
+
+
+def test_spec_emits_one_entry_per_layer():
+    """KV cache groups temporarily disabled: every layer is FullAttentionSpec
+    regardless of layer_types entry. Reverts to one uniform spec until the
+    bounded-sliding-cache decode bug is fixed."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    layers = ["sliding_attention"] * 5 + ["full_attention"] + ["sliding_attention"] * 5
+    spec = HybridAttentionForCausalLM.get_kv_cache_spec(_make_vllm_config(layers))
+
+    assert len(spec) == len(layers)
+    for i in range(len(layers)):
+        name = f"model.layers.{i}.self_attn"
+        assert name in spec
+        assert isinstance(spec[name], FullAttentionSpec)
+        assert not isinstance(spec[name], SlidingWindowSpec)
+
+
+def test_spec_gemma3_27b_pattern():
+    """All layers are FullAttentionSpec while kv cache groups are disabled."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    pattern = ["sliding_attention"] * 5 + ["full_attention"]
+    layers = pattern * 10  # 60 layers
+    spec = HybridAttentionForCausalLM.get_kv_cache_spec(_make_vllm_config(layers))
+
+    full_count = sum(isinstance(v, FullAttentionSpec) for v in spec.values())
+    sliding_count = sum(isinstance(v, SlidingWindowSpec) for v in spec.values())
+    assert full_count == 60
+    assert sliding_count == 0
+
+
+def test_spec_gpt_oss_alternating_pattern():
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    layers = ["sliding_attention", "full_attention"] * 12  # 24 layers
+    spec = HybridAttentionForCausalLM.get_kv_cache_spec(_make_vllm_config(layers))
+
+    full_count = sum(isinstance(v, FullAttentionSpec) for v in spec.values())
+    sliding_count = sum(isinstance(v, SlidingWindowSpec) for v in spec.values())
+    assert full_count == 24
+    assert sliding_count == 0
+
+
+def test_spec_uniform_full_attention_still_works():
+    """All-full layer_types → single-type config; spec generation still succeeds."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    spec = HybridAttentionForCausalLM.get_kv_cache_spec(_make_vllm_config(["full_attention"] * 4))
+    assert all(isinstance(v, FullAttentionSpec) for v in spec.values())
+    assert not any(isinstance(v, SlidingWindowSpec) for v in spec.values())
+
+
+def test_spec_propagates_kv_heads_and_head_size():
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    spec = HybridAttentionForCausalLM.get_kv_cache_spec(
+        _make_vllm_config(["full_attention", "sliding_attention"], num_kv_heads=4, head_size=256)
+    )
+
+    for layer_spec in spec.values():
+        assert layer_spec.num_kv_heads == 4
+        assert layer_spec.head_size == 256
+        assert layer_spec.block_size == 64
+        assert layer_spec.dtype == torch.bfloat16
+
+
+def test_spec_missing_layer_types_raises():
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    cfg = _make_vllm_config(["full_attention"])
+    cfg.model_config.hf_config.text_config.layer_types = None
+
+    with pytest.raises(ValueError, match="layer_types"):
+        HybridAttentionForCausalLM.get_kv_cache_spec(cfg)
+
+
+def test_spec_unknown_layer_type_raises():
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    cfg = _make_vllm_config(["full_attention", "rotary_chunked_xyz"])
+
+    with pytest.raises(ValueError, match="Unsupported layer_type"):
+        HybridAttentionForCausalLM.get_kv_cache_spec(cfg)
+
+
+def test_subclass_must_override_prefill_and_decode():
+    """The base class's prefill_forward / decode_forward are explicit
+    NotImplementedError stubs — subclasses must provide model-specific
+    routing that consumes ``page_tables_per_group``."""
+    from models.tt_transformers.tt.generator_vllm import HybridAttentionForCausalLM
+
+    instance = HybridAttentionForCausalLM.__new__(HybridAttentionForCausalLM)
+
+    with pytest.raises(NotImplementedError, match="prefill_forward"):
+        instance.prefill_forward()
+    with pytest.raises(NotImplementedError, match="decode_forward"):
+        instance.decode_forward()

code/models/tt_transformers/tests/test_interleaved_to_sharded.py

@@ -0,0 +1,89 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.model_config import ModelArgs
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_decoder_inference(mesh_device, reset_seeds):
+    model_args = ModelArgs(mesh_device, cache_hf=True)
+    state_dict = torch.load(model_args.consolidated_weights_path, map_location=torch.device("cpu"))
+
+    # Ref model needs partial state dict, but our models use full state dict keys as cached weight names
+    first_layer_prefix = model_args.get_state_dict_prefix("TransformerBlock", 0)
+    partial_state_dict = {
+        k[len(first_layer_prefix) :]: v for k, v in state_dict.items() if (k.startswith(first_layer_prefix))
+    }
+    reference_model = model_args.reference_decoder()
+    reference_model.load_state_dict(partial_state_dict)
+
+    generation_length = 10
+
+    seqlen = 1
+    batch = model_args.max_batch_size
+
+    tt_ccl = TT_CCL(mesh_device)
+    for i in range(generation_length):
+        logger.info(f"[Decoder] Generating token {i}")
+
+        # input = torch.randn(1, 32, 4096)
+        pt_decode_input = (torch.rand(batch, seqlen, model_args.dim) * 2) - 1
+        tt_decode_input = pt_decode_input.clone()
+
+        decode_input = model_args.prepare_residual_tensor_decode(
+            tt_decode_input,
+            ttnn.L1_MEMORY_CONFIG,
+        )
+
+        dim = 2048
+
+        attn_input_grid = ttnn.CoreGrid(y=2, x=8)
+        mem_cfg = ttnn.create_sharded_memory_config(
+            (
+                32,
+                dim // attn_input_grid.num_cores,
+            ),
+            attn_input_grid,
+            ttnn.ShardStrategy.WIDTH,
+            ttnn.ShardOrientation.ROW_MAJOR,
+            use_height_and_width_as_shard_shape=True,
+        )
+
+        # Run TT model
+        tt_out = ttnn.experimental.all_gather_async(
+            decode_input,
+            persistent_output_buffer=None,
+            dim=3,
+            multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(),
+            num_links=1,
+            topology=model_args.ccl_topology(),
+            memory_config=mem_cfg,
+            barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(),
+            chunks_per_sync=10,
+            num_workers_per_link=2,
+            num_buffers_per_channel=2,
+        )
+
+        debug_max = lambda t: ttnn.to_torch(
+            t, mesh_composer=ttnn.ConcatMeshToTensor(model_args.mesh_device, dim=-1)
+        ).max()
+        logger.info(f"decode_input max: {debug_max(decode_input)=}, {decode_input.memory_config()=}")
+        logger.info(f"tt_out max: {debug_max(tt_out)=}, {tt_out.memory_config()=}")

code/models/tt_transformers/tests/test_llama90b_decoder_json.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""Llama 3.2 90B decoder JSON must match ModelOptimizations baseline except LI_QKV_DECODE (issue #36378)."""
+from pathlib import Path
+
+import pytest
+
+from models.tt_transformers.tt.model_config import (
+    DecodersPrecision,
+    MathFidelitySetting,
+    ModelOptimizations,
+    OpGroup,
+    TensorGroup,
+    parse_decoder_json,
+)
+
+_MN = "Llama-3.2-90B-Instruct"
+_TT_ROOT = Path(__file__).resolve().parents[1]
+_ACC = _TT_ROOT / "model_params" / _MN / "accuracy_decoder_config.json"
+_PERF = _TT_ROOT / "model_params" / _MN / "performance_decoder_config.json"
+
+
+def _tensor_map(m: ModelOptimizations):
+    return {
+        k.name: m.tensor_dtype_settings[k].name
+        for k in TensorGroup
+        if k in m.tensor_dtype_settings and m.tensor_dtype_settings[k] is not None
+    }
+
+
+def _op_map(m: ModelOptimizations):
+    return {k.name: m.op_fidelity_settings[k].name for k in OpGroup if k in m.op_fidelity_settings}
+
+
+def _assert_json_matches_baseline_except_qkv_decode(json_path: Path, opt_fn):
+    baseline = opt_fn(_MN)
+    b_tensor = _tensor_map(baseline)
+    b_op = _op_map(baseline)
+    loaded = parse_decoder_json(json_path, default_optimization=opt_fn)
+    n = len(loaded.decoder_optimizations)
+    assert n == 80
+    for di in range(n):
+        j = loaded.decoder_optimizations[di]
+        j_tensor = _tensor_map(j)
+        j_op = _op_map(j)
+        assert j_tensor == b_tensor, f"decoder {di} tensor_dtype mismatch"
+        for ok, vb in b_op.items():
+            jv = j_op.get(ok)
+            if ok == OpGroup.LI_QKV_DECODE.name:
+                assert (
+                    vb == MathFidelitySetting.HIFI2.name and jv == MathFidelitySetting.HIFI2_NOL1ACC.name
+                ), f"decoder {di} LI_QKV_DECODE expected HIFI2->HIFI2_NOL1ACC, got {vb}->{jv}"
+            else:
+                assert jv == vb, f"decoder {di} op {ok} baseline={vb} json={jv}"
+
+
+@pytest.mark.skipif(not _ACC.is_file(), reason="accuracy_decoder_config.json not present")
+def test_llama90b_accuracy_decoder_json_matches_baseline_except_qkv_decode():
+    _assert_json_matches_baseline_except_qkv_decode(_ACC, ModelOptimizations.accuracy)
+
+
+@pytest.mark.skipif(not _PERF.is_file(), reason="performance_decoder_config.json not present")
+def test_llama90b_performance_decoder_json_matches_baseline_except_qkv_decode():
+    _assert_json_matches_baseline_except_qkv_decode(_PERF, ModelOptimizations.performance)
+
+
+def test_llama90b_decoders_precision_without_json_layers_match_accuracy_baseline():
+    """When no JSON is used, DecodersPrecision repeats optimization_level(model_name) per layer."""
+    base = ModelOptimizations.accuracy(_MN)
+    dp = DecodersPrecision(80, _MN, base)
+    for i in range(80):
+        assert dp.decoder_optimizations[i]._full_name == base._full_name

code/models/tt_transformers/tests/test_lm_head.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.common import Mode
+from models.tt_transformers.tt.lm_head import LMHead
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.prefetcher import Prefetcher
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "use_prefetcher",
+    ([False]),
+)
+@pytest.mark.parametrize(
+    "seq_len",
+    (32,),
+)
+@pytest.mark.parametrize(
+    "batch_size",
+    (1,),
+)
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_lm_head_inference(seq_len, batch_size, mesh_device, use_prefetcher, reset_seeds):
+    dtype = ttnn.bfloat8_b
+
+    prefetcher = Prefetcher(mesh_device, num_tensors=0, num_layers=1) if use_prefetcher else None
+
+    if use_prefetcher:
+        prefetcher.init(mode=Mode.DECODE)
+
+    model_args = ModelArgs(
+        mesh_device, max_batch_size=batch_size, max_seq_len=seq_len, cache_hf=True, prefetcher=prefetcher
+    )
+    model_args.n_layers = 1
+
+    state_dict = model_args.load_state_dict()
+
+    state_dict_prefix = model_args.get_state_dict_prefix("", None)
+    # Ref model needs partial state dict, but our models use full state dict keys as cached weight names
+    partial_state_dict = {
+        "weight": state_dict[f"{state_dict_prefix}output.weight"],
+    }
+
+    model_args.WEIGHTS_DTYPE = dtype
+    reference_model = model_args.reference_lm_head()
+    reference_model.load_state_dict(partial_state_dict)
+
+    tt_ccl = TT_CCL(mesh_device)
+    tt_model = LMHead(
+        args=model_args,
+        mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
+        dtype=dtype,
+        state_dict=state_dict,
+        state_dict_prefix=state_dict_prefix,
+        weight_cache_path=model_args.weight_cache_path(dtype),
+        max_columns_per_device=model_args.max_columns_per_device_lm_head,
+        prefetcher=prefetcher,
+    )
+
+    torch_input = torch.randn(1, 1, seq_len, model_args.dim, dtype=torch.bfloat16)
+    reference_output = reference_model(torch_input)
+    tt_input = ttnn.from_torch(
+        torch_input,
+        device=mesh_device,
+        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(None, None), mesh_shape=model_args.cluster_shape),
+        dtype=ttnn.bfloat8_b,
+        memory_config=model_args.get_lm_head_input_mem_config(Mode.PREFILL, prefetcher),
+        layout=ttnn.TILE_LAYOUT,
+    )
+    tt_output = tt_model(tt_input)
+    tt_output_torch = ttnn.to_torch(
+        tt_output,
+        mesh_composer=ttnn.ConcatMesh2dToTensor(
+            mesh_device, model_args.cluster_shape, dims=(3, 1) if model_args.is_galaxy else (1, 3)
+        ),
+    )
+    tt_output_torch = tt_output_torch[:, 0:1, :, : model_args.vocab_size]
+
+    pcc_required = 0.99
+    passing, pcc_message = comp_pcc(reference_output, tt_output_torch, pcc_required)
+
+    logger.info(comp_allclose(reference_output, tt_output_torch))
+    logger.info(f"PCC: {pcc_message}")
+    if passing:
+        logger.info("LM_Head Passed!")
+    else:
+        logger.warning("LM_Head Failed!")
+
+    assert passing, f"LM_Head output does not meet PCC requirement {pcc_required}: {pcc_message}."

code/models/tt_transformers/tests/test_lm_head_cache.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+from models.tt_transformers.tt.lm_head import _lm_head_cache_file_name
+
+
+def test_galaxy_2d_lm_head_uses_distinct_cache_key():
+    cache_root = Path("weights")
+    common = {
+        "dummy_weights": False,
+        "num_splits": 1,
+        "shard": 0,
+        "width": 131072,
+        "mode": 0,
+        "mesh_shape": (8, 4),
+    }
+
+    flat = _lm_head_cache_file_name(cache_root, galaxy_2d=False, **common)
+    galaxy_2d = _lm_head_cache_file_name(cache_root, galaxy_2d=True, **common)
+
+    assert flat.name == "output_lm_head_1_split_shard_0_131072_mode_0"
+    assert galaxy_2d.name == "output_lm_head_1_split_shard_0_131072_mode_0_galaxy_2d_8x4_v1"
+    assert flat != galaxy_2d
+
+
+def test_dummy_lm_head_does_not_use_weight_cache():
+    assert (
+        _lm_head_cache_file_name(
+            Path("weights"),
+            dummy_weights=True,
+            num_splits=1,
+            shard=0,
+            width=131072,
+            mode=0,
+            galaxy_2d=True,
+            mesh_shape=(8, 4),
+        )
+        is None
+    )

code/models/tt_transformers/tests/test_mha_wo_sharding.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Verify that the wo output projection weight is correctly sharded for MHA
+models (n_heads == n_kv_heads) on multi-device meshes.
+
+Regression test for: ShardTensor2dMesh producing incorrect weight shapes
+when n_heads == n_kv_heads on N300 (1x2 mesh).
+"""
+
+import os
+
+import pytest
+import torch
+
+import ttnn
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_mha_wo_sharding(mesh_device):
+    """On a multi-device mesh, wo sharded via ShardTensorToMesh(dim=2) must
+    produce per-device shape [1, 1, n_heads*head_dim // num_devices, dim]
+    and the subsequent ttnn.linear must not crash."""
+
+    num_devices = mesh_device.get_num_devices()
+    if num_devices < 2:
+        pytest.skip("MHA sharding test requires >= 2 devices")
+
+    n_heads = 16
+    head_dim = 128
+    dim = 2048
+    seq_len = 128
+    n_local_heads = n_heads // num_devices
+    qkv_inner = n_local_heads * head_dim
+
+    pt_wo = torch.randn(1, 1, n_heads * head_dim, dim)
+    wo = ttnn.as_tensor(
+        pt_wo,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=2),
+    )
+
+    expected_sharded_dim = (n_heads * head_dim) // num_devices
+    assert wo.shape[-2] == expected_sharded_dim, f"wo dim 2 should be {expected_sharded_dim}, got {wo.shape[-2]}"
+
+    attn_output = ttnn.as_tensor(
+        torch.randn(1, 1, seq_len, qkv_inner),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+    )
+
+    output = ttnn.linear(
+        attn_output,
+        wo,
+        dtype=ttnn.bfloat16,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    assert output.shape[-1] == dim, f"output width should be {dim}, got {output.shape[-1]}"
+    assert output.shape[-2] == seq_len, f"output seq_len should be {seq_len}, got {output.shape[-2]}"
+
+    ttnn.deallocate(wo)
+    ttnn.deallocate(attn_output)
+    ttnn.deallocate(output)

code/models/tt_transformers/tests/test_mlp.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+from models.tt_transformers.tests.test_utils import get_ref_model_dype
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.common import Mode
+from models.tt_transformers.tt.mlp import MLP
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.prefetcher import Prefetcher
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "use_prefetcher",
+    ([False]),
+)
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "seq_len",
+    (64 * 1024, 32 * 1024, 512, 32),
+)
+@pytest.mark.parametrize(
+    "batch_size",
+    (1,),
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_mlp_inference(seq_len, batch_size, mesh_device, reset_seeds, ensure_gc, use_prefetcher):
+    dtype = ttnn.bfloat8_b
+    mode = Mode.DECODE if seq_len <= 32 else Mode.PREFILL
+
+    # Setup prefetcher (FF1, FF2, FF3 weights are prefetched)
+    num_tensors = 3 if mode == Mode.DECODE else 0
+    prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None
+
+    if use_prefetcher:
+        prefetcher.init(mode)
+
+    model_args = ModelArgs(
+        mesh_device,
+        max_batch_size=batch_size,
+        max_seq_len=128,
+        cache_hf=True,
+        prefetcher=prefetcher,
+    )
+    model_args.n_layers = 1
+    state_dict = model_args.load_state_dict()
+
+    # Ref model needs partial state dict, but our models use full state dict keys as cached weight names
+    first_layer_prefix = model_args.get_state_dict_prefix("MLP", 0)
+    partial_state_dict = {
+        k[len(first_layer_prefix) + 1 :]: v for k, v in state_dict.items() if (k.startswith(first_layer_prefix))
+    }
+
+    reference_model = model_args.reference_mlp()
+    reference_model.load_state_dict(partial_state_dict)
+    if model_args.is_90b:
+        # float32 ~3x faster than bfloat16.
+        # bfloat16 fails on CI (32k and 64k seq_len) with "This test seems to have hung... Timing out test case"
+        reference_model.to(torch.float32)
+
+    tt_ccl = TT_CCL(mesh_device)
+    tt_model = MLP(
+        mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
+        args=model_args,
+        state_dict=state_dict,
+        weight_cache_path=model_args.weight_cache_path(dtype),
+        layer_num=0,
+        dtype=dtype,
+        model_config=model_args.get_model_config(),
+        prefetcher=prefetcher,
+    )
+
+    # Run prefetcher if it is used
+    if prefetcher is not None and mode == Mode.DECODE:
+        prefetcher.prefetch()
+        prefetcher.run()
+
+    torch_input = torch.randn(
+        1, 1, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
+    )
+    reference_output = reference_model(torch_input)
+
+    tt_input = ttnn.from_torch(
+        torch_input,
+        device=mesh_device,
+        mesh_mapper=ttnn.ShardTensor2dMesh(
+            mesh_device,
+            dims=(None, 3) if model_args.is_galaxy else (None, None),
+            mesh_shape=model_args.cluster_shape,
+        ),  # When both dims are None, the mapper used is `ReplicateTensorToMesh`
+        dtype=ttnn.bfloat8_b,
+        memory_config=model_args.get_mlp_input_mem_config(mode, prefetcher),
+        layout=ttnn.TILE_LAYOUT,
+    )
+    logger.info("Run MLP")
+    tt_output = tt_model(tt_input, mode)
+
+    tt_output_torch = ttnn.to_torch(
+        tt_output,
+        mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),
+    )
+
+    tt_output_torch = tt_output_torch[:, :1, :, :]
+
+    pcc_required = 0.99
+    passing, pcc_message = comp_pcc(reference_output, tt_output_torch, pcc_required)
+
+    logger.info(comp_allclose(reference_output, tt_output_torch))
+    logger.info(f"PCC: {pcc_message}")
+    if passing:
+        logger.info("MLP Passed!")
+    else:
+        logger.warning("MLP Failed!")
+
+    assert passing, f"MLP output does not meet PCC requirement {pcc_required}: {pcc_message}."

code/models/tt_transformers/tests/test_model_config_utils.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from models.tt_transformers.tt.model_config import compute_padded_vocab_size, should_pad_sampling_logits_to_power_of_2
+
+
+@pytest.mark.parametrize(
+    ("vocab_size", "num_devices", "expected"),
+    [
+        (151936, 1, 151936),
+        (151936, 4, 151936),
+        (151936, 8, 152064),
+        (151936, 32, 152576),
+        (32001, 2, 32064),
+    ],
+)
+def test_compute_padded_vocab_size(vocab_size, num_devices, expected):
+    padded_vocab_size = compute_padded_vocab_size(vocab_size, num_devices)
+
+    assert padded_vocab_size == expected
+    assert padded_vocab_size >= vocab_size
+    assert padded_vocab_size % (32 * num_devices) == 0
+    assert (padded_vocab_size // num_devices) % 32 == 0
+
+
+def test_compute_padded_vocab_size_rejects_invalid_num_devices():
+    with pytest.raises(ValueError, match="num_devices must be >= 1"):
+        compute_padded_vocab_size(32000, 0)
+
+
+@pytest.mark.parametrize(
+    ("base_model_name", "padded_vocab_size", "sampling_splits", "expected"),
+    [
+        ("Llama-3.1-70B", 128256, 4, True),
+        ("Llama-3.1-70B", 131072, 4, False),
+        ("Llama-3.1-8B", 128256, 4, False),
+    ],
+)
+def test_should_pad_sampling_logits_to_power_of_2(base_model_name, padded_vocab_size, sampling_splits, expected):
+    assert should_pad_sampling_logits_to_power_of_2(base_model_name, padded_vocab_size, sampling_splits) is expected
+
+
+def test_should_pad_sampling_logits_to_power_of_2_rejects_invalid_sampling_splits():
+    with pytest.raises(ValueError, match="sampling_splits must be >= 1"):
+        should_pad_sampling_logits_to_power_of_2("Llama-3.1-70B", 128256, 0)

code/models/tt_transformers/tests/test_model_prefill.py

@@ -0,0 +1,322 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import bz2
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_pcc
+from models.tt_transformers.tt.common import PagedAttentionConfig, create_tt_model
+from models.tt_transformers.tt.generator import Generator
+from models.tt_transformers.tt.model_config import DecodersPrecision
+
+
+@torch.no_grad()
+@pytest.mark.timeout(900)
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("use_prefetcher", ([False]))
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+# Model and attention prefill tests should run both with and without paged attention to debug any issues that may occur with default attention
+@pytest.mark.parametrize(
+    "paged_attention",
+    (
+        True,
+        # False,
+    ),
+    ids=(
+        "paged_attention",
+        # "default_attention",
+    ),
+)
+@pytest.mark.parametrize(
+    "page_params",
+    [{"page_block_size": 32, "page_max_num_blocks": 1024}],
+)
+@pytest.mark.parametrize(
+    "seq_len",
+    (128, 256, 3072, 4096, 8192, 16384, 32768),
+    ids=["128", "256", "3k", "4k", "8k", "16k", "32k"],
+)
+@pytest.mark.parametrize(
+    "max_seq_len",
+    (128 * 1024,),
+    ids=[
+        "max128k",
+    ],
+)
+@pytest.mark.parametrize(
+    "optimizations",
+    [
+        lambda model_args: DecodersPrecision.performance(model_args.n_layers, model_args.model_name),
+        lambda model_args: DecodersPrecision.accuracy(model_args.n_layers, model_args.model_name),
+    ],
+    ids=["performance", "accuracy"],
+)
+@pytest.mark.parametrize(
+    "num_layers",
+    (1, None),
+    ids=["1layer", "all_layers"],
+)
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
+def test_model_inference(
+    paged_attention,
+    page_params,
+    optimizations,
+    seq_len,
+    max_seq_len,
+    num_layers,
+    mesh_device,
+    reset_seeds,
+    ensure_gc,
+    is_ci_env,
+    request,
+    use_prefetcher,
+):
+    test_id = request.node.callspec.id
+    use_hf_rope = request.config.getoption("--use_hf_rope")
+    if is_ci_env:
+        if "accuracy" in test_id:
+            pytest.skip("CI test only runs performance mode to reduce CI pipeline load")
+
+        # TODO: Save ref outputs to avoid running reference model for large seq_len
+        if seq_len > 8192:
+            pytest.skip("CI test only runs up to 8192 seq_len to avoid out of ram issues for ref model")
+        if use_hf_rope:
+            if num_layers != 1 and seq_len != 256:
+                pytest.skip("When HF rope is used CI only runs full model for 256 seq len to reduce CI pipeline load")
+
+        elif num_layers != 1 and seq_len != 4096:
+            pytest.skip("CI only runs full model for 4k seq len to reduce CI pipeline load")
+
+        hf_model_env = os.getenv("HF_MODEL", "")
+        if ("Llama" in hf_model_env) and ("Vision" in hf_model_env) and (num_layers is None):
+            pytest.skip("Skipping Llama Vision full model test: no CrossAttention functionality in this test.")
+
+    run_ref_pt = True  # Flag to run reference PyTorch model and compare PCC
+    dtype = ttnn.bfloat8_b
+    batch_size = 1  # For prefill we only support batch_size = 1
+
+    # Use instruct weights instead of general weights
+    instruct = True
+
+    paged_attention_config = (
+        PagedAttentionConfig(
+            block_size=page_params["page_block_size"],
+            max_num_blocks=page_params["page_max_num_blocks"],
+        )
+        if paged_attention
+        else None
+    )
+
+    # Load TTNN model
+    logger.info(f"Loading TT model...")
+    model_args, tt_model, tt_kv_cache, state_dict = create_tt_model(
+        mesh_device,
+        instruct=instruct,
+        max_batch_size=batch_size,
+        optimizations=optimizations,
+        max_seq_len=max_seq_len,
+        paged_attention_config=paged_attention_config,
+        dtype=dtype,
+        num_layers=num_layers,
+        use_prefetcher=use_prefetcher,
+        use_hf_rope=use_hf_rope,
+    )
+
+    if (
+        model_args.base_model_name.startswith("Mistral-")
+        or model_args.base_model_name.startswith("Qwen3-")
+        or model_args.base_model_name.startswith("Phi-3-mini-")
+        or model_args.base_model_name.startswith("phi-4")
+    ):
+        # TODO: Per layer KV cache fetching is not implemented for all models
+        # See issue https://github.com/tenstorrent/tt-metal/issues/19806"
+        cache_pcc = False
+    else:
+        cache_pcc = True
+
+    # This sets the minimum PCC for each iteration based on optimization mode
+    # TODO: See issue https://github.com/tenstorrent/tt-metal/issues/19806
+    perf_out_pcc_map = {"Mistral-7B-Instruct-v0.3": 0.73}
+    acc_out_pcc_map = {
+        "Mistral-7B-Instruct-v0.3": 0.75,
+        "Phi-3-mini-128k-instruct": 0.89,
+    }
+    kv_cache_pcc_map = {"Mistral-7B-Instruct-v0.3": 0.75}
+
+    if num_layers == 1:
+        expec_out_pcc = 0.97
+        expec_kv_cache_pcc = 0.99
+    else:
+        if "accuracy" in test_id:
+            default_expec_out_pcc = 0.91  # TODO Look on improving PCC
+            expec_out_pcc = acc_out_pcc_map.get(model_args.model_name, default_expec_out_pcc)
+        else:  # performance mode
+            assert "performance" in test_id
+            default_expec_out_pcc = 0.869  # TODO Look on improving PCC
+            expec_out_pcc = perf_out_pcc_map.get(model_args.model_name, default_expec_out_pcc)
+
+        default_expec_kv_cache_pcc = 0.88
+        expec_kv_cache_pcc = kv_cache_pcc_map.get(model_args.model_name, default_expec_kv_cache_pcc)
+
+    processor = model_args.processor
+    tokenizer = model_args.tokenizer
+    generator = Generator([tt_model], [model_args], mesh_device, processor=processor, tokenizer=tokenizer)
+    logger.info("Finished loading TT model.")
+
+    # Create page table if paged attention is enabled
+    if paged_attention:
+        # Implied shuffling of blocks
+        permutation = torch.randperm(paged_attention_config.max_num_blocks)
+        # Page table which maps virtual blocks to physical
+        reverse_permutation = torch.argsort(permutation)
+        page_table = reverse_permutation.reshape(
+            model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size
+        )
+    else:
+        page_table = None
+
+    # Load prompt
+    current_file_path = os.path.abspath(__file__)
+    current_file_dir = os.path.dirname(current_file_path)
+    prompt_file = os.path.join(current_file_dir, "tale-of-two-cities.txt.bz2")
+    with bz2.open(prompt_file, "rt", encoding="utf-8") as f:
+        prompt = f.read()
+    encoded_prompt = model_args.encode_prompt(prompt, instruct=instruct)[:seq_len]
+    logger.info(f"Prompt length: {len(encoded_prompt)} tokens")
+
+    # Load reference model
+    if run_ref_pt:
+        logger.info("Loading reference model...")
+        # On a warm ttnn cache, create_tt_model defers the HF load and returns a dataless
+        # placeholder state_dict. The host reference embedding below still needs the REAL weights,
+        # so reload them here when the state_dict is a warm-cache placeholder (either flavor) or
+        # genuinely empty. Branch on is_placeholder, NOT truthiness: tt_transformers'
+        # _PlaceholderStateDict is falsy but the shared CachedStateDict is truthy, so `if not
+        # state_dict` alone would silently stop reloading -- and build this torch reference from
+        # torch.empty garbage -- if the loaders are ever collapsed onto the shared class. (#45400)
+        if getattr(state_dict, "is_placeholder", False) or not state_dict:
+            state_dict = model_args.load_state_dict()
+        state_dict_prefix = model_args.get_state_dict_prefix("", None)
+        reference_model = model_args.reference_transformer(load_checkpoint=True)
+        # Embedding on host
+        embd = model_args.reference_embedding()
+        if model_args.is_llama_vision():
+            weight = torch.cat(
+                [
+                    state_dict[f"{state_dict_prefix}tok_embeddings.weight"],
+                    state_dict[f"{state_dict_prefix}learnable_embedding.weight"],
+                ],
+                dim=0,
+            )
+        else:
+            weight = state_dict[f"{state_dict_prefix}tok_embeddings.weight"]
+        embd.load_state_dict({"emb.weight": weight})
+        logger.info("Finished loading reference model.")
+
+    # Select the first token from the prompt for initial decoding
+    encoded_prompt_tensor = torch.tensor(encoded_prompt)  # [:,0]
+    tt_prefill_input = encoded_prompt_tensor.unsqueeze(0)
+    prompt_lens = [seq_len]
+    start_pos = 0
+
+    # Run TT model
+    logger.info(f"Running TT model...")
+    tt_output_torch = generator.prefill_forward_text(
+        tt_prefill_input,
+        page_table=page_table,
+        kv_cache=[tt_kv_cache],
+        prompt_lens=prompt_lens,
+    )
+    logger.info(f"Finished running TT model.")
+
+    if run_ref_pt:
+        # Run reference model
+        logger.info(f"Running reference model...")
+        pt_prefill_input = embd(encoded_prompt_tensor).view(batch_size, seq_len, -1)
+        ref_output = reference_model(pt_prefill_input, start_pos)
+        ref_output = ref_output[:, -1:, :]  # Get last token since TT model only returns the last token
+        logger.info(f"Finished running reference model.")
+
+        # Measure PCC if also running reference model
+        all_tests_pass = True
+
+        # Check output pcc
+        passing, pcc_message = comp_pcc(ref_output, tt_output_torch, expec_out_pcc)
+        logger.info(f"Output PCC: {pcc_message}")
+        if not passing:
+            all_tests_pass = False
+            logger.warning(f"Output PCC {pcc_message} is lower than {expec_out_pcc}")
+
+        # Compare KV caches
+        if cache_pcc:
+            for i in range(model_args.n_layers):
+                pytorch_layer_present = [
+                    reference_model.cache_k[i].clone().permute(0, 2, 1, 3),  # [batch_size, n_kv_heads, seq, head_dim]
+                    reference_model.cache_v[i].clone().permute(0, 2, 1, 3),  # [batch_size, n_kv_heads, seq, head_dim]
+                ]
+
+                tt_layer_present = []
+                if paged_attention:
+                    for layer_past in tt_model.layers[i].attention.layer_past:
+                        tt_layer_present.append(
+                            ttnn.to_torch(
+                                layer_past,
+                                mesh_composer=ttnn.ConcatMesh2dToTensor(
+                                    mesh_device,
+                                    dims=(1, 3) if model_args.is_galaxy else (0, 1),
+                                    mesh_shape=model_args.cluster_shape,
+                                ),
+                            )[reverse_permutation][:, : model_args.n_kv_heads, :, : model_args.head_dim]
+                            .reshape(
+                                model_args.max_batch_size,
+                                paged_attention_config.max_num_blocks // model_args.max_batch_size,
+                                model_args.n_kv_heads,
+                                paged_attention_config.block_size,
+                                model_args.head_dim,
+                            )
+                            .transpose(1, 2)
+                            .reshape(model_args.max_batch_size, model_args.n_kv_heads, -1, model_args.head_dim)[
+                                :batch_size, ...
+                            ]
+                        )
+                else:
+                    for layer_past in tt_model.layers[i].attention.layer_past_list[0]:
+                        tt_layer_present.append(
+                            ttnn.to_torch(
+                                layer_past,
+                                mesh_composer=ttnn.ConcatMesh2dToTensor(
+                                    mesh_device,
+                                    dims=(1, 0) if model_args.is_galaxy else (0, 1),
+                                    mesh_shape=model_args.cluster_shape,
+                                ),
+                            )
+                        )
+
+                for j, (cache_pt, cache_tt) in enumerate(zip(pytorch_layer_present, tt_layer_present)):
+                    cache_length_to_check = seq_len
+                    cache_pt = cache_pt[:, :, 0:cache_length_to_check, :]
+                    cache_tt = cache_tt[:, :, 0:cache_length_to_check, :]
+                    pcc_passed, output_pcc = comp_pcc(cache_pt, cache_tt, expec_kv_cache_pcc)
+                    kv_str = "K" if j == 0 else "V"
+                    logger.info(f"[layer={i+1}] {kv_str} cache PCC: {output_pcc}")
+                    if not pcc_passed:
+                        all_tests_pass = False
+                        logger.warning(f"[layer={i+1}] {kv_str} PCC {output_pcc} is lower than {expec_kv_cache_pcc}")
+
+        if all_tests_pass:
+            logger.info("All PCC checks passed!")
+        else:
+            assert all_tests_pass, f"PCC is lower than expected for some of the outputs. Check warnings!"

code/models/tt_transformers/tests/test_rope_utils.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from models.tt_transformers.tt.rope import get_batch_size_per_device_group
+
+
+@pytest.mark.parametrize(
+    ("batch_size", "use_qk_fused", "num_devices", "mesh_shape", "expected"),
+    [
+        (1, True, 32, (8, 4), 2),
+        (1, False, 32, (8, 4), 1),
+        (32, True, 32, (8, 4), 16),
+        (1, True, 1, (), 2),
+    ],
+)
+def test_get_batch_size_per_device_group(batch_size, use_qk_fused, num_devices, mesh_shape, expected):
+    assert get_batch_size_per_device_group(batch_size, use_qk_fused, num_devices, mesh_shape, 1) == expected

code/models/tt_transformers/tests/test_utils.py

@@ -0,0 +1,439 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from collections import defaultdict
+
+import pandas as pd
+import torch
+from loguru import logger
+
+from models.tt_transformers.tt.model_config import HfAttentionWrapper, HfDecoderWrapper, HfModelWrapper
+
+
+def _extract_dtype_from_state_dict(model):
+    """Helper to extract dtype from model's state_dict."""
+    try:
+        state_dict = model.state_dict()
+        for key, param in state_dict.items():
+            if "weight" in key:
+                print(f"get_ref_model_dype: key={key}, dtype={param.dtype}")
+                return param.dtype
+    except Exception as e:
+        pass
+    return None
+
+
+def get_ref_model_dype(ref_model, model_name):
+    default_dype = torch.float32
+
+    if ref_model is None and model_name is None:
+        return default_dype
+
+    try:
+        models_to_check = []
+        if isinstance(ref_model, HfAttentionWrapper):
+            models_to_check.append(ref_model.attention)
+        elif isinstance(ref_model, HfDecoderWrapper):
+            models_to_check.append(ref_model.decoder)
+        elif isinstance(ref_model, HfModelWrapper):
+            models_to_check.append(ref_model.model)
+        else:
+            models_to_check = [ref_model]
+
+        # Try all models until one works
+        for model in models_to_check:
+            if model is not None:
+                dtype = _extract_dtype_from_state_dict(model)
+                if dtype is not None:
+                    return dtype
+
+    except Exception as e:
+        pass
+
+    # try hardcoded dtypes
+    if model_name and isinstance(model_name, str):
+        model_name_lower = model_name.lower()
+        if "mistral-7b" in model_name_lower:
+            return torch.bfloat16
+        if "llama" in model_name_lower:
+            return torch.bfloat16
+        if "phi-3-mini" in model_name_lower or "phi-4" in model_name_lower:
+            return torch.bfloat16
+
+    return default_dype
+
+
+### UTIL FUNCTIONS FOR DEVICE PERF
+def build_duration_dict(raw_dict, column_name):
+    """Build a dictionary of op codes to list of durations."""
+    op_code_dict = {}
+    for entry in raw_dict:
+        if column_name not in entry:
+            logger.warning(f"Warning: {entry} does not have column {column_name}")
+        op_code = entry["OP CODE"]
+        duration = entry[column_name]
+        if op_code not in op_code_dict:
+            op_code_dict[op_code] = []
+        op_code_dict[op_code].append(duration)
+    return op_code_dict
+
+
+def build_duration_per_instance_dict(input_dict, num_layers):
+    """Build a dictionary of op codes to list of durations per instance."""
+    per_instance_dict = {}
+    for op_code in input_dict:
+        num_ops_with_op_code = len(input_dict[op_code])
+        num_instances = num_ops_with_op_code // num_layers
+        if num_ops_with_op_code % num_layers != 0:
+            logger.warning(
+                f"Warning: {op_code} has {num_ops_with_op_code} ops, not a multiple of {num_layers} layers. Skipping per-instance analysis for this op."
+            )
+            continue  # Skip this op_code instead of asserting
+        for iteration_id in range(num_layers):
+            for instance_id in range(num_instances):
+                op_code_with_id = f"{op_code}_{instance_id}"
+                if op_code_with_id not in per_instance_dict:
+                    per_instance_dict[op_code_with_id] = []
+                per_instance_dict[op_code_with_id].append(
+                    input_dict[op_code][iteration_id * num_instances + instance_id]
+                )
+    return per_instance_dict
+
+
+def merge_device_rows(df):
+    """
+    Merges device rows from a DataFrame into a single row per device.
+
+    Args:
+        df: A DataFrame containing measurements.
+
+    Returns:
+        A DataFrame with merged rows.
+    """
+    block_by_device = defaultdict(list)
+
+    for _, row in df.iterrows():
+        op_name = row["OP CODE"]
+        op_type = row["OP TYPE"]
+
+        if op_type == "tt_dnn_device":
+            device_id = int(row["DEVICE ID"])
+            block_by_device[device_id].append((op_name, row.to_dict()))
+
+    device_ids = sorted(block_by_device.keys())
+    merged_blocks = []
+    global_index = 0
+    while max(len(block_by_device[device_id]) for device_id in device_ids) > 0:
+        blocks = []
+        op_name = None
+        missing_devices = []
+        for device_id in device_ids:
+            if not len(block_by_device[device_id]):
+                logger.warning(f"Warning: Device {device_id} is missing operation {op_name} at index {global_index}")
+                continue
+            if op_name is None:
+                op_name = block_by_device[device_id][0][0]
+            elif op_name != block_by_device[device_id][0][0]:
+                missing_devices.append(device_id)
+                continue
+
+            blocks.append(block_by_device[device_id].pop(0))
+
+        if missing_devices:
+            logger.warning(
+                f"Warning: {op_name} at index {global_index} not present in CSV for {len(missing_devices)} devices {missing_devices} - do not trust data for this op or directly subsequent ops with the same name"
+            )
+
+        if not blocks:
+            break
+
+        if "AllGather" in op_name or "ReduceScatter" in op_name or "AllReduce" in op_name or "Matmul_RS" in op_name:
+            # For collective ops, take the average duration over all rows within a block
+            device_kernel_durations = [
+                d["DEVICE KERNEL DURATION [ns]"]
+                for _, d in blocks
+                if "DEVICE KERNEL DURATION [ns]" in d and not math.isnan(d["DEVICE KERNEL DURATION [ns]"])
+            ]
+
+            average_duration = (
+                sum(device_kernel_durations) / len(device_kernel_durations) if device_kernel_durations else float("nan")
+            )
+            # Use the first block's data but update its duration with the average
+            base_block = blocks[0][1].copy()
+            base_block["DEVICE KERNEL DURATION [ns]"] = average_duration
+            merged_blocks.append(base_block)
+        else:
+            # For non-collective ops, take the row with maximum duration
+            max_duration_block = max(blocks, key=lambda x: x[1]["DEVICE KERNEL DURATION [ns]"])
+            merged_blocks.append(max_duration_block[1])
+
+        global_index += 1
+
+    return pd.DataFrame(merged_blocks)
+
+
+def process_measurements(df, num_layers):
+    """
+    Given a Dataframe containing op device perf measurements, return the average, min, and max durations per instance on kerne
+    dispatch, and first to last start.
+
+    Args:
+        df: A DataFrame containing measurements.
+        num_layers: The number of layers in the model.
+
+    Returns:
+        A dictionary of aggregated values.
+        - kernel_duration_per_instance_aggregate_dict: A dictionary of aggregated kernel durations per instance.
+        - dispatch_duration_per_instance_aggregate_dict: A dictionary of aggregated dispatch durations per instance.
+        - first_to_last_start_per_instance_aggregate_dict: A dictionary of aggregated first to last start durations per instance.
+    """
+    raw_dict = df[
+        ["OP CODE", "DEVICE KERNEL DURATION [ns]", "OP TO OP LATENCY [ns]", "DEVICE KERNEL FIRST TO LAST START [ns]"]
+    ].to_dict(orient="records")
+
+    # Kernel duration
+    kernel_duration_dict = build_duration_dict(raw_dict, "DEVICE KERNEL DURATION [ns]")
+    kernel_duration_per_instance_dict = build_duration_per_instance_dict(kernel_duration_dict, num_layers)
+    kernel_duration_per_instance_aggregate_dict = {
+        "avg": aggregate_per_instance_dict(kernel_duration_per_instance_dict, lambda v: sum(v) / len(v)),
+        "min": aggregate_per_instance_dict(kernel_duration_per_instance_dict, min),
+        "max": aggregate_per_instance_dict(kernel_duration_per_instance_dict, max),
+    }
+
+    # Dispatch duration
+    dispatch_duration_dict = build_duration_dict(raw_dict, "OP TO OP LATENCY [ns]")
+    dispatch_duration_per_instance_dict = build_duration_per_instance_dict(dispatch_duration_dict, num_layers)
+    dispatch_duration_per_instance_aggregate_dict = {
+        "avg": aggregate_per_instance_dict(dispatch_duration_per_instance_dict, lambda v: sum(v) / len(v)),
+        "min": aggregate_per_instance_dict(dispatch_duration_per_instance_dict, min),
+        "max": aggregate_per_instance_dict(dispatch_duration_per_instance_dict, max),
+    }
+    # First to last start
+    first_to_last_start_dict = build_duration_dict(raw_dict, "DEVICE KERNEL FIRST TO LAST START [ns]")
+    first_to_last_start_per_instance_dict = build_duration_per_instance_dict(first_to_last_start_dict, num_layers)
+    first_to_last_start_per_instance_aggregate_dict = {
+        "avg": aggregate_per_instance_dict(first_to_last_start_per_instance_dict, lambda v: sum(v) / len(v)),
+        "min": aggregate_per_instance_dict(first_to_last_start_per_instance_dict, min),
+        "max": aggregate_per_instance_dict(first_to_last_start_per_instance_dict, max),
+    }
+
+    return (
+        kernel_duration_per_instance_aggregate_dict,
+        dispatch_duration_per_instance_aggregate_dict,
+        first_to_last_start_per_instance_aggregate_dict,
+    )
+
+
+def print_dict(input_dict, dict_name):
+    # print dict as a readable python dict
+    logger.info(f"\n{dict_name} = {{")
+    for op_code_with_id in input_dict:
+        logger.info(f'"{op_code_with_id}": {input_dict[op_code_with_id]},')
+    logger.info("}")
+
+
+def aggregate_per_instance_dict(input_dict, agg_fn, default=0):
+    """
+    Aggregates a dictionary of values by a given function.
+
+    Args:
+        input_dict: A dictionary of values to aggregate.
+        agg_fn: A function to aggregate the values.
+        default: The default value to return if the dictionary is empty.
+
+    Returns:
+        A dictionary of aggregated values.
+    """
+    result = {}
+    for key, values in input_dict.items():
+        clean_values = [v if v is not None else 0 for v in values]
+        result[key] = agg_fn(clean_values) if clean_values else default
+    return result
+
+
+def find_repeated_runs(ops, num_runs):
+    """
+    Find the starting index of repeated operation runs in a list.
+
+    This function scans through a list of operations (`ops`) to find the
+    first index (`left`) such that the remaining portion of the list,
+    `ops[left:]`, can be evenly divided into `num_runs` contiguous segments
+    (runs), all of which are identical.
+    """
+
+    def check_ops(left):
+        n = len(ops) - left
+        if n % num_runs != 0:
+            return False  # Can't evenly split
+
+        run_length = n // num_runs
+        first = ops[left : left + run_length]
+        for i in range(1, num_runs):
+            if ops[left + i * run_length : left + (i + 1) * run_length] != first:
+                return False
+        return True
+
+    left = 0
+    while left < len(ops):
+        if check_ops(left):
+            return left
+        left += 1
+    return -1  # return -1 if not found
+
+
+def find_repeated_block(ops, min_repeat=2):
+    """
+    Detect a repeating block (pattern) of operations within a list.
+
+    This function scans through the list of operations `ops` to find a contiguous
+    sub-sequence (block) that repeats consecutively at least `min_repeat` times.
+    It returns information about the prefix (head) before the repeated region,
+    the size and count of the repeated block, and the suffix (tail) after it.
+
+    The function assumes that each block represents a "layer" or
+    repeating structure (e.g., neural network layer operations).
+    It tries multiple possible block sizes (starting from 10) to identify
+    the first valid repeated pattern.
+
+    """
+    n = len(ops)
+    for block_size in range(10, n // min_repeat + 1):  # ignore tiny blocks
+        for start in range(n - 2 * block_size):
+            block = ops[start : start + block_size]
+            next_block = ops[start + block_size : start + 2 * block_size]
+
+            if block == next_block:
+                # Found a repeating pattern
+                # Extend it as far as it repeats
+                i = start
+                while i + block_size <= n and ops[i : i + block_size] == block:
+                    i += block_size
+                repeat_count = (i - start) // block_size
+
+                head = ops[:start]
+                tail = ops[i:]
+                return {
+                    "num_head_ops": len(head),
+                    "num_layer_block_ops": len(block),
+                    "num_layers": repeat_count,
+                    "num_tail_ops": len(tail),
+                }
+    # No repetition found
+    return {
+        "num_head_ops": len(ops),
+        "num_layer_block_ops": 0,
+        "num_layers": 0,
+        "num_tail_ops": len(ops),
+    }
+
+
+def split_compile_and_trace(
+    df: pd.DataFrame,
+    mode: str = "prefill",
+    num_runs: int = 1,
+    num_layers: int = None,
+):
+    """
+    Split a concatenated ops DataFrame into compile and runtime-trace segments,
+    and further partition those into first layer, mid layers, and model tail DataFrames.
+
+    The ops CSV typically contains three consecutive phases: compile, capture/trace,
+    and runtime trace. When an extra sampling compile pass is present (to enable
+    random sampling), it contributes a fixed number of rows that should not be used
+    to determine the thirds split.
+
+    Parameters:
+        df:                the input DataFrame (all ops)
+        mode:              the mode of the test (prefill or decode)
+        num_runs:          number of runs in the CSV (typically 3: compile, capture, trace)
+        num_layers:        number of core layers to partition (required for further splits)
+
+    Returns:
+        (
+            df_model_compilation, df_model_trace,
+            df_first_layer_compilation, df_first_layer_trace,
+            df_mid_layers_compilation, df_mid_layers_trace,
+            df_model_tail_compilation, df_model_tail_trace
+        )
+        Any of the additional outputs may be None if slicing arguments are not provided.
+    """
+
+    # Finds the first index such that ops[left:] contains num_runs of identical blocks of ops
+    first_run_start = find_repeated_runs(df["OP CODE"].tolist(), num_runs)
+    adjusted_len = (len(df) - first_run_start) // num_runs  # The number of ops in each run
+    first_run_end = first_run_start + adjusted_len
+    last_run_start = len(df) - adjusted_len
+    df_model_compilation = df[first_run_start:first_run_end]
+    df_model_trace = df[last_run_start:]
+
+    # Find the head and tail of the repeating region in the model compilation/ trace region of ops
+    head_tail_ops = find_repeated_block(df_model_compilation["OP CODE"].tolist(), num_layers)
+
+    # [op_start_index:op_end_index] = all core layers region
+    op_start_index = head_tail_ops["num_head_ops"]
+    op_end_index = len(df_model_compilation) - head_tail_ops["num_tail_ops"]
+    df_layers_compilation = df_model_compilation[op_start_index:op_end_index]
+    df_layers_trace = df_model_trace[op_start_index:op_end_index]
+
+    # First layer: always first 'len/num_layers'
+    split_point = int(len(df_layers_compilation) / num_layers)
+    df_first_layer_compilation = df_layers_compilation[:split_point]
+    df_first_layer_trace = df_layers_trace[:split_point]
+
+    # Mid layers: remainder of layers region
+    if num_layers > 1:
+        df_mid_layers_compilation = df_layers_compilation[split_point:]
+        df_mid_layers_trace = df_layers_trace[split_point:]
+    else:
+        df_mid_layers_compilation = None
+        df_mid_layers_trace = None
+
+    # Model tail ops (e.g. lmhead/sampling): [tail_start_index:]
+    if op_end_index is not None:
+        df_model_tail_compilation = df_model_compilation[op_end_index:]
+        df_model_tail_trace = df_model_trace[op_end_index:]
+    else:
+        df_model_tail_compilation = None
+        df_model_tail_trace = None
+
+    return (
+        df_model_compilation,
+        df_model_trace,
+        df_first_layer_compilation,
+        df_first_layer_trace,
+        df_mid_layers_compilation,
+        df_mid_layers_trace,
+        df_model_tail_compilation,
+        df_model_tail_trace,
+    )
+
+
+def verify_value_within_margin(value, target, margin, op_code_with_id, perf_type):
+    upper_limit = target + margin * target
+    lower_limit = target - margin * target
+
+    passing = True
+
+    if value > upper_limit:
+        passing = False
+        logger.warning(
+            f"{op_code_with_id} {perf_type}: {value} ns is larger than target "
+            f"({target}) ns, difference: "
+            f"{abs(value - upper_limit)} ns, margin: "
+            f"{margin}, "
+            f"relative margin to pass would be: "
+            f"{(abs(target - value) / target) if target != 0 else -1}"
+        )
+    elif value < lower_limit:
+        passing = False
+        logger.warning(
+            f"{op_code_with_id} {perf_type}: {value} ns is smaller than target "
+            f"({target}) ns, difference: "
+            f"{abs(value - lower_limit)} ns, margin: "
+            f"{margin}, "
+            f"relative margin to pass would be: "
+            f"{(abs(target - value) / target) if target != 0 else -1}"
+        )
+    return passing

code/models/tt_transformers/tests/test_vllm_kv_cache.py

@@ -0,0 +1,141 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the vLLM-side KV cache allocator helpers in
+``generator_vllm.py``.
+
+Verifies the new per-layer entry point (``allocate_vllm_kv_cache_per_layer``)
+and that the legacy uniform-shape entry point (``allocate_vllm_kv_cache``)
+still delegates to it bit-for-bit.
+
+Real ttnn allocation requires a mesh device, so this test mocks
+``ttnn.as_tensor`` / ``ttnn.ReplicateTensorToMesh`` and the ``dp_model``
+handles. We verify call structure and shape routing, not the resulting
+tensor contents.
+"""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def dp_model():
+    """One submesh handle whose optimizations return None (so the allocator
+    falls back to the bfloat8_b default — keeps the test independent of
+    the model's optimization config table)."""
+    submesh = MagicMock()
+    args = MagicMock()
+    args.optimizations = None  # Force the bfloat8_b fallback path.
+    model = MagicMock()
+    model.mesh_device = submesh
+    model.args = args
+    return [model]
+
+
+def _make_ttnn_mock():
+    ttnn_mock = MagicMock()
+    ttnn_mock.as_tensor.side_effect = lambda *a, **kw: ("tt-tensor", kw.get("dtype"), kw.get("cache_file_name"))
+    ttnn_mock.bfloat8_b = "bfloat8_b-sentinel"
+    ttnn_mock.bfloat16 = "bfloat16-sentinel"
+    return ttnn_mock
+
+
+def test_per_layer_allocates_one_kv_pair_per_unique_tensor(dp_model):
+    """Each unique ``tensor_idx`` allocates one (k, v) pair; layers that
+    share a ``tensor_idx`` reuse the same handles."""
+    from models.tt_transformers.tt import generator_vllm
+
+    # Layers 0, 1, 2 all use tensor_idx=0,1,2 respectively → three buffers.
+    per_layer = [
+        ((4, 2, 32, 64), torch.bfloat16, 0),
+        ((4, 2, 32, 64), torch.bfloat16, 1),
+        ((4, 2, 32, 64), torch.bfloat16, 2),
+    ]
+
+    with patch.object(generator_vllm, "ttnn", new=_make_ttnn_mock()) as ttnn_mock:
+        kv_cache = generator_vllm.allocate_vllm_kv_cache_per_layer(
+            per_layer, dp_model=dp_model, tt_cache_path=Path("/tmp/tt-test-cache")
+        )
+
+    # One submesh, three layers, two tensors per layer (k, v) = 6 calls.
+    assert ttnn_mock.as_tensor.call_count == 6
+    assert len(kv_cache) == 1  # one submesh
+    assert len(kv_cache[0]) == 3  # three layers
+    assert all(len(layer) == 2 for layer in kv_cache[0])  # k, v
+
+
+def test_shared_tensor_idx_reuses_one_buffer(dp_model):
+    """Layers sharing a ``tensor_idx`` (HMA tensor sharing) point at the
+    same underlying ttnn handles and only one allocation runs per
+    ``tensor_idx``."""
+    from models.tt_transformers.tt import generator_vllm
+
+    # Layers 0 and 2 share tensor 0; layer 1 has its own tensor 1.
+    per_layer = [
+        ((4, 2, 32, 64), torch.bfloat16, 0),
+        ((4, 2, 32, 64), torch.bfloat16, 1),
+        ((4, 2, 32, 64), torch.bfloat16, 0),
+    ]
+
+    with patch.object(generator_vllm, "ttnn", new=_make_ttnn_mock()) as ttnn_mock:
+        kv_cache = generator_vllm.allocate_vllm_kv_cache_per_layer(
+            per_layer, dp_model=dp_model, tt_cache_path=Path("/tmp/tt-test-cache")
+        )
+
+    # 2 unique tensor_idx values × 2 (k, v) = 4 allocations.
+    assert ttnn_mock.as_tensor.call_count == 4
+    # Layers 0 and 2 must reference the *same* handle list.
+    assert kv_cache[0][0] is kv_cache[0][2]
+    assert kv_cache[0][0] is not kv_cache[0][1]
+
+
+def test_per_layer_keys_cache_filename_on_tensor_idx(dp_model):
+    """Cache filenames must distinguish independent buffers even when
+    shapes are identical, so on-disk caches can't collide across layers
+    that don't share a ``tensor_idx``."""
+    from models.tt_transformers.tt import generator_vllm
+
+    per_layer = [
+        ((4, 2, 32, 64), torch.bfloat16, 0),
+        ((4, 2, 32, 64), torch.bfloat16, 1),
+    ]
+
+    with patch.object(generator_vllm, "ttnn", new=_make_ttnn_mock()) as ttnn_mock:
+        generator_vllm.allocate_vllm_kv_cache_per_layer(
+            per_layer, dp_model=dp_model, tt_cache_path=Path("/tmp/tt-test-cache")
+        )
+
+    cache_filenames = [str(call.kwargs["cache_file_name"]) for call in ttnn_mock.as_tensor.call_args_list]
+    assert sum("_t0" in f for f in cache_filenames) == 2
+    assert sum("_t1" in f for f in cache_filenames) == 2
+
+
+def test_legacy_uniform_shape_delegates_to_per_layer(dp_model):
+    """The legacy ``allocate_vllm_kv_cache`` must produce identical output to
+    calling ``allocate_vllm_kv_cache_per_layer`` with a per-layer triple
+    list (each layer its own ``tensor_idx``), so existing single-group
+    callers keep working unchanged."""
+    from models.tt_transformers.tt import generator_vllm
+
+    shape = (4, 2, 32, 64)
+    dtype = torch.bfloat16
+    num_layers = 3
+
+    with patch.object(generator_vllm, "ttnn", new=_make_ttnn_mock()) as ttnn_mock:
+        legacy = generator_vllm.allocate_vllm_kv_cache(
+            shape, dtype, num_layers, dp_model=dp_model, tt_cache_path=Path("/tmp/c")
+        )
+    legacy_call_count = ttnn_mock.as_tensor.call_count
+
+    with patch.object(generator_vllm, "ttnn", new=_make_ttnn_mock()) as ttnn_mock:
+        per_layer = generator_vllm.allocate_vllm_kv_cache_per_layer(
+            [(shape, dtype, i) for i in range(num_layers)],
+            dp_model=dp_model,
+            tt_cache_path=Path("/tmp/c"),
+        )
+    per_layer_call_count = ttnn_mock.as_tensor.call_count
+
+    assert legacy_call_count == per_layer_call_count
+    assert len(legacy[0]) == len(per_layer[0]) == num_layers

code/models/tt_transformers/tests/test_warm_cache_marker.py

@@ -0,0 +1,437 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Lightweight (CPU-only, no device, no model download) tests for the warm ttnn weight-cache
+detector generalized into tt_transformers (issue #45400, generalizes GPT-OSS PR #48531).
+
+These exercise the *real* ModelArgs.weight_cache_is_complete / mark_weight_cache_complete /
+placeholder_state_dict logic by binding the unbound methods to a tiny stub whose
+weight_cache_path points at a tmp dir -- so we validate marker round-trip, the shape/dtype
+manifest, staleness rejection, the force-load override, the .tensorbin belt-and-suspenders
+check, and the dataless placeholder state_dict without constructing a full ModelArgs.
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from models.tt_transformers.tt.model_config import ModelArgs
+
+DTYPE = "bfp8"  # opaque here: the stub's weight_cache_path ignores it; marker stores str(dtype)
+
+SAMPLE_SD = {
+    "tok_embeddings.weight": torch.zeros(4, 8, dtype=torch.bfloat16),
+    "layers.0.attention.wo.weight": torch.zeros(8, 8, dtype=torch.float32),
+}
+
+
+class _FakeArgs:
+    """Minimal stand-in exposing exactly what the marker methods touch, with the real methods
+    bound so the production logic is under test."""
+
+    WEIGHT_CACHE_MARKER = ModelArgs.WEIGHT_CACHE_MARKER
+    WEIGHT_CACHE_FORMAT_VERSION = ModelArgs.WEIGHT_CACHE_FORMAT_VERSION
+    _weight_cache_identity = ModelArgs._weight_cache_identity
+    # Bound too: _weight_cache_identity calls it on self, so leaving it out made every gate call
+    # raise AttributeError -- and nothing ran this file to notice. (#45400 review, finding B1)
+    _weight_cache_build_variant = ModelArgs._weight_cache_build_variant
+    weight_cache_is_complete = ModelArgs.weight_cache_is_complete
+    mark_weight_cache_complete = ModelArgs.mark_weight_cache_complete
+    placeholder_state_dict = ModelArgs.placeholder_state_dict
+
+    def __init__(self, cache_dir, model_name="Test-Model-8B", n_layers=32, mesh_shape=(1, 8)):
+        self._cache_dir = Path(cache_dir)
+        self.model_name = model_name
+        self.n_layers = n_layers
+        self.dummy_weights = False
+        self.is_mixture_of_experts = False
+        self.mesh_device = SimpleNamespace(shape=mesh_shape)
+        # Everything _weight_cache_build_variant reads. get_tensor_dtype deliberately lives on
+        # self.optimizations (a DecodersPrecision in production), NOT on the args object: the
+        # variant helper shipped calling self.get_tensor_dtype and threw AttributeError on every
+        # model, and a stub that put the method on the args would have kept agreeing with that
+        # bug. Mirror the production shape so the stub can only pass against correct code.
+        self.prefetcher = None
+        self.max_batch_size = 1
+        self.use_fused_all_gather_matmul = False
+        self.use_hf_rope = False
+        self.optimizations = SimpleNamespace(
+            get_tensor_dtype=lambda decoder_id, tensor, prefetcher=False: "DataType.BFLOAT8_B"
+        )
+
+    def weight_cache_path(self, dtype):
+        return self._cache_dir
+
+
+def _touch_tensorbin(cache_dir):
+    (Path(cache_dir) / "some.weight.tensorbin").write_bytes(b"\x00")
+
+
+@pytest.fixture(autouse=True)
+def _clear_force_env(monkeypatch):
+    monkeypatch.delenv("TT_TRANSFORMERS_FORCE_MODEL_LOAD", raising=False)
+
+
+def test_cold_cache_is_incomplete(tmp_path):
+    args = _FakeArgs(tmp_path)
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_mark_then_complete_roundtrip(tmp_path):
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)  # a real build writes tensor files alongside the marker
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    assert args.weight_cache_is_complete(DTYPE) is True
+
+    # Marker payload includes the shape/dtype manifest.
+    meta = json.loads(marker_path(tmp_path, args._weight_cache_build_variant()).read_text())
+    assert meta["model_name"] == "Test-Model-8B"
+    assert meta["n_layers"] == 32
+    assert meta["mesh_shape"] == "(1, 8)"
+    assert meta["format_version"] == ModelArgs.WEIGHT_CACHE_FORMAT_VERSION
+    assert meta["weights"]["tok_embeddings.weight"] == [[4, 8], "torch.bfloat16"]
+
+
+def test_marker_without_manifest_is_incomplete(tmp_path):
+    # A marker with no weight manifest (e.g. an old v1-style write) can't back a warm build.
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    args.mark_weight_cache_complete(DTYPE)  # no state_dict -> weights == {}
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_marker_without_tensorbin_is_incomplete(tmp_path):
+    args = _FakeArgs(tmp_path)
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    # No .tensorbin present -> belt-and-suspenders check fails.
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_force_env_disables_skip(tmp_path, monkeypatch):
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    assert args.weight_cache_is_complete(DTYPE) is True  # warm...
+    monkeypatch.setenv("TT_TRANSFORMERS_FORCE_MODEL_LOAD", "1")
+    assert args.weight_cache_is_complete(DTYPE) is False  # ...but forced to cold-load
+
+
+@pytest.mark.parametrize(
+    "mutate",
+    [
+        pytest.param({"format_version": 999}, id="stale-format"),
+        pytest.param({"model_name": "Other-Model"}, id="wrong-model"),
+        pytest.param({"n_layers": 1}, id="partial-build"),
+        pytest.param({"mesh_shape": "(2, 4)"}, id="wrong-mesh"),
+        pytest.param({"weights": {}}, id="empty-manifest"),
+    ],
+)
+def test_stale_marker_rejected(tmp_path, mutate):
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    marker = marker_path(tmp_path, args._weight_cache_build_variant())
+    meta = json.loads(marker.read_text())
+    meta.update(mutate)
+    marker.write_text(json.dumps(meta))
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_corrupt_marker_is_incomplete(tmp_path):
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    marker_path(tmp_path, args._weight_cache_build_variant()).write_text("{ not json")
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_placeholder_state_dict_is_dataless_and_falsy(tmp_path):
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+
+    sd = args.placeholder_state_dict(DTYPE)
+    # Falsy so reference-building callers (`if not state_dict`) load real weights instead.
+    assert not sd
+    assert len(sd) == 2
+    assert set(sd.keys()) == set(SAMPLE_SD.keys())
+    # Reconstructs correct shape/dtype without any real data.
+    emb = sd["tok_embeddings.weight"]
+    assert tuple(emb.shape) == (4, 8)
+    assert emb.dtype == torch.bfloat16
+    assert sd["layers.0.attention.wo.weight"].dtype == torch.float32
+
+
+# ---------------------------------------------------------------------------
+# models/common/weight_cache.py -- the shared helper used by the forked loaders.
+# The tests above bind the tt_transformers ModelArgs methods; these cover the shared
+# module's own behaviour: sidecar capture/rejection, per-file completeness, component
+# matching, atomic publish, and the CachedStateDict contract. (#45400 review)
+# ---------------------------------------------------------------------------
+
+from models.common.weight_cache import (  # noqa: E402
+    HOST_WEIGHTS_SIDECAR,
+    WEIGHT_CACHE_MARKER,
+    CachedStateDict,
+    build_cached_state_dict,
+    mark_weight_cache_complete,
+    marker_path,
+    normalize_mesh_shape,
+    weight_cache_is_complete,
+)
+
+SHARED_ID = dict(model_name="unit/test-model", n_layers=2, mesh_shape=(1, 8))
+
+
+def _seed(tmp_path, *, components=None, is_host_weight=None, sd=None):
+    """Write a tensorbin then mark the cache complete, mimicking a real cold build."""
+    _touch_tensorbin(tmp_path)
+    mark_weight_cache_complete(
+        tmp_path, sd if sd is not None else SAMPLE_SD, components=components, is_host_weight=is_host_weight, **SHARED_ID
+    )
+
+
+def test_shared_marker_roundtrip(tmp_path):
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False
+    _seed(tmp_path)
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is True
+
+
+def test_shared_marker_rejects_missing_tensorbin(tmp_path):
+    """A recorded tensorbin that later disappears must force a cold load -- otherwise as_tensor
+    regenerates it from the placeholder and writes garbage into the cache permanently."""
+    _seed(tmp_path)
+    for f in tmp_path.glob("*.tensorbin"):
+        f.unlink()
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False
+
+
+def test_shared_marker_finds_tensorbins_in_subdirs(tmp_path):
+    """Forked loaders nest per-layer weights (qwen36 layers.N/, gemma4 layer_N/)."""
+    sub = tmp_path / "layers.0"
+    sub.mkdir()
+    (sub / "wq_dtype_BFLOAT8_B_layout_TILE.tensorbin").write_bytes(b"x")
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, **SHARED_ID)
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is True
+    (sub / "wq_dtype_BFLOAT8_B_layout_TILE.tensorbin").unlink()
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False
+
+
+def test_components_subset_matching(tmp_path):
+    """A text-only seed must not certify a build that also needs the vision tower; the reverse
+    (vision seed satisfying a text-only build) is fine."""
+    _seed(tmp_path, components=["text"])
+    assert weight_cache_is_complete(tmp_path, components=["text"], **SHARED_ID) is True
+    assert weight_cache_is_complete(tmp_path, components=["text", "vision"], **SHARED_ID) is False
+
+    _seed(tmp_path, components=["text", "vision"])
+    assert weight_cache_is_complete(tmp_path, components=["text"], **SHARED_ID) is True
+
+
+def test_sidecar_capture_and_corruption_rejected(tmp_path):
+    _seed(tmp_path, is_host_weight=lambda k: k == "tok_embeddings.weight")
+    assert (tmp_path / HOST_WEIGHTS_SIDECAR).is_file()
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is True
+
+    sd = build_cached_state_dict(tmp_path)
+    # The captured host weight is served REAL; everything else is a dataless placeholder.
+    assert torch.equal(sd["tok_embeddings.weight"], SAMPLE_SD["tok_embeddings.weight"])
+
+    # A torn sidecar must degrade to a cold load, not crash every later run.
+    (tmp_path / HOST_WEIGHTS_SIDECAR).write_bytes(b"not a torch file")
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False
+
+
+def test_no_temp_files_left_behind(tmp_path):
+    _seed(tmp_path, is_host_weight=lambda k: k == "tok_embeddings.weight")
+    assert not list(tmp_path.glob("*.tmp*")), "atomic publish must not leave temp files"
+
+
+def test_mesh_shape_encoding_is_writer_agnostic(tmp_path):
+    """Both marker writers must encode the mesh identically or each rejects the other's marker."""
+
+    class _MeshShape:
+        def __init__(self, dims):
+            self._dims = dims
+
+        def __iter__(self):
+            return iter(self._dims)
+
+        def __str__(self):
+            return f"MeshShape({list(self._dims)})"
+
+    assert normalize_mesh_shape(_MeshShape((1, 8))) == normalize_mesh_shape((1, 8))
+    _seed(tmp_path)
+    identity = dict(SHARED_ID, mesh_shape=_MeshShape((1, 8)))
+    assert weight_cache_is_complete(tmp_path, **identity) is True
+
+
+def test_cached_state_dict_contract():
+    manifest = {k: [list(v.shape), str(v.dtype)] for k, v in SAMPLE_SD.items()}
+    real = SAMPLE_SD["tok_embeddings.weight"]
+    sd = CachedStateDict(manifest, {"tok_embeddings.weight": real})
+
+    # Truthy and flagged (the tt_transformers placeholder is falsy -- callers must branch on the
+    # attribute, not truthiness).
+    assert sd
+    assert sd.is_placeholder is True
+
+    # Membership must not materialize a tensor.
+    assert "layers.0.attention.wo.weight" in sd
+    assert "nope" not in sd
+    assert sd.get("nope") is None
+
+    # Mutable: loaders setdefault KV-shared weights.
+    sd["extra"] = torch.zeros(2)
+    assert "extra" in sd
+    del sd["extra"]
+    assert "extra" not in sd
+
+    del sd["tok_embeddings.weight"]
+    assert "tok_embeddings.weight" not in sd
+    try:
+        sd["tok_embeddings.weight"]
+        raise AssertionError("a deleted key must raise KeyError, even when the sidecar still has it")
+    except KeyError:
+        pass
+
+
+def test_build_variant_must_match_exactly(tmp_path):
+    """Build options that change an as_tensor cache FILENAME (prefetcher, precision) are matched
+    exactly, not as a superset: a different variant needs DIFFERENT files, and any it is missing
+    would be regenerated from the placeholder rather than cold-loaded."""
+    perf = {"prefetcher": False, "precision": "aaaaaaaaaaaa"}
+    _touch_tensorbin(tmp_path)
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, build_variant=perf, **SHARED_ID)
+
+    assert weight_cache_is_complete(tmp_path, build_variant=perf, **SHARED_ID) is True
+    # different precision config
+    assert (
+        weight_cache_is_complete(tmp_path, build_variant={"prefetcher": False, "precision": "bbbb"}, **SHARED_ID)
+        is False
+    )
+    # prefetcher flips the dtypes and adds ring-matmul splits
+    assert (
+        weight_cache_is_complete(tmp_path, build_variant={"prefetcher": True, "precision": "aaaaaaaaaaaa"}, **SHARED_ID)
+        is False
+    )
+    # a caller that records no variant must not be satisfied by one that did
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False
+
+
+# ---------------------------------------------------------------------------
+# Review-fix coverage (#45400 findings B3 / R1 / R3).
+# ---------------------------------------------------------------------------
+
+import models.common.weight_cache as _wc  # noqa: E402
+
+
+def test_variant_markers_coexist(tmp_path):
+    """One marker file PER build variant: two variants sharing a cache dir must not evict each
+    other's marker. The live case is the Llama CI job running eval-32 with and without the DRAM
+    prefetcher against one instruct cache -- a single exactly-matched marker made each leg's seed
+    clobber the other's, so both cold-loaded forever with nothing going red. (finding B3)"""
+    no_pf = {"prefetcher": False, "precision": "aaaaaaaaaaaa"}
+    with_pf = {"prefetcher": True, "precision": "aaaaaaaaaaaa"}
+    _touch_tensorbin(tmp_path)
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, build_variant=no_pf, **SHARED_ID)
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, build_variant=with_pf, **SHARED_ID)
+
+    # Both warm at once -- the second seed did not evict the first.
+    assert weight_cache_is_complete(tmp_path, build_variant=no_pf, **SHARED_ID) is True
+    assert weight_cache_is_complete(tmp_path, build_variant=with_pf, **SHARED_ID) is True
+    assert marker_path(tmp_path, no_pf) != marker_path(tmp_path, with_pf)
+    # A variant nobody seeded stays cold.
+    assert (
+        weight_cache_is_complete(tmp_path, build_variant={"prefetcher": False, "precision": "b"}, **SHARED_ID) is False
+    )
+    # And each variant's builder reads its own manifest.
+    sd = build_cached_state_dict(tmp_path, build_variant=with_pf)
+    assert set(sd.keys()) == set(SAMPLE_SD.keys())
+
+
+def test_unverifiable_variant_fails_closed(tmp_path):
+    """A build variant that could not be computed must never certify or match a cache: the gate
+    returns False and mark refuses to write, so the run cold-loads instead of risking a
+    placeholder persist under an unchecked filename set. (finding R3)"""
+    bad = {"unverifiable": True, "error": "RuntimeError: boom"}
+    _touch_tensorbin(tmp_path)
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, build_variant=bad, **SHARED_ID)
+    assert not list(tmp_path.glob(f"{WEIGHT_CACHE_MARKER}*")), "unverifiable variant must not write a marker"
+    # Even with a marker forged at the matching path, the gate rejects the request side.
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, **SHARED_ID)  # legit .none marker
+    assert weight_cache_is_complete(tmp_path, build_variant=bad, **SHARED_ID) is False
+
+
+def test_modelargs_variant_error_disables_skip(tmp_path):
+    """ModelArgs path: if computing the precision signature raises, the sentinel flows through
+    identity -> gate -> False (cold load), and marking is refused -- instead of the old behaviour
+    of collapsing to a match-anything 'unknown'. (finding R3)"""
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+
+    def _boom(decoder_id, tensor, prefetcher=False):
+        raise RuntimeError("precision config unavailable")
+
+    args.optimizations = SimpleNamespace(get_tensor_dtype=_boom)
+    variant = args._weight_cache_build_variant()
+    assert variant.get("unverifiable") is True
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    assert not list(Path(tmp_path).glob(f"{WEIGHT_CACHE_MARKER}*"))
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_hf_rope_is_part_of_the_variant(tmp_path):
+    """load_state_dict permutes QKV differently per rope mode under the SAME cache filenames, so a
+    marker seeded in one mode must not certify the other. (finding R2)"""
+    args = _FakeArgs(tmp_path)
+    _touch_tensorbin(tmp_path)
+    args.mark_weight_cache_complete(DTYPE, SAMPLE_SD)
+    assert args.weight_cache_is_complete(DTYPE) is True
+    args.use_hf_rope = True
+    assert args.weight_cache_is_complete(DTYPE) is False
+
+
+def test_sidecar_loaded_once_per_warm_run(tmp_path, monkeypatch):
+    """The completeness gate's validation load must be reused by build_cached_state_dict, not
+    repeated -- the sidecar can be multi-GB on NAS. The builder consumes the memoized entry so the
+    tensors are not pinned afterwards. (finding R1)"""
+    _seed(tmp_path, is_host_weight=lambda k: k == "tok_embeddings.weight")
+
+    real_load = torch.load
+    calls = []
+
+    def counting_load(*a, **k):
+        calls.append(a[0] if a else k.get("f"))
+        return real_load(*a, **k)
+
+    monkeypatch.setattr(torch, "load", counting_load)
+    _wc._SIDECAR_CACHE.clear()
+
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is True
+    sd = build_cached_state_dict(tmp_path)
+    assert torch.equal(sd["tok_embeddings.weight"], SAMPLE_SD["tok_embeddings.weight"])
+    assert len(calls) == 1, f"sidecar torch.load'ed {len(calls)}x per warm run, expected 1"
+    assert not _wc._SIDECAR_CACHE, "builder must consume the memoized sidecar entry"
+
+
+def test_mark_survives_non_oserror_write_failure(tmp_path, monkeypatch):
+    """mark_weight_cache_complete only records completion; a failure to record must warn, not
+    raise. torch.save on a read-only mount raises RuntimeError from the C++ serializer (not
+    OSError), which crashed every read-only cold run of the sidecar models right after a
+    successful build. (finding R5, Gemma-4-E2B bh_p150, run 32511945147)"""
+    _touch_tensorbin(tmp_path)
+
+    def _ro_save(*a, **k):
+        raise RuntimeError(
+            "[enforce fail at inline_container.cc:747] . open file failed with strerror: Read-only file system"
+        )
+
+    monkeypatch.setattr(torch, "save", _ro_save)
+    # Must not raise, and must not publish a marker that claims a sidecar it could not write.
+    mark_weight_cache_complete(tmp_path, SAMPLE_SD, is_host_weight=lambda k: k == "tok_embeddings.weight", **SHARED_ID)
+    assert weight_cache_is_complete(tmp_path, **SHARED_ID) is False

image/index.json

@@ -0,0 +1 @@
+{"schemaVersion":2,"mediaType":"application/vnd.oci.image.index.v1+json","manifests":[{"mediaType":"application/vnd.oci.image.index.v1+json","digest":"sha256:c4805783cb377cf99c36c43a26f9cf41f4605732b8dd244d9a11211a9d290708","size":856,"annotations":{"io.containerd.image.name":"docker.io/tt-model/qwen3.8-flash-next-p300x2:c4805783cb37","org.opencontainers.image.ref.name":"c4805783cb37"}}]}

image/manifest.json

@@ -0,0 +1 @@
+[{"Config":"blobs/sha256/2d4ec0ee7e1d8e2afa34d4462fd36e238adc97abc08950df6826e1ba6ae51b0a","RepoTags":["tt-model/qwen3.8-flash-next-p300x2:c4805783cb37"],"Layers":["blobs/sha256/d544298cabd50e7c86bfef1e52b67f01db6b3a57bfecfe37a851873dee83e52a","blobs/sha256/926bdef455064f157521bce0256de58abdbaab511c6c27e6159d549865f318fa","blobs/sha256/5ab8bdd2eb9167ebfb3bcfbf8c8a130f606bdc70c07b847c0a7fdaa502049043","blobs/sha256/df309695d498e77bd9a50ea7f3050e22efd1eae88c6611dbc89606e7e414ea78","blobs/sha256/e59b07f819fd19b7b5194adc92ccb10272747a70622188426f6ea64ba4775dbd","blobs/sha256/22aac323a4902588424acb82548cb3942084b85d5002631163f0732023598ff6","blobs/sha256/bf170f51420a62eeb56b8f6afb7a86e8f78ff820ccbbf7c9f543dcf76230aab5","blobs/sha256/bfe3bd79fd494c00040a96308433fad0dfaa9c8599707dcfa728cf9aac71c63e","blobs/sha256/e6675335cae939ad7a92ef21ad2b5afcf94f773a947383d7be56ac2af240a2ca","blobs/sha256/0741ec9077e1cca75a2621bc81f592f783376ffcbcd3d28b448336bdd57df894","blobs/sha256/cb9d04479d4428f62ea6193b0afa5c626ad8c5cb85af4bb3707b628747c4f303","blobs/sha256/a8d0a837d932862db7e85f056a9739d2fd6e1a44e4187266c268a6ca3eb2147c","blobs/sha256/2c83578d6015b0ec43dc4f1776e4ee53bd60491c15de2ad76c01b7f4fb6cb1c1","blobs/sha256/be22a17157f88c31873fddab027afdabbc75acbd737e4e6832f5a27147888915","blobs/sha256/5380a77208b9d0a514bc798ae2f366739a45192c19ac74005fea7eb2432a55a5","blobs/sha256/0babaa50033ef09eb22eeb3a79677044a82f7fc85b37ae8ca184d47f27c5c150","blobs/sha256/4637c6b0dd8ff29e9c09c442bcc91a26734b5df20edbcdce665497db5460bb65","blobs/sha256/5ff7f4251980e087f0e4491f352436860694fd3cd5cf13ce6b88c87e9da9ea10","blobs/sha256/4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1","blobs/sha256/d9efe2bed67066cab065e775a92326ebe261059ff9dc356d11fd5a51f778285d","blobs/sha256/532e399e6e08e45df8593c1496794824dcc34b86681a0dde5d8d166b02b5615f","blobs/sha256/a697d5857e3141dfa04762d522373c42fae90b16056bc40371802fa4574d5b25"]}]

image/oci-layout

@@ -0,0 +1 @@
+{"imageLayoutVersion":"1.0.0"}

requirements.lock

@@ -0,0 +1,423 @@
+accelerate==1.7.0
+aiofiles==25.1.0
+aiohappyeyeballs==2.7.1
+aiohttp==3.14.3
+aiosignal==1.4.0
+alabaster==0.7.16
+annotated-doc==0.0.5
+annotated-types==0.8.0
+anthropic==0.122.0
+anyio==4.14.2
+apache-tvm-ffi==0.1.9
+argcomplete==3.7.2
+argon2-cffi==25.1.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+astor==0.8.1
+asttokens==3.0.2
+async-lru==2.3.0
+attrs==26.1.0
+audioread==3.1.0
+av==18.0.0
+babel==2.18.0
+backcall==0.2.0
+backports-zstd==1.6.0
+beautifulsoup4==4.15.0
+beautifultable==1.1.0
+bert-score==0.3.12
+black==26.3.1
+blake3==1.0.9
+bleach==6.4.0
+blessed==1.48.0
+blinker==1.9.0
+blobfile==3.1.0
+bokeh==3.8.2
+boto3==1.43.68
+botocore==1.43.68
+breathe==4.35.0
+brotli==1.2.0
+build==0.10.0
+cachetools==7.1.7
+cbor2==6.1.4
+certifi==2026.7.22
+cffi==2.1.1
+cfgv==3.5.0
+charset-normalizer==3.4.9
+choreographer==1.2.1
+clang-format==19.1.4
+click==8.1.7
+cloudpickle==3.1.2
+comm==0.2.3
+compressed-tensors==0.17.0
+contourpy==1.3.3
+coverage==7.15.4
+cryptography==50.0.0
+cycler==0.12.1
+dacite==1.7.0
+dash==2.15.0
+dash-core-components==2.0.0
+dash-html-components==2.0.0
+dash-table==5.0.0
+datasets==2.21.0
+debugpy==1.8.21
+decorator==5.3.1
+defusedxml==0.7.1
+deprecated==1.3.1
+deprecation==2.1.0
+depyf==0.20.0
+detect-installer==0.1.0
+diffusers==0.38.0
+dill==0.3.8
+diskcache==5.6.3
+distlib==0.4.3
+distro==1.9.0
+dnspython==2.8.0
+docopt==0.6.2
+docstring-parser==0.18.0
+docutils==0.18.1
+editorconfig==0.17.1
+efficientnet-pytorch==0.7.1
+einops==0.6.1
+email-validator==2.3.0
+enlighten==1.12.4
+evaluate==0.4.0
+executing==2.2.1
+fastapi==0.136.3
+fastapi-cli==0.0.32
+fastapi-cloud-cli==0.23.0
+fastar==0.11.0
+faster-fifo==1.5.2
+fastjsonschema==2.22.1
+fastnumbers==5.1.1
+fastsafetensors==0.3.3
+filelock==3.32.2
+flask==3.0.3
+fonttools==4.63.0
+fqdn==1.5.1
+frozenlist==1.8.0
+fsspec==2023.9.2
+ftfy==6.1.1
+future==1.0.0
+gitdb==4.0.12
+gitpython==3.1.55
+glob2==0.7
+googleapis-common-protos==1.75.0
+graphql-core==3.2.11
+graphviz==0.21
+grpcio==1.83.0
+h11==0.16.0
+h2==4.4.1
+hf-xet==1.6.0
+hpack==4.2.0
+httpcore==1.0.9
+httpcore2==2.10.0
+httptools==0.8.0
+httpx==0.28.1
+httpx2==2.10.0
+huggingface-hub==1.16.1
+humanize==4.16.0
+humming-kernels==0.1.6
+hypercorn==0.18.0
+hyperframe==6.1.0
+identify==2.6.19
+idna==3.18
+ijson==3.5.1
+imageio==2.37.4
+imagesize==2.0.0
+importlib-metadata==9.0.0
+inflate64==1.0.4
+iniconfig==2.3.0
+interegular==0.3.3
+ipykernel==7.3.0
+ipython==8.12.3
+ipywidgets==8.1.1
+isoduration==20.11.0
+itsdangerous==2.2.0
+jaraco-classes==3.4.0
+jaraco-context==6.1.2
+jaraco-functools==4.6.0
+jedi==0.20.0
+jeepney==0.9.0
+jinja2==3.1.6
+jinxed==2.1.0
+jiter==0.16.0
+jiwer==3.0.5
+jmespath==1.1.0
+joblib==1.5.3
+jsbeautifier==1.14.7
+json5==0.15.0
+jsonlines==4.0.0
+jsonpointer==3.1.1
+jsonschema==4.26.0
+jsonschema-specifications==2025.9.1
+jupyter-client==8.9.1
+jupyter-core==5.9.1
+jupyter-events==0.12.1
+jupyter-lsp==2.3.1
+jupyter-server==2.20.0
+jupyter-server-terminals==0.5.4
+jupyterlab==4.5.9
+jupyterlab-pygments==0.3.0
+jupyterlab-server==2.28.0
+jupyterlab-widgets==3.0.16
+kagglehub==0.2.9
+kaleido==1.2.0
+keyring==25.7.0
+kiwisolver==1.5.0
+lark==1.2.2
+lazy-loader==0.5
+librosa==0.10.0
+lightning-utilities==0.15.3
+llguidance==1.7.6
+llvmlite==0.47.0
+lm-format-enforcer==0.11.3
+logistro==2.0.1
+loguru==0.6.0
+lxml==6.1.0
+markdown-it-py==3.0.0
+markupsafe==3.0.3
+matplotlib==3.10.9
+matplotlib-inline==0.2.2
+mcp==2.0.0
+mcp-types==2.0.0
+mdit-py-plugins==0.6.1
+mdurl==0.1.2
+medpy==0.5.2
+mistral-common==1.11.7
+mistune==3.3.4
+ml-dtypes==0.5.4
+model-hosting-container-standards==0.1.16
+mongoengine==0.24.2
+more-itertools==11.1.0
+motor==3.5.3
+mpmath==1.3.0
+msgpack==1.2.1
+msgspec==0.21.1
+multidict==6.7.1
+multiprocess==0.70.14
+multivolumefile==0.2.3
+mypy==1.9.0
+mypy-extensions==1.1.0
+myst-parser==3.0.0
+narwhals==2.24.0
+nbclient==0.11.0
+nbconvert==7.17.1
+nbformat==5.11.0
+nbsphinx==0.9.3
+nest-asyncio==1.6.0
+nest-asyncio2==1.7.2
+networkx==3.1
+nh3==0.3.6
+ninja==1.13.0
+nodeenv==1.10.0
+notebook-shim==0.2.4
+numba==0.65.0
+numpy==1.26.4
+open-clip-torch==2.26.1
+openai==3.1.0
+openai-harmony==0.0.8
+opencv-python-headless==4.11.0.86
+opentelemetry-api==1.44.0
+opentelemetry-exporter-otlp==1.44.0
+opentelemetry-exporter-otlp-proto-common==1.44.0
+opentelemetry-exporter-otlp-proto-grpc==1.44.0
+opentelemetry-exporter-otlp-proto-http==1.44.0
+opentelemetry-proto==1.44.0
+opentelemetry-sdk==1.44.0
+opentelemetry-semantic-conventions==0.65b0
+opentelemetry-semantic-conventions-ai==0.5.1
+orjson==3.11.9
+outlines-core==0.2.14
+packaging==26.3
+pandas==2.3.3
+pandoc==2.3
+pandocfilters==1.5.1
+parso==0.8.7
+partial-json-parser==0.2.1.1.post7
+pathspec==1.1.1
+peft==0.19.1
+pexpect==4.9.0
+pickleshare==0.7.5
+pillow==12.3.0
+pkginfo==1.12.1.2
+platformdirs==3.11.0
+plotly==5.18.0
+pluggy==1.6.0
+plumbum==2.0.2
+ply==3.11
+pooch==1.9.0
+pprintpp==0.4.0
+pre-commit==3.5.0
+prefixed==0.9.0
+priority==2.0.0
+prometheus-client==0.26.0
+prometheus-fastapi-instrumentator==8.1.0
+prompt-toolkit==3.0.52
+propcache==0.5.2
+protobuf==5.29.6
+psutil==7.2.2
+psycopg2-binary==2.9.11
+ptyprocess==0.7.0
+pure-eval==0.2.3
+py-cpuinfo==9.0.0
+py7zr==1.1.3
+pyarrow==21.0.0
+pybase64==1.5.0
+pybcj==1.0.8
+pycapnp==2.0.0
+pycountry==26.2.16
+pycparser==3.0
+pycryptodomex==3.23.0
+pydantic==2.13.4
+pydantic-core==2.46.4
+pydantic-extra-types==2.11.1
+pydantic-settings==2.15.0
+pydash==8.0.6
+pyelftools==0.33
+pygments==2.20.0
+pyjwt==2.13.0
+pymongo==4.8.0
+pyparsing==3.3.2
+pyppmd==1.3.1
+pyproject-hooks==1.2.0
+pyro5==5.16
+pytest==9.0.3
+pytest-benchmark==5.2.3
+pytest-cov==7.0.0
+pytest-github-actions-annotate-failures==0.3.0
+pytest-repeat==0.9.4
+pytest-split==0.11.0
+pytest-timeout==2.4.0
+python-dateutil==2.9.0.post0
+python-discovery==1.5.1
+python-dotenv==1.2.2
+python-json-logger==4.1.0
+python-multipart==0.0.32
+pytokens==0.4.1
+pytz==2026.3.post1
+pyyaml==6.0.3
+pyzmq==27.1.0
+quack-kernels==0.5.0
+qwen-vl-utils==0.0.14
+rapidfuzz==3.14.5
+rapidyaml==0.12.0
+rarfile==4.5
+readme-renderer==43.0
+referencing==0.37.0
+regex==2026.7.19
+requests==2.34.2
+requests-toolbelt==1.0.0
+responses==0.18.0
+retrying==1.4.2
+rfc3339-validator==0.1.4
+rfc3986==2.0.0
+rfc3986-validator==0.1.1
+rfc3987-syntax==1.1.0
+rich==15.0.0
+rich-toolkit==0.20.3
+rignore==0.8.1
+rpds-py==2026.6.3
+s3transfer==0.19.2
+safetensors==0.8.0
+scikit-image==0.26.0
+scikit-learn==1.9.0
+scipy==1.17.1
+seaborn==0.13.2
+secretstorage==3.5.0
+send2trash==2.1.0
+sentencepiece==0.2.1
+sentry-sdk==2.68.0
+serpent==1.43
+setproctitle==1.3.7
+setuptools==80.10.2
+shellingham==1.5.4
+simpleitk==2.5.6
+simplejson==4.1.1
+six==1.17.0
+smmap==5.0.3
+sniffio==1.3.1
+snowballstemmer==3.1.1
+sortedcontainers==2.4.0
+soundfile==0.14.0
+soupsieve==2.9.2
+soxr==1.1.0
+sphinx==7.1.2
+sphinx-rtd-theme==1.3.0
+sphinxcontrib-applehelp==2.0.0
+sphinxcontrib-devhelp==2.0.0
+sphinxcontrib-email==0.3.5
+sphinxcontrib-htmlhelp==2.1.0
+sphinxcontrib-jquery==4.1
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==2.0.0
+sphinxcontrib-serializinghtml==2.0.0
+sse-starlette==3.4.8
+sseclient-py==1.9.0
+stack-data==0.6.3
+starlette==1.6.0
+strawberry-graphql==0.138.1
+supervisor==4.3.0
+sympy==1.14.0
+tabulate==0.9.0
+tblib==3.2.2
+tenacity==9.1.4
+termcolor==3.3.0
+terminado==0.18.1
+texttable==1.7.0
+threadpoolctl==3.6.0
+tifffile==2026.3.3
+tiktoken==0.7.0
+tilelang==0.1.9
+timm==1.0.28
+tinycss2==1.5.1
+tokenizers==0.23.1
+tokenspeed-mla==0.1.2
+tokenspeed-triton==3.8.10.post20260721
+toolz==0.12.0
+torch==2.11.0+cpu
+torch-c-dlpack-ext==0.1.5
+torch-fidelity==0.3.0
+torchmetrics==1.9.0
+torchvision==0.26.0+cpu
+tornado==6.5.8
+tqdm==4.66.3
+traitlets==5.16.1
+transformers==5.16.0
+triton==3.7.1
+truststore==0.10.4
+tt-exalens==0.3.27
+tt-perf-report==1.2.8
+tt-umd==0.9.9
+twine==4.0.2
+typer==0.27.1
+typing-extensions==4.16.0
+typing-inspection==0.4.4
+tzdata==2026.3
+tzlocal==5.4.4
+universal-analytics-python3==1.1.1
+uri-template==1.3.0
+urllib3==2.7.0
+uv==0.9.26
+uvicorn==0.52.3
+uvloop==0.22.1
+virtualenv==21.7.4
+watchfiles==1.2.0
+wcwidth==0.8.2
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+websockets==16.0
+werkzeug==3.0.6
+wheel==0.47.0
+widgetsnbextension==4.0.15
+wrapt==2.3.0
+wsproto==1.3.2
+xgrammar==0.2.3
+xlsxwriter==3.0.8
+xmltodict==1.0.4
+xxhash==3.8.1
+xyzservices==2026.3.0
+yamllint==1.32.0
+yarl==1.24.5
+z3-solver==4.15.4.0
+zipp==4.1.0