Add files using upload-large-folder tool

README.md

@@ -0,0 +1,60 @@
+---
+tags:
+- blackhole
+- p300x2
+- tt-model-cache
+- tt-model-container
+- vllm-plugin
+---
+
+# nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026
+
+UNPUBLISHED TEST VARIANT on stock vLLM 0.26.0 + the standalone vllm-tt-plugin, with NVIDIA's recommended nemotron_v3 reasoning parser and qwen3_coder tool parser. NVIDIA Nemotron 3.5 Lightning 30B-A3B (NVFP4 checkpoint): a hybrid Mamba-Transformer mixture-of-experts reasoning model with a 1,048,576-token context, ported to Tenstorrent Blackhole and served through vLLM on a QuietBox2 (2x p300c, 4 chips).
+
+Runs on **p300x2** (mesh `P300x2`) — 1,048,576-token context, up to 32 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  ndaly/Nemotron-3.5-Lightning-30B-A3B-NVFP4-tt-p300x2-vllm026 --with-weights
+tt-model serve ndaly/Nemotron-3.5-Lightning-30B-A3B-NVFP4-tt-p300x2-vllm026
+```
+
+`pull --with-weights` downloads the Docker image and the [`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4) weights at `6dbbd757ea75a8ece6e0702872e3ae53f9987728` (into your HF cache; they are not in the image). `serve` starts an OpenAI-compatible server on port 20000 (or the next free port, if that one is busy); the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs `Application startup complete`.
+
+### What was validated
+Full 1M-token context on a QuietBox2 (P300x2 mesh, 4 Blackhole chips). Release validation
+(2026-08-19, kernel 7.0.0-28): TPOT 15.4 ms (~65 tok/s/user) single-stream, TTFT 109 ms,
+22/22 spec-test conformance. This container was re-verified on the same box on 2026-09-09
+(kernel 7.0.0-30): identity, deterministic generation and the separate `reasoning` field
+all pass; TTFT 110 ms; single-stream TPOT measured 42 ms (~24 tok/s/user) — identical to a
+host-native run of the bring-up environment on that day, so the gap is the host, not the
+container. Reasoning is returned in a separate `reasoning` field (`deepseek_r1` parser); with
+`chat_template_kwargs: {"enable_thinking": false}` the parser has no `</think>` to split on
+and the answer lands in `reasoning` with empty `content`.
+
+### Try it
+```bash
+curl -s localhost:20000/v1/chat/completions -H 'Content-Type: application/json' -d '{
+  "model": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4",
+  "messages": [{"role": "user", "content": "Merge two sorted lists in Python."}],
+  "max_tokens": 2048, "temperature": 0}'
+```
+The first boot converts the NVFP4 weights to device layout and JIT-compiles kernels (~12 min
+on a QuietBox2); both are cached under `~/.cache/tt-model/<name>/` so later boots take ~4 min.
+After any non-clean stop, run `tt-smi -r all` before serving again.
+
+## 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.26.0`](https://github.com/vllm-project/vllm/releases/tag/v0.26.0) |
+| vllm-tt-plugin | a local checkout — commit not published |
+| `code/` digest | `5787b8dfce9d7e38` (sha256, first 16 hex digits) |
+| built | 2026-09-11T17:47:54+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/demos/llama31_8B_demo.py

@@ -0,0 +1,823 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+TTTv2 MLP1D Demo Script for Llama 3.1-8B-Instruct
+
+Demonstrates the TTTv2 MLP1D module by replacing the MLP layers in Llama 3.1-8B-Instruct model.
+
+Measurements:
+- Top-1/Top-5 token accuracy (like `ci-token-matching`)
+- tok/s/u and TTFT performance (like `batch-1` and `batch-32`)
+
+Supports:
+- `performance` and `accuracy` optimization modes
+- N150 (1x1), N300 (1x2), T3K (1x8) topologies
+
+Usage:
+    # Token accuracy test on N150
+    MESH_DEVICE=N150 HF_MODEL=meta-llama/Llama-3.1-8B-Instruct \
+    python_env/bin/pytest models/common/demos/llama31_8B_demo.py -k "token-accuracy" -v
+
+    # Batch-1 latency test on N300
+    MESH_DEVICE=N300 HF_MODEL=meta-llama/Llama-3.1-8B-Instruct \
+    python_env/bin/pytest models/common/demos/llama31_8B_demo.py -k "batch-1" --performance -v
+
+    # Batch-32 throughput test on T3K
+    MESH_DEVICE=T3K HF_MODEL=meta-llama/Llama-3.1-8B-Instruct \
+    python_env/bin/pytest models/common/demos/llama31_8B_demo.py -k "batch-32" --performance -v
+
+"""
+
+import json
+import os
+import re
+import subprocess
+import sys
+from contextlib import contextmanager
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.sampling import SamplingParams
+from models.perf.benchmarking_utils import BenchmarkProfiler
+from models.tt_transformers.tt.common import (
+    PagedAttentionConfig,
+    create_tt_model,
+    preprocess_inputs_prefill,
+    sample_host,
+)
+from models.tt_transformers.tt.generator import Generator
+from models.tt_transformers.tt.model_config import DecodersPrecision
+
+# =============================================================================
+# Constants and Expected Metrics
+# =============================================================================
+
+# Expected accuracy metrics from centralized targets for Llama-3.1-8B (top1, top5 only)
+# Performance metrics (tok_s_u, ttft_ms) are collected dynamically by running
+# simple_text_demo.py with the corresponding test case to get real baseline values.
+# Higher is better for tok_s_u, top1, top5
+# Lower is better for ttft_ms
+EXPECTED_METRICS = {
+    "performance": {
+        "N150": {"top1": 90, "top5": 97, "tok_s_u": 28.3, "ttft_ms": 104},
+        "N300": {"top1": 90, "top5": 97, "tok_s_u": 44.2, "ttft_ms": 67},
+        "T3K": {"top1": 90, "top5": 98, "tok_s_u": 64.3, "ttft_ms": 53},
+    },
+    "accuracy": {
+        "N150": {"top1": 96, "top5": 100, "tok_s_u": 25.2, "ttft_ms": 138},
+        "N300": {"top1": 96, "top5": 100, "tok_s_u": 38.8, "ttft_ms": 79},
+        "T3K": {"top1": 97, "top5": 100, "tok_s_u": 60.8, "ttft_ms": 81},
+    },
+}
+
+# Tolerance for performance validation (5%)
+PERF_TOLERANCE = 0.05
+
+
+# =============================================================================
+# Baseline Collection from simple_text_demo.py
+# =============================================================================
+
+
+def collect_baseline_from_simple_text_demo(
+    device_name: str,
+    batch_size: int,
+    opt_mode: str,
+) -> dict | None:
+    """
+    Run simple_text_demo.py to collect real baseline performance metrics (tok_s_u, ttft_ms).
+
+    This provides an accurate comparison by measuring TTTv1 MLP performance
+    on the same hardware under the same conditions, rather than relying on
+    potentially outdated markdown reference values.
+
+    Args:
+        device_name: Device name (N150, N300, T3K)
+        batch_size: Batch size (1 or 32)
+        opt_mode: Optimization mode (performance or accuracy)
+
+    Returns:
+        Dict with baseline metrics {tok_s_u, ttft_ms} or None if collection fails
+    """
+    # Build the test ID for simple_text_demo
+    batch_id = f"batch-{batch_size}"
+    test_filter = f"{opt_mode}-{batch_id}"
+
+    # Get HF_MODEL from environment
+    hf_model = os.environ.get("HF_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
+
+    logger.info(f"Collecting baseline from simple_text_demo.py ({test_filter})...")
+
+    # Build pytest command (use sys.executable for CI compatibility)
+    cmd = [
+        sys.executable,
+        "-m",
+        "pytest",
+        "models/tt_transformers/demo/simple_text_demo.py",
+        "-k",
+        test_filter,
+        "-v",
+        "--tb=short",
+    ]
+
+    env = os.environ.copy()
+    env["MESH_DEVICE"] = device_name
+    env["HF_MODEL"] = hf_model
+    # Unset CI to ensure tests run normally (CI=true may skip or alter test behavior)
+    env.pop("CI", None)
+
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=600,  # 10 minute timeout
+            env=env,
+            cwd=os.getcwd(),
+        )
+
+        output = result.stdout + result.stderr
+
+        # Parse the output for performance metrics
+        # Looking for lines like:
+        # "Average Time to First Token (TTFT): 68.09ms"
+        # "Average speed: 57.89ms @ 17.28 tok/s/user (17.28 tok/s throughput)"
+        ttft_match = re.search(r"Average Time to First Token \(TTFT\):\s*([\d.]+)ms", output)
+        speed_match = re.search(r"Average speed:.*@\s*([\d.]+)\s*tok/s/user", output)
+
+        if ttft_match and speed_match:
+            baseline = {
+                "ttft_ms": float(ttft_match.group(1)),
+                "tok_s_u": float(speed_match.group(1)),
+            }
+            logger.info(f"Baseline collected: {baseline['tok_s_u']:.2f} tok/s/u, TTFT {baseline['ttft_ms']:.2f}ms")
+            return baseline
+        else:
+            logger.warning(f"Could not parse baseline metrics from simple_text_demo output")
+            if "PASSED" not in output and "passed" not in output:
+                logger.warning(f"Test may have failed. Exit code: {result.returncode}")
+            return None
+
+    except subprocess.TimeoutExpired:
+        logger.warning("Baseline collection timed out after 10 minutes")
+        return None
+    except Exception as e:
+        logger.warning(f"Baseline collection failed: {e}")
+        return None
+
+
+# =============================================================================
+# Token Accuracy Helper
+# =============================================================================
+
+
+class TokenAccuracy:
+    """Helper class for measuring token accuracy against reference data."""
+
+    def __init__(self, model_name: str):
+        self.gt_pos = -1
+        self.store_predicted_tokens = []
+        reference_data_file = os.path.join("models/tt_transformers/tests/reference_outputs/", model_name) + ".refpt"
+        if not os.path.exists(reference_data_file):
+            raise FileNotFoundError(f"Reference data file not found: {reference_data_file}")
+
+        logger.info(f"Loading reference data from {reference_data_file}")
+        reference_data = torch.load(reference_data_file)
+        reference_tokens = reference_data["reference_tokens"]
+        split_point = reference_tokens.shape[-1] // 2
+        self.input_prompt = reference_tokens[0, :split_point]
+        self.reference_tokens = reference_tokens[0, split_point:]
+        self.top5_tokens = reference_data["top5_tokens"][split_point - 1 :, :]
+        self.maxindex = len(self.reference_tokens) - 1
+
+    def prepare_ref_tokens(self, tokenizer):
+        """Decode input prompt for reference."""
+        return tokenizer.decode(self.input_prompt.tolist())
+
+    def collect_predicted_tokens(self, tokens):
+        """Collect predicted tokens and return the next reference token (teacher forcing)."""
+        self.store_predicted_tokens.append(tokens)
+        self.gt_pos += 1
+        return self.reference_tokens[min(self.gt_pos, self.maxindex)].unsqueeze(-1).unsqueeze(-1)
+
+    def compute_accuracy(self) -> tuple[float, float]:
+        """Compute Top-1 and Top-5 accuracy."""
+        count = 0
+        count_t5 = 0
+        matching_sz = min(len(self.reference_tokens), len(self.store_predicted_tokens))
+        for i in range(matching_sz):
+            if self.top5_tokens[i, 0].item() == self.store_predicted_tokens[i]:
+                count += 1
+            if self.store_predicted_tokens[i] in self.top5_tokens[i, :]:
+                count_t5 += 1
+        accuracy_top1 = count / matching_sz
+        accuracy_top5 = count_t5 / matching_sz
+
+        return accuracy_top1 * 100, accuracy_top5 * 100
+
+
+# =============================================================================
+# MLP1D Patching Context Manager
+# =============================================================================
+
+
+@contextmanager
+def patch_mlp_with_mlp1d():
+    """
+    Context manager that patches decoder.MLP with MLP1D.from_model_args.
+
+    The signatures now match, enabling drop-in replacement without adapters.
+    """
+    from models.common.modules.mlp.mlp_1d import MLP1D
+    from models.tt_transformers.tt import decoder
+
+    # Store original MLP class
+    original_MLP = decoder.MLP
+
+    # Replace with MLP1D.from_model_args - signatures match!
+    decoder.MLP = MLP1D.from_model_args
+
+    try:
+        yield
+    finally:
+        # Restore original MLP class
+        decoder.MLP = original_MLP
+
+
+def create_tt_model_with_mlp1d(
+    mesh_device,
+    instruct: bool,
+    max_batch_size: int,
+    optimizations,
+    max_seq_len: int,
+    paged_attention_config,
+    dtype,
+    state_dict=None,
+    num_layers=None,
+):
+    """
+    Wrapper around create_tt_model that patches MLP → MLP1D.from_model_args.
+
+    With matching signatures, this is now a simple class swap without adapters.
+    """
+    with patch_mlp_with_mlp1d():
+        model_args, model, tt_kv_cache, state_dict = create_tt_model(
+            mesh_device,
+            instruct=instruct,
+            max_batch_size=max_batch_size,
+            optimizations=optimizations,
+            max_seq_len=max_seq_len,
+            paged_attention_config=paged_attention_config,
+            dtype=dtype,
+            state_dict=state_dict,
+            num_layers=num_layers,
+        )
+
+    return model_args, model, tt_kv_cache, state_dict
+
+
+# =============================================================================
+# Performance Validation
+# =============================================================================
+
+
+def validate_metrics(measured: dict, expected: dict, tolerance: float = PERF_TOLERANCE) -> tuple[bool, str]:
+    """
+    Validate measured metrics against expected values with tolerance.
+
+    For performance tests: expected comes from simple_text_demo.py baseline (tok_s_u, ttft_ms).
+    For accuracy tests: expected comes from centralized model targets (top1, top5).
+
+    Returns:
+        (passed, message) tuple
+    """
+    failures = []
+    for metric, expected_val in expected.items():
+        if expected_val is None:
+            continue
+        measured_val = measured.get(metric)
+        if measured_val is None:
+            continue
+
+        # Higher is better for tok_s_u, top1, top5
+        # Lower is better for ttft_ms
+        if metric == "ttft_ms":
+            if measured_val > expected_val * (1 + tolerance):
+                failures.append(
+                    f"{metric}: measured {measured_val:.2f} > expected {expected_val * (1 + tolerance):.2f}"
+                )
+        else:
+            if measured_val < expected_val * (1 - tolerance):
+                failures.append(
+                    f"{metric}: measured {measured_val:.2f} < expected {expected_val * (1 - tolerance):.2f}"
+                )
+
+    if failures:
+        return False, "; ".join(failures)
+    return True, "All metrics within tolerance"
+
+
+def get_device_name_from_mesh_shape(mesh_shape: tuple[int, int]) -> str:
+    """Map mesh_shape to device name."""
+    mapping = {
+        (1, 1): "N150",
+        (1, 2): "N300",
+        (1, 8): "T3K",
+    }
+    return mapping.get(mesh_shape, f"Unknown({mesh_shape})")
+
+
+# =============================================================================
+# Test Parametrization
+# =============================================================================
+
+
+def _get_mesh_shape():
+    """Get mesh shape from MESH_DEVICE environment variable."""
+    shape_map = {
+        "N150": (1, 1),
+        "N300": (1, 2),
+        "T3K": (1, 8),
+    }
+    return shape_map.get(os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids()))
+
+
+# Test configurations
+TEST_CONFIGS = [
+    # Token accuracy test (like ci-token-matching)
+    pytest.param(
+        {
+            "name": "token-accuracy",
+            "batch_size": 1,
+            "max_generated_tokens": 500,
+            "measure_accuracy": True,
+        },
+        id="token-accuracy",
+    ),
+    # Latency test (like batch-1)
+    pytest.param(
+        {
+            "name": "batch-1-latency",
+            "batch_size": 1,
+            "max_generated_tokens": 200,
+            "measure_accuracy": False,
+        },
+        id="batch-1",
+    ),
+    # Throughput test (like batch-32)
+    pytest.param(
+        {
+            "name": "batch-32-throughput",
+            "batch_size": 32,
+            "max_generated_tokens": 200,
+            "measure_accuracy": False,
+        },
+        id="batch-32",
+    ),
+]
+
+# Global cache for baseline metrics (collected before any test opens mesh_device)
+_baseline_cache: dict[str, dict] = {}
+_baseline_collected: bool = False
+
+
+def _get_baseline_cache_key(device_name: str, batch_size: int, opt_mode: str) -> str:
+    """Generate cache key for baseline metrics."""
+    return f"{device_name}-batch{batch_size}-{opt_mode}"
+
+
+def _collect_all_baselines():
+    """
+    Collect all baseline metrics ONCE at module load time, before any tests run.
+
+    This ensures baselines are collected when no mesh_device is open.
+    """
+    global _baseline_collected
+    if _baseline_collected:
+        return
+
+    device_name = os.environ.get("MESH_DEVICE")
+    assert device_name is not None, "MESH_DEVICE environment variable is not set"
+    logger.info(f"=== Collecting baselines from simple_text_demo.py for {device_name} ===")
+
+    # Collect baselines for batch_size x opt_mode combinations
+    # CI runs batch-32, so skip batch-1 when CI=true
+    is_ci = os.environ.get("CI", "").lower() in ("true", "1", "yes")
+    batch_sizes = [32] if is_ci else [1, 32]
+
+    for batch_size in batch_sizes:
+        for opt_mode in ["performance", "accuracy"]:
+            cache_key = _get_baseline_cache_key(device_name, batch_size, opt_mode)
+            logger.info(f"Collecting baseline for {cache_key}...")
+            baseline = collect_baseline_from_simple_text_demo(
+                device_name=device_name,
+                batch_size=batch_size,
+                opt_mode=opt_mode,
+            )
+            if baseline:
+                _baseline_cache[cache_key] = baseline
+                logger.info(f"  ✓ {cache_key}: {baseline['tok_s_u']:.2f} tok/s/u, {baseline['ttft_ms']:.2f}ms TTFT")
+            else:
+                logger.warning(f"  ✗ {cache_key}: baseline collection failed")
+
+    _baseline_collected = True
+    logger.info("=== Baseline collection complete ===")
+
+
+def _get_cached_baseline(device_name: str, batch_size: int, opt_mode: str) -> dict | None:
+    """Get cached baseline metrics. Must be called after baselines are collected."""
+    cache_key = _get_baseline_cache_key(device_name, batch_size, opt_mode)
+    return _baseline_cache.get(cache_key)
+
+
+# Collect all baselines at module import time, before pytest sets up any fixtures.
+# This ensures baselines are collected when no mesh_device is open.
+_collect_all_baselines()
+
+
+# =============================================================================
+# Main Test Function
+# =============================================================================
+
+
+@pytest.mark.parametrize(
+    "device_params",
+    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "trace_region_size": 50000000, "num_command_queues": 1}],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "mesh_device",
+    [_get_mesh_shape()],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "optimizations",
+    [
+        pytest.param(
+            lambda args: DecodersPrecision.performance(args.n_layers, args.model_name),
+            id="performance",
+        ),
+        pytest.param(
+            lambda args: DecodersPrecision.accuracy(args.n_layers, args.model_name),
+            id="accuracy",
+        ),
+    ],
+)
+@pytest.mark.parametrize("test_config", TEST_CONFIGS)
+def test_mlp1d_llama_demo(
+    mesh_device: ttnn.MeshDevice,
+    optimizations,
+    test_config: dict,
+    is_ci_env: bool,
+    request,
+):
+    """
+    Demo test for MLP1D replacement in Llama 3.1-8B-Instruct.
+
+    This test:
+    1. Creates the TT model with MLP layers replaced by MLP1D
+    2. Runs prefill and decode inference
+    3. Measures performance (TTFT, tok/s/u) and/or accuracy (Top-1, Top-5)
+    4. Validates results against centralized targets
+    """
+    test_id = request.node.callspec.id
+
+    if is_ci_env and "batch-1" in test_id:
+        pytest.skip("CI only runs batch-32 and token-accuracy tests")
+
+    # Skip TG devices - MLP1D does not support Galaxy
+    mesh_shape = mesh_device.shape
+    if mesh_shape[0] > 1 or mesh_shape[1] > 8:
+        pytest.skip("MLP1D does not support TG/Galaxy devices")
+
+    # Get test parameters
+    test_name = test_config["name"]
+    batch_size = test_config["batch_size"]
+    max_generated_tokens = test_config["max_generated_tokens"]
+    measure_accuracy = test_config["measure_accuracy"]
+
+    # Get optimization mode from test ID
+    opt_mode = "performance" if "performance" in test_id else "accuracy"
+
+    # Get device name
+    device_name = get_device_name_from_mesh_shape(tuple(mesh_shape))
+
+    logger.info(f"=== MLP1D Demo: {test_name} ===")
+    logger.info(f"Device: {device_name}, Optimization: {opt_mode}, Batch: {batch_size}")
+
+    # Configuration
+    instruct = not measure_accuracy  # Use non-instruct for token accuracy (matches ci-token-matching)
+    max_seq_len = 1024
+    paged_attention = True
+    page_params = {"page_block_size": 32, "page_max_num_blocks_per_dp": 1024}
+    dtype = ttnn.bfloat8_b
+
+    # Paged attention config
+    paged_attention_config = (
+        PagedAttentionConfig(
+            block_size=page_params["page_block_size"],
+            max_num_blocks=page_params["page_max_num_blocks_per_dp"],
+        )
+        if paged_attention
+        else None
+    )
+
+    # Create model with MLP1D replacement
+    logger.info("Creating TT model with MLP1D replacement...")
+    model_args, model, tt_kv_cache, _ = create_tt_model_with_mlp1d(
+        mesh_device=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,
+    )
+
+    # Validate model was created successfully
+    assert model is not None, "Failed to create model"
+    assert model_args is not None, "Failed to get model args"
+
+    # Wrap in lists for API compatibility (like prepare_generator_args does for data parallel)
+    model_args_list = [model_args]
+    model_list = [model]
+    tt_kv_cache_list = [tt_kv_cache]
+
+    tokenizer = model_args.tokenizer
+
+    # Initialize profiler
+    profiler = BenchmarkProfiler()
+    profiler.start("run")
+
+    # Setup token accuracy if needed
+    token_acc = None
+    if measure_accuracy:
+        try:
+            token_acc = TokenAccuracy(model_name=model_args.model_name)
+        except FileNotFoundError as e:
+            pytest.skip(f"Reference data not found: {e}")
+
+    # Prepare input prompts
+    if measure_accuracy and token_acc:
+        input_prompts = [token_acc.prepare_ref_tokens(tokenizer)]
+    else:
+        # Use the same 128-token prompts as simple_text_demo.py for fair perf comparison
+        prompt_file = "models/tt_transformers/demo/sample_prompts/input_data_questions_prefill_128.json"
+        with open(prompt_file, "r") as f:
+            prompts_data = json.load(f)
+        # Repeat/truncate prompts to match batch_size
+        input_prompts = [prompts_data[i % len(prompts_data)]["prompt"] for i in range(batch_size)]
+
+    # Create page table for paged attention
+    page_table = None
+    if paged_attention_config:
+        permutation = torch.randperm(paged_attention_config.max_num_blocks)
+        reverse_permutation = torch.argsort(permutation)
+        page_table = reverse_permutation.reshape(batch_size, paged_attention_config.max_num_blocks // batch_size)
+
+    # Preprocess inputs
+    logger.info("Preprocessing inputs...")
+    (
+        input_tokens_prefill_pt,
+        encoded_prompts,
+        decoding_pos,
+        prefill_lens,
+    ) = preprocess_inputs_prefill(
+        input_prompts, tokenizer, model_args_list, instruct, max_generated_tokens, max_prefill_len=max_seq_len
+    )
+    input_tokens_prefill_pt = torch.stack(input_tokens_prefill_pt).view(batch_size, -1)
+
+    # Create generator
+    generator = Generator(model_list, model_args_list, mesh_device, tokenizer=tokenizer)
+
+    # --- Prefill Phase ---
+    logger.info("Starting prefill warmup...")
+    profiler.start("compile_prefill")
+    logits = generator.prefill_forward_text(
+        input_tokens_prefill_pt,
+        page_table=page_table,
+        kv_cache=tt_kv_cache_list,
+        prompt_lens=decoding_pos,
+    )
+    profiler.end("compile_prefill")
+
+    logger.info("Starting prefill inference...")
+    profiler.start("inference_prefill")
+    logits = generator.prefill_forward_text(
+        input_tokens_prefill_pt,
+        page_table=page_table,
+        kv_cache=tt_kv_cache_list,
+        prompt_lens=decoding_pos,
+    )
+    prefilled_token = torch.argmax(logits, dim=-1)
+    profiler.end("inference_prefill")
+
+    # --- Decode Phase ---
+    # Keep track of generated outputs to print out every iteration
+    all_outputs = [encoded_prompts[b][: prefill_lens[b]] for b in range(batch_size)]
+    for user in range(batch_size):
+        user_tok = int(prefilled_token[user].item())
+        all_outputs[user].append(user_tok)
+
+    logger.info("Starting decode loop...")
+    current_pos = torch.tensor([decoding_pos[b] for b in range(batch_size)])
+    out_tok = prefilled_token
+    user_done = [False] * batch_size
+
+    # Sampling params (argmax for deterministic results)
+    sampling_params = {"temperature": 0, "top_p": 0.08, "top_k": 32}
+    device_sampling_params = (
+        SamplingParams(
+            temperature=sampling_params["temperature"],
+            top_k=sampling_params["top_k"],
+            top_p=sampling_params["top_p"],
+        )
+        if model_list[0]._supports_on_device_sampling
+        else None
+    )
+
+    profiler.start("inference_decode")
+
+    for iteration in range(max_generated_tokens):
+        if iteration == 0:
+            profiler.start("compile_decode")
+        else:
+            profiler.start(f"inference_decode_time_{iteration}")
+
+        # Teacher forcing for token accuracy
+        if token_acc:
+            out_tok[0] = token_acc.collect_predicted_tokens(out_tok[0].item())
+
+        # Decode forward
+        logits, _ = generator.decode_forward(
+            out_tok,
+            current_pos,
+            enable_trace=not measure_accuracy,  # Disable trace for accuracy (teacher forcing)
+            page_table=page_table,
+            kv_cache=tt_kv_cache_list,
+            sampling_params=device_sampling_params,
+            prompt_tokens=input_tokens_prefill_pt,
+            output_tokens=out_tok,
+        )
+
+        # Get next token
+        if device_sampling_params is not None:
+            out_tok = logits.unsqueeze(1)
+        else:
+            _, out_tok = sample_host(
+                logits,
+                temperature=sampling_params["temperature"],
+                top_p=sampling_params["top_p"],
+                on_host=True,
+            )
+
+        if iteration == 0:
+            profiler.end("compile_decode")
+            decode_iteration_time = profiler.get_duration("compile_decode")
+        else:
+            profiler.end(f"inference_decode_time_{iteration}")
+            decode_iteration_time = profiler.get_duration(f"inference_decode_time_{iteration}")
+
+        # Log per-iteration performance
+        tokens_per_second_per_user = 1 / decode_iteration_time if decode_iteration_time > 0 else 0
+        logger.debug(
+            f"Iteration {iteration}: {1000 * decode_iteration_time:.0f}ms @ {tokens_per_second_per_user:.1f} tok/s/user"
+        )
+
+        current_pos += 1
+
+        # Save output token to print out later and check for EOS
+        for user in range(batch_size):
+            user_tok = out_tok[user].item()
+            if user_tok not in tokenizer.stop_tokens and not user_done[user]:
+                all_outputs[user].append(user_tok)
+            else:
+                if not measure_accuracy:  # For accuracy tests, keep decoding (teacher forcing)
+                    user_done[user] = True
+                    logger.trace(f"[User {user}] Finished decoding at iteration {iteration}")
+                else:
+                    all_outputs[user].append(user_tok)
+
+        # Print out generated outputs for each user (debug level)
+        for user in range(batch_size):
+            text = "".join(tokenizer.decode(all_outputs[user]))
+            if len(text) > 100:
+                text = "..." + text[-97:]
+            text = text.replace("\n", " ")
+            logger.debug(f"[User {user}] {text}")
+
+        if all(user_done):
+            break
+
+    profiler.end("inference_decode")
+
+    # Final print of generated text
+    logger.info("Finished decoding, printing the final outputs...\n")
+    for user, output in enumerate(all_outputs):
+        text = tokenizer.decode(output)
+        # Find where prompt ends and generated text begins
+        prompt_text = input_prompts[user] if user < len(input_prompts) else ""
+        # Strip leading/trailing newlines from generated portion
+        text_after_prompt = text[len(prompt_text) :] if text.startswith(prompt_text) else text
+        short_prompt = (
+            (prompt_text[:100] + "\n<long prompt not printed in full>\n" + prompt_text[-100:])
+            if len(prompt_text) > 200
+            else prompt_text
+        )
+        logger.info(f"\n==USER {user} - PROMPT\n{short_prompt}\n==USER {user} - OUTPUT\n{text_after_prompt.strip()}\n")
+    profiler.end("run")
+
+    # --- Compute Metrics ---
+    compile_prefill_time = profiler.get_duration("compile_prefill")
+    compile_decode_time = profiler.get_duration("compile_decode")
+    total_inference_prefill_time = profiler.get_duration("inference_prefill")
+
+    # Sum decode times (skip compile iteration)
+    total_inference_decode_time = 0
+    num_decode_iterations = min(iteration + 1, max_generated_tokens)
+    for i in range(1, num_decode_iterations):
+        total_inference_decode_time += profiler.get_duration(f"inference_decode_time_{i}")
+
+    # Calculate performance metrics
+    avg_ttft_ms = (total_inference_prefill_time / batch_size) * 1000
+    tok_s_u = (num_decode_iterations - 1) / total_inference_decode_time if num_decode_iterations > 1 else 0
+
+    # Compute accuracy if applicable
+    top1_acc = None
+    top5_acc = None
+    if token_acc:
+        top1_acc, top5_acc = token_acc.compute_accuracy()
+
+    # --- Report Results ---
+    logger.info("")
+    logger.info("=== Performance Metrics ===")
+    logger.info(f"Device: {device_name}")
+    logger.info(f"Mode: {opt_mode}")
+    logger.info(f"Compile prefill: {compile_prefill_time:.2f}s")
+    logger.info(f"Compile decode: {compile_decode_time:.2f}s")
+    logger.info(f"TTFT: {avg_ttft_ms:.2f}ms")
+    logger.info(f"Decode speed: {tok_s_u:.2f} tok/s/u")
+    if top1_acc is not None:
+        logger.info(f"Top-1 accuracy: {top1_acc:.2f}%")
+    if top5_acc is not None:
+        logger.info(f"Top-5 accuracy: {top5_acc:.2f}%")
+
+    # Build measured metrics dict
+    measured = {
+        "ttft_ms": avg_ttft_ms,
+        "tok_s_u": tok_s_u,
+    }
+    if top1_acc is not None:
+        measured["top1"] = top1_acc
+    if top5_acc is not None:
+        measured["top5"] = top5_acc
+
+    # --- Validate Against Baseline ---
+    # Performance metrics (tok_s_u, ttft_ms): MUST come from simple_text_demo.py baseline.
+    # Accuracy metrics (top1, top5): centralized targets source for token-accuracy tests.
+    # Baselines are collected at module load time, before any test opens mesh_device.
+    baseline_metrics = _get_cached_baseline(device_name, batch_size, opt_mode)
+
+    if measure_accuracy:
+        # Token-accuracy tests: use centralized targets for top1/top5 only
+        baseline_source = "models/model_targets.yaml"
+        perf_md_metrics = EXPECTED_METRICS.get(opt_mode, {}).get(device_name, {})
+        expected_for_validation = {k: v for k, v in perf_md_metrics.items() if k in ("top1", "top5")}
+        logger.info("Token-accuracy test runs without trace - validating top1/top5 from model_targets.yaml only")
+    else:
+        # Performance tests: MUST use baseline from simple_text_demo.py
+        assert baseline_metrics is not None, (
+            f"Baseline collection from simple_text_demo.py failed for {device_name}/{opt_mode}/batch-{batch_size}. "
+            "Performance metrics (tok_s_u, ttft_ms) require a valid baseline measurement."
+        )
+        assert "tok_s_u" in baseline_metrics, f"Baseline missing 'tok_s_u': {baseline_metrics}"
+        assert "ttft_ms" in baseline_metrics, f"Baseline missing 'ttft_ms': {baseline_metrics}"
+        baseline_source = "simple_text_demo.py (TTTv1)"
+        expected_for_validation = baseline_metrics
+
+    if expected_for_validation:
+        logger.info(f"Comparing against baseline from: {baseline_source}")
+        if baseline_metrics:
+            logger.info(f"  Baseline tok_s_u: {baseline_metrics.get('tok_s_u', 'N/A'):.2f}")
+            logger.info(f"  Baseline ttft_ms: {baseline_metrics.get('ttft_ms', 'N/A'):.2f}")
+            logger.info(f"  MLP1D tok_s_u:    {tok_s_u:.2f}")
+            logger.info(f"  MLP1D ttft_ms:    {avg_ttft_ms:.2f}")
+
+            # Calculate relative performance
+            if baseline_metrics.get("tok_s_u"):
+                rel_perf = (tok_s_u / baseline_metrics["tok_s_u"]) * 100
+                logger.info(f"  Relative performance: {rel_perf:.1f}% of baseline")
+
+        passed, message = validate_metrics(measured, expected_for_validation)
+        if passed:
+            logger.info(f"✓ All metrics within {PERF_TOLERANCE * 100:.0f}% of {baseline_source} targets")
+        else:
+            logger.warning(f"✗ Metrics outside tolerance: {message}")
+            # Don't fail the test on performance regression - just warn
+            # This allows the demo to run even if performance varies
+
+    logger.info(f"=== MLP1D Demo Complete ===")

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,487 @@
+# 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
+
+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
+        if torch.any(golden.bool()) != torch.any(calculated.bool()):
+            return 0.0
+
+        # 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 case
+        if golden.numel() == 1:
+            return float(torch.equal(golden, calculated))
+
+        # If both tensors are constant
+        if torch.max(golden) == torch.min(golden) and torch.max(calculated) == torch.min(calculated):
+            return torch.isclose(torch.max(golden), torch.max(calculated)).item()
+
+        # 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(),
+        )
+        # Remove correlation coefficient with self (typically always 1.0)
+        mask = np.ones(cal_pcc.shape, dtype=bool)
+        np.fill_diagonal(mask, 0)
+        cal_pcc = np.min(cal_pcc[mask])
+
+        if isinstance(cal_pcc, np.ma.core.MaskedConstant):
+            return 1.0
+
+        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/readiness_check/check_degenerate_output.py

@@ -0,0 +1,495 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Machine check for degenerate generated text in readiness artifacts.
+
+This is a runner-side verification, not a quality judgement. It flags
+mechanically broken generation — the kind produced by decode-loop bugs
+(stale token/position feedback, stale trace inputs) — while staying
+agnostic to model style, checkpoint type (base vs instruct), and content.
+
+The key signal is adjacent-token duplication: a healthy model, base or
+instruct, does not emit nearly every word twice while the text continues
+to advance ("the the difference difference between between ..."). Phrase
+-level looping, by contrast, is common in base checkpoints under greedy
+decoding and is reported only as advisory.
+
+Checked artifacts (discovered under one or more roots):
+
+  - ``readiness_vllm/vllm_qualitative_outputs.json`` written by
+    ``run_vllm_server`` (list of {prompt, greedy_completion,
+    sampled_completion}).
+  - ``autoregressive_meta.json`` written by ``run_autoregressive``
+    ({hf: {token_ids}, tt: {token_ids}, ...}) plus the sibling
+    ``tt_completion.txt`` when present.
+
+Exit codes: 0 = clean, 1 = advisory findings only, 2 = critical findings,
+3 = checker-internal error (never a model verdict).
+
+Invoke by file path rather than ``python -m`` so the package ``__init__``
+(which imports torch) is not pulled in:
+
+    python models/common/readiness_check/check_degenerate_output.py \\
+        [--model-dir models/autoports/<model>] [--hf-model <hf-model-id>] \\
+        [--root models/autoports] [--scope all|vllm|autoregressive] \\
+        [--missing-artifacts advisory|critical] [--json report.json]
+
+Scoping precedence: explicit paths, then ``--model-dir``, then ``--hf-model``
+fuzzy resolution against autoport directory names, then the unscoped
+``--root``. Scoped runs ensure stale artifacts from another model on the
+same machine can neither pass nor fail this check. Unreadable artifacts are
+critical findings (the stage owns its evidence); exit 3 is reserved for
+checker/environment faults.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Iterable, Sequence
+
+# A completion shorter than this many words is too small for a stable
+# duplication rate; it is skipped rather than judged.
+MIN_WORDS_FOR_DUPLICATION = 20
+# Critical: fraction of consecutive word pairs that are identical.
+# Calibration on experiment archives: healthy serving outputs measure
+# <= 0.8% (max over completions); the known token-feedback bug measures
+# 54% mean / 94% max. Anything above 10% is mechanical, not stylistic.
+ADJACENT_DUPLICATION_CRITICAL = 0.10
+# Advisory: fraction of the completion covered by repeats of its most
+# common trigram. Base checkpoints legitimately loop phrases under greedy
+# decoding, so this never fails the check on its own.
+TRIGRAM_LOOP_ADVISORY = 0.50
+MIN_WORDS_FOR_LOOP = 50
+# Advisory: a near-empty completion when many tokens were requested.
+NEAR_EMPTY_CHARS = 5
+
+_WORD_RE = re.compile(r"\w+", re.UNICODE)
+
+
+@dataclass
+class Finding:
+    severity: str  # "critical" | "advisory"
+    artifact: str
+    label: str
+    metric: str
+    value: float
+    threshold: float
+    detail: str
+
+
+@dataclass
+class Report:
+    findings: list[Finding] = field(default_factory=list)
+    measured: list[dict[str, Any]] = field(default_factory=list)
+
+    @property
+    def exit_code(self) -> int:
+        if any(f.severity == "critical" for f in self.findings):
+            return 2
+        if self.findings:
+            return 1
+        return 0
+
+
+def words_of(text: str) -> list[str]:
+    return [w.lower() for w in _WORD_RE.findall(text)]
+
+
+def adjacent_duplication(tokens: Sequence[Any]) -> float:
+    """Fraction of consecutive token pairs that are identical."""
+    if len(tokens) < 2:
+        return 0.0
+    dup = sum(1 for a, b in zip(tokens, tokens[1:]) if a == b)
+    return dup / (len(tokens) - 1)
+
+
+def trigram_loop_fraction(tokens: Sequence[Any]) -> float:
+    """Fraction of the sequence covered by non-overlapping repeats of its most common trigram."""
+    if len(tokens) < 3:
+        return 0.0
+    counts: dict[tuple[Any, ...], int] = {}
+    for i in range(len(tokens) - 2):
+        gram = tuple(tokens[i : i + 3])
+        counts[gram] = counts.get(gram, 0) + 1
+    top_gram = max(counts, key=counts.get)  # type: ignore[arg-type]
+    covered = 0
+    i = 0
+    while i <= len(tokens) - 3:
+        if tuple(tokens[i : i + 3]) == top_gram:
+            covered += 3
+            i += 3
+        else:
+            i += 1
+    return covered / len(tokens)
+
+
+def check_completion(
+    report: Report,
+    *,
+    artifact: Path,
+    label: str,
+    text: str | None,
+    token_ids: Sequence[int] | None = None,
+) -> None:
+    """Apply degeneracy metrics to one completion (text and/or token ids)."""
+    tokens: Sequence[Any] | None = None
+    source = None
+    if text is not None and text.strip():
+        tokens = words_of(text)
+        source = "words"
+    elif token_ids:
+        # No text available: token-id duplication is a weaker but still
+        # meaningful signal (repeated layout tokens make it noisier).
+        tokens = list(token_ids)
+        source = "token_ids"
+
+    measured: dict[str, Any] = {"artifact": str(artifact), "label": label, "source": source}
+
+    if text is not None and len(text.strip()) < NEAR_EMPTY_CHARS:
+        report.findings.append(
+            Finding(
+                severity="advisory",
+                artifact=str(artifact),
+                label=label,
+                metric="near_empty_completion",
+                value=float(len(text.strip())),
+                threshold=float(NEAR_EMPTY_CHARS),
+                detail="Completion is empty or whitespace; verify EOS handling and sampler output.",
+            )
+        )
+        measured["near_empty"] = True
+        report.measured.append(measured)
+        return
+
+    if tokens is None:
+        report.measured.append(measured)
+        return
+
+    dup = adjacent_duplication(tokens)
+    loop = trigram_loop_fraction(tokens)
+    measured.update(
+        {"num_tokens": len(tokens), "adjacent_duplication": round(dup, 4), "trigram_loop_fraction": round(loop, 4)}
+    )
+    report.measured.append(measured)
+
+    if len(tokens) >= MIN_WORDS_FOR_DUPLICATION and dup > ADJACENT_DUPLICATION_CRITICAL:
+        report.findings.append(
+            Finding(
+                severity="critical",
+                artifact=str(artifact),
+                label=label,
+                metric="adjacent_duplication",
+                value=round(dup, 4),
+                threshold=ADJACENT_DUPLICATION_CRITICAL,
+                detail=(
+                    "Adjacent-token duplication this high while the text advances is a "
+                    "decode-loop input bug signature (stale token/position/trace feedback), "
+                    "not a model-quality property. Compare the same prompt against the HF "
+                    "reference before classifying this as a model limitation."
+                ),
+            )
+        )
+    elif len(tokens) >= MIN_WORDS_FOR_LOOP and loop > TRIGRAM_LOOP_ADVISORY:
+        report.findings.append(
+            Finding(
+                severity="advisory",
+                artifact=str(artifact),
+                label=label,
+                metric="trigram_loop_fraction",
+                value=round(loop, 4),
+                threshold=TRIGRAM_LOOP_ADVISORY,
+                detail=(
+                    "Completion is dominated by one repeating phrase. This can be normal for "
+                    "base checkpoints under greedy decoding; verify against the HF reference "
+                    "on the same prompt."
+                ),
+            )
+        )
+
+
+def _load_artifact(report: Report, path: Path) -> Any | None:
+    """Parse an artifact, recording unreadable files as critical findings.
+
+    A required artifact that cannot be read is failed stage evidence, not a
+    checker fault, so it goes through the normal finding path rather than
+    exiting 3.
+    """
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except Exception as exc:  # noqa: BLE001
+        report.findings.append(
+            Finding(
+                severity="critical",
+                artifact=str(path),
+                label="artifact parse",
+                metric="unreadable_artifact",
+                value=0.0,
+                threshold=1.0,
+                detail=f"Artifact could not be read or parsed ({exc}). Regenerate it from the documented runner.",
+            )
+        )
+        return None
+
+
+def check_vllm_qualitative(report: Report, path: Path) -> None:
+    items = _load_artifact(report, path)
+    if items is None:
+        return
+    for i, item in enumerate(items):
+        prompt = str(item.get("prompt", ""))[:60]
+        for key in ("greedy_completion", "sampled_completion"):
+            if key in item:
+                check_completion(
+                    report,
+                    artifact=path,
+                    label=f"prompt[{i}] {key} ({prompt!r})",
+                    text=item.get(key) or "",
+                )
+
+
+def check_autoregressive_meta(report: Report, path: Path) -> None:
+    meta = _load_artifact(report, path)
+    if meta is None:
+        return
+    tt = meta.get("tt", {})
+    text_path = path.parent / "tt_completion.txt"
+    text = text_path.read_text(encoding="utf-8") if text_path.exists() else None
+    check_completion(
+        report,
+        artifact=path,
+        label="tt free-running completion",
+        text=text,
+        token_ids=tt.get("token_ids"),
+    )
+    hf_ids = meta.get("hf", {}).get("token_ids")
+    tt_ids = tt.get("token_ids")
+    if hf_ids and tt_ids:
+        match = sum(1 for a, b in zip(hf_ids, tt_ids) if a == b)
+        report.measured.append(
+            {
+                "artifact": str(path),
+                "label": "hf/tt token agreement (informational)",
+                "matching_tokens": match,
+                "compared_tokens": min(len(hf_ids), len(tt_ids)),
+            }
+        )
+
+
+def discover(roots: Iterable[Path], scope: str) -> tuple[list[Path], list[Path]]:
+    vllm_files: list[Path] = []
+    meta_files: list[Path] = []
+    for root in roots:
+        if root.is_file():
+            if root.name == "autoregressive_meta.json":
+                meta_files.append(root)
+            else:
+                vllm_files.append(root)
+            continue
+        if scope in ("all", "vllm"):
+            vllm_files.extend(sorted(root.rglob("vllm_qualitative_outputs.json")))
+        if scope in ("all", "autoregressive"):
+            meta_files.extend(sorted(root.rglob("autoregressive_meta.json")))
+    return vllm_files, meta_files
+
+
+def _squash(text: str) -> str:
+    return re.sub(r"[^a-z0-9]+", "", text.lower())
+
+
+def resolve_model_dirs(root: Path, hf_model: str) -> tuple[list[Path], str]:
+    """Find the autoport directory (or directories) belonging to hf_model.
+
+    Agents choose their own autoport directory names (observed variants for
+    one model include `llama31_8b_instruct` and `meta_llama_Llama_3_1_8B_Instruct`),
+    so match on alphanumeric-squashed containment between the directory path
+    relative to root and the HF model id, in either direction.
+    """
+    if not root.is_dir():
+        return [], f"{root} does not exist"
+    target = _squash(hf_model)
+    markers = ("tt", "doc", "readiness_vllm")
+    candidates = [
+        path
+        for path in sorted(root.rglob("*"))
+        if path.is_dir()
+        and len(path.relative_to(root).parts) <= 3
+        and any((path / marker).is_dir() for marker in markers)
+    ]
+    matches = []
+    for path in candidates:
+        squashed = _squash(str(path.relative_to(root)))
+        if squashed and (squashed in target or target in squashed):
+            matches.append(path)
+    # Prefer the deepest matching directories; a parent that only matched
+    # because its child did adds nothing but duplicate scanning.
+    matches = [m for m in matches if not any(other != m and other.is_relative_to(m) for other in matches)]
+    if matches:
+        return matches, ""
+    if candidates:
+        listing = ", ".join(str(c.relative_to(root)) for c in candidates)
+        return [], f"no autoport directory under {root} matches {hf_model!r} (found: {listing})"
+    return [], f"no autoport directories found under {root}"
+
+
+class _Parser(argparse.ArgumentParser):
+    def error(self, message: str) -> None:  # noqa: D102 - argparse override
+        # argparse exits 2 on misuse, which stage gates would misread as a
+        # critical model verdict. CLI misuse is a checker/environment fault.
+        self.print_usage(sys.stderr)
+        print(f"{self.prog}: error: {message}", file=sys.stderr)
+        raise SystemExit(3)
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = _Parser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "paths",
+        nargs="*",
+        type=Path,
+        help="Artifact files or directories to scan. Default: --model-dir, --hf-model resolution, or --root.",
+    )
+    parser.add_argument(
+        "--root",
+        type=Path,
+        default=Path("models/autoports"),
+        help="Directory scanned when no paths are given (default: models/autoports).",
+    )
+    parser.add_argument(
+        "--model-dir",
+        type=Path,
+        help="Explicit autoport directory to scope the scan to. Takes precedence over --hf-model.",
+    )
+    parser.add_argument(
+        "--hf-model",
+        help="Scope the scan to the autoport directory matching this HF model id (fuzzy fallback when --model-dir is not given).",
+    )
+    parser.add_argument("--scope", choices=("all", "vllm", "autoregressive"), default="all")
+    parser.add_argument(
+        "--missing-artifacts",
+        choices=("advisory", "critical"),
+        default="advisory",
+        help=(
+            "Severity when no generation artifacts (or no matching model "
+            "directory) are found. Stage gates should use 'critical': by "
+            "stages 5+, missing generation evidence is a failed completion "
+            "requirement, not a soft warning."
+        ),
+    )
+    parser.add_argument("--json", type=Path, help="Write the machine-readable report here.")
+    args = parser.parse_args(argv)
+
+    report = Report()
+
+    if args.paths:
+        roots: list[Path] = args.paths
+    elif args.model_dir:
+        if args.model_dir.is_dir():
+            roots = [args.model_dir]
+        else:
+            roots = []
+            report.findings.append(
+                Finding(
+                    severity=args.missing_artifacts,
+                    artifact=str(args.model_dir),
+                    label="model directory resolution",
+                    metric="missing_model_dir",
+                    value=0.0,
+                    threshold=1.0,
+                    detail=(
+                        f"--model-dir {args.model_dir} does not exist. The stage requires "
+                        "generation evidence under the model's autoport directory."
+                    ),
+                )
+            )
+    elif args.hf_model:
+        roots, why_empty = resolve_model_dirs(args.root, args.hf_model)
+        if not roots:
+            report.findings.append(
+                Finding(
+                    severity=args.missing_artifacts,
+                    artifact=str(args.root),
+                    label="model directory resolution",
+                    metric="missing_model_dir",
+                    value=0.0,
+                    threshold=1.0,
+                    detail=(
+                        f"{why_empty}. The stage requires generation evidence under the "
+                        "model's autoport directory; create it at the documented location."
+                    ),
+                )
+            )
+        else:
+            print(f"scoped to: {', '.join(str(r) for r in roots)}")
+    else:
+        roots = [args.root]
+
+    vllm_files, meta_files = discover(roots, args.scope)
+
+    if roots and not vllm_files and not meta_files and not report.findings:
+        report.findings.append(
+            Finding(
+                severity=args.missing_artifacts,
+                artifact=", ".join(str(r) for r in roots),
+                label="artifact discovery",
+                metric="missing_artifacts",
+                value=0.0,
+                threshold=1.0,
+                detail=(
+                    f"No generation artifacts found (scope={args.scope}). If this stage "
+                    "requires generation evidence, produce it at the documented location "
+                    "before reporting complete."
+                ),
+            )
+        )
+    for path in vllm_files:
+        check_vllm_qualitative(report, path)
+    for path in meta_files:
+        check_autoregressive_meta(report, path)
+
+    for m in report.measured:
+        compact = {k: v for k, v in m.items() if k not in ("artifact", "label")}
+        print(f"measured: {m.get('label')} [{m.get('artifact')}] {compact}")
+
+    if report.findings:
+        print(f"\n{len(report.findings)} finding(s):")
+        for f in report.findings:
+            print(f"\n[{f.severity.upper()}] {f.metric}={f.value} (threshold {f.threshold})")
+            print(f"  artifact: {f.artifact}")
+            print(f"  where:    {f.label}")
+            print(f"  {f.detail}")
+    else:
+        print("\nNo degenerate output detected.")
+
+    if args.json:
+        args.json.parent.mkdir(parents=True, exist_ok=True)
+        args.json.write_text(
+            json.dumps(
+                {
+                    "findings": [asdict(f) for f in report.findings],
+                    "measured": report.measured,
+                    "exit_code": report.exit_code,
+                },
+                indent=2,
+            ),
+            encoding="utf-8",
+        )
+
+    return report.exit_code
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except SystemExit:
+        raise
+    except Exception:  # noqa: BLE001 - exit 3 marks checker-internal errors, never a model verdict
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(3)

code/models/common/readiness_check/run_prefill_check.py

@@ -0,0 +1,266 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Run prefill-based readiness check against a reference file.
+
+Validates TT model prefill accuracy by:
+  1. Loading the generator from `<model_dir>/tt/generator.py`
+  2. Running `prefill_forward(return_all_logits=True)` on full sequences
+  3. Comparing predictions against reference top-K at each position
+  4. Reporting top-1, top-5, and top-K accuracy
+
+Use alongside `run_teacher_forcing.py` to validate both prefill and decode paths.
+
+CLI:
+    python -m models.common.readiness_check.run_prefill_check \\
+        --model-dir models/autoports/<model_name> \\
+        --reference references/<model>.refpt \\
+        --mesh-device N150
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import sys
+from pathlib import Path
+from typing import Any, Dict, List
+
+import torch
+
+from models.common.readiness_check.contract import (
+    BUILD_GENERATOR_FUNCTION_NAME,
+    GENERATOR_MODULE_RELPATH,
+    BuildGeneratorFn,
+    Generator,
+)
+from models.common.readiness_check.mesh_device import (
+    add_mesh_device_args,
+    close_readiness_mesh_device,
+    open_readiness_mesh_device,
+)
+from models.common.readiness_check.schema import Reference, ReferenceEntry, load_reference
+
+
+def _import_build_generator(model_dir: Path) -> BuildGeneratorFn:
+    """
+    Load `<model_dir>/tt/generator.py` and return its `build_generator` function.
+    """
+    generator_path = model_dir / GENERATOR_MODULE_RELPATH
+    if not generator_path.exists():
+        raise FileNotFoundError(
+            f"Expected generator at {generator_path}. The readiness check requires "
+            f"<model_dir>/{GENERATOR_MODULE_RELPATH} to exist and expose `{BUILD_GENERATOR_FUNCTION_NAME}`."
+        )
+
+    module_name = f"_readiness_generator_{model_dir.resolve().name}"
+    spec = importlib.util.spec_from_file_location(module_name, generator_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load spec for {generator_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+
+    fn = getattr(module, BUILD_GENERATOR_FUNCTION_NAME, None)
+    if fn is None or not callable(fn):
+        raise AttributeError(
+            f"{generator_path} does not export a callable `{BUILD_GENERATOR_FUNCTION_NAME}`. "
+            f"See models/common/readiness_check/contract.py."
+        )
+    return fn  # type: ignore[return-value]
+
+
+def _run_one_entry_prefill(
+    *,
+    generator: Generator,
+    entry: ReferenceEntry,
+    reference: Reference,
+) -> Dict[str, Any]:
+    """
+    Run batch prefill check for one entry.
+
+    Args:
+        generator: TT generator instance
+        entry: Reference entry with prompt, generated tokens, and top-K predictions
+        reference: Reference object (for metadata)
+
+    Returns:
+        Accuracy dict with top1, top5, top100, matches, and total
+    """
+    prompt_tokens = entry.prompt_tokens[0].tolist()  # [P]
+    gen_tokens = entry.generated_tokens[0].tolist()  # [G]
+    topk_reference = entry.topk_tokens  # [G, K]
+
+    # Concatenate prompt + generated for full sequence
+    full_sequence = prompt_tokens + gen_tokens
+    prompt_len = len(prompt_tokens)
+    gen_len = len(gen_tokens)
+    full_len = prompt_len + gen_len
+
+    # Convert to tensor [1, full_len]
+    tokens_tensor = torch.tensor([full_sequence], dtype=torch.long)
+
+    # Prepare KV cache and page table (implementation-specific)
+    # This is a simplified version - real implementations need proper setup
+    # For now, we'll call prefill_forward with minimal setup and let the
+    # generator handle defaults through **kwargs
+
+    # Call prefill with return_all_logits=True to get logits at all positions
+    # Note: This requires the generator implementation to support return_all_logits
+    try:
+        # Most generators will need proper page_table and kv_cache setup
+        # This is a placeholder - real usage requires model-specific initialization
+        import ttnn
+
+        # Allocate dummy page table and kv cache
+        # Real implementations should use generator's initialization methods
+        batch_size = 1
+        max_blocks = 1024
+        page_table = torch.arange(max_blocks).reshape(batch_size, max_blocks)
+        page_table_tt = ttnn.from_torch(
+            page_table,
+            device=generator.mesh_device if hasattr(generator, "mesh_device") else None,
+            dtype=ttnn.int32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+        )
+
+        # Placeholder kv_cache - real implementation needs proper initialization
+        kv_cache = None  # Generator should handle None gracefully or we need to init properly
+
+        logits = generator.prefill_forward(
+            tokens=tokens_tensor,
+            page_table=page_table_tt,
+            kv_cache=kv_cache,
+            prompt_lens=[full_len],
+            return_all_logits=True,
+        )
+    except TypeError as e:
+        if "return_all_logits" in str(e):
+            raise NotImplementedError(
+                f"Generator {type(generator).__name__} does not support return_all_logits parameter. "
+                "Please update the generator to support the new contract."
+            ) from e
+        raise
+
+    # logits should be [1, full_len, vocab_size]
+    if logits.dim() == 2:
+        # [batch, vocab] - only got last position, not all positions
+        raise RuntimeError(
+            f"Generator returned logits with shape {logits.shape}, but return_all_logits=True "
+            "should return shape [batch, seq_len, vocab]. Generator may not support return_all_logits."
+        )
+
+    # Extract logits at positions that predict gen_tokens
+    # logits[0, i] predicts token at position i+1
+    # So logits[0, prompt_len-1:prompt_len+gen_len-1] predicts gen_tokens
+    prediction_logits = logits[0, prompt_len - 1 : prompt_len + gen_len - 1, :]  # [gen_len, vocab]
+
+    # Get top-1 predictions (argmax)
+    tt_predictions = torch.argmax(prediction_logits, dim=-1).cpu()  # [gen_len]
+
+    # Compare against reference
+    matches_top1 = 0
+    matches_top5 = 0
+    matches_topk = 0
+    k_cols = min(5, topk_reference.shape[1])
+
+    for i in range(gen_len):
+        tt_pred = int(tt_predictions[i].item())
+        ref_topk = topk_reference[i]
+
+        if tt_pred == int(ref_topk[0].item()):
+            matches_top1 += 1
+        if tt_pred in ref_topk[:k_cols].tolist():
+            matches_top5 += 1
+        if tt_pred in ref_topk.tolist():
+            matches_topk += 1
+
+    return {
+        "top1": matches_top1 / gen_len,
+        "top5": matches_top5 / gen_len,
+        "top100": matches_topk / gen_len,
+        "matches_top1": matches_top1,
+        "matches_top5": matches_top5,
+        "matches_top100": matches_topk,
+        "total": gen_len,
+        "k": reference.k,
+    }
+
+
+def _format_row(label: str, stats: Dict[str, Any]) -> str:
+    return (
+        f"{label:<20} "
+        f"top1={stats['top1']:.3f} ({stats['matches_top1']}/{stats['total']})  "
+        f"top5={stats['top5']:.3f} ({stats['matches_top5']}/{stats['total']})  "
+        f"top{stats['k']}={stats['top100']:.3f} ({stats['matches_top100']}/{stats['total']})"
+    )
+
+
+def run_prefill_check(
+    *,
+    model_dir: Path,
+    reference_path: Path,
+    mesh_device,
+    build_kwargs: Dict[str, Any] | None = None,
+) -> List[Dict[str, Any]]:
+    """
+    Programmatic entry point. Builds the generator, runs batch prefill
+    over all reference entries, and returns the per-entry accuracy dicts.
+    """
+    build_kwargs = build_kwargs or {}
+    build_generator = _import_build_generator(model_dir)
+    generator: Generator = build_generator(model_dir=model_dir, mesh_device=mesh_device, **build_kwargs)
+
+    reference = load_reference(reference_path)
+    per_entry: List[Dict[str, Any]] = []
+
+    try:
+        for entry_idx, entry in enumerate(reference.entries):
+            if entry_idx > 0:
+                generator.reset()
+            stats = _run_one_entry_prefill(generator=generator, entry=entry, reference=reference)
+            per_entry.append(stats)
+            print(_format_row(f"entry[{entry_idx}]", stats))
+    finally:
+        teardown = getattr(generator, "teardown", None)
+        if callable(teardown):
+            teardown()
+
+    total = sum(s["total"] for s in per_entry)
+    if total:
+        agg = {
+            "top1": sum(s["matches_top1"] for s in per_entry) / total,
+            "top5": sum(s["matches_top5"] for s in per_entry) / total,
+            "top100": sum(s["matches_top100"] for s in per_entry) / total,
+            "matches_top1": sum(s["matches_top1"] for s in per_entry),
+            "matches_top5": sum(s["matches_top5"] for s in per_entry),
+            "matches_top100": sum(s["matches_top100"] for s in per_entry),
+            "total": total,
+            "k": per_entry[0]["k"],
+        }
+        print(_format_row("AGGREGATE", agg))
+
+    return per_entry
+
+
+def _main() -> None:
+    parser = argparse.ArgumentParser(description="Run the batch prefill readiness check against a reference file.")
+    parser.add_argument("--model-dir", type=Path, required=True, help="Path to the model directory.")
+    parser.add_argument("--reference", type=Path, required=True, help="Path to the .refpt reference file.")
+    add_mesh_device_args(parser)
+    args = parser.parse_args()
+
+    mesh_device = open_readiness_mesh_device(args.mesh_device, args.fabric_config)
+    try:
+        run_prefill_check(
+            model_dir=args.model_dir.resolve(),
+            reference_path=args.reference.resolve(),
+            mesh_device=mesh_device,
+        )
+    finally:
+        close_readiness_mesh_device(mesh_device, args.fabric_config)
+
+
+if __name__ == "__main__":
+    _main()

code/models/common/readiness_check/schema.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Reference file schema for the model-readiness check.
+
+A reference file pins, for one or more prompts, the HuggingFace teacher's
+top-K next-token predictions at every generated position. The TT model
+under test is then run with teacher forcing against this file and scored
+on top-1 / top-5 / top-K hit rate.
+
+On-disk format (single torch.save'd dict):
+
+    {
+        "format_version": "readiness_v1",
+        "k": int,
+        "hf_model_id": str,
+        "token_ids_meta": {"bos_id": int|None, "eos_id": int, "pad_id": int|None},
+        "entries": [
+            {
+                "prompt_text": str,
+                "prompt_tokens":    Tensor [1, P]   int64,
+                "generated_tokens": Tensor [1, G]   int64,
+                "topk_tokens":      Tensor [G, K]   int32,
+                "tf_prompt_len":    int,            # == P, stored for clarity
+            },
+            ...
+        ],
+    }
+
+`topk_tokens[i]` is the teacher's top-K prediction for the token at
+position `tf_prompt_len + i`, conditioned on positions `[0 .. tf_prompt_len + i - 1]`.
+Top-1 is column 0; top-5 is columns 0..4; top-K is the full row.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import torch
+
+FORMAT_VERSION = "readiness_v1"
+
+
+@dataclass
+class ReferenceEntry:
+    prompt_text: str
+    prompt_tokens: torch.Tensor  # [1, P] int64
+    generated_tokens: torch.Tensor  # [1, G] int64
+    topk_tokens: torch.Tensor  # [G, K] int32
+    tf_prompt_len: int
+
+    def __post_init__(self) -> None:
+        if self.prompt_tokens.dim() != 2 or self.prompt_tokens.shape[0] != 1:
+            raise ValueError(f"prompt_tokens must be [1, P], got {tuple(self.prompt_tokens.shape)}")
+        if self.generated_tokens.dim() != 2 or self.generated_tokens.shape[0] != 1:
+            raise ValueError(f"generated_tokens must be [1, G], got {tuple(self.generated_tokens.shape)}")
+        if self.topk_tokens.dim() != 2:
+            raise ValueError(f"topk_tokens must be [G, K], got {tuple(self.topk_tokens.shape)}")
+        if self.topk_tokens.shape[0] != self.generated_tokens.shape[1]:
+            raise ValueError(
+                f"topk_tokens rows ({self.topk_tokens.shape[0]}) must equal "
+                f"generated_tokens length ({self.generated_tokens.shape[1]})"
+            )
+        if self.tf_prompt_len != self.prompt_tokens.shape[1]:
+            raise ValueError(
+                f"tf_prompt_len ({self.tf_prompt_len}) must equal prompt length " f"({self.prompt_tokens.shape[1]})"
+            )
+
+    @property
+    def num_generated(self) -> int:
+        return int(self.generated_tokens.shape[1])
+
+    @property
+    def k(self) -> int:
+        return int(self.topk_tokens.shape[1])
+
+
+@dataclass
+class Reference:
+    k: int
+    hf_model_id: str
+    entries: List[ReferenceEntry]
+    token_ids_meta: Dict[str, Optional[int]] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.entries:
+            raise ValueError("Reference must contain at least one entry")
+        for idx, entry in enumerate(self.entries):
+            if entry.k != self.k:
+                raise ValueError(f"Entry {idx} has k={entry.k}, expected {self.k} from Reference.k")
+
+
+def save_reference(reference: Reference, path: Path | str) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload: Dict[str, Any] = {
+        "format_version": FORMAT_VERSION,
+        "k": reference.k,
+        "hf_model_id": reference.hf_model_id,
+        "token_ids_meta": reference.token_ids_meta,
+        "entries": [
+            {
+                "prompt_text": e.prompt_text,
+                "prompt_tokens": e.prompt_tokens.to(torch.int64).cpu().contiguous(),
+                "generated_tokens": e.generated_tokens.to(torch.int64).cpu().contiguous(),
+                "topk_tokens": e.topk_tokens.to(torch.int32).cpu().contiguous(),
+                "tf_prompt_len": int(e.tf_prompt_len),
+            }
+            for e in reference.entries
+        ],
+    }
+    torch.save(payload, path)
+    return path
+
+
+def load_reference(path: Path | str) -> Reference:
+    path = Path(path)
+    payload = torch.load(path, weights_only=False)
+    fmt = payload.get("format_version")
+    if fmt != FORMAT_VERSION:
+        raise ValueError(f"Unsupported reference format_version={fmt!r} in {path}; " f"expected {FORMAT_VERSION!r}")
+    raw_entries = payload.get("entries")
+    if not isinstance(raw_entries, list) or not raw_entries:
+        raise ValueError(f"Reference {path} missing non-empty 'entries' list")
+    entries = [
+        ReferenceEntry(
+            prompt_text=str(e["prompt_text"]),
+            prompt_tokens=e["prompt_tokens"],
+            generated_tokens=e["generated_tokens"],
+            topk_tokens=e["topk_tokens"],
+            tf_prompt_len=int(e["tf_prompt_len"]),
+        )
+        for e in raw_entries
+    ]
+    return Reference(
+        k=int(payload["k"]),
+        hf_model_id=str(payload["hf_model_id"]),
+        entries=entries,
+        token_ids_meta=dict(payload.get("token_ids_meta") or {}),
+    )

code/models/common/readiness_check/teacher_forcing.py

@@ -0,0 +1,182 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Teacher-forcing accuracy harness driven by a `readiness_v1` reference file.
+
+Usage:
+
+    acc = TokenAccuracy("llama31_8b.refpt")
+    prompt_ids = acc.get_prompt_token_ids(user_idx=0)
+
+    # Prefill the TT model with prompt_ids, then decode one token at a time:
+    for _ in range(acc.num_gt_tokens(user_idx=0)):
+        tt_pred = run_tt_decode(...)            # int token id
+        forced_next = acc.collect_predicted_tokens(tt_pred, user_idx=0)
+        feed_token_to_tt(forced_next)            # teacher forcing
+
+    print(acc.compute_accuracy(user_idx=0))
+    # {'top1': ..., 'top5': ..., 'top100': ..., 'matches_top1': ..., ...}
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+
+from models.common.readiness_check.schema import Reference, load_reference
+
+
+class TokenAccuracy:
+    """
+    Per-user teacher-forcing helper backed by a `readiness_v1` reference file.
+
+    Each `entry` in the reference corresponds to one prompt and is exposed as
+    a separate "user." For multi-user vLLM-style decode, callers pass
+    `user_idx` to route per-user state.
+    """
+
+    def __init__(self, reference: str | Path | Reference) -> None:
+        if isinstance(reference, Reference):
+            self.reference = reference
+        else:
+            self.reference = load_reference(reference)
+
+        self.entry_count = len(self.reference.entries)
+        self.k = self.reference.k
+
+        # Flattened per-user views.
+        self.prompt_1d: List[torch.Tensor] = [
+            e.prompt_tokens[0].to(torch.long).contiguous() for e in self.reference.entries
+        ]
+        self.gt_gen_1d: List[torch.Tensor] = [
+            e.generated_tokens[0].to(torch.long).contiguous() for e in self.reference.entries
+        ]
+        self.topk: List[torch.Tensor] = [e.topk_tokens.to(torch.long).contiguous() for e in self.reference.entries]
+        self.tf_prompt_len: List[int] = [int(e.tf_prompt_len) for e in self.reference.entries]
+
+        meta = self.reference.token_ids_meta or {}
+        self.eos_id: Optional[int] = int(meta["eos_id"]) if meta.get("eos_id") is not None else None
+
+        # Per-user runtime state.
+        self.pred_tokens: List[List[int]] = [[] for _ in range(self.entry_count)]
+        self.cursor: List[int] = [0] * self.entry_count
+
+    # Bounds + reset.
+
+    @property
+    def num_entries(self) -> int:
+        return self.entry_count
+
+    def _check_idx(self, user_idx: int) -> int:
+        idx = int(user_idx)
+        if idx < 0 or idx >= self.entry_count:
+            raise IndexError(f"user_idx={idx} out of range for {self.entry_count} entry(ies)")
+        return idx
+
+    def reset(self) -> None:
+        for buf in self.pred_tokens:
+            buf.clear()
+        for i in range(self.entry_count):
+            self.cursor[i] = 0
+
+    # Accessors.
+
+    def get_prompt_token_ids(self, user_idx: int = 0) -> List[int]:
+        return self.prompt_1d[self._check_idx(user_idx)].tolist()
+
+    def num_gt_tokens(self, user_idx: int = 0) -> int:
+        return int(self.gt_gen_1d[self._check_idx(user_idx)].numel())
+
+    def num_pred_tokens(self, user_idx: int = 0) -> int:
+        return len(self.pred_tokens[self._check_idx(user_idx)])
+
+    def get_predicted_tokens(self, user_idx: int = 0) -> List[int]:
+        return list(self.pred_tokens[self._check_idx(user_idx)])
+
+    # Core teacher-forcing interface.
+
+    def collect_predicted_tokens(self, tt_pred_token: int, *, user_idx: int = 0) -> int:
+        """
+        Record TT's predicted token for the next generated position of
+        `user_idx` and return the ground-truth token to force into the
+        next TT decode step.
+
+        Once the ground-truth sequence is exhausted, returns `eos_id` if
+        known, else the last ground-truth token.
+        """
+        idx = self._check_idx(user_idx)
+        gt_gen = self.gt_gen_1d[idx]
+        cursor = self.cursor[idx]
+
+        self.pred_tokens[idx].append(int(tt_pred_token))
+
+        if cursor >= int(gt_gen.numel()):
+            if self.eos_id is not None:
+                return int(self.eos_id)
+            return int(gt_gen[-1].item())
+
+        forced = int(gt_gen[cursor].item())
+        self.cursor[idx] = cursor + 1
+        return forced
+
+    def compute_accuracy(self, user_idx: int = 0) -> Dict[str, float]:
+        """
+        Top-1 / top-5 / top-K hit rate for `user_idx`:
+
+          For generated step i:
+            - top-1 = topk_tokens[i, 0]
+            - top-5 = topk_tokens[i, 0:5]
+            - top-K = topk_tokens[i, :]
+
+        Returns a dict with float ratios and integer counts:
+            {top1, top5, top100, matches_top1, matches_top5, matches_top100, total}
+        (`top100` is named for the K=100 default; if the reference has a
+        different K, it is still keyed `top100` for consistency. `k` is
+        also returned so callers can verify.)
+        """
+        idx = self._check_idx(user_idx)
+        preds = self.pred_tokens[idx]
+        gt_gen = self.gt_gen_1d[idx]
+        topk = self.topk[idx]
+
+        total = min(len(preds), int(gt_gen.numel()))
+        if total == 0:
+            return {
+                "top1": 0.0,
+                "top5": 0.0,
+                "top100": 0.0,
+                "matches_top1": 0,
+                "matches_top5": 0,
+                "matches_top100": 0,
+                "total": 0,
+                "k": self.k,
+            }
+
+        matches_top1 = 0
+        matches_top5 = 0
+        matches_topk = 0
+        k_cols = min(5, topk.shape[1])
+
+        for i in range(total):
+            row = topk[i]
+            tt_pred = int(preds[i])
+            if tt_pred == int(row[0].item()):
+                matches_top1 += 1
+            if tt_pred in row[:k_cols].tolist():
+                matches_top5 += 1
+            if tt_pred in row.tolist():
+                matches_topk += 1
+
+        return {
+            "top1": matches_top1 / total,
+            "top5": matches_top5 / total,
+            "top100": matches_topk / total,
+            "matches_top1": matches_top1,
+            "matches_top5": matches_top5,
+            "matches_top100": matches_topk,
+            "total": total,
+            "k": self.k,
+        }

code/models/common/readiness_check/vllm_prompts.txt

@@ -0,0 +1,11 @@
+Write a haiku about machine learning.
+
+Explain the difference between supervised and unsupervised learning in simple terms.
+
+Complete this story: Once upon a time, in a faraway kingdom, there lived a curious young inventor who discovered
+
+What are the three laws of thermodynamics?
+
+Translate the following to French: "Hello, how are you today?"
+
+Write a Python function to calculate the Fibonacci sequence.

code/models/common/rmsnorm.py

@@ -0,0 +1,203 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+import ttnn
+from models.common.lightweightmodule import LightweightModule
+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
+
+        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 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
+
+        # 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)
+        norm = self._distributed_rmsnorm if distributed else ttnn.rms_norm
+        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"
+
+        x = 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
+    ):
+        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"
+
+        # Run distributed rmsnorm part 1
+        tt_stats = ttnn.rms_norm_pre_all_gather(inp, compute_kernel_config=compute_kernel_config, dtype=ttnn.bfloat16)
+        # 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=ttnn.DRAM_MEMORY_CONFIG,
+            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,
+        )
+        tt_stats.deallocate(True)
+
+        return tt_out

code/models/common/sampling/README.md

@@ -0,0 +1,113 @@
+# Sampling Module Overview
+
+The `models.common.sampling` package bundles everything needed to run on-device
+sampling (top-k / top-p / temperature/ seed) plus presence/frequency/repetition
+penalties with optional trace capture.
+
+## Key Components
+- `SamplingGenerator`: high-level class that owns both `TTSampling` and
+  `TTPenalties`, exposes helper methods to reset sampling parameters, penalties,
+  prompt/output state, and to run sampling with or without trace capture.
+- `format_sampling_params`: utility that pads/clamps sampling parameters to the
+  hardware-friendly layout expected by `TTSampling`.
+- `LogProbsCalculator`: computes per-token log-probabilities across a sharded
+  vocabulary using numerically stable log-softmax (global max / sum-exp
+  reduction across devices).
+
+## Quick Start
+```python
+from models.common.sampling import SamplingGenerator, format_sampling_params
+
+sampling = SamplingGenerator(args=args, mesh_device=mesh_device, tt_ccl=tt_ccl)
+
+params = format_sampling_params(user_params, max_batch_size=32)
+sampling.reset_sampling_params(params)
+
+sampling.reset_seed(seed)
+
+sampling.reset_prompt_tokens(prompt_tokens)   # torch tensor shaped [B, S]
+sampling.reset_output_state(output_tokens)
+
+tt_tokens = sampling.sample(
+    tt_logits,
+    tt_out_tok=tt_out_buffer,
+)
+```
+
+`SamplingGenerator.sample()` accepts `enable_trace=True` to record/replay
+sampling traces.
+
+## File Map
+
+| File | Purpose |
+|---|---|
+| `generator.py` | `SamplingGenerator` orchestrator; `SamplingParams`; `format_sampling_params`; `broadcast_sampling_params`; `chunk_sampling_params`; `SeedManager` |
+| `tt_sampling.py` | `TTSampling` — on-device top-k/top-p/temp with multi-device all-gather |
+| `tt_penalties.py` | `TTPenalties` — presence / frequency / repetition penalties |
+| `tt_log_probs.py` | `LogProbsCalculator` — log-softmax across sharded vocabulary |
+| `_utils.py` | Shared helpers: `clamp`, `is_default_value`, `filter_none`, `split_list` |
+
+## Required `args` Attributes
+
+```python
+vocab_size: int           # actual vocabulary size (unpadded)
+cluster_shape: tuple      # (rows, cols) of the device mesh, e.g. (4, 8)
+```
+
+Optional (with defaults):
+
+```python
+padded_vocab_size: int    # tile-aligned total vocab; defaults to vocab_size
+max_batch_size: int       # per sampling row; default 32
+max_top_k: int            # default 32
+sampling_dp: int          # >1 for multi-row DP; default 1
+sub_core_grids            # CoreRangeSet or None
+model_config: dict        # keys: GALAXY_NUM_LINKS, DECODE_SAMPLING_INPUT_MEMCFG, SAMPLING_AG_CONFIG
+```
+
+## `data_parallel` vs `sampling_dp`
+
+These are different concepts and should not be mixed:
+
+- **`data_parallel`** lives above this package. It means multiple TT model
+  instances / submeshes process different requests in parallel.
+- **`sampling_dp`** lives inside this package. It means one TT model instance
+  has multiple independent sampling groups, usually one per mesh row.
+
+For `sampling_dp > 1`:
+- logits are still computed per sampling group
+- but sampling params, seeds, and penalty state are flattened to
+  `max_batch_size * sampling_dp`
+- those flattened host tensors are then row-sharded onto the device
+
+Decode already follows this contract by using `chunk_sampling_params(...)`
+plus `apply_decode_state(...)`.
+
+## Param Distribution API
+
+**`SamplingParams`**: Canonical dataclass for sampling parameters (temp, top_k, top_p, penalties, seed, log_probs). Import from `models.common.sampling`. vLLM has its own duck-type-compatible `TTSamplingParams`.
+
+**`broadcast_sampling_params(params, idx, slot_len=32)`**: Expand a single user's params to fill `slot_len` slots. Used during prefill.
+
+**`chunk_sampling_params(params, sampling_dp)`**: Split a SamplingParams into `sampling_dp` pieces. List fields split evenly; scalars replicated. Works with duck-typed objects (vLLM).
+
+**`SamplingGenerator.apply_prefill_state(...)`**: Reset params, seeds, prompt tokens, and output state for a prefill request.
+
+**`SamplingGenerator.apply_decode_state(chunks, ...)`**: Format/merge params and apply for one model instance. Handles both simple (1 chunk) and row-sharded (multiple chunks) cases. Does NOT advance seeds — callers manage `seed_manager.get_new_values()` separately.
+
+## Pitfalls
+
+**`padded_vocab_size` vs `vocab_size`**: TTSampling device offsets for global token IDs must use the padded vocab size to match how the LM head shards logits across devices. Using unpadded `vocab_size` for offsets shifts token IDs from devices 1+ and produces garbled output.
+
+**`sampling_dp`**: When >1, k/p/temp tensors must have length `max_batch_size * sampling_dp` and are row-sharded via `ShardTensor2dMesh(dims=(0, None))`. Use `chunk_sampling_params` + `apply_decode_state` to distribute params across mesh rows.
+
+**Batched prefill + on-device sampling**: This path is only valid when the
+runtime prefill compute layout matches the sampling-group layout. If a model
+uses `sampling_dp > 1` but does not expose a row-sharded batched-prefill input
+contract, batched prefill must fall back to sequential prefill for correctness.
+
+**Trace invalidation**: Changing `force_argmax_sampling` state invalidates captured traces. Force-argmax is triggered when callers pass k=1, p=1.0, temp=1.0 (note: p=1.0 means "no top-p filtering", distinct from the internal initialization default of p=0). `SamplingGenerator.reset_sampling_params` handles this.
+
+## Future Work
+
+- Consolidate DeepSeek's minimal `SamplingParams` (in `models/demos/deepseek_v3/tt/generator.py`) to use the common one

code/models/common/sampling/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from .tt_sampling import TTSampling
+from .tt_penalties import TTPenalties, apply_penalties
+from .tt_log_probs import LogProbsCalculator, LogProbsResult
+from .generator import (
+    SamplingGenerator,
+    SamplingParams,
+    SAMPLING_PARAM_FIELDS,
+    format_sampling_params,
+    broadcast_sampling_params,
+    chunk_sampling_params,
+    SeedManager,
+)
+from ._utils import split_list
+
+__all__ = [
+    "TTSampling",
+    "TTPenalties",
+    "apply_penalties",
+    "LogProbsCalculator",
+    "LogProbsResult",
+    "SamplingGenerator",
+    "SamplingParams",
+    "SAMPLING_PARAM_FIELDS",
+    "format_sampling_params",
+    "broadcast_sampling_params",
+    "chunk_sampling_params",
+    "SeedManager",
+    "split_list",
+]

code/models/common/sampling/_utils.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from loguru import logger
+
+
+def clamp(value, min_value, max_value):
+    if value < min_value:
+        return min_value
+    elif value > max_value:
+        return max_value
+    return value
+
+
+def is_default_value(values, default):
+    """Check if values match a default, handling None, scalar, and iterable inputs."""
+    if values is None:
+        return True
+    if isinstance(values, (int, float)):
+        return values == default
+    return all(value == default for value in values)
+
+
+def filter_none(kwargs: dict) -> dict:
+    return {k: v for k, v in kwargs.items() if v is not None}
+
+
+def split_list(lst, n):
+    """Split list into n equal parts."""
+    chunk_size = len(lst) // n
+    return [list(lst[i * chunk_size : (i + 1) * chunk_size]) for i in range(n)]
+
+
+def compact_debug_list(values, max_items=12):
+    if values is None:
+        return None
+    if hasattr(values, "reshape") and hasattr(values, "tolist"):
+        values = values.reshape(-1).tolist()
+    elif isinstance(values, tuple):
+        values = list(values)
+    elif not isinstance(values, list):
+        values = list(values) if isinstance(values, range) else [values]
+    if len(values) <= max_items:
+        return values
+    half = max(1, max_items // 2)
+    return {"len": len(values), "head": values[:half], "tail": values[-half:]}
+
+
+def is_llama33_70b_model(args) -> bool:
+    if isinstance(args, list):
+        args = args[0] if args else None
+    if args is None:
+        return False
+
+    fields_to_check = (
+        "model_name",
+        "base_model_name",
+        "model_base_path",
+        "model_cache_path",
+        "tokenizer_path",
+        "CKPT_DIR",
+        "LLAMA_DIR",
+        "hf_model",
+        "HF_MODEL",
+    )
+    for field in fields_to_check:
+        value = getattr(args, field, None)
+        if value is None:
+            continue
+        normalized = str(value).lower().replace("_", "-")
+        if "llama" in normalized and "3.3-70b" in normalized:
+            return True
+    return False
+
+
+def log_sampling_debug(enabled, message, **kwargs):
+    if not enabled:
+        return
+    compact = {key: value for key, value in kwargs.items() if value is not None}
+    logger.info(f"SamplingDBG {message}: {compact}")
+
+
+def is_power_of_2(n):
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def upper_power_of_2(n: int) -> int:
+    if n <= 1:
+        return 1
+    return 1 << (n - 1).bit_length()

code/models/common/sampling/generator.py

@@ -0,0 +1,932 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import random
+import secrets
+from dataclasses import dataclass, fields, replace
+from typing import List, Optional
+
+import torch
+from loguru import logger
+
+import ttnn
+
+from ._utils import clamp
+from ._utils import compact_debug_list as _compact_debug_list
+from ._utils import is_default_value, is_llama33_70b_model
+from ._utils import log_sampling_debug as _log_sampling_debug
+from ._utils import split_list
+from .tt_penalties import TTPenalties
+from .tt_sampling import TTSampling
+
+MAX_UINT32 = 2**32 - 1
+# MAX_UINT32 is reserved as the device skip sentinel; keep real seeds in a bounded positive range.
+DEVICE_SEED_MAX = 1_000_000
+_UINT64_MASK = (1 << 64) - 1
+
+
+def _hash_request_seed_to_device_seed(seed: int, counter: int) -> int:
+    """Derive a stable per-token device seed from a request seed.
+
+    The device sampling op accepts bounded positive seeds, while vLLM
+    request seeds can be any integer and must be reproducible regardless
+    of batch slot. Hashing (request seed, token counter) gives each token
+    a deterministic but well-mixed device seed without relying on mutable
+    per-slot RNG state. The constants below are the SplitMix64 finalizer.
+    """
+    value = (int(seed) & _UINT64_MASK) ^ ((int(counter) + 0x9E3779B97F4A7C15) & _UINT64_MASK)
+    value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & _UINT64_MASK
+    value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & _UINT64_MASK
+    value = (value ^ (value >> 31)) & _UINT64_MASK
+    return (value % DEVICE_SEED_MAX) + 1
+
+
+@dataclass(frozen=True)
+class SamplingParams:
+    """
+    Sampling parameters for on-device greedy decoding / sampling.
+
+    Used by Generator decode/prefill functions. vLLM has its own duck-type-compatible
+    TTSamplingParams (in vllm/worker/tt_model_runner.py) that works with the same
+    format_sampling_params / chunk_sampling_params functions.
+    """
+
+    temperature: float | list[float]
+    top_k: int | list[int]
+    top_p: float | list[float]
+    presence_penalty: float | list[float] = 0.0
+    frequency_penalty: float | list[float] = 0.0
+    repetition_penalty: float | list[float] = 1.0
+    seed: int | list[int] | None = None
+    enable_log_probs: bool | list[bool] = False
+    num_logprobs: int | list[int] = 0
+
+
+SAMPLING_PARAM_FIELDS = tuple(f.name for f in fields(SamplingParams))
+
+
+@dataclass(frozen=True)
+class _TraceKey:
+    penalties_on: bool
+    log_probs_on: bool
+    force_argmax: bool
+
+
+class SamplingGenerator:
+    """
+    High-level sampling helper that owns both `TTSampling` and `TTPenalties`
+    modules and optionally manages TTNN trace capture/execution for sampling.
+
+    Typical usage:
+        generator = SamplingGenerator(args=args, mesh_device=mesh_device, tt_ccl=tt_ccl)
+        generator.reset_sampling_params(k=..., p=..., temp=...)
+        tokens = generator.sample(logits, enable_trace=True)
+    """
+
+    _DEFAULT_PENALTIES = {
+        "presence": 0.0,
+        "frequency": 0.0,
+        "repetition": 1.0,
+    }
+
+    def __init__(
+        self,
+        *,
+        args,
+        mesh_device,
+        tt_ccl,
+        cq_id: int = 0,
+    ):
+        self.mesh_device = mesh_device
+        self.cq_id = cq_id
+        self.args = args
+        self._sampling_debug_enabled = is_llama33_70b_model(args)
+        self.sub_core_grids = getattr(args, "sub_core_grids", None)
+        self.tt_sampling = TTSampling(mesh_device=mesh_device, tt_ccl=tt_ccl, args=args)
+        self.tt_penalties = TTPenalties(mesh_device=mesh_device, args=args)
+
+        self._penalties_active = False
+
+        self._trace_states: dict[_TraceKey, dict] = {}
+        seed_batch_size = self.tt_sampling.max_batch_size * self.tt_sampling._sampling_dp
+        self.seed_manager = SeedManager(
+            self.tt_sampling,
+            max_batch_size=seed_batch_size,
+        )
+
+    def _new_trace_state(self):
+        return {"id": None, "input": None, "output": None, "kwargs": {}}
+
+    def _trace_slot(self, penalties_on: bool, log_probs_on: bool, force_argmax: bool):
+        key = _TraceKey(penalties_on=penalties_on, log_probs_on=log_probs_on, force_argmax=force_argmax)
+        slot = self._trace_states.get(key)
+        if slot is None:
+            slot = self._new_trace_state()
+            self._trace_states[key] = slot
+        return key, slot
+
+    def reset_trace(self):
+        """
+        Drop any cached trace metadata for both penalties/no-penalties and log-probs/no-log-probs paths.
+        """
+        for key, slot in self._trace_states.items():
+            if slot["id"] is None:
+                continue
+            logger.debug(
+                f"Resetting sampling trace (penalties={key.penalties_on}, log_probs={key.log_probs_on}, force_argmax={key.force_argmax}, trace_id={slot['id']})"
+            )
+            try:
+                ttnn.release_trace(self.mesh_device, slot["id"])
+            except Exception as e:
+                logger.warning(f"Failed to release trace {slot['id']} : {e}")
+                continue
+        self._trace_states.clear()
+
+    def reset_prompt_tokens(self, prompt_tokens):
+        if not self._penalties_active:
+            return
+        self.tt_penalties.reset_prompt_tokens(prompt_tokens)
+
+    def reset_output_state(self, tokens=None):
+        if not self._penalties_active:
+            return
+        self.tt_penalties.reset_output_tokens(tokens)
+
+    # ---------------------------------------------------------------------
+    # Prefill / decode state helpers
+    # ---------------------------------------------------------------------
+    def apply_prefill_state(
+        self,
+        *,
+        sampling_params,
+        prompt_tokens: torch.Tensor | None,
+        empty_slots: list[int],
+        replicate_seeds: bool = True,
+    ):
+        """Prepare sampling state for a prefill request.
+
+        Resets params, seeds, prompt tokens, and output state in the correct order.
+        """
+        self.reset_sampling_params(sampling_params, empty_slots=empty_slots)
+        seed = getattr(sampling_params, "seed", None)
+        # assert on condition that seed is not None
+        assert seed is not None, "sampling_params must be formatted (seed should be a list, not None)"
+        self.seed_manager.reset_seed(seed, empty_slots)
+        self.seed_manager.get_new_values(empty_slots, replicate_seeds=replicate_seeds)
+        if prompt_tokens is not None:
+            self.reset_prompt_tokens(prompt_tokens)
+        self.reset_output_state()
+
+    def apply_decode_state(
+        self,
+        sampling_params_chunks: list,
+        *,
+        reset_batch: bool = False,
+        prompt_tokens: torch.Tensor | None = None,
+        output_tokens: torch.Tensor | None = None,
+    ):
+        """Format, merge (if row-sharded), and apply sampling params for one model instance.
+
+        Args:
+            sampling_params_chunks: List of SamplingParams assigned to this instance.
+                Length-1 for simple cases; >1 for row-sharded (sampling_dp > data_parallel).
+            reset_batch: Also reset prompt tokens and output state (first decode step).
+            prompt_tokens: Prompt tokens for penalty tracking.
+            output_tokens: Output tokens for penalty tracking.
+
+        Does NOT call ``seed_manager.get_new_values()`` — callers manage seed
+        advancement separately since generators call it at different points.
+        """
+        chunks_per_model = len(sampling_params_chunks)
+
+        max_batch_size = self.tt_sampling.max_batch_size
+
+        if chunks_per_model == 1:
+            formatted_params = format_sampling_params(sampling_params_chunks[0], max_batch_size)
+            self.reset_sampling_params(formatted_params)
+        else:
+            # Row-sharded case: format each chunk to max_batch_size, concatenate.
+            # After (0, None) sharding each row gets its own chunk of max_batch_size entries.
+            # Both TTSampling and TTPenalties use the same concatenated params.
+            formatted_chunks = [format_sampling_params(chunk, max_batch_size) for chunk in sampling_params_chunks]
+            concat_fields = {}
+            for field in SAMPLING_PARAM_FIELDS:
+                lists = [getattr(fc, field) for fc in formatted_chunks]
+                if all(v is None for v in lists):
+                    concat_fields[field] = None
+                else:
+                    concat_fields[field] = sum((v if isinstance(v, list) else [v] for v in lists), [])
+            formatted_params = SamplingParams(**concat_fields)
+            self.reset_sampling_params(formatted_params)
+
+        if reset_batch:
+            self.reset_prompt_tokens(prompt_tokens)
+            self.reset_output_state(output_tokens)
+
+    # ---------------------------------------------------------------------
+    # Sampling helpers
+    # ---------------------------------------------------------------------
+    def reset_sampling_params(self, sampling_params, empty_slots: list[int] | None = None):
+        old_force_argmax_sampling = self.tt_sampling.force_argmax_sampling
+        num_logprobs = getattr(sampling_params, "num_logprobs", None)
+        self.tt_sampling.reset_params(
+            k=sampling_params.top_k,
+            p=sampling_params.top_p,
+            temp=sampling_params.temperature,
+            enable_log_probs=sampling_params.enable_log_probs,
+            num_logprobs=num_logprobs,
+            empty_slots=empty_slots,
+        )
+        if self.tt_sampling.force_argmax_sampling != old_force_argmax_sampling:
+            self.reset_trace()
+
+        old_penalties_active = self._penalties_active
+        self._penalties_active = not (
+            is_default_value(sampling_params.presence_penalty, self._DEFAULT_PENALTIES["presence"])
+            and is_default_value(sampling_params.frequency_penalty, self._DEFAULT_PENALTIES["frequency"])
+            and is_default_value(sampling_params.repetition_penalty, self._DEFAULT_PENALTIES["repetition"])
+        )
+        if (
+            not self.tt_sampling.force_argmax_sampling
+            or self._penalties_active
+            or self._penalties_active != old_penalties_active
+        ):
+            self.tt_penalties.reset_params(
+                sampling_params.presence_penalty, sampling_params.frequency_penalty, sampling_params.repetition_penalty
+            )
+        self._log_probs_active = self.tt_sampling.log_probs_calculator.enable_log_probs
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SamplingGenerator reset params",
+            empty_slots=_compact_debug_list(empty_slots),
+            force_argmax=self.tt_sampling.force_argmax_sampling,
+            force_argmax_changed=self.tt_sampling.force_argmax_sampling != old_force_argmax_sampling,
+            penalties_active=self._penalties_active,
+            log_probs_active=self._log_probs_active,
+            temperature=_compact_debug_list(sampling_params.temperature),
+            top_k=_compact_debug_list(sampling_params.top_k),
+            top_p=_compact_debug_list(sampling_params.top_p),
+            presence_penalty=_compact_debug_list(sampling_params.presence_penalty),
+            frequency_penalty=_compact_debug_list(sampling_params.frequency_penalty),
+            repetition_penalty=_compact_debug_list(sampling_params.repetition_penalty),
+            seed=_compact_debug_list(getattr(sampling_params, "seed", None)),
+        )
+
+    def _validate_trace_inputs(self, slot, logits: ttnn.Tensor, tt_out_tok: Optional[ttnn.Tensor]):
+        if slot["input"] is None or slot["output"] is None:
+            raise RuntimeError("Trace metadata missing. Call capture_trace first.")
+
+        if logits is not slot["input"]:
+            raise ValueError(
+                "The provided logits tensor does not match the tensor used during trace capture. "
+                "Call `reset_trace()` before tracing with new tensors."
+            )
+        if isinstance(slot["output"], tuple):
+            if tt_out_tok is not None and tt_out_tok is not slot["output"][0]:
+                raise ValueError(
+                    "The provided output tensor does not match the tensor used during trace capture. "
+                    "Call `reset_trace()` before tracing with new tensors."
+                )
+        else:
+            if tt_out_tok is not None and tt_out_tok is not slot["output"]:
+                raise ValueError(
+                    "The provided output tensor does not match the tensor used during trace capture. "
+                    "Call `reset_trace()` before tracing with new tensors."
+                )
+
+    def _run_sampling(
+        self,
+        logits,
+        *,
+        penalties_on: bool,
+        tt_out_tok: Optional[ttnn.Tensor],
+    ):
+        if penalties_on:
+            logits = self.tt_penalties.apply(logits)
+        tt_tokens, tt_log_probs = self.tt_sampling(logits, tt_out_tok=tt_out_tok)
+        return tt_tokens, tt_log_probs
+
+    def capture_trace(
+        self,
+        logits: ttnn.Tensor,
+        *,
+        tt_out_tok: Optional[ttnn.Tensor] = None,
+        skip_precompile: bool = False,
+    ) -> ttnn.Tensor:
+        """
+        Capture a trace of the sampling pipeline for the given configuration.
+        """
+        penalties_on = self._penalties_active
+        log_probs_on = getattr(self, "_log_probs_active", False)
+        force_argmax = self.tt_sampling.force_argmax_sampling
+
+        key, slot = self._trace_slot(penalties_on, log_probs_on, force_argmax)
+
+        if not skip_precompile:
+            logger.debug(
+                f"Pre-compiling sampling path before trace capture (penalties={penalties_on},log_probs_on={log_probs_on},force_argmax={force_argmax})"
+            )
+            self._run_sampling(
+                logits,
+                penalties_on=penalties_on,
+                tt_out_tok=tt_out_tok,
+            )
+
+        trace_id = ttnn.begin_trace_capture(self.mesh_device, cq_id=self.cq_id)
+        sampled = self._run_sampling(
+            logits,
+            penalties_on=penalties_on,
+            tt_out_tok=tt_out_tok,
+        )
+        ttnn.end_trace_capture(self.mesh_device, trace_id, cq_id=self.cq_id)
+        ttnn.synchronize_device(self.mesh_device)
+
+        if tt_out_tok is not None:
+            if isinstance(sampled, tuple):
+                output = (tt_out_tok, sampled[-1])
+            else:
+                output = (tt_out_tok, sampled)
+        else:
+            output = sampled
+
+        slot["id"] = trace_id
+        slot["input"] = logits
+        slot["output"] = output
+        slot["kwargs"] = {"tt_out_tok": tt_out_tok}
+
+        return slot["output"]
+
+    def _execute_trace(self, key: _TraceKey) -> ttnn.Tensor:
+        slot = self._trace_states.get(key)
+        if slot is None:
+            raise RuntimeError("Trace has not been captured yet.")
+        if slot["id"] is None or slot["output"] is None:
+            raise RuntimeError("Trace has not been captured yet.")
+
+        ttnn.execute_trace(self.mesh_device, slot["id"], cq_id=self.cq_id, blocking=False)
+        return slot["output"]
+
+    def sample(
+        self,
+        logits: ttnn.Tensor,
+        *,
+        enable_trace: bool = True,
+        tt_out_tok: Optional[ttnn.Tensor] = None,
+        skip_precompile: bool = False,
+    ) -> ttnn.Tensor:
+        """
+        Convenience wrapper that either runs the sampling module directly or
+        replays a captured trace.
+        """
+
+        penalties_on = self._penalties_active
+        log_probs_on = getattr(self, "_log_probs_active", False)
+        force_argmax = self.tt_sampling.force_argmax_sampling
+        # Explicit request seeds update a persistent seed tensor every token;
+        # run them directly so trace replay cannot observe stale seed state.
+        use_internal_trace = enable_trace and not self.seed_manager.has_active_request_seed()
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SamplingGenerator sample",
+            enable_trace=enable_trace,
+            use_internal_trace=use_internal_trace,
+            penalties_on=penalties_on,
+            log_probs_on=log_probs_on,
+            force_argmax=force_argmax,
+            logits_shape=list(logits.shape),
+            tt_out_tok_shape=list(tt_out_tok.shape) if tt_out_tok is not None else None,
+        )
+
+        if not use_internal_trace:
+            tt_out = self._run_sampling(
+                logits,
+                penalties_on=penalties_on,
+                tt_out_tok=tt_out_tok,
+            )
+        else:
+            key, slot = self._trace_slot(penalties_on, log_probs_on, force_argmax)
+            if slot["id"] is None:
+                return self.capture_trace(
+                    logits,
+                    tt_out_tok=tt_out_tok,
+                    skip_precompile=skip_precompile,
+                )
+
+            self._validate_trace_inputs(slot, logits, tt_out_tok)
+            tt_out = self._execute_trace(key)
+
+        if penalties_on and tt_out is not None:
+            if isinstance(tt_out, tuple):
+                self.tt_penalties.update_output_tokens(tt_out[0])
+            else:
+                self.tt_penalties.update_output_tokens(tt_out)
+        return tt_out
+
+
+def format_sampling_params(sampling_params, max_batch_size):
+    """
+    Format sampling parameters for on-device use.
+
+    Converts scalar fields to lists, pads all lists to ``max_batch_size``,
+    inverts temperature, clamps top-p/top-k, and normalises penalties.
+
+    Returns a **new** SamplingParams — the input is never mutated.
+    """
+    if not isinstance(sampling_params.temperature, List):
+        update_dict = {field.name: [getattr(sampling_params, field.name)] for field in fields(sampling_params)}
+        sampling_params = replace(sampling_params, **update_dict)
+
+    target_len = max_batch_size
+    assert target_len % 32 == 0, f"Sampling batch size must be a multiple of 32, got {target_len}"
+
+    # Defaults used when padding short lists to target_len
+    defaults = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": 1,
+        "presence_penalty": 0.0,
+        "frequency_penalty": 0.0,
+        "repetition_penalty": 1.0,
+        "seed": None,
+        "num_logprobs": 0,
+        "enable_log_probs": False,
+    }
+
+    def _pad(lst, name):
+        """Return a new list padded to target_len with the default for *name*."""
+        if len(lst) >= target_len:
+            return list(lst)
+        return list(lst) + [defaults[name]] * (target_len - len(lst))
+
+    # Pad core sampling fields (scalar→list already done above)
+    temperature = _pad(sampling_params.temperature, "temperature")
+    top_p = _pad(sampling_params.top_p, "top_p")
+    top_k = _pad(sampling_params.top_k, "top_k")
+
+    # enable_log_probs / num_logprobs: scalar → broadcast to all users.
+    # Multi-element list → pad with default (False/0) for inactive slots.
+    # Single-element list (from scalar→list conversion) → broadcast to all.
+    def _broadcast_pad(lst, name):
+        if not isinstance(lst, list):
+            return [lst] * target_len
+        if len(lst) == 1:
+            return lst * target_len
+        return _pad(lst, name)
+
+    enable_log_probs = _broadcast_pad(sampling_params.enable_log_probs, "enable_log_probs")
+    if getattr(sampling_params, "num_logprobs", None) is not None:
+        num_logprobs = _broadcast_pad(sampling_params.num_logprobs, "num_logprobs")
+    else:
+        num_logprobs = None
+
+    # Normalise and pad penalty / seed fields (may still be None/scalar)
+    def _normalise_and_pad(name):
+        value = getattr(sampling_params, name, None)
+        if value is None:
+            lst = [defaults[name]]
+        elif isinstance(value, List):
+            lst = list(value)
+        else:
+            lst = [value]
+        return _pad(lst, name)
+
+    presence_penalty = _normalise_and_pad("presence_penalty")
+    frequency_penalty = _normalise_and_pad("frequency_penalty")
+    repetition_penalty = _normalise_and_pad("repetition_penalty")
+    seed = _normalise_and_pad("seed")
+
+    # Clamp / transform values in the new lists (no mutation of the input)
+    TOP_P_MIN = 0.0
+    TOP_P_MAX = 1.0
+
+    for i in range(len(temperature)):
+        top_p[i] = clamp(top_p[i], TOP_P_MIN, TOP_P_MAX)
+
+        if temperature[i] == 0:
+            temperature[i] = 1.0
+            top_k[i] = 1
+            # Device sampling treats p=0 as a first-token cutoff; with k=1
+            # this is the compact argmax representation for greedy rows.
+            top_p[i] = 0.0
+        else:
+            temperature[i] = 1 / temperature[i]
+
+        # top_k contract: TT sampling supports up to 32 today.
+        # k < 1 means "no restriction" → max (32); k > 32 → capped to 32.
+        if top_k[i] < 1:
+            top_k[i] = 32
+        if top_k[i] > 32:
+            top_k[i] = 32
+
+        if repetition_penalty[i] == 0:
+            repetition_penalty[i] = defaults["repetition_penalty"]
+
+    kwargs = dict(
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        presence_penalty=presence_penalty,
+        frequency_penalty=frequency_penalty,
+        repetition_penalty=repetition_penalty,
+        seed=seed,
+    )
+    # Only include logprobs fields if the input dataclass has them
+    # (vLLM's TTSamplingParams may not have these fields)
+    input_fields = {f.name for f in fields(sampling_params)}
+    if "num_logprobs" in input_fields:
+        kwargs["num_logprobs"] = num_logprobs
+    if "enable_log_probs" in input_fields:
+        kwargs["enable_log_probs"] = enable_log_probs
+
+    return replace(sampling_params, **kwargs)
+
+
+def broadcast_sampling_params(
+    formatted_sampling_params,
+    idx: int,
+    slot_len: int = 32,
+):
+    """
+    Create a new SamplingParams where each list field is broadcast to a full list of length
+    ``slot_len``, taking the value from ``idx``. Does not mutate the input.
+    """
+    kwargs = {}
+    for f in fields(formatted_sampling_params):
+        value = getattr(formatted_sampling_params, f.name)
+        value_is_list = isinstance(value, List)
+        if value_is_list:
+            chosen = value[idx] if idx < len(value) else value[0]
+        else:
+            chosen = value
+        if value_is_list:
+            # Preserve list fields as lists even when the selected value is None.
+            kwargs[f.name] = [chosen] * slot_len
+        elif chosen is None:
+            kwargs[f.name] = None
+        else:
+            kwargs[f.name] = [chosen] * slot_len
+    return SamplingParams(**kwargs)
+
+
+def chunk_sampling_params(sampling_params, sampling_dp: int) -> list:
+    """
+    Chunk a SamplingParams (or duck-type-compatible object) into ``sampling_dp`` pieces.
+
+    List fields are split evenly (length must be divisible by ``sampling_dp``).
+    Scalar fields are replicated to all chunks.  Falls back to dataclass defaults
+    for missing attributes so that vLLM's TTSamplingParams works transparently.
+
+    Returns a list of SamplingParams.
+    """
+    if sampling_dp == 1:
+        return [sampling_params]
+
+    chunked_fields = {}
+    for field_name in SAMPLING_PARAM_FIELDS:
+        try:
+            val = getattr(sampling_params, field_name)
+        except AttributeError:
+            if hasattr(SamplingParams, field_name):
+                val = getattr(SamplingParams, field_name)
+            else:
+                raise
+        if isinstance(val, list):
+            assert (
+                len(val) % sampling_dp == 0
+            ), f"Sampling param '{field_name}' length {len(val)} not divisible by sampling_dp {sampling_dp}"
+            chunked_fields[field_name] = split_list(val, sampling_dp)
+        else:
+            chunked_fields[field_name] = [val] * sampling_dp
+
+    return [
+        SamplingParams(**{field: chunked_fields[field][i] for field in SAMPLING_PARAM_FIELDS})
+        for i in range(sampling_dp)
+    ]
+
+
+class SeedManager:
+    """Manages per-user RNG seeds for on-device sampling.
+
+    Tracks which users have explicit seeds set (``_seed_active``) and avoids
+    unnecessary host-to-device copies during decode when no seeds are active.
+    """
+
+    def __init__(self, tt_sampling, max_batch_size=32):
+        self.max_batch_size = max_batch_size
+        self.seeds = [None for _ in range(max_batch_size)]
+        self.seed_counters = [0 for _ in range(max_batch_size)]
+        # Pre-allocate RNG objects; actual request seeds are set via reset_seed().
+        self.rngs = [random.Random(secrets.randbits(64)) for _ in range(max_batch_size)]
+        self.tt_sampling = tt_sampling
+        self._sampling_debug_enabled = getattr(tt_sampling, "_sampling_debug_enabled", False)
+        # True when at least one user slot has a non-None request seed.
+        self._seed_active = False
+        # Set to True by reset_seed() so the next get_new_values() pushes
+        # fresh values to the device. When _seed_active is True this pushes
+        # per-user seeds; when False it pushes varied per-user
+        # values to diversify the device RNG state. Cleared after the push.
+        self._reseted = False
+        # When True, the next get_new_values() must push MAX_UINT32 (SKIP) so
+        # the device transitions from rand_tile_init to rand_tile advance.
+        self._needs_skip = False
+        # True only for the most recent get_new_values() call when at least
+        # one active slot used an explicit request seed.
+        self._active_request_seed = False
+        # Mesh mapper for sharding seeds across rows when sampling_dp > 1.
+        if tt_sampling._sampling_dp > 1:
+            self._seed_mapper = ttnn.ShardTensor2dMesh(
+                tt_sampling.mesh_device, dims=tt_sampling._param_dims, mesh_shape=tt_sampling.cluster_shape
+            )
+        else:
+            self._seed_mapper = None
+
+    def _next_unseeded_rng_seed(self) -> int:
+        return secrets.randbits(64)
+
+    def _next_unseeded_device_seed(self) -> int:
+        return secrets.randbelow(DEVICE_SEED_MAX) + 1
+
+    def _next_device_seed_from_rng(self, rng: random.Random) -> int:
+        return rng.randint(1, DEVICE_SEED_MAX)
+
+    def _next_device_seed_for_slot(self, slot: int) -> int:
+        request_seed = self.seeds[slot]
+        if request_seed is None:
+            return self._next_device_seed_from_rng(self.rngs[slot])
+        device_seed = _hash_request_seed_to_device_seed(int(request_seed), self.seed_counters[slot])
+        self.seed_counters[slot] += 1
+        return device_seed
+
+    def _seed_from_slot_params(self, seeds, slot: int):
+        if seeds is None:
+            return None
+        if isinstance(seeds, torch.Tensor):
+            flat = seeds.reshape(-1)
+            if slot < 0 or slot >= flat.numel():
+                return None
+            seed = flat[slot]
+        elif isinstance(seeds, list):
+            if slot < 0 or slot >= len(seeds):
+                return None
+            seed = seeds[slot]
+        else:
+            seed = seeds
+
+        if seed is None:
+            return None
+        if isinstance(seed, torch.Tensor):
+            if seed.numel() == 0:
+                return None
+            seed = seed.reshape(-1)[0].item()
+        return int(seed)
+
+    def reset_seed_from_slots(self, seeds, user_ids):
+        """Reset decode seed state from slot-indexed sampling params."""
+        if user_ids is None:
+            user_ids = range(self.max_batch_size)
+        for user in user_ids:
+            slot = int(user)
+            seed = self._seed_from_slot_params(seeds, slot)
+            self.seeds[slot] = seed
+            self.seed_counters[slot] = 0
+            if seed is None:
+                self.rngs[slot].seed(self._next_unseeded_rng_seed())
+            else:
+                self.rngs[slot].seed(int(seed))
+        self._seed_active = any(s is not None for s in self.seeds)
+        self._reseted = True
+
+    def reset_seed_from_slots_if_needed(self, seeds, user_ids) -> bool:
+        """Reset only active slots whose slot-indexed seed changed."""
+        if user_ids is None:
+            user_ids = range(self.max_batch_size)
+        reset_slots = []
+        for user in user_ids:
+            slot = int(user)
+            if self._seed_from_slot_params(seeds, slot) != self.seeds[slot]:
+                reset_slots.append(slot)
+        if not reset_slots:
+            return False
+        self.reset_seed_from_slots(seeds, reset_slots)
+        return True
+
+    def align_seed_counters_to_positions(self, seeds, user_ids, positions, offset: int = 1):
+        """Make explicit-seed decode independent of persistent slot lifetime.
+
+        vLLM can temporarily remove running requests from the persistent batch
+        while admitting another prefill batch, then re-add them in different
+        slots. For explicit request seeds, deriving the per-token device seed
+        from the absolute decode position keeps the stream reproducible even
+        when the Python-side slot counter was reset or moved.
+        """
+        if positions is None:
+            return
+        if user_ids is None:
+            user_ids = range(self.max_batch_size)
+
+        if isinstance(positions, torch.Tensor):
+            flat_positions = positions.reshape(-1)
+
+            def _position(slot):
+                if slot < 0 or slot >= flat_positions.numel():
+                    return None
+                pos = flat_positions[slot]
+                return int(pos.item())
+
+        elif isinstance(positions, list):
+
+            def _position(slot):
+                if slot < 0 or slot >= len(positions):
+                    return None
+                return int(positions[slot])
+
+        else:
+
+            def _position(_slot):
+                return int(positions)
+
+        for user in user_ids:
+            slot = int(user)
+            seed = self._seed_from_slot_params(seeds, slot)
+            if seed is None:
+                continue
+            position = _position(slot)
+            if position is None or position < 0:
+                continue
+            self.seed_counters[slot] = max(0, position + offset)
+
+    def has_active_request_seed(self) -> bool:
+        return self._active_request_seed
+
+    def _debug_state(self, slots=None):
+        if slots is None:
+            slots = range(self.max_batch_size)
+        state = []
+        for slot in slots:
+            slot = int(slot)
+            if slot < 0 or slot >= self.max_batch_size:
+                continue
+            seed = self.seeds[slot]
+            if seed is not None:
+                state.append((slot, seed))
+        return _compact_debug_list(state)
+
+    def apply_slot_remap(self, remap):
+        """Reindex RNG state after batch condense.
+
+        ``remap`` is a 1-D int tensor of length ``max_batch_size`` where
+        ``remap[i] = j`` means slot *i* now holds the request that was
+        previously at slot *j*. Identity entries (``remap[i] == i``) are
+        no-ops. Only non-identity entries trigger a move.
+        """
+        if not self._seed_active:
+            return
+        moves = [(int(remap[i]), i) for i in range(len(remap)) if int(remap[i]) != i]
+        if not moves:
+            _log_sampling_debug(
+                self._sampling_debug_enabled, "SeedManager slot remap identity", seed_active=self._seed_active
+            )
+            return
+        # Snapshot the state we're about to overwrite.
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SeedManager slot remap",
+            moves=_compact_debug_list(moves),
+            state_before=self._debug_state(),
+        )
+        old_seeds = list(self.seeds)
+        old_counters = list(self.seed_counters)
+        old_rngs = list(self.rngs)
+        moved_sources = {old_slot for old_slot, _ in moves}
+        moved_destinations = {new_slot for _, new_slot in moves}
+        for old_slot, new_slot in moves:
+            self.seeds[new_slot] = old_seeds[old_slot]
+            self.seed_counters[new_slot] = old_counters[old_slot]
+            # copy.copy preserves internal RNG state but creates an
+            # independent object so the old slot reference does not alias
+            # the new one.
+            self.rngs[new_slot] = copy.copy(old_rngs[old_slot])
+        for old_slot in moved_sources - moved_destinations:
+            self.seeds[old_slot] = None
+            self.seed_counters[old_slot] = 0
+        self._seed_active = any(s is not None for s in self.seeds)
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SeedManager slot remap done",
+            seed_active=self._seed_active,
+            state_after=self._debug_state(),
+        )
+
+    def reset_seed(self, seeds, user_ids):
+        """Update RNG state for the given user slots after a prefill.
+
+        Args:
+            seeds: Seed values in request order. Accepts a list, tensor, scalar,
+                or None (treated as all unseeded).
+            user_ids: Batch slot indices being prefilled.
+        """
+        user_ids = [int(user) for user in user_ids]
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SeedManager reset prefill",
+            user_ids=_compact_debug_list(user_ids),
+            requested_seeds=_compact_debug_list(seeds),
+            state_before=self._debug_state(user_ids),
+        )
+        for i, user in enumerate(user_ids):
+            slot = int(user)
+            seed = self._seed_from_slot_params(seeds, i)
+            self.seeds[slot] = seed
+            self.seed_counters[slot] = 0
+            if seed is None:
+                self.rngs[slot].seed(self._next_unseeded_rng_seed())
+            else:
+                self.rngs[slot].seed(int(seed))
+        self._seed_active = any(s is not None for s in self.seeds)
+        self._reseted = True
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SeedManager reset prefill done",
+            seed_active=self._seed_active,
+            state_after=self._debug_state(user_ids),
+        )
+
+    def get_new_values(self, empty_slots=None, replicate_seeds=False):
+        """Generate and push new seed values to the device.
+
+        **Seeded path** (``_seed_active=True``):
+        Advances each active slot seed state and copies the new values to
+        the device every step. Explicit request seeds produce slot-independent
+        device seeds derived from the request seed and the slot counter. Some
+        decode callers align that counter to the absolute token position so
+        vLLM batch-layout changes cannot reset a request's random stream.
+
+        **Unseeded path** (``_seed_active=False``):
+        Uses a three-state machine to ensure each user gets a unique device
+        RNG state without redundant host-to-device copies during decode:
+
+          State 1 - **init** (``_reseted=True``):
+            Push varied per-user values from system entropy.
+
+          State 2 - **transition** (``_needs_skip=True``):
+            Push MAX_UINT32 (SKIP) so the device stops reinitializing and
+            starts advancing via rand_tile().
+
+          State 3 - **steady** (both flags clear):
+            Early-return with no device copy.
+        """
+        if empty_slots is None:
+            empty_slots = list(range(self.max_batch_size))
+        else:
+            empty_slots = [int(slot) for slot in empty_slots]
+        empty_slot_set = set(empty_slots)
+        self._active_request_seed = any(self.seeds[i] is not None for i in empty_slot_set)
+
+        if not self._seed_active:
+            self._active_request_seed = False
+            if self._reseted:
+                new_seeds = [self._next_unseeded_device_seed() for _ in range(self.max_batch_size)]
+                self._needs_skip = True
+            elif self._needs_skip:
+                new_seeds = [MAX_UINT32] * self.max_batch_size
+                self._needs_skip = False
+            else:
+                # State 3 (steady): device already has SKIP, rand_tile
+                # advances on its own, so no host-to-device copy is needed.
+                _log_sampling_debug(
+                    self._sampling_debug_enabled,
+                    "SeedManager seed update skipped",
+                    active_slots=_compact_debug_list(empty_slots),
+                    seed_active=self._seed_active,
+                    reseted=self._reseted,
+                    needs_skip=self._needs_skip,
+                )
+                return
+        else:
+            new_seeds = [
+                self._next_device_seed_for_slot(i) if i in empty_slot_set else MAX_UINT32
+                for i in range(self.max_batch_size)
+            ]
+            if replicate_seeds:
+                assert len(empty_slots) == 1, "Cannot replicate seeds if empty_slots is not length 1"
+                new_seeds = self.max_batch_size * [new_seeds[empty_slots[0]]]
+
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "SeedManager seed update",
+            active_slots=_compact_debug_list(empty_slots),
+            replicate_seeds=replicate_seeds,
+            seed_active=self._seed_active,
+            reseted=self._reseted,
+            needs_skip=self._needs_skip,
+            new_device_seeds=_compact_debug_list(new_seeds),
+            state_after_counter_advance=self._debug_state(empty_slots),
+        )
+
+        new_seed_tt = ttnn.from_torch(
+            torch.tensor(new_seeds), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_mapper=self._seed_mapper
+        )
+        ttnn.copy_host_to_device_tensor(new_seed_tt, self.tt_sampling.seeds_tt_tensor)
+        self._reseted = False

code/models/common/sampling/sampling_params.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SamplingParams:
+    """
+    Used in Generator decode forward functions for greedy decoding / sampling on device.
+    The same data class exists in vLLM at vllm/v1/worker/tt_model_runner.py.
+    """
+
+    temperature: float | list[float]
+    top_k: int | list[int]
+    top_p: float | list[float]
+    presence_penalty: float | list[float] = 0.0
+    frequency_penalty: float | list[float] = 0.0
+    repetition_penalty: float | list[float] = 1.0
+    seed: int | list[int] | None = None
+    enable_log_probs: bool | list[bool] = False
+    num_logprobs: int | list[int] = 0

code/models/common/sampling/tt_log_probs.py

@@ -0,0 +1,713 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.sampling._utils import filter_none
+
+# Maximum number of top logprobs that can be requested (OpenAI API limit)
+MAX_TOP_LOGPROBS = 20
+# Number of top logprobs computed on device (gathered top-k from all devices)
+DEVICE_TOP_K = 32
+
+
+@dataclass
+class LogProbsResult:
+    """Result of log-probs calculation for a batch.
+
+    Contains logprobs and global indices for the gathered top-k tokens across all
+    devices.  The sampled token is always part of the gathered top-k (it was selected
+    from them by ttnn.sampling), so its logprob can be looked up by matching its
+    index in ``topk_indices``.
+
+    Attributes:
+        topk_logprobs: Optional ttnn.Tensor of shape (1, 1, batch_size, DEVICE_TOP_K)
+            containing logprobs for the gathered top-k tokens.
+        topk_indices: Optional ttnn.Tensor of shape (1, 1, batch_size, DEVICE_TOP_K)
+            containing global vocabulary indices for the gathered top-k tokens.
+        topk_logprobs_host: Optional ttnn.Tensor on host of shape (1, 1, batch_size, DEVICE_TOP_K)
+            containing logprobs for the gathered top-k tokens after moving to host.
+        topk_indices_host: Optional ttnn.Tensor on host of shape (1, 1, batch_size, DEVICE_TOP_K)
+            containing global vocabulary indices for the gathered top-k tokens after moving to host.
+    """
+
+    topk_logprobs: Optional[ttnn.Tensor]
+    topk_indices: Optional[ttnn.Tensor]
+    topk_logprobs_host: Optional[ttnn.Tensor]
+    topk_indices_host: Optional[ttnn.Tensor]
+
+    def cpu(self, blocking: bool = True) -> "LogProbsResult":
+        """Transfer device tensors to host CPU.
+
+        Args:
+            blocking: If True, wait for the host tensor to be ready before returning.
+
+        Returns:
+            A new LogProbsResult so that each iteration gets its own host
+            tensor references (avoids race with trace-captured singletons).
+        """
+        return LogProbsResult(
+            topk_logprobs=None,
+            topk_indices=None,
+            topk_logprobs_host=self.topk_logprobs.cpu(blocking=blocking),
+            topk_indices_host=self.topk_indices.cpu(blocking=blocking),
+        )
+
+    def extract_user(self, user_batch_idx: int) -> "LogProbsResult":
+        """Extract a single user's top-K logprobs from a batched host result.
+
+        Args:
+            user_batch_idx: Index of the user within the batch dimension.
+
+        Returns:
+            A new LogProbsResult with host tensors sliced to shape (1, DEVICE_TOP_K).
+        """
+        lp_torch = ttnn.to_torch(ttnn.get_device_tensors(self.topk_logprobs_host)[0])
+        idx_torch = ttnn.to_torch(ttnn.get_device_tensors(self.topk_indices_host)[0])
+        return LogProbsResult(
+            topk_logprobs=None,
+            topk_indices=None,
+            topk_logprobs_host=lp_torch[0, 0, user_batch_idx, :].unsqueeze(0),
+            topk_indices_host=idx_torch[0, 0, user_batch_idx, :].unsqueeze(0),
+        )
+
+    def to_torch_pair(self):
+        """Convert host tensors to a (logprobs, indices) pair of torch tensors.
+
+        Returns:
+            Tuple of (logprobs_tensor, indices_tensor), each of shape (DEVICE_TOP_K,).
+        """
+        lp = self.topk_logprobs_host
+        idx = self.topk_indices_host
+        if isinstance(lp, torch.Tensor):
+            return lp.reshape(-1)[:DEVICE_TOP_K].float(), idx.reshape(-1)[:DEVICE_TOP_K].int()
+        return ttnn.to_torch(lp).reshape(-1)[:DEVICE_TOP_K].float(), ttnn.to_torch(idx).reshape(-1)[:DEVICE_TOP_K].int()
+
+
+def reformat_logprobs(output_log_probs, batch_size):
+    """Convert a list of per-user logprobs into the format expected by vLLM.
+
+    Args:
+        output_log_probs: List of LogProbsResult, torch.Tensor, float, int, or None per user.
+        batch_size: Total batch size.
+
+    Returns:
+        For new path (LogProbsResult): tuple of (topk_lp[B, DEVICE_TOP_K], topk_idx[B, DEVICE_TOP_K])
+        For old path: flat tensor of shape [B]
+    """
+    has_logprobs_result = any(isinstance(lp, LogProbsResult) for lp in output_log_probs)
+    if has_logprobs_result:
+        all_lp = torch.zeros(batch_size, DEVICE_TOP_K, dtype=torch.float32)
+        all_idx = torch.zeros(batch_size, DEVICE_TOP_K, dtype=torch.int32)
+        for i, lp in enumerate(output_log_probs):
+            if isinstance(lp, LogProbsResult) and lp.topk_logprobs_host is not None:
+                all_lp[i], all_idx[i] = lp.to_torch_pair()
+        return (all_lp, all_idx)
+    else:
+        flat_lp = torch.ones(batch_size, dtype=torch.float32)
+        for i, lp in enumerate(output_log_probs):
+            if lp is not None and isinstance(lp, (torch.Tensor, float, int)):
+                flat_lp[i] = float(lp)
+        return flat_lp
+
+
+class LogProbsCalculator:
+    """
+    Class to calculate log-probs for a given logits tensor and indices tensor.
+
+    Supports two modes:
+    - Old mode (backward compat): calculate_log_probs() returns single sampled-token logprob
+    - New mode (gpt-oss-120b): calculate_topk_log_probs() returns top-32 logprobs + indices
+
+    Args:
+        mesh_device: MeshDevice to use for all-gather operations
+        sub_core_grids: Sub-core grid configuration for operations (optional)
+        tt_ccl: CCL object for distributed operations (optional)
+        batch_size: Maximum batch size for log-probs calculation (default: 32)
+        use_topk_logprobs: If True, allocate tensors for new top-K path instead of old path
+    """
+
+    def __init__(
+        self,
+        mesh_device: ttnn.MeshDevice,
+        sub_core_grids: ttnn.CoreRangeSet = None,
+        tt_ccl=None,
+        batch_size: int = 32,
+        use_topk_logprobs: bool = False,
+    ):
+        self.global_max = None
+        self.global_exp_sum = None
+        self.mesh_device = mesh_device
+        self.enable_log_probs = False  # default to False
+        # Per-user boolean array tracking which users have logprobs enabled
+        self.logprobs_enabled = [False] * batch_size
+        # Per-user integer array tracking how many top logprobs each user requested (0-20)
+        self.num_logprobs = [0] * batch_size
+        # Flag: True when at least one user needs top-k logprobs (num_logprobs > 0)
+        self.topk_logprobs_needed = False
+        self.cluster_shape = list(mesh_device.shape)
+        self.sub_core_grids = sub_core_grids
+        self.tt_ccl = tt_ccl
+        self.batch_size = batch_size
+        self._use_topk_logprobs = use_topk_logprobs
+        self.common_args = filter_none(
+            {
+                "sub_core_grids": sub_core_grids,
+            }
+        )
+
+        # CCL introspection (same pattern as TTSampling)
+        self._line_all_gather = getattr(self.tt_ccl, "line_all_gather", None)
+        self._line_all_gather_supports_buffer_key = False
+        if callable(self._line_all_gather):
+            try:
+                sig = inspect.signature(self._line_all_gather)
+                params = sig.parameters
+                self._line_all_gather_supports_buffer_key = "buffer_key" in params or any(
+                    p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
+                )
+            except (TypeError, ValueError):
+                logger.warning("Unable to inspect line_all_gather signature; assuming no buffer_key support.")
+
+        num_devices = self.mesh_device.get_num_devices()
+
+        # Determine the TP dimension: for 2D meshes, logits are sharded across
+        # the larger dimension (TP). For 1D or single-device, use all devices.
+        if self.cluster_shape[0] > 1 and self.cluster_shape[1] > 1:
+            # 2D mesh: TP axis is the larger dimension
+            tp_axis = 0 if self.cluster_shape[0] >= self.cluster_shape[1] else 1
+            num_devices_for_sharding = self.cluster_shape[tp_axis]
+            self._all_gather_cluster_axis = tp_axis
+        elif num_devices > 1:
+            # 1D mesh
+            num_devices_for_sharding = num_devices
+            self._all_gather_cluster_axis = None
+        else:
+            # Single device
+            num_devices_for_sharding = num_devices
+            self._all_gather_cluster_axis = None
+
+        self.num_devices_for_sharding = num_devices_for_sharding
+
+        # Initialize tensors based on mode
+        # Old path tensors (backward compat for non-gpt-oss models)
+        self.mask = None
+        self.output_tensor = None
+        # New path tensors (gpt-oss-120b top-K logprobs)
+        self.topk_logprobs_output = None
+        self.topk_indices_output = None
+
+        if use_topk_logprobs:
+            # New path: allocate top-K output tensors
+            self.topk_logprobs_output = ttnn.as_tensor(
+                torch.zeros(1, 1, batch_size, DEVICE_TOP_K),
+                dtype=ttnn.bfloat16,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                device=self.mesh_device,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+            )
+            self.topk_indices_output = ttnn.as_tensor(
+                torch.zeros(1, 1, batch_size, DEVICE_TOP_K, dtype=torch.int32),
+                dtype=ttnn.uint32,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                device=self.mesh_device,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+            )
+        else:
+            # Old path: allocate mask and output tensors
+            mask_tensor = (
+                torch.arange(num_devices_for_sharding).unsqueeze(1).expand(num_devices_for_sharding, batch_size)
+            )
+
+            if self.cluster_shape[0] > 1 and self.cluster_shape[1] > 1:
+                dims = (0, None) if self._all_gather_cluster_axis == 0 else (None, 0)
+                mesh_mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=dims, mesh_shape=self.cluster_shape)
+            elif num_devices > 1:
+                mesh_mapper = ttnn.ShardTensorToMesh(self.mesh_device, dim=0)
+            else:
+                mesh_mapper = ttnn.ReplicateTensorToMesh(self.mesh_device)
+
+            self.mask = ttnn.as_tensor(
+                mask_tensor,
+                dtype=ttnn.bfloat16,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                device=self.mesh_device,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                preprocess=lambda x: x.to(torch.bfloat16),
+                mesh_mapper=mesh_mapper,
+            )
+            self.output_tensor = ttnn.as_tensor(
+                torch.ones(1, 1, 1, batch_size),
+                dtype=ttnn.bfloat16,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                device=self.mesh_device,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+            )
+
+    def _perform_all_gather(self, tensor: ttnn.Tensor, dim: int, num_links: int, buffer_key: str = None):
+        if callable(self._line_all_gather):
+            kwargs = {
+                "dim": dim,
+                "num_links": num_links,
+                "memory_config": tensor.memory_config(),
+                "cluster_axis": self._all_gather_cluster_axis,
+            }
+            if self._line_all_gather_supports_buffer_key and buffer_key is not None:
+                kwargs["buffer_key"] = buffer_key
+            return self._line_all_gather(tensor, **kwargs)
+
+        return ttnn.all_gather(
+            tensor,
+            dim=dim,
+            num_links=num_links,
+            memory_config=tensor.memory_config(),
+            cluster_axis=self._all_gather_cluster_axis,
+            topology=ttnn.Topology.Linear,
+        )
+
+    def set_log_probs_mode(
+        self,
+        enable_log_probs: bool | list[bool] = False,
+        num_logprobs: int | list[int] | None = None,
+        empty_slots: list[int] | None = None,
+    ):
+        """Set logprobs mode for the current batch.
+
+        Args:
+            enable_log_probs: Boolean or per-user boolean list. If any user has logprobs
+                enabled, the entire batch runs logprobs computation.
+            num_logprobs: Integer or per-user integer list (0-20). Specifies how many
+                top logprobs to return per user. 0 means sampled token logprob only.
+                Values > 0 trigger top-k logprobs computation on device.
+            empty_slots: Optional list of batch indices at which to apply the new
+                logprobs settings. When provided, only those positions are updated
+                and the rest of the batch retains its previous values.
+        """
+        if empty_slots is not None:
+            # Partial update: only modify the specified batch positions
+            if isinstance(enable_log_probs, list):
+                for i, slot in enumerate(empty_slots):
+                    self.logprobs_enabled[slot] = enable_log_probs[i]
+            else:
+                for slot in empty_slots:
+                    self.logprobs_enabled[slot] = enable_log_probs
+
+            if num_logprobs is not None:
+                if isinstance(num_logprobs, list):
+                    for i, slot in enumerate(empty_slots):
+                        self.num_logprobs[slot] = num_logprobs[i]
+                else:
+                    for slot in empty_slots:
+                        self.num_logprobs[slot] = num_logprobs
+        else:
+            # Full batch update
+            if isinstance(enable_log_probs, list):
+                self.logprobs_enabled = list(enable_log_probs)
+            else:
+                self.logprobs_enabled = [enable_log_probs] * self.batch_size
+
+            if num_logprobs is not None:
+                if isinstance(num_logprobs, list):
+                    self.num_logprobs = list(num_logprobs)
+                else:
+                    self.num_logprobs = [num_logprobs] * self.batch_size
+            else:
+                self.num_logprobs = [0] * self.batch_size
+
+        # Recompute derived flags from the full arrays
+        self.enable_log_probs = any(self.logprobs_enabled)
+        # Top-K computation is needed whenever logprobs are enabled (even with
+        # num_logprobs=0) because the sampled token's logprob is extracted from
+        # the top-K results in the new path.
+        self.topk_logprobs_needed = self.enable_log_probs
+
+    def _compute_global_stats(
+        self,
+        logits_tensor: ttnn.Tensor,
+    ):
+        """
+        To calculate log-probs, we need to calculate the global max and global sum(exp(logits - global_max)) for each chip.
+        This is done by all-gathering the max and sum(exp(logits - global_max)) for each chip and then taking the max and sum of the gathered tensors.
+        log-prob formula: log-prob(x) = logits(x) - global_max - log(sum(exp(logits - global_max)))
+
+        Args:
+            logits_tensor (ttnn.Tensor): Logits as model output (1, 1, batch_size, vocab_size_per_device)
+        """
+        # Calculate local max
+        local_max_tensor = ttnn.max(logits_tensor, dim=-1, keepdim=True, **self.common_args)
+
+        gathered_max_tensors = self._perform_all_gather(
+            local_max_tensor,
+            dim=1,
+            num_links=1,
+            buffer_key="LOGPROBS_MAX_REDUCTION",
+        )
+        # Convert to ROW_MAJOR_LAYOUT due to memory clobbering which affects all ttnn.reshape ops with TILE_LAYOUT
+        gathered_max_tensors = ttnn.to_layout(gathered_max_tensors, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        ttnn.deallocate(local_max_tensor)
+        D = self.num_devices_for_sharding
+        B = gathered_max_tensors.shape[2]
+        gathered_max_tensors = ttnn.reshape(gathered_max_tensors, (1, 1, D, B), **self.common_args)
+        gathered_max_tensors = ttnn.to_layout(gathered_max_tensors, ttnn.TILE_LAYOUT, **self.common_args)
+
+        self.global_max = ttnn.max(gathered_max_tensors, dim=2, keepdim=True, **self.common_args)
+
+        global_max_to_subtract = ttnn.to_layout(self.global_max, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        global_max_to_subtract = ttnn.reshape(global_max_to_subtract, (1, 1, B, 1), **self.common_args)
+        global_max_to_subtract = ttnn.to_layout(global_max_to_subtract, ttnn.TILE_LAYOUT, **self.common_args)
+
+        # Calculate stable local sum-exp using subtract of global-max from each local logit
+        subtracted_tensor = ttnn.subtract(logits_tensor, global_max_to_subtract, **self.common_args)
+        exp_tensor = ttnn.exp(subtracted_tensor, **self.common_args)
+        ttnn.deallocate(global_max_to_subtract)
+        ttnn.deallocate(subtracted_tensor)
+        sum_exp_tensor = ttnn.sum(exp_tensor, dim=-1, keepdim=True, **self.common_args)
+        ttnn.deallocate(exp_tensor)
+
+        gathered_sum_exp_tensors = self._perform_all_gather(
+            sum_exp_tensor,
+            dim=1,
+            num_links=1,
+            buffer_key="LOGPROBS_SUM_EXP_REDUCTION",
+        )
+        gathered_sum_exp_tensors = ttnn.to_layout(gathered_sum_exp_tensors, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        ttnn.deallocate(sum_exp_tensor)
+        B_sum = gathered_sum_exp_tensors.shape[2]
+        gathered_sum_exp_tensors = ttnn.reshape(gathered_sum_exp_tensors, (1, 1, D, B_sum), **self.common_args)
+        gathered_sum_exp_tensors = ttnn.to_layout(gathered_sum_exp_tensors, ttnn.TILE_LAYOUT, **self.common_args)
+
+        self.global_exp_sum = ttnn.sum(gathered_sum_exp_tensors, dim=2, keepdim=True, **self.common_args)
+        ttnn.deallocate(gathered_sum_exp_tensors)
+
+    def _is_supported(self):
+        """Check if logprobs computation is supported on this device configuration."""
+        num_devices = self.mesh_device.get_num_devices()
+        if num_devices not in (8, 32):
+            return False
+        if self.num_devices_for_sharding < 2:
+            return False
+        return True
+
+    # -----------------------------------------------------------------------
+    # Old path (backward compat for non-gpt-oss models)
+    # -----------------------------------------------------------------------
+
+    def _prepare_relevant_logits(self, logits_tensor: ttnn.Tensor, global_idx_tensor: ttnn.Tensor):
+        """
+        Prepare global idx tensor with correct values on all devices.
+        """
+        size_per_device = logits_tensor.shape[-1]
+
+        # convert global_idx_tensor to ttnn.TILE_LAYOUT
+        global_idx_tilized_tensor = ttnn.to_layout(global_idx_tensor, ttnn.TILE_LAYOUT, **self.common_args)
+
+        # TODO: Raise an issue on this since for UINT_32 ttnn.div produces incorrect output (all zeros)
+        global_idx_tilized_tensor = ttnn.typecast(global_idx_tilized_tensor, ttnn.float32, **self.common_args)
+
+        # Get chip_id for each user based on global_idx values in global_idx_tensor
+        chip_ids_tensor = ttnn.div(
+            global_idx_tilized_tensor,
+            size_per_device,
+            rounding_mode="floor",
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            **self.common_args,
+        )
+
+        # Get local index for each user based on global_idx values in global_idx_tensor
+        remainder_tensor = ttnn.remainder(
+            global_idx_tilized_tensor,
+            size_per_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            **self.common_args,
+        )
+
+        # Convert remainder_tensor to int32
+        remainder_tensor = ttnn.typecast(remainder_tensor, ttnn.uint32, **self.common_args)
+        # convert to ROW_MAJOR_LAYOUT due to memory clobbering which affects all ttnn.reshape ops with TILE_LAYOUT
+        remainder_tensor = ttnn.to_layout(remainder_tensor, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        batch_vol = remainder_tensor.shape[2] * remainder_tensor.shape[3]
+        remainder_tensor = ttnn.reshape(remainder_tensor, (1, 1, batch_vol, 1), **self.common_args)
+        remainder_tensor = ttnn.to_layout(remainder_tensor, ttnn.TILE_LAYOUT, **self.common_args)
+
+        # Get logits for each user on each chip based on local index
+        selected_logits_tensor = ttnn.gather(logits_tensor, dim=3, index=remainder_tensor, **self.common_args)
+
+        # convert to ROW_MAJOR_LAYOUT due to memory clobbering which affects all ttnn.reshape ops with TILE_LAYOUT
+        selected_logits_tensor = ttnn.to_layout(selected_logits_tensor, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        batch_vol_s = selected_logits_tensor.shape[2] * selected_logits_tensor.shape[3]
+        selected_logits_tensor = ttnn.reshape(selected_logits_tensor, (1, 1, 1, batch_vol_s), **self.common_args)
+        selected_logits_tensor = ttnn.to_layout(selected_logits_tensor, ttnn.TILE_LAYOUT, **self.common_args)
+        # Compare mask to chip_ids tensor and select correct positions for each user on all chips inplace
+        ttnn.eq_(chip_ids_tensor, self.mask, **self.common_args)
+
+        # Multiply selected_logits_tensor with chip_ids_tensor to get expected logits for each user
+        selected_logits_tensor = ttnn.multiply(selected_logits_tensor, chip_ids_tensor, **self.common_args)
+
+        # All gather logits across all devices
+        selected_logits_tensor = self._perform_all_gather(
+            selected_logits_tensor,
+            dim=1,
+            num_links=1,
+            buffer_key="LOGPROBS_LOGITS",
+        )
+
+        selected_logits_tensor = ttnn.to_layout(selected_logits_tensor, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        D_s = self.num_devices_for_sharding
+        B_g = selected_logits_tensor.shape[3]
+        selected_logits_tensor = ttnn.reshape(selected_logits_tensor, (1, 1, D_s, B_g), **self.common_args)
+        selected_logits_tensor = ttnn.to_layout(selected_logits_tensor, ttnn.TILE_LAYOUT, **self.common_args)
+
+        # Apply sum over device dimension to get logits for each user on all chips
+        selected_logits_tensor = ttnn.sum(selected_logits_tensor, dim=2, keepdim=True, **self.common_args)
+
+        return selected_logits_tensor
+
+    def _calculate_log_probs(self, sampled_logits_tensor: ttnn.Tensor):
+        """
+        Calculate log-probs for a given logits tensor with formula:
+        log-prob(x) = logits(x) - global_max - log(global_exp_sum)
+        """
+        out = ttnn.subtract(sampled_logits_tensor, self.global_max, **self.common_args)
+        log_global_exp_sum = ttnn.log(self.global_exp_sum, **self.common_args)
+        # Subtract and put result to self.output_tensor
+        ttnn.subtract(out, log_global_exp_sum, output_tensor=self.output_tensor, **self.common_args)
+
+    def calculate_log_probs(
+        self,
+        logits_tensor: ttnn.Tensor,
+        indices_tensor: ttnn.Tensor,
+    ):
+        """
+        Calculate log-probs for a given logits tensor and indices tensor.
+        Returns None if log-probs are not requested, not supported, or the device count is not 8 or 32.
+        (Old path — backward compat for non-gpt-oss models)
+        """
+        if not self.enable_log_probs:
+            return None
+
+        if not self._is_supported():
+            return None
+
+        # Calculating log-probs requires bfloat16 precision for near-stable sum-exp calculation
+        if logits_tensor.dtype == ttnn.bfloat8_b:
+            logits_tensor = ttnn.typecast(logits_tensor, ttnn.bfloat16, **self.common_args)
+
+        # Compute global max and global sum(exp(logits - global_max)) for each chip
+        self._compute_global_stats(logits_tensor)
+
+        # Prepare relevant logits for each user on each chip
+        relevant_logits = self._prepare_relevant_logits(logits_tensor, indices_tensor)
+
+        # Calculate log-probs for each user on each chip and stores in self.output_tensor
+        self._calculate_log_probs(relevant_logits)
+
+        return self.output_tensor
+
+    # -----------------------------------------------------------------------
+    # New path (gpt-oss-120b top-K logprobs)
+    # -----------------------------------------------------------------------
+
+    def _calculate_topk_log_probs_from_values(self, topk_values: ttnn.Tensor):
+        """Compute logprobs for gathered top-k values using pre-computed global stats.
+
+        Applies the log-softmax formula: logprob = logit - global_max - log(global_exp_sum)
+        to each of the gathered top-k values.
+
+        Args:
+            topk_values: Gathered top-k values tensor of shape (1, 1, batch_size, num_topk).
+        """
+        B = topk_values.shape[2]
+
+        # Reshape global_max from (1,1,1,B) to (1,1,B,1) for broadcasting with (1,1,B,K)
+        global_max_bcast = ttnn.to_layout(self.global_max, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        global_max_bcast = ttnn.reshape(global_max_bcast, (1, 1, B, 1), **self.common_args)
+        global_max_bcast = ttnn.to_layout(global_max_bcast, ttnn.TILE_LAYOUT, **self.common_args)
+
+        # Compute log(global_exp_sum) and reshape for broadcasting
+        log_global_exp_sum = ttnn.log(self.global_exp_sum, **self.common_args)
+        log_global_exp_sum_bcast = ttnn.to_layout(log_global_exp_sum, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        log_global_exp_sum_bcast = ttnn.reshape(log_global_exp_sum_bcast, (1, 1, B, 1), **self.common_args)
+        log_global_exp_sum_bcast = ttnn.to_layout(log_global_exp_sum_bcast, ttnn.TILE_LAYOUT, **self.common_args)
+        ttnn.deallocate(log_global_exp_sum)
+
+        # Apply log-softmax formula: logprob = logit - global_max - log(global_exp_sum)
+        ttnn.subtract(topk_values, global_max_bcast, output_tensor=self.topk_logprobs_output, **self.common_args)
+        ttnn.subtract(
+            self.topk_logprobs_output,
+            log_global_exp_sum_bcast,
+            output_tensor=self.topk_logprobs_output,
+            **self.common_args,
+        )
+        ttnn.deallocate(global_max_bcast)
+        ttnn.deallocate(log_global_exp_sum_bcast)
+
+        return self.topk_logprobs_output
+
+    def calculate_topk_log_probs(
+        self,
+        logits_tensor: ttnn.Tensor,
+        topk_values: ttnn.Tensor,
+        topk_global_indices: ttnn.Tensor,
+        sub_core_grid_topk: ttnn.CoreRangeSet = None,
+    ) -> LogProbsResult | None:
+        """Calculate logprobs for the gathered top-k tokens in a single pass.
+
+        Args:
+            logits_tensor: Full logits tensor, sharded across devices.
+            topk_values: Gathered top-k values from all devices. Shape: (1,1,B,256).
+            topk_global_indices: Global vocabulary indices for the gathered top-k. Shape: (1,1,B,256).
+            sub_core_grid_topk: Sub-core grid for topk operation.
+
+        Returns:
+            LogProbsResult with top-32 logprobs and indices, or None if disabled/unsupported.
+        """
+        if not self.enable_log_probs:
+            return None
+
+        if not self._is_supported():
+            return None
+
+        # Ensure bfloat16 precision for numerical stability
+        if logits_tensor.dtype == ttnn.bfloat8_b:
+            logits_tensor = ttnn.typecast(logits_tensor, ttnn.bfloat16, **self.common_args)
+
+        # Compute global stats — large intermediates allocated and freed here
+        self._compute_global_stats(logits_tensor)
+
+        # Narrow 256 -> 32 topk values and indices
+        topk_local_values, topk_local_indices = ttnn.topk(
+            topk_values,
+            k=DEVICE_TOP_K,
+            dim=-1,
+            sub_core_grids=sub_core_grid_topk,
+        )
+
+        # ttnn.gather requires uint32 indices in TILE_LAYOUT
+        topk_local_indices = ttnn.typecast(topk_local_indices, ttnn.uint32, **self.common_args)
+        topk_local_indices = ttnn.to_layout(topk_local_indices, ttnn.ROW_MAJOR_LAYOUT, **self.common_args)
+        topk_local_indices = ttnn.to_layout(topk_local_indices, ttnn.TILE_LAYOUT, **self.common_args)
+        ttnn.gather(
+            topk_global_indices, dim=-1, index=topk_local_indices, out=self.topk_indices_output, **self.common_args
+        )
+        ttnn.deallocate(topk_local_indices)
+
+        # Ensure topk_values is bfloat16 for consistent computation
+        if topk_local_values.dtype != ttnn.bfloat16:
+            topk_local_values = ttnn.typecast(topk_local_values, ttnn.bfloat16, **self.common_args)
+
+        # Single-pass logprob computation for all gathered top-k tokens
+        self._calculate_topk_log_probs_from_values(topk_local_values)
+        ttnn.deallocate(topk_local_values)
+
+        return LogProbsResult(
+            topk_logprobs=self.topk_logprobs_output,
+            topk_indices=self.topk_indices_output,
+            topk_logprobs_host=None,
+            topk_indices_host=None,
+        )
+
+    # -----------------------------------------------------------------------
+    # Host transfer helpers
+    # -----------------------------------------------------------------------
+
+    def _build_mesh_composer(self):
+        """Build the appropriate mesh composer for transferring tensors from device to host."""
+        if self.cluster_shape[0] > 1 and self.cluster_shape[1] > 1:
+            return ttnn.ConcatMesh2dToTensor(
+                self.mesh_device,
+                dims=(0, 1),
+                mesh_shape=self.cluster_shape,
+            )
+        else:
+            return ttnn.ConcatMeshToTensor(self.mesh_device, dim=0)
+
+    def transfer_logprobs_to_host(
+        self,
+        log_probs_result: LogProbsResult | None,
+        sampled_token_ids: torch.Tensor,
+        num_logprobs_per_user: list[int] | None = None,
+    ) -> list[dict | None]:
+        """Move logprobs from device to host and build per-user response objects.
+
+        tt-metal path only (standalone demos/tests). vLLM does NOT call this.
+
+        Args:
+            log_probs_result: LogProbsResult from calculate_topk_log_probs.
+            sampled_token_ids: Host tensor of sampled token IDs, shape (batch_size,).
+            num_logprobs_per_user: Per-user count of top logprobs to return (0-20).
+                If None, uses self.num_logprobs.
+
+        Returns:
+            List of length batch_size. Each element is None for users with
+            logprobs disabled, otherwise a dict with returned_token and top_logprobs.
+        """
+        if log_probs_result is None:
+            return [None] * self.batch_size
+
+        if num_logprobs_per_user is None:
+            num_logprobs_per_user = self.num_logprobs
+
+        mesh_composer = self._build_mesh_composer()
+
+        topk_logprobs_host = ttnn.to_torch(
+            log_probs_result.topk_logprobs_host
+            if log_probs_result.topk_logprobs_host is not None
+            else log_probs_result.topk_logprobs,
+            mesh_composer=mesh_composer,
+        )
+        topk_indices_host = ttnn.to_torch(
+            log_probs_result.topk_indices_host
+            if log_probs_result.topk_indices_host is not None
+            else log_probs_result.topk_indices,
+            mesh_composer=mesh_composer,
+        )
+        # Remove replicas
+        topk_logprobs_host = topk_logprobs_host[0, 0, ...].float()
+        topk_indices_host = topk_indices_host[0, 0, ...].to(torch.int32)
+
+        results: list[dict | None] = []
+        for user_idx in range(self.batch_size):
+            if not self.logprobs_enabled[user_idx]:
+                results.append(None)
+                continue
+
+            sampled_id = int(sampled_token_ids[user_idx].item())
+            user_logprobs = topk_logprobs_host[user_idx]
+            user_indices = topk_indices_host[user_idx]
+
+            # Extract sampled token logprob by matching its ID in the top-k
+            match_mask = user_indices == sampled_id
+            if match_mask.any():
+                sampled_logprob = float(user_logprobs[match_mask][0].item())
+            else:
+                logger.warning(f"Sampled token {sampled_id} not found in top-k for user {user_idx}")
+                sampled_logprob = float("nan")
+
+            n = num_logprobs_per_user[user_idx] if user_idx < len(num_logprobs_per_user) else 0
+
+            results.append(
+                {
+                    "returned_token": {
+                        "token_idx": sampled_id,
+                        "logprob": sampled_logprob,
+                    },
+                    "top_logprobs": {
+                        "token_indices": user_indices[:n].tolist(),
+                        "logprobs": user_logprobs[:n].tolist(),
+                        "num_logprobs": n,
+                        "logprobs_formatted": False,
+                    },
+                }
+            )
+
+        return results

code/models/common/sampling/tt_penalties.py

@@ -0,0 +1,390 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""
+On-device penalties module with persistent buffers, mirroring TTSampling.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import torch
+
+import ttnn
+from models.common.lightweightmodule import LightweightModule
+from models.common.sampling._utils import compact_debug_list as _compact_debug_list
+from models.common.sampling._utils import is_llama33_70b_model
+from models.common.sampling._utils import log_sampling_debug as _log_sampling_debug
+
+
+@dataclass
+class PenaltyContext:
+    prompt_mask: ttnn.Tensor
+    output_mask: ttnn.Tensor
+    output_counts: ttnn.Tensor
+    output_counts_gathered: ttnn.Tensor
+    presence_penalties: ttnn.Tensor
+    frequency_penalties: ttnn.Tensor
+    repetition_penalties: ttnn.Tensor
+    inverse_repetition_penalties: ttnn.Tensor
+    sub_core_grids: Any | None = None
+
+
+def _tokens_debug_summary(tokens: torch.Tensor) -> dict[str, Any]:
+    valid = tokens >= 0
+    row_lengths = valid.sum(dim=1).tolist()
+    flat_valid = tokens[valid]
+    if flat_valid.numel() == 0:
+        unique_count = 0
+        duplicate_count = 0
+        head = []
+    else:
+        unique_count = int(torch.unique(flat_valid).numel())
+        duplicate_count = int(flat_valid.numel() - unique_count)
+        head = flat_valid[:16].tolist()
+    return {
+        "shape": list(tokens.shape),
+        "dtype": str(tokens.dtype),
+        "valid_tokens": int(valid.sum().item()),
+        "unique_tokens": unique_count,
+        "duplicate_tokens": duplicate_count,
+        "row_lengths": _compact_debug_list(row_lengths),
+        "head_tokens": head,
+    }
+
+
+def apply_penalties(logits: ttnn.Tensor, context: Optional[PenaltyContext]) -> ttnn.Tensor:
+    if context is None:
+        return logits
+
+    op_kwargs = {"sub_core_grids": context.sub_core_grids} if context.sub_core_grids else {}
+    # presence
+    presence_term = ttnn.multiply(
+        ttnn.typecast(context.output_mask, ttnn.bfloat16, **op_kwargs), context.presence_penalties, **op_kwargs
+    )
+    presence_term_bf16 = ttnn.typecast(presence_term, ttnn.bfloat16, **op_kwargs)
+    logits = ttnn.subtract(logits, presence_term_bf16, output_tensor=logits, **op_kwargs)
+    presence_term_bf16.deallocate()
+
+    # frequency
+    output_counts_bf16 = ttnn.typecast(context.output_counts, ttnn.bfloat16, **op_kwargs)
+
+    freq_term = ttnn.multiply(output_counts_bf16, context.frequency_penalties, **op_kwargs)
+
+    freq_term_bf16 = ttnn.typecast(freq_term, ttnn.bfloat16, **op_kwargs)
+    logits = ttnn.subtract(logits, freq_term_bf16, output_tensor=logits, **op_kwargs)
+    freq_term_bf16.deallocate()
+
+    # repetition
+
+    # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
+
+    combined_mask_int32 = ttnn.add(context.prompt_mask, context.output_mask, **op_kwargs)
+    combined_mask = ttnn.typecast(combined_mask_int32, ttnn.bfloat16, **op_kwargs)
+    combined_mask_int32.deallocate()
+    penalties = ttnn.where(combined_mask, context.repetition_penalties, 1.0, **op_kwargs)
+    inverse_penalties = ttnn.where(combined_mask, context.inverse_repetition_penalties, 1.0, **op_kwargs)
+    combined_mask.deallocate()
+
+    # If logits are >0, divide by penalty, otherwise multiply by penalty.
+    logits_bf16 = ttnn.typecast(logits, ttnn.bfloat16, **op_kwargs)
+    logits_gt1 = ttnn.gt(logits_bf16, 0, **op_kwargs)
+    scaling = ttnn.where(logits_gt1, inverse_penalties, penalties, **op_kwargs)
+    logits_gt1.deallocate()
+    penalties.deallocate()
+    inverse_penalties.deallocate()
+    logits = ttnn.multiply(logits, scaling, output_tensor=logits, **op_kwargs)
+    scaling.deallocate()
+
+    return logits
+
+
+class TTPenalties(LightweightModule):
+    """
+    Penalty module with persistent device tensors, similar to TTSampling.
+    """
+
+    def __init__(self, mesh_device, args):
+        super().__init__()
+        self.mesh_device = mesh_device
+        self._sampling_debug_enabled = is_llama33_70b_model(args)
+        self.cluster_shape = mesh_device.shape
+        # Floor at 32 so that ROW_MAJOR [batch, vocab] buffers passed to
+        # ttnn.tilize always have physical_volume divisible by TILE_HW
+        # (32*32 = 1024).  32 * V is 1024-aligned for any 32-aligned V.
+        self.max_batch_size = max(getattr(args, "max_batch_size", 32), 32)
+
+        padded_vocab_size = getattr(args, "padded_vocab_size", None)
+        self.vocab_size = padded_vocab_size if padded_vocab_size is not None else args.vocab_size
+
+        self.sub_core_grids = getattr(args, "sub_core_grids", None)
+        self._op_kwargs = {"sub_core_grids": self.sub_core_grids} if self.sub_core_grids else {}
+
+        # sampling_dp > 1 when multiple mesh rows each sample independently
+        # (e.g. GPT-OSS on [4,8] Galaxy: 4 rows × 32 users = 128 total)
+        self._sampling_dp = getattr(args, "sampling_dp", 1)
+
+        # When rows are used for data parallelism (sampling_dp > 1), vocab
+        # must be sharded along columns; otherwise pick the larger dimension.
+        if self._sampling_dp > 1:
+            num_devices = mesh_device.shape[-1]
+        else:
+            num_devices = max(mesh_device.shape[-1], mesh_device.shape[-2])
+        self.num_devices = num_devices
+        # Total batch across all rows. Host tensors use this size; after
+        # (0, ...) sharding each row gets max_batch_size entries.
+        self._total_batch = self.max_batch_size * self._sampling_dp
+
+        # shard vocab size over larger cluster dim
+        if mesh_device.shape[-1] == self.num_devices:
+            shard_dims = (None, 1)
+            shard_dims_slice = (None, 0)
+        else:
+            shard_dims = (1, None)
+            shard_dims_slice = (0, None)
+
+        # For row-sharded mode (sampling_dp > 1), also shard the batch dimension
+        # across mesh rows so each row gets its own per-user penalty state.
+        if self._sampling_dp > 1:
+            assert (
+                mesh_device.shape[-1] == self.num_devices
+            ), "Row-sharded penalties require vocab sharding along mesh columns"
+            shard_dims = (0, 1)  # batch across rows, vocab across cols
+            shard_dims_gathered = (0, None)  # batch across rows, vocab replicated
+            shard_dims_bf16 = (0, None)  # per-row penalty params
+            per_row_batch = self.max_batch_size  # NOT divided: each row gets max_batch_size
+        else:
+            shard_dims_gathered = (None, None)
+            shard_dims_bf16 = None
+            per_row_batch = self.max_batch_size
+
+        self.per_row_batch_size = per_row_batch
+        self._shard_dims_gathered = shard_dims_gathered
+
+        self.prompt_mask = self._alloc_int_buffer(shard_dims=shard_dims)
+        self.output_mask = self._alloc_int_buffer(shard_dims=shard_dims)
+        self.output_counts_gathered = self._alloc_int_buffer(shard_dims=shard_dims_gathered)
+        self.output_counts = self._alloc_int_buffer(shard_dims=shard_dims)
+        self._shard_dims_mask = shard_dims
+        self.decode_src = self._alloc_int_buffer(
+            host=torch.ones(self._total_batch, 1), shard_dims=shard_dims_gathered, layout=ttnn.ROW_MAJOR_LAYOUT
+        )
+        self.zeros = self._alloc_int_buffer(shard_dims=shard_dims_gathered, layout=ttnn.ROW_MAJOR_LAYOUT)
+        self.presence_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16)
+        self.frequency_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16)
+        self.repetition_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16)
+        self.inverse_repetition_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16)
+
+        vocab_per_dev = self.vocab_size // self.num_devices
+        d = torch.arange(self.num_devices, dtype=torch.int32)
+
+        # [0, 0, 0, vocab_per_dev, 0, 2*vocab_per_dev, ...]
+        start_1d = torch.empty(2 * self.num_devices, dtype=torch.int32)
+        start_1d[0::2] = 0
+        start_1d[1::2] = d * vocab_per_dev
+
+        # [batch, vocab_per_dev, batch, 2*vocab_per_dev, ...]
+        end_1d = torch.empty(2 * self.num_devices, dtype=torch.int32)
+        end_1d[0::2] = per_row_batch  # per-row batch size, exclusive
+        end_1d[1::2] = (d + 1) * vocab_per_dev  # exclusive
+
+        self.slice_start = ttnn.from_torch(
+            start_1d,
+            device=self.mesh_device,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims_slice, mesh_shape=self.cluster_shape),
+        )
+        self.slice_end = ttnn.from_torch(
+            end_1d,
+            device=self.mesh_device,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims_slice, mesh_shape=self.cluster_shape),
+        )
+
+    def _alloc_int_buffer(self, shard_dims, host=None, layout=ttnn.TILE_LAYOUT):
+        if host is None:
+            host = torch.zeros((self._total_batch, self.vocab_size), dtype=torch.int32)
+        return ttnn.from_torch(
+            host,
+            dtype=ttnn.int32,
+            layout=layout,
+            device=self.mesh_device,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape),
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    def _alloc_bf16_buffer(self, shard_dims=None):
+        host = torch.zeros((self._total_batch, 1), dtype=torch.float32)
+        if shard_dims is not None:
+            return ttnn.from_torch(
+                host,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                device=self.mesh_device,
+                mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape),
+            )
+        return ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.mesh_device)
+
+    def _copy_host_to_device(self, dst: ttnn.Tensor, src: torch.Tensor):
+        if self._sampling_dp > 1:
+            # For row-sharded buffers, create a properly sharded host tensor
+            # so copy_host_to_device_tensor writes per-row shards correctly.
+            mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=(0, None), mesh_shape=self.cluster_shape)
+            src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None, mesh_mapper=mapper)
+        else:
+            src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None)
+        ttnn.copy_host_to_device_tensor(src_tt, dst)
+
+    def _copy_int_host_to_device(self, dst: ttnn.Tensor, src: torch.Tensor, shard_dims):
+        mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape)
+        src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None, mesh_mapper=mapper)
+        ttnn.copy_host_to_device_tensor(src_tt, dst)
+
+    def _token_counts_host(self, tokens_2d: torch.Tensor) -> torch.Tensor:
+        valid = (tokens_2d >= 0) & (tokens_2d < self.vocab_size)
+        token_ids = torch.where(valid, tokens_2d, torch.zeros_like(tokens_2d)).to(torch.int64)
+        counts = torch.zeros((self._total_batch, self.vocab_size), dtype=torch.int32)
+        counts.scatter_add_(1, token_ids, valid.to(torch.int32))
+        return counts
+
+    def reset_params(self, presence: List[float], frequency: List[float], repetition: List[float]):
+        presence_tensor = self._pad_params(presence)
+        frequency_tensor = self._pad_params(frequency)
+        repetition_tensor = self._pad_params(repetition)
+        inverse_repetition_tensor = 1 / repetition_tensor
+
+        self._copy_host_to_device(self.presence_penalties, presence_tensor)
+        self._copy_host_to_device(self.frequency_penalties, frequency_tensor)
+        self._copy_host_to_device(self.repetition_penalties, repetition_tensor)
+        self._copy_host_to_device(self.inverse_repetition_penalties, inverse_repetition_tensor)
+
+    def _pad_params(self, values: List[float]) -> torch.Tensor:
+        tensor = torch.tensor(values, dtype=torch.float32)
+        if tensor.numel() < self._total_batch:
+            pad_value = tensor[-1] if tensor.numel() > 0 else torch.tensor(0.0)
+            pad = pad_value.repeat(self._total_batch - tensor.numel())
+            tensor = torch.cat([tensor, pad])
+        elif tensor.numel() > self._total_batch:
+            tensor = tensor[: self._total_batch]
+        return tensor.view(self._total_batch, 1)
+
+    def _pad_batch_to_max(self, tokens_2d: torch.Tensor, pad_value: int) -> torch.Tensor:
+        """Pad/truncate first dim to _total_batch."""
+        if tokens_2d.dim() != 2:
+            raise ValueError(f"Expected 2D tensor [B, S], got {tokens_2d.shape}")
+        B, S = tokens_2d.shape
+        if B < self._total_batch:
+            pad = torch.full((self._total_batch - B, S), pad_value, dtype=tokens_2d.dtype)
+            return torch.cat([tokens_2d, pad], dim=0)
+        if B > self._total_batch:
+            return tokens_2d[: self._total_batch]
+        return tokens_2d
+
+    def reset_prompt_tokens(self, prompt_tokens: torch.Tensor):
+        prompt_tokens_2d = prompt_tokens.reshape(-1, prompt_tokens.shape[-1])
+        prompt_tokens_2d = self._pad_batch_to_max(prompt_tokens_2d, pad_value=-1)
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "TTPenalties reset prompt tokens",
+            tokens=_tokens_debug_summary(prompt_tokens_2d),
+        )
+
+        # Build reset masks on host to avoid device scatter_add races on
+        # duplicate prompt token ids (common in penalty tests/prompts).
+        prompt_counts = self._token_counts_host(prompt_tokens_2d)
+        prompt_mask = (prompt_counts > 0).to(torch.int32)
+        self._copy_int_host_to_device(self.prompt_mask, prompt_mask, self._shard_dims_mask)
+
+    def reset_output_tokens(self, tokens=None):
+        # ALWAYS reset output buffers to zero first (this is the core accuracy fix from issue #35731)
+        # This ensures penalty statistics are cleared between prefill and decode phases
+        self.output_mask = ttnn.mul(self.output_mask, 0, output_tensor=self.output_mask, **self._op_kwargs)
+        self.output_counts = ttnn.mul(self.output_counts, 0, output_tensor=self.output_counts, **self._op_kwargs)
+        self.output_counts_gathered = ttnn.mul(
+            self.output_counts_gathered, 0, output_tensor=self.output_counts_gathered, **self._op_kwargs
+        )
+
+        # THEN optionally repopulate if tokens are provided
+        if tokens is not None:
+            tokens_2d = tokens.reshape(-1, tokens.shape[-1])
+            tokens_2d = self._pad_batch_to_max(tokens_2d, pad_value=-1)
+            _log_sampling_debug(
+                self._sampling_debug_enabled,
+                "TTPenalties reset output tokens",
+                tokens=_tokens_debug_summary(tokens_2d),
+            )
+            output_counts = self._token_counts_host(tokens_2d)
+            output_mask = (output_counts > 0).to(torch.int32)
+            self._copy_int_host_to_device(self.output_counts_gathered, output_counts, self._shard_dims_gathered)
+            self._copy_int_host_to_device(self.output_counts, output_counts, self._shard_dims_mask)
+            self._copy_int_host_to_device(self.output_mask, output_mask, self._shard_dims_mask)
+        else:
+            _log_sampling_debug(self._sampling_debug_enabled, "TTPenalties reset output tokens", tokens=None)
+
+    def update_output_tokens(self, new_tokens):
+        # Reshape decode token to [batch, 1] for scatter_add.
+        # Non-row-sharded: token shape is [1,1,1,batch] → shape[-1]==batch, shape[-2]==1
+        # Row-sharded:     token shape is [1,1,batch,1] → shape[-2]==batch, shape[-1]==1
+        batch = self.per_row_batch_size
+        if (new_tokens.shape[-1] == batch and new_tokens.shape[-2] == 1) or (
+            new_tokens.shape[-2] == batch and new_tokens.shape[-1] == 1
+        ):
+            new_tokens = ttnn.reshape(new_tokens, [batch, 1], **self._op_kwargs)
+            src = self.decode_src
+        else:
+            src = self._alloc_int_buffer(
+                host=torch.ones(self._total_batch, new_tokens.shape[-1]),
+                shard_dims=self._shard_dims_gathered,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+            )
+        self.token_bin_counts_and_mask(
+            new_tokens=new_tokens,
+            counts=self.output_counts_gathered,
+            src=src,
+            counts_sliced=self.output_counts,
+            mask=self.output_mask,
+        )
+
+    def token_bin_counts_and_mask(self, new_tokens, src, counts=None, mask=None, counts_sliced=None):
+        counts_new = ttnn.scatter_add(self.zeros, 1, new_tokens, src, **self._op_kwargs)
+
+        new_tokens.deallocate()
+        # need to use use_low_perf because llama galaxy runs out of L1 otherwise
+        counts_new = ttnn.tilize(
+            counts_new, **self._op_kwargs, use_low_perf=True if self.sub_core_grids is not None else False
+        )
+        if counts:
+            counts = ttnn.add(counts, counts_new, output_tensor=counts, **self._op_kwargs)
+        else:
+            counts = counts_new
+        counts_sliced = ttnn.slice(
+            counts,
+            self.slice_start,
+            self.slice_end,
+            output_tensor=counts_sliced,
+            slice_dim=1,
+            num_devices=self.num_devices,
+            **self._op_kwargs,
+        )
+
+        mask = ttnn.gt(counts_sliced, 0, output_tensor=mask, **self._op_kwargs)
+        return counts, mask
+
+    def apply(self, tt_logits: ttnn.Tensor) -> ttnn.Tensor:
+        if tt_logits is None:
+            return tt_logits
+        context = PenaltyContext(
+            prompt_mask=self.prompt_mask,
+            output_mask=self.output_mask,
+            output_counts=self.output_counts,
+            output_counts_gathered=self.output_counts_gathered,
+            presence_penalties=self.presence_penalties,
+            frequency_penalties=self.frequency_penalties,
+            repetition_penalties=self.repetition_penalties,
+            inverse_repetition_penalties=self.inverse_repetition_penalties,
+            sub_core_grids=self.sub_core_grids,
+        )
+        original_shape = tt_logits.shape
+        reshaped = ttnn.reshape(tt_logits, (-1, original_shape[-1]))
+        apply_penalties(reshaped, context)
+        return ttnn.reshape(reshaped, original_shape)

code/models/common/sampling/tt_sampling.py

@@ -0,0 +1,655 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+import sys
+
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.lightweightmodule import LightweightModule
+from models.common.sampling._utils import compact_debug_list as _compact_debug_list
+from models.common.sampling._utils import is_default_value, is_llama33_70b_model, is_power_of_2
+from models.common.sampling._utils import log_sampling_debug as _log_sampling_debug
+from models.common.sampling._utils import upper_power_of_2
+from models.common.sampling.tt_log_probs import LogProbsCalculator
+
+
+class TTSampling(LightweightModule):
+    """
+    On-device sampling module supporting top-k, top-p, and temperature-based sampling.
+
+    This class implements high-performance on-device sampling that can work across different
+    model implementations by accepting configuration parameters rather than assuming specific
+    args structures.
+
+    Multi-device sampling works by partitioning the vocabulary across devices. Each device
+    computes top-k locally on its vocabulary partition, then all-gather operations combine
+    the results across devices to perform global top-k selection before final sampling.
+
+    Args:
+        mesh_device: The device or MeshDevice for computations
+        tt_ccl: CCL object for distributed operations (supports both line_all_gather and tt_all_gather)
+        vocab_size: Vocabulary size of the model
+        padded_vocab_size: Padded vocabulary size (must be divisible by num devices)
+        max_batch_size: Maximum batch size supported
+        max_top_k: Maximum number of top-k tokens to consider
+        cluster_shape: Shape of the device cluster (rows, cols)
+        sampling_all_gather_axis: Axis to all-gather over in 2D meshes (0=rows, 1=cols, default: 0)
+        sub_core_grids: Sub-core grid configuration for operations
+        sub_core_grid_topk: Sub-core grid configuration specifically for top-k operations
+        start_core: Starting core coordinate for sampling operations
+        num_gather_links: Number of links to use for all-gather operations (optional)
+        sampling_memory_config: Memory configuration for sampling tensors (optional)
+        k, p, temp: Initial sampling parameters (tensors of size max_batch_size)
+
+    Note:
+        Uses persistent buffers when CCL supports line_all_gather (llama3_70b_galaxy),
+        otherwise uses standard all_gather where the CCL API handles memory allocation (tt-transformers).
+    """
+
+    def _is_force_argmax_sampling(self, k, p, temp):
+        """Detect whether all users request deterministic greedy decoding.
+
+        When every user in the batch has k=1 (top-1), p=0.0 or p=1.0 (no top-p filter),
+        and temp=1.0 (no temperature scaling), we can skip the full top-k / top-p /
+        temperature / RNG pipeline and use a single all-gather + argmax instead.
+        This is significantly faster because argmax needs only one all-gather of the
+        full logits tensor vs. three gathers (values, indices, sampled tokens) in the
+        normal path.
+
+        Note: callers may represent greedy rows with p=1.0, while the
+        device argmax-style representation uses p=0.0.
+        The model config must also set allow_force_argmax=True for this to activate.
+
+        Changing this state between decode steps invalidates captured traces, so
+        SamplingGenerator maintains separate trace slots keyed by force_argmax.
+        """
+        return (
+            self._allow_force_argmax_sampling
+            and is_default_value(k, 1)
+            and (is_default_value(p, 1.0) or is_default_value(p, 0.0))
+            and is_default_value(temp, 1.0)
+        )
+
+    def _select_topk_indices_dtype(self, per_device_vocab_size: int, multi_step_reduction: bool):
+        # if vocab is larger than uint16 max, return uint32 for indices
+        if per_device_vocab_size > torch.iinfo(torch.uint16).max:
+            return ttnn.uint32
+
+        # if vocab size is missaligned with tile size and multi-step reduction is used, we need uint32 because of slice op compatibility
+        if multi_step_reduction and (per_device_vocab_size // 2) % ttnn.TILE_SIZE != 0:
+            return ttnn.uint32
+
+        return ttnn.uint16
+
+    @property
+    def force_argmax_sampling(self) -> bool:
+        return self._force_argmax_sampling
+
+    def __init__(
+        self,
+        mesh_device,
+        tt_ccl,
+        args,
+        k=None,
+        p=None,
+        temp=None,
+    ):
+        super().__init__()
+        self.mesh_device = mesh_device
+        self._sampling_debug_enabled = is_llama33_70b_model(args)
+        # Multi-step reduction is supported only on single device
+        self.multi_step_reduction = list(mesh_device.shape) == [1, 1]
+        self.tt_ccl = tt_ccl
+        self._line_all_gather = getattr(self.tt_ccl, "line_all_gather", None)
+        self._line_all_gather_supports_buffer_key = False
+        self._line_all_gather_supports_dtype = False
+        self.pad_to_power_of_2 = getattr(args, "pad_logits_to_power_of_2", False)
+        if callable(self._line_all_gather):
+            try:
+                line_all_gather_sig = inspect.signature(self._line_all_gather)
+                line_all_gather_params = line_all_gather_sig.parameters
+                self._line_all_gather_supports_buffer_key = "buffer_key" in line_all_gather_params or any(
+                    param.kind == inspect.Parameter.VAR_KEYWORD for param in line_all_gather_params.values()
+                )
+                self._line_all_gather_supports_dtype = "dtype" in line_all_gather_params or any(
+                    param.kind == inspect.Parameter.VAR_KEYWORD for param in line_all_gather_params.values()
+                )
+            except (TypeError, ValueError):
+                logger.warning("Unable to inspect line_all_gather signature; assuming no buffer_key or dtype support.")
+
+        padded_vocab_size = getattr(args, "padded_vocab_size", None)
+        self.padded_vocab_size = padded_vocab_size if padded_vocab_size is not None else args.vocab_size
+        self.vocab_size = args.vocab_size
+        # Round up to the next tile boundary (32) — device tensors must be tile-aligned.
+        raw_batch = getattr(args, "max_batch_size", 32)
+        self.max_batch_size = max(32, ((raw_batch + 31) // 32) * 32)
+        self.max_top_k = getattr(args, "max_top_k", 32)
+        self.cluster_shape = args.cluster_shape
+
+        self.sampling_all_gather_axis = getattr(args, "sampling_all_gather_axis", 0)
+        self.sub_core_grids = getattr(args, "sub_core_grids", None)
+        self.sub_core_grid_topk = getattr(args, "sub_core_grid_topk", None)
+        self.start_core = getattr(args, "start_core", ttnn.CoreCoord(0, 0))
+        self._sampling_sub_core_grids = (
+            ttnn.num_cores_to_corerangeset_in_subcoregrids(
+                self.start_core, self.max_batch_size, self.sub_core_grids, row_wise=True
+            )
+            if self.sub_core_grids is not None
+            else None
+        )
+
+        # sampling_dp > 1 when multiple mesh groups each sample users independently
+        # (e.g. GPT-OSS on [4,8]: 4 rows × 32 users; Llama Galaxy on [8,4]: 4 cols × 8 users)
+        self._sampling_dp = getattr(args, "sampling_dp", 1)
+        # Shard params along the non-all-gather axis; replicate along the all-gather axis
+        if self._sampling_dp > 1:
+            if self.sampling_all_gather_axis == 0:
+                self._param_dims = (None, 0)  # shard along cols
+            else:
+                self._param_dims = (0, None)  # shard along rows
+        else:
+            self._param_dims = (None, None)
+
+        if hasattr(args, "model_config") and "GALAXY_NUM_LINKS" in args.model_config:
+            # Calculate num_gather_links based on model config
+            max_num_gather_links = args.model_config["GALAXY_NUM_LINKS"]
+            self.num_gather_links = (
+                args.max_top_k // 32 if args.max_top_k // 32 <= max_num_gather_links else max_num_gather_links
+            )
+        else:
+            self.num_gather_links = 1
+        if hasattr(args, "model_config") and "DECODE_SAMPLING_INPUT_MEMCFG" in args.model_config:
+            self.sampling_memory_config = args.model_config["DECODE_SAMPLING_INPUT_MEMCFG"]
+        else:
+            self.sampling_memory_config = ttnn.DRAM_MEMORY_CONFIG
+
+        # Force argmax sampling
+        if hasattr(args, "model_config") and "SAMPLING_AG_CONFIG" in args.model_config:
+            # The model config may describe the fastest full-size Galaxy path, but
+            # the actual CCL shape is resolved from the runtime mesh below.
+            sampling_ag_config = args.model_config["SAMPLING_AG_CONFIG"]
+            self._allow_force_argmax_sampling = sampling_ag_config["allow_force_argmax"]
+            self.num_argmax_gather_links = sampling_ag_config["num_links"]
+            self.argmax_chunks_per_sync = sampling_ag_config.get("chunks_per_sync", 10)
+            self.argmax_num_workers_per_link = 1
+            self.ag_topology = sampling_ag_config["topology"]
+        else:
+            self._allow_force_argmax_sampling = False
+            self.num_argmax_gather_links = self.num_gather_links
+            self.argmax_chunks_per_sync = 10
+            self.argmax_num_workers_per_link = 1
+            self.ag_topology = ttnn.Topology.Linear
+
+        # Set defaults for sampling parameters if not provided
+        # Default: k=1 (top-1), p=0 (effectively argmax), temp=1 (no temperature scaling)
+        # When p=0, the sampling operation will select the token with highest probability (argmax)
+        total_param_size = self.max_batch_size * self._sampling_dp
+        if k is None:
+            k = torch.ones(total_param_size)
+        if p is None:
+            p = torch.zeros(total_param_size)
+        if temp is None:
+            temp = torch.ones(total_param_size)
+
+        self._force_argmax_sampling = self._is_force_argmax_sampling(k, p, temp)
+
+        # Create sampling parameter tensors on device
+        # When _sampling_dp > 1, dims=(0, None) shards the [128] tensor across 4 rows → [32] per row
+        self.k_tensor = ttnn.from_torch(
+            k,
+            device=self.mesh_device,
+            dtype=ttnn.uint32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=self._param_dims, mesh_shape=self.cluster_shape),
+        )
+        self.p_tensor = ttnn.from_torch(
+            p,
+            device=self.mesh_device,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=self._param_dims, mesh_shape=self.cluster_shape),
+        )
+        self.temp_tensor = ttnn.from_torch(
+            temp,
+            device=self.mesh_device,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=self._param_dims, mesh_shape=self.cluster_shape),
+        )
+
+        # Create device offset indices for global indexing
+        self._create_indices_tensors()
+        # Log-probs tensor to store the log-probs for the batch
+        self.tt_log_probs = None
+        self.log_probs_calculator = LogProbsCalculator(
+            self.mesh_device,
+            self.sub_core_grids,
+            self.tt_ccl,
+            batch_size=self.max_batch_size,
+            use_topk_logprobs=getattr(args, "use_topk_logprobs", False),
+        )
+
+        # Seeds tensor: one RNG slot per user across all rows.
+        # When sampling_dp > 1, shard across rows so each row gets its own slice.
+        # user_ids tensor: core routing only (32 per row, replicated).
+        self.seeds_tt_tensor = ttnn.from_torch(
+            torch.arange(total_param_size).to(torch.uint32),
+            device=self.mesh_device,
+            dtype=ttnn.uint32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=self._param_dims, mesh_shape=self.cluster_shape)
+            if self._sampling_dp > 1
+            else None,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+        self.user_ids_tt_tensor = ttnn.as_tensor(
+            torch.arange(self.max_batch_size).to(torch.uint32),
+            dtype=ttnn.uint32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=self.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    def _create_indices_tensors(self):
+        """Create the indices tensors needed for distributed top-k operations."""
+        # Create indices tensor for device offsets
+        # For multi-step reduction, we use reduce over 2 steps in a single device
+        if self.multi_step_reduction:
+            num_devices_in_mesh = 2
+        else:
+            # If the mesh is effectively 1D, use the non-singleton dimension.
+            # If the mesh is 2D, use the configured gather axis.
+            if 1 in self.cluster_shape:
+                num_devices_in_mesh = max(self.cluster_shape[0], self.cluster_shape[1])
+            else:
+                assert self.sampling_all_gather_axis in (
+                    0,
+                    1,
+                ), f"sampling_all_gather_axis must be 0 or 1 for 2D meshes, got {self.sampling_all_gather_axis}"
+                num_devices_in_mesh = self.cluster_shape[self.sampling_all_gather_axis]
+        indices_device_offsets = torch.ones(
+            1, 1, self.max_batch_size, self.max_top_k * num_devices_in_mesh, dtype=torch.int64
+        )
+        # padded_per_device: tile-aligned width matching actual logit tensors (for indices tensor)
+        padded_per_device = self.padded_vocab_size // num_devices_in_mesh
+
+        for device_id in range(num_devices_in_mesh):
+            indices_device_offsets[:, :, :, device_id * self.max_top_k : (device_id + 1) * self.max_top_k] = (
+                device_id * padded_per_device
+            )
+        self.tt_indices_device_offsets = ttnn.from_torch(
+            indices_device_offsets,
+            device=self.mesh_device,
+            dtype=ttnn.int32,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=(None, None), mesh_shape=self.cluster_shape),
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+        # Create local indices tensor for top-k operations (must match logit width)
+        indices_tensor_torch = torch.zeros(1, 1, self.max_batch_size, padded_per_device, dtype=torch.int32)
+        for i in range(padded_per_device):
+            indices_tensor_torch[:, :, :, i] = i
+
+        # pad to power of 2 if needed
+        if self.pad_to_power_of_2 and not is_power_of_2(indices_tensor_torch.shape[-1]):
+            padded_value = upper_power_of_2(indices_tensor_torch.shape[-1])
+            indices_tensor_torch = torch.nn.functional.pad(
+                indices_tensor_torch,
+                (0, padded_value - indices_tensor_torch.shape[-1]),  # pad only last dim
+                mode="constant",
+                value=-1,  # invalid index to ensure that the padding values are not used
+            )
+
+        indices_dtype = self._select_topk_indices_dtype(padded_per_device, self.multi_step_reduction)
+        self.tt_indices_tensor = ttnn.from_torch(
+            indices_tensor_torch,
+            dtype=indices_dtype,
+            layout=ttnn.Layout.TILE,
+            device=self.mesh_device,
+            mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=(None, None), mesh_shape=self.cluster_shape),
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    def _perform_all_gather(self, tensor, dim, cluster_axis, memory_config, num_links, buffer_key=None, dtype=None):
+        """
+        Flexible all-gather that works across different CCL implementations.
+
+        - If `tt_ccl` exposes `line_all_gather`, prefer it (enables persistent buffer usage on some stacks).
+        - Otherwise fall back to `ttnn.all_gather`.
+        """
+        if callable(self._line_all_gather):
+            # Some implementations accept `buffer_key` (for persistent buffers), others may not.
+            line_all_gather_kwargs = {
+                "dim": dim,
+                "cluster_axis": cluster_axis,
+                "memory_config": memory_config,
+                "num_links": num_links,
+            }
+            if self._line_all_gather_supports_buffer_key and buffer_key is not None:
+                line_all_gather_kwargs["buffer_key"] = buffer_key
+            if self._line_all_gather_supports_dtype and dtype is not None:
+                line_all_gather_kwargs["dtype"] = dtype
+            return self._line_all_gather(tensor, **line_all_gather_kwargs)
+
+        return ttnn.all_gather(
+            tensor,
+            dim=dim,
+            num_links=num_links,
+            memory_config=memory_config,
+            cluster_axis=cluster_axis,
+            topology=ttnn.Topology.Linear,
+        )
+
+    def _get_sampling_cluster_axis(self):
+        if self.mesh_device.get_num_devices() <= 1:
+            return None
+        # 1D submeshes should use the default CCL axis; forcing axis 1 can make
+        # smaller Galaxy DP groups request routes outside the submesh.
+        if 1 in self.cluster_shape:
+            return None
+        return self.sampling_all_gather_axis
+
+    def _get_force_argmax_all_gather_config(self, cluster_axis):
+        num_links = self.num_argmax_gather_links
+        if hasattr(self.tt_ccl, "get_num_links"):
+            # Clamp the tuned config to the links available on the actual submesh.
+            num_links = min(num_links, self.tt_ccl.get_num_links(cluster_axis))
+
+        topology = self.ag_topology
+        # Ring is available for T3K-like 8-device groups; smaller DP groups need
+        # linear routing to avoid wraparound routes such as D0 -> D12.
+        if self.mesh_device.get_num_devices() < 8:
+            topology = ttnn.Topology.Linear
+
+        return max(1, num_links), topology
+
+    def reset_params(
+        self,
+        k,
+        p,
+        temp,
+        enable_log_probs: bool | list[bool] = None,
+        num_logprobs: int | list[int] = None,
+        empty_slots: list[int] | None = None,
+    ):
+        """Update sampling parameters (k, p, temperature, logprobs) dynamically."""
+        self._force_argmax_sampling = self._is_force_argmax_sampling(k, p, temp)
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "TTSampling reset params",
+            force_argmax=self._force_argmax_sampling,
+            empty_slots=_compact_debug_list(empty_slots),
+            top_k=_compact_debug_list(k),
+            top_p=_compact_debug_list(p),
+            temperature=_compact_debug_list(temp),
+            enable_log_probs=_compact_debug_list(enable_log_probs),
+            num_logprobs=_compact_debug_list(num_logprobs),
+            sampling_dp=self._sampling_dp,
+        )
+        if not self._force_argmax_sampling:
+            # When _sampling_dp > 1, create multi-device host tensors so
+            # copy_host_to_device_tensor writes per-row shards correctly.
+            if self._sampling_dp > 1:
+                mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=self._param_dims, mesh_shape=self.cluster_shape)
+            else:
+                mapper = None
+
+            self.k_tensor_new = ttnn.from_torch(
+                torch.tensor(k),
+                device=None,
+                dtype=ttnn.uint32,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                mesh_mapper=mapper,
+            )
+            self.p_tensor_new = ttnn.from_torch(
+                torch.tensor(p),
+                device=None,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                mesh_mapper=mapper,
+            )
+            self.temp_tensor_new = ttnn.from_torch(
+                torch.tensor(temp),
+                device=None,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                mesh_mapper=mapper,
+            )
+
+            ttnn.copy_host_to_device_tensor(self.k_tensor_new, self.k_tensor)
+            ttnn.copy_host_to_device_tensor(self.p_tensor_new, self.p_tensor)
+            ttnn.copy_host_to_device_tensor(self.temp_tensor_new, self.temp_tensor)
+
+        self.log_probs_calculator.set_log_probs_mode(
+            enable_log_probs, num_logprobs=num_logprobs, empty_slots=empty_slots
+        )
+
+    def forward(
+        self,
+        x: ttnn.Tensor,
+        tt_out_tok: ttnn.Tensor = None,
+    ):
+        """
+        Perform on-device sampling on logits tensor.
+        The logits are sharded over the devices in the cluster.
+        We perform local top-k on each device, then all-gather the top-k values and indices across all devices.
+        We then convert the gathered values and indices to the appropriate format, add the device offsets to get the global vocabulary indices,
+        and perform the actual sampling with top-k, top-p, and temperature.
+
+        Args:
+            x: Input logits tensor
+            tt_out_tok: Optional output tensor to write results to
+
+        Returns:
+            Sampled token indices tensor
+        """
+        _log_sampling_debug(
+            self._sampling_debug_enabled,
+            "TTSampling forward",
+            force_argmax=self._force_argmax_sampling,
+            logits_shape=list(x.shape),
+            tt_out_tok_shape=list(tt_out_tok.shape) if tt_out_tok is not None else None,
+            max_top_k=self.max_top_k,
+            multi_step_reduction=self.multi_step_reduction,
+            sampling_dp=self._sampling_dp,
+        )
+        if self._force_argmax_sampling:
+            logger.info("Forcing argmax sampling")
+            # Gather the output across all devices and untilize the tensor (for argmax)
+            num_devices = self.mesh_device.get_num_devices()
+            if num_devices > 1:
+                cluster_axis = self._get_sampling_cluster_axis()
+                num_links, topology = self._get_force_argmax_all_gather_config(cluster_axis)
+                logger.debug(
+                    f"Force argmax sampling all-gather: cluster_axis={cluster_axis}, "
+                    f"num_links={num_links}, topology={topology}"
+                )
+                x = ttnn.experimental.all_gather_async(
+                    x,
+                    persistent_output_buffer=None,
+                    dim=3,
+                    multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis),
+                    num_links=num_links,
+                    memory_config=x.memory_config(),
+                    cluster_axis=cluster_axis,
+                    topology=topology,
+                    barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis),
+                    chunks_per_sync=self.argmax_chunks_per_sync,
+                    num_workers_per_link=self.argmax_num_workers_per_link,
+                    num_buffers_per_channel=2,
+                )
+            x_untilized = ttnn.untilize(x, use_multicore=True)
+            tt_out_tok = ttnn.argmax(
+                x_untilized,
+                dim=-1,
+                output_tensor=tt_out_tok,
+                keepdim=False,
+            )
+            # Argmax path: logprobs not supported (force-argmax is disabled
+            # when logprobs are enabled via format_sampling_params guard).
+            self.tt_log_probs = None
+            return tt_out_tok, self.tt_log_probs
+
+        # Convert to bfloat16 for top-k operations (typecast is no-op if already bfloat16)
+        x_bf16 = ttnn.typecast(x, dtype=ttnn.bfloat16, sub_core_grids=self.sub_core_grids)
+
+        if self.multi_step_reduction:
+            x_bf16_list = ttnn.split(x_bf16, x_bf16.shape[-1] // 2, dim=3)
+            indices_tensor_list = ttnn.split(self.tt_indices_tensor, self.tt_indices_tensor.shape[-1] // 2, dim=3)
+            topk_values_list = []
+            topk_indices_list = []
+
+            for i in range(len(x_bf16_list)):
+                topk_values, topk_indices = ttnn.topk(
+                    x_bf16_list[i],
+                    k=self.max_top_k,
+                    dim=-1,
+                    sub_core_grids=self.sub_core_grid_topk,
+                    indices_tensor=indices_tensor_list[i],
+                )
+                topk_values_list.append(topk_values)
+                topk_indices_list.append(topk_indices)
+                x_bf16_list[i].deallocate()
+                indices_tensor_list[i].deallocate()
+
+            topk_values_gathered_bf16_interleaved = ttnn.concat(topk_values_list, dim=3)
+            topk_indices_gathered = ttnn.concat(topk_indices_list, dim=3)
+
+            for i in range(len(topk_indices_list)):
+                ttnn.deallocate(topk_values_list[i])
+                ttnn.deallocate(topk_indices_list[i])
+
+        else:
+            # apply padding to the input tensor if needed
+            # if number is not power of 2, pad to upper power of 2
+            # pad only last dimension with float::min value to upper_power_of_2
+            # This is necessary to use full optimization in the topk operation.
+            if self.pad_to_power_of_2 and not is_power_of_2(x_bf16.shape[-1]):
+                padded_value = upper_power_of_2(x_bf16.shape[-1])
+                x_bf16 = ttnn.pad(
+                    x_bf16,
+                    [(0, 0), (0, 0), (0, 0), (0, padded_value - x_bf16.shape[-1])],
+                    value=-sys.float_info.max,
+                    sub_core_grids=self.sub_core_grids,
+                )
+            # Perform local top-k on each device
+            topk_values, topk_indices = ttnn.topk(
+                x_bf16,
+                k=self.max_top_k,
+                dim=-1,
+                sub_core_grids=self.sub_core_grid_topk,
+                indices_tensor=self.tt_indices_tensor,
+            )
+
+            # For 1D meshes use `cluster_axis=None`. For 2D meshes, use the configured gather axis.
+            sampling_cluster_axis = self._get_sampling_cluster_axis()
+
+            # Gather top-k values across all devices
+            topk_values_gathered = self._perform_all_gather(
+                topk_values,
+                dim=3,
+                cluster_axis=sampling_cluster_axis,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                num_links=self.num_gather_links,
+                buffer_key="SAMPLING_VALUES",
+            )
+
+            ttnn.deallocate(topk_values)
+
+            # Convert gathered values to appropriate format
+            if self.sampling_memory_config != ttnn.DRAM_MEMORY_CONFIG:
+                topk_values_gathered_bf16 = ttnn.to_memory_config(
+                    topk_values_gathered,
+                    memory_config=self.sampling_memory_config,
+                    dtype=ttnn.bfloat16,
+                )
+                topk_values_gathered_bf16_interleaved = ttnn.to_memory_config(
+                    topk_values_gathered_bf16, memory_config=ttnn.DRAM_MEMORY_CONFIG
+                )
+                ttnn.deallocate(topk_values_gathered_bf16)
+            else:
+                topk_values_gathered_bf16_interleaved = topk_values_gathered
+
+            # Gather top-k indices across all devices
+            topk_indices_gathered = self._perform_all_gather(
+                topk_indices,
+                dim=3,
+                cluster_axis=sampling_cluster_axis,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                num_links=self.num_gather_links,
+                buffer_key="SAMPLING_INDICES",
+                dtype=ttnn.uint16,
+            )
+            ttnn.deallocate(topk_indices)
+
+        # Convert indices to appropriate data types
+
+        topk_indices_gathered_int32 = ttnn.typecast(
+            topk_indices_gathered, dtype=ttnn.int32, sub_core_grids=self.sub_core_grids
+        )
+
+        if self.sampling_memory_config != ttnn.DRAM_MEMORY_CONFIG:
+            topk_indices_gathered_int32_sharded = ttnn.to_memory_config(
+                topk_indices_gathered_int32, self.sampling_memory_config
+            )
+            ttnn.deallocate(topk_indices_gathered_int32)
+        else:
+            topk_indices_gathered_int32_sharded = topk_indices_gathered_int32
+
+        # Add device offsets to get global vocabulary indices
+        topk_global_indices = ttnn.add(
+            self.tt_indices_device_offsets,
+            topk_indices_gathered_int32_sharded,
+            dtype=ttnn.uint32,
+            memory_config=self.sampling_memory_config,
+        )
+
+        ttnn.deallocate(topk_indices_gathered_int32_sharded)
+
+        topk_global_indices_interleaved = ttnn.to_memory_config(topk_global_indices, ttnn.DRAM_MEMORY_CONFIG)
+
+        # Untilize indices for sampling operation
+        topk_global_indices_interleaved_untilised = ttnn.untilize(
+            topk_global_indices_interleaved, use_multicore=True, sub_core_grids=self.sub_core_grids
+        )
+        ttnn.manual_seed(
+            seeds=self.seeds_tt_tensor,
+            user_ids=self.user_ids_tt_tensor,
+            sub_core_grids=self._sampling_sub_core_grids,
+        )
+        # Perform the actual sampling with top-k, top-p, and temperature
+        tt_out_tok = ttnn.sampling(
+            topk_values_gathered_bf16_interleaved,
+            topk_global_indices_interleaved_untilised,
+            k=self.k_tensor,
+            p=self.p_tensor,
+            temp=self.temp_tensor,
+            sub_core_grids=self._sampling_sub_core_grids,
+            output_tensor=tt_out_tok,
+        )
+
+        # Compute logprobs if enabled
+        if self.log_probs_calculator.enable_log_probs and self.log_probs_calculator._use_topk_logprobs:
+            # New path: top-K logprobs for gpt-oss-120b
+            self.tt_log_probs = self.log_probs_calculator.calculate_topk_log_probs(
+                logits_tensor=x,
+                topk_values=topk_values_gathered_bf16_interleaved,
+                topk_global_indices=topk_global_indices_interleaved,
+                sub_core_grid_topk=self.sub_core_grid_topk,
+            )
+        elif self.log_probs_calculator.enable_log_probs:
+            # Old path: single sampled-token logprob
+            self.tt_log_probs = self.log_probs_calculator.calculate_log_probs(x, tt_out_tok)
+        else:
+            self.tt_log_probs = None
+
+        ttnn.deallocate(topk_values_gathered_bf16_interleaved)
+        ttnn.deallocate(topk_global_indices_interleaved)
+        ttnn.deallocate(topk_global_indices_interleaved_untilised)
+
+        return tt_out_tok, self.tt_log_probs

code/models/common/tensor_utils.py

@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tensor utility functions for TTTv2 modules.
+"""
+
+import json
+import re
+
+import torch
+
+import ttnn
+
+# Standard tile size - hardware constant
+TILE_SIZE = ttnn.TILE_SIZE  # 32
+
+
+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/tests/conftest.py

@@ -0,0 +1,265 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import contextlib
+import fcntl
+import os
+import time
+
+import pytest
+
+import ttnn
+
+# ==============================================================================
+# Device Lock - Coordinates exclusive access to TT devices across processes
+# ==============================================================================
+
+_TT_DEVICE_LOCK_PATH = os.environ.get("TT_DEVICE_LOCK_PATH", "/tmp/tt_device.lock")
+_TT_DEVICE_LOCK_TIMEOUT = float(os.environ.get("TT_DEVICE_LOCK_TIMEOUT", "60"))  # 1 min default
+
+
+class DeviceLockTimeout(Exception):
+    """Raised when acquiring the device lock times out."""
+
+
+# todo)) the UMD already provides a lock mechanism -- use it instead of this?
+@contextlib.contextmanager
+def tt_device_lock(lock_path: str = _TT_DEVICE_LOCK_PATH, timeout: float = _TT_DEVICE_LOCK_TIMEOUT):
+    """
+    Context manager for exclusive access to TT devices.
+
+    Uses flock for cross-process coordination. Blocks until lock is acquired
+    or timeout is reached.
+
+    Usage:
+        with tt_device_lock():
+            mesh = ttnn.open_mesh_device(...)
+            # ... do work ...
+            ttnn.close_mesh_device(mesh)
+
+    Debug stuck locks with: lsof /tmp/tt_device.lock
+
+    Environment variables:
+        TT_DEVICE_LOCK_PATH: Override lock file path (default: /tmp/tt_device.lock)
+        TT_DEVICE_LOCK_TIMEOUT: Override timeout in seconds (default: 300)
+    """
+    lock_dir = os.path.dirname(lock_path)
+    if lock_dir and not os.path.exists(lock_dir):
+        os.makedirs(lock_dir, exist_ok=True)
+
+    lock_file = open(lock_path, "a+")  # open the file in append mode to avoid truncation race condition among processes
+    start_time = time.monotonic()
+    lock_acquired = False
+
+    try:
+        # Poll for lock with timeout
+        logged_waiting = False
+        while True:
+            try:
+                fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                lock_acquired = True
+                break
+            except BlockingIOError:
+                pass  # Lock held by another process
+
+            if not logged_waiting:
+                print(f"[tt_device_lock] Waiting for device lock (held by another process)...")
+                print(f"[tt_device_lock] Debug with: lsof {lock_path}")
+                logged_waiting = True
+
+            if time.monotonic() - start_time >= timeout:
+                lock_file.close()
+                raise DeviceLockTimeout(
+                    f"Timed out after {timeout}s waiting for device lock. " f"Check: lsof {lock_path}"
+                )
+
+            time.sleep(1)  # sleep for 1 second to avoid busy-waiting
+
+        if logged_waiting:
+            print(f"[tt_device_lock] Lock acquired after {time.monotonic() - start_time:.1f}s")
+
+        # Write PID for debugging
+        lock_file.truncate(0)  # clear the file
+        lock_file.write(f"{os.getpid()}\n")
+        lock_file.flush()
+
+        yield
+
+    finally:
+        if lock_acquired:
+            fcntl.flock(lock_file, fcntl.LOCK_UN)
+        lock_file.close()
+
+
+def pytest_collection_modifyitems(config, items):
+    """Deselect tests where ttnn_mesh_device fixture doesn't match mesh_shape param.
+
+    This enables tests to use cross-product parametrization (all meshes × all cases)
+    while only running the valid combinations, without noisy skip messages.
+    """
+    selected = []
+    deselected = []
+
+    for item in items:
+        if not hasattr(item, "callspec"):
+            selected.append(item)
+            continue
+
+        params = item.callspec.params
+        fixture_mesh = params.get("ttnn_mesh_device")
+        required_mesh = params.get("mesh_shape")
+
+        # Keep test if no mesh_shape param or if meshes match
+        if required_mesh is None or fixture_mesh == required_mesh:
+            selected.append(item)
+        else:
+            deselected.append(item)
+
+    items[:] = selected
+    if deselected:
+        config.hook.pytest_deselected(items=deselected)
+
+
+@pytest.fixture(scope="module")
+def ttnn_mesh_device(request):
+    """Create and yield a mesh device for a given mesh shape, cleanup on teardown."""
+    if not hasattr(request, "param"):
+        pytest.skip(f"{__file__}: mesh_device fixture called without parametrization")
+
+    if ttnn.device.is_blackhole():
+        pytest.skip(f"{__file__}: Blackhole device is not supported for this test yet")
+
+    # request.param is either a Sequence of ints or a dict with fabric_config and etc.
+    params = getattr(request, "param", tuple())
+    if isinstance(params, tuple):
+        mesh_shape = params
+        updated_params = dict()
+    else:
+        try:
+            updated_params = params.copy()
+            mesh_shape = updated_params.pop("mesh_shape")
+        except Exception as e:
+            pytest.skip(f"{__file__}: mesh_shape is required: {e}")
+
+    # Pre-check: if no devices at all, skip without invoking C++ open.
+    # Some environments can throw here (e.g. transient driver/UMD issues); treat as "device unavailable".
+    try:
+        num_pcie = ttnn.get_num_pcie_devices()
+    except Exception as e:
+        pytest.skip(f"{__file__}: Unable to query TT devices on this system: {e}")
+
+    if isinstance(num_pcie, int) and num_pcie == 0:
+        pytest.skip(f"{__file__}: No TT devices detected on this system")
+
+    # Pre-check: skip shapes that cannot fit into the SystemMesh to avoid native exceptions
+    sys_desc = ttnn._ttnn.multi_device.SystemMeshDescriptor()  # type: ignore[attr-defined]
+    sys_shape = tuple(sys_desc.shape())
+    req_shape = tuple(mesh_shape)
+    allowed = _allowed_req_shapes_for_system(sys_shape)
+    if req_shape not in allowed:
+        pytest.skip(
+            f"{__file__}: Requested mesh {req_shape} unsupported on system {sys_shape}. "
+            f"Allowed for this system: {allowed}"
+        )
+
+    parent_shape = _pick_parent_shape_for_submesh(sys_shape, req_shape)
+
+    # config fabric config
+    fabric_config = updated_params.pop("fabric_config", None)
+    if parent_shape == (1, 1):
+        # Single device does not need fabric config.
+        pass
+    else:
+        # Provide default fabric config for the mesh we actually open (full system mesh).
+        num_devices = parent_shape[0] * parent_shape[1]
+        if fabric_config is None:
+            if num_devices >= 8:
+                fabric_config = ttnn.FabricConfig.FABRIC_1D_RING
+            else:
+                fabric_config = ttnn.FabricConfig.FABRIC_1D
+        # set all other input arguments to default values by top-level conftest.py
+        ttnn.set_fabric_config(
+            fabric_config, ttnn.FabricReliabilityMode.STRICT_INIT, None, ttnn.FabricTensixConfig.DISABLED
+        )
+
+    # config dispatch core to default values by conftest.py
+    updated_params["dispatch_core_config"] = ttnn.DispatchCoreConfig(type=None, axis=None, fabric_tensix_config=None)
+
+    # If a test requests a submesh of a larger system mesh (e.g. request 2x4 on a 8x4 system),
+    # fabric cannot be initialized on only the subset of devices. In that case, open the full
+    # system mesh first, then return the "first" submesh. We intentionally rely on the default
+    # offset behavior here (i.e. no explicit offset selection).
+    parent_device = None
+    submesh_device = None
+
+    # Acquire exclusive lock to prevent concurrent device access across processes
+    with tt_device_lock():
+        try:
+            if req_shape != parent_shape:
+                parent_device = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(parent_shape), **updated_params)
+                submesh_device = parent_device.create_submesh(ttnn.MeshShape(req_shape))
+                yield submesh_device
+            else:
+                parent_device = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(parent_shape), **updated_params)
+                yield parent_device
+        except Exception as e:
+            pytest.skip(f"{__file__}: Mesh device unavailable or unsupported for this configuration: {e}")
+        finally:
+            if submesh_device is not None:
+                ttnn.close_mesh_device(submesh_device)
+            if parent_device is not None:
+                ttnn.close_mesh_device(parent_device)
+            if fabric_config:
+                ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED)
+            del parent_device
+
+
+def _allowed_req_shapes_for_system(sys_shape: tuple[int, int]) -> set[tuple[int, int]]:
+    # todo)) Different cluster has potentially different physical interconnects (in terms of number of links, topology, etc.).
+    #        Thus, a tuple of ints may not be enough to fingerprint the parent/system mesh device. We need to use a more sophisticated fingerprinting mechanism so we can base the allowed list of (sub)mesh shapes on the parent/system mesh device.
+    # [INFO] The most robust way to identify the underlying system is to use ttnn.cluster.get_cluster_type(), which returns a ClusterType enum that precisely identifies your hardware configuration. cluster.cpp:16-37
+
+    _CANDIDATE_REQ_SHAPES = {
+        (1, 1): ((1, 1),),
+        (1, 2): ((1, 2), (1, 1)),
+        (2, 4): ((2, 4), (1, 8), (1, 4), (1, 2), (1, 1)),
+        (8, 4): ((8, 4), (4, 8), (1, 8), (1, 4), (1, 2), (1, 1)),
+        # [INFO] add more system shapes here
+    }
+
+    allowed: set[tuple[int, int]] = set()
+
+    if sys_shape in _CANDIDATE_REQ_SHAPES:
+        for mesh_shape in _CANDIDATE_REQ_SHAPES[sys_shape]:
+            allowed.add(mesh_shape)
+
+    return allowed
+
+
+def _pick_parent_shape_for_submesh(system_shape: tuple[int, int], requested_shape: tuple[int, int]) -> tuple[int, int]:
+    # For multi-device workloads we always open the full system mesh (fabric cannot be launched on a subset),
+    # but we may choose the *orientation* of the full mesh such that the requested submesh fits with the
+    # default offset (i.e. "first submesh").
+    if requested_shape == (1, 1):
+        return (1, 1)
+
+    # If the request uses all devices, treat it as a "full-mesh view" shape and open the parent mesh in that view.
+    # This enables shapes like (1,32) on a system whose SystemMeshDescriptor reports (8,4).
+    system_num_devices = system_shape[0] * system_shape[1]
+    requested_num_devices = requested_shape[0] * requested_shape[1]
+    if requested_num_devices == system_num_devices:
+        return requested_shape
+
+    if requested_shape[0] <= system_shape[0] and requested_shape[1] <= system_shape[1]:
+        return system_shape
+
+    rotated = (system_shape[1], system_shape[0])
+    if requested_shape[0] <= rotated[0] and requested_shape[1] <= rotated[1]:
+        return rotated
+
+    # No orientation can fit this request without an explicit offset / mapping.
+    pytest.skip(
+        f"{__file__}: Requested submesh {requested_shape} does not fit within system mesh {system_shape} "
+        f"(or its rotated view {rotated}) with default offset."
+    )

code/models/common/tests/requirements.txt

@@ -0,0 +1 @@
+pytest-cov

code/models/common/tests/setup.cfg

@@ -0,0 +1,12 @@
+[coverage:run]
+branch = True
+source = models/common/modules
+
+[coverage:report]
+fail_under = 80
+show_missing = True
+exclude_lines =
+    pragma: no cover
+    if TYPE_CHECKING:
+    raise NotImplementedError
+    @abstractmethod

code/models/common/tests/test_auto_compose.py

@@ -0,0 +1,416 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for automatic composition of multi-device sharded tensors using TensorTopology.
+
+This test module validates that the auto-composition logic correctly infers
+MeshToTensor composers from a sharded ttnn.Tensor's topology metadata.
+
+It validates both host-sharded and device-sharded cases.
+"""
+
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+import ttnn
+from models.common.auto_compose import _infer_mesh_composer_from_topology, to_torch_auto_compose
+
+# ======================================================================================
+# Test Parameters (for device-dependent tests)
+# ======================================================================================
+
+
+_DEVICE_TEST_MARKS = {
+    "ttnn_mesh_device": pytest.mark.parametrize(
+        "ttnn_mesh_device",
+        [
+            (1, 1),  # single device # [INFO] apply auto_compose on single device would incur error in c++ code
+            (1, 2),  # 1D mesh, 2 devices
+            (1, 4),  # 1D mesh, 4 devices
+            (1, 8),  # 1D mesh, 8 devices
+            (2, 4),  # 2D mesh, 8 devices
+            (4, 8),  # 2D mesh, 32 devices
+            (8, 4),  # 2D mesh, 32 devices
+        ],
+        ids=[
+            "1x1",
+            "1x2",
+            "1x4",
+            "1x8",
+            "2x4",
+            "4x8",
+            "8x4",
+        ],
+        indirect=True,
+    ),
+    "layout_dtype": pytest.mark.parametrize(
+        "layout,dtype",
+        [
+            (ttnn.ROW_MAJOR_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat8_b),  # bfloat8_b only works with TILE_LAYOUT
+            (ttnn.TILE_LAYOUT, ttnn.bfloat4_b),  # bfloat4_b only works with TILE_LAYOUT
+        ],
+        ids=["row_major_bf16", "tile_bf16", "tile_bf8b", "tile_bf4b"],
+    ),
+}
+
+
+# ======================================================================================
+# Helper Functions
+# ======================================================================================
+
+
+def _make_known_pattern(num_chunks: int) -> torch.Tensor:
+    """
+    Produces shape [num_chunks, 1, 3, 1] with per-chunk distinct values.
+    Chunk i contains [i*1, i*2, i*3].
+    """
+    rows = []
+    for i in range(num_chunks):
+        rows.append(torch.tensor([[[i * 1.0], [i * 2.0], [i * 3.0]]]).transpose(0, 1))  # [1,3,1]
+    data = torch.stack(rows, dim=0)  # [num_chunks,1,3,1]
+    return data.to(torch.bfloat16)
+
+
+def _make_arange_dtype(
+    shape: tuple[int, ...], dtype: torch.dtype = torch.bfloat16, min_value: float = 0, max_value: float = 100
+) -> torch.Tensor:
+    """Create a deterministic tensor with arange data and specified dtype."""
+    numel = 1
+    for s in shape:
+        numel *= s
+    # Generate values from min_value to max_value with step of 1
+    values = torch.arange(min_value, max_value + 1, dtype=dtype)
+    # Randomly sample indices (with replacement) to fill the tensor
+    indices = torch.randint(0, len(values), size=(numel,))
+    data = values[indices].reshape(shape)
+    return data
+
+
+def _pos_dim(dim: int, rank: int) -> int:
+    """Convert possibly-negative dim to positive index for given rank."""
+    return dim % rank
+
+
+def _get_hw_shard_unit() -> int:
+    """
+    Hardware-related shard unit threshold (default 32).
+    Override via env var TT_TEST_SHARD_UNIT for future hardware.
+    """
+    try:
+        return int(os.environ.get("TT_TEST_SHARD_UNIT", "32"))
+    except Exception:
+        return 32
+
+
+def _build_and_compose_sharded(
+    torch_in: torch.Tensor,
+    device: ttnn.MeshDevice | None,
+    layout,
+    ttnn_mesh_device: ttnn.MeshDevice,
+    shard_dim: int,
+    dtype: torch.dtype = ttnn.bfloat16,
+) -> tuple[ttnn.Tensor, torch.Tensor, torch.Tensor]:
+    """Build sharded tensor and compose it back to torch."""
+    tt_sharded = ttnn.from_torch(
+        torch_in,
+        device=device,
+        dtype=dtype,
+        layout=layout,
+        mesh_mapper=ttnn.ShardTensorToMesh(ttnn_mesh_device, dim=shard_dim),
+    )
+    torch_auto = to_torch_auto_compose(tt_sharded, device=ttnn_mesh_device if device is None else None)
+    torch_ref = ttnn.to_torch(tt_sharded, mesh_composer=ttnn.ConcatMeshToTensor(ttnn_mesh_device, dim=shard_dim))
+    return tt_sharded, torch_auto, torch_ref
+
+
+# ======================================================================================
+# Device-Dependent Tests (require mesh device fixture)
+# ======================================================================================
+
+
+@pytest.mark.parametrize("storage", ["host", "device"])  # where the sharded tensor lives
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_sharded_1d_basic(ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, storage: str) -> None:
+    """Basic 1D sharding auto-composition for both host and device storage."""
+    num_devices = ttnn_mesh_device.get_num_devices()
+
+    # Input tensor of shape [num_devices, 1, 3, 1]
+    torch_in = _make_known_pattern(num_devices)
+
+    # Build sharded tensor on host or device along dim=0 and compose back
+    device = None if storage == "host" else ttnn_mesh_device
+    _, torch_auto, torch_ref = _build_and_compose_sharded(
+        torch_in, device, layout, ttnn_mesh_device, shard_dim=0, dtype=dtype
+    )
+
+    if dtype == ttnn.bfloat16:
+        assert torch.equal(torch_ref, torch_in), "Explicit composer mismatch"
+        assert torch.equal(torch_auto, torch_in), "Auto-composer mismatch"
+    else:
+        # For quantized dtypes, compare auto vs explicit composed results
+        assert torch.equal(torch_auto, torch_ref), "Auto vs explicit composer mismatch for quantized dtype"
+
+
+@pytest.mark.parametrize("storage", ["host", "device"])  # where the replicated tensor lives
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_replicate_1d_basic(ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, storage: str) -> None:
+    """Replicated 1D distribution should compose to identity for host and device storage."""
+    # Any shape works; replication does not change global shape
+    # ttnn.from_torch perform naive quantization to lower dtypes -- work on existing exponent and mantissa values
+    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
+    min_value, max_value = -7, 7
+    torch_in = _make_arange_dtype((2, 3, 4, 5), dtype=torch.float32, min_value=min_value, max_value=max_value)
+
+    device = None if storage == "host" else ttnn_mesh_device
+    tt_replicated = ttnn.from_torch(
+        torch_in,
+        device=device,
+        dtype=dtype,
+        layout=layout,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(ttnn_mesh_device),
+    )
+
+    # Auto-composition should detect full replication and yield identity
+    if device is None:
+        torch_auto = to_torch_auto_compose(tt_replicated, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_replicated)
+
+    assert torch.equal(torch_auto, torch_in)
+
+
+# --------------------------------------------------------------------------------------
+# Shard various tensor dims on 1D meshes
+# --------------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dim", [0, 1, 2, -1])
+@pytest.mark.parametrize("storage", ["host", "device"])  # where the sharded tensor lives
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_sharded_various_dims(ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, dim: int, storage: str) -> None:
+    num_devices = ttnn_mesh_device.get_num_devices()
+
+    rank = 4
+    axis = _pos_dim(dim, rank)
+    shape = [2, 3, 4, 1]
+    shape[axis] = num_devices
+    torch_in = _make_arange_dtype(tuple(shape))
+
+    device = None if storage == "host" else ttnn_mesh_device
+    _, torch_auto, torch_ref = _build_and_compose_sharded(torch_in, device, layout, ttnn_mesh_device, dim, dtype=dtype)
+
+    if dtype == ttnn.bfloat16:
+        assert torch.equal(torch_ref, torch_in)
+        assert torch.equal(torch_auto, torch_in)
+    else:
+        assert torch.equal(torch_auto, torch_ref)
+
+
+# --------------------------------------------------------------------------------------
+# Coverage for 2D mesh sharding: shard-shard and replicate-shard
+# --------------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dims_pair", [(0, 1), (0, -1), (1, -1)])
+@pytest.mark.parametrize("storage", ["host", "device"])  # where the sharded tensor lives
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_sharded_2d_basic(
+    ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, dims_pair: tuple[int, int], storage: str
+) -> None:
+    mesh_shape = tuple(ttnn_mesh_device.shape)
+    if len(mesh_shape) != 2 and torch.prod(torch.tensor(mesh_shape)).item() <= 1:
+        pytest.skip("Requires a 2D mesh with both dims > 1")
+
+    rank = 4
+    d0 = _pos_dim(dims_pair[0], rank)
+    d1 = _pos_dim(dims_pair[1], rank)
+    assert d0 != d1, "Shard dims for 2D sharding must be distinct"
+
+    shape = [2, 3, 4, 5]
+    shape[d0] = mesh_shape[0]
+    shape[d1] = mesh_shape[1]
+    torch_in = _make_arange_dtype(tuple(shape))
+
+    mapper = ttnn.ShardTensor2dMesh(ttnn_mesh_device, mesh_shape=mesh_shape, dims=(dims_pair[0], dims_pair[1]))
+    device = None if storage == "host" else ttnn_mesh_device
+    tt_sharded = ttnn.from_torch(torch_in, device=device, dtype=dtype, layout=layout, mesh_mapper=mapper)
+
+    if device is None:
+        torch_auto = to_torch_auto_compose(tt_sharded, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_sharded)
+    composer = ttnn.ConcatMesh2dToTensor(ttnn_mesh_device, mesh_shape=mesh_shape, dims=(dims_pair[0], dims_pair[1]))
+    torch_ref = ttnn.to_torch(tt_sharded, mesh_composer=composer)
+
+    if dtype == ttnn.bfloat16:
+        assert torch.equal(torch_ref, torch_in)
+        assert torch.equal(torch_auto, torch_in)
+    else:
+        assert torch.equal(torch_auto, torch_ref)
+
+
+@pytest.mark.parametrize(
+    "dims_pair",
+    [
+        (None, -1),  # replicate along mesh dim 0, shard along last tensor dim
+        (1, None),  # shard along tensor dim 1 on mesh dim 0, replicate mesh dim 1
+    ],
+)
+@pytest.mark.parametrize("storage", ["host", "device"])  # host vs device sharded tensor
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_sharded_2d_with_replicate(
+    ttnn_mesh_device: ttnn.MeshDevice,
+    layout,
+    dtype,
+    dims_pair: tuple[object, object],
+    storage: str,
+) -> None:
+    # None indicates replicate axis
+    replicate_axis = [i for i, d in enumerate(dims_pair) if d is None][0]
+    mesh_shape = tuple(ttnn_mesh_device.shape)
+    if len(mesh_shape) != 2 and torch.prod(torch.tensor(mesh_shape)).item() <= 1:
+        pytest.skip("Requires a 2D mesh with at least one dim > 1 to observe replication")
+
+    rank = 4
+    # Determine which tensor axis is sharded (the non-None entry)
+    shard_dim = [d for d in dims_pair if d is not None][0]
+    shard_axis = _pos_dim(shard_dim, rank)
+    shape = [2, 3, 4, 5]
+    # Set size along sharded axis rounded up to a multiple of the other mesh dim
+    other_mesh_dim = mesh_shape[1 - replicate_axis]
+    shape[shard_axis] = ((shape[shard_axis] + other_mesh_dim - 1) // other_mesh_dim) * other_mesh_dim
+
+    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
+    torch_in = _make_arange_dtype(tuple(shape), dtype=torch.float32, min_value=-7, max_value=7)
+
+    mapper = ttnn.ShardTensor2dMesh(ttnn_mesh_device, mesh_shape=mesh_shape, dims=dims_pair)  # type: ignore[arg-type]
+    device = None if storage == "host" else ttnn_mesh_device
+    tt_sharded = ttnn.from_torch(torch_in, device=device, dtype=dtype, layout=layout, mesh_mapper=mapper)
+
+    if device is None:
+        torch_auto = to_torch_auto_compose(tt_sharded, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_sharded)
+
+    assert torch.equal(torch_auto, torch_in)
+
+
+# --------------------------------------------------------------------------------------
+# Tensor shape categories around hardware threshold (e.g., 32)
+# --------------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("category", ["lt", "eq", "gt"])  # per-shard length relative to threshold
+@pytest.mark.parametrize("storage", ["host", "device"])  # where the sharded tensor lives
+@_DEVICE_TEST_MARKS["ttnn_mesh_device"]
+@_DEVICE_TEST_MARKS["layout_dtype"]
+def test_sharded_shape_thresholds(
+    ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, category: str, storage: str
+) -> None:
+    num_devices = ttnn_mesh_device.get_num_devices()
+
+    unit = _get_hw_shard_unit()
+    if category == "lt":
+        per_shard = max(1, unit - 1)
+    elif category == "eq":
+        per_shard = unit
+    else:
+        per_shard = unit + 1
+
+    shard_dim = -1  # test last dimension as sharded axis (rank=4)
+    rank = 4
+    axis = _pos_dim(shard_dim, rank)
+    # Global size across sharded dim = per_shard_len * num_devices
+    shape = [2, 3, 4, 5]
+    shape[axis] = per_shard * num_devices
+    torch_in = _make_arange_dtype(tuple(shape))
+
+    device = None if storage == "host" else ttnn_mesh_device
+    _, torch_auto, torch_ref = _build_and_compose_sharded(
+        torch_in, device, layout, ttnn_mesh_device, shard_dim, dtype=dtype
+    )
+
+    if dtype == ttnn.bfloat16:
+        assert torch.equal(torch_ref, torch_in)
+        assert torch.equal(torch_auto, torch_in)
+    else:
+        assert torch.equal(torch_auto, torch_ref)
+
+
+# --------------------------------------------------------------------------------------
+# Test coverage for auto_compose
+# --------------------------------------------------------------------------------------
+
+
+def test_to_torch_auto_compose_exception_handler():
+    """Test the exception handler in to_torch_auto_compose (lines 38-40)."""
+    mock_tensor = MagicMock(spec=ttnn.Tensor)
+    mock_topology = MagicMock()
+    mock_tensor.tensor_topology.return_value = mock_topology
+    mock_topology.placements.return_value = [ttnn.PlacementShard(0)]
+    mock_topology.distribution_shape.return_value = [2]
+
+    mock_device = MagicMock(spec=ttnn.MeshDevice)
+    mock_tensor.device.return_value = mock_device
+
+    with patch("ttnn.create_mesh_composer", return_value="fake_composer"):
+        with patch("ttnn.to_torch", side_effect=RuntimeError("Mock failure")):
+            with pytest.raises(RuntimeError, match="Mock failure"):
+                to_torch_auto_compose(mock_tensor)
+
+
+def test_to_torch_auto_compose_no_device_error():
+    """Test RuntimeError when tensor is on host and no device is provided/available (lines 102-104)."""
+    mock_tensor = MagicMock(spec=ttnn.Tensor)
+    mock_topology = MagicMock()
+    mock_tensor.tensor_topology.return_value = mock_topology
+    mock_topology.placements.return_value = [ttnn.PlacementShard(0)]
+    mock_topology.distribution_shape.return_value = [2]
+
+    # Tensor on host
+    mock_tensor.device.return_value = None
+
+    with patch("ttnn.GetDefaultDevice", return_value=None):
+        with pytest.raises(RuntimeError, match="Tensor is on host and no mesh_device provided"):
+            to_torch_auto_compose(mock_tensor)
+
+
+def test_infer_composer_1d_sharded_mock():
+    """
+    Use mocking to hit the 1D sharded paths (lines 113, 125-131)
+    if real 1D meshes are hard to come by.
+    """
+    mock_tensor = MagicMock(spec=ttnn.Tensor)
+    mock_topology = MagicMock()
+    mock_tensor.tensor_topology.return_value = mock_topology
+
+    # Case 1: 1D Sharded
+    mock_topology.placements.return_value = [ttnn.PlacementShard(0)]
+    mock_topology.distribution_shape.return_value = [2]
+
+    mock_device = MagicMock(spec=ttnn.MeshDevice)
+    mock_device.shape.dims.return_value = 1
+    mock_tensor.device.return_value = mock_device
+
+    with patch("ttnn.create_mesh_composer") as mock_create:
+        mock_create.return_value = "fake_composer"
+        composer = _infer_mesh_composer_from_topology(mock_tensor)
+        assert composer == "fake_composer"
+        mock_create.assert_called_once()
+
+    # Case 2: 1D Replicated
+    mock_topology.placements.return_value = [ttnn.PlacementReplicate()]
+    mock_topology.distribution_shape.return_value = [2]
+    composer = _infer_mesh_composer_from_topology(mock_tensor)
+    assert composer is None

code/models/common/tests/test_device_lock.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for tt_device_lock cross-process coordination."""
+
+import multiprocessing
+import os
+import time
+
+import pytest
+
+from models.common.tests.conftest import DeviceLockTimeout, tt_device_lock
+
+
+class TestDeviceLock:
+    """Tests for tt_device_lock context manager."""
+
+    def test_lock_acquires_and_releases(self, tmp_path):
+        """Basic acquire/release works."""
+        lock_path = str(tmp_path / "test.lock")
+
+        with tt_device_lock(lock_path=lock_path, timeout=5):
+            assert os.path.exists(lock_path)
+            with open(lock_path) as f:
+                content = f.read().strip()
+            assert content == str(os.getpid())
+
+    def test_lock_blocks_concurrent_access(self, tmp_path):
+        """Second process waits while first holds lock."""
+        lock_path = str(tmp_path / "test.lock")
+        result_file = str(tmp_path / "results.txt")
+
+        def worker(worker_id, hold_time):
+            """Acquire lock, record timestamp, hold for hold_time, release."""
+            with tt_device_lock(lock_path=lock_path, timeout=30):
+                with open(result_file, "a") as f:
+                    f.write(f"{worker_id}:acquired:{time.time()}\n")
+                time.sleep(hold_time)
+                with open(result_file, "a") as f:
+                    f.write(f"{worker_id}:released:{time.time()}\n")
+
+        # Start worker 1, let it acquire lock
+        p1 = multiprocessing.Process(target=worker, args=(1, 2))
+        p1.start()
+        time.sleep(0.3)  # Give p1 time to acquire
+
+        # Start worker 2, should block
+        p2 = multiprocessing.Process(target=worker, args=(2, 0.1))
+        p2.start()
+
+        p1.join()
+        p2.join()
+
+        # Parse results
+        with open(result_file) as f:
+            lines = f.read().strip().split("\n")
+
+        events = []
+        for line in lines:
+            parts = line.split(":")
+            events.append((int(parts[0]), parts[1], float(parts[2])))
+
+        # Worker 1 should acquire before worker 2
+        next(e for e in events if e[0] == 1 and e[1] == "acquired")
+        w1_release = next(e for e in events if e[0] == 1 and e[1] == "released")
+        w2_acquire = next(e for e in events if e[0] == 2 and e[1] == "acquired")
+
+        # Worker 2 should only acquire AFTER worker 1 releases
+        assert (
+            w2_acquire[2] >= w1_release[2]
+        ), f"Worker 2 acquired at {w2_acquire[2]} but worker 1 released at {w1_release[2]}"
+
+    def test_lock_timeout_raises(self, tmp_path):
+        """Timeout raises DeviceLockTimeout."""
+        lock_path = str(tmp_path / "test.lock")
+
+        def holder():
+            """Hold lock for a long time."""
+            with tt_device_lock(lock_path=lock_path, timeout=60):
+                time.sleep(10)
+
+        # Start holder
+        p = multiprocessing.Process(target=holder)
+        p.start()
+        time.sleep(0.3)  # Let it acquire
+
+        # Try to acquire with short timeout - should raise
+        try:
+            with pytest.raises(DeviceLockTimeout):
+                with tt_device_lock(lock_path=lock_path, timeout=1):
+                    pass
+        finally:
+            p.terminate()
+            p.join()
+
+    def test_lock_file_created_if_missing(self, tmp_path):
+        """Lock file is created if it doesn't exist."""
+        lock_path = str(tmp_path / "subdir" / "test.lock")
+        assert not os.path.exists(lock_path)
+
+        with tt_device_lock(lock_path=lock_path, timeout=5):
+            assert os.path.exists(lock_path)
+
+    def test_sequential_locks_work(self, tmp_path):
+        """Multiple sequential lock/unlock cycles work."""
+        lock_path = str(tmp_path / "test.lock")
+
+        for i in range(3):
+            with tt_device_lock(lock_path=lock_path, timeout=5):
+                with open(lock_path) as f:
+                    assert f.read().strip() == str(os.getpid())

code/models/common/tests/test_distribute_as.py

@@ -0,0 +1,240 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for distributing a torch.Tensor using the topology of an existing TTNN tensor.
+
+Follows the same style as test_auto_compose.py, but in the opposite
+direction: given a reference distributed tensor, verify that
+`from_torch_dist_as(torch_tensor, ref_tensor)` produces a TTNN tensor
+with the same topology and data that composes back to the original
+torch input.
+"""
+
+import pytest
+import torch
+
+import ttnn
+from models.common.auto_compose import extract_tensor_topology_info, to_torch_auto_compose
+from models.common.distribute_as import from_torch_dist_as
+
+# ======================================================================================
+# Test Parameters
+# ======================================================================================
+
+
+pytestmark = [
+    pytest.mark.parametrize(
+        "ttnn_mesh_device",
+        [
+            (1, 1),  # single device
+            (1, 2),  # 1D mesh, 2 devices
+            (1, 8),  # 1D mesh, 8 devices
+            (2, 4),  # 2D mesh, 8 devices
+        ],
+        ids=[
+            "1x1",
+            "1x2",
+            "1x8",
+            "2x4",
+        ],
+        indirect=True,
+    ),
+    pytest.mark.parametrize(
+        "layout,dtype",
+        [
+            (ttnn.ROW_MAJOR_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat8_b),  # bfloat8_b only works with TILE_LAYOUT
+            (ttnn.TILE_LAYOUT, ttnn.bfloat4_b),  # bfloat4_b only works with TILE_LAYOUT
+        ],
+        ids=["row_major_bf16", "tile_bf16", "tile_bf8b", "tile_bf4b"],
+    ),
+    pytest.mark.parametrize(
+        "memory_config",
+        [ttnn.DRAM_MEMORY_CONFIG, ttnn.L1_MEMORY_CONFIG],
+        ids=["dram", "l1"],
+    ),
+]
+
+
+# ======================================================================================
+# Helpers
+# ======================================================================================
+
+
+def _make_arange_dtype(
+    shape: tuple[int, ...], dtype: torch.dtype = torch.bfloat16, min_value: float = 0, max_value: float = 100
+) -> torch.Tensor:
+    """Create a deterministic tensor with arange data and specified dtype."""
+    numel = 1
+    for s in shape:
+        numel *= s
+    # Generate values from min_value to max_value with step of 1
+    values = torch.arange(min_value, max_value + 1, dtype=dtype)
+    # Randomly sample indices (with replacement) to fill the tensor
+    indices = torch.randint(0, len(values), size=(numel,))
+    data = values[indices].reshape(shape)
+    return data
+
+
+def _pos_dim(dim: int, rank: int) -> int:
+    return dim % rank
+
+
+def _topology_signature(tt: ttnn.Tensor) -> tuple[tuple[str, tuple[int, ...]], tuple[int, ...]]:
+    placements, dist_shape = extract_tensor_topology_info(tt)
+    sig_p = []
+    for p in placements:
+        if isinstance(p, ttnn.PlacementShard):
+            sig_p.append(("shard", (p.dim,)))
+        else:
+            assert isinstance(p, ttnn.PlacementReplicate)
+            sig_p.append(("replicate", ()))
+    return tuple(sig_p), tuple(dist_shape)
+
+
+# ======================================================================================
+# Tests
+# ======================================================================================
+
+
+@pytest.mark.parametrize("storage", ["host", "device"])
+def test_sharded_1d(
+    ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, memory_config: ttnn.MemoryConfig, storage: str
+) -> None:
+    is_host_ref = storage == "host"
+    num_devices = ttnn_mesh_device.get_num_devices()
+
+    # Reference topology: host-sharded along dim 0
+    shape = (num_devices, 1, 3, 1)
+    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
+    max_value = 6  # allows for range of value to grow to 7 for new_input
+    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)
+    ref_tt = ttnn.from_torch(
+        ref_input,
+        device=None if is_host_ref else ttnn_mesh_device,
+        dtype=dtype,
+        layout=layout,
+        memory_config=memory_config,
+        mesh_mapper=ttnn.ShardTensorToMesh(ttnn_mesh_device, dim=0),
+    )
+
+    # Distribute a different tensor using the same topology
+    new_input = ref_input + 1  # change values to avoid accidental equality
+    if is_host_ref:
+        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
+    else:
+        tt_as = from_torch_dist_as(new_input, ref_tt)
+
+    # Verify topology matches reference
+    expected_topology = _topology_signature(ref_tt)
+    assert _topology_signature(tt_as) == expected_topology
+    assert tt_as.dtype == dtype
+    assert tt_as.memory_config() == memory_config
+
+    # Verify round-trip via auto-composition matches the source
+    if is_host_ref:
+        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_as)
+    assert torch.equal(torch_auto, new_input)
+
+
+@pytest.mark.parametrize("storage", ["host", "device"])
+@pytest.mark.parametrize("dim", [0, 1, 2, -1])
+def test_sharded_various_dims(
+    ttnn_mesh_device: ttnn.MeshDevice,
+    layout,
+    dtype,
+    memory_config: ttnn.MemoryConfig,
+    dim: int,
+    storage: str,
+) -> None:
+    is_host_ref = storage == "host"
+    num_devices = ttnn_mesh_device.get_num_devices()
+
+    rank = 4
+    axis = _pos_dim(dim, rank)
+    shape = [2, 3, 4, 1]
+    shape[axis] = num_devices
+    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
+    max_value = 6  # allows for range of value to grow to 7 for new_input
+    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)
+
+    ref_tt = ttnn.from_torch(
+        ref_input,
+        device=None if is_host_ref else ttnn_mesh_device,
+        dtype=dtype,
+        layout=layout,
+        memory_config=memory_config,
+        mesh_mapper=ttnn.ShardTensorToMesh(ttnn_mesh_device, dim=dim),
+    )
+
+    new_input = ref_input + 1  # change values to avoid accidental equality
+    if is_host_ref:
+        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
+    else:
+        tt_as = from_torch_dist_as(new_input, ref_tt)
+
+    assert _topology_signature(tt_as) == _topology_signature(ref_tt)
+    assert tt_as.dtype == dtype
+    assert tt_as.memory_config() == memory_config
+    if is_host_ref:
+        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_as)
+    assert torch.equal(torch_auto, new_input)
+
+
+@pytest.mark.parametrize("storage", ["host", "device"])
+@pytest.mark.parametrize("dims_pair", [(0, 1), (0, -1), (1, -1)])
+def test_sharded_2d(
+    ttnn_mesh_device: ttnn.MeshDevice,
+    layout,
+    dtype,
+    memory_config: ttnn.MemoryConfig,
+    dims_pair: tuple[int, int],
+    storage: str,
+) -> None:
+    is_host_ref = storage == "host"
+    mesh_shape = tuple(ttnn_mesh_device.shape)
+    if len(mesh_shape) != 2 and torch.prod(torch.tensor(mesh_shape)).item() <= 1:
+        pytest.skip("Requires a 2D mesh with both dims > 1")
+
+    rank = 4
+    d0 = _pos_dim(dims_pair[0], rank)
+    d1 = _pos_dim(dims_pair[1], rank)
+    assert d0 != d1
+
+    shape = [2, 3, 4, 5]
+    shape[d0] = mesh_shape[0]
+    shape[d1] = mesh_shape[1]
+    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
+    max_value = 6  # allows for range of value to grow to 7 for new_input
+    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)
+
+    mapper = ttnn.ShardTensor2dMesh(ttnn_mesh_device, mesh_shape=mesh_shape, dims=(dims_pair[0], dims_pair[1]))
+    ref_tt = ttnn.from_torch(
+        ref_input,
+        device=None if is_host_ref else ttnn_mesh_device,
+        dtype=dtype,
+        layout=layout,
+        memory_config=memory_config,
+        mesh_mapper=mapper,
+    )
+
+    new_input = ref_input + 1  # change values to avoid accidental equality
+    if is_host_ref:
+        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
+    else:
+        tt_as = from_torch_dist_as(new_input, ref_tt)
+
+    assert _topology_signature(tt_as) == _topology_signature(ref_tt)
+    assert tt_as.dtype == dtype
+    assert tt_as.memory_config() == memory_config
+    if is_host_ref:
+        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
+    else:
+        torch_auto = to_torch_auto_compose(tt_as)
+    assert torch.equal(torch_auto, new_input)

code/models/common/tests/test_executor_parity.py

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests for eager == traced executor output parity.
+
+The real contract is "same inputs, same outputs, regardless of executor".
+These tests verify that EagerLLMExecutor and TracedLLMExecutor produce
+identical outputs for the same model and inputs.
+"""
+
+import pytest
+import torch
+
+import ttnn
+from models.common.models.executor import EagerLLMExecutor, TensorSpec, TracedLLMExecutor
+from models.common.models.module_input_validation import ConfigMismatch
+
+
+class TestTensorSpec:
+    """Unit tests for TensorSpec dataclass."""
+
+    def test_tensor_spec_creation(self):
+        """Test TensorSpec can be created with expected fields."""
+        spec = TensorSpec(
+            shape=(32, 1, 128256),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            memory_config=None,
+        )
+        assert spec.shape == (32, 1, 128256)
+        assert spec.dtype == ttnn.bfloat16
+        assert spec.layout == ttnn.TILE_LAYOUT
+        assert spec.memory_config is None
+
+    def test_tensor_spec_from_tensor_requires_device(self):
+        """Test TensorSpec.from_tensor requires an allocated device tensor."""
+        # This test documents expected behavior - from_tensor needs a real device tensor
+        # Can't test without device, but we can test the method signature exists
+        assert hasattr(TensorSpec, "from_tensor")
+        assert callable(TensorSpec.from_tensor)
+
+
+class TestConfigMismatch:
+    """Unit tests for ConfigMismatch dataclass."""
+
+    def test_config_mismatch_creation(self):
+        """Test ConfigMismatch can be created with expected fields."""
+        mismatch = ConfigMismatch(
+            module_name="layer[0].attention",
+            expected_memcfg=ttnn.DRAM_MEMORY_CONFIG,
+            actual_memcfg=ttnn.L1_MEMORY_CONFIG,
+        )
+        assert mismatch.module_name == "layer[0].attention"
+        assert mismatch.expected_memcfg == ttnn.DRAM_MEMORY_CONFIG
+        assert mismatch.actual_memcfg == ttnn.L1_MEMORY_CONFIG
+
+
+class TestBoundaryAssertions:
+    """Tests for boundary assertions on public methods."""
+
+    @pytest.fixture
+    def mock_model(self):
+        """Create a minimal mock model for testing assertions."""
+
+        class MockModel:
+            vocab_size = 128256
+            n_layers = 32
+            num_devices = 1
+            sampling = None
+            model_args = None
+
+            def set_kv_cache(self, kv_cache):
+                pass
+
+        return MockModel()
+
+    @pytest.fixture
+    def mock_mesh_device(self):
+        """Return None as placeholder - real device tests need hardware."""
+        return None
+
+    def test_prefill_forward_rejects_1d_tokens(self, mock_model):
+        """Test that prefill_forward rejects 1D token tensor."""
+        # Can't fully test without device, but can test assertion message format
+        tokens_1d = torch.tensor([1, 2, 3])
+        assert tokens_1d.dim() == 1
+
+        # The assertion should fail with a clear message
+        with pytest.raises(AssertionError, match=r"tokens must be \[batch_size, seq_len\], got 1D"):
+            # We can't call the actual method without a device, but we can test the assertion logic
+            assert tokens_1d.dim() == 2, f"tokens must be [batch_size, seq_len], got {tokens_1d.dim()}D"
+
+    def test_prefill_forward_accepts_2d_tokens(self):
+        """Test that 2D tokens pass the assertion."""
+        tokens_2d = torch.tensor([[1, 2, 3]])
+        assert tokens_2d.dim() == 2  # Should pass
+
+    def test_decode_forward_rejects_2d_tokens(self):
+        """Test that decode_forward rejects 2D token tensor."""
+        tokens_2d = torch.tensor([[1, 2, 3]])
+        assert tokens_2d.dim() == 2
+
+        with pytest.raises(AssertionError, match=r"tokens must be \[batch_size\], got 2D"):
+            assert tokens_2d.dim() == 1, f"tokens must be [batch_size], got {tokens_2d.dim()}D"
+
+    def test_decode_forward_accepts_1d_tokens(self):
+        """Test that 1D tokens pass the decode assertion."""
+        tokens_1d = torch.tensor([1, 2, 3])
+        assert tokens_1d.dim() == 1  # Should pass
+
+    def test_page_table_must_be_2d(self):
+        """Test that page_table must be 2D when provided."""
+        page_table_1d = torch.tensor([0, 1, 2])
+        assert page_table_1d.dim() == 1
+
+        with pytest.raises(AssertionError, match=r"page_table must be \[batch_size, max_blocks\], got 1D"):
+            assert page_table_1d.dim() == 2, f"page_table must be [batch_size, max_blocks], got {page_table_1d.dim()}D"
+
+    def test_start_pos_must_be_1d(self):
+        """Test that start_pos must be 1D."""
+        start_pos_2d = torch.tensor([[0, 1, 2]])
+        assert start_pos_2d.dim() == 2
+
+        with pytest.raises(AssertionError, match=r"start_pos must be \[batch_size\], got 2D"):
+            assert start_pos_2d.dim() == 1, f"start_pos must be [batch_size], got {start_pos_2d.dim()}D"
+
+
+class TestOutputSpecCapture:
+    """Tests for output spec capture during compile."""
+
+    def test_eager_executor_has_output_spec_attributes(self):
+        """Test that EagerLLMExecutor has output spec attributes."""
+
+        class MockModel:
+            vocab_size = 128256
+            n_layers = 32
+            num_devices = 1
+            sampling = None
+            model_args = None
+
+            def set_kv_cache(self, kv_cache):
+                pass
+
+        # Can't fully instantiate without device, but can check class structure
+        assert hasattr(EagerLLMExecutor, "__init__")
+
+    def test_traced_executor_has_output_spec_attributes(self):
+        """Test that TracedLLMExecutor has output spec attributes."""
+        assert hasattr(TracedLLMExecutor, "__init__")
+
+
+class TestValidateModuleConfigs:
+    """Tests for the _validate_module_configs context manager."""
+
+    def test_validate_module_configs_method_exists(self):
+        """Test that _validate_module_configs method exists on EagerLLMExecutor."""
+        assert hasattr(EagerLLMExecutor, "_validate_module_configs")
+
+    def test_compile_accepts_validate_configs_flag(self):
+        """Test that compile() accepts validate_configs parameter."""
+        import inspect
+
+        sig = inspect.signature(EagerLLMExecutor.compile)
+        params = sig.parameters
+        assert "validate_configs" in params
+        assert params["validate_configs"].default is False

code/models/common/tests/test_executor_trace_prefill.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from models.common.models import executor as executor_module
+
+
+def test_easy_trace_prefill_replay_copies_only_mutable_inputs(monkeypatch):
+    engine = executor_module.TracedLLMExecutor.__new__(executor_module.TracedLLMExecutor)
+    engine.mesh_device = "mesh"
+    engine.trace_id_prefill = {128: 7}
+    engine.trace_inputs_prefill = {
+        128: ("device_tokens", "device_cos", "device_sin", "device_page_table", None),
+    }
+    engine.trace_output_prefill = {128: "trace_output"}
+
+    monkeypatch.setattr(
+        engine,
+        "_prepare_prefill_trace_inputs_host",
+        lambda tokens, page_table, last_token_idx: (
+            "host_tokens",
+            "device_cos_slice",
+            "device_sin_slice",
+            "host_page_table",
+            None,
+        ),
+    )
+
+    copy_calls = []
+    monkeypatch.setattr(
+        executor_module,
+        "copy_host_to_device",
+        lambda host_tensors, device_tensors: copy_calls.append((host_tensors, device_tensors)),
+    )
+
+    execute_calls = []
+    monkeypatch.setattr(
+        executor_module.ttnn,
+        "execute_trace",
+        lambda mesh_device, trace_id, cq_id, blocking: execute_calls.append((mesh_device, trace_id, cq_id, blocking)),
+    )
+
+    result = engine._easy_trace_prefill(
+        tokens="tokens",
+        page_table="page_table",
+        user_id=0,
+        last_token_idx=127,
+        prefill_seq_len=128,
+    )
+
+    assert result == "trace_output"
+    assert copy_calls == [
+        (
+            ("host_tokens", "host_page_table", None),
+            ("device_tokens", "device_page_table", None),
+        )
+    ]
+    assert execute_calls == [("mesh", 7, 0, False)]

code/models/common/tests/test_lazy_weight.py

@@ -0,0 +1,988 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for LazyWeight module.
+"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+import ttnn
+from models.common.modules.lazy_weight import (
+    LazyWeight,
+    _auto_pad_for_sharded_tiles,
+    _from_torch_and_dump,
+    resolve_lazy_weight,
+)
+
+# ============================================================================
+# Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def mock_tensor():
+    """Create a mock torch tensor for testing."""
+    return torch.randn(1, 32, 128)
+
+
+@pytest.fixture
+def mock_tensor_2d():
+    """Create a 2D mock torch tensor for testing."""
+    return torch.randn(64, 128)
+
+
+@pytest.fixture
+def mock_mesh_device():
+    """Create a mock MeshDevice for unit tests."""
+    device = MagicMock(spec=ttnn.MeshDevice)
+    device.get_num_devices.return_value = 8
+    device.id.return_value = "test_device_0"
+    return device
+
+
+@pytest.fixture
+def mock_memory_config():
+    """Create a mock MemoryConfig."""
+    return ttnn.DRAM_MEMORY_CONFIG
+
+
+@pytest.fixture
+def mock_mesh_mapper_config():
+    """Create a mock MeshMapperConfig for sharding."""
+    return ttnn.MeshMapperConfig(
+        placements=[ttnn.PlacementShard(-1)],
+        mesh_shape_override=ttnn.MeshShape([8]),
+    )
+
+
+# ============================================================================
+# Test LazyWeight Construction
+# ============================================================================
+
+
+class TestLazyWeightConstruction:
+    """Tests for LazyWeight dataclass construction and __post_init__."""
+
+    def test_construction_with_tensor(self, mock_tensor):
+        """Test basic construction with a tensor."""
+        lw = LazyWeight(source=mock_tensor)
+        assert lw.source is mock_tensor
+        assert lw.dtype == ttnn.bfloat16
+        assert lw.pad_value == 0.0
+        assert lw.device is None
+        assert lw.mesh_mapper_config is None
+        assert lw.memory_config is None
+        assert lw.layout is None
+        assert lw._value is None
+
+    def test_construction_with_all_parameters(
+        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
+    ):
+        """Test construction with all parameters specified."""
+        cache_dir = Path("/tmp/cache")
+        lw = LazyWeight(
+            source=mock_tensor,
+            cache_dir_weight_name=(cache_dir, "test_weight"),
+            pad_value=1.0,
+            dtype=ttnn.bfloat4_b,
+            device=mock_mesh_device,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        assert lw.source is mock_tensor
+        assert lw.cache_dir_weight_name == (cache_dir, "test_weight")
+        assert lw.pad_value == 1.0
+        assert lw.dtype == ttnn.bfloat4_b
+        assert lw.device is mock_mesh_device
+        assert lw.mesh_mapper_config is mock_mesh_mapper_config
+        assert lw.memory_config is mock_memory_config
+        assert lw.layout == ttnn.TILE_LAYOUT
+
+    def test_post_init_validates_shape(self):
+        """Test that __post_init__ validates tensor shape."""
+        mock_source = MagicMock()
+        mock_source.shape = None
+
+        with pytest.raises(AssertionError, match="source must have a shape"):
+            LazyWeight(source=mock_source)
+
+    def test_post_init_validates_empty_shape(self):
+        """Test that __post_init__ rejects empty shape."""
+        mock_source = MagicMock()
+        mock_source.shape = ()
+
+        with pytest.raises(AssertionError, match="source must have a shape"):
+            LazyWeight(source=mock_source)
+
+
+# ============================================================================
+# Test is_resolved Method
+# ============================================================================
+
+
+class TestIsResolved:
+    """Tests for LazyWeight.is_resolved() method."""
+
+    def test_not_resolved_missing_device(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
+        """Test is_resolved returns False when device is missing."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        assert not lw.is_resolved()
+
+    def test_not_resolved_missing_layout(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test is_resolved returns False when layout is missing."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+        )
+        assert not lw.is_resolved()
+
+    def test_not_resolved_missing_memory_config(self, mock_tensor, mock_mesh_device):
+        """Test is_resolved returns False when memory_config is missing."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        assert not lw.is_resolved()
+
+    def test_is_resolved_with_all_required_fields(
+        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
+    ):
+        """Test is_resolved returns True when all required fields are set."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        assert lw.is_resolved()
+
+    def test_is_resolved_with_none_mesh_mapper_config(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test is_resolved returns False when mesh_mapper_config is None (replicated case requires it to be explicitly set)."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=None,  # None means not resolved per is_resolved check
+        )
+        # mesh_mapper_config is in required_fields so None means not resolved
+        assert not lw.is_resolved()
+
+
+# ============================================================================
+# Test _get_fingerprint Method
+# ============================================================================
+
+
+class TestGetFingerprint:
+    """Tests for LazyWeight._get_fingerprint() method."""
+
+    def test_fingerprint_includes_source_shape(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint includes source shape."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "srcshape_1_32_128" in fingerprint
+
+    def test_fingerprint_includes_non_default_memory_config(self, mock_tensor, mock_mesh_device):
+        """Test that fingerprint includes memory_config hash when non-default."""
+        # Use L1_MEMORY_CONFIG which is different from DRAM_MEMORY_CONFIG
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=ttnn.L1_MEMORY_CONFIG,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "memcfg_" in fingerprint
+
+    def test_fingerprint_includes_dtype(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint includes dtype."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            dtype=ttnn.bfloat16,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "dtype_BFLOAT16" in fingerprint
+
+    def test_fingerprint_includes_layout(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint includes layout."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "layout_TILE" in fingerprint
+
+    def test_fingerprint_includes_pad_value_when_non_default(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint includes pad_value when non-default."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            pad_value=1.5,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "pad_1.5" in fingerprint
+
+    def test_fingerprint_excludes_pad_value_when_default(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint excludes pad_value when default (0.0)."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            pad_value=0.0,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "pad_" not in fingerprint
+
+    def test_fingerprint_includes_device_id(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint includes device id."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "device_" in fingerprint
+
+    def test_fingerprint_device_without_id_method(self, mock_tensor, mock_memory_config):
+        """Test fingerprint uses 'single' when device has no id() method."""
+        device = MagicMock()
+        del device.id  # Remove id attribute
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        fingerprint = lw._get_fingerprint()
+        assert "device_single" in fingerprint
+
+    def test_fingerprint_changes_with_different_dtype(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint changes with different dtype."""
+        lw1 = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            dtype=ttnn.bfloat16,
+        )
+        lw2 = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            dtype=ttnn.bfloat4_b,
+        )
+        assert lw1._get_fingerprint() != lw2._get_fingerprint()
+
+
+# ============================================================================
+# Test _get_mesh_mapper_fingerprint Method
+# ============================================================================
+
+
+class TestGetMeshMapperFingerprint:
+    """Tests for LazyWeight._get_mesh_mapper_fingerprint() method."""
+
+    def test_fingerprint_replicated_when_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that fingerprint is 'replicated' when mesh_mapper_config is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=None,
+        )
+        fingerprint = lw._get_mesh_mapper_fingerprint()
+        assert fingerprint == "replicated"
+
+    def test_fingerprint_with_mesh_mapper_config(
+        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
+    ):
+        """Test that fingerprint includes mapper hash when config is set."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=mock_mesh_mapper_config,
+        )
+        fingerprint = lw._get_mesh_mapper_fingerprint()
+        assert fingerprint.startswith("mapper_")
+        # Hash should be 12 characters
+        assert len(fingerprint) == len("mapper_") + 12
+
+    def test_different_configs_produce_different_fingerprints(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that different mesh_mapper_configs produce different fingerprints."""
+        config1 = ttnn.MeshMapperConfig(
+            placements=[ttnn.PlacementShard(-1)],
+            mesh_shape_override=ttnn.MeshShape([8]),
+        )
+        config2 = ttnn.MeshMapperConfig(
+            placements=[ttnn.PlacementShard(-2)],
+            mesh_shape_override=ttnn.MeshShape([8]),
+        )
+        lw1 = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=config1,
+        )
+        lw2 = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=config2,
+        )
+        assert lw1._get_mesh_mapper_fingerprint() != lw2._get_mesh_mapper_fingerprint()
+
+
+# ============================================================================
+# Test _get_cache_fill_path Method
+# ============================================================================
+
+
+class TestGetCacheFillPath:
+    """Tests for LazyWeight._get_cache_fill_path() method."""
+
+    def test_returns_none_when_cache_dir_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that _get_cache_fill_path returns None when cache_dir is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        result = lw._get_cache_fill_path(cache_dir=None, weight_name="test")
+        assert result is None
+
+    def test_returns_none_when_weight_name_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that _get_cache_fill_path returns None when weight_name is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        result = lw._get_cache_fill_path(cache_dir=Path("/tmp/cache"), weight_name=None)
+        assert result is None
+
+    def test_returns_path_with_fingerprint(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that _get_cache_fill_path returns correct path with fingerprint."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        cache_dir = Path("/tmp/cache")
+        result = lw._get_cache_fill_path(cache_dir=cache_dir, weight_name="my_weight")
+        assert result is not None
+        assert result.parent == cache_dir
+        assert result.name.startswith("my_weight_")
+        assert result.suffix == ".tensorbin"
+
+
+# ============================================================================
+# Test padded_shape Property
+# ============================================================================
+
+
+class TestPaddedShape:
+    """Tests for LazyWeight.padded_shape property."""
+
+    def test_padded_shape_raises_when_not_resolved(self, mock_tensor, mock_mesh_device):
+        """Test that padded_shape raises assertion when not resolved."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            # Missing memory_config and layout
+        )
+        with pytest.raises(AssertionError, match="LazyWeight must be resolved"):
+            _ = lw.padded_shape
+
+    def test_padded_shape_raises_when_device_none(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
+        """Test that padded_shape raises ValueError when device is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=None,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        # First it will fail is_resolved check
+        with pytest.raises(AssertionError, match="LazyWeight must be resolved"):
+            _ = lw.padded_shape
+
+    def test_padded_shape_no_padding_when_replicated(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that padded_shape returns original shape when replicated (mesh_mapper_config is None)."""
+        # Need to set mesh_mapper_config to something to pass is_resolved - use a mock
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=None,
+        )
+        # Override is_resolved to return True for this test
+        with patch.object(lw, "is_resolved", return_value=True):
+            result = lw.padded_shape
+            assert result == tuple(mock_tensor.shape)
+
+    def test_padded_shape_no_padding_single_device(self, mock_tensor, mock_memory_config, mock_mesh_mapper_config):
+        """Test that padded_shape returns original shape for single device."""
+        single_device = MagicMock(spec=ttnn.MeshDevice)
+        single_device.get_num_devices.return_value = 1
+
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=single_device,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        result = lw.padded_shape
+        assert result == tuple(mock_tensor.shape)
+
+    def test_padded_shape_pads_for_sharding(self, mock_memory_config):
+        """Test that padded_shape correctly pads for sharding across multiple devices."""
+        # Create tensor with shape that needs padding: dim -1 is 100
+        # With 8 devices, shard_dim = 100/8 = 12.5, needs to be tile aligned (32)
+        # padded_shard = ceil(12.5/32)*32 = 32
+        # padded_hidden = 32 * 8 = 256
+        tensor = torch.randn(1, 32, 100)
+        device = MagicMock(spec=ttnn.MeshDevice)
+        device.get_num_devices.return_value = 8
+
+        config = ttnn.MeshMapperConfig(
+            placements=[ttnn.PlacementShard(-1)],
+            mesh_shape_override=ttnn.MeshShape([8]),
+        )
+
+        lw = LazyWeight(
+            source=tensor,
+            device=device,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        result = lw.padded_shape
+        # Original: (1, 32, 100), padded dim -1 should be 256 (32*8 where 32 is tile-aligned shard)
+        assert result == (1, 32, 256)
+
+    def test_padded_shape_raises_valueerror_when_device_none_but_resolved(
+        self, mock_tensor, mock_memory_config, mock_mesh_mapper_config
+    ):
+        """Test that padded_shape raises ValueError when device is None but is_resolved mocked True."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=None,
+            memory_config=mock_memory_config,
+            mesh_mapper_config=mock_mesh_mapper_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        # Override is_resolved to return True to reach the ValueError branch
+        with patch.object(lw, "is_resolved", return_value=True):
+            with pytest.raises(ValueError, match="device must be set to compute padded_shape"):
+                _ = lw.padded_shape
+
+
+# ============================================================================
+# Test get_device_weight Method
+# ============================================================================
+
+
+class TestGetDeviceWeight:
+    """Tests for LazyWeight.get_device_weight() method."""
+
+    def test_raises_when_device_none(self, mock_tensor, mock_memory_config):
+        """Test that get_device_weight raises when device is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=None,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        with pytest.raises(ValueError, match="device must be provided"):
+            lw.get_device_weight()
+
+    def test_raises_when_layout_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that get_device_weight raises when layout is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=None,
+        )
+        with pytest.raises(ValueError, match="layout must be provided"):
+            lw.get_device_weight()
+
+    def test_raises_when_memory_config_none(self, mock_tensor, mock_mesh_device):
+        """Test that get_device_weight raises when memory_config is None."""
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=None,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        with pytest.raises(ValueError, match="memory_config must be provided"):
+            lw.get_device_weight()
+
+    def test_returns_cached_value_on_second_call(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that get_device_weight returns cached _value on second call."""
+        mock_ttnn_tensor = MagicMock()
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+        # Manually set _value to simulate previous call
+        lw._value = mock_ttnn_tensor
+
+        result = lw.get_device_weight()
+        assert result is mock_ttnn_tensor
+
+    def test_get_device_weight_replicated_flow(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test get_device_weight with replicated tensor (no sharding) - mocked."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+            mock_ttnn.replicate_tensor_to_mesh_mapper.return_value = mock_mapper
+            mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+
+            lw = LazyWeight(
+                source=mock_tensor,
+                device=mock_mesh_device,
+                memory_config=mock_memory_config,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=None,  # Replicated
+            )
+
+            result = lw.get_device_weight()
+
+            # Verify replicate_tensor_to_mesh_mapper was called
+            mock_ttnn.replicate_tensor_to_mesh_mapper.assert_called_once_with(mock_mesh_device)
+            assert result is mock_ttnn_tensor
+
+    def test_get_device_weight_sharded_flow(
+        self, mock_tensor, mock_mesh_device, mock_memory_config, mock_mesh_mapper_config
+    ):
+        """Test get_device_weight with sharded tensor - mocked."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+            mock_ttnn.create_mesh_mapper.return_value = mock_mapper
+            mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+
+            lw = LazyWeight(
+                source=mock_tensor,
+                device=mock_mesh_device,
+                memory_config=mock_memory_config,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=mock_mesh_mapper_config,
+            )
+
+            result = lw.get_device_weight()
+
+            # Verify create_mesh_mapper was called with the config
+            mock_ttnn.create_mesh_mapper.assert_called_once_with(mock_mesh_device, mock_mesh_mapper_config)
+            assert result is mock_ttnn_tensor
+
+    def test_get_device_weight_with_cache_dir_weight_name(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test get_device_weight extracts cache_dir and weight_name from tuple."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_dir = Path(tmpdir)
+
+            with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+                mock_ttnn.replicate_tensor_to_mesh_mapper.return_value = mock_mapper
+                mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+                mock_ttnn.StorageType.HOST = ttnn.StorageType.HOST
+                mock_ttnn_tensor.storage_type.return_value = ttnn.StorageType.HOST
+
+                lw = LazyWeight(
+                    source=mock_tensor,
+                    cache_dir_weight_name=(cache_dir, "my_weight"),
+                    device=mock_mesh_device,
+                    memory_config=mock_memory_config,
+                    layout=ttnn.TILE_LAYOUT,
+                    mesh_mapper_config=None,
+                )
+
+                result = lw.get_device_weight()
+                assert result is not None
+
+    def test_get_device_weight_cache_hit(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test get_device_weight loads from cache when file exists."""
+        mock_ttnn_tensor = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_dir = Path(tmpdir)
+
+            # Create LazyWeight and compute cache path WITHOUT mocking
+            lw = LazyWeight(
+                source=mock_tensor,
+                cache_dir_weight_name=(cache_dir, "cached_weight"),
+                device=mock_mesh_device,
+                memory_config=mock_memory_config,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=None,
+            )
+
+            # Get expected cache path and create fake cache file
+            cache_path = lw._get_cache_fill_path(cache_dir, "cached_weight")
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            cache_path.touch()
+
+            # Now patch ttnn.load_tensor for the get_device_weight call
+            with patch("models.common.modules.lazy_weight.ttnn.load_tensor") as mock_load:
+                mock_load.return_value = mock_ttnn_tensor
+
+                result = lw.get_device_weight()
+
+                # Verify load_tensor was called (cache hit)
+                mock_load.assert_called_once()
+                assert result is mock_ttnn_tensor
+
+
+# ============================================================================
+# Test _auto_pad_for_sharded_tiles Function
+# ============================================================================
+
+
+class TestAutoPadForSharding:
+    """Tests for _auto_pad_for_sharded_tiles() function."""
+
+    def test_no_padding_when_shapes_match(self, mock_tensor):
+        """Test that no padding is applied when shapes match."""
+        result = _auto_pad_for_sharded_tiles(mock_tensor, tuple(mock_tensor.shape))
+        assert result is mock_tensor
+
+    def test_padding_applied_when_shapes_differ(self):
+        """Test that padding is applied when shapes differ."""
+        tensor = torch.randn(1, 32, 100)
+        padded_shape = (1, 32, 128)
+
+        result = _auto_pad_for_sharded_tiles(tensor, padded_shape)
+        assert result.shape == padded_shape
+        # Original data should be preserved
+        assert torch.equal(result[:, :, :100], tensor)
+
+    def test_padding_with_custom_pad_value(self):
+        """Test padding with custom pad value."""
+        tensor = torch.zeros(1, 32, 100)
+        padded_shape = (1, 32, 128)
+
+        result = _auto_pad_for_sharded_tiles(tensor, padded_shape, pad_value=1.0)
+        assert result.shape == padded_shape
+        # Padded region should have pad_value
+        assert torch.allclose(result[:, :, 100:], torch.ones(1, 32, 28))
+
+
+# ============================================================================
+# Test _from_torch_and_dump Function
+# ============================================================================
+
+
+class TestFromTorchAndDump:
+    """Tests for _from_torch_and_dump() function."""
+
+    def test_no_cache_when_cache_file_none(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that no caching occurs when cache_file_name is None."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+            mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+
+            result = _from_torch_and_dump(
+                tensor=mock_tensor,
+                device=mock_mesh_device,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                memory_config=mock_memory_config,
+                mesh_mapper=mock_mapper,
+                is_replicated=False,
+                pad_value=0.0,
+                cache_file_name=None,
+            )
+
+            mock_ttnn.from_torch.assert_called_once()
+            assert result is mock_ttnn_tensor
+
+    def test_cache_dump_when_cache_file_provided(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that tensor is cached when cache_file_name is provided."""
+        mock_ttnn_tensor = MagicMock()
+        mock_ttnn_tensor.storage_type.return_value = ttnn.StorageType.HOST
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_file = Path(tmpdir) / "test_cache.tensorbin"
+
+            with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+                mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+                mock_ttnn.StorageType.HOST = ttnn.StorageType.HOST
+
+                _from_torch_and_dump(
+                    tensor=mock_tensor,
+                    device=mock_mesh_device,
+                    dtype=ttnn.bfloat16,
+                    layout=ttnn.TILE_LAYOUT,
+                    memory_config=mock_memory_config,
+                    mesh_mapper=None,
+                    is_replicated=True,
+                    pad_value=0.0,
+                    cache_file_name=cache_file,
+                )
+
+                # Verify from_torch was called with device=None for caching
+                call_kwargs = mock_ttnn.from_torch.call_args.kwargs
+                assert call_kwargs.get("device") is None
+
+                # Verify dump_tensor_flatbuffer was called
+                mock_ttnn._ttnn.tensor.dump_tensor_flatbuffer.assert_called_once()
+
+    def test_replicated_uses_none_mesh_mapper(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that replicated tensors use None mesh_mapper in from_torch."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+            mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+
+            _from_torch_and_dump(
+                tensor=mock_tensor,
+                device=mock_mesh_device,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                memory_config=mock_memory_config,
+                mesh_mapper=mock_mapper,
+                is_replicated=True,
+                pad_value=0.0,
+                cache_file_name=None,
+            )
+
+            call_kwargs = mock_ttnn.from_torch.call_args.kwargs
+            assert call_kwargs.get("mesh_mapper") is None
+
+    def test_sharded_uses_mesh_mapper(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that sharded tensors use the provided mesh_mapper."""
+        mock_ttnn_tensor = MagicMock()
+        mock_mapper = MagicMock()
+
+        with patch("models.common.modules.lazy_weight.ttnn") as mock_ttnn:
+            mock_ttnn.from_torch.return_value = mock_ttnn_tensor
+
+            _from_torch_and_dump(
+                tensor=mock_tensor,
+                device=mock_mesh_device,
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                memory_config=mock_memory_config,
+                mesh_mapper=mock_mapper,
+                is_replicated=False,
+                pad_value=0.0,
+                cache_file_name=None,
+            )
+
+            call_kwargs = mock_ttnn.from_torch.call_args.kwargs
+            assert call_kwargs.get("mesh_mapper") is mock_mapper
+
+
+# ============================================================================
+# Test resolve_lazy_weight Function
+# ============================================================================
+
+
+class TestResolveLazyWeight:
+    """Tests for resolve_lazy_weight() function."""
+
+    def test_resolves_none_fields(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that resolve_lazy_weight fills in None fields."""
+        lw = LazyWeight(source=mock_tensor)
+
+        resolved = resolve_lazy_weight(
+            lw,
+            device=mock_mesh_device,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+
+        assert resolved.device is mock_mesh_device
+        assert resolved.memory_config is mock_memory_config
+        assert resolved.layout == ttnn.TILE_LAYOUT
+
+    def test_does_not_override_existing_fields(self, mock_tensor, mock_mesh_device, mock_memory_config):
+        """Test that resolve_lazy_weight does not override non-None fields."""
+        original_device = MagicMock(spec=ttnn.MeshDevice)
+        lw = LazyWeight(
+            source=mock_tensor,
+            device=original_device,
+            dtype=ttnn.bfloat4_b,
+        )
+
+        resolved = resolve_lazy_weight(
+            lw,
+            device=mock_mesh_device,
+            dtype=ttnn.bfloat16,
+            memory_config=mock_memory_config,
+            layout=ttnn.TILE_LAYOUT,
+        )
+
+        # Original non-None fields should be preserved
+        assert resolved.device is original_device
+        assert resolved.dtype == ttnn.bfloat4_b
+        # None fields should be filled
+        assert resolved.memory_config is mock_memory_config
+        assert resolved.layout == ttnn.TILE_LAYOUT
+
+    def test_returns_new_instance(self, mock_tensor, mock_mesh_device):
+        """Test that resolve_lazy_weight returns a new LazyWeight instance."""
+        lw = LazyWeight(source=mock_tensor)
+        resolved = resolve_lazy_weight(lw, device=mock_mesh_device)
+
+        assert resolved is not lw
+        assert resolved.device is mock_mesh_device
+        assert lw.device is None  # Original unchanged
+
+    def test_with_no_overrides(self, mock_tensor):
+        """Test resolve_lazy_weight with no kwargs returns equivalent instance."""
+        lw = LazyWeight(source=mock_tensor, dtype=ttnn.bfloat16)
+        resolved = resolve_lazy_weight(lw)
+
+        assert resolved is not lw
+        assert resolved.source is lw.source
+        assert resolved.dtype == lw.dtype
+
+
+# ============================================================================
+# Integration Tests (require device - use parametrized fixture)
+# ============================================================================
+
+
+@pytest.mark.parametrize(
+    "ttnn_mesh_device",
+    [
+        (1, 1),
+    ],
+    ids=[
+        "1x1",
+    ],
+    indirect=True,
+)
+class TestLazyWeightIntegration:
+    """Integration tests for LazyWeight that require actual TTNN device."""
+
+    def test_get_device_weight_replicated(self, ttnn_mesh_device: ttnn.MeshDevice):
+        """Test get_device_weight with replicated tensor (no mesh_mapper_config)."""
+        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)
+
+        lw = LazyWeight(
+            source=tensor,
+            device=ttnn_mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper_config=None,  # Replicated
+        )
+
+        result = lw.get_device_weight()
+        assert result is not None
+        # Subsequent calls should return cached value
+        result2 = lw.get_device_weight()
+        assert result2 is result
+
+    def test_get_device_weight_with_cache(self, ttnn_mesh_device: ttnn.MeshDevice):
+        """Test get_device_weight with caching enabled."""
+        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_dir = Path(tmpdir)
+
+            lw = LazyWeight(
+                source=tensor,
+                cache_dir_weight_name=(cache_dir, "test_weight"),
+                device=ttnn_mesh_device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=None,
+            )
+
+            # First call should create cache
+            result1 = lw.get_device_weight()
+            assert result1 is not None
+
+            # Verify cache file was created
+            cache_files = list(cache_dir.glob("*.tensorbin"))
+            assert len(cache_files) == 1
+            assert cache_files[0].name.startswith("test_weight_")
+
+    def test_get_device_weight_cache_hit(self, ttnn_mesh_device: ttnn.MeshDevice):
+        """Test get_device_weight loads from cache on second instance."""
+        tensor = torch.randn(1, 1, 32, 64, dtype=torch.bfloat16)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_dir = Path(tmpdir)
+
+            # First LazyWeight creates cache
+            lw1 = LazyWeight(
+                source=tensor,
+                cache_dir_weight_name=(cache_dir, "test_weight"),
+                device=ttnn_mesh_device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=None,
+            )
+            lw1.get_device_weight()
+
+            # Second LazyWeight should load from cache
+            lw2 = LazyWeight(
+                source=tensor,
+                cache_dir_weight_name=(cache_dir, "test_weight"),
+                device=ttnn_mesh_device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                layout=ttnn.TILE_LAYOUT,
+                mesh_mapper_config=None,
+            )
+            result2 = lw2.get_device_weight()
+
+            assert result2 is not None
+
+
+if __name__ == "__main__":
+    # Run unit tests that don't require device
+    pytest.main([__file__, "-v", "-k", "not Integration"])

code/models/common/tests/test_metrics.py

@@ -0,0 +1,584 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for TTNN-native metric functions - NO MOCKS.
+Tests numerical correctness of metrics computed with actual TTNN ops.
+
+Run with: pytest test_ttnn_metrics_numerical.py -v
+Requires: TTNN hardware/installation, pytest
+"""
+
+import pytest
+import torch
+
+import ttnn
+
+# Import metric functions from the validation framework
+from models.common.metrics import comp_allclose, compute_max_abs_error, compute_mean_abs_error, compute_pcc
+
+pytestmark = [
+    pytest.mark.parametrize(
+        "ttnn_mesh_device",
+        [
+            (1, 1),
+            # todo)) currently device metrics are only supported on 1x1 mesh device, which is experimental feature.
+            # we could add support for more mesh shapes later if there is demand.
+            # (1, 2),
+            # (1, 8),
+            # (2, 4),
+        ],
+        ids=[
+            "1x1",
+            # "1x2",
+            # "1x8",
+            # "2x4",
+        ],
+        indirect=True,
+    ),
+    pytest.mark.parametrize(
+        "layout,dtype",
+        [
+            (ttnn.ROW_MAJOR_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat16),
+            (ttnn.TILE_LAYOUT, ttnn.bfloat8_b),  # quantized dtypes only supported with TILE_LAYOUT
+            (ttnn.TILE_LAYOUT, ttnn.bfloat4_b),  # quantized dtypes only supported with TILE_LAYOUT
+        ],
+        ids=["row_major_bf16", "tile_bf16", "tile_bf8b", "tile_bf4b"],
+    ),
+]
+
+
+def _quantize_like_ttnn(x: torch.Tensor, device: ttnn.MeshDevice, dtype, layout):
+    """Round-trip through TTNN to obtain torch tensor quantized like given dtype.
+
+    For bf16, returns input unchanged. For bf8b/bf4b, uses TILE_LAYOUT only.
+    """
+    x_t = ttnn.from_torch(x, device=device, dtype=dtype, layout=layout)
+    return ttnn.to_torch(x_t)
+
+
+# Test case definitions
+# Format: (name, tensor_a_fn, tensor_b_fn, max_spec, mean_spec, cosine_spec)
+# Each spec is a tuple: (expected_value, tolerance) or None to skip expected check
+TEST_CASES = [
+    # Identical tensors
+    pytest.param(
+        "identical_random",
+        lambda: torch.randn(32, 64, dtype=torch.bfloat16),
+        lambda t: t.clone(),
+        (0.0, 1e-6),  # max: expect ~0, tight tolerance
+        (0.0, 1e-6),  # mean: expect ~0, tight tolerance
+        id="identical_random",
+    ),
+    # Known differences
+    pytest.param(
+        "known_diff_0.5",
+        lambda: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.bfloat16),
+        lambda t: t + torch.tensor([[0.0, 0.0, 0.5], [0.0, 0.0, 0.0]], dtype=torch.bfloat16),
+        (0.5, 0.02),  # max: expect 0.5
+        None,  # mean: skip expected check, just compare TTNN vs PyTorch
+        id="known_diff_max_0.5",
+    ),
+    pytest.param(
+        "uniform_diff_0.5",
+        lambda: torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16),
+        lambda t: t + 1,
+        (1, 0.02),  # max: expect 0.5
+        (1, 0.02),  # mean: expect 0.5 (all elements differ by same amount)
+        id="uniform_diff_0.5",
+    ),
+    # Orthogonal vectors
+    pytest.param(
+        "orthogonal",
+        lambda: torch.tensor([[1.0, 0.0]], dtype=torch.bfloat16),
+        lambda t: torch.tensor([[0.0, 1.0]], dtype=torch.bfloat16),
+        None,  # max: skip expected check
+        None,  # mean: skip expected check
+        id="orthogonal_vectors",
+    ),
+    # Opposite vectors
+    pytest.param(
+        "opposite",
+        lambda: torch.ones(1, 16, dtype=torch.bfloat16),
+        lambda t: -torch.ones(1, 16, dtype=torch.bfloat16),
+        None,  # max: skip expected check
+        None,  # mean: skip expected check
+        id="opposite_vectors",
+    ),
+    # Large tensors
+    pytest.param(
+        "large_128x256",
+        lambda: torch.randn(128, 256, dtype=torch.bfloat16),
+        lambda t: t + 0.1,
+        None,  # max: skip expected check
+        None,  # mean: skip expected check
+        id="large_128x256",
+    ),
+    # Edge cases - all zeros
+    pytest.param(
+        "all_zeros",
+        lambda: torch.zeros(16, 16, dtype=torch.bfloat16),
+        lambda t: t.clone(),
+        (0.0, 1e-6),  # max: expect 0
+        (0.0, 1e-6),  # mean: expect 0
+        id="edge_all_zeros",
+    ),
+    # Edge cases - all ones
+    pytest.param(
+        "all_ones",
+        lambda: torch.ones(16, 16, dtype=torch.bfloat16),
+        lambda t: t.clone(),
+        (0.0, 1e-6),  # max: expect 0
+        (0.0, 1e-6),  # mean: expect 0
+        id="edge_all_ones",
+    ),
+]
+
+
+@pytest.mark.parametrize("name,tensor_a_fn,tensor_b_fn,max_spec,mean_spec", TEST_CASES)
+def test_abs_metrics_vs_pytorch(ttnn_mesh_device, layout, dtype, name, tensor_a_fn, tensor_b_fn, max_spec, mean_spec):
+    """
+    Unified test for all metrics against PyTorch ground truth.
+    Tests various tensor configurations and verifies TTNN metrics match PyTorch.
+    Each metric spec is (expected_value, tolerance) or None.
+    """
+    # Generate tensors
+    torch_a = tensor_a_fn()
+    torch_b = tensor_b_fn(torch_a) if callable(tensor_b_fn) else tensor_b_fn
+
+    # Compute PyTorch ground truth (use TTNN-like quantized tensors for quantized dtypes)
+    torch_a_q = _quantize_like_ttnn(torch_a, ttnn_mesh_device, dtype, layout)
+    torch_b_q = _quantize_like_ttnn(torch_b, ttnn_mesh_device, dtype, layout)
+    max_error_torch = (torch_a_q - torch_b_q).abs().max().item()
+    mean_error_torch = (torch_a_q - torch_b_q).abs().mean().item()
+
+    # Convert to TTNN
+    ttnn_a = ttnn.from_torch(torch_a, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    ttnn_b = ttnn.from_torch(torch_b, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    # Compute with TTNN (ttnn vs ttnn)
+    max_error_ttnn = compute_max_abs_error(ttnn_a, ttnn_b)
+    mean_error_ttnn = compute_mean_abs_error(ttnn_a, ttnn_b)
+
+    # Mixed-mode checks (torch vs ttnn and ttnn vs torch)
+    # Use quantized torch views for mixed checks to fairly compare quantized dtypes
+    max_error_mixed_torch_ttnn = compute_max_abs_error(torch_a_q, ttnn_b)
+    mean_error_mixed_torch_ttnn = compute_mean_abs_error(torch_a_q, ttnn_b)
+    max_error_mixed_ttnn_torch = compute_max_abs_error(ttnn_a, torch_b_q)
+    mean_error_mixed_ttnn_torch = compute_mean_abs_error(ttnn_a, torch_b_q)
+
+    # Default tolerance for TTNN vs PyTorch comparison (bf16 precision)
+    default_tolerance = 0.02
+
+    # Verify max_abs_error
+    if max_spec is not None:
+        expected_max, tolerance_max = max_spec
+        assert (
+            abs(max_error_ttnn - expected_max) < tolerance_max
+        ), f"max_abs_error: expected {expected_max}, got {max_error_ttnn}"
+        # Check TTNN matches PyTorch
+        assert (
+            abs(max_error_ttnn - max_error_torch) < tolerance_max
+        ), f"max_abs_error TTNN vs PyTorch: {max_error_ttnn} vs {max_error_torch}"
+        # Mixed variants also match expected and PyTorch
+        assert (
+            abs(max_error_mixed_torch_ttnn - expected_max) < tolerance_max
+        ), f"max_abs_error (torch,ttnn): expected {expected_max}, got {max_error_mixed_torch_ttnn}"
+        assert (
+            abs(max_error_mixed_ttnn_torch - expected_max) < tolerance_max
+        ), f"max_abs_error (ttnn,torch): expected {expected_max}, got {max_error_mixed_ttnn_torch}"
+        assert (
+            abs(max_error_mixed_torch_ttnn - max_error_torch) < tolerance_max
+        ), f"max_abs_error (torch,ttnn) vs PyTorch: {max_error_mixed_torch_ttnn} vs {max_error_torch}"
+        assert (
+            abs(max_error_mixed_ttnn_torch - max_error_torch) < tolerance_max
+        ), f"max_abs_error (ttnn,torch) vs PyTorch: {max_error_mixed_ttnn_torch} vs {max_error_torch}"
+    else:
+        # No expected value, just check TTNN matches PyTorch with default tolerance
+        assert (
+            abs(max_error_ttnn - max_error_torch) < default_tolerance
+        ), f"max_abs_error TTNN vs PyTorch: {max_error_ttnn} vs {max_error_torch}"
+        assert (
+            abs(max_error_mixed_torch_ttnn - max_error_torch) < default_tolerance
+        ), f"max_abs_error (torch,ttnn) vs PyTorch: {max_error_mixed_torch_ttnn} vs {max_error_torch}"
+        assert (
+            abs(max_error_mixed_ttnn_torch - max_error_torch) < default_tolerance
+        ), f"max_abs_error (ttnn,torch) vs PyTorch: {max_error_mixed_ttnn_torch} vs {max_error_torch}"
+
+    # Verify mean_abs_error
+    if mean_spec is not None:
+        expected_mean, tolerance_mean = mean_spec
+        assert (
+            abs(mean_error_ttnn - expected_mean) < tolerance_mean
+        ), f"mean_abs_error: expected {expected_mean}, got {mean_error_ttnn}"
+        # Check TTNN matches PyTorch
+        assert (
+            abs(mean_error_ttnn - mean_error_torch) < tolerance_mean
+        ), f"mean_abs_error TTNN vs PyTorch: {mean_error_ttnn} vs {mean_error_torch}"
+        # Mixed variants also match expected and PyTorch
+        assert (
+            abs(mean_error_mixed_torch_ttnn - expected_mean) < tolerance_mean
+        ), f"mean_abs_error (torch,ttnn): expected {expected_mean}, got {mean_error_mixed_torch_ttnn}"
+        assert (
+            abs(mean_error_mixed_ttnn_torch - expected_mean) < tolerance_mean
+        ), f"mean_abs_error (ttnn,torch): expected {expected_mean}, got {mean_error_mixed_ttnn_torch}"
+        assert (
+            abs(mean_error_mixed_torch_ttnn - mean_error_torch) < tolerance_mean
+        ), f"mean_abs_error (torch,ttnn) vs PyTorch: {mean_error_mixed_torch_ttnn} vs {mean_error_torch}"
+        assert (
+            abs(mean_error_mixed_ttnn_torch - mean_error_torch) < tolerance_mean
+        ), f"mean_abs_error (ttnn,torch) vs PyTorch: {mean_error_mixed_ttnn_torch} vs {mean_error_torch}"
+    else:
+        # No expected value, just check TTNN matches PyTorch with default tolerance
+        assert (
+            abs(mean_error_ttnn - mean_error_torch) < default_tolerance
+        ), f"mean_abs_error TTNN vs PyTorch: {mean_error_ttnn} vs {mean_error_torch}"
+        assert (
+            abs(mean_error_mixed_torch_ttnn - mean_error_torch) < default_tolerance
+        ), f"mean_abs_error (torch,ttnn) vs PyTorch: {mean_error_mixed_torch_ttnn} vs {mean_error_torch}"
+        assert (
+            abs(mean_error_mixed_ttnn_torch - mean_error_torch) < default_tolerance
+        ), f"mean_abs_error (ttnn,torch) vs PyTorch: {mean_error_mixed_ttnn_torch} vs {mean_error_torch}"
+
+
+# check pcc computed on device against known good pcc values on host
+def test_pcc_ttnn_native(ttnn_mesh_device, layout, dtype):
+    """
+    Test TTNN-native PCC computation with actual TTNN tensors on device.
+
+    Verifies that:
+    1. PCC computes correctly on-device using TTNN ops
+    2. Results match the robust CPU/numpy implementation
+    3. Handles various correlation patterns (perfect, high, negative)
+    """
+    print("\nTest: TTNN-native PCC computation")
+
+    # Test case 1: Perfect positive correlation
+    torch.manual_seed(42)
+    a_torch = torch.randn(32, 32).bfloat16()
+    b_torch = a_torch.clone()
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    # Mixed-mode PCC (ttnn vs torch and torch vs ttnn)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Perfect correlation - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn >= 0.999, f"Perfect correlation should be ~1.0, got {pcc_ttnn}"
+    assert abs(pcc_ttnn - pcc_torch) < 0.01, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    # Mixed should match PyTorch within same tolerance
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.01, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.01, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 2: High correlation with small noise
+    torch.manual_seed(42)
+    a_torch = torch.randn(32, 64).bfloat16()
+    b_torch = a_torch + torch.randn(32, 64).bfloat16() * 0.01
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  High correlation - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn >= 0.95, f"High correlation should be >= 0.95, got {pcc_ttnn}"
+    # Allow more tolerance here due to bfloat16 and noise
+    assert abs(pcc_ttnn - pcc_torch) < 0.05, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.05, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.05, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 3: Negative correlation
+    torch.manual_seed(42)
+    a_torch = torch.randn(32, 32).bfloat16()
+    b_torch = -a_torch + torch.randn(32, 32).bfloat16() * 0.1
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Negative correlation - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn < -0.8, f"Negative correlation should be < -0.8, got {pcc_ttnn}"
+    assert abs(pcc_ttnn - pcc_torch) < 0.1, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.1, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.1, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 4: Larger tensor (128x256)
+    torch.manual_seed(42)
+    a_torch = torch.randn(128, 256).bfloat16()
+    b_torch = a_torch + torch.randn(128, 256).bfloat16() * 0.05
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Large tensor - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn >= 0.90, f"Large tensor PCC should be >= 0.90, got {pcc_ttnn}"
+    assert abs(pcc_ttnn - pcc_torch) < 0.1, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.1, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.1, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 5: Large noise - should produce low PCC (< 1.0)
+    torch.manual_seed(42)
+    a_torch = torch.randn(32, 32).bfloat16()
+    # Add significant noise (50% of signal strength)
+    b_torch = a_torch + torch.randn(32, 32).bfloat16() * 0.5
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Large noise (0.5x) - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn < 0.99, f"Large noise should reduce PCC below 0.99, got {pcc_ttnn}"
+    assert pcc_ttnn > 0.50, f"PCC should still show some correlation (>0.50), got {pcc_ttnn}"
+    assert abs(pcc_ttnn - pcc_torch) < 0.1, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.1, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.1, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 6: Very large noise - should produce very low PCC
+    torch.manual_seed(42)
+    a_torch = torch.randn(32, 32).bfloat16()
+    # Add massive noise (2x signal strength) - correlation should be weak
+    b_torch = a_torch + torch.randn(32, 32).bfloat16() * 2.0
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Very large noise (2.0x) - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn < 0.80, f"Very large noise should reduce PCC below 0.80, got {pcc_ttnn}"
+    assert abs(pcc_ttnn - pcc_torch) < 0.15, f"TTNN vs PyTorch mismatch: {pcc_ttnn} vs {pcc_torch}"
+    assert abs(pcc_mixed_1 - pcc_torch) < 0.15, f"PCC (ttnn,torch) vs PyTorch mismatch: {pcc_mixed_1} vs {pcc_torch}"
+    assert abs(pcc_mixed_2 - pcc_torch) < 0.15, f"PCC (torch,ttnn) vs PyTorch mismatch: {pcc_mixed_2} vs {pcc_torch}"
+
+    print("  ✓ TTNN-native PCC correctly detects varying correlation strengths!")
+
+
+def test_pcc_constant_tensors(ttnn_mesh_device, layout, dtype):
+    """
+    Test that TTNN-native PCC correctly handles constant tensors.
+
+    Constant tensors are an edge case where variance is zero. The implementation
+    should detect this and return 1.0 if both constants are equal, 0.0 otherwise.
+    """
+    print("\nTest: PCC with constant tensors")
+
+    # Test case 1: Same constant value
+    a_torch = torch.ones(32, 32).bfloat16() * 5.0
+    b_torch = torch.ones(32, 32).bfloat16() * 5.0
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    # Mixed-mode PCC computation
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Same constant (5.0) - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn == 1.0, f"Same constant tensors should have PCC=1.0, got {pcc_ttnn}"
+    assert pcc_torch == 1.0, f"PyTorch should also return 1.0, got {pcc_torch}"
+    assert pcc_mixed_1 == pcc_torch, f"PCC (ttnn,torch) should equal PyTorch: {pcc_mixed_1} vs {pcc_torch}"
+    assert pcc_mixed_2 == pcc_torch, f"PCC (torch,ttnn) should equal PyTorch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 2: Different constant values
+    a_torch = torch.ones(32, 32).bfloat16() * 5.0
+    b_torch = torch.ones(32, 32).bfloat16() * 3.0
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  Different constants (5.0 vs 3.0) - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn == 0.0, f"Different constant tensors should have PCC=0.0, got {pcc_ttnn}"
+    # PyTorch returns True/False which becomes 1.0/0.0
+    assert pcc_torch in [0.0, 1.0], f"PyTorch should return 0.0 or 1.0, got {pcc_torch}"
+    assert pcc_mixed_1 == pcc_torch, f"PCC (ttnn,torch) should equal PyTorch: {pcc_mixed_1} vs {pcc_torch}"
+    assert pcc_mixed_2 == pcc_torch, f"PCC (torch,ttnn) should equal PyTorch: {pcc_mixed_2} vs {pcc_torch}"
+
+    # Test case 3: All zeros
+    a_torch = torch.zeros(32, 32).bfloat16()
+    b_torch = torch.zeros(32, 32).bfloat16()
+
+    a_ttnn = ttnn.from_torch(a_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_ttnn = ttnn.from_torch(b_torch, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_ttnn, b_ttnn)
+    a_q = _quantize_like_ttnn(a_torch, ttnn_mesh_device, dtype, layout)
+    b_q = _quantize_like_ttnn(b_torch, ttnn_mesh_device, dtype, layout)
+    pcc_torch = compute_pcc(a_q, b_q)
+    pcc_mixed_1 = compute_pcc(a_ttnn, b_q)
+    pcc_mixed_2 = compute_pcc(a_q, b_ttnn)
+
+    print(f"  All zeros - TTNN: {pcc_ttnn:.6f}, PyTorch: {pcc_torch:.6f}")
+    assert pcc_ttnn == 1.0, f"Zero tensors should have PCC=1.0, got {pcc_ttnn}"
+    assert pcc_torch == 1.0, f"PyTorch should also return 1.0, got {pcc_torch}"
+    assert pcc_mixed_1 == pcc_torch, f"PCC (ttnn,torch) should equal PyTorch: {pcc_mixed_1} vs {pcc_torch}"
+    assert pcc_mixed_2 == pcc_torch, f"PCC (torch,ttnn) should equal PyTorch: {pcc_mixed_2} vs {pcc_torch}"
+
+    print("  ✓ TTNN-native PCC correctly handles constant tensors!")
+
+
+def test_pcc_all_nan_and_mixed_nan(ttnn_mesh_device, layout, dtype):
+    """
+    TTNN-native PCC should mirror CPU semantics for NaN cases:
+    - both all-NaN -> 1.0
+    - mixed NaN presence -> 0.0
+    """
+    print("\nTest: PCC NaN edge cases")
+
+    a_nan = torch.full((32, 32), float("nan"), dtype=torch.float32)
+    b_nan = torch.full((32, 32), float("nan"), dtype=torch.float32)
+    a_num = torch.zeros(32, 32, dtype=torch.float32)
+
+    a_nan_t = ttnn.from_torch(a_nan, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_nan_t = ttnn.from_torch(b_nan, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    a_num_t = ttnn.from_torch(a_num, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    # Both all-NaN -> 1.0
+    pcc_ttnn = compute_pcc(a_nan_t, b_nan_t)
+    a_nan_q = _quantize_like_ttnn(a_nan, ttnn_mesh_device, dtype, layout)
+    b_nan_q = _quantize_like_ttnn(b_nan, ttnn_mesh_device, dtype, layout)
+    pcc_cpu = compute_pcc(a_nan_q, b_nan_q)
+    # Mixed-mode PCC computation
+    pcc_mixed_1 = compute_pcc(a_nan_t, b_nan_q)
+    pcc_mixed_2 = compute_pcc(a_nan_q, b_nan_t)
+    print(f"  both NaN - TTNN: {pcc_ttnn}, CPU: {pcc_cpu}")
+    assert pcc_ttnn == 1.0
+    assert pcc_cpu == 1.0
+    assert pcc_mixed_1 == 1.0
+    assert pcc_mixed_2 == 1.0
+
+    # Mixed NaN presence -> 0.0
+    pcc_ttnn = compute_pcc(a_nan_t, a_num_t)
+    a_num_q = _quantize_like_ttnn(a_num, ttnn_mesh_device, dtype, layout)
+    pcc_cpu = compute_pcc(a_nan_q, a_num_q)
+    # Mixed-mode PCC computation
+    pcc_mixed_1 = compute_pcc(a_nan_t, a_num_q)
+    pcc_mixed_2 = compute_pcc(a_nan_q, a_num_t)
+    print(f"  mixed NaN - TTNN: {pcc_ttnn}, CPU: {pcc_cpu}")
+    assert pcc_ttnn == 0.0
+    assert pcc_cpu == 0.0
+    assert pcc_mixed_1 == 0.0
+    assert pcc_mixed_2 == 0.0
+
+    print("  ✓ TTNN-native PCC correctly handles NaN cases!")
+
+
+def test_pcc_zero_vs_nonzero(ttnn_mesh_device, layout, dtype):
+    """
+    One tensor all-zero and the other non-zero -> PCC = 0.0 on both TTNN and CPU.
+    """
+    print("\nTest: PCC zero vs non-zero")
+
+    a_zero = torch.zeros(32, 32, dtype=torch.float32)
+    b_nonzero = torch.ones(32, 32, dtype=torch.float32)
+
+    a_zero_t = ttnn.from_torch(a_zero, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_nonzero_t = ttnn.from_torch(b_nonzero, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+
+    pcc_ttnn = compute_pcc(a_zero_t, b_nonzero_t)
+    a_zero_q = _quantize_like_ttnn(a_zero, ttnn_mesh_device, dtype, layout)
+    b_nonzero_q = _quantize_like_ttnn(b_nonzero, ttnn_mesh_device, dtype, layout)
+    pcc_cpu = compute_pcc(a_zero_q, b_nonzero_q)
+    # Mixed-mode PCC computation
+    pcc_mixed_1 = compute_pcc(a_zero_t, b_nonzero_q)
+    pcc_mixed_2 = compute_pcc(a_zero_q, b_nonzero_t)
+    print(f"  zero vs non-zero - TTNN: {pcc_ttnn}, CPU: {pcc_cpu}")
+    assert pcc_ttnn == 0.0
+    assert pcc_cpu == 0.0
+    assert pcc_mixed_1 == 0.0
+    assert pcc_mixed_2 == 0.0
+    print("  ✓ TTNN-native PCC correctly handles zero vs non-zero!")
+
+
+def test_comp_allclose_ttnn_native(ttnn_mesh_device, layout, dtype):
+    """TTNN-native tests for comp_allclose using on-device ops."""
+
+    # Exact equality
+    a_t = ttnn.from_torch(torch.randn(2, 4, dtype=torch.bfloat16), device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_t = a_t
+    passed, msg = comp_allclose(a_t, b_t)
+    assert passed, f"TTNN equality should pass. Got: {msg}"
+
+    # Fail with tight tolerance
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16)
+    b = a + 0.5
+    a_t = ttnn.from_torch(a, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_t = ttnn.from_torch(b, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    passed, msg = comp_allclose(a_t, b_t, rtol=1e-6, atol=1e-6)
+    assert not passed and "Allclose check failed" in msg
+
+    # Pass with relaxed tolerance
+    passed, msg = comp_allclose(a_t, b_t, rtol=0.2, atol=0.6)
+    assert passed, f"Expected pass with relaxed tolerance. Got: {msg}"
+
+    # NaN equal
+    a = torch.tensor([float("nan"), 1.0], dtype=torch.float32)
+    b = torch.tensor([float("nan"), 1.0], dtype=torch.float32)
+    a_t = ttnn.from_torch(a, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_t = ttnn.from_torch(b, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    passed, _ = comp_allclose(a_t, b_t)
+    assert passed, "TTNN: Both NaNs at same positions should pass"
+
+    # Inf same sign pass, different sign fail
+    a = torch.tensor([float("inf"), -float("inf"), 2.0], dtype=torch.float32)
+    b = torch.tensor([float("inf"), -float("inf"), 2.0], dtype=torch.float32)
+    a_t = ttnn.from_torch(a, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    b_t = ttnn.from_torch(b, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    passed, _ = comp_allclose(a_t, b_t)
+    assert passed
+
+    b = torch.tensor([float("inf"), float("inf"), 2.0], dtype=torch.float32)
+    b_t = ttnn.from_torch(b, device=ttnn_mesh_device, dtype=dtype, layout=layout)
+    passed, msg = comp_allclose(a_t, b_t)
+    assert not passed and "Allclose check failed" in msg
+
+
+if __name__ == "__main__":
+    # Run with pytest for better output
+    pytest.main([__file__, "-v", "--tb=short"])

code/models/common/tests/test_module_input_validation.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+from models.common.models import module_input_validation as validation
+
+
+class FakeTensor:
+    def __init__(self, memory_config="dram"):
+        self._memory_config = memory_config
+        self.memory_config_calls = 0
+
+    def is_allocated(self):
+        return True
+
+    def memory_config(self):
+        self.memory_config_calls += 1
+        return self._memory_config
+
+
+class FakeConfig:
+    prefill_input_memcfg = "dram"
+
+
+class FakeModule:
+    config = FakeConfig()
+
+    def prefill_forward(self, x):
+        return x
+
+
+class FakeModel:
+    def __init__(self):
+        self.module = FakeModule()
+
+
+def iter_named_modules(model):
+    yield "module", model.module
+
+
+def test_validate_module_input_configs_reads_memory_config_when_active(monkeypatch):
+    monkeypatch.setattr(validation, "ttnn", SimpleNamespace(Tensor=FakeTensor))
+    model = FakeModel()
+    tensor = FakeTensor()
+
+    with validation.validate_module_input_configs(
+        model=model,
+        iter_named_modules=iter_named_modules,
+        mode="prefill",
+    ):
+        assert model.module.prefill_forward(tensor) is tensor
+
+    assert tensor.memory_config_calls == 1
+
+
+def test_suspend_module_input_validation_bypasses_memory_config(monkeypatch):
+    monkeypatch.setattr(validation, "ttnn", SimpleNamespace(Tensor=FakeTensor))
+    model = FakeModel()
+    tensor = FakeTensor()
+
+    with validation.validate_module_input_configs(
+        model=model,
+        iter_named_modules=iter_named_modules,
+        mode="prefill",
+    ):
+        with validation.suspend_module_input_validation():
+            assert model.module.prefill_forward(tensor) is tensor
+
+    assert tensor.memory_config_calls == 0
+
+
+def test_suspend_module_input_validation_restores_after_nested_contexts(monkeypatch):
+    monkeypatch.setattr(validation, "ttnn", SimpleNamespace(Tensor=FakeTensor))
+    model = FakeModel()
+    tensor = FakeTensor()
+
+    with validation.validate_module_input_configs(
+        model=model,
+        iter_named_modules=iter_named_modules,
+        mode="prefill",
+    ):
+        with validation.suspend_module_input_validation():
+            with validation.suspend_module_input_validation():
+                model.module.prefill_forward(tensor)
+            model.module.prefill_forward(tensor)
+        model.module.prefill_forward(tensor)
+
+    model.module.prefill_forward(tensor)
+
+    assert tensor.memory_config_calls == 1

code/models/common/tests/test_sampling.py

@@ -0,0 +1,598 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ttnn
+from models.common.sampling import (
+    LogProbsCalculator,
+    SamplingParams,
+    SeedManager,
+    broadcast_sampling_params,
+    format_sampling_params,
+)
+from models.common.sampling.tt_log_probs import MAX_TOP_LOGPROBS, LogProbsResult
+from models.common.utility_functions import comp_pcc
+
+
+# ---------------------------------------------------------------------------
+# Helper: simulate per-device top-k gather (mirrors TTSampling behaviour)
+# ---------------------------------------------------------------------------
+def _simulate_gathered_topk(torch_logits, num_devices, top_k=32):
+    """Simulate the per-device top-k + all-gather that TTSampling performs.
+
+    Args:
+        torch_logits: Full logits tensor, shape (1, 1, B, V).
+        num_devices: Number of TP devices.
+        top_k: Per-device top-k count.
+
+    Returns:
+        gathered_values: (1, 1, B, num_devices * top_k) raw logit values.
+        gathered_indices: (1, 1, B, num_devices * top_k) global vocab indices.
+    """
+    V = torch_logits.shape[-1]
+    shard_size = V // num_devices
+    all_values = []
+    all_indices = []
+    for d in range(num_devices):
+        shard = torch_logits[:, :, :, d * shard_size : (d + 1) * shard_size]
+        vals, local_idx = torch.topk(shard, top_k, dim=-1)
+        global_idx = local_idx + d * shard_size
+        all_values.append(vals)
+        all_indices.append(global_idx)
+    gathered_values = torch.cat(all_values, dim=-1)
+    gathered_indices = torch.cat(all_indices, dim=-1)
+    return gathered_values, gathered_indices
+
+
+# ===========================================================================
+# Top-K logprobs tests (TG Galaxy only)
+# ===========================================================================
+
+# Common TG Galaxy device parametrization for all new tests
+TG_SHAPE = [1, 1, 32, 8 * 16032]  # Llama on TG with 8-chip TP sharded vocab
+TG_DEVICE_PARAMS = {
+    "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING,
+    "dispatch_core_axis": ttnn.DispatchCoreAxis.COL,
+}
+TG_MESH_SHAPE = (8, 4)
+TG_SUB_CORE_GRIDS = ttnn.CoreRangeSet(
+    [
+        ttnn.CoreRange(ttnn.CoreCoord(1, 0), ttnn.CoreCoord(3, 9)),
+        ttnn.CoreRange(ttnn.CoreCoord(5, 0), ttnn.CoreCoord(6, 9)),
+    ]
+)
+TG_NUM_TP_DEVICES = 8  # TP dimension for Galaxy
+
+
+def _make_host_only_seed_manager(max_batch_size=4):
+    return SeedManager(SimpleNamespace(_sampling_dp=1), max_batch_size=max_batch_size)
+
+
+def test_seed_manager_seed_params_do_not_fallback_to_slot_zero():
+    seed_manager = _make_host_only_seed_manager()
+
+    assert seed_manager._seed_from_slot_params([11], 0) == 11
+    assert seed_manager._seed_from_slot_params([11], 1) is None
+    assert seed_manager._seed_from_slot_params(torch.tensor([22]), 1) is None
+    assert seed_manager._seed_from_slot_params(33, 3) == 33
+
+
+def test_seed_counter_position_alignment_skips_out_of_bounds_slots():
+    seed_manager = _make_host_only_seed_manager()
+
+    seed_manager.align_seed_counters_to_positions([101, None, 303], [0, 2], [5], offset=1)
+
+    assert seed_manager.seed_counters == [6, 0, 0, 0]
+
+
+def test_broadcast_sampling_params_preserves_none_list_fields():
+    params = SamplingParams(temperature=[1.0, 1.0], top_k=[1, 1], top_p=[1.0, 1.0], seed=[None, 42])
+
+    broadcast = broadcast_sampling_params(params, 0, slot_len=4)
+
+    assert broadcast.seed == [None, None, None, None]
+
+
+def test_format_sampling_params_uses_device_argmax_sentinel_for_greedy_rows():
+    params = format_sampling_params(
+        SamplingParams(temperature=0.0, top_k=32, top_p=0.95),
+        max_batch_size=32,
+    )
+
+    assert params.temperature[0] == 1.0
+    assert params.top_k[0] == 1
+    assert params.top_p[0] == 0.0
+
+
+def _skip_if_not_galaxy(mesh_device):
+    """Skip test if not running on TG Galaxy (32 devices)."""
+    if mesh_device.get_num_devices() != 32:
+        pytest.skip(f"Test requires TG Galaxy (32 devices), got {mesh_device.get_num_devices()}")
+
+
+def _push_topk_test_tensors_to_tg(torch_tensor, gathered_values, gathered_indices, mesh_device):
+    """Push logits, topk values, and topk indices to a TG Galaxy mesh device."""
+    logits_tt = ttnn.from_torch(
+        torch_tensor,
+        device=mesh_device,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(-1, None), mesh_shape=list(mesh_device.shape)),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    topk_values_tt = ttnn.from_torch(
+        gathered_values,
+        device=mesh_device,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    topk_indices_tt = ttnn.from_torch(
+        gathered_indices.to(torch.int32),
+        device=mesh_device,
+        dtype=ttnn.int32,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    return logits_tt, topk_values_tt, topk_indices_tt
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        [1, 1, 32, 8 * 18992],  # Qwen3 on T3K
+    ],
+)
+@pytest.mark.parametrize(
+    "device_params",
+    [
+        ({"fabric_config": ttnn.FabricConfig.FABRIC_1D}),
+    ],
+    indirect=["device_params"],
+    ids=["fabric_linear"],
+)
+def test_log_probs_calculation(shape, mesh_device):
+    seed = 1234
+    torch.manual_seed(seed)
+
+    log_probs_calculator = LogProbsCalculator(mesh_device)
+
+    torch_tensor = torch.randn(shape)
+    # shuffle the tensor in last 2 dimensions
+    for i in range(shape[-2]):
+        torch_tensor[:, :, i, :] = torch_tensor[:, :, i, torch.randperm(shape[-1])]
+
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+    indices_tensor = argmax_tensor.reshape(
+        argmax_tensor.shape[0], argmax_tensor.shape[1], argmax_tensor.shape[-1], argmax_tensor.shape[-2]
+    )
+    # Push inputs to device
+    logits_tensor = ttnn.from_torch(
+        torch_tensor,
+        device=mesh_device,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=-1),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    ttnn_indices_tensor = ttnn.from_torch(
+        indices_tensor,
+        device=mesh_device,
+        dtype=ttnn.int32,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    log_probs_calculator.set_log_probs_mode(True)
+    tt_log_probs = log_probs_calculator.calculate_log_probs(logits_tensor, ttnn_indices_tensor)
+    log_probs_tt_host = ttnn.to_torch(tt_log_probs, mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=3))
+    log_probs_tt_host = log_probs_tt_host[:, :, :1, :32]
+
+    # Calculate log-probs for each user on each chip using torch
+    log_probs_torch = F.log_softmax(torch_tensor.float(), dim=-1)
+    log_probs_torch_argmax = torch.gather(log_probs_torch, dim=-1, index=argmax_tensor)
+    log_probs_torch_argmax = torch.reshape(log_probs_torch_argmax, (1, 1, 1, 32))
+
+    passing, pcc = comp_pcc(log_probs_torch_argmax, log_probs_tt_host, pcc=0.99)
+    print(f"pcc={pcc}")
+
+    assert passing, f"Assertion failed, PCC={pcc}"
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        [1, 1, 32, 8 * 18992],  # Qwen3 on T3K
+    ],
+)
+@pytest.mark.parametrize(
+    "device_params",
+    [
+        ({"fabric_config": ttnn.FabricConfig.FABRIC_1D}),
+    ],
+    indirect=["device_params"],
+    ids=["fabric_linear"],
+)
+def test_log_probs_returns_none_when_disabled(shape, mesh_device):
+    """Test that calculate_log_probs returns None when enable_log_probs is False."""
+    log_probs_calculator = LogProbsCalculator(mesh_device)
+
+    torch_tensor = torch.randn(shape)
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+    indices_tensor = argmax_tensor.reshape(
+        argmax_tensor.shape[0], argmax_tensor.shape[1], argmax_tensor.shape[-1], argmax_tensor.shape[-2]
+    )
+
+    logits_tensor = ttnn.from_torch(
+        torch_tensor,
+        device=mesh_device,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=-1),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    ttnn_indices_tensor = ttnn.from_torch(
+        indices_tensor,
+        device=mesh_device,
+        dtype=ttnn.int32,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    # Log probs disabled (default) - should return None
+    log_probs_calculator.set_log_probs_mode(False)
+    result = log_probs_calculator.calculate_log_probs(logits_tensor, ttnn_indices_tensor)
+    assert result is None, f"Expected None when log_probs disabled, got {type(result)}"
+
+    # Log probs enabled - should return a tensor (not None)
+    log_probs_calculator.set_log_probs_mode(True)
+    num_devices = mesh_device.get_num_devices()
+    result = log_probs_calculator.calculate_log_probs(logits_tensor, ttnn_indices_tensor)
+    if num_devices in (8, 32) and log_probs_calculator.num_devices_for_sharding >= 2:
+        assert result is not None, "Expected tensor when log_probs enabled on supported device"
+    else:
+        assert result is None, "Expected None on unsupported device count"
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        [1, 1, 32, 8 * 16032],  # llama on TG with 8 chips sharded vocab
+    ],
+)
+@pytest.mark.parametrize(
+    "device_params",
+    [
+        (
+            {
+                "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING,
+                "dispatch_core_axis": ttnn.DispatchCoreAxis.COL,
+            }
+        ),
+    ],
+    indirect=True,
+    ids=["fabric_linear"],
+)
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        (8, 4),
+    ],
+    indirect=True,
+)
+def test_log_probs_with_sub_core_grids_on_galaxy(shape, mesh_device):
+    seed = 1234
+    torch.manual_seed(seed)
+
+    sub_core_grids = ttnn.CoreRangeSet(
+        [
+            ttnn.CoreRange(ttnn.CoreCoord(1, 0), ttnn.CoreCoord(3, 9)),
+            ttnn.CoreRange(ttnn.CoreCoord(5, 0), ttnn.CoreCoord(6, 9)),
+        ]
+    )
+    log_probs_calculator = LogProbsCalculator(mesh_device, sub_core_grids)
+
+    torch_tensor = torch.randn(shape)
+    # shuffle the tensor in last 2 dimensions
+    for i in range(shape[-2]):
+        torch_tensor[:, :, i, :] = torch_tensor[:, :, i, torch.randperm(shape[-1])]
+
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+    indices_tensor = argmax_tensor.reshape(
+        argmax_tensor.shape[0], argmax_tensor.shape[1], argmax_tensor.shape[-1], argmax_tensor.shape[-2]
+    )
+
+    if mesh_device.get_num_devices() == 8:
+        mesh_mapper = ttnn.ShardTensorToMesh(mesh_device, dim=-1)
+    elif mesh_device.get_num_devices() == 32:
+        mesh_mapper = ttnn.ShardTensor2dMesh(mesh_device, dims=(-1, None), mesh_shape=list(mesh_device.shape))
+    else:
+        raise ValueError(f"Unsupported number of devices: {mesh_device.get_num_devices()}")
+
+    logits_tensor = ttnn.from_torch(
+        torch_tensor,
+        device=mesh_device,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=mesh_mapper,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    ttnn_indices_tensor = ttnn.from_torch(
+        indices_tensor,
+        device=mesh_device,
+        dtype=ttnn.int32,
+        layout=ttnn.TILE_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+    log_probs_calculator.set_log_probs_mode(True)
+    tt_log_probs = log_probs_calculator.calculate_log_probs(logits_tensor, ttnn_indices_tensor)
+    log_probs_tt_host = ttnn.to_torch(tt_log_probs, mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=3))
+    # slice from (1,1,32,256) -> (1,1,1,32)
+    log_probs_tt_host = log_probs_tt_host[:, :, :1, :32]
+
+    log_probs_torch = F.log_softmax(torch_tensor.float(), dim=-1)
+    log_probs_torch_argmax = torch.gather(log_probs_torch, dim=-1, index=argmax_tensor)
+    log_probs_torch_argmax = torch.reshape(log_probs_torch_argmax, (1, 1, 1, 32))
+
+    passing, pcc = comp_pcc(log_probs_torch_argmax, log_probs_tt_host, pcc=0.99)
+    print(f"pcc={pcc}")
+
+    assert passing, f"Assertion failed, PCC={pcc}"
+
+
+# ===========================================================================
+# New top-K logprobs tests (TG Galaxy only)
+# ===========================================================================
+
+
+@pytest.mark.parametrize("shape", [TG_SHAPE])
+@pytest.mark.parametrize("device_params", [TG_DEVICE_PARAMS], indirect=True, ids=["tg"])
+@pytest.mark.parametrize("mesh_device", [TG_MESH_SHAPE], indirect=True)
+def test_top_k_log_probs_on_galaxy(shape, mesh_device):
+    """Top-K logprobs PCC check on TG Galaxy (32-device 2D mesh)."""
+    _skip_if_not_galaxy(mesh_device)
+    torch.manual_seed(1234)
+    batch_size = shape[2]
+
+    calc = LogProbsCalculator(mesh_device, TG_SUB_CORE_GRIDS, batch_size=batch_size, use_topk_logprobs=True)
+
+    torch_tensor = torch.randn(shape)
+    for i in range(batch_size):
+        torch_tensor[:, :, i, :] = torch_tensor[:, :, i, torch.randperm(shape[-1])]
+
+    log_probs_torch = F.log_softmax(torch_tensor.to(torch.float16), dim=-1)
+    gathered_values, gathered_indices = _simulate_gathered_topk(torch_tensor, TG_NUM_TP_DEVICES)
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+
+    logits_tt, topk_values_tt, topk_indices_tt = _push_topk_test_tensors_to_tg(
+        torch_tensor, gathered_values, gathered_indices, mesh_device
+    )
+
+    calc.set_log_probs_mode([True] * batch_size, num_logprobs=[5] * batch_size)
+    result = calc.calculate_topk_log_probs(logits_tt, topk_values_tt, topk_indices_tt)
+
+    assert result is not None, "Expected LogProbsResult, got None"
+    assert isinstance(result, LogProbsResult)
+
+    host_results = calc.transfer_logprobs_to_host(result, argmax_tensor.squeeze())
+
+    composer = calc._build_mesh_composer()
+    topk_logprobs_host = ttnn.to_torch(result.topk_logprobs, mesh_composer=composer)
+    topk_logprobs_host = topk_logprobs_host[0, 0, ...]
+    topk_indices_host = ttnn.to_torch(result.topk_indices, mesh_composer=composer)
+    topk_indices_host = topk_indices_host[0, 0, ...].long()
+
+    expected_logprobs = torch.gather(
+        log_probs_torch.squeeze(0).squeeze(0),
+        dim=-1,
+        index=topk_indices_host,
+    )
+
+    passing, pcc = comp_pcc(expected_logprobs, topk_logprobs_host, pcc=0.99)
+    print(f"Galaxy top-K logprobs PCC={pcc}")
+    assert passing, f"Galaxy top-K logprobs PCC failed: {pcc}"
+
+    for user_idx in range(batch_size):
+        r = host_results[user_idx]
+        assert r is not None
+        sampled_id = argmax_tensor[0, 0, user_idx, 0].item()
+        torch_lp = log_probs_torch[0, 0, user_idx, sampled_id].item()
+        assert abs(r["returned_token"]["logprob"] - torch_lp) < 0.05
+
+
+@pytest.mark.parametrize("shape", [TG_SHAPE])
+@pytest.mark.parametrize("device_params", [TG_DEVICE_PARAMS], indirect=True, ids=["tg"])
+@pytest.mark.parametrize("mesh_device", [TG_MESH_SHAPE], indirect=True)
+def test_top_k_log_probs_returns_none_when_not_needed(shape, mesh_device):
+    """calculate_topk_log_probs returns None when disabled."""
+    _skip_if_not_galaxy(mesh_device)
+    batch_size = shape[2]
+    calc = LogProbsCalculator(mesh_device, TG_SUB_CORE_GRIDS, batch_size=batch_size, use_topk_logprobs=True)
+
+    torch_tensor = torch.randn(shape)
+    gathered_values, gathered_indices = _simulate_gathered_topk(torch_tensor, TG_NUM_TP_DEVICES)
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+
+    logits_tt, topk_values_tt, topk_indices_tt = _push_topk_test_tensors_to_tg(
+        torch_tensor, gathered_values, gathered_indices, mesh_device
+    )
+
+    calc.set_log_probs_mode(False, num_logprobs=0)
+    result = calc.calculate_topk_log_probs(logits_tt, topk_values_tt, topk_indices_tt)
+    assert result is None, "Expected None when logprobs disabled"
+
+    calc.set_log_probs_mode(True, num_logprobs=0)
+    assert calc.topk_logprobs_needed  # needed for sampled token logprob
+    result = calc.calculate_topk_log_probs(logits_tt, topk_values_tt, topk_indices_tt)
+    assert result is not None, "Expected LogProbsResult when logprobs enabled"
+
+    sampled_ids = argmax_tensor.squeeze()
+    host_results = calc.transfer_logprobs_to_host(result, sampled_ids)
+    assert len(host_results) == batch_size
+    for i in range(batch_size):
+        r = host_results[i]
+        assert r is not None
+        assert r["returned_token"]["token_idx"] == int(sampled_ids[i].item())
+        assert len(r["top_logprobs"]["token_indices"]) == 0
+
+
+@pytest.mark.parametrize("shape", [TG_SHAPE])
+@pytest.mark.parametrize("device_params", [TG_DEVICE_PARAMS], indirect=True, ids=["tg"])
+@pytest.mark.parametrize("mesh_device", [TG_MESH_SHAPE], indirect=True)
+def test_per_user_logprobs_enabled(shape, mesh_device):
+    """Mixed per-user logprobs: only even users enabled."""
+    _skip_if_not_galaxy(mesh_device)
+    torch.manual_seed(42)
+    batch_size = shape[2]
+
+    calc = LogProbsCalculator(mesh_device, TG_SUB_CORE_GRIDS, batch_size=batch_size, use_topk_logprobs=True)
+
+    torch_tensor = torch.randn(shape)
+    for i in range(batch_size):
+        torch_tensor[:, :, i, :] = torch_tensor[:, :, i, torch.randperm(shape[-1])]
+
+    log_probs_torch = F.log_softmax(torch_tensor.to(torch.float16), dim=-1)
+    gathered_values, gathered_indices = _simulate_gathered_topk(torch_tensor, TG_NUM_TP_DEVICES)
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+
+    logits_tt, topk_values_tt, topk_indices_tt = _push_topk_test_tensors_to_tg(
+        torch_tensor, gathered_values, gathered_indices, mesh_device
+    )
+
+    enable_log_probs = [i % 2 == 0 for i in range(batch_size)]
+    num_logprobs_list = [5 if i % 2 == 0 else 0 for i in range(batch_size)]
+    calc.set_log_probs_mode(enable_log_probs, num_logprobs=num_logprobs_list)
+
+    result = calc.calculate_topk_log_probs(logits_tt, topk_values_tt, topk_indices_tt)
+    assert result is not None
+
+    sampled_ids = argmax_tensor.squeeze()
+    host_results = calc.transfer_logprobs_to_host(result, sampled_ids)
+
+    for i in range(batch_size):
+        if enable_log_probs[i]:
+            assert host_results[i] is not None
+            sampled_id = int(sampled_ids[i].item())
+            torch_lp = log_probs_torch[0, 0, i, sampled_id].item()
+            assert abs(host_results[i]["returned_token"]["logprob"] - torch_lp) < 0.05
+        else:
+            assert host_results[i] is None
+
+
+@pytest.mark.parametrize("shape", [TG_SHAPE])
+@pytest.mark.parametrize("device_params", [TG_DEVICE_PARAMS], indirect=True, ids=["tg"])
+@pytest.mark.parametrize("mesh_device", [TG_MESH_SHAPE], indirect=True)
+def test_set_log_probs_mode_validation(shape, mesh_device):
+    """Verify set_log_probs_mode internal state."""
+    _skip_if_not_galaxy(mesh_device)
+    batch_size = shape[2]
+    calc = LogProbsCalculator(mesh_device, TG_SUB_CORE_GRIDS, batch_size=batch_size, use_topk_logprobs=True)
+
+    calc.set_log_probs_mode(True)
+    assert calc.enable_log_probs is True
+    assert all(calc.logprobs_enabled)
+    assert calc.topk_logprobs_needed  # needed for sampled token logprob
+
+    calc.set_log_probs_mode(True, num_logprobs=5)
+    assert calc.topk_logprobs_needed is True
+    assert all(n == 5 for n in calc.num_logprobs)
+
+    enable_list = [True, False, True] + [False] * (batch_size - 3)
+    num_lp_list = [10, 0, 3] + [0] * (batch_size - 3)
+    calc.set_log_probs_mode(enable_list, num_logprobs=num_lp_list)
+    assert calc.enable_log_probs is True
+    assert calc.topk_logprobs_needed is True
+    assert calc.logprobs_enabled == enable_list
+    assert calc.num_logprobs == num_lp_list
+
+    calc.set_log_probs_mode(False, num_logprobs=0)
+    assert calc.enable_log_probs is False
+
+    calc.set_log_probs_mode(True, num_logprobs=0)
+    assert calc.enable_log_probs is True
+    assert calc.topk_logprobs_needed  # needed for sampled token logprob
+
+    calc.set_log_probs_mode(False, num_logprobs=0)
+    calc.set_log_probs_mode([True, True], num_logprobs=[10, 15], empty_slots=[2, 5])
+    assert calc.logprobs_enabled[2] is True
+    assert calc.logprobs_enabled[5] is True
+    assert calc.logprobs_enabled[0] is False
+    assert calc.num_logprobs[2] == 10
+    assert calc.num_logprobs[5] == 15
+
+    calc.set_log_probs_mode(False, num_logprobs=0)
+    calc.set_log_probs_mode(True, num_logprobs=7, empty_slots=[0, 3, 4])
+    assert all(calc.logprobs_enabled[i] for i in [0, 3, 4])
+    assert calc.logprobs_enabled[1] is False
+
+    calc.set_log_probs_mode([True], num_logprobs=[20], empty_slots=[1])
+    assert calc.logprobs_enabled[1] is True
+    assert calc.num_logprobs[0] == 7
+    assert calc.num_logprobs[1] == 20
+
+
+@pytest.mark.parametrize("shape", [TG_SHAPE])
+@pytest.mark.parametrize("device_params", [TG_DEVICE_PARAMS], indirect=True, ids=["tg"])
+@pytest.mark.parametrize("mesh_device", [TG_MESH_SHAPE], indirect=True)
+def test_top_k_logprobs_pcc_torch_vs_tt(shape, mesh_device):
+    """Compare host (PyTorch bfloat16) vs device (bfloat16) logprobs for full batch."""
+    _skip_if_not_galaxy(mesh_device)
+    torch.manual_seed(9999)
+    batch_size = shape[2]
+    requested_logprobs = MAX_TOP_LOGPROBS
+
+    calc = LogProbsCalculator(mesh_device, TG_SUB_CORE_GRIDS, batch_size=batch_size, use_topk_logprobs=True)
+
+    torch_tensor = torch.randn(shape).to(torch.bfloat16)
+    for i in range(batch_size):
+        torch_tensor[:, :, i, :] = torch_tensor[:, :, i, torch.randperm(shape[-1])]
+
+    log_probs_torch = F.log_softmax(torch_tensor, dim=-1, dtype=torch.bfloat16)
+    gathered_values, gathered_indices = _simulate_gathered_topk(torch_tensor, TG_NUM_TP_DEVICES)
+    argmax_tensor = torch.argmax(torch_tensor, dim=-1, keepdim=True)
+
+    logits_tt, topk_values_tt, topk_indices_tt = _push_topk_test_tensors_to_tg(
+        torch_tensor, gathered_values, gathered_indices, mesh_device
+    )
+
+    calc.set_log_probs_mode([True] * batch_size, num_logprobs=[requested_logprobs] * batch_size)
+
+    result = calc.calculate_topk_log_probs(logits_tt, topk_values_tt, topk_indices_tt)
+    assert result is not None
+
+    sampled_ids = argmax_tensor.squeeze()
+    host_results = calc.transfer_logprobs_to_host(result, sampled_ids)
+
+    for user in range(batch_size):
+        r = host_results[user]
+        assert r is not None
+
+        device_sampled_lp = r["returned_token"]["logprob"]
+        token_idx = r["returned_token"]["token_idx"]
+        torch_sampled_lp = log_probs_torch[0, 0, user, token_idx].item()
+        assert abs(device_sampled_lp - torch_sampled_lp) < 0.05
+
+        top_indices = r["top_logprobs"]["token_indices"]
+        top_lps_device = torch.tensor(r["top_logprobs"]["logprobs"], dtype=torch.float32)
+        assert len(top_indices) == requested_logprobs
+
+        top_lps_torch = log_probs_torch[0, 0, user, top_indices].float()
+        passing, pcc = comp_pcc(top_lps_torch.unsqueeze(0), top_lps_device.unsqueeze(0), pcc=0.98)
+        assert passing, (
+            f"User {user} top-{requested_logprobs} logprobs PCC failed: {pcc}\n"
+            f"  device: {top_lps_device[:5].tolist()}...\n"
+            f"  torch:  {top_lps_torch[:5].tolist()}..."
+        )

code/models/common/tests/test_utils.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from models.common.tests.utils import stable_model_seed
+
+
+def test_stable_model_seed_deterministic() -> None:
+    assert stable_model_seed("llama-3") == stable_model_seed("llama-3")
+
+
+def test_stable_model_seed_distinct() -> None:
+    assert stable_model_seed("llama-3") != stable_model_seed("mistral-7b")
+
+
+def test_stable_model_seed_uint32_range() -> None:
+    seed = stable_model_seed("llama-3")
+    assert 0 <= seed < 2**32

code/models/common/tests/test_validation_tools.py

@@ -0,0 +1,622 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+These tests demonstrate how to use @validate_against to compare TTNN implementations
+against reference PyTorch implementations with automatic metrics collection.
+"""
+
+import pytest
+import torch
+
+import ttnn
+from models.common.auto_compose import to_torch_auto_compose
+from models.common.metrics import compute_pcc_host
+from models.common.validation_tools import (
+    Metric,
+    MetricSpec,
+    clear_validation_results,
+    compare_to_torch,
+    compare_to_ttnn,
+    enable_validation,
+    get_validation_registry,
+)
+
+# [INFO] the purpose of this test is to validate the validation framework itself,
+# which does not care about the mesh shape or tensor layout; we have other test files on those topics.
+pytestmark = [
+    pytest.mark.parametrize(
+        "ttnn_mesh_device",
+        [
+            (1, 1),
+        ],
+        ids=[
+            "1x1",
+        ],
+        indirect=True,
+    ),
+]
+
+# ============================================================================
+# Example 1: Validating RMSNorm against PyTorch reference
+# ============================================================================
+
+
+def torch_rms_norm(x, weight, eps=1e-6):
+    """Reference PyTorch implementation of RMS normalization"""
+    variance = x.pow(2).mean(-1, keepdim=True)
+    x = x * torch.rsqrt(variance + eps)
+    return weight * x
+
+
+class HostValidatedRMSNorm:
+    """RMS Normalization with validation decorator using old input_map pattern"""
+
+    def __init__(self, weight: torch.Tensor, eps: float, device):
+        self.eps = eps
+        self.weight = ttnn.from_torch(
+            weight.unsqueeze(0).unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+        )
+
+    @compare_to_torch(
+        reference_fn=torch_rms_norm,
+        input_to_torch=lambda self, x: (
+            # [INFO] produce input args to torch_rms_norm as a tuple
+            (to_torch_auto_compose(x), to_torch_auto_compose(self.weight)),
+            # [INFO] produce input kwargs to torch_rms_norm as a dict
+            {"eps": self.eps},
+        ),
+        metric_tolerances={
+            Metric.MAX_ABS_ERROR: 1e-2,
+            Metric.MEAN_ABS_ERROR: 1e-3,
+            "pcc": 0.99,  # can use enum or their string values
+        },
+        enabled=True,
+    )
+    def __call__(self, x):
+        # x shape: [1, seq_len, hidden_size]
+        x_squared = ttnn.mul(x, x)
+        mean_x_squared = ttnn.mean(x_squared, dim=-1, keepdim=True)
+        rms = ttnn.sqrt(ttnn.add(mean_x_squared, self.eps))
+        x_normed = ttnn.mul(x, ttnn.reciprocal(rms))
+        return ttnn.mul(x_normed, self.weight)
+
+
+class DeviceValidatedRMSNorm:
+    """RMS Normalization - ultra-clean pattern: NO conversions needed!"""
+
+    def __init__(self, weight: torch.Tensor, eps: float, device):
+        self.eps = eps
+        self.weight_torch = weight  # Keep for reference
+        self.weight = ttnn.from_torch(
+            weight.unsqueeze(0).unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+        )
+        self.device = device
+
+    def _reference_impl(self, x):
+        """Reference implementation - mocking a TTNN reference implementation for testing"""
+        # Convert TTNN to torch for reference computation
+        x_torch = ttnn.to_torch(x).squeeze(0)
+        result_torch = torch_rms_norm(x_torch, self.weight_torch, self.eps)
+        # Convert back to TTNN to match __call__ output type
+        return ttnn.from_torch(
+            result_torch.unsqueeze(0), device=self.device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+        )
+
+    # [INFO] this decorator is useful when the reference function is a TTNN-native function.
+    #        currently, it is experimental and requires the reference function has same-ordered
+    #        inputs as the decorated function.
+    @compare_to_ttnn(
+        reference_fn=lambda self, x: self._reference_impl(x),
+        # [INFO] passing `metric_tolerances` is optional; if not provided, the default tolerances will be used:
+        # metric_tolerances={
+        #     Metric.MAX_ABS_ERROR: 1e-2,
+        #     Metric.PCC: 0.99,
+        # },
+    )
+    def __call__(self, x):
+        # x shape: [1, seq_len, hidden_size]
+        x_squared = ttnn.mul(x, x)
+        mean_x_squared = ttnn.mean(x_squared, dim=-1, keepdim=True)
+        rms = ttnn.sqrt(ttnn.add(mean_x_squared, self.eps))
+        x_normed = ttnn.mul(x, ttnn.reciprocal(rms))
+        return ttnn.mul(x_normed, self.weight)
+
+    @compare_to_ttnn(
+        reference_fn=lambda self, x: self._reference_impl(x),
+    )
+    def _call_torch__(self, x):
+        # copied __call__ code below and converted to torch tensor to mock a function under test that returns a torch tensor
+        # x shape: [1, seq_len, hidden_size]
+        x_squared = ttnn.mul(x, x)
+        mean_x_squared = ttnn.mean(x_squared, dim=-1, keepdim=True)
+        rms = ttnn.sqrt(ttnn.add(mean_x_squared, self.eps))
+        x_normed = ttnn.mul(x, ttnn.reciprocal(rms))
+        return to_torch_auto_compose(ttnn.mul(x_normed, self.weight))
+
+
+def test_validation_rmsnorm_host_and_device(ttnn_mesh_device: ttnn.MeshDevice):
+    registry = get_validation_registry()
+
+    hidden_size = 64
+    batch_size = 1
+    seq_len = 8
+
+    weight = torch.randn(hidden_size, dtype=torch.bfloat16)
+
+    # Device-validated RMSNorm
+    rms_device = DeviceValidatedRMSNorm(weight, eps=1e-6, device=ttnn_mesh_device)
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=torch.bfloat16)
+    x_tt = ttnn.from_torch(x.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    _ = rms_device(x_tt)
+
+    _ = rms_device._call_torch__(x_tt)
+
+    # Host-validated RMSNorm
+    rms_host = HostValidatedRMSNorm(weight, eps=1e-6, device=ttnn_mesh_device)
+    x2 = torch.randn(batch_size, seq_len, hidden_size, dtype=torch.bfloat16)
+    x2_tt = ttnn.from_torch(x2.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    _ = rms_host(x2_tt)
+
+    assert len(registry.results) >= 3
+    # Expect all validations to pass
+    assert registry.results[0].metrics[Metric.MAX_ABS_ERROR].passed
+    assert registry.results[1].metrics[Metric.MAX_ABS_ERROR].passed
+    assert registry.results[2].metrics[Metric.MAX_ABS_ERROR].passed
+    assert registry.results[2].metrics[Metric.MEAN_ABS_ERROR].passed
+    assert registry.results[0].metrics[Metric.PCC].passed
+    assert registry.results[1].metrics[Metric.PCC].passed
+    assert registry.results[2].metrics[Metric.PCC].passed
+
+
+# ============================================================================
+# Example 2: Validating matrix multiplication
+# ============================================================================
+
+
+@compare_to_torch(
+    reference_fn=torch.matmul,
+    # [INFO] when reference function accepts inputs in the same order as the decorated function,
+    #        we can omit input_to_torch; the mapping will be inferred automatically.
+    metric_tolerances={
+        Metric.MAX_ABS_ERROR: 1.5e-1,
+        Metric.PCC: 0.99,
+    },
+)
+def ttnn_matmul(a, b):
+    """TTNN matrix multiplication with validation"""
+    return ttnn.matmul(a, b)
+
+
+# make a test case to show how to directly use auto_compose to convert ttnn to torch
+@compare_to_torch(
+    reference_fn=torch.matmul,
+    # [INFO] this is a simple example of input remapping.
+    input_to_torch=lambda a, b: (to_torch_auto_compose(b), to_torch_auto_compose(a)),
+    metric_tolerances={
+        Metric.MAX_ABS_ERROR: 1.5e-1,
+        Metric.PCC: 0.99,
+    },
+)
+def ttnn_matmul_reverse(a, b):
+    """TTNN matrix multiplication with validation"""
+    return ttnn.matmul(b, a)
+
+
+def test_validation_matmul(ttnn_mesh_device: ttnn.MeshDevice):
+    registry = get_validation_registry()
+
+    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=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    _ = ttnn_matmul(a_tt, b_tt)
+    _ = ttnn_matmul_reverse(b_tt, a_tt)
+
+    # Expect two validations recorded and both passed
+    assert len(registry.results) >= 2
+    assert registry.results[-1].passed
+    assert registry.results[-2].passed
+
+
+# ============================================================================
+# Example 3: Custom metrics and complex mappings
+# ============================================================================
+
+
+def custom_attention_reference(q, k, v, scale):
+    """Reference attention computation"""
+    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+    attn_weights = torch.nn.functional.softmax(scores, dim=-1)
+    return torch.matmul(attn_weights, v)
+
+
+@compare_to_torch(
+    reference_fn=custom_attention_reference,
+    # [INFO]{ when reference function accepts inputs in the same order as the decorated function,
+    # we can omit input_to_torch; it will be inferred automatically as if the following code were written:
+    # input_to_torch=lambda q, k, v, scale: (
+    #     to_torch_auto_compose(q),
+    #     to_torch_auto_compose(k),
+    #     to_torch_auto_compose(v),
+    #     scale,
+    # ),
+    # [INFO]}
+    metric_tolerances={
+        Metric.MAX_ABS_ERROR: 0.1,
+        Metric.MEAN_ABS_ERROR: 0.02,
+        Metric.PCC: 0.99,
+    },
+)
+def ttnn_attention(q, k, v, scale):
+    """Simplified attention with validation"""
+    scores = ttnn.matmul(q, ttnn.transpose(k, -2, -1))
+    scores = ttnn.mul(scores, scale)
+    attn_weights = ttnn.softmax(scores, dim=-1)
+    return ttnn.matmul(attn_weights, v)
+
+
+def test_validation_attention(ttnn_mesh_device: ttnn.MeshDevice):
+    m, n, dk, dv = 8, 8, 16, 16
+    q = torch.randn(1, m, dk, dtype=torch.bfloat16)
+    k = torch.randn(1, n, dk, dtype=torch.bfloat16)
+    v = torch.randn(1, n, dv, dtype=torch.bfloat16)
+
+    q_tt = ttnn.from_torch(q.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    k_tt = ttnn.from_torch(k.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    v_tt = ttnn.from_torch(v.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    registry = get_validation_registry()
+    before = len(registry.results)
+    scale = 1.0 / (dk**0.5)
+    _ = ttnn_attention(q_tt, k_tt, v_tt, scale)
+    assert len(registry.results) == before + 1
+    test_result = registry.results[-1]
+    # expect the test to pass max_abs_error and pcc checks
+    assert test_result.metrics[Metric.MAX_ABS_ERROR].passed
+    assert test_result.metrics[Metric.PCC].passed
+    assert test_result.metrics[Metric.MEAN_ABS_ERROR].passed
+
+
+# ============================================================================
+# Example 4: Validating from_torch checkpoint
+# ============================================================================
+
+
+@compare_to_torch(
+    reference_fn=lambda tensor, device: tensor,
+    output_to_torch=to_torch_auto_compose,
+    metric_tolerances={
+        Metric.MAX_ABS_ERROR: 0.015,
+        Metric.MEAN_ABS_ERROR: 0.01,
+        Metric.PCC: 0.99,
+    },
+)
+def from_torch_checkpoint(tensor: torch.Tensor, device: ttnn.MeshDevice):
+    """Return TTNN tensor created via from_torch from a checkpoint tensor."""
+    return ttnn.from_torch(tensor, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+
+def test_validation_checkpoint_from_torch(ttnn_mesh_device: ttnn.MeshDevice):
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    # Simulated checkpoint tensor (e.g., a weight matrix)
+    rows, cols = 32, 128
+    weight = torch.randn(rows, cols, dtype=torch.float32)
+
+    # Validate a direct from_torch call via the decorated function
+    _ = from_torch_checkpoint(weight, ttnn_mesh_device)
+
+    # Ensure a result was recorded and it passed
+    assert len(registry.results) == before + 1
+    assert registry.results[-1].passed
+
+
+# ============================================================================
+# Example 5: Validating with MetricSpec
+# ============================================================================
+
+
+@compare_to_torch(
+    reference_fn=torch.matmul,
+    metric_tolerances={
+        "pcc_host": MetricSpec(tolerance=0.99, higher_is_better=True, compute_fn=compute_pcc_host),
+    },
+)
+def ttnn_matmul_metric_spec(a, b):
+    return ttnn.matmul(a, b)
+
+
+def test_validation_matmul_metric_spec(ttnn_mesh_device: ttnn.MeshDevice):
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    m, n, k = 8, 10, 6
+    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=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    _ = ttnn_matmul_metric_spec(a_tt, b_tt)
+
+    assert len(registry.results) == before + 1
+    assert registry.results[-1].passed
+
+
+# ============================================================================
+# Example 6: Validating with non-decorator use of compare_to_torch
+#            between a class instance and a reference class instance!
+# ============================================================================
+
+
+def test_validation_non_decorator_class_vs_class_torch(ttnn_mesh_device: ttnn.MeshDevice):
+    """Validate a callable class against a reference class using non-decorator style."""
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    # Simple linear layer implemented in TTNN (__call__) vs Torch reference (forward)
+    m, n, k = 8, 10, 6
+    x = torch.randn(1, m, k, dtype=torch.bfloat16)
+    w = torch.randn(1, k, n, dtype=torch.bfloat16)
+
+    class TTLinear:
+        def __init__(self, weight: torch.Tensor, device: ttnn.MeshDevice):
+            # Weight expected as [1, k, n]; add device batch dim for TTNN tensor
+            self.weight = ttnn.from_torch(
+                weight.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+            )
+
+        def __call__(self, inp):
+            return ttnn.matmul(inp, self.weight)
+
+    class TorchLinearRef:
+        def __init__(self, weight: torch.Tensor):
+            self.weight = weight
+
+        def forward(self, inp: torch.Tensor):
+            return torch.matmul(inp, self.weight)
+
+    # Instantiate both implementations
+    layer = TTLinear(w, ttnn_mesh_device)
+    ref_layer = TorchLinearRef(w)
+
+    # Convert input to TTNN tensor (add device batch dim)
+    x_tt = ttnn.from_torch(x.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    # Non-decorator usage: wrap the unbound __call__ so we can pass (self, x)
+    validated_call = compare_to_torch(
+        reference_fn=lambda self, inp: ref_layer.forward(inp),
+        metric_tolerances={
+            Metric.MAX_ABS_ERROR: 1.5e-1,
+            Metric.PCC: 0.99,
+        },
+    )(TTLinear.__call__)
+
+    _ = validated_call(layer, x_tt)
+
+    assert len(registry.results) == before + 1
+    assert registry.results[-1].passed
+
+
+# ============================================================================
+# Example 7: Validating with non-decorator use of compare_to_ttnn
+#            between a class instance (return torch tensor) and a reference class instance
+#            (return TTNN tensor)
+# NOTE: This use of compare_to_ttnn could come in handy in situations where a module instance
+#       within torch implementation is being replaced by a TTNN module instance and
+#       we want to check the output of the TTNN module instance against the output of
+#       the torch module instance during end2end testing.
+# ============================================================================
+
+
+def test_validation_non_decorator_class_vs_class_ttnn(ttnn_mesh_device: ttnn.MeshDevice):
+    """Validate a callable TTNN class against a TTNN reference class using non-decorator style."""
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    # Simple linear layer implemented in TTNN (__call__) vs TTNN reference (forward)
+    m, n, k = 8, 10, 6
+    x = torch.randn(1, m, k, dtype=torch.bfloat16)
+    w = torch.randn(1, k, n, dtype=torch.bfloat16)
+
+    class TTLinear:
+        def __init__(self, weight: torch.Tensor, device: ttnn.MeshDevice):
+            # Weight expected as [1, k, n]; add device batch dim for TTNN tensor
+            self.weight = ttnn.from_torch(
+                weight.unsqueeze(0), device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT
+            )
+
+        def __call__(self, inp):
+            return ttnn.matmul(inp, self.weight)
+
+    class TorchLinearRef:
+        def __init__(self, weight: torch.Tensor, device: ttnn.MeshDevice):
+            self.weight = weight
+
+        def forward(self, inp):
+            return torch.matmul(inp, self.weight)
+
+    # Instantiate both implementations
+    layer = TTLinear(w, ttnn_mesh_device)
+    ref_layer = TorchLinearRef(w, ttnn_mesh_device)
+
+    # # Convert input to TTNN tensor (add device batch dim)
+    # x_tt = ttnn.from_torch(x.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    # Non-decorator usage: wrap the unbound __call__ so we can pass (self, x)
+    validated_call = compare_to_ttnn(
+        reference_fn=lambda inp: layer(inp),
+        input_to_ttnn=lambda self, inp: (
+            ttnn.from_torch(inp, device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT),
+        ),
+        metric_tolerances={
+            Metric.MAX_ABS_ERROR: 1.5e-1,
+            Metric.PCC: 0.99,
+        },
+    )(TorchLinearRef.forward)
+
+    ttnn.SetDefaultDevice(ttnn_mesh_device)
+    _ = validated_call(ref_layer, x.unsqueeze(0))
+
+    assert len(registry.results) == before + 1
+    assert registry.results[-1].passed
+
+
+def test_return_reference_output_torch(ttnn_mesh_device: ttnn.MeshDevice):
+    """Demonstrate return_reference_output=True returns the reference (torch) output.
+
+    The decorator computes torch.matmul on host for reference, then returns that
+    reference result converted back to a TTNN tensor distributed like the impl output.
+    """
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    m, n, k = 8, 10, 6
+    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=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    ref_host = torch.ones(m, n, dtype=torch.bfloat16)
+
+    @compare_to_torch(
+        # mock a reference function that returns a torch tensor with the same shape as the decorated function output
+        reference_fn=lambda a, b: ref_host,
+        metric_tolerances={
+            Metric.MAX_ABS_ERROR: 1,  # outrageous tolerance to confirm the mock
+            Metric.PCC: 0.99,
+        },
+        return_reference_output=True,
+    )
+    def _impl_matmul(a, b):
+        return ttnn.matmul(a, b)
+
+    # Call impl; returned value should be the reference result (distributed as impl output)
+    out_tt = _impl_matmul(a_tt, b_tt)
+
+    # Registry records one validation
+    assert len(registry.results) == before + 1
+    assert not registry.results[-1].metrics[Metric.MAX_ABS_ERROR].passed
+    assert registry.results[-1].metrics[Metric.PCC].passed
+
+    # Convert both outputs to host and verify numerical equivalence
+    out_host = to_torch_auto_compose(out_tt)
+    assert torch.allclose(out_host, ref_host)
+
+
+# ============================================================================
+# Additional test functions
+# ============================================================================
+
+
+def test_validation_enable_disable(ttnn_mesh_device: ttnn.MeshDevice):
+    a = torch.randn(1, 8, 8, dtype=torch.bfloat16)
+    b = torch.randn(1, 8, 8, dtype=torch.bfloat16)
+    a_tt = ttnn.from_torch(a.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    registry = get_validation_registry()
+    recorded_after_enable = len(registry.results)
+    # Disable validation: should not record
+    enable_validation(False)
+    _ = ttnn_matmul(a_tt, b_tt)
+    assert len(registry.results) == recorded_after_enable
+
+    # Re-enable for subsequent tests
+    enable_validation(True)
+
+
+def test_validation_non_decorator_host(ttnn_mesh_device: ttnn.MeshDevice):
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    m, n, k = 8, 10, 6
+    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=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    def _matmul(a, b):
+        return ttnn.matmul(a, b)
+
+    validated_matmul = compare_to_torch(
+        reference_fn=torch.matmul,
+        metric_tolerances={
+            Metric.MAX_ABS_ERROR: 1.5e-1,
+            Metric.PCC: 0.99,
+        },
+    )(_matmul)
+
+    _ = validated_matmul(a_tt, b_tt)
+
+    assert len(registry.results) == before + 1
+    assert registry.results[-1].passed
+
+
+def test_validation_raises_on_reference_exception(ttnn_mesh_device: ttnn.MeshDevice):
+    """When raise_exceptions=True, reference exceptions should propagate and not record results."""
+    registry = get_validation_registry()
+    before = len(registry.results)
+
+    a = torch.randn(1, 8, 8, dtype=torch.bfloat16)
+    b = torch.randn(1, 8, 8, dtype=torch.bfloat16)
+    a_tt = ttnn.from_torch(a.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+    b_tt = ttnn.from_torch(b.unsqueeze(0), device=ttnn_mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
+
+    def _ref_raises(a, b):
+        pass
+
+    # [INFO] make a mismatched signature on reference function to force the reference function to raise an exception!
+    @compare_to_torch(reference_fn=lambda a, b, c: _ref_raises(a, b), raise_exceptions=True)
+    def _matmul(a, b):
+        return ttnn.matmul(a, b)
+
+    with pytest.raises(TypeError) as e:
+        _ = _matmul(a_tt, b_tt)
+    assert "missing 1 required positional argument: 'c'" in str(e.value)
+    assert len(registry.results) == before + 1
+
+    # [INFO] make a mismatched signature on output_to_torch to force the reference function to raise an exception!
+    @compare_to_torch(reference_fn=lambda a, b: ..., output_to_torch=lambda x, y: ..., raise_exceptions=True)
+    def _matmul_too(a, b):
+        return ttnn.matmul(a, b)
+
+    with pytest.raises(TypeError) as e:
+        _ = _matmul_too(a_tt, b_tt)
+    assert "missing 1 required positional argument: 'y'" in str(e.value)
+    assert len(registry.results) == before + 2
+
+    # [INFO] make a mismatched signature on input_to_torch to force the reference function to raise an exception!
+    @compare_to_torch(reference_fn=lambda a, b: ..., input_to_torch=lambda x: ..., raise_exceptions=True)
+    def _matmul_three(a, b):
+        return ttnn.matmul(a, b)
+
+    with pytest.raises(TypeError) as e:
+        _ = _matmul_three(a_tt, b_tt)
+    assert "takes 1 positional argument but 2 were given" in str(e.value)
+    assert len(registry.results) == before + 3
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _print_validation_report_after_module(request):
+    # Runs once after all tests in this module finish
+    yield
+    registry = get_validation_registry()
+    reporter = request.config.pluginmanager.get_plugin("terminalreporter")
+    reporter.write_line("Printing validation report after yield")
+    registry.print_report(verbose=True)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _clear_validation_results_before_module():
+    clear_validation_results()
+
+
+# todo)) add an example where the reference is a file name that we can load output tensors from!
+# idea: validating against tensors -- give us a file name and we can compare our tensor under test against that!
+# - a filename for the tensor

code/models/common/tests/utils.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared helpers for models/common tests."""
+
+from __future__ import annotations
+
+import zlib
+
+
+def stable_model_seed(model_name: str) -> int:
+    """Stable 32-bit seed derived from model name.
+
+    Python's built-in hash is randomized per process, which breaks reproducibility
+    across runs and can mismatch on-disk cached weights. Use CRC32 instead.
+
+    NOTE: Avoid a single hardcoded global seed (e.g., 1234) for all models; a
+    per-model stable seed keeps caches distinct and reduces correlated RNG paths.
+    """
+    return zlib.crc32(model_name.encode("utf-8")) & 0xFFFFFFFF

code/models/common/utility_functions.py

@@ -0,0 +1,1216 @@
+# 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):
+    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
+    if torch.any(golden.bool()) != torch.any(calculated.bool()):
+        logger.error("One tensor is all zero")
+        return False, 0.0
+
+    # For now, mask all infs and nans 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:
+        # Check if dtype is FP8 - they don't support isinf/isneginf/masked_fill operations
+        is_fp8 = golden.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+
+        if is_fp8:
+            # Convert FP8 to float32 for comparison since FP8 doesn't support many operations
+            golden = golden.to(torch.float32)
+            calculated = calculated.to(torch.float32)
+
+        golden = golden.clone()
+        calculated = calculated.clone()
+
+        # Mask NaN and inf values
+        golden[
+            torch.logical_or(
+                torch.isnan(golden),
+                torch.logical_or(torch.isinf(golden), torch.isneginf(golden)),
+            )
+        ] = 0
+        calculated[
+            torch.logical_or(
+                torch.isnan(calculated),
+                torch.logical_or(torch.isinf(calculated), torch.isneginf(calculated)),
+            )
+        ] = 0
+
+    if torch.equal(golden, calculated):
+        return True, 1.0
+
+    if golden.dtype == torch.bfloat16:
+        golden = golden.type(torch.float32)
+        calculated = calculated.type(torch.float32)
+    cal_pcc = np.min(
+        np.ma.corrcoef(
+            np.ma.masked_invalid(torch.squeeze(golden).detach().numpy()).flatten(),
+            np.ma.masked_invalid(torch.squeeze(calculated).detach().numpy()).flatten(),
+        )
+    )
+
+    if isinstance(cal_pcc, np.ma.core.MaskedConstant):
+        return True, 1.0
+
+    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)} = |{calculated[ulp_index_tuple]} - {golden[ulp_index_tuple]}| / {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)
+        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_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):
+    assert stride_h == stride_w
+    assert filter_pyt_nchw_tensor.shape[2] == filter_pyt_nchw_tensor.shape[3]
+    # Fold activation for unity stride
+    # Pad channel size to 4. This is to make sure L1 read addresses are 16 bit aligned
+    C = _nearest_y(filter_pyt_nchw_tensor.shape[1], 4)
+    # 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

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/demos/gpt_oss/tt/ccl.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import ttnn
+
+
+class CCLManager:
+    def __init__(self, mesh_device, num_links, topology=ttnn.Topology.Ring):
+        self.mesh_device = mesh_device
+        self.num_links = num_links
+        self.topology = topology
+
+        # Cache for ping pong buffers: key = (shape_tuple, dim, mesh_axis), value = [buffer1, buffer2]
+        self._ping_pong_buffer_cache = {}
+        self._ping_pong_buffer_indices = {}
+
+        # Setup semaphores
+        self._init_subdevice()
+
+        # Initialize semaphores for reduce scatter and all gather
+        self._init_semaphores()
+        self.rs_ping_pong_idx = 0
+        self.ag_ping_pong_idx = 0
+        self.barrier_idx = 0
+
+    def _init_subdevice(self):
+        compute_grid_size = ttnn.CoreCoord(8, 8)
+        self.ccl_cores = ttnn.CoreRangeSet(
+            {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(compute_grid_size.x - 1, compute_grid_size.y - 1))}
+        )
+
+        _worker_sub_device = ttnn.SubDevice(
+            [
+                self.ccl_cores,
+            ]
+        )
+        self.ccl_sub_device_id = ttnn.SubDeviceId(0)
+
+    def _init_semaphores(self):
+        # Initialize semaphores for reduce scatter ping pong
+        rs_n_sems = 3 * 2  # 3 semaphores * 2 for ping pong
+        self.rs_ping_pong_semaphores = [
+            ttnn.create_global_semaphore(self.mesh_device, self.ccl_cores, 0) for _ in range(rs_n_sems)
+        ]
+
+        # Initialize semaphores for all gather ping pong
+        ag_n_sems = 2 * 2  # 2 semaphores * 2 for ping pong (2 buffers)
+        self.ag_ping_pong_semaphores = [
+            ttnn.create_global_semaphore(self.mesh_device, self.ccl_cores, 0) for _ in range(ag_n_sems)
+        ]
+
+        # Initialize barrier semaphores
+        barrier_ns_sems = 2 * 1
+        self.barrier_semaphore = [
+            ttnn.create_global_semaphore(self.mesh_device, self.ccl_cores, 0) for _ in range(barrier_ns_sems)
+        ]
+
+    def get_rs_ping_pong_semaphore(self):
+        """
+        Get semaphores for reduce scatter ping pong operations.
+
+        Returns:
+            List of 3 semaphores for the current ping pong cycle
+        """
+        cur_idx = self.rs_ping_pong_idx
+        n_sems = 3
+        self.rs_ping_pong_idx = (cur_idx + 1) % 2
+        return self.rs_ping_pong_semaphores[cur_idx * n_sems : (cur_idx + 1) * n_sems]
+
+    def get_ag_ping_pong_semaphore(self):
+        """
+        Get semaphores for all gather ping pong operations.
+
+        Returns:
+            List of 3 semaphores for the current ping pong cycle
+        """
+        cur_idx = self.ag_ping_pong_idx
+        n_sems = 2
+        self.ag_ping_pong_idx = (cur_idx + 1) % 2
+        return self.ag_ping_pong_semaphores[cur_idx * n_sems : (cur_idx + 1) * n_sems]
+
+    def get_barrier_semaphore(self):
+        """
+        Get semaphores for barrier operations.
+        """
+        cur_idx = self.barrier_idx
+        self.barrier_idx = (cur_idx + 1) % 2
+        return self.barrier_semaphore[cur_idx]
+
+    def reset_global_semaphores(self):
+        """Reset all global semaphores to 0"""
+        for sem in self.rs_ping_pong_semaphores:
+            ttnn.reset_global_semaphore_value(sem, 0)
+        for sem in self.ag_ping_pong_semaphores:
+            ttnn.reset_global_semaphore_value(sem, 0)

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:745c72341b2d4058b7a08d6e8ed95d0e493586d5b4d3b78d75fc973d5fe8db51","size":856,"annotations":{"io.containerd.image.name":"docker.io/tt-model/nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026:745c72341b2d","org.opencontainers.image.ref.name":"745c72341b2d"}}]}

image/manifest.json

@@ -0,0 +1 @@
+[{"Config":"blobs/sha256/f2b4df69449c2b17e203e5c5c45a7eafb9e17691b52baee2b55f054f2b69b0a3","RepoTags":["tt-model/nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026:745c72341b2d"],"Layers":["blobs/sha256/e51aee9c82ec5dd5ba2add49c45c6d85d460512757e2615b69bcdf9469c7cb58","blobs/sha256/f95188533c275c88f64d0ea2af69d9c5bd5f6e8ae6228a72f7ca1a474391490e","blobs/sha256/44958d2683cf8259e11d1fc5969eebacbd9182cba8c61a944d02ad6f9d0208f5","blobs/sha256/f4403345261d6a0a456e52ecb4bfcefafc8f60fe769f7a158c108d85f14a8442","blobs/sha256/0c3f1a4dac205920b50ede37e2a71089fc73c112353f1817ba049a573e3d85b0","blobs/sha256/7a4a65a66726b5cae2adac1ecf91f6ef3279e817aa14905130b096af0713543b","blobs/sha256/a0a726a404e7bde4fd5f62328276f46da44372697b7f29f475ab73c4f3981c02","blobs/sha256/1d138841c7fb11420920117c66dc923e100debfad7e2715f74e3b402c7a2b42f","blobs/sha256/2e0cfc89ff175fb9daec4a421637369dce62441303743f043c88fa8e4d5c84ca","blobs/sha256/3a03bff8c5a73a353bad8ae751e0703b43b10a63ccc7bb88c69f2559a04b3c25","blobs/sha256/45d1acbd1c6ba16132bd0ed960cdd7869ce831966b46cd7054c5c6db89f50d7c","blobs/sha256/633e13d0cbce3c46913bfa2b768a5b2d8485589423834ca7e36e06be81df2b5f","blobs/sha256/b7ff18b6e28d541b44e6b3859161b646143687aee783a5140bba7c82eb6f8dc3","blobs/sha256/87f9f907633999cc2166f3a11c05ef6c65bcdba7d728b5f46e54e8ae8e35a655","blobs/sha256/bc24046f95cf467afdbb275d4d7115fa2898918e49ef5a38b6f9d4e0369e7931","blobs/sha256/7de6f03cf8ad41347dfcd688ad0e2fd2b9f9bda48f10225e50418c2f2003c515","blobs/sha256/7f224de3c4cc07456c8a346f6d71cd8270921b772536691313ded51c82ce5b09","blobs/sha256/625235333d0c70632afb385283d5a06a49496d65f369b50a319656e103fe2281","blobs/sha256/4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1","blobs/sha256/66f86dec5bd03c37f55c9103d282115cfd2097437c2967808baa3532558a39bb","blobs/sha256/9ad78b76ad8377c5e1e7270d222f4ba7ebb2a4d8c6f8144fc03a36d317b2b852","blobs/sha256/c274f728a5846ebb5617e94318b4b25ed1df2e5aa7d63af334e0e4dd4c937424"]}]

image/oci-layout

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

requirements.lock

@@ -0,0 +1,215 @@
+agent-detector==2.0.0
+aiohappyeyeballs==2.7.1
+aiohttp==3.14.3
+aiosignal==1.4.0
+annotated-doc==0.0.5
+annotated-types==0.8.0
+anthropic==1.5.0
+anyio==4.15.1
+apache-tvm-ffi==0.1.10
+astor==0.8.1
+attrs==26.1.0
+blake3==1.0.9
+cachetools==7.1.8
+cbor2==6.1.4
+certifi==2026.7.22
+cffi==2.1.1
+cfgv==3.5.0
+charset-normalizer==3.5.1
+click==8.5.0
+cloudpickle==3.1.2
+compressed-tensors==0.17.0
+contourpy==1.3.3
+cryptography==50.0.1
+cuda-bindings==13.4.1
+cuda-core==1.2.0
+cuda-pathfinder==1.8.1
+cuda-python==13.4.1
+cuda-tile==1.6.0
+cycler==0.12.1
+depyf==0.20.0
+detect-installer==0.2.1
+dill==0.4.1
+distlib==0.4.3
+distro==1.9.0
+dnspython==2.8.0
+docstring_parser==0.18.0
+einops==0.8.2
+elastic-transport==9.4.2
+elasticsearch==9.5.1
+email-validator==2.3.0
+fastapi==0.136.3
+fastapi-cli==0.0.32
+fastapi-cloud-cli==0.26.0
+fastar==0.12.0
+fastsafetensors==0.4.0
+filelock==3.32.6
+flashinfer-python==0.6.14
+fonttools==4.65.0
+frozenlist==1.8.0
+fsspec==2026.7.0
+googleapis-common-protos==1.75.3
+graphviz==0.21
+grpcio==1.83.1
+h11==0.16.0
+hf-xet==1.6.0
+httpcore==1.0.9
+httpcore2==2.12.0
+httptools==0.8.0
+httpx==0.28.1
+httpx2==2.12.0
+huggingface_hub==1.16.1
+humming-kernels==0.1.10
+identify==2.6.19
+idna==3.19
+ijson==3.5.1
+interegular==0.3.3
+Jinja2==3.1.6
+jiter==0.16.0
+jmespath==1.1.0
+jsonschema==4.26.0
+jsonschema-specifications==2025.9.1
+kiwisolver==1.5.1
+lark==1.2.2
+linkify-it-py==2.2.0
+llguidance==1.7.6
+llvmlite==0.47.0
+lm-format-enforcer==0.11.3
+loguru==0.6.0
+markdown-it-py==4.2.0
+MarkupSafe==3.0.3
+matplotlib==3.11.1
+mcp==2.2.0
+mcp-types==2.2.0
+mdit-py-plugins==0.6.1
+mdurl==0.1.2
+mistral_common==1.11.7
+ml_dtypes==0.6.0
+model-hosting-container-standards==0.1.16
+mpmath==1.3.0
+msgspec==0.21.1
+multidict==6.8.0
+networkx==3.6.1
+ninja==1.13.2
+nodeenv==1.10.0
+numba==0.65.0
+numpy==2.3.5
+nvidia-cuda-cccl==13.3.4.2.1
+nvidia-cuda-crt==13.4.59
+nvidia-cuda-nvcc==13.4.59
+nvidia-cuda-nvdisasm==13.4.49
+nvidia-cuda-nvrtc==13.4.59
+nvidia-cuda-runtime==13.4.49
+nvidia-cudnn-frontend==1.28.0
+nvidia-cutlass-dsl==4.6.0
+nvidia-cutlass-dsl-libs-base==4.6.0
+nvidia-cutlass-dsl-libs-core==4.6.0
+nvidia-cutlass-dsl-libs-cu12==4.6.0
+nvidia-cutlass-dsl-libs-cu13==4.6.0
+nvidia-ml-py==13.610.43
+nvidia-nvvm==13.4.59
+nvtx==0.2.15
+openai==3.13.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
+outlines_core==0.2.14
+packaging==26.3
+pandas==3.0.5
+partial-json-parser==0.2.1.1.post7
+pillow==12.3.0
+platformdirs==4.11.8
+pre_commit==4.6.2
+prometheus-fastapi-instrumentator==8.1.0
+prometheus_client==0.26.0
+propcache==0.5.2
+protobuf==6.33.6
+psutil==7.2.2
+py-cpuinfo==9.0.0
+pybase64==1.5.0
+pycountry==26.2.16
+pycparser==3.0
+pydantic==2.13.5
+pydantic-extra-types==2.11.1
+pydantic-settings==2.15.0
+pydantic_core==2.46.5
+pyelftools==0.33
+Pygments==2.21.0
+PyJWT==2.14.0
+pyluwen==0.9.0
+pynvvideocodec==2.0.4
+pyparsing==3.3.2
+python-dateutil==2.9.0.post0
+python-discovery==1.6.0
+python-dotenv==1.2.3
+python-json-logger==4.2.0
+python-multipart==0.0.32
+PyYAML==6.0.3
+pyzmq==27.2.0
+quack-kernels==0.6.3
+referencing==0.37.0
+regex==2026.9.10
+requests==2.34.2
+rich==15.0.0
+rich-toolkit==0.20.5
+rignore==0.8.1
+rpds-py==2026.6.3
+safetensors==0.8.0
+seaborn==0.13.2
+sentencepiece==0.2.2
+sentry-sdk==2.69.1
+setproctitle==1.3.7
+setuptools==80.10.2
+setuptools-scm==8.1.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+sse-starlette==3.4.11
+starlette==1.6.0
+supervisor==4.3.0
+sympy==1.14.0
+tabulate==0.10.0
+tblib==3.2.2
+textual==8.2.8
+tiktoken==0.14.0
+tilelang==0.1.9
+tokenizers==0.22.2
+tokenspeed-mla==0.1.8
+tokenspeed-triton==3.8.10.post20260906
+tomli==2.4.1
+torch==2.11.0+cpu
+torch_c_dlpack_ext==0.1.5
+torchcodec==0.16.0+cpu
+torchvision==0.26.0+cpu
+tqdm==4.70.1
+transformers==5.10.2
+triton==3.8.0
+truststore==0.10.4
+tt-smi==6.5.0
+tt-tools-common==1.6.0
+tt-umd==0.9.10
+ttnn==0.65.2.dev7097+gcfd2056ecb4
+ttnn==0.65.1rc17.dev7096+ggcfd2056ecb4
+typer==0.27.2
+typing-inspection==0.4.4
+typing_extensions==4.16.0
+urllib3==2.7.0
+uvicorn==0.52.4
+uvloop==0.22.1
+virtualenv==21.7.9
+vllm==0.26.0+empty
+vllm-tt-plugin==0.1.0
+watchfiles==1.2.0
+websockets==17.1
+wheel==0.48.0
+xgrammar==0.2.6
+yarl==1.24.5
+z3-solver==4.15.4.0

tt_kernel_manifest.json

@@ -0,0 +1,150 @@
+{
+  "schema_version": "5.1",
+  "name": "nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026",
+  "tt_metal_version": "0.65.2.dev7097+gcfd2056ecb4",
+  "arch": "blackhole",
+  "device_count": 4,
+  "producer": {
+    "tt_kernel_version": "0.1.0",
+    "created_at": "2026-09-11T17:52:26.396571+00:00",
+    "hostname": "tt-quietbox"
+  },
+  "weights": {
+    "repo_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4",
+    "revision": "6dbbd757ea75a8ece6e0702872e3ae53f9987728",
+    "allow_patterns": null,
+    "ignore_patterns": null,
+    "repo_type": "model"
+  },
+  "mesh": null,
+  "entrypoint": null,
+  "resources": null,
+  "capabilities": null,
+  "env": {},
+  "bundled": null,
+  "deps": null,
+  "container": {
+    "image": {
+      "registry": "hf",
+      "repository": "nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026",
+      "tag": "tt-model/nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026:745c72341b2d",
+      "digest": "sha256:745c72341b2d4058b7a08d6e8ed95d0e493586d5b4d3b78d75fc973d5fe8db51"
+    },
+    "kind": "vllm-plugin",
+    "runtime": {
+      "vllm": {
+        "version": "0.26.0"
+      },
+      "plugin": {
+        "path": "/home/ttuser/nemo35lightning/vllm-tt-plugin",
+        "sha": "52db64380a1564c2c16959b55c2cd7378ecbe9f0",
+        "dirty": false
+      },
+      "extra_models_dir": "models/autoports/nvidia_nvidia_nemotron_3_5_lightning_30b_a3b_nvfp4/vllm_bundle",
+      "wheels": [
+        "/home/ttuser/nemo35lightning/wheels_pins/torch-2.11.0+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/transformers-5.10.2-py3-none-any.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/huggingface_hub-1.16.1-py3-none-any.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/numpy-2.3.5-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl",
+        "/home/ttuser/nemo35lightning/wheels_pins/loguru-0.6.0-py3-none-any.whl"
+      ],
+      "lock": "requirements.lock"
+    },
+    "serve": {
+      "hardware": "p300x2",
+      "mesh_device": "P300x2",
+      "port": null,
+      "max_model_len": 1048576,
+      "max_num_seqs": 32,
+      "block_size": 64,
+      "server_timeout": null,
+      "capabilities": {
+        "tool_parser": "qwen3_coder",
+        "reasoning_parser": "nemotron_v3"
+      },
+      "additional_config": {
+        "tt": {
+          "sample_on_device_mode": "all",
+          "trace_region_size": 1073741824,
+          "fabric_config": "FABRIC_1D_RING"
+        }
+      },
+      "args": [
+        [
+          "--hf-overrides",
+          "{\"quantization_config\":null}"
+        ],
+        [
+          "--revision",
+          "6dbbd757ea75a8ece6e0702872e3ae53f9987728"
+        ],
+        [
+          "--tokenizer-revision",
+          "6dbbd757ea75a8ece6e0702872e3ae53f9987728"
+        ]
+      ],
+      "env": {
+        "TT_METAL_VISIBLE_DEVICES": "0,1,2,3",
+        "ARCH_NAME": "blackhole",
+        "TORCHDYNAMO_DISABLE": "1",
+        "VLLM_RPC_TIMEOUT": "900000",
+        "NEMOTRON_WEIGHT_CACHE": "/cache/nemotron_weights",
+        "NEMOTRON_REVISION": "6dbbd757ea75a8ece6e0702872e3ae53f9987728"
+      }
+    },
+    "serve_profiles": [
+      {
+        "hardware": null,
+        "mesh_device": null,
+        "port": null,
+        "max_model_len": null,
+        "max_num_seqs": null,
+        "block_size": null,
+        "server_timeout": null,
+        "capabilities": null,
+        "additional_config": {},
+        "args": [],
+        "env": {},
+        "name": "default",
+        "description": null
+      }
+    ],
+    "default_profile": null,
+    "code_dir": "code",
+    "verify": [
+      "import vllm; assert vllm.__version__.startswith('0.26.'), vllm.__version__",
+      "from vllm.reasoning import ReasoningParserManager as R; R.get_reasoning_parser('nemotron_v3')",
+      "from vllm.tool_parsers import ToolParserManager as T; T.get_tool_parser('qwen3_coder')",
+      "import vllm_tt_plugin, pathlib; assert '/tt-metal/' not in vllm_tt_plugin.__file__ and 'site-packages' in vllm_tt_plugin.__file__, vllm_tt_plugin.__file__",
+      "from models.autoports.nvidia_nvidia_nemotron_3_5_lightning_30b_a3b_nvfp4.tt.generator_vllm import NemotronHForCausalLM; assert NemotronHForCausalLM",
+      "from pathlib import Path; p = Path('/opt/tt-metal/models/autoports/nvidia_nvidia_nemotron_3_5_lightning_30b_a3b_nvfp4/doc/datatype_sweep/selected_precision_config.json'); assert p.is_file(), 'selected_precision_config.json missing: the model would silently serve at code-default precision'",
+      "from models.autoports.nvidia_nvidia_nemotron_3_5_lightning_30b_a3b_nvfp4.tt import precision; r = precision.resolve_policies(); assert r[2]['config_id'] != 'code-defaults', r[2]",
+      "import transformers, tokenizers, safetensors, numpy, torch; assert transformers.__version__ == '5.10.2', transformers.__version__; assert numpy.__version__ == '2.3.5', numpy.__version__; assert torch.__version__ == '2.11.0+cpu', torch.__version__"
+    ],
+    "built": {
+      "image": "tt-model/nemotron-3.5-lightning-30b-a3b-nvfp4-tt-p300x2-vllm026:745c72341b2d",
+      "repo": "ndaly/Nemotron-3.5-Lightning-30B-A3B-NVFP4-tt-p300x2-vllm026",
+      "tt_model_version": "0.1.0",
+      "created_at": "2026-09-11T17:47:54+00:00",
+      "tt_metal": {
+        "sha": "cfd2056ecb4601bd8b136b4ef1f89c6329f2eaef",
+        "describe": "v0.74.0-dev20260622-202-gcfd2056ecb-dirty",
+        "dirty": true,
+        "scm_version": "0.65.2.dev7097+gcfd2056ecb4",
+        "mode": "local",
+        "remote": "https://github.com/tenstorrent/tt-metal.git",
+        "branch": "agentic-research/hous/multigoal-claude",
+        "pushed": false
+      },
+      "code_sha256": "5787b8dfce9d7e38e43fdb6e6bc7cbfcb93b48d48719c9e825261a5624c2d151",
+      "plugin": {
+        "sha": "52db64380a1564c2c16959b55c2cd7378ecbe9f0",
+        "path": "/home/ttuser/nemo35lightning/vllm-tt-plugin",
+        "dirty": false
+      },
+      "image_digest": "sha256:745c72341b2d4058b7a08d6e8ed95d0e493586d5b4d3b78d75fc973d5fe8db51"
+    }
+  }
+}