Loading post…
ModalCharles FryeExplainer
Making FlashAttention-4 faster for inference
Summary
FlashAttention-4, originally designed for pre-training, required modifications to efficiently handle LLM inference workloads dominated by the memory bandwidth-limited decode phase. Modal’s contributions include adjusting parallelism strategies from query parallelism to key/value parallelism, implementing split KV to parallelize across KV tiles, and adding a heuristic for optimal splits. They added support for FP8 inputs, arbitrary KV page sizes via cp.async loads, and reduced wasted work for small query sizes by using a single query tile. The kernel also gained extended GQA packing for irregular head ratios, using cp.async for query tiles to bypass TMA alignment restrictions. These changes, measured in throughput improvements up to 3x for single-token decode and up to 2.4x for small page sizes, demonstrate the leverage of flexible parallelism and tile-based programming for inference optimization.
Context
Modal, an infrastructure provider for AI inference, contributed optimizations to the open-source FlashAttention-4 kernel to improve performance for large language model inference, particularly the decode phase. The original kernel was oriented toward pre-training workloads with large query tiles and regular memory access patterns, but inference workloads are dominated by variable batch sizes, sequence lengths, and KV cache retrieval. Modal’s changes focus on adjusting parallelism strategies, supporting irregular memory accesses, and adding lower-precision input support.
Approach / What changed
The optimizations include: adjusting parallelism from query parallelism to key/value parallelism and implementing split KV to parallelize across KV tiles; supporting irregular memory accesses via cp.async loads instead of TMA for arbitrary KV page sizes; adding FP8 input support (e4m3 and e5m2); reducing wasted work for small query sizes by operating on a single query tile; and extending GQA packing for irregular head ratios using cp.async for query tiles to bypass TMA alignment restrictions. A heuristic determines optimal split KV count based on SM count and sequence length.
Takeaways
Adjusting parallelism strategy (e.g., from query to KV parallelism) provides higher leverage than optimizing a fixed parallel component, as per Amdahl’s Law, because it moves work between sequential and parallel components.
Using cp.async loads instead of TMA allows supporting arbitrary KV page sizes, improving KV cache efficiency at the cost of lower arithmetic and memory throughput, which is a beneficial trade-off for many inference workloads.
Split KV parallelizes across KV tiles but introduces an extra reduction step, so a heuristic based on SM count and sequence length determines the optimal number of splits to avoid performance degradation.
FAQ
Why does FlashAttention-4 need changes for inference?
The original kernel was designed for pre-training with large query tiles and regular memory access, but inference is dominated by the decode phase with variable batch sizes, sequence lengths, and KV cache retrieval, requiring new parallelism and memory access strategies.
How does FP8 support improve inference performance?
FP8 reduces memory and arithmetic bandwidth demand by half, leading to faster inference and smaller KV caches. The speedup is less than 2x because the softmax operation still operates at higher precision on CUDA Cores.
What is the trade-off of using cp.async loads over TMA?
cp.async allows arbitrary KV page sizes for better cache efficiency but has lower arithmetic and memory throughput compared to TMA, making it a trade-off that prioritizes KV cache efficiency over raw throughput.
Source and freshness
- Source
- modal.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.