M4.4 - FlashAttention-2 Work Partitioning and the Decode Boundary

0x00 From IO-Aware Exactness To The Work-Partition Problem

M4.3 established a result that is easy to state and difficult to implement: exact dense attention does not need to materialize its complete score and probability matrices in HBM. Tiled execution can stream Q, K, and V, carry compact row state, and recompute local score tiles in backward.

That result removes a major memory-schedule bottleneck. It does not prove that the remaining kernel uses the GPU efficiently.

Owner question: after FlashAttention-1 removes global S/P traffic, how does FlashAttention-2 change state coordinates and assign work across thread blocks and warps so that exact attention spends more time on useful matrix operations?

This question marks a historical shift. FlashAttention-1 asks how exact attention can avoid unnecessary HBM movement. FlashAttention-2 starts from that IO-aware algorithm and asks why its realized throughput still trails highly optimized matrix multiplication. The answer is not one new attention equation. It is a set of connected scheduling decisions:

FA1 result:
    compact exact row state
    tiled HBM/on-chip schedule
    no global score or probability matrix

remaining FA1-era pressure:
    repeated non-matmul normalization work
    too few independent thread blocks for some prefill shapes
    shared-memory traffic caused by warp ownership

FA2 response:
    retain an unnormalized output numerator
    expose query-row blocks as independent forward work
    split Q across warps so one warp owns an output slice
    assign backward K/V columns to CTAs and reduce shared dQ updates

The durable thesis is:

exactness determines the state that must exist
IO awareness determines where that state may live
work partitioning determines who owns it and whether the GPU stays busy

M4.4 therefore continues the M4 accounting chain rather than restarting the attention derivation:

M4.1: count tensors, FLOPs, bytes, and persistent state
M4.2: map those ledgers to GPU execution and resource limits
M4.3: derive the exact online state and reconstruct the FA1 IO schedule
M4.4: derive the FA2 state and work-ownership changes

One-token decode appears here only as the point where FA2’s query-row axis collapses. The exact K/V-split reduction and runtime backend contract belong to M4.5, while M4.6 owns the Hopper and Blackwell pipeline branch.

0x00A Claim Status, Owner Contract, And Scope

Claim and evidence classes

• Source-local. A claim, explanatory order, or visual idea preserved from local article A19. It is not automatically treated as technically correct.
• Externally verified. A mathematical, algorithmic, or historical claim checked against the FlashAttention-2 paper or another named primary source.
• Version-scoped implementation. Behavior verified in FlashAttention tag v2.0.0 at the historical commit pinned below. It is not a claim about the current repository or every backend.
• Synthetic contract test. A repository-authored oracle or estimator used to test a stated invariant. It is not a GPU measurement.
• Source-reported measurement. A number reproduced from a paper under that source’s hardware, shape, metric, and software conditions.
• Rewrite interpretation. A teaching or boundary decision that connects formulas, schedules, and neighboring modules.
• Unmeasured. A runtime, dispatch, or profiler claim for which this article has no local execution artifact.

What M4.4 owns, consumes, and defers

• Consumes from M4.1: tensor axes, FLOP/byte ledgers, persistent-state accounting, and declared units.
• Consumes from M4.2: thread-block/warp hierarchy, occupancy and wave vocabulary, shared-memory/register pressure, and the distinction between logical traffic and measured transactions.
• Consumes from M4.3: exact online-softmax algebra, the attention numerator, FA1 forward/backward schedules, compact LSE state, and the semantic/runtime evidence boundary.
• Owns in M4.4: FA2’s normalized-to-unnormalized state-coordinate change; non-matmul work ledger; query-block launch geometry; split-K versus split-Q warp ownership; FA2 backward ownership; causal prefill and MQA/GQA mapping; historical v2.0.0 source-to-kernel and training integration; and the paper measurement contract.
• Defers to M4.5: one/few-query K/V splitting, exact partial-state merge, paged-cache addressing, mask programs, SDPA/backend admission, and profiler proof.
• Defers to M4.6: Hopper TMA/WGMMA scheduling, FA3 ping-pong and FP8 paths, Blackwell TMEM/2-CTA MMA, and FA4 pipelines.
• Defers to M5: request scheduling, KV allocation and paging, batching, eviction, and serving-level latency/throughput policy.

Historical evidence contract

The mechanism authority is the FlashAttention-2 paper and the implementation surface released with FlashAttention v2.0.0:

paper:
    arXiv:2307.08691

historical implementation:
    tag:    v2.0.0
    commit: 4f285b354796fb17df8636485b9a04df3ebbb7dc

current repository main:
    useful for later drift and runtime questions
    not evidence for what the historical FA2 release did

The article traces the pinned Python API, C++ admission path, launch specializations, forward/backward kernels, and Triton reference. It does not substitute a current high-level API call for that historical source trace.

Scope constraints

This is a source-linked kernel reconstruction, not a CUDA code dump.
This is not a locally reproduced benchmark unless an artifact says so.
This does not re-prove M4.3's online-state algebra.
This does not derive M4.5's Flash-Decoding merge or backend dispatch.
This does not derive M4.6's FA3/FA4 hardware pipelines.

Background knowledge: a fast algorithm still needs a balanced GPU work decomposition

A GPU kernel exposes work at several levels: thread blocks are scheduled across streaming multiprocessors, warps cooperate inside a block, and specialized units execute matrix operations efficiently when shapes and data movement align. Removing HBM intermediates does not guarantee that this hierarchy stays occupied.

algorithm question:
    how much arithmetic and data movement are required?

partition question:
    which independent tiles can run on different workers?

warp question:
    which operands are shared and which partial results need reduction?

shape question:
    do prefill and decode expose the same parallel axes?

FlashAttention-2 improves the partition of exact tiled attention, especially for long query sequences. One-token decode then exposes a different shape. M4.4 ends at that ownership change instead of treating one implementation as universal.

0x01 What FlashAttention-2 Improves

The FlashAttention-2 paper starts from a very specific observation.

FlashAttention already reduces HBM traffic and memory. But the v1 kernel still did not run as close to optimized GEMM efficiency as the authors wanted. The paper reports that FlashAttention reached only 25-40% of theoretical maximum FLOPs/s in its framing, while optimized matrix multiplication was much closer to the hardware ceiling.

The paper identifies three main improvements:

1. reduce non-matmul FLOPs
2. parallelize attention even within one head across thread blocks
3. partition work between warps to reduce shared-memory communication

The paper reports roughly 2x speedup over FlashAttention and 50-73% of theoretical maximum FLOPs/s on A100 for its tested settings. Treat those numbers correctly:

source-reported
hardware-specific
shape-specific
implementation-specific

The stable lesson is not a fixed multiplier. The stable lesson is:

avoiding HBM traffic is necessary,
but work placement still determines GPU utilization.

The work-axis ledger is:

FA1-style simple prefill grid:
    work blocks ~= B * H_q

FA2-style prefill grid:
    work blocks ~= B * H_q * num_q_blocks

one-token decode without KV splitting:
    work blocks ~= B * H_q

Flash-Decoding-style decode:
    work blocks ~= B * H_q * num_kv_chunks
    plus a final reduction over chunks

The exact grid layout depends on implementation, but the pressure is stable:

When B and H_q are small, adding a sequence axis can decide whether the GPU is busy.

0x02 Non-Matmul FLOPs Are Not Like GEMM FLOPs

Attention contains two large matrix products:

score product: Q @ K.T
value product: P @ V

It also contains row maxima, masks, subtraction, exponentials, row sums, reciprocals or division, frame conversion, and accumulator addition. Those operations do not all use the matrix-multiply path. A total FLOP count can therefore hide the instruction mix that limits a real kernel.

The FlashAttention-2 paper makes this argument with one A100 coordinate. It quotes theoretical peaks of 312 TFLOPs/s for FP16/BF16 matrix multiplication and 19.5 TFLOPs/s for FP32 non-matmul work, then uses their 16x ratio as a motivation to remove scalar work from the inner loop. These are source-reported, hardware-specific peak rates, not a claim that every maximum, exponential, reciprocal, or addition takes exactly sixteen times the latency of one matrix-multiply FLOP.

The useful question is consequently narrower:

Which normalizations follow from the mathematical result, and which follow only from the coordinate system chosen for the loop-carried state?

0x02A One Dense Row Defines Two Equivalent State Coordinates

M4.3 proved the general mergeable state. M4.4 consumes that result and reconstructs the particular coordinate change used by FA2. Let A be one non-empty set of visible key positions for a single query row. With scores s_j and value vectors v_j, define:

$$m_A=\max_{j\in A}s_j,$$ $$l_A=\sum_{j\in A}\exp(s_j-m_A), \qquad u_A=\sum_{j\in A}\exp(s_j-m_A)v_j.$$

The exact normalized attention output over A is:

$$O_A=\frac{u_A}{l_A}.$$

Renderer-safe state definitions:

m_A = max_{j in A} s_j
l_A = sum_{j in A} exp(s_j - m_A)
u_A = sum_{j in A} exp(s_j - m_A) * v_j
O_A = u_A / l_A

The two valid coordinate systems are therefore:

Coordinate	Stored row state	Meaning
normalized output	(m_A,l_A,O_A)	output has already been divided by the current denominator
unnormalized numerator	(m_A,l_A,u_A)	numerator and denominator remain in the same maximum frame

They represent the same row because u_A=l_A O_A. FA2 does not change the softmax distribution. It carries u instead of repeatedly carrying a newly normalized O.

0x02B Two Blocks Must First Share One Maximum Frame

Let A and B be disjoint K/V blocks summarized independently. Their local maxima define different exponential coordinates. Choose the union maximum:

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

Every contribution from A must be multiplied by alpha, and every contribution from B by beta, before they can be added. For the denominator:

$$\begin{aligned} l &=\sum_{j\in A\cup B}\exp(s_j-m) \\ &=\alpha\sum_{j\in A}\exp(s_j-m_A) +\beta\sum_{j\in B}\exp(s_j-m_B) \\ &=\alpha l_A+\beta l_B. \end{aligned}$$

The numerator has the same change of coordinates:

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

The normalized result is then:

$$O=\frac{\alpha u_A+\beta u_B}{\alpha l_A+\beta l_B}.$$

Substitute u_A=l_A O_A and u_B=l_B O_B:

$$O = \frac{\alpha l_A}{l}O_A + \frac{\beta l_B}{l}O_B.$$

Renderer-safe two-block 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
O = u / l

# The same row in normalized coordinates:
O = (alpha * l_A / l) * O_A
  + (beta  * l_B / l) * O_B

This derivation identifies the non-negotiable operation: maximum-frame rescaling remains necessary in either coordinate. Delaying the denominator division is valid. Dropping alpha or beta is not.

0x02C The Sequential FA1-Style And FA2 Updates Are The Same Row

In a sequential kernel, A is all previously visited blocks and B is the current score tile S_blk. Instead of first building a local summary for B, write the current block directly in the new maximum frame:

$$m_{new}=\max\left(m_{old},\operatorname{rowmax}(S_{blk})\right), \qquad \alpha=\exp(m_{old}-m_{new}),$$ $$\widetilde P_{blk}=\exp(S_{blk}-m_{new}),$$ $$l_{new}=\alpha l_{old}+\operatorname{rowsum}(\widetilde P_{blk}).$$

A normalized working state must update both already-visited mass and the new value contribution under l_new:

$$O_{new} = \frac{\alpha l_{old}}{l_{new}}O_{old} + \frac{\widetilde P_{blk}V_{blk}}{l_{new}}.$$

The second division can legally be placed on either side of the value product:

$$\left(\frac{\widetilde P_{blk}}{l_{new}}\right)V_{blk} = \frac{\widetilde P_{blk}V_{blk}}{l_{new}}.$$

That is a real design alternative. Scaling the probability tile performs work over its B_c columns; scaling the matrix-product result performs work over its d_v output columns. Fusion, register layout, and tile shape determine which physical schedule is preferable.

FA2 instead stores u_old=l_old O_old and updates:

$$u_{new}=\alpha u_{old}+\widetilde P_{blk}V_{blk}.$$

Renderer-safe sequential coordinate updates:

normalized:
    O_new = (alpha * l_old / l_new) * O_old
          + (P_tilde @ V_blk) / l_new

unnormalized:
    u_new = alpha * u_old + P_tilde @ V_blk

Substituting u_old=l_old O_old into u_new/l_new recovers the normalized update exactly:

$$\frac{u_{new}}{l_{new}} = \frac{\alpha l_{old}}{l_{new}}O_{old} + \frac{\widetilde P_{blk}V_{blk}}{l_{new}}.$$

The removed work is repeated denominator normalization. The retained work is the maximum, exponentiation, row reduction, and rescaling of the old numerator into the new maximum frame.

0x02D Induction, Final Normalization, And Backward LSE

After visiting blocks 1..j, the unnormalized recurrence maintains the invariant:

$$m^{(j)}=\max_{t\le j}\max_{k\in t}s_k,$$ $$l^{(j)}=\sum_{t\le j}\sum_{k\in t}\exp(s_k-m^{(j)}),$$ $$u^{(j)}=\sum_{t\le j}\sum_{k\in t}\exp(s_k-m^{(j)})v_k.$$

The first valid block establishes the base case. For the induction step, choose m^(j)=max(m^(j-1),m_blk), move the old state and current block into that common frame, and add them. The two-block derivation then proves the invariant for block j. No intermediate division is part of that invariant.

After T_c blocks, one reciprocal normalizes the final vector:

$$O=\frac{u^{(T_c)}}{l^{(T_c)}}.$$

The same denominator gives the row log-sum-exp saved for backward:

$$LSE=m^{(T_c)}+\log l^{(T_c)}.$$

For any reconstructed score s_j, its exact normalized probability before dropout is:

$$p_j=\exp(s_j-LSE).$$

Renderer-safe finalization:

O   = u_final / l_final
LSE = m_final + log(l_final)
p_j = exp(s_j - LSE)       # backward reconstruction before dropout

Saving LSE rather than both m and l is the paper’s second state change. It replaces two row scalars with one row scalar; it does not save a dense probability matrix.

0x02E Worked Two-Block Trace From The Executable Oracle

The repository fixture uses:

block A:
    scores = [1.0, -1.0]
    values = [[1.0, 2.0], [3.0, -1.0]]

block B:
    scores = [2.0, 0.5, -2.0]
    values = [[-1.0, 4.0], [2.0, 0.0], [5.0, -3.0]]

Its generated states are easier to audit as vertical ledgers than as one wide table:

Block A

m = 1.000000
l = 1.135335
u = [1.406006, 1.864665]
O = [1.238406, 1.642391]

Block B

m = 2.000000
l = 1.241446
u = [-0.462161, 3.945053]
O = [-0.372277, 3.177789]

Union A+B

m = 2.000000
l = 1.659112
u = [0.055079, 4.631025]
O = [0.033198, 2.791267]

The dense row, normalized online merge, and unnormalized online merge agree to within 8.89e-16 in the generated binary64 report. The state-coordinate residual max_abs(u-l*O) is also 8.89e-16. This is a synthetic contract test, not a kernel accuracy or performance measurement.

The reproducible artifacts are:

docs/numinens-evidence/m4-4-flashattention2-work-partition/
    state_coordinate_fixture.json
    state_coordinate_oracle.py
    state_coordinate_report.json
    test_state_coordinate_oracle.py

0x02F Empty Rows And Finite Precision Are Separate Contracts

An empty block is the merge identity. An all-masked row is not a normal state with m=-inf, because evaluating exp(-inf - -inf) is indeterminate. A safe state machine carries a validity flag:

empty state:
    valid = false
    m = -inf
    l = 0
    u = 0

merge(empty, state) = state
merge(state, empty) = state

all-masked final row:
    use an explicit API/backend output policy
    do not infer it by evaluating the normal recurrence

The synthetic oracle reports zero payload and LSE=null for an all-masked row so the test remains serializable. That is a fixture policy, not a universal framework contract.

The exact-real-arithmetic coordinate proof is not a promise of bitwise equality. The oracle includes a deliberately ill-conditioned binary32 row with values [16777216, 1, -16777216, 1] and equal scores. Two legal merge orders produce 0.25 and 0.5, while the binary64 dense reference is 0.5. The large gap is intentionally constructed to make non-associativity visible; it is not a claim about typical attention error.

mathematical contract:
    normalized and unnormalized coordinates represent the same exact row

finite-precision contract:
    operation order, accumulation dtype, exp implementation, and fusion may
    change rounding; compare under declared tolerances, not bitwise identity

0x02G Declared Non-Matmul Operation Ledger

The following is a logical scalar ledger per valid query row, not an instruction trace. Let:

T_c = number of K/V blocks
B_c = visible score positions per full K/V block
d_v = value/output width

counting convention:
    one division = one reciprocal + elementwise multiplies

excluded here:
    QK.T and PV matrix products
    masks and tail predicates
    loads and stores
    vector instructions and compiler fusion

Work retained by both normalized and unnormalized coordinates is:

Logical scalar work	Count per row	T_c=32, B_c=128, d_v=128
score shifts	T_c B_c	4,096
score exponentials	T_c B_c	4,096
within-block max comparisons	T_c(B_c-1)	4,064
within-block row-sum additions	T_c(B_c-1)	4,064
running-max comparisons	T_c-1	31
old-frame exponentials alpha	T_c-1	31
denominator frame multiply/add pairs	T_c-1	31 pairs
old output/numerator vector scaling	(T_c-1)d_v	3,968

The state-specific alternatives are:

Normalized state, scale the probability tile

reciprocals:                            T_c
new-contribution normalization muls:   T_c * B_c
old-mass coefficient formation muls:   2 * (T_c - 1)
final output scaling muls:              0

Normalized state, scale the PV result

reciprocals:                            T_c
new-contribution normalization muls:   T_c * d_v
old-mass coefficient formation muls:   2 * (T_c - 1)
final output scaling muls:              0

FA2 unnormalized numerator

reciprocals:                            1
new-contribution normalization muls:   0
old-mass coefficient formation muls:   0
final output scaling muls:              d_v

For the declared 32 x 128 x 128 fixture, the unnormalized schedule removes:

31 repeated row reciprocals
4,096 inner-loop normalization multiplications
62 scalar multiplications used to form normalized old-mass coefficients

and adds:
128 final output-scaling multiplications

net state-specific multiplication delta under either displayed normalized
placement for this equal-width fixture: 4,030 fewer logical multiplies

B_c=d_v=128 makes the two normalized placements equal in this example. They are not equal for every tile shape. More importantly, these logical counts are not physical instruction counts, cycle estimates, or measured speedups. Reciprocal approximations, vectorization, epilogue fusion, compiler scheduling, and register layout can change how the work is realized. Later sections will trace the historical kernel and paper schedule before making implementation claims.

The cost-model translation is now precise:

FA2 removes repeated denominator normalization.
FA2 does not remove stable maxima, exponentials, row reductions, or old-frame
rescaling.
One final normalization and one final LSE evaluation remain.

0x03 Algorithm 1 Is A Tensor Storage Contract

The recurrence in 0x02 answers a mathematical question: what state is necessary to merge the next K/V block exactly? Algorithm 1 answers a different question: where does each tensor live while that recurrence executes, and how long must it remain live?

That distinction is essential. The following source-local image, preserved from A19, is a crop of the paper algorithm. It is useful as the primary schedule reference, but its square N x d notation, abstract “on-chip SRAM,” and real-valued tensors are not a production CUDA storage specification.

Source-local A19 crop of FlashAttention-2 Algorithm 1. The small annotation on lines 9-10 came with the original A19 asset. The primary paper, not the annotation, is authoritative; the crop is not a literal CUDA transcript or local runtime evidence.

0x03A Paper Algorithm, Generalized Trace, And Historical Kernel Are Three Contracts

This section keeps three layers separate:

paper Algorithm 1:
    Q,K,V in R^(N x d)
    abstract HBM and on-chip SRAM
    query-block outer, K/V-block inner

generalized teaching trace:
    separate S_q and S_k
    true tail extents b_r_i and b_c_j
    explicit d_h and d_v
    explicit causal, padding, and all-masked predicates

historical implementation trace:
    FlashAttention tag v2.0.0
    commit 4f285b354796fb17df8636485b9a04df3ebbb7dc
    concrete FP16/BF16 storage, FP32 accumulators, padding, and epilogue

The generalized trace is a rewrite interpretation of the paper schedule. The pinned source trace is version-scoped implementation evidence. Neither should be presented as a description of current FlashAttention main, a framework’s present dispatcher, or a profiler capture.

For one batch item and one query head, define:

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

The true extents of block i and block j are:

$$b_{r,i}=\min(B_r,S_q-iB_r), \qquad b_{c,j}=\min(B_c,S_k-jB_c).$$

Renderer-safe block contract:

T_r = ceil(S_q / B_r)
T_c = ceil(S_k / B_c)

b_r_i = min(B_r, S_q - i*B_r)
b_c_j = min(B_c, S_k - j*B_c)

Q_i:         [b_r_i, d_h]
K_j, V_j:    [b_c_j, d_h] in the paper's equal-width case
S_ij:        [b_r_i, b_c_j]
P_tilde_ij:  [b_r_i, b_c_j]
m_i, l_i:    [b_r_i]
u_i, O_i:    [b_r_i, d_v]
LSE_i:       [b_r_i]

The paper uses one d for Q, K, V, and O. Naming d_h and d_v in the teaching ledger exposes which product owns each width; it does not claim that the pinned v2.0.0 API accepted unequal K and V widths.

0x03B Lines 1-2: Partition Views, Not Yet Tensor Movement

Source

Q [S_q,d_h], K [S_k,d_h], V [S_k,d_v] in HBM

Destination and shape

logical views Q_i, K_j, V_j, O_i, and LSE_i
no score or probability tile exists yet

Dtype and lifetime

paper: real-valued tensors
v2.0.0 storage: Q/K/V share FP16 or BF16
metadata: launch/block coordinates
lifetime: launch scope

Partitioning is an address calculation. It does not require copying all blocks into a new global tensor. A final block may have fewer than B_r query rows or fewer than B_c key rows; calling the logical tail a full block would already change the accounting model.

0x03C Lines 3-5: One Worker Becomes The Query-Row Owner

The outer loop selects query block i. Paper line 4 moves Q_i from HBM to on-chip storage. Line 5 creates the identity state:

source:
    Q_i in HBM
    identity constants -inf and 0

destination:
    Q_i on chip                 [b_r_i,d_h]
    m_i = -inf                 [b_r_i]
    l_i = 0                    [b_r_i]
    u_i = 0                    [b_r_i,d_v]

lifetime:
    Q_i, m_i, l_i, u_i survive the complete inner K/V loop

The paper writes the working output accumulator with an O-like symbol. This article calls it u_i because 0x02 proved that it is the unnormalized numerator until line 12. Calling it the final normalized output inside the loop would erase the state-coordinate change that FA2 introduced.

In the pinned kernel, Q_i may live in shared memory, register fragments, or a shared-memory region later reused by K, depending on the launch specialization. The stable ownership statement is narrower:

Q_i and row state remain local to the query-block owner;
they are not read from and written to HBM after every K/V block.

0x03D Lines 6-10: K/V Stream, Score And Probability Fragments Expire

Each inner iteration has three different lifetimes.

K/V block, line 7

source:       K_j,V_j in HBM
destination:  on-chip K/V tiles
shape:        [b_c_j,d_h] and [b_c_j,d_v]
lifetime:     one inner iteration, possibly overlapped by pipelining

