base_IIXIV / fla /ops /linear_attn /utils.py
mainline777's picture
Duplicate from silx-ai/Quasar-Preview
41865df
Raw
History Blame Contribute Delete
190 Bytes
import torch
def normalize_output(q: torch.Tensor, k: torch.Tensor, o: torch.Tensor) -> torch.Tensor:
k = k.cumsum(1)
z = (q * k).sum(-1, keepdim=True)
return o / (z + 1e-10)