A benchmark comparing standard SDPA against Flash-Decoding-style chunked attention at the decode step of autoregressive LLM generation.
Closes a pair with flash-attention-benchmark:
flash-attention-benchmark: prefill attention (compute-bound, many query tokens) flash-decoding-bench: decode attention (bandwidth-bound, q_len=1)
For full methodology and design decisions see design.md.
FlashAttention parallelizes over query tokens. At decode time, q_len=1. There is nothing to parallelize over. The GPU is underutilized.
Flash Decoding fixes this by parallelizing over the KV cache dimension instead. This benchmark measures where that strategy wins.
naive Standard q @ k^T -> softmax -> @ v sdpa torch.scaled_dot_product_attention (fused kernel reference) sequential Flash-Decoding-style chunks, sequential Python loop parallel Flash-Decoding-style chunks, concurrent CUDA streams auto variants Auto-select chunk size by KV length
Hardware: RTX 2070, PyTorch 2.13
auto_parallel vs SDPA (bs=1, q_len=1):
kv=512: 0.10x stream overhead dominates kv=1024: 0.23x kv=2048: 0.46x kv=4096: 0.95x crossover kv=8192: 1.38x parallel chunked wins kv=16384: 1.36x
auto_sequential vs SDPA (bs=1, q_len=1):
kv=512: 0.90x kv=16384: 1.13x only wins at extreme KV length
SDPA wins for short KV lengths. The fused kernel reads KV in a single efficient pass. Chunk overhead dominates when KV is small.
Parallel chunked decode wins above approximately 8K KV tokens. At kv=8192: 1.38x speedup vs SDPA for bs=1. This is the regime where KV-parallelism pays off.
Sequential chunked decode only wins at extreme KV lengths. Python loop overhead limits speedup to 1.13x even at 16K tokens. This is why Flash Decoding requires a fused CUDA kernel.
Batch size 4 does not benefit from chunked decode. With batched requests, SDPA already achieves good GPU utilization. Flash Decoding benefit is most valuable at batch_size=1.
Measured gains are a lower bound on real Flash Decoding. A fused CUDA kernel moves the crossover to shorter KV lengths and achieves higher peak speedup. Production systems (FlashInfer, vLLM, SGLang) implement exactly this.
flash-decoding-bench/ ├── src/ │ ├── init.py │ ├── config.py │ ├── tensors.py │ ├── kernels.py │ ├── benchmark.py │ └── analysis.py ├── results/ ├── plots/ ├── LICENSE ├── design.md ├── README.md ├── requirements.txt └── run.py
python3 -m venv venv source venv/bin/activate pip install -r requirements.txt python run.py
Outputs:
results/results.csv results/summary.txt plots/decode_vs_kv.png plots/chunk_sweep.png plots/q_len_scaling.png
MIT License. See LICENSE for details.
Joao Felipe De Souza 2026