File size: 3,992 Bytes
0957e22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | # coding=utf-8
# AETHER-30B-11Attn — Auto-extracted from new_attentions.py
import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
# =============================================================================
# 1. MLA (Multi-head Latent Attention) — DeepSeek-V2/V3
# =============================================================================
class MLAAttention(nn.Module):
"""Multi-head Latent Attention (low-rank Q/KV decomposition).
DeepSeek-V2 paper: Q/KV을 low-rank로 분해 → KV cache 1/14 절감.
Q: hidden → q_lora_rank → split (rope + nope) → multi-head
KV: hidden → kv_lora_rank → split (rope + nope) → K (rope+nope), V (nope)
"""
def __init__(self, config, layer_idx: int = 0):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.q_lora_rank = getattr(config, "mla_q_lora_rank", 768)
self.kv_lora_rank = getattr(config, "mla_kv_lora_rank", 256)
self.qk_rope_head_dim = getattr(config, "mla_qk_rope_head_dim", 64)
self.qk_nope_head_dim = getattr(config, "mla_qk_nope_head_dim", 64)
self.v_head_dim = getattr(config, "mla_v_head_dim", 128)
qk_head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim
# Q low-rank
self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False)
self.q_a_norm = nn.LayerNorm(self.q_lora_rank, eps=1e-5)
self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * qk_head_dim, bias=False)
# KV low-rank (compressed)
self.kv_a_proj = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
self.kv_a_norm = nn.LayerNorm(self.kv_lora_rank, eps=1e-5)
self.kv_b_proj = nn.Linear(
self.kv_lora_rank,
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
bias=False,
)
# Output
self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, self.hidden_size, bias=False)
def forward(self, hidden_states, attention_mask=None, position_ids=None,
past_key_value=None, use_cache=False, **kwargs):
bsz, q_len, _ = hidden_states.shape
qk_head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim
# Q
q = self.q_a_proj(hidden_states)
q = self.q_a_norm(q)
q = self.q_b_proj(q).view(bsz, q_len, self.num_heads, qk_head_dim).transpose(1, 2)
# KV
kv = self.kv_a_proj(hidden_states)
kv_compressed, k_rope = kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
kv_compressed = self.kv_a_norm(kv_compressed)
kv_b = self.kv_b_proj(kv_compressed).view(
bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
).transpose(1, 2)
k_nope, v = kv_b.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
# Combine k (broadcast k_rope to all heads)
k_rope_expanded = k_rope.view(bsz, q_len, 1, self.qk_rope_head_dim).expand(
-1, -1, self.num_heads, -1
).transpose(1, 2)
k = torch.cat([k_nope, k_rope_expanded], dim=-1) # (bsz, n_h, seq, qk_head_dim)
# SDPA
attn_out = F.scaled_dot_product_attention(
q, k, v,
attn_mask=(attention_mask.to(q.dtype) if attention_mask is not None else None),
dropout_p=0.0, is_causal=(attention_mask is None and q_len > 1),
)
attn_out = attn_out.transpose(1, 2).contiguous().view(bsz, q_len, -1)
return self.o_proj(attn_out), past_key_value
# =============================================================================
# 2. Mamba2 — State Space Model (placeholder simplification)
__all__ = ['MLAAttention']
|