"""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)