code/models/common/modules/rope/rope_1d.py

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

"""
TTTv2-style Rotary Position Embedding (RoPE) setup for 1D-topology devices.

RotarySetup1D takes pre-computed cos/sin rotation matrices (as LazyWeight) and
provides efficient position-based lookup at runtime.

Core API:
  - decode_forward(rot_idxs): Look up cos/sin rotation matrices from ttnn.Tensor indices
  - prefill_forward(start_pos, seq_len): Slice cos/sin matrices for prefill
  - get_both_trans_mats(): Decode + prefill transformation matrices (read-only)

Backward-compatible API (will retire with TTTv1):
  - get_rot_idxs(position_idxs): Convert torch position indices to ttnn tensor
  - get_rot_mats(position_idxs): Combined get_rot_idxs + decode_forward

Helper:
  - prepare_rot_idxs(config, position_idxs, on_host): Standalone torch→ttnn index conversion

Note: torch is used at construction time for building transformation matrices.
      The decode_forward() and prefill_forward() methods use only ttnn ops.
"""

from dataclasses import dataclass, replace
from typing import Any, Dict, List, Optional, Tuple, Union

import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight
from models.common.tensor_utils import TILE_SIZE, get_rot_transformation_mat, nearest_32

# =============================================================================
# Config dataclass
# =============================================================================


@dataclass
class Rope1DConfig:
    """
    Configuration for RotarySetup1D.

    Simple usage:
        rope = RotarySetup1D(cos_lw, sin_lw, max_batch_size=32)

    Power API:
        config = Rope1DConfig(cos_lw, sin_lw, max_batch_size=32,
                              device=mesh_device, use_qk_fused=True)
        rope = RotarySetup1D.from_config(config)
    """

    # Required: cos/sin embedding lookup tables as LazyWeight
    # source shape: [1, 1, max_seq_len, head_dim]
    cos_matrix: LazyWeight
    sin_matrix: LazyWeight

    # Required scalars
    max_batch_size: int

    # Optional (derived from weights if None)
    head_dim: int | None = None  # derived from cos_matrix.source.shape[-1]
    device: Any | None = None  # derived from cos_matrix.device
    use_qk_fused: bool = False
    datatype: ttnn.DataType = ttnn.bfloat16

    # Derived (resolved by _resolve_rope_config, None = auto)
    batch_size_per_device_group: int | None = None
    core_grid: Any | None = None
    batch_grid: Any | None = None
    decode_trans_mat_mem_config: Optional[ttnn.MemoryConfig] = None
    cos_sin_shard_mem_config: Optional[ttnn.MemoryConfig] = None

    # Internal: transformation matrix LazyWeights (built by _resolve_rope_config)
    _decode_trans_mat: Optional[LazyWeight] = None
    _prefill_trans_mat: Optional[LazyWeight] = None


# =============================================================================
# RotarySetup1D
# =============================================================================


