初始化项目，由ModelHub XC社区提供模型
Model: pathcosmos/frankenstallm Source: Original Platform
This commit is contained in:
ModelHub XC
263
source/model/attention.py
Normal file
263
source/model/attention.py
Normal file
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
Multi-Head (and Grouped-Query) Attention with optional FlashAttention-2 backend.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .config import LMConfig
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Optional FlashAttention import
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
from flash_attn import flash_attn_func # type: ignore[import]
|
||||
HAS_FLASH_ATTN = True
|
||||
except ImportError:
|
||||
HAS_FLASH_ATTN = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Optional TransformerEngine import (FP8 support)
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
import transformer_engine.pytorch as te # type: ignore[import]
|
||||
HAS_TE = True
|
||||
except ImportError:
|
||||
te = None # type: ignore[assignment]
|
||||
HAS_TE = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Rotary embedding helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def apply_rotary_emb(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Apply rotary positional embeddings to query or key tensor.
|
||||
|
||||
Args:
|
||||
x: (B, T, H, D_head)
|
||||
cos: (T, D_head // 2) — from RotaryEmbedding.forward
|
||||
sin: (T, D_head // 2) — from RotaryEmbedding.forward
|
||||
|
||||
Returns:
|
||||
Tensor with the same shape as *x*, rotated.
|
||||
"""
|
||||
d = x.shape[-1]
|
||||
half_d = d // 2
|
||||
|
||||
x1 = x[..., :half_d] # (B, T, H, D//2)
|
||||
x2 = x[..., half_d:] # (B, T, H, D//2)
|
||||
|
||||
# Broadcast cos/sin from (T, D//2) → (1, T, 1, D//2)
|
||||
cos = cos.unsqueeze(0).unsqueeze(2) # (1, T, 1, D//2)
|
||||
sin = sin.unsqueeze(0).unsqueeze(2) # (1, T, 1, D//2)
|
||||
|
||||
rotated = torch.cat(
|
||||
[x1 * cos - x2 * sin, x1 * sin + x2 * cos],
|
||||
dim=-1,
|
||||
)
|
||||
return rotated.to(x.dtype)
|
||||
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-Head Attention
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MultiHeadAttention(nn.Module):
|
||||
"""Multi-head (or grouped-query) causal self-attention.
|
||||
|
||||
Supports:
|
||||
- Standard MHA: n_kv_heads == n_heads
|
||||
- GQA / MQA: n_kv_heads < n_heads (must evenly divide n_heads)
|
||||
|
||||
Attention backend:
|
||||
- FlashAttention-2 when available and config.use_flash_attn is True
|
||||
- Vanilla scaled dot-product otherwise (causal mask via upper-triangular)
|
||||
"""
|
||||
|
||||
def __init__(self, config: LMConfig) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.n_heads = config.n_heads
|
||||
self.n_kv_heads = config.n_kv_heads # resolved in __post_init__
|
||||
self.head_dim = config.d_model // config.n_heads
|
||||
self.d_model = config.d_model
|
||||
self.dropout = config.dropout
|
||||
self.use_flash = config.use_flash_attn
|
||||
|
||||
# Number of query-head groups per KV head
|
||||
self.n_rep = self.n_heads // self.n_kv_heads
|
||||
|
||||
# Projections ----------------------------------------------------
|
||||
# Select Linear implementation: te.Linear (FP8) or nn.Linear (BF16)
|
||||
_Linear = te.Linear if (config.use_fp8 and HAS_TE) else nn.Linear
|
||||
|
||||
# Fused QKV projection: single GEMM (d_model → q_dim + k_dim + v_dim)
|
||||
# For GQA 24:8 with head_dim=128: 3072 + 1024 + 1024 = 5120
|
||||
self._q_dim = self.n_heads * self.head_dim # e.g. 24 * 128 = 3072
|
||||
self._kv_dim = self.n_kv_heads * self.head_dim # e.g. 8 * 128 = 1024
|
||||
self.qkv_proj = _Linear(
|
||||
config.d_model,
|
||||
self._q_dim + 2 * self._kv_dim, # 3072 + 2*1024 = 5120
|
||||
bias=config.bias,
|
||||
)
|
||||
self.out_proj = _Linear(
|
||||
config.d_model,
|
||||
config.d_model,
|
||||
bias=config.bias,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# KV-head expansion for GQA
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""Expand KV heads to match the number of query heads.
|
||||
|
||||
Args:
|
||||
x: (B, T, n_kv_heads, head_dim)
|
||||
n_rep: repetition factor
|
||||
|
||||
Returns:
|
||||
(B, T, n_kv_heads * n_rep, head_dim)
|
||||
"""
|
||||
if n_rep == 1:
|
||||
return x
|
||||
B, T, n_kv, D = x.shape
|
||||
return x.repeat_interleave(n_rep, dim=2)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forward
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: (B, T, C)
|
||||
cos: (T, head_dim // 2) — from RotaryEmbedding
|
||||
sin: (T, head_dim // 2) — from RotaryEmbedding
|
||||
|
||||
Returns:
|
||||
(B, T, C)
|
||||
"""
|
||||
B, T, C = x.shape
|
||||
|
||||
# --- Fused QKV projection (single GEMM) --------------------------------
|
||||
qkv = self.qkv_proj(x) # (B, T, q_dim + 2*kv_dim)
|
||||
q, k, v = qkv.split([self._q_dim, self._kv_dim, self._kv_dim], dim=-1)
|
||||
q = q.view(B, T, self.n_heads, self.head_dim)
|
||||
k = k.view(B, T, self.n_kv_heads, self.head_dim)
|
||||
v = v.view(B, T, self.n_kv_heads, self.head_dim)
|
||||
|
||||
# FlashAttention-2 and rotary embedding require bf16/fp16.
|
||||
# te.Linear with MXFP8 may emit FP8-format output tensors; cast if needed.
|
||||
if q.dtype not in (torch.float16, torch.bfloat16):
|
||||
q = q.to(torch.bfloat16)
|
||||
k = k.to(torch.bfloat16)
|
||||
v = v.to(torch.bfloat16)
|
||||
|
||||
# --- Rotary embeddings -----------------------------------------------
|
||||
q = apply_rotary_emb(q, cos, sin)
|
||||
k = apply_rotary_emb(k, cos, sin)
|
||||
|
||||
# --- Attention -------------------------------------------------------
|
||||
if self.use_flash and HAS_FLASH_ATTN and x.is_cuda:
|
||||
attn_out = self._flash_attention(q, k, v, B, T)
|
||||
else:
|
||||
attn_out = self._standard_attention(q, k, v, B, T)
|
||||
|
||||
# --- Output projection -----------------------------------------------
|
||||
# attn_out: (B, T, C)
|
||||
return self.out_proj(attn_out)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# FlashAttention-2 path
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _flash_attention(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
B: int,
|
||||
T: int,
|
||||
) -> torch.Tensor:
|
||||
"""Run FlashAttention-2.
|
||||
|
||||
flash_attn_func expects inputs in (B, T, H, D) layout and returns
|
||||
(B, T, H, D). FlashAttention-2 natively supports GQA via head count
|
||||
mismatch (q has n_heads, k/v have n_kv_heads) — no KV expansion needed.
|
||||
"""
|
||||
dropout_p = self.dropout if self.training else 0.0
|
||||
|
||||
# flash_attn_func: (B, T, H, D) → (B, T, H, D)
|
||||
# GQA is handled natively: q=(B,T,n_heads,D), k/v=(B,T,n_kv_heads,D)
|
||||
out = flash_attn_func(q, k, v, dropout_p=dropout_p, causal=True)
|
||||
|
||||
# Reshape (B, T, n_heads, head_dim) → (B, T, C)
|
||||
return out.reshape(B, T, self.n_heads * self.head_dim)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Standard (fallback) attention path
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _standard_attention(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
B: int,
|
||||
T: int,
|
||||
) -> torch.Tensor:
|
||||
"""Vanilla scaled dot-product causal attention.
|
||||
|
||||
Softmax is computed in float32 for numerical stability.
|
||||
"""
|
||||
# Expand KV heads for GQA
|
||||
k = self._repeat_kv(k, self.n_rep) # (B, T, n_heads, head_dim)
|
||||
v = self._repeat_kv(v, self.n_rep) # (B, T, n_heads, head_dim)
|
||||
|
||||
# (B, T, H, D) → (B, H, T, D)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
scale = math.sqrt(self.head_dim)
|
||||
|
||||
# Scaled dot-product: (B, H, T, T)
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) / scale
|
||||
|
||||
# Causal mask: fill upper triangle (excluding diagonal) with -inf
|
||||
causal_mask = torch.triu(
|
||||
torch.ones(T, T, device=q.device, dtype=torch.bool), diagonal=1
|
||||
)
|
||||
scores = scores.masked_fill(causal_mask, float("-inf"))
|
||||
|
||||
# Softmax in fp32, then cast back
|
||||
attn_weights = F.softmax(scores.float(), dim=-1).to(q.dtype)
|
||||
|
||||
if self.training and self.dropout > 0.0:
|
||||
attn_weights = F.dropout(attn_weights, p=self.dropout)
|
||||
|
||||
# Weighted sum: (B, H, T, D)
|
||||
out = torch.matmul(attn_weights, v)
|
||||
|
||||
# (B, H, T, D) → (B, T, H, D) → (B, T, C)
|
||||
out = out.transpose(1, 2).contiguous().reshape(B, T, self.d_model)
|
||||
return out
|
||||
Reference in New Issue
Block a user