Kernel visualization: Flash Attention 1

Flash Attention builds on the ideas behind online softmax. While online softmax requires two fetches from HBM, Flash Attention only requires one. Just like online softmax, the output matrix can be written as a recurrence formula:

Kernel implementation

The visualization shows how data moves from HBM to a single block in SRAM, how the intermediate values are calculated, and finally how the output is moved back to HBM. Each block operates identically on different tiles.

# Drawing helpers: memory slots in HBM and SRAM, and the actions the kernel below takes on them.
# Each action is recorded into the current step; `show_steps` draws the steps.
import json
import math

INK, GRAY = '#27272a', '#3f3f46'  # text; outlines

SLOTS = []  # every slot, in drawing order: HBM first, so that data moving into SRAM passes over it

class Slot:
    """A place in HBM or SRAM that holds one tile, value or register, drawn as a w by h box centered at (x, y).
    `data`: what it holds at the start. `kind`: the palette its updates use. An SRAM slot is also drawn empty
    (`outline`), with `caption`, by default its name, under it."""
    def __init__(self, name, memory, x, y, w, h, data=None, font=14, kind=None, outline=None, caption=None):
        self.name, self.memory, self.x, self.y, self.w, self.h = name, memory, x, y, w, h
        self.initial, self.font, self.kind = data, font, kind
        self.outline = memory == 'sram' if outline is None else outline
        self.caption = name if caption is None else caption
        SLOTS.append(self)

def data(label='', fill='none', line=GRAY, level=0):
    """What a slot holds, as drawn: a label, a fill and an outline color. `level` counts its updates."""
    return dict(label=label, fill=fill, line=line, level=level)

def as_list(x):
    return list(x) if isinstance(x, (list, tuple)) else [x]

# ---- The actions ----

def new_figure(block=None, caption='Setup ready'):
    """Start recording. `block`: (square, slots) of the CUDA block SRAM shows from the start, or None."""
    global FRAMES, state, pending, block_now, transfers
    state = {s: s.initial for s in SLOTS}
    block_now, pending, transfers = block, {}, {'fetches': 0, 'stores': 0}
    FRAMES = [dict(label='start', caption=caption, events={}, state=dict(state), block=block)]

def content(slot):
    """What `slot` holds now, this step's actions included."""
    return pending[slot][2] if slot in pending else state[slot]

def move_data(src, dst):
    """Copy the data in `src` to `dst`, slots or lists of slots, e.g. a tile from HBM to SRAM. Data without a label
    takes the label `dst` had. One call is one transfer."""
    for a, b in zip(as_list(src), as_list(dst)):
        moved = dict(content(a))
        moved['label'] = moved['label'] or (content(b) or data())['label']
        pending[b] = ('move', a, moved, 0, 1)
    a, b = as_list(src)[0], as_list(dst)[0]
    if a.memory != b.memory:
        transfers['fetches' if a.memory == 'hbm' else 'stores'] += 1

def change(slots, one_by_one=False, **fields):
    """Change the data in `slots` in place: `fields` replace its label, fill, line or level.
    With `one_by_one`, the slots change one after another, e.g. the cells of a matrix."""
    slots = as_list(slots)
    for n, s in enumerate(slots):
        new = {**(content(s) or data()), **fields}
        pending[s] = ('change', None, new, n if one_by_one else 0, len(slots) if one_by_one else 1)

def run_block(square, slots):
    """From this step, SRAM shows the shared memory of the CUDA block at `square` of the SRAM grid: `slots`, empty."""
    global block_now
    block_now = (square, tuple(slots))
    for s in slots:
        state[s] = s.initial

def step(label, caption=None):
    """End a step. `label` names it in the step list; `caption`, by default `label`, is shown on the canvas."""
    for s, event in pending.items():
        state[s] = event[2]
    FRAMES.append(dict(label=label, caption=caption or label, events=dict(pending), state=dict(state), block=block_now))
    pending.clear()

# ---- Drawing: the whole figure is one Diagram node, and the step list changes it ----

