Hello, I'm Seongjae Park from OptAI.

This article explains the core ideas of FlashDecoding++, which was proposed for GPUs. The details of the device port will follow in the next article. Here, we focus on the principles behind its goal: computing softmax quickly and with efficient memory use over a long KV cache.

1. Why FlashDecoding++ for On-Device Inference?

During decoding, attention repeats the same operations for every token. The Query length remains fixed at 1, while Key/Value grows with the KV cache. Score computation and softmax therefore incur an O(K) cost at every step, proportional to the context length K. On both GPUs and NPUs, as this region grows, the time spent reading data increasingly dominates computation.

Before Flash-Decoding: a single Query scans the entire KV cache sequentially

*During decoding, a Query of length 1 must sequentially read the entire, continually growing Key/Value sequence.*

Flash-Decoding: split the KV cache into chunks, compute in parallel, and combine the results

*Flash-Decoding partitions the KV cache into multiple splits for parallel processing. Combining the partial results requires coordinating softmax statistics across blocks. (Source: PyTorch Blog — Flash-Decoding for long-context inference)*

FlashDecoding++ addresses this problem in three ways. Our focus is asynchronous softmax, which contributes most directly to on-device attention.

FlashDecoding++ inference dataflow overview and its three solutions

*Overview of the FlashDecoding++ inference dataflow. (a) Prefill primarily uses GEMM, while decode primarily uses GEMV/Flat GEMM. (b) Asynchronous softmax based on a unified maximum, (c) Flat GEMM optimization with double buffering, and (d) heuristic dataflow: the three bottlenecks FlashDecoding++ redesigns and their respective solutions. (Source: Hong et al., arXiv:2311.01282)*

2. Three Bottlenecks

(1) Synchronized partial Softmax. Splitting softmax into multiple blocks to process a long sequence in parallel requires synchronization to share and update the maximum and sum from each block before computing the full probabilities. This process alone introduces roughly 20% overhead in the attention computation.

(2) Low utilization in Flat GEMM. Decode-stage matrices have a flat shape with a very small M. Existing libraries increase their size by padding with 0 to improve hardware utilization, but the resulting unnecessary operations can account for more than half of the computation.

(3) Losses from static dataflow. The optimal kernel varies with the input shape and hardware configuration. A single static execution path cannot flexibly handle different GEMM shapes, resulting in substantial performance losses.

3. The Core Solution: Asynchronous Softmax with a Unified Maximum

FlashDecoding++ sets a statistically safe unified max value φ in advance so that each block can compute independently. This removes inter-block synchronization and leaves only a simple sum at the end, yielding speedups of 1.18× for prefill and 1.14× for decoding.

Conventional online softmax tracks the maximum during execution.

softmax(x_i) = exp(x_i - max_j x_j) / Σ_j exp(x_j - max_j x_j)

The unified-maximum approach replaces the running max with a constant φ fixed in advance.

softmax(x_i) = exp(x_i - φ) / Σ_j exp(x_j - φ)

Each block no longer needs to find the maximum again. However, if an x_i exceeds the threshold defined by φ and creates a risk of overflow, the block stops its asynchronous computation and falls back to conventional online softmax for safe recomputation. In other words, φ is chosen by observing the distribution beforehand; values large enough to exceed its range trigger the safe fallback path.

4. The Other Two Solutions

Flat GEMM optimization with double buffering. Minimizing padding and introducing double buffering allow computation to overlap with data reads. This latency hiding helps work around the memory bottleneck.

Heuristic dataflow adapted to hardware resources. The linear-layer computations in an LLM ultimately fall into four groups: Q/K/V generation, O projection, and the two FFN operations. The weight dimensions K·N are constant; only M, representing sequence length and batch, varies. A lookup table prepared offline selects kernels for different ranges of M, recovering losses from static dataflow.

Of these three solutions, asynchronous softmax contributes most directly to on-device attention. Deploying the algorithm on actual hardware, however, makes the results highly sensitive to quantization and execution-unit allocation. The next article covers that part of the story.

References
  1. Hong, K., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
  2. Dao, T., et al. (2023). Flash-Decoding for long-context inference. PyTorch Blog.
  3. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.