mla_components.py

"""MLA attention blocks used by TaoNet."""

import math

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


class RotaryEmbedding(nn.Module):
    """Rotary position embedding with optional YaRN scaling."""

    def __init__(
        self,
        dim,
        rope_scale=40.0,
        max_seq_length=1024,
        yarn_enabled=False,
        yarn_original_max_seq_length=None,
        yarn_alpha=1.0,
    ):
        super().__init__()
        if dim % 2 != 0:
            raise ValueError("RotaryEmbedding requires an even dimension")
        self.dim = dim
        self.rope_scale = rope_scale
        self.max_seq_length = max_seq_length
        self.yarn_enabled = yarn_enabled
        self.yarn_original_max_seq_length = (
            yarn_original_max_seq_length if yarn_original_max_seq_length is not None else max_seq_length
        )
        self.yarn_alpha = yarn_alpha

        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def _apply_yarn_scaling(self, freqs, seq_len):
        if not self.yarn_enabled:
            return freqs
        original_max_seq_length = self.yarn_original_max_seq_length
        if seq_len <= original_max_seq_length:
            return freqs
        target_scale = self.max_seq_length / original_max_seq_length
        current_ratio = seq_len / original_max_seq_length
        progress = min(current_ratio / target_scale, 1.0)
        scale_factor = 1.0 + (target_scale - 1.0) * (progress ** (1.0 / self.yarn_alpha))
        return freqs / scale_factor

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq) / self.rope_scale
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        freqs = self._apply_yarn_scaling(freqs, seq_len)
        return torch.cat((freqs, freqs), dim=-1)


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary(x, cos, sin):
    if cos.dim() == 2:
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
    cos = cos[..., : x.shape[-1]]
    sin = sin[..., : x.shape[-1]]
    x_rot = x[..., : cos.shape[-1]]
    x_base = x[..., cos.shape[-1] :]
    x_rot = (x_rot * cos) + (rotate_half(x_rot) * sin)
    if x_base.shape[-1] > 0:
        return torch.cat([x_rot, x_base], dim=-1)
    return x_rot


class DeepSeekMLA(nn.Module):
    """DeepSeek-style multi-head latent attention."""

    def __init__(
        self,
        d_model,
        d_latent_kv,
        n_heads,
        d_rope,
        dropout=0.1,
        gqa_groups=1,
        rope_scale=40.0,
        max_seq_length=1024,
        yarn_enabled=False,
        yarn_original_max_seq_length=None,
        yarn_alpha=1.0,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_latent_kv = d_latent_kv
        self.n_heads = n_heads
        self.d_rope = d_rope
        self.gqa_groups = gqa_groups

        if d_model % n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if d_latent_kv % n_heads != 0:
            raise ValueError("d_latent_kv must be divisible by n_heads")

        self.d_head_full = d_model // n_heads
        self.d_head_latent = d_latent_kv // n_heads

        self.norm = nn.LayerNorm(d_model)
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_latent_kv, bias=False)
        self.v_proj = nn.Linear(d_model, d_latent_kv, bias=False)
        self.rotary = RotaryEmbedding(
            d_rope,
            rope_scale=rope_scale,
            max_seq_length=max_seq_length,
            yarn_enabled=yarn_enabled,
            yarn_original_max_seq_length=yarn_original_max_seq_length,
            yarn_alpha=yarn_alpha,
        )
        self.out_proj = nn.Linear(d_latent_kv, d_model, bias=False)
        self.head_weights = nn.Parameter(torch.ones(n_heads))
        self.attn_dropout = nn.Dropout(dropout)
        self.proj_dropout = nn.Dropout(dropout)

    def forward(self, x, attention_mask=None):
        batch_size, seq_len, _ = x.shape
        x_norm = self.norm(x)

        q = self.q_proj(x_norm)
        k = self.k_proj(x_norm)
        v = self.v_proj(x_norm)

        q = q.view(batch_size, seq_len, self.n_heads, self.d_head_full).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_head_latent).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_head_latent).transpose(1, 2)

        if self.d_rope > 0:
            rotary_emb = self.rotary(seq_len, x.device)
            cos = torch.cos(rotary_emb).unsqueeze(0).unsqueeze(0)
            sin = torch.sin(rotary_emb).unsqueeze(0).unsqueeze(0)
            q_rope = apply_rotary(q[..., : self.d_rope], cos, sin)
            q = torch.cat([q_rope, q[..., self.d_rope :]], dim=-1)
            k_rope = apply_rotary(k[..., : self.d_rope], cos, sin)
            k = torch.cat([k_rope, k[..., self.d_rope :]], dim=-1)

        q_for_attn = q[..., : self.d_head_latent]
        attn_mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                attn_mask_bool = attention_mask.bool().unsqueeze(1).unsqueeze(1)
            else:
                attn_mask_bool = attention_mask.bool()

        dropout_p = self.attn_dropout.p if self.training else 0.0
        out_heads = F.scaled_dot_product_attention(
            q_for_attn,
            k,
            v,
            attn_mask=attn_mask_bool,
            dropout_p=dropout_p,
            scale=None,
        )
        out_concat = out_heads.transpose(1, 2).reshape(batch_size, seq_len, self.d_latent_kv)
        out = self.out_proj(out_concat)
        return self.proj_dropout(out)


class AttentionBlock(nn.Module):
    """Attention block with SwiGLU feed-forward network."""

    def __init__(
        self,
        d_model,
        d_latent_kv,
        n_heads,
        d_rope,
        d_ff,
        dropout=0.1,
        gqa_groups=1,
        rope_scale=40.0,
        max_seq_length=1024,
        yarn_enabled=False,
        yarn_original_max_seq_length=None,
        yarn_alpha=1.0,
    ):
        super().__init__()
        self.mla = DeepSeekMLA(
            d_model,
            d_latent_kv,
            n_heads,
            d_rope,
            dropout,
            gqa_groups,
            rope_scale=rope_scale,
            max_seq_length=max_seq_length,
            yarn_enabled=yarn_enabled,
            yarn_original_max_seq_length=yarn_original_max_seq_length,
            yarn_alpha=yarn_alpha,
        )
        self.ff_norm = nn.LayerNorm(d_model)
        self.ff_gate = nn.Linear(d_model, d_ff, bias=False)
        self.ff_value = nn.Linear(d_model, d_ff, bias=False)
        self.ff_out = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attention_mask=None):
        attn_out = self.mla(x, attention_mask)
        x = x + self.dropout(attn_out)
        ff_norm = self.ff_norm(x)
        ff_gate = self.ff_gate(ff_norm)
        ff_value = self.ff_value(ff_norm)
        ff_out = ff_value * F.silu(ff_gate)
        ff_out = self.ff_out(ff_out)
        return x + self.dropout(ff_out)