Score and probability fragments, lines 8-9

S_ij = tau * Q_i @ K_j.T       [b_r_i,b_c_j]
P_tilde_ij                     [b_r_i,b_c_j]

lifetime:
    create for this block
    apply sequence/causal predicates before rowmax
    update m_i and l_i
    feed the value product
    discard

tau is the score scale, M4.3’s 1/sqrt(d_h). The paper omits it for clarity (its footnote 2) and writes line 8 as S_ij=Q_iK_j^T; this trace keeps it because backward must recompute the same scaled scores from LSE_i.

Loop-carried state, lines 9-10

m_i: [b_r_i]       running maximum
l_i: [b_r_i]       running denominator in the m_i frame
u_i: [b_r_i,d_v]   running value numerator in the same frame

lifetime:
    every visited K/V block for this query owner

Line 10 is the crucial ownership update:

u_i <- exp(m_old - m_new) * u_i + P_tilde_ij @ V_j

S_ij and P_tilde_ij do not become global S and P matrices. They are temporary fragments whose only purpose is to update the compact row state.

0x03E Lines 11-17: Normalize And Commit Once

After all visible K/V blocks have been merged:

O_i   = u_i / l_i
LSE_i = m_i + log(l_i)

The epilogue then writes:

O_i:   [b_r_i,d_v] -> HBM once
LSE_i: [b_r_i]     -> HBM once

The lifetimes now diverge:

u_i, m_i, l_i:
    die after epilogue finalization

O_i:
    persists as the layer output

LSE_i:
    persists when backward must reconstruct probabilities

This is the storage consequence of the algebra from 0x02: carry (m,l,u) on chip, normalize once, and save one row scalar LSE rather than repeatedly materializing normalized output state between K/V blocks.

0x03F Executable Tail-Aware Schedule

The paper’s 17 lines can be expanded into a tail-aware teaching schedule. This is still not CUDA source:

def fa2_forward_owner(Q, K, V, B_r, B_c, causal_policy, tau):
    S_q, d_h = Q.shape
    S_k, d_k = K.shape
    assert d_h == d_k

    for i in parallel_range(ceil_div(S_q, B_r)):
        q0 = i * B_r
        q1 = min(q0 + B_r, S_q)
        Q_i = predicated_load(Q, rows=(q0, q1), width=d_h)

        m_i = full([q1 - q0], -inf, dtype=fp32)
        l_i = zeros([q1 - q0], dtype=fp32)
        u_i = zeros([q1 - q0, V.shape[-1]], dtype=fp32)
        valid_i = zeros([q1 - q0], dtype=bool)

        for j in visible_kv_tiles(i, S_q, S_k, B_r, B_c, causal_policy):
            k0 = j * B_c
            k1 = min(k0 + B_c, S_k)
            K_j = predicated_load(K, rows=(k0, k1), width=d_h)
            V_j = predicated_load(V, rows=(k0, k1), width=V.shape[-1])

            S_ij = tau * fp32_accumulate(Q_i @ K_j.T)
            allowed = sequence_and_causal_predicate(q0, q1, k0, k1)
            S_ij = where(allowed, S_ij, -inf)  # before rowmax

            block_valid = any(allowed, axis=-1)
            # Invalid rows merge the identity; they do not evaluate exp(-inf - -inf).
            m_new, l_new, u_new = merge_valid_rows(
                m_i, l_i, u_i, S_ij, V_j, block_valid
            )
            m_i, l_i, u_i = m_new, l_new, u_new
            valid_i |= block_valid

        O_i, LSE_i = finalize_under_declared_all_masked_policy(
            m_i, l_i, u_i, valid_i
        )
        predicated_store(O[q0:q1], cast_output_dtype(O_i))
        predicated_store(LSE[q0:q1], LSE_i)  # FP32 in v2.0.0

The valid_i branch is explicit because the algebraic identity and a backend’s all-masked output policy are separate contracts. The historical kernel has its own finite-state handling and early exits; this teaching code does not claim line-for-line source equivalence.

0x03G Query, Key, Head-Width, And Causal Tails Are Different Predicates

A kernel can have four boundaries in one tile:

query tail:
    q_index < S_q

key/value tail:
    k_index < S_k

head-width tail:
    feature_index < materialized_d

causal boundary:
    allowed(q_index, k_index) under the declared alignment policy

The independent-tail fixture makes all three widths visible:

S_q = 257, S_k = 193, d_h = 78
B_r = 128, B_c = 64

query block rows:  [128, 128, 1]
K/V block rows:    [64, 64, 64, 1]

logical width:                 78
v2.0.0 wrapper storage width:  ceil(78/8)*8  = 80
kernel specialization width:   ceil(80/32)*32 = 96

The pinned wrapper physically pads Q/K/V from 78 to 80 values per row. It then dispatches a 96-wide specialization. The kernel predicates feature lanes against params.d=80, and the wrapper slices the returned output back to 78. Those sixteen specialization lanes are not automatically sixteen extra HBM values per row; predicates can suppress their global loads and stores.

The same caution applies to sequence rounding. v2.0.0 computes query and key workspace extents rounded to 128, including the optional returned softmax/dropout debug tensor. That does not mean the ordinary Q/K/V tensors are copied to full 128-row blocks before every kernel call.

For square causal attention, declare allowed(q,k) iff k <= q. A tile is:

full:
    k_end - 1 <= q_start

skipped:
    k_start > q_end - 1

partial:
    neither condition holds; mask cells before rowmax

With S_q=S_k=256, B_r=128, and B_c=64, the executable fixture finds:

2 full tiles
4 partial tiles
2 skipped tiles

6 visited tile rectangles contain 49,152 score slots
32,896 cells are causally allowed

The distinction matters: partial tile rectangles still contain masked cells. The pinned forward kernel bounds n_block_max, iterates K/V blocks in reverse, and separates masking iterations from mask-free iterations. Reverse order is a source-level scheduling choice, not a different online-softmax result. Sections 0x05.1 and 0x06 own the full causal and variable-length geometry; this section establishes only the predicate and lifetime contract.

0x03H Dtype Is A Pipeline, Not One Label On The Operation

Algorithm 1 writes tensors over the reals. The pinned v2.0.0 source realizes that algorithm with several dtype roles:

Object or stage	v2.0.0 dtype role	Source-inspected consequence
Q, K, V global/shared storage	FP16 or BF16, all equal	tensor-core input elements
QK^T score fragment	FP32 accumulator	score maximum and exponential state start from FP32 accumulation
m, l, and u/acc_o	FP32	loop-carried numerical state stays in ElementAccum=float
shifted P_tilde math	FP32	exponentials and row sums update FP32 state
P_tilde entering P @ V	converted to FP16/BF16	second tensor-core product uses low-precision multiplicands with FP32 accumulation
final normalization	FP32	reciprocal and output-accumulator scaling occur before output conversion
O write	same FP16/BF16 dtype as input	epilogue converts FP32 acc_o to Element
LSE write	FP32	wrapper allocates softmax_lse with at::kFloat
row/feature predicates	boolean/control state	invalid rows and feature lanes are not committed

These are not inferred from one paper equation. They are visible in the pinned flash_api.cpp, kernel_traits.h, and flash_fwd_kernel.h.

The dtype contract also explains why a paper-level capacity count in scalar elements cannot be translated directly into shared-memory bytes. Q/K/V storage, score/output accumulators, row scalars, predicate state, and double-buffered tiles do not all have one width or one residency.

0x03I Dropout Adds Replay State Without Changing The Softmax Denominator

Training adds a second correctness contract. Attention dropout acts on the normalized probability path, but the stable maximum and denominator describe the undropped softmax distribution:

1. compute undropped m, l, and shifted probability numerator
2. generate the dropout decision for probability elements
3. apply the mask before P @ V
4. apply reciprocal keep-probability scaling
5. regenerate the same decisions during backward

The v2.0.0 Python autograd wrapper saves an opaque CUDA RNG state when dropout_p > 0, restores it for backward, and then restores the caller’s later RNG state. The kernel derives Philox seed/offset coordinates from batch, head, lane, and attention-block location so forward and backward can traverse work differently while regenerating the same mask.

Normal training does not require a dense probability matrix to persist. The historical API can optionally return an S_dmask diagnostic tensor, but in v2.0.0 return_softmax is accepted only when dropout is active. That optional debug path must not be added to the ordinary Algorithm 1 traffic ledger.

0x03J FA1 And FA2 Move Different Objects

M4.3’s FA1 Algorithm 1 trace established the K/V-outer paper schedule. FA2 swaps loop ownership. For a non-causal square case, let:

b_in:   bytes per Q/K/V element
b_out:  bytes per O element
b_s:    bytes per m/l/LSE scalar

Under the declared logical schedule model, FA1 moves:

$$\begin{aligned} M_{FA1}={}&2S_kd\,b_{in} +T_cS_qd\,b_{in} +2T_cS_qd\,b_{out}\\ &+4T_cS_q\,b_s +S_qd\,b_{out} +2S_q\,b_s. \end{aligned}$$

The final two terms are the paper’s initial writes of O, l, and m. FA2 moves:

$$M_{FA2} =S_qd\,b_{in} +2T_rS_kd\,b_{in} +S_qd\,b_{out} +S_q\,b_s.$$

Renderer-safe traffic formulas:

FA1 K/V-outer:
    K,V reads once             = 2*S_k*d*b_in
    Q reads per K/V block      = T_c*S_q*d*b_in
    O reads+writes per visit   = 2*T_c*S_q*d*b_out
    m,l reads+writes per visit = 4*T_c*S_q*b_s
    initialization writes      = S_q*d*b_out + 2*S_q*b_s

FA2 query-outer:
    Q reads once               = S_q*d*b_in
    K,V stream per Q block     = 2*T_r*S_k*d*b_in
    O writes once              = S_q*d*b_out
    LSE writes once            = S_q*b_s

For the fixture:

S_q=S_k=256, d=64
B_r=128 -> T_r=2
B_c=64  -> T_c=4
b_in=b_out=2 bytes, b_s=4 bytes

FA1 logical total including initialization: 509,952 bytes
FA2 logical total:                          197,632 bytes
declared difference:                       -312,320 bytes

FA2 does not reduce every term. It reloads K/V for each query-block owner, while FA1’s idealized K/V-outer trace reuses each K/V block across query blocks. FA2 wins this fixture’s ledger by eliminating repeated Q loads and repeated HBM round trips for O,m,l, while also exposing independent query-block work. Different tile sizes, cache behavior, GQA reuse, causal skipping, and launch placement can change physical traffic.

The independent-tail case reports another pair of coordinates:

logical valid-payload FA2 bytes at d=78:     261,860
v2.0.0 materialized-width bytes at d=80:    268,548
extra bytes in that wrapper-width model:      6,688

The 96-wide specialization is reported as a tile-slot envelope, not as traffic. Predication, vector load policy, cache-line coalescing, compiler lowering, and cache residency decide which slots become physical transactions.

0x03K Executable Schedule Evidence And Failure Lab

The checked fixture and oracle are:

docs/numinens-evidence/m4-4-flashattention2-work-partition/
    algorithm1_schedule_fixture.json
    algorithm1_schedule_oracle.py
    algorithm1_schedule_report.json
    test_algorithm1_schedule_oracle.py

They verify:

every Algorithm 1 line group names source, destination, shape, dtype, lifetime,
and tail behavior

the 257 x 193 fixture retains every logical score rectangle exactly once

the causal fixture separates visited rectangles from actually allowed cells

the d=78 logical, d=80 materialized, and d=96 specialization widths stay distinct

FA1/FA2 traffic follows the declared loop order and byte widths

all generated evidence remains unmeasured

Common failures now become mechanically identifiable:

Failure	Broken contract	Repair
call the in-loop u_i a normalized output	erases the delayed denominator division	name unnormalized numerator and final epilogue separately
apply a tail or causal mask after rowmax	forbidden scores can change m and l	predicate scores before rowmax and rowsum
count all B_r x B_c x d_kernel slots as HBM bytes	conflates logical shape, wrapper storage, specialization, and transactions	report each coordinate separately
say FA2 reads every K/V element once	ignores one stream per query-block owner	include T_r in the logical K/V term
include optional returned S_dmask in ordinary inference traffic	turns a diagnostic dropout path into baseline state	count it only under its explicit API option
infer current backend behavior from v2.0.0	historical source is not current dispatch evidence	keep implementation claims commit-scoped

Run the schedule checks with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_algorithm1_schedule_oracle.py

This oracle is a synthetic contract test. It does not build the historical extension, run a GPU kernel, inspect SASS, collect Nsight counters, or measure latency.

Source: FlashAttention-2, Figure 2. This is a source-reported partition schematic reproduced for technical commentary, not a literal CUDA grid or warp map and not a local measurement. Algorithm 1’s query-owner schedule explains the forward half; section 0x06.1 reconstructs the different backward write-conflict graph.

The formula and implementation are therefore not separate worlds. The formula determines the row state that must survive. The paper schedule chooses its owner. The historical kernel chooses dtypes, predicates, buffering, traversal order, and output conversion. Profiling is still required to determine what the device actually moved.

0x04 Launch Geometry Is A Resource Contract

Section 0x03 established one independent owner for each query-row block. That is an algorithmic parallelism statement. A GPU launch needs a stricter chain:

logical tensor axes
    -> launched CTA coordinates
    -> active versus early-return CTAs
    -> threads and warps inside one CTA
    -> shared-memory and compiled-register demand
    -> resident CTAs per SM
    -> rounded waves across the device

M4.2 owns the generic CUDA residency model. M4.4 now applies that model to the historical FA2 forward kernel. It does not repeat the full CUDA programming model, and it does not substitute a paper tile for compiled resource metadata.

0x04A The Historical Grid Has Three Independent Axes

The paper writes T_r independent query blocks. The pinned v2.0.0 launch template realizes them as the x dimension of a three-dimensional CUDA grid:

const int num_m_block =
    (params.seqlen_q + Kernel_traits::kBlockM - 1)
    / Kernel_traits::kBlockM;

dim3 grid(num_m_block, params.b, params.h);

kernel<<<grid,
         Kernel_traits::kNThreads,
         Kernel_traits::kSmemSize,
         stream>>>(params);

The mapping is exact for that historical source:

blockIdx.x -> query-row tile m_block
blockIdx.y -> batch item bidb
blockIdx.z -> query head bidh

Therefore:

$$T_r=\left\lceil\frac{S_q}{B_r}\right\rceil$$ $$N_{\mathrm{CTA,launch}}=T_r B H_q.$$

Renderer-safe launch formulas:

T_r          = ceil_div(S_q, B_r)
grid         = (T_r, B, H_q)
launched CTA = T_r * B * H_q

For B=1, H_q=32, S_q=4096, and B_r=128:

grid = (32, 1, 32)
N_CTA,launch = 32 * 1 * 32 = 1,024

At S_q=1, the same specialization still launches one query tile per query head:

grid = (1, 1, 32)
N_CTA,launch = 32

That is the exact arithmetic behind the later decode boundary. It is not yet an occupancy statement: a grid can expose 1,024 CTAs while resource limits allow only a small number to reside on each SM.

The head axis remains H_q, not H_kv. In the kernel, query head bidh maps to K/V head bidh / h_h_k_ratio. GQA therefore changes K/V indexing and reuse, but this v2.0.0 launch still creates one owner per query head and query tile.

The visual is a source-derived ownership reconstruction, not a literal lane map. The exact CUTE fragments depend on the selected trait; compiled register allocation still requires a build artifact.

0x04B Launched, Active, Full, And Useful Are Different Counts

For fixed-length input, every valid query block has work unless the key sequence is empty. Packed variable-length batches are different. v2.0.0 launches from the declared maximum seqlen_q, then the kernel checks the actual sequence lengths:

if (m_block * kBlockM >= binfo.actual_seqlen_q
    || binfo.actual_seqlen_k == 0) {
    return;
}

So four populations must not be conflated:

launched CTA: grid coordinate exists
active CTA:   survives the kernel-entry early return
full CTA:     owns B_r valid query rows
tail CTA:     owns fewer than B_r valid query rows

The executable packed fixture uses:

B=3, H_q=8, B_r=128
declared maximum S_q=512
actual S_q=[512,257,0]
actual S_k=[512,193,256]

The launch is:

grid = (4,3,8)
launched CTAs = 4 * 3 * 8 = 96

But only:

active CTAs after early return = 56
early-return CTAs              = 40
full query-tile CTAs           = 48
tail query-tile CTAs           = 8

The active CTAs expose 7,168 query-row slots, of which 6,152 are valid. The remaining 1,016 rows are predicates inside the eight tail CTAs. A profiler can still see launched blocks and instructions associated with checks; a tensor ledger should count only valid rows. Both coordinates are useful, but they answer different questions.

0x04C One CTA Is A Compile-Time Trait, Not Just A Query Tile

The v2.0.0 forward trait turns four template choices into a block:

Flash_fwd_kernel_traits<
    kHeadDim,
    kBlockM,
    kBlockN,
    kNWarps,
    Is_Q_in_regs,
    Share_Q_K_smem,
    Element>

The thread count is derived, not independently tuned:

$$T_{\mathrm{block}}=32N_{\mathrm{warps}}.$$

4 warps -> 128 threads
8 warps -> 256 threads

The selected trait also fixes:

B_r: rows of Q and O owned by the CTA
B_c: K/V rows consumed per inner-loop iteration
d:   compile-time head-dimension specialization
Q/K/V shared-memory layouts and swizzles
CUTE tensor-core partition across warps
128-bit global-memory copy layout

The active v2.0.0 branches are not one universal tile:

Historical branch	(d,B_r,B_c,warps)	Threads	Source-selection reason
d=64, no dropout	(64,128,128,4)	128	selected active branch; comments reject slower alternatives for the tested case
d=128, SM80 no dropout	(128,128,64,4)	128	A100-class branch
d=128, SM86/SM89 non-causal	(128,128,32,4)	128	source comment targets two CTAs/SM with 48 KiB SMEM
d=128, SM86/SM89 causal	(128,64,64,4)	128	square causal tile selected in source
d=224, high-SMEM branch	(224,128,64,8)	256	selected only when the opt-in per-block limit admits 112 KiB
d=224, fallback	(224,64,64,4)	128	lower per-block shared-memory demand
d=256, A100 branch	(256,128,64,8)	256	source chooses 128 KiB on A100-class resources

These are version-scoped implementation facts from commit 4f285b354796fb17df8636485b9a04df3ebbb7dc, not recommendations for a current backend. The branch comments contain historical tuning observations; they do not establish performance on another compiler, GPU, sequence shape, or FlashAttention release.

0x04D Shared Memory Follows Tensor Lifetimes

For two-byte FP16/BF16 elements and the active default Share_Q_K_smem=false, the forward trait computes:

$$S_Q=B_r d b_e,$$ $$S_{KV}=2B_c d b_e,$$ $$S_{\mathrm{block}}=S_Q+S_{KV}.$$

Renderer-safe shared-memory derivation:

S_Q     = B_r * d * bytes_element
S_KV    = 2 * B_c * d * bytes_element
S_block = S_Q + S_KV                 # default active branches

S_KV contains one K tile and one V tile. It is not the full sequence. S_Q contains the query-owner tile. The score fragment, probability fragment, row maximum, denominator, and output numerator are predominantly accumulator or register state in the source organization rather than additional persistent shared-memory matrices.

At the epilogue, the kernel writes the output accumulator through an sO shared-memory view. That view reuses the already allocated dynamic arena after the K/V loop; it is not added to kSmemSize as another simultaneously live B_r x d allocation. This is why a storage formula must follow lifetimes, not sum every tensor name that appears anywhere in the kernel.

If Share_Q_K_smem=true, the trait instead uses:

$$S_{\mathrm{block}}=\max(S_Q,S_{KV}),$$

but the source contract forces Q into registers before K/V reuse that storage. Shared-memory saving therefore exchanges one resource for register pressure; it does not make the Q tile disappear.

The checked active configurations produce:

Configuration	S_Q	S_KV	S_block
d=64, 128x128, 4 warps	16 KiB	32 KiB	48 KiB
d=128, 128x64, 4 warps	32 KiB	32 KiB	64 KiB
d=128, 128x32, 4 warps	32 KiB	16 KiB	48 KiB
d=224, 128x64, 8 warps	56 KiB	56 KiB	112 KiB
d=224, 64x64, 4 warps	28 KiB	56 KiB	84 KiB
d=256, 128x64, 8 warps	64 KiB	64 KiB	128 KiB

0x04E Feasibility Comes Before Occupancy

A candidate cannot be assigned an occupancy percentage until it is a legal launch. At minimum, check:

threads per CTA <= architectural per-block limit
dynamic SMEM per CTA <= opt-in per-block limit
head dimension satisfies tensor/copy-layout divisibility
thread count satisfies Q/K/V and P copy-layout divisibility
at least one CTA fits all SM resources after compilation

The historical launch path performs two notable CUDA calls:

if (smem_size >= 48 * 1024) {
    cudaFuncSetAttribute(
        kernel,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        smem_size);
}

cudaOccupancyMaxActiveBlocksPerMultiprocessor(
    &ctas_per_sm,
    kernel,
    Kernel_traits::kNThreads,
    smem_size);

The first opts the kernel into a larger dynamic shared-memory allocation. The second asks CUDA about the compiled kernel function on the selected device, using its thread count and dynamic SMEM. CUDA defines the result as the maximum active blocks per SM for that function and launch configuration. The current CUDA programming guidance also makes the missing dependency explicit: compiled register and shared-memory usage are required for occupancy calculation.

The v2.0.0 source stores ctas_per_sm only as a local diagnostic variable; the checked source does not emit it into an article, log, or result tensor. This repository did not build the extension or execute that API. Any numeric resident CTA count below is therefore either a static ceiling or an explicitly declared sensitivity scenario.

For the device limits used in the oracle:

A100 / compute capability 8.0:
    108 SMs
    64 resident warps/SM
    2,048 resident threads/SM
    32 resident CTAs/SM architectural limit
    65,536 32-bit registers/SM
    164 KiB shared memory/SM
    163 KiB maximum opt-in shared memory/CTA

generic compute capability 8.6 profile:
    48 resident warps/SM
    1,536 resident threads/SM
    16 resident CTAs/SM architectural limit
    65,536 32-bit registers/SM
    100 KiB shared memory/SM
    99 KiB maximum opt-in shared memory/CTA

The A100 SM count comes from NVIDIA’s A100 architecture whitepaper. The SM resource limits come from NVIDIA’s Ampere Tuning Guide. They describe hardware ceilings, not this kernel’s achieved occupancy.

0x04F A Static Ceiling Is Not The Compiled Residency

Specializing M4.2’s general resource model to a selected FA2 CTA gives:

$$N_{\mathrm{static}} = \min\left( N_{\mathrm{block,arch}}, \left\lfloor\frac{T_{SM}}{T_{block}}\right\rfloor, \left\lfloor\frac{W_{SM}}{W_{block}}\right\rfloor, \left\lfloor\frac{S_{SM}}{S_{block}}\right\rfloor \right).$$

This deliberately excludes registers. With a compiled R_thread and a declared allocation model, a first-order register term is:

$$N_R=\left\lfloor\frac{R_{SM}}{R_{thread}T_{block}}\right\rfloor.$$

Actual allocation rounds registers at architecture-specific granularities, so the CUDA occupancy API or occupancy calculator remains authoritative for the compiled binary. The source’s C++ template tells us accumulator shapes and ElementAccum=float; it does not tell us the final allocated registers per thread after compilation.