class RotarySetup1D(LightweightModule):
    """
    Rotary position embedding setup for non-TG (1D) devices.

    Takes pre-computed cos/sin lookup tables as LazyWeight and provides
    efficient position-based rotation matrix lookup at runtime.

    Simple API:
        rope = RotarySetup1D(cos_lw, sin_lw, max_batch_size=32)
        rot_idxs = prepare_rot_idxs(rope.config, position_idxs)
        [cos, sin] = rope.decode_forward(rot_idxs)
        [cos, sin] = rope.prefill_forward(start_pos=0, seq_len=128)

    Power API:
        config = Rope1DConfig(cos_lw, sin_lw, max_batch_size=32, use_qk_fused=True)
        rope = RotarySetup1D.from_config(config)

    Backward-compatible API:
        rope.get_rot_mats(position_idxs)  # accepts torch.Tensor or ttnn.Tensor
        rope.get_rot_idxs(position_idxs, on_host=True)
    """

    def __init__(
        self,
        cos_matrix: LazyWeight,
        sin_matrix: LazyWeight,
        max_batch_size: int,
    ) -> None:
        """
        Simple API — takes cos/sin as LazyWeight, derives all other config.

        Args:
            cos_matrix: LazyWeight with source shape [1, 1, max_seq_len, head_dim].
            sin_matrix: LazyWeight with source shape [1, 1, max_seq_len, head_dim].
            max_batch_size: Maximum batch size for decode mode.
        """
        super().__init__()
        config = Rope1DConfig(
            cos_matrix=cos_matrix,
            sin_matrix=sin_matrix,
            max_batch_size=max_batch_size,
        )
        self.config = _resolve_rope_config(config)
        self._device_weights_loaded = False

    @classmethod
    def from_config(cls, config: Rope1DConfig):
        """Power API — construct from config dataclass with full customization."""
        instance = object.__new__(cls)
        super(RotarySetup1D, instance).__init__()
        instance.config = _resolve_rope_config(config)
        instance._device_weights_loaded = False
        return instance

    def load_device_weights(self):
        """Load cos/sin and transformation matrices to device. Called lazily on first use."""
        if self._device_weights_loaded:
            return

        cfg = self.config

        self.cos_matrix = cfg.cos_matrix.get_device_weight()
        self.sin_matrix = cfg.sin_matrix.get_device_weight()
        self.transformation_mat = cfg._decode_trans_mat.get_device_weight()
        self.transformation_mat_prefill = cfg._prefill_trans_mat.get_device_weight()

        self._device_weights_loaded = True

    # ---- Core API ----

    def decode_forward(self, rot_idxs: ttnn.Tensor) -> List[ttnn.Tensor]:
        """Look up cos/sin rotation matrices for given position indices (decode mode).

        Args:
            rot_idxs: ttnn.Tensor of shape [1, batch], dtype uint32.

        Returns:
            [cos, sin] — height-sharded rotation matrices.
        """
        self.load_device_weights()
        cfg = self.config

        # Send to device if needed
        if rot_idxs.device != cfg.device:
            rot_idxs = ttnn.to_device(rot_idxs, cfg.device, memory_config=ttnn.DRAM_MEMORY_CONFIG)

        # Embedding lookup: [1, batch] → [1, batch, head_dim]
        cos = ttnn.embedding(rot_idxs, self.cos_matrix, layout=ttnn.TILE_LAYOUT)
        sin = ttnn.embedding(rot_idxs, self.sin_matrix, layout=ttnn.TILE_LAYOUT)

        # Reshape: [1, batch, head_dim] → [1, 1, batch, head_dim] → [1, batch, 1[32], head_dim]
        cos = ttnn.unsqueeze_to_4D(cos)
        sin = ttnn.unsqueeze_to_4D(sin)
        cos = ttnn.transpose(cos, 1, 2)
        sin = ttnn.transpose(sin, 1, 2)

        # Trim to actual batch size.
        # When batch_size_per_device_group is not a multiple of TILE_SIZE (e.g., batch=1),
        # the embedding produces nearest_32(batch) rows that need trimming.
        # When tile-aligned (e.g., batch=32), this is a no-op slice.
        cos = cos[:, : cfg.batch_size_per_device_group, :, :]
        sin = sin[:, : cfg.batch_size_per_device_group, :, :]

        # Shard to batch cores
        cos = ttnn.interleaved_to_sharded(cos, cfg.cos_sin_shard_mem_config)
        sin = ttnn.interleaved_to_sharded(sin, cfg.cos_sin_shard_mem_config)

        return [cos, sin]

    def prefill_forward(self, start_pos: int, seq_len: int, pad_to: int | None = None) -> List[ttnn.Tensor]:
        """Slice cos/sin matrices for prefill mode.

        This replaces the duplicated cos/sin slicing logic found in every model's
        prepare_inputs_prefill(). Models no longer need to reach into
        rope_setup.cos_matrix / sin_matrix directly.

        Args:
            start_pos: Starting position in the sequence.
            seq_len: Number of positions to slice (typically the padded input length S).
            pad_to: If set, zero-pad the sequence dim to this length (for SDPA alignment).
                    Ignored when pad_to <= seq_len.

        Returns:
            [cos_slice, sin_slice] — interleaved DRAM tensors,
            shape [1, 1, seq_len (or pad_to), head_dim].
        """
        self.load_device_weights()

        mat_len = self.cos_matrix.shape[2]
        end_pos = start_pos + seq_len
        assert mat_len >= end_pos, f"Requested range [{start_pos}:{end_pos}) exceeds cos/sin table length {mat_len}"

        cos_slice = self.cos_matrix[:, :, start_pos:end_pos, :]
        sin_slice = self.sin_matrix[:, :, start_pos:end_pos, :]

        if pad_to is not None and pad_to > seq_len:
            pad_len = pad_to - seq_len
            padding = [(0, 0)] * 4
            padding[2] = (0, pad_len)
            cos_slice = ttnn.pad(cos_slice, padding=padding, value=0.0)
            sin_slice = ttnn.pad(sin_slice, padding=padding, value=0.0)

        return [cos_slice, sin_slice]

    def forward(self, mode: str, **kwargs) -> List[ttnn.Tensor]:
        """Dispatch to decode_forward or prefill_forward based on mode.

        Args:
            mode: "decode" or "prefill".
            **kwargs: Forwarded to the underlying method.
                decode:  rot_idxs (ttnn.Tensor)
                prefill: start_pos (int), seq_len (int), pad_to (int | None)
        """
        if mode == "decode":
            return self.decode_forward(**kwargs)
        elif mode == "prefill":
            return self.prefill_forward(**kwargs)
        else:
            raise ValueError(f"Unknown mode: {mode!r}. Expected 'decode' or 'prefill'.")

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

    # ---- Backward-compatible API (will retire with TTTv1) ----

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

    def get_rot_mats(
        self,
        position_idxs: "Union[torch.Tensor, ttnn.Tensor]",
        return_rot_idxs: bool = False,
    ) -> "Union[List[ttnn.Tensor], Tuple[List[ttnn.Tensor], ttnn.Tensor]]":
        """Look up cos/sin from torch or ttnn position indices.

        For new code, prefer: rot_idxs = prepare_rot_idxs(...); rope.decode_forward(rot_idxs)
        """
        import torch

        if isinstance(position_idxs, torch.Tensor):
            rot_idxs = self.get_rot_idxs(position_idxs)
        else:
            rot_idxs = position_idxs

        cos_sin = self.decode_forward(rot_idxs)

        if return_rot_idxs:
            return cos_sin, rot_idxs
        return cos_sin

    # ---- Factory for TTTv1 backward compatibility ----

    @classmethod
    def from_model_args(cls, device: Any, args, model_name: str = "unknown"):
        """Factory method for backward compatibility with ModelArgs.

        Args:
            device: MeshDevice or single device.
            args: Model arguments (ModelArgs instance).
            model_name: Model name for logging (default "unknown").
        """
        if args.is_galaxy:
            raise ValueError("RotarySetup1D cannot be used for Galaxy devices.")

        from models.tt_transformers.tt.common import rope_scaling_model_factory

        rope_scaling = None
        if hasattr(args, "rope_scaling_params") and args.rope_scaling_params is not None:
            rope_scaling = rope_scaling_model_factory(
                args.rope_scaling_params, getattr(args, "original_max_context_len", None)
            )

        # Compute cos/sin torch tensors using TTTv1 reference
        from models.tt_transformers.tt.rope import compute_gather_cos_sin

        cos_torch, sin_torch = compute_gather_cos_sin(
            dhead=args.head_dim, end=2 * args.max_seq_len, theta=args.rope_theta, rope_scaling=rope_scaling
        )

        # Wrap in LazyWeight
        cos_lw = LazyWeight(source=cos_torch, device=device)
        sin_lw = LazyWeight(source=sin_torch, device=device)

        config = Rope1DConfig(
            cos_matrix=cos_lw,
            sin_matrix=sin_lw,
            max_batch_size=args.max_batch_size,
            head_dim=args.head_dim,
            device=device,
            use_qk_fused=getattr(args, "use_qk_fused", False),
            datatype=ttnn.bfloat16,
        )
        return cls.from_config(config)


