Github 倉庫：https://github.com/One-sixth/flash-linear-attention-pytorch

flash-linear-attention-pytorch

純 Pytorch 實(shí)現(xiàn) TransnormerLLM 中快速線性注意力算子。
用于學(xué)習(xí)目的。
如果你希望用于訓(xùn)練模型，你可能要修改為 CUDA 或 Triton 的實(shí)現(xiàn)，不然會很慢。

注意

這個(gè)算子有精度問題，誤差較大，是正常的。
這是因?yàn)樽⒁饬仃嚊]有激活函數(shù)，導(dǎo)致注意力矩陣的值很大。
在使用 float16 類型時(shí)需要特別小心。

這是一個(gè)簡單的緩解方法：限制 q 和 k 的值，從而減少float16溢出的可能性。

q = q / q.norm(-1, keepdim=True)
k = k / k.norm(-1, keepdim=True)
o = linear_attention(q, k, v, m)

使用方法

import torch
from flash_linear_attention_ops import flash_linear_attention, normal_linear_attentionbatch_size = 16
seq_len = 1024
dim = 64
n_head = 12
device = 'cuda'
dtype = torch.float32Q = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
K = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
V = torch.randn(batch_size, n_head, seq_len, dim, requires_grad=True, dtype=dtype, device=device)
M = torch.randint(0, 2, (1, 1, seq_len, seq_len), device=device, dtype=dtype)O_flash = flash_linear_attention(Q, K, V, M)
O_normal = normal_linear_attention(Q, K, V, M)print('O_flash.shape', O_flash.shape)
print('O_normal.shape', O_normal.shape)print('O diff', (O_flash - O_normal).abs().max().item())

參考引用

https://github.com/OpenNLPLab/TransnormerLLM
https://github.com/shreyansh26/FlashAttention-PyTorch