The corresponding warp-occupancy ceiling is:

$$\mathrm{occupancy}_{\mathrm{static}} = \frac{N_{\mathrm{static}}W_{block}}{W_{SM}}.$$

Using the declared A100 limits, while explicitly excluding compiled registers:

Historical configuration	Static CTA ceiling	First limiter	Static warp-occupancy ceiling
d=64, 128x128, 4 warps, 48 KiB	3	shared memory	18.75%
d=128, 128x64, 4 warps, 64 KiB	2	shared memory	12.50%
d=224, 128x64, 8 warps, 112 KiB	1	shared memory	12.50%
d=224, 64x64, 4 warps, 84 KiB	1	shared memory	6.25%
d=256, 128x64, 8 warps, 128 KiB	1	shared memory	12.50%

For the d=64 case, preserving three CTAs under the unrounded register model would require no more than floor(65536/(3*128)) = 170 registers per thread. The fixture does not assume that condition. For d=128, the analogous simple threshold for two CTAs is 256, but the hardware maximum is 255; allocation rounding and the compiled function still belong to the CUDA calculator rather than prose inference.

The evidence ladder is:

source traits:
    tile, warps, threads, dynamic SMEM

compiled metadata:
    registers/thread, static SMEM, spills, target architecture

occupancy API/calculator:
    theoretical resident CTAs and warps for that binary and device

profiler:
    achieved active/eligible/issued warps, stall reasons, throughput

This section has the first layer and a static device-limit calculation. It has not collected the last three layers locally.

0x04G Resource Trade-Offs Explain The Configuration Matrix

Changing a tile changes several constraints at once:

increase B_r:
    fewer grid CTAs
    more Q/output accumulator state per CTA
    more Q shared memory

increase B_c:
    fewer K/V loop iterations
    more K/V shared memory
    larger score/probability fragments

increase warps:
    more threads and possible intra-CTA parallel work
    more resident warps per CTA
    different fragment ownership and possible register pressure

share Q/K shared memory:
    lower dynamic SMEM
    Q must survive in registers

The source itself records rejected alternatives. For d=64, a 256x64 tile was reported 85% slower at sequence length 2K in the source comment because of register spilling. For d=128, commented eight-warp alternatives were reported 28% slower for the same source-local test shape. These are historical tuning notes, not portable benchmark claims, but they expose the mechanism:

more rows or more warps can increase parallel work
while simultaneously increasing live fragments, register demand, or spills

The d=224 branch queries maximum opt-in shared memory before choosing between 112 KiB / 8 warps and 84 KiB / 4 warps. The d=256 branch queries both per-block and per-SM shared memory before choosing a 128 KiB A100-oriented tile or a smaller-row alternative intended to permit two CTAs on hardware with a larger SM budget. M4.6 owns the later Hopper pipeline consequences; M4.4 uses this branch only to show that hardware resources participate in dispatch.

0x04H Device Waves Are Quantized

Suppose a compiled occupancy result admits N_res CTAs per SM on a device with N_SM SMs. The simultaneous CTA-slot envelope is:

$$C_{\mathrm{CTA}}=N_{SM}N_{res}.$$

For N_launch query-owner CTAs, an equal-duration wave model gives:

$$N_{\mathrm{waves}} = \left\lceil\frac{N_{launch}}{C_{\mathrm{CTA}}}\right\rceil.$$

The last wave contains:

$$N_{last} = N_{launch}-(N_{waves}-1)C_{\mathrm{CTA}},$$

and its slot fill is:

$$f_{last}=\frac{N_{last}}{C_{\mathrm{CTA}}}.$$

Renderer-safe wave formulas:

CTA slots      = N_SM * resident_CTA_per_SM
rounded waves  = ceil_div(launched_CTA, CTA_slots)
last-wave CTA  = launched_CTA - (waves - 1) * CTA_slots
last-wave fill = last-wave CTA / CTA_slots

Use the paper’s A100-class device count, the d=128,128x64,4-warp source configuration, and a declared sensitivity scenario of two resident CTAs per SM. This scenario matches the static SMEM ceiling, but no compiled occupancy result was collected here:

N_SM = 108
resident CTA/SM scenario = 2
C_CTA = 216 slots
B=1, H_q=32, B_r=128

S_q	Query blocks	Launched CTAs	Rounded waves	Last-wave CTAs	Last-wave fill	Rounded slot fill
1	1	32	1	32	14.81%	14.81%
768	6	192	1	192	88.89%	88.89%
769	7	224	2	8	3.70%	51.85%
2,048	16	512	3	80	37.04%	79.01%
4,096	32	1,024	5	160	74.07%	94.81%

The first full 216-slot wave requires:

$$T_r\geq\left\lceil\frac{216}{1\cdot32}\right\rceil=7,$$

so the smallest sequence length exposing at least that many query blocks is:

$$S_q^{min}=(7-1)128+1=769.$$

That same 769 case crosses a quantization boundary: 768 launches 192 CTAs and fits in one wave; one additional valid query row creates a seventh tile per head, 224 CTAs, and a second wave with only eight CTAs in this envelope.

If the compiled kernel admitted only one CTA per SM, the same 4,096-token case would have 108 slots and ten rounded waves rather than five. This sensitivity is why wave arithmetic must consume a compiled residency result, not infer it from tile names.

Wave count is not latency. It ignores overlap, scheduling order, unequal CTA duration, cache state, memory-system contention, and instruction dependencies. It is a launch-capacity diagnostic.

0x04I Causal CTAs Break The Equal-Duration Wave Model

For square causal attention, query block i visits only K/V blocks that can contain visible keys. In the pinned kernel:

n_block_max = min(
    ceil_div(S_k, B_c),
    ceil_div((i+1)*B_r, B_c)
)

With S_q=S_k=4096, B_r=128, and B_c=64, the 32 query blocks perform:

K/V iterations per query block:
2, 4, 6, ..., 64

minimum / maximum: 2 / 64
longest-to-shortest ratio: 32x

causal total per head:    1,056 K/V-tile iterations
non-causal total per head: 2,048 K/V-tile iterations
ratio:                     51.5625%

The CTA resource footprint is chosen at launch; the early block does not return most of its registers or shared memory merely because its causal loop is short. Yet its duration can be much shorter than the final query block. CUDA does not promise grid scheduling order, so grouping CTAs into equal-duration waves is no longer a timeline prediction. A trace or profiler is required to observe actual tail behavior.

Packed inputs add another distortion. In the fixture’s two-CTA-per-SM scenario, the grid launches 96 CTAs into 216 slots, a nominal 44.44% one-wave fill. Only 56 CTAs survive entry checks, a 25.93% useful-CTA fill before query-row and K/V tails are considered. Reporting only the grid count would hide that gap.

0x04J Executable Resource Evidence And Failure Closure

The launch/resource evidence package adds:

launch_resource_fixture.json
launch_resource_oracle.py
launch_resource_report.json
test_launch_resource_oracle.py

The oracle checks:

the exact v2.0.0 grid and trait-derived thread count
Q plus K/V shared-memory formulas and lifetime aliasing
per-block feasibility under declared CC 8.0 and CC 8.6 limits
static CTA ceilings without inventing compiled registers
one-CTA and two-CTA residency sensitivity
the 768 -> 769 wave boundary
packed launched/active/full/tail CTA populations
causal 2 -> 64 inner-loop iteration imbalance

Run it with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_launch_resource_oracle.py

This is a synthetic contract test. It does not compile v2.0.0, inspect ptxas register metadata, execute cudaOccupancyMaxActiveBlocksPerMultiprocessor, collect Nsight counters, or measure latency.

Failure	Broken coordinate	Repair
call B H_q ceil(S_q/B_r) occupancy	grid population is not resident warps	compute compiled per-SM residency first
call kSmemSize the complete resource footprint	registers and allocation granularity are missing	add compiled metadata and occupancy API result
add Q, K, V, and epilogue O SMEM as simultaneously live	ignores arena reuse after the K/V loop	derive peak from lifetimes in kernel_traits.h and flash_fwd_kernel.h
infer register count from accumulator tensor shape	compiler allocation, temporary liveness, and spills are hidden	use ptxas --resource-usage or a compiled function attribute
call a two-CTA scenario measured occupancy	turns a declared sensitivity into a runtime fact	label the scenario and preserve the unmeasured status
report grid CTAs as useful work for packed input	ignores entry returns and tail rows	report launched, active, full, tail, and valid-row coordinates
treat rounded waves as elapsed time	assumes equal CTA duration and known scheduling	combine waves with trace/profiler evidence
assume causal halves every CTA equally	hides a 2-to-64 K/V-loop range in the fixture	retain per-query-block work

This completes the inter-CTA launch and resource contract. The next section moves inside one CTA and asks how its warps divide Q, K, V, score, and output fragments.

0x05 Warp Partitioning Is An Output-Ownership Contract

Section 0x04 stopped at the CTA boundary. For one grid coordinate, the historical forward kernel has already selected a query tile, batch item, and query head. The remaining question is smaller in scale but not in consequence:

Which warp owns each query row, which operands are shared, and how many producers must be merged before one output row is valid?

The FlashAttention-2 paper compares two answers. Its older illustrated plan splits K/V across warps; its replacement splits Q. Calling those plans “split-K” and “split-Q” is only shorthand. Their actual difference is the output conflict graph:

split-K:
    W warps contribute to the same output row
    -> W producers, then an inter-warp reduction

split-Q:
    one warp owns each output row and scans all K/V tiles for it
    -> one producer, no inter-warp output-row reduction

That distinction recovers the mechanism explained in A19 while correcting one dangerous simplification. The paper says there is no need for communication between warps to add the illustrated partial outputs. It does not say that the complete CTA stops using shared memory, barriers, or lane-level reductions.

0x05A Figure 3 Is A Reduction Graph, Not A Literal Lane Map

Source: FlashAttention-2, Figure 3. This source-reported conceptual comparison gives the left plan overlapping output ownership and the right plan disjoint query/output slices. It does not specify the CuTe lane layout of every specialization and is not local runtime evidence.

For one query tile Q_i, one K/V tile, and W warps, the two plans can be written as an ownership table:

Plan	Operand visible to every warp	Per-warp operand	Producers per output row	Required output-row action
illustrated split-K	Q_i	disjoint K_w,V_w slices	W	write partials, synchronize, reduce
illustrated split-Q	staged K_j,V_j	disjoint Q_w rows and their row state	1	write the owned row; no cross-warp row merge

The paper’s performance argument is then precise: after FA1 removes the large HBM materialization of S and P, FA2 removes an avoidable on-chip partial output reduction from the forward warp plan.

0x05B Split-K Gives One Row Several Producers

Let the key indices be partitioned into disjoint sets J_0,...,J_{W-1}. For query row i, warp w sees only its key/value slice and can form a stable local state:

$$m_i^{(w)}=\max_{j\in J_w} S_{ij},$$ $$l_i^{(w)}=\sum_{j\in J_w}e^{S_{ij}-m_i^{(w)}},$$ $$u_i^{(w)}=\sum_{j\in J_w}e^{S_{ij}-m_i^{(w)}}V_j.$$

Every u_i^(w) targets the same final output row. The states are also in different maximum frames, so averaging the locally normalized outputs is not valid. They first need one common frame:

$$m_i=\max_w m_i^{(w)},$$ $$l_i=\sum_w e^{m_i^{(w)}-m_i}l_i^{(w)},$$ $$u_i=\sum_w e^{m_i^{(w)}-m_i}u_i^{(w)}, \qquad O_i=\frac{u_i}{l_i}.$$

Renderer-safe merge contract:

m_i = max_w m_i^(w)
l_i = sum_w exp(m_i^(w) - m_i) * l_i^(w)
u_i = sum_w exp(m_i^(w) - m_i) * u_i^(w)
O_i = u_i / l_i

If globally normalized P already existed, the value product alone could be written as O_i = sum_w P[i,J_w] @ V[J_w]. Fused attention does not begin with that global P; its row maximum and denominator are part of the state that must be reconciled. The exact merge above makes the hidden dependency explicit. It is the same associative state algebra derived in M4.3, applied here to simultaneous warp producers rather than sequential K/V tiles.

The executable fixture uses four query rows, eight keys, a value width of three, and four two-key partitions. Across all rows:

maximum error, exact split-K state merge vs dense: 2.22e-16
maximum error, average of four local outputs:       1.1794
output producers per row:                           4
binary merges required per row:                     3

For the first row, the result is:

dense output:       [0.947517, 0.664354, -0.216300]
exact merged output:[0.947517, 0.664354, -0.216300]
naive local average:[0.729781, 0.831874,  0.006329]

The toy is not a model of FA1’s physical shared-memory layout. It proves the semantic conflict: partitioning the reduction axis gives several partial row states that cannot independently become final output rows.

0x05C Split-Q Gives Every Output Row One Warp Owner

Now partition the query rows into disjoint sets I_w. Warp w owns:

Q[I_w, :]
row maxima m[I_w]
row denominators l[I_w]
unnormalized outputs u[I_w, :]
final O[I_w, :] and LSE[I_w]

The CTA cooperatively stages each K/V tile, and every warp consumes that tile for its own rows. The per-warp recurrence is:

for warp in warps:
    rows = I_w
    m, l, u = empty_row_state(rows)

    for K_blk, V_blk in shared_kv_stream:
        S = Q[rows] @ K_blk.T
        m, l, u = online_update(m, l, u, S, V_blk)

    O[rows] = u / l[:, None]

For I_0,...,I_3 = {0},{1},{2},{3} in the fixture, each output row has one producer and matches the dense reference exactly. There is nothing to combine with another warp because no other warp owns that row.

The reconstruction has a different job from paper Figure 3. Figure 3 records the historical design comparison. This visual exposes the exact state merge, the number of output producers, and the boundary of the communication claim.

0x05D Reduction Topology Changes Depth, Not The Ownership Conflict

With W partial states, an exact binary reduction needs W-1 merges per output row regardless of topology:

$$N_{\mathrm{merge}}=B_r(W-1).$$

The topology changes the critical path:

serial collector: depth = W - 1
balanced tree:    depth = ceil(log2(W)) for a power-of-two tree

For B_r=128, d_v=128, W=4, and FP32 (m,l,u) state:

state width                         = d_v + 2 = 130 FP32 values
partial row states                  = 128 * 4 = 512
binary state merges                 = 128 * 3 = 384
non-owner semantic payload floor    = 128 * 3 * 130 * 4 B
                                    = 199,680 B = 195 KiB
all partial states written once     = 128 * 4 * 130 * 4 B
                                    = 266,240 B = 260 KiB
serial depth / balanced-tree depth  = 3 / 2

This is a semantic payload envelope, not a profiler count and not a claim that historical FA1 physically wrote this exact (m,l,u) representation. Bank layout, vectorization, staging reuse, reduction topology, and whether a component remains in registers determine real shared-memory transactions. The ledger exists to show why “only on-chip” is not the same as “free.”

Split-Q removes those 384 cross-warp output-state merge edges for the illustrated row tile. It does not remove its operand movement:

one staged K/V tile at B_c=64, d=128, FP16/BF16:
    2 * 64 * 128 * 2 B = 32 KiB

one Q tile at B_r=128, d=128, FP16/BF16:
    128 * 128 * 2 B = 32 KiB

one epilogue O staging tile at the same shape:
    128 * 128 * 2 B = 32 KiB

Those are logical tensor payloads for one selected trait, not measured bank transactions. K/V are staged cooperatively once per CTA iteration; saying “every warp sees K/V” must not be misread as four independent HBM loads.

0x05E The Pinned v2.0.0 Trait Places Warps Along Query Rows

The paper establishes the design. The historical source pin shows one concrete realization. In kernel_traits.h, the forward tensor-core layout is:

using TiledMma = TiledMMA<
    typename Base::MMA_Atom_Arch,
    Layout<Shape<Int<kNWarps>, _1, _1>>,
    typename Base::ValLayoutMNK>;

The warp replication is on the first MMA dimension, M, which corresponds to query/output rows. It is not on the N key-column dimension. That source fact is the implementation-level split-Q contract.

The same forward kernel gives a source-visible causal-mask row seed:

row_seed = (tidx / 32) * 16 + (tidx % 32) / 4;
row_stride = kNWarps * 16;

For the active A100 d=128, B_r=128, B_c=64, four-warp specialization:

Warp	Unique query/output rows derived from the source layout	Row count
0	0-15, 64-79	32
1	16-31, 80-95	32
2	32-47, 96-111	32
3	48-63, 112-127	32

The non-contiguous row bands matter. Figure 3’s four clean horizontal slices are an ownership illustration, not a promise that every specialization stores one contiguous quarter per warp. For B_r=64 with four warps, the same source formula gives contiguous 16-row ranges. For B_r=128 with eight warps, it also gives 16 rows per warp. All three checked specializations cover the complete row tile with zero cross-warp row overlap.

At lane granularity, four lanes cooperate over columns for each row seed. In the B_r=128, four-warp case, lanes 0-3 of warp 0 touch row set {0,8,64,72} across the repeated MMA-M positions. The pinned softmax.h reduces each row maximum and sum with Allreduce<4>. That is a four-lane intra-warp reduction, not the cross-warp partial-output reduction that split-Q removes.

0x05F One K/V Iteration Has A Fragment-Level Ownership Trace

The pinned forward source can be followed without pretending CUTE tensors are ordinary contiguous arrays:

// All CTA threads partition cooperative global -> shared copies.
auto copy_qkv = GmemTiledCopyQKV{}.get_thread_slice(tidx);
tQgQ = copy_qkv.partition_S(gQ);  tQsQ = copy_qkv.partition_D(sQ);
tKgK = copy_qkv.partition_S(gK);  tKsK = copy_qkv.partition_D(sK);
tVgV = copy_qkv.partition_S(gV);  tVsV = copy_qkv.partition_D(sV);

// The TiledMma gives each thread fragments in the warp's M-row ownership.
auto thr_mma = tiled_mma.get_thread_slice(tidx);
tSrQ = thr_mma.partition_fragment_A(sQ);
tSrK = thr_mma.partition_fragment_B(sK);
tOrV = thr_mma.partition_fragment_B(sV_transposed);
acc_o = partition_fragment_C(tiled_mma, Shape<B_r, d>{});

for each visible K/V tile:
    wait_for_shared_K();
    acc_s = Q_fragment @ K_fragment.T;       // FP32 score fragment
    mask_invalid_score_elements(acc_s);
    wait_before_reusing_KV_shared_storage();
    P, m, l = row_softmax_and_rescale(acc_s, m, l, acc_o);
    acc_o += P_fragment @ V_fragment;        // owned output rows only

normalize_owned_rows(acc_o, l);

The important tensor facts are:

Object	Logical tile	Source-visible location during the step	Ownership
sQ	[B_r,d]	dynamic shared memory; optionally copied into registers	CTA-staged, warp-row consumed
sK,sV	[B_c,d] each	dynamic shared memory	CTA-staged, visible to all warp row owners
acc_s	per-thread fragment of [B_r,B_c]	FP32 registers	query-row fragment plus key-column lanes
scores_max,scores_sum	row-state fragments	FP32 registers	same warp as the output row
tOrP	converted probability fragment	registers	transient input to P-V MMA
acc_o	per-thread fragment of [B_r,d]	FP32 registers across K/V iterations	disjoint warp-row ownership

There is no shared-memory write of four overlapping acc_o row states followed by a cross-warp add. The output accumulator remains with its row owner while K/V tiles stream.

0x05G No Output Reduction Does Not Mean No CTA Communication

The pinned source still contains several coordination mechanisms:

1. GmemTiledCopyQKV partitions Q/K/V copies across all CTA threads.
2. cp_async_wait<...>() and __syncthreads() protect K/V shared-memory production, consumption, and reuse in the inner loop.
3. Row softmax uses four-lane Allreduce<4> operations for maxima and sums.
4. The epilogue converts FP32 acc_o fragments, writes them to shared sO, synchronizes, then repartitions sO for a coalesced global-memory write.
5. When Q and K share the same shared-memory arena, another synchronization protects the lifetime transition before sO reuses that arena.

So the precise claim is:

removed:
    cross-warp reduction of overlapping forward output-row partials

retained:
    cooperative operand staging
    shared-memory lifetime barriers
    lane-level row reductions
    epilogue shared-memory permutation and synchronization

This distinction is also why counting __syncthreads() occurrences cannot by itself measure the optimization. The semantic purpose, dependency path, and bytes protected by a barrier matter more than the source token count.

0x05H Alternative Choices Move Pressure Rather Than Erase It

Choice	Advantage	New or retained pressure	Appropriate boundary
split-K inside one CTA	exposes parallel key-reduction work	overlapping row state and reduction	historical comparison in this section
split-Q inside one CTA	disjoint output rows; no output-state merge	all owners consume staged K/V; rows per warp shrink with more warps	FA2 forward path
four versus eight warps	more warp-level work can cover larger tiles	more threads, registers, and resource pressure per CTA	0x04 trait matrix
smaller B_r	more CTAs and simpler row ownership	less reuse and more launch/tail overhead	0x04 wave sensitivity
split K/V across CTAs	recovers parallelism when query rows collapse	global partial-state merge and scheduler/backend contract	M4.5
producer/consumer warp specialization	overlaps movement and compute on newer hardware	explicit pipeline stages and architecture-specific synchronization	M4.6

Split-Q is therefore not “the universally best axis.” It is the forward prefill answer when enough query rows exist inside the CTA. At one-token decode, that row-parallel surface collapses; M4.5 owns the exact K/V-split alternative and its partial-state merge rather than duplicating it here.

0x05I Executable Ownership Evidence And Failure Closure

The ownership evidence is reproducible with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_warp_partition_oracle.py

The checked artifacts are:

warp_partition_fixture.json
warp_partition_oracle.py
warp_partition_report.json
test_warp_partition_oracle.py

They verify exact split-K state merging, the failure of naive local normalization, split-Q single-producer ownership, the declared communication envelope, and source-derived row coverage for three v2.0.0 traits. This remains a synthetic contract test plus version-scoped source inspection. It does not emulate CUTE register fragments, provide a compiled register map, compile CUDA, count physical shared-memory transactions, or measure latency.

Failure	Broken contract	Repair
average locally normalized split-K outputs	local partitions use different maxima and denominators	merge (m,l,u) in one maximum frame
call each split-K partial a final output row	W warps write the same semantic row	expose producer multiplicity and reduction
say every split-Q warp reloads K/V from HBM	confuses shared visibility with global duplication	trace cooperative CTA staging into sK/sV
say FA2 has no inter-warp communication	generalizes one removed reduction to the whole CTA	inventory remaining staging, barriers, lane reductions, and epilogue
draw Figure 3 as a literal lane map	hides specialization-dependent CUTE layouts	derive row sets from the pinned trait and source indices
call the 195/260 KiB envelope measured traffic	converts semantic state volume into hardware transactions	require compiled kernel and profiler counters
reuse forward split-Q ownership for backward	ignores dQ/dK/dV write conflicts	derive Algorithm 2 separately in 0x06.1

The key result is now stronger than “split Q is faster”:

FA2 changes the number of producers per forward output row from many to one.
That removes one exact reduction graph while preserving the CTA communication
needed to feed tensor cores and commit the owned rows.

Did The New Ownership Actually Move Throughput?

The work-partition argument is useful only if it survives measurement. The FlashAttention-2 paper benchmarks forward plus backward across causal and non-causal attention, two head dimensions, and several sequence lengths:

Source: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, Figure 4. These are paper-reported TFLOPs/s on one A100 80GB SXM4 benchmark configuration; they compare the implementations available to the authors at publication time and are not a current backend ranking.

Across the four panels, the purple FlashAttention-2 bars stay above the FlashAttention-1 and Triton baselines. The important connection to the previous sections is causal, not merely chronological:

FA1:
    fixes the dominant global-memory schedule

FA2:
    then exposes more thread-block work
    and reduces intra-block communication

measurement:
    checks whether those ownership changes raise achieved throughput

The chart still does not select a kernel for a new workload. Head dimension, mask mode, dtype, GPU generation, backward requirements, and framework dispatch remain part of that decision.

Backward Uses A Different Ownership Problem

Forward can assign one query block to one program because that program owns the corresponding O_i rows. Backward has three different output axes:

dQ_i += dS_ij @ K_j
dK_j += dS_ij.T @ Q_i
dV_j += P_ij.T @ dO_i

No two-dimensional tile can be the unique owner of all three. A query-tile owner can commit dQ_i directly but collides on dK_j/dV_j; a K/V-tile owner can commit dK_j/dV_j directly but collides on dQ_i.

The paper chooses K/V-column workers and states that their dQ contributions communicate through atomic adds. The historical source turns that sentence into a preprocessing kernel, an FP32 atomic scratch tensor, a column-owner main kernel, a conversion kernel, and a separate grouped-head reduction. Section 0x06.1 derives that full contract after causal and grouped coordinates have been established.

The immediate lesson is narrower: “split Q in forward” is not a complete training schedule. Backward changes the write-conflict graph.

0x05.1 Causal Tile Geometry Is A Coordinate Contract

A19 correctly introduced causal masking as schedule work: whole-future tiles can be skipped, old-history tiles need no element mask, and only the boundary band needs element predicates. The missing step was to derive those classes from integer coordinates and to state which diagonal convention the kernel implements.

0x05.1A One Cell Predicate Determines Every Tile Class

Let a query tile i and key tile j use half-open logical ranges:

$$q\in[q_0,q_1),\qquad k\in[k_0,k_1),$$

with:

$$q_0=iB_r,\qquad q_1=\min(S_q,(i+1)B_r),$$ $$k_0=jB_c,\qquad k_1=\min(S_k,(j+1)B_c).$$

Use an alignment offset delta and declare a score cell visible when:

$$k\le q+\delta.$$

Renderer-safe coordinate contract:

q tile: [q0, q1)
k tile: [k0, k1)
visible(q, k) iff k <= q + delta

top-left diagonal:    delta = 0
bottom-right diagonal: delta = S_k - S_q

For square self-attention, S_q = S_k, so both conventions reduce to k <= q. Rectangular attention is where version and API semantics become observable.

0x05.1B Full, Partial, And Skipped Follow From Range Extremes

A tile is fully visible only when its largest key index is visible to its smallest query index:

$$k_1-1\le q_0+\delta.$$

It is fully future only when its smallest key index is later than its largest query index:

$$k_0>(q_1-1)+\delta.$$

Everything else intersects the causal boundary and needs an element mask:

def classify_causal_tile(q0, q1, k0, k1, delta):
    if k0 > (q1 - 1) + delta:
        return "skipped_future"
    if (k1 - 1) <= q0 + delta:
        return "full_visible"
    return "partial_mask"

This is more general than the shorthand j < i, j = i, j > i used for aligned square tiles. If B_r != B_c, one query tile can intersect several key tiles along the diagonal. If either sequence has a tail, tile ordinal alone no longer states how many valid rows or columns exist.

0x05.1C Tail Shape And Causal Class Are Independent Predicates

Four questions must remain separate:

is the query tile short because q1 == S_q?
is the key tile short because k1 == S_k?
is the whole tile causally visible, partial, or future?
which head and batch item own this tile?

A tail tile can be fully visible. For S_q=S_k=257, B_r=128, and B_c=64, the last coordinate is one cell:

query range: [256,257)
key range:   [256,257)
causal class: full_visible, because 256 <= 256
tail class: both query-tail and key-tail

Calling every tail a diagonal partial tile would conflate storage bounds with attention semantics.

0x05.1D The Pinned Kernel Culls Whole-Future Tiles Before QK^T

The historical v2.0.0 forward kernel computes:

n_block_max = ceil_div(actual_seqlen_k, kBlockN);
if (Is_causal) {
    n_block_max = min(
        n_block_max,
        ceil_div((m_block + 1) * kBlockM, kBlockN)
    );
}

For its top-left diagonal, query block m_block therefore visits only:

$$N_{K,\mathrm{visit}}(i)= \min\left( \left\lceil\frac{S_k}{B_c}\right\rceil, \left\lceil\frac{(i+1)B_r}{B_c}\right\rceil \right).$$

Future blocks beyond that limit never enter the QK^T loop. This is actual work culling, not merely multiplying computed probabilities by zero later.

The source then traverses K/V blocks in reverse. It applies element masking to:

ceil(B_r / B_c) near-diagonal iterations in causal mode
the final K/V tail iteration in non-causal mode

and runs older, fully visible history tiles through an unmasked loop. In the mask helper, every lane-derived score coordinate is set to -INFINITY when:

col_idx >= min(actual_seqlen_k, row_idx + 1)

The mask therefore executes before softmax_rescale_o computes the local maximum and denominator. If future scores enter either reduction, masking the output afterward cannot repair the denominator.

0x05.1E A Square 4096-Token Ledger Shows Tile Quantization

The executable fixture uses:

S_q = S_k = 4096
B_r = 128
B_c = 64
query tiles = 32
key tiles   = 64

The dense tile rectangle has 32 * 64 = 2,048 coordinates. Exact classification gives:

Tile class	Count	Kernel consequence
fully visible	992	run QK, softmax update, and PV without causal element predicates
partial causal	64	run the tile and set future score cells to -inf before row reductions
fully future	992	skip the complete QK/PV tile

The source-visible K/V-loop counts are:

2, 4, 6, ..., 64
sum: 1,056 visited K/V tiles per query head
mask-loop iterations: 64

At cell granularity:

dense rectangle:                  16,777,216 cells
causally allowed:                  8,390,656 cells
visited rectangle after tile skip: 8,650,752 cells
visited but element-masked:          260,096 cells

The allowed count is exactly 4096 * 4097 / 2. The difference between 8,650,752 visited rectangle cells and 8,390,656 allowed cells is the quantized diagonal-band tax. “Causal is half the FLOPs” is therefore an asymptotic intuition, not a literal kernel ledger.

For the 257-token dual-tail fixture, the 3 x 5 grid contains:

7 full-visible tiles
4 semantic partial tiles
4 skipped-future tiles
11 source-visited tiles
6 source mask-loop iterations

There are six mask-loop iterations but only four semantically partial tiles because the source keeps a fixed two-tile causal mask band and the final tail coordinate is fully visible. Schedule masking work and semantic tile class are related, not identical.

The reconstruction is informed by the FlashAttention-2 causal-masking note, A19’s causal introduction, pinned source, and causal_gqa_report.json. It is not a measured instruction trace.

0x05.1F v2.0 And v2.1 Give Rectangular Causal Calls Different Meanings

The historical kernel pin uses delta=0: the triangle begins at the top-left. The official FlashAttention changelog states that version 2.1 changed rectangular causal masks to bottom-right alignment:

$$\delta=S_k-S_q.$$

For S_q=2, S_k=5:

v2.0 top-left:       v2.1 bottom-right:
1 0 0 0 0            1 1 1 1 0
1 1 0 0 0            1 1 1 1 1

For S_q=5, S_k=2:

v2.0 top-left:       v2.1 bottom-right:
1 0                  0 0
1 1                  0 0
1 1                  0 0
1 1                  1 0
1 1                  1 1

This is a semantic API change, not a floating-point implementation detail. The square training case is unchanged; rectangular prefix/decode-style calls must name the alignment convention in tests and runtime records.

0x05.1G Source-Reported Causal Speedup Is Not A Geometry Theorem

The FlashAttention-2 paper reports roughly 1.7-1.8x causal speedup over its non-causal attention path for the discussed settings. That result is source-reported, not derivable from “half the matrix” alone. The measured ratio also includes:

diagonal-band QK/PV work
element predicates and row reductions
CTA launch and wave quantization
non-matmul scalar work
memory and synchronization behavior
the paper's exact shapes, dtypes, kernels, and A100 hardware

The geometric oracle proves cell and tile counts. It does not predict the paper’s runtime.

0x06 Prefill MQA/GQA Is A Head-Ownership Contract

Prefill supplies many query rows, but MQA and GQA change which K/V head those rows read. Let:

Q: [B, S_q, H_q,  d_h]
K: [B, S_k, H_kv, d_h]
V: [B, S_k, H_kv, d_v]
O: [B, S_q, H_q,  d_v]

The output and LSE remain query-head objects. K/V storage uses H_kv heads.

0x06A Divisibility Produces The Direct Head Map

The admissibility condition is:

$$H_q\bmod H_{kv}=0.$$

Define the contiguous group size:

$$g=\frac{H_q}{H_{kv}}.$$

Then query head h_q reads:

$$h_{kv}(h_q)=\left\lfloor\frac{h_q}{g}\right\rfloor.$$

Renderer-safe form:

assert H_q % H_kv == 0
g = H_q // H_kv
kv_head = q_head // g

For H_q=16, H_kv=4, and g=4:

Q heads  0  1  2  3 -> KV head 0
Q heads  4  5  6  7 -> KV head 1
Q heads  8  9 10 11 -> KV head 2
Q heads 12 13 14 15 -> KV head 3

MQA is the endpoint H_kv=1; all query heads map to K/V head zero. MHA is the other endpoint H_kv=H_q; g=1 and every query head maps to itself.

0x06B One Historical Docstring Reverses The Divisibility Sentence

At commit 4f285b3, the Python examples correctly map six query heads onto two K/V heads, but their sentence says the number of K/V heads must be divisible by the number of query heads. The pinned C++ admission check is the reverse:

TORCH_CHECK(
    num_heads % num_heads_k == 0,
    "Number of heads in key/value must divide number of heads in query"
);

The example, C++ check, and kernel address calculation agree on H_q % H_kv == 0. This is why source reconstruction cannot treat comments, wrapper text, and executable guards as interchangeable evidence.

0x06C Forward Does Not Materialize Repeated K/V Heads

The historical launch grid remains:

grid.x = ceil_div(S_q, B_r)  # query tile
grid.y = B                   # batch item
grid.z = H_q                 # query/output head

Inside one CTA, bidh = blockIdx.z is a query head. The wrapper stores:

h_h_k_ratio = H_q / H_kv

and the forward kernel forms K/V addresses with:

kv_head = bidh / params.h_h_k_ratio;
row_offset_k = ... + kv_head * params.k_head_stride;
row_offset_v = ... + kv_head * params.v_head_stride;

The tensor trace is therefore:

for batch in range(B):
    for q_head in range(H_q):
        kv_head = q_head // g
        for q_tile in range(ceil_div(S_q, B_r)):
            Q_i = Q[batch, q_tile, q_head]
            K_group = K[batch, :, kv_head]
            V_group = V[batch, :, kv_head]
            O[batch, q_tile, q_head] = exact_tiled_attention(Q_i, K_group, V_group)

No [B,S_k,H_q,d_h] K/V tensor is required in forward. The executable parity fixture compares this direct map with explicit K/V repetition for H_q=6, H_kv=2, g=3 and obtains zero binary64 error. Replacing contiguous division with the shape-valid but wrong map q_head % H_kv produces maximum output error 0.11 in the same fixture.

0x06D Reduced Storage Does Not Automatically Produce Grouped CTA Reuse

Consider:

B = 2
S_q = S_k = 257
B_r = 128
B_c = 64
H_q = 16

Every MHA, GQA, or MQA case launches:

$$P_{\mathrm{fwd}}= B H_q\left\lceil\frac{S_q}{B_r}\right\rceil =2\cdot16\cdot3=96$$

programs, because the grid owns query/output heads, not stored K/V heads.

Design	H_kv	g	Forward programs	Unique (batch,kv_head,q_tile) coordinates	Q-head CTAs sharing each coordinate
MHA	16	1	96	96	1
GQA	4	4	96	24	4
MQA	1	16	96	6	16

Across five K/V tiles, the query-head CTA schedule performs 480 logical K/V tile visits in all three cases. The unique grouped coordinates are 480, 120, and 30 respectively. That ratio identifies a reuse opportunity. It is not a claim that GQA causes exactly four HBM reads or MQA sixteen: L2/cache residency, CTA order, and later fused kernels determine physical traffic.

The historical FA2 mapping is a simple, correct alternative:

one CTA owns one query head
several CTAs may address the same stored K/V head
no cross-CTA output reduction is needed

A head-group-fused alternative can place several query heads that share one K/V head in the same work unit, trading a more complex tile and more live query state for stronger explicit K/V reuse. M5.7 owns that grouped-kernel and persistent-cache analysis.

0x06E Backward Must Reduce The Implicitly Shared Heads

Forward uses one stored K/V head as input to g independent query-head outputs. By the chain rule, every one of those output paths contributes to the same learned K/V head:

$$dK_{h_{kv}}= \sum_{h_q:\,h_{kv}(h_q)=h_{kv}} dK^{\mathrm{expanded}}_{h_q},$$ $$dV_{h_{kv}}= \sum_{h_q:\,h_{kv}(h_q)=h_{kv}} dV^{\mathrm{expanded}}_{h_q}.$$

The pinned C++ wrapper makes the ownership conflict explicit. When H_kv != H_q, it allocates expanded gradient buffers shaped by H_q, lets the backward path produce query-head-indexed contributions, reshapes them as:

dK_expanded: [B, S_k, H_q, d_h]
    -> [B, S_k, H_kv, g, d_h]

dV_expanded: [B, S_k, H_q, d_v]
    -> [B, S_k, H_kv, g, d_v]

and sums the group axis. For the oracle’s six-query-head fixture:

expanded dK rows by Q head:
    [1,2], [2,4], [3,6], [4,8], [5,10], [6,12]

reduced dK by KV head:
    KV0 <- Q0+Q1+Q2 = [6,12]
    KV1 <- Q3+Q4+Q5 = [15,30]

Forward index sharing and backward gradient summation are two sides of one trained ownership contract. Section 0x06.1 derives Algorithm 2’s Q/K/V block owners, atomics, and reduction alternatives; this section establishes the final MQA/GQA head-group reduction that composes with that schedule.

This visual reconstructs the mechanism described by FlashAttention-2 and A19’s MQA/GQA section. MQA’s bandwidth motivation comes from Shazeer; GQA’s trained middle point comes from Ainslie et al. Checkpoint conversion, quality, cache bytes, decode reuse, and tensor-parallel placement remain in M5.7.

0x06F Causal, Grouped, And Packed Coordinates Compose Per Sequence

For variable-length batches, BlockInfo derives actual query and K/V extents from cumulative sequence offsets before the kernel computes tile bounds. The complete prefill owner key is therefore:

(batch/packed-sequence, query head, query tile)
    -> K/V head = query_head // g
    -> actual S_q and S_k for this sequence
    -> causal alignment and tile class

Using a global maximum sequence length for every packed item would launch early-return CTAs and could apply the wrong causal geometry if actual lengths are ignored. Using H_kv as the output-head grid would omit query-head outputs. The sequence and head maps are independent axes that meet only in the final address calculation.

0x06G Executable Geometry And Ownership Evidence

The causal/GQA contract is backed by:

causal_gqa_fixture.json
causal_gqa_oracle.py
causal_gqa_report.json
test_causal_gqa_oracle.py

The evidence class is synthetic contract test plus version-scoped source inspection. It proves:

tile classification from integer ranges
historical source visit and mask-band counts
v2.0/v2.1 rectangular mask matrices
direct grouped forward parity against explicit repetition
failure of a wrong but shape-valid head map
expanded-to-grouped dK/dV reduction
prefill-to-decode query-program collapse

It does not compile the extension, observe cache hits, count HBM transactions, or measure speedup.

Failure	Broken contract	Repair
mask scores after rowmax	forbidden scores already changed maximum and denominator	predicate score fragments before online-softmax update
classify tiles only by i and j	wrong for unequal tile sizes, tails, or rectangular alignment	derive from half-open integer ranges
assume every source mask-loop tile is semantically partial	fixed mask band and tails can include fully visible work	report semantic class and schedule path separately
run rectangular causal attention without naming alignment	v2.0 and v2.1 compute different masks	pin version and delta convention
use q_head % H_kv	shape-valid wrong K/V ownership	enforce contiguous q_head // g map
launch output programs over H_kv	missing query-head outputs	retain H_q as output-owner axis unless a fused kernel owns all grouped queries
call direct indexing guaranteed K/V reuse	storage reduction is mistaken for traffic measurement	inspect chosen schedule and profiler bytes
skip grouped gradient sum	K/V receives only one query path’s gradient	reduce expanded dK/dV across group axis

0x06.1 Algorithm 2 Backward Is A Reduction-Topology Contract

The forward pass has one semantic output family, O, indexed by query row. Backward must produce three families with two incompatible ownership axes:

dQ follows query rows
dK follows key rows
dV follows value rows, which share the key-row axis

This section owns the prefill/training backward schedule. It derives the probability reconstruction, maps Algorithm 2 to tensors and lifetimes, traces the active historical v2.0.0 launch, distinguishes atomicity from numerical determinism, and compares alternative owner axes. It does not claim that this historical source is the current backend, and it does not move one-token cache backward into M4.4.

0x06.1A D_i = dO_i dot O_i Eliminates A Materialized Softmax Reduction

Start with one head, no dropout, and the forward pass’s score scale tau (not the maximum-frame rescale factor alpha of 0x02):

$$S=\tau QK^\top,\qquad P=\operatorname{softmax}(S),\qquad O=PV.$$

The direct derivatives are:

$$dP=dOV^\top,$$ $$dQ=\tau dSK,\qquad dK=\tau dS^\top Q,\qquad dV=P^\top dO.$$

The only non-obvious term is dS. For one query row, the softmax Jacobian is:

$$\frac{\partial P_j}{\partial S_k} =P_j(\mathbf{1}_{j=k}-P_k).$$

Contracting that Jacobian with the incoming row dP gives:

$$dS_j =P_j\left(dP_j-\sum_k P_kdP_k\right).$$

The row scalar appears to require reconstructed P and dP over the complete key axis. Associativity exposes a cheaper saved-output identity:

$$\begin{aligned} \sum_k P_kdP_k &=\sum_k P_k(dO\cdot V_k)\\ &=dO\cdot\left(\sum_k P_kV_k\right)\\ &=dO\cdot O. \end{aligned}$$

Define:

$$D_i=dO_i\cdot O_i.$$

Then every score tile can reconstruct its local derivative with the same row scalar:

$$dS_{ij}=P_{ij}(dP_{ij}-D_i).$$

Renderer-safe contract:

S      = tau * Q @ K.T
P_ij   = exp(S_ij - LSE_i)
dP_ij  = dot(dO_i, V_j)
D_i    = dot(dO_i, O_i) = sum_j P_ij * dP_ij
dS_ij  = P_ij * (dP_ij - D_i)

dQ = tau * dS @ K
dK = tau * dS.T @ Q
dV = P.T @ dO

O and rowwise LSE are therefore sufficient saved forward statistics for this reconstruction. The quadratic S, P, dP, and dS objects remain logical tensors; each tile is recomputed and expires on chip. Dropout adds Philox replay and scaling to this path, but it does not change the owner-axis conflict derived below.

0x06.1B Algorithm 2 Names A Column Owner Before It Names A CUDA Grid

Let:

T_r = ceil(S_q / B_r) query tiles
T_c = ceil(S_k / B_c) K/V tiles

Algorithm 2 places the K/V tile loop outside the query-tile loop. Its storage contract is:

Algorithm stage	Tensor view	Typical source dtype	Lifetime and owner
partition inputs	Q_i,O_i,dO_i:[B_r,d], K_j,V_j:[B_c,d]	element dtype	logical views; no movement yet
saved statistics	LSE_i,D_i:[B_r]	FP32	query-row state loaded for each visible pair
recomputed score state	S_ij,P_ij,dP_ij,dS_ij:[B_r,B_c]	FP32 accumulation with element fragments for tensor-core paths	one tile iteration
local K/V gradients	acc_dK_j,acc_dV_j:[B_c,d]	FP32 accumulators	live across all visible query tiles for column j
query contribution	Delta dQ_i^(j):[B_r,d]	FP32	one (i,j) pair; must join other columns
final outputs	dQ:[S_q,d], dK,dV:[S_k,d]	requested element dtype	HBM outputs after reduction/conversion

Paper lines 5-18 can be read as one transaction:

choose K/V column j
    load K_j and V_j once
    initialize dK_j and dV_j once
    stream every visible query tile i
        reconstruct P_ij and dS_ij
        accumulate dK_j and dV_j locally
        contribute to dQ_i
    commit dK_j and dV_j once

That outer-loop choice makes one worker the natural owner of dK_j/dV_j. The paper’s sequential pseudocode reads, updates, and writes dQ_i for every column. The parallel implementation cannot use an ordinary read-modify-write, because multiple column workers can reach the same dQ_i concurrently.

0x06.1C The Gradient Equations Define The Producer Graph

At tile granularity define:

$$\Delta dQ_i^{(j)}=\tau dS_{ij}K_j,$$ $$\Delta dK_j^{(i)}=\tau dS_{ij}^{\top}Q_i, \qquad \Delta dV_j^{(i)}=P_{ij}^{\top}dO_i.$$

The complete gradients are reductions over opposite axes:

$$dQ_i=\sum_{j\in\mathcal{V}(i)}\Delta dQ_i^{(j)},$$ $$dK_j=\sum_{i\in\mathcal{V}(j)}\Delta dK_j^{(i)}, \qquad dV_j=\sum_{i\in\mathcal{V}(j)}\Delta dV_j^{(i)}.$$

V(i) is the set of K/V tiles visible to query tile i; V(j) is the set of query tiles that can see K/V tile j. The conflict cannot be removed by naming one two-dimensional score tile as owner:

Main CTA owner	Direct single-writer outputs	Cross-CTA reduction
K/V column (j,b,h_q)	expanded dK_j, expanded dV_j	dQ_i across K/V columns
query row (i,b,h_q)	dQ_i	dK_j, dV_j across query rows and grouped query heads
one complete head (b,h_q)	expanded dQ,dK,dV	grouped-head sum only

This is a conservation rule: changing the owner axis moves the reduction. It does not make all three outputs single-writer at once.

0x06.1D The Active Historical Source Is A Three-Kernel Column-Owner Pipeline

At commit 4f285b3, run_flash_bwd unconditionally calls run_flash_bwd_seqk_parallel. For one specialization, the launch sequence is:

1. preprocess
   grid = (ceil(S_q/B_r), B, H_q)
   compute D = rowsum(dO * O)
   clear FP32 dQaccum

2. main column-owner kernel
   grid = (ceil(S_k/B_c), B, H_q)
   one CTA owns (K/V tile j, batch b, query head h_q)
   accumulate expanded dK_j/dV_j locally
   atomicAdd every dQ_i contribution into FP32 dQaccum

3. dQ conversion
   grid = (ceil(S_q/B_r), B, H_q)
   load FP32 dQaccum, apply final scale, cast, write dQ

