9 / 11 Kernels (CUDA / Triton / Pallas)

Rescale The Running Max

In FlashAttention's online-softmax tiling over K/V blocks, when a new tile produces a local max m_new that exceeds the running max m_old, the existing accumulator O and running denominator l must be corrected by multiplying both by:

A

[exp(m_new) / exp(m_old)] squared

B

No correction is needed for O, only l needs rescaling

C

exp(m_old - m_new)

D

exp(m_new - m_old)

E

1 / m_new