Kernels:
Trusted publisher
File size: 2,668 Bytes
1779e8b f81bfcb 1779e8b f81bfcb 1779e8b f81bfcb 1779e8b 60c6a88 1779e8b 568865e 60c6a88 74f8aef 60c6a88 568865e 60c6a88 1779e8b 568865e f81bfcb 1779e8b 568865e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | ---
library_name: kernels
license: bsd-3-clause
---
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](https://github.com/huggingface/kernels). This card was automatically generated.
## How to use
```python
# 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`.
## Benchmarks
Benchmarking script is available for this kernel. Run `kernels benchmark kernels-community/flash-attn2 --version 4`.
|