Module Ir.Online_softmax

The online-softmax attention rewrite (gh-ocannl-483): the first member of the algebraic-rewrite tier over lowered code (Rewrites), a pattern-directed substitution that CHANGES the computation where the schedule transforms only rearrange it.

The composed attention softmax (q * k^T) * v lowers to a max-reduction and a sum-reduction over the key axis with two elementwise nests between them, then a reduction over the key axis against v. Two of its intermediates are seq^2-shaped and get materialized: the scores are read by both reductions, and the probabilities are read once per value-width iteration of the final reduction. This pass rewrites four shapes, all gated by the config key online_softmax (default off, on in the approximate profile), the third and the fourth by a second key each:

• The (max, sum-of-exp-shifted-by-max) reduction pair over one axis becomes ONE Ir.Low_level.t.Scan_loop per row carrying the running max and the rescaled running sum -- the online-softmax recurrence, Ir.Low_level.t.Scan_loop's founding use case -- writing both trajectories to the original nodes so every downstream reader is unaffected. The chain is recognized only at one precision, the scores', and the carried pair lives at that precision widened to f32 (f64 stays f64): under narrow scores the state does not round per step, which is the one difference from the composed form beyond summation order. The elementwise nests in between stay as the definitions of their nodes. This reassociates the normalizer's summation, which is why the pass is a numerics policy rather than an optimizer decision: results move within rounding.
• A reduction consuming the normalized probabilities (through elementwise nests only) has the loops the probabilities do not index moved innermost, behind one read of the probability cell into a scope local. Bitwise exact: every moved loop indexes the reduction's target, so no per-cell summation order changes; the read multiplicity that forced the probabilities into a seq^2 buffer is gone, and the virtualizer inlines their whole chain into that one read.

• The composed backward of that attention, in a routine that also holds the rewritten forward (the training step of Train.grad_update), gated separately by online_softmax_backward (default off, on in approximate; gh-ocannl-1002). Anchored on the normalizer and the value pass the first two shapes matched -- O and v; dO from the reduction dP += dO * v; the chain from dP through e.grad, l.grad and n.grad = e.grad * e to the score gradient dS the two contractions against the scores' operands q and k read -- it replaces every seq, seq gradient buffer with a per-row D = sum (dO * O) into a minted node and three nests recomputing each pair's p, dp and ds = chain (p * (dp - D)) into scope locals (the recovered elementwise chain -- the mask's where, the scale -- applies to ds only, so a finite mask fill keeps its probabilities and their dV): one over the query rows accumulating q.grad, one over the keys accumulating k.grad, one over the keys accumulating v.grad from p alone, each owning what it writes. The per-cell summation orders of the three gradients are the composed ones; ds reassociates the composed (dP / l + dl) * e, hence the numerics gate. It declines whole -- the backward stays composed -- on anything it cannot prove: no normalizer in the routine, a composed max gradient, another reader of a consumed gradient node, a requested (materialized) intermediate, code the census cannot see in the span, a dead or partial loop, mixed precisions along either chain, or an elementwise step between e / l and the value pass (active dropout).

• With a positive block size B in online_softmax_block (gh-ocannl-1003), the first two shapes give way, per attention, to ONE pass: per query row, a scan over the key blocks of B carrying the row's running max and sum, whose body computes the block's scores into a minted tile (the score contraction and its elementwise chain -- scale, mask), the new max, the rescale factor, the exponentials in place, the rescaled sum, and the row's output numerator in a second tile, rescaled and accumulating p * v -- FlashAttention-2's block recurrence; the last key block writes the output O = U / l. The tiles are carried through memory by the scan body, sound because a write in a scan body is never a virtualization candidate and schedule transforms do not enter a scan body. The scores are read once, so no seq, seq buffer is left at any recompute cap. The value contraction is reassociated too ((sum p * v) / l for sum (p / l) * v), within the same numerics gate. A key count not a multiple of B leaves a guarded tail, padded with -inf scores. The final (m, l) land in their nodes when a later statement reads them, and the elementwise definitions reading them move behind the fold, so the composed or fused backward reads what it reads after the two-pass form. It declines -- the two-pass form applies -- on a value pass that does not read e / l directly (active dropout), a reader of the row state it cannot move, scores that are not a contraction of two reads, or anything else the two-pass contract refuses. The pass emits the scalar form; on a GPU with f32 matrix units the default schedule renders its two contractions as block Tile_mmas, one lane per query row (Schedule.optop.Fold_mma).

What stays as it was in the two-pass form: the score reduction q * k^T keeps its own placement decision -- the recompute cap virtualize_max_inline_reduction decides whether it is replayed at its two read sites (the scan and the hoisted read) or stored once. Recomputing is flash attention's memory trade (no seq, seq buffer at all); storing was faster at every cell measured on cc and Metal (benchmarks/report-gh483-online-softmax.md, seq 128 to 1024), and a training step's backward reads the scores again through cross-routine splicing. (The block fold computes its scores into the tile from q and k and reads nothing else of the score chain, which it removes where nothing else reads it.)

The pass runs at lowering, ahead of the analyses (Rewrites.rewrite.apply, from Assignments.lower, over the raw lowered code), so the traced store and the placements see the rewritten routine and no analysis digest needs the knob: the code itself carries the decision (Code_borne). It is idempotent, as the tier's fixpoint requires: its output has no max-reduction nest left to match -- so no normalizer to anchor a backward or a fold on either -- and a hoisted reduction is no longer a single nest.