File size: 6,167 Bytes
2db32a1 | 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 | # coding=utf-8
"""Gated Cross-Layer Attention (GCLA).
Implements Section VII of the Wiola paper. Each decoder layer performs:
* GQA self-attention with SRPE on the local sequence, and
* cross-attention to compressed summaries of up to ``Lambda`` preceding
layers (supplied by the model as ``context_summaries``),
blended by a scalar gate ``beta = sigmoid(phi)`` and modulated by a sigmoid
output gate ``G``:
O = (1 - beta) * O_self + beta * O_ctx
A = (G * concat(O)) W_O
The context tensor is provided *per query position* as a causal cumulative mean
of prior-layer outputs (see :class:`WiolaModel`), so a cached incremental decode
reproduces a full forward pass exactly.
"""
import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from .srpe import SpiralRotaryEmbedding, apply_srpe
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
"""Expand [B, H_kv, T, d] -> [B, H_kv*n_rep, T, d] (GQA)."""
if n_rep == 1:
return x
b, kv, t, d = x.shape
x = x[:, :, None, :, :].expand(b, kv, n_rep, t, d)
return x.reshape(b, kv * n_rep, t, d)
class GatedCrossLayerAttention(nn.Module):
def __init__(self, config, layer_idx: int):
super().__init__()
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.num_kv_heads = config.num_key_value_heads
self.head_dim = config.head_dim
self.n_rep = self.num_heads // self.num_kv_heads
self.lookback = config.gcla_lookback
q_dim = self.num_heads * self.head_dim
kv_dim = self.num_kv_heads * self.head_dim
self.q_proj = nn.Linear(self.hidden_size, q_dim, bias=False)
self.k_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
self.v_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
self.o_proj = nn.Linear(q_dim, self.hidden_size, bias=False)
# Cross-layer context projections.
self.k_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
self.v_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
# Sigmoid output gate.
self.gate_proj = nn.Linear(self.hidden_size, q_dim, bias=False)
# Scalar blend gate beta = sigmoid(phi).
self.blend_logit = nn.Parameter(torch.tensor(float(config.gcla_gate_init)))
self.srpe = SpiralRotaryEmbedding(
head_dim=self.head_dim,
max_position_embeddings=config.max_position_embeddings,
theta=config.srpe_theta,
spiral_divisor=config.srpe_spiral_divisor,
radial_amplitude=config.srpe_radial_amplitude,
radial_frequency=config.srpe_radial_frequency,
)
def forward(
self,
hidden_states: torch.Tensor, # [B, T, d]
position_ids: torch.Tensor, # [B, T]
attention_mask: Optional[torch.Tensor], # [B, 1, T, T_k] additive
context_summaries: Optional[torch.Tensor] = None, # [B, T, Lambda, d]
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
use_cache: bool = False,
):
bsz, q_len, _ = hidden_states.shape
q = self.q_proj(hidden_states)
k = self.k_proj(hidden_states)
v = self.v_proj(hidden_states)
q = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = v.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
# SRPE rotation on q/k (per head).
cos, sin = self.srpe(hidden_states, position_ids)
q, k = apply_srpe(q, k, cos, sin)
# Concatenate cached KV (incremental decoding).
# Concatenate cached KV (incremental decoding).
if (
past_key_value is not None
and past_key_value[0] is not None # ← guard against placeholder None
and past_key_value[1] is not None
):
past_k, past_v = past_key_value
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
present = (k, v) if use_cache else None
# GQA expansion.
k_rep = repeat_kv(k, self.n_rep)
v_rep = repeat_kv(v, self.n_rep)
scale = 1.0 / math.sqrt(self.head_dim)
scores = torch.matmul(q, k_rep.transpose(-1, -2)) * scale # [B,H,T,T_k]
if attention_mask is not None:
scores = scores + attention_mask
attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
o_self = torch.matmul(attn, v_rep) # [B,H,T,dh]
# Cross-layer context attention.
if context_summaries is not None and context_summaries.shape[2] > 0:
lam = context_summaries.shape[2]
k_ctx = self.k_ctx_proj(context_summaries) # [B,T,Lam,kv_dim]
v_ctx = self.v_ctx_proj(context_summaries)
k_ctx = k_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
v_ctx = v_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
# Expand kv heads to full head count.
k_ctx = k_ctx.repeat_interleave(self.n_rep, dim=3) # [B,T,Lam,H,dh]
v_ctx = v_ctx.repeat_interleave(self.n_rep, dim=3)
# scores[b,h,t,l] = q[b,h,t,:] . k_ctx[b,t,l,h,:]
ctx_scores = torch.einsum("bhtd,btlhd->bhtl", q, k_ctx) * scale
ctx_attn = F.softmax(ctx_scores, dim=-1, dtype=torch.float32).to(q.dtype)
o_ctx = torch.einsum("bhtl,btlhd->bhtd", ctx_attn, v_ctx)
beta = torch.sigmoid(self.blend_logit)
o = (1.0 - beta) * o_self + beta * o_ctx
else:
o = o_self
# Merge heads.
o = o.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.head_dim)
# Sigmoid output gate.
g = torch.sigmoid(self.gate_proj(hidden_states))
o = self.o_proj(g * o)
return o, present
|