FLASH ATTENTION & PAGED ATTENTION

Picture this.You just rented an NVIDIA A100. It costs three dollars an hour. It has 80 gigabytes of memory and can do nineteen trillion math operations per second. You feel unstoppable.

You feed it a prompt: “Summarize the entire history of the Roman Empire in one paragraph.” The model starts strong. “The Roman Empire, one of history’s greatest civilizations…” Then, at token 3,000, something terrible happens.

The fan screams. The memory bar turns red. The generation slows to a crawl. What should take milliseconds now takes seconds. Then, at token 8,000, it dies.

Out of memory.

You stare at the screen. How? You have 80 GB. The model weights are only 13 GB. Where did the other 67 gigabytes go?

The answer lies in a dirty secret that almost killed the entire LLM revolution. And the only thing standing between us and that death was two clever, borderline-insane engineering hacks.

This is the story of Flash Attention and Paged Attention.

The O(n²) Monster

Let’s say you’re running a transformer model. You type a sentence: “The cat sat on the mat.” Seven words.

The model doesn’t read this like you do. It converts every word into a vector — a list of numbers, say 128 numbers per word. That’s your Query (Q), Key (K), and Value (V).

For attention, the model needs to compare every word with every other word.

Word 1 looks at Word 1, Word 2, Word 3… all the way to Word 7. Word 2 looks at Word 1, Word 2, Word 3… all the way to Word 7.

For 7 words, that’s 7 × 7 = 49 comparisons. No big deal.

But now imagine you’re processing a legal contract. 4,096 tokens. That’s 4,096 × 4,096 = 16,777,216 comparisons.

Still fine. The GPU laughs at this.

Now imagine you’re processing a book. 32,768 tokens. That’s 1,073,741,824 comparisons. Over one billion.

The GPU stops laughing.

The Dirty Secret

Here’s where it gets sinister.

You would think the problem is the math. One billion multiplications sounds scary. But remember: this GPU does nineteen trillion operations per second. One billion is a drop in the bucket. The math isn’t the villain.

The villain is memory.

Let me show you exactly what standard attention does. Step by step. Watch closely, because this is where the crime happens.

Step 1: The GPU loads the Query matrix and Key matrix from memory. It multiplies them. This creates a giant score matrix — let’s call it S. For 4,096 tokens, S is a 4,096 × 4,096 grid of numbers. That’s 16.7 million floats. At 4 bytes each, that’s 67 megabytes.

and this is the murder weapon it writes S back to memory.

Step 2: It loads S back from memory, runs softmax on it, and creates another giant matrix P. P is also 4,096 × 4,096. Another 67 megabytes written to memory.

Step 3: It loads P and the Value matrix, multiplies them, and finally gets the output.

So for one attention layer, for one sequence, the GPU wrote 134 megabytes of temporary data to memory. Data it didn’t even need to keep.

Now stack that 32 times (for a 32-layer model). And do it for a batch of 16 sequences at once.

The memory bandwidth the speed at which the GPU can read and write to its own RAM chokes. It’s like trying to drain a swimming pool through a drinking straw. The compute cores sit idle, twiddling their thumbs, waiting for data to arrive.

The GPU is bored and the memory is dying.

This is the bottleneck that almost ended the long-context revolution. You could buy the most expensive GPU on the planet, and it would still crawl on long documents because it was spending 90% of its time moving data around instead of doing math.

Flash Attention — The Heist

The Mastermind

In 2022, a team at Stanford looked at this mess and asked a dangerous question:

What if we just… never wrote the score matrix to memory?

Sounds impossible, right? Softmax needs the entire row. You need all the scores to compute the probabilities. How can you compute softmax without having all the numbers?

The answer is a mathematical trick called online softmax. And it’s the kind of idea that makes you angry you didn’t think of it first.

The Heist Plan

Imagine you’re a bank robber. Your crew needs to count all the money in a vault, find the biggest stack, and divide everything proportionally. But the vault is huge, and you only have a tiny bag.

