Kernels:

kernels-community
/

flash-attn2

Trusted publisher

This is the repository card of kernels-community/flash-attn2 that has been pushed on the Hub. It was built to be used with the kernels library. This card was automatically generated.

How to use

# make sure `kernels` is installed: `pip install -U kernels`
from kernels import get_kernel

kernel_module = get_kernel("kernels-community/flash-attn2", version=4)
flash_attn_func = kernel_module.flash_attn_func

flash_attn_func(...)

Available functions

flash_attn_func
flash_attn_kvpacked_func
flash_attn_qkvpacked_func
flash_attn_varlen_func
flash_attn_varlen_kvpacked_func
flash_attn_varlen_qkvpacked_func
flash_attn_with_kvcache
flash_attn_interface
fwd
varlen_fwd
bwd
varlen_bwd
fwd_kvcache

Low-level ops

flash_attn_interface is part of the public API and mirrors the upstream flash_attn.flash_attn_interface module. Besides the high-level functions, these four functions are supported entry points:

flash_attn_interface._flash_attn_forward
flash_attn_interface._flash_attn_backward
flash_attn_interface._flash_attn_varlen_forward
flash_attn_interface._flash_attn_varlen_backward

The upstream _wrapped_* aliases of these four functions are available as well. They are Torch custom ops with upstream's signatures. They have no autograd, always return the softmax LSE, and are traceable by torch.compile. Use them when you need the LSE or drive the backward pass yourself, for instance in ring or context-parallel attention. The names keep their upstream leading underscore so code written against the flash_attn package ports unchanged.

Dropout is supported. The CUDA build uses the Torch stable ABI and cannot read the CUDA generator itself, so the forward is handed a generator and reserves its Philox seed and offset through the dispatcher, exactly as the non-stable build did. Masks are therefore bit-identical to earlier kernel versions and to the pip flash_attn package for the same torch.manual_seed, including under activation checkpointing and CUDA graph capture.

Torch 2.10 to 2.12 cannot pass a generator through the stable ABI. There the kernel draws its seed and offset from the default generator instead, which is reproducible under torch.manual_seed but not bit-identical to other Torch versions, and the gen argument of fwd/varlen_fwd raises a ValueError.

return_softmax requires dropout_p > 0.