FlashAttention Tiling: Inspect Online Softmax

FlashAttention Tiling: Inspect Online Softmax: experiment guide

This experiment implements the online-softmax idea used in the FlashAttention forward pass on a deliberately small, single-head fixture. Query, key and value vectors each have two coordinates. Sequence length selects two through eight rows, while Matrix fixture seed selects one of twenty reproducible integer-based fixtures. The Q, K and V arrays appear in the numerical table and full-precision inspector. They are generated examples, not trained projections or a model checkpoint. A seed selects different data; it does not improve attention quality. The implementation uses JavaScript numbers and does not execute a GPU kernel or emulate reduced-precision arithmetic.

Worked details 1

Tile edge sets both the query-block height and key-block width, from one through four. The algorithm visits query blocks in order, then visits all key blocks for each query block. Partial edge tiles are clipped to the actual sequence length. For each allowed query-key pair it computes Score multiplier times the dot product of Q and K divided by the square root of two. Score multiplier changes attention sharpness by scaling these logits; zero makes all allowed logits equal. Causal mask excludes keys whose index exceeds the query index. Turning it off permits every key. The initial causal diagonal always leaves at least one valid key in a completed row, so no final softmax row has an empty denominator.

Worked details 2

For each query row the model retains a running maximum m, a normalizer l and an unnormalized weighted-value vector u. Initially m is negative infinity, l is zero and u is zero. When a tile arrives, the new maximum is the larger of the old maximum and the tile scores. Old l and u are multiplied by exp(old m minus new m), then receive the tile exponentials and their weighted values. The initial rescaling factor is explicitly zero. The normalized output is u divided by l. If every key in a tile row is masked, that row remains unchanged; the code does not subtract negative infinity from negative infinity. These cases appear explicitly in the inspector.

Worked details 3

At the opening settings, sequence length four and tile edge two produce four tile visits. Causal attention evaluates ten query-key pairs. The V rows are (1.5, minus 1.5), (1.5, minus 1), (1.5, minus 0.5) and (1.5, zero). Row zero can attend only to itself, so its final output is exactly (1.5, minus 1.5). The other final second coordinates are approximately minus 1.165119, minus 0.724822 and minus 1.021037. Change tile edge to one: there are sixteen visits, but the final attention result agrees up to floating-point reduction error. A tile edge of three produces clipped edge tiles while preserving the same result. The top-right causal tile contains only masked pairs and leaves both affected rows unchanged.

Worked details 4

Set Score multiplier to zero with the causal mask enabled. Every allowed key receives equal weight, so the outputs are prefix means of V: their second coordinates are minus 1.5, minus 1.25, minus 1 and minus 0.75. Disable the causal mask and every output becomes the mean of all four V rows, namely (1.5, minus 0.75). Change the seed to two to see a different matrix fixture and different weighted outputs. A multiplier of four hundred exercises large logits while the running-maximum subtraction keeps the exponential arguments nonpositive. It does not guarantee identical bit patterns across tile orders or simulate every numerical issue in production kernels.

Worked details 5

Play reveals the tile trace over sixty teaching seconds. Pause freezes the current tile; Next tile advances exactly one visit; Seek moves to an inspected tile while pausing. Restart returns to the initial empty statistics without altering parameters. Reset parameters restores the opening sequence, tile, seed, multiplier and causal mask. Speed selects quarter, normal or four times the teaching rate and changes no arithmetic. Editing any model input recomputes the history and pauses playback. Blank, fractional count and out-of-range inputs produce a clear validation error and disable Play and Next tile. Gray crossed cells are masked, the current tile has an orange border, and the table provides the numerical alternative to every output label.

Worked details 6

The displayed peak rectangular score tile counts at most min(sequence length, tile edge) squared scalars, compared with sequence length squared for a dense score matrix. This is a score-scratch comparison only. It excludes Q, K, V, output, running accumulators and implementation overhead. The browser deliberately retains the entire teaching trace for inspection, so its own total memory is not the working-memory bound of an optimized attention kernel. The lesson demonstrates tiled forward arithmetic and online normalization; it does not implement a backward pass, dropout, multiple heads, SRAM allocation, actual HBM traffic or measured speedup. Numerical agreement with dense attention is the acceptance criterion for this model.