# =============================================================================
# Config resolver
# =============================================================================


def _resolve_rope_config(config: Rope1DConfig) -> Rope1DConfig:
    """Resolve None fields in Rope1DConfig to sensible defaults.

    Fills in:
      - head_dim: derived from cos_matrix.source.shape[-1]
      - device: derived from cos_matrix.device if None
      - batch_size_per_device_group, core_grid, batch_grid
      - decode_trans_mat_mem_config, cos_sin_shard_mem_config
      - cos_matrix/sin_matrix LazyWeight fields: device, memory_config, layout, dtype
      - _decode_trans_mat, _prefill_trans_mat: internal transformation matrix LazyWeights
    """
    to_set = {}

    # --- Phase 1: Foundational fields ---

    head_dim = config.head_dim
    if head_dim is None:
        head_dim = config.cos_matrix.source.shape[-1]
        to_set["head_dim"] = head_dim

    device = config.device
    if device is None:
        device = config.cos_matrix.device
    if device is None:
        device = ttnn.GetDefaultDevice()
    if config.device is None:
        to_set["device"] = device

    assert device is not None, "device must be available at this point!"

    # --- Phase 2: Derived scalars ---

    batch_size_per_device_group = config.batch_size_per_device_group
    if batch_size_per_device_group is None:
        batch_size_per_device_group = config.max_batch_size * 2 if config.use_qk_fused else config.max_batch_size
        to_set["batch_size_per_device_group"] = batch_size_per_device_group

    core_grid = config.core_grid
    if core_grid is None:
        core_grid = device.compute_with_storage_grid_size()
        to_set["core_grid"] = core_grid

    batch_grid = config.batch_grid
    if batch_grid is None:
        batch_grid = ttnn.num_cores_to_corerangeset(batch_size_per_device_group, core_grid, row_wise=True)
        to_set["batch_grid"] = batch_grid

    # --- Phase 3: Memory configs ---

    if config.decode_trans_mat_mem_config is None:
        to_set["decode_trans_mat_mem_config"] = ttnn.create_sharded_memory_config(
            shape=(TILE_SIZE, TILE_SIZE),
            core_grid=batch_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )

    if config.cos_sin_shard_mem_config is None:
        to_set["cos_sin_shard_mem_config"] = ttnn.create_sharded_memory_config(
            shape=(TILE_SIZE, head_dim),
            core_grid=batch_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )

    # --- Phase 4: Resolve LazyWeights ---

    to_set["cos_matrix"] = resolve_lazy_weight(
        config.cos_matrix,
        device=device,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper_config=None,  # replicated
        layout=ttnn.TILE_LAYOUT,
        dtype=config.datatype,
    )
    to_set["sin_matrix"] = resolve_lazy_weight(
        config.sin_matrix,
        device=device,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper_config=None,  # replicated
        layout=ttnn.TILE_LAYOUT,
        dtype=config.datatype,
    )

    # --- Phase 5: Transformation matrix LazyWeights ---

    decode_trans_mat_mem_config = to_set.get("decode_trans_mat_mem_config", config.decode_trans_mat_mem_config)

    assert config._decode_trans_mat is None, "_decode_trans_mat is internal and should not be set by the user"
    decode_trans_mat_source = get_rot_transformation_mat(dhead=TILE_SIZE).repeat(
        1,
        1,
        batch_size_per_device_group,
        1,
    )
    to_set["_decode_trans_mat"] = LazyWeight(
        source=decode_trans_mat_source,
        device=device,
        memory_config=decode_trans_mat_mem_config,
        mesh_mapper_config=None,  # replicated
        layout=ttnn.TILE_LAYOUT,
        dtype=config.datatype,
    )

    assert config._prefill_trans_mat is None, "_prefill_trans_mat is internal and should not be set by the user"
    prefill_trans_mat_source = get_rot_transformation_mat(dhead=TILE_SIZE)
    to_set["_prefill_trans_mat"] = LazyWeight(
        source=prefill_trans_mat_source,
        device=device,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper_config=None,  # replicated
        layout=ttnn.TILE_LAYOUT,
        dtype=config.datatype,
    )

    return replace(config, **to_set)


