Split the KV, Not the Query

At autoregressive decode time with batch size 1 and a very long KV cache, a plain FlashAttention kernel leaves the GPU badly underutilized even though the same kernel is compute-efficient during prefill. Flash-Decoding fixes this by splitting the KV sequence dimension into chunks processed in parallel, then combining the partial results using log-sum-exp statistics. Why does splitting specifically along the KV dimension solve the utilization problem?

It converts the memory-bound decode step into a compute-bound one by routing the reduction through tensor cores instead of CUDA cores.

It removes the softmax normalization step entirely, replacing exp() with a cheaper linear approximation to cut per-token latency.

At single-token decode the query dimension has size 1, so there's no parallelism to exploit there; splitting along KV length instead creates enough independent tiles to fill the GPU's streaming multiprocessors, even though the underlying workload stays memory-bandwidth-bound rather than becoming compute-bound.

Splitting along KV length reduces the total number of floating-point operations required for the attention computation, directly speeding up what is fundamentally a compute-bound step.

It avoids reading the KV cache from HBM more than once per generated token by fusing every chunk into a single on-chip tile that fits in shared memory regardless of context length.