4. wrapper-level GQA/MQA closure when H_q != H_kv
   reshape expanded dK/dV over [H_kv, g]
   sum the g query-head contributions

The three CUDA launches are ordered on the same stream. The preprocessing launch establishes zeroed atomic destinations before the main grid, and the conversion launch reads them only after the main grid has completed. No CTA in the main grid waits for another CTA or assumes which one runs first.

The diagram reconstructs source ownership and tensor state. It is not a CUDA timeline, a physical-traffic trace, or a current-backend claim.

0x06.1E One Active CTA Reconstructs Five Matmul-Shaped Paths

For main-program coordinate (n_block,b,h_q), the historical kernel computes:

j       = n_block
h_kv    = floor(h_q / g)
K_j,V_j = K/V[b, j, h_kv]

It retains FP32 acc_dk and acc_dv across visible query blocks. For each query block, processed in reverse order, it performs the five backward matmul-shaped operations counted by the paper:

1. recompute S_ij = tau * Q_i @ K_j.T
2. compute    dP_ij = dO_i @ V_j.T
3. accumulate dV_j += P_ij.T @ dO_i
4. accumulate dQ_i += tau * dS_ij @ K_j
5. accumulate dK_j += tau * dS_ij.T @ Q_i

The softmax tile is reconstructed from S_ij and LSE_i; causal and tail predicates run before invalid entries can contribute. P and dS use shared memory as layout bridges between MMA paths. sP and sdQ alias storage at different lifetimes, so the barriers around their reuse are correctness barriers, not decorative synchronization.

Conceptual source-faithful pseudocode:

def backward_column_owner(params, b, q_head, kv_tile):
    kv_head = q_head // params.group_size
    K_j = load_K(b, kv_head, kv_tile)
    V_j = load_V(b, kv_head, kv_tile)
    acc_dK = fp32_zeros([B_c, d])
    acc_dV = fp32_zeros([B_c, d])

    for q_tile in visible_query_tiles_reverse(kv_tile):
        Q_i, dO_i = load_main_query_state(b, q_head, q_tile)
        LSE_i = load_lse(b, q_head, q_tile)
        D_i = load_D(b, q_head, q_tile)

        S_ij = tau * Q_i @ K_j.T
        S_ij = apply_tail_and_causal_mask_before_exp(S_ij)
        P_ij = exp(S_ij - LSE_i[:, None])  # undropped
        # Z_ij: 0/1 mask, each entry kept with probability keep
        # (Z_ij = 1 and keep = 1 without dropout)
        Z_ij = replay_optional_dropout_mask(b, q_head, q_tile, kv_tile)

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

        acc_dV += (P_ij * Z_ij / keep).T @ dO_i
        acc_dK += tau * dS_ij.T @ Q_i
        # The scratch omits tau, which conversion applies; dS_ij already
        # carries 1/keep, so conversion must not apply it again.
        atomic_add_fp32(dQaccum[i], dS_ij @ K_j)

    store_element_dtype(dK_expanded[b, q_head, kv_tile], acc_dK)
    store_element_dtype(dV_expanded[b, q_head, kv_tile], acc_dV)

Under dropout the two paths need different probabilities. dV needs the dropped P_ij*Z_ij/keep that formed O, while dS needs the undropped P_ij with the masked and rescaled dP_ij. With this dP_ij, D_i=dot(dO_i,O_i)=sum_j P_ij*dP_ij still holds, because O is the dropped output. v2.0.0 places the 1/keep factor differently: it marks dropped elements in the sign bit of the recomputed P, masks dP without 1/keep, and scales D by keep in the preprocessing kernel. Its loop therefore produces keep*dS_ij and keep*dV_j. It applies 1/keep once at the end: together with tau to dK and, during conversion, to dQaccum, and alone to dV. Both placements give the same gradients.

The paper estimates backward FLOPs as 2.5x forward because backward has five matmuls including score recomputation while forward has two. That is a paper-reported FLOP convention, not a runtime theorem: exponentials, masks, dropout replay, atomics, conversions, launch overhead, and grouped reduction are outside the simple 5 / 2 ratio.

0x06.1F Atomic Add Solves Lost Updates, Not Reduction Ordering

An ordinary update is unsafe:

dQaccum[row, col] += contribution

Suppose two column CTAs both read zero and independently compute +0.75 and -0.20. If their ordinary stores race and the second writer wins:

correct sum:  0.55
stored value: -0.20
absolute error from lost update: 0.75

The source instead selects GmemTiledCopydQaccumAtomicAdd and issues scalar FP32 atomicAdd operations for dQaccum. CUDA defines atomicAdd as one atomic read-modify-write at an address, so neither contribution is lost. Atomicity does not impose a fixed order across CTAs. The CUDA programming model permits thread blocks to run in any order, and FP32 addition is not associative. The oracle’s explicit order example gives:

float32((1e20 + -1e20) + 1) = 1
float32((1e20 + 1) + -1e20) = 0

Therefore:

atomicity contract:
    every producer contributes exactly once to one address

numerical-order contract:
    arrival order is not fixed by the grid

determinism contract:
    bitwise repeatability is not established by atomicAdd

A deterministic alternative can materialize one partial tensor per producer and reduce those partials in a fixed tree. That replaces atomic contention and unordered arrival with more scratch storage, more writes, a reduction launch, and a prescribed summation tree. It is a design alternative, not the active historical path.

0x06.1G Causality Makes Producer Multiplicity Triangular

Use a square four-by-four tile grid with B_r=B_c=128. Under top-left causal semantics, the visible tile pairs are:

q0 <- k0
q1 <- k0,k1
q2 <- k0,k1,k2
q3 <- k0,k1,k2,k3

The active column-owner main grid still launches one CTA per K/V tile and query head. Its logical dQ producer multiplicity is:

Query tile	Column CTAs contributing to its dQ rows
q0	1
q1	2
q2	3
q3	4

Each expanded dK_j/dV_j tile still has one main-kernel writer: that owner loops over all query tiles that can see j. The inactive row-owner alternative reverses the triangular pattern. Before grouped heads are considered, K/V tile producer counts become [4,3,2,1].

The diagonal tile still needs element predicates. Tile-level producer counts say which CTAs may contribute to an output tile; they do not say every score cell within the tile is visible or that every atomic experiences equal contention.

0x06.1H GQA Adds A Head Reduction After The Tile Reduction

The active main grid uses H_q, not H_kv:

grid = (T_c, B, H_q)

For GQA group size g=H_q/H_kv, query heads in one group read the same stored K/V head but follow different query paths. To keep main-kernel writes single-owner, the wrapper supplies element-dtype expanded buffers:

dK_expanded: [B,S_k,H_q,d]
dV_expanded: [B,S_k,H_q,d]

The main CTA (j,b,h_q) writes its own expanded slice directly. After the kernel, the wrapper performs:

[B,S_k,H_q,d]
    -> reshape [B,S_k,H_kv,g,d]
    -> sum axis g
    -> [B,S_k,H_kv,d]

Two independent reductions now exist:

tile-axis reduction:
    K/V-column CTAs atomicAdd into dQaccum across j

head-group reduction:
    expanded dK/dV query-head slices sum across g

Calling both simply “the backward reduction” hides their different axes, dtypes, synchronization mechanisms, and storage costs.

0x06.1I The Source Contains A Row-Owner Alternative, But Does Not Dispatch It

The pinned source also contains a sequence-Q-parallel helper and launcher. Its dispatch and required dKaccum/dVaccum wrapper allocations are commented out at this pin, so this is a dormant implementation branch rather than a ready runtime choice. If completed and selected, it would:

clear FP32 dKaccum and dVaccum over (K/V tile, batch, H_kv)
launch one main CTA over (query tile, batch, H_q)
compute dQ directly for that query tile
atomicAdd dK/dV contributions into grouped K/V scratch
convert FP32 dKaccum/dVaccum to output dtype

For the oracle coordinate:

B=1, H_q=8, H_kv=2, g=4
S_q=S_k=512, B_r=B_c=128

the non-causal owner counts are:

Schedule	Direct owner	Atomic producer multiplicity
active column owner	expanded dK/dV	each dQ tile receives 4 column CTAs
row-owner alternative	dQ	each final grouped dK/dV tile receives 4 query tiles * g=4, or 16 CTA paths

For the causal coordinate, the row-owner K/V multiplicities become [16,12,8,4]. That path can avoid an explicit expanded-head reduction by accumulating directly into H_kv scratch, but it moves atomics onto two output families and combines sequence and head-group contention.

The run_flash_bwd body calls the sequence-K path unconditionally at this historical pin; the sequence-Q dispatch is commented out. Its presence is evidence for a considered implementation alternative, not evidence that the alternative ran in published benchmarks.

The broader design space is:

Design	Parallel surface	Reduction cost	Main trade-off
one CTA per complete query head	B*H_q	no tile-axis global atomics	poor occupancy for small batch/head count and long sequences
K/V-column CTAs	B*H_q*T_c	FP32 atomic dQ, then optional GQA group sum	active historical FA2 choice
query-row CTAs	B*H_q*T_r	FP32 atomic dK/dV, naturally targets H_kv scratch	source-present inactive alternative
per-producer partial tensors plus fixed reduction	configurable	extra storage, writes, and reduction kernels	deterministic order is possible

This comparison also prevents a terminology error. FA2 avoids a split-K warp partition for an MMA output tile, but its active backward schedule still splits the sequence-K axis across CTAs and reduces their dQ contributions atomically. Warp-level split-K and a grid-level sequence-K split are different hierarchy levels.

0x06.1J Scratch Memory Is Part Of The Algorithm Contract

For:

B=1, H_q=32, H_kv=8, g=4
S_q=S_k=4096, d=128
B_r=B_c=128
element dtype=2 bytes, accumulator dtype=4 bytes

the active historical source-level ledger is:

Object	Shape	Width	Size
dQaccum	[1,32,4096,128]	FP32	64 MiB
D / softmax_d	[1,32,4096]	FP32	0.5 MiB
dK_expanded	[1,4096,32,128]	element dtype	32 MiB
dV_expanded	[1,4096,32,128]	element dtype	32 MiB
one final grouped dK or dV	[1,4096,8,128]	element dtype	8 MiB

The declared backward scratch subset is therefore:

$$64+0.5+32+32=128.5\ \text{MiB}.$$

declared backward scratch subset = 128.5 MiB

It excludes Q/K/V/O/dO, saved LSE, final gradients, allocator effects, shared memory, registers, and any framework reduction workspace. “FlashAttention uses linear memory” does not mean “backward has no material scratch.” These objects are linear in sequence length, but their head and width multipliers matter.

The same coordinate launches:

1,024 preprocess programs
1,024 column-owner main programs
1,024 dQ-conversion programs
plus the grouped dK/dV reductions

For non-causal attention, each dQ element has 32 K/V-tile producers. The oracle’s one-logical-add-per-element ledger contains 536,870,912 FP32 dQ element updates. Top-left causal tiling reduces the corresponding full-tile ledger to 276,824,064. These counts describe logical atomic update slots in the reconstructed schedule. They do not measure coalescing, cache sectors, atomic serialization, contention latency, or executed tail instructions.

0x06.1K Executable Backward Evidence And Failure Closure

The backward-ownership contract is reproducible with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_backward_ownership_oracle.py

The checked artifacts are:

backward_ownership_fixture.json
backward_ownership_oracle.py
backward_ownership_report.json
test_backward_ownership_oracle.py

They prove binary64 reference-gradient parity for column- and row-owner schedules, causal and grouped producer multiplicities, the declared scratch ledger, a lost-update failure, and an FP32 order-sensitivity example. Pinned source inspection separately establishes the active launcher, grids, atomic target, conversion path, direct expanded dK/dV commits, and grouped wrapper sum.

They do not compile the extension, count executed atomics, measure contention, profile memory traffic, or establish bitwise determinism.

Failure	Broken contract	Repair
parallel column CTAs use ordinary dQ += partial	read-modify-write races lose contributions	clear FP32 scratch, then use device-wide atomic adds
convert each partial to FP16/BF16 before global accumulation	repeated low-precision rounding changes the reduction	accumulate shared dQ state in FP32, then convert once
call atomic addition deterministic	atomicity is confused with a fixed summation order	report numerical tolerance or use a fixed reduction tree
let two column CTAs write final dK_j/dV_j	K/V ownership was not preserved	keep all visible query tiles inside one column owner
launch column owners over H_kv in GQA	query-head gradient paths disappear	launch over H_q, then reduce expanded K/V gradients across g
skip the preprocessing clear	atomic scratch begins from undefined data	preserve launch ordering and explicit zero initialization
call source-present row ownership active	helper code is mistaken for selected dispatch	inspect run_flash_bwd, not only available templates
treat 5/2 as measured slowdown	matmul FLOPs omit scalar, atomic, conversion, and launch work	keep the ratio as a paper FLOP convention and benchmark the full path
say FA2 eliminates split-K everywhere	warp-level work partition is conflated with grid-level sequence splitting	name hierarchy, owner, target, and reduction separately

0x06.2 One Backward CTA Is A Storage-Lifetime Contract

Section 0x06.1 established the grid-level owner: one active main CTA owns one K/V column tile for one (batch,query_head) coordinate. That is not yet an implementation trace. Inside that CTA, five MMA paths need incompatible output shapes, P and dS need shared-memory layout bridges, dK and dV remain live across the complete query loop, asynchronous copies need publication edges, and several logical tensors deliberately reuse the same byte interval.

This section owns those CTA internals for the pinned v2.0.0 source. It derives declared dynamic shared memory, source-visible barrier hazards, CuTe atom replication, and source-named accumulator-fragment pressure. It does not claim a compiled register count, spill count, bank-conflict count, achieved occupancy, or barrier latency. Those require a built kernel and profiler evidence.

0x06.2A The Active Template Disables Two Source Branches

The active sequence-K launcher instantiates the main helper as:

compute_dq_dk_dv_1colblock<
    Kernel_traits,
    Is_dropout,
    Is_causal,
    false,  // Is_first
    false,  // Is_last
    true    // Seq_parallel
>(params, bidb, bidh, n_block);

Those booleans change the live storage graph:

Template fact	Consequence in the active main CTA
Is_first=false	D=rowsum(dO*O) is loaded from the preprocessing kernel; this CTA does not use the source branch that recomputes and stores it
Is_last=false	the shared sdQ output path is inactive; each query-tile contribution is atomically added to FP32 dQaccum
Seq_parallel=true	the CTA directly commits its expanded dK/dV column after the reverse query loop

In other words, D is produced by the preprocessing kernel, while this active main CTA consumes it.

The source still declares storage for branches used by other template combinations. Reading every declared pointer as simultaneously live would overstate the active path. In particular, sdQ points at the same shared address as sP, but the active Is_last=false main CTA never materializes a shared dQ tile there.

0x06.2B Five MMA Paths Need Three Warp-Replication Layouts

For one visible query tile, let:

Q,dO : [B_r,d]
K,V  : [B_c,d]
S,P,dP,dS : [B_r,B_c]

The CTA executes five matmul-shaped paths:

Path	Contraction	FP32 C-accumulator shape	Lifetime
score	Q @ K.T	[B_r,B_c]	one query-tile iteration; transformed into probability state
probability gradient	dO @ V.T	[B_r,B_c]	one iteration; transformed pointwise into dS
value gradient	P.T @ dO	[B_c,d]	acc_dV persists across all query tiles
query contribution	dS @ K	[B_r,d]	one iteration; then atomic-adds to global FP32 scratch
key gradient	dS.T @ Q	[B_c,d]	acc_dK persists across all query tiles

One warp arrangement is not assumed to fit all three output geometries. The pinned trait declares:

TiledMmaSdP: (AtomLayoutMSdP, warps/AtomLayoutMSdP, 1)
TiledMmadKV: (AtomLayoutNdKV, warps/AtomLayoutNdKV, 1)
TiledMmadQ:  (AtomLayoutMdQ,  warps/AtomLayoutMdQ,  1)

For the high-shared-memory d=128 specialization (B_r=64,B_c=128,W=8), the three CuTe replication shapes are:

score and dP: (2,4,1)
dK and dV:    (4,2,1)
dQ:           (2,4,1)

These tuples describe CTA-level replication of an MMA atom. They are not a literal lane-to-element map. Such a map also depends on the selected MMA atom, copy atoms, fragment layouts, dtype, and compiler lowering.

0x06.2C Shared Memory Is Derived From Simultaneous Lifetimes

Let one element occupy b_e bytes. The Q/dO prefix contains either one or two Q stages plus one dO tile:

$$E_{QdO}=B_r d \begin{cases} 2,&\text{single-buffer Q}\cr 3,&\text{double-buffer Q}. \end{cases}$$

The remaining logical tile sizes, in elements, are:

$$E_{KV}=2B_cd,\qquad E_{dS}=B_rB_c,\qquad E_P=B_rB_c.$$

When V remains in shared memory, all four intervals are retained in the main arena:

$$S_{\mathrm{shared\ V}} =b_e(E_{QdO}+E_{KV}+E_{dS}+E_P).$$

When Is_V_in_regs=true, every thread captures its V fragment before dS is published. The former V interval can then become the start of sdS, so the post-Q/dO arena is the larger of the original K/V pair and the K plus derived state:

$$S_{\mathrm{V\ in\ regs}} =b_e\left[E_{QdO}+\max\left(E_{KV},B_cd+E_{dS}+E_P\right)\right].$$

Renderer-safe formulas:

QdO elements = B_r*d*(2 if single-buffer Q else 3)
KV elements  = 2*B_c*d
dS elements  = B_r*B_c
P elements   = B_r*B_c

shared-V total = bytes * (QdO + KV + dS + P)
V-in-reg total = bytes * (QdO + max(KV, B_c*d + dS + P))

For element width 2 bytes:

Specialization	Q/dO prefix	Post-prefix arena	Dynamic shared memory
d=128, B_r=64, B_c=128, double Q, shared V	48 KiB	K/V 64 + dS 16 + P 16 = 96 KiB	144 KiB
d=128, B_r=64, B_c=64, double Q, V in registers	48 KiB	max(K/V 32, K 16 + dS 8 + P 8) = 32 KiB	80 KiB
d=256, B_r=B_c=64, double Q, shared V	96 KiB	80 KiB	176 KiB
same d=256 tile, single Q, shared V	64 KiB	80 KiB	144 KiB

The lower-memory branches do not perform the same schedule with free savings. The d=128 fallback also shrinks B_c and keeps V in registers. The d=256 fallback removes one Q stage, reducing overlap and introducing a loop-carried reuse barrier.

0x06.2D An Alias Is Valid Only After Its Last Reader

The source maps pointers into one dynamic shared-memory allocation. In the double-buffer, shared-V case, the conceptual address order is:

[ Q stage 0 ][ Q stage 1 ][ dO ][ K ][ V ][ dS ][ P / inactive sdQ ]

With V captured in registers, sdS starts where sV started:

[ Q stage 0 ][ Q stage 1 ][ dO ][ K ][ V then dS ][ P / inactive sdQ ]

After the reverse query loop, the epilogue reinterprets the K/V-and-derived arena:

old [ K ][ V / dS ][ P ]
new [ dK ][ dV ]

Three different alias contracts are present:

1. sV -> sdS is legal only after every warp has copied its V fragment into registers.
2. sP <-> sdQ is a compile-time alias, but the active main CTA does not use the sdQ lifetime.
3. sK/sV/derived state -> sdK/sdV is legal only after the last loop consumer has finished and before global epilogue copies begin.

An address equality is therefore not enough to explain correctness. The proof needs a last-reader edge, a barrier if readers span the CTA, and a new producer lifetime after that edge.

0x06.2E Fence, Wait, And CTA Barrier Have Different Jobs

The pinned kernel uses asynchronous global-to-shared copies. Its relevant sequence is conceptually:

copy_async(Q_or_next_Q, sQ_stage);
copy_async(K, sK);
copy_async(dO, sdO);
copy_async(optional_V, sV);
cp_async_fence();

cp_async_wait<0>();
__syncthreads();
consume_shared_operands();

cp_async_fence() commits or orders the current asynchronous-copy group; it is not by itself a CTA-wide publication barrier. The consumer path waits for the required asynchronous groups and then uses __syncthreads() so all threads cross the shared-memory producer/consumer edge. NVIDIA’s CUDA documentation likewise separates asynchronous copy completion from thread-block synchronization.

The active source hazards are:

Producer -> consumer	Reused storage	Required source edge	Failure without the edge
async Q/K/dO/V copies -> MMA operand readers	sQ,sK,sdO,sV	cp_async_wait then __syncthreads()	consumers can observe incomplete tiles
shared V -> per-thread V fragment	sV, then sdS when V is in registers	wait plus CTA barrier	dS can overwrite V before all warps capture it
fragment dS stores -> dQ/dK operand readers	sdS	CTA barrier	MMA reads a partially published derivative tile
dV readers of current dO -> next dO async copy	sdO	CTA barrier	next dO overwrites rows still consumed by dV
dK readers of Q -> next Q copy in single-buffer mode	sole sQ stage	CTA barrier, then next copy group	next Q overwrites the only stage too early
loop consumers -> dK/dV epilogue producers	aliased K/V arena	CTA barrier when required by branch	epilogue overwrites still-live operands
fragment dK/dV stores -> global epilogue readers	sdK,sdV	CTA barrier	global copy observes incomplete CTA-produced tiles

This is also why “FA2 removed a reduction” must not become “FA2 removed CTA synchronization.” Split-Q removes one output-state reduction topology in the forward kernel; operand publication and storage reuse still require barriers.

0x06.2F Barrier Count Follows Hazards, Not The Number Of Matmuls

Let T be the number of full visible query tiles processed by one active column CTA, with no early exit. The source-path block-barrier executions are:

async operand publication: T
dS publication:            T
dO reuse protection:       T
dK/dV epilogue publish:    1
V-register capture:        +1 when V is in registers
V-reg epilogue alias:      +1 when V is in registers
single-Q-stage reuse:      +(T-1) when Q is not double buffered

Therefore:

Q staging	V location	Source-path barrier formula	T=32
double buffer	shared memory	3T+1	97
double buffer	registers	3T+3	99
single buffer	shared memory	4T	128
single buffer	registers	4T+2	130

These are dynamic executions along the reconstructed full path, not emitted SASS instruction counts or measured stalls. Causal culling, sequence tails, template branches, and an empty column can change T or bypass work. A larger count also does not directly imply proportionally larger latency: the cost depends on arrival skew, outstanding dependencies, occupancy, and hardware.

0x06.2G Accumulator Pressure Can Be Bounded Before Compiling, Not Measured

With W warps, the CTA has N_t=32W threads. For the source’s MMA output fragments, the average logical FP32 C-accumulator elements per thread are:

$$n_S=n_{dP}=\frac{B_rB_c}{N_t},\qquad n_{dQ}=\frac{B_rd}{N_t},\qquad n_{dK}=n_{dV}=\frac{B_cd}{N_t}.$$

acc_dK and acc_dV persist across every query tile. A source-named live-slot ledger is therefore:

$$R_{S}=n_{dK}+n_{dV}+n_S,$$ $$R_{dP}=n_{dK}+n_{dV}+n_S+n_{dP},$$ $$R_{dQ}=n_{dK}+n_{dV}+n_{dP}+n_{dQ}.$$

Define the named-fragment peak:

$$R_{\mathrm{named}}=\max(R_S,R_{dP},R_{dQ}).$$

For three eight-warp specializations:

Traits	nS,nP,nQ,nK,nV FP32 values/thread	Persistent dK+dV	R_named
d=128, 64x128, shared V	32,32,32,64,64	128	192
d=128, 64x64, V in registers	16,16,32,32,32	64	112
d=256, 64x64, shared V	16,16,64,64,64	128	208

The first row names 49,152 FP32 fragment slots across the CTA (192*256). That is a tensor-fragment liveness ledger, not a claim that ptxas allocates exactly 192 registers per thread. It excludes A/B operand fragments, fragment copies, LSE/D scalars, predicates, RNG state, addresses, compiler reuse, allocation granularity, and spills.

For the low-memory d=128 branch, V-in-registers adds 32 element-dtype V values per thread. At FP16/BF16 width that is only a packing lower bound of 16 32-bit slots; actual compiled allocation can be larger. The branch reduces the named FP32 peak by shrinking B_c, but then adds the V operand lifetime. Only ptxas --resource-usage, local memory diagnostics, and a GPU profile can close that trade.

0x06.2H The Specialization Matrix Moves Pressure Between Resources

The historical launcher contains a resource policy, not one universal tile:

Head width and branch	Tile	V state	Q staging	Declared shared memory
d=64, high-shared-memory path	128x128	shared	double	144 KiB
d=64, fallback	64x128	registers	double	72 KiB
d=96, no dropout	64x128	registers	double	92 KiB
d=96, dropout path	64x128	shared	double	116 KiB
d=128, high-shared-memory path	64x128	shared	double	144 KiB
d=128, fallback	64x64	registers	double	80 KiB
d=160, high/fallback	64x64	shared	double/single	116/96 KiB
d=192, high/fallback	64x64	shared/registers	double/single	136/96 KiB
d=256, high/fallback	64x64	shared	double/single	176/144 KiB

Moving V into registers can shrink the shared arena but lengthens a register operand lifetime and adds alias barriers. Removing double buffering saves one B_r*d stage but loses the safe next-Q prefetch window. Shrinking B_c reduces score, probability, K/V, and persistent-gradient fragments, but creates more K/V-column CTAs and therefore more grid-level dQ producers for the same sequence. Section 0x06.1 owns that last reduction effect; this section owns why the local tile changed.

Dropout can also select a different storage branch at the same head width. That is a reminder that algorithmic features alter RNG and probability-state lifetimes, not merely a scalar after the matmul.

0x06.2I The Epilogue Is A Final Alias Transition

After the last query tile, the CTA no longer needs K/V operands, P, or dS. The source converts FP32 acc_dK/acc_dV fragments to the output element dtype and stages them in the old operand arena:

finish all query-tile consumers
    -> protect the aliased arena where required
store converted dK fragments into sdK over old sK
store converted dV fragments into sdV over old sV/derived-state storage
    -> __syncthreads()
cooperatively copy sdK and sdV to expanded global output slices

The final barrier is a producer/consumer barrier inside one CTA: fragment owners publish shared rows, and global-copy threads consume them. It is not the grid-level GQA reduction and not the FP32 dQ atomic reduction. Those happen at different hierarchy levels and in different storage.

0x06.2J Executable CTA Evidence And Failure Closure

The CTA storage-lifetime contract is reproducible with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_backward_cta_oracle.py

The checked artifacts are:

backward_cta_fixture.json
backward_cta_oracle.py
backward_cta_report.json
test_backward_cta_oracle.py

They derive all 15 declared historical specialization rows, exact shared-memory formulas and alias offsets, three CTA-level MMA replication shapes, barrier formulas for a full 32-query-tile path, and source-named accumulator-fragment pressure. Pinned source inspection establishes the active template flags and synchronization sites. The oracle does not compile the extension or emulate CuTe instructions.

Failure	Broken contract	Repair
treat every declared shared pointer as active	inactive template branches inflate or confuse the live set	begin with Is_first, Is_last, and Seq_parallel
overwrite V with dS before all warps capture it	shared alias begins before the old lifetime ends	wait for V, synchronize, then start the sdS lifetime
call cp_async_fence a publication barrier	copy-group commit is confused with completion and CTA visibility	preserve cp_async_wait and the block barrier before consumers
remove the dS barrier because one warp produced its rows	dQ/dK operand readers can belong to other warps	synchronize the CTA after shared dS publication
single-buffer Q without a reuse edge	next async copy races the current dK consumer	add the loop-carried barrier or retain two Q stages
report R_named as registers/thread	tensor-fragment arithmetic is confused with compiler allocation	label it a source-level pressure ledger and inspect ptxas
declare V-in-registers always cheaper	shared-memory savings hide register lifetime and extra barriers	compare compiled spills, occupancy, and timing for the selected branch
overwrite the K/V arena before the loop ends	epilogue aliases live operands	start sdK/sdV only after the last reader and required barrier
infer barrier latency from 3T+1	execution count is confused with stall cost	profile barrier stalls and arrival skew on the target GPU

0x06.3 Backward Arithmetic Intensity Needs Two Byte Ledgers

Section 0x06.1 established who repeats data across CTAs. Section 0x06.2 established what one CTA retains on chip. Neither result is a physical HBM measurement. This section connects them with an explicit source-logical I/O model, derives the five-path matmul work, and then specifies the profiler experiment needed to replace that model with hardware counters.

The distinction matters because three quantities are often called simply “traffic” or “arithmetic intensity”:

source-logical bytes:
    valid tensor elements requested or committed by the reconstructed schedule

physical hierarchy bytes:
    sectors transferred through L1, L2, DRAM, and local memory on one GPU run

semantic tensor size:
    unique elements in Q, K, V, gradients, and saved state

The first is derived here. The second remains unmeasured. The third is not a schedule: a tensor can be logically read many times by different owners even though it has one stored copy.

0x06.3A Five Matmuls Give 10 B_r B_c d FLOPs Per Scheduled Tile Pair

Use the conventional accounting rule that one fused multiply-add is two FLOPs. Each of the five contractions derived in 0x06.2 performs a matrix product with B_r B_c d multiply-adds:

$$F_{\mathrm{one\ path}}=2B_rB_cd.$$

Therefore one scheduled score-tile pair carries:

$$F_{\mathrm{pair}}=5(2B_rB_cd)=10B_rB_cd.$$

This is the derivation behind the paper’s 2.5x backward-to-forward matmul ratio: backward has five matmul-shaped paths where forward has two. It does not count maxima, exponentials, probability reconstruction, dropout replay, atomics, casts, predicates, address arithmetic, or launch overhead.

For causal and tail-sensitive accounting, define the visited tile set T and two score areas:

$$A_{\mathrm{scheduled}} =\sum_{(i,j)\in\mathcal{T}}B_rB_c,$$ $$A_{\mathrm{useful}} =\sum_{(i,j)\in\mathcal{T}} \sum_{q\in Q_i}\sum_{k\in K_j}\mathbf{1}[k\le q].$$

Then, for batch B and H_q query heads:

$$F_{\mathrm{scheduled}}=10BH_qdA_{\mathrm{scheduled}}, \qquad F_{\mathrm{useful}}=10BH_qdA_{\mathrm{useful}}.$$

The scheduled numerator includes masked cells inside a visited diagonal tile and padded slots in a tail tile. The useful numerator does not. Neither is an instruction counter: tensor-core lowering, predication, and compiler behavior still require a compiled kernel.

0x06.3B Column Ownership Repeats Query State, Not K/V State

Let r_i be the number of valid query rows in tile i, and define the source-logical row-visit count:

$$R=\sum_{(i,j)\in\mathcal{T}}r_i.$$

Use b_e bytes for Q/K/V/O/dO and final gradients, and b_a bytes for FP32 LSE, D, and dQaccum. The active column-owner main grid has the following ledger:

$$B_{KV}=2BH_qS_kdb_e,$$ $$B_{Q,dO}=2BH_qRdb_e,$$ $$B_{LSE,D}=2BH_qRb_a,$$ $$B_{\Delta dQ}=BH_qRdb_a,$$ $$B_{dK,dV}=2BH_qS_kdb_e.$$

The sum is the main-kernel logical byte ledger:

$$B_{\mathrm{main}} =B_{KV}+B_{Q,dO}+B_{LSE,D}+B_{\Delta dQ}+B_{dK,dV}.$$

The terms follow directly from ownership:

1. One column CTA loads its K and V tile once, so the complete K/V axis is loaded once per query head.
2. That CTA loops over visible query tiles, so Q, dO, LSE, and D are repeated according to R.
3. Every visited pair contributes one FP32 dQ tile update.
4. The owner writes its expanded dK/dV column once.

H_q, rather than H_kv, appears in the K/V term because the historical grid launches one column owner per query head. GQA storage sharing creates a cache reuse opportunity; it does not prove that physical K/V bytes collapse by the group size.

B_Delta dQ counts four payload bytes per logical FP32 update. A hardware atomic is a read-modify-write operation whose requests, sectors, cache path, replay, and serialization are not represented by those four bytes. Calling this term measured atomic traffic would be wrong.

Also note what is absent: the active main CTA does not read O. The preprocessing kernel consumes O and dO to produce D; the main kernel consumes Q, dO, LSE, and D.

0x06.3C The Pipeline Boundary Adds Fixed And Grouped Traffic

The historical backward operation is not one kernel. Its complete logical ledger includes the surrounding stages from 0x06.1.

Preprocessing reads O and dO, writes D, and clears FP32 dQaccum:

$$B_{\mathrm{pre}} =2BH_qS_qdb_e +BH_qS_qb_a +BH_qS_qdb_a.$$

The dQ conversion kernel reads FP32 scratch and writes the requested output dtype:

$$B_{\mathrm{convert}} =BH_qS_qd(b_a+b_e).$$

For GQA group size g=H_q/H_{kv}>1, the wrapper reads two expanded gradient tensors and writes two grouped tensors:

$$B_{\mathrm{GQA}} =2BH_qS_kdb_e+2BH_{kv}S_kdb_e.$$

For MHA, where g=1, that separate wrapper reduction is absent. The complete source-logical pipeline ledger is:

$$B_{\mathrm{pipeline}} =B_{\mathrm{pre}}+B_{\mathrm{main}} +B_{\mathrm{convert}}+B_{\mathrm{GQA}}.$$

The preprocess dot products, dQ scaling, and grouped additions can be reported separately, but adding only those operations still does not produce a complete scalar FLOP count. The exponential and pointwise derivative path remains outside this matmul convention.

0x06.3D A 4,096-Token Ledger Is Dominated By Repeated Rows And dQ Updates

Use the same training coordinate as 0x06.1:

B=1, H_q=32, H_kv=8, g=4
S_q=S_k=4096, d=128
B_r=64, B_c=128
b_e=2 bytes, b_a=4 bytes

For non-causal attention, each of 64 query tiles visits all 32 K/V tiles:

visited pairs per query head = 64 * 32 = 2,048
row visits per query head     = 4,096 * 32 = 131,072
scheduled score slots         = useful cells = 16,777,216
scheduled five-path work      = 687.195 GFLOPs

The byte ledger is:

Stage or main term	Logical bytes
preprocess: O/dO reads, D write, dQaccum clear	128.5 MiB
main: K/V loads	64 MiB
main: repeated Q/dO loads	2,048 MiB
main: repeated LSE/D loads	32 MiB
main: logical FP32 dQ update payload	2,048 MiB
main: expanded dK/dV stores	64 MiB
dQ conversion	96 MiB
GQA expanded-read and grouped-write reduction	80 MiB
complete pipeline	4,560.5 MiB

Under the declared convention:

$$I_{\mathrm{logical,pipeline}} =\frac{687.194767\times 10^9\ \mathrm{FLOPs}} {4{,}782{,}030{,}848\ \mathrm{bytes}} =143.704\ \mathrm{FLOPs/byte}.$$

Renderer-safe result: 143.704 FLOPs/byte.

This number is reproducible, but it is not a roofline point measured from DRAM. The denominator includes one four-byte logical atomic payload rather than physical atomic traffic and assumes no cache reuse between CTAs.

For top-left causal attention, the high-memory specialization visits 1,056 tile pairs per query head:

row visits                         = 67,584
useful causal score cells          = 8,390,656
scheduled score slots              = 8,650,752
masked diagonal slots              = 260,096
scheduled work                     = 354.335 GFLOPs
main logical bytes                 = 2,256.5 MiB
complete pipeline logical bytes    = 2,561 MiB
scheduled pipeline intensity       = 131.948 FLOPs/byte
useful pipeline intensity          = 127.981 FLOPs/byte

Causality reduces repeated main-loop work and bytes together, but preprocess, conversion, and grouped reduction remain fixed at this shape. That fixed floor is why the complete-pipeline intensity falls even though the main loop does less work.

0x06.3E Tile Changes Can Move Different Ledgers In Opposite Directions

The historical trait sweep prevents four common but invalid shortcuts.

Smaller B_c: less diagonal waste, more repeated destinations

For d=128, the fallback changes B_c from 128 to 64 and moves V into registers. Under the causal coordinate:

Trait	Row visits/head	Diagonal/tail excess	Pipeline bytes	Scheduled intensity
64x128, shared V	67,584	3.10%	2,561 MiB	131.948 FLOPs/byte
64x64, V in registers	133,120	1.54%	4,625 MiB	71.957 FLOPs/byte

The smaller key tile wastes fewer masked slots, so scheduled matmul work falls from 354.335 to 348.966 GFLOPs. But there are twice as many K/V-column owners, nearly doubling repeated Q/dO/LSE/D reads and logical dQ update payloads. The resulting logical intensity is much lower.

This does not prove the larger tile is faster. The fallback exists because the larger tile may be infeasible under the device shared-memory envelope, and V-in-registers changes compiled register pressure. Feasibility precedes the traffic comparison.

Smaller B_r: more loop iterations can preserve the same bytes

The d=64 high-memory and fallback traits use 128x128 and 64x128 tiles. The fallback has twice as many causal tile-pair iterations per query head (1,056 versus 528), but each iteration contains half as many rows. For the divisible 4,096-token coordinate, scheduled work and every source-logical byte term are identical. Section 0x06.2 still predicts more source-path barrier executions.

Therefore iteration count is neither byte count nor runtime.

Q buffering and V placement are not HBM terms by themselves

The d=256 double- and single-Q-buffer traits use the same 64x64 geometry. Both produce exactly 9,217 MiB and 72.214 scheduled FLOPs/byte in the causal pipeline ledger. The d=192 double/shared-V and single/V-register traits are also identical in logical work and bytes.

Their differences are on-chip lifetimes, overlap, barriers, compiled registers, possible spills, residency, and timing. Any reported HBM difference must come from measured cache/spill behavior, not from relabeling 0x06.2’s shared-memory savings as DRAM savings.

0x06.3F Modeled And Measured Arithmetic Intensity Are Different Coordinates

The source model uses:

$$I_{\mathrm{logical}} =\frac{F_{\mathrm{scheduled\ matmul}}} {B_{\mathrm{source\ logical}}}.$$

A physical DRAM coordinate would use a reconciled measured numerator and measured off-chip denominator:

$$I_{\mathrm{DRAM}} =\frac{F_{\mathrm{measured\ or\ reconciled}}} {B_{\mathrm{DRAM\ read}}+B_{\mathrm{DRAM\ write}}}.$$

Nsight Compute defines roofline arithmetic intensity as work divided by memory traffic, in FLOPs per byte. Its achieved point, memory boundary, compute boundary, and ridge point are useful only when the work and traffic conventions match the selected roofline. A tensor-core mixed-precision numerator, a scalar FP32 numerator, and a shape-derived 10d numerator are not interchangeable.

Likewise, L2 bytes answer a different question from DRAM bytes. A K/V tile can be requested independently by four grouped query heads while some requests hit in L2. The logical ledger remains four requests; the DRAM counter may show less off-chip traffic. That is evidence of hierarchy reuse, not a contradiction.

No bound label follows from I_logical alone. Barrier skew, atomics, scalar pipelines, occupancy, register spills, bank conflicts, and launch tails can limit a kernel before either a tensor-core or DRAM roof is reached.

0x06.3G The Profiler Contract Is Per Stage, Not Only The Main Kernel

backward_profile_manifest.json freezes the experiment before any counters are collected. Every hardware observation is currently null. A valid run must record:

Contract	Required evidence
build identity	exact source commit, local patch hash, CUDA/compiler versions, flags
device identity	GPU name and UUID, compute capability, driver, clocks, power, persistence
workload identity	tensor shapes, values hash, dtype, scale, causal/dropout policy, seed
launch identity	resolved kernel names and launch ordinals for all four stages
timing	warmups and repeated CUDA-event or framework timing outside profiler replay
launch resources	registers/thread, dynamic shared memory, residency, achieved warps
hierarchy traffic	DRAM read/write, L2 traffic, local spill traffic, shared conflicts
pipeline behavior	tensor/scalar utilization, atomics, barrier stalls, source hot spots

The collection sequence is intentionally staged:

1. Time the complete backward range without Nsight Compute.
2. Resolve exact metric names with ncu --query-metrics for this GPU/tool pair.
3. Capture LaunchStats and Occupancy per kernel stage.
4. Capture SpeedOfLight, roofline, compute, and memory sections.
5. Capture scheduler, warp-state, instruction, and source counters separately.
6. Reconcile every measured numerator and denominator before comparing intensity.

Nsight Compute may need multiple replay passes, serializes profiled launches, and restores written memory for kernel replay. That makes profiler duration a diagnostic coordinate, not a substitute for the separately timed application range. The exact metric names are also tool- and architecture-dependent; the manifest records semantic questions and requires --query-metrics instead of hard-coding one architecture’s counter vocabulary.

0x06.3H Owner Alternatives Move The Byte Terms

The active column owner loads K/V once and repeats query state; a query-row owner would load Q/dO/LSE/D once and repeat K/V state. It would also move atomics from dQ to dK/dV. A fixed-tree design would replace atomic updates with partial-gradient stores and a deterministic reduction pass.

Those are different algorithms at the global-memory boundary:

Owner choice	Repeated operand	Global reduction state	Additional cost
active K/V column	Q, dO, LSE, D	FP32 dQ atomics	expanded-head dK/dV sum for GQA
dormant query row	K and V	FP32 dK/dV atomics	combines sequence and group producers
fixed partial tree	chosen operand depends on owner	materialized partial tensors	partial writes, fixed reduction kernels, workspace

The source-logical model can derive all three once their launch and storage contracts are frozen. It cannot choose the fastest one without physical cache, atomic, residency, and timing evidence. The historical source chooses the column path; the row path remains source-present but inactive at the pin.

0x06.3I Executable I/O Evidence And Failure Closure

The backward I/O contract is reproducible with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_backward_io_oracle.py

The checked artifacts are:

backward_io_fixture.json
backward_io_oracle.py
backward_io_report.json
backward_profile_manifest.json
test_backward_io_oracle.py

The oracle sweeps non-causal and top-left-causal work over all 15 historical backward trait rows. It derives visited tile pairs, valid row visits, useful cells, scheduled slots, five-path work, four pipeline stages, three intensity coordinates, and four specialization comparisons. It does not compile or profile FlashAttention.

Failure	Broken contract	Repair
count unique tensor size as kernel traffic	CTA ownership and repeated reads disappear	sum logical requests over the visited owner graph
use H_kv for main K/V loads	active grid over H_q is ignored	count one K/V column owner per query head; measure cache reuse separately
count four payload bytes as physical atomic traffic	read-modify-write, sectors, replay, and cache path are hidden	profile atomics and hierarchy bytes
use useful causal cells as executed tensor work	diagonal tile computation and tails disappear	report useful cells and scheduled slots separately
say smaller tiles always reduce work	repeated owners and destinations can increase	derive R, scheduled slots, and stage bytes for each trait
turn shared-memory savings into HBM savings	on-chip capacity is confused with off-chip transfer	preserve equal logical bytes for identical tile geometry
classify the kernel from modeled FLOPs/byte	physical bytes, scalar pipelines, stalls, and occupancy are absent	collect a reconciled per-stage profile
time only the main kernel	preprocessing, conversion, and GQA reduction vanish	report each stage and the complete backward range
compare Nsight replay duration with normal execution	collection perturbation becomes benchmark data	time outside the profiler and record replay mode separately
import M4.5 split-state traffic	prefill backward is mixed with one/few-query decode	keep decode K/V splitting and exact state merge in M4.5

0x06.4 One Name Does Not Identify One FlashAttention Contract

“FlashAttention” can name a mathematical invariant, a paper pseudocode loop, a CUDA launch, a Triton program, a source-present alternative, or a later API. Those objects are related, but they do not prove the same claims. Treating them as one implementation creates exactly the contradictions that appeared in the original A19 account: FA1 can seem both sequence-serial and sequence-parallel; FA2 can seem to eliminate all warp communication while its source still uses barriers; and a helper present in a header can be mistaken for a dispatched kernel.

The comparison therefore starts with an evidence rule:

paper algorithm
    proves the stated invariant, abstract owner, and paper claim

pinned source tree
    proves which code and alternatives exist at that commit

pinned launch trace
    proves which source-present path is active for an admitted call

later README or release note
    proves later API drift, not historical behavior

synthetic oracle
    proves only its declared algebra or ownership model

This is not administrative bookkeeping. The evidence coordinate is part of the technical claim.

0x06.4A The Chronology Contains More Than FA1 Then FA2

The FA1 paper established exact, IO-aware tiled attention. The FA2 paper then changed three different coordinates: non-matmul state updates, thread-block parallelism, and warp work partitioning. But the paper itself gives an important historical qualification: the query-outer loop swap and sequence parallelism were first suggested and implemented by Phil Tillet in Triton. FA2 did not invent every idea that later appeared inside the FA2 package; it assembled and extended them into a more complete work-partition design.

The relevant artifacts are:

Artifact	What it establishes	What it does not establish
FA1 paper Algorithm 1	normalized online state and a K/V-outer abstract schedule	the v1.0.9 CUDA grid or warp map
paper-attributed Phil Tillet Triton lineage	first query-outer loop swap and sequence-parallel implementation, according to FA2	the exact code later shipped as the FA2 Triton reference
pinned FA1 v1.0.9	optional query-sequence splits, rounded extents, scratch, and selectable backward paths	that the paper pseudocode was copied literally into CUDA
FA2 paper Algorithms 1 and 2	unnormalized numerator, query-row forward owners, split-Q warps, and K/V-column backward owners	the active v2.0.0 specialization or launch helper
pinned FA2 v2.0.0 CUDA	active grids, traits, tails, barriers, atomics, scratch, GQA mapping, and epilogues	current repository or framework behavior
source-present v2.0.0 alternatives	a Triton reference and a dormant CUDA row-owner backward helper exist	that either replaces the active CUDA dispatch
later repository boundary	v2.1 rectangular causal drift and a v2.2 short-query path	what the paper or v2.0.0 did

This chronology resolves one inherited sentence that needs narrowing. The FA2 paper describes its first-version baseline as exposing B * H thread blocks. That is a valid description of the baseline used in the paper’s comparison. It is not a safe universal statement about every pre-FA2 implementation: the pinned FA1 v1.0.9 launch can already split the query sequence to expose more work, at the cost of temporary output state and another reduction path.

Long sequence length also does not change a model’s architectural head count. It often forces a smaller batch under a memory budget, so the available B * H surface can become too small. The runtime problem is reduced concurrent work, not heads disappearing from the architecture.

0x06.4B Paper Pseudocode Is A Proof-Bearing Schedule, Not A CUDA Transcript

