FlashAttention Online Softmax Block-wise Algorithm

• Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
• Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
• Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

💻 Industrial Code Implementation

import numpy as np

def online_softmax_block_simulation(
    q: np.ndarray,   # (S, D)
    k: np.ndarray,   # (S, D)
    v: np.ndarray,   # (S, D)
    block_size: int = 2
) -> np.ndarray:
    """
    FlashAttention 核心 Online Softmax 分块算法的纯 Python 标量模拟。
    展示如何在不保存全量 S*S 矩阵的情况下精确计算 Attention 输出。
    """
    S, D = q.shape
    d_k = D
    
    # 最终输出累加器与统计量跟踪
    O = np.zeros((S, D), dtype=np.float64)
    m = np.full((S, 1), -np.inf)  # 当前已知最大值
    l = np.zeros((S, 1))          # 当前累加分母 sum(exp)
    
    # 外循环：遍历 Key 和 Value 的 Block
    for j_start in range(0, S, block_size):
        j_end = min(j_start + block_size, S)
        K_block = k[j_start:j_end, :]  # (Bc, D)
        V_block = v[j_start:j_end, :]  # (Bc, D)
        
        # 内循环：遍历 Query 的 Block
        for i_start in range(0, S, block_size):
            i_end = min(i_start + block_size, S)
            Q_block = q[i_start:i_end, :]  # (Br, D)
            
            # 1. 计算块内局部点积: (Br, D) @ (D, Bc) -> (Br, Bc)
            S_ij = np.matmul(Q_block, K_block.T) / np.sqrt(d_k)
            
            # 2. 块内当前统计量
            m_prev = m[i_start:i_end, :]
            l_prev = l[i_start:i_end, :]
            O_prev = O[i_start:i_end, :]
            
            # 局部最大值
            m_ij = np.max(S_ij, axis=-1, keepdims=True)
            m_new = np.maximum(m_prev, m_ij)
            
            # 3. 缩放系数
            p_prev = np.exp(m_prev - m_new)
            p_curr = np.exp(S_ij - m_new)
            l_curr = np.sum(p_curr, axis=-1, keepdims=True)
            l_new = l_prev * p_prev + l_curr
            
            # 4. 增量更新输出累加器
            # O_new = (O_prev * (l_prev * p_prev) + p_curr @ V_block) / l_new
            O_new = (O_prev * (l_prev * p_prev) + np.matmul(p_curr, V_block)) / l_new
            
            # 回写全局状态 (实际在 SRAM 中完成)
            m[i_start:i_end, :] = m_new
            l[i_start:i_end, :] = l_new
            O[i_start:i_end, :] = O_new
            
    return O

import numpy as np
S, D = 4, 8
q = np.random.randn(S, D)
k = np.random.randn(S, D)
v = np.random.randn(S, D)
# 1. 经典标准算法
scores = np.matmul(q, k.T) / np.sqrt(D)
probs = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
probs /= np.sum(probs, axis=-1, keepdims=True)
std_out = np.matmul(probs, v)

# 2. Online Softmax 分块算法 (block_size=2)
flash_out = online_softmax_block_simulation(q, k, v, block_size=2)

assert np.allclose(std_out, flash_out, atol=1e-5), "分块结果与标准注意力不一致！"
print("✓ FlashAttention Online Softmax 分块模拟自测通过")