def shown(s, frame):
    """Whether `frame` shows slot s: HBM always, SRAM while its CUDA block is the one shown."""
    return s.memory == 'hbm' or (frame['block'] is not None and s in frame['block'][1])

def item(x=0, y=0, w=0, h=0, fill='none', stroke='none', label='', font=12, wrap=0, line=False):
    """One shape of the diagram: a box with (x, y) its bottom-left corner, or a line from (x, y) to (x + w, y + h).
    A box of size 0 draws only its label."""
    return dict(x=round(x, 4), y=round(y, 4), w=round(w, 4), h=round(h, 4), fill=fill, stroke=stroke,
                label=label, font=font, wrap=wrap, line=line, color=INK)

ROLES = ('outline', 'under', 'top', 'ring', 'caption')  # each slot's items, in drawing order

def slot_items(s, k, phase):
    """Slot s's items at step k, by role: its empty outline and its caption, what it held before the step ('under'),
    what it holds ('top'), and a ring while it changes in place. `phase` is 'start', 'move' or 'end': how the step
    begins, where its moves go, and how it ends."""
    frame, items = FRAMES[k], dict.fromkeys(ROLES, item())
    if not shown(s, frame):
        return items
    bottom = s.y - s.h / 2

    def holding(d, cx=s.x, cy=s.y, h=s.h):
        return item(cx - s.w / 2, cy - s.h / 2, s.w, h, d['fill'], d['line'], d['label'], s.font) if d else item()

    if s.outline:
        items['outline'] = item(s.x - s.w / 2, bottom, s.w, s.h, stroke=GRAY)
        items['caption'] = item(s.x, bottom - 0.3, label=s.caption, font=min(s.font, 12))
    top, event = frame['state'][s], frame['events'].get(s)
    under = FRAMES[k - 1]['state'][s] if event else None
    if event is None or phase == 'end':
        items['top'] = holding(top)
    elif event[0] == 'move':  # the copy flies in over what the slot held
        start = (event[1].x, event[1].y) if phase == 'start' else (s.x, s.y)
        items['under'], items['top'] = holding(under), holding(top, *start)
    else:  # a change in place: the new data fills the slot from the bottom, under its new label
        old = holding(under) if under else item(s.x - s.w / 2, bottom, s.w, s.h)
        items['under'] = dict(old, label=top['label'], font=s.font)
        items['top'] = dict(holding(top, h=0 if phase == 'start' else s.h), label='')
    if event is not None and event[0] == 'change':
        items['ring'] = item(s.x - s.w / 2 - 0.06, bottom - 0.06, s.w + 0.12, s.h + 0.12, stroke=INK)
    return items

def zoom_ends(block, side):
    """The zoom line on `side` (0 left, 1 right) of `block`: from a bottom corner of its square in the SRAM grid
    to the top corner of its slots on that side."""
    (bx, by), slots = block
    edge = min(s.x - s.w / 2 for s in slots) if side == 0 else max(s.x + s.w / 2 for s in slots)
    return (bx + 0.3 * side, by), (edge, max(s.y + s.h / 2 for s in slots))

def zoom_items(k, phase):
    """The two zoom lines at step k. They grow when SRAM is first shown, and slide when the block shown changes."""
    now = FRAMES[k]['block']
    before = FRAMES[k - 1]['block'] if k else now
    lines = []
    for side in (0, 1):
        if now is None or now[0] is None:
            lines.append(item(line=True))
            continue
        start, end = zoom_ends(now, side)
        if phase == 'start' and (before is None or before[0] is None):
            end = start
        elif phase == 'start' and before != now:
            start, end = zoom_ends(before, side)
        lines.append(item(*start, end[0] - start[0], end[1] - start[1], stroke=GRAY, line=True))
    return lines

def rect(left, bottom, w, h, **colors):
    """A rectangle. Every argument but the colors can be a node."""
    return pie.rectangle(pie.point(left, bottom, hidden=True), w, h, **colors)