All of the compared forward paths compute the same dense row:

$$O_i = \frac{\sum_j \exp(S_{ij}-m_i)V_j} {\sum_j \exp(S_{ij}-m_i)}.$$

They differ in which state coordinate survives, which loop is outer, and which worker owns the final row:

Contract	Persistent row state	Global owner	Commit boundary
FA1 paper Algorithm 1	normalized O_i, maximum m_i, denominator l_i	K/V block is outer in the abstract loops	updated row state crosses the abstract outer iterations
query-outer Triton lineage	online row state	one program owns a query-row block	row owner commits after streaming K/V
FA1 v1.0.9 source	online state plus split-dependent scratch	batch/head/split work unit	optional partial output and reduction state
FA2 paper Algorithm 1	unnormalized numerator u_i, m_i, l_i	one worker owns one query-row block	normalize once, then commit O and LSE
FA2 v2.0.0 CUDA	fragmented FP32 online state	grid (query_tile,batch,query_head)	register fragments use an epilogue before global store

The semantic identity

u_i = l_i * O_i
O_i = u_i / l_i

allows normalized and unnormalized coordinates to represent the same row. It does not imply the same instruction order, register layout, scratch allocation, or bit pattern. Likewise, “query outer” does not uniquely determine the CUDA kernel. Two query-row owners may use different tile sizes, numbers of warps, tail predicates, shared-memory aliases, dropout replay coordinates, and epilogues.

The safe statement is therefore:

paper Algorithm 1:
    proof-bearing abstract schedule

v2.0.0 CUDA forward:
    one version-scoped materialization of that invariant and ownership idea

v2.0.0 Triton reference:
    a separate implementation path with its own feature and correctness scope

The historical Triton file shipped in v2.0.0 explicitly says it starts from Phil Tillet’s implementation, but that file is not a frozen copy of the first credited implementation. It adds and changes features, carries experimental warnings, and differs from the CUDA path. Historical lineage and executable identity must remain separate.

0x06.4C Global Loop Order Trades Output Traffic For Operand Repetition

For a paper-style K/V-outer schedule, the same query-row state is updated across column blocks. For a query-row owner, that row state can remain private while the worker streams every required K/V block. Let T_r and T_c be query- and key-tile counts. A deliberately simplified ownership ledger is:

K/V-outer paper schedule:
    K/V tile requests:        T_c
    row-state update visits:  T_r * T_c

query-row owner schedule:
    row-state final owners:   T_r
    logical K/V requests:     T_r * T_c

This is not a physical DRAM formula. Cache hierarchy and CTA scheduling may serve repeated logical K/V requests from L2. The ledger identifies the cost that moved:

loop swap reduces:
    repeated global row-state commits

loop swap transfers or retains:
    repeated K/V requests across query-row owners
    per-owner register and shared-memory state
    enough query tiles to occupy the machine

Sequence parallelism adds another choice. Exposing query tiles as independent CTAs can rescue a small B * H launch surface, but each additional owner has resource state and a position in a quantized device wave. The improvement is not “more parallelism is always faster”; it is “the launch gains schedulable work while its repeated operands, resources, and tails remain affordable.”

0x06.4D Warp Partition Changes The Producer Set

Define the warp producer set for output row r as

$$\mathcal{P}_O(r)= \{w:\text{warp }w\text{ contributes to row }r\}.$$

The Figure 3 split-K design has

$$|\mathcal{P}_O(r)|=W,$$

so partial states from W warps require an exact maximum-frame merge. A split-Q design partitions query rows, giving

$$|\mathcal{P}_O(r)|=1.$$

That removes the overlapping output-row reduction. It does not prove that the CTA performs no inter-warp communication. The pinned CUDA path still has cooperative K/V staging, asynchronous-copy waits, CTA barriers, lane-group max/sum reductions, and an output epilogue. “No need for communication” in the paper’s split-Q discussion is precise only when scoped to communication for combining overlapping output slices.

This alternative also moves resource pressure. Split-Q duplicates or partitions different fragments, changes the MMA replication layout, and may increase register pressure even while reducing shared-memory reduction traffic. Sections 0x04 and 0x05 derived why the final choice still depends on head width, tile geometry, compiled registers, and device capacity.

0x06.4E Backward Ownership Chooses Which Gradient Needs A Reduction

The gradient equations do not name one launch:

dQ_i += dS_ij @ K_j
dK_j += dS_ij.T @ Q_i
dV_j += P_ij.T @ dO_i

The owner graph does. With K/V-column ownership, one CTA owns column j, so dK_j and dV_j are direct outputs. Every visible column owner contributes to dQ_i:

$$\mathcal{P}_{dQ}(i)= \{j:(i,j)\text{ is a visited tile pair}\}.$$

The active v2.0.0 pipeline therefore accumulates FP32 dQ atomically. A query-row owner reverses the graph: dQ_i becomes direct, while dK_j and dV_j receive contributions from many row owners. Atomics have not vanished; they moved.

A third alternative materializes one partial tensor per producer and reduces those partials through a fixed tree:

$$G = \operatorname{TreeReduce}(G^{(0)},\ldots,G^{(p-1)}).$$

That replaces concurrent atomic ordering with explicit workspace, partial writes, reduction kernels, and tree depth. It can make the reduction topology repeatable, but it is not free and still performs finite-precision additions.

Backward owner	Direct output	Reduced output	Cost transferred
K/V column	dK_j, dV_j	dQ_i	repeated Q/dO/LSE/D and FP32 dQ atomics
query row	dQ_i	dK_j, dV_j	repeated K/V and K/V-gradient atomics
chosen owner + fixed tree	owner’s direct gradients	materialized partial gradients	workspace, extra traffic, and reduction launches

The row-owner CUDA helper is present in the pinned source, but the active run_flash_bwd launch calls the sequence-K column-owner path. Source presence proves a design alternative was encoded; only the launch trace proves dispatch.

0x06.4F Causal Alignment And GQA Show Why Version Pins Matter

Some alternatives preserve the same semantic output and move cost. Others change the API contract itself.

For rectangular causal attention, v2.0.0 uses top-left alignment:

$$\operatorname{visible}_{2.0}(q,k) \iff k\le q.$$

The repository documents bottom-right alignment from v2.1:

$$\operatorname{visible}_{2.1+}(q,k) \iff k\le q+(S_k-S_q).$$

When S_q != S_k, these predicates expose different cells. This is semantic version drift, not a faster materialization of one fixed mask. A current README can establish that the later contract exists; it cannot be projected backward onto the FA2 paper or v2.0.0.

GQA provides a different cost-transfer example. With

$$g=H_q/H_{kv},\qquad h_{kv}=\left\lfloor h_q/g\right\rfloor,$$

forward avoids storing repeated K/V heads. But the pinned forward grid remains over H_q, not H_kv: each query head still owns an output and merely indexes its grouped K/V head. Backward then sums the implicitly expanded dK/dV contributions over each query-head group. Reduced storage does not imply a reduced query-head program count or a free backward path.

0x06.4G The Design Matrix Records Costs Removed, Retained, And Moved

No alternative is evaluated by a single “better” column:

Transition	Reduced	Retained or transferred	Correctness boundary
normalized O to unnormalized u	repeated reciprocal and row normalization	max, exp, denominator sum, numerator frame rescale	same real-arithmetic row; finite order differs
K/V outer to query outer	repeated row-state commits	repeated logical K/V requests and row-owner resources	same online merge invariant
B*H to sequence-parallel surface	low available CTA count	more owners, wave quantization, scratch depending on split	partial states need an exact merge
split-K to split-Q warps	overlapping output-state merge	operand staging, barriers, lane reductions, epilogue	one output-row producer in split-Q
column to row backward owner	one gradient’s atomics	atomics and repeated operands move to another gradient	same gradient equations
atomics to fixed tree	concurrent reduction order	workspace, partial traffic, reduction kernels	fixed topology, not exact real arithmetic
MHA to GQA storage	stored K/V heads	H_q output owners and grouped backward reduction	divisibility and head map
top-left to bottom-right causality	nothing: this is not only an optimization	visible cells change for rectangular calls	versioned API semantics

The choice cannot be ranked until a workload and machine are frozen. A column-owner backward can win when K/V reuse and direct dK/dV ownership outweigh atomic dQ; a row owner can win under a different cache, shape, or contention regime. A fixed tree can be justified by reproducibility or bounded atomic contention despite extra traffic. The paper supplies design reasoning; the pinned source supplies a candidate schedule; only measurement at a declared coordinate supplies a runtime ranking.

0x06.4H Corrected Claims From A19

The original article contains valuable explanatory diagrams and the right high-level motivation, but four statements need stronger contracts here:

Compressed statement	Contract-safe replacement
GEMM means a generative matrix operation	GEMM means general matrix-matrix multiplication
long sequences mean fewer heads	long sequences often force smaller batches; architectural head count is unchanged
FA1 only parallelizes over batch and heads	the FA2 paper’s baseline does; Phil Tillet’s Triton lineage and pinned FA1 source already expose sequence-parallel forms
split-Q removes warp communication	split-Q removes the overlapping output-row reduction; other synchronization and exchange remain

One more correction is methodological: “FlashAttention-2 implementation” is not a unique coordinate. The paper, active CUDA path, Triton reference, dormant helpers, and later release paths are a family. Any implementation claim must name which member it describes.

0x06.4I Executable Contract Resolution And Failure Closure

The paper/source comparison is reproducible with:

python3 -m unittest -v \
  docs/numinens-evidence/m4-4-flashattention2-work-partition/test_paper_contract_oracle.py

The checked artifacts are:

paper_contract_fixture.json
paper_contract_oracle.py
paper_contract_report.json
test_paper_contract_oracle.py

The fixture contains seven evidence coordinates, eight design transitions, and eight deliberately conflated claims. The oracle requires each transition to name an invariant, the cost reduced, the cost retained or transferred, and its exactness boundary. It does not compile either historical extension, prove that a dormant helper dispatches, or rank alternatives by runtime.

def resolve_implementation_claim(claim, evidence):
    require_version_pin(evidence)
    require_authority_coordinate(evidence)  # paper, source, launch, or later drift

    if claim.says_path_is_active:
        require_pinned_launch_trace(evidence)

    if claim.compares_designs:
        require_named_invariant(claim)
        require_removed_cost(claim)
        require_transferred_cost(claim)

    if claim.ranks_runtime:
        require_frozen_workload_hardware_and_measurement(evidence)

Failure	Broken contract	Repair
paper pseudocode is read as the CUDA loop nest	proof schedule and materialization collapse	name paper and source coordinates separately
historical credit is assigned only to FA2	prior Triton lineage disappears	preserve the paper’s Phil Tillet attribution
pinned FA1 split support is ignored	“FA1 only has B*H work” becomes universal	scope that sentence to the paper baseline
source helper is called active	source presence substitutes for launch evidence	trace the selected launcher and template flags
split-Q is called communication-free	retained barriers and reductions disappear	name the removed output-owner conflict only
logical K/V requests are called DRAM traffic	cache service is assumed	preserve logical and measured byte coordinates
GQA program count is changed to H_kv	output ownership is confused with storage	keep the grid over H_q and group only K/V addresses
later causal semantics are backported	a different mask is attributed to v2.0.0	pin the API version and visibility predicate
alternative is ranked without a profile	cost transfer becomes a speed claim	freeze shape, traits, GPU, build, and measured stages

0x06.5 A Throughput Number Is A Numerator-Duration Contract

The source Figure 4 above answers a real empirical question: under the paper’s benchmark coordinate, did FA2’s work partition complete attention forward plus backward faster than the compared implementations? It does not directly measure tensor-core instructions, scalar instructions, DRAM traffic, or the fraction of wall time spent inside one kernel. Those claims require different observations.

Every reported throughput in this section is therefore treated as a pair:

$$R_{reported}=\frac{F_{declared}}{t_{observed}}.$$

F_declared is an accounting numerator chosen by the paper. t_observed is elapsed time. A bar can move because elapsed time moved even though the declared numerator is identical, and two bars with different numerators cannot be compared as if they measured the same operation set.

The evidence classes remain separate:

paper bar or table cell:
    source-reported measurement

FLOP reconstruction, ratio, or implied duration:
    analytically derived from the paper convention and rounded source value

future local run:
    unmeasured until a complete build/workload/timing manifest is filled

0x06.5A Figure 4 Freezes Tokens, Not Batch Size

Section 4.1 of the FA2 paper declares:

GPU:                  one A100 80GB SXM4
sequence length S:    512, 1k, 2k, 4k, 8k, 16k
tokens per batch:     B * S = 16,384
hidden width:         H * d_h = 2,048
head shapes:          (H=32, d_h=64) or (H=16, d_h=128)
mask modes:           noncausal and causal
phase in Figure 4:    forward + backward

Thus batch size falls as sequence length rises:

S	512	1,024	2,048	4,096	8,192	16,384
B=16,384/S	32	16	8	4	2	1

This sweep is designed to expose the exact pressure discussed in 0x04: long sequences produce more query tiles but fewer batch owners. It is not a same-batch sequence-length sweep, and it is not one fixed tensor shape.

The two head configurations also have the same hidden width. Therefore the paper numerator below is identical for d_h=64 and d_h=128 at the same sequence and mask mode. Different TFLOPs/s between those two panels means different elapsed time under one declared amount of model work; it does not mean the wider-head panel was assigned more numerator FLOPs.

0x06.5B Two Forward Matmuls Produce The Paper Numerator

For square attention with

Q, K, V: [B, H, S, d_h]
scores:  [B, H, S, S]

the QK^T contraction contains B H S^2 d_h multiply-adds. Counting one multiply and one add as two FLOPs gives

$$F_{QK}=2BHS^2d_h.$$

The PV contraction has the same shape count:

$$F_{PV}=2BHS^2d_h.$$

The paper prints 4 * S^2 * d_h * H, which is the per-sequence expression. Figure 4 times a batch, so extending that printed formula over B sequences gives the dense forward numerator

$$F_{fwd}=F_{QK}+F_{PV}=4BHS^2d_h.$$

This derivation counts the two matmul-shaped paths. It does not add maxima, exponentials, denominator sums, masking, dropout, casts, memory operations, or launch overhead to the numerator. Those operations still consume the measured duration. The resulting TFLOPs/s is best read as equivalent useful matmul throughput under the paper’s convention, not a hardware instruction count.

For causal attention the paper divides that numerator by two:

$$F_{fwd,causal}^{paper}=2BHS^2d_h.$$

This is a convenient half-dense convention. An exact valid-cell ledger that includes the diagonal would instead use S(S+1)/2 visible pairs:

$$F_{fwd,causal}^{valid} =4BHd_h\frac{S(S+1)}{2} =2BHd_hS(S+1).$$

The valid-cell expression exceeds the paper numerator by 2 B H d_h S, and neither expression is a scheduled tile/instruction count. 0x05.1 owns the separate full/partial/skipped tile geometry.

The paper then uses the two-forward-matmul versus five-backward-matmul argument:

$$F_{bwd}^{paper}=\frac{5}{2}F_{fwd}^{paper},$$

so Figure 4’s combined numerator is

$$F_{fwd+bwd}^{paper} =\left(1+\frac{5}{2}\right)F_{fwd}^{paper} =\frac{7}{2}F_{fwd}^{paper}.$$

The 2.5x is an operation-count convention derived from matmul paths. It is not a claim that backward latency must be exactly 2.5x forward latency.

0x06.5C A Worked Figure 4 Row Recovers Implied Duration, Not Raw Timing

At S=4096, the fixed-token rule gives B=4. Both head configurations have H d_h=2048, so noncausal forward plus backward contains

$$F_{fwd+bwd}^{paper} =\frac{7}{2}(4\cdot4\cdot4096^2\cdot2048) =1{,}924{,}145{,}348{,}608\ \text{FLOPs}.$$

The causal numerator is exactly half under the paper convention:

noncausal combined numerator:  1.924145 TFLOPs
causal combined numerator:     0.962073 TFLOPs

Inverting t=F/R against the rounded FA2 labels in Figure 4 gives:

Panel at S=4096	Source label	Paper numerator	Implied duration
noncausal, d_h=64	171 TFLOPs/s	1.924145 TFLOPs	11.252 ms
noncausal, d_h=128	196 TFLOPs/s	1.924145 TFLOPs	9.817 ms
causal, d_h=64	156 TFLOPs/s	0.962073 TFLOPs	6.167 ms
causal, d_h=128	173 TFLOPs/s	0.962073 TFLOPs	5.561 ms

These durations are derived from rounded chart labels. They are useful for checking arithmetic, but they are not recovered timer samples and must not be given extra precision in a benchmark report.

0x06.5D The 73 Percent, 63 Percent, 65 Percent, And 72 Percent Claims Differ

The paper cites 312 TFLOPs/s as the A100 FP16/BF16 dense tensor-core peak. Dividing several source-reported maxima by that one peak yields superficially similar percentages, but the numerators and measured scopes differ:

Source coordinate	Reported maximum	Peak fraction	What was timed
Figure 4 forward + backward	203 TFLOPs/s	65.1%	attention forward plus backward
Figure 5 forward	227 TFLOPs/s	72.8%	attention forward
Section 4 rounded forward statement	230 TFLOPs/s	73.7%	rounded prose summary, reported as 73%
Figure 6 backward	196 TFLOPs/s	62.8%	attention backward
Table 1 end-to-end model	225 TFLOPs/s/GPU	72.1%	distributed GPT training under a model-FLOP numerator

The paper’s “up to 73%” attention claim belongs to the forward result, not the Figure 4 combined maximum. The 72% MFU claim belongs to the end-to-end model numerator. Neither percentage is an Nsight measurement of tensor-pipe occupancy, and neither says that 72-73% of every cycle issued a tensor-core instruction.

0x06.5E Figure Names Are Not Sufficient Baseline Pins

Figure 4 names PyTorch, FlashAttention, xFormers, a Triton FlashAttention path, and FlashAttention-2. It does not publish a complete environment record for every bar. The paper gives the GPU, shape sweep, and numerator, but the prose does not close all of these fields:

exact dtype assigned to each plotted bar
input distribution and seed
dropout probability
warmup and timed repeat count
reported timing statistic and synchronization protocol
CUDA, PyTorch, compiler, xFormers, and Triton versions
clock and power policy

The pinned v2.0.0 repository contains a generic benchmark helper and a one-shape varlen example. That example uses B=64, S=1024, hidden width 1024, 16 heads, FP16, dropout 0.1, and 30 repeats. It is not the Figure 4 sweep, so its defaults cannot be backfilled into the paper as missing metadata.

This creates a precise reproducibility status:

paper result:
    source-reported and accounting-reconstructable

exact raw benchmark reproduction:
    blocked by open environment and timing metadata

version-scoped local reconstruction:
    possible after a new manifest chooses and records every open field

0x06.5F Table 1 Uses A Different Model-FLOP Numerator

Section 4.2 reports eight-GPU GPT training. Its per-sequence model numerator is

$$F_{model}(S) =6SN+12Ld_{model}S^2,$$

where N is the paper’s parameter count, L is layer count, and d_model is hidden width. The first term is the familiar parameterized-matmul training proxy: approximately 2SN forward work plus 4SN backward work. The second retains sequence-dependent attention work.

Although these are causal GPT models, the paper explicitly does not halve the attention term. It keeps the literature convention for consistency. This means the Table 1 numerator and Figure 4 causal numerator make different choices:

Figure 4 causal attention:
    halve the dense attention numerator

Table 1 causal model training:
    keep the full 12 * L * d_model * S^2 attention term

For global sequence batch G, W GPUs, and observed step duration t_step, the per-GPU reported rate has the form

$$R_{model/GPU} =\frac{G F_{model}(S)}{W t_{step}}.$$

That rate covers the complete training step under the model numerator. It is not an attention-kernel TFLOPs/s result, and communication, optimizer work, non-matmul primitives, and idle time can consume duration without appearing in the numerator.

The paper’s Table 1 reports:

Model and context	Without FlashAttention	FA1	FA2	FA2 / no-Flash	FA2 / FA1
GPT3-1.3B, 2k	142	189	196	1.380x	1.037x
GPT3-1.3B, 8k	72	170	220	3.056x	1.294x
GPT3-2.7B, 2k	149	189	205	1.376x	1.085x
GPT3-2.7B, 8k	80	175	225	2.813x	1.286x

All first three numeric columns are source-reported TFLOPs/s/GPU. The ratio columns are derived from the rounded table cells.

The table exposes an internal reporting wrinkle: 220/72 = 3.056x, while the paper prose says “up to 2.8x” versus no FlashAttention. The 225/80 = 2.813x 2.7B/8k row matches that prose after rounding, but it is not the largest ratio obtainable from the printed cells. We preserve both source coordinates. We do not silently replace the prose with 3.1x, nor suppress the table-derived ratio to make it match the headline.

The maximum model rate does reconcile cleanly with the paper’s rounded MFU:

$$\frac{225}{312}=0.72115385\approx72\%.$$

0x06.5G Repository Documentation Also Has Measurement Generations

At the same v2.0.0 pin, the top-level README advertises 225 TFLOPs/s and 72% model FLOPs utilization, matching the paper. The nested training README still says 189 TFLOPs/s and 60.6% for its older training account. A git tag does not force every documentation page to describe one measurement campaign.

The pinned configs do provide useful implementation coordinates: the base run uses eight GPUs, BF16, and 2k sequences; separate 8k configs change local and global batch. They do not prove that every Table 1 value came from an untouched config at this exact commit. Configuration presence, source-reported result, and executable reproduction remain three different claims.

0x06.5H Local Measurement Needs A Comparison Gate

A new benchmark should begin with the checked-in benchmark_measurement_manifest.json, not with a copied throughput formula:

def reconcile(local_run, paper_contract):
    require_source_commit(local_run)
    require_gpu_driver_cuda_compiler(local_run)
    require_dtype_shape_mask_dropout(local_run)
    require_warmup_repeats_statistic_sync(local_run)
    require_selected_implementation_per_case(local_run)

    paper_flops = paper_numerator(local_run.shape, local_run.phase)
    paper_equivalent_tflops = paper_flops / local_run.elapsed_seconds / 1e12

    return {
        "paper_equivalent_tflops": paper_equivalent_tflops,
        "latency_distribution": local_run.latency_distribution,
        "comparison_status": close_or_list_open_differences(local_run, paper_contract),
    }

Latency remains the primary direct observation. Paper-equivalent TFLOPs/s is a derived view over the same timing samples. Peak memory, OOM status, selected kernel, profiler counters, and numerical error are separate result columns.

The current local manifest deliberately contains zero timing samples and zero profiler reports. Therefore this article makes no local A100 performance claim and no current-backend ranking. Nsight replay duration must not be substituted for application timing, and a TFLOPs/s bar cannot be inverted into DRAM bytes, occupancy, or stall reasons.

0x06.5I Executable Accounting Evidence And Failure Closure

The benchmark-accounting contract is reproduced with:

cd docs/numinens-evidence/m4-4-flashattention2-work-partition
python3 -m unittest -v test_benchmark_accounting_oracle.py

The checked artifacts are:

benchmark_accounting_fixture.json
benchmark_accounting_oracle.py
benchmark_accounting_report.json
benchmark_measurement_manifest.json
test_benchmark_accounting_oracle.py

The oracle reconstructs all Figure 4 FA2 paper numerators and implied durations, checks fixed tokens and hidden width, reconciles the phase-specific peak percentages, derives every Table 1 ratio, preserves the 2.8x versus 3.056x discrepancy, and requires the local measurement coordinate to remain empty. It does not recover raw timer samples, execute CUDA, validate the transcribed source bars against hidden source data, or claim that a current backend reproduces the historical result.

