Introduction

Every modern language model — GPT, LLaMA, Gemini — is built on the same core operation: self-attention. For each token in a sequence, attention computes a weighted blend of all the other tokens, letting every word "look at" every other word at once.

The catch is the cost. With NN tokens you need an N×NN \times N matrix of scores — quadratic in both time and memory. For short sequences that is fine. For sequences of thousands of tokens, that matrix alone can fill and overflow the fast on-chip memory of a GPU.

In 2022, Tri Dao and colleagues at Stanford published FlashAttention: an algorithm that computes the exact same result as standard attention while never materializing the full N×NN \times N matrix in fast memory. The trick is not an approximation. It is an IO-aware kernel — a piece of code that reasons about where data lives (slow HBM vs. fast SRAM) and moves it as little as possible.

The IO-Aware Kernel

The demo below visualises FlashAttention's core idea: tile the attention computation so that each tile fits in fast on-chip memory (SRAM), accumulate the result with an online softmax, then discard the tile before loading the next one.

<!-- FlashAttention demo: compare HBM accesses in standard vs. flash attention -->
<div class="controls">
  <label>Sequence length N = <strong id="nval">8</strong>
    <input type="range" id="nslider" min="4" max="16" step="4" value="8">
  </label>
  <label>Tile size T = <strong id="tval">4</strong>
    <input type="range" id="tslider" min="2" max="8" step="2" value="4">
  </label>
</div>
<div class="mode-btns">
  <button id="btn-standard" class="active" type="button">Standard Attention</button>
  <button id="btn-flash" type="button">FlashAttention</button>
</div>
<canvas id="canvas" width="480" height="280"></canvas>
<div id="stats" class="stats"></div>
<div class="action-btns">
  <button id="btn-run" type="button">Run animation</button>
  <button id="btn-reset" type="button" class="ghost">Reset</button>
</div>
<p class="caption" id="caption">Press ’Run animation’ to see each memory access step.</p>

