File size: 6,029 Bytes
4d74557 | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """
Full Multi-Head Attention Layer for FORGE's global attention blocks.
Uses FlashAttention-3 when available, with GQA and RoPE.
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from .dgn_norm import build_norm
from .rotary_embedding import RotaryEmbedding
try:
from flash_attn import flash_attn_func
FLASH_ATTN = True
except ImportError:
FLASH_ATTN = False
class GlobalMHALayer(nn.Module):
"""
Global Multi-Head Attention with:
- GQA (Grouped Query Attention) with n_kv_heads < n_heads
- RoPE positional encoding (long-context base 500k)
- FlashAttention-3 for maximum GPU utilization
- KV cache support for efficient inference
- Pre-norm with DGN
"""
def __init__(
self,
d_model: int = 1280,
n_heads: int = 16,
n_kv_heads: int = 4,
head_dim: int = 80,
rope_base: float = 500_000.0,
rope_scaling_type: Optional[str] = None,
rope_scaling_factor: float = 1.0,
max_seq_len: int = 4096,
dropout: float = 0.0,
norm_type: str = "dgn",
dgn_n_groups: int = 16,
norm_eps: float = 1e-6,
use_flash_attention: bool = True,
layer_idx: int = 0,
):
super().__init__()
assert d_model == n_heads * head_dim, \
f"d_model={d_model} must equal n_heads*head_dim={n_heads*head_dim}"
assert n_heads % n_kv_heads == 0, \
f"n_heads={n_heads} must be divisible by n_kv_heads={n_kv_heads}"
self.d_model = d_model
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = head_dim
self.n_rep = n_heads // n_kv_heads # GQA repetition factor
self.dropout = dropout
self.use_flash = use_flash_attention and FLASH_ATTN
self.layer_idx = layer_idx
self.scale = head_dim ** -0.5
# Pre-norm
self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
# Q/K/V projections
self.q_proj = nn.Linear(d_model, n_heads * head_dim, bias=False)
self.k_proj = nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
self.v_proj = nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
self.o_proj = nn.Linear(d_model, d_model, bias=False)
# RoPE
self.rope = RotaryEmbedding(
head_dim=head_dim,
max_seq_len=max_seq_len,
base=rope_base,
scaling_type=rope_scaling_type,
scaling_factor=rope_scaling_factor,
)
# KV cache for inference
self._kv_cache_k: Optional[torch.Tensor] = None
self._kv_cache_v: Optional[torch.Tensor] = None
self._kv_cache_len: int = 0
self._init_weights()
def _init_weights(self):
nn.init.normal_(self.q_proj.weight, std=0.02)
nn.init.normal_(self.k_proj.weight, std=0.02)
nn.init.normal_(self.v_proj.weight, std=0.02)
# Output projection: scale by 1/sqrt(2 * n_layers) for residual depth
nn.init.normal_(self.o_proj.weight, std=0.02)
def clear_kv_cache(self):
self._kv_cache_k = None
self._kv_cache_v = None
self._kv_cache_len = 0
def forward(
self,
x: torch.Tensor, # (B, T, d_model)
attention_mask: Optional[torch.Tensor] = None, # (B, T) or (B, 1, T, T)
position_ids: Optional[torch.Tensor] = None,
use_cache: bool = False,
) -> Tuple[torch.Tensor, Optional[Tuple]]:
B, T, _ = x.shape
residual = x
x_norm = self.norm(x)
# Project
Q = self.q_proj(x_norm) # (B, T, n_heads * head_dim)
K = self.k_proj(x_norm) # (B, T, n_kv_heads * head_dim)
V = self.v_proj(x_norm) # (B, T, n_kv_heads * head_dim)
# Reshape to (B, H, T, D)
Q = rearrange(Q, 'b t (h d) -> b h t d', d=self.head_dim)
K = rearrange(K, 'b t (h d) -> b h t d', d=self.head_dim)
V = rearrange(V, 'b t (h d) -> b h t d', d=self.head_dim)
# RoPE
Q, K = self.rope.apply_rotary(Q, K, position_ids)
# KV cache concat (inference)
if use_cache:
if self._kv_cache_k is not None:
K = torch.cat([self._kv_cache_k, K], dim=2)
V = torch.cat([self._kv_cache_v, V], dim=2)
self._kv_cache_k = K
self._kv_cache_v = V
self._kv_cache_len = K.shape[2]
# GQA: expand K,V from n_kv_heads to n_heads
if self.n_rep > 1:
K = K.repeat_interleave(self.n_rep, dim=1)
V = V.repeat_interleave(self.n_rep, dim=1)
# Attention
if self.use_flash:
# FlashAttention: expects (B, T, H, D)
Q_fa = rearrange(Q, 'b h t d -> b t h d').contiguous()
K_fa = rearrange(K, 'b h t d -> b t h d').contiguous()
V_fa = rearrange(V, 'b h t d -> b t h d').contiguous()
out = flash_attn_func(
Q_fa, K_fa, V_fa,
dropout_p=self.dropout if self.training else 0.0,
causal=True,
softmax_scale=self.scale,
) # (B, T, H, D)
out = rearrange(out, 'b t h d -> b t (h d)')
else:
# Standard PyTorch SDPA (fused, memory efficient)
out = F.scaled_dot_product_attention(
Q, K, V,
attn_mask=attention_mask,
dropout_p=self.dropout if self.training else 0.0,
is_causal=(attention_mask is None),
scale=self.scale,
) # (B, H, T, D)
out = rearrange(out, 'b h t d -> b t (h d)')
out = self.o_proj(out) # (B, T, d_model)
return residual + out, (K, V) if use_cache else None |