162 lines
5.2 KiB
Python
162 lines
5.2 KiB
Python
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
# For a list of all contributors, visit:
|
|
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
|
|
|
|
def naive_recurrent_gated_delta_rule(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
g: torch.Tensor,
|
|
scale: float = None,
|
|
initial_state: torch.Tensor = None,
|
|
output_final_state: bool = False,
|
|
):
|
|
"""
|
|
Reference PyTorch implementation of recurrent gated delta rule.
|
|
|
|
Args:
|
|
q: [B, T, H, K]
|
|
k: [B, T, H, K]
|
|
v: [B, T, H, V]
|
|
beta: [B, T, H]
|
|
g: [B, T, H]
|
|
scale: float, optional
|
|
initial_state: [B, H, K, V], optional
|
|
output_final_state: bool
|
|
|
|
Returns:
|
|
o: [B, T, H, V]
|
|
final_state: [B, H, K, V] if output_final_state else None
|
|
"""
|
|
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
|
|
B, H, T, K, V = *k.shape, v.shape[-1]
|
|
o = torch.zeros(B, H, T, V).to(v)
|
|
h = torch.zeros(B, H, K, V).to(v)
|
|
if initial_state is not None:
|
|
h = initial_state.to(torch.float32)
|
|
if scale is None:
|
|
scale = 1 / (q.shape[-1] ** 0.5)
|
|
q = q * scale
|
|
|
|
for i in range(T):
|
|
b_q = q[:, :, i]
|
|
b_k = k[:, :, i]
|
|
b_v = v[:, :, i].clone()
|
|
h = h.clone() * g[:, :, i].exp()[..., None, None]
|
|
b_beta = beta[:, :, i]
|
|
b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
|
|
b_v = b_v * b_beta[..., None]
|
|
h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
|
|
o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
|
|
|
|
if not output_final_state:
|
|
h = None
|
|
o = o.transpose(1, 2).contiguous()
|
|
return o, h
|
|
|
|
|
|
def naive_chunk_gated_delta_rule(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
g: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
chunk_size: int = 64,
|
|
scale: float = None,
|
|
initial_state: torch.Tensor = None,
|
|
output_final_state: bool = False,
|
|
):
|
|
"""
|
|
Reference PyTorch implementation of chunk gated delta rule.
|
|
|
|
Args:
|
|
q: [B, T, H, K]
|
|
k: [B, T, H, K]
|
|
v: [B, T, H, V]
|
|
g: [B, T, H]
|
|
beta: [B, T, H]
|
|
chunk_size: int
|
|
scale: float, optional
|
|
initial_state: [B, H, K, V], optional
|
|
output_final_state: bool
|
|
|
|
Returns:
|
|
o: [B, T, H, V]
|
|
final_state: [B, H, K, V] if output_final_state else None
|
|
"""
|
|
BT = chunk_size
|
|
if scale is None:
|
|
scale = 1 / (q.shape[-1] ** 0.5)
|
|
|
|
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
|
|
|
|
T = q.shape[-2]
|
|
pad_len = (BT - (T % BT)) % BT
|
|
if pad_len > 0:
|
|
q = F.pad(q, (0, 0, 0, pad_len))
|
|
k = F.pad(k, (0, 0, 0, pad_len))
|
|
v = F.pad(v, (0, 0, 0, pad_len))
|
|
beta = F.pad(beta, (0, pad_len))
|
|
g = F.pad(g, (0, pad_len))
|
|
|
|
q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
|
|
decay = g
|
|
chunk_size = BT
|
|
b, h, l, d_k = q.shape
|
|
d_v = v.shape[-1]
|
|
q = q * scale
|
|
v = v * beta[..., None]
|
|
k_beta = k * beta[..., None]
|
|
assert l % chunk_size == 0
|
|
|
|
# note that diagonal is masked.
|
|
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
|
|
q, k, v, k_beta, decay = map(
|
|
lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
|
|
[q, k, v, k_beta, decay.unsqueeze(-1)],
|
|
)
|
|
decay = decay.squeeze(-1).cumsum(-1)
|
|
decay_exp = decay.exp()[..., None]
|
|
L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
|
|
attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
|
|
for i in range(1, chunk_size):
|
|
attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
|
|
attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
|
|
attn = attn
|
|
k_cumsum = attn @ v
|
|
k_cumdecay = attn @ (k_beta * decay_exp)
|
|
v = k_cumsum
|
|
|
|
S = k.new_zeros(b, h, d_k, d_v)
|
|
if initial_state is not None:
|
|
S = initial_state.to(torch.float32)
|
|
|
|
o = torch.zeros_like(v)
|
|
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
|
|
for i in range(0, l // chunk_size):
|
|
q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
|
|
attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
|
|
v_prime = (k_cumdecay[:, :, i]) @ S
|
|
v_new = v_i - v_prime
|
|
o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
|
|
o[:, :, i] = o_inter + attn @ v_new
|
|
S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
|
|
[..., None]).transpose(-1, -2) @ v_new
|
|
if not output_final_state:
|
|
S = None
|
|
|
|
# unpad
|
|
o = rearrange(o, 'b h n c d -> b h (n c) d')
|
|
o = o[:, :, :T]
|
|
o = o.transpose(1, 2)
|
|
return o, S
|