Hugging Face's logo

sunandbean
/
diffrhythm-1.2-p100a

Text-to-Audio
ONNX
blackhole
diffrhythm
music-generation
p100a
tenstorrent
tt-dit-server
tt-metal
tt-model-cache
tt-model-catalog
tt-model-container
tt-nn
ttnn
Model card Files
xet
 Community
diffrhythm-1.2-p100a / code /models /common /tests /modules /attention /test_attention_1d.py
sunandbean's picture
sunandbean
Add files using upload-large-folder tool
656477e verified
Raw History Blame Contribute Delete
134 kB
# 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})")