M4.3 - IO-Aware Attention and FlashAttention-1: Online Softmax, Tiling, and Exactness

M4.3 - IO-Aware Attention and FlashAttention-1: Online Softmax, Tiling, and Exactness

0x00 From Exact Attention To An IO-Aware Algorithm

M4.1 separated a Transformer into parameter, state, activation, operation, transfer, and K/V ledgers. M4.2 then showed why those logical ledgers do not determine one GPU runtime: the selected kernel still decides which objects cross HBM, which tiles remain on chip, and how much parallel work reaches the SMs.

M4.3 applies those contracts to one exact attention algorithm. Its owner question is:

How can exact dense attention stream score and value tiles without globally materializing the score or probability matrices, and what proof, saved state, IO schedule, and backward contract make that implementation correct?

The distinction inside that question is the foundation of the article:

quadratic score relations:
    part of exact dense attention

global [S, S] score and probability tensors in HBM:
    one execution plan, not a mathematical requirement

FlashAttention-1 did not make dense attention linear-time. It changed which intermediate objects had to survive outside the tile that produced them. Understanding that change requires more than seeing the final kernel pseudocode. It requires reconstructing the sequence of ideas that made the schedule possible.

0x00A Two Responses To The Long-Sequence Attention Problem

The compact one-head equation hides two different problems:

$$O = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_h}}\right)V.$$

First, dense attention evaluates a quadratic set of query-key relations. Second, a conventional program may also materialize the resulting score and probability matrices in HBM. Those are related costs, but they are not the same object.

One line of research changed the relation itself. Sparse patterns, low-rank approximations, kernelized or linear attention, and other approximations can reduce arithmetic or storage under their own model and quality contracts. The motivation section of the FlashAttention paper makes a narrower systems observation: a better asymptotic label does not by itself guarantee wall-clock speedup. Kernel overhead, unsupported operations, hardware utilization, memory movement, and the realized workload still matter. This article does not audit those approximate methods; it follows the other branch.

The second branch keeps the exact dense relation and asks whether its conventional storage schedule is necessary.

The enabling ideas arrived in stages:

1. Safe softmax subtracts a row maximum before exponentiation. It protects the denominator calculation, but the conventional algorithm still obtains row-wide statistics in separate passes.
2. The 2018 Online Normalizer paper showed how to update the running maximum and normalization sum together. The compact state (m,l) represents every score seen so far, but softmax alone still needs a later pass to emit every normalized probability.
3. Rabe and Staats showed that exact attention and self-attention do not intrinsically require a stored quadratic attention matrix. Their practical accelerator algorithm used query/key chunking, stable summaries, and checkpointing to exchange recomputation for memory.
4. FlashAttention-1 joined the online normalization state to tiled matrix multiplication, fused masking, softmax, optional dropout, and value accumulation, and reconstructed local score/probability tiles during backward. Its changed object was not the dense attention relation; it was the HBM schedule around that relation.

This history prevents two common collapses:

exact compact-memory attention != automatically an IO-optimal GPU kernel
online softmax                 != a complete attention implementation

Online normalization supplies the state algebra. Exact chunked attention shows that the global matrix is avoidable. FlashAttention-1 turns those facts into an IO-aware GPU schedule.

0x00B The Algorithm Staircase

Each step below removes one dependency while preserving a contract needed by the next step.

The original A18 - FlashAttention article correctly built its explanation through this progression rather than presenting FlashAttention as an isolated trick. M4.3 preserves that explanatory method, but pins each claim to the primary papers, separates mathematical state from implementation schedule, and leaves generic GPU execution to M4.2.

0x00C Why The Question Matters Numerically

For full-sequence prefill or training-style attention, one conventional dense schedule materializes score and probability objects:

S_score: [B, H_q, S, S]
P_attn:  [B, H_q, S, S]

For the running MHA fixture:

B = 1, S = 4096, H_q = H_kv = 32, d_h = 128
b_qkv = b_score = b_prob = b_out = 2 bytes

each Q/K/V/O payload is 32 MiB, while each full S_score/P_attn payload is 1 GiB. A separate-stage program that writes and rereads both quadratic objects therefore schedules 4 GiB of visible S/P payload movement before counting Q/K/V/O, masks, dropout state, workspace, cache transactions, or replay.

declared stored payload width != accumulator width != saved-state width
logical stage traffic        != measured DRAM transactions

0x02 derives the complete fixture and the symbolic byte ledger. The preview here establishes only the scale of the materialization problem.

FlashAttention answers the owner question this way:

The dense attention relation is required.
The global S/P matrices in HBM are not.

0x00D Background Knowledge: Exact Attention Can Have Different Execution Plans

GPU memory is hierarchical. Registers and on-chip shared memory are small and fast; HBM is much larger but more expensive to access. A mathematically compact expression can therefore be slow when its program writes large intermediates to HBM and reads them back for the next operation.

FlashAttention does not approximate or sparsify the dense equation studied in this article. It changes the schedule:

standard materialized plan:
    write score matrix -> read for softmax -> write probabilities -> read for PV

IO-aware tiled plan:
    load Q/K/V tiles -> maintain online-attention state on chip -> accumulate output

The enabling background is tiled matrix multiplication plus a mergeable softmax recurrence that summarizes earlier key tiles without retaining every score. M4.2 owns the complete hardware hierarchy, tiling, fusion, and profiler vocabulary. M4.3 owns the proof that this particular compact state still represents the exact dense row.

0x01 Claim Status, Owner Contract, And Scope

Claim and evidence classes

The historical implementation anchor for this article is the official FlashAttention repository at tag v1.0.9, commit 6d48e14a. The evidence manifest retains the full commit identifier. Current repository HEAD and current framework SDPA behavior are separate contracts. No local FA1 GPU benchmark is claimed here.

What M4.3 owns, consumes, and defers

Repeated vocabulary is allowed; repeated derivation is not. M4.3 can consume the statement that HBM movement matters, then derive which attention objects move. M4.4 can consume the final online state and FA1 schedule, then derive a different work partition without re-proving softmax.

Scope constraints

This is not a generic CUDA programming guide.
This is not a survey of approximate attention methods.
This is not a current-framework benchmark report.
This is not proof that every SDPA call runs a FlashAttention kernel.

The target is narrower and deeper:

derive the invariant that makes tiled exact attention possible
prove what each compact row-state field represents
trace the actual FA1 paper schedule separately from teaching pseudocode
show which HBM objects disappear and which state replaces them
derive the backward reconstruction and numerical contract
version implementation and runtime claims

0x02 Logical Graph, Materialized Schedule, And Dtype Contract

The attention equation names values and dependencies. It does not require each named value to become a globally stored tensor. This section fixes the semantic graph first, then declares one deliberately materialized baseline schedule.

0x02A The Semantic Graph Includes Mask And Dropout Order

Let g(h) map query head h to its K/V head. For MHA, g(h)=h; the running fixture uses MHA. For each batch item, query head, query position, and key position:

$$S^{\mathrm{raw}}_{b,h,i,j} = \tau\left\langle Q_{b,h,i,:}, K_{b,g(h),j,:} \right\rangle, \qquad \tau=\frac{1}{\sqrt{d_h}}.$$

Masking changes which scores participate in the row denominator:

$$S=\operatorname{mask}\!\left(S^{\mathrm{raw}}\right), \qquad P_{b,h,i,j} = \frac{\exp(S_{b,h,i,j})} {\sum_k\exp(S_{b,h,i,k})}.$$

During training with dropout probability p, let D be a Bernoulli keep mask. Dropout is applied to normalized probabilities, not to logits before the denominator:

$$P^{\mathrm{use}} = \begin{cases} D\odot P/(1-p), & \text{training with dropout},\\ P, & \text{dropout disabled or inference}, \end{cases}$$ $$O_{b,h,i,:}=\sum_j P^{\mathrm{use}}_{b,h,i,j}V_{b,g(h),j,:}.$$

Renderer-safe semantic fallback:

S_raw = tau * Q @ K.T
S     = mask(S_raw)                 # before row normalization
P     = softmax_row(S)
P_use = dropout(P)                  # training only; identity otherwise
O     = P_use @ V

The logical shapes are:

These shapes describe values. They do not prove that S, P, or P_use owns an HBM allocation.

0x02B One Declared Materialized Forward Schedule

The baseline below is a no-dropout visible-forward schedule with three separate global stages. It may fuse mask application into the score or softmax stage; either choice preserves the semantic rule that masking precedes row normalization.

stage 1: read Q,K       -> compute and mask S -> write S
stage 2: read S         -> row softmax        -> write P
stage 3: read P,V       -> matrix product     -> write O

This is one comparison schedule, not a claim that every non-FlashAttention implementation materializes both matrices. A framework can fuse some stages, retain an intermediate in cache, choose another representation, or recompute. Those choices define different transfer rows.

0x02C Birth, Last Use, And Storage Are Different Questions

For the declared forward schedule:

Training can extend the lifetime of P, a log-normalizer, mask metadata, or RNG state into backward. This forward baseline does not pretend those objects are free; it defers their exact saved-state and replay contract to 0x0D and the historical source trace.

0x02D Declare Numerical Roles Instead Of One Universal b

One phrase such as “bf16 attention” is insufficient for a byte ledger. The following symbols identify storage or execution roles independently:

The running example will declare b_qkv=b_score=b_prob=b_out=2 and b_acc=b_state=4. This is a teaching fixture, not evidence that all historical or current kernels use those exact internal representations. Stored payload, operator input dtype, multiply precision, reduction precision, and saved-state dtype remain separate contracts.

0x02E Symbolic Payload And Transfer Ledger

Define element counts:

$$\begin{aligned} E_Q &= BH_qN_qd_h, & E_O &= BH_qN_qd_h,\\ E_K &= BH_{kv}N_kd_h, & E_V &= BH_{kv}N_kd_h,\\ E_S &= BH_qN_qN_k, & E_P &= BH_qN_qN_k. \end{aligned}$$

The globally stored payload of an object is its element count times its declared stored width. The visible transfer ledger for the three-stage schedule is therefore:

$$\begin{aligned} \mathcal B_{\mathrm{materialized}} ={}&E_Qb_{qkv}+E_Kb_{qkv} +2E_Sb_{score}\\ &+2E_Pb_{prob}+E_Vb_{qkv}+E_Ob_{out}. \end{aligned}$$

The factors of two are not dtype multipliers. They are one write plus one later read for each globally materialized quadratic object.

B_materialized =
    read(Q) + read(K)
  + write(S) + read(S)
  + write(P) + read(P)
  + read(V) + write(O)

If a materialized mask, dropout mask, saved probability, workspace, spill, or another copy crosses the boundary, add a named directional row. Do not hide it inside b_score or an unexplained constant.

0x02F Running MHA Fixture, Reconciled

Declare:

B = 1, H_q = H_kv = 32
N_q = N_k = 4096, d_h = 128
b_qkv = b_score = b_prob = b_out = 2 bytes
b_acc = b_state = 4 bytes

Then:

The score/probability cycle is:

$$2(1\ \mathrm{GiB})+2(1\ \mathrm{GiB})=4\ \mathrm{GiB}.$$

One visible read of each Q/K/V payload plus one O write contributes 128 MiB, so:

$$\mathcal B_{\mathrm{materialized}} =4\ \mathrm{GiB}+128\ \mathrm{MiB} =4.125\ \mathrm{GiB/layer}.$$

This is not peak memory: S and P need not be live simultaneously if the allocator can release or reuse storage. It is not training saved-state memory, and it is not measured DRAM traffic. It is the visible complete-payload transfer sum implied by the declared separate-stage schedule.

0x02G Four Byte Ledgers Must Not Collapse

M4.1 established four different questions; M4.3 applies them to attention:

stored payload != live residency
live residency != scheduled logical transfer
scheduled logical transfer != measured transactions

M4.3 will derive which attention rows disappear under FA1. M4.2 remains the owner of profiler interpretation, caches, sectors, launch evidence, and the mapping from a logical schedule to physical transactions.

0x03 Paper Figure 1 And The Inherited Hardware Boundary

M4.2 already owns the GPU hierarchy, residency constraints, coalescing, occupancy, tiling tradeoffs, and Roofline evidence. M4.3 consumes only this boundary:

an intermediate that survives a global stage boundary needs durable storage;
a fused schedule can consume a temporary tile before that tile becomes a
global score or probability matrix.

Source: Dao et al., FlashAttention, arXiv v2, Figure 1. The left panel’s capacities and bandwidths and the right panel’s GPT-2 attention timings are source-reported for the paper’s A100-era setup. The middle panel supplies the durable algorithmic evidence used here: K/V and Q blocks cross the HBM/on-chip boundary, local score work is consumed inside the fused schedule, and output blocks return to HBM. The drawing is not a production buffer-allocation map, an assertion that one block owns one SM, or a local benchmark.

0x04 Why Matmul Tiles Cleanly

Matrix multiplication already has the tiling pattern FlashAttention wants to reuse.

For:

C = A @ B

a tiled kernel can compute one C block by streaming over the shared inner dimension:

for each k tile:
    load A_i,k
    load B_k,j
    C_i,j += A_i,k @ B_k,j

This works because partial sums are associative:

C_partial_new = C_partial_old + A_tile @ B_tile

The accumulator only needs the current partial sum. It does not need the full product matrix to exist in HBM.

Attention wants the same execution shape:

stream K/V blocks
update O block
avoid storing S_score and P_attn

But attention adds the obstruction that plain matmul does not have:

softmax is row-normalized across every allowed key in the row.

0x05 Why The Denominator Blocks A Naive Tile Schedule

Take one non-empty score row x = [x_1, ..., x_V] after masking. For now, assume it contains at least one finite value and no NaN or +inf. Native softmax defines one distribution over the whole row:

$$y_i = \operatorname{softmax}(x)_i = \frac{\exp(x_i)} {\sum_{j=1}^{V}\exp(x_j)}$$

The numerator for position i is local. The denominator is not: every allowed position contributes to it. If a key tile contains only x_B, that tile cannot emit its final probabilities until it knows the denominator shared with every other tile.

This is why the following tempting schedule is wrong:

softmax(x_B1) @ V_B1
+ softmax(x_B2) @ V_B2
+ softmax(x_B3) @ V_B3

Each local softmax gives its own block probability mass 1. Three non-empty blocks therefore contribute three independently normalized distributions. Exact dense attention requires one row distribution whose total mass is 1. 0x09 gives a numerical counterexample.

The obstruction is narrower than “softmax cannot be tiled.” The actual requirement is:

maintain one row denominator while score tiles arrive in sequence

Native, safe, and online softmax all compute that same denominator. They differ in numerical coordinates, retained state, and how many times the row must be loaded.

0x05A Declared Pass And Access Model

The pass counts below use the model from Milakov and Gimelshein’s online-normalizer paper:

x resides in slow/global memory
scalar reduction state can remain in registers or local on-chip storage
y is written once to slow/global memory
the final output pass recomputes exponentials instead of storing them
one input-row scan means V input loads
one output emission means V output stores

Under this model, “three accesses per value” means two loads of x_i plus one store of y_i. It is a logical array-access count. It is not a claim about cache lines, coalesced transactions, compiler-generated instructions, or profiler counters; M4.2 owns those physical distinctions.

This declared model matters. A schedule that stores every exponential would have another payload and another access pattern. A fused consumer that never writes y would have a different boundary. The comparison here isolates the normalizer itself.

0x06 Native And Safe Softmax Share One Row Distribution

0x06A Native Softmax: Two Scans, Three Accesses

The direct algorithm first computes the raw exponential sum, then reloads each value to emit its normalized output:

def native_softmax(x):
    d = 0.0
    for x_i in x:                 # scan 1: V loads
        d += exp(x_i)

    y = []
    for x_i in x:                 # scan 2: V loads
        y.append(exp(x_i) / d)    # V stores
    return y

Renderer-safe formula:

d_native = sum_j exp(x_j)
y_i      = exp(x_i) / d_native

In exact real arithmetic this is the definition. In finite arithmetic, exponentiating the original coordinates is fragile. A sufficiently large positive value can make exp(x_i) overflow; a row of sufficiently negative values can make every exponential underflow to zero. Either event can make the final division non-finite or undefined.

For example, direct exponentiation of the finite row [1000, 999, -1000] overflows in common floating-point formats even though its probability distribution is well-defined.

0x06B Translation Invariance Produces Safe Softmax

In exact real arithmetic, softmax does not change when the same finite constant c is subtracted from every position. Factor the shared scale out of both numerator and denominator:

$$\begin{aligned} \frac{\exp(x_i-c)}{\sum_j\exp(x_j-c)} &= \frac{\exp(x_i)\exp(-c)} {\sum_j\exp(x_j)\exp(-c)} \\ &= \frac{\exp(x_i)\exp(-c)} {\exp(-c)\sum_j\exp(x_j)} \\ &= \frac{\exp(x_i)}{\sum_j\exp(x_j)}. \end{aligned}$$

Choose the row maximum as the reference:

$$m = \max_j x_j, \qquad d = \sum_j \exp(x_j-m).$$

Then every shifted exponent satisfies:

$$x_j-m \le 0,$$

and at least one maximum term contributes exp(0) = 1. The row [1000, 999, -1000] becomes [0, -1, -2000], with:

d_safe = exp(0) + exp(-1) + exp(-2000)
       ~= 1.367879

Subtracting the maximum prevents positive-exponent overflow for finite inputs. It does not make finite precision exact: negligible tail terms can underflow to zero, additions can round, and low-precision exponentials can differ by implementation.

0x06C Safe Softmax: Three Scans, Four Accesses

The classical safe implementation discovers the maximum before it can accumulate the shifted denominator:

def safe_softmax(x):
    m = -inf
    for x_i in x:                     # scan 1: V loads
        m = max(m, x_i)

    d = 0.0
    for x_i in x:                     # scan 2: V loads
        d += exp(x_i - m)

    y = []
    for x_i in x:                     # scan 3: V loads
        y.append(exp(x_i - m) / d)    # V stores
    return y

Renderer-safe formula:

m      = max_j x_j
d_safe = sum_j exp(x_j - m)
y_i    = exp(x_i - m) / d_safe

The extra stability pass creates the opportunity that online softmax exploits:

Under the declared model, native and online softmax require three accesses per value; safe softmax requires four accesses per value. The difference is one eliminated input scan, not a claim that online normalization performs no additional scalar arithmetic.

Online softmax does not remove the row-wide distribution. It fuses discovery of the maximum and stable denominator into the first scan, then uses the second scan to emit probabilities.

0x07 Online Normalizer: Change The Reference Maximum While Streaming

0x07A A Running Maximum Alone Is Not Enough

Suppose the first j-1 values have been processed under maximum m_(j-1). Their stable denominator is:

$$d_{j-1} = \sum_{k=1}^{j-1}\exp(x_k-m_{j-1}).$$

If x_j raises the maximum, the stored denominator is expressed in the wrong coordinates. Keeping its number unchanged would mix terms shifted by two different maxima. The old sum must be converted to the new reference maximum before the incoming term is added.

0x07B Derive The Scalar Recurrence From One Denominator

Start with the prefix invariant:

$$m_{j-1}=\max_{1\le k<j}x_k, \qquad d_{j-1}=\sum_{k=1}^{j-1}\exp(x_k-m_{j-1}).$$

After reading x_j, choose one maximum for the enlarged prefix:

$$m_j=\max(m_{j-1},x_j).$$

Every old term must now use m_j. Insert and subtract the old maximum inside its exponent:

$$\begin{aligned} \sum_{k=1}^{j-1}\exp(x_k-m_j) &= \sum_{k=1}^{j-1} \exp(x_k-m_{j-1}+m_{j-1}-m_j) \\ &= \exp(m_{j-1}-m_j) \sum_{k=1}^{j-1}\exp(x_k-m_{j-1}) \\ &= \exp(m_{j-1}-m_j)d_{j-1}. \end{aligned}$$

The incoming value contributes exp(x_j-m_j) under that same maximum. Therefore:

$$d_j = d_{j-1}\exp(m_{j-1}-m_j) + \exp(x_j-m_j).$$

Renderer-safe recurrence:

m_j = max(m_(j-1), x_j)

d_j = d_(j-1) * exp(m_(j-1) - m_j)
    + exp(x_j - m_j)

Every rescaling exponent now has a reason:

m_(j-1) - m_j <= 0    re-express old terms under the new maximum
x_j       - m_j <= 0    express the incoming term under that maximum

There are only two cases:

0x07C The Prefix Invariant Is Inductive

For the first finite value, initialize:

m_1 = x_1
d_1 = 1

because the stable denominator of a one-element prefix is exp(x_1-x_1) = 1. This base case matters: starting the loop at x_2 with d = 0 silently drops the first element.

Assume the invariant holds through j-1. The recurrence above converts every old contribution from reference m_(j-1) to reference m_j, then adds the new contribution under m_j. Thus:

$$d_j=\sum_{k=1}^{j}\exp(x_k-m_j).$$

At j=V, the state is exactly the safe-softmax state:

m_V = max_j x_j
d_V = sum_j exp(x_j - m_V)

The second scan can therefore emit the same real-arithmetic distribution as safe softmax:

$$y_i=\frac{\exp(x_i-m_V)}{d_V}.$$

0x07D Executable Scalar Reference

The safest pedagogical reference initializes from the first unmasked element instead of evaluating an indeterminate sentinel expression. That element must be finite: a NaN or +inf goes to the exceptional-value policy wherever it appears in the row, and only a row whose scores are all masked -inf reaches the all-masked policy:

def online_normalizer(x):
    if len(x) == 0:
        return empty_row_policy()

    first = first_index(i for i in range(len(x)) if x[i] != -inf)
    if first is None:
        return all_masked_row_policy()
    if not isfinite(x[first]):
        return exceptional_value_policy(x[first])

    m = x[first]
    d = 1.0

    for x_j in x[first + 1:]:
        if x_j == -inf:               # masked contribution is zero
            continue
        if not isfinite(x_j):
            return exceptional_value_policy(x_j)
        m_new = max(m, x_j)
        d = d * exp(m - m_new) + exp(x_j - m_new)
        m = m_new

    return m, d


def online_softmax(x):
    m, d = online_normalizer(x)       # scan 1
    return [exp(x_i - m) / d for x_i in x]  # scan 2