Standard attention says: “Drag every stack out of the vault, line them up in the parking lot, then do the math.”

Flash Attention says: “Bring a small stack to your bag, update your running totals, put it back, and never let the parking lot see a single bill.”

Here’s how the math actually works.

Softmax for a row looks like this:

​Where m is the maximum value in the row. You need the max to keep the exponentials from exploding.

The trick: you don’t need the global max all at once. You can process chunks.

Chunk 1: You see scores [2.0, 3.0, 1.0]. Local max = 3.0. Local sum = e^(2−3)+e^(3−3)+e^(1−3) = 0.368 + 1.0 + 0.135 = 1.503.

Chunk 2: You see scores [5.0, 2.0]. Oh no. New max = 5.0. The old calculations used the wrong max.

But here’s the beautiful part: you can rescale everything.

New global max = 5.0. Old sum needs to be rescaled by e^(3.0−5.0)=e^(−2)=0.135 . New sum = (1.503 × 0.135) + e^(5−5) + e^(2−5) = 0.203 + 1.0 + 0.050 = 1.253.

You keep two running numbers: the current max, and the current rescaled sum. Every time you see a bigger number, you rescale your accumulator. By the end, you have the exact same result as if you’d seen everything at once.

The Real Workflow (Step-by-Step)

Let’s trace through Flash Attention with real dimensions. You’re processing a sequence of 4,096 tokens. Head dimension is 64. Your GPU’s fast on-chip memory (SRAM) can hold a 64 × 64 tile comfortably.

The Setup:

• Q, K, V are each 4096 × 64 matrices.
• SRAM can hold blocks of size, say, 64 × 64.
• We tile Q into 64 blocks of size 64 × 64.
• We tile K and V into 64 blocks of size 64 × 64.

The Loop:

For each block of Q (say, Q_tile covering tokens 0-63):
    Initialize: running_max = -infinity, running_sum = 0, running_output = 0

    For each block of K and V (K_tile, V_tile covering tokens 0-63, then 64-127, etc.):
        1. Load Q_tile into SRAM
        2. Load K_tile into SRAM
        3. Compute S_partial = Q_tile @ K_tile^T  (64 × 64 block of scores)
        4. Find local_max in S_partial
        5. Compute local_sum = sum(exp(S_partial - local_max))
        6. RESCALE:
           - If local_max > running_max:
               running_sum = running_sum * exp(running_max - local_max) + local_sum
               running_max = local_max
           - Else:
               local_sum = local_sum * exp(local_max - running_max)
               running_sum = running_sum + local_sum
        7. Compute partial output: O_partial = exp(S_partial - running_max) @ V_tile
        8. Accumulate into running_output, rescaling previous values if needed
        9. Drop S_partial from memory. It never touches HBM.

    10. Finalize: O_tile = running_output / running_sum
    11. Write only O_tile (64 × 64) to HBM

The Result:

Standard Attention memory traffic: Read Q, K, V. Write S (409⁶²). Write P (409⁶²). Write O. Total: roughly 134 MB of writes for the attention matrix alone.

Flash Attention memory traffic: Read Q, K, V. Write O. Zero writes for S and P. Total: roughly 2 MB of writes.

That’s a 67× reduction in memory traffic.

Why It Feels Like Magic

Your original intuition was: “If a task takes 5 minutes, doing it in chunks should still take 5 minutes.”

You’re right about the math. The multiplications still happen. But imagine the task isn’t doing math it’s carrying boxes between two rooms.

Standard attention carries 134 boxes to the slow room (HBM) and back. Flash Attention carries 2 boxes.

The actual “work” the math happens in the fast room (SRAM) where carrying is free. Flash Attention doesn’t reduce work. It reduces carrying.

Real Numbers

On an A100 GPU:

• HBM bandwidth: ~2 TB/s
• SRAM bandwidth: ~19 TB/s
• Math throughput: 312 TFLOPS

For a 4096×4096 attention with head dim 64:

• Math required: ~2 billion FLOPs
• Time if math-bound: 0.006 milliseconds
• Time if memory-bound (standard): ~0.1 milliseconds
• Time with Flash Attention: ~0.008 milliseconds

Flash Attention brings you back to being limited by math, not memory. That’s the win.

Part III: Paged Attention — The Landlord

A Different Crisis

Flash Attention solved the speed problem. But there was another monster lurking: wasted space.

Imagine you start a company. You rent an office building. But here’s the lease: you must pay for 100 desks, even if you only have 3 employees. And if your friend wants to start a company too, they must rent a separate 100-desk building. Even if you both do the exact same work for the first month.

That’s how standard LLM inference worked.

When you send a prompt to a model, the GPU pre-allocates a giant contiguous block of memory for the KV cache. If the model supports up to 8,192 tokens, it reserves space for 8,192 keys and 8,192 values even if your prompt is only 10 words long.

For one user, this is annoying. For a chatbot serving thousands of users, it’s catastrophic. You have 80 GB of memory, but 70 GB of it is empty, reserved desks.

Worse: if User A generates 50 tokens and User B generates 51 tokens, their reserved blocks are different sizes. You can’t shuffle them together efficiently. The memory becomes Swiss cheese full of holes.

Enter the Landlord

The team behind vLLM looked at this and asked: What if we stopped pre-renting buildings and started renting by the desk?

They stole an idea from your computer’s operating system: virtual memory paging.

Instead of one giant block, you cut the KV cache into small, fixed-size chunks called pages. Maybe 16 tokens per page. You keep a giant pool of free pages in GPU memory. When a sequence needs space, you hand it a page. When it needs more, you hand it another. The pages don’t have to be next to each other. They can be scattered anywhere.

But the model expects a continuous stream of K and V vectors. How do you feed it a jigsaw puzzle?

You lie.

You maintain a Block Table a simple lookup list that says:

• Token 0–15 → Physical Page #7
• Token 16–31 → Physical Page #42
• Token 32–47 → Physical Page #103

The model thinks it’s looking at one long, beautiful, continuous memory block. But under the hood, the GPU kernel is bouncing around physical memory like a pinball, grabbing pages from wherever they live.

The Real Workflow: A Chatbot Serving Three Users

Let’s walk through a real scenario. You’re running a customer service LLM. Three users show up at the same time.

System Setup:

• Page size: 16 tokens
• GPU memory: 80 GB
• Free page pool: 50,000 pages available

User 1 arrives. Prompt: “How do I reset my password?” (6 tokens)

Engine Action:
1. Allocate 1 physical page from free pool → Page #100
2. Compute K, V for all 6 prompt tokens
3. Write K[0:6], V[0:6] into Page #100 (slots 0-5 used, slots 6-15 empty)
4. Update Block Table for User 1:
   Logical Block 0 → Physical Page #100
5. Remaining capacity in Page #100: 10 tokens

User 2 arrives. Prompt: “How do I reset my password? I forgot my email too.” (11 tokens)

Engine Action:
1. Allocate 1 physical page → Page #205
2. Compute K, V for 11 prompt tokens
3. Write K[0:11], V[0:11] into Page #205
4. Update Block Table for User 2:
   Logical Block 0 → Physical Page #205
5. Remaining capacity: 5 tokens

User 3 arrives. Prompt: “How do I reset my password?” (6 tokens)

Engine Action:
1. Allocate 1 physical page → Page #301
2. Write K, V into Page #301
3. Block Table for User 3:
   Logical Block 0 → Physical Page #301

Now all three users start generating tokens autoregressively.

Generation Step 1: User 1 produces token #6

Engine Action:
1. Check Block Table: User 1's next token index is 6
2. Token 6 belongs to Logical Block 0 (tokens 0-15)
3. Logical Block 0 → Physical Page #100
4. There are 10 empty slots left in Page #100
5. Compute K, V for new token
6. Write into Page #100 at slot 6
7. GPU Attention Kernel receives:
   - New query vector q_new
   - Block Table: [100]
   - Context length: 7
