flash-attn3-12864

flash-attn3-12864

FlashAttention-3 (Hopper, sm90) with two head-dim pairs added that the upstream build does not instantiate: (d_qk=64, d_v=128) forward + backward, and (d_qk=128, d_v=64). Everything else is unmodified upstream kernels-community/flash-attn3.

Without these instantiations, attention at d_qk=64 / d_v=128 falls back to the padded (64,256) template — roughly double the O-accumulator and V smem traffic for the same math. The (64,128) forward here runs a retuned 128x128 tile; the backward carries a TMA transaction-bytes fix required for unequal head dims (details in src_patches/FA3_BWD_ANALYSIS.md).

Use

from kernels import get_kernel
fa3 = get_kernel("devenpzak/flash-attn3-12864", version=1)
# drop-in: same flash_attn_interface API as kernels-community/flash-attn3

Ships one variant: torch-stable-abi29-cu128-x86_64-linux (torch >= 2.10 cu128, stable ABI, needs libcudart.so.13 available — --cudart shared).

Rebuild from source

Everything needed to reproduce the binary bit-for-bit is in src_patches/:

1. Kernel source: commit 8a730d96c37560ccf1d3e09f7bbccdc886818f33 of the flash-attn3 kernel (that commit was force-pushed out of the HF repo's history; it is recoverable from the github.com/huggingface/kernels-community monorepo, flash-attn3/ subtree). CUTLASS v4.0.0.
2. Apply fa3_full_64x128_from_built_src.patch (patch -p0 inside flash-attn/): the complete (64,128) fwd+bwd instantiations, dispatcher branches, tile retune, and the TmaTransactionBytesdO fix — 33 files.
3. Add the 24 hdim64_128 TU entries to build.toml (ordering mirrors the patch's own 128_64 block), then apply fa3_12864_from_64x128.patch: the (128,64) support — 28 files.
4. Build per FA3_FWD_BUILD.md §4: hand-rolled from build.toml (nvcc 13.0, gcc 13.x, sm90 TUs at compute_90a, binding TU with -DTORCH_EXTENSION_NAME=_flash_attn3_cuda_8a730d9 -DPy_LIMITED_API=0x03090000), link -shared --cudart shared against libtorch/libcudart. 451 TUs total; note the ptxas 13.2 ICE (C7907) on ~28 fwd TUs — build those with nvcc 13.0, which build.toml's cuda.maxver = "13.0" already pins.
5. fa3_build_cuda.sh is the compile driver (expects a compile_jobs.txt, one nvcc line per TU).

Built .so sha256: 9beee90dbef9c9c221e677dad9f932108b70c46aa35841fc0862281d5b84c09b.

FA3_FWD_BUILD.md documents the forward instantiation, dispatcher, and tile retune; FA3_BWD_ANALYSIS.md the backward analysis and the unequal-head-dim fix.