Why LLM Prefill Costs Scale Quadratically — Even After FlashAttention

FlashAttention enabled and a 32K-context prefill still runs hundreds of times slower than a 512-token one — that's disorienting the first time you see it. The natural reaction is "I thought FlashAttention fixed the memory problem." But what FlashAttention reduces is HBM traffic; what it cannot reduce is the number of matrix dot products. These are independent bottlenecks, which is why IO optimization alone cannot fully resolve the latency problem of long-context prefill.

Where Does O(n²) Come From?

Self-attention cost breaks down along two axes: FLOPs and HBM traffic.

For sequence length n, head dimension d, and number of heads h, the QKᵀ matmul requires 2n²d FLOPs. After the softmax, the AV matmul costs another 2n²d. The total is 4n²dh FLOPs — O(n²) in n.

HBM traffic is a separate calculation. A standard attention implementation writes the full n×n matrix S = QKᵀ to HBM, then reads it back for the softmax. In FP16, reads and writes combined come to 4n²h bytes. At large n, this term dominates the linear-in-n cost of loading Q, K, and V.

Combining the two bottlenecks into a single metric gives Arithmetic Intensity (AI):

AI = FLOPs / bytes = 2n²d / 4n² = d/2   (FP16, QKᵀ kernel)

With d=128, AI = 64 ops/byte. The A100 SXM roofline is 312 TFLOPS ÷ 2000 GB/s = 156 ops/byte. Since 64 < 156, the QKᵀ step is IO-bound — the standard attention implementation is bottlenecked by HBM reads and writes. That is the bottleneck FlashAttention targets.

What FlashAttention Does and Doesn't Reduce

FlashAttention-2 (Dao et al., 2023) never materializes the S matrix to HBM. It tiles the computation into SRAM-sized blocks, completes the softmax on-chip, and only writes the final output O back to HBM. As a result, HBM traffic drops from O(n²) to O(n·d) — linear in n.

FLOPs are unchanged. Tiling QKᵀ into blocks does not alter the total number of dot products.

Once the IO bottleneck is removed, the GPU is held back by raw compute. The effective AI rises after FlashAttention is applied, which means the GPU can enter a compute-bound regime:

prefill compute time ≈ 4n²dh / peak_TFLOPS
prefill IO time      ≈ n × d × h × 4 bytes / bandwidth   (Q,K,V + O load/store)

As n grows, compute time scales as n² while IO time scales linearly in n. Beyond some crossover point, compute time dominates, and in that regime no further IO reduction from FlashAttention helps — the FLOPs themselves are the bottleneck.

Measurements on an A100 80GB with Llama-2-7B and batch size 1:

With FlashAttention ON, the 512→32768 range is still about 640× slower. The gap relative to the theoretical n² scaling (64×) is explained by the GPU shifting into compute-bound territory above 16K, with KV cache memory-management overhead added on top. The speedup ratio growing from 1.6× to 3.7× at longer sequences reflects how severe IO waste becomes in the OFF state — not that FlashAttention becomes more effective at longer sequences.

What Chunked Prefill Actually Does

Sarathi-Serve (Agrawal et al., 2024, USENIX OSDI) proposes Chunked Prefill, which does not change the total O(n²) compute. Splitting a long prefill request into chunks of size c and processing them over n/c steps reduces per-step compute to O(c·n), but the total remains O(n²).

The benefit is in scheduling. In a standard scheduler, decode requests wait until prefill completes. A 32K prefill taking 11 seconds stalls the entire decode queue. Chunked Prefill inserts decode requests into the batch at every chunk boundary, which raises GPU utilization and reduces decode queuing.

To enable it in vLLM:

vllm serve meta-llama/Llama-2-7b-hf \
  --enable-chunked-prefill \
  --max-num-batched-tokens 512

The default chunk size is 512 tokens. In vLLM V1 it is enabled by default when applicable. Smaller chunks reduce decode wait time and improve TPOT (time per output token), but at each chunk boundary there is overhead from partial KV cache loading and the scheduler, which can raise p50 TTFT for individual requests. The Sarathi-Serve paper identifies chunk sizes in the 256–512 range as the sweet spot between TTFT p95 and throughput, reporting a 2.6× serving capacity improvement over vLLM on a single A100 with Mistral-7B.

Documentation around Chunked Prefill often says it "reduces TTFT" — more precisely, it suppresses tail latency at p95 and above. Treating it as a lever for lowering p50 TTFT and deploying it with that expectation will produce results where p50 TTFT gets worse, not better.

Ring Attention Partitions the O(n²) Work

Handling 128K- or 1M-token contexts that don't fit on a single GPU requires distributing the sequence dimension across multiple GPUs. Ring Attention (Liu et al., 2023) assigns n/k tokens to each of k GPUs, reducing each GPU's attention compute to O((n/k)²). Doubling the GPU count theoretically cuts per-GPU compute to one quarter.

The catch is that each GPU must attend over K and V blocks from other GPUs, making cross-device communication unavoidable. For Llama-2-7B (d=128, h=32) with n=32K and k=2, the KV blocks that must be exchanged total roughly 512 MB. Over NVLink (A100 bidirectional ~600 GB/s) the transfer takes ~1.7ms; over PCIe Gen4 (unidirectional 32 GB/s) it is ~16ms. Since a 32K prefill with FlashAttention ON takes ~11.5 seconds, the communication is completely hidden by compute under NVLink. Under PCIe, the sequence length needs to reach 64K–128K before distributing across GPUs beats a single GPU.

The Ring Attention paper notes that communication can be pipeline-overlapped with compute, but the practical effectiveness of that overlap depends on the hardware topology and batch size. This is why adding more GPUs does not linearly accelerate prefill.

A Single Framework for the Three Strategies

In practice the order is roughly as follows. FlashAttention should always be on — it gives clear gains in the IO-bound regime and causes no harm in the compute-bound regime. Chunked Prefill is added for mixed prefill/decode workloads to improve TPOT and suppress tail latency. Ring Attention is considered only for ultra-long contexts that exceed single-GPU memory and compute limits, and only when NVLink interconnects are available.

Because prefill cost scales as n², the maximum context length is directly tied to the TTFT p95 SLO. Even if the average context is 4K, a p99 request at 32K stalls the entire decode queue for as long as that request takes. Chunked Prefill softens the impact, but the right first step is setting max_model_len based on the SLO budget.