8. Kernel loads all 7 K, V vectors from Page #100
9. Computes attention over all 7 positions
10. Returns output

Generation Step 10: User 1 produces token #15 (the 16th token)

Engine Action:
1. Token index 15 belongs to Logical Block 0
2. Physical Page #100 now has exactly 1 slot left (slot 15)
3. Write K, V into slot 15
4. Page #100 is now FULL

Generation Step 11: User 1 produces token #16 (the 17th token)

Engine Action:
1. Token index 16 belongs to Logical Block 1 (tokens 16-31)
2. Check Block Table: User 1 has no mapping for Logical Block 1
3. Allocate NEW physical page from free pool → Page #450
4. Update Block Table:
   Logical Block 0 → Physical Page #100
   Logical Block 1 → Physical Page #450   [NEW]
5. Write K, V for token 16 into Page #450, slot 0
6. GPU Kernel now receives Block Table: [100, 450]
7. Kernel loads:
   - All 16 vectors from Page #100
   - 1 vector from Page #450
   - Total context: 17 tokens
8. Computes full attention over all 17 positions

Notice: the model attended to all 17 previous tokens. Paged Attention doesn’t skip tokens or select a subset. It loads the entire history. It just stores that history in scattered pages.

The Plot Twist: Sharing Is Caring

Now for the move that breaks the game.

User 2 and User 3 both started with the same prompt: “How do I reset my password?”

In standard inference, their KV caches are in separate contiguous blocks. The same 6 K, V vectors are stored twice.

With Paged Attention and copy-on-write, the engine can do this:

User 1 Block Table:
  Logical 0 → Physical #100 (tokens: "How do I reset my password?")

User 2 Block Table:
  Logical 0 → Physical #100 (SAME PAGE — shared!)
  Logical 1 → Physical #205 (unique continuation)

User 3 Block Table:
  Logical 0 → Physical #100 (SAME PAGE — shared!)

All three users point to the same physical page for their prompt. The memory for that prompt is stored once, not three times.

When User 2 generates a unique continuation, they get their own new pages. The shared page stays read-only until someone tries to write to it.

Real impact: In production, 60–80% of prompts share prefixes (system prompts, documentation, common questions). Paged Attention turns that redundancy into free memory.

The Exact Bottleneck

Paged Attention doesn’t make one user faster. If you’re alone on the GPU, your tokens generate at the same speed.

What it does is let you pack more users onto the same GPU.

Table

The bottleneck was never compute. It was memory fragmentation and inability to share. Paged Attention turns a GPU memory crisis into a throughput paradise.

Part IV: The Convergence -When Both Tricks Meet

Here’s where it gets beautiful.

Modern inference engines like vLLM, TensorRT-LLM, and TGI use both hacks together.

Flash Attention handles the math inside each attention operation, making sure the GPU isn’t drowning in memory traffic.

Paged Attention handles the storage of K and V across time, making sure the GPU isn’t drowning in wasted space.

Together, they let you run models on sequences of 100,000+ tokens, serving dozens of users simultaneously, on hardware that would have choked on a single long document three years ago.

Epilogue: The Physics Still Win

These aren’t magic. They’re engineering.

Flash Attention can’t create more SRAM. If your sequence is so long that even the tiled blocks don’t fit in fast memory, you’re back to being slow.

Paged Attention can’t create more GPU RAM. If you have a million users, you’ll still run out of pages.

But what these two ideas did was shift the frontier. They turned problems that looked like brick walls into problems that look like speed bumps.

The next time you use ChatGPT, Claude, or any long-context model, remember: somewhere in a data center, a GPU is running a heist and playing landlord at the same time. And because of that, you get your answer in seconds instead of minutes.

The math didn’t change. The physics didn’t change. The engineers just stopped carrying so many boxes.