❯ what does this mean "Attention mask + max-reduce (first pass of softmax)."

⏺ Softmax is computed in two passes for numerical stability:

Pass 1 (this test): Find the max across each row
max_j = max_over_j(scores[b, i, j] + mask[i, j])

Pass 2: Subtract max, exponentiate, sum, divide
softmax[b,i,j] = exp(scores[b,i,j] + mask[i,j] - max_j[b,i]) / Σ_j exp(...)

The mask is typically a causal mask (upper triangle = -inf, lower = 0.0) so that future tokens get zeroed out after exponentiation. You add it before taking the max so that masked positions don't affect the max value, which would cause
numerical issues in pass 2.

This test shows pass 1: applying the mask and computing the row-wise max, fused into one kernel.

❯ "Softmax is computed in two passes for numerical stability" where does this usually happen?

⏺ In a transformer's self-attention block, right after computing the attention scores Q @ K^T / sqrt(d) and before multiplying by V:

scores = Q @ K^T / sqrt(d_k) # [B, heads, S, S]
scores = scores + mask # causal mask
attn_weights = softmax(scores) # along last dim
output = attn_weights @ V

The two-pass softmax (sometimes called "online softmax" or "safe softmax") happens inside that softmax call. Every deep learning framework does it — PyTorch's F.softmax, JAX's jax.nn.softmax, etc. If you just computed exp(x) / sum(exp(x))
naively, large values in x would overflow exp() to inf. Subtracting the max first keeps everything in a safe numerical range:

pass 1: m = max(x)
pass 2: softmax = exp(x - m) / Σ exp(x - m)

Flash Attention (Dao et al.) fuses both passes and the matmuls into a single tiled kernel by maintaining running max/sum statistics — that's the "online" trick that avoids materializing the full S x S attention matrix.

❯ why is the mask added instead of multiplied - i thought attention was multiplicative

⏺ The attention mechanism itself is multiplicative (Q times K^T), but the mask is additive. They work in different spaces:

• The mask is applied in log space (before softmax)
• Adding -inf in log space is equivalent to multiplying by 0 in probability space

softmax(scores + (-inf)) = exp(-inf) / Σ... = 0
softmax(scores + 0) = exp(scores) / Σ... = unchanged

So mask = 0 means "keep" and mask = -inf means "zero out". If you multiplied the scores by 0 instead, you'd get softmax(0) = 1/S (uniform) rather than 0 — the masked position would still get attention weight.