FlashDecoding: How Sequence Partitioning Solved the Memory Bandwidth Bottleneck in LLM Generation

In large language model serving, execution divides into two distinct operational regimes: prompt prefill and autoregressive token generation (decoding). While FlashAttention transformed prefill throughput by eliminating High Bandwidth Memory (HBM) round-trips for intermediate attention matrices, standard FlashAttention algorithms encounter a severe hardware utilization bottleneck during decoding.

Autoregressive decoding processes only a single query token per active sequence ($S_q = 1$) while scanning key-value caches that span thousands or tens of thousands of past tokens ($S_{kv} \gg 1$). Under standard parallelization schemes, this dynamic leaves modern GPUs heavily underutilized.

To overcome this hardware stall, researchers from Stanford CRFM and Meta AI introduced FlashDecoding (also known as Split-KV attention). By adding the sequence length of keys and values as a dedicated parallelization dimension, FlashDecoding saturates GPU compute cores, unlocks full memory bandwidth, and delivers up to 8x faster per-token generation in long-context workloads.

The Generation Wall: Why Autoregressive Decoding Stalls GPUs

The core difference between prefill and decoding lies in arithmetic intensity, defined as the ratio of compute operations (FLOPs) to memory traffic (bytes accessed).

During prompt prefill, the query tensor contains the entire input sequence ($S_q \ge 1024$). Computing the attention matrix involves large matrix-matrix multiplications (GEMMs). Because data loaded from HBM into on-chip SRAM is reused across multiple tokens, arithmetic intensity is high, allowing GPU Tensor Cores to operate near peak computational throughput.

Prefill Attention:
Q (S_q x d) @ K^T (d x S_kv) -> High arithmetic intensity (Compute-bound GEMM)

Decoding Attention:
Q (1 x d)   @ K^T (d x S_kv) -> Low arithmetic intensity (Memory-bandwidth-bound GEMV)

During decoding, the query is a single vector ($S_q = 1$). Computing attention scores against the past context is no longer a dense matrix multiplication; it is a matrix-vector product (GEMV). For every generated token, the GPU must stream the entire accumulated Key ($K$) and Value ($V$) tensors from HBM into SRAM, perform a minimal number of multiply-accumulate steps with the single query vector, and discard the loaded weights.

The arithmetic intensity drops by two to three orders of magnitude. Under these conditions, performance is entirely constrained by memory bandwidth.

The Parallelization Failure of Standard Attention

Standard FlashAttention algorithms schedule work across two primary dimensions:

1. Batch size ($B$)
2. Number of query attention heads ($H$)

Each thread block (CUDA Cooperative Thread Array, or CTA) is assigned a single (batch_idx, head_idx) pair and processes the entire historical sequence length sequentially.

This scheduling strategy fails when serving latency-critical requests at small batch sizes:

• Hardware Underutilization: Modern data center GPUs contain substantial parallel processing resources. An NVIDIA A100 features 108 Streaming Multiprocessors (SMs), while an NVIDIA H100 provides 132 SMs. In an interactive scenario with batch size $B = 1$ and a model with 32 attention heads (such as Llama-3-8B), standard FlashAttention launches only 32 thread blocks. As a result, 75% to 80% of the GPU Streaming Multiprocessors sit completely idle.
• Serial Context Iteration: The few active SMs must sequentially load the entire KV cache block by block across 8k, 32k, or 128k tokens. Memory bandwidth is bounded by the memory throughput of individual SMs rather than the aggregate memory bus of the entire chip.

The FlashDecoding Architecture: Split-KV Parallelism

Introduced by Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov in late 2023, FlashDecoding resolves this bottleneck by decomposing the KV sequence length dimension into independent parallel splits.

Instead of launching $B \times H$ thread blocks, FlashDecoding launches:

$$\text{Total Thread Blocks} = B \times H \times M$$

where $M$ represents the number of sequence splits.

FlashDecoding divides the computation into two decoupled phases: parallel local attention evaluation and a final reduction step.

Phase 1: Local Partial Attention Computation

