# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import math from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple, Union import torch from torch import nn import ttnn from models.common.lightweightmodule import LightweightModule from models.common.utility_functions import nearest_32 from models.tt_transformers.tt.common import RopeScaling, gather_cos_sin, get_rot_transformation_mat from models.tt_transformers.tt.prefetcher import Prefetcher from ttnn import replicate_tensor_to_mesh_mapper def get_batch_size_per_device_group( batch_size: int, use_qk_fused: bool, num_devices: int, mesh_shape: Tuple[int, ...], shard_batch_to_mesh_dim: int ) -> int: batch_size_per_device_group = ( max(batch_size // mesh_shape[shard_batch_to_mesh_dim], 1) if num_devices == 32 else batch_size ) return batch_size_per_device_group * (2 if use_qk_fused else 1) # Copied from DeepseekV3RotaryEmbedding: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py#L114 class RotaryEmbedding(nn.Module): def __init__(self, dim: int, max_position_embeddings: int, base: float, device: Optional[Any] = None) -> None: super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) # Build here to make `torch.jit.trace` work. self._set_cos_sin_cache( seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype(), ) self.max_seq_len_cached = None @staticmethod def permute_to_meta_format(cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: # Undo the HF permute 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) cos = cos.unsqueeze(0).unsqueeze(0) # [1, 1, max_seq_len, dim] sin = sin.unsqueeze(0).unsqueeze(0) # [1, 1, max_seq_len, dim] return cos, sin def _set_cos_sin_cache(self, seq_len: int, device: Any, dtype: torch.dtype) -> None: self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq.to(t.device)) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() self.register_buffer("freqs_cis", torch.complex(cos.float(), sin.float()), persistent=False) cos, sin = self.permute_to_meta_format(cos, sin) self.register_buffer("cos_cached", cos.to(dtype), persistent=False) self.register_buffer("sin_cached", sin.to(dtype), persistent=False) def forward(self, x: torch.Tensor, seq_len: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: # x: [bs, num_attention_heads, seq_len, head_size] if seq_len is None: seq_len = x.shape[-2] # Get sequence length from input tensor if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype), ) class ScaledRotaryEmbedding(RotaryEmbedding, ABC): def __init__( self, dim: int, max_position_embeddings: int, base: float, factor: float, device: Optional[Any] = None, ) -> None: self.scaling_factor = factor super().__init__(dim, max_position_embeddings, base, device) @abstractmethod def apply_scaling(self, freqs: torch.Tensor) -> torch.Tensor: pass def _set_cos_sin_cache(self, seq_len: int, device: Any, dtype: torch.dtype) -> None: self.max_seq_len_cached = seq_len freqs = 1.0 / (self.base ** (torch.arange(0, self.dim, 2)[: (self.dim // 2)].float() / self.dim)) t = torch.arange(seq_len * 2.0) freqs = self.apply_scaling(freqs) freqs = torch.outer(t, freqs).float() cos = torch.cos(freqs) sin = torch.sin(freqs) self.register_buffer("freqs_cis", torch.complex(cos.float(), sin.float()), persistent=False) cos, sin = gather_cos_sin(torch.arange(seq_len), cos, sin) self.register_buffer("cos_cached", cos.to(dtype), persistent=False) self.register_buffer("sin_cached", sin.to(dtype), persistent=False) # Copied from DeepseekV3YarnRotaryEmbedding: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py#L262 class YarnRotaryEmbedding(RotaryEmbedding): def __init__( self, dim: int, max_position_embeddings: int, base: float, factor: float, original_max_position_embeddings: int, beta_fast: float, beta_slow: float, mscale: float, mscale_all_dim: float, truncate: bool = True, device: Optional[Any] = None, ) -> None: self.scaling_factor = factor self.original_max_position_embeddings = original_max_position_embeddings self.beta_fast = beta_fast self.beta_slow = beta_slow self.mscale = mscale self.mscale_all_dim = mscale_all_dim self.truncate = truncate super().__init__(dim, max_position_embeddings, base, device) # Inverse dim formula to find dim based on number of rotations @staticmethod def yarn_find_correction_dim(num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float: return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) # Find dim range bounds based on rotations @staticmethod def yarn_find_correction_range( low_rot: float, high_rot: float, dim: int, base: float, max_position_embeddings: int, truncate: bool = True ) -> Tuple[float, float]: low = YarnRotaryEmbedding.yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) high = YarnRotaryEmbedding.yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) if truncate: low = math.floor(low) high = math.ceil(high) return max(low, 0), min(high, dim - 1) @staticmethod def yarn_get_mscale(scale: float, mscale: float) -> float: if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 @staticmethod def yarn_linear_ramp_mask(min: float, max: float, dim: int) -> torch.Tensor: if min == max: max += 0.001 # Prevent singularity linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) ramp_func = torch.clamp(linear_func, 0, 1) return ramp_func def _set_cos_sin_cache(self, seq_len: int, device: Any, dtype: torch.dtype) -> None: self.max_seq_len_cached = seq_len dim = self.dim freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) freq_inter = 1.0 / ( self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) ) low, high = YarnRotaryEmbedding.yarn_find_correction_range( self.beta_fast, self.beta_slow, dim, self.base, self.original_max_position_embeddings, self.truncate, ) inv_freq_mask = 1.0 - YarnRotaryEmbedding.yarn_linear_ramp_mask(low, high, dim // 2).to( device=device, dtype=torch.float32 ) inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask self.register_buffer("inv_freq", inv_freq, persistent=False) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) _mscale = float( YarnRotaryEmbedding.yarn_get_mscale(self.scaling_factor, self.mscale) / YarnRotaryEmbedding.yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) ) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * _mscale sin = emb.sin() * _mscale cos, sin = self.permute_to_meta_format(cos, sin) self.register_buffer("cos_cached", cos.to(dtype), persistent=False) self.register_buffer("sin_cached", sin.to(dtype), persistent=False) class LinearScaledRotaryEmbedding(ScaledRotaryEmbedding): def __init__( self, dim: int, max_position_embeddings: int, base: float, factor: float, device: Optional[Any] = None ) -> None: super().__init__(dim, max_position_embeddings, base, factor, device) def apply_scaling(self, freqs: torch.Tensor) -> torch.Tensor: return freqs / self.scaling_factor class LlamaRotaryEmbedding(ScaledRotaryEmbedding): def __init__( self, dim: int, max_position_embeddings: int, base: float, factor: float, original_max_position_embeddings: int, low_freq_factor: float, high_freq_factor: float, device: Optional[Any] = None, ) -> None: self.orig_context_len = original_max_position_embeddings self.low_freq_factor = low_freq_factor self.high_freq_factor = high_freq_factor super().__init__(dim, max_position_embeddings, base, factor, device) def apply_scaling(self, freqs: torch.Tensor) -> torch.Tensor: # Llama-3.x specific scaling # Values obtained from grid search low_freq_wavelen = self.orig_context_len / self.low_freq_factor high_freq_wavelen = self.orig_context_len / self.high_freq_factor new_freqs = [] for freq in freqs: wavelen = 2 * math.pi / freq if wavelen < high_freq_wavelen: new_freqs.append(freq) elif wavelen > low_freq_wavelen: new_freqs.append(freq / self.scaling_factor) else: assert low_freq_wavelen != high_freq_wavelen smooth = (self.orig_context_len / wavelen - self.low_freq_factor) / ( self.high_freq_factor - self.low_freq_factor ) new_freqs.append((1 - smooth) * freq / self.scaling_factor + smooth * freq) return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device) class Phi3RotaryEmbedding(ScaledRotaryEmbedding): def __init__( self, dim: int, max_position_embeddings: int, base: float, original_max_position_embeddings: int, long_factor: List[int], short_factor: List[int], device: Optional[Any] = None, ) -> None: self.orig_context_len = original_max_position_embeddings self.long_factor = long_factor self.short_factor = short_factor scale = 1024 * 128 / self.orig_context_len # Specific for Phi-3-mini-128k if scale <= 1.0: scaling_factor = 1.0 else: scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.orig_context_len)) super().__init__(dim, max_position_embeddings, base, scaling_factor, device) def apply_scaling(self, freqs: torch.Tensor) -> torch.Tensor: if self.max_seq_len_cached > self.orig_context_len: ext_factors = torch.tensor(self.long_factor, dtype=torch.float32) else: ext_factors = torch.tensor(self.short_factor, dtype=torch.float32) assert freqs.shape[-1] == ext_factors.shape[-1] return freqs / ext_factors def _set_cos_sin_cache(self, seq_len: int, device: Any, dtype: torch.dtype) -> None: self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) inv_freq_shape = torch.arange(0, self.dim, 2).float().to(device) / self.dim self.inv_freq = 1.0 / (self.base**inv_freq_shape) self.inv_freq = self.apply_scaling(self.inv_freq) freqs = torch.outer(t, self.inv_freq.to(t.device)) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.scaling_factor sin = emb.sin() * self.scaling_factor cos, sin = self.permute_to_meta_format(cos, sin) self.register_buffer("cos_cached", cos.to(dtype), persistent=False) self.register_buffer("sin_cached", sin.to(dtype), persistent=False) def rotary_embedding_factory( dim: int, max_position_embeddings: int, base: float, rope_scaling: Optional[RopeScaling] = None, device: Optional[Any] = None, ) -> Union[RotaryEmbedding, ScaledRotaryEmbedding]: if rope_scaling is None: return RotaryEmbedding(dim, max_position_embeddings, base, device) else: if rope_scaling.rope_type.value == "linear": rotary_embedding = LinearScaledRotaryEmbedding elif rope_scaling.rope_type.value == "llama3": rotary_embedding = LlamaRotaryEmbedding elif rope_scaling.rope_type.value == "yarn": rotary_embedding = YarnRotaryEmbedding elif rope_scaling.rope_type.value == "longrope": rotary_embedding = Phi3RotaryEmbedding else: raise ValueError(f"Invalid rope_scaling: {rope_scaling}") return rotary_embedding( dim=dim, max_position_embeddings=max_position_embeddings, base=base, **rope_scaling.model_dump(exclude_none=True), ) def compute_freqs_cis( dhead: int, end: int, theta: float, rope_scaling: Optional[RopeScaling] ) -> Tuple[torch.Tensor, torch.Tensor]: rotary_embedding = rotary_embedding_factory( dim=dhead, max_position_embeddings=end // 2, base=theta, rope_scaling=rope_scaling ) return rotary_embedding.freqs_cis def compute_gather_cos_sin( dhead: int, end: int, theta: float, rope_scaling: Optional[RopeScaling] ) -> Tuple[torch.Tensor, torch.Tensor]: rotary_embedding = rotary_embedding_factory( dim=dhead, max_position_embeddings=end // 2, base=theta, rope_scaling=rope_scaling ) return rotary_embedding.cos_cached, rotary_embedding.sin_cached def get_rot_mats( head_dim: int, device: Any, seq_len: int, theta: float, rope_scaling: Optional[RopeScaling], datatype: Any = ttnn.bfloat16, rot_mats_layout: ttnn.Layout = ttnn.TILE_LAYOUT, nope: bool = False, ) -> List[ttnn.Tensor]: cos_matrix, sin_matrix = compute_gather_cos_sin( dhead=head_dim, end=2 * seq_len, theta=theta, rope_scaling=rope_scaling, ) if nope: # NoPE layers (e.g. EXAONE-4.x full-attention layers): rotary must be the # identity. x*cos + rotate_half(x)*sin with cos=1, sin=0 is exactly x, for # every rotary op variant, so keep the shapes/layout and neutralize content. cos_matrix = torch.ones_like(cos_matrix) sin_matrix = torch.zeros_like(sin_matrix) cos_matrix = ttnn.from_torch( cos_matrix, device=device, layout=rot_mats_layout, dtype=datatype, mesh_mapper=replicate_tensor_to_mesh_mapper(device), ) sin_matrix = ttnn.from_torch( sin_matrix, device=device, layout=rot_mats_layout, dtype=datatype, mesh_mapper=replicate_tensor_to_mesh_mapper(device), ) return [cos_matrix, sin_matrix] def get_rot_mats_hf( head_dim: int, device: Any, seq_len: int, theta: float, rope_scaling: Optional[RopeScaling], datatype: Any = ttnn.bfloat16, layout: ttnn.Layout = ttnn.TILE_LAYOUT, ) -> List[ttnn.Tensor]: """Generate HF-format cos/sin matrices (no Meta permutation). Returns cos/sin in HF format: [c0, c1, ..., c_{d/2-1}, c0, c1, ..., c_{d/2-1}] Instead of Meta format: [c0, c0, c1, c1, ...] Args: layout: Device tensor layout. Decode caches for :class:`HfRotarySetup` use ``ROW_MAJOR``; prefill uses the default ``TILE`` layout. """ from models.tt_transformers.tt.common import precompute_freqs # Generate HF-format cos/sin directly # precompute_freqs returns cos/sin in shape [seq_len, head_dim//2] cos_freqs, sin_freqs = precompute_freqs( head_dim, seq_len * 2, # Generate for 2*seq_len to match compute_gather_cos_sin behavior theta, rope_scaling.factor if rope_scaling else None, rope_scaling.original_max_position_embeddings if rope_scaling else None, rope_scaling.rope_type.value if rope_scaling else "llama3", ) # HF format: concat freqs with itself [c0, c1, ..., c_{d/2-1}, c0, c1, ..., c_{d/2-1}] # cos_freqs and sin_freqs are [seq_len*2, head_dim//2], we need [seq_len, head_dim] # Take first seq_len rows and duplicate cos_hf = torch.cat([cos_freqs[:seq_len], cos_freqs[:seq_len]], dim=-1) # [seq_len, head_dim] sin_hf = torch.cat([sin_freqs[:seq_len], sin_freqs[:seq_len]], dim=-1) # [seq_len, head_dim] # Add batch dimensions: [1, 1, seq_len, head_dim] cos_hf = cos_hf.unsqueeze(0).unsqueeze(0) sin_hf = sin_hf.unsqueeze(0).unsqueeze(0) cos_matrix = ttnn.from_torch( cos_hf, device=device, layout=layout, dtype=datatype, mesh_mapper=replicate_tensor_to_mesh_mapper(device), ) sin_matrix = ttnn.from_torch( sin_hf, device=device, layout=layout, dtype=datatype, mesh_mapper=replicate_tensor_to_mesh_mapper(device), ) return [cos_matrix, sin_matrix] class HfRotarySetupOld(LightweightModule): """Legacy HF rope setup: HF-format cos/sin caches for ``ttnn.experimental.rotary_embedding``. Prefer :class:`HfRotarySetup` with ``ttnn.experimental.rotary_embedding_hf`` for production. """ def __init__( self, device: Any, batch_size: int, head_dim: int, max_seq_len: int, rope_theta: float, rope_scaling: Optional[RopeScaling] = None, use_qk_fused: bool = False, datatype: ttnn.DataType = ttnn.bfloat16, shard_batch_to_mesh_dim: Optional[int] = 1, # Those are kept for API compatibility with RotarySetup prefetcher: Optional[Prefetcher] = None, ) -> None: super().__init__() if use_qk_fused: raise NotImplementedError("use_qk_fused") self.batch_size = batch_size self.head_dim = head_dim self.max_seq_len = max_seq_len self.device = device # Generate the cos/sin matrices in HF format (no Meta permutation) # Generate for max_seq_len to allow slicing in prepare_inputs_prefill self.cos_matrix, self.sin_matrix = get_rot_mats_hf( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, ) self.cos_matrix_prefill, self.sin_matrix_prefill = get_rot_mats_hf( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, ) # Store 2D versions for embedding lookup (trace-compatible slicing) # Reshape from [1, 1, max_seq_len, head_dim] to [max_seq_len, head_dim] self.cos_matrix_2d = ttnn.reshape(self.cos_matrix, (max_seq_len, head_dim)) self.sin_matrix_2d = ttnn.reshape(self.sin_matrix, (max_seq_len, head_dim)) self.transformation_mat = None self.transformation_mat_prefill = None def get_rot_idxs(self, position_idxs: torch.Tensor, on_host: bool = False) -> ttnn.Tensor: assert isinstance(position_idxs, torch.Tensor), "Position ids must be a torch tensor" assert len(position_idxs.shape) == 1, "position idxs must be a [batch] tensor" batch = position_idxs.shape[0] position_idxs = position_idxs.reshape(1, batch) # [1, 1, 1, batch] assert position_idxs.shape == (1, batch), "position idxs must be a [1, batch] tensor" assert torch.min(position_idxs) >= 0, "position idxs must be non-negative" # Add padding if needed pad_size = nearest_32(batch) - batch position_idxs = torch.nn.functional.pad(position_idxs, (0, pad_size), "constant", 0) if on_host: # If tensor is on host, don't pass a mesh mapper if single-device rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) else: # On device rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) return rot_idxs def get_rot_mats( self, position_idxs: Union[torch.Tensor, ttnn.Tensor], return_rot_idxs: bool = False ) -> List[ttnn.Tensor]: """Get rotation matrices (cos/sin) for HF-style RoPE, one row per batch slot. Decode attention slices ``cos[:, :, b:b+1, :]`` / ``sin`` to ``[1, 1, 1, head_dim]`` and calls ``ttnn.experimental.rotary_embedding(..., token_idx=0)`` per batch index (trace-safe fixed loop). Prefill uses full-sequence cos/sin from separate tensors. Args: position_idxs: Per-batch positions. Device ``ttnn.Tensor`` ``[1, batch_padded]`` (``uint32``, same padding as ``get_rot_idxs``) for trace, or 1D / ``[1, batch]`` ``torch.Tensor`` (processed via ``get_rot_idxs``). Each batch slot must appear explicitly in the index tensor; there is no special case for a single position replicated across ``batch_size`` or for a Python ``int``. return_rot_idxs: If True, also return ``position_idxs`` unchanged. Returns: ``[cos, sin]`` with shape ``[1, 1, batch_padded, head_dim]``. """ if isinstance(position_idxs, ttnn.Tensor): rot_idx = position_idxs if len(rot_idx.shape) == 1: rot_idx = ttnn.unsqueeze(rot_idx, 0) cos_emb = ttnn.embedding(rot_idx, self.cos_matrix_2d, layout=ttnn.TILE_LAYOUT) sin_emb = ttnn.embedding(rot_idx, self.sin_matrix_2d, layout=ttnn.TILE_LAYOUT) cos_sliced = ttnn.unsqueeze_to_4D(cos_emb) sin_sliced = ttnn.unsqueeze_to_4D(sin_emb) elif isinstance(position_idxs, torch.Tensor): idx_1d = position_idxs.reshape(-1) rot_idx = self.get_rot_idxs(idx_1d) cos_emb = ttnn.embedding(rot_idx, self.cos_matrix_2d, layout=ttnn.TILE_LAYOUT) sin_emb = ttnn.embedding(rot_idx, self.sin_matrix_2d, layout=ttnn.TILE_LAYOUT) cos_sliced = ttnn.unsqueeze_to_4D(cos_emb) sin_sliced = ttnn.unsqueeze_to_4D(sin_emb) else: raise TypeError(f"position_idxs must be torch.Tensor or ttnn.Tensor, got {type(position_idxs)}") if return_rot_idxs: return [cos_sliced, sin_sliced], position_idxs return [cos_sliced, sin_sliced] def get_both_trans_mats(self) -> Dict[str, ttnn.Tensor]: return {"decode": self.transformation_mat, "prefill": self.transformation_mat_prefill} class HfRotarySetup(LightweightModule): """HF rope setup for ``ttnn.experimental.rotary_embedding_hf`` (decode and prefill). Decode cos/sin caches use ``ROW_MAJOR`` layout for ``ttnn.embedding`` row gather; prefill uses ``TILE`` layout via :func:`get_rot_mats_hf`. See :class:`HfRotarySetupOld` for the legacy ``rotary_embedding`` path. """ def __init__( self, device: Any, batch_size: int, head_dim: int, max_seq_len: int, rope_theta: float, rope_scaling: Optional[RopeScaling] = None, use_qk_fused: bool = False, datatype: ttnn.DataType = ttnn.bfloat16, shard_batch_to_mesh_dim: Optional[int] = 1, # Kept for API compatibility prefetcher: Optional[Prefetcher] = None, ) -> None: super().__init__() if use_qk_fused: raise NotImplementedError("use_qk_fused") self.batch_size = batch_size self.original_batch_size = batch_size self.head_dim = head_dim self.device = device self.is_mesh_device = isinstance(device, ttnn._ttnn.multi_device.MeshDevice) self.prefetcher = prefetcher self.num_devices = device.get_num_devices() if self.is_mesh_device else 1 if self.num_devices == 32: self.batch_size_per_device_group = max( self.original_batch_size // list(device.shape)[shard_batch_to_mesh_dim], 1 ) else: self.batch_size_per_device_group = self.original_batch_size # Match RotarySetup: Wormhole Galaxy reports (8, 9) storage grid; decode rope shards use (8, 8). self.core_grid = ( device.compute_with_storage_grid_size() if ttnn.get_arch_name() == "blackhole" else ttnn.CoreCoord(8, 8) ) # Decode: ROW_MAJOR cache for embedding lookup (same numerics as prefill via get_rot_mats_hf). self.cos_matrix, self.sin_matrix = get_rot_mats_hf( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, layout=ttnn.ROW_MAJOR_LAYOUT, ) self.cos_matrix_prefill, self.sin_matrix_prefill = get_rot_mats_hf( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, ) self.transformation_mat = None self.transformation_mat_prefill = None def get_rot_idxs(self, position_idxs: torch.Tensor, on_host: bool = False) -> ttnn.Tensor: assert isinstance(position_idxs, torch.Tensor), "Position ids must be a torch tensor" assert len(position_idxs.shape) == 1, "position idxs must be a [batch] tensor" batch = position_idxs.shape[0] position_idxs = position_idxs.reshape(1, batch) # [1, batch] assert position_idxs.shape == (1, batch), "position idxs must be a [1, batch] tensor" assert torch.min(position_idxs) >= 0, "Position idxs must be non-negative" if on_host: rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) else: rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) return rot_idxs def get_rot_mats( self, position_idxs: Union[torch.Tensor, ttnn.Tensor], return_rot_idxs: bool = False, ) -> List[ttnn.Tensor]: """Get rotation matrices (cos/sin) for decode mode with per-batch positions. This method extracts cos/sin values from the cache for each batch element's position. Returns tensors shaped [1, batch, 1, head_dim] for use with rotary_embedding_hf. The cos/sin tensors are placed in HEIGHT_SHARDED memory with the same core grid layout that nlp_create_qkv_heads_decode uses for Q/K, ensuring the sharded rotary_embedding_hf kernel reads matching data on each core. Args: position_idxs: [batch] tensor of positions, one per batch element return_rot_idxs: If True, also return the processed rotation indices. Returns: List of [cos, sin] tensors, each shaped [1, batch, 1, head_dim]. If return_rot_idxs=True, returns ([cos, sin], rot_idxs). """ device = self.device if isinstance(position_idxs, torch.Tensor): rot_idxs = self.get_rot_idxs(position_idxs) else: rot_idxs = position_idxs assert len(rot_idxs.shape) == 2 and rot_idxs.shape[0] == 1, "rot_idxs must be a [1, batch] tensor" if rot_idxs.device != device: rot_idxs = ttnn.to_device(rot_idxs, device, memory_config=ttnn.DRAM_MEMORY_CONFIG) embedding_layout = ttnn.TILE_LAYOUT cos = ttnn.embedding( rot_idxs, self.cos_matrix, layout=embedding_layout, memory_config=ttnn.DRAM_MEMORY_CONFIG ) # [1, batch, head_dim] sin = ttnn.embedding( rot_idxs, self.sin_matrix, layout=embedding_layout, memory_config=ttnn.DRAM_MEMORY_CONFIG ) # [1, batch, head_dim] cos = ttnn.unsqueeze_to_4D(cos) # [1, 1, batch, head_dim] sin = ttnn.unsqueeze_to_4D(sin) # [1, 1, batch, head_dim] cos = ttnn.transpose(cos, 1, 2) # [1, batch, 1(padded to 32), head_dim] sin = ttnn.transpose(sin, 1, 2) # [1, batch, 1(padded to 32), head_dim] batch_decode = self.batch_size_per_device_group num_cores = min(batch_decode, self.core_grid.x * self.core_grid.y) batch_grid = ttnn.num_cores_to_corerangeset(num_cores, self.core_grid, row_wise=True) mem_config = ttnn.create_sharded_memory_config( shape=(ttnn.TILE_SIZE, self.head_dim), core_grid=batch_grid, strategy=ttnn.ShardStrategy.HEIGHT, orientation=ttnn.ShardOrientation.ROW_MAJOR, use_height_and_width_as_shard_shape=True, ) if batch_decode % ttnn.TILE_SIZE != 0: cos = cos[:, :batch_decode, :, :] sin = sin[:, :batch_decode, :, :] cos = ttnn.interleaved_to_sharded(cos, mem_config) sin = ttnn.interleaved_to_sharded(sin, mem_config) if return_rot_idxs: return [cos, sin], rot_idxs return [cos, sin] def get_both_trans_mats(self) -> Dict[str, ttnn.Tensor]: return {"decode": self.transformation_mat, "prefill": self.transformation_mat_prefill} class RotarySetup(LightweightModule): def __init__( self, device: Any, batch_size: int, head_dim: int, max_seq_len: int, rope_theta: float, rope_scaling: Optional[RopeScaling] = None, use_qk_fused: bool = False, datatype: ttnn.DataType = ttnn.bfloat16, shard_batch_to_mesh_dim: Optional[int] = 1, prefetcher: Optional[Prefetcher] = None, nope: bool = False, ) -> None: super().__init__() self.use_qk_fused = use_qk_fused self.original_batch_size = batch_size self.prefetcher = prefetcher # NOTE: If qk fused ops (rotary embedding + paged cache update) are used # we need to double the batch size in order to replicate the transformation matrix on double the batch size number of cores self.doubled_batch_size = self.original_batch_size * 2 if use_qk_fused else self.original_batch_size self.head_dim = head_dim self.device = device self.is_mesh_device = isinstance(device, ttnn._ttnn.multi_device.MeshDevice) self.num_devices = device.get_num_devices() if self.is_mesh_device else 1 self.batch_size_per_device_group = get_batch_size_per_device_group( self.original_batch_size, use_qk_fused, self.num_devices, tuple(device.shape) if self.is_mesh_device else (), shard_batch_to_mesh_dim, ) # Always use (8, 8) on wormhole (compute_with_storage_grid_size returns (8, 9) on Galaxy) self.core_grid = ( device.compute_with_storage_grid_size() if ttnn.get_arch_name() == "blackhole" else ttnn.CoreCoord(8, 8) ) self.start_core = ttnn.CoreCoord(1, 0) # Generate the cos/sin matrices needed for ttnn.embedding op self.cos_matrix, self.sin_matrix = get_rot_mats( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, rot_mats_layout=ttnn.ROW_MAJOR_LAYOUT, nope=nope, ) self.cos_matrix_prefill, self.sin_matrix_prefill = get_rot_mats( head_dim=head_dim, device=device, seq_len=max_seq_len, theta=rope_theta, rope_scaling=rope_scaling, datatype=datatype, rot_mats_layout=ttnn.TILE_LAYOUT, nope=nope, ) def get_batch_grid(batch_size, core_grid, start_core, batch_size_per_device_group, prefetcher): if ttnn.get_arch_name() == "blackhole": if prefetcher is not None: return ttnn.num_cores_to_corerangeset_in_subcoregrids( start_core, batch_size_per_device_group, prefetcher.all_worker_cores_range_set, row_wise=True, ) else: # Use batch_size (which is doubled_batch_size for fused QK) to determine the number of cores if batch_size % 32 == 0: return ttnn.CoreGrid(y=8, x=8) return ttnn.num_cores_to_corerangeset(batch_size, core_grid, row_wise=True) else: return ttnn.num_cores_to_corerangeset(batch_size, core_grid, row_wise=True) self.batch_grid = get_batch_grid( self.batch_size_per_device_group, self.core_grid, self.start_core, self.batch_size_per_device_group, self.prefetcher, ) # Generate the transformation matrix trans_mat = get_rot_transformation_mat(dhead=ttnn.TILE_SIZE).repeat( 1, 1, self.batch_size_per_device_group, 1, # 1, 1, num_cores, 1 ) # Repeat across all cores on device 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=replicate_tensor_to_mesh_mapper(device), ) # TODO: Colman, should this be TILE_SIZE or head_dim? Why should it be different for prefill and decode? prefill_trans_mat_torch = get_rot_transformation_mat(dhead=head_dim) self.transformation_mat_prefill = ttnn.from_torch( prefill_trans_mat_torch, device=device, layout=ttnn.TILE_LAYOUT, dtype=datatype, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper=replicate_tensor_to_mesh_mapper(device), ) def get_trans_mat_on_sub_core_grids( self, x: ttnn.Tensor, sub_core_grids: ttnn.CoreRangeSet, use_qk_fused: bool = False ) -> ttnn.Tensor: # Reshape the cos/sin matrices to the sub-core grids # This function is only used when sub core grids is enabled via prefetcher orig_mem_cfg = x.memory_config() x_dram = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) x_reshaped = ttnn.reshape(x_dram, [1, 64 if use_qk_fused else 32, 1, 128], sub_core_grids=sub_core_grids) x_restored = ttnn.to_memory_config(x_reshaped, orig_mem_cfg) return x_restored def get_both_trans_mats(self) -> Dict[str, ttnn.Tensor]: assert self.transformation_mat is not None, "Transformation matrix not initialized" assert self.transformation_mat_prefill is not None, "Prefill Transformation matrix not initialized" 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: assert isinstance(position_idxs, torch.Tensor), "Position ids must be a torch tensor" assert len(position_idxs.shape) == 1, "position idxs must be a [batch] tensor" if self.use_qk_fused: # NOTE: For fused QK ops (rotary embedding + paged cache update), we intentionally double the batch dimension so that # the rotary indices can be used for Q and K tensors each. position_idxs = position_idxs.repeat(2) assert ( position_idxs.shape[0] == self.batch_size_per_device_group ), "Position idxs must be the same as the batch size per device group" batch = position_idxs.shape[0] position_idxs = position_idxs.reshape(1, batch) # [1, 1, 1, batch] assert position_idxs.shape == (1, batch), "Position idxs must be a [1, batch] tensor" assert torch.min(position_idxs) >= 0, "Position idxs must be non-negative" # Add padding if needed pad_size = nearest_32(batch) - batch position_idxs = torch.nn.functional.pad(position_idxs, (0, pad_size), "constant", 0) if on_host: # If tensor is on host, don't pass a mesh mapper if single-device rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) else: # On device rot_idxs = ttnn.as_tensor( position_idxs, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, device=self.device, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper=replicate_tensor_to_mesh_mapper(self.device), ) return rot_idxs def get_rot_mats( self, position_idxs: Union[torch.Tensor, ttnn.Tensor], return_rot_idxs: bool = False, ) -> List[ttnn.Tensor]: """Get rotation matrices (cos/sin) for specific positions using Meta-style RoPE. This method is designed for use with standard rotary embedding operations (e.g., rotary_embedding_llama). It slices the cos/sin cache by position indices and returns batch-specific, sharded rotation matrices. NOTE: This behaves differently from HfRotarySetupOld.get_rot_mats() due to different underlying RoPE implementations: - RotarySetup (this class): Uses Meta-style RoPE with embedding-based position slicing. Returns cos/sin matrices sliced by position_idxs and sharded across batch dimension. - HfRotarySetupOld: Legacy HF-style RoPE (``ttnn.experimental.rotary_embedding``) which expects the full cos/sin cache. Returns the raw, unsliced cache matrices. Args: position_idxs: Position indices to slice rotation matrices. Can be torch.Tensor or ttnn.Tensor. return_rot_idxs: If True, also return the processed rotation indices. Returns: List of [cos, sin] tensors sliced and sharded for the given positions. If return_rot_idxs=True, returns ([cos, sin], rot_idxs). """ device = self.device # If position_idxs is a torch tensor, get the TTNN version of it if isinstance(position_idxs, torch.Tensor): rot_idxs = self.get_rot_idxs(position_idxs) else: rot_idxs = position_idxs assert len(rot_idxs.shape) == 2 and rot_idxs.shape[0] == 1, "rot_idxs must be a [1, batch] tensor" # Send the idxs to device if rot_idxs.device != device: rot_idxs = ttnn.to_device(rot_idxs, device, memory_config=ttnn.DRAM_MEMORY_CONFIG) embedding_layout = ttnn.TILE_LAYOUT if self.prefetcher is not None: trans_mat_core_grids = ttnn.num_cores_to_corerangeset_in_subcoregrids( self.start_core, self.batch_size_per_device_group, self.prefetcher.all_worker_cores_range_set, row_wise=True, ) mem_config = ttnn.create_sharded_memory_config( shape=(ttnn.TILE_SIZE, self.head_dim), core_grid=trans_mat_core_grids, strategy=ttnn.ShardStrategy.HEIGHT, orientation=ttnn.ShardOrientation.ROW_MAJOR, use_height_and_width_as_shard_shape=True, ) else: trans_mat_core_grids = None mem_config = ttnn.DRAM_MEMORY_CONFIG cos = ttnn.embedding( rot_idxs, self.cos_matrix, layout=embedding_layout, memory_config=mem_config ) # [1, batch, head_dim] sin = ttnn.embedding( rot_idxs, self.sin_matrix, layout=embedding_layout, memory_config=mem_config ) # [1, batch, head_dim] if self.batch_size_per_device_group % ttnn.TILE_SIZE == 0 and self.prefetcher is not None: cos = self.get_trans_mat_on_sub_core_grids(cos, trans_mat_core_grids, use_qk_fused=self.use_qk_fused) sin = self.get_trans_mat_on_sub_core_grids(sin, trans_mat_core_grids, use_qk_fused=self.use_qk_fused) cos = ttnn.unsqueeze_to_4D(cos) # [1, 1, batch, head_dim] sin = ttnn.unsqueeze_to_4D(sin) # [1, 1, batch, head_dim] if self.prefetcher is None: cos = ttnn.transpose(cos, 1, 2) # [1, batch, 1[32], head_dim] sin = ttnn.transpose(sin, 1, 2) # [1, batch, 1[32], head_dim] 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, ) 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, :, :] cos = ttnn.interleaved_to_sharded( cos, mem_config ) # [1, 1 (= batch / shard_num_cores), 1[32], self.head_dim] sin = ttnn.interleaved_to_sharded( sin, mem_config ) # [1, 1 (= batch / shard_num_cores), 1[32], self.head_dim] if return_rot_idxs: return [cos, sin], rot_idxs return [cos, sin]