FlashAttention 与 SDPA Backend

先从 reference 语义出发

缩放点积注意力 定义：

dense reference 会逻辑上形成 $S=QK^{\mathsf T}$ 与 probability 矩阵，shape 为 $[B,H_q,T,S]$。FlashAttention 不改变这个数学函数，也不是近似稀疏 attention；它通过 SRAM/shared-memory tiling 和 IO-aware 调度，避免把完整 scores/probabilities 反复写入 HBM。

Online softmax 的三个状态

固定一个 Query row，把 K/V 按列分 tiles。处理完旧 tiles 后维护：

• running max $m$；
• 指数和 $\ell=\sum_j\exp(s_j-m)$；
• 未归一化输出累计量 $o=\sum_j\exp(s_j-m)v_j$。

新 tile 分数为 $s^{new}$，其最大值为 $m_{tile}$：

最后输出：

关键是 $m$ 增大时，旧状态必须乘 $\exp(m-m')$ 重缩放。只更新最大值而不缩放旧 $\ell,o$ 会得到错误分布。Causal mask 可在每个 score tile 内把不可见位置设为 $-\infty$，不需要完整全局 mask 矩阵。

图中省略了 Query tiling、多个 warps/CTAs、shared-memory bank/layout、反向传播重计算、split-KV、序列并行与具体 GPU 指令。o 的 shape 实际为每个 Query row 的 [D] 向量，而不是一个标量。

可运行 online-softmax reference

实验问题：只分 K/V tiles、维护 $m,\ell,o$ 的教学实现，是否与一次性 dense causal softmax 对齐？当前环境下 PyTorch SDPA profiler 实际看到了哪个 backend event？

import math, platform
import torch
import torch.nn.functional as F

torch.manual_seed(71)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")

def dense_attention(q,k,v):
    scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    t,s=q.shape[-2],k.shape[-2]
    mask=torch.triu(torch.full((t,s),-torch.inf,device=q.device),1)
    return torch.softmax(scores+mask,-1)@v

def tiled_online_attention(q,k,v,block_n=3):
    b,h,t,d=q.shape; source=k.shape[-2]
    m=torch.full((b,h,t),-torch.inf,device=q.device,dtype=q.dtype)
    ell=torch.zeros_like(m); acc=torch.zeros_like(q)
    query_pos=torch.arange(t,device=q.device)[:,None]
    for start in range(0,source,block_n):
        end=min(start+block_n,source)
        scores=q@k[...,start:end,:].transpose(-2,-1)/math.sqrt(d)
        key_pos=torch.arange(start,end,device=q.device)[None,:]
        scores=scores.masked_fill(key_pos>query_pos,-torch.inf)
        tile_max=scores.max(-1).values; m_new=torch.maximum(m,tile_max)
        old_scale=torch.exp(m-m_new)
        p=torch.exp(scores-m_new[...,None])
        p=torch.where(torch.isfinite(scores),p,torch.zeros_like(p))
        ell=ell*old_scale+p.sum(-1)
        acc=acc*old_scale[...,None]+p@v[...,start:end,:]
        m=m_new
    return acc/ell[...,None],(m,ell,acc)

def backend_evidence():
    if DEVICE.type!="cuda": return "CPU: math/reference backend"
    from torch.profiler import ProfilerActivity,profile
    q=torch.randn(1,8,512,64,device=DEVICE,dtype=torch.float16)
    with profile(activities=[ProfilerActivity.CPU,ProfilerActivity.CUDA]) as prof:
        F.scaled_dot_product_attention(q,q,q,is_causal=True)
    torch.cuda.synchronize()
    selected=[e.key for e in prof.key_averages() if "flash_attention" in e.key]
    return selected[0] if selected else "no flash_attention event observed"

def main():
    q=torch.randn(1,2,7,4,device=DEVICE,dtype=torch.float32)
    k=torch.randn_like(q); v=torch.randn_like(q)
    dense=dense_attention(q,k,v)
    tiled,(m,ell,acc)=tiled_online_attention(q,k,v,3)
    torch.testing.assert_close(tiled,dense,rtol=1e-5,atol=1e-6)
    assert m.shape==ell.shape==(1,2,7) and acc.shape==q.shape
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={q.dtype} Q=K=V={tuple(q.shape)} tile_n=3")
    print(f"state_m={tuple(m.shape)} state_ell={tuple(ell.shape)} state_acc={tuple(acc.shape)}")
    print(f"online_vs_dense_max_abs_diff={(tiled-dense).abs().max():.3e}")
    print(f"observed_sdpa_backend_event={backend_evidence()}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/flash_attention_online_softmax.py。2026-08-13 在 PyTorch 2.13.0+cu130、RTX 3080 上，online 与 dense 最大绝对差为 1.192e-07；对 [1,8,512,64] FP16 causal SDPA 的 profiler 观察到 aten::_scaled_dot_product_flash_attention。

Profiler 名称只是本环境、本 shape、本版本证据。部署时应记录 GPU capability、PyTorch/CUDA 版本、dtype、shape、mask/dropout，并用 profiler 或 backend controls 复核；不能把本页结果承诺给 CPU、不同 GPU 或任意 head dimension。

显存与 IO 成本

dense reference 的 score/probability 逻辑中间量是 $O(BH_qTS)$ 元素。FlashAttention 仍做 dense attention 所需的主 FLOPs，但让 Q/K/V tiles 在片上存储中复用，并只把必要状态和输出写回 HBM，从而把大量 $T\times S$ 中间矩阵 IO 降掉。峰值 allocated 是否显著下降取决于框架是否真的选择 fused backend、allocator 缓存与测量窗口。

examples/flash_attention_memory.py 用完全相同的 FP16 causal Q/K/V:[1,8,4096,64]，先分别 warmup 强制 SDPBackend.MATH 与 SDPBackend.FLASH_ATTENTION，再证明两条路径输出对齐，最后清空 allocator cache、记录各自 baseline、reset peak 并隔离测量。可直接运行：

source /home/nh/workshop/mylab/bin/activate
python -W error examples/flash_attention_memory.py

2026-08-13 在 PyTorch 2.13.0+cu130、RTX 3080（compute capability 8.6）上的顺序实测为：

该 shape 的 Q/K/V 总 payload 为 12 MiB，输出为 4 MiB；若以 FP16 计算，一张逻辑 score 矩阵 [1,8,4096,4096] 就是 256 MiB。Math 路径的峰值不能简单解释为“恰好保存了一张 score”：具体实现还可能使用更高精度中间量、mask、softmax workspace 和其他临时 tensor。Flash 与 Math 最大绝对差为 0.001953，在先定义的 rtol=atol=3e-3 内通过；低精度 backend 的累计顺序不同，因此这里验证的是容差内语义一致，而非 bitwise 一致。

reserved 增量也不等于本次算子的 workspace：baseline reserved segment 内的空闲块可以直接承载新 allocation，所以本次 Flash 的 peak reserved 增量（2 MiB）小于仍存活输出的 4 MiB。跨实现比较时应优先同时报告 baseline、absolute peak、above-baseline peak 与 live-output payload，不能只摘一个 allocator 数字。

它的收益在长序列、合适 dtype/head dimension、支持的 GPU 与足够工作量下更明显。短 decode、很小 tensor、unsupported mask/dropout 或 backend fallback 时，launch、addressing、KV 读取和其他层可能主导；不能写成“FlashAttention 必然更快”。

与 PagedAttention、KV Cache 的组合

• KV Cache 解决历史 K/V 跨 decode step 持久化；FlashAttention 不负责 cache 生命周期。
• PagedAttention 解决逻辑 token 到非连续物理 blocks 的映射；FlashAttention 解决 attention tile 的 IO/softmax 累计。
• 一个 paged decode kernel 可以边查 block table、边加载 K/V tiles、边执行 online softmax。因此“paged”描述地址管理，“flash”描述 IO-aware attention 算法，两者可同时成立。

失败场景与生产实现差距

• 教学代码只 tile K/V，仍由 PyTorch 张量操作物化 tile scores；它比生产 kernel 慢，不是性能实现。
• 真正 kernel 需要选择 block sizes、warp 数、shared-memory layout、vectorized load、causal 边界与 split-KV 合并策略。
• 低精度下不同 backend 的累加顺序不同，数值不必 bitwise 相同；应先定义可接受容差和下游语义。
• GQA、ragged batch、paged KV、sliding window、ALiBi 等支持矩阵取决于实现版本，API 可接受参数不等于每个 fast path 都支持。

上下游关系

• 缩放点积注意力 是必须保持的 dense reference 语义。
• MHA/MQA/GQA 决定 head mapping；原生 GQA fast path 要避免物理扩展 KV heads。
• GEMM/GEMV 与算术强度 提供 compute/IO 分析框架；FlashAttention 的核心收益来自减少 HBM IO，而非减少 dense attention 的数学依赖。
• Backend Dispatch/Fallback用同一输入对比默认 Flash、强制 math 与不支持 dtype，说明“API 可调用”不等于“命中本页 fast path”。
• Kernel Backend、Triton Blocking 与 Persistent Kernel解释 tile/program/warp、片上资源与 persistent scheduling；这些是实现 FlashAttention kernel 时必须选择、但不属于 attention 数学语义的层次。

参考资料

• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Dao, FlashAttention-2
• PyTorch — scaled_dot_product_attention
• PyTorch — torch.nn.attention backend controls