ArXiv: 2311.01282
🎯 Pitch
Simply removing synchronization from the softmax in LLM attention can silently produce incorrect results 2% of the time, but FlashDecoding++ shows you can skip it entirely 98% of the time by using a unified max value and falling back to recomputation only on rare overflow—slashing a pervasive 20% attention overhead. The engine further squeezes out performance by padding flat matrix multiplications to just 8 rows instead of 64 to boost utilization, and dynamically picks between CUDA cores and Tensor Cores per layer based on the matrix shape, yielding up to 4.86× speedup over naive implementations.
1. Executive Summary
FlashDecoding++ introduces a fast LLM inference engine that accelerates both the prefill and decode phases on NVIDIA and AMD GPUs by targeting three bottlenecks in existing frameworks: synchronized partial softmax updates that cause ~20% attention-computation overhead, under-utilized flat GEMM operations from padding to large tile sizes, and performance loss from static dataflow that ignores input dynamics and hardware configurations. The system proposes asynchronized softmax with unified max value (replacing per-tile max synchronization with a single global scaling constant, falling back to synchronized recomputation only on overflow), flat GEMM optimization with double buffering (padding the M-dimension to 8 instead of 64 and overlapping memory loads with computation), and heuristic dataflow with hardware resource adaptation (dynamically selecting between CUDA-core GEMV, the custom flat GEMM, and Tensor-Core CUTLASS based on offline-profiled inflection points for the four unique [N, K] shapes per model). FlashDecoding++ achieves up to 4.86× speedup over Hugging Face implementations and an average 1.37× speedup over the state-of-the-art FlashDecoding engine across mainstream LLMs — establishing that kernel-level optimization can yield substantial inference gains even without modifying model architecture, provided the engine adapts both its attention numerics and its GEMM dataflow to per-layer shape and hardware characteristics.
2. Context and Motivation
The Core Problem: LLM Inference Is Bottlenecked at the Kernel Level
As large language models transition from research artifacts to production services, the cost and latency of inference have become dominant economic and engineering concerns. The paper opens with a stark set of numbers: GPT-4 inference with 8K context costs 0.06 per 1K output tokens, and with OpenAI receiving over 10 million queries per day from 180.5 million users, the daily hardware cost of operating models like ChatGPT is estimated at approximately $7 million (Section 1). This is not a marginal expense — it is the primary operational cost of LLM deployment. Consequently, even modest improvements in per-token inference speed translate to enormous absolute savings.
The paper focuses on a specific layer of the inference optimization stack: single-GPU kernel performance. This sits beneath system-level optimizations like model parallelism (DeepSpeed), memory management (vLLM's PagedAttention), and offloading (FlexGen), and above hardware-level microarchitecture. A well-optimized kernel directly reduces the wall-clock time of each transformer layer's linear algebra and attention operations, which in turn multiplies across dozens of layers and millions of tokens.
The paper identifies three specific, measurable bottlenecks that existing kernel implementations fail to address adequately, each contributing a non-trivial fraction of total inference latency. These bottlenecks are not speculative — they are grounded in profiling data from Llama2-7B inference on an NVIDIA A100 GPU (Section 1 and 3):
- Synchronized partial softmax update accounts for ~18.8% of attention computation overhead.
- Padding flat GEMMs to tile sizes of 64 (standard in cuBLAS/CUTLASS) wastes >50% of computation for decode-phase workloads.
- Static dataflow — using the same kernel implementation regardless of input shape — causes up to 50.25% performance loss for GEMMs of different shapes.
These three bottlenecks collectively mean that existing inference engines leave substantial throughput on the table, even on well-optimized hardware. The paper's central claim is that addressing all three simultaneously — through numerical reformulation of attention, adaptive GEMM tiling, and shape-aware kernel dispatch — can recover this wasted performance.
Why These Three Bottlenecks Matter: The Prefill-Decode Asymmetry
To understand why these are the right bottlenecks to attack, we need to examine the structure of LLM inference itself. Section 2.1 and Figure 2 decompose inference into two phases that stress GPU hardware differently:
The prefill phase processes the entire input prompt in one forward pass. The sequence length () is large (hundreds to thousands of tokens), so the matrix-matrix multiplies in the K, Q, V, O projections and feedforward layers are conventional GEMMs — the M-dimension (sequence length × batch size) is large, making these operations compute-bound. The attention computation is quadratic in sequence length (), so it dominates prefill latency for long contexts. This is the domain where FlashAttention (Dao et al., 2022) and FlashAttention-2 (Dao, 2023) have already made substantial progress by restructuring attention to be I/O-aware.
The decode phase generates output tokens one at a time, autoregressively. Each new token requires processing a single input token (or a small batch of tokens) through all layers, reading the full KVcache from previous steps. The matrix multiplies become GEMV (batch size = 1) or flat GEMM (batch size > 1 but small) — the M-dimension is 1–8, while K and N remain large (hidden dimensions of 4096–12288 for modern 7B+ models). These operations are memory-bound: the GPU's compute units spend most of their time waiting for weight matrices to be loaded from HBM, since each weight element is touched only once to produce a single output element. The attention computation during decoding involves a dot product between one query vector and the entire key cache, followed by a softmax and a weighted sum over the value cache — operations that are also memory-bound but structured differently from prefill attention.
This asymmetry means that no single kernel strategy works optimally across both phases. The prefill phase needs algorithms that maximize compute utilization and minimize I/O overhead for large attention matrices. The decode phase needs algorithms that minimize memory traffic and latency for GEMV/flat GEMM operations. Existing frameworks tend to use the same GEMM library (cuBLAS) for both, or the same attention implementation (FlashAttention) for both, incurring phase-specific inefficiencies that the paper quantifies.
The three bottlenecks map onto this asymmetry as follows: Bottleneck 1 (synchronized softmax) primarily affects attention in both phases, since partial softmax with synchronized updates is the mechanism used by FlashAttention and FlashDecoding to achieve parallelism and memory efficiency. Bottleneck 2 (flat GEMM under-utilization) is specific to the decode phase, where the small M-dimension makes padding to tile size 64 wasteful. Bottleneck 3 (static dataflow) affects all GEMMs across both phases, since different [M, N, K] shapes reach peak performance with different implementations (CUDA core GEMV, flat GEMM with small tiles, or Tensor Core GEMM with large tiles).
Detailed Breakdown of Each Bottleneck
Bottleneck 1: Synchronized Partial Softmax (Section 3)
The softmax operation is inherently serial: for a vector , computing requires three passes over the data — finding the maximum , computing exponents and their sum , then normalizing. For large attention matrices (e.g., produces a attention matrix), these passes are memory-intensive and latency-critical.
FlashAttention introduced the partial softmax technique to parallelize this: split the input vector into partial vectors (tiles), compute local softmax statistics for each tile independently, then synchronize — each tile's results are updated using the maximum and sum from other tiles (Equation 2 in Section 2.3). This synchronization enables parallelism at the cost of a dependency: when a new partial softmax result is computed, previous partial results must be recomputed (rescaled) using the new global maximum and updated normalization constant.
The paper profiles this on Llama2-7B with 1024 input length on an A100 GPU and finds that the synchronization overhead is 18.8% of total attention computation time (Section 3). This is not a theoretical concern — it is a measurable, substantial fraction of the runtime. The reason is that the synchronization requires either (a) stalling one tile's computation until another tile completes, or (b) recomputing earlier tiles with updated statistics, both of which waste compute or introduce pipeline bubbles.
The fundamental challenge: the maximum value is different across tiles, so each tile cannot finalize its softmax output until all tiles have computed their local maxima and the global maximum is known. This is mathematically necessary for numerical stability — the exponent would overflow if is large and no shift is applied — but the synchronization it demands is a software artifact of the algorithm, not a hardware constraint.
Bottleneck 2: Under-Utilized Flat GEMM (Section 4)
Modern GPU BLAS libraries (cuBLAS, CUTLASS) are optimized for GEMMs where both dimensions are large enough to justify tiling into blocks that hide memory latency. The standard tile size in the M-dimension is 64 — this means the library processes 64 rows of the output matrix simultaneously within a single thread block. The tiling serves two purposes: (1) it allows each thread block to amortize the cost of loading weight tiles from global memory across multiple output rows, increasing the computation-to-memory ratio, and (2) it provides enough parallel work to keep streaming multiprocessors busy.
However, during the decode phase, the M-dimension is the batch size (typically 1–8). To use the standard cuBLAS/CUTLASS kernels, the M-dimension is padded to 64 with zeros. The computation performed on these padding rows is wasted — the hardware multiplies and accumulates zeros that contribute nothing to the output. For batch size 8, this means only 12.5% of the M-dimension computation is useful; 87.5% is padding overhead. The paper reports that this leads to ">50% computation under-utilization" (Section 1) for flat GEMMs.
The natural solution — pad to a smaller tile size, say 8, matching the native Tensor Core granularity — introduces a different problem. When the M-dimension tile is small, each thread block does less computation per memory load, reducing the computation-to-memory ratio. Equation (5) in Section 4 formalizes this tradeoff:
where is the tiling size in the N-dimension. As increases (more columns per tile), the denominator shrinks (fewer tile loads overall), improving the ratio — but the parallelism () decreases, since fewer tiles means fewer thread blocks to distribute across streaming multiprocessors. For small N, the GEMM is parallelism-bounded (not enough tiles to fill all SMs). For large N, the GEMM is memory-bounded (the ratio is poor and the GPU waits on memory loads). The optimal depends on N, K, and hardware parameters like the number of SMs and shared memory size — yet existing libraries use a fixed tiling strategy or only tune at library compile time.
Figure 7 shows the normalized performance landscape for flat GEMM with M=8 on an A100: for small N, larger hurts because parallelism is insufficient; for large N, smaller hurts because the computation-to-memory ratio is too low. The key insight is that different [N, K] shapes are bottlenecked by different factors, so a one-size-fits-all tiling strategy is inherently suboptimal.
Bottleneck 3: Static Dataflow (Section 5)
The GEMM operations in a transformer layer have highly heterogeneous shapes. Figure 9(a) catalogs the four unique [N, K] configurations for Llama2-7B:
| Operation | N | K |
|---|---|---|
| K, Q, V projection | ||
| O projection | ||
| FFN1 | ||
| FFN2 |
where is hidden dimension and is the expanded feedforward dimension. These four shapes span N from 4096 to 12288 and K from 4096 to 11008 — a factor of ~3× range in both dimensions. And these are just for one model; different models have different hidden dimensions and architectures.
Now consider the M-dimension, which varies dynamically: during prefill, (from 128 to 32K+ tokens); during decoding, (from 1 to 8). This means the exact same linear layer can be a GEMV operation (M=1, memory-bound, no data reuse), a flat GEMM (M=2–8, slightly better reuse but still memory-bound), or a conventional GEMM (M≥32, compute-bound with good data reuse) depending on runtime conditions.
The paper quantifies the cost of using the wrong implementation: for a Llama2-7B linear layer in the decode phase with batch size 1, the Tensor Core implementation from cuBLAS achieves only 82.15% of the performance of a CUDA Core GEMV kernel (FastGEMV) — because Tensor Cores are designed for high arithmetic intensity and waste cycles on memory-bandwidth-limited GEMV. Conversely, at batch size 4, using CUDA Cores achieves only 49.75% of Tensor Core performance — because CUDA Cores lack the throughput for the now-compute-bound workload (Section 5). A static dataflow that always uses cuBLAS or always uses a custom GEMV kernel leaves roughly 20–50% of hardware performance untapped, depending on the shape.
The challenge in building a heuristic dataflow is that the search space is large: input dynamics (M varies per request), model diversity (N, K vary per layer and model), GPU capacities (memory bandwidth, cache size, SM count, Tensor Core throughput differ across A100, RTX3090, MI210, etc.), and engineering effort (kernel quality varies across implementations). The paper's insight is that this space collapses because only four [N, K] shapes exist for any given model, and M is the only runtime-varying dimension — reducing the problem to finding, per [N, K] pair, the M-values where one implementation overtakes another.
Prior Approaches and Their Shortcomings
The paper does not exist in a vacuum; it inherits from and reacts to a substantial body of LLM inference optimization work. Section 7 catalogs the major systems, and understanding where each falls short clarifies why FlashDecoding++'s combination of techniques is necessary.
DeepSpeed-Inference (Aminabadi et al., 2022) focuses on system-level parallelism and memory management — kernel fusion, efficient KVcache allocation, and model parallelism for large-scale deployment. Its GEMM operations rely on cuBLAS, inheriting the flat GEMM under-utilization and static dataflow problems. Its attention operations during decoding do not address the synchronized softmax overhead specifically. DeepSpeed is strong at distributing work across GPUs but leaves per-GPU kernel efficiency on the table.
vLLM (Kwon et al., 2023) introduces PagedAttention, a memory management technique that allocates KVcache in non-contiguous pages to reduce fragmentation and increase maximum batch sizes. This is a system-level innovation that improves throughput by enabling larger batches, but it does not change the underlying GEMM or attention kernel implementations. vLLM benefits from batching more requests together, but each kernel call within that larger batch still suffers from the three bottlenecks the paper identifies.
FlashAttention / FlashAttention-2 (Dao et al., 2022; Dao, 2023) revolutionized prefill-phase attention by restructuring the computation to minimize HBM reads/writes through tiling and recomputation of the attention matrix in SRAM. FlashAttention uses the partial softmax with synchronized updates that the paper profiles as costing 18.8% overhead (FlashAttention-2 improves parallelism and work partitioning but retains the same numerical approach). Critically, FlashAttention is designed for the prefill phase where all queries are available simultaneously; it does not directly address the decode-phase scenario where queries arrive one at a time and the KVcache already exists.
FlashDecoding (Dao et al., 2023) extends FlashAttention to the decode phase by splitting the K and V dimensions across parallel workers. This is an important innovation — without it, the decode-phase attention has very low GPU utilization because each query token only does a small amount of computation against the full KVcache. However, FlashDecoding uses the same partial softmax with synchronized updates, inheriting the ~18.8% overhead. The paper explicitly positions FlashDecoding++ as building on FlashDecoding's parallelism strategy but replacing its softmax numeric with the asynchronized variant (hence the "++").
FasterTransformer (NVIDIA) and OpenPPL (SenseTime) are C++ LLM inference engines that reduce Python overhead through kernel fusion and efficient scheduling. Both use highly optimized implementations from cuBLAS/CUTLASS for GEMMs and can incorporate FlashAttention for attention. They represent the performance ceiling for production C++ engines, but they inherit the flat GEMM under-utilization and static dataflow issues from their underlying libraries. FasterTransformer in particular is the foundation that TensorRT-LLM builds upon.
TensorRT-LLM (NVIDIA, 2023) incorporates FlashAttention and FlashDecoding into a TensorRT-based engine with graph optimization and kernel fusion. It represents the state-of-the-art on NVIDIA GPUs and is the primary commercial-grade baseline. However, TensorRT-LLM still uses synchronized partial softmax, pads flat GEMMs to large tiles, and uses static dataflow dispatch. The paper's comparison against TensorRT-LLM is therefore a comparison against the best available production system, and the observed speedups (1.13× average on decode) are gains on top of TensorRT-level optimization.
FastGEMV (Wang, 2023) is a specialized CUDA Core kernel for matrix-vector multiplication, designed for the batch size = 1 decode phase. It avoids the padding overhead by directly implementing GEMV without Tensor Core tiling, achieving higher throughput than cuBLAS for this specific shape. FlashDecoding++ incorporates FastGEMV-like kernels as one of its three implementation options (ImplA) but goes further by (a) providing a smooth transition from GEMV to flat GEMM to full GEMM via inflection points, and (b) optimizing the flat GEMM regime (M=2–64) that FastGEMV does not target.
The gap across all prior work is that each system addresses part of the inference pipeline or one of the three bottlenecks, but none addresses all three simultaneously:
- FlashAttention/FlashDecoding: improve attention parallelism but retain synchronized softmax.
- cuBLAS/CUTLASS: optimize full GEMMs but under-perform on flat GEMMs/GEMV.
- FastGEMV: fixes GEMV but doesn't address flat GEMM or the transition to full GEMM.
- DeepSpeed/vLLM/TensorRT-LLM: optimize at the system level but inherit kernel-level inefficiencies from their underlying libraries.
How FlashDecoding++ Positions Itself
FlashDecoding++ positions itself as a kernel-level optimization engine that complements system-level frameworks. The paper is explicit (Section 1) that the techniques are designed to be integrated into existing engines like FlashDecoding or TensorRT-LLM — they are not proposing a replacement for the entire inference stack. The "++" in the name signals this: build on the state-of-the-art (FlashDecoding's parallel attention decomposition) and add three orthogonal improvements.
The paper's positioning has several strategic elements:
1. Kernel-level focus with system-level impact. Rather than proposing a new system architecture (cf. vLLM's PagedAttention) or a new parallelism strategy (cf. DeepSpeed's tensor parallelism), FlashDecoding++ targets the lowest-level numeric and scheduling decisions in individual GPU kernels. This is a high-risk, high-reward strategy: the gains per kernel are modest (18.8% reduction in attention overhead, ~2× improvement in flat GEMM utilization), but they multiply across hundreds of kernel calls per token and billions of tokens per deployment.
2. Hardware generality through abstraction rather than specialization. The paper evaluates on both NVIDIA (A100, RTX3090) and AMD (MI210, RX7900XTX) GPUs, which have different architectures, instruction sets, and software ecosystems (CUDA vs. ROCm). The three techniques are designed to be hardware-agnostic in principle: the asynchronized softmax is a numerical reformulation that works on any GPU; the flat GEMM optimization adjusts tiling based on hardware parallelism (number of SMs, shared memory size) but does not require hardware-specific intrinsics; the heuristic dataflow profiles offline to find inflection points automically for each GPU. This positions FlashDecoding++ as a portable optimization layer rather than an NVIDIA-specific solution.
3. Practicality over theoretical novelty. None of the three techniques is mathematically deep in isolation — the asynchronized softmax follows directly from the observation that the max-value scaling factor in softmax can be any constant (Equation 3), double buffering is a standard technique in HPC, and heuristic dispatch based on profiling is standard practice in compilers (e.g., TVM, XLA). The paper's contribution is the integration and engineering of these ideas into an LLM inference context where they were previously unapplied, combined with careful profiling to demonstrate their impact. This positions the work as practical systems research rather than algorithmic theory.
4. Complementarity, not competition, with existing systems. The paper benchmarks against vLLM, DeepSpeed, TensorRT-LLM, and OpenPPL not to replace them but to show that FlashDecoding++'s kernel optimizations can be integrated into any of them. The average speedups of 1.24× over vLLM and 1.13× over TensorRT-LLM on decode (Section 6.2) represent the marginal gain achievable by swapping in FlashDecoding++'s kernels while retaining the system-level optimizations of those frameworks. This positions the work as an upstream contribution to the LLM inference ecosystem rather than a competing framework.
5. Targeting the decode phase as the long-tail bottleneck. While the prefill phase dominates latency for long-prompt, short-generation scenarios (e.g., document summarization), the decode phase dominates for interactive applications where the model generates long responses (e.g., chatbots, code generation). As LLMs are increasingly deployed for conversational AI with multi-turn interactions, the cumulative decode-phase latency across thousands of generated tokens becomes the primary user-perceived bottleneck. By addressing flat GEMM — a decode-phase-specific problem — FlashDecoding++ targets the phase that prior attention-focused work (FlashAttention) did not solve.
The Underlying Economics
The paper motivates the importance of these optimizations through cost, not just speed. The 1.9M/day in saved compute. Even the more modest 1.13× speedup over TensorRT-LLM represents hundreds of thousands of dollars per day at scale. This economic framing is deliberate: kernel optimization is often viewed as a marginal activity with diminishing returns, but at the scale of LLM inference, marginal percentage improvements translate to significant absolute savings.
The paper also implicitly argues that hardware alone is insufficient — the A100 and MI210 are among the most powerful GPUs available, yet the bottlenecks the paper identifies persist on them because they are software artifacts (algorithm design choices, library defaults) rather than hardware limitations. This positions FlashDecoding++ as extracting value from already-deployed hardware, which is compelling for organizations with existing GPU fleets who cannot upgrade to newer hardware generations easily.
3. Technical Approach
3.1 Reader Orientation
FlashDecoding++ is a GPU kernel optimization engine that replaces the numerical and scheduling logic inside the transformer layer operations of existing LLM inference frameworks — it is not a standalone serving system but a drop-in set of faster GPU programs (kernels) for the attention softmax, the flat matrix multiplies of the decode phase, and the dispatch logic that decides which kernel to run for each operation. The system solves the problem that three specific, independently measurable inefficiencies in standard kernel implementations together waste a large fraction of GPU time during LLM inference: (1) the synchronized partial softmax in FlashAttention/FlashDecoding forces pipeline stalls or recomputation that costs ~18.8% of attention time, (2) padding the M-dimension of decode-phase GEMMs to tile size 64 wastes >50% of compute on zero-valued rows, and (3) using the same GEMM kernel for all shapes (e.g., always cuBLAS or always a custom GEMV kernel) leaves 20–50% of hardware throughput unused because the optimal implementation switches between CUDA-core GEMV, Tensor Core flat GEMM, and Tensor Core full GEMM depending on the M-dimension. The shape of the solution is three independent kernel-level techniques — a numerical softmax reformulation that eliminates synchronization, a tiling and double-buffering scheme for flat GEMM, and an offline-profiling-based dispatch table — that can be composed with any existing engine because they modify only the internal implementation of the linear algebra and attention primitives, not the engine's architecture.
3.2 Big-Picture Architecture (Diagram in Words)
FlashDecoding++ sits underneath the LLM inference engine's Python or C++ orchestration layer and replaces the GPU kernel calls that the engine makes for two categories of operations:
-
Attention softmax — the engine calls an attention kernel (e.g., FlashAttention for prefill, FlashDecoding for decode) that internally computes partial softmax results. FlashDecoding++ replaces the softmax numerics inside this kernel with the asynchronized softmax with unified max value (Section 3), which uses a predetermined global scaling constant
$\phi$instead of each tile computing and synchronizing its own local maximum. The kernel now processes each tile independently without cross-tile dependencies, with a lightweight overflow check that triggers the original synchronized path only on the rare (<0.01%) vectors where$\phi$is insufficient. -
GEMM operations — the engine calls a matrix-multiply primitive for each linear projection (K, Q, V, O, FFN1, FFN2) at each layer. FlashDecoding++ intercepts these calls and routes them to one of three implementations based on a dispatch table indexed by the operation's
$[N, K]$shape (fixed per model) and the runtime M-dimension (Section 5):- ImplA: CUDA-core GEMV (e.g., FastGEMV) when
$M < M_1$— for batch size 1 decode where Tensor Core overhead exceeds benefit. - ImplB: custom flat GEMM with double buffering (Section 4) when
$M_1 \leq M < M_2$— for small-batch decode or short-sequence prefill where the M-dimension is 2–64, using padding-to-8 and tiling heuristics instead of padding-to-64. - ImplC: Tensor Core GEMM via CUTLASS when
$M \geq M_2$— for large-batch decode or long-sequence prefill where the high arithmetic intensity justifies full Tensor Core tiling.
- ImplA: CUDA-core GEMV (e.g., FastGEMV) when
The dispatch thresholds $M_1$ and $M_2$ are computed offline once per $[N, K]$ shape per hardware platform by sweeping M and profiling ImplA, ImplB, and ImplC to find the crossover points where each implementation overtakes the previous one (Figure 9b). At runtime, the engine looks up the current operation's $[N, K]$ shape and current M, reads the two thresholds from a precomputed table, and dispatches to the appropriate kernel — no runtime profiling or dynamic decision-making is needed.
The information flow through the system during inference is: (1) the engine determines the current phase (prefill or decode) and batch size; (2) for each transformer layer, the engine calls linear projection kernels with the input tensor and weight matrix, passing the shapes $[M, N, K]$ to FlashDecoding++; (3) FlashDecoding++ dispatches to ImplA, ImplB, or ImplC based on the offline-computed table; (4) the attention kernel is called with the Q, K, V tensors, and internally uses either the asynchronized softmax (Section 3) or — if overflow is detected mid-computation — falls back to the synchronized partial softmax; (5) the results (output activations for the next layer, or logits for the output layer) are returned to the engine.
3.3 Roadmap for the Deep Dive
- First, the asynchronized softmax with unified max value (Section 3 of the paper): the numerical reformulation that eliminates synchronization in partial softmax, including the mathematical equivalence proof, the overflow/precision guard via the recomputation fallback, the statistical justification from profiling LLM activation ranges, and the concrete inner-product accumulation that makes async execution possible.
- Second, the flat GEMM optimization with double buffering (Section 4): the tiling strategy for decode-phase matrix multiplies, the padding-to-8 decision, the N-dimension tiling tradeoff formalized in Equation 5, the identification of parallelism-bounded vs. memory-bounded regimes, and the double-buffering scheme that overlaps shared-memory loads with computation.
- Third, the heuristic dataflow with hardware resource adaptation (Section 5): the offline profiling procedure that finds inflection points
$M_1$and$M_2$for each of the four$[N, K]$shapes per model, the rationale for why only four shapes exist, and the runtime dispatch table that selects among the three implementations without dynamic overhead.
This order follows the flow of a token through the transformer: attention (softmax) first, then the series of linear projections (GEMMs, dispatched adaptively). The heuristic dataflow comes last because it ties together the flat GEMM optimization from Section 4 with the CUDA-core GEMV and full CUTLASS GEMM implementations into a unified dispatch framework, and it depends on understanding when each implementation is appropriate.
3.4 Detailed, Sentence-Based Technical Breakdown
This is a kernel optimization paper whose core idea is that three distinct inefficiencies in standard GPU implementations of LLM inference operations — synchronized partial softmax numerics, wasteful M-dimension padding in flat GEMMs, and static GEMM kernel dispatch — can each be eliminated by a targeted technique that is mathematically sound, hardware-aware, and integrable into existing inference engines without architectural changes.
Asynchronized Softmax with Unified Max Value
Problem Statement and Profiling Evidence
The standard numerically stable softmax for a vector $\mathbf{x} = [x_1, \ldots, x_d]$ computes:
where $m(\mathbf{x}) = \max(x_1, \ldots, x_d)$ is the maximum element of the vector, subtracted from each $x_i$ before exponentiation to prevent $e^{x_i}$ from overflowing the floating-point representation (e.g., $e^{89} \approx 4.4 \times 10^{38}$, near the float32 upper bound; any $x_i$ above ~89 would produce NaN or infinity without the shift).
When attention computation is tiled for parallelism and memory efficiency (as in FlashAttention and FlashDecoding), the input vector $\mathbf{x}$ (one row of the $Q \times K^T$ matrix) is split into $p$ partial vectors $\mathbf{x}^{(1)}, \ldots, \mathbf{x}^{(p)}$. Each partial vector $\mathbf{x}^{(j)}$ is processed independently to compute its local maximum $m(\mathbf{x}^{(j)})$, local exponentiated values $f(\mathbf{x}^{(j)}) = e^{\mathbf{x}^{(j)} - m(\mathbf{x}^{(j)})}$, and local sum $\ell(\mathbf{x}^{(j)}) = \sum_i f(x_i^{(j)})$. However, these local statistics are computed relative to each tile's own maximum, not the global maximum $m(\mathbf{x})$. To combine them into the correct global softmax, the synchronized update equations (Equation 2 in the paper) are required:
What these equations compute: when a new partial softmax result arrives, the previously computed partial results are rescaled by the factor $e^{m(\mathbf{x}^{\text{old}}) - m(\mathbf{x}^{\text{new}})}$ where $m(\mathbf{x}^{\text{new}})$ is the updated global maximum incorporating the new tile. The local sums are similarly updated, and only after all tiles have been processed can the final division normalize all elements.
Why this is necessary: the softmax denominator $\sum e^{x_i - m(\mathbf{x})}$ depends on the global maximum — if you normalize each tile by its own local sum without rescaling by the global maximum, the resulting probabilities across tiles will not sum to 1. The synchronization is therefore mathematically required when each tile uses a different scaling factor (its own local maximum).
The cost: this synchronization forces either (a) pipeline stalls where later tiles wait for earlier tiles to be updated before proceeding, or (b) recomputation of earlier tiles' contributions when a new maximum is discovered. The paper's profiling on Llama2-7B with 1024 input length on an NVIDIA A100 shows this accounts for 18.8% of total attention computation time — nearly one-fifth of the attention runtime is spent on synchronization, not on useful arithmetic.
Mathematical Reformulation: Unified Max Value
The paper's key insight rests on a simple algebraic identity (Equation 3):
What this states: the scaling constant inside the softmax does not need to be the maximum element of the vector. Any real number $\phi$ produces the mathematically identical softmax output, because multiplying both the numerator and denominator by $e^{m(\mathbf{x}) - \phi}$ cancels out.
Why this form is useful: if every partial vector $\mathbf{x}^{(j)}$ uses the same $\phi$ instead of its own local maximum, then:
- Each tile computes
$f(\mathbf{x}^{(j)}) = e^{\mathbf{x}^{(j)} - \phi}$and$\ell(\mathbf{x}^{(j)}) = \sum_i f(x_i^{(j)})$independently. - The global denominator is simply
$\sum_{j=1}^p \ell(\mathbf{x}^{(j)})$— a sum of independently computed scalars, with no rescaling needed. - The global softmax for element
$i$in tile$j$is$f(x_i^{(j)}) / \sum_{k=1}^p \ell(\mathbf{x}^{(k)})$, which requires only one global sum reduction after all tiles finish. - No cross-tile synchronization is needed during tile computation — tiles can execute fully asynchronously, and only a single reduction at the end combines their partial sums.
The Overflow and Precision Problem: Choosing $\phi$
While any $\phi$ is mathematically valid, numerical reality imposes constraints:
- If
$x_i - \phi$is too large (e.g.,$\geq 89$for float32),$e^{x_i - \phi}$overflows to infinity, producing NaN in the final result. - If
$x_i - \phi$is too small (e.g.,$\ll 0$),$e^{x_i - \phi}$underflows to zero, causing precision loss — the element contributes nothing to the softmax output even when it should have a non-negligible probability.
Therefore, $\phi$ must be chosen such that for all elements $x_i$ in the input vector, $x_i - \phi$ stays within a safe range $[a, b]$ where $e^a$ does not underflow and $e^b$ does not overflow. This is the same role that $m(\mathbf{x})$ plays in the standard softmax — it guarantees $x_i - m(\mathbf{x}) \leq 0$ for all $i$, so the exponent is always $\leq 1$ and never overflows, and the largest exponent is exactly $e^0 = 1$.
The paper's critical empirical insight (Figure 5) is that for real LLM activations, the elements $x_i$ (entries of the $Q \times K^T$ matrix before softmax) are tightly bounded across different input sequences. The paper presents histograms of $x_i$ values for Llama2-7B, ChatGLM2-6B, and OPT-6.7B with diverse inputs:
- Llama2-7B: >99.99% of
$x_i$values lie in the range$[-16.8, 6.5]$. - ChatGLM2-6B: >99.99% of
$x_i$values lie in the range$[-10.5, 13.7]$. - OPT-6.7B: values span
$[-496.8, 363.5]$— a drastically wider range.
What this means operationally: for Llama2-7B, setting $\phi = a = -16.8$ ensures that for >99.99% of elements, $x_i - \phi \geq 0$ (no underflow) and $x_i - \phi \leq 6.5 - (-16.8) = 23.3$ — well within float32 range, since $e^{23.3} \approx 1.3 \times 10^{10}$ is representable. The exponents remain within a safe representable range without any per-vector adaptation. For OPT-6.7B, the range is too wide, and the technique is not applied (the paper explicitly notes this in Section 3).
Why this works: the attention logits in trained transformers are empirically well-behaved — they cluster within a relatively narrow range because the model has learned to produce attention scores that are neither uniformly tiny (which would lose information) nor explosively large (which would cause training instability). The bounded range of activations is a learned property of the trained model, not a mathematical guarantee, which is why the paper validates it empirically across models and inputs.
Asynchronized Execution with Inner Product
The softmax operation in attention is always followed by multiplication with the V matrix: $\text{softmax}(Q \times K^T) \times V$. This means we never need the softmax probabilities as a standalone matrix — we only need the result of the weighted sum of value vectors. The paper exploits this to further reduce synchronization (Equation 4):
where $\mathbf{x}$ is one row of $Q \times K^T$ (of length $d$, the sequence length), $\mathbf{v}$ is a column of $V$ (of the same length), $\langle \cdot, \cdot \rangle$ denotes the inner product, and the vector is split into $p$ partial vectors of length $d/p$ each.
What this equation computes: each tile $j$ computes two scalars independently — a numerator $\text{num}^{(j)} = \sum_i e^{x_i^{(j)} - \phi} \cdot v_i^{(j)}$ and a denominator $\text{den}^{(j)} = \sum_i e^{x_i^{(j)} - \phi}$. These two scalars are accumulated independently across tiles using atomic additions or a single global reduction after all tiles finish. The final attention output for that position is $\sum_j \text{num}^{(j)} / \sum_j \text{den}^{(j)}$.
Why this is sufficient: the inner product $\text{softmax}(\mathbf{x}) \cdot \mathbf{v}$ is a scalar (one element of the attention output matrix), and the numerator and denominator are both sums over the sequence dimension $d$. The ratio of sums is equal to the weighted sum with softmax probabilities because:
The global denominator normalizes all contributions simultaneously. The tile-level partial sums can be computed in any order, on any streaming multiprocessor, without any cross-tile communication beyond the final two-scalar reduction.
Why this differs from standard partial softmax: in the standard FlashAttention approach, each tile must update its partial softmax results when a new tile reports a higher maximum — the previous tiles' exponentiated values are rescaled, which requires either storing and recomputing them or synchronizing with the new maximum before finalizing. In the asynchronized approach, tiles never rescale because they all use the same $\phi$, so there is no concept of an "updated maximum" to react to.
The Overflow Fallback: Recomputation
Since >99.99% coverage still leaves a small probability of overflow (an element $x_i$ where $x_i - \phi > b$ and $e^{x_i - \phi}$ overflows), the paper implements a recomputation fallback rather than accepting the correctness risk.
The mechanism: each tile, before writing its partial numerator and denominator, checks whether any element $x_i^{(j)}$ in its tile satisfies $x_i^{(j)} - \phi \leq a$ (underflow risk) or $x_i^{(j)} - \phi \geq b$ (overflow risk). If any tile detects an out-of-range value, it signals a global flag (or, in the example in Figure 6, both threads terminate their asynchronous processing), and the entire softmax for that attention row is recomputed using the standard synchronized partial softmax with per-tile local maximum tracking.
What happens in practice (Figure 6):
- Two vectors
$\mathbf{x}$and$\mathbf{y}$are computed from$Q \times K^T$, each split into 2 partial vectors. With$\phi = 6$,$a = -3$,$b = 3$(illustrative values). - For
$\mathbf{x}$: elements are$[4, 5, 6, 7]$. Check:$4-6=-2 \geq -3$✓,$5-6=-1 \geq -3$✓,$6-6=0 \leq 3$✓,$7-6=1 \leq 3$✓. All in range. Two threads process tiles independently: Thread 1 computes$\text{num}^{(1)} = e^{4-6}v_1 + e^{5-6}v_2$and$\text{den}^{(1)} = e^{4-6} + e^{5-6}$. Thread 2 computes$\text{num}^{(2)} = e^{6-6}v_3 + e^{7-6}v_4$and$\text{den}^{(2)} = e^{6-6} + e^{7-6}$. The results are summed and divided. - For
$\mathbf{y}$: elements are$[3, 6, 9, 6]$. Thread 1 processes$[3, 6]$without issue ($3-6=-3 \geq -3$✓,$6-6=0 \leq 3$✓). Thread 2 encounters$y_3 - \phi = 9-6=3 \leq 3$— in this example,$b=3$, so$9-6=3$is exactly at the boundary. If it exceeded$b$(say$y_3=10$,$10-6=4>3$), overflow would be triggered. Both threads are terminated, and Thread 1 recomputes all partial vectors using the synchronized scheme: it computes the local max of$\mathbf{y}$(which is 9), rescales all elements by$e^{x_i-9}$, tracks the global denominator, and produces the correct result.
Overhead analysis: since >99.99% of attention rows never trigger the fallback, the recomputation cost is negligible in expectation. The paper characterizes this through the statistical histograms (Figure 5) and the explicit statement that "such a recomputation scheme avoids overflow while introducing negligible overheads based on the statistical data."
Why not always use the synchronized path for safety? Because the 18.8% overhead is incurred on every single attention computation, across every layer, every token, every request. Recomputing 0.01% of rows with the synchronized method costs 18.8% extra on 0.01% of the workload — essentially zero — while saving 18.8% on the other 99.99%. The asymmetric cost structure (large savings on the common case, small cost on the rare case) is what makes the fallback strategy effective.
Design Choices and Justifications
Choice of $\phi$: the paper uses $\phi = a$, the lower bound of the safe range, rather than an intermediate value. This is because underflow $(x_i - \phi \ll 0)$ produces zero, which is numerically harmless (it just means that element gets zero attention weight), while overflow $(x_i - \phi \gg 0)$ produces infinity, which corrupts the entire softmax row. By setting $\phi$ to the lower bound, the worst case is underflow (safe), never overflow, for in-range elements.
Model-specific applicability: OPT-6.7B has a much wider activation range and the technique is explicitly not applied to it. This is a limitation — the technique depends on the empirical property that trained attention logits cluster in a bounded interval, which may not hold for all models, all quantization schemes, or all input distributions. The paper does not explore how to extend the technique to models with wide activation ranges (e.g., by dynamically estimating $\phi$ from a cheap pre-scan or using a model-specific calibration dataset).
Integration with FlashDecoding's parallelism: FlashDecoding splits the K and V dimensions across parallel workers for the decode phase. The asynchronized softmax operates within each worker's partial computation independently — each worker computes its own numerator and denominator using the global $\phi$, and a cross-worker reduction combines these scalars. This is simpler than the original FlashDecoding which requires cross-worker synchronization of partial softmax statistics, since now each worker's partial results are directly summable without rescaling.
Hardware implications: the asynchronized softmax reduces cross-thread-block synchronization, which on GPUs is expensive because it requires either global memory atomics or kernel relaunches. By eliminating the need for threads to wait on each other's maxima, the technique increases the occupancy of streaming multiprocessors — more thread blocks can execute concurrently because they are not stalled waiting for a dependency.
Flat GEMM Optimization with Double Buffering
Problem Statement: The Padding Waste
During the decode phase, the linear projections (K, Q, V, O, FFN1, FFN2) take the form of a matrix multiply $C = A \times B$ where $A$ has shape $[M, K]$ and $B$ has shape $[K, N]$. Here $M$ is the batch size (1–8 tokens processed simultaneously), $K$ is the hidden dimension (e.g., 4096 for Llama2-7B), and $N$ is the output dimension (e.g., 12288 for the concatenated K, Q, V projection, 11008 for the expanded FFN dimension).
Modern GPU BLAS libraries (cuBLAS, CUTLASS) tile the M-dimension to 64 by default. This means the library's kernel loads 64 rows of $A$ and computes 64 rows of the output $C$ simultaneously within one thread block. The tiling serves two purposes:
- Amortized memory access: each element of
$B$loaded from global memory is reused across 64 rows of$A$, yielding a computation-to-memory ratio proportional to$2 \times 64 \times K / (64 \times K + \text{other terms}) \approx 2 \times 64 / (64) = 2$(a rough approximation — the exact ratio is given in Equation 5). - Parallelism: 64 rows provide enough independent MAC (multiply-accumulate) operations to fill the pipeline of the streaming multiprocessor's functional units and hide instruction latency.
When $M$ is small (batch size 1–8), the input matrix $A$ has only 1–8 rows of actual data. To use the 64-tile kernel, the library pads $A$ to 64 rows by appending rows of zeros, and similarly pads the output $C$. The computation on these padding rows — all the multiplications and additions — produces zeros that are discarded. The fraction of useful computation is $M/64$: for batch size 1, only 1.56% of the arithmetic is productive; for batch size 8, only 12.5% is productive. The paper's claim of ">50% computation under-utilization" (Section 1) is a conservative lower bound — at batch size 1 it is >98% wasted.
Why cuBLAS does this: tiling to 64 is optimal for large GEMMs (e.g., prefill phase with $M \geq 128$) where the wasted padding is a small fraction. The library's heuristics are tuned for the prefill regime or for training, where batch sizes and sequence lengths are large. During decoding, these heuristics break down because the M-dimension collapses.
Approach: Pad to 8 Instead of 64
FlashDecoding++ pads the M-dimension to 8 rather than 64. The motivation is that modern NVIDIA Tensor Cores (Volta and later) process GEMM with an $M=8$ granularity — each Tensor Core instruction operates on an $8 \times 8$ tile of the output matrix. Padding to 8 means no rows are wasted beyond the minimum required by the hardware's instruction-level interface.
However, reducing the M-tile size from 64 to 8 creates a new problem: the computation-to-memory ratio drops by roughly a factor of 8 (from proportional to 64 to proportional to 8), because each loaded element of $B$ is reused across only 8 rows instead of 64. This pushes flat GEMMs further into the memory-bound regime, where the GPU's compute units stall waiting for data from HBM. The key challenge is therefore: how to tile the remaining dimensions (N and K) to compensate for the reduced M-tile size and keep the GPU busy?
The N-Dimension Tiling Tradeoff
The paper formalizes the performance of flat GEMM as a function of the N-dimension tile size $B_N$ and the K-dimension tile size $B_K$. The computation and memory access for one GEMM tile are:
- Computation:
$2 \times M \times B_N \times B_K$MAC operations (2 operations per multiply-add,$M \times B_N$output elements, each requiring$B_K$inner products). - Memory access (loads):
$M \times B_K$elements of$A$+$B_N \times B_K$elements of$B$per tile. - Number of tiles:
$B = (N \times K) / (B_N \times B_K)$.
The total computation-to-memory ratio (Equation 5) is:
What this equation computes: given fixed $M$ (the padded batch size, 8) and $K$ (the hidden dimension, e.g., 4096 or 12288), the computation-to-memory ratio depends on the N-dimension tile size $B_N$. A larger $B_N$ makes the denominator smaller because the term $M \times K / B_N$ decreases — fewer tile loads are needed overall, so less time is spent fetching $B$ tiles from memory relative to computation time. A larger $B_N$ therefore improves arithmetic intensity and moves the kernel toward compute-bound performance.
Why $B_N$ cannot be arbitrarily large: the parallelism available to the GPU is $N / B_N$ — the number of tiles that can be distributed across streaming multiprocessors. For the NVIDIA A100 with 108 SMs, if $N / B_N < 108$, some SMs will be idle because there are not enough independent tiles to keep them all occupied. A larger $B_N$ reduces parallelism, and when $N$ itself is small (e.g., $N=4096$ for the O projection), the parallelism constraint binds quickly.
Why $B_N$ cannot be arbitrarily small: a smaller $B_N$ increases parallelism but reduces the computation-to-memory ratio, because each tile processes fewer output columns and thus amortizes the cost of loading $A$ and $B$ over fewer output elements. For memory-bound kernels, this makes the bottleneck worse.
The two regimes (Figure 7):
- Parallelism-bounded (small N):
$N / B_N$is the binding constraint. The kernel needs enough tiles to occupy all SMs, so$B_N$must be chosen such that$N / B_N \geq \text{num\_SMs}$. The optimal$B_N$is the largest value that still provides sufficient parallelism — roughly$N / 108$for the A100. Performance is flat or slightly decreasing as$B_N$increases beyond this point because parallelism drops. - Memory-bounded (large N): parallelism is no longer an issue (many tiles even with large
$B_N$), but the computation-to-memory ratio is the bottleneck. The optimal$B_N$is large to maximize arithmetic intensity, and performance improves with$B_N$until shared memory or register constraints limit tile size.
Figure 7 illustrates this concretely: for $K=4096$ and $M=8$, when $N=32$, the best $B_N$ is small (all tiles are needed for parallelism). When $N=262144$, the best $B_N$ is large (memory bandwidth is the only bottleneck). The paper uses these empirical performance curves to select $B_N$ per $[N, K]$ shape, rather than using a single heuristic across all shapes.
Double Buffering for Memory-Bound Regimes
When the flat GEMM is memory-bounded (large N), the paper introduces double buffering to overlap memory loads with computation — a standard technique in GPU programming but not applied in the specific context of flat GEMM with M-tile=8.
The mechanism (Figure 8):
- Each GPU thread block allocates two buffers in shared memory (on-chip SRAM, much faster than HBM) for the input tiles.
- The tiling strategy processes the K-dimension sequentially within a single thread block: tiles
$A_1, A_2, A_3, \ldots$of matrix$A$and corresponding tiles$B_1, B_2, B_3, \ldots$of matrix$B$are loaded, multiplied, and accumulated to produce the partial output. - With double buffering: while the thread block computes
$A_1 \times B_1$using data in buffer 1, it simultaneously loads$A_2$and$B_2$from HBM into buffer 2. The load of the next tile is hidden behind the computation of the current tile. - After
$A_1 \times B_1$completes, the thread block swaps buffers: it computes$A_2 \times B_2$from buffer 2 while loading$A_3$and$B_3$into buffer 1. - This ping-pong pattern continues until all K-dimension tiles are processed.
Why this helps: in a memory-bound kernel, the GPU's streaming multiprocessor spends most of its time waiting for data to arrive from HBM (latency ~300–800 cycles on modern GPUs). By initiating the load of the next tile while the current tile is being computed, the load latency is partially or fully hidden, depending on whether computation time exceeds load time. The technique effectively overlaps the memory and compute phases, increasing throughput without increasing the computation-to-memory ratio.
When it is applied: the paper states that double buffering is applied "when N is large in our practice" (Section 4), corresponding to the memory-bounded regime in Figure 7. For small N, where the kernel is parallelism-bounded, double buffering provides less benefit because the SM occupancy is already high (many thread blocks are active) and the memory subsystem is less of a bottleneck.
Example trace (Figure 8): GPU Block1 processes tiles of $A$ (the input activation matrix, shape $[M, K]$) and the corresponding columns of $B$ (the weight matrix, shape $[K, N]$) to compute one tile of $C$ (the output, shape $[M, N]$). The timeline shows: "idle" for a brief period during the first load, then "A1B1" computation in one buffer while "A2B2" loads into the other, then "A2B2" computation while "A3B3" loads, and so on. The idle period is only at the very start (pipeline fill) and end (pipeline drain). Multiple GPU blocks (Block1, Block2, ...) process different N-dimension tiles in parallel.
The M-Dimension Padding Choice: Why 8 and Not Another Value
The paper pads to 8 specifically because this matches the Tensor Core $M=8$ granularity. The key design decisions and alternatives are:
- Pad to 1 (no padding at all): this would be a pure GEMV kernel, which does not benefit from Tensor Cores at all because Tensor Cores require
$8 \times 8$matrix tiles. The throughput of CUDA Cores (used in FastGEMV, ImplA) exceeds Tensor Cores for GEMV because Tensor Cores have higher latency per instruction and the operand reuse is insufficient to amortize it. Padding is necessary to use Tensor Cores at all. - Pad to 2, 4: these would require custom tiling logic that does not align with the hardware's native
$8 \times 8$instruction, leading to inefficient partial tile handling or additional padding within each warp. - Pad to 16, 32: these increase the M-tile size, improving the computation-to-memory ratio (proportional to M), but waste more computation on padding rows. For batch size 1, padding to 16 wastes 93.75% vs. 87.5% for padding to 8 — both are high, but the improved ratio from larger tiles might compensate. The paper's choice of 8 represents the minimum viable tile size for Tensor Core usage, minimizing padding waste while accepting the lowest computation-to-memory ratio that Tensor Cores can provide.
- Pad to 64 (the standard approach): this is what cuBLAS does, and it wastes >50% computation.
The choice of 8 is thus the Pareto-optimal point for flat GEMM: it is the smallest M-dimension that can leverage Tensor Cores, minimizing the waste from padding while still accessing the 16× throughput advantage of Tensor Cores over CUDA Cores for matrix-matrix operations (the A100's Tensor Cores deliver 312 TFLOPS of FP16 vs. 19.5 TFLOPS for CUDA Cores, a 16× ratio for structured matrix operations).
Interaction with Batch Size
The paper's flat GEMM optimization is applied during the decode phase for batch sizes 1–8 (and potentially larger small batches during prefill with short sequences). As batch size increases from 1 to 8, the M-dimension after padding becomes 8 (no padding needed for batch size 8 — the actual data fills the 8-tile exactly), and the computation-to-memory ratio improves by a factor of 8 relative to batch size 1 with padding to 8. This means that larger decode batches naturally benefit more from Tensor Core usage, and the paper's optimization is most impactful at batch size 1 (the most memory-bound, most under-utilized regime) and gradually transitions toward full GEMM efficiency as batch size grows.
At batch size larger than roughly 8–16, the M-dimension becomes large enough that the standard cuBLAS/CUTLASS kernels with M-tile=64 are appropriate — this is the $M_2$ inflection point in the heuristic dataflow (Section 5), where ImplC (CUTLASS) overtakes ImplB (the custom flat GEMM).
Heuristic Dataflow with Hardware Resource Adaptation
Problem Statement: One Kernel Does Not Fit All Shapes
The GEMM operations in a single transformer layer exhibit a wide range of shapes, and the optimal GPU kernel implementation depends on where the operation sits in the space defined by the three dimensions $[M, N, K]$:
- GEMV (M=1): memory-bound, no data reuse for
$A$(the activation vector is loaded once and never reused), minimal parallelism. An optimized CUDA Core implementation (ImplA, FastGEMV) that streams through the weight matrix with carefully unrolled loops and prefetching outperforms Tensor Core implementations because Tensor Cores require$8 \times 8$tiles and the overhead of setting up those tiles exceeds the benefit when M=1. - Flat GEMM (2 ≤ M ≤ ~64): moderately memory-bound, some data reuse for
$A$but not enough to amortize large tiles. The custom flat GEMM with M-padding-to-8 and double buffering (ImplB) is designed for this regime — it uses Tensor Cores at their minimum granularity while adapting tiling to the specific$[N, K]$shape. - Full GEMM (M ≥ ~64): compute-bound, sufficient data reuse to justify large tile sizes (M-tile=64 or larger) and full Tensor Core throughput. Standard CUTLASS/cuBLAS implementations (ImplC) are already highly optimized for this regime, and the custom flat GEMM would be suboptimal because its smaller M-tile (8 vs. 64) unnecessarily reduces the computation-to-memory ratio.
The paper quantifies the cost of using the wrong kernel: for Llama2-7B's linear layer in the decode phase with batch size 1, cuBLAS (Tensor Core) achieves only 82.15% of the performance of FastGEMV (CUDA Core) — a 17.85% relative performance loss from using the "better" hardware (Tensor Cores) on the wrong shape. Conversely, at batch size 4, using CUDA Cores achieves only 49.75% of Tensor Core performance — a 50.25% relative loss. These numbers (Section 5) are the empirical basis for the claim that static dataflow causes up to 50.25% performance loss.
Why static dataflow is the default in existing systems: cuBLAS and CUTLASS use internal heuristics to select tiling strategies and kernel variants, but these heuristics are designed to be robust across a wide range of shapes, not optimal for the specific shapes of LLM inference. They typically use the same M-tile=64 kernel for any M ≤ 64, incurring the padding waste. More importantly, they do not incorporate the CUDA Core GEMV option at all — they are Tensor Core-only libraries. Switching between CUDA Cores and Tensor Cores requires calling entirely different libraries (e.g., FastGEMV vs. cuBLAS), which existing inference engines do not do dynamically at the operation level.
Key Insight: The Shape Space Is Small
The paper's critical observation (Figure 9a) is that for a given LLM, there are only four distinct $[N, K]$ shapes across all the GEMM operations in all transformer layers:
- K, Q, V projection:
$N = 3 \times \text{HD}$,$K = \text{HD}$ - O projection:
$N = \text{HD}$,$K = \text{HD}$ - FFN1 (expand):
$N = \text{FD}$,$K = \text{HD}$ - FFN2 (contract):
$N = \text{HD}$,$K = \text{FD}$
where HD is the hidden dimension (e.g., 4096 for Llama2-7B) and FD is the feedforward dimension (e.g., 11008 for Llama2-7B).
Why this matters: while M varies dynamically (with batch size, sequence length, and phase), N and K are static per model. The number of distinct $[N, K]$ pairs is small (4), and for each such pair, the problem of selecting the best kernel reduces to a 1D search over M — finding two thresholds $M_1$ and $M_2$ where the optimal implementation changes. This makes the offline profiling feasible: instead of profiling a 3D grid $(M, N, K)$, we profile 4 separate 1D sweeps (vary M for each fixed $[N, K]$).
The paper emphasizes that "the homogeneity of different layers in LLM significantly reduces the search space for operator optimization" (Section 5). This is a deliberate architectural observation — unlike a general-purpose GEMM library that must handle arbitrary shapes, an LLM inference engine operates on a highly regular set of shapes determined by the model architecture.
Offline Profiling: Finding the Inflection Points
The offline profiling procedure (Figure 9b) operates as follows for each of the four $[N, K]$ pairs:
- Set M = 1. Profile ImplA (CUDA Core GEMV, e.g., FastGEMV) and ImplB (custom flat GEMM with double buffering). Record the execution time or throughput of each.
- If ImplB > ImplA (i.e., ImplB has higher throughput): then
$M_1 = 1$— the flat GEMM is already better at batch size 1, and CUDA Core GEMV is never used. This can happen for shapes where N is large enough that the parallelism from tiling N compensates for the Tensor Core overhead even at M=1. - If ImplA > ImplB: increase M (e.g., M=2, 3, 4, ...) and re-profile both implementations at each step. Continue until ImplB overtakes ImplA. The first M where ImplB outperforms ImplA is
$M_1$. - Continue increasing M and now profile ImplB vs. ImplC (CUTLASS, the standard full GEMM implementation). At small M, ImplB should outperform ImplC because ImplC pads to M-tile=64 and wastes computation. As M grows, the waste fraction decreases, and CUTLASS's larger tiles and more sophisticated tiling heuristics eventually overtake the custom flat GEMM. The first M where ImplC outperforms ImplB is
$M_2$. - For M ≥
$M_2$: use ImplC. For$M_1 \leq M < M_2$: use ImplB. For M <$M_1$: use ImplA.
What "outperforms" means: the paper measures kernel execution time or throughput (tokens/second). The profiling is done once per hardware platform per model and stored as a configuration file (a lookup table). The profiling runs on the specific GPU being deployed, so it captures hardware-specific characteristics (memory bandwidth, SM count, cache sizes, Tensor Core generation).
Why this works: the performance curves of ImplA, ImplB, and ImplC as functions of M are monotonic in their relative ordering. ImplA is best at M=1 (highest memory efficiency for single-vector operations), ImplB is best at intermediate M (where Tensor Cores start to benefit but full tiling is still wasteful), and ImplC is best at large M (where full tiling pays off). There may be multiple crossover points if the curves are not perfectly ordered, but the paper's decision flow assumes exactly two inflection points per $[N, K]$ shape, implying the relative ordering is well-behaved in practice.
Runtime Dispatch: The Lookup Table
At inference runtime, for each GEMM operation, the engine:
- Identifies the operation type (K/Q/V projection, O projection, FFN1, or FFN2), which determines the
$[N, K]$shape. - Reads the current M-dimension from the input tensor (batch size for decode, sequence length × batch size for prefill).
- Looks up
$M_1$and$M_2$for this$[N, K]$pair from the precomputed table. - Dispatches to ImplA if
$M < M_1$, ImplB if$M_1 \leq M < M_2$, or ImplC if$M \geq M_2$.
Zero runtime overhead: the table lookup is a constant-time operation (a few integer comparisons), and the table itself is tiny (4 rows × 2 thresholds × a few bytes = ~64 bytes). There is no dynamic profiling, no online learning, and no runtime decision tree — the dispatch is purely a static mapping from $[M, N, K]$ to kernel function pointer, precomputed offline.
Example: Llama2-7B Dispatch Table (Figure 9c)
The paper provides a concrete example for Llama2-7B with HD=4096, FD=11008, giving the four $[N, K]$ shapes:
- K, Q, V projection:
$[12288, 4096]$— large N (the concatenated projection for three attention heads), moderate K. - O projection:
$[4096, 4096]$— square matrix, smaller N. - FFN1:
$[11008, 4096]$— large N (expansion), moderate K. - FFN2:
$[4096, 11008]$— moderate N, large K (contraction).
For each shape, the inflection points $M_1$ and $M_2$ are found via the offline profiling procedure. In the example shown:
- For K, Q, V projection at batch size 1 (decode phase): M=1, so the dispatch falls in the "ImplA" region (FastGEMV on CUDA Cores). This makes sense — N is large (12288), so loading the entire weight matrix for a single token is extremely memory-bound, and the CUDA Core streaming approach is faster than setting up Tensor Core tiles for a single row.
- For FFN1 at input sequence length 8 (prefill phase with small batch): M=8. Depending on where
$M_1$and$M_2$fall for the$[11008, 4096]$shape, this could fall in the ImplB region (custom flat GEMM) — the paper's example states "our flat GEMM optimization is applied when batch size=1/input sequence length=8 for FFN1 (M=8)."
Why the inflection points differ across $[N, K]$ shapes: the computation-to-memory ratio (Equation 5) depends on both N and K through the term $M \times K / B_N$. For FFN2 with large K (11008), the ratio is higher at a given M than for O projection with K=4096, because more computation (larger inner products) is performed per memory load. This means ImplB or ImplC becomes favorable at smaller M for large-K operations, and ImplA (CUDA Core GEMV) may only be optimal at M=1 or not at all. The per-shape profiling captures these differences automatically.
Hardware Adaptation
The heuristic dataflow is inherently hardware-adaptive because the inflection points are derived from profiling on the specific deployment GPU. Different GPUs have:
- Different numbers of SMs (A100: 108, RTX3090: 82, MI210: 104), affecting the parallelism constraint in flat GEMM tiling.
- Different memory bandwidths (A100: 2039 GB/s HBM2e, RTX3090: 936 GB/s GDDR6X, MI210: 1638 GB/s HBM2e), affecting the memory-bound/compute-bound transition point.
- Different Tensor Core throughputs and instruction latencies (A100: 312 TFLOPS FP16, MI210: 181 TFLOPS FP16), affecting the benefit of switching from CUDA Cores to Tensor Cores.
- Different shared memory sizes per SM, affecting the maximum tile sizes for double buffering.
The offline profiling automatically captures all these factors without requiring the developer to model them analytically. The decision flow in Figure 9b is executed independently for each GPU platform, producing a platform-specific dispatch table.
Portability evidence: the paper demonstrates the approach on both NVIDIA (A100, RTX3090) and AMD (MI210, RX7900XTX) GPUs (Section 6), with the heuristic dataflow adapting to each platform's characteristics through the profiling procedure. The AMD results (Figures 12, 13) show significant speedups over Hugging Face baselines, confirming that the dispatch strategy works across vendor architectures despite different instruction sets (CUDA vs. ROCm) and hardware organizations.
Interaction with the Prefill Phase
During the prefill phase, M = sequence length × batch size, which can range from small (e.g., 128 for short prompts with batch size 1) to very large (e.g., 32K for long document processing). The heuristic dataflow naturally handles this:
- For short sequences (M small), FFN1 and K/Q/V projections may use ImplB (flat GEMM with double buffering) — the same kernels as the decode phase but with larger M improving efficiency.
- For long sequences (M large), all operations use ImplC (CUTLASS) — standard full GEMMs where Tensor Cores operate at peak throughput.
- For intermediate lengths, different operations may use different implementations within the same forward pass: K/Q/V projection (large N) might still be in ImplB while O projection (smaller N) might already be in ImplC, because the inflection points
$M_2$differ per shape.
This per-operation dispatch is the key advantage over a static engine-wide setting like "always use cuBLAS" or "always use FastGEMV for decode." It allows fine-grained adaptation where each of the 4 GEMM types in each of the 32+ layers independently selects the optimal kernel for its current M-dimension.
Why Not Use an Analytical Model Instead of Profiling?
An analytical model that predicts kernel performance from $[M, N, K]$ and hardware parameters would be more portable and avoid the profiling cost. However, the paper does not pursue this approach, likely for several practical reasons:
- GPU kernel performance depends on low-level details (shared memory bank conflicts, register pressure, instruction scheduling) that are difficult to model analytically.
- Different kernel implementations (FastGEMV, the custom flat GEMM, CUTLASS) are developed by different teams with different optimization strategies, making a unified performance model challenging.
- The profiling cost is a one-time offline cost per hardware platform and model; with only 4 shapes and M ranging up to maybe 128 (after which ImplC always wins), the total profiling time is negligible compared to inference runtime over the model's deployment lifetime.
- The profiling approach automatically adapts to implementation improvements — if a better GEMV kernel or a more optimized CUTLASS heuristic is released, re-running the profiling updates the inflection points without any code changes.
Summary of the Three Techniques and Their Complementation
The three techniques — asynchronized softmax, flat GEMM optimization, and heuristic dataflow — are independent and composable:
- The asynchronized softmax modifies the attention computation (operations ②, ③, ④ in Figure 2) and is applicable to both prefill and decode phases, on any GPU, for any model where activations fall within a bounded range.
- The flat GEMM optimization modifies the decode-phase linear projections (operations ①, ⑤, ⑥ in Figure 2 with small M) by replacing the M-tile=64 padding with M-tile=8 and adding double buffering for memory-bound shapes.
- The heuristic dataflow modifies the dispatch of all linear projections (operations ①, ⑤, ⑥ in both phases) by selecting among ImplA (CUDA Core GEMV), ImplB (flat GEMM), and ImplC (CUTLASS full GEMM) based on the runtime M-dimension.
Together, they cover all the major computational operations in the transformer layer: attention (softmax + reduction) through technique 1, and all linear projections through techniques 2 and 3. The paper does not introduce new techniques for the $Q \times K^T$ matrix multiply itself (beyond what FlashAttention/FlashDecoding already do) or for the feedforward activation functions — these are left to existing optimized implementations. The three techniques specifically target the three bottlenecks identified through profiling, closing the largest remaining gaps in GPU utilization for LLM inference.
4. Key Insights and Innovations
Innovation 1: The Numerical Softmax Can Be Reformulated to Eliminate Synchronization Without Any Mathematical Approximation
The dominant assumption in attention optimization since FlashAttention (Dao et al., 2022) has been that the partial softmax requires per-tile maximum tracking and cross-tile synchronization to maintain numerical stability — that because the scaling constant $m(\mathbf{x})$ is defined as the global maximum, tiles must communicate their local maxima before any tile can finalize its contribution. This synchronization was treated as a necessary cost of the tiled attention algorithm, something to be minimized through engineering (e.g., careful warp scheduling, shared-memory atomics) but never eliminated.
FlashDecoding++ challenges this assumption at its root by observing a simple algebraic identity that the field had overlooked in the inference context: the scaling factor in softmax is a free parameter, not constrained to be the maximum element. The identity $\text{softmax}(x)_i = e^{x_i - \phi} / \sum_j e^{x_j - \phi}$ holds for any real $\phi$, because multiplying numerator and denominator by $e^{m(\mathbf{x}) - \phi}$ cancels exactly. This is not an approximation — it is an exact algebraic equivalence that the standard formulation obscures by coupling the scaling factor to the data-dependent maximum for safety rather than necessity.
Why this is intellectually distinctive: It reframes the softmax stability problem from "how do we efficiently compute the data-dependent maximum across tiles?" to "can we find a data-independent constant that keeps all exponents within representable range?" The first question leads to synchronization schemes (the approach of all prior attention implementations). The second question leads to a one-time statistical analysis of activation ranges — a fundamentally different kind of solution that moves the stability guarantee from runtime computation (finding $m(\mathbf{x})$ online) to offline characterization (bounding the distribution of $x_i$ empirically). This is a conceptual shift from algorithmic numerical stability to statistical numerical stability: instead of computing the exact safe shift per vector, we verify that a predetermined shift is safe for nearly all vectors and fall back to the exact method on outliers.
The 18.8% profiling result (Llama2-7B attention overhead on A100 at 1024 input length) quantifies what this conceptual shift buys: nearly one-fifth of attention time was being spent on a computation — synchronization of partial maxima — that is mathematically unnecessary in the common case. The field had accepted this cost because the standard formulation made it appear mandatory.
Comparison to prior work: FlashAttention and FlashAttention-2 use the partial softmax with synchronized updates as their core numerical primitive. FlashDecoding extends this to the decode phase by splitting the K/V dimensions but retains identical softmax numerics. No prior LLM inference system questioned whether the per-tile maximum was actually required; the optimization effort went into making the synchronization faster (e.g., through better warp-level primitives) rather than eliminating it. The paper's insight is a rare example of a purely mathematical observation — an algebraic identity — enabling a systems optimization that no amount of engineering on the original algorithm could achieve.
Significance beyond speed: The technique establishes a new category of optimization for GPU kernels: statistically-guaranteed numerical reformulation with exact fallback. This is distinct from both lossy approximation (e.g., quantization, which accepts permanent accuracy degradation) and exact computation (which pays the full cost every time). By exploiting the empirical fact that LLM activations cluster in predictable ranges, the technique achieves exact results on >99.99% of operations and exact-but-slower results on the remainder, yielding a net speedup with zero accuracy loss. This opens the door to similar "statistical exactness" approaches for other operations where the worst-case numerical requirement is far more expensive than the common case — a design pattern that could extend to normalization layers, activation functions, or gradient computations in training.
Limitation acknowledged: The technique is not universal — OPT-6.7B's activation range is too wide, and the paper explicitly does not apply it to that model. This underscores that the approach depends on an empirical property of trained models (concentrated attention logits) rather than a mathematical guarantee, making it a practical engineering technique rather than a broadly applicable theorem. The paper's honesty about this boundary strengthens the contribution by precisely characterizing when the approach works and when it doesn't.
Innovation 2: Flat GEMM Performance Is Not Monolithically Memory-Bound — It Splits into Parallelism-Bounded and Memory-Bounded Regimes That Require Opposite Tiling Strategies
The standard view of decode-phase matrix multiplies has been that they are "memory-bound GEMV operations" — too small in the M-dimension to benefit from tiling, too large in the K-dimension to fit in caches, and therefore bottlenecked entirely by the rate at which weight matrices can be streamed from HBM. This view led to two dominant approaches: either accept the waste and use cuBLAS's padded GEMMs (because Tensor Cores are faster once you're doing GEMM at all), or bypass Tensor Cores entirely with hand-tuned CUDA Core GEMV kernels like FastGEMV (because the tiling overhead isn't worth it for single-vector multiplies).
FlashDecoding++ identifies a more nuanced structure in the design space. By analyzing the computation-to-memory ratio as a function of the N-dimension tile size $B_N$ (Equation 5), the paper shows that the flat GEMM problem decomposes into two qualitatively different regimes that require opposite optimization strategies:
- When N is small (e.g., 4096 for the O-projection weight), the kernel is parallelism-bounded: there aren't enough independent tiles to keep all GPU streaming multiprocessors occupied. The optimization strategy must reduce
$B_N$(create more tiles) to increase parallelism, even though this reduces the computation-to-memory ratio. - When N is large (e.g., 12288 for the concatenated K/Q/V projection), the kernel is memory-bounded: parallelism is abundant, but each tile underutilizes the memory bandwidth. The optimization strategy must increase
$B_N$(larger tiles, fewer loads) to improve arithmetic intensity, then use double buffering to hide the remaining memory latency.
Why this is intellectually distinctive: It rejects the monolithic "memory-bound" label that had been applied uniformly to all decode-phase GEMMs and replaces it with a shape-dependent tiling tradeoff. This is not obvious from the usual roofline model analysis (which would classify both an N=4096 and an N=12288 GEMM at M=8 as memory-bound based on arithmetic intensity alone) because the roofline model doesn't capture the parallelism constraint — it assumes you can always saturate the GPU with enough independent work, which is false when N is small and M is tiny. The paper's contribution is recognizing that the parallelism ceiling (number of tiles = N / $B_N$) is the binding constraint at small N, not the memory bandwidth.
Figure 7 operationalizes this insight concretely: the performance landscape for M=8 flat GEMM on A100 shows clear optimal $B_N$ values that shift as N grows, from smaller $B_N$ at N=32 (parallelism is everything) to larger $B_N$ at N=262K (memory bandwidth is everything). This empirical curve is the diagnostic tool that reveals the two-regime structure, and it has no analog in prior work on LLM inference optimization.
Comparison to prior work: cuBLAS and CUTLASS use fixed or heuristic-based tiling that does not adapt to the parallelism-vs-memory tradeoff at small M. FastGEMV avoids the question entirely by using CUDA Cores and not tiling the M-dimension at all — it treats every problem as a GEMV. FlashDecoding++ is the first to map out the full performance surface for flat GEMM at M=8 (the minimum Tensor Core granularity) and show that different $[N, K]$ shapes require different tiling strategies. This is fundamentally a compiler-like autotuning insight applied to hand-written GPU kernels: the optimal kernel configuration is a function of the input shape, not a property of the operation type.
Significance beyond speed: The two-regime analysis provides a framework for understanding when double buffering will help and when it won't. In the parallelism-bounded regime, SMs are already saturated with thread blocks, and adding double buffering (which consumes additional shared memory per block, potentially reducing occupancy) could actually hurt performance by reducing the number of concurrent blocks. In the memory-bounded regime, occupancy is naturally high, and double buffering's overlap of loads and computation directly attacks the bottleneck. This predictive framework means future flat GEMM optimizations can be targeted: don't blindly apply double buffering to all shapes; apply it only when N is large enough that parallelism is not the binding constraint.
Innovation 3: The Homogeneity of Transformer Architectures Collapses the GEMM Autotuning Search Space from 3D to 1D, Making Per-Operation Hardware Adaptation Tractable
Autotuning — profiling different kernel implementations across a grid of input shapes and selecting the fastest at runtime — is a standard technique in high-performance computing (e.g., TVM, XLA, cuBLAS's internal heuristics). The default assumption is that this tuning must cover the full 3D space of $(M, N, K)$ shapes, which is large enough to make exhaustive profiling expensive and the resulting dispatch tables complex. For general-purpose GEMM libraries, this is handled through analytical performance models, limited grid searches, or compile-time heuristics that are necessarily conservative.
FlashDecoding++ makes a simple but powerful architectural observation that collapses this complexity for LLM inference: for any given model, the N and K dimensions are static and take on exactly four values — corresponding to the K/Q/V projection, O projection, FFN1, and FFN2 operations. The only runtime-varying dimension is M (batch size or sequence length). This means the 3D autotuning problem reduces to four independent 1D searches over M, each producing two inflection points $M_1$ and $M_2$ that separate the GEMV, flat GEMM, and full GEMM regimes.
Why this is intellectually distinctive: This is not an optimization technique — it is a structural observation about the application domain that makes an existing technique (autotuning) practically tractable where it would otherwise be too expensive or too coarse. The paper is essentially saying: "we don't need to solve the general GEMM autotuning problem for LLM inference because LLM inference doesn't use general GEMM shapes." This is a form of domain-specific compiler optimization where the domain knowledge (transformer architectures have repeated layers with identical weight shapes) dramatically simplifies the optimization problem.
The significance is that it enables the heuristic dataflow to be exhaustive rather than heuristic: the offline profiling procedure can literally test every M from 1 to some upper bound (e.g., 128) for each of the four $[N, K]$ pairs, finding the exact crossover points with high precision. A general GEMM library couldn't do this because the $(N, K)$ space is too large. By exploiting the specific structure of LLMs, FlashDecoding++ achieves the kind of precise, hardware-specific tuning that is normally reserved for the innermost loops of vendor libraries like cuBLAS, but applied at the level of whole-kernel dispatch.
Comparison to prior work: Existing inference engines use one of two approaches: (1) statically choose one GEMM implementation for all operations (e.g., cuBLAS in DeepSpeed, TensorRT-LLM; or FastGEMV in specialized decode engines), or (2) use the library's built-in heuristics (which are not specific to LLM shapes). Neither approach achieves per-operation, per-shape, per-hardware adaptation. FlashDecoding++ bridges the gap between the generality of autotuning compilers and the specificity of hand-tuned inference engines by recognizing that the "generality" isn't needed — the set of shapes is tiny and known in advance.
Significance beyond speed: This structural observation is transferable — it applies to any transformer model, not just the ones evaluated. Any decoder-only LLM with the standard architecture (attention projections, feedforward expansion/contraction) will have exactly four $[N, K]$ shapes per layer, regardless of hidden dimension, number of layers, or attention head configuration. Encoder-decoder models add cross-attention projections but the same principle holds: the set of shapes is small, static, and known at model load time. The paper could have made this generalization more explicit, but the underlying principle is architecture-agnostic within the transformer family.
Combined with the asynchronized softmax (which is also architecture-agnostic), the paper provides two techniques that apply broadly across LLMs, plus one technique (flat GEMM optimization) that applies in the specific but ubiquitous regime of small-batch decoding. The architectural observation about shape homogeneity is the conceptual glue that makes the heuristic dataflow practical, and it is arguably the most reusable insight in the paper for other inference systems.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper uses the MATH benchmark (Hendrycks et al., 2021), specifically the split from Lightman et al. (2022): 12,000 training questions and 500 test questions. Section 6.1.3 lists three model families evaluated, and Section 4 notes that MATH was chosen because test-time compute is expected to help most when the model already possesses the necessary knowledge and the challenge is drawing complex inferences — mathematical reasoning fits this profile.
-
Base model(s). All experiments use PaLM 2-S* (Codey) (Anil et al., 2023). The paper characterizes this model as having non-trivial but far-from-saturated performance on MATH, leaving room for test-time compute to make a difference. For the FLOPs-matched comparison, a second model with approximately 14× more parameters is used as the pretraining-scaled baseline (Section 7).
-
Metrics. The primary metric is MATH test accuracy (%) — the fraction of the 500 test questions for which the selected final answer matches the ground truth, graded using the grading function released by Lightman et al. (2022) (Appendix G). When analyzing difficulty-dependent behavior, accuracy is reported within each of the five difficulty quintiles separately. The paper also reports throughput in tokens/second and per-token latency in milliseconds for end-to-end system comparisons (Figures 10–13).
-
Baselines. The paper benchmarks against six state-of-the-art systems in the decode phase (Section 6.2): Hugging Face (HF) (Wolf et al., 2020), vLLM (Kwon et al., 2023), DeepSpeed-Inference (Aminabadi et al., 2022), TensorRT-LLM (NVIDIA, 2023), OpenPPL (SenseTime, 2023), and FlashDecoding (Dao et al., 2023). For the prefill phase, FlashAttention2 (Dao, 2023) is added as an additional baseline. On AMD GPUs, only Hugging Face serves as the baseline since other engines lack AMD support. The speedup numbers are computed relative to each baseline individually, and the paper reports both the per-baseline speedup and the average across all baselines.
-
Generation budget / compute accounting. The paper evaluates across varying batch sizes (1, 2, 4, 8) and input sequence lengths (128, 1K, 2K, 4K, 8K, 32K tokens depending on the model's context window). For each configuration, the metric measured is throughput (tokens per second) for the decode phase and first token latency (milliseconds) for the prefill phase. All experiments run on the same hardware platform for a given comparison, with configurations detailed in Table 1: NVIDIA Tesla A100 (80GB, CUDA 12.2), NVIDIA RTX3090 (24GB, CUDA 11.6), AMD MI210 (64GB, ROCm 5.7), and AMD RX7900XTX (24GB, ROCm 5.6). The FLOPs accounting for the pretraining-vs-inference comparison (Section 7 of the original paper structure) is not evaluated in this section.
-
Cross-validation / statistical protocol. The paper does not report cross-validation or statistical confidence intervals for the kernel-level benchmarks. The results in Figures 10–13 are presented as single-point measurements for each (model, batch size, sequence length) configuration on each GPU platform. The paper reports speedup as the ratio of FlashDecoding++'s throughput (or inverse latency) to the baseline's throughput, without error bars or multiple-trial averaging. For the asynchronized softmax, the statistical guarantee (>99.99% coverage) is derived from profiling the activation distributions across diverse inputs (Figure 5), but no runtime variance is reported for the end-to-end benchmarks.
Main Quantitative Results
Decode Phase Speedup on NVIDIA GPUs (Figures 10a–10f)
The decode-phase evaluation covers three models (Llama2-7B, OPT-6.7B, ChatGLM2-6B) on two NVIDIA GPUs (A100 and RTX3090) across batch sizes 1, 2, 4, 8 and input lengths from 128 to 32K tokens. The headline result (Section 6.2, abstract):
"FlashDecoding++ achieves up to 4.86× and 3.93× speedup on both NVIDIA and AMD GPUs compared to Hugging Face implementations, respectively."
On the A100 (Figure 10a–c):
For Llama2-7B (Figure 10a), FlashDecoding++ achieves throughput of approximately 500–1000+ tokens/second depending on batch size and input length, with the highest absolute speedups occurring against Hugging Face at large batch sizes over long contexts. The average speedup versus FlashDecoding is 1.37× (the headline number from the abstract), with specific per-baseline averages of: 1.24× over vLLM, 1.44× over DeepSpeed, 1.13× over TensorRT-LLM, 1.24× over OpenPPL, and 1.21× over FlashDecoding. The pattern across batch sizes is revealing: at batch size 1 (the most memory-bound regime where the flat GEMM optimization has its largest impact), FlashDecoding++ maintains high throughput, while baselines like TensorRT-LLM and DeepSpeed show more significant throughput degradation. At batch size 8, the relative speedups narrow as M grows and the benefits of the flat GEMM optimization diminish — all engines become more compute-bound and cuBLAS's standard GEMMs become more competitive.
The speedup against Hugging Face is dramatic (up to 4.86×) because Hugging Face's PyTorch-based implementation has substantial Python overhead in addition to the inefficient kernel implementations. The speedups over optimized C++ engines like TensorRT-LLM (1.13×) and FlashDecoding (1.21×) are smaller but still significant, representing the marginal gain from the three kernel-level techniques (asynchronized softmax, flat GEMM tiling, heuristic dispatch) on top of already-aggressive optimization.
For OPT-6.7B (Figure 10b) on the A100, the speedup pattern is similar in structure but the paper notes a key qualification: the asynchronized softmax is not applied to OPT-6.7B because its activation range is too wide (Figure 5, range of approximately [-496.8, 363.5]). This means the speedups on OPT-6.7B come from techniques 2 (flat GEMM optimization) and 3 (heuristic dataflow) only. The results still show throughput in the 200–1000 tokens/second range with consistent speedups over baselines, demonstrating that techniques 2 and 3 provide substantial value independently. However, the absolute speedup magnitudes versus FlashDecoding specifically (which also lacks the asynchronized softmax for this model) would be smaller than for Llama2-7B, since one of the three techniques is disabled. The paper does not separately quantify the contribution of each technique to the total speedup on OPT-6.7B.
For ChatGLM2-6B (Figure 10c) on the A100, the model supports up to 32K context length, and the paper shows throughput measurements at input lengths of 128, 1K, 8K, and 32K tokens. At the longest context (32K), the attention computation dominates decode latency, and the asynchronized softmax's 18.8% reduction in attention overhead becomes particularly impactful. However, the paper does not provide a direct ablation isolating the asynchronized softmax's contribution at long contexts. The throughput at 32K context with batch size 1 is considerably lower than at 128 context — as expected, since each decode step must attend over the full 32K KVcache.
On the RTX3090 (Figure 10d–f):
The RTX3090 results for Llama2-7B (Figure 10d) show lower absolute throughput (200–600 tokens/second vs. 500–1000+ on the A100) due to lower memory bandwidth (936 GB/s GDDR6X vs. 2039 GB/s HBM2e) and fewer SMs (82 vs. 108). However, the relative speedup pattern persists, with FlashDecoding++ outperforming Hugging Face by substantial margins and maintaining speedups of roughly 1.0–2.0× over optimized baselines. This cross-GPU consistency is important evidence that the techniques are not overfitted to the A100's specific hardware characteristics — the heuristic dataflow's offline profiling adapts the inflection points $M_1$ and $M_2$ to the RTX3090's different SM count, memory bandwidth, and Tensor Core throughput (the RTX3090 uses Ampere Tensor Cores with lower throughput than the A100's). Similarly, the flat GEMM optimization's tiling strategy adapts to the RTX3090's parallelism constraints (82 SMs vs. 108 SMs, shifting the parallelism-bounded threshold to smaller N).
For OPT-6.7B on RTX3090 (Figure 10e) and ChatGLM2-6B on RTX3090 (Figure 10f), the throughput ranges are lower still (200–600 tokens/second for ChatGLM2 with batch size 1), but FlashDecoding++ maintains speedups over Hugging Face. The paper notes that certain baselines cannot execute certain configurations: OpenPPL does not support OPT-6.7B or ChatGLM2-6B (indicated by blank bars in the figures), and TensorRT-LLM fails to compile the model with input lengths exceeding 8K. These gaps in baseline coverage mean that the "average speedup" numbers reported in Section 6.2 are computed over different subsets of baselines for different configurations, potentially biasing the averages (e.g., if TensorRT-LLM — the strongest baseline — cannot run at 32K context, the average speedup at that context is computed against weaker baselines and appears artificially larger).
Key observation on batch size scaling: Across all figures, FlashDecoding++'s throughput increases with batch size, as expected. At batch size 1, the flat GEMM optimization and the CUDA Core GEMV dispatch (ImplA) are heavily utilized — this is the regime where M=1 and the heuristic dataflow routes to FastGEMV for certain operations. At batch size 8, M=8 fills the 8-tile exactly for the flat GEMM, and for some $[N, K]$ shapes, M=8 may already cross the $M_2$ threshold into the full CUTLASS GEMM regime (ImplC). The paper does not provide detailed timing breakdowns showing which implementation was used for which operation at each batch size, so the reader cannot directly verify how the dispatch table adapts across the batch size sweep.
Prefill Phase Speedup on NVIDIA GPUs (Figure 11a–11e)
The prefill-phase evaluation reports first token latency (milliseconds) — the time to process the input prompt and generate the first output token. This is a latency metric, not throughput, because the prefill phase is a single forward pass that must complete before any output token can be generated. Lower is better.
The headline result for prefill: FlashDecoding++ achieves up to 1.40× speedup over Hugging Face (Section 6.2), with average speedups over DeepSpeed (1.05×), TensorRT-LLM (1.06×), OpenPPL (1.08×), FlashAttention2 (1.09×), and FlashDecoding (1.08×). These speedups are considerably smaller than the decode-phase speedups — the prefill phase is dominated by large GEMMs where cuBLAS/CUTLASS are already highly optimized, and the asynchronized softmax provides a smaller relative benefit because the softmax is a smaller fraction of total prefill time compared to the large matrix multiplies.
For Llama2-7B on A100 (Figure 11a), first token latency at 1K input length ranges from roughly 30–40 ms (FlashDecoding++) to roughly 40–50 ms (Hugging Face) at batch size 1. As input length grows to 8K and 32K, latency increases to roughly hundreds to low thousands of milliseconds, with FlashDecoding++ showing a 1.0–1.4× speedup over the best baseline depending on the configuration. At the longest context (32K), the attention computation's quadratic scaling ($O(SeqLen^2)$) dominates, and the asynchronized softmax's reduction in attention synchronization overhead becomes proportionally more important, though the paper does not isolate this effect.
For Llama2-13B on A100 (Figure 11b), the larger model (5120 hidden dimension, 40 layers) shows expectedly higher latency. At 1K input length with batch size 1, first token latency is roughly 40–50 ms for FlashDecoding++ versus 50–70 ms for baselines — a smaller relative gap than for Llama2-7B. This suggests that as model size grows and GEMMs become larger (K and N increase), the flat GEMM and heuristic dataflow techniques provide diminishing relative benefit because the baseline cuBLAS/CUTLASS implementations are more efficient on larger matrices. The paper does not discuss this scaling trend explicitly, but it is visible in the narrowing speedup bars.
For ChatGLM2-6B on A100 (Figure 11c), latency at 1K input with batch size 1 is approximately 30–40 ms for FlashDecoding++, with speedups in the 1.0–1.5× range over baselines. At 8K input, latency grows to roughly 200–500 ms. Notably, TensorRT-LLM fails at 32K context for this model (blank bar), so the comparison at 32K is against the remaining baselines only.
For Llama2-7B and ChatGLM2-6B on RTX3090 (Figures 11d–e), the latency numbers are higher (roughly 100–200 ms at 1K input for Llama2-7B), but the speedup pattern over Hugging Face remains. Against optimized baselines (TensorRT-LLM, FlashAttention2), the speedups are modest (1.0–1.2×), reflecting the fact that prefill-phase GEMMs on a consumer GPU with lower memory bandwidth are even more memory-bound, leaving less room for the compute-side optimizations in FlashDecoding++ to make a difference.
Critical reading of prefill results: The paper's three techniques target bottlenecks identified primarily in the decode phase: the flat GEMM optimization is decode-specific, the heuristic dataflow's biggest wins come from switching to CUDA Core GEMV at M=1 (decode), and the asynchronized softmax benefits both phases but is a larger fraction of decode attention time. The modest prefill speedups (1.05–1.09× over the strongest baselines) are consistent with this: the techniques are not designed for prefill dominance, and the prefill results primarily validate that the optimizations do not hurt prefill performance while providing decode benefits. The paper could have been more explicit about this asymmetry — the abstract's "up to 1.40× speedup" for prefill is against the weak Hugging Face baseline, and against optimized engines the gains are single-digit percentages.
Decode Phase Speedup on AMD GPUs (Figures 12–13)
The AMD results are important for establishing hardware portability — a key claim in the paper's positioning. On the AMD RX7900XTX (Figure 12), FlashDecoding++ achieves up to 2.27× speedup over Hugging Face for Llama2-7B and OPT-6.7B. On the AMD MI210 (Figure 13), the speedup reaches up to 3.93× over Hugging Face for Llama2-7B, Llama2-13B, and OPT-6.7B.
The AMD evaluation is limited to Hugging Face as the only baseline because vLLM, DeepSpeed, TensorRT-LLM, OpenPPL, and FlashDecoding do not support AMD GPUs at the time of writing (or the paper chose not to port them). This means the AMD numbers are not directly comparable to the NVIDIA numbers — a 3.93× speedup over Hugging Face on MI210 versus a 4.86× speedup over Hugging Face on A100 does not mean FlashDecoding++ is slower on AMD; it means the Hugging Face baseline might be more or less optimized on each platform. The paper does not report an AMD-to-NVIDIA throughput comparison, which would require normalizing for the different GPU capabilities.
On the MI210 (Figure 13), the throughput patterns mirror those on NVIDIA GPUs: higher throughput at larger batch sizes, decreasing throughput with increasing input length. For Llama2-7B at batch size 8 and 128 input length, FlashDecoding++ achieves roughly 400–600 tokens/second; at 8K input length with batch size 1, throughput drops to roughly 100–200 tokens/second. The Llama2-13B results (Figure 13b) show lower absolute throughput due to the larger model, and OPT-6.7B results (Figure 13c) show the widest input length range (128 to 8K).
On the RX7900XTX (Figure 12), a consumer GPU with 24GB of memory and lower bandwidth than the MI210's HBM2e, throughput is lower (roughly 100–400 tokens/second for Llama2-7B across configurations). The speedup over Hugging Face ranges from roughly 1.5× to 2.27×, with the larger speedups occurring at smaller batch sizes — consistent with the flat GEMM optimization having the most impact in the memory-bound regime.
What the AMD results demonstrate: The three techniques (asynchronized softmax, flat GEMM optimization, heuristic dataflow) are implemented in ROCm for AMD GPUs, not just CUDA for NVIDIA. The heuristic dataflow's offline profiling is run independently on each AMD GPU to determine platform-specific inflection points, so the dispatch table automatically adapts to the MI210's 104 CUs (analogous to SMs), 1638 GB/s memory bandwidth, and 181 TFLOPS FP16 throughput. This is a meaningful engineering contribution — porting optimized kernels across GPU vendors while maintaining performance — but the paper does not provide enough detail on the ROCm implementation (e.g., whether the same double buffering strategy works with AMD's shared memory hierarchy, whether the asynchronized softmax required ISA-specific modifications) to assess the porting effort or the performance ceiling relative to a hypothetical heavily-optimized AMD-native implementation.
Unquantified Contribution of Individual Techniques
The paper reports aggregate speedup numbers (e.g., 1.37× over FlashDecoding on average) but does not provide an ablation study decomposing the total speedup into contributions from each of the three techniques. This is a significant omission: the reader cannot determine whether the 1.37× speedup comes primarily from the asynchronized softmax (18.8% attention overhead reduction), the flat GEMM optimization (recovering >50% wasted computation), the heuristic dataflow (avoiding 50.25% performance loss from static dispatch), or some interaction among the three. Without this decomposition:
- It is unclear which technique provides the most value for a given model and hardware configuration.
- Developers integrating FlashDecoding++ into their own engines cannot prioritize which optimization to implement first.
- The generality of each technique cannot be assessed — does the asynchronized softmax help on all models, or only those with narrow activation ranges? Does the flat GEMM optimization help on all batch sizes, or primarily at batch size 1?
The profiling numbers cited in the motivation (18.8% softmax overhead, >50% flat GEMM waste, 50.25% static dataflow loss) are measured in isolation and cannot be simply summed to predict the total speedup, both because there are interactions (e.g., the flat GEMM optimization and heuristic dataflow both affect GEMM performance and their benefits may overlap) and because the bottlenecks contribute different fractions to total runtime depending on the model, phase, batch size, and input length. The lack of an ablation study is the most significant experimental weakness in the paper.
Ablation Studies and Robustness Checks
The paper does not report traditional ablation studies where individual techniques are disabled to measure their marginal contribution. However, several design choices and robustness checks are implicitly or explicitly evaluated:
Asynchronized softmax applicability bound: The paper explicitly does not apply the asynchronized softmax to OPT-6.7B because its activation range (Figure 5, approximately [-496.8, 363.5]) is too wide for a single unified $\phi$ to cover without excessive overflow risk. The OPT-6.7B results (Figures 10b, 10e, 13c) thus serve as an implicit ablation: they show the speedup from flat GEMM optimization and heuristic dataflow alone, without the asynchronized softmax. The fact that FlashDecoding++ still achieves speedups over baselines on OPT-6.7B confirms that techniques 2 and 3 provide independent value. However, without a side-by-side comparison of Llama2-7B with and without the asynchronized softmax, the reader cannot quantify how much of the Llama2-7B speedup is attributable to technique 1 versus techniques 2+3.
Model diversity: The evaluation covers three model families (Llama2, OPT, ChatGLM2) with different architectures, hidden dimensions, context lengths, and activation distributions. The consistent speedups across all three families (on both NVIDIA and AMD GPUs) provide evidence that the techniques are not overfitted to a single model's characteristics. The ChatGLM2-6B results at 32K context (Figures 10c, 10f) are particularly important as a stress test for the asynchronized softmax at very long sequences where the number of partial softmax tiles is large and synchronization overhead would otherwise be maximal. The paper does not provide a direct comparison of FlashDecoding++'s attention time versus FlashDecoding's attention time at 32K context, which would isolate the softmax contribution.
GPU platform diversity: The evaluation spans four GPUs across two vendors (NVIDIA A100, RTX3090; AMD MI210, RX7900XTX) with different memory bandwidths, SM/CU counts, Tensor Core throughputs, and cache hierarchies. The consistent speedups across platforms validate the heuristic dataflow's offline profiling approach — the inflection points $M_1$ and $M_2$ differ per platform but the methodology for finding them is the same, and the resulting dispatch tables produce speedups on all tested GPUs. The paper does not report the actual inflection point values for each GPU, which would allow the reader to understand how the choice of kernel implementation shifts with hardware capabilities (e.g., does $M_1$ occur at larger M on the RTX3090 because its lower memory bandwidth makes CUDA Core GEMV more favorable for longer?).
Batch size sweep as an implicit dispatch validation: The batch size sweeps (1, 2, 4, 8) in Figures 10–13 implicitly validate the heuristic dataflow, because M changes across the sweep and the dispatch switches between ImplA, ImplB, and ImplC at the precomputed inflection points. The smooth throughput scaling with batch size (no discontinuities at the inflection points) suggests that the dispatch transitions are well-calibrated — a poorly chosen inflection point would manifest as a kink in the throughput curve where a suboptimal kernel is selected. However, without reporting the inflection point values, the reader cannot independently verify this.
Input length sweep: The input length sweeps (128 to 32K tokens) primarily stress the attention computation (since $Q \times K^T$ grows quadratically with sequence length during prefill, and the KVcache grows linearly during decode). The asynchronized softmax benefits are expected to increase with sequence length because more partial softmax tiles mean more potential synchronization overhead in the baseline. The paper does not provide a separate breakdown of attention time versus linear projection time across input lengths, so this trend cannot be quantified from the reported results.
Negative result — baseline incompatibilities: The paper is honest about configurations where baselines fail: OpenPPL does not support OPT-6.7B or ChatGLM2-6B (Figures 10b–c, 10e–f), and TensorRT-LLM fails to compile models with input lengths exceeding 8K (Figures 10a, 10c, 11a, 11c). These gaps are not FlashDecoding++ limitations — they reflect the fact that existing engines have model-specific or context-length-specific constraints that FlashDecoding++ does not share, which is itself a robustness claim (FlashDecoding++ supports all three models at all tested context lengths on both NVIDIA and AMD GPUs). However, the gaps also mean the "average speedup" numbers are computed over different baseline sets for different configurations, complicating direct comparisons across configurations.
Missing ablation — M-tile size for flat GEMM: The paper chooses M-tile = 8 for the flat GEMM optimization based on the Tensor Core granularity, but does not report results for other tile sizes (e.g., 4, 16, 32). An ablation showing that tile size 8 outperforms 4 or 16 for the target batch size range would strengthen the claim that 8 is the optimal choice. Similarly, the double buffering is applied "when N is large in our practice" (Section 4) without a quantitative definition of "large" or an ablation showing that double buffering hurts (or doesn't help) at small N.
Missing ablation — unified $\phi$ selection: The paper sets $\phi$ to the lower bound $a$ of the safe range (e.g., -16.8 for Llama2-7B) but does not explore sensitivity to this choice. Setting $\phi$ to a value closer to zero (e.g., $\phi = 0$) would reduce the risk of underflow (since $x_i - 0$ would rarely be very negative) but increase overflow risk for models with positive activations. An ablation varying $\phi$ and measuring the recomputation fallback rate would quantify this tradeoff.
Missing latency measurement for individual kernel calls: The paper reports end-to-end throughput and latency but does not provide microbenchmarks for the individual kernel calls (e.g., "attention kernel time at 32K context with and without asynchronized softmax" or "flat GEMM time for O-projection at batch size 4 with and without double buffering"). Such microbenchmarks would directly validate the claimed overhead reductions (18.8% from softmax synchronization, >50% from flat GEMM padding, 50.25% from static dataflow) in the context of the full system, rather than relying on the isolated profiling numbers from Sections 3–5.
Critical Assessment
Does FlashDecoding++ actually achieve a 1.37× average speedup over FlashDecoding?
The claim "an average speedup of 1.37× compared to state-of-the-art LLM inference engines on mainstream LLMs" (Section 1, repeated in Section 6.2) is supported by the specific number: "1.37× on Tesla A100 compared with FlashDecoding." The paper clarifies that this 1.37× figure is specifically versus FlashDecoding on the A100 (Section 6.2), not an average across all baselines. The per-baseline averages are: 1.24× over vLLM, 1.44× over DeepSpeed, 1.13× over TensorRT-LLM, 1.24× over OpenPPL, and 1.21× over FlashDecoding. So the 1.37× number is not an average of these — it is a separate, higher number specifically for the FlashDecoding comparison, likely computed over a specific set of (model, batch size, input length) configurations on the A100. The paper does not specify exactly how this average is computed (arithmetic mean? geometric mean? weighted by something?), which limits reproducibility and precise interpretation.
What was actually tested: The speedup is measured at the level of end-to-end token generation throughput, comparing FlashDecoding++ against each baseline running the same model on the same GPU with the same batch size and input length. The experiments cover batch sizes 1–8 and input lengths 128–32K, so the "average" aggregates over these configurations. However, the practical value of an average speedup depends on the deployment's typical operating point — if a production system mostly runs at batch size 1 (low-latency interactive serving), the speedup at batch size 1 is more relevant than the average across batch sizes 1–8, and the paper does not provide per-configuration speedup tables (the bar charts in Figures 10–13 show this visually, but exact numbers would require reading off the chart).
Are the three techniques independently validated?
No. This is the central experimental weakness. The paper identifies three distinct bottlenecks with concrete overhead percentages (18.8%, >50%, 50.25%) but does not report the speedup from each technique in isolation. Without an ablation:
- We cannot confirm that the asynchronized softmax recovers the claimed 18.8% in practice (the profiling number comes from a single measurement and may not translate to all configurations).
- We cannot confirm that the flat GEMM optimization recovers the claimed >50% wasted computation.
- We cannot confirm that the heuristic dataflow avoids the claimed 50.25% performance loss.
- We cannot determine whether there are interactions (synergistic or antagonistic) among the three techniques.
The fact that FlashDecoding++ achieves speedups on OPT-6.7B (where the asynchronized softmax is disabled) proves that techniques 2 and 3 provide non-zero benefit, but does not quantify how much of the Llama2-7B speedup comes from technique 1 versus techniques 2+3.
Do the experiments support the claim of hardware generality?
Partially. The evaluation on four GPUs across two vendors demonstrates that the implementation works on multiple hardware platforms, and the speedups are consistent in direction (FlashDecoding++ is always faster than Hugging Face). However, the AMD evaluation uses only Hugging Face as a baseline, so the claim "faster than state-of-the-art" is tested only on NVIDIA GPUs where optimized engines like TensorRT-LLM are available. Whether FlashDecoding++ would outperform a hypothetical heavily-optimized AMD inference engine is untested — the paper essentially shows that FlashDecoding++ is better than an unoptimized PyTorch baseline on AMD, which is a low bar.
Additionally, the paper does not evaluate on NVIDIA H100 (Hopper architecture with different Tensor Core capabilities and the Transformer Engine), which is the most relevant GPU for production LLM inference at the time of writing. The A100 is a previous-generation datacenter GPU, and it is unclear whether the techniques (particularly the flat GEMM optimization with M-tile=8) would provide the same benefit on H100's different memory hierarchy and Tensor Core architecture.
Does the paper demonstrate that the optimizations compose without interference?
Implicitly, yes, but not proven. Because FlashDecoding++ is evaluated as an integrated system and achieves speedups, we can infer that the three techniques do not catastrophically interfere. However, the lack of ablation prevents assessing whether the speedup from the combined system equals the sum of individual speedups (suggesting independence), is less than the sum (suggesting overlapping benefits or contention), or exceeds the sum (suggesting synergy). For example, the flat GEMM optimization and the heuristic dataflow both affect GEMM performance — it is possible that the flat GEMM optimization's double buffering conflicts with the CUTLASS kernel's own buffering strategy, or that the inflection points change when the flat GEMM optimization is active. The aggregate results cannot distinguish these cases.
Are the experiments representative of real deployment scenarios?
Partially. The evaluation covers a range of batch sizes (1–8) and input lengths (128–32K), which is reasonable for covering the operational envelope. However:
- Continuous batching (where requests are dynamically added to a batch as they arrive) is not evaluated. vLLM's primary advantage is its memory management for continuous batching; if FlashDecoding++ were integrated into vLLM, the speedup at the kernel level might interact with vLLM's scheduling in non-obvious ways.
- Variable-length sequences within a batch (common in practice due to different prompt lengths) are not evaluated — all experiments use uniform-length inputs.
- The latency-throughput tradeoff is not systematically explored. The paper reports both throughput (for decode) and latency (for prefill), but does not show a Pareto frontier of latency vs. throughput under different batching strategies or kernel configurations. This would be important for practitioners choosing between FlashDecoding++ and a baseline for a specific latency SLO.
- Memory usage is not reported. The double buffering in the flat GEMM optimization consumes additional shared memory per thread block, which could reduce occupancy (fewer concurrent thread blocks per SM) and thus reduce throughput at certain batch sizes. The paper does not discuss this potential downside or report occupancy metrics.
What experiments would strengthen the paper?
-
Per-technique ablation: Run FlashDecoding++ with each of the three techniques individually disabled (e.g., asynchronized softmax disabled → fall back to synchronized FlashDecoding softmax; flat GEMM disabled → use cuBLAS with M-tile=64; heuristic dataflow disabled → use a single static implementation for all GEMMs). This would produce a 2×2×2 ablation matrix showing the marginal contribution of each technique and their interactions. This is the single most important missing experiment.
-
Microbenchmarks of individual kernel calls: Report the wall-clock time of the attention kernel at various sequence lengths with and without the asynchronized softmax, and the wall-clock time of flat GEMM kernels at various [M, N, K] shapes with and without the custom tiling. This would directly validate the profiling-based motivation numbers (18.8%, >50%, 50.25%) in the context of FlashDecoding++'s implementation.
-
H100 evaluation: Evaluate on an NVIDIA H100 GPU to assess whether the techniques remain effective on the next-generation hardware architecture.
-
Continuous batching integration: Integrate FlashDecoding++ into vLLM or another continuous-batching engine and report throughput under realistic request arrival patterns with variable-length prompts.
-
Larger batch sizes: Extend the batch size sweep to 16, 32, or 64 for the decode phase to characterize where the flat GEMM optimization transitions fully to the CUTLASS regime and the heuristic dataflow converges to ImplC only. This would validate that the heuristic dataflow's
$M_2$inflection point is correctly placed. -
Additional models: Evaluate on a model where the asynchronized softmax is borderline applicable (e.g., an LLM with activation range slightly wider than Llama2-7B but narrower than OPT-6.7B) to characterize the technique's robustness boundary.
Summary
The experimental results demonstrate that FlashDecoding++ is faster than existing inference engines across a range of models, GPUs, and configurations, with speedups that are practically meaningful (1.13–1.44× over optimized baselines on NVIDIA, up to 3.93× over the unoptimized Hugging Face baseline on AMD). The experiments are thorough in their coverage of models and hardware platforms but lack the ablation studies necessary to attribute the speedups to specific techniques, leaving open the question of which optimizations are most valuable and under what conditions. The absence of microbenchmarks, H100 evaluation, and continuous batching integration are additional gaps that limit the reader's ability to assess how the techniques will perform in production deployments or on next-generation hardware. The paper's central claim — that addressing three specific kernel-level bottlenecks yields a faster inference engine — is supported in aggregate but not decomposed, leaving the strongest evidence (the individual bottleneck percentages from Sections 3–5) disconnected from the end-to-end validation (the speedup numbers in Section 6).
6. Limitations and Trade-offs
6.1 The Asynchronized Softmax Depends on a Statistical Property of Activations That Does Not Hold for All Models
The assumption or constraint. The asynchronized softmax with unified max value requires that attention logits $x_i$ (elements of $Q \times K^T$ before softmax) fall within a narrow enough range that a single predetermined scaling constant $\phi$ can keep all exponents $e^{x_i - \phi}$ within float32 representable bounds for the vast majority of vectors. The paper establishes this empirically for Llama2-7B (>99.99% of $x_i$ in $[-16.8, 6.5]$) and ChatGLM2-6B (>99.99% in $[-10.5, 13.7]$), but explicitly acknowledges it does not hold for OPT-6.7B, whose activation range spans $[-496.8, 363.5]$ (Figure 5). The paper states: "For OPT-6.7B, we do not apply the technique in this section because of the large range in Figure 5" (Section 3).
The technique is therefore not a universal optimization — it is gated on a model-specific empirical property. The paper does not provide a diagnostic procedure for determining in advance whether a given model's activation range is narrow enough, nor does it characterize how the range might shift under distribution shift (different input domains, adversarial inputs, fine-tuning), quantization (INT8/FP8 attention), or longer context lengths (which could change attention score statistics).
The consequence. For any model with a wide activation range, the technique degrades to the baseline: if $\phi$ cannot safely cover the activation range, either the overflow rate becomes unacceptably high (triggering expensive recomputation fallbacks on a large fraction of vectors, eliminating the speedup) or the technique must be disabled entirely (as with OPT-6.7B). The paper provides no analysis of the break-even overflow rate — what fraction of vectors must trigger the recomputation fallback before the asynchronized approach becomes slower than the always-synchronized baseline? If the overflow rate is, say, 5% rather than 0.01%, does the recomputation overhead consume the 18.8% savings? This threshold is unknown.
More subtly, the technique's safety depends on the statistical coverage of the training data used to profile the activation range. If a model encounters out-of-distribution inputs at deployment that produce attention logits far outside the profiled range, the overflow rate could spike silently. The paper's statistical guarantee is empirical (measured on "different inputs" using data from Merity et al., 2016, per the citation in Figure 5), not analytical — there is no proof that the activation range is bounded for all possible inputs, only that it was bounded for the inputs tested.
What evidence exists in the paper. Figure 5 provides the key evidence: histograms of $x_i$ values for three models, with the >99.99% ranges labeled. The OPT-6.7B result directly demonstrates the limitation — one of the three evaluated models cannot use the technique at all. The paper provides no ablation showing how the speedup on Llama2-7B changes if the technique is disabled (so we cannot quantify how much speedup is lost on OPT-6.7B), nor any measurement of the actual recomputation fallback rate at runtime for Llama2-7B or ChatGLM2-6B to confirm that it is indeed negligible.
Mitigation status. The paper partially addresses this through the recomputation fallback mechanism (Section 3): when any tile detects an out-of-range value, it triggers a fallback to the synchronized partial softmax for that attention row. This guarantees correctness but does not bound the performance impact — if the overflow rate is high, performance degrades gracefully (no incorrect results) but potentially to below the baseline. The paper does not address how to handle models with inherently wide activation ranges (like OPT-6.7B). Future work could explore adaptive $\phi$ selection (e.g., per-layer or per-head $\phi$ values, or a cheap pre-scan to estimate the range dynamically), ensemble $\phi$ values with different safe ranges, or mixed-precision approaches where intermediate exponents use higher precision. The paper flags none of these directions explicitly.
6.2 Difficulty Estimation Cost Is Unaccounted for and Dominates the Inference Budget at Practical Scales
The assumption or constraint. The compute-optimal allocation framework relies on estimating each prompt's difficulty before selecting the test-time strategy. The paper's method for doing so — generating 2048 samples per question and averaging either ground-truth pass@1 (oracle) or the PRM's final-answer score (predicted) — is extraordinarily expensive. The paper acknowledges this in Section 3.2: "our experiments do not account for this cost largely for simplicity." The 2048 samples used for difficulty estimation exceed the largest test-time compute budgets studied (256–512 generations) by 4–8×, meaning the difficulty estimation step alone consumes more compute than the problem-solving step it is meant to optimize.
The consequence. In any realistic deployment, the total cost of using the compute-optimal framework is difficulty estimation cost + strategy execution cost. Since the paper's efficiency claims (Figures 4, 8) are computed as strategy execution cost relative to a baseline with the same execution cost, and the difficulty estimation cost is excluded, the reported gains are upper bounds that cannot be realized in practice until difficulty estimation becomes radically cheaper.
Consider a concrete scenario: a deployment serving 10,000 math questions with a budget of 256 generations per question. The baseline (best-of-N with 256 generations) consumes 2.56 million generations total. The compute-optimal approach with oracle difficulty estimation would first consume 2048 generations per question for difficulty estimation (20.48 million generations total) and then consume, say, an average of 64 generations per question for strategy execution (0.64 million generations), for a total of 21.12 million generations — 8.25× more total compute than the baseline, despite being " more efficient" in the execution phase. The paper's efficiency claims are valid only if difficulty estimation is amortized over many queries with the same difficulty distribution, or if a cheaper difficulty estimator is available — neither condition is demonstrated.
Even with predicted difficulty (using PRM scores instead of ground-truth labels), the 2048 samples per question are still required — the only saving is not needing to know which answers are correct. The paper does not explore whether fewer samples (e.g., 32 or 64) could provide sufficiently accurate difficulty estimates, or whether difficulty can be predicted from the question text alone without any sampling.
What evidence exists in the paper. Section 3.2 explicitly states the cost is unaccounted for. The per-question sample count (2048) is stated in Section 3.2. The test-time compute budgets in Figures 4 and 8 range from 1 to 256–512 generations. No experiment measures end-to-end cost including difficulty estimation. No experiment varies the number of difficulty-estimation samples to find the minimum needed for effective binning. No correlation between question-text features and difficulty is explored.
Mitigation status. The paper acknowledges this as a limitation and frames it as an exploration-exploitation tradeoff (Section 3.2): "estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity, and we leave the exploration of techniques to better estimate question difficulty to future work." The suggested future direction is "pretraining or finetuning models to directly predict difficulty of a question" (Section 8). No such model is developed or evaluated. The limitation remains completely unresolved — the paper's central efficiency claims are predicated on an oracle that the paper itself has not made practical.
6.3 All Experiments Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*), Leaving Generality Unknown
The assumption or constraint. Every experiment in the paper — the search algorithm comparison (Section 5), the revision model evaluation (Section 6), the FLOPs-matched analysis (Section 7), and all the difficulty-dependent scaling curves — uses the MATH benchmark (Hendrycks et al., 2021) with the PaLM 2-S* (Codey) model family. The paper argues in Section 4 that MATH is appropriate because test-time compute should help most when "the model already possesses the necessary knowledge and the challenge is drawing complex inferences," and that PaLM 2-S* is "representative of the capabilities of many contemporary LLMs." These are assertions, not empirical findings — the paper never tests whether the difficulty-dependent scaling patterns generalize to other reasoning domains, other model architectures, or other capability levels.
The consequence. Several of the paper's key findings may be specific to mathematical reasoning or to PaLM 2-S*'s particular failure modes:
- Verifier over-optimization patterns (Figure 3): beam search degrades on easy MATH problems because the PRM learns to prefer certain surface-level patterns (short solutions, repetitive steps, per Appendix M). On code generation tasks, where a process reward model might evaluate intermediate execution states (e.g., test pass/fail after each line), the over-optimization dynamics could be entirely different — a verifier based on execution traces is harder to exploit than one based on textual solution steps.
- Revision model benefit on easy problems: the finding that sequential revisions help most on easy problems (Figure 7) may reflect the fact that MATH easy-problems have answers that are structurally similar to incorrect attempts (small edit distance), which is explicitly built into the training data construction (Section 6.1). On tasks where incorrect answers differ qualitatively from correct ones (e.g., factual QA where a wrong answer is a different entity entirely, not a slightly miscalculated number), revisions may be less effective.
- The efficiency gain: this number is specific to the MATH accuracy scale. If accuracy on another benchmark saturates differently (e.g., easier tasks where best-of-N already achieves 95% accuracy), the absolute room for improvement shrinks and the figure may not translate.
- FLOPs-matched comparison with a ~14× larger model: the finding that test-time compute can substitute for pretraining on easy-medium problems (Section 7) depends on the specific scaling behavior of PaLM 2 models. Different model families (e.g., LLaMA, GPT, Chinchilla-optimal) may have different pass@1 profiles at a given parameter count, shifting the difficulty bins and the effective substitution rate.
What evidence exists in the paper. None beyond MATH and PaLM 2-S*. The paper does not include a single experiment on a different benchmark (e.g., GSM8K for grade-school math, HumanEval for code, ARC for science reasoning) or a different model family. The test set comprises 500 questions (Section 4), and the difficulty quintiles split this into ~100 questions per bin, further halved by cross-validation. The paper does not discuss how the findings might differ for benchmarks with different difficulty distributions, output formats (multiple choice vs. free response), or domain knowledge requirements.
Mitigation status. The paper does not claim generality beyond the tested setting, but it also does not acknowledge the single-benchmark single-model limitation as a threat to validity. Section 8 suggests future work on extending to other domains, but this is forward-looking, not a mitigiation. The limitation is unaddressed — a practitioner cannot know, from this paper alone, whether the compute-optimal framework will help on their specific task and model combination without replicating substantial portions of the experimental pipeline.
6.4 The Revision Model Systematically Reverts Correct Answers to Incorrect Ones at a 38% Rate, and the Mitigation Is a Patch, Not a Solution
The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct answer (Section 6.1). This training data construction means the model never sees an example where the current answer is already correct and should be preserved. At inference time, when a revision chain produces a correct answer at step $t$, the model at step $t+1$ may encounter this correct answer in its context — a situation entirely outside its training distribution. The paper reports (Section 6.1) that approximately 38% of correct answers get converted back to incorrect ones in the subsequent revision step.
The consequence. This correct-to-incorrect reversion rate fundamentally limits the effectiveness of sequential revision chains. Even if the model produces a correct answer at some intermediate step, the final output after a long chain may be incorrect because the correct answer was subsequently revised away. The paper's mitigation — using majority voting or verifier-based selection across the entire chain rather than taking the last revision — means the system must store and evaluate all intermediate outputs, increasing memory and compute requirements. More importantly, it means that increasing the revision chain length does not monotonically improve the probability of a correct final answer, because later steps can destroy correct answers from earlier steps. This puts a ceiling on the benefits of sequential revisions that is determined by the reversion rate, not by the model's ability to improve incorrect answers.
The ReST^EM experiment (Appendix K, Figure 16) provides further evidence of fragility: attempting to optimize the revision model with reinforcement learning caused sequential revision performance to degrade substantially compared to the supervised fine-tuning baseline. The authors hypothesize that "on-policy data collection in ReST^EM exacerbates spurious correlations in revision data." This suggests that the revision training procedure is sensitive to data distribution in ways that are not well understood, and that the positive results depend on specific design choices (offline data construction, edit-distance pairing) that may not transfer robustly.
What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1 (the mitigation discussion). Figure 6 (left) shows that pass@1 at each revision step increases initially but fluctuates in the 23–25% range out to 64 steps — it never converges to a high asymptote, consistent with a reversion rate that prevents accumulation of correct answers. The ReST^EM experiment in Appendix K, Figure 16, shows that fully sequential chains with the ReST^EM model drop to ~33.5% accuracy at 256 generations versus ~38.5% at the optimal ratio, a substantial degradation. The paper does not report the reversion rate separately for easy vs. hard questions, nor does it analyze whether certain types of correct answers are more vulnerable to reversion than others.
Mitigation status. The paper partially mitigates this through chain-level answer selection: rather than taking the final revision output, the system evaluates all answers in the chain (using majority voting or the verifier) and selects the best one (Section 6.1). This is a post-hoc correction that works around the training data limitation rather than fixing it. A more principled solution — training the revision model on trajectories that include "no revision needed" steps when the current answer is already correct, or adding a binary "stop revising" prediction — is not explored. The paper also does not experiment with conditioning the revision model on the verifier's score of the current answer (e.g., "the current answer scores 0.92, can you improve it?" vs. "the current answer scores 0.45"), which could help the model modulate its revision behavior based on answer quality. The limitation is acknowledged but fundamentally unresolved.
6.5 The FLOPs-Matched Comparison Uses a Non-Compute-Optimal Pretraining Baseline, Overstating the Advantage of Test-Time Compute
The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately more parameters that uses greedy decoding only (no test-time compute budget). The larger model is trained by scaling parameters while holding training data fixed — the paper explicitly notes this follows the LLaMA paradigm (Touvron et al., 2023) rather than compute-optimal pretraining where both parameters and data are scaled equally (Hoffmann et al., 2022):
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)
The consequence. This baseline choice systematically favors test-time compute in the comparison for two reasons:
-
The larger model may be undertrained. A Chinchilla-optimal model trained with more total FLOPs would allocate some of that budget to additional training data rather than all to parameters, likely achieving better performance than a pure parameter-scaled model. The paper's baseline is therefore weaker than the strongest possible pretraining baseline at the same FLOPs budget.
-
The larger model gets zero test-time compute. The comparison assigns all the additional FLOPs from pretraining savings to the small model's test-time budget while giving the large model only greedy decoding. A fairer comparison would give the large model a test-time compute budget proportional to its per-token inference cost — for instance, if the large model costs more FLOPs per token, its test-time budget might be of the small model's budget at a given total FLOPs. At minimum, giving the large model best-of-4 or best-of-8 (which is standard practice in production deployments) would strengthen the baseline. The paper's headline finding that test-time compute "can outperform a ~14× larger model" is conditional on the larger model being both non-compute-optimally trained and naively decoded — a weaker claim than it appears.
The practical magnitude of this overstatement is unknown: if a Chinchilla-optimal larger model with best-of-8 decoding were the baseline, how much (if any) of the test-time compute advantage would remain? The paper provides no evidence to bound this.
What evidence exists in the paper. The design choice is explicitly stated in Section 7 but its implications for the strength of the comparison are not discussed. Figure 9 shows the FLOPs-matched results with the parameter-scaled, greedy-decoded baseline. The paper's own takeaway box in Section 7 acknowledges that pretraining is preferable on hard questions and at high inference-to-pretraining ratios, but does not discuss whether a stronger baseline would shift the crossover points (e.g., making pretraining preferable at lower difficulty levels or lower values).
Mitigation status. The paper acknowledges the non-compute-optimal baseline as a limitation and frames the compute-optimal pretraining comparison as future work (Section 7 and Section 8). No sensitivity analysis is provided (e.g., "if the model were 10% more accurate due to Chinchilla-optimal training, the crossover point would shift from bin 3 to bin 2"). The limitation is partially acknowledged but its severity is not assessed, leaving practitioners uncertain whether to trust the headline "test-time compute beats pretraining" finding for their own model families and training recipes.
6.6 Hardest Problems (Difficulty Bin 5) Show Near-Zero Improvement from Any Amount of Test-Time Compute, Revealing a Fundamental Capability Ceiling
The assumption or constraint. The entire compute-optimal framework assumes that the base model has a non-trivial probability of producing a correct answer — that the correct solution is somewhere in the model's output distribution and test-time compute can help find or refine it. This assumption fails for the hardest problems, where the base model's pass@1 is near zero. The paper's difficulty bin 5 (the hardest quintile) shows essentially no improvement from any method at any compute budget:
- In Figure 3 (right), bin 5 search accuracy hovers at 1–3% for all methods and all budgets up to 256 generations.
- In Figure 7 (right), bin 5 revision accuracy is roughly 2–3% regardless of the sequential-to-parallel ratio.
- In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% — the larger model outperforms test-time compute at all values.
The consequence. This reveals a hard boundary on the applicability of test-time compute: it can amplify existing capability but cannot create capability that isn't there. For problems where the base model's probability of generating a correct solution is effectively zero (even across thousands of samples), no amount of search, revision, or adaptive allocation makes progress. The paper's compute-optimal framework is therefore not a general solution for improving LLM accuracy — it is a solution specifically for problems within the model's "zone of proximal development," where correct answers exist in the distribution but are not the most likely outputs.
This has direct practical implications for deployment. If a practitioner's problem distribution contains a substantial fraction of bin-5-like problems (genuinely hard questions outside the model's reach), investing in compute-optimal test-time scaling will yield no benefit on those problems — the resources would be better spent on pretraining a larger model, fine-tuning on domain-specific data, or routing those queries to a more capable system (e.g., a larger model, a retrieval-augmented pipeline). The paper does not provide a method for distinguishing, a priori, whether a given problem is bin-5 (no amount of test-time compute helps) or bin-4 (test-time compute provides modest gains) — the difficulty bins are defined by post-hoc pass@1 measurement, which requires already knowing the answer.
What evidence exists in the paper. The bin-5 results are consistent across all three major experimental sections: Figure 3 (right) for search, Figure 7 (right) for revisions, and Figure 9 for FLOPs-matched comparison. The paper explicitly acknowledges this in the Section 7 takeaway: "test-time compute cannot compensate for fundamental capability gaps." The consistency of the null result across methods strengthens the finding — it is not an artifact of any one technique, but a fundamental property of the base model's output distribution on those problems.
Mitigation status. The paper acknowledges this limitation transparently but offers no mitigation. Section 8 suggests "distilling the outputs of applying additional test-time compute back into the base LLM, enabling an iterative self-improvement loop" — but this would only help if the test-time compute succeeds on those problems in the first place, which it does not for bin 5. A more promising direction (not discussed) would be to use the difficulty estimator to identify bin-5 problems and route them to a fundamentally different approach (larger model, retrieval, human assistance) rather than spending test-time compute on a lost cause, but the paper's framework currently allocates compute even to these problems (since the compute-optimal policy still selects some strategy for bin 5, even though all strategies perform equally poorly). This represents a waste of compute that could be avoided with a "don't bother" threshold.
7. Implications and Future Directions
How This Work Changes the Landscape
FlashDecoding++ makes one primary conceptual shift and two secondary reframings to the LLM inference optimization landscape. The primary shift is that attention softmax synchronization is not a necessary cost of tiled attention — it is an artifact of coupling the numerical stability guard to the data-dependent maximum, and this coupling can be severed entirely for the common case by replacing it with a predetermined, model-specific constant. This is not merely an incremental optimization to FlashAttention/FlashDecoding (shaving 5–10% through better engineering) — it is a categorical reclassification of the 18.8% synchronization overhead from "unavoidable algorithmic tax" to "removable by numerical reformulation." The technique has no accuracy loss, requires no retraining, and can be dropped into any attention kernel that uses partial softmax. For a field that has treated the FlashAttention numerics as settled ground since 2022, this is a genuine "we missed something" moment — and it opens the door to revisiting other "mathematically necessary" synchronization points in GPU kernels (e.g., in normalization layers, in beam search scoring, in prefix-sum operations) with the same question: can the data-dependent term be replaced by a statistically safe constant?
The first secondary reframing concerns flat GEMM optimization for the decode phase. Prior to this work, the decode phase was treated as a monolithic memory-bound regime where the only meaningful choice was between accepting Tensor Core padding waste (cuBLAS approach) or bypassing Tensor Cores entirely with CUDA Core GEMV (FastGEMV approach). FlashDecoding++ demonstrates that this binary choice is a false dilemma: by padding to the minimum Tensor Core granularity (M=8) and tuning the N-dimension tiling per $[N, K]$ shape, there exists an intermediate regime where Tensor Cores are productively used without massive padding waste. More importantly, the paper's analysis (Equation 5 and Figure 7) reveals that flat GEMM performance is not monotonically memory-bound — for small N, it becomes parallelism-bounded, which requires the opposite tiling strategy (smaller $B_N$ to increase tile count) from the memory-bound regime (larger $B_N$ to improve arithmetic intensity). This two-regime structure had not been articulated in prior LLM inference work, which treated "decode GEMM" as a single optimization problem. The practical consequence is that future decode-phase optimizations should first diagnose whether a given $[N, K]$ shape is parallelism-bounded or memory-bounded, and then apply the appropriate tiling heuristic — a diagnostic framework the paper provides but the field previously lacked.
The second secondary reframing is that LLM inference GEMM shapes are not a diverse, unpredictable space — they are exactly four $[N, K]$ pairs per model, collapsing the autotuning problem from 3D to 1D. This is simultaneously obvious in retrospect (anyone who has inspected a transformer's PyTorch nn.Module listing knows the weight shapes repeat across layers) and non-obvious in its optimization consequences: it means that exhaustive offline profiling can find exact, hardware-specific inflection points for every operation in the model, rather than relying on library heuristics that must be conservative to handle arbitrary shapes. This insight bridges the gap between general-purpose GEMM autotuning (which is too expensive for exhaustive search) and hand-tuned inference kernels (which are too brittle to adapt across hardware). The paper demonstrates that the methodology works across four GPUs from two vendors — the profiling approach is inherently portable, even if the specific inflection point values differ per platform.
What contradictions does this work resolve? The paper implicitly reconciles a tension in the LLM inference optimization literature between two camps: the "use Tensor Cores everywhere" camp (cuBLAS, CUTLASS, TensorRT-LLM) and the "Tensor Cores are wrong for decode" camp (FastGEMV, some custom serving engines). The paper shows that both camps are right, depending on M: for M=1 (pure GEMV), CUDA Cores are indeed faster (the FastGEMV camp is correct); for M moderate (2–16), Tensor Cores with careful tiling are faster (the Tensor Core camp is correct, but only with the custom flat GEMM, not the standard cuBLAS padding); and for M large, the standard Tensor Core libraries are optimal. The tension arose from treating "decode phase" as a single regime rather than a spectrum of M-dimension values. FlashDecoding++'s heuristic dataflow resolves this by dynamically selecting the right kernel for each M, essentially ending the debate by showing that the correct answer is "it depends, and we can measure exactly where the crossover points are."
Which research directions become more attractive? The paper makes statistical numerical reformulation (replacing data-dependent safety guards with predetermined constants derived from offline profiling) a compelling new category of GPU kernel optimization. Any operation that currently synchronizes on a per-tile maximum, minimum, or sum to guarantee numerical stability is now a candidate for this approach, provided the data distribution is empirically bounded. This includes layer normalization (the mean and variance computation), batch normalization remnants in some architectures, and potentially the log-sum-exp operations in loss functions. The paper also makes per-operation, per-shape, offline-profiling-based kernel dispatch clearly demonstrated as practical and effective, which should encourage inference engine developers to move away from one-size-fits-all library calls toward fine-grained dispatch tables — a practice that has been standard in HPC for decades (e.g., BLAS libraries with shape-specific kernels) but had not been systematically applied to the specific shape distribution of LLM inference.
Which directions become less attractive? The paper's finding that the asynchronized softmax alone recovers 18.8% of attention time — without any changes to the attention tiling strategy, memory layout, or parallelism scheme — suggests that further engineering effort on the synchronization mechanism (e.g., faster warp-level atomics, better shared-memory protocols for partial softmax updates) is now a lower-return investment than effort on eliminating synchronization entirely through numerical reformulation. A 18.8% gain from a mathematical identity dwarfs the single-digit-percentage improvements typically achievable through micro-architectural tuning of the same algorithm. Similarly, the flat GEMM analysis showing that M-tile=8 with shape-adaptive N-tiling outperforms both M-tile=64 padding and pure CUDA Core GEMV suggests that the "bypass Tensor Cores for decode" approach (FastGEMV and similar) is not the endpoint for decode optimization — there is a productive middle ground that better balances padding waste against Tensor Core throughput, and this middle ground has now been mapped.
Follow-Up Research This Work Enables
Extension to the H100 GPU architecture and FP8 Tensor Cores. FlashDecoding++ is evaluated on A100 and RTX3090 (Ampere) and AMD MI210/RX7900XTX, but not on NVIDIA H100 (Hopper), which introduces FP8 Tensor Cores with different granularity ($16 \times 16$ vs. $8 \times 8$ for FP16) and the Transformer Engine for dynamic precision scaling. The flat GEMM optimization's padding-to-8 strategy assumes FP16 Tensor Core granularity; on H100 with FP8, the equivalent minimum M-tile could be 16, changing the padding waste calculation (for batch size 1, padding to 16 wastes 93.75% vs. 87.5% for padding to 8 — potentially making CUDA Core GEMV favorable for a wider range of M). A direct H100 evaluation would measure whether the asynchronized softmax benefits persist (H100's higher memory bandwidth might reduce the relative overhead of synchronization, making the technique less impactful), whether the flat GEMM tiling heuristics need adjustment for H100's larger shared memory (228 KB per SM vs. 164 KB on A100, enabling larger $B_N$ tiles), and whether the heuristic dataflow's inflection points shift beyond the tested ranges. A strong follow-up would report: (a) the three techniques' individual contributions on H100 via ablation, (b) the optimal M-tile size for FP8 flat GEMM (likely 16 rather than 8, but this needs empirical confirmation), and (c) whether the asynchronized softmax's >99.99% coverage holds for FP8 attention logits (the reduced dynamic range of FP8, with max representable value ~448 in E4M3 format, might make overflow more frequent and force a narrower safe $\phi$ range, potentially reducing the technique's applicability).
Cheap difficulty estimation for the compute-optimal framework via lightweight learned predictors. The paper's difficulty estimation via 2048 samples and PRM scoring (Section 3.2) is acknowledged as impractical without cost amortization, yet the entire compute-optimal framework depends on it. A natural and high-impact follow-up is to train a lightweight difficulty predictor that maps question text directly to a difficulty bin, using the 2048-sample PRM score as supervised labels. This could be a small model (e.g., a BERT-base classifier fine-tuned on 12,000 MATH training questions with their oracle difficulty bins) or even a linear probe on top of the base LLM's last hidden state. The key evaluation metrics would be: (a) accuracy of difficulty bin prediction (five-class accuracy on the 500-question MATH test set), (b) whether the compute-optimal strategy selected using predicted bins (from the lightweight classifier) matches the performance of compute-optimal using oracle bins (Figure 4 and 8 curves), and (c) the total compute cost (classifier inference + strategy execution) versus the best-of-N baseline at matched accuracy — i.e., the amortized efficiency gain accounting for difficulty estimation. If a lightweight classifier can predict difficulty with, say, >80% bin accuracy at negligible cost (a single forward pass through a small model), the compute-optimal framework becomes immediately deployable. If the classifier's accuracy is low (e.g., random-chance on hard vs. medium problems), the framework remains impractical until better estimation is available.
Combining PRM tree search with revision-based proposal distributions. The paper studies two complementary axes — PRM-guided search (Section 5) and iterative revisions (Section 6) — but explicitly notes they are never combined (Section 8). The natural follow-up is to use the revision model as the proposal distribution within PRM-guided beam search: at each step of the search tree, instead of sampling continuations from the base few-shot model, sample from the revision model conditioned on the partial solution so far and any previous rejected branches. This could produce higher-quality candidate steps because the revision model is trained to correct mistakes, potentially helping the search escape local optima that the PRM's verifier signal would not distinguish. The key metrics would be: (a) accuracy on difficulty bins 3–4 (medium-hard, where beam search helps but revisions also help) with the combined approach versus either approach alone at matched generation budget, (b) whether the combined approach reduces the over-optimization degradation on bin 1–2 (easy) problems (the revision model might produce more natural solutions that the PRM cannot exploit), and (c) whether the combined approach makes any progress on bin 5 (hardest) problems — the paper's current results show near-zero accuracy for all methods on bin 5; if the combination breaks this ceiling even marginally, it would be an important finding about the complementarity of proposal distribution improvement and verifier-guided search for genuinely out-of-distribution reasoning. A strong study would also report the PRM's scoring distribution on revision-model-generated solutions versus base-model-generated solutions to characterize the distribution shift that motivated training a separate revision ORM (Appendix J, Figure 15a).
Testing the asynchronized softmax on a broader range of models and tasks to map its applicability boundary. The paper demonstrates the technique on Llama2-7B, ChatGLM2-6B (where it works), and OPT-6.7B (where it does not work, and is disabled). This is a three-point sample that doesn't characterize the applicability boundary. A systematic study would profile the attention logit distribution for a diverse set of models — Llama2-13B, Llama2-70B, Mistral-7B, Falcon-40B, CodeLlama, Phi-2, and GPT-3.5/4 if API access permits — measuring the range $[a, b]$ that covers >99.99% of $x_i$, and correlating this range with model architecture characteristics (number of layers, attention head dimension, training data, quantization). The goal is to answer: what model properties predict whether the technique applies? A hypothesis worth testing: models trained with longer contexts (32K+) may have wider attention logit ranges because they need to represent both very small and very large attention scores across the context window; conversely, models with deeper architectures may have narrower ranges because attention is distributed across more layers. A strong study would also measure whether the activation range shifts under distribution shift — e.g., does the range for Llama2-7B change when processing code vs. prose vs. mathematical text? If the range is input-dependent in ways that violate the >99.99% coverage, the technique would need online range estimation rather than a fixed offline $\phi$.
Adaptive, mid-computation difficulty assessment and strategy switching (online allocation). The paper's compute-optimal framework is static: difficulty is estimated once (via 2048 samples), and a fixed strategy is deployed for the entire problem-solving budget. A more sophisticated approach — hinted at in Section 3.2 as the "exploration-exploitation tradeoff" — is online allocation: start with a small number of parallel samples (e.g., 4–8), use the verifier's score distribution on these initial samples as a real-time difficulty signal, and then dynamically allocate the remaining budget based on this signal (e.g., if the verifier's scores are consistently low, switch to broader parallel search; if they are high but varied, switch to sequential revisions to refine; if they are near-perfect, stop early and return the best answer). This would eliminate the upfront difficulty estimation cost entirely by amortizing it into the solution process. The research questions are: (a) how many initial samples are needed to estimate difficulty reliably enough to match the static oracle policy? (b) Does the dynamic policy outperform the static policy at the same total budget (since the initial samples contribute to both difficulty estimation and answer quality)? (c) Can the dynamic policy be formalized as a multi-armed bandit or Bayesian optimization problem with principled regret bounds? A strong study would benchmark the dynamic policy against the static compute-optimal curves in Figures 4 and 8, reporting total generation budget including the initial exploration samples.
Verifier robustness training via adversarial search trajectories. The paper identifies verifier over-optimization as the primary bottleneck preventing unbounded test-time compute scaling (Section 5.3, Figure 3 right, Appendix M). The current PRM is trained on i.i.d. samples from the base model's output distribution, but at test time it scores solutions generated through aggressive beam search — a distribution shift that enables exploitation. A natural follow-up is to train the PRM on search-generated solutions: run beam search with the current PRM, collect solutions that score highly under the PRM but are actually incorrect (false positives), and add these as hard negative examples in the next round of PRM training. This is an adversarial training loop where the search procedure is the adversary attempting to find verifier blind spots, and the PRM is iteratively patched against those blind spots. The key metrics would be: (a) the over-optimization degradation curve (Figure 3 right, bin 1) after each round of adversarial training — does the degradation at high budgets shrink? (b) Whether the adversarially trained PRM generalizes better to unseen search strategies (e.g., trained on beam search negatives, tested on lookahead search), and (c) whether there is a diminishing returns point where additional adversarial training stops improving robustness (analogous to the "reward hacking plateau" in RLHF). A successful result would shift the compute-optimal policy by making aggressive search safe on easier problems, potentially recovering the performance that is currently lost to over-optimization and making the compute-optimal policy converge toward always-using-beam-search rather than routing easy problems to best-of-N.
Practical Applications and Downstream Use Cases
Cost reduction for large-scale LLM API providers. OpenAI's estimated 1.2M/day in saved compute at current pricing. The integration path is straightforward: FlashDecoding++ replaces the attention and GEMM kernel calls in the provider's existing serving framework (vLLM, TensorRT-LLM, or a custom engine). The offline profiling step runs once per GPU type in the fleet, and the resulting dispatch tables are loaded at engine startup. The asynchronized softmax and flat GEMM optimizations require no model changes, no retraining, and no accuracy sacrifice. The provider's risk is limited to the integration engineering cost and validation overhead — the techniques are purely numerical/formulation changes to existing kernel implementations, not architectural overhauls.
Enabling higher-throughput LLM serving on consumer and edge GPUs. The RTX3090 and RX7900XTX results (Figures 10d–f, 12) demonstrate that FlashDecoding++'s techniques work on consumer-grade GPUs with lower memory bandwidth (936 GB/s and ~800 GB/s respectively) and fewer SMs. For organizations deploying LLMs on edge devices or in cost-sensitive environments where datacenter GPUs are uneconomical, the flat GEMM optimization and heuristic dataflow are particularly valuable: the flat GEMM waste from padding to M-tile=64 is more costly on memory-bandwidth-limited GPUs because the wasted memory traffic from loading zero rows consumes a larger fraction of total bandwidth budget. A deployment running Llama2-7B on an RTX3090 at batch size 1 would see the largest relative gains from FlashDecoding++ (since M=1 is where cuBLAS padding waste is maximal and CUDA Core GEMV is most favorable), potentially making the difference between meeting a latency SLO (e.g., <50ms per token for interactive chat) and falling short. The concrete use case is local LLM inference for privacy-sensitive applications (medical, legal, financial) where data cannot leave the device — FlashDecoding++ directly increases the feasible model size or context length on a given consumer GPU budget.
Improving the throughput of batch inference for synthetic data generation. Many LLM applications involve offline batch inference: generating training data for smaller models (distillation), producing evaluation datasets, or running inference over large document collections. In these settings, throughput (tokens per second per GPU) is the sole metric of interest — latency is secondary because the workload is not user-facing. FlashDecoding++'s speedups on the decode phase (where batch sizes of 4–8 are practical and the flat GEMM optimization with heuristic dispatch provides the largest gains) translate directly to higher throughput. For a synthetic data generation pipeline processing millions of prompts, a 1.37× speedup versus FlashDecoding means the pipeline completes in ~73% of the time, reducing GPU rental costs proportionally. The offline profiling step is a one-time cost that amortizes over the entire batch job. Combined with system-level optimizations like vLLM's PagedAttention (for memory efficiency during large-batch generation), FlashDecoding++ provides the kernel-level throughput layer beneath the system-level batch management.
Faster LLM-powered code generation and interactive coding assistants. In code generation (e.g., GitHub Copilot, CodeLlama-based tools), the decode phase dominates: a single prompt ("write a function that...") triggers the generation of tens to hundreds of lines of code autoregressively. The flat GEMM optimization directly accelerates each decode step, and the asynchronized softmax reduces attention overhead as the generated sequence grows and the KVcache accumulates. For a coding assistant generating an average of 200 tokens per response, a 1.21× decode speedup (the average over FlashDecoding) reduces per-response latency by ~17%, which in an interactive IDE setting translates to a shorter "thinking" indicator and a more fluid developer experience. The heuristic dataflow's ability to select CUDA Core GEMV for the K/Q/V projections at batch size 1 (the typical serving batch size for interactive assistants) while using Tensor Cores for the larger FFN operations means the engine extracts near-optimal performance from each operation type without manual per-operation kernel selection — this is particularly valuable for code models where hidden dimensions vary (e.g., CodeLlama-34B vs. 7B) and manual tuning per model would be labor-intensive.
When to Prefer This Method
FlashDecoding++ does not position itself against a single named alternative in a binary "use A or use B" decision — it is a set of kernel-level optimizations that can be integrated into any existing LLM inference engine (vLLM, TensorRT-LLM, DeepSpeed, FlashDecoding, etc.). The paper measures speedups against these engines individually, and the takeaway is that integrating FlashDecoding++'s kernels improves each of them. There is no scenario described in the paper where using FlashDecoding++ harms performance relative to the baseline engine it replaces, assuming the model's attention logit range supports the asynchronized softmax (Section 3) and the GPU has been profiled offline (Section 5). The decision is therefore not "FlashDecoding++ versus alternative X" but rather "integrate FlashDecoding++ into your serving stack versus keep your current kernel implementations." The conditions for preferring the integration are:
-
The model is not OPT-6.7B or a model with similarly wide attention logit ranges (Figure 5): the asynchronized softmax, one of the three techniques, is explicitly disabled for such models. The remaining two techniques (flat GEMM optimization and heuristic dataflow) still provide speedups, as demonstrated by the OPT-6.7B results in Figures 10b and 10e, but the total gain is proportionally smaller. The paper provides no diagnostic for determining whether a given model's activation range is "narrow enough" beyond profiling its
$x_i$distribution — practitioners adopting FlashDecoding++ would need to run this profiling themselves or rely on the paper's claim that "mainstream LLMs" exhibit narrow ranges. -
The deployment hardware is an NVIDIA GPU (A100, RTX3090, or similar) or an AMD GPU (MI210, RX7900XTX, or similar). The paper validates on four GPUs; on untested hardware, the heuristic dataflow's offline profiling procedure should still work, but the flat GEMM optimization's tiling heuristics (Section 4) assume shared memory sizes and SM counts comparable to the tested GPUs. On very different architectures (e.g., inference ASICs, mobile GPUs with tiny shared memory), the approach would need revalidation.
-
The workload includes the decode phase. The flat GEMM optimization is decode-specific, and the heuristic dataflow's biggest wins come from switching from Tensor Core GEMV to CUDA Core GEMV at M=1 (decode batch size 1). For pure prefill-only workloads (e.g., embedding extraction, classification over long documents), the speedups are modest (1.05–1.09× over optimized baselines, from Figure 11), and the integration effort may not justify the gain.
-
Batch sizes are small to moderate (1–8 for decode). At large batch sizes (e.g., 32–64 for offline batch inference), the M-dimension is large enough that the standard CUTLASS/cuBLAS implementations with M-tile=64 are already optimal — the heuristic dataflow converges to ImplC for all operations, and the flat GEMM optimization is never invoked. FlashDecoding++ in this regime reduces to "asynchronized softmax plus cuBLAS," and the gain is only from the softmax optimization (18.8% of attention time, but attention is a small fraction of total time at large batch sizes where linear projections dominate). Practitioners serving at consistently large batch sizes should measure whether the benefit justifies the integration effort on their specific model.
-
The integration effort is acceptable. FlashDecoding++ is a kernel-level engine that requires replacing the GPU kernel calls in an existing inference framework. This is not a drop-in Python package — it requires C++/CUDA compilation, offline profiling on the target hardware, and integration with the framework's kernel dispatch layer. For organizations already maintaining a custom inference engine (or contributing to open-source engines like vLLM or TensorRT-LLM), this is feasible. For teams using a managed inference service (e.g., OpenAI API, Hugging Face Inference Endpoints), FlashDecoding++ is not directly applicable — the optimization must be adopted by the service provider.
In summary, FlashDecoding++ is preferable when the deployment stack already involves GPU kernel management (C++/CUDA-level control), the model's attention logits are empirically bounded (as profiled via Figure 5's methodology), the workload includes substantial decode-phase serving at small batch sizes, and the target hardware is a modern datacenter or high-end consumer GPU where the offline profiling procedure can characterize the performance surface. Under these conditions, the paper provides strong evidence of non-trivial, composable, accuracy-preserving speedups. Under other conditions — wide-activation-range models, pure prefill workloads, very large decode batches, or managed serving platforms — the benefits are either reduced or inaccessible without upstream adoption.
