Tiny VoltKernel visualization: Flash Decoding
Generating text with a transformer happens in two stages. First comes prefill, where the whole prompt goes through at once, so attention has hundreds or thousands of query rows to work on. Then comes decode, where tokens come out one at a time and each step has exactly one query row per head.
Flash Attention divides its work across query rows. That fits prefill well and decode badly, because in decode there is only one row to hand out. One thread block ends up walking the entire KV cache by itself while the rest of the GPU waits, and since every generated token reads that whole cache, this is where decode spends its time.
Flash decoding keeps the same math and changes what gets divided up: the K/V sequence instead of the query rows.
Standard FlashAttention parallelizes over query tiles (the outer loop). At decode time there is only one query row per (batch, head), so that outer axis is width-1 — most of the GPU sits idle while a single thread block streams through the entire sequence.
Flash decoding fixes this by parallelizing over the sequence itself. Split along the sequence dimension into chunks and give each chunk to its own thread block. Each block runs FlashAttention's online softmax over its own slice and writes a partial . A second combine kernel merges them.
For a single query row , per split covering indices in :
Each split stores unnormalized (no per-split division) and the combine kernel applies a single global at the end. The per-split weight rescales each unnormalized partial so its max aligns with — the same rescaling FlashAttention already does across inner tiles, lifted one level up across splits.
Visualizing the steps in kernel
There are two kernels. The first launches one thread block per batch, head and split, so the splits run side by side; the second merges what they produce.
template<int Bc, int D, int S> // S = split size along K/V
__global__ void fd_split(const float* Q, const float* K, const float* V,
float* O_partial, float* m_partial, float* l_partial,
int N, float scale, int n_splits) {
int s = blockIdx.x, bh = blockIdx.y, tid = threadIdx.x;
Q += bh*D; // 1 query row per head
K += bh*N*D; V += bh*N*D;
O_partial += (bh*n_splits + s)*D;
m_partial += bh*n_splits + s;
l_partial += bh*n_splits + s;
int n_start = s*S, n_end = min(n_start + S, N);
__shared__ float Qi[D], Kj[Bc][D], Vj[Bc][D], Sij[Bc];
__shared__ float mi, li, Oi[D];
}
Visualizing operations in a single block (split s = 0, K/V range [0, S))
This block sees only its own slice of K and V, and runs the ordinary online softmax over it, so it never learns anything about the rest of the sequence.
Qi = Q; Oi = 0; mi = -inf; li = 0:
// s = 0
for (int d = tid; d < D; d += blockDim.x) Qi[d] = Q[d];
for (int d = tid; d < D; d += blockDim.x) Oi[d] = 0.f;
if (tid == 0) { mi = -INFINITY; li = 0.f; }
__syncthreads();
0.0 Inner tile j = n_start
Kj, Vj = K[j], V[j]:
// s = 0, j = n_start (= 0)
int j = n_start;
int cols = min(Bc, n_end - j);
for (int x = tid; x < cols*D; x += blockDim.x) {
Kj[x/D][x%D] = K[(j + x/D)*D + x%D];
Vj[x/D][x%D] = V[(j + x/D)*D + x%D];
}
__syncthreads();
Sij = Qi @ Kj.T:
// s = 0, j = 0
for (int c = tid; c < cols; c += blockDim.x) {
float s_ = 0.f;
for (int k = 0; k < D; k++) s_ += Qi[k]*Kj[c][k];
Sij[c] = s_*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = alpha * li + exp(Sij - mi).sum(-1)Oi = alpha * Oi + softmax(Sij) @ Vj:
// s = 0, j = 0 — one thread does the row-wise online softmax update
if (tid == 0) {
float m_ij = -INFINITY;
for (int c = 0; c < cols; c++) m_ij = fmaxf(m_ij, Sij[c]);
float m_new = fmaxf(mi, m_ij);
float alpha = expf(mi - m_new);
float lsum = 0.f;
for (int c = 0; c < cols; c++) { Sij[c] = expf(Sij[c] - m_new); lsum += Sij[c]; }
li = alpha*li + lsum;
mi = m_new;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < cols; c++) pv += Sij[c]*Vj[c][d];
Oi[d] = alpha*Oi[d] + pv;
}
}
__syncthreads();
0.1 Inner tile j = n_start + Bc
Kj, Vj = K[j], V[j]:
// s = 0, j = n_start + Bc
int j = n_start + Bc;
int cols = min(Bc, n_end - j);
for (int x = tid; x < cols*D; x += blockDim.x) {
Kj[x/D][x%D] = K[(j + x/D)*D + x%D];
Vj[x/D][x%D] = V[(j + x/D)*D + x%D];
}
__syncthreads();
Sij = Qi @ Kj.T:
// s = 0, j = 1
for (int c = tid; c < cols; c += blockDim.x) {
float s_ = 0.f;
for (int k = 0; k < D; k++) s_ += Qi[k]*Kj[c][k];
Sij[c] = s_*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = alpha * li + exp(Sij - mi).sum(-1)Oi = alpha * Oi + softmax(Sij) @ Vj:
// s = 0, j = 1
if (tid == 0) {
float m_ij = -INFINITY;
for (int c = 0; c < cols; c++) m_ij = fmaxf(m_ij, Sij[c]);
float m_new = fmaxf(mi, m_ij);
float alpha = expf(mi - m_new);
float lsum = 0.f;
for (int c = 0; c < cols; c++) { Sij[c] = expf(Sij[c] - m_new); lsum += Sij[c]; }
li = alpha*li + lsum;
mi = m_new;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < cols; c++) pv += Sij[c]*Vj[c][d];
Oi[d] = alpha*Oi[d] + pv;
}
}
__syncthreads();
0. Write partial
Nothing is divided here. Oi is still a running total, and mi and li are what the next kernel needs to finish it, so all three are written out.
O_partial[0], l_partial[0], m_partial[0] = Oi, li, mi:
// s = 0
for (int d = tid; d < D; d += blockDim.x) O_partial[d] = Oi[d];
if (tid == 0) { *m_partial = mi; *l_partial = li; }
Same operations in another block (split s = 1, K/V range [S, 2S), runs in parallel with split 0)
| Operation | Action |
|---|---|
Qi = Q; Oi = 0; mi = -inf; li = 0 |
|
Kj, Vj = K[j], V[j] (j = n_start) |
|
Sij = Qi @ Kj.T |
|
mi, li, Oi = ... |
|
Kj, Vj = K[j], V[j] (j = n_start + Bc) |
|
Sij = Qi @ Kj.T |
|
mi, li, Oi = ... |
|
O_partial[1], l_partial[1], m_partial[1] = Oi, li, mi |
Split 1 runs the identical kernel body as split 0 on its own SM, differing only in blockIdx.x = 1, so n_start = S.
Combine kernel — one thread block per (batch, head)
After both split blocks finish, a second kernel merges their partials.
In the steps below, the per-split m_partial, l_partial, and O_partial tiles are shown sliding from HBM into the combine block, which represents a real global-memory read. But the combine kernel does not store these tiles in shared memory: they flow through registers, and only the reduced scalars m_global, l_global (and the acc output) ever live in __shared__.
m_global = max_s m_partial[s]:
template<int D>
__global__ void fd_combine(const float* O_partial, const float* m_partial,
const float* l_partial, float* O, int n_splits) {
int bh = blockIdx.x, tid = threadIdx.x;
O_partial += bh*n_splits*D;
m_partial += bh*n_splits;
l_partial += bh*n_splits;
O += bh*D;
__shared__ float m_global, l_global;
if (tid == 0) {
float m_g = -INFINITY;
for (int s = 0; s < n_splits; s++) m_g = fmaxf(m_g, m_partial[s]);
m_global = m_g;
}
__syncthreads();
l_global = sum_s exp(m_s - m_global) * l_s:
if (tid == 0) {
float l_g = 0.f;
for (int s = 0; s < n_splits; s++)
l_g += expf(m_partial[s] - m_global) * l_partial[s];
l_global = l_g;
}
__syncthreads();
acc = sum_s exp(m_s - m_global) * O_partial[s]:
__shared__ float acc[D];
for (int d = tid; d < D; d += blockDim.x) {
float a = 0.f;
for (int s = 0; s < n_splits; s++) {
float w = expf(m_partial[s] - m_global);
a += w * O_partial[s*D + d];
}
acc[d] = a;
}
__syncthreads();
O = acc / l_global(write back to HBM):
for (int d = tid; d < D; d += blockDim.x) {
O[d] = acc[d] / l_global;
}
}
In Flash Attention, one thread block read the whole KV cache by itself. In flash decoding, a few dozen blocks each read a different chunk of it. The chunks do not overlap and together they cover the sequence, so the number of bytes read from HBM is exactly the same in both cases.
The gain is in how fast those bytes arrive. A single block cannot ask for data quickly enough to use all of the GPU's memory bandwidth, so most of it goes unused while that block waits. A few dozen blocks asking at once can use it, and the same read finishes sooner.