In the first phase, the sequence length $S_{kv}$ is partitioned into $M$ equal chunks of size $B_c$ (typically 256 or 512 tokens). Each thread block is assigned a distinct tuple (batch_idx, head_idx, split_idx) and executes in parallel:

1. Chunk Ingestion: The thread block loads its designated chunk of Key ($K^{(m)}$) and Value ($V^{(m)}$) tensors from HBM into SRAM.
2. Local Score Computation: It calculates the local unnormalized attention scores $S^{(m)} = Q (K^{(m)})^T / \sqrt{d}$.
3. Local Softmax Statistics: Using the online softmax algorithm, the thread block computes:

• Local maximum score: $m^{(m)} = \max_{j} S_{j}^{(m)}$
• Local softmax normalizer: $l^{(m)} = \sum_{j} \exp(S_{j}^{(m)} - m^{(m)})$

1. Partial Value Accumulation: It computes the partial output vector:

$$O^{(m)} = \sum_{j} \exp(S_{j}^{(m)} - m^{(m)}) V_{j}^{(m)}$$

1. Intermediate Storage: Each thread block writes its partial output vector $O^{(m)} \in \mathbb{R}^{d}$ and its scalar statistics $(m^{(m)}, l^{(m)})$ to a temporary intermediate global buffer in HBM.

Because $M$ can be scaled up to 32, 64, or 128 depending on context length, the total number of blocks easily reaches 512 to 2048, achieving 100% SM occupancy even at batch size 1.

Phase 2: Log-Sum-Exp Reduction Kernel

Because the softmax function is non-linear due to exponentiation, partial outputs cannot simply be summed together. In standard attention, scaling requires global normalization across all tokens.

FlashDecoding leverages the mathematical properties of the Log-Sum-Exp (LSE) identity to rescale and merge the $M$ partial accumulators in a lightweight secondary kernel (flash_fwd_combine):

1. Global Maximum Identification: The reduction kernel identifies the overall maximum score across all splits:

$$m^{\text{global}} = \max_{m=1}^M m^{(m)}$$

1. Rescaling Factors: Each split's normalizer is adjusted relative to the global maximum:

$$\alpha^{(m)} = \exp(m^{(m)} - m^{\text{global}})$$

1. Global Normalizer Aggregation: The global denominator is computed:

$$l^{\text{global}} = \sum_{m=1}^M l^{(m)} \cdot \alpha^{(m)}$$

1. Normalized Output Synthesis: The final attention output vector $O \in \mathbb{R}^{d}$ is synthesized by weighting each partial output:

$$O = \frac{1}{l^{\text{global}}} \sum_{m=1}^M O^{(m)} \cdot \alpha^{(m)}$$

Because $M$ is small relative to the sequence length (typically $16 \le M \le 64$), the reduction kernel processes negligible data volume, executing in under 20 microseconds. The compute and memory latency gains from parallelizing the KV cache reads dwarf the small reduction overhead.

Standard Decoding:
[ SM 0 ] -> Processes KV Tokens 0 to 32,768 sequentially (Slow, Low Occupancy)
[ SM 1..131 ] -> IDLE

FlashDecoding (Split-KV, M=32):
[ SM 0  ] -> Processes KV Tokens 0..1023     -> Outputs (O^(0),  m^(0),  l^(0))
[ SM 1  ] -> Processes KV Tokens 1024..2047  -> Outputs (O^(1),  m^(1),  l^(1))
...
[ SM 31 ] -> Processes KV Tokens 31744..32768-> Outputs (O^(31), m^(31), l^(31))
                                      |
                      [ flash_fwd_combine Kernel ]
                                      v
                             Final Token Output O

Dynamic Split Heuristics and GQA Packing

Determining the optimal number of splits $M$ requires balancing GPU occupancy against reduction overhead.

Heuristic Selection

In production implementations across vLLM, SGLang, and TensorRT-LLM, the partition count $M$ is selected dynamically based on hardware properties and sequence dimensions:

$$M = \min\left(\left\lceil \frac{S_{kv}}{\text{block\_size}} \right\rceil, \max\left(1, \left\lfloor \frac{\text{num\_SMs} \times \text{target\_waves}}{B \times H} \right\rfloor\right)\right)$$