Failure	Broken contract	Repair
compare TFLOPs/s without naming numerator	equal units hide different operation sets	publish phase, FLOP formula, causal rule, and duration
call Figure 4’s 203 the 73% result	combined and forward charts are conflated	bind 73% to Figure 5 forward
call Table 1’s 225 attention-kernel speed	model and attention numerators collapse	preserve the model-FLOP equation and end-to-end scope
use exact triangular cells for a paper ratio	numerator convention changes after measurement	reproduce the paper half-dense numerator, report alternatives separately
treat implied duration as raw timing	rounded bar labels acquire false precision	label it analytically inverted and rounded
infer dtype or timing defaults from a nearby script	source example is projected into the paper	list open metadata and call the run a reconstruction
turn 220/72 into the prose’s 2.8x	internal source discrepancy is hidden	retain 3.056x table ratio and 2.8x headline
divide by 312 and call it hardware utilization	model-equivalent work is confused with issued operations	call it paper peak fraction or MFU under the named numerator
rank a current backend against historical labels	versions, hardware, and dispatch drift disappear	pin both coordinates and measure the same manifest
infer physical traffic from a throughput bar	no byte counter was observed	use the separate profiler contract in 0x06.3G
import short-query decode measurements	prefill/training and decode schedules mix	keep K/V splitting and cache backends in M4.5

0x07 Prefill And Decode Expose Different Parallel Surfaces

For prefill:

Q: [B, H_q, S_q, d_h], S_q >> 1
K: [B, H_kv, S_k, d_h]
V: [B, H_kv, S_k, d_v]

The dense score semantics are [B,H_q,S_q,S_k], although M4.3’s online schedule never materializes that full object in HBM. FA2’s historical forward grid exposes:

$$P_{\mathrm{prefill}}= B H_q\left\lceil\frac{S_q}{B_r}\right\rceil.$$

With B=1, H_q=32, S_q=4096, and B_r=128:

query tiles: 32
forward programs: 1 * 32 * 32 = 1,024

GQA with H_kv=8 does not change that number in this kernel. It changes which K/V head each of the 32 query-head owners addresses.

At one-token decode:

Q:       [B, H_q,  1, d_h]
K_cache: [B, H_kv, T, d_h]
V_cache: [B, H_kv, T, d_v]
score semantics: [B, H_q, 1, T]

The same query-owner grid becomes:

$$P_{\mathrm{decode,no\ split}}=BH_q.$$

For the same B=1, H_q=32, H_kv=8 coordinate:

prefill programs: 1,024
decode programs:      32

Using eight K/V heads as the launch count would make the situation worse and would also omit output heads unless each program explicitly fused four query heads. Program count is not achieved occupancy or latency, but the 32x loss of the query-tile axis is an exact ownership fact.

0x07A Rectangular Causality Is Also A Decode API Boundary

Set S_q=1, S_k=8. Under historical top-left alignment:

delta = 0
the only query row sees key 0

Under bottom-right alignment:

delta = 7
the query row sees keys 0 through 7

The oracle records one allowed key versus eight. A one-token query over a long prefix therefore needs offset-aware causal semantics or a specialized cache API; simply passing causal=True to code with the wrong historical alignment changes the model operation before performance is considered.

The official changelog separates the evolution cleanly:

v2.0: top-left rectangular causal alignment
v2.1: bottom-right rectangular causal alignment
v2.2: specialized inference path for very short Q, including K/V splitting

M4.4 pins v2.0.0 to explain FA2’s historical work partition. It does not use current repository behavior as evidence for that artifact.

The high-level call can remain visually stable:

model(input_ids, past_key_values=cache)

but query length, causal offset, K/V storage, and kernel ownership have changed.

0x08 Decode Boundary: Query Parallelism Collapses

M4.4 stops after proving three boundary facts:

1. S_q=1 removes the query-tile multiplier.
2. GQA keeps H_q output owners even though only H_kv K/V heads are stored.
3. A short-Q/long-K causal call needs the correct positional alignment.

The cached sequence is the remaining long axis. Splitting it can expose more work, but no split then owns the global softmax denominator. That creates the partial-state reduction derived in M4.5, together with:

K/V-axis split selection
exact partial-state merge
contiguous or paged cache addressing
decode-specific GQA indexing and reuse
mask programs and FlashMask interval representation
SDPA/backend eligibility and profiler evidence
decode bandwidth and split-count tradeoffs

M5.7 separately owns the model and serving consequences:

checkpoint conversion and uptraining
persistent K/V bytes
logical broadcast versus materialized repetition
head-group-fused decode kernels
tensor-parallel placement and replication

This boundary is deliberate. Repeating M4.5’s exact merge or M5.7’s cache ledger here would create competing owners for the same mechanism. M4.4’s job is to show exactly why those mechanisms become necessary.

0x09 What FA2 Does Not Solve

FA2 is a major kernel improvement, but it does not erase all attention costs.

It does not remove the dense query-key relation:

prefill still has dense allowed pairs

It does not remove the KV cache:

decode still reads cached K/V

It does not guarantee the best kernel for one-query decode:

query-block parallelism collapses when Q length is 1

It does not make framework dispatch irrelevant:

unsupported masks, dtypes, head dimensions, or GQA layouts can change backend

It does not replace serving-level memory management:

many requests still need KV allocation, paging, batching, and eviction policy

This is the reason M4 points to M5 instead of trying to solve serving inside one kernel article.

0x0A Handoff To M4.6: FA3 And FA4 Hardware Pipelines

FlashAttention-3 is not the subject of this FA2 article, but it confirms the same direction. M4.6 now reconstructs that hardware branch through both FA3 on Hopper and FA4 on Blackwell.

The FA3 work targets Hopper features:

WGMMA
TMA
asynchrony
FP8
incoherent processing to reduce quantization error

The important M4 lesson is not the details of Hopper programming. It is the continuity:

FA1: avoid S/P HBM materialization
FA2: improve work partitioning and utilization
FA3: exploit newer hardware scheduling and precision features

Attention optimization keeps moving down the stack:

math expression
memory schedule
work partition
hardware pipeline
framework dispatch

M4.6 turns that continuity into a quantitative pipeline study: Hopper TMA/WGMMA producer-consumer schedules, FA3 pingpong and register pressure, Blackwell TMEM and 2-CTA MMA, FA4 forward/backward resource cycles, numerical approximations, deterministic reduction, and profiler evidence.

0x0B Handoff To M4.5, M4.6, And M5

M4.4 has stayed inside one attention call and shown the transition from FA2 prefill work to the decode problem.

It answered:

How does FA2 reduce repeated non-matmul normalization work?
What state does the unnormalized accumulator retain?
How does FA2 expose query-block work during prefill?
How do warp ownership and backward ownership differ?
Why does the softmax derivative reduce to D_i = dot(dO_i,O_i)?
Why do K/V-column CTAs directly own dK/dV but atomically reduce dQ?
How do causal tiling and GQA add independent reduction axes?
Which row-owner and fixed-tree alternatives move the reduction elsewhere?
What FP32 and expanded-head scratch does the historical backward path allocate?
Which backward template branches are active in the historical sequence-K main CTA?
How do V, dS, P, dK, and dV legally reuse shared-memory intervals?
Why do V placement and Q buffering move pressure among shared memory, registers, and barriers?
Why does shrinking B_c reduce diagonal waste while increasing repeated query-state and dQ-update traffic?
Why can Q buffering or V placement change on-chip feasibility without changing source-logical HBM bytes?
How must scheduled FLOPs, logical bytes, hierarchy counters, and per-stage time be reconciled before claiming a roofline bound?
Why are FA1 paper pseudocode, pinned FA1 source, FA2 paper pseudocode, and pinned FA2 source different evidence coordinates?
Which query-outer and sequence-parallel ideas does the FA2 paper credit to Phil Tillet's Triton lineage?
Why does source presence not prove that a row-owner backward helper is dispatched?
Which costs move when state coordinates, global loop owners, warp producers, or backward owners change?
Why must rectangular causal alignment and grouped-head behavior be version-scoped?
How do integer tile ranges determine full, partial, skipped, and tail work?
How does a query-head owner address grouped K/V without materializing repeats?
Why does grouped forward indexing imply a backward head-group reduction?
Why does one-token decode collapse the query-tile launch axis?
Where does one-token decode require a new kernel contract?

M5 will move one level outward.

M4.5 now owns the deeper decode/backend branch:

exact K/V split-reduction derivation
FlashMask interval representation
paged/GQA-aware kernel indexing
backend eligibility and profiler evidence
FA2 / Flash-Decoding / FlashMask / FA3 comparison

That extra article prevents the decode and mask implementation details from being compressed into a short FA2 epilogue.

M4.6 owns the hardware-generation branch:

Hopper TMA/WGMMA pipeline ownership
FA3 GEMM-softmax overlap and FP8 layout
Blackwell TMEM and 2-CTA MMA
FA4 forward/backward feeds-and-speeds derivations
architecture-specific compilation, scheduling, and numerical evidence

Serving is not one isolated attention call. It is many requests:

arriving at different times
with different prompt lengths
generating for different lengths
sharing limited GPU memory
competing for KV cache space
batched and rebatched by a scheduler

PagedAttention/vLLM frames the serving pressure directly: KV cache memory is large, grows and shrinks dynamically, and inefficient management wastes memory through fragmentation and duplication.

That is the next module’s problem:

M4: cost and kernel boundary
M5: serving scheduler and KV pressure

0x0C Publish-Gate Checks

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

Gate	Status
FA2 is presented as work partitioning, not a new attention formula	Covered in 0x01 through 0x05
Source-reported speedups are qualified	0x01
Non-matmul work is separated from GEMM-shaped work	0x02
Normalized and unnormalized coordinates are derived from one denominator	0x02A-0x02D, executable oracle
The unnormalized FA2 accumulator and final normalization are explicit	0x02C-0x02E
Removed reciprocals are separated from retained maxima, exponentials, sums, and frame scaling	0x02G, generated operation ledger
Empty/all-masked rows and finite-precision ordering have explicit contracts	0x02F, executable oracle
FA1/FA2 loop and state ownership are made explicit	0x03, 0x04, work-ownership visual
FA2 launch geometry and loop-order alternatives are traced numerically	0x04
Split-K partial states and exact maximum-frame merge are derived	0x05B, executable oracle
Split-Q gives each forward output row one warp producer	0x05C, 0x05E, source-derived row map
Removed output reduction is separated from retained CTA communication	0x05D-0x05G, paired visuals
Warp layouts are version-scoped rather than inferred literally from Figure 3	0x05E, v2.0.0 trait and mask trace
Backward D_i, dQ, dK, and dV ownership is explicit	0x05
Causal classes are derived from half-open integer ranges rather than tile ordinals	0x05.1A-0x05.1C, executable oracle
Semantic partial tiles are separated from the historical source mask-loop band	0x05.1D-0x05.1E, causal-coordinate visual
Rectangular top-left and bottom-right causal alignment are version-scoped	0x05.1F, 0x07A, executable masks
GQA divisibility, contiguous head mapping, and direct K/V indexing are explicit	0x06A-0x06C, pinned C++ and kernel source
Stored K/V heads, query-head programs, and physical traffic are separate ledgers	0x06D, ownership visual
MQA/GQA backward gradients reduce the query-head group axis	0x06E, executable reduction fixture
Packed sequence, causal, and grouped-head coordinates compose without replacing one another	0x06F
D_i = dot(dO_i,O_i) is derived from the softmax Jacobian rather than asserted	0x06.1A, executable backward oracle
Algorithm 2 tensor shape, dtype, lifetime, and owner are explicit	0x06.1B-0x06.1E, source-pipeline visual
The active historical preprocessing, column-owner, conversion, and GQA-reduction stages are separated	0x06.1D, pinned launch and C++ wrapper source
Atomic correctness is separated from floating-point reduction order and bitwise determinism	0x06.1F, CUDA contract, executable failure trace
Causal and grouped producer multiplicities are quantified	0x06.1G-0x06.1I, ownership-topology visual and oracle
Source-present row ownership is not mislabeled as active dispatch	0x06.1I, pinned launch source
Backward scratch and logical atomic-update ledgers have declared scope	0x06.1J-0x06.1K, executable oracle
Active backward template flags are separated from dormant source branches	0x06.2A, pinned main-helper instantiation
Five MMA output geometries and their three CTA replication layouts are traced	0x06.2B, responsive CTA-pipeline visual
Shared-memory capacity follows simultaneous lifetimes and aliases	0x06.2C-0x06.2D, executable CTA oracle
Async-copy fence, wait, and block barrier have distinct contracts	0x06.2E-0x06.2F, CUDA documentation and hazard ledger
Named fragment pressure is not mislabeled as compiled registers	0x06.2G-0x06.2H, specialization matrix and unmeasured boundary
The dK/dV epilogue alias begins only after loop consumers finish	0x06.2I-0x06.2J, source trace and failure closure
Backward scheduled and useful FLOPs are derived from the same visited tile set	0x06.3A, executable I/O oracle
Unique tensor size is separated from source-logical requests and physical hierarchy traffic	0x06.3B-0x06.3C, profiler contract
Preprocess, main, dQ conversion, and grouped-head reduction bytes are accounted separately	0x06.3C-0x06.3D, generated stage ledger
The logical atomic term is labeled as payload rather than physical read-modify-write traffic	0x06.3B, 0x06.3I
Tile-width, row-height, Q-buffering, and V-placement sensitivities do not cross accounting coordinates	0x06.3E, responsive sensitivity visual
Logical and measured arithmetic intensity use explicit, non-interchangeable numerators and denominators	0x06.3F, Nsight Compute measurement boundary
Profiling is staged per kernel and its replay duration is not used as benchmark timing	0x06.3G, unmeasured profile manifest
Owner alternatives identify which operand and gradient state become repeated or reduced	0x06.3H, design-alternative ledger
Paper, pinned source, active launch, and later repository drift are separate evidence coordinates	0x06.4A-0x06.4B, responsive contract-stack visual
Phil Tillet’s credited Triton lineage is preserved without equating it with a later reference file	0x06.4A-0x06.4B, FA2 paper and pinned source
Global loop and sequence-parallel alternatives identify both removed and transferred costs	0x06.4C, contract oracle
Split-Q is scoped to removing the overlapping output-row merge, not all CTA communication	0x06.4D, producer-set derivation
Column-owner, row-owner, atomic, and fixed-tree backward alternatives preserve explicit reduction graphs	0x06.4E, responsive cost-transfer visual
Source-present alternatives are not labeled active without a launch trace	0x06.4E, pinned v2.0.0 launcher
Rectangular causality is treated as semantic version drift and GQA storage is separated from program ownership	0x06.4F, pinned version predicates
A19 terminology, head-count, FA1 parallelism, and warp-communication simplifications are corrected	0x06.4H
Design alternatives are not ranked without a frozen workload, hardware, build, and measurement coordinate	0x06.4G-0x06.4I, executable contract oracle
Figure 4 forward, backward, and combined numerators are derived from the two-versus-five matmul convention	0x06.5B-0x06.5C, executable benchmark-accounting oracle
The paper half-dense causal numerator is separated from exact valid cells and scheduled tile work	0x06.5B, numerator-stack visual
Figure 4/5/6 attention rates and Table 1 model MFU retain different phases and numerators	0x06.5D-0x06.5F, phase reconciliation ledger
Table 1 ratios preserve the printed 3.056x maximum and the prose’s separate 2.8x claim	0x06.5F, checked report
Missing paper environment fields are not backfilled from a nearby benchmark example	0x06.5E, source-completeness contract
Source-reported bars, analytical inversions, local timings, and profiler counters remain separate evidence classes	0x06.5H-0x06.5I, unmeasured manifest and comparison-gate visual
Prefill and decode tensor shapes and program surfaces are not conflated	0x07, prefill/decode boundary visual
Decode collapse is derived over H_q, not incorrectly over H_kv	0x07, executable program-count fixture
Decode details are handed off rather than duplicated	0x08, 0x0B
M4 hands off to M5 without pretending kernels solve serving	0x09, 0x0B
M4.5 receives the full decode, mask, and backend implementation boundary	0x0B
M4.6 receives the Hopper/Blackwell hardware-pipeline boundary	0x0A, 0x0B
Every required formula family is present under one declared accounting coordinate	formula_contract.json, article_audit_report.json
Every local post and fragment target resolves to an existing owner	final publication audit
Every responsive SVG pair has stable dimensions, title, description, semantic alt text, and a pinned digest	visual_provenance_report.json
All four directly reused source figures have source, figure number, digest, rights boundary, and anti-claims	visual provenance contract
No exact long prose paragraph is duplicated from adjacent M4/M5 owners or A19	final publication audit, 240-character threshold
Synthetic, source-reported, version-scoped, and unmeasured claims remain separate	source manifest and final audit

0x0CA Formula Contract

The final audit does not treat “the page contains math” as formula closure. It registers the formula families by the coordinate each one describes:

exact semantic identities:
    online state, common-maximum merge, final normalization, gradients

paper algorithm and source-logical models:
    launch population, storage lifetimes, tile visits, logical bytes

source-reported accounting conventions:
    attention FLOP numerators, model FLOPs, peak fractions

runtime measurements:
    elapsed time, hierarchy bytes, occupancy, stalls, compiled resources

formula_contract.json requires the state, schedule, launch, ownership, causal/GQA, backward, I/O, benchmark, and decode-boundary families. The offline audit checks balanced and non-empty display-math blocks and confirms that every registered expression is present. This is a notation and coverage check, not a second mathematical proof; the executable oracles remain the numerical evidence for their declared fixtures. The publication snapshot contains 126 display-math blocks; that count is a regression floor, not a depth target or a substitute for checking the registered formula families.

0x0CB Source And Evidence Contract

source_manifest.json separates local article owners, primary papers and official documentation, historical implementation pins, the later repository boundary, repository-authored reconstructions, deterministic artifacts, and unmeasured runtime fields. A source URL in the bibliography is not enough: the audit also requires its claim status and role, and requires full commit pins for implementation claims.

Four source figures are directly reused: A19’s annotated crop of paper Algorithm 1 and the paper’s Figures 2, 3, and 4. Their digests, figure numbers, source URLs, commentary purposes, rights boundaries, and anti-claims are recorded separately from the 21 repository-authored responsive visual pairs. visual_coverage_contract.json enumerates that complete semantic surface. They support the paper schedule, abstract block/warp partition contrasts, and the source-reported benchmark coordinate. They do not prove literal CUDA launch geometry, lane maps, current dispatch paths, or local performance results. They are retained for attributed technical commentary; no general reuse license is claimed. This evidence contract does not turn an abstract paper schedule into a literal CUDA transcript.

0x0CC Render And Accessibility Contract

The static visual audit requires every desktop/mobile SVG to provide fixed width, height, and viewBox, role="img", one non-empty <title>, one non-empty <desc>, and aria-labelledby="title desc". The article must give every image non-empty semantic alt text, select the mobile asset at (max-width: 640px), and contain no orphaned reconstruction or missing target. The complete reconstructed surface is 21 responsive SVG pairs; the four source rasters retain their original files and scale through article CSS rather than pretending to have repository-authored mobile reconstructions.

The browser acceptance matrix is separate:

desktop: 1440 x 1000, light and dark
mobile:   390 x  844, light and dark

required:
    all formulas typeset without MathJax errors
    correct responsive source selected
    no broken image
    no document or article horizontal overflow
    wide tables and code blocks remain intentionally scrollable on mobile

Static provenance cannot prove browser layout, and one browser run cannot pin future CSS behavior. Both checks are therefore retained as different gates.

0x0CD Boundary And Duplication Contract

The final boundary audit verifies outbound ownership links to M4.1-M4.3, M4.5-M4.6, and M5.7, plus reciprocal handoffs from M4.3, M4.5, and M4.6. It also resolves local fragment links against the target headings and rejects exact long-paragraph duplication against the adjacent owner corpus and A19 at the declared threshold of 240 normalized characters.

That duplicate check is intentionally narrow. Zero exact overlap does not prove zero conceptual repetition. The article still needs the owner table and handoff sections to explain why a repeated term is context rather than a second competing derivation.

0x0CE Final Publication Closure

The reproducible closure commands are:

cd docs/numinens-evidence/m4-4-flashattention2-work-partition
python3 audit_visual_provenance.py --output visual_provenance_report.json
python3 audit_article.py --output article_audit_report.json --fail-on-mismatch
python3 -m unittest discover -v

The audit can establish formula presence, source and artifact identity, responsive visual metadata, local-link integrity, boundary closure, and exact paragraph non-duplication. It cannot turn the empty benchmark and profiler manifests into measurements. Historical extension build status, local GPU timings, physical hierarchy traffic, compiled registers, spills, stalls, occupancy, and current-backend ranking remain explicitly unmeasured.

0xFF References

• Local source: Exploring the Transformer Series (19) - FlashAttention v2 and its upgrade
• 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 source: M4.3 - IO-Aware Attention and FlashAttention-1
• M4.5 - Decode Attention Kernels, Mask Programs, and Backend Boundaries
• M4.6 - FlashAttention-3/4 Hardware Pipelines
• M5.7 - MQA/GQA and KV-Head Sharing
• Tri Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Dao-AILab, FlashAttention v1.0.9 historical source at commit 6d48e14a
• Dao-AILab, FlashAttention v2.0.0 historical source at commit 4f285b354796fb17df8636485b9a04df3ebbb7dc
• Dao-AILab, FlashAttention v2.0.0 benchmark utility
• Dao-AILab, FlashAttention v2.0.0 one-shape varlen benchmark example
• Dao-AILab, FlashAttention v2.0.0 top-level performance README
• Dao-AILab, FlashAttention v2.0.0 training README and older throughput account
• Dao-AILab, FlashAttention v2.0.0 Triton reference path and its lineage/feature warnings
• Dao-AILab, FlashAttention v2.0.0 Python interface and historical MQA/GQA docstring
• Dao-AILab, FlashAttention v2.0.0 C++ admission, head ratio, and grouped backward reduction
• Dao-AILab, FlashAttention v2.0.0 forward kernel traits
• Dao-AILab, FlashAttention v2.0.0 forward kernel
• Dao-AILab, FlashAttention v2.0.0 backward launch selection and conversion pipeline
• Dao-AILab, FlashAttention v2.0.0 backward column-owner, row-owner, atomic, and epilogue source
• Dao-AILab, FlashAttention v2.0.0 softmax and mask helpers
• Dao-AILab, Current FlashAttention README at commit 00756db for later causal-alignment and short-query inference drift
• Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need
• Joshua Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• NVIDIA, CUDA Runtime Occupancy API
• NVIDIA, Ampere GPU Architecture Tuning Guide: Occupancy
• NVIDIA, CUDA C++ Programming Guide: Compute Capabilities
• NVIDIA, CUDA Programming Guide: Atomics
• NVIDIA, CUDA Programming Model: thread-block scheduling order
• NVIDIA, CUDA Programming Guide: Cooperative Groups and asynchronous data copies
• NVIDIA, CUDA Programming Guide: thread-block synchronization
• NVIDIA, CUTLASS CuTe MMA atom-layout documentation
• NVIDIA, Nsight Compute Profiling Guide
• NVIDIA, NVIDIA A100 Tensor Core GPU Architecture
• Tri Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision