model_v2.py

# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.

"""Full definition of a decoder-only transformer-based language model, all of it in this single file.

Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
"""

import math
from typing import Any, Optional, Tuple
from typing import Dict, List, TypeVar

import torch
import torch.nn as nn
import torch.nn.functional as F

from typing_extensions import Self
from functools import partial

from litgpt.config_v2 import Config
from mamba_ssm.modules.mamba_simple import Mamba
from mamba_ssm.modules.mamba_simple import Block as MBBlock
from einops import einsum, rearrange

class GPT(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        assert config.padded_vocab_size is not None
        self.config = config

        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias)

        if config.rnn_type == "Mamba":
            self.transformer = nn.ModuleDict(
                dict(
                    wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                    h=nn.ModuleList(MambaBlock(config, layer_idx=layer_idx) for layer_idx in range(config.n_layer * 2)),
                    ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
                )
            )
        else:
            if config.use_mod:
                self.transformer = nn.ModuleDict(
                    dict(
                        wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                        h=nn.ModuleList(Block(config) if _ % 2 == 0 else MoD(config, Block(config)) for _ in range(config.n_layer)),
                        ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
                    )
                )
            else:
                self.transformer = nn.ModuleDict(
                    dict(
                        wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                        h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
                        ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
                    )
                )

        self.max_seq_length = self.config.block_size
        self.mask_cache: Optional[torch.Tensor] = None

    @property
    def max_seq_length(self) -> int:
        return self._max_seq_length

    @max_seq_length.setter
    def max_seq_length(self, value: int) -> None:
        """
        When doing inference, the sequences used might be shorter than the model's context length.
        This allows setting a smaller number to avoid allocating unused memory
        """
        if value > self.config.block_size:
            raise ValueError(f"Cannot attend to {value}, block size is only {self.config.block_size}")
        self._max_seq_length = value
        if not hasattr(self, "cos"):
            # first call
            cos, sin = self.rope_cache()
            self.register_buffer("cos", cos, persistent=False)
            self.register_buffer("sin", sin, persistent=False)
        # override
        elif value != self.cos.size(0):
            self.cos, self.sin = self.rope_cache(device=self.cos.device)
        # the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know
        # if the kv cache is expected

    def reset_parameters(self) -> None:
        # Trigger resetting the rope-cache
        self.cos, self.sin = self.rope_cache(device=self.cos.device)

    def _init_weights(self, module: nn.Module) -> None:
        """Meant to be used with `gpt.apply(gpt._init_weights)`."""
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx: torch.Tensor, input_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
        T = idx.size(1)
        if self.max_seq_length < T:
            raise ValueError(f"Cannot forward sequence of length {T}, max seq length is only {self.max_seq_length}.")

        if self.config.rnn_type is None:
            if input_pos is not None:  # use the kv cache
                cos = self.cos.index_select(0, input_pos)
                sin = self.sin.index_select(0, input_pos)
                if self.mask_cache is None:
                    raise TypeError("You need to call `gpt.set_kv_cache()`")
                mask = self.mask_cache.index_select(2, input_pos)
            else:
                cos = self.cos[:T]
                sin = self.sin[:T]
                mask = None

            x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
            if self.config.scale_embeddings:
                x = x * (self.config.n_embd**0.5)

            for block in self.transformer.h:
                x = block(x, cos, sin, mask, input_pos)
            x = self.transformer.ln_f(x)
            return self.lm_head(x)  # (b, t, vocab_size)
        
        else:
            x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
            if self.config.scale_embeddings:
                x = x * (self.config.n_embd**0.5)
            h_s = None
            for block in self.transformer.h:
                x, h_s = block(x, h_s)
            x = self.transformer.ln_f(x)
            return self.lm_head(x)  # (b, t, vocab_size)

    @classmethod
    def from_name(cls, name: str, **kwargs: Any) -> Self:
        return cls(Config.from_name(name, **kwargs))

    def rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        return build_rope_cache(
            seq_len=self.max_seq_length,
            n_elem=self.config.rope_n_elem,
            device=device,
            condense_ratio=self.config.rope_condense_ratio,
            base=self.config.rope_base,
        )

    def set_kv_cache(
        self,
        batch_size: int,
        rope_cache_length: Optional[int] = None,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None:
        if rope_cache_length is None:
            rope_cache_length = self.cos.size(-1)
        max_seq_length = self.max_seq_length

        # initialize the kv cache for all blocks
        for block in self.transformer.h:
            block.attn.kv_cache = block.attn.build_kv_cache(
                batch_size, max_seq_length, rope_cache_length, device, dtype
            )

        if self.mask_cache is None or self.mask_cache.size(3) != max_seq_length:
            # passing `attn_mask` to SDPA disables the flash implementation. since we only need the mask
            # for the kv-cache support (only during inference), we only create it in that situation
            self.mask_cache = build_mask_cache(max_seq_length, device)

    def clear_kv_cache(self) -> None:
        self.mask_cache = None
        for block in self.transformer.h:
            block.attn.kv_cache = None

# =============================================================
# Block
# =============================================================

class Block(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        if not config.parallel_residual and config.shared_attention_norm:
            raise NotImplementedError(
                "No checkpoint amongst the ones we support uses this configuration"
                " (non-parallel residual and shared attention norm)."
            )

        self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
        self.attn = CausalSelfAttention(config)
        self.norm_2 = None if config.shared_attention_norm else config.norm_class(config.n_embd, eps=config.norm_eps)
        self.mlp = config.mlp_class(config)

        self.config = config

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        input_pos: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Non-parallel residual       Parallel residual
           ┌─ x                     ┌─ x ────────────┐             Note: if `shared_attention_norm` is True,
           │  ↓                     │  ↓             ↓                   the output from `norm_1` is reused
           │  norm_1                │  norm_1  ───►  norm_2
           │  ↓                     │  ↓             ↓
           │  attn                  │  attn          mlp
           │  ↓                     │  ↓             │
        ┌─ └► +                     └► + ◄───────────┘
        │     norm_2
        │     ↓
        │     mlp
        │     ↓
        └───► +
        """

        x_normed = self.norm_1(x)
        attention_output = self.attn(x_normed, cos, sin, mask, input_pos)

        if self.config.parallel_residual:
            x_normed = x_normed if self.config.shared_attention_norm else self.norm_2(x)
            x = self.mlp(x_normed) + attention_output + x
        else:
            x = attention_output + x
            x = self.mlp(self.norm_2(x)) + x
        return x


class CausalSelfAttention(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        shape = (config.n_head + 2 * config.n_query_groups) * config.head_size
        # key, query, value projections for all heads, but in a batch
        self.attn = nn.Linear(config.n_embd, shape, bias=config.bias)
        # output projection
        # if `head_size` is explicitly specified in the config, `n_emd` might not be equal to `head_size * n_head`
        self.proj = nn.Linear(config.head_size * config.n_head, config.n_embd, bias=config.bias)
        # disabled by default
        self.kv_cache: Optional[KVCache] = None

        self.use_infini_transformer = config.use_infini_transformer
        if self.use_infini_transformer:
            # ============ memory params
            self.beta = nn.Parameter(torch.randn(1))
            if config.init_learnable_memory_state:
                self.init_mem = nn.Parameter(torch.zeros(1, config.n_head, config.head_size, config.head_size))
                self.init_z = nn.Parameter(torch.ones(1, config.n_head, config.head_size))
            else:
                self.init_mem = None
                self.init_z = None
            # ============

            self.segment_len = config.segment_len
            self.use_delta_update_memory = config.use_delta_update_memory

        self.config = config

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        input_pos: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:

        if self.use_infini_transformer:
            return self.forward_infini(x, cos, sin, mask, input_pos)

        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        qkv = self.attn(x)

        # assemble into a number of query groups to support MHA, MQA and GQA together (see `config.n_query_groups`)
        q_per_kv = self.config.n_head // self.config.n_query_groups
        total_qkv = q_per_kv + 2  # each group has 1+ queries, 1 key, and 1 value
        qkv = qkv.view(B, T, self.config.n_query_groups, total_qkv, self.config.head_size)
        qkv = qkv.permute(0, 2, 3, 1, 4)  # (B, n_query_groups, total_qkv, T, hs)

        # split batched computation into three
        q, k, v = qkv.split((q_per_kv, 1, 1), dim=2)

        # maybe repeat k and v if for the non multi-head attention cases
        # training: flash attention requires it
        # inference: multi-query would require a full kv cache so avoid it to limit its memory usage
        if self.config.n_query_groups != self.config.n_head and (input_pos is None or self.config.n_query_groups != 1):
            k = k.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
            v = v.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)

        q = q.reshape(B, -1, T, self.config.head_size)  # (B, nh_q, T, hs)
        k = k.reshape(B, -1, T, self.config.head_size)  # (B, nh_k, T, hs)
        v = v.reshape(B, -1, T, self.config.head_size)  # (B, nh_v, T, hs)

        q_roped = apply_rope(q[..., : self.config.rope_n_elem], cos, sin)
        k_roped = apply_rope(k[..., : self.config.rope_n_elem], cos, sin)
        q = torch.cat((q_roped, q[..., self.config.rope_n_elem :]), dim=-1)
        k = torch.cat((k_roped, k[..., self.config.rope_n_elem :]), dim=-1)

        if input_pos is not None:
            if not isinstance(self.kv_cache, KVCache):
                raise TypeError("You need to call `gpt.set_kv_cache()`")
            k, v = self.kv_cache(input_pos, k, v)

        y = self.scaled_dot_product_attention(q, k, v, mask)

        y = y.reshape(B, T, self.config.head_size * self.config.n_head)  # re-assemble all head outputs side by side

        # output projection
        return self.proj(y)


    def forward_infini(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        input_pos: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, full_T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        M = torch.zeros(1, self.config.n_head, self.config.head_size, self.config.head_size).to(x.device)
        z = torch.ones(1, self.config.n_head, self.config.head_size).to(x.device)
        if self.init_mem is not None and self.init_z is not None:
            M = self.init_mem
            z = self.init_z

        num_segments, rem = divmod(full_T, self.segment_len)
        num_segments += 1 if rem > 0 else 0

        out = []
        for ix in range(num_segments):
            ix_begin = ix * self.segment_len
            ix_end = min(ix_begin + self.segment_len, x.size(1))
            T = ix_end - ix_begin

            x_ = x[:, ix_begin:ix_end, :]

            cos_ = cos[ix_begin:ix_end]
            sin_ = sin[ix_begin:ix_end]

            qkv = self.attn(x_)

            # assemble into a number of query groups to support MHA, MQA and GQA together (see `config.n_query_groups`)
            q_per_kv = self.config.n_head // self.config.n_query_groups
            total_qkv = q_per_kv + 2  # each group has 1+ queries, 1 key, and 1 value
            qkv = qkv.view(B, T, self.config.n_query_groups, total_qkv, self.config.head_size)
            qkv = qkv.permute(0, 2, 3, 1, 4)  # (B, n_query_groups, total_qkv, T, hs)

            # split batched computation into three
            q, k, v = qkv.split((q_per_kv, 1, 1), dim=2)

            # maybe repeat k and v if for the non multi-head attention cases
            # training: flash attention requires it
            # inference: multi-query would require a full kv cache so avoid it to limit its memory usage
            if self.config.n_query_groups != self.config.n_head and (input_pos is None or self.config.n_query_groups != 1):
                k = k.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
                v = v.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)

            q = q.reshape(B, -1, T, self.config.head_size)  # (B, nh_q, T, hs)
            k = k.reshape(B, -1, T, self.config.head_size)  # (B, nh_k, T, hs)
            v = v.reshape(B, -1, T, self.config.head_size)  # (B, nh_v, T, hs)

            q_roped = apply_rope(q[..., : self.config.rope_n_elem], cos_, sin_)
            k_roped = apply_rope(k[..., : self.config.rope_n_elem], cos_, sin_)
            q = torch.cat((q_roped, q[..., self.config.rope_n_elem :]), dim=-1)
            k = torch.cat((k_roped, k[..., self.config.rope_n_elem :]), dim=-1)

            # ============================== retrieve from memory
            # shape: (batch_size, num_heads, segment_len, dim)
            memory_output = self._retrieve_from_memory(q, M, z)
            # Update memory with current segment's key and value states
            M, z = self._update_memory(k, v, M, z, use_delta=self.use_delta_update_memory)
            # ==============================

            if input_pos is not None:
                if not isinstance(self.kv_cache, KVCache):
                    raise TypeError("You need to call `gpt.set_kv_cache()`")
                k, v = self.kv_cache(input_pos, k, v)

            y = self.scaled_dot_product_attention(q, k, v, mask)

            # ============================== long term injection
            memory_output = memory_output.transpose(1, 2)
            y = self._long_term_injection(y, memory_output)
            # ==============================

            y = y.reshape(B, T, self.config.head_size * self.config.n_head)  # re-assemble all head outputs side by side
            # output projection
            out.append(self.proj(y))

        # Return concatenated full sequence from buffer
        out = torch.concat(out, dim=1)

        return out

    def scaled_dot_product_attention(
        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        scale = 1.0 / math.sqrt(self.config.head_size)
        y = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale, is_causal=mask is None
        )
        return y.transpose(1, 2)

    def build_kv_cache(
        self,
        batch_size: int,
        max_seq_length: int,
        rope_cache_length: Optional[int] = None,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> "KVCache":
        heads = 1 if self.config.n_query_groups == 1 else self.config.n_head
        v_shape = (batch_size, heads, max_seq_length, self.config.head_size)
        if rope_cache_length is None:
            if self.config.rotary_percentage != 1.0:
                raise TypeError("Please pass the `rope_cache_length=gpt.cos.size(-1)` value")
            k_shape = v_shape
        else:
            k_shape = (
                batch_size,
                heads,
                max_seq_length,
                rope_cache_length + self.config.head_size - self.config.rope_n_elem,
            )
        return KVCache(k_shape, v_shape, device=device, dtype=dtype)

    def _retrieve_from_memory(self, Q, M, z):
        if M is None:
            return torch.zeros_like(Q)

        # Retrieve context from compressive memory using linear attention (Eq. 3)
        sigma_q = F.elu(Q) + 1
        M_s_1 = torch.matmul(sigma_q, M)
        Z_s_1 = torch.matmul(sigma_q, z.unsqueeze(-1) + 1e-8)
        A_mem = M_s_1 / Z_s_1
        return A_mem

    def _update_memory(self, K, V, M, z, use_delta=False):
        sigma_k = F.elu(K) + 1
        if M is not None:
            if use_delta:
                V_retrieved = torch.matmul(sigma_k, M) / (torch.matmul(sigma_k, z.unsqueeze(-1) + 1e-8))
                updated_M = M + torch.matmul(sigma_k.transpose(-2, -1), V - V_retrieved)
            else:
                updated_M = M + torch.matmul(sigma_k.transpose(-2, -1), V)
        else:
            updated_M = torch.matmul(sigma_k.transpose(-2, -1), V)

        if z is not None:
            updated_z = z + sigma_k.sum(dim=-2)
        else:
            updated_z = sigma_k.sum(dim=-2)

        M = updated_M
        z = updated_z
        return M, z

    def _long_term_injection(self, A_dot, A_mem):
        beta = torch.sigmoid(self.beta)
        A = beta * A_mem + (1 - beta) * A_dot
        return A


class MambaBlock(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()

        ssm_cfg = {}
        factory_kwargs = {"device": None, "dtype": None}
        mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)

        self.block = MBBlock(
            config.n_embd,
            mixer_cls,
            fused_add_norm=False,
            residual_in_fp32=False,
        )
        self.block.layer_idx = layer_idx

        self.config = config

    def forward(self, x, residual=None, inference_params=None):
        hidden_states, residual = self.block(x, residual=residual, inference_params=inference_params)

        return hidden_states, residual


# =============================================================
# MLP
# =============================================================

class GptNeoxMLP(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        self.fc = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
        self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)

        self.config = config

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc(x)
        x = torch.nn.functional.gelu(x, approximate=self.config.gelu_approximate)
        return self.proj(x)


class LLaMAMLP(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        self.fc_1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
        self.fc_2 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
        self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)

        self.config = config

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_fc_1 = self.fc_1(x)
        x_fc_2 = self.fc_2(x)
        x = torch.nn.functional.silu(x_fc_1) * x_fc_2
        return self.proj(x)


class GemmaMLP(LLaMAMLP):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_fc_1 = self.fc_1(x)
        x_fc_2 = self.fc_2(x)
        x = torch.nn.functional.gelu(x_fc_1, approximate=self.config.gelu_approximate) * x_fc_2
        return self.proj(x)


class LLaMAMoE(nn.Module):
    def __init__(self, config: Config) -> None:
        super().__init__()
        self.gate = nn.Linear(config.n_embd, config.n_expert, bias=False)
        self.experts = nn.ModuleList(LLaMAMLP(config) for _ in range(config.n_expert))

        self.config = config

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Derived from: https://github.com/mistralai/mistral-src/blob/b46d6/moe_one_file_ref.py#L203-L219
        See also figure 1 in https://arxiv.org/abs/2211.15841
        """
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
        x = x.view(-1, C)  # (B*T, C)
        router = self.gate(x)  # (B*T, n_expert)
        probs, indices = torch.topk(router, self.config.n_expert_per_token)  # (B*T, n_expert_per_token)
        probs = probs.softmax(dim=1, dtype=torch.float).to(dtype=x.dtype)
        masks = indices.unsqueeze(-1) == torch.arange(self.config.n_expert, device=x.device)
        masks = masks.permute(2, 0, 1)  # (n_expert, B*T, n_expert_per_token)
        y = torch.zeros_like(x)  # (B*T, C)
        for mask, expert in zip(masks, self.experts):
            token_idx, expert_idx = torch.where(mask)
            y[token_idx] += probs[token_idx, expert_idx, None] * expert(x[token_idx])
        return y.view(B, T, C)


class TokenRouter(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.weight_predictor = nn.Linear(embed_dim, 1)

    def forward(self, x):
        weights = self.weight_predictor(x).squeeze(
            -1
        )  # [batch_size, seq_len]
        return weights


class MoD(nn.Module):
    def __init__(self, config, block):
        super().__init__()
        self.router = TokenRouter(config.n_embd)
        self.block = block
        self.capacity = config.capacity

    def forward(self, 
                x: torch.Tensor,
                cos: torch.Tensor,
                sin: torch.Tensor,
                mask: Optional[torch.Tensor] = None,
                input_pos: Optional[torch.Tensor] = None,):
        b, s, d = x.shape
        weights = self.router(x)

        # Compute B-th percentil for router weightsto determine the capacity threshold
        k = int(self.capacity * s)
        top_k_values, _ = torch.topk(weights, k, dim=1, sorted=True)
        threshold = top_k_values[:, -1]

        # Determine which tokens exceed the threshold shape: b,s
        selected_mask = weights > threshold.unsqueeze(-1) 
        # Process onlys elected tokens through the block
        masked_input = x * selected_mask[:, :, None]

        processed_tokens = self.block(
                                masked_input,
                                cos,
                                sin,
                                mask,
                                input_pos
                            )
        output = processed_tokens + (
            x * (~selected_mask).unsqueeze(-1).to(x.dtype)
        )

        return output


def build_rope_cache(
    seq_len: int, n_elem: int, device: Optional[torch.device] = None, base: int = 10000, condense_ratio: int = 1
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Enhanced Transformer with Rotary Position Embedding.

    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
    transformers/rope/__init__.py. MIT License:
    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
    """
    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))

    # Create position indexes `[0, 1, ..., seq_len - 1]`
    seq_idx = torch.arange(seq_len, device=device) / condense_ratio

    # Calculate the product of position index and $\theta_i$
    idx_theta = torch.outer(seq_idx, theta).repeat(1, 2)

    return torch.cos(idx_theta), torch.sin(idx_theta)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    head_size = x.size(-1)
    x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
    x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
    rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
    roped = (x * cos) + (rotated * sin)
    return roped.to(dtype=x.dtype)


class KVCache(nn.Module):
    def __init__(
        self,
        k_shape: Tuple[int, int, int, int],
        v_shape: Tuple[int, int, int, int],
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None:
        super().__init__()
        self.register_buffer("k", torch.zeros(k_shape, device=device, dtype=dtype), persistent=False)
        self.register_buffer("v", torch.zeros(v_shape, device=device, dtype=dtype), persistent=False)

    def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # move the buffer to the activation dtype for when AMP is used
        self.k = self.k.to(k.dtype)
        self.v = self.v.to(v.dtype)
        # update the cache
        k = self.k.index_copy_(2, input_pos, k)
        v = self.v.index_copy_(2, input_pos, v)
        return k, v

    def reset_parameters(self) -> None:
        torch.nn.init.zeros_(self.k)
        torch.nn.init.zeros_(self.v)


def build_mask_cache(max_seq_length: int, device: Optional[torch.device] = None) -> torch.Tensor:
    ones = torch.ones((max_seq_length, max_seq_length), device=device, dtype=torch.bool)
    return torch.tril(ones).unsqueeze(0).unsqueeze(0)


class RMSNorm(torch.nn.Module):
    """Root Mean Square Layer Normalization.

    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
    """

    def __init__(self, size: int, dim: int = -1, eps: float = 1e-6, add_unit_offset: bool = False) -> None:
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(size))
        self.eps = eps
        self.dim = dim
        self.add_unit_offset = add_unit_offset

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        # NOTE: the original RMSNorm paper implementation is not equivalent
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        x_normed = x_normed.to(dtype=dtype)
        if self.add_unit_offset:
            # Gemma model requires a unit offset
            # https://github.com/google/gemma_pytorch/blob/main/gemma/model.py#L176
            return x_normed * (1 + self.weight)
        return x_normed * self.weight

    def reset_parameters(self) -> None:
        torch.nn.init.ones_(self.weight)