b11140
Pre-releaseCUDA: enable sparse-fa for dsv4 prefill (again) (#29298)
-
CUDA: enable sparse-fa for dsv4 prefill (again)
-
CUDA: unroll the query loop of the sparse mask scan
The query loop of flash_attn_mask_to_sparse_indices has a runtime trip
count, which keeps the unrolled scan over the values of a lane from
issuing its loads together. Template the kernel on ncols1 so the loop
is bounded at compile time: batch one decodes compile to straight line
code and the scan drops from 46 to 17 us at 49k columns on sparse
decode shapes.
- CUDA: pick the out of bounds check of the sparse mask scan in host code
The query loop of the ncols1 == 8 scan keeps a runtime bound and an
early exit, so it does not unroll past its first iteration. Template the
kernel on whether the last group of queries is partial, decided on the
host from n_queries, and hoist the column bound out of the loop: the
loop becomes straight line code and the batched sparse op at 49k
context drops from 586 to 244 us.
Co-authored-by: Pascal admin@serveurperso.com
Website:
Attestations:
macOS/iOS:
- macOS Apple Silicon (arm64)
- macOS Apple Silicon (arm64, KleidiAI enabled) DISABLED
- macOS Intel (x64)
- iOS XCFramework
Linux:
- Ubuntu x64 (CPU)
- Ubuntu arm64 (CPU)
- Ubuntu s390x (CPU)
- Ubuntu x64 (Vulkan)
- Ubuntu arm64 (Vulkan)
- Ubuntu x64 (CUDA 12) - CUDA 12.8 libraries
- Ubuntu x64 (CUDA 13) - CUDA 13.4 libraries
- Ubuntu arm64 (CUDA 13) - CUDA 13.4 libraries
- Ubuntu x64 (ROCm 10.0)
- Ubuntu x64 (OpenVINO)
- Ubuntu x64 (SYCL FP32)
- Ubuntu x64 (SYCL FP16)
- Linux arm64 (Snapdragon: CPU, Adreno GPU, Hexagon NPU) - setup guide
Android:
Windows:
- Windows x64 (CPU)
- Windows arm64 (CPU)
- Windows arm64 (OpenCL Adreno)
- Windows x64 (CUDA 12) - CUDA 12.4 DLLs
- Windows x64 (CUDA 13) - CUDA 13.4 DLLs
- Windows arm64 (CUDA 13) - CUDA 13.4 DLLs
- Windows x64 (Vulkan)
- Windows x64 (OpenVINO)
- Windows x64 (SYCL)
- Windows x64 (ROCm 10.0)
openEuler:
- DISABLED
- openEuler x86 (310p)
- openEuler x86 (910b, ACL Graph)
- openEuler aarch64 (310p)
- openEuler aarch64 (910b, ACL Graph)
UI: