Flash-Decoding: Mathematical Foundations, Split-KV Parallelism, Online Softmax Reduction, and Serving Latency Optimization

Large language model inference consists of two distinct computational regimes: the compute-bound prefill phase and the memory-bandwidth-bound autoregressive decoding phase. While the prefill phase processes prompt tokens in parallel across dense matrix multiplication units, generation proceeds token by token. At each generation step, a single newly generated query vector must attend to the entire historical Key-Value (KV) cache stored in High-Bandwidth Memory (HBM).

As context windows scale from 4,000 tokens to 32,000, 128,000, and beyond, attention over the accumulated KV cache becomes the dominant latency bottleneck during token generation. Standard exact attention kernels such as FlashAttention-1 and FlashAttention-2 effectively eliminate high-bandwidth memory round-trips during training and prefill by tiling along the query sequence dimension. However, when applied to single-token decoding where the query length is exactly one ($Q = 1$), FlashAttention suffers from a structural parallelism collapse that leaves modern GPU hardware largely idle.

Flash-Decoding resolves this hardware underutilization by introducing parallelism along the KV sequence dimension. By partitioning the historical KV cache into independent chunks, processing attention across chunks concurrently over all available Streaming Multiprocessors (SMs), and executing an exact two-pass online softmax reduction, Flash-Decoding achieves up to an 8x end-to-end speedup and up to a 50x attention kernel acceleration for long-context generation without mathematical approximation.

The Single-Token Generation Bottleneck

During autoregressive generation, a transformer layer receives a single input token at time step $t$. The model projects this token into query, key, and value vectors:

$$q_t \in \mathbb{R}^{1 \times d}, \quad k_t \in \mathbb{R}^{1 \times d}, \quad v_t \in \mathbb{R}^{1 \times d}$$

The new key and value vectors are appended to the accumulated KV cache:

$$K = [k_1; k_2; \dots; k_t] \in \mathbb{R}^{N \times d}, \quad V = [v_1; v_2; \dots; v_t] \in \mathbb{R}^{N \times d}$$

where $N$ denotes the total sequence length and $d$ denotes the head dimension. The standard scaled dot-product attention output vector $o_t \in \mathbb{R}^{1 \times d}$ is defined as:

$$o_t = \text{Softmax}\left(\frac{q_t K^T}{\sqrt{d}}\right) V$$

In standard attention implementations, evaluating this equation requires reading $K$ and $V$ from GPU global memory (HBM) into on-chip Static Random-Access Memory (SRAM), computing dot products, writing intermediate attention scores back to HBM, evaluating softmax normalizers across all $N$ tokens, and performing a second round-trip read from HBM to compute the weighted sum over $V$. The memory traffic scales linearly with context length $N$, transferring $2 \times N \times d \times \text{sizeof(dtype)}$ bytes per attention head at every generated token.

Why FlashAttention Collapses During Autoregressive Decoding

FlashAttention optimizes the prefill phase by tiling the query sequence $Q \in \mathbb{R}^{N_q \times d}$ into blocks of size $B_r$ and the key/value sequence into blocks of size $B_c$. FlashAttention assigns each query tile to an independent GPU thread block.

Because the prefill sequence length $N_q$ is large, the total number of parallel thread blocks scheduled across the GPU grid is:

$$\text{Grid Size}_{\text{prefill}} = \text{Batch Size} \times \text{Num Heads} \times \left\lceil \frac{N_q}{B_r} \right\rceil$$

On an NVIDIA A100 GPU containing 108 Streaming Multiprocessors (SMs), a prefill batch with $B=4$, $H=32$, and $N_q=2048$ with block size $B_r=128$ generates $4 \times 32 \times 16 = 2,048$ thread blocks. This workload saturates all 108 SMs with ample occupancy.

During autoregressive generation, however, $N_q = 1$. The ceiling term $\lceil N_q / B_r \rceil$ collapses to 1. The total thread block grid size becomes strictly:

$$\text{Grid Size}_{\text{decode}} = \text{Batch Size} \times \text{Num Query Heads}$$

When serving long-context requests, batch sizes are constrained by KV-cache memory capacity. For a single inference stream ($B=1$) with Grouped-Query Attention (GQA) utilizing 8 key-value heads or 16 query heads, FlashAttention generates only 16 thread blocks. When scheduled onto an A100 GPU with 108 SMs, 92 SMs remain completely unallocated. The GPU operates at under 15% hardware occupancy regardless of how long the context length $N$ grows. The single thread block handling each head is forced to iterate sequentially over all $N / B_c$ key-value tiles, creating a severe latency bottleneck.

Mathematical Formulation of Split-KV Parallelism

Flash-Decoding overcomes this limitation by partitioning the key and value sequence dimension into $S$ non-overlapping chunks.

Let the total context length $N$ be divided into $S$ uniform partitions:

$$N = \sum_{s=1}^{S} N_s$$

For each partition $s \in \{1, 2, \dots, S\}$, define the local key and value slices as:

$$K^{(s)} = K_{(s-1)N_s : s N_s, :} \in \mathbb{R}^{N_s \times d}, \quad V^{(s)} = V_{(s-1)N_s : s N_s, :} \in \mathbb{R}^{N_s \times d}$$

The GPU scheduling grid expands by a factor of $S$:

$$\text{Grid Size}_{\text{Flash-Decoding}} = \text{Batch Size} \times \text{Num Query Heads} \times S$$

By setting $S = \lceil N / \text{SplitSize} \rceil$ where $\text{SplitSize}$ is typically chosen between 128 and 256 tokens, a context length of $N = 32,768$ yields $S = 128$ splits. Even at batch size $B=1$ with 16 query heads, the kernel dispatches $16 \times 128 = 2,048$ independent thread blocks across the GPU, saturating all SMs.

Each split $s$ independently executes the FlashAttention kernel on its local slice $K^{(s)}$ and $V^{(s)}$. It calculates:

1. Local unnormalized attention logits:

$$S^{(s)} = \frac{q K^{(s)T}}{\sqrt{d}} \in \mathbb{R}^{1 \times N_s}$$

1. Local maximum logit scalar for numerical stability:

$$m^{(s)} = \max_{j \in [1, N_s]} S^{(s)}_j$$

1. Local unnormalized exponential sum:

$$\ell^{(s)} = \sum_{j=1}^{N_s} \exp(S^{(s)}_j - m^{(s)})$$

1. Local partial attention output vector:

$$O^{(s)} = \frac{1}{\ell^{(s)}} \sum_{j=1}^{N_s} \exp(S^{(s)}_j - m^{(s)}) V^{(s)}_j \in \mathbb{R}^{1 \times d}$$

Along with the partial output vector $O^{(s)}$, each thread block writes two scalars per attention head to global memory: the local row maximum $m^{(s)}$ and the log-sum-exp value $L^{(s)} = m^{(s)} + \ln(\ell^{(s)})$.

The Rescaling and Online Softmax Reduction Algorithm

Because softmax normalizers depend on global sums across the entire sequence, individual split outputs $O^{(s)}$ cannot simply be averaged. Flash-Decoding applies the mathematical properties of online softmax to perform an exact reduction over the $S$ partial outputs.

Let $m$ denote the true global maximum logit across all $S$ splits:

$$m = \max_{s \in \{1, \dots, S\}} m^{(s)}$$

To combine the local exponential sums $\ell^{(s)}$ under the global maximum $m$, each local sum is rescaled by the exponential difference $\exp(m^{(s)} - m)$:

$$\ell = \sum_{s=1}^{S} \ell^{(s)} \exp(m^{(s)} - m)$$

The exact global attention output vector $O \in \mathbb{R}^{1 \times d}$ is then computed as the normalized weighted combination of all partial outputs:

$$O = \frac{1}{\ell} \sum_{s=1}^{S} \left( \ell^{(s)} \exp(m^{(s)} - m) \right) O^{(s)}$$

Alternatively, expressing the combination in terms of log-sum-exp scalars $L^{(s)} = m^{(s)} + \ln(\ell^{(s)})$, the global log-sum-exp is:

$$L = \ln \left( \sum_{s=1}^{S} \exp(L^{(s)}) \right)$$

and the output weight for split $s$ is:

$$\alpha^{(s)} = \exp(L^{(s)} - L)$$

$$O = \sum_{s=1}^{S} \alpha^{(s)} O^{(s)}$$

Because this formulation preserves exact algebraic equivalence with full-sequence softmax, Flash-Decoding introduces zero quantization error, zero precision drift, and zero approximation noise.

Hardware Execution Topology and Memory Traffic

Flash-Decoding executes in two sequential GPU kernel launches:

[KV Cache in HBM: Length N]
       │
       ├── Split 1: KV[0 : N/S]     ───> SM Threadblock 1  ───> Partial (O^(1), m^(1), l^(1))
       ├── Split 2: KV[N/S : 2N/S]  ───> SM Threadblock 2  ───> Partial (O^(2), m^(2), l^(2))
       ├── ...                      ───> ...               ───> ...
       └── Split S: KV[(S-1)N/S : N]───> SM Threadblock S  ───> Partial (O^(S), m^(S), l^(S))
                                                                       │
                                                                       ▼
                                                         [Kernel 2: Reduction Stage]
                                                         SM Threadblock reads S partials
                                                         Rescales via exp(m^(s) - m)
                                                                       │
                                                                       ▼
                                                         Final Output Vector O in HBM

The memory overhead introduced by the intermediate buffers is negligible:

• Intermediate outputs: $B \times H \times S \times d \times \text{sizeof(dtype)}$ bytes.
• Intermediate metadata: $B \times H \times S \times 2 \times 4$ bytes (for FP32 scalars $m^{(s)}$ and $\ell^{(s)}$).

For a representative configuration with $B=1$, $H=32$, $S=128$, $d=128$, and 16-bit floats, the intermediate buffer requires only $1 \times 32 \times 128 \times 128 \times 2 = 1.05 \text{ MB}$, fitting comfortably within GPU L2 cache (which spans 40 MB to 50 MB on modern enterprise accelerators).

The second reduction kernel reads $S$ vectors of size $d$, which requires negligible computation compared to reading the full KV cache.

Benchmarks and Performance Characteristics

Empirical evaluations across CodeLlama-34B and Llama-2 architectures on NVIDIA A100 and H100 GPUs demonstrate distinct performance regimes:

• Short Context Regimes ($N \le 1,024$ tokens): Standard FlashAttention and Flash-Decoding exhibit comparable latency because memory bandwidth over small caches is minimal and kernel launch overhead dominates.
• Medium Context Regimes ($1,024 < N \le 8,192$ tokens): FlashAttention execution time begins scaling linearly with context length due to single-SM serialization. Flash-Decoding latency remains virtually flat as additional SMs are engaged to absorb the split workloads.
• Long Context Regimes ($N \ge 16,384$ tokens to $65,536$ tokens): Flash-Decoding demonstrates up to an 8x reduction in total token generation time and up to a 50x speedup in the attention sub-layer relative to FlashAttention-2 at batch size 1.

Scaled multi-head attention execution time micro-benchmarks on NVIDIA A100-SXM4-80GB (FP16, 16 query heads, head dimension $d=128$, batch size $B=1$):

• Context length 512 tokens: PyTorch Eager 3,151 microseconds, FlashAttention-2 366 microseconds, Flash-Decoding 67 microseconds.
• Context length 4,096 tokens: PyTorch Eager 3,157 microseconds, FlashAttention-2 401 microseconds, Flash-Decoding 57 microseconds.
• Context length 16,384 tokens: PyTorch Eager 3,223 microseconds, FlashAttention-2 582 microseconds, Flash-Decoding 58 microseconds.
• Context length 32,768 tokens: PyTorch Eager 3,224 microseconds, FlashAttention-2 1,156 microseconds, Flash-Decoding 60 microseconds.
• Context length 65,536 tokens: PyTorch Eager 1,335 microseconds, FlashAttention-2 2,300 microseconds, Flash-Decoding 64 microseconds.

Modern Evolutions: FlashDecoding++, FlashInfer, and Multi-Block Attention

The split-KV paradigm pioneered by Flash-Decoding has been integrated into production LLM serving runtimes through several advanced variants:

1. FlashDecoding++ (MLSys 2024): Traditional Flash-Decoding requires partial output synchronization across all splits. FlashDecoding++ introduces flat-gemm optimization and unified maximum approximation using statistical prior bounds, eliminating dynamic maximum synchronization across splits.
2. Multi-Block Attention (NVIDIA TensorRT-LLM): TensorRT-LLM implements Multi-Block Attention (MBA), which integrates split-KV tiling directly into custom paged memory managers with runtime heuristic switching between single-block and multi-block execution modes based on batch occupancy.
3. FlashInfer: FlashInfer provides customized CUDA and Cutlass templates for split-KV decoding across ragged, paged, and heterogeneous batch topologies, dynamically balancing split chunk sizes when different requests in a batch have vastly different context lengths.

Production Implementation Guidelines

When deploying Flash-Decoding in high-throughput inference engines such as vLLM, SGLang, and TensorRT-LLM, engineers should adhere to the following operational practices:

• Dynamic Heuristic Dispatching: Use Flash-Decoding when the total parallel work factor $B \times H_{\text{query}} < \text{Num SMs}$. When large batch sizes already saturate the GPU ($B \times H_{\text{query}} \ge 2 \times \text{Num SMs}$), revert to standard non-split decoding kernels to avoid the reduction kernel synchronization overhead.
• Tuning Split Granularity: Set the split chunk size $N_s$ between 128 and 256 tokens. Choosing $N_s < 64$ creates excessive reduction overhead, while $N_s > 512$ degrades SM occupancy on intermediate sequence lengths.
• Coordination with Paged KV-Caches: Ensure split boundaries align with block boundaries in paged memory allocators (e.g., 16 or 32 tokens per block) to prevent memory misalignment during asynchronous global loads.

Sources

• PyTorch Technical Blog: Flash-Decoding for long-context inference
• Tri Dao et al.: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (NeurIPS 2022)
• Tri Dao: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024)
• Hong et al.: FlashDecoding++: Faster Large Language Model Inference on GPUs (MLSys 2024)
• FlashAttention GitHub Repository (Dao-AILab)