• When $B \times H \ge \text{num\_SMs}$, the GPU is already fully occupied. $M$ defaults to 1, reverting to standard FlashAttention without reduction overhead.
• When $B \times H < \text{num\_SMs}$, $M$ increases proportionally to ensure every SM has at least one to two active warps to hide memory access latency.

Interaction with Grouped-Query Attention (GQA)

Under Grouped-Query Attention (used in Llama 3, Mistral, and Qwen models), multiple query heads share a single Key-Value head (e.g., an 8:1 query-to-KV head ratio).

FlashDecoding implementations exploit this structure through GQA packing. A single thread block loads a chunk of $K$ and $V$ once into shared memory, and simultaneously computes dot products against all 8 associated query heads. This amortizes the HBM load bandwidth across multiple query heads, multiplying the effective arithmetic intensity during the decode phase.

Evolution: FlashDecoding++, FlashInfer, and FA-3

Following the release of the initial FlashDecoding algorithm, several production extensions emerged:

FlashDecoding++

Developed by researchers from Tsinghua University and Shanghai Jiao Tong University, FlashDecoding++ introduced three architectural refinements:

• Unified Softmax with Pre-Determined Maxima: Replaces dynamic local max tracking with partial statistical bounds, reducing inter-warp synchronization.
• Flat GEMM Optimization: Bypasses padded tensor representations for variable-length batch serving, preventing redundant compute on padding tokens.
• Double Buffering and Async Copy: Maximizes overlap between asynchronous memory copy (cp.async) and Tensor Core execution on Ampere and Ada Lovelace GPUs.

Integration in PagedAttention and FlashInfer

Production inference engines do not store KV caches in contiguous memory buffers; they use paged memory allocators (PagedAttention). Libraries like FlashInfer implement Paged Split-KV kernels, where the sequence partition logic indexes over non-contiguous physical page tables in GPU memory while maintaining split reduction mechanics.

FlashAttention-3 and FlashAttention-4

In FlashAttention-3, Tri Dao and Jay Shah integrated Split-KV directly with NVIDIA Hopper Tensor Memory Accelerator (TMA) hardware asynchronous copy features and FP8 tensor core accumulation, allowing decoding attention to maintain near-constant latency curves up to 64k context lengths.

Performance Benchmarks and Serving Economics

The implementation of Split-KV sequence partitioning fundamentally altered the latency profile of long-context LLM serving:

Attention Latency Comparison (NVIDIA A100-80GB SXM4, Batch Size 1, Head Dim 128, 32 Heads):

Sequence Length   Standard Attn   FlashAttention-2   FlashDecoding (Split-KV)   Speedup
512 tokens        0.32 ms         0.28 ms            0.28 ms                    1.0x (M=1)
4,096 tokens      2.10 ms         1.85 ms            0.62 ms                    3.0x
16,384 tokens     8.40 ms         7.20 ms            1.35 ms                    5.3x
32,768 tokens     17.10 ms        14.80 ms           2.10 ms                    7.0x
65,536 tokens     35.60 ms        30.20 ms           3.90 ms                    7.7x

(Indicative single-token generation latencies on NVIDIA A100-80GB SXM4, Batch Size = 1, Head Dim = 128, 32 Heads).

Before FlashDecoding, single-token generation latency scaled linearly with context length, making interactive chatbots and real-time coding assistants unviable at context lengths exceeding 16k tokens.

By restructuring attention across the sequence dimension, FlashDecoding ensures that GPU execution scales with available hardware parallelism rather than sequential sequence length, establishing the foundational architecture for modern long-context inference engines.

Sources

• Stanford CRFM: Flash-Decoding for Long-Context Inference
• Dao-AILab: FlashAttention-2 GitHub Repository
• FlashDecoding++: Faster Large Language Model Inference with Asynchronous, Flat GEMM and Unified Softmax (arXiv:2311.01282)
• Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023 / arXiv:2309.06180)
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision
• FlashInfer: High-Performance GPU Kernel Library for LLM Serving