def show_steps(name, background, caption_at=(0, -5.6), caption_width=8, unit=50):
    """Draw the recorded steps as one Diagram node, with a step list that plays them."""
    pie.clear()
    pie.show_canvas('right')
    pie.scalar(0, out='gridOpacity')
    pie.scalar(unit, out='unit')
    background()
    squares = list(dict.fromkeys(f['block'][0] for f in FRAMES if f['block'] and f['block'][0]))

    def snapshot(k, phase):
        """Every item at step k: the filled square of the block shown, the zoom lines, each role of every slot, the caption."""
        block = FRAMES[k]['block']
        fills = [item(*q, 0.3, 0.3, GRAY, GRAY) if block and block[0] == q else item(*q) for q in squares]
        slots = [slot_items(s, k, phase) for s in SLOTS]
        caption = item(*caption_at, label=FRAMES[k]['caption'], wrap=caption_width)
        return fills + zoom_items(k, phase) + [items[role] for role in ROLES for items in slots] + [caption]

    def diff(a, b):
        """[index, changed fields] of each item that differs from `a` to `b`."""
        return [[i, {key: v for key, v in new.items() if old[key] != v}] for i, (old, new) in enumerate(zip(a, b)) if old != new]

    first_top = len(squares) + 2 + ROLES.index('top') * len(SLOTS)  # the index of the first slot's 'top' item
    steps, last = [dict(label=FRAMES[0]['label'])], snapshot(0, 'end')
    for k in range(1, len(FRAMES)):
        start, move, end = (snapshot(k, phase) for phase in ('start', 'move', 'end'))
        events = FRAMES[k]['events']
        stagger = [[first_top + i, events[s][3], events[s][4]] for i, s in enumerate(SLOTS) if s in events and events[s][4] > 1]
        steps.append(dict(label=FRAMES[k]['label'], start=diff(last, start), move=diff(start, move), end=diff(move, end), stagger=stagger))
        last = end
    figure = pie.diagram(json.dumps(snapshot(0, 'end')), out=f'{name}_figure')
    pie.ui.step_list(figure, steps=steps)

# ---- This page: where the slots are, and the kernel's actions on them ----

TILE = {'Q': ('#99f6e4', '#115e59'), 'K': ('#fecaca', '#b91c1c'), 'V': ('#f4f4f5', GRAY)}  # (fill, outline)
# The carries, in the hues of the equations: (outline, fill once one K, V tile is folded in, fill once both are)
CARRY = {'O': ('#b45309', '#fde68a', '#f59e0b'), 'l': ('#047857', '#a7f3d0', '#34d399'), 'm': ('#1d4ed8', '#bfdbfe', '#60a5fa')}
S_LINE, S_FILLS = '#7c3aed', ('#ddd6fe', '#a78bfa')  # Sij: each new Sij takes the other fill

def hbm_tiles(name, x, w, font=14):
    """Tiles 0 and 1 of matrix `name` in HBM."""
    fill, line = TILE[name] if name in TILE else ('none', CARRY[name][0])
    return [Slot(f'{name}{n}', 'hbm', x, y, w, 1.2, data(f'{name}{n}', fill, line), font) for n, y in enumerate((3.625, 2.375))]

Q, K, V = hbm_tiles('Q', -3.15, 1.6), hbm_tiles('K', -1.35, 1.6), hbm_tiles('V', 0.45, 1.6)
O, l, m = hbm_tiles('O', 2.25, 1.6), hbm_tiles('l', 3.375, 0.25, 10), hbm_tiles('m', 3.825, 0.25, 10)

# The shared memory of one CUDA block
Qi = Slot('Qi', 'sram', -3.7, -3.0, 1.6, 1.2)
Kj = Slot('Kj', 'sram', -1.95, -3.0, 1.6, 1.2)
Sij = Slot('Sij', 'sram', -0.4, -3.0, 1.2, 1.2)
Sij.cells = [Slot(f'Sij[{r}][{c}]', 'sram', -0.85 + 0.3 * c, -2.55 - 0.3 * r, 0.3, 0.3, outline=False)
             for r in range(4) for c in range(4)]
Vj = Slot('Vj', 'sram', 1.15, -3.0, 1.6, 1.2)
mi = Slot('mi', 'sram', 2.225, -3.0, 0.25, 1.2, font=10, kind='m')
li = Slot('li', 'sram', 2.625, -3.0, 0.25, 1.2, font=10, kind='l')
Oi = Slot('Oi', 'sram', 3.7, -3.0, 1.6, 1.2, kind='O')
SRAM = [s for s in SLOTS if s.memory == 'sram']

def compute(S):
    """Compute matrix S, e.g. Sij = Qi @ Kj.T: its cells fill one after another."""
    level = (content(S.cells[0]) or data())['level'] + 1
    change(S.cells, one_by_one=True, fill=S_FILLS[level % 2], line=S_LINE, level=level)

def update(*carries):
    """The online softmax update of mi, li and Oi with one more K, V tile: each turns a shade darker."""
    for s in carries:
        level = (content(s) or data())['level'] + 1
        line, *shades = CARRY[s.kind]
        change(s, fill=shades[min(level, len(shades)) - 1], line=line, level=level)

def background():
    """What doesn't change: the names of HBM, SRAM and the matrices, and the SRAM grid of CUDA blocks."""
    pie.text('HBM', 0, 5, 14, stroke=INK)
    for label, x in (('Q', -3.15), ('K', -1.35), ('V', 0.45), ('O', 2.25), ('L', 3.375), ('m', 3.825)):
        pie.text(label, x, 1.5, stroke=INK)
    pie.text('SRAM', 0, 0.6, 14, stroke=INK)
    for y in (0.05, -0.35):
        for x in (-0.55, -0.15, 0.25):
            rect(x, y, 0.3, 0.3, stroke=GRAY)

# The kernel: K and V on the outer loop, Q on the inner loop
new_figure(block=((-0.15, -0.35), SRAM))  # SRAM shows one CUDA block: the filled square of the SRAM grid
Tr = Tc = 2  # blocks of Q rows; tiles of K and V

for j in range(Tc):
    move_data([K[j], V[j]], [Kj, Vj])                      # Kj, Vj = K[j], V[j]
    step(f'{j}. load K{j}, V{j}', f'[Outer loop: {j}] load K{j}, V{j}')
    for i in range(Tr):
        loop = f'[Outer loop: {j}, inner loop: {i}]'
        move_data([Q[i], m[i], l[i], O[i]], [Qi, mi, li, Oi])  # Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]
        step(f'{j}.{i} load Q{i}, m{i}, l{i}, O{i}', f'{loop} load Q{i}, m{i}, l{i}, O{i}')
        compute(Sij)                                       # Sij = Qi @ Kj.T
        step(f'{j}.{i} fill Sij', f'{loop} fill Sij')
        update(mi, li, Oi)                                 # mi, li, Oi = rescale, then add this tile
        step(f'{j}.{i} update mi, li and Oi', f'{loop} update mi, li and Oi')
        move_data([mi, li, Oi], [m[i], l[i], O[i]])        # m[i], L[i], O[i] = mi, li, Oi
        step(f'{j}.{i} write m{i}, l{i}, O{i}', f'{loop} write m{i}, l{i}, O{i}')

show_steps('fa1', background)

Run the two editors in order: the first defines the drawing and the actions on memory, the second is the kernel. Then pick a step. Drag the step list by its ⋮⋮ grip onto the canvas to keep it in view while you read.

template<int Br, int Bc, int D>
__global__ void fa1(const float* Q, const float* K, const float* V,
                    float* O, float* l, float* m, int N, float scale) {
    int bh = blockIdx.x, tid = threadIdx.x;
    Q += bh*N*D; K += bh*N*D; V += bh*N*D; O += bh*N*D;
    l += bh*N;   m += bh*N;

    __shared__ float Qi[Br][D], Oi[Br][D], Kj[Bc][D], Vj[Bc][D];
    __shared__ float Sij[Br][Bc], mi[Br], li[Br];

    int Tr = (N+Br-1)/Br, Tc = (N+Bc-1)/Bc;

0. Outer iteration j = 0

One K/V tile is loaded into SRAM and stays there while every Q block is swept against it.

• Kj, Vj = K[j], V[j]

// j = 0
for (int x = tid; x < Bc*D; x += blockDim.x) {
    Kj[x/D][x%D] = K[(0*Bc + x/D)*D + x%D];
    Vj[x/D][x%D] = V[(0*Bc + x/D)*D + x%D];
}
__syncthreads();

0.0 Inner iteration i = 0

It carries empty mi, li, Oi from HBM.

• Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]

// j = 0, i = 0
for (int x = tid; x < Br*D; x += blockDim.x) {
    Qi[x/D][x%D] = Q[(0*Br + x/D)*D + x%D];
    Oi[x/D][x%D] = O[(0*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[0*Br+tid]; li[tid] = l[0*Br+tid]; }
__syncthreads();

• Sij = Qi @ Kj.T

// j = 0, i = 0
for (int x = tid; x < Br*Bc; x += blockDim.x) {
    int r = x/Bc, c = x%Bc;
    float s = 0.f;
    for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
    Sij[r][c] = s*scale;
}
__syncthreads();

• mi = max(mi, Sij.max(-1))
• li = rescale(li) + exp(Sij - mi).sum(-1)
• Oi = rescale(Oi) + softmax(Sij) @ Vj

// j = 0, i = 0
if (tid < Br) {
    float m_ij = -INFINITY;
    for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
    float m_new = fmaxf(mi[tid], m_ij);
    float alpha = expf(mi[tid] - m_new);
    float lsum = 0.f;
    for (int c = 0; c < Bc; c++) {
        Sij[tid][c] = expf(Sij[tid][c] - m_new);
        lsum += Sij[tid][c];
    }
    float l_new = alpha*li[tid] + lsum;
    for (int d = 0; d < D; d++) {
        float pv = 0.f;
        for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
        Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
    }
    mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();

• m[i], L[i], O[i] = mi, li, Oi

// j = 0, i = 0
for (int x = tid; x < Br*D; x += blockDim.x)
    O[(0*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[0*Br+tid] = mi[tid]; l[0*Br+tid] = li[tid]; }
__syncthreads();

0.1 Inner iteration i = 1

It carries empty mi, li, Oi from HBM

• Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]

// j = 0, i = 1
for (int x = tid; x < Br*D; x += blockDim.x) {
    Qi[x/D][x%D] = Q[(1*Br + x/D)*D + x%D];
    Oi[x/D][x%D] = O[(1*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[1*Br+tid]; li[tid] = l[1*Br+tid]; }
__syncthreads();

• Sij = Qi @ Kj.T

// j = 0, i = 1
for (int x = tid; x < Br*Bc; x += blockDim.x) {
    int r = x/Bc, c = x%Bc;
    float s = 0.f;
    for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
    Sij[r][c] = s*scale;
}
__syncthreads();

• mi = max(mi, Sij.max(-1))
• li = rescale(li) + exp(Sij - mi).sum(-1)
• Oi = rescale(Oi) + softmax(Sij) @ Vj

// j = 0, i = 1
if (tid < Br) {
    float m_ij = -INFINITY;
    for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
    float m_new = fmaxf(mi[tid], m_ij);
    float alpha = expf(mi[tid] - m_new);
    float lsum = 0.f;
    for (int c = 0; c < Bc; c++) {
        Sij[tid][c] = expf(Sij[tid][c] - m_new);
        lsum += Sij[tid][c];
    }
    float l_new = alpha*li[tid] + lsum;
    for (int d = 0; d < D; d++) {
        float pv = 0.f;
        for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
        Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
    }
    mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();

• m[i], L[i], O[i] = mi, li, Oi

// j = 0, i = 1
for (int x = tid; x < Br*D; x += blockDim.x)
    O[(1*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[1*Br+tid] = mi[tid]; l[1*Br+tid] = li[tid]; }
__syncthreads();

1. Outer iteration j=1

K0, V0 are dropped and the next K/V tile takes their place, and the same Q blocks come round again, now carrying the partial results left behind by j = 0.

• Kj, Vj = K[j], V[j]

// j = 1
for (int x = tid; x < Bc*D; x += blockDim.x) {
    Kj[x/D][x%D] = K[(1*Bc + x/D)*D + x%D];
    Vj[x/D][x%D] = V[(1*Bc + x/D)*D + x%D];
}
__syncthreads();

1.0 Inner iteration i = 0

It carries partial (not empty) mi, li, Oi from HBM

• Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]

// j = 1, i = 0
for (int x = tid; x < Br*D; x += blockDim.x) {
    Qi[x/D][x%D] = Q[(0*Br + x/D)*D + x%D];
    Oi[x/D][x%D] = O[(0*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[0*Br+tid]; li[tid] = l[0*Br+tid]; }
__syncthreads();

• Sij = Qi @ Kj.T

// j = 1, i = 0
for (int x = tid; x < Br*Bc; x += blockDim.x) {
    int r = x/Bc, c = x%Bc;
    float s = 0.f;
    for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
    Sij[r][c] = s*scale;
}
__syncthreads();

• mi = max(mi, Sij.max(-1))
• li = rescale(li) + exp(Sij - mi).sum(-1)
• Oi = rescale(Oi) + softmax(Sij) @ Vj

// j = 1, i = 0
if (tid < Br) {
    float m_ij = -INFINITY;
    for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
    float m_new = fmaxf(mi[tid], m_ij);
    float alpha = expf(mi[tid] - m_new);
    float lsum = 0.f;
    for (int c = 0; c < Bc; c++) {
        Sij[tid][c] = expf(Sij[tid][c] - m_new);
        lsum += Sij[tid][c];
    }
    float l_new = alpha*li[tid] + lsum;
    for (int d = 0; d < D; d++) {
        float pv = 0.f;
        for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
        Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
    }
    mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();

• m[i], L[i], O[i] = mi, li, Oi

// j = 1, i = 0
for (int x = tid; x < Br*D; x += blockDim.x)
    O[(0*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[0*Br+tid] = mi[tid]; l[0*Br+tid] = li[tid]; }
__syncthreads();

1.1 Inner iteration i = 1

It carries partial mi, li, Oi from HBM

• Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]

// j = 1, i = 1
for (int x = tid; x < Br*D; x += blockDim.x) {
    Qi[x/D][x%D] = Q[(1*Br + x/D)*D + x%D];
    Oi[x/D][x%D] = O[(1*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[1*Br+tid]; li[tid] = l[1*Br+tid]; }
__syncthreads();

• Sij = Qi @ Kj.T

// j = 1, i = 1
for (int x = tid; x < Br*Bc; x += blockDim.x) {
    int r = x/Bc, c = x%Bc;
    float s = 0.f;
    for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
    Sij[r][c] = s*scale;
}
__syncthreads();

• mi = max(mi, Sij.max(-1))
• li = rescale(li) + exp(Sij - mi).sum(-1)
• Oi = rescale(Oi) + softmax(Sij) @ Vj

// j = 1, i = 1
if (tid < Br) {
    float m_ij = -INFINITY;
    for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
    float m_new = fmaxf(mi[tid], m_ij);
    float alpha = expf(mi[tid] - m_new);
    float lsum = 0.f;
    for (int c = 0; c < Bc; c++) {
        Sij[tid][c] = expf(Sij[tid][c] - m_new);
        lsum += Sij[tid][c];
    }
    float l_new = alpha*li[tid] + lsum;
    for (int d = 0; d < D; d++) {
        float pv = 0.f;
        for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
        Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
    }
    mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();

• m[i], L[i], O[i] = mi, li, Oi

// j = 1, i = 1
for (int x = tid; x < Br*D; x += blockDim.x)
    O[(1*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[1*Br+tid] = mi[tid]; l[1*Br+tid] = li[tid]; }
__syncthreads();

Counting trips across HBM:

• $K$ and $V$ are read once
• $Q$, $O$, $m$ and $l$ are read twice, once per tile of $K$ and $V$, with $O$, $m$ and $l$ written back twice as well.
  • Why are $m$ and $l$ moved? Every step but the last leaves $O$ unfinished, and $m$ and $l$ are the running max and running sum the next step needs to carry on. Only after the final tile of $K$ and $V$ is $O$ the real answer.

Flash Attention 2 improves on it by putting $Q$ on the outer loop, so each block of $O$ is written just once.