attention-kernel-optimization
verified443b357b-ee83-4b54-a1d4-552c77cb0243
Understand and apply the GPU-kernel optimizations that make long-context inference possible — FlashAttention 1/2/3, PagedAttention, Flash-Decoding, and FlashInfer — with when to reach for each.
Metadata
Skill file
# Attention-Kernel Optimization for LLM Inference
Use when long-context inference is your bottleneck — serving 32K+ context, high
throughput with varying sequence lengths, or latency-sensitive streaming — and you
need to understand *which* kernel-level optimization your stack is using (or should
use), and what's happening under the hood. These techniques are the reason 128K+
context is practical today.
## Why standard attention kills long context
Naïve attention computes the full N×N attention matrix and writes it to HBM
(high-bandwidth memory). For an N=128K sequence with head_dim=128, that's ~132 GB
of intermediate data that must be read/written to HBM — far exceeding GPU memory
bandwidth. The bottleneck isn't compute, it's *data movement*. Every optimization
below attacks this memory bandwidth problem, not FLOPs.
## FlashAttention (1 → 2 → 3)
- **FA-1 (2022)**: IO-aware tiling — stream tiles of Q, K, V that fit in SRAM,
compute attention locally, and keep the softmax rescaling correct with online
normalization. Memory is O(N·d) instead of O(N²). The key insight is recomputing
intermediate results in SRAM rather than writing them to HBM and reading back.
- **FA-2 (2023)**: Better parallelism across warps, reduced non-matmul FLOPs, ~2×
faster than FA-1 on long sequences.
- **FA-3 (2024)**: Built for Hopper GPUs (H100). Uses asynchronous pipelining
(overlapping GEMM and softmax), FP8 low-precision support, and Tensor Core
acceleration. Achieves 740 TFLOPS/s — 75% of H100's theoretical peak, up from
35% utilization in FA-2. For sequences ≥1K tokens, FA-3 even beats vendor
hand-tuned libraries (cuDNN).
## Inference-specific: Flash-Decoding and FlashInfer
Training kernels parallelize across batch×heads — fine for training where query
length is large. Inference has query length = 1 (one token at a time), so this
parallelism doesn't work. Two solutions:
- **Flash-Decoding**: Split the K/V sequence across blocks, compute attention
in parallel, then reduce. Designed specifically for the q_len=1 case.
- **FlashInfer**: Builds on FA-3 and adds KV-cache-aware inference scheduling,
reducing inter-token latency by 29-69% vs standard backends.
## PagedAttention (the vLLM revolution)
PagedAttention, pioneered in vLLM, treats the KV cache as *pages* — fixed-size
blocks that can be allocated non-contiguously — instead of one giant contiguous
tensor. This eliminates fragmentation and allows memory sharing across requests
(prefix caching). Result: near-zero memory waste, 2-4× throughput for serving, and
the ability to batch requests with wildly different sequence lengths efficiently.
PagedAttention alone is responsible for most of vLLM's throughput advantage over
naive serving setups.
## What you control (the practical knob)
Most teams don't write attention kernels. But you choose the serving framework that
wraps the right ones:
- **vLLM**: PagedAttention + FA-2/3 + FlashInfer for KV-cache scheduling.
- **SGLang**: Similar stack; RadixAttention for prefix caching.
- **Text-Generation-Inference (TGI)**: Flash-Attention integration.
- **llama.cpp**: Has its own optimized attention for CPU/GPU.
The practical lever: `--enable-prefix-caching` (shared KV cache for identical
prefixes, massive win when many requests share a system prompt), `--max-model-len`,
and choosing the right attention backend.
## Pitfalls
- Assuming the latest kernel (FA-3) helps on all hardware — it's Hopper-specific.
FA-2 is the right call on Ampere (A100/A10).
- Ignoring prefix caching when your system prompt is large and shared — easily the
largest single throughput optimization.
- Believing attention optimization solves all latency — for short sequences,
attention is not the bottleneck (the MLP layers and linear projections dominate).
- Using training-optimized kernels for inference without Flash-Decoding — the
q_len=1 case needs different parallelism.
## Verify
- Profile your inference with `nvidia-smi` or a framework's built-in metrics before
and after switching backends; measure tokens/sec/gpu and inter-token latency.
- For a shared-prompt workload, benchmark throughput with and without prefix caching.
- Check the attention backend your serving framework is actually using (vLLM logs it
at startup).Attached files
No attached files.