# =============================================================================
# Standalone helper function
# =============================================================================


def prepare_rot_idxs(
    config: Rope1DConfig,
    position_idxs: "torch.Tensor",
    on_host: bool = False,
) -> ttnn.Tensor:
    """Convert torch position indices to a TTNN rotation index tensor.

    This is the standalone equivalent of RotarySetup1D.get_rot_idxs().
    Useful for pre-allocating host tensors for trace mode.

    Args:
        config: Resolved Rope1DConfig (device must be set).
        position_idxs: 1D torch tensor of position indices, shape [batch].
        on_host: If True, return tensor on host (for later copy_host_to_device_tensor).

    Returns:
        ttnn tensor of shape [1, nearest_32(batch)], dtype uint32.
    """
    import torch

    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"
    assert config.device is not None, "config.device must be set (use a resolved config)"
    assert config.batch_size_per_device_group is not None, "config must be resolved (_resolve_rope_config)"

    if config.use_qk_fused:
        position_idxs = position_idxs.repeat(2)
        assert (
            position_idxs.shape[0] == config.batch_size_per_device_group
        ), "Position idxs must match batch_size_per_device_group"

    batch = position_idxs.shape[0]
    position_idxs = position_idxs.reshape(1, batch)
    assert torch.min(position_idxs) >= 0, "Position idxs must be non-negative"

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

    device = config.device

    if on_host:
        rot_idxs = ttnn.as_tensor(
            position_idxs,
            dtype=ttnn.uint32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(device),
        )
    else:
        rot_idxs = ttnn.as_tensor(
            position_idxs,
            dtype=ttnn.uint32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            device=device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(device),
        )

    return rot_idxs