Naive Attention Kernel

Before understanding FlashAttention, you must understand what it fixes. Standard attention materializes the full $S \times S$ attention matrix in HBM — creating an $O(S^2)$ memory bottleneck that makes long-context serving impossible without algorithmic innovation.

---

1. Standard Attention: Step by Step

For a single head with query $Q$, key $K$, value $V$ all of shape $(S, d_k)$:

Step 1: S = Q × K^T              → (S, dk) × (dk, S) = (S, S)    [GEMM, write to HBM]
Step 2: S = S / sqrt(dk)          → (S, S)                         [element-wise, read+write HBM]
Step 3: S = S + mask               → (S, S)                         [element-wise, read+write HBM]
Step 4: P = softmax(S, dim=-1)    → (S, S)                         [reduction + element-wise, read+write HBM]
Step 5: O = P × V                 → (S, S) × (S, dk) = (S, dk)    [GEMM, read S×S from HBM]

The Memory Problem

The attention matrix $S$ has $S^2$ elements. For multi-head attention with batch:

$$\text{Attention matrix size} = B \times N_h \times S^2 \times \text{bytes}$$

Sequence length	Attention matrix (B=1, Nh=32, FP16)	vs H100 memory
S = 2,048	256 MB	Fits easily
S = 8,192	4 GB	~5% of HBM
S = 32,768	64 GB	80% of HBM
S = 131,072	1 TB	Doesn’t fit

At 128K context, the attention matrix alone exceeds GPU memory. This is why naive attention cannot serve long-context models.

---

2. HBM Traffic Analysis

Count every byte read from and written to HBM:

Step	Operation	HBM reads	HBM writes
QK^T	GEMM	$2 \times S \times d_k$ (Q and K)	$S^2$ (scores)
Scale	Element-wise	$S^2$	$S^2$
Mask	Element-wise	$S^2$ (scores) + $S^2$ (mask)	$S^2$
Softmax	Reduction + normalize	$S^2$ (3 passes)	$S^2$
PV	GEMM	$S^2$ (P) + $S \times d_k$ (V)	$S \times d_k$ (O)

Total HBM traffic: ~$7S^2 + 4Sd_k$ bytes (FP16). The $S^2$ terms dominate.

For S=4096, dk=128: $7 \times 4096^2 \times 2 + 4 \times 4096 \times 128 \times 2 \approx 236$ MB per head. With 32 heads: 7.5 GB of HBM traffic just for attention.

Arithmetic Intensity

$$\text{FLOPs} = 2S^2 d_k + 2S^2 d_k + O(S^2) \approx 4S^2 d_k$$ $$\text{AI} = \frac{4S^2 d_k}{14 S^2} \approx \frac{4 \times 128}{14} \approx 37$$

With H100’s ops:byte ratio of 295, naive attention is severely memory-bound (AI = 37 << 295). The GPU spends most of its time reading and writing the attention matrix, not computing.

---

3. Naive CUDA Implementation

// Step 1: QK^T
// Use cuBLAS: cublasSgemm(S, S, dk, Q, K^T, scores)

// Step 2-4: Scale + Mask + Softmax (fused for fewer HBM trips)
__global__ void scaleMaskSoftmax(float* scores, float* mask, int S, float scale) {
    int row = blockIdx.x;
    int tid = threadIdx.x;

    // Scale and mask
    for (int j = tid; j < S; j += blockDim.x) {
        scores[row * S + j] = scores[row * S + j] * scale + mask[row * S + j];
    }
    __syncthreads();

    // Online softmax (single pass for max + sum)
    float m = -INFINITY, s = 0.0f;
    for (int j = tid; j < S; j += blockDim.x) {
        float val = scores[row * S + j];
        float new_m = fmaxf(m, val);
        s = s * expf(m - new_m) + expf(val - new_m);
        m = new_m;
    }
    // Reduce m, s across block...
    __syncthreads();

    // Normalize
    for (int j = tid; j < S; j += blockDim.x) {
        scores[row * S + j] = expf(scores[row * S + j] - m) / s;
    }
}

// Step 5: PV
// Use cuBLAS: cublasSgemm(S, dk, S, P, V, output)

Even with fused scale+mask+softmax, we still:

1. Materialize the S×S matrix in HBM after QK^T
2. Read it back for softmax
3. Read it again for PV multiplication

The S×S matrix is written once and read twice — 3 complete passes over an $O(S^2)$ tensor.

---

4. Why This Can’t Be Fixed Without Algorithmic Change

You might think: “just fuse everything into one kernel.” But:

• QK^T is a GEMM that produces the $S \times S$ output
• Softmax requires the entire row to compute (needs global max and sum)
• PV is a GEMM that reads the $S \times S$ attention weights

You can’t fuse two GEMMs with a softmax between them using standard approaches — the softmax needs the complete row before PV can start.

FlashAttention’s insight: Use tiling to break this dependency. Process Q in tiles, compute partial attention scores, use online softmax to maintain running statistics, and accumulate the output — all without ever writing the full $S \times S$ matrix.

This is covered in the next two chapters.

---

5. When Naive Attention Is Fine

Despite its limitations, naive attention is used when:

Scenario	Why naive is OK
Short sequences (S < 2K)	S×S matrix is small (< 32 MB), fits easily
Profiling/debugging	Easier to understand, debug, and verify correctness
Non-LLM models	Vision (ViT) with S=197 patches — S² is tiny
Reference implementation	Gold standard for checking FlashAttention correctness

For LLM inference with S > 4K, FlashAttention is mandatory.

---

Self-Check

• Calculate the attention matrix size for given B, Nh, S in FP16
• Count the total HBM reads and writes for standard 5-step attention
• Calculate arithmetic intensity and explain why naive attention is memory-bound
• Explain why you can’t simply “fuse” QK^T + softmax + PV into one kernel
• Explain when naive attention is acceptable vs when FlashAttention is required

---

References

• Attention Is All You Need (Vaswani et al., 2017)
• Self-Attention Does Not Need O(n²) Memory (Rabe & Staats, 2021)
• FlashAttention (Dao et al., 2022) — motivating analysis