-inf is the algebraic maximum identity, but it must not be replaced with a guessed finite sentinel such as -1000. A literal implementation initialized as m = -inf, d = 0 also needs an identity-safe first-element branch: an all-masked first value can otherwise form -inf - (-inf), which is NaN in floating-point arithmetic.

0x07E The Scalar Update Already Contains The Block Update

Let processed set A and incoming block B each summarize their own values:

$$m_A=\max_{j\in A}x_j, \qquad d_A=\sum_{j\in A}\exp(x_j-m_A),$$ $$m_B=\max_{j\in B}x_j, \qquad d_B=\sum_{j\in B}\exp(x_j-m_B).$$

Choose the one denominator reference for their union:

$$m_{A\cup B}=\max(m_A,m_B).$$

The same coordinate conversion gives:

$$d_{A\cup B} = \exp(m_A-m_{A\cup B})d_A + \exp(m_B-m_{A\cup B})d_B.$$

Renderer-safe block merge:

m = max(m_A, m_B)
d = exp(m_A - m) * d_A
  + exp(m_B - m) * d_B

The scalar recurrence is the special case where B contains one value, so m_B=x_j and d_B=1. The next state-operator section will turn this observation into a tested binary merge and check identity, associativity in real arithmetic, permutation invariance, and finite-precision tolerances. Those stronger properties are not assumed here.

0x07F Numerical Range And Exceptional Rows

For a non-empty prefix containing at least one finite value:

$$1 \le d_j \le j.$$

The lower bound follows because at least one maximum contributes exp(0)=1. The upper bound follows because every shifted exponent is non-positive, so every contribution is at most 1. The online recurrence therefore avoids exponential overflow for this finite-input contract.

That statement has boundaries:

very negative, non-maximum contributions can underflow to zero
addition and exponentials still round in finite precision
different legal merge orders can produce slightly different rounded results
-inf masked entries contribute zero once a finite maximum exists
an empty or all-masked row needs an explicit semantic output policy
NaN and +inf inputs need an explicit runtime policy

In particular, an all-masked row is not repaired by writing m = -inf: subtracting that maximum from a masked -inf score is indeterminate. The state-operator audit owns the executable oracle, empty-state identity, all-masked policy, permutation tests, and tolerance contract.

Online normalization has now removed one stable-softmax scan, but it still needs a second scan to emit every probability. Attention offers another opportunity: instead of storing or emitting those probabilities, carry their weighted value numerator alongside the denominator.

0x08 Online State Algebra: From Sets To Merge Trees

The scalar recurrence proves that a denominator can move between maximum coordinates. Tiled attention needs a stronger statement: independently summarized blocks must be mergeable in any legal partition tree without changing the dense row represented by the state.

From this point, write the denominator d as l to match the FlashAttention paper. Also attach a vector payload v_j \in \mathbb{R}^{d_v} to each allowed score. The resulting unnormalized numerator is called u; later conceptual kernel pseudocode uses the variable name o for the same role.

0x08A A State Is A Canonical Summary Of One Allowed Index Set

For a non-empty set of allowed, finite score positions A, define:

$$\operatorname{state}(A)=(m_A,\ell_A,u_A,n_A),$$

with:

$$m_A=\max_{j\in A}x_j,$$ $$\ell_A=\sum_{j\in A}\exp(x_j-m_A),$$ $$u_A=\sum_{j\in A}\exp(x_j-m_A)v_j, \qquad n_A=|A|.$$

Renderer-safe state definition:

state(A) = (m_A, l_A, u_A, n_A)

m_A = max_{j in A} x_j
l_A = sum_{j in A} exp(x_j - m_A)
u_A = sum_{j in A} exp(x_j - m_A) * v_j
n_A = number of allowed positions represented exactly once

m_A defines the numerical coordinate system. l_A carries scalar probability mass under that coordinate. u_A carries the same mass multiplied by a vector payload. The normalized weighted result is u_A/l_A, but the next state-representation section owns the full derivation of attention output, normalized-state alternatives, and LSE storage.

The count is not needed by the softmax formula. It is included in the proof and oracle contract so duplicated or omitted partition elements cannot pass unnoticed.

0x08B Closure Comes From One Common Maximum

Let A and B be disjoint summaries. Choose one coordinate for their union:

$$m=\max(m_A,m_B), \qquad \alpha=\exp(m_A-m), \qquad \beta=\exp(m_B-m).$$

The denominator of the union is:

$$\begin{aligned} \ell_{A\cup B} &=\sum_{j\in A}\exp(x_j-m) +\sum_{j\in B}\exp(x_j-m) \\ &=\alpha\ell_A+\beta\ell_B. \end{aligned}$$

The numerator uses the same change of coordinates, component by component:

$$\begin{aligned} u_{A\cup B} &=\sum_{j\in A}\exp(x_j-m)v_j +\sum_{j\in B}\exp(x_j-m)v_j \\ &=\alpha u_A+\beta u_B. \end{aligned}$$

Therefore the binary merge is:

$$\operatorname{state}(A)\oplus\operatorname{state}(B) = \left( m, \alpha\ell_A+\beta\ell_B, \alpha u_A+\beta u_B, n_A+n_B \right).$$

Renderer-safe merge:

m     = max(m_A, m_B)
alpha = exp(m_A - m)
beta  = exp(m_B - m)

l = alpha * l_A + beta * l_B
u = alpha * u_A + beta * u_B
n = n_A + n_B

This proves closure: the merged tuple is not merely shaped like a state. Expanding its fields recovers the defining sums over A union B under the correct union maximum.

0x08C The Mathematical Identity Needs An Executable Branch

The empty set has the algebraic summary:

$$e_{d_v}=(-\infty,0,\mathbf{0}_{d_v},0).$$

It should satisfy:

$$e\oplus S=S\oplus e=S.$$

But substituting two empty states into the general formula asks a floating-point program to evaluate exp(-inf - -inf). That subtraction is NaN. The algebraic identity therefore needs an executable validity contract:

def merge(left, right):
    if not left.valid:
        return right
    if not right.valid:
        return left

    m = max(left.m, right.m)
    alpha = exp(left.m - m)
    beta = exp(right.m - m)

    return State(
        valid=True,
        m=m,
        l=alpha * left.l + beta * right.l,
        u=alpha * left.u + beta * right.u,
        count=left.count + right.count,
    )

An empty tile and a fully masked tile both produce valid=False, m=-inf, l=0, u=0, and count=0 in the local oracle. Their merge takes the identity branch and never evaluates an indeterminate exponent.

A fully masked row is a semantic policy boundary because mathematical softmax has no allowed support there. The checked-in synthetic oracle returns an invalid state and a zero vector. That is a declared test policy, not a universal framework or kernel guarantee.

0x08D Commutativity Follows From Symmetric Coordinates

For disjoint A and B, both operand orders choose the same maximum m. They merely exchange alpha with beta and exchange the two terms in each sum. In exact real arithmetic:

$$\operatorname{state}(A)\oplus\operatorname{state}(B) = \operatorname{state}(B)\oplus\operatorname{state}(A).$$

This contract assumes a true partition: every allowed position appears once. Merging overlapping summaries would count the overlap twice. Commutativity does not repair an invalid partition manifest.

0x08E Associativity Is A Common-Denominator Statement

Take three pairwise-disjoint sets A, B, and C. Instead of expanding two nested binary merges, choose their final common maximum directly:

$$M=\max(m_A,m_B,m_C).$$

Every possible grouping must represent:

$$\ell_{ABC} = \exp(m_A-M)\ell_A +\exp(m_B-M)\ell_B +\exp(m_C-M)\ell_C,$$ $$u_{ABC} = \exp(m_A-M)u_A +\exp(m_B-M)u_B +\exp(m_C-M)u_C.$$

Closure says the intermediate merge represents its set union. Applying closure again therefore makes both groupings represent the same canonical state of A union B union C:

$$(S_A\oplus S_B)\oplus S_C = S_A\oplus(S_B\oplus S_C).$$

This is associativity in exact real arithmetic. It is more precise than saying either “softmax is associative” or “softmax is not associative.” Final normalized probability vectors are not directly concatenated and added, but the compact unnormalized set summary has an associative merge.

0x08F Partition And Permutation Invariance

Let disjoint blocks B_1,...,B_R cover every allowed position exactly once. Their common representation is:

$$M=\max_r m_{B_r},$$ $$\ell=\sum_{r=1}^{R}\exp(m_{B_r}-M)\ell_{B_r}, \qquad u=\sum_{r=1}^{R}\exp(m_{B_r}-M)u_{B_r}.$$

Any sequential fold, balanced tree, or block permutation built from oplus reaches this same expression in exact arithmetic. That gives the implementation freedom needed by tiled and parallel schedules without changing the dense mathematical row.

0x08G Finite Precision Is Not The Algebra

The proof above uses real-number addition and multiplication. A finite implementation rounds exponentials, products, sums, and final division. Reassociation can therefore change low bits even when every tree is semantically legal.

The local oracle reports output spread across three legal trees, not runtime or throughput. Its binary32 and binary16 modes round every modeled scalar operation through the named IEEE format; they do not reproduce a particular GPU’s accumulator width, instruction fusion, exponential approximation, or reduction topology.

The tiny float64 differences are ordinary reassociation rounding, not a violation of the exact-arithmetic proof. The intentionally adversarial binary16 fixture uses equal scores and payloads [2048, 1, -2048, 1, 2048, 1, -2048, 1]. Its exact output is 0.5, while the modeled sequential, balanced, and permuted trees produce 0.125, 0.25, and 0.5. This is a sensitivity witness, not a claim that production attention commonly sees that error.

0x08H Executable Oracle Contract

The repository includes:

docs/numinens-evidence/m4-3-io-aware-attention/
    online_state_fixture.json
    online_state_oracle.py
    online_state_report.json
    test_online_state_oracle.py

The tests require:

dense float64 output ~= sequential merge ~= balanced merge ~= permuted merge
left identity == right identity == original valid state
empty and fully masked tiles never evaluate exp(-inf - -inf)
all-masked row follows the declared zero-output test policy
non-multiple final blocks retain every allowed index exactly once
finite-precision order spread is recorded rather than denied
the checked-in report is reproducible from the fixture

Run the contract with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-3-io-aware-attention/test_online_state_oracle.py

The oracle is synthetic_contract_test evidence. It proves the checked equations and policies under its declared arithmetic model. It is not a FlashAttention benchmark, a framework conformance suite, or evidence about one GPU backend.

0x09 Online Attention Numerator And State Alternatives

Online softmax can stream the denominator, but a probability vector still has one entry per key. Attention does not need that vector as an externally visible result. It needs its contraction with the value rows:

$$O=PV=\sum_j p_jv_j.$$

That contraction is what turns a two-pass online normalizer into a one-pass online attention recurrence. This section reconstructs the progression in A18 images 31—36 and 38—44, while separating three objects that are easy to conflate:

working state:
    updated while score/value blocks stream

normalized attention output O:
    the final weighted value vector for one query row

forward-saved state:
    compact data retained so backward can reconstruct probabilities

The derivation follows the exact-attention target in FlashAttention Algorithm 1. The normalized-output recurrence is not a different attention approximation. It is the same numerator and denominator written in different coordinates.

0x09A The Value Numerator Falls Out Of The Dense Attention Equation

For one query row, absorb the scale, mask, and any additive score bias into allowed finite scores s_j. Before dropout:

$$p_j=\frac{\exp(s_j)}{\sum_k\exp(s_k)}, \qquad O=\sum_jp_jv_j.$$

Choose the stable row maximum m=max_j s_j and multiply the numerator and denominator by exp(-m):

$$O = \frac{ \sum_j\exp(s_j-m)v_j }{ \sum_j\exp(s_j-m) } =\frac{u}{\ell}.$$

Renderer-safe derivation:

l = sum_j exp(s_j - m)
u = sum_j exp(s_j - m) * v_j
O = u / l

This is why the denominator and numerator must use the same maximum coordinate. l measures total unnormalized probability mass; every component of u carries that same mass after multiplication by a value component. A running maximum without l cannot normalize the row, and (m,l) without u still cannot emit attention output in one traversal.

For a set A, the relation to the ordinary partition function is:

$$Z_A=\sum_{j\in A}e^{s_j}=e^{m_A}\ell_A, \qquad N_A=\sum_{j\in A}e^{s_j}v_j=e^{m_A}u_A.$$

The exponential factor cancels only because both use the same m_A:

$$O_A=\frac{N_A}{Z_A}=\frac{u_A}{\ell_A}.$$

0x09B The Multi-Pass Bridge Makes The Remaining Dependency Visible

A18 first extends two-pass online softmax into multi-pass attention. In a conceptual row traversal, the first pass discovers the final normalizer; the second revisits each score/value position to accumulate output:

# Pass 1: score stream -> final (m, l)
m, l = online_normalizer(scores)

# Pass 2: revisit or recompute each score, then consume V
O = zeros(d_v)
for score, value in zip(scores, values):
    p = exp(score - m) / l
    O += p * value

This is exact and need not materialize a global probability matrix. It still has a global phase boundary: no probability in pass 2 is final until pass 1 has finished. If scores were not saved, they must be recomputed from Q and K; if they were saved globally, the quadratic intermediate has returned.

Here, “one pass” describes the logical score/value stream for one row. It does not claim that a production FA1 kernel reads every HBM tensor exactly once. Loop order, tile ownership, cache behavior, and writeback belong to the next algorithm-schedule loop.

0x09C Unnormalized State Delays Division Until Finalization

Let disjoint summaries A and B have states (m_A,l_A,u_A) and (m_B,l_B,u_B). The previous state-algebra section proved the common-maximum merge:

$$m=\max(m_A,m_B), \quad \alpha=e^{m_A-m}, \quad \beta=e^{m_B-m},$$ $$\ell=\alpha\ell_A+\beta\ell_B, \qquad u=\alpha u_A+\beta u_B.$$

The attention output represented by the merged state is therefore:

$$O_{A\cup B} = \frac{\alpha u_A+\beta u_B} {\alpha\ell_A+\beta\ell_B}.$$

For one arriving score/value pair (s_i,v_i), its singleton state has m_B=s_i, l_B=1, and u_B=v_i. The element recurrence is:

m_i     = max(m_(i-1), s_i)
alpha_i = exp(m_(i-1) - m_i)
beta_i  = exp(s_i - m_i)

l_i = alpha_i * l_(i-1) + beta_i
u_i = alpha_i * u_(i-1) + beta_i * v_i

O = u_N / l_N

The vector accumulator stays unnormalized. There is no vector division inside the recurrence, only one final component-wise division by l_N. This is a useful mathematical and reference-code representation; it is not a claim that every optimized kernel literally stores a variable named u.

0x09D Normalized State Carries The Partial Output Directly

Suppose each local state stores its normalized partial output:

$$O_A=\frac{u_A}{\ell_A}, \qquad O_B=\frac{u_B}{\ell_B}.$$

Substitute u_A=l_A O_A and u_B=l_B O_B into the unnormalized union output:

$$O_{A\cup B} = \frac{ \alpha\ell_AO_A+\beta\ell_BO_B }{ \alpha\ell_A+\beta\ell_B }.$$

Define each block’s mass in the union distribution:

$$\gamma_A=\frac{\alpha\ell_A}{\ell}, \qquad \gamma_B=\frac{\beta\ell_B}{\ell}.$$

Then:

$$\gamma_A+\gamma_B=1, \qquad O_{A\cup B}=\gamma_AO_A+\gamma_BO_B.$$

Renderer-safe normalized merge:

left_mass  = alpha * l_A
right_mass = beta  * l_B
l          = left_mass + right_mass

gamma_A = left_mass  / l
gamma_B = right_mass / l
gamma_A + gamma_B = 1

O = (left_mass * O_A + right_mass * O_B) / l

This is the normalized-output form used by the paper-level FA1 schedule: old output rows are rescaled by their retained global mass, and the current tile’s local P_tilde @ V contribution is rescaled into the same denominator. The state still needs m and l; O alone does not reveal how much probability mass its partition owns.

For a singleton arriving value, the recurrence becomes:

$$O_i = \frac{ \alpha_i\ell_{i-1}O_{i-1}+\beta_iv_i }{\ell_i}.$$

Compared with (m,l,u), this form performs normalization during every update. It exposes the current partial output directly, but it introduces vector scaling by normalized partition weights at each merge. Neither representation is universally superior; schedule, accumulator dtype, fusion, and downstream use decide which is convenient.

0x09E LSE Is The Logarithm Of The Same Partition Function

The row log-sum-exp is not a fourth normalization rule. It is the ordinary partition function represented in log coordinates:

$$\operatorname{LSE} =\log\sum_je^{s_j} =m+\log\ell.$$

Consequently, a local probability can be reconstructed as:

$$p_j =\exp(s_j-\operatorname{LSE}) =\frac{\exp(s_j-m)}{\ell}.$$

Renderer-safe LSE equivalence:

LSE = m + log(l)
p_j = exp(s_j - LSE)
    = exp(s_j - m) / l

Normalized partial outputs can also be merged directly in log-partition coordinates. Let L_A and L_B denote the two local LSE values:

$$L=\operatorname{logaddexp}(L_A,L_B),$$ $$O_{A\cup B} =e^{L_A-L}O_A+e^{L_B-L}O_B.$$

The two coefficients are again the partitions’ global probability masses. This (LSE,O) form is especially natural when independent workers return normalized partial outputs, as in later split-K/decode reductions. A tiled attention kernel may still prefer (m,l,O) internally because it already computes tile maxima and shifted denominator sums. Mathematical equivalence does not imply equal instruction count, rounding path, or hardware efficiency.

0x09F Working State And Saved State Have Different Jobs

The online update and backward checkpoint should not be described as one universal tuple.

For the running fixture with B=1, H=32, and S=4096, one fp32 scalar per query row/head is 0.5 MiB; saving separate fp32 (m,l) is 1.0 MiB. The output tensor O exists in either case and must not be counted as if LSE replaced it.

LSE is also not a complete training replay contract:

LSE:
    restores the undropped softmax normalizer

still required when applicable:
    causal/padding/additive-mask semantics
    sequence lengths or indexing metadata
    deterministic dropout RNG seed/counter or equivalent mask replay state

Softmax normalization occurs before attention-probability dropout. The saved LSE therefore describes the undropped distribution; it does not encode which probabilities dropout later kept or how they were scaled.

0x09G Executable Equivalence Contract

This section extends the preceding synthetic oracle rather than introducing a second unrelated reference. For every fixture row it now evaluates:

dense float64 attention
two-pass online-normalizer then value accumulation
one-pass unnormalized (m,l,u) merge
normalized (m,l,O) merge
log-domain (LSE,O) merge

The current suite contains 16 tests. Its local_normalization_trap fixture also records the partition masses used in 0x09H.

For the three-block article fixture, all four streamed forms target:

O = [3.554972349613208, 1.2032242390332164]
LSE = 3.19076575192487

Their maximum float64 absolute errors against the dense reference are:

These are tolerance-based semantic checks under Python float64, not bitwise kernel parity. Binary32 and binary16 modes remain arithmetic sensitivity models, not GPU measurements.

0x09H Why Naive Block Softmax Is Wrong

The local-normalization error is not a tiny floating-point effect. It is a different computation.

Let one row be split into two key blocks:

x_B1 = [0, 0]
x_B2 = [10]

Let scalar values be:

V_B1 = [1, 1]
V_B2 = [100]

Wrong local-block attention:

softmax([0, 0]) @ [1, 1] + softmax([10]) @ [100]
= 1 + 100
= 101

Correct global attention:

$$\frac{ e^0 \cdot 1 + e^0 \cdot 1 + e^{10}\cdot 100 } { e^0 + e^0 + e^{10} } \approx 99.99$$

The generated oracle makes the missing quantity explicit:

The masses sum to one, and the weighted contributions sum to 99.9910116300. Adding local outputs gives 101 because it assigns each block unit mass. Averaging them would give 50.5, which is also wrong. A legal merge must carry either (m,l), LSE, or equivalent partition-mass information.

The wrong answer gave each block its own full probability mass. FlashAttention does not do that. It streams blocks while maintaining one row denominator.

This is the reason the article cannot skip the online-softmax proof. The proof is the difference between:

tiling attention exactly

and:

inventing a blockwise approximation

0x0A FlashAttention-1 Algorithm 1: Tensor And Loop Trace

The online-state proof does not choose a hardware schedule. It permits many legal traversal orders because every disjoint key partition can be merged into the same row state. This section now separates three objects that are often collapsed:

1. mathematical invariant:
       merge exact row summaries in any legal partition order

2. FlashAttention-1 paper schedule:
       K/V block j outer; Q/O/l/m block i inner

3. historical CUDA source schedule:
       concrete thread blocks, warps, buffers, and launch constraints

The first object was proved in 0x08-0x09. The second is reconstructed here from FlashAttention Algorithm 1. The third requires the pinned v1.0.9 source-to-kernel trace and remains a later implementation loop. A paper algorithm is not CUDA source, and a teaching loop is not either of them.

0x0AA Reference Invariant Versus Paper Schedule

The row-owned teaching schedule keeps one Q_i block and its unnormalized state (m_i,l_i,u_i) on chip while streaming all K/V blocks:

for i in Q blocks:
    load Q_i
    initialize m_i=-inf, l_i=0, u_i=0

    for j in K/V blocks:
        load K_j, V_j
        S_ij = tau * Q_i @ K_j.T
        mask S_ij before rowmax
        merge tile state into (m_i,l_i,u_i)

    O_i = u_i / l_i
    write O_i and the selected compact saved state

That code is an executable description of the invariant. It is not Algorithm 1. The paper reverses those loops so one K/V block is reused while all Q/O/state blocks visit it:

for j in K/V blocks; for i in Q/O/state blocks

Both traversals target the same dense row. They assign different blocks to the long-lived on-chip role, so they also imply different reloads, writebacks, and parallel ownership.

The Q-block-owner view is useful preparation for M4.4, but this article does not relabel it as FlashAttention-2. A recurrence, a loop order, and a production kernel are separate contracts.

0x0AB Algorithm 1 Line By Line

The paper states Q,K,V in R^(N x d) and an on-chip SRAM capacity M. Its idealized block rule is:

$$B_c=\left\lceil\frac{M}{4d}\right\rceil, \qquad B_r=\min\left( \left\lceil\frac{M}{4d}\right\rceil,d \right).$$

The labels matter:

Q_i, O_i: [B_r, d]
K_j, V_j: [B_c, d]
S_ij:     [B_r, B_c]
m_i,l_i: [B_r]

One early bullet in A18 swaps B_r and B_c; its later tensor walkthrough uses the correct assignment. M4.3 follows the primary paper: row blocks use B_r; column/key blocks use B_c.

Paper-faithful tensor pseudocode, with the article’s scale and mask hook made explicit, is:

# O, l, m are initialized in HBM.
for j in range(T_c):
    K_j = load_hbm(K[j])                  # [b_c_j, d]
    V_j = load_hbm(V[j])                  # [b_c_j, d]

    for i in range(T_r):
        Q_i = load_hbm(Q[i])              # [b_r_i, d]
        O_i = load_hbm(O[i])              # [b_r_i, d]
        l_i = load_hbm(l[i])              # [b_r_i]
        m_i = load_hbm(m[i])              # [b_r_i]

        S_ij = tau * Q_i @ K_j.T          # [b_r_i, b_c_j]
        S_ij = apply_mask_before_rowmax(S_ij, i, j)

        valid_blk = row_has_allowed_score(S_ij)
        # Rows with valid_blk=False merge the identity state and keep
        # their previous O_i,l_i,m_i values unchanged.

        m_tilde = rowmax(S_ij)            # [b_r_i]
        P_tilde = exp(S_ij - m_tilde[:, None])
        l_tilde = rowsum(P_tilde)         # [b_r_i]

        m_new = maximum(m_i, m_tilde)
        alpha = exp(m_i     - m_new)
        beta  = exp(m_tilde - m_new)
        l_new = alpha * l_i + beta * l_tilde

        O_new = (
            (alpha * l_i)[:, None] * O_i
            + beta[:, None] * (P_tilde @ V_j)
        ) / l_new[:, None]

        write_hbm(O[i], O_new)
        write_hbm(l[i], l_new)
        write_hbm(m[i], m_new)

The paper’s main Algorithm 1 displays the dense finite baseline without the scale, mask, or dropout lines. Its implementation discussion and appendix add optional masking and dropout. M4.3 keeps the semantic contract from 0x02A: scale belongs in the score, masking occurs before row maximum, and probability dropout occurs after undropped normalization. A fully masked tile also needs the explicit identity branch from 0x08C; evaluating exp(-inf - -inf) is not a safe implementation.

The normalized O update on paper line 12 is not a new formula. Define:

$$w_{\mathrm{old}} =e^{m_i-m_i^{\mathrm{new}}}\ell_i, \qquad w_{\mathrm{tile}} =e^{\widetilde m_{ij}-m_i^{\mathrm{new}}}\widetilde\ell_{ij}.$$

Then l_new=w_old+w_tile, and:

$$O_i^{\mathrm{new}} = \frac{ w_{\mathrm{old}}O_i +e^{\widetilde m_{ij}-m_i^{\mathrm{new}}} \widetilde P_{ij}V_j } {\ell_i^{\mathrm{new}}}.$$

This is exactly the normalized-state merge derived in 0x09D. Algorithm 1 turns that algebra into a repeated HBM/on-chip schedule.

0x0AC Non-Multiple Tail Tiles Are Part Of The Algorithm

The paper writes square N x d inputs. The synthetic trace uses separate N_q and N_k only to expose row and column axes and both kinds of tail; it does not change the online-state recurrence:

N_q = 7
N_k = 10
d   = 4
B_r = 3
B_c = 4

T_r = ceil(7 / 3)  = 3
T_c = ceil(10 / 4) = 3

The partitions are:

With j outer and i inner, the nine score tiles occur in this order:

Those tiles address 70 = N_q*N_k score positions exactly once. The tail is not silently discarded and is not logically a full padded tile. A production kernel may use predication or padded storage; that is a different physical traffic contract.

0x0AD One Query Row Through Three Key Tiles

The checked fixture extracts query row 6 and gives it this score row before masking:

[0, -2, 3, 1 | -1, 4, 2, 7 | 5, 6]

Only key indices 0..6 are allowed. Therefore key 7 is masked in the middle tile and the final tail tile is fully masked. Masking happens before each tile maximum.

The first tile produces the unnormalized value state:

m = 3
l = 1.191860298604
u = [1.914452, 1.277409, 0.142073, -0.814878]

The second tile has a larger maximum. The old state must change coordinates:

alpha_old = exp(3 - 4) = 0.367879441171
beta_tile = exp(4 - 4) = 1

m = 4
l = 1.580534130840
u = [3.839623, 0.598530, 0.930406, 1.706961]

The third tile is fully masked, so it follows the executable identity branch instead of evaluating an indeterminate exponential. Final normalization gives:

O = [2.429320, 0.378688, 0.588666, 1.079990]

The streamed float64 result differs from the dense reference by at most 4.440892098500626e-16. The trace reports (m,l,u) because that makes each numerator update auditable; Algorithm 1 stores the equivalent normalized partial O=u/l after each tile.

0x0AE HBM Movement Under The Paper Schedule

For the same non-multiple fixture, the synthetic report counts logical scalar words at the paper’s HBM/on-chip boundary. It counts true tail extents, not padded transactions:

Therefore:

nested-loop reads    = 290 scalar words
nested-loop writes   = 126 scalar words
nested-loop transfer = 416 scalar words

line-2 initialization writes:
    O: 28, l: 7, m: 7 = 42 scalar words

total including initialization = 458 scalar words

The paper schedule loads each logical K/V element once in this idealized serial trace, but it reloads every Q/O/state row for each of the three K/V blocks. The row-owned conceptual schedule flips that choice:

The second number is not a benchmark and does not prove a universally better kernel. The fixture is tiny, the block sizes were selected to expose tails rather than optimize a device, and the two schedules induce different capacity, parallelism, reuse, and writeback requirements. Its purpose is to make one fact impossible to miss:

changing loop order changes the traffic ledger

The original paper’s theorem analyzes its own K/V-outer schedule. M4.4 owns the later FA2 work-partition change. Neither count is a cache-line, DRAM-sector, or profiler transaction count.

0x0AF The Tile-Capacity Contract Is An Idealized Model

Algorithm 1’s M is an idealized on-chip capacity measured in scalar elements for its two-level IO model. The proof requires these object families to fit:

$$B_c d=O(M), \qquad B_r d=O(M), \qquad B_rB_c=O(M).$$

A useful shape ledger is:

K_j,V_j:     2 * B_c * d scalar slots
Q_i,O_i:     2 * B_r * d scalar slots
S_ij/P_ij:       B_r * B_c scalar slots, potentially reused
row state:       O(B_r) scalar slots

These terms explain the paper’s asymptotic choice: B_c grows like M/d, and B_r is also bounded by d so the score tile fits. They are not an exact sum for a production kernel. Real launch admissibility must distinguish bytes from elements, stored from accumulator widths, shared memory from registers, double buffering, reduction scratch, occupancy, and architecture limits.

If one uniform stored width b_elem were actually applicable, the first unit conversion would be:

M_words = floor(M_bytes / b_elem)

Mixed dtypes and register-resident accumulators already break that shortcut. M4.2 owns the physical-residency calculation; this section owns only the paper model and tensor constraints.

0x0AG Failure Lab And Executable Contract

The reproducible schedule bundle is:

docs/numinens-evidence/m4-3-io-aware-attention/
    fa1_schedule_fixture.json
    fa1_schedule_oracle.py
    fa1_schedule_report.json
    test_fa1_schedule_oracle.py

It checks block shapes, j-outer visit order, exact score-position coverage, logical movement counts, the masked row trace, dense float64 parity, and the article’s generated values. Run it with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-3-io-aware-attention/test_fa1_schedule_oracle.py

This bundle is synthetic_contract_test evidence. It does not claim that the historical FA1 CUDA code built locally, that a current framework dispatched FA1, or that 416 words correspond to measured DRAM transactions.

The local score and shifted-probability tiles remain temporary:

S_ij:       [b_r_i, b_c_j]
P_tilde_ij: [b_r_i, b_c_j]

They exist long enough to update one exact row state, then disappear. That is the central systems idea:

store compact row state, not global score/probability matrices

0x0B What The Online Schedule Removes From The Baseline

0x02 already derived the 4.125 GiB visible complete-payload ledger for the declared materialized MHA fixture. The online schedule changes the status of the two quadratic objects rather than changing their semantic role:

For the same fixture, one fp32 logsumexp value per row/head has payload:

1 * 32 * 4096 * 4 bytes
= 524,288 bytes
= 0.5 MiB

If saving both m and l as fp32:

1.0 MiB

That comparison does not say total FlashAttention traffic is 0.5 MiB. Q/K/V/O blocks and row state still move according to the paper schedule, and some blocks are reloaded. It says the compact saved state replaces the need to preserve full global S/P matrices.

The original FlashAttention paper supplies a measured version of the same argument. In its A100/GPT-2-medium configuration, FlashAttention performs more floating-point work because backward recomputes local intermediates, yet the reported HBM read/write volume falls from 40.3 GB to 4.4 GB and the forward-plus-backward runtime falls from 41.7 ms to 7.3 ms. The middle panel also shows the limit of the slogan: increasing block size reduces HBM accesses and runtime only until another resource becomes limiting.

Source: Dao et al., FlashAttention, arXiv v2, Figure 2. The left and middle panels use the paper’s A100 setup with sequence length 1024, head dimension 64, 16 heads, and batch size 64; the right panel is a block-sparse sequence-length-4096 experiment. These values validate the IO argument in those configurations, not a universal speedup or traffic constant.

Do not overclaim the FlashAttention traffic as:

read Q/K/V exactly once and write O exactly once

Real traffic depends on tile sizes, loop order, cache behavior, mask/dropout path, hardware, backend, and implementation. The stable claim is narrower and stronger:

the global S/P write-read cycle disappears

0x0C IO Complexity: Derive The Transfer Count Before Taking A Limit

Dense exact attention still performs dense pairwise arithmetic:

QK^T: Theta(N^2 d)
PV:   Theta(N^2 d)

IO complexity asks a different question: how many scalar elements cross between a slow memory and a fast memory of capacity M? The answer depends on the schedule and on what M means. This section uses the two-level scalar-word model of the FA1 paper. It does not substitute aggregate GPU SRAM bytes for M, and it does not call logical words DRAM transactions.

0x0CA Block Counts Are The Missing Link

For square one-head attention, let:

Q,K,V,O: [N,d]
M:       fast-memory capacity in scalar elements
B_c:     key/value rows per outer tile
B_r:     query/output rows per inner tile

Algorithm 1 chooses, up to ceilings and constant factors:

$$B_c=\Theta(M/d), \qquad B_r=\Theta(\min(M/d,d)).$$

The numbers of blocks are:

$$T_c=\left\lceil\frac{N}{B_c}\right\rceil, \qquad T_r=\left\lceil\frac{N}{B_r}\right\rceil.$$

Renderer-safe form:

T_c = ceil(N / B_c)
T_r = ceil(N / B_r)

The K/V-outer schedule loads each K/V row once, but performs T_c complete passes over Q, O, and row state. That repeated-pass count, not the recurrence alone, produces the paper’s IO bound.

0x0CB Forward Logical-Word Ledger

Under the exact instruction-level model used by the local oracle, Algorithm 1 moves:

Therefore:

$$W_{\mathrm{fwd}} =2Nd+3T_cNd+4T_cN+Nd+2N.$$

For the reproducible sensitivity fixture:

N = 4096
d = 64
M = 65,536 scalar elements

B_c = ceil(M / (4d)) = 256
B_r = min(B_c,d) = 64
T_c = 16
T_r = 64

the declared ledger is:

K/V reads                        524,288
Q/O reads plus O writes       12,582,912
m/l reads plus writes            262,144
O/m/l initialization             270,336
----------------------------------------
forward total                 13,639,680 scalar words

The initialization row follows the paper pseudocode. An optimized kernel may fuse or elide initialization traffic, which would define a different constant ledger without changing the asymptotic result.

0x0CC Backward Has The Same Limit But A Different Constant Ledger

Algorithm 4 keeps K/V and the partial dK/dV tile in the outer loop. For every K/V block it rereads Q, O, dO, dQ, m, and l, then writes the updated dQ. Its declared logical-word ledger is:

Thus:

$$W_{\mathrm{bwd}} =7Nd+5T_cNd+2T_cN.$$

For the same N,d,M fixture:

K/V reads                         524,288
Q/O/dO/dQ reads + dQ writes    20,971,520
m/l reads                         131,072
dK/dV final writes                524,288
gradient initialization           786,432
-----------------------------------------
backward total                 22,937,600 scalar words

These formulas count neither RNG-state payload nor mask metadata, workspace, padding, cache lines, sectors, write allocation, or repeated physical loads. They are a bridge from Algorithm 1/4 to asymptotics, not a profiler model.

0x0CD Recover The FA1 Theorems

Since B_c=Theta(M/d):

$$T_c = \Theta\!\left(\frac{Nd}{M}\right).$$

The repeated Q/O term therefore becomes:

$$T_cNd = \Theta\!\left(\frac{N^2d^2}{M}\right).$$

Keeping the unavoidable input/output floor visible gives:

$$W_{\mathrm{FA1}} = \Theta\!\left( Nd+\frac{N^2d^2}{M} \right).$$

W_FA1 = Theta(Nd + N^2 d^2 / M)

The FA1 paper states d <= M <= Nd; in that interval the second term is at least Nd, so Theorems 2 and 5 report the compact form:

$$W_{\mathrm{FA1}} = \Theta\!\left(\frac{N^2d^2}{M}\right)$$

for its forward and backward schedules. Standard materialized attention has Theta(Nd + N^2) IO in the same model. Neither statement changes the dense Theta(N^2d) arithmetic count.

0x0CE What “Optimal” Actually Means

The original paper’s Proposition 3 is sometimes paraphrased too strongly. It rules out one algorithm that is little-o of N^2d^2/M for every M in [d,Nd]; its contradiction uses the endpoint M=Theta(Nd) and the unavoidable Omega(Nd) input/output size. That is not a pointwise matching lower bound at every cache size.

Later work closes that distinction. For exact attention computed with standard matrix multiplication, Saha and Ye derive:

$$\Theta\!\left( \min\left\{ \frac{N^2d}{\sqrt{M}}, \frac{N^2d^2}{M} \right\} \right).$$

This creates two regimes:

large cache: M >= d^2
    FA-style N^2 d^2 / M is tight

small cache: M < d^2
    N^2 d / sqrt(M) is smaller
    a matrix-multiplication/materialization schedule can beat FA1 asymptotically

For backward with standard matrix multiplication, Li et al. derive the tighter all-regime expression:

$$\Theta\!\left( \min\left\{ \frac{N^2d^2+Nd^3}{M}, \frac{N^2d+Nd^2}{\sqrt{M}} \right\} \right).$$

For common N >= d shapes in the large-cache regime, the leading term reduces to the familiar N^2d^2/M. The additional rectangular terms matter when the shape assumptions change. The durable claim is therefore conditional:

FA1's schedule is IO-optimal in the large-cache regime under the cited model.
It is not the pointwise optimal schedule for every possible M and shape.

0x0CF Sensitivity Is A Staircase Before It Is An Asymptote

The exact block formula includes ceilings. Consequently, a small change to M may leave T_c unchanged, then remove an entire pass when B_c crosses a divisor boundary. The generated fixture records these ratios relative to its base case:

Changing B_r without changing B_c does not change the paper’s dominant HBM term: each K/V tile still triggers one complete Q/O traversal. It does change inner-loop count, local score shape, reduction work, launch geometry, and feasible resource usage. Those production effects belong to M4.2 and M4.4, not to this scalar-word theorem.

Element width also changes the legal comparison:

M_elements = floor(M_bytes / bytes_per_element)

Mixed Q/K/V, accumulator, row-state, and scratch widths invalidate a single uniform conversion. A production tile manifest must account for each buffer, register pressure, shared memory, occupancy, and alignment separately.

0x0CG Executable Contract And Failure Lab

io_complexity_oracle.py regenerates the block choices, exact forward and backward word ledgers, cache regimes, later tight bounds, and sensitivity ratios in io_complexity_report.json. The report is a synthetic_contract_test; it contains no GPU measurement.

M4.4 starts where this derivation stops: after the global S/P round trip is removed, work partitioning, warp communication, non-matmul instructions, and decode shape become the next bottlenecks.

0x0D FA1 Algorithm 4: Backward And Recomputation Trace

Training backward is not a footnote to the memory claim. It determines what forward must save, what backward may reconstruct, and which stochastic and masking decisions must replay exactly.

0x0DA Saved State Is Smaller Than Recomputed State

The full Appendix B forward algorithm saves:

O: [N_q,d]
m: [N_q]
l: [N_q]
R: pseudo-random generator state when dropout is active

alongside the original Q/K/V inputs or whatever the surrounding autograd graph retains. A modern implementation may save LSE instead of separate m,l, but the reconstruction identity is the same:

LSE_i = m_i + log(l_i)
P_ij  = exp(S_ij - LSE_i)

Backward reconstructs only local quadratic tiles:

S_ij: [B_r,B_c] = tau * Q_i @ K_j.T
P_ij: [B_r,B_c] = exp(mask(S_ij) - m_i) / l_i
Z_ij: [B_r,B_c] = replayed inverted-dropout multipliers

The local tiles die after contributing to dQ, dK, and dV. No global [N_q,N_k] score, probability, dropout-mask, dP, or dS matrix is required in HBM.

0x0DB Derive The Gradient With Mask, Dropout, And Scale Intact

Let Z contain the inverted-dropout multiplier: 0 for a dropped entry and 1/(1-p) for a kept entry. Forward computes:

$$P^{\mathrm{drop}}=P\odot Z, \qquad O=P^{\mathrm{drop}}V.$$

The value gradient and the gradient entering dropout are:

$$dV=(P^{\mathrm{drop}})^T dO, \qquad dP^{\mathrm{drop}}=dO V^T.$$

Replay the same Z, not a fresh mask:

$$dP=dP^{\mathrm{drop}}\odot Z.$$

For one row, the softmax Jacobian gives:

$$dS_{ij} =P_{ij}\left(dP_{ij}-\sum_k P_{ik}dP_{ik}\right).$$

The contraction looks global, but the saved output removes that dependency:

$$\begin{aligned} D_i &=\sum_k P_{ik}dP_{ik}\\ &=\sum_k P_{ik}Z_{ik}(dO_i^T V_k)\\ &=dO_i^T\left(\sum_k P_{ik}Z_{ik}V_k\right)\\ &=dO_i^T O_i. \end{aligned}$$

Renderer-safe contract:

D_i = rowsum(dO_i * O_i)

dPdrop_ij = dO_i @ V_j.T
dP_ij     = dPdrop_ij * Z_ij
dS_ij     = P_ij * (dP_ij - D_i[:,None])

dV_j += Pdrop_ij.T @ dO_i
dQ_i += tau * dS_ij @ K_j
dK_j += tau * dS_ij.T @ Q_i

The tau factor belongs on dQ and dK because S=tau QK^T; it does not belong on dV. Masked entries have P_ij=0, so their dS entries are zero. This remains true only if the same mask coordinates are applied before reconstructing P.

0x0DC Algorithm 4 Line By Line

The paper’s backward schedule is K/V-outer, matching its forward traversal:

Paper-faithful pseudocode with explicit tails is:

for j0, j1 in kv_ranges(N_k, B_c):
    K_j = load(K[j0:j1, :])
    V_j = load(V[j0:j1, :])
    dK_j = zeros([j1-j0, d])
    dV_j = zeros([j1-j0, d])

    for i0, i1 in query_ranges(N_q, B_r):
        Q_i  = load(Q[i0:i1, :])
        O_i  = load(O[i0:i1, :])
        dO_i = load(dO[i0:i1, :])
        dQ_i = load(dQ[i0:i1, :])
        m_i, l_i = load(m[i0:i1]), load(l[i0:i1])

        S_ij = tau * Q_i @ K_j.T
        S_ij = reapply_same_mask(S_ij, i0, j0)
        P_ij = exp(S_ij - m_i[:, None]) / l_i[:, None]

        Z_ij = replay_dropout(R, i0, j0, shape=P_ij.shape)
        Pdrop_ij = P_ij * Z_ij
        dPdrop_ij = dO_i @ V_j.T
        dP_ij = dPdrop_ij * Z_ij
        D_i = rowsum(dO_i * O_i)
        dS_ij = P_ij * (dP_ij - D_i[:, None])

        dV_j += Pdrop_ij.T @ dO_i
        dQ_i += tau * dS_ij @ K_j
        dK_j += tau * dS_ij.T @ Q_i
        store(dQ[i0:i1, :], dQ_i)

    store(dK[j0:j1, :], dK_j)
    store(dV[j0:j1, :], dV_j)

This is Algorithm 4’s tensor contract, not a claim about current CUDA source, thread ownership, vectorization, or accumulation dtype. The pinned historical source trace remains a separate evidence task.

0x0DD A Non-Multiple Gradient Trace

The executable fixture uses:

N_q=4, N_k=5, d=3, B_r=2, B_c=3
query blocks: [2,3], [2,3]
K/V blocks:   [3,3], [2,3]
score visits: [2,3], [2,3], [2,2], [2,2]

Its mask is causal-like for keys 0..3; key 4 is never allowed. The first query block against the final K/V tail is fully masked, while the second has one allowed coordinate. The fixed dropout multiplier matrix stands in for a deterministically replayed RNG stream.

The generated row contractions are:

D = [-0.066667, 0.607831, -0.130259, -0.592483]

and the invariants are:

max |D - rowsum(P*dP)|       = 1.110223e-16
max masked |dS|              = 0
unused key-4 max |dK|,|dV|   = 0, 0
tiled vs dense dQ/dK/dV error = 0, 0, 0
finite-difference max error   = 6.364323e-11

Under the same logical scalar-word model used in 0x0C, the fixture moves:

K/V reads                              30
Q/O/dO/dQ reads                        96
m/l reads                              16
dQ writes                              24
dK/dV final writes                     30
dQ/dK/dV initialization                42
-----------------------------------------
backward total = 238 scalar words

That count is not a cache-line trace. It deliberately excludes RNG-state payload, mask metadata, workspace, transactions, and source-level kernel details.

0x0DE Gradient Ownership Is The Runtime Problem Left Behind

The K/V-outer paper schedule gives one outer owner all query contributions to one dK/dV tile. That makes dK/dV finalization local, but dQ is shared across all K/V tiles and must be resumed and written after each outer visit:

outer K/V ownership:
    local dK_j,dV_j reduction
    repeated dQ_i HBM read-modify-write

query ownership alternative:
    local dQ_i reduction
    dK_j,dV_j need cross-owner reduction or atomics

Neither ownership choice changes the calculus. It changes synchronization, traffic constants, available parallelism, and reduction placement. FA2’s repartitioning belongs to M4.4; this article establishes the dependency that forces the choice.

0x0DF Replay, Integration, And Failure Closure

Saving an RNG seed is not automatically sufficient. Backward must reconstruct the same random draw for each logical (query,key) coordinate. Changing tile order, vector width, batch/head mapping, or counter advancement without a coordinate-stable RNG contract can silently generate a different dropout mask.

Similarly, the mask contract includes all coordinate transforms:

causal offset
padding lengths
packed-sequence boundaries
additive bias coordinates
query/key tail extents

fa1_backward_oracle.py regenerates the forward saved state, dense analytic backward, tiled Algorithm 4 backward, finite-difference checks, visit order, and logical movement in fa1_backward_report.json. This is a synthetic_contract_test, not a GPU benchmark.

0x0DG The Paper Measures Both Sides Of The Claim

The IO argument predicts two observable consequences: avoiding global S and P storage should reduce memory growth, and reducing HBM traffic can reduce runtime even though backward recomputes local tiles. The original paper plots those consequences together.

Source: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Figure 3. The left panel is source-reported forward-plus-backward runtime; the right panel is attention memory. The models, baselines, kernels, and hardware belong to the paper’s benchmark setup, so the reusable evidence is the joint runtime-and-memory trend, not a universal speedup factor.

left:  extra recomputation does not prevent lower measured runtime
right: compact saved state avoids quadratic attention activation storage

The figure does not say dense attention became linear-time, and it does not measure the local fixture above. It shows that one exact equation can have a materially different training-memory schedule and measured runtime.

Historical v1.0.9 Source Trace: Python To CUDA

The paper algorithms explain why exact attention can stream. They do not by themselves specify the tensor layouts admitted by a library, what autograd saves, which shapes a binding rejects, how a launch is specialized, or what a historical implementation tested. This section follows those contracts through FlashAttention tag v1.0.9, pinned to commit 6d48e14a6c2f551db96f0badc658a6279a929df3.

The evidence class is version_scoped_implementation: source-inspected, not locally executed. The archive and selected source spans are hash-addressed in fa1_source_trace.json. The historical extension was not built here, its CUDA tests were not run here, and no kernel timing or profiler counter is inferred.

The Build Surface Fixes The Historical Vocabulary

At this pin, setup.py builds one flash_attn_cuda extension from the binding file csrc/flash_attn/fmha_api.cpp, head-dimension-specific forward and backward translation units, and the block-sparse paths. The build checks for CUDA 11 or newer and emits architecture targets including sm75 and sm80, with sm90 added when the detected toolkit is new enough. Those facts matter because a path copied from a later source tree, such as flash_api.cpp, would not describe this revision.

The static call chain is:

public unpadded Python API
    -> torch.autograd.Function
        -> flash_attn_cuda.fwd / flash_attn_cuda.bwd
            -> fmha_api.cpp admission + allocation
                -> head-dimension / causal / dropout / split launcher
                    -> forward or backward CTA device routine
                        -> dense-reference and race-oriented pytest checks

That chain is an implementation graph, not a performance result. Every arrow means “this source calls or configures the next layer at this commit.”

The Python API Is Already A Tensor And Replay Contract

The historical public functions are explicitly unpadded. They carry token rows plus cumulative sequence lengths rather than a padded [B,S,H,d] tensor:

The wrapper supplies 1/sqrt(d) when no softmax scale is passed and forwards causal, dropout, split, and generator state into the extension. Its autograd classes save:

Q, K, V
forward output O
float softmax LSE
cu_seqlens metadata
dropout RNG state

In compact form, the phase boundary is Q, K, V, O, float LSE, RNG state.

That list connects the source directly to the Algorithm 4 derivation: backward does not retrieve a durable global probability matrix. It receives the compact normalizer and replay state needed to reconstruct local work. In deterministic mode, the wrapper requests num_splits=1; it does not claim that every possible GPU operation is bitwise deterministic.

The optional return_attn_probs output has a deliberately narrow contract. The public docstrings mark it as testing-only and warn that the returned values may use a different scaling. Treating that diagnostic tensor as a stable attention-probability API would erase a source-level warning.

The C++ Binding Turns Mathematical Shapes Into Admission Rules

The binding does more than forward pointers. Its forward checks and allocations make the historical kernel domain explicit:

The rounded extents select and size kernel work; they do not change the logical sequence lengths stored in cu_seqlens. The binding allocates fp32 LSE, an RNG state tensor containing a Philox seed and offset, an optional score diagnostic, and, on looping paths, a temporary fp32 output. Backward similarly allocates a rowwise softmax_d quantity and may allocate fp32 dQ scratch when split work requires another accumulation path.

This distinction is important:

logical attention state:
    Q, K, V, O, LSE, mask/dropout semantics

binding and launch state:
    padded kernel extents, split count, scratch buffers,
    architecture/dtype admission, Philox coordinates

The first line is close to the paper proof. The second is necessary to run one specific implementation.

The Forward Source Preserves The Invariant Without Copying Algorithm 1 Literally

The forward launch template chooses compile-time causal, dropout, and diagnostic-output variants. It also uses an occupancy-based heuristic for splitting the query sequence: more splits can expose more parallel work, but they introduce extra output traffic and reduction state. Therefore num_splits is not a free speedup parameter.

Inside the pinned fmha_fprop_kernel_1xN.h, the source-level order can be summarized as:

map one batch/head/split work unit
load sequence-local Q and K/V fragments
form a QK score fragment
apply causal and tail masking before normalization
merge the fragment into online max/sum state
derive dropout draws from stable sequence coordinates
contract the surviving probabilities with V
rescale or combine partial output state
write O and float LSE; preserve Philox seed/offset

The mathematical invariant is the same one proved in 0x07—0x09. The materialized schedule is not the paper’s nested loops copied into CUDA. A CTA contains fragment layouts, warp reductions, rounded extents, optional sequence splits, and temporary output state. Saying “the source implements Algorithm 1” is valid only at the invariant level, not as a line-for-line scheduling claim.

Backward Makes Recompute And Ownership Concrete

The backward launcher can choose a regular path or a sequence-parallel path. The split path performs the rowwise dO dot O work separately and changes how dQ is accumulated; the tradeoff is more parallelism versus temporary state, atomics or another reduction, and potentially different finite-precision order.

The pinned fmha_dgrad_kernel_1xN_loop.h restores the RNG seed and offset, reconstructs local score/probability work from Q/K/V, LSE, O, and dO, then accumulates dQ/dK/dV. This source trace supports the article’s recomputation contract. It does not imply that the source uses exactly the same ownership schedule as paper Algorithm 4 for every split and head dimension.

The Historical Tests Define Relative Numerical Expectations

The dense reference in tests/test_flash_attn.py upcasts the dense computation, applies mask before softmax, reuses the dropout mask, and then contracts with V. One QKV-packed test family crosses:

2 dtypes when the architecture admits bf16
2 causal modes
6 head dimensions: 128, 64, 80, 40, 32, 16
11 sequence lengths: 97 through 2048, including boundary-adjacent tails
2 dropout probabilities: 0.0 and 0.17

2 * 2 * 6 * 11 * 2 = 528 nominal combinations

Architecture and memory guards can skip cases, so 528 is not a recorded pass count. The source compares FlashAttention error with a low-precision PyTorch baseline rather than asserting one universal absolute tolerance: forward and the testing attention diagnostic are bounded relative to that baseline, while the gradient allowance is wider when sequence splitting changes accumulation. The test also checks the observed dropout fraction. A separate race-oriented family repeats output and gradient comparisons ten times and acknowledges that split dQ accumulation can be nondeterministic while the tested dK/dV path is stable.

These tests were inspected but not executed in this repository. They establish what the historical authors chose to test, not that those checks pass on the current machine.

Source-Trace Publish Contract

The reproducible inspection manifest is:

docs/numinens-evidence/m4-3-io-aware-attention/
    fa1_source_trace.json
    verify_fa1_source_trace.py
    test_fa1_source_trace.py

Given an extracted copy of the pinned source and its downloaded archive, verify all recorded bytes with:

python3 verify_fa1_source_trace.py \
  --source-root /path/to/flash-attention-6d48e14a \
  --archive /path/to/6d48e14a.tar.gz

The manifest pins the archive hash, file hashes, selected line-span hashes, call-chain order, test axes, and anti-claims. Static inspection supports the historical API, saved-state, admission, launch, kernel, and test contracts above. It does not prove that current SDPA dispatches this kernel, that current FlashAttention has the same source organization, or that any locally measured runtime belongs to v1.0.9.

0x0E Masks, Causality, And Dropout

Dense attention is rarely just:

softmax(QK^T) V

A real kernel often handles:

scale by 1 / sqrt(d_h)
causal mask
padding mask
attention bias
dropout during training

Mask placement is part of the invariant.

Masked positions must be removed before:

rowmax
exp
rowsum
value accumulation

For a causal full prefill, query row i cannot attend to key positions j > i. Mathematically:

score[i, j] = -inf for j > i

That statement changes the allowed score pattern. It does not by itself prove how much work a specific kernel skips. Some kernels skip upper-triangular tiles. Some operate on tiles with masked elements. Some have special causal traversal. The implementation determines the actual work.

Dropout is similarly nontrivial. If dropout is applied to attention probabilities during training, backward must preserve the same stochastic semantics. A fused kernel can avoid saving the full dropout mask only if it saves enough RNG state or metadata to regenerate the same effect.

So the implementation contract is:

same masked attention semantics
same dropout semantics
no global S/P matrix

0x0F Backend Boundary: SDPA Is Not Proof

The paper, the historical implementation, and a current framework call expose three different interfaces:

paper contract:
    an exact IO-aware schedule exists under the stated memory model

historical v1.0.9 contract:
    one pinned Python/C++/CUDA implementation admits named layouts and shapes

current framework contract:
    one semantic SDPA call is dispatched by the installed framework build

None is evidence for the next one without another observation. In particular, the current PyTorch documentation says that torch.nn.functional.scaled_dot_product_attention can select among a FlashAttention-2 implementation, a memory-efficient implementation, and a C++ math implementation on CUDA. A successful call proves the requested attention semantics completed. It does not identify which implementation executed, and it does not mean the historical FA1 kernel traced above ran.

0x0FA Freeze The Semantic Call Before Diagnosing Dispatch

The current API takes tensors with logical axes:

Q: [B, ..., H_q, L, E]
K: [B, ..., H_kv, S, E]
V: [B, ..., H_kv, S, E_v]
output: [B, ..., H_q, L, E_v]

The runtime manifest must also freeze storage dtype, device, layout and strides, mask kind and shape, is_causal, dropout_p, scale, and whether GQA is enabled. Those fields affect both semantics and backend eligibility.

import torch.nn.functional as F

out = F.scaled_dot_product_attention(
    q, k, v,
    attn_mask=attn_mask,
    dropout_p=(dropout_p if training else 0.0),
    is_causal=is_causal,
    scale=scale,
    enable_gqa=enable_gqa,
)

The explicit evaluation branch matters: PyTorch documents that SDPA applies dropout according to dropout_p even when the surrounding module is in evaluation mode. The boolean-mask convention also matters: True means that a coordinate participates in SDPA, the inverse of MultiheadAttention.key_padding_mask. Under the documented API, supplying both attn_mask and is_causal=True is an error.

Grouped-query attention adds another admission boundary. The documented constraints include:

H_q % H_kv == 0
number of K heads == number of V heads

and current support is backend-dependent. M4.3 does not re-derive GQA—M5.7 owns that architecture—but a backend report that omits H_q and H_kv cannot explain why one case was admitted and another fell back.

0x0FB Reference Parity And Backend Identity Are Orthogonal

There are two questions:

semantic question:
    does the output match a declared reference under declared tolerances?

execution question:
    which backend and device kernel produced that output?

A math-backend result can pass numerical parity without exercising an IO-aware fused kernel. A fused backend can execute while differing slightly from the math reference because operation order and accumulator policy differ. Therefore neither answer substitutes for the other.

The evidence ladder is:

This separation is why a parity test and a profiler trace belong in the same evidence bundle but answer different questions.

0x0FC Backend Forcing Is An Admission Test

PyTorch’s current preferred control is torch.nn.attention.sdpa_kernel. A single-backend context removes silent fallback to another enabled SDPA backend:

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    flash_out = F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=attn_mask,
        dropout_p=0.0,
        is_causal=is_causal,
        enable_gqa=enable_gqa,
    )

If the fused implementation cannot run, the documentation says the runtime raises a warning explaining why. The diagnostic must preserve that warning or exception together with the complete input manifest. Changing dtype, removing a mask, making a tensor contiguous, or changing head width to make the warning disappear creates a new workload case; it is not a repair to the old measurement.

Forcing is stronger than an automatic call, but it is still an admission observation. It says nothing about latency until warmup, synchronization, raw samples, and the rest of the benchmark protocol are recorded.

0x0FD Profiler Evidence Closes The Execution Question

The profiler must surround the same forced call and include CUDA activity:

from torch.profiler import ProfilerActivity, profile

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
) as prof:
    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        flash_out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
    torch.cuda.synchronize()

kernel_symbols = sorted({
    event.name
    for event in prof.events()
    if "flash" in event.name.lower()
    or "scaled_dot_product" in event.name.lower()
})

The exact symbol vocabulary is version-specific. The report should preserve the raw trace or event table rather than encode one timeless substring as the definition of FlashAttention. The official profiler documentation also warns that CUDA device traces depend on CUPTI availability; when CUDA exists but CUPTI does not, timing may appear in a table while the JSON trace lacks CUDA events. That limitation belongs in the evidence record.

M4.2 owns the interpretation of physical transactions, cache counters, sectors, and rooflines. M4.3 only requires enough profiler evidence to stop a semantic API call from masquerading as a backend identity claim.

0x0FE Timing Is A New Experiment, Not A Side Effect Of Dispatch

A publishable local runtime claim needs the shared numinens.kernel-evidence.v1 bundle from M4.2. At minimum it records:

framework and kernel-library versions
GPU name, architecture, count, driver, and CUDA version
complete Q/K/V workload and stride manifest
requested backend, selected backend, diagnostics, and kernel symbols
reference, tolerances, maximum error, and pass/fail
warmup count and at least 20 synchronized raw timing samples
profiler command, report path, and measured metrics

Median or p95 without the raw samples is not enough to audit the summary. A timing without numerical parity can benchmark the wrong answer. A timing without profiler identity can benchmark a fallback. A kernel symbol without a workload manifest cannot be generalized to another shape.

The executable diagnostic in this article deliberately stops before making a claim it cannot support:

python3 docs/numinens-evidence/m4-3-io-aware-attention/sdpa_backend_diagnostic.py \
  --observed-at 2026-07-27 \
  --trace-output /path/to/sdpa_profile_trace.json \
  --output docs/numinens-evidence/m4-3-io-aware-attention/runtime_observation_unmeasured.json

On the repository host, the probe observed macOS on arm64, Python 3.9.6, no importable PyTorch installation, and no CUDA runtime path. Its result remains:

evidence class: unmeasured
semantic SDPA execution: not run
forced FlashAttention backend: not run
CUDA profiler trace: not run
kernel timing: not run
historical FA1 build/tests: not run

That is useful evidence: it closes the temptation to fill the missing fields with paper numbers or results from another machine. On a compatible CUDA host, the same diagnostic can perform a float64 math-reference comparison, force the current FlashAttention backend, collect profiler symbols, and capture raw CUDA event timings. Publication still requires promoting those observations into the shared validated kernel bundle.

0x0FF Failure Closure

The debugging order is therefore not merely a checklist. It is a chain of non-interchangeable claims:

freeze semantics and tensor manifest
    -> establish reference parity
        -> force backend and retain diagnostics
            -> capture profiler kernel identity
                -> collect synchronized raw timing
                    -> compare only within the declared environment

0x10 What FlashAttention Is Not

FlashAttention is easiest to keep correct by ruling out the common wrong explanations.

It is not sparse attention by default.

Dense FlashAttention still evaluates the dense allowed query-key relation.

It is not low-rank attention.

It does not replace QK^T with a low-rank approximation.

It is not linear attention.

Dense pairwise query-key interactions remain.
Time is still O(S^2 d) for dense attention.

It is not a quality tradeoff.

The attention target is exact dense attention,
modulo floating-point order and precision differences.

It is not automatically faster for every input.

shape, dtype, mask, dropout, hardware, head dimension,
sequence length, and backend dispatch all matter

The precise statement is:

FlashAttention is an IO-aware exact attention schedule.
It reduces HBM traffic and activation memory by avoiding global S/P materialization.

0x11 Oracle-Generated Three-Block Trace

The worked row below is generated from online_state_fixture.json, then checked against the dense float64 reference by test_online_state_oracle.py. It uses:

scores = [0, -2, 3, 1, -1]

values = [
    [ 1, 0],
    [ 0, 2],
    [ 4, 1],
    [ 2, 3],
    [-1, 2],
]

blocks = [[0, 1], [2, 3], [4]]

Each block first computes its own canonical state:

The first merge raises the reference maximum from 0 to 3:

alpha = exp(0 - 3) ~= 0.049787
beta  = exp(3 - 3) = 1

state(B0 union B1):
    m = 3.000000
    l = 1.191860
    u = [4.320458, 1.419482]

The last block has a lower maximum, so the existing state keeps scale 1 while B2 is rescaled by exp(-1-3):

state(B0 union B1 union B2):
    m = 3.000000
    l = 1.210176
    u = [4.302142, 1.456113]

u / l = [3.554972, 1.203224]

The dense reference computes:

m_dense = 3.000000
l_dense = 1.2101759374922962
u_dense = [4.302141995952355, 1.4561130214854774]
O_dense = [3.554972349613208, 1.2032242390332164]

The sequential merge differs from the dense float64 output by at most 8.881784e-16. Balanced and permuted trees also pass the 1e-12 output tolerance. That tolerance is part of the executable contract; the article does not call the floating-point paths bitwise identical.

0x12 Publish-Gate Checks

0x12A Visual Provenance Is Part Of The Claim

A diagram is not evidence merely because it resembles a paper figure. M4.3 uses three distinct visual classes:

The mobile paper variants do not redraw curves, values, or annotations. They crop the original panels, scale them uniformly, and stack them so labels remain readable at the 390 px acceptance viewport. The original desktop raster stays the provenance anchor and both files are hash-pinned.

The reconstructed diagrams are not independent support for their formulas. Their underlying paper section, source pin, or executable fixture remains the claim owner. Conversely, the three paper rasters are not allowed to inherit a local benchmark status from their placement in this article.

paper or fixture establishes claim
    -> semantic-purpose ledger chooses one visual job
        -> desktop/mobile assets expose that job
            -> alt text states the same job without relying on color
                -> hash and article-reference audit closes provenance

visual_coverage_contract.json maps the twelve required semantic jobs onto the complete visual surface. audit_visual_provenance.py regenerates visual_provenance_report.json; the checked report contains 17 responsive reconstructions, 3 source figures, 20 total visuals, and no unresolved structural error. No A18 screenshot raster is reused: A18 supplies explanatory lineage, while the article either cites a primary-paper figure directly or reconstructs the mechanism in a responsive repository-owned asset.

0x12B Gate Ledger

This article should pass these checks before it is treated as publication-quality:

0x12C Final Publication Contract

The publish gate is executable rather than ceremonial. Four surfaces must close together:

formula contract
    balanced non-empty math blocks
    stable notation and required derivation landmarks
    checked fixture values and oracle reports

source contract
    every primary, official, implementation, and local owner source is named
    source-reported measurements remain separated from local observations
    checked artifacts retain their recorded digests

render contract
    local post links and image targets resolve
    desktop and mobile picture variants have non-empty alternative text
    code fences, equations, tables, and figures stay inside the article surface

boundary contract
    M4.3 owns FA1 exactness, schedule, backward reconstruction, and IO analysis
    M4.1/M4.2 remain prerequisite owners
    M4.4-M4.6, M5.1-M5.8, and M8.1 retain their deferred derivations

audit_article.py checks the static publication surface, all marker families, source coverage, evidence hashes, responsive assets, and adjacent-article ownership. It also rejects exact long-paragraph duplication across the local M4.1, M4.2, and M4.4 comparison corpus. Its result does not replace the numerical oracles or a browser render: it verifies that those independent contracts remain connected to the article being published.

0x13 Bridge To FlashAttention-2

M4.3 solved one problem:

the global S/P HBM round trip was not required

After that, the bottleneck moves. A kernel can still have:

too much non-matmul scalar work
too little parallel work for the available SMs
poor warp partitioning
shared-memory communication overhead
causal tile imbalance
backend dispatch limitations

The next article is about that transfer of bottleneck.

M4.4 asks:

Once FlashAttention-1 fixes the HBM schedule,
how does FlashAttention-2 improve work partitioning?
And why does one-token decode become a different kernel problem?

0xFF References

• Local source: Exploring the Transformer Series (18) - FlashAttention
• Local source: M4.1 - Transformer Cost Model
• Local source: M4.2 - GPU Execution Model for Transformer Kernels
• Local handoff: M4.4 - FlashAttention-2 Work Partitioning and the Decode Boundary
• Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv v2
• Milakov and Gimelshein, Online normalizer calculation for softmax, arXiv v2
• Rabe and Staats, Self-attention Does Not Need O(n^2) Memory, arXiv v3
• Barna Saha and Christopher Ye, The I/O Complexity of Attention, or How Optimal is FlashAttention?, arXiv v1
• Xiaoyu Li et al., Fine-grained Attention I/O Complexity: Comprehensive Analysis for Backward Passes, arXiv v1
• Official FlashAttention repository, historical v1.0.9 source at commit 6d48e14a
• PyTorch documentation, torch.nn.functional.scaled_dot_product_attention
• PyTorch documentation, sdpa_kernel backend control
• PyTorch documentation, torch.profiler activity and trace contract
• NVIDIA documentation, CUDA Programming Guide