/* Styles for the FlashAttention IO-aware tiling demo */
* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; color: #222; margin: 0; }
.controls { display: flex; gap: 1.4rem; flex-wrap: wrap; margin-bottom: .6rem; }
label { font-size: .85rem; color: #444; display: flex; align-items: center; gap: .5rem; }
input[type=range] { width: 90px; }
.mode-btns { display: flex; gap: .5rem; margin-bottom: .5rem; }
button { font: 600 13px system-ui; padding: .4rem .85rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 7px; cursor: pointer; }
button.ghost { background: #fff; color: #1d3557; }
button.active { background: #e63946; border-color: #c0303c; }
canvas { display: block; border: 1px solid #cdd9e3; border-radius: 8px;
         background: #f8fafc; max-width: 100%; }
.stats { font-size: .85rem; margin: .4rem 0; min-height: 1.3em; color: #1d3557; font-weight: 600; }
.action-btns { display: flex; gap: .5rem; margin-top: .4rem; }
.caption { font-size: .82rem; color: #555; margin: .5rem 0 0; line-height: 1.5; }

const fmt = (s, v) => s.replace(/%(\w+)%/g, (_, k) => String(v[k] ?? ''));
// FlashAttention IO-aware tiling visualizer
// Shows how standard attention makes O(N^2) HBM round-trips while FlashAttention tiles to SRAM.

const canvas = document.getElementById('canvas');
const ctx = canvas.getContext('2d');

// Demo state: sequence length N, tile size T, current animation mode
let N = 8, T = 4, mode = 'standard';
let animFrame = null, animStep = 0, animSteps = [];

const COLORS = {
  q: '#4a90d9',   // query matrix color
  k: '#e67e22',   // key matrix color
  v: '#27ae60',   // value matrix color
  s: '#9b59b6',   // score matrix color
  hbm_write: '#e63946',  // HBM write highlight
  hbm_read:  '#f4a261',  // HBM read highlight
  sram:      '#2ecc71',  // SRAM tile highlight
  out:       '#1d3557',  // output matrix color
  tile:      'rgba(230,57,70,0.18)',  // active tile overlay
  bg:        '#f8fafc',
  border:    '#cdd9e3',
  text:      '#222',
};

function getN(){ return parseInt(document.getElementById('nslider').value); }
function getT(){ return parseInt(document.getElementById('tslider').value); }

// Build the list of memory-access steps for a given mode
function buildSteps(n, t, m){
  const steps = [];
  const tClamped = Math.min(t, n);
  if(m === 'standard'){
    // Step 1: read Q, K, V from HBM
    steps.push({ type: 'hbm_read', label: 'read_qkv', rows: n });
    // Step 2: write N×N score matrix S to HBM
    steps.push({ type: 'hbm_write', label: 'write_s', rows: n });
    // Step 3: read S back from HBM for softmax
    steps.push({ type: 'hbm_read', label: 'read_s', rows: n });
    // Step 4: write softmax(S) = P to HBM
    steps.push({ type: 'hbm_write', label: 'write_p', rows: n });
    // Step 5: read P back from HBM
    steps.push({ type: 'hbm_read', label: 'read_p', rows: n });
    // Step 6: write output O to HBM
    steps.push({ type: 'hbm_write', label: 'write_o', rows: n });
  } else {
    // FlashAttention: loop over tiles, each tile stays in SRAM
    const numTiles = Math.ceil(n / tClamped);
    for(let i = 0; i < numTiles; i++){
      // Load one tile of Q/K/V into SRAM
      steps.push({ type: 'tile_sram', label: 'load_tile', tile: i, numTiles, rows: tClamped });
      // Compute partial softmax + accumulate O in SRAM
      steps.push({ type: 'tile_compute', label: 'compute_tile', tile: i, numTiles, rows: tClamped });
      // Write only the O chunk back to HBM
      steps.push({ type: 'tile_write_o', label: 'write_tile_o', tile: i, numTiles, rows: tClamped });
    }
  }
  return steps;
}

// Count total HBM element reads + writes across all steps
function countHBM(steps){
  let reads = 0, writes = 0;
  for(const s of steps){
    if(s.type === 'hbm_read') reads += s.rows * s.rows;
    if(s.type === 'hbm_write') writes += s.rows * s.rows;
    if(s.type === 'tile_sram') reads += s.rows * s.rows;
    if(s.type === 'tile_write_o') writes += s.rows * s.rows;
  }
  return { reads, writes };
}

// Render the canvas: HBM grid on left, SRAM grid on right
function draw(step){
  const W = canvas.width, H = canvas.height;
  ctx.clearRect(0, 0, W, H);

  const n = N, t = Math.min(T, n);
  const cellSize = Math.min(Math.floor((W * 0.42) / n), 24);
  const gridW = cellSize * n, gridH = cellSize * n;
  const leftX = 16, topY = Math.floor((H - gridH) / 2);
  const rightX = W - gridW - 16;

  // Clear canvas background
  ctx.fillStyle = COLORS.bg;
  ctx.fillRect(0, 0, W, H);

  // Draw HBM / SRAM section labels
  ctx.font = 'bold 11px system-ui';
  ctx.fillStyle = COLORS.text;
  ctx.textAlign = 'center';
  ctx.fillText('HBM ' + step.hbmLabel, leftX + gridW / 2, topY - 10);
  ctx.fillText('SRAM', rightX + gridW / 2, topY - 10);

  // Draw the HBM matrix grid with highlights
  for(let r = 0; r < n; r++){
    for(let c = 0; c < n; c++){
      const x = leftX + c * cellSize, y = topY + r * cellSize;
      let fill = '#e8eef3';
      if(step.hbmHighlight){
        if(step.hbmHighlight === 'all') fill = COLORS.hbm_read;
        else if(step.hbmHighlight === 'write_s') fill = COLORS.hbm_write;
        else if(step.hbmHighlight === 'write_p') fill = COLORS.s;
        else if(step.hbmHighlight === 'write_o') fill = COLORS.out;
      }
      ctx.fillStyle = fill;
      ctx.fillRect(x + 1, y + 1, cellSize - 2, cellSize - 2);
      ctx.strokeStyle = COLORS.border;
      ctx.strokeRect(x + 1, y + 1, cellSize - 2, cellSize - 2);
    }
  }

  // Draw the SRAM tile grid with active-tile highlight
  for(let r = 0; r < n; r++){
    for(let c = 0; c < n; c++){
      const x = rightX + c * cellSize, y = topY + r * cellSize;
      let fill = '#e8eef3';
      if(step.sramHighlight){
        const inTile = step.tileRow !== undefined
          ? (r >= step.tileRow && r < step.tileRow + t)
          : false;
        if(inTile) fill = COLORS.sram;
      }
      ctx.fillStyle = fill;
      ctx.fillRect(x + 1, y + 1, cellSize - 2, cellSize - 2);
      ctx.strokeStyle = COLORS.border;
      ctx.strokeRect(x + 1, y + 1, cellSize - 2, cellSize - 2);
    }
  }

  // Draw a directional arrow between HBM and SRAM
  if(step.arrowDir){
    const ay = topY + gridH / 2;
    const ax1 = step.arrowDir === 'right' ? leftX + gridW + 4 : rightX - 4;
    const ax2 = step.arrowDir === 'right' ? rightX - 4 : leftX + gridW + 4;
    ctx.strokeStyle = step.arrowDir === 'right' ? COLORS.sram : COLORS.hbm_read;
    ctx.lineWidth = 2.2;
    ctx.beginPath();
    ctx.moveTo(ax1, ay);
    ctx.lineTo(ax2, ay);
    const dir = step.arrowDir === 'right' ? 1 : -1;
    ctx.lineTo(ax2 - dir * 8, ay - 5);
    ctx.moveTo(ax2, ay);
    ctx.lineTo(ax2 - dir * 8, ay + 5);
    ctx.stroke();
    ctx.lineWidth = 1;
  }
}

// Map an animation step to its visual parameters
function stepToVis(s, n, t){
  if(s.type === 'hbm_read')
    return { hbmLabel: '(S, read)', hbmHighlight: 'all', sramHighlight: false, arrowDir: 'right' };
  if(s.type === 'hbm_write' && s.label === 'write_s')
    return { hbmLabel: '(writing S)', hbmHighlight: 'write_s', sramHighlight: false, arrowDir: 'left' };
  if(s.type === 'hbm_write' && s.label === 'write_p')
    return { hbmLabel: '(writing P=softmax(S))', hbmHighlight: 'write_p', sramHighlight: false, arrowDir: 'left' };
  if(s.type === 'hbm_write' && s.label === 'write_o')
    return { hbmLabel: '(writing O)', hbmHighlight: 'write_o', sramHighlight: false, arrowDir: 'left' };
  if(s.type === 'tile_sram')
    return { hbmLabel: '(idle)', hbmHighlight: null, sramHighlight: true, tileRow: s.tile * t, arrowDir: 'right' };
  if(s.type === 'tile_compute')
    return { hbmLabel: '(idle)', hbmHighlight: null, sramHighlight: true, tileRow: s.tile * t, arrowDir: null };
  if(s.type === 'tile_write_o')
    return { hbmLabel: '(write O chunk)', hbmHighlight: 'write_o', sramHighlight: false, tileRow: s.tile * t, arrowDir: 'left' };
  return { hbmLabel: '', hbmHighlight: null, sramHighlight: false, arrowDir: null };
}

const STEP_LABELS = {
  read_qkv:    'Read Q, K, V from HBM',
  write_s:     'Write S = QKᵀ to HBM',
  read_s:      'Read S from HBM',
  write_p:     'Write P = softmax(S) to HBM',
  read_p:      'Read P from HBM',
  write_o:     'Write O = PV to HBM',
  load_tile:   'Load tile into SRAM',
  compute_tile:'Compute tile in SRAM (online softmax)',
  write_tile_o:'Write O chunk to HBM',
};

function updateStats(){
  const n = getN(), t = getT();
  const std = countHBM(buildSteps(n, t, 'standard'));
  const fl  = countHBM(buildSteps(n, t, 'flash'));
  const stdTotal = std.reads + std.writes;
  const flTotal  = fl.reads + fl.writes;
  const ratio = (stdTotal / flTotal).toFixed(1);
  document.getElementById('stats').textContent =
    (mode === 'standard'
      ? fmt('Standard attention: %stdTotal% HBM element accesses (reads + writes)', { stdTotal: stdTotal })
      : fmt('FlashAttention: %flTotal% HBM element accesses — %ratio%× fewer than standard', { flTotal: flTotal, ratio: ratio }));
}

function reset(){
  if(animFrame){ cancelAnimationFrame(animFrame); animFrame = null; }
  animStep = 0; animSteps = [];
  N = getN(); T = getT();
  const vis = stepToVis({ type: 'hbm_read', label: 'read_qkv', rows: N }, N, Math.min(T,N));
  draw({ ...vis, hbmHighlight: null, arrowDir: null, hbmLabel: '' });
  document.getElementById('caption').textContent = 'Press ’Run animation’ to see each memory access step.';
  updateStats();
}

function runAnimation(){
  if(animFrame){ cancelAnimationFrame(animFrame); animFrame = null; }
  N = getN(); T = getT();
  animSteps = buildSteps(N, Math.min(T, N), mode);
  animStep = 0;
  function tick(){
    if(animStep >= animSteps.length){ animFrame = null; return; }
    const s = animSteps[animStep];
    const vis = stepToVis(s, N, Math.min(T, N));
    draw(vis);
    const lbl = STEP_LABELS[s.label] || s.label;
    document.getElementById('caption').textContent =
      (mode === 'standard' ? 'Standard attention step:' : 'FlashAttention step:') + ' ' + lbl;
    updateStats();
    animStep++;
    animFrame = setTimeout(tick, 650);
  }
  tick();
}

document.getElementById('nslider').addEventListener('input', function(){
  document.getElementById('nval').textContent = this.value; reset();
});
document.getElementById('tslider').addEventListener('input', function(){
  document.getElementById('tval').textContent = this.value; reset();
});
document.getElementById('btn-standard').addEventListener('click', function(){
  mode = 'standard';
  document.getElementById('btn-standard').classList.add('active');
  document.getElementById('btn-flash').classList.remove('active');
  reset();
});
document.getElementById('btn-flash').addEventListener('click', function(){
  mode = 'flash';
  document.getElementById('btn-flash').classList.add('active');
  document.getElementById('btn-standard').classList.remove('active');
  reset();
});
document.getElementById('btn-run').addEventListener('click', runAnimation);
document.getElementById('btn-reset').addEventListener('click', reset);

reset();
Result
Console

Notice what never happens: the full N×NN \times N matrix of attention scores is never written to slow GPU memory (HBM). Standard attention writes it once and reads it back; FlashAttention rewrites each tile entirely in fast SRAM and throws it away. The final output is the same — bit-for-bit identical — because the online softmax accumulates running statistics that let it rescale on the fly.

The Real Complexity

Standard benchmarks measure FLOPs. But on modern GPUs, arithmetic is cheap — memory bandwidth is the real bottleneck.

  • Standard attention reads the QQ, KK, VV matrices, writes the N×NN \times N score matrix SS to HBM, reads it back to apply softmax, writes the softmax result, reads it again to multiply by VV. Total HBM accesses: O(N2)O(N^2).
  • FlashAttention tiles the computation. For a tile block size Br×BcB_r \times B_c, it keeps the current tile entirely in SRAM, accumulates partial softmax statistics, and writes only the output OO back to HBM. Total HBM accesses: O(N2d/M)O(N^2 d / M), where dd is head dimension and MM is SRAM size — sublinear in NN relative to the naïve approach when MM is large.
  • Recomputation on the backward pass: FlashAttention does not store the N×NN \times N attention matrix for backpropagation. Instead it recomputes it from the tiles during the backward pass. This trades extra FLOPs for drastically less memory — and FLOPs are cheap, HBM reads are not.
  • The result is exact: unlike sparse or approximate attention methods, FlashAttention produces the numerically identical output. No information is lost; only the order of operations changes.

Tri Dao's 2022 paper showed 2–4× wall-clock speedup over standard PyTorch attention on A100 GPUs, and 5–20× less memory for long sequences. The 2023 FlashAttention-2 paper pushed it further with better parallelism, reaching up to 9× speedup. This is what allows modern LLMs to handle context windows of 100k+ tokens — it is the same reason transformer training became practical at scale.

Where It Matters

FlashAttention is now a default ingredient in virtually every serious transformer implementation:

  • LLM training: GPT-4, LLaMA 2/3, Mistral, and most large-scale models use FlashAttention to cut training costs significantly. Without it, training on sequences longer than a few thousand tokens would require impractically large GPU clusters.
  • Long-context inference: 100k-token context windows (Anthropic Claude, GPT-4 Turbo) are only feasible because attention no longer blows up GPU memory at those lengths.
  • Protein and molecular structure: AlphaFold-style models fold thousands of residues; the same IO-aware trick lets them run on hardware that would otherwise run out of memory.
  • Video and multimodal models: treating video frames as tokens means sequences of tens of thousands; FlashAttention makes the attention pass tractable.
  • Speculative decoding and fine-tuning: techniques like LoRA and QLoRA that fine-tune large models on consumer GPUs depend heavily on the memory savings FlashAttention provides.

The insight generalises: whenever you have a computation that produces a large intermediate matrix that is immediately consumed and discarded, tiling and recomputation can replace it with a series of passes over fast memory. The matrix multiplication literature has known this for decades; FlashAttention applied the idea to the attention operator specifically.

Conclusion

FlashAttention is a reminder that the bottleneck is rarely where you think it is. The naive analysis says attention is slow because of O(N2)O(N^2) multiplications. The real analysis says attention is slow because those multiplications force O(N2)O(N^2) round-trips between slow and fast memory — and that those round-trips can be slashed without losing a single bit of accuracy.

Tiling, online softmax, and recomputation on the backward pass are not exotic tricks. They are classical ideas from high-performance computing, applied precisely to the right operation at the right moment. The payoff is enormous: context windows have grown from a few thousand tokens to hundreds of thousands, and transformer models that once required dozens of GPUs now fit on a handful.

The deeper lesson: when you understand the memory hierarchy, you can redesign algorithms that look fixed — and sometimes the redesign changes everything.

Share this article

Pick a channel — or use your device's native share sheet.

Comments

Loading comments...

https://www.kipuhub.com/en/article/flash-attention/Content licensed under CC BY-NC 4.0.