code/models/tt_transformers/tt/rope.py

# 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]