Blame for code/models/common/tests/modules/attention/test_attention_1d.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Tests for the Attention1D module (1D mesh topology: N150, N300, T3K).

This test suite verifies:
1. Unit tests for config dataclasses (no device needed)
2. Attention1D class matches HuggingFace/Meta reference model
3. Attention1D correctly rejects TG/Galaxy devices
4. Sliding window attention works correctly (seq_len > window_size)

Test coverage notes:
- Paged attention: Tested via (page_block_size, chunk_size) parameter combinations.
- Chunked prefill: Tested via paged-chunked variant. Requires paged=True and mode="prefill".
- Variants: non-paged, paged, paged-chunked (3 combinations per test case).
"""

import inspect
import os
import time
from dataclasses import dataclass, replace
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest
import torch
from loguru import logger
from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM

from models.common.utility_functions import hf_cache_layer_kv, hf_cache_num_layers

# transformers 5.x moved no_init_weights to transformers.initialization; fall back
# to the old location for transformers < 5.x.
try:
    from transformers.initialization import no_init_weights
except ImportError:
    from transformers.modeling_utils import no_init_weights

import ttnn
from models.common.auto_compose import to_torch_auto_compose
from models.common.modules.attention import attention_1d as attention_1d_module
from models.common.modules.attention.attention_1d import Attention1D, Attention1DConfig, _resolve_attention1d_config
from models.common.modules.lazy_weight import LazyWeight
from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1DConfig
from models.common.modules.tt_ccl import TT_CCL
from models.common.tensor_utils import (
    get_rot_transformation_mat,
    nearest_32,
    zeros_like_kv_cache,
    zeros_like_paged_cache,
)
from models.common.tests.utils import stable_model_seed
from models.common.utility_functions import comp_allclose, comp_pcc

# 1D module suites target the T3K; skip when the host system is a Galaxy.
pytestmark = pytest.mark.usefixtures("skip_on_galaxy_system")

# =============================================================================
# RoPE Helper Functions (replaces TTTv1 rope imports)
# =============================================================================


def _permute_to_meta_format(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Convert HuggingFace RoPE format to Meta format for TTNN compatibility.

    HF stores cos/sin with shape [batch, seq_len, head_dim] where head_dim is interleaved.
    Meta format expects [1, 1, seq_len, head_dim] with doubled values.
    """
    # Handle different HF output shapes
    if len(cos.shape) == 3:
        cos = cos.squeeze(0)  # [seq_len, head_dim]
        sin = sin.squeeze(0)

    # Undo the HF permute: take first half and duplicate
    cos = cos[:, : cos.shape[1] // 2]
    cos = torch.stack((cos, cos), dim=-1).flatten(-2)

    sin = sin[:, : sin.shape[1] // 2]
    sin = torch.stack((sin, sin), dim=-1).flatten(-2)

    # Add batch dimensions: [1, 1, seq_len, head_dim]
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)

    return cos, sin


def get_cos_sin_from_hf(
    rotary_emb,
    seq_len: int,
    head_dim: int,
    dtype: torch.dtype = torch.bfloat16,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Extract cos/sin rotation matrices from HuggingFace rotary_emb module.

    Args:
        rotary_emb: HuggingFace rotary embedding module (e.g., LlamaRotaryEmbedding)
        seq_len: Maximum sequence length
        head_dim: Head dimension
        dtype: Output dtype

    Returns:
        (cos, sin) tensors in Meta format with shape [1, 1, seq_len, head_dim]
    """
    # Create dummy input for HF's rotary_emb forward
    x = torch.zeros(1, 1, seq_len, head_dim, dtype=dtype)
    position_ids = torch.arange(seq_len).unsqueeze(0)

    # HF rotary_emb.forward() returns (cos, sin)
    with torch.no_grad():
        cos_hf, sin_hf = rotary_emb(x, position_ids)

    # Convert to Meta format
    cos_meta, sin_meta = _permute_to_meta_format(cos_hf.float(), sin_hf.float())

    return cos_meta.to(dtype), sin_meta.to(dtype)


def get_rot_mats_from_hf(
    rotary_emb,
    seq_len: int,
    head_dim: int,
    device: ttnn.MeshDevice,
    dtype: ttnn.DataType = ttnn.bfloat16,
) -> list[ttnn.Tensor]:
    """
    Create TTNN rotation matrices from HuggingFace rotary_emb.

    Replaces `get_rot_mats` from models.tt_transformers.tt.rope.
    """
    cos_meta, sin_meta = get_cos_sin_from_hf(rotary_emb, seq_len * 2, head_dim)

    cos_tt = ttnn.from_torch(
        cos_meta,
        device=device,
        layout=ttnn.TILE_LAYOUT,
        dtype=dtype,
        mesh_mapper=ttnn.ReplicateTensorToMesh(device),
    )
    sin_tt = ttnn.from_torch(
        sin_meta,
        device=device,
        layout=ttnn.TILE_LAYOUT,
        dtype=dtype,
        mesh_mapper=ttnn.ReplicateTensorToMesh(device),
    )

    return [cos_tt, sin_tt]


class RotarySetupHelper:
    """
    Simplified RotarySetup for testing - extracts rotation matrices from HuggingFace's
    rotary_emb instead of computing from scratch. Replaces TTTv1's RotarySetup class.
    """

    def __init__(
        self,
        device: ttnn.MeshDevice,
        batch_size: int,
        head_dim: int,
        max_seq_len: int,
        rotary_emb,  # HuggingFace rotary embedding module
        use_qk_fused: bool = False,
        datatype: ttnn.DataType = ttnn.bfloat16,
    ):
        self.device = device
        self.head_dim = head_dim
        self.use_qk_fused = use_qk_fused
        self.batch_size = batch_size
        self.doubled_batch_size = batch_size * 2 if use_qk_fused else batch_size

        is_mesh = isinstance(device, ttnn.MeshDevice)
        num_devices = device.get_num_devices() if is_mesh else 1

        if num_devices == 32:
            self.batch_size_per_device_group = max(self.doubled_batch_size // device.shape[1], 1)
        else:
            self.batch_size_per_device_group = self.doubled_batch_size

        self.core_grid = device.compute_with_storage_grid_size()
        self.batch_grid = ttnn.num_cores_to_corerangeset(self.doubled_batch_size, self.core_grid, row_wise=True)

        # Get cos/sin from HuggingFace rotary_emb
        self.cos_matrix, self.sin_matrix = get_rot_mats_from_hf(rotary_emb, max_seq_len, head_dim, device, datatype)

        # Create transformation matrices
        trans_mat = get_rot_transformation_mat().repeat(1, 1, self.doubled_batch_size, 1)
        trans_mat_mem_config = ttnn.create_sharded_memory_config(
            shape=(ttnn.TILE_SIZE, ttnn.TILE_SIZE),
            core_grid=self.batch_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )
        self.transformation_mat = ttnn.from_torch(
            trans_mat,
            device=device,
            layout=ttnn.TILE_LAYOUT,
            dtype=datatype,
            memory_config=trans_mat_mem_config,
            mesh_mapper=(
                ttnn.ShardTensor2dMesh(
                    device,
                    dims=(None, 2) if (num_devices == 32 and batch_size > 1) else (None, None),
                    mesh_shape=list(device.shape),
                )
                if is_mesh
                else None
            ),
        )

        # Prefill transformation matrix
        prefill_trans_mat = get_rot_transformation_mat()
        if head_dim != ttnn.TILE_SIZE:
            prefill_trans_mat = torch.zeros(1, 1, head_dim, head_dim)
            base_mat = get_rot_transformation_mat()
            prefill_trans_mat[:, :, : ttnn.TILE_SIZE, : ttnn.TILE_SIZE] = base_mat

        self.transformation_mat_prefill = ttnn.from_torch(
            prefill_trans_mat,
            device=device,
            layout=ttnn.TILE_LAYOUT,
            dtype=datatype,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ReplicateTensorToMesh(device) if is_mesh else None,
        )

    def get_both_trans_mats(self) -> dict[str, ttnn.Tensor]:
        """Return transformation matrices for decode and prefill."""
        return {"decode": self.transformation_mat, "prefill": self.transformation_mat_prefill}

    def get_rot_idxs(self, position_idxs: torch.Tensor, on_host: bool = False) -> ttnn.Tensor:
        """Convert position indices to TTNN tensor."""

        if self.use_qk_fused:
            position_idxs = position_idxs.repeat(2)

        batch = position_idxs.shape[0]
        position_idxs = position_idxs.reshape(1, batch)

        # Pad to tile boundary
        pad_size = nearest_32(batch) - batch
        position_idxs = torch.nn.functional.pad(position_idxs, (0, pad_size), "constant", 0)

        rot_idxs = ttnn.as_tensor(
            position_idxs,
            dtype=ttnn.uint32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            device=None if on_host else self.device,
            memory_config=None if on_host else ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.device),
        )
        return rot_idxs

    def get_rot_mats(self, position_idxs: torch.Tensor) -> list[ttnn.Tensor]:
        """Get rotation matrices for given position indices."""
        rot_idxs = self.get_rot_idxs(position_idxs)

        if rot_idxs.device != self.device:
            rot_idxs = ttnn.to_device(rot_idxs, self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG)

        cos = ttnn.embedding(rot_idxs, self.cos_matrix, layout=ttnn.TILE_LAYOUT)
        sin = ttnn.embedding(rot_idxs, self.sin_matrix, layout=ttnn.TILE_LAYOUT)

        cos = ttnn.unsqueeze_to_4D(cos)
        sin = ttnn.unsqueeze_to_4D(sin)

        cos = ttnn.transpose(cos, 1, 2)
        sin = ttnn.transpose(sin, 1, 2)

        if self.batch_size_per_device_group % ttnn.TILE_SIZE != 0:
            cos = cos[:, : self.batch_size_per_device_group, :, :]
            sin = sin[:, : self.batch_size_per_device_group, :, :]

        mem_config = ttnn.create_sharded_memory_config(
            shape=(ttnn.TILE_SIZE, self.head_dim),
            core_grid=self.batch_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )

        cos = ttnn.interleaved_to_sharded(cos, mem_config)
        sin = ttnn.interleaved_to_sharded(sin, mem_config)

        return [cos, sin]


# =============================================================================
# HfAttentionWrapper (replaces TTTv1 model_config import)
# =============================================================================


class HfAttentionWrapper:
    """
    Wrapper for HuggingFace attention modules with KV cache support.
    Provides a consistent interface for running HF attention as reference.
    """

    def __init__(self, attention, head_dim: int, rotary_emb):
        from transformers import DynamicCache

        self.attention = attention
        self.past_key_value = DynamicCache()
        self.head_dim = head_dim
        self.rotary_emb = rotary_emb
        self._uses_past_key_values = "past_key_values" in inspect.signature(attention.forward).parameters

    def forward(self, x: torch.Tensor, start_pos: int, mask=None):
        """Run attention forward pass using rotary_emb directly."""
        position_ids = torch.tensor([list(range(start_pos, start_pos + x.shape[1]))] * x.shape[0])

        if mask is not None:
            while len(mask.shape) < 4:
                mask = mask.unsqueeze(0)

        if self.rotary_emb is not None:
            position_embeddings = self.rotary_emb(x, position_ids)
            cache_kwargs = (
                {"past_key_values": self.past_key_value}
                if self._uses_past_key_values
                else {"past_key_value": self.past_key_value, "use_cache": True}
            )
            output, *_ = self.attention(x, position_embeddings=position_embeddings, attention_mask=mask, **cache_kwargs)
        else:
            cache_kwargs = (
                {"past_key_values": self.past_key_value}
                if self._uses_past_key_values
                else {"past_key_value": self.past_key_value, "use_cache": True}
            )
            outputs = self.attention(x, position_ids=position_ids, attention_mask=mask, **cache_kwargs)
            output = outputs[0]
            if not self._uses_past_key_values and len(outputs) > 2:
                self.past_key_value = outputs[2]
        return output

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def reset_cache(self):
        """Reset KV cache for new sequence."""
        from transformers import DynamicCache

        self.past_key_value = DynamicCache()

    @property
    def cache_k(self) -> torch.Tensor:
        """Get key cache in shape [batch, seq_len, n_kv_heads, head_dim]."""
        if hf_cache_num_layers(self.past_key_value) == 0:
            return torch.zeros(0)
        # DynamicCache stores as [batch, n_heads, seq_len, head_dim]
        # Transpose to [batch, seq_len, n_heads, head_dim]
        return hf_cache_layer_kv(self.past_key_value, 0)[0].transpose(1, 2)

    @property
    def cache_v(self) -> torch.Tensor:
        """Get value cache in shape [batch, seq_len, n_kv_heads, head_dim]."""
        if hf_cache_num_layers(self.past_key_value) == 0:
            return torch.zeros(0)
        # DynamicCache stores as [batch, n_heads, seq_len, head_dim]
        # Transpose to [batch, seq_len, n_heads, head_dim]
        return hf_cache_layer_kv(self.past_key_value, 0)[1].transpose(1, 2)


# =============================================================================
# PagedAttentionConfig (replaces TTTv1 common import)
# =============================================================================


@dataclass
class PagedAttentionConfig:
    """Configuration for paged attention."""

    block_size: int = 64
    max_num_blocks: int = 2048


# =============================================================================
# Weight extraction helpers
# =============================================================================


def _reverse_permute(tensor, n_heads, dim1, dim2):
    """Convert HuggingFace Q/K weights to Meta format for RoPE compatibility.

    HuggingFace stores Q/K weights in a format optimized for their attention implementation,
    while Meta format is required for TTNN's RoPE implementation.
    """
    return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2)


def _reverse_permute_1d(tensor):
    """Convert the last dim from separate real/imaginary (r1,r2,i1,i2,...) to interleaved (r1,i1,r2,i2,...)"""
    shape = tensor.shape
    dim = shape[-1]
    assert dim % 2 == 0, "Last dimension must be even"
    reals = tensor[..., : dim // 2]
    imags = tensor[..., dim // 2 :]
    interleaved = torch.stack((reals, imags), dim=-1).flatten(start_dim=len(shape) - 1)
    return interleaved


def get_attention_weights_from_ref_model(
    reference_attn, num_devices: int = 1
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
    """
    Extract attention weights from a reference attention module in TTNN layout.

    Applies reverse_permute to Q and K weights to convert from HuggingFace format
    to Meta format, which is required for TTNN's RoPE implementation.

    Returns:
        (wqkv, wo, q_norm, k_norm, wqkv_bias) tensors in TTNN layout
    """
    # Phi-3 / Phi-4 ship a FUSED qkv_proj (single Linear) instead of separate q/k/v projections.
    # Split it into Q/K/V rows here so the rest of the pipeline is architecture-agnostic.
    fused_qkv = hasattr(reference_attn, "qkv_proj") and not hasattr(reference_attn, "q_proj")
    if fused_qkv:
        cfg = reference_attn.config
        _n_heads = cfg.num_attention_heads
        _n_kv = getattr(cfg, "num_key_value_heads", _n_heads)
        _hd = (
            getattr(reference_attn, "head_dim", None) or getattr(cfg, "head_dim", None) or (cfg.hidden_size // _n_heads)
        )
        _q = _n_heads * _hd
        _kv = _n_kv * _hd
        qkv_w = reference_attn.qkv_proj.weight  # (n_heads*hd + 2*n_kv*hd, dim), order Q|K|V
        wq_raw = qkv_w[:_q]  # (n_heads * head_dim, dim)
        wk_raw = qkv_w[_q : _q + _kv]  # (n_kv_heads * head_dim, dim)
        wv_raw = qkv_w[_q + _kv : _q + 2 * _kv]  # (n_kv_heads * head_dim, dim)
        wo_raw = reference_attn.o_proj.weight  # (dim, n_heads * head_dim)
    else:
        # Get raw weights from HF module
        wq_raw = reference_attn.q_proj.weight  # (n_heads * head_dim, dim)
        wk_raw = reference_attn.k_proj.weight  # (n_kv_heads * head_dim, dim)
        wv_raw = reference_attn.v_proj.weight  # (n_kv_heads * head_dim, dim)
        wo_raw = reference_attn.o_proj.weight  # (dim, n_heads * head_dim)

    # Compute head_dim from weight shapes
    dim = wq_raw.shape[1]
    n_heads_times_head_dim = wq_raw.shape[0]
    n_kv_heads_times_head_dim = wk_raw.shape[0]

    # For head_dim calculation, we need n_heads. Use the ratio of Q/K sizes.
    # Q: (n_heads * head_dim, dim), K: (n_kv_heads * head_dim, dim)
    # If n_heads == n_kv_heads (no GQA), just use q shape
    # Otherwise, we need to infer from config or assume head_dim from common values
    if hasattr(reference_attn, "head_dim"):
        head_dim = reference_attn.head_dim
    elif hasattr(reference_attn, "config") and hasattr(reference_attn.config, "head_dim"):
        head_dim = reference_attn.config.head_dim
    else:
        # Common head_dim values for LLaMA models
        head_dim = 128 if n_heads_times_head_dim >= 4096 else 64

    n_heads = n_heads_times_head_dim // head_dim
    n_kv_heads = n_kv_heads_times_head_dim // head_dim

    # Apply reverse_permute to convert HF format to Meta format for RoPE compatibility
    # This transformation is critical for Q and K weights
    wq_meta = _reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim)
    wk_meta = _reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim)
    # V and O don't need permutation
    wv_meta = wv_raw
    wo_meta = wo_raw

    # Transpose to TTNN layout: (dim, out_features)
    wq = wq_meta.T  # (dim, n_heads * head_dim)
    wk = wk_meta.T  # (dim, n_kv_heads * head_dim)
    wv = wv_meta.T  # (dim, n_kv_heads * head_dim)
    wo = wo_meta.T  # (n_heads * head_dim, dim)

    # Build combined QKV weight
    # Shape: (1, 1, dim, qkv_size_per_device * num_devices)
    qkv_list = []
    for i in range(num_devices):
        wq_chunk = torch.chunk(wq, num_devices, dim=1)[i]
        wk_chunk = torch.chunk(wk, num_devices, dim=1)[i]
        wv_chunk = torch.chunk(wv, num_devices, dim=1)[i]
        qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1)
        qkv_list.append(qkv)

    wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0)

    # WO weight: (1, 1, n_heads * head_dim, dim)
    wo = wo.unsqueeze(0).unsqueeze(0)

    # Q/K norm weights (optional, e.g., for Qwen models)
    # These also need reverse_permute_1d transformation
    q_norm = None
    k_norm = None
    if hasattr(reference_attn, "q_norm") and reference_attn.q_norm is not None:
        q_norm = _reverse_permute_1d(reference_attn.q_norm.weight)
    if hasattr(reference_attn, "k_norm") and reference_attn.k_norm is not None:
        k_norm = _reverse_permute_1d(reference_attn.k_norm.weight)

    # QKV bias (optional, e.g., for Qwen2/Qwen2.5 models)
    # Bias also needs the same chunking/concat pattern as weights
    wqkv_bias = None
    if not fused_qkv and hasattr(reference_attn.q_proj, "bias") and reference_attn.q_proj.bias is not None:
        bq_raw = reference_attn.q_proj.bias  # (n_heads * head_dim,)
        bk_raw = reference_attn.k_proj.bias  # (n_kv_heads * head_dim,)
        bv_raw = reference_attn.v_proj.bias  # (n_kv_heads * head_dim,)

        # Apply reverse_permute to Q and K biases (same as weights)
        bq_meta = _reverse_permute_1d(bq_raw.view(n_heads, head_dim)).view(-1)
        bk_meta = _reverse_permute_1d(bk_raw.view(n_kv_heads, head_dim)).view(-1)
        bv_meta = bv_raw  # V doesn't need permutation

        # Build combined QKV bias with chunking for multi-device
        qkv_bias_list = []
        for i in range(num_devices):
            bq_chunk = torch.chunk(bq_meta, num_devices, dim=0)[i]
            bk_chunk = torch.chunk(bk_meta, num_devices, dim=0)[i]
            bv_chunk = torch.chunk(bv_meta, num_devices, dim=0)[i]
            qkv_bias = torch.cat([bq_chunk, bk_chunk, bv_chunk], dim=-1)
            qkv_bias_list.append(qkv_bias)

        wqkv_bias = torch.cat(qkv_bias_list, dim=-1)

    return wqkv, wo, q_norm, k_norm, wqkv_bias


# ============================================================================
# Weight Caching - Avoid expensive torch.randn_like() per test
# ============================================================================

_CACHED_ATTN_WEIGHTS: dict[str, dict[str, torch.Tensor]] = {}


def _init_weight_scaled_normal(param: torch.Tensor, name: str) -> torch.Tensor:
    """Initialize a weight tensor using scaled normal distribution.

    Uses a scaling factor based on typical pretrained LLM weight statistics
    (std ~0.02) rather than pure randn (std=1.0) which can cause
    extreme activations and numerical issues in softmax.

    For 2D weights (linear layers): Uses normal with std = 0.02
    For 1D weights (biases, norms): Uses appropriate initialization

    NOTE: Uses torch.randn() instead of torch.randn_like() to ensure
    the global RNG state (set via torch.manual_seed) is respected.
    torch.randn_like() may not use the global RNG consistently.
    """
    if param.dim() >= 2:
        # Linear layer weights: scaled normal initialization
        # std=0.02 matches typical pretrained transformer weights
        return torch.randn(param.shape, dtype=param.dtype, device=param.device) * 0.02
    elif "norm" in name.lower() or "weight" in name.lower():
        # Norm weights (e.g., q_norm.weight, k_norm.weight): use ones
        return torch.ones_like(param)
    else:
        # Biases: use small random values
        return torch.randn(param.shape, dtype=param.dtype, device=param.device) * 0.01


def _get_or_init_attn_weights(model_name: str, reference_attn) -> None:
    """Initialize attention weights once per model, cache and reuse across tests.

    Uses scaled normal initialization (std=0.02) for better numerical conditioning
    compared to pure random noise (std=1.0). This helps maintain reasonable
    activation magnitudes through the attention computation.

    NOTE: Uses a deterministic seed per model to ensure reproducible weights
    regardless of test execution order or Python process hash randomization.
    This prevents flaky tests where PCC varies based on which tests ran first.
    """
    if model_name not in _CACHED_ATTN_WEIGHTS:
        logger.info(f"\033[33m[cache miss]\033[0m Initializing weights for {model_name}")
        _CACHED_ATTN_WEIGHTS[model_name] = {}
        # Use deterministic seed based on model name to ensure reproducible weights
        # regardless of test execution order
        seed = stable_model_seed(model_name)
        rng_state = torch.get_rng_state()
        torch.manual_seed(seed)
        with torch.no_grad():
            for name, param in reference_attn.named_parameters():
                _CACHED_ATTN_WEIGHTS[model_name][name] = _init_weight_scaled_normal(param, name)
        torch.set_rng_state(rng_state)  # Restore original RNG state
    else:
        logger.info(f"\033[32m[cache hit]\033[0m Reusing cached weights for {model_name}")

    # Load cached weights into model
    with torch.no_grad():
        for name, param in reference_attn.named_parameters():
            param.copy_(_CACHED_ATTN_WEIGHTS[model_name][name])


# ============================================================================
# Unit Tests - No device required
# ============================================================================


def test_attention_1d_config_creation():
    """Test that Attention1DConfig dataclass can be created with explicit values."""
    mock_mesh_device = MagicMock()
    mock_tt_ccl = MagicMock()
    mock_wqkv = MagicMock()
    mock_wo = MagicMock()

    # Create config with explicit values
    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=mock_wo,
        mesh_device=mock_mesh_device,
        tt_ccl=mock_tt_ccl,
        dim=4096,
        n_heads=32,
        n_kv_heads=8,
        head_dim=128,
        max_batch_size=64,
        topology=ttnn.Topology.Ring,
    )

    # Verify explicit values are preserved
    assert config.wqkv is mock_wqkv
    assert config.wo is mock_wo
    assert config.mesh_device is mock_mesh_device
    assert config.tt_ccl is mock_tt_ccl
    assert config.dim == 4096
    assert config.n_heads == 32
    assert config.n_kv_heads == 8
    assert config.head_dim == 128
    assert config.max_batch_size == 64
    assert config.topology == ttnn.Topology.Ring


def test_attention_1d_config_defaults():
    """Test that Attention1DConfig has sensible defaults."""
    # Minimal creation - only required fields
    config = Attention1DConfig(wqkv=MagicMock(), wo=MagicMock())

    # Check defaults
    assert config.max_batch_size == 32
    assert config.max_seq_len == 128 * 1024
    assert config.use_vllm_paged_kv_cache is False
    assert config.kv_cache_dtype == ttnn.bfloat8_b
    assert config.use_qk_fused is False
    assert config.num_reduce_scatter_links is None
    assert config.num_all_gather_links is None

    # Optional fields default to None
    assert config.mesh_device is None
    assert config.tt_ccl is None
    assert config.dim is None
    assert config.n_heads is None
    assert config.n_kv_heads is None
    assert config.q_norm_config is None
    assert config.k_norm_config is None


def test_attention_1d_config_power_user_overrides():
    """Test that Attention1DConfig accepts power-user overrides for program configs."""
    mock_prg_config = MagicMock()
    mock_mem_config = MagicMock()

    config = Attention1DConfig(
        wqkv=MagicMock(),
        wo=MagicMock(),
        decode_xqkv_prg_config=mock_prg_config,
        decode_sdpa_prg_config=mock_prg_config,
        decode_attn_output_prg_config=mock_prg_config,
        decode_residual_memcfg=mock_mem_config,
        activation_dtype=ttnn.bfloat16,
        scale=0.08838834764831843,  # 1/sqrt(128)
    )

    # User-provided overrides should be preserved
    assert config.decode_xqkv_prg_config is mock_prg_config
    assert config.decode_sdpa_prg_config is mock_prg_config
    assert config.decode_attn_output_prg_config is mock_prg_config
    assert config.decode_residual_memcfg is mock_mem_config
    assert config.activation_dtype == ttnn.bfloat16
    assert config.scale == pytest.approx(0.08838834764831843)


def test_attention_1d_happy_path_signature():
    """Test that Attention1D.__init__ accepts the happy path signature (weights + dimensions).

    Note: This is a unit test that verifies the API signature without a device.
    Full integration testing of Attention1D creation happens in device tests.
    """
    import inspect

    # Verify the __init__ signature has the expected parameters
    sig = inspect.signature(Attention1D.__init__)
    params = list(sig.parameters.keys())

    # Expected: self, wqkv, wo, n_heads, n_kv_heads, head_dim
    assert "wqkv" in params, "wqkv should be a parameter"
    assert "wo" in params, "wo should be a parameter"
    assert "n_heads" in params, "n_heads should be a required parameter"
    assert "n_kv_heads" in params, "n_kv_heads should be a required parameter"
    assert "head_dim" in params, "head_dim should be a required parameter"

    # Verify n_heads, n_kv_heads, head_dim have no defaults (are required)
    for param_name in ["n_heads", "n_kv_heads", "head_dim"]:
        param = sig.parameters[param_name]
        assert param.default is inspect.Parameter.empty, f"{param_name} should be required (no default)"


def test_attention_1d_resolve_requires_n_heads(expect_error):
    """Test that _resolve_attention1d_config raises ValueError when n_heads is missing."""
    mock_source = MagicMock()
    mock_source.shape = (4096, 1536)

    mock_wqkv = MagicMock(spec=LazyWeight)
    mock_wqkv.source = mock_source
    mock_wqkv.device = None

    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=MagicMock(spec=LazyWeight),
        n_heads=None,  # Missing!
        n_kv_heads=8,
        head_dim=128,
    )

    with expect_error(ValueError, "n_heads must be provided"):
        _resolve_attention1d_config(config)


def test_attention_1d_resolve_requires_n_kv_heads(expect_error):
    """Test that _resolve_attention1d_config raises ValueError when n_kv_heads is missing."""
    mock_source = MagicMock()
    mock_source.shape = (4096, 1536)

    mock_wqkv = MagicMock(spec=LazyWeight)
    mock_wqkv.source = mock_source
    mock_wqkv.device = None

    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=MagicMock(spec=LazyWeight),
        n_heads=32,
        n_kv_heads=None,  # Missing!
        head_dim=128,
    )

    with expect_error(ValueError, "n_kv_heads must be provided"):
        _resolve_attention1d_config(config)


def test_attention_1d_resolve_requires_head_dim(expect_error):
    """Test that _resolve_attention1d_config raises ValueError when head_dim is missing."""
    mock_source = MagicMock()
    mock_source.shape = (4096, 1536)

    mock_wqkv = MagicMock(spec=LazyWeight)
    mock_wqkv.source = mock_source
    mock_wqkv.device = None

    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=MagicMock(spec=LazyWeight),
        n_heads=32,
        n_kv_heads=8,
        head_dim=None,  # Missing!
    )

    with expect_error(ValueError, "head_dim must be provided"):
        _resolve_attention1d_config(config)


def test_attention_1d_resolve_kv_cache_tensor_passthrough():
    """Test that _resolve_attention1d_config passes through raw ttnn.Tensor KV cache entries."""
    mock_source = MagicMock()
    mock_source.shape = (4096, 1536)

    mock_wqkv = MagicMock(spec=LazyWeight)
    mock_wqkv.source = mock_source
    mock_wqkv.device = None

    # Simulate pre-allocated ttnn.Tensor KV cache (e.g., from vLLM).
    # Use plain MagicMock (NOT spec=LazyWeight) so isinstance(_, LazyWeight) is False.
    mock_cache_k = MagicMock()
    mock_cache_v = MagicMock()

    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=MagicMock(spec=LazyWeight),
        n_heads=32,
        n_kv_heads=8,
        head_dim=128,
        max_batch_size=1,
        max_seq_len=128,
        kv_cache=(mock_cache_k, mock_cache_v),
    )

    # Should not crash — previously would fail with dataclasses.replace on non-dataclass
    try:
        resolved = _resolve_attention1d_config(config)
    except (ValueError, AssertionError) as e:
        # Only device-related errors are acceptable, not kv_cache errors
        assert "kv_cache" not in str(e).lower(), f"Unexpected kv_cache error: {e}"
        return

    # If resolution succeeded fully, verify the tensors were passed through as-is
    assert resolved.kv_cache[0] is mock_cache_k
    assert resolved.kv_cache[1] is mock_cache_v


def test_attention_1d_resolve_rejects_sliding_window_with_paged(expect_error):
    """Test that _resolve_attention1d_config rejects sliding_window + paged_attention_config."""
    mock_source = MagicMock()
    mock_source.shape = (4096, 1536)

    mock_wqkv = MagicMock(spec=LazyWeight)
    mock_wqkv.source = mock_source
    mock_wqkv.device = None

    config = Attention1DConfig(
        wqkv=mock_wqkv,
        wo=MagicMock(spec=LazyWeight),
        n_heads=32,
        n_kv_heads=8,
        head_dim=128,
        max_batch_size=1,
        max_seq_len=128,
        sliding_window=4096,
        paged_attention_config=PagedAttentionConfig(block_size=64, max_num_blocks=2048),
    )

    with expect_error(ValueError, "sliding_window"):
        _resolve_attention1d_config(config)


class _AttentionPrefillTensor:
    def __init__(self, shape, dtype):
        self.shape = shape
        self.dtype = dtype


@pytest.mark.parametrize("use_runtime_tensor", [False, True])
def test_attention_prefill_selects_scalar_or_tensor_chunk_start_api(monkeypatch, use_runtime_tensor):
    bfloat16 = object()
    bfloat8_b = object()
    chunked_sdpa = MagicMock(return_value=_AttentionPrefillTensor((1, 32, 128, 128), bfloat8_b))
    xqkv = _AttentionPrefillTensor((1, 1, 128, 6144), bfloat16)
    q_heads = _AttentionPrefillTensor((1, 32, 128, 128), bfloat16)
    k_heads = _AttentionPrefillTensor((1, 8, 128, 128), bfloat16)
    v_heads = _AttentionPrefillTensor((1, 8, 128, 128), bfloat16)
    output = _AttentionPrefillTensor((1, 1, 128, 4096), bfloat8_b)
    # TTNN operation fakes retain backend-specific overload arguments.
    fake_ttnn = SimpleNamespace(
        DRAM_MEMORY_CONFIG=object(),
        bfloat16=bfloat16,
        bfloat8_b=bfloat8_b,
        deallocate=MagicMock(),
        experimental=SimpleNamespace(
            nlp_concat_heads=MagicMock(side_effect=lambda tensor, **_kwargs: tensor),
            nlp_create_qkv_heads=MagicMock(return_value=(q_heads, k_heads, v_heads)),
            rotary_embedding_llama=MagicMock(side_effect=lambda tensor, *_args, **_kwargs: tensor),
        ),
        linear=MagicMock(side_effect=[xqkv, output]),
        reshape=MagicMock(side_effect=lambda tensor, *_args, **_kwargs: tensor),
        transformer=SimpleNamespace(chunked_scaled_dot_product_attention=chunked_sdpa),
        typecast=MagicMock(side_effect=lambda tensor, **_kwargs: tensor),
    )
    prefill_sdpa_prg_config = MagicMock(return_value="sdpa-program")
    cfg = SimpleNamespace(
        activation_dtype=None,
        head_dim=128,
        li_o_prefill_compute_kernel_cfg=object(),
        li_qkv_prefill_compute_kernel_cfg=object(),
        mesh_device=SimpleNamespace(get_num_devices=MagicMock(return_value=1)),
        min_kv_prefill_shard_seqlen=256,
        n_heads=32,
        n_kv_heads=8,
        paged_attention_config=SimpleNamespace(block_size=32),
        prefill_sdpa_prg_config=prefill_sdpa_prg_config,
        prefill_wo_prg_config=MagicMock(return_value="wo-program"),
        prefill_xqkv_prg_config=MagicMock(return_value="qkv-program"),
        scale=0.125,
        sdpa_prefill_compute_kernel_cfg=object(),
        sliding_window=None,
        transformation_mat_prefill=object(),
        use_minimal_qkv_matmul=MagicMock(return_value=False),
        use_minimal_wo_matmul=MagicMock(return_value=False),
        wo_prefill_len_cutoff=1024,
    )
    cache_dtype = object()
    attention = SimpleNamespace(
        _all_gather_before_wo_prefill=MagicMock(side_effect=lambda tensor: tensor),
        _kv_fill_prefill=MagicMock(),
        _reduce_after_wo_prefill=MagicMock(side_effect=lambda tensor: tensor),
        config=cfg,
        k_norm=None,
        kv_cache=(
            _AttentionPrefillTensor((1, 8, 4096, 128), cache_dtype),
            _AttentionPrefillTensor((1, 8, 4096, 128), cache_dtype),
        ),
        load_device_weights=MagicMock(),
        q_norm=None,
        wo=object(),
        wqkv=object(),
        wqkv_bias_prefill=None,
    )
    x = _AttentionPrefillTensor((1, 1, 128, 4096), bfloat16)
    page_table = object()
    chunk_start_idx_tensor = object() if use_runtime_tensor else None
    monkeypatch.setattr(attention_1d_module, "ttnn", fake_ttnn)
    # This loader is a generic bypass, not the API under test.
    monkeypatch.setattr(attention_1d_module, "_load_input_device_tensor", lambda tensor, *_args, **_kwargs: tensor)

    Attention1D.prefill_forward(
        attention,
        x,
        (object(), object()),
        page_table=page_table,
        chunk_start_idx=96,
        chunk_start_idx_tensor=chunk_start_idx_tensor,
    )

    sdpa_kwargs = chunked_sdpa.call_args.kwargs
    if use_runtime_tensor:
        assert sdpa_kwargs["chunk_start_idx_tensor"] is chunk_start_idx_tensor
        assert "chunk_start_idx" not in sdpa_kwargs
        prefill_sdpa_prg_config.assert_called_once_with(128, 32)
    else:
        assert sdpa_kwargs["chunk_start_idx"] == 96
        assert "chunk_start_idx_tensor" not in sdpa_kwargs
        prefill_sdpa_prg_config.assert_called_once_with(128, 96)


# ============================================================================
# Model name constants
# ============================================================================

LLAMA_8B = "meta-llama/Llama-3.1-8B-Instruct"
LLAMA_70B = "meta-llama/Llama-3.3-70B-Instruct"
LLAMA_1B = "meta-llama/Llama-3.2-1B-Instruct"
LLAMA_3B = "meta-llama/Llama-3.2-3B-Instruct"
LLAMA_11B = "meta-llama/Llama-3.2-11B-Vision-Instruct"
LLAMA_90B = "meta-llama/Llama-3.2-90B-Vision-Instruct"
MISTRAL_7B = "mistralai/Mistral-7B-Instruct-v0.3"
MIXTRAL_8X7B = "mistralai/Mixtral-8x7B-v0.1"
QWEN2_7B = "Qwen/Qwen2-7B-Instruct"
QWEN25_7B = "Qwen/Qwen2.5-7B-Instruct"
QWEN25_72B = "Qwen/Qwen2.5-72B-Instruct"
QWEN25_CODER_32B = "Qwen/Qwen2.5-Coder-32B-Instruct"
DEEPSEEK_R1_14B = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
QWEN3_32B = "Qwen/Qwen3-32B"
PHI4 = "microsoft/phi-4"  # Phi3 architecture: fused qkv_proj, no bias, no q/k norm, GQA 40/10

_slow = pytest.mark.slow


# ============================================================================
# Test cases from attn_1d_performance.csv - hardcoded as pytest parameters
# ============================================================================


def _list_test_cases() -> list[pytest.param]:
    """
    Hardcoded test cases from attn_1d_performance.csv.

    Parameters: mesh_shape, seq_len, batch_size, mode, x_dtype, wqkv_dtype, hf_model_name, pcc

    batch_size semantics:
    - Prefill: batch_size=1 (single user prefill), input shape (1, 1, seq_len, dim)
    - Decode: batch_size=32 (continuous batching), input shape (1, 1, batch_size, dim)
    """
    # fmt: off
    return [
        # === Fast tests (minimal coverage set) ===
        # Single device (1x1)
        pytest.param((1, 1), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-128-1B"),
        pytest.param((1, 1), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-decode-32-1B"),
        pytest.param((1, 1), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-8192-1B"),
        # Dual device (1x2)
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-128-8B"),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-decode-32-8B"),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.99, id="1x2-decode-32-11B"),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.95, id="1x2-decode-32-Qwen2.5-7B", marks=pytest.mark.skip(reason="Disabled: see #45980")),
        # Multi-device (1x8)
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-128-8B"),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-decode-32-8B"),

        # === Slow tests (full coverage from models sweep) ===
        # --- Llama-3.2-1B on N150 (1x1) ---
        pytest.param((1, 1), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-1024-1B", marks=_slow),
        pytest.param((1, 1), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-2048-1B", marks=_slow),
        pytest.param((1, 1), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-4096-1B", marks=_slow),
        pytest.param((1, 1), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-16384-1B", marks=_slow),
        pytest.param((1, 1), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x1-prefill-32768-1B", marks=_slow),

        # --- Llama-3.2-3B on N150 (1x1) ---
        pytest.param((1, 1), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-prefill-128-3B", marks=_slow),
        pytest.param((1, 1), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-prefill-1024-3B", marks=_slow),
        pytest.param((1, 1), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-prefill-2048-3B", marks=_slow),
        pytest.param((1, 1), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-prefill-4096-3B", marks=_slow),
        pytest.param((1, 1), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-prefill-8192-3B", marks=_slow),
        pytest.param((1, 1), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x1-decode-32-3B", marks=_slow),

        # --- Llama-3.1-8B on N150 (1x1) ---
        pytest.param((1, 1), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x1-prefill-128-8B", marks=_slow),
        pytest.param((1, 1), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x1-prefill-1024-8B", marks=_slow),
        pytest.param((1, 1), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x1-prefill-2048-8B", marks=_slow),
        pytest.param((1, 1), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x1-prefill-4096-8B", marks=_slow),
        pytest.param((1, 1), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x1-decode-32-8B", marks=_slow),

        # --- Mistral-7B on N150 (1x1) ---
        pytest.param((1, 1), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x1-prefill-128-Mistral-7B", marks=_slow),
        pytest.param((1, 1), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x1-prefill-1024-Mistral-7B", marks=_slow),
        pytest.param((1, 1), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x1-prefill-2048-Mistral-7B", marks=_slow),
        pytest.param((1, 1), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x1-prefill-4096-Mistral-7B", marks=_slow),
        pytest.param((1, 1), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x1-decode-32-Mistral-7B", marks=_slow),

        # --- Llama-3.2-1B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-128-1B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-1024-1B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-2048-1B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-4096-1B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-8192-1B", marks=_slow),
        pytest.param((1, 2), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-16384-1B", marks=_slow),
        pytest.param((1, 2), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-prefill-32768-1B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x2-decode-32-1B", marks=_slow),

        # --- Llama-3.2-3B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-128-3B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-1024-3B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-2048-3B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-4096-3B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-8192-3B", marks=_slow),
        pytest.param((1, 2), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-16384-3B", marks=_slow),
        pytest.param((1, 2), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-prefill-32768-3B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x2-decode-32-3B", marks=_slow),

        # --- Llama-3.1-8B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-128-8B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-1024-8B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-2048-8B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-4096-8B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-8192-8B", marks=_slow),
        pytest.param((1, 2), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-16384-8B", marks=_slow),
        pytest.param((1, 2), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-prefill-32768-8B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x2-decode-32-8B", marks=_slow),

        # --- Llama-3.2-11B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.99, id="1x2-prefill-128-11B", marks=_slow),
        # NOTE: 11B 1024+ prefill has lower PCC (0.9845) due to vision model complexity
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-1024-11B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-2048-11B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-4096-11B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-8192-11B", marks=_slow),
        pytest.param((1, 2), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-16384-11B", marks=_slow),
        pytest.param((1, 2), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x2-prefill-32768-11B", marks=_slow),

        # --- Mistral-7B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x2-prefill-128-Mistral-7B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x2-prefill-1024-Mistral-7B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x2-prefill-2048-Mistral-7B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x2-prefill-4096-Mistral-7B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x2-decode-32-Mistral-7B", marks=_slow),

        # --- Phi-4 on N300 (1x2) --- fused qkv_proj split, GQA 40/10, head_dim 128, no bias/qk-norm.
        # decode-32 is validated on N300 for both standard and paged SDPA-decode (the earlier
        # num_output_cores limit for Phi-4's 10 KV heads no longer trips); batch-1 decode is
        # additionally covered end-to-end by the M5 token-accuracy demo (97-99% top-1).
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, PHI4, 0.99, id="1x2-prefill-128-Phi-4", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, PHI4, 0.99, id="1x2-decode-32-Phi-4", marks=_slow),

        # --- Qwen2-7B on N300 (1x2) ---
        # NOTE: Qwen2-7B has Q/K biases causing numerical precision issues
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN2_7B, 0.98, id="1x2-prefill-128-Qwen2-7B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN2_7B, 0.97, id="1x2-prefill-1024-Qwen2-7B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN2_7B, 0.97, id="1x2-prefill-2048-Qwen2-7B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN2_7B, 0.97, id="1x2-prefill-4096-Qwen2-7B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, QWEN2_7B, 0.99, id="1x2-decode-32-Qwen2-7B", marks=_slow),

        # --- Qwen2.5-7B on N300 (1x2) ---
        # NOTE: Qwen2.5-7B has large Q/K biases causing numerical precision issues
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.98, id="1x2-prefill-128-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-1024-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-2048-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-4096-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-8192-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-16384-Qwen2.5-7B", marks=_slow),
        pytest.param((1, 2), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_7B, 0.97, id="1x2-prefill-32768-Qwen2.5-7B", marks=_slow),
        # NOTE: Qwen2.5-7B has lower PCC for prefill+decode due to Q/K biases + RoPE interaction.
        # TTTv1's test_attention.py also shows ~0.984 min PCC. With 128-token prefill, accumulated
        # numerical error in SDPA over the larger KV cache causes further degradation.
        # See models/common/tests/modules/attention/low_pcc_notes.md for detailed analysis

        # --- DeepSeek-R1-14B on N300 (1x2) ---
        pytest.param((1, 2), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-prefill-128-DeepSeek-R1-14B", marks=_slow),
        pytest.param((1, 2), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-prefill-1024-DeepSeek-R1-14B", marks=_slow),
        pytest.param((1, 2), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-prefill-2048-DeepSeek-R1-14B", marks=_slow),
        pytest.param((1, 2), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-prefill-4096-DeepSeek-R1-14B", marks=_slow),
        pytest.param((1, 2), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-prefill-8192-DeepSeek-R1-14B", marks=_slow),
        pytest.param((1, 2), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, DEEPSEEK_R1_14B, 0.99, id="1x2-decode-32-DeepSeek-R1-14B", marks=_slow),

        # --- Llama-3.2-1B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-128-1B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-1024-1B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-2048-1B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-4096-1B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-8192-1B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-16384-1B", marks=_slow),
        pytest.param((1, 8), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-prefill-32768-1B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_1B, 0.99, id="1x8-decode-32-1B", marks=_slow),

        # --- Llama-3.2-3B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-128-3B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-1024-3B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-2048-3B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-4096-3B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-8192-3B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-16384-3B", marks=_slow),
        pytest.param((1, 8), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-prefill-32768-3B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_3B, 0.99, id="1x8-decode-32-3B", marks=_slow),

        # --- Llama-3.1-8B on T3K (1x8) ---
        pytest.param((1, 8), 256, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-256-8B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-1024-8B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-2048-8B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-4096-8B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-8192-8B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-16384-8B", marks=_slow),
        pytest.param((1, 8), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-prefill-32768-8B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_8B, 0.99, id="1x8-decode-32-8B", marks=_slow),

        # --- Llama-3.2-11B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.99, id="1x8-prefill-128-11B", marks=_slow),
        # NOTE: 11B 1024+ prefill has lower PCC (0.9844) due to vision model complexity
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-1024-11B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-2048-11B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-4096-11B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-8192-11B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-16384-11B", marks=_slow),
        pytest.param((1, 8), 32768, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.98, id="1x8-prefill-32768-11B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_11B, 0.99, id="1x8-decode-32-11B", marks=_slow),

        # --- Llama-3.3-70B on T3K (1x8) ---
        # NOTE: 70B has slightly lower PCC (0.997) due to model size and multi-device communication
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_70B, 0.99, id="1x8-prefill-128-70B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_70B, 0.99, id="1x8-prefill-1024-70B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_70B, 0.99, id="1x8-prefill-2048-70B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_70B, 0.99, id="1x8-prefill-4096-70B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_70B, 0.97, id="1x8-decode-32-70B", marks=_slow),

        # --- Llama-3.2-90B on T3K (1x8) ---
        # NOTE: 90B has slightly lower PCC (0.995-0.996) due to model size
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_90B, 0.99, id="1x8-prefill-128-90B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, LLAMA_90B, 0.99, id="1x8-decode-32-90B", marks=_slow),

        # --- Mistral-7B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x8-prefill-128-Mistral-7B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x8-prefill-1024-Mistral-7B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x8-prefill-2048-Mistral-7B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x8-prefill-4096-Mistral-7B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, MISTRAL_7B, 0.99, id="1x8-decode-32-Mistral-7B", marks=_slow),

        # --- Mixtral-8x7B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MIXTRAL_8X7B, 0.99, id="1x8-prefill-128-Mixtral-8x7B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MIXTRAL_8X7B, 0.99, id="1x8-prefill-1024-Mixtral-8x7B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MIXTRAL_8X7B, 0.99, id="1x8-prefill-2048-Mixtral-8x7B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, MIXTRAL_8X7B, 0.99, id="1x8-prefill-4096-Mixtral-8x7B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, MIXTRAL_8X7B, 0.99, id="1x8-decode-32-Mixtral-8x7B", marks=_slow),

        # --- Qwen2.5-72B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-128-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-1024-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-2048-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-4096-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-8192-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-prefill-16384-Qwen2.5-72B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_72B, 0.99, id="1x8-decode-32-Qwen2.5-72B", marks=_slow),

        # --- Qwen2.5-Coder-32B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_CODER_32B, 0.99, id="1x8-prefill-128-Qwen2.5-Coder-32B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_CODER_32B, 0.99, id="1x8-prefill-1024-Qwen2.5-Coder-32B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_CODER_32B, 0.99, id="1x8-prefill-2048-Qwen2.5-Coder-32B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_CODER_32B, 0.99, id="1x8-prefill-4096-Qwen2.5-Coder-32B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, QWEN25_CODER_32B, 0.99, id="1x8-decode-32-Qwen2.5-Coder-32B", marks=_slow),
        # BF16 weights
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_CODER_32B, 0.99, id="1x8-prefill-128-Qwen2.5-Coder-32B-bf16", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN25_CODER_32B, 0.99, id="1x8-prefill-1024-Qwen2.5-Coder-32B-bf16", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat16, QWEN25_CODER_32B, 0.99, id="1x8-decode-32-Qwen2.5-Coder-32B-bf16", marks=_slow),

        # --- Qwen3-32B on T3K (1x8) ---
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-128-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-1024-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 2048, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-2048-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 4096, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-4096-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 8192, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-8192-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 16384, 1, "prefill", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-prefill-16384-Qwen3-32B", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat8_b, QWEN3_32B, 0.99, id="1x8-decode-32-Qwen3-32B", marks=_slow),
        # BF16 weights
        pytest.param((1, 8), 128, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN3_32B, 0.99, id="1x8-prefill-128-Qwen3-32B-bf16", marks=_slow),
        pytest.param((1, 8), 1024, 1, "prefill", ttnn.bfloat16, ttnn.bfloat16, QWEN3_32B, 0.99, id="1x8-prefill-1024-Qwen3-32B-bf16", marks=_slow),
        pytest.param((1, 8), 32, 32, "decode", ttnn.bfloat16, ttnn.bfloat16, QWEN3_32B, 0.99, id="1x8-decode-32-Qwen3-32B-bf16", marks=_slow),
    ]
    # fmt: on


# ============================================================================
# Integration Tests - Require device
# ============================================================================


@torch.no_grad()
@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [(1, 1), (1, 2), (1, 8)],
    ids=["1x1", "1x2", "1x8"],
    indirect=True,
)
@pytest.mark.parametrize(
    "mesh_shape,seq_len,batch_size,mode,act_dtype,wqkv_dtype,hf_model_name,pcc",
    _list_test_cases(),
)
@pytest.mark.parametrize(
    "page_block_size,chunk_size",
    [
        (None, None),  # standard (non-paged, non-chunked)
        (64, None),  # paged only (non-chunked)
        (64, 4096),  # paged + chunked
    ],
    ids=["standard", "paged", "paged-chunked"],
)
@pytest.mark.parametrize("num_decode_iterations", [10])
def test_attention_1d_vs_reference(
    ttnn_mesh_device: ttnn.MeshDevice,
    mesh_shape,
    seq_len,
    batch_size,
    mode,
    act_dtype,
    wqkv_dtype,
    hf_model_name,
    pcc,
    page_block_size,
    chunk_size,
    num_decode_iterations,
):
    """
    Test Attention1D constructed via from_config (direct API) executes successfully.

    This test uses the TTTv2 pattern:
    1. Load HF config directly via AutoConfig.from_pretrained
    2. Extract weights from HF reference model
    3. Create LazyWeight objects
    4. Build Attention1DConfig with explicit parameters
    5. Create Attention1D via from_config (NOT from_model_args)
    6. Run forward pass and verify execution

    Note: Full numerical comparison against HF reference is done in
    test_attention_1d_vs_reference_from_model_args which uses existing infrastructure
    that handles HF API differences.
    """
    # Skip if mesh_shape doesn't match device
    device_shape = (ttnn_mesh_device.shape[0], ttnn_mesh_device.shape[1])
    if device_shape != mesh_shape:
        pytest.skip(f"Test requires {mesh_shape} mesh, got {device_shape}")

    # Chunked prefill only applies to prefill mode
    chunked_prefill = chunk_size is not None
    if chunked_prefill and mode != "prefill":
        pytest.skip("Chunked prefill only applies to prefill mode")

    # Chunked prefill requires seq_len > chunk_size
    # TTTv1 uses chunk_size = N * 1024 where N ranges from 4-128 depending on model/device
    # Default of 4096 matches TTTv1's minimum (4 * 1024)
    if chunked_prefill and seq_len <= chunk_size:
        pytest.skip(f"Chunked prefill requires seq_len > chunk_size ({chunk_size})")

    # batch_size is now a test parameter:
    # - Prefill: typically batch_size=1, input shape (1, 1, seq_len, dim)
    # - Decode: typically batch_size=32 (continuous batching), input shape (1, 1, batch, dim)
    #   current_pos has shape (batch_size,) - one position per user
    # Note: For decode tests, seq_len parameter is unused (kept for parameterization compatibility)

    # max_seq_len: minimum allocation for KV cache
    # - Prefill: exactly seq_len tokens written to cache
    # - Decode: num_decode_iterations tokens (starts from position 0, no prefill)
    # Round up to alignment (SDPA kernel requires multiples of 32, paged attention requires page_block_size)
    if mode == "prefill":
        max_seq_len = seq_len
    else:
        max_seq_len = num_decode_iterations
    # Round up: use page_block_size if paged, otherwise 32 (SDPA tile alignment)
    alignment = page_block_size if page_block_size is not None else 32
    max_seq_len = ((max_seq_len + alignment - 1) // alignment) * alignment
    num_devices = ttnn_mesh_device.get_num_devices()

    # Seed for reproducibility
    seed = 1234
    torch.manual_seed(seed)

    # Load HF config directly (no ModelArgs)
    hf_config = AutoConfig.from_pretrained(hf_model_name)

    # Handle multimodal models (Mllama, LLaVA, etc.) which nest text config under .text_config
    is_multimodal = hasattr(hf_config, "text_config") and hf_config.text_config is not None
    cfg = hf_config.text_config if is_multimodal else hf_config
    cfg.num_hidden_layers = 1  # Only need 1 layer for testing

    dim = cfg.hidden_size
    n_heads = cfg.num_attention_heads
    n_kv_heads = getattr(cfg, "num_key_value_heads", n_heads)
    # Use explicit head_dim if available (e.g., Qwen3 models), else calculate from dim/n_heads
    # Note: some configs have head_dim=None explicitly, so we use `or` to fallback
    head_dim = getattr(cfg, "head_dim", None) or (dim // n_heads)
    sliding_window = getattr(cfg, "sliding_window", None)

    # Load HF model structure without weights, then initialize with random weights
    # This avoids slow network downloads and ensures reproducible deterministic testing
    if is_multimodal:
        # For multimodal models, import and use the specific model class
        from transformers import MllamaForConditionalGeneration

        with no_init_weights():
            # MllamaForConditionalGeneration uses _from_config (internal method) instead of from_config
            hf_model = MllamaForConditionalGeneration._from_config(hf_config, torch_dtype=torch.bfloat16)
        # Mllama has layers directly at language_model.layers (not language_model.model.layers).
        # transformers 5.x nests the text model under hf_model.model.language_model.
        text_model = hf_model.language_model if hasattr(hf_model, "language_model") else hf_model.model.language_model
        first_layer = text_model.layers[0]
        rotary_emb = getattr(text_model, "rotary_emb", None)
    else:
        with no_init_weights():
            hf_model = AutoModelForCausalLM.from_config(hf_config, torch_dtype=torch.bfloat16)
        first_layer = hf_model.model.layers[0]
        rotary_emb = getattr(hf_model.model, "rotary_emb", None)

    # Get reference attention from first layer
    reference_attn = first_layer.self_attn

    # Initialize attention weights deterministically (cached for speed across test cases)
    _get_or_init_attn_weights(hf_model_name, reference_attn)

    # Wrap in HfAttentionWrapper for consistent KV cache and RoPE handling (local class)
    reference_wrapper = HfAttentionWrapper(reference_attn, head_dim, rotary_emb)

    # Extract attention weights in TTNN layout
    wqkv_torch, wo_torch, q_norm_torch, k_norm_torch, wqkv_bias_torch = get_attention_weights_from_ref_model(
        reference_attn, num_devices
    )

    # Create LazyWeights with caching enabled
    # Cache keys include model name + seed to avoid mismatched cached weights across runs
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    cache_dir = Path(os.getenv("TT_CACHE_PATH", "model_cache/attention_1d"))
    model_seed = stable_model_seed(hf_model_name)
    model_cache_prefix = f"{hf_model_name.replace('/', '_').replace('-', '_')}_seed{model_seed:08x}"

    # QKV weight: shard on dim=-1
    lazy_wqkv = LazyWeight(
        source=wqkv_torch,
        dtype=wqkv_dtype,
        cache_dir_weight_name=(cache_dir, f"{model_cache_prefix}_layer0_wqkv"),
    )

    # WO weight: shard on dim=-2
    lazy_wo = LazyWeight(
        source=wo_torch,
        dtype=wqkv_dtype,
        cache_dir_weight_name=(cache_dir, f"{model_cache_prefix}_layer0_wo"),
    )

    # Q/K norm configs (optional) - using RMSNorm1DConfig composition pattern
    q_norm_config = None
    k_norm_config = None
    if q_norm_torch is not None:
        lazy_q_norm = LazyWeight(
            source=q_norm_torch.unsqueeze(0).unsqueeze(0).unsqueeze(0),
            dtype=ttnn.bfloat16,
            cache_dir_weight_name=(cache_dir, f"{model_cache_prefix}_layer0_q_norm"),
        )
        q_norm_config = RMSNorm1DConfig(
            weight=lazy_q_norm,
            mesh_device=ttnn_mesh_device,
            eps=1e-5,
            decode_in_sharded=False,  # Q/K heads are interleaved
            decode_out_sharded=False,
            prefill_distributed=False,
        )
    if k_norm_torch is not None:
        lazy_k_norm = LazyWeight(
            source=k_norm_torch.unsqueeze(0).unsqueeze(0).unsqueeze(0),
            dtype=ttnn.bfloat16,
            cache_dir_weight_name=(cache_dir, f"{model_cache_prefix}_layer0_k_norm"),
        )
        k_norm_config = RMSNorm1DConfig(
            weight=lazy_k_norm,
            mesh_device=ttnn_mesh_device,
            eps=1e-5,
            decode_in_sharded=False,  # Q/K heads are interleaved
            decode_out_sharded=False,
            prefill_distributed=False,
        )

    # Create TT_CCL for multi-device
    tt_ccl = TT_CCL(ttnn_mesh_device) if num_devices > 1 else None

    # Determine topology
    if num_devices == 1:
        topology = None
    elif num_devices == 2:
        topology = ttnn.Topology.Linear
    else:
        topology = ttnn.Topology.Ring

    # Setup paged attention config and page table if enabled
    paged_attention_config = None
    page_table_tt = None
    page_table = None
    reverse_permutation = None

    if page_block_size is not None:
        # Paged attention parameters (use local PagedAttentionConfig)
        # Each user needs ceil(max_seq_len / block_size) blocks
        blocks_per_user = (max_seq_len + page_block_size - 1) // page_block_size
        max_num_blocks = max(128, blocks_per_user * batch_size)

        paged_attention_config = PagedAttentionConfig(
            block_size=page_block_size,
            max_num_blocks=max_num_blocks,
        )

        # Create page table: random permutation simulates block allocation
        permutation = torch.randperm(max_num_blocks)
        reverse_permutation = torch.argsort(permutation)
        page_table = reverse_permutation.reshape(batch_size, max_num_blocks // batch_size)
        page_table_tt = ttnn.from_torch(
            page_table,
            device=ttnn_mesh_device,
            dtype=ttnn.int32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
        )

    # Determine if we need the power user path (from_config) or can use the simple API
    # Power user path is needed for models with special features that aren't in the simple API
    has_special_features = (
        q_norm_config is not None
        or k_norm_config is not None
        or wqkv_bias_torch is not None
        or sliding_window is not None
        or paged_attention_config is not None
    )

    if has_special_features:
        # Power user path: use from_config() for models with special features
        config = Attention1DConfig(
            wqkv=lazy_wqkv,
            wo=lazy_wo,
            n_heads=n_heads,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            max_batch_size=batch_size,
            max_seq_len=max_seq_len,
            q_norm_config=q_norm_config,
            k_norm_config=k_norm_config,
            wqkv_bias=LazyWeight(source=wqkv_bias_torch) if wqkv_bias_torch is not None else None,
            sliding_window=sliding_window,
            paged_attention_config=paged_attention_config,
        )
        tt_model = Attention1D.from_config(config)
    else:
        # Happy path: simple API for basic Llama-style models
        tt_model = Attention1D(
            wqkv=lazy_wqkv,
            wo=lazy_wo,
            n_heads=n_heads,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            max_batch_size=batch_size,
            max_seq_len=max_seq_len,
        )

    # Verify config is properly resolved (works for both happy path and from_config)
    assert tt_model.config.is_resolved(), "Config should be resolved after Attention1D creation"
    assert tt_model.config.dim == dim
    assert tt_model.config.n_heads == n_heads
    assert tt_model.config.n_kv_heads == n_kv_heads
    assert tt_model.config.head_dim == head_dim

    if mode == "prefill":
        _run_prefill_test(
            tt_model=tt_model,
            reference_wrapper=reference_wrapper,
            ttnn_mesh_device=ttnn_mesh_device,
            mesh_shape=mesh_shape,
            batch_size=batch_size,
            seq_len=seq_len,
            dim=dim,
            n_heads=n_heads,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            act_dtype=act_dtype,
            page_block_size=page_block_size,
            chunked_prefill=chunked_prefill,
            chunk_size=chunk_size,
            paged_attention_config=paged_attention_config,
            page_table=page_table,
            page_table_tt=page_table_tt,
            pcc=pcc,
        )
    else:
        _run_decode_test(
            tt_model=tt_model,
            reference_wrapper=reference_wrapper,
            ttnn_mesh_device=ttnn_mesh_device,
            mesh_shape=mesh_shape,
            batch_size=batch_size,
            dim=dim,
            n_heads=n_heads,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            max_seq_len=max_seq_len,
            act_dtype=act_dtype,
            page_block_size=page_block_size,
            page_table_tt=page_table_tt,
            num_decode_iterations=num_decode_iterations,
            pcc=pcc,
        )


def _run_prefill_test(
    tt_model,
    reference_wrapper,
    ttnn_mesh_device,
    mesh_shape,
    batch_size,
    seq_len,
    dim,
    n_heads,
    n_kv_heads,
    head_dim,
    act_dtype,
    page_block_size,
    chunked_prefill,
    chunk_size,
    paged_attention_config,
    page_table,
    page_table_tt,
    pcc,
):
    """Run prefill test and compare against HuggingFace reference."""
    pt_attention_input = torch.randn(batch_size, seq_len, dim, dtype=torch.bfloat16)

    # Prepare TT input for prefill
    tt_input = ttnn.from_torch(
        pt_attention_input.unsqueeze(0),  # [1, batch, seq_len, dim]
        device=ttnn_mesh_device,
        dtype=act_dtype,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
    )

    # Get rot_mats for prefill (cos/sin matrices) - extract from HF rotary_emb
    if chunked_prefill:
        # Chunked prefill: process sequence in chunks
        # Compute full rotation matrices once (covering all positions 0 to seq_len)
        full_cos, full_sin = get_cos_sin_from_hf(
            reference_wrapper.rotary_emb,
            seq_len=seq_len,
            head_dim=head_dim,
        )

        num_chunks = seq_len // chunk_size
        tt_outputs_chunked = []

        for chunk_idx in range(num_chunks):
            chunk_start = chunk_idx * chunk_size
            chunk_end = chunk_start + chunk_size

            # Extract chunk input
            pt_chunk = pt_attention_input[:, chunk_start:chunk_end, :]

            # Prepare TT input for this chunk
            tt_chunk = ttnn.from_torch(
                pt_chunk.unsqueeze(0),
                device=ttnn_mesh_device,
                dtype=act_dtype,
                layout=ttnn.TILE_LAYOUT,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
            )

            # Chunk page table: subset of pages for this chunk
            block_size = paged_attention_config.block_size
            chunk_page_table_pt = page_table[:, chunk_start // block_size : chunk_end // block_size]
            chunk_page_table_tt = ttnn.from_torch(
                chunk_page_table_pt,
                device=ttnn_mesh_device,
                dtype=ttnn.int32,
                layout=ttnn.ROW_MAJOR_LAYOUT,
                mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
            )

            # Slice rotation matrices for this chunk's positions
            chunk_cos = full_cos[:, :, chunk_start:chunk_end, :]
            chunk_sin = full_sin[:, :, chunk_start:chunk_end, :]

            chunk_cos_tt = ttnn.from_torch(
                chunk_cos,
                device=ttnn_mesh_device,
                layout=ttnn.TILE_LAYOUT,
                dtype=act_dtype,
                mesh_mapper=ttnn.ReplicateTensorToMesh(ttnn_mesh_device),
            )
            chunk_sin_tt = ttnn.from_torch(
                chunk_sin,
                device=ttnn_mesh_device,
                layout=ttnn.TILE_LAYOUT,
                dtype=act_dtype,
                mesh_mapper=ttnn.ReplicateTensorToMesh(ttnn_mesh_device),
            )
            chunk_rot_mats = [chunk_cos_tt, chunk_sin_tt]

            # Run chunked prefill
            tt_chunk_out = tt_model.forward(
                tt_chunk,
                None,
                chunk_rot_mats,
                mode="prefill",
                page_table=page_table_tt,
                chunk_page_table=chunk_page_table_tt,
                chunk_start_idx=chunk_start,
            )

            # Collect output
            tt_chunk_out_torch = to_torch_auto_compose(tt_chunk_out)
            tt_chunk_output = tt_chunk_out_torch[:, 0:1, :chunk_size, :dim].view(batch_size, chunk_size, dim)
            tt_outputs_chunked.append(tt_chunk_output)

        # Concatenate all chunk outputs
        tt_output_torch = torch.cat(tt_outputs_chunked, dim=1)
    else:
        # Standard prefill: single forward pass - use HF rotary_emb
        rot_mats = get_rot_mats_from_hf(
            reference_wrapper.rotary_emb,
            seq_len=seq_len,
            head_dim=head_dim,
            device=ttnn_mesh_device,
        )

        # Run TT model - verify forward pass executes without error
        tt_out = tt_model.forward(
            tt_input,
            None,  # current_pos not used in prefill
            rot_mats,
            mode="prefill",
            page_table=page_table_tt,
        )

        # Convert output to torch and verify shape
        tt_out = to_torch_auto_compose(tt_out)
        tt_output_torch = tt_out[:, 0:1, :seq_len, :dim].view(batch_size, seq_len, dim)
    ttnn.SetDefaultDevice(None)

    # Verify output shape and content
    assert tt_output_torch.shape == (
        batch_size,
        seq_len,
        dim,
    ), f"Expected shape {(batch_size, seq_len, dim)}, got {tt_output_torch.shape}"
    assert not torch.isnan(tt_output_torch).any(), "Output contains NaN values"
    assert not torch.isinf(tt_output_torch).any(), "Output contains Inf values"

    # Run reference HuggingFace attention using HfAttentionWrapper
    # Note: freqs_cis_i is None because HfAttentionWrapper uses rotary_emb directly
    with torch.no_grad():
        reference_output = reference_wrapper(pt_attention_input, start_pos=0, mask=None)

    # Compare TT output with reference using PCC
    passing, pcc_message = comp_pcc(reference_output, tt_output_torch.to(reference_output.dtype), pcc)
    logger.info(f"  PCC comparison: {pcc_message}")
    logger.info(comp_allclose(reference_output, tt_output_torch.to(reference_output.dtype)))
    assert passing, f"Prefill PCC failed: {pcc_message} (expected >= {pcc})"

    # Note: KV cache comparison is skipped because:
    # - HF's DynamicCache stores K/V after applying HF-format RoPE
    # - TT's Attention1D stores K/V after applying Meta-format RoPE
    # These are different rotary embedding formats, so cached values won't match.
    # The output comparison above is the meaningful correctness check.
    logger.info("  KV cache validation: SKIPPED (HF/TT use different RoPE formats in cache)")

    paged_str = f"paged(block_size={page_block_size})" if page_block_size is not None else "non-paged"
    chunked_str = "chunked" if chunked_prefill else "non-chunked"
    logger.info(
        f"test_attention_1d_vs_reference (from_config): PASSED for mode=prefill, seq_len={seq_len}, {paged_str}, {chunked_str}"
    )
    logger.info(f"  Config: dim={dim}, n_heads={n_heads}, n_kv_heads={n_kv_heads}, head_dim={head_dim}")
    logger.info(f"  Output shape: {tt_output_torch.shape}, dtype: {tt_output_torch.dtype}")


def _run_decode_test(
    tt_model,
    reference_wrapper,
    ttnn_mesh_device,
    mesh_shape,
    batch_size,
    dim,
    n_heads,
    n_kv_heads,
    head_dim,
    max_seq_len,
    act_dtype,
    page_block_size,
    page_table_tt,
    num_decode_iterations,
    pcc,
):
    """
    Run decode-only test starting from position 0 (TTTv1 style).

    This is a pure unit test for decode_forward:
    - No prefill step - KV cache builds incrementally during decode
    - Runs num_decode_iterations decode steps at positions 0, 1, 2, ...
    - Compares each iteration against HuggingFace reference

    For batch_size > 1 (continuous batching):
    - All users decode with identical inputs at the same position
    - Verifies batching mechanism doesn't introduce numerical errors

    Note: prefill→decode transition is tested separately in test_attention_1d_prefill_decode_transition.
    """
    # Create decode-specific RotarySetupHelper using HF rotary_emb (reused across iterations)
    decode_rope_setup = RotarySetupHelper(
        ttnn_mesh_device,
        batch_size,
        head_dim,
        max_seq_len,
        reference_wrapper.rotary_emb,
        use_qk_fused=False,
    )

    # Decode iterations starting from position 0
    # KV cache builds incrementally: position 0, 1, 2, ... (TTTv1 style)
    all_iterations_passing = True
    min_pcc_across_iterations = 1.0

    for decode_iter in range(num_decode_iterations):
        current_pos_value = decode_iter  # Start from 0, not after prefill

        # Create identical input for all users (for comparison against single HF reference)
        # TTNN decode expects shape (1, 1, batch_size, dim) - seq_len=1, batch in 3rd dim
        pt_decode_single = torch.randn(1, 1, dim, dtype=torch.bfloat16)
        pt_decode_batched = pt_decode_single.unsqueeze(2).expand(1, 1, batch_size, dim).contiguous()

        tt_decode_input = ttnn.from_torch(
            pt_decode_batched,
            device=ttnn_mesh_device,
            dtype=act_dtype,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
        )

        # Convert to decode input memory config
        tt_decode_input = ttnn.to_memory_config(tt_decode_input, tt_model.config.decode_input_memcfg)

        # Position for decode: all users at the same position
        position_idxs = torch.full((batch_size,), current_pos_value, dtype=torch.long)
        decode_rot_mats = decode_rope_setup.get_rot_mats(position_idxs)

        current_pos = ttnn.from_torch(
            position_idxs,
            device=ttnn_mesh_device,
            dtype=ttnn.int32,
            mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(ttnn_mesh_device),
        )

        # Run TT decode (batched)
        tt_out = tt_model.forward(
            tt_decode_input,
            current_pos,
            decode_rot_mats,
            mode="decode",
            page_table=page_table_tt,
        )

        # Convert output to torch
        tt_out = to_torch_auto_compose(tt_out)

        # Extract TT output: shape (1, 1, batch_size, dim) -> (batch_size, 1, dim)
        tt_output_torch = tt_out[:, 0:1, :batch_size, :dim].view(batch_size, 1, dim)

        # Run HuggingFace decode using wrapper (single user reference)
        with torch.no_grad():
            reference_output = reference_wrapper(pt_decode_single, start_pos=current_pos_value, mask=None)

        # Verify output content
        assert tt_output_torch.numel() > 0, f"Output is empty at iteration {decode_iter}"
        assert not torch.isnan(tt_output_torch).any(), f"Output contains NaN at iteration {decode_iter}"
        assert not torch.isinf(tt_output_torch).any(), f"Output contains Inf at iteration {decode_iter}"

        # Compare EACH user's TT output with the single HF reference
        for user_idx in range(batch_size):
            user_output = tt_output_torch[user_idx : user_idx + 1]
            passing, pcc_value = comp_pcc(reference_output, user_output.to(reference_output.dtype), pcc)
            if isinstance(pcc_value, (int, float)):
                min_pcc_across_iterations = min(min_pcc_across_iterations, float(pcc_value))
            if not passing:
                logger.warning(f"  Iteration {decode_iter}, User {user_idx} PCC failed: {pcc_value}")
                all_iterations_passing = False

    ttnn.SetDefaultDevice(None)

    logger.info(
        f"  Decode iterations: {num_decode_iterations}, min_pcc={min_pcc_across_iterations:.6f} across all iterations"
    )
    assert all_iterations_passing, f"Decode PCC failed (min_pcc={min_pcc_across_iterations:.6f}, expected >= {pcc})"

    paged_str = f"paged(block_size={page_block_size})" if page_block_size is not None else "non-paged"
    logger.info(
        f"test_attention_1d_vs_reference (from_config): PASSED for mode=decode, "
        f"batch_size={batch_size}, {paged_str}, iterations={num_decode_iterations}"
    )
    logger.info(f"  Config: dim={dim}, n_heads={n_heads}, n_kv_heads={n_kv_heads}, head_dim={head_dim}")


# =============================================================================
# Integration Test: Prefill → Decode Transition
# =============================================================================


@torch.no_grad()
@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [
        (1, 1),  # single device
        (1, 2),  # 1D mesh, 2 devices
        (1, 8),  # 1D mesh, 8 devices
    ],
    ids=["1x1", "1x2", "1x8"],
    indirect=True,
)
@pytest.mark.parametrize(
    "page_block_size",
    [None, 64],  # Test both non-paged and paged attention
    ids=["non-paged", "paged"],
)
def test_attention_1d_prefill_decode_transition(ttnn_mesh_device: ttnn.MeshDevice, page_block_size):
    """
    Integration test for prefill→decode transition.

    This test verifies that KV cache state is correctly maintained across the
    prefill→decode boundary. Unlike the unit tests which test prefill and decode
    in isolation, this test verifies the handoff between modes.

    **Integration Point: KV Cache**
    The KV cache is explicitly created and passed to the model config, making
    the integration point crystal clear:
    - Prefill writes keys/values to positions 0..prefill_seq_len-1
    - Decode reads from those positions and writes to prefill_seq_len+

    Test flow:
    1. Create explicit KV cache tensors (the integration point)
    2. Run prefill → populates KV cache positions 0..N-1
    3. Run decode at positions N, N+1, ... → reads from cache, writes new positions
    4. Compare outputs against HuggingFace reference
    """
    # Minimal test parameters - we're testing transition, not model variations
    hf_model_name = "meta-llama/Llama-3.2-1B-Instruct"
    prefill_seq_len = 128  # Must be divisible by 128
    num_decode_after_prefill = 5
    batch_size = 32  # Realistic continuous batching scenario
    pcc = 0.98

    seed = 42
    torch.manual_seed(seed)

    # Load HF config
    hf_config = AutoConfig.from_pretrained(hf_model_name)
    cfg = hf_config.text_config if hasattr(hf_config, "text_config") else hf_config
    cfg.num_hidden_layers = 1

    dim = cfg.hidden_size
    n_heads = cfg.num_attention_heads
    n_kv_heads = getattr(cfg, "num_key_value_heads", n_heads)
    head_dim = getattr(cfg, "head_dim", None) or (dim // n_heads)

    # Calculate max_seq_len with proper alignment
    max_seq_len = prefill_seq_len + num_decode_after_prefill
    alignment = page_block_size if page_block_size is not None else 32
    max_seq_len = ((max_seq_len + alignment - 1) // alignment) * alignment

    mesh_shape = ttnn_mesh_device.shape
    num_devices = ttnn_mesh_device.get_num_devices()
    n_local_kv_heads = n_kv_heads // num_devices

    # Load HF model
    with no_init_weights():
        hf_model = AutoModelForCausalLM.from_config(hf_config, torch_dtype=torch.bfloat16)
    first_layer = hf_model.model.layers[0]
    rotary_emb = getattr(hf_model.model, "rotary_emb", None)

    reference_attn = first_layer.self_attn
    _get_or_init_attn_weights(hf_model_name, reference_attn)
    reference_wrapper = HfAttentionWrapper(reference_attn, head_dim, rotary_emb)

    # Extract weights
    wqkv_torch, wo_torch, q_norm_torch, k_norm_torch, wqkv_bias_torch = get_attention_weights_from_ref_model(
        reference_attn, num_devices
    )

    act_dtype = ttnn.bfloat16
    wqkv_dtype = ttnn.bfloat8_b
    lazy_wqkv = LazyWeight(source=wqkv_torch, dtype=wqkv_dtype, cache_dir_weight_name=None)
    lazy_wo = LazyWeight(source=wo_torch, dtype=wqkv_dtype, cache_dir_weight_name=None)

    # Setup paged attention if enabled
    paged_attention_config = None
    page_table_tt = None
    if page_block_size is not None:
        # Each user needs ceil(max_seq_len / block_size) blocks
        blocks_per_user = (max_seq_len + page_block_size - 1) // page_block_size
        max_num_blocks = max(128, blocks_per_user * batch_size)
        paged_attention_config = PagedAttentionConfig(block_size=page_block_size, max_num_blocks=max_num_blocks)

        permutation = torch.randperm(max_num_blocks)
        reverse_permutation = torch.argsort(permutation)
        page_table = reverse_permutation.reshape(batch_size, max_num_blocks // batch_size)
        page_table_tt = ttnn.from_torch(
            page_table,
            device=ttnn_mesh_device,
            dtype=ttnn.int32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
        )

    # =========================================================================
    # INTEGRATION POINT: Explicitly create KV cache tensors
    # =========================================================================
    # The KV cache is the shared state between prefill and decode.
    # By creating it explicitly here, we make the integration point visible:
    # - Prefill will WRITE to this cache (positions 0..prefill_seq_len-1)
    # - Decode will READ from this cache and WRITE new positions
    # =========================================================================
    if paged_attention_config is not None:
        cache_k = zeros_like_paged_cache(paged_attention_config, n_local_kv_heads, head_dim)
        cache_v = zeros_like_paged_cache(paged_attention_config, n_local_kv_heads, head_dim)
    else:
        cache_k = zeros_like_kv_cache(batch_size, n_local_kv_heads, max_seq_len, head_dim)
        cache_v = zeros_like_kv_cache(batch_size, n_local_kv_heads, max_seq_len, head_dim)

    # Wrap as LazyWeight for config
    kv_cache = (LazyWeight(source=cache_k), LazyWeight(source=cache_v))

    # Build config with explicit KV cache
    topology = ttnn.Topology.Ring if num_devices > 1 else None
    tt_ccl = TT_CCL(ttnn_mesh_device) if num_devices > 1 else None

    config = Attention1DConfig(
        wqkv=lazy_wqkv,
        wo=lazy_wo,
        mesh_device=ttnn_mesh_device,
        tt_ccl=tt_ccl,
        topology=topology,
        dim=dim,
        n_heads=n_heads,
        n_kv_heads=n_kv_heads,
        head_dim=head_dim,
        max_batch_size=batch_size,
        max_seq_len=max_seq_len,
        scale=head_dim**-0.5,
        use_vllm_paged_kv_cache=False,
        paged_attention_config=paged_attention_config,
        kv_cache=kv_cache,  # <-- EXPLICIT KV CACHE: the integration point
        wqkv_dtype=wqkv_dtype,
        wo_dtype=wqkv_dtype,
        activation_dtype=act_dtype,
    )

    tt_model = Attention1D.from_config(config)

    # =========================================================================
    # Step 1: PREFILL - populates KV cache positions 0..prefill_seq_len-1
    # =========================================================================
    # Run prefill per-user (continuous batching style) to avoid compute grid limits.
    # Each user gets the same input for easy comparison.
    # =========================================================================
    pt_prefill_input_single = torch.randn(1, prefill_seq_len, dim, dtype=torch.bfloat16)
    prefill_rot_mats = get_rot_mats_from_hf(rotary_emb, prefill_seq_len, head_dim, ttnn_mesh_device)

    # Collect outputs for all users
    tt_prefill_outputs = []
    for user_id in range(batch_size):
        tt_prefill_input = ttnn.from_torch(
            pt_prefill_input_single.unsqueeze(0),  # (1, 1, prefill_seq_len, dim)
            device=ttnn_mesh_device,
            dtype=act_dtype,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
        )
        tt_prefill_out = tt_model.forward(
            tt_prefill_input, None, prefill_rot_mats, user_id=user_id, mode="prefill", page_table=page_table_tt
        )
        tt_prefill_out = to_torch_auto_compose(tt_prefill_out)
        tt_prefill_outputs.append(tt_prefill_out[:, 0:1, :prefill_seq_len, :dim].view(1, prefill_seq_len, dim))

    # Stack outputs: (batch_size, prefill_seq_len, dim)
    tt_prefill_output = torch.cat(tt_prefill_outputs, dim=0)

    # HF reference prefill (single user, same input)
    with torch.no_grad():
        ref_prefill_output = reference_wrapper(pt_prefill_input_single, start_pos=0, mask=None)

    # Compare first user's output (all users have same input, should match)
    passing, pcc_msg = comp_pcc(ref_prefill_output, tt_prefill_output[0:1].to(ref_prefill_output.dtype), pcc)
    logger.info(f"  Prefill PCC: {pcc_msg}")
    assert passing, f"Prefill failed: {pcc_msg}"

    # =========================================================================
    # Step 2: DECODE - uses prefill output as input (realistic autoregressive flow)
    # =========================================================================
    # In real autoregressive generation:
    # - First decode input = last position of prefill output
    # - Subsequent decode inputs = previous decode output
    # This tests both KV cache integration AND data flow between modes.
    # =========================================================================
    decode_rope_setup = RotarySetupHelper(
        ttnn_mesh_device, batch_size, head_dim, max_seq_len, rotary_emb, use_qk_fused=False
    )

    # First decode input: last position of prefill output (shape: batch, 1, dim)
    pt_decode_input = tt_prefill_output[:, -1:, :].clone()  # (batch_size, 1, dim)
    ref_decode_input = ref_prefill_output[:, -1:, :].clone()  # For HF reference

    for decode_iter in range(num_decode_after_prefill):
        current_pos_value = prefill_seq_len + decode_iter

        # Prepare TT input: (batch, 1, dim) -> (1, 1, batch, dim) for TTNN decode format
        pt_decode_batched = pt_decode_input.transpose(0, 1).unsqueeze(0)  # (1, 1, batch_size, dim)

        tt_decode_input = ttnn.from_torch(
            pt_decode_batched,
            device=ttnn_mesh_device,
            dtype=act_dtype,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
        )
        tt_decode_input = ttnn.to_memory_config(tt_decode_input, tt_model.config.decode_input_memcfg)

        position_idxs = torch.full((batch_size,), current_pos_value, dtype=torch.long)
        decode_rot_mats = decode_rope_setup.get_rot_mats(position_idxs)

        current_pos = ttnn.from_torch(
            position_idxs,
            device=ttnn_mesh_device,
            dtype=ttnn.int32,
            mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(ttnn_mesh_device),
        )

        tt_decode_out = tt_model.forward(
            tt_decode_input, current_pos, decode_rot_mats, mode="decode", page_table=page_table_tt
        )
        tt_decode_out = to_torch_auto_compose(tt_decode_out)
        tt_decode_output = tt_decode_out[:, 0:1, :batch_size, :dim].view(batch_size, 1, dim)

        # HF reference decode (using same input flow)
        with torch.no_grad():
            ref_decode_output = reference_wrapper(ref_decode_input, start_pos=current_pos_value, mask=None)

        # Compare first user's output (all users have identical input, should produce identical output)
        passing, pcc_msg = comp_pcc(ref_decode_output, tt_decode_output[0:1].to(ref_decode_output.dtype), pcc)
        logger.info(f"  Decode[pos={current_pos_value}] PCC: {pcc_msg}")
        assert passing, f"Decode at position {current_pos_value} failed: {pcc_msg}"

        # Next decode input = this decode output (autoregressive flow)
        pt_decode_input = tt_decode_output.clone()
        ref_decode_input = ref_decode_output.clone()

    ttnn.SetDefaultDevice(None)

    paged_str = f"paged(block_size={page_block_size})" if page_block_size is not None else "non-paged"
    logger.info(
        f"test_attention_1d_prefill_decode_transition: PASSED ({paged_str}, "
        f"prefill={prefill_seq_len}, decode={num_decode_after_prefill})"
    )


# =============================================================================
# Focused Blackhole hardware gates
# =============================================================================


def _attention_gate_kernel(fidelity, *, approximate, fp32):
    return ttnn.WormholeComputeKernelConfig(
        math_fidelity=fidelity,
        math_approx_mode=approximate,
        fp32_dest_acc_en=fp32,
        packer_l1_acc=True,
    )


def _build_synthetic_attention_gate(mesh_device, *, paged: bool, is_blackhole: bool):
    """Build a reduced Llama layer with explicit common-config requests."""
    torch.manual_seed(2026)
    num_devices = mesh_device.get_num_devices()
    # Preserve Llama-3.1-8B attention geometry while reducing unrelated MLP
    # and vocabulary allocations in the host reference model.
    dim = 4096
    n_heads = 32
    n_kv_heads = 8
    head_dim = 128
    max_batch_size = 1
    max_seq_len = 320
    hf_config = LlamaConfig(
        vocab_size=256,
        hidden_size=dim,
        intermediate_size=512,
        num_hidden_layers=1,
        num_attention_heads=n_heads,
        num_key_value_heads=n_kv_heads,
        head_dim=head_dim,
        max_position_embeddings=max_seq_len,
        torch_dtype=torch.bfloat16,
    )
    reference_model = LlamaForCausalLM(hf_config).to(torch.bfloat16)
    reference_attention = reference_model.model.layers[0].self_attn
    rotary_emb = reference_model.model.rotary_emb
    reference = HfAttentionWrapper(reference_attention, head_dim, rotary_emb)
    wqkv, wo, _, _, _ = get_attention_weights_from_ref_model(reference_attention, num_devices)

    page_config = PagedAttentionConfig(block_size=64, max_num_blocks=32) if paged else None
    common = Attention1DConfig(
        wqkv=LazyWeight(source=wqkv, dtype=ttnn.bfloat8_b),
        wo=LazyWeight(source=wo, dtype=ttnn.bfloat8_b),
        mesh_device=mesh_device,
        tt_ccl=TT_CCL(mesh_device) if num_devices > 1 else None,
        topology=ttnn.Topology.Ring if num_devices > 1 else None,
        dim=dim,
        n_heads=n_heads,
        n_kv_heads=n_kv_heads,
        head_dim=head_dim,
        max_batch_size=max_batch_size,
        max_seq_len=max_seq_len,
        scale=head_dim**-0.5,
        paged_attention_config=page_config,
        kv_cache_dtype=ttnn.bfloat8_b,
        wqkv_dtype=ttnn.bfloat8_b,
        wo_dtype=ttnn.bfloat8_b,
        activation_dtype=ttnn.bfloat16,
        prefill_qkv_minimal_matmul=is_blackhole,
    )
    ordinary_slots = {
        name: _attention_gate_kernel(
            ttnn.MathFidelity.HiFi2,
            approximate=is_blackhole,
            fp32=is_blackhole,
        )
        for name in (
            "li_qkv_decode_compute_kernel_cfg",
            "sdpa_decode_compute_kernel_cfg",
            "li_o_decode_compute_kernel_cfg",
            "li_qkv_prefill_compute_kernel_cfg",
            "li_o_prefill_compute_kernel_cfg",
        )
    }
    common = replace(
        common,
        **ordinary_slots,
        sdpa_prefill_compute_kernel_cfg=_attention_gate_kernel(ttnn.MathFidelity.HiFi4, approximate=False, fp32=True),
        prefill_qkv_grid=(8, 10) if is_blackhole else (8, 8),
        dram_shard_grid_width=mesh_device.dram_grid_size().x if is_blackhole else 8,
        decode_create_qkv_head_grid=ttnn.CoreGrid(y=4, x=8) if is_blackhole else None,
        decode_transformation_core_grid=(
            ttnn.CoreCoord(8, 8) if is_blackhole else mesh_device.compute_with_storage_grid_size()
        ),
    )
    model = Attention1D.from_config(common)
    assert model.config.prefill_qkv_grid == ((8, 10) if is_blackhole else (8, 8))
    if is_blackhole:
        assert model.config.decode_create_qkv_head_grid.x == 8
        assert model.config.decode_create_qkv_head_grid.y == 4
    else:
        assert model.config.decode_create_qkv_head_grid is None
    assert model.config.use_minimal_qkv_matmul(256) is is_blackhole
    assert not model.config.use_minimal_qkv_matmul(128)
    for slot in (
        "li_qkv_decode_compute_kernel_cfg",
        "sdpa_decode_compute_kernel_cfg",
        "li_o_decode_compute_kernel_cfg",
        "li_qkv_prefill_compute_kernel_cfg",
        "sdpa_prefill_compute_kernel_cfg",
        "li_o_prefill_compute_kernel_cfg",
    ):
        assert getattr(model.config, slot) is not None
    return model, reference, rotary_emb, page_config


def _run_attention_standard_gate(model, reference, rotary_emb, mesh_device, *, mode):
    dim = model.config.dim
    mesh_shape = tuple(mesh_device.shape)
    if mode == "prefill":
        _run_prefill_test(
            tt_model=model,
            reference_wrapper=reference,
            ttnn_mesh_device=mesh_device,
            mesh_shape=mesh_shape,
            batch_size=1,
            seq_len=256,
            dim=dim,
            n_heads=model.config.n_heads,
            n_kv_heads=model.config.n_kv_heads,
            head_dim=model.config.head_dim,
            act_dtype=ttnn.bfloat16,
            page_block_size=None,
            chunked_prefill=False,
            chunk_size=None,
            paged_attention_config=None,
            page_table=None,
            page_table_tt=None,
            pcc=0.98,
        )
    else:
        _run_decode_test(
            tt_model=model,
            reference_wrapper=reference,
            ttnn_mesh_device=mesh_device,
            mesh_shape=mesh_shape,
            batch_size=1,
            dim=dim,
            n_heads=model.config.n_heads,
            n_kv_heads=model.config.n_kv_heads,
            head_dim=model.config.head_dim,
            max_seq_len=model.config.max_seq_len,
            act_dtype=ttnn.bfloat16,
            page_block_size=None,
            page_table_tt=None,
            num_decode_iterations=2,
            pcc=0.98,
        )


@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [
        pytest.param((1, 1), id="p150-1x1"),
        pytest.param(
            {"mesh_shape": (1, 4), "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING},
            id="p150x4-1x4",
        ),
    ],
    indirect=True,
)
@pytest.mark.parametrize("mode", ["prefill", "decode"])
def test_attention_1d_blackhole_common_config_standard_correctness_cache_and_timing(
    request, ttnn_mesh_device, require_blackhole_mesh_device, mode
):
    """Correctness/cache gate; synchronized timings are evidence without a threshold."""
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    request.addfinalizer(lambda: ttnn.SetDefaultDevice(None))
    model, reference, rotary_emb, _ = _build_synthetic_attention_gate(ttnn_mesh_device, paged=False, is_blackhole=True)
    ttnn_mesh_device.enable_program_cache()
    ttnn_mesh_device.clear_program_cache()
    request.addfinalizer(ttnn_mesh_device.disable_and_clear_program_cache)

    _run_attention_standard_gate(model, reference, rotary_emb, ttnn_mesh_device, mode=mode)
    ttnn.synchronize_device(ttnn_mesh_device)
    cache_entries = ttnn_mesh_device.num_program_cache_entries()
    assert cache_entries > 0

    timings_ms = []
    for _ in range(2):
        reference.reset_cache()
        start = time.perf_counter()
        _run_attention_standard_gate(model, reference, rotary_emb, ttnn_mesh_device, mode=mode)
        ttnn.synchronize_device(ttnn_mesh_device)
        timings_ms.append((time.perf_counter() - start) * 1000)
        assert ttnn_mesh_device.num_program_cache_entries() == cache_entries
    logger.info(
        "BH Attention1D standard measurement mode={} mesh={}: warm-cache mean={:.3f} ms, samples={}",
        mode,
        tuple(ttnn_mesh_device.shape),
        sum(timings_ms) / len(timings_ms),
        timings_ms,
    )


def _run_attention_paged_transition_gate(model, reference, rotary_emb, page_config, mesh_device):
    dim = model.config.dim
    mesh_shape = tuple(mesh_device.shape)
    seq_len = 256
    torch_input = torch.randn(1, seq_len, dim, dtype=torch.bfloat16)
    page_table = torch.randperm(page_config.max_num_blocks).reshape(1, -1)
    page_table_tt = ttnn.from_torch(
        page_table,
        device=mesh_device,
        dtype=ttnn.int32,
        layout=ttnn.ROW_MAJOR_LAYOUT,
        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(None, None), mesh_shape=mesh_shape),
    )
    tt_input = ttnn.from_torch(
        torch_input.unsqueeze(0),
        device=mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(None, None), mesh_shape=mesh_shape),
    )
    prefill_rot = get_rot_mats_from_hf(rotary_emb, seq_len, model.config.head_dim, mesh_device)
    tt_prefill = model.forward(tt_input, None, prefill_rot, mode="prefill", page_table=page_table_tt)
    tt_prefill_torch = to_torch_auto_compose(tt_prefill)[:, 0, :seq_len, :dim]
    reference_prefill = reference(torch_input, start_pos=0)
    passing, message = comp_pcc(reference_prefill, tt_prefill_torch.to(reference_prefill.dtype), 0.98)
    assert passing, f"Blackhole paged prefill PCC failed: {message}"

    decode_input = torch.randn(1, 1, dim, dtype=torch.bfloat16)
    tt_decode_input = ttnn.from_torch(
        decode_input.unsqueeze(0),
        device=mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(None, None), mesh_shape=mesh_shape),
    )
    tt_decode_input = ttnn.to_memory_config(tt_decode_input, model.config.decode_input_memcfg)
    rope = RotarySetupHelper(mesh_device, 1, model.config.head_dim, model.config.max_seq_len, rotary_emb)
    position = torch.tensor([seq_len], dtype=torch.long)
    current_pos = ttnn.from_torch(
        position,
        device=mesh_device,
        dtype=ttnn.int32,
        mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device),
    )
    tt_decode = model.forward(
        tt_decode_input,
        current_pos,
        rope.get_rot_mats(position),
        mode="decode",
        page_table=page_table_tt,
    )
    tt_decode_torch = to_torch_auto_compose(tt_decode)[:, 0, :1, :dim].view(1, 1, dim)
    reference_decode = reference(decode_input, start_pos=seq_len)
    passing, message = comp_pcc(reference_decode, tt_decode_torch.to(reference_decode.dtype), 0.98)
    assert passing, f"Blackhole paged decode PCC failed: {message}"


@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [
        pytest.param((1, 1), id="p150-1x1"),
        pytest.param(
            {"mesh_shape": (1, 4), "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING},
            id="p150x4-1x4",
        ),
    ],
    indirect=True,
)
def test_attention_1d_blackhole_common_config_paged_prefill_decode_transition_cache_and_timing(
    request, ttnn_mesh_device, require_blackhole_mesh_device
):
    """Paged transition gate with synchronized timing evidence and stable cache count."""
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    request.addfinalizer(lambda: ttnn.SetDefaultDevice(None))
    model, reference, rotary_emb, page_config = _build_synthetic_attention_gate(
        ttnn_mesh_device, paged=True, is_blackhole=True
    )
    ttnn_mesh_device.enable_program_cache()
    ttnn_mesh_device.clear_program_cache()
    request.addfinalizer(ttnn_mesh_device.disable_and_clear_program_cache)

    _run_attention_paged_transition_gate(model, reference, rotary_emb, page_config, ttnn_mesh_device)
    ttnn.synchronize_device(ttnn_mesh_device)
    cache_entries = ttnn_mesh_device.num_program_cache_entries()
    assert cache_entries > 0

    reference.reset_cache()
    start = time.perf_counter()
    _run_attention_paged_transition_gate(model, reference, rotary_emb, page_config, ttnn_mesh_device)
    ttnn.synchronize_device(ttnn_mesh_device)
    elapsed_ms = (time.perf_counter() - start) * 1000
    assert ttnn_mesh_device.num_program_cache_entries() == cache_entries
    logger.info(
        "BH Attention1D paged transition measurement mesh={}: warm-cache elapsed={:.3f} ms",
        tuple(ttnn_mesh_device.shape),
        elapsed_ms,
    )


@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [pytest.param((1, 1), id="n150-1x1")],
    indirect=True,
)
@pytest.mark.parametrize("mode", ["prefill", "decode"])
def test_attention_1d_wormhole_common_config_correctness_cache_and_timing(request, ttnn_mesh_device, mode):
    """Focused WH correctness/cache gate using all six explicit compute slots."""
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    request.addfinalizer(lambda: ttnn.SetDefaultDevice(None))
    model, reference, rotary_emb, _ = _build_synthetic_attention_gate(ttnn_mesh_device, paged=False, is_blackhole=False)
    ttnn_mesh_device.enable_program_cache()
    ttnn_mesh_device.clear_program_cache()
    request.addfinalizer(ttnn_mesh_device.disable_and_clear_program_cache)

    _run_attention_standard_gate(model, reference, rotary_emb, ttnn_mesh_device, mode=mode)
    ttnn.synchronize_device(ttnn_mesh_device)
    cache_entries = ttnn_mesh_device.num_program_cache_entries()
    assert cache_entries > 0

    timings_ms = []
    for _ in range(2):
        reference.reset_cache()
        start = time.perf_counter()
        _run_attention_standard_gate(model, reference, rotary_emb, ttnn_mesh_device, mode=mode)
        ttnn.synchronize_device(ttnn_mesh_device)
        timings_ms.append((time.perf_counter() - start) * 1000)
        assert ttnn_mesh_device.num_program_cache_entries() == cache_entries
    logger.info(
        "WH Attention1D standard measurement mode={} mesh={}: warm-cache mean={:.3f} ms, samples={}",
        mode,
        tuple(ttnn_mesh_device.shape),
        sum(timings_ms) / len(timings_ms),
        timings_ms,
    )


@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [pytest.param((1, 1), id="n150-1x1")],
    indirect=True,
)
def test_attention_1d_wormhole_common_config_paged_prefill_decode_transition_cache_and_timing(
    request, ttnn_mesh_device
):
    """Focused WH paged transition gate with stable program-cache evidence."""
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    request.addfinalizer(lambda: ttnn.SetDefaultDevice(None))
    model, reference, rotary_emb, page_config = _build_synthetic_attention_gate(
        ttnn_mesh_device, paged=True, is_blackhole=False
    )
    ttnn_mesh_device.enable_program_cache()
    ttnn_mesh_device.clear_program_cache()
    request.addfinalizer(ttnn_mesh_device.disable_and_clear_program_cache)

    _run_attention_paged_transition_gate(model, reference, rotary_emb, page_config, ttnn_mesh_device)
    ttnn.synchronize_device(ttnn_mesh_device)
    cache_entries = ttnn_mesh_device.num_program_cache_entries()
    assert cache_entries > 0

    reference.reset_cache()
    start = time.perf_counter()
    _run_attention_paged_transition_gate(model, reference, rotary_emb, page_config, ttnn_mesh_device)
    ttnn.synchronize_device(ttnn_mesh_device)
    elapsed_ms = (time.perf_counter() - start) * 1000
    assert ttnn_mesh_device.num_program_cache_entries() == cache_entries
    logger.info(
        "WH Attention1D paged transition measurement mesh={}: warm-cache elapsed={:.3f} ms",
        tuple(ttnn_mesh_device.shape),
        elapsed_ms,
    )


@torch.no_grad()
@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [
        (1, 1),  # single device
        (1, 2),  # 1D mesh, 2 devices
        (1, 8),  # 1D mesh, 8 devices
    ],
    ids=["1x1", "1x2", "1x8"],
    indirect=True,
)
@pytest.mark.parametrize("seq_len", (512, 32))
def test_attention_1d_vs_reference_from_model_args(ttnn_mesh_device: ttnn.MeshDevice, seq_len):
    """
    Test that Attention1D class created via from_model_args matches reference model.

    This test validates backward compatibility with the ModelArgs factory method
    and performs numerical PCC comparison against the reference attention.
    """
    from models.tt_transformers.tests.test_utils import get_ref_model_dype
    from models.tt_transformers.tt.ccl import TT_CCL
    from models.tt_transformers.tt.common import precompute_freqs
    from models.tt_transformers.tt.model_config import Mode, ModelArgs
    from models.tt_transformers.tt.rope import RotarySetup, get_rot_mats

    # Use HF_MODEL env var if set, otherwise use appropriate default based on device count
    # Multi-device requires larger models (dim >= 4096) for proper sharding
    num_devices = ttnn_mesh_device.get_num_devices()
    env_model = os.environ.get("HF_MODEL")

    if not env_model:
        # Set default model based on device configuration
        if num_devices == 1:
            # Small model works for single device
            default_model = "meta-llama/Llama-3.2-1B"
        else:
            # Multi-device needs larger model with dim >= 4096
            default_model = "meta-llama/Llama-3.1-8B-Instruct"
        os.environ["HF_MODEL"] = default_model
        logger.info(f"HF_MODEL not set, using default: {default_model}")

    dtype = ttnn.bfloat8_b
    pcc = 0.98
    batch_size = 1
    mode = "decode" if seq_len <= 32 else "prefill"

    model_args = ModelArgs(ttnn_mesh_device, max_batch_size=batch_size, max_seq_len=2048, cache_hf=True)
    model_args.n_layers = 1

    if model_args.is_galaxy:
        pytest.skip("Attention1D test only runs on non-TG devices")

    # Verify model dimensions are sufficient for multi-device
    if num_devices > 1 and model_args.dim < 4096:
        pytest.skip(f"Model dim={model_args.dim} too small for {num_devices} devices - use 8B+ models")

    state_dict = model_args.load_state_dict()

    # Load reference attention model
    first_layer_prefix = model_args.get_state_dict_prefix("Attention", 0) + "."
    partial_state_dict = {
        k[len(first_layer_prefix) :]: v for k, v in state_dict.items() if k.startswith(first_layer_prefix)
    }
    reference_model = model_args.reference_attention()
    reference_model.load_state_dict(partial_state_dict)

    # Setup RoPE transformation matrices
    rope_setup = RotarySetup(
        ttnn_mesh_device,
        batch_size,
        model_args.head_dim,
        model_args.max_seq_len,
        model_args.rope_theta,
        model_args.rope_scaling,
        model_args.use_qk_fused,
    )
    transformation_mats = rope_setup.get_both_trans_mats()

    # Precompute freqs_cis for reference model
    cos, sin = precompute_freqs(
        model_args.head_dim,
        model_args.max_seq_len * 2,
        model_args.rope_theta,
        model_args.rope_scaling.factor if model_args.rope_scaling else None,
        model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None,
        model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3",
    )
    freqs_cis = torch.complex(cos, sin)

    # Cache path
    def topology_aware_cache_path(dtype):
        if model_args.instruct:
            return model_args.model_cache_path / {
                ttnn.bfloat16: f"tensor_cache_instruct_bf16_{ttnn_mesh_device.shape}",
                ttnn.bfloat8_b: f"tensor_cache_instruct_bfp8_{ttnn_mesh_device.shape}",
            }.get(dtype, f"tensor_cache_instruct_{ttnn_mesh_device.shape}")
        return model_args.model_cache_path / {
            ttnn.bfloat16: f"tensor_cache_bf16_{ttnn_mesh_device.shape}",
            ttnn.bfloat8_b: f"tensor_cache_bfp8_{ttnn_mesh_device.shape}",
        }.get(dtype, f"tensor_cache_{ttnn_mesh_device.shape}")

    weight_cache_path = topology_aware_cache_path(dtype)

    # Create TT_CCL for multi-device
    tt_ccl = TT_CCL(ttnn_mesh_device) if num_devices > 1 else None

    # Create Attention1D via from_model_args
    tt_model = Attention1D.from_model_args(
        mesh_device=ttnn_mesh_device,
        tt_ccl=tt_ccl,
        args=model_args,
        state_dict=state_dict,
        weight_cache_path=weight_cache_path,
        layer_num=0,
        transformation_mats=transformation_mats,
        use_paged_kv_cache=False,
    )

    # Verify the model was created successfully
    assert tt_model is not None
    assert tt_model.config.is_resolved()
    assert tt_model.config.dim == model_args.dim
    assert tt_model.config.n_heads == model_args.n_heads
    assert tt_model.config.n_kv_heads == model_args.n_kv_heads

    logger.info(f"test_attention_1d_vs_reference_from_model_args: Testing mode={mode}, seq_len={seq_len}")

    if mode == "prefill":
        # Prefill mode test
        pt_attention_input = torch.randn(
            batch_size, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
        )

        # Prepare TT input
        tt_input = ttnn.from_torch(
            pt_attention_input.unsqueeze(0),  # [1, batch, seq_len, dim]
            device=ttnn_mesh_device,
            dtype=dtype,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                ttnn_mesh_device, dims=(None, None), mesh_shape=model_args.cluster_shape
            ),
        )

        # Get rot_mats for prefill
        rot_mats = get_rot_mats(
            head_dim=model_args.head_dim,
            device=ttnn_mesh_device,
            seq_len=seq_len,
            theta=model_args.rope_theta,
            rope_scaling=model_args.rope_scaling,
        )

        # Run TT model
        tt_out = tt_model.forward(
            tt_input,
            None,  # current_pos not used in prefill
            rot_mats,
            mode="prefill",
        )

        # Convert TT output to torch
        tt_out = to_torch_auto_compose(tt_out)
        tt_output_torch = tt_out[:, 0:1, :seq_len, : model_args.dim].view(batch_size, seq_len, model_args.dim)

        # Run reference model
        freqs_cis_slice = freqs_cis[:seq_len]
        reference_output = reference_model(pt_attention_input, start_pos=0, freqs_cis_i=freqs_cis_slice, mask=None)

        # Compare with PCC
        passing, pcc_message = comp_pcc(reference_output, tt_output_torch.to(reference_output.dtype), pcc)
        logger.info(f"  PCC comparison: {pcc_message}")
        logger.info(comp_allclose(reference_output, tt_output_torch.to(reference_output.dtype)))
        assert passing, f"Prefill PCC failed: {pcc_message} (expected >= {pcc})"

        logger.info(f"test_attention_1d_vs_reference_from_model_args: PASSED for mode={mode}, seq_len={seq_len}")

    else:
        # Decode mode test - requires prefill first to populate KV cache
        prefill_seq_len = 128
        pt_prefill_input = torch.randn(
            batch_size,
            prefill_seq_len,
            model_args.dim,
            dtype=get_ref_model_dype(reference_model, model_args.model_name),
        )

        # Prefill TT model
        tt_prefill_input = ttnn.from_torch(
            pt_prefill_input.unsqueeze(0),
            device=ttnn_mesh_device,
            dtype=dtype,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                ttnn_mesh_device, dims=(None, None), mesh_shape=model_args.cluster_shape
            ),
        )

        prefill_rot_mats = get_rot_mats(
            head_dim=model_args.head_dim,
            device=ttnn_mesh_device,
            seq_len=prefill_seq_len,
            theta=model_args.rope_theta,
            rope_scaling=model_args.rope_scaling,
        )

        _ = tt_model.forward(
            tt_prefill_input,
            None,
            prefill_rot_mats,
            user_id=0,
            mode="prefill",
        )

        # Prefill reference model (to populate its KV cache state)
        freqs_cis_prefill = freqs_cis[:prefill_seq_len]
        _ = reference_model(pt_prefill_input, start_pos=0, freqs_cis_i=freqs_cis_prefill, mask=None)

        # Decode pass
        current_pos = torch.tensor([prefill_seq_len])
        current_pos_tensor = ttnn.from_torch(
            current_pos,
            device=ttnn_mesh_device,
            dtype=ttnn.int32,
            mesh_mapper=ttnn.ReplicateTensorToMesh(ttnn_mesh_device),
        )

        pt_decode_input = torch.randn(
            batch_size, 1, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
        )

        # Prepare decode input for TT
        attention_input = model_args.prepare_residual_tensor_decode(
            pt_decode_input.clone(),
            model_args.get_attn_input_mem_config(Mode.DECODE),
            force_replicated=True,
        )

        decode_rot_mats = rope_setup.get_rot_mats(current_pos)

        # Run TT decode
        tt_out = tt_model.forward(
            attention_input,
            current_pos_tensor,
            decode_rot_mats,
            mode="decode",
        )

        # Convert TT output
        tt_out = to_torch_auto_compose(tt_out)
        tt_output_torch = tt_out[:, 0:1, :batch_size, : model_args.dim].view(batch_size, 1, model_args.dim)

        # Run reference decode
        freqs_cis_i = freqs_cis[prefill_seq_len, :].unsqueeze(0)
        reference_output = reference_model(
            pt_decode_input, start_pos=prefill_seq_len, freqs_cis_i=freqs_cis_i, mask=None
        )

        # Compare with PCC
        passing, pcc_message = comp_pcc(reference_output, tt_output_torch.to(reference_output.dtype), pcc)
        logger.info(f"  PCC comparison: {pcc_message}")
        logger.info(comp_allclose(reference_output, tt_output_torch.to(reference_output.dtype)))
        assert passing, f"Decode PCC failed: {pcc_message} (expected >= {pcc})"

        logger.info(f"test_attention_1d_vs_reference_from_model_args: PASSED for mode={mode}, seq_len={seq_len}")


def test_attention_1d_rejects_galaxy(expect_error):
    """Test that Attention1D.from_model_args rejects Galaxy/TG devices."""
    # Mock args with is_galaxy = True
    mock_args = MagicMock()
    mock_args.is_galaxy = True

    with expect_error(ValueError, "cannot be used for Galaxy"):
        Attention1D.from_model_args(
            mesh_device=MagicMock(),
            tt_ccl=MagicMock(),
            args=mock_args,
            state_dict={},
            weight_cache_path=None,
            layer_num=0,
            transformation_mats={},
        )


@torch.no_grad()
@pytest.mark.parametrize(
    "ttnn_mesh_device",
    [(1, 1)],
    ids=["1x1"],
    indirect=True,
)
@pytest.mark.parametrize(
    "sliding_window,seq_len,pcc",
    [
        pytest.param(64, 128, 0.99, id="sw64-seq128"),
        pytest.param(64, 256, 0.99, id="sw64-seq256"),
        pytest.param(128, 256, 0.99, id="sw128-seq256"),
    ],
)
def test_attention_1d_sliding_window(
    ttnn_mesh_device: ttnn.MeshDevice,
    sliding_window: int,
    seq_len: int,
    pcc: float,
):
    """
    Test Attention1D with sliding window attention.

    This test explicitly verifies that sliding window masking works correctly by:
    1. Using seq_len > sliding_window to ensure the window is actually applied
    2. Creating a reference implementation with sliding window causal mask
    3. Comparing TT output against the masked reference

    The sliding window limits attention to the last `sliding_window` tokens,
    which affects which KV entries are attended to during SDPA.
    """
    # Use Llama-3.2-1B as base model (fast to load, no sliding window by default)
    hf_model_name = LLAMA_1B

    mesh_shape = tuple(ttnn_mesh_device.shape)
    num_devices = ttnn_mesh_device.get_num_devices()
    batch_size = 1
    max_seq_len = max(512, seq_len * 2)

    torch.manual_seed(42)

    # Load HuggingFace model
    hf_config = AutoConfig.from_pretrained(hf_model_name)
    hf_model = AutoModelForCausalLM.from_pretrained(hf_model_name, torch_dtype=torch.bfloat16)
    reference_attn = hf_model.model.layers[0].self_attn
    rotary_emb = getattr(hf_model.model, "rotary_emb", None)

    dim = hf_config.hidden_size
    n_heads = hf_config.num_attention_heads
    n_kv_heads = hf_config.num_key_value_heads
    head_dim = dim // n_heads

    # Get weights from reference model
    wqkv_torch, wo_torch, _, _, _ = get_attention_weights_from_ref_model(reference_attn, num_devices)

    # Create TT model with sliding window
    ttnn.SetDefaultDevice(ttnn_mesh_device)
    lazy_wqkv = LazyWeight(source=wqkv_torch, dtype=ttnn.bfloat8_b, cache_dir_weight_name=None)
    lazy_wo = LazyWeight(source=wo_torch, dtype=ttnn.bfloat8_b, cache_dir_weight_name=None)

    # Note: kv_cache is auto-created by config resolution
    config = Attention1DConfig(
        wqkv=lazy_wqkv,
        wo=lazy_wo,
        mesh_device=ttnn_mesh_device,
        dim=dim,
        n_heads=n_heads,
        n_kv_heads=n_kv_heads,
        head_dim=head_dim,
        max_batch_size=batch_size,
        max_seq_len=max_seq_len,
        scale=head_dim**-0.5,
        sliding_window=sliding_window,  # Enable sliding window
        use_vllm_paged_kv_cache=False,
        activation_dtype=ttnn.bfloat16,
    )

    tt_model = Attention1D.from_config(config)

    rot_mats = get_rot_mats_from_hf(
        rotary_emb,
        seq_len=seq_len,
        head_dim=head_dim,
        device=ttnn_mesh_device,
    )

    # Prepare input
    pt_input = torch.randn(batch_size, seq_len, dim, dtype=torch.bfloat16)

    tt_input = ttnn.from_torch(
        pt_input.unsqueeze(0),  # [1, batch, seq_len, dim]
        device=ttnn_mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=ttnn.ShardTensor2dMesh(ttnn_mesh_device, dims=(None, None), mesh_shape=mesh_shape),
    )

    # Run TT model with sliding window
    tt_out = tt_model.forward(
        tt_input,
        None,  # current_pos not used in prefill
        rot_mats,
        mode="prefill",
    )

    tt_out_torch = to_torch_auto_compose(tt_out)
    tt_output = tt_out_torch[:, 0:1, :seq_len, :dim].view(batch_size, seq_len, dim)

    # Create reference with sliding window causal mask
    # Sliding window mask: position i can only attend to positions max(0, i - sliding_window + 1) to i
    def create_sliding_window_mask(seq_len: int, sliding_window: int) -> torch.Tensor:
        """Create a sliding window causal attention mask for HuggingFace attention."""
        # HF expects mask shape: (batch, 1, seq_len, seq_len) where 0 = attend, -inf = mask
        mask = torch.zeros(seq_len, seq_len, dtype=torch.bfloat16)

        for i in range(seq_len):
            # Position i can attend to [max(0, i - sliding_window + 1), i]
            for j in range(seq_len):
                if j > i:
                    # Causal: can't attend to future
                    mask[i, j] = torch.finfo(torch.bfloat16).min
                elif j < i - sliding_window + 1:
                    # Sliding window: can't attend beyond window
                    mask[i, j] = torch.finfo(torch.bfloat16).min

        return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)

    sliding_mask = create_sliding_window_mask(seq_len, sliding_window)

    # Setup reference attention using local HfAttentionWrapper
    reference_wrapper = HfAttentionWrapper(reference_attn, head_dim, rotary_emb)

    # Run reference WITH sliding window mask
    ref_output_sw = reference_wrapper(pt_input, start_pos=0, mask=sliding_mask)

    # Also run reference WITHOUT sliding window (full attention) for comparison
    # Need a fresh wrapper since it caches KV
    reference_wrapper_full = HfAttentionWrapper(reference_attn, head_dim, rotary_emb)
    ref_output_full = reference_wrapper_full(pt_input, start_pos=0, mask=None)

    ttnn.SetDefaultDevice(None)

    # Basic sanity checks
    assert not torch.isnan(tt_output).any(), "TT output contains NaN"
    assert not torch.isinf(tt_output).any(), "TT output contains Inf"
    assert tt_output.shape == (batch_size, seq_len, dim), f"Shape mismatch: {tt_output.shape}"

    logger.info(f"Sliding window test: sw={sliding_window}, seq_len={seq_len}")

    # Compare TT output (with sliding window) against HF reference (with sliding window mask)
    passing_sw, pcc_sw = comp_pcc(ref_output_sw, tt_output.to(ref_output_sw.dtype), pcc)
    logger.info(f"  PCC TT vs HF (both with sliding window): {pcc_sw}")
    assert passing_sw, f"Sliding window PCC failed: {pcc_sw} (expected >= {pcc})"

    # Verify sliding window actually has an effect by comparing to full attention
    if seq_len > sliding_window:
        # The sliding window reference should differ from full attention
        ref_diff = (ref_output_sw - ref_output_full).abs().mean()
        logger.info(f"  HF sliding window vs full attention diff: {ref_diff:.6f}")
        assert ref_diff > 1e-4, "HF sliding window mask had no effect on reference"

        # TT output should also differ from full attention
        tt_diff = (tt_output - ref_output_full).abs().mean()
        logger.info(f"  TT sliding window vs HF full attention diff: {tt_diff:.6f}")
        assert tt_diff > 1e-4, "TT sliding window had no effect"

        # The TT-vs-reference-sliding-window PCC should be higher than TT-vs-full-attention
        _, pcc_full = comp_pcc(ref_output_full, tt_output.to(ref_output_full.dtype), 0.0)
        logger.info(f"  PCC TT (sliding) vs HF (full): {pcc_full}")

        # Compare PCC values
        pcc_sw_val = float(pcc_sw)
        pcc_full_val = float(pcc_full)
        assert pcc_sw_val > pcc_full_val, (
            f"TT sliding window should match HF sliding window better than HF full attention. "
            f"PCC(TT,HF_sw)={pcc_sw_val:.4f} should be > PCC(TT,HF_full)={pcc_full_val:.4f}"
        )

    logger.info(f"test_attention_1d_sliding_window: PASSED (sw={sliding_window}, seq_len={seq_len})")