cd0d9e1a91
Bug: if l_i > 0 branch in Triton is invalid (compiled as constexpr). Also: p = exp(scores - m_i_new) computed after m_i_new update was using the wrong reference max (should subtract m_ij first, then rescale). Fix: Adapted exactly from prefix_prefill.py's proven-correct pattern: p = exp(scores - m_ij) # probs relative to chunk max l_ij = sum(p) # chunk sum m_i_new = max(m_i, m_ij) # new running max alpha = exp(m_i - m_i_new) # old accumulator rescale beta = exp(m_ij - m_i_new) # new chunk rescale l_i_new = alpha*l_i + beta*l_ij acc = acc*(alpha*l_i/l_i_new) + (p*beta/l_i_new) @ V This is the Flash Attention online softmax tiling algorithm. Same math as CCCL's parallel_reduce with compound accumulators.