This repository contains PyTorch implementations of the infinite-setting penumbral and umbral cone attention operators introduced in https://arxiv.org/abs/2306.00392.
Both files should be self-contained, and can be used as replacements for dot product attention.
If you have a Volta or newer NVIDIA GPU and are using PyTorch 2.0+, you may wish to torch.compile these implementations to get significant speedups.
If you have an Ampere or newer NVIDIA GPU, you may also want to turn on TF32 for internal matrix multiplications, which can give additional speedups (see https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html).