手写分块注意力：中间矩阵从 128MB 压到 128KB，输出偏差 1.42e-07

手写分块注意力：中间矩阵从 128MB 压到 128KB，输出偏差 1.42e-07

序列长度 4096、单头、float32，标准写法要先算出一张 4096×4096 的打分矩阵，占 128 MB；改成分块计算之后，同时在内存里的中间张量只有 128 KB，缩小 1024 倍，进程峰值内存从 367 MB 降到 45 MB。两者的输出最大偏差 1.42e-07，属于 float32 累加的正常范围。

结论先说：分块换来的是内存，不是速度，也不需要付出精度。我这台 2 核机器上分块反而略快（139 ms 对 95 ms），原因是省掉了两次 128 MB 的大块分配；但真正让 GPU 上快出一个数量级的，是片上 SRAM 复用和算子融合，纯 CPU 上循环写出来的版本拿不到这部分收益。

标准式写法卡在哪

打分矩阵的账很好算：单头 S 是 N×N 个 float32，N=4096 时 64 MB，softmax 之后再来一份 P，又是 64 MB，峰值 128 MB。序列翻倍，这个数字翻四倍。按公式推一下：N=16384 单头 1 GB，32 头一起算是 32 GB；N=32768 就是 128 GB。这就是长上下文先在显存上撞墙的原因，跟算力是否够用是两件事。

这套机制在生产里的成熟实现是 GitHub 上的 flash-attention 仓库，BSD-3-Clause 协议，star 约 2.5 万，最新标签是 4.0.0 的 beta 系列（star 数与标签名取自该仓库公开元数据）。它靠 CUDA 核函数把分块与在线 softmax 融合在一起。我没装它，下面全部用 numpy 手写一版，好处是每一层的误差和内存都能单独量。

基准写法就这么几行：

s=q @ k.T# N x N，先落在内存里s*=1.0/np.sqrt(d)s-=s.max(axis=-1,keepdims=True)p=np.exp(s)p/=p.sum(axis=-1,keepdims=True)out=p @ v

在线 softmax 的三个游标

分块版本的思路是：不算出整行 softmax，而是在遍历 K/V 块的过程中，用三个游标把这行的结果一点点凑出来。行最大值 m、行分母 l、以及还没除分母的分子累加 acc。每读到一个新块，如果块内最大值比旧的大，就把已经累加的 acc 和 l 按 exp(m_旧 − m_新) 整体缩小一次，再并入新块的结果。

acc=np.zeros((n,d),dtype=np.float32)# 分子（未除分母）m=np.full(n,-np.inf,dtype=np.float32)# 每行当前的最大值l=np.zeros(n,dtype=np.float32)# 每行当前的分母scale=1.0/np.sqrt(d)foriinrange(0,n,B):qb=q[i:i+B]forjinrange(0,n,B):kb,vb=k[j:j+B],v[j:j+B]s=(qb @ kb.T)*scale# 只有 B x B，不再有 N x Nm_new=np.maximum(m[i:i+B],s.max(axis=-1))alpha=np.exp(m[i:i+B]-m_new)# 旧的累加要按新最大值修正p=np.exp(s-m_new[:,None])l[i:i+B]=l[i:i+B]*alpha+p.sum(axis=1)acc[i:i+B]=acc[i:i+B]*alpha[:,None]+p @ vb m[i:i+B]=m_new out=acc/l[:,None]

那个 alpha 修正项是唯一容易写错的地方：忘了它，输出不会报错，只是数值悄悄偏掉。我第一轮实现就漏了 acc 的修正，在 N=2048 上对照标准写法，最大偏差从 1.94e-07 涨到 4.75e-01。

实测：内存、误差、块大小

跑法：随机正态的 q/k/v，d=64，单头，float32，每个配置独立进程跑 5 次取中位数，中间张量字节数由数组自己报，峰值内存读进程高水位。

分块那一列的中间张量是常数，因为缓冲区大小只跟块大小有关，跟序列长度无关。朴素那一列按 N² 增长，所以在 512 上只差 16 倍，到 4096 就差 1024 倍。峰值内存更能说明问题：朴素版在 4096 上比 512 上多占 322 MB，分块版只多 5 MB。

块大小本身也有取舍。固定 N=4096 扫一遍：

块太小（32）会被 Python 层的循环开销吃掉，278 ms 里有大半花在几千次小块矩阵乘的调度上；块太大（512）缓冲区回到 2 MB，收益又开始缩水。这台的拐点在 256 附近。精度上五个块大小的偏差都在 1.3e-07 到 1.6e-07 之间，看不出系统性差异。

再确认一下误差到底来自哪。N=1024 加跑一个 float64 参考实现做对照，最大偏差 6.26e-07（朴素 fp32 对 fp64）和 1.83e-07（分块 fp32 对 fp64），分块版本甚至更贴近双精度参考，因为它始终在用修正后的最大值做减法，exp 的输入更小、抵消更少。

N=4096 B=128 d=64 单头 fp32 中间张量 128.0 MB -> 128 KB (1024x) 峰值内存 367 MB -> 45 MB 耗时 139 ms -> 95 ms 最大偏差 1.42e-07

我怎么用它

分块写完，第一件想清楚的事是它解决什么。我的场景是本地跑小模型做长文档处理，显存先于算力不够，所以这版代码值钱的地方是内存曲线拉平，能把可处理的序列长度从几千提到几万。速度上不要有期待，CPU 上循环调度本身就贵，纯 Python 写的那一版没有融合算子，别拿它跟 CUDA 实现比。

用的时候有三点注意。一是别用 Python 逐元素去算 exp，我把小块的 exp 和行求和换成极简循环跑了一次，N=2048 上从 19.7 ms 变成 1485 ms，慢 75 倍。二是这个实现在 N 不能整除块大小时要把最后一块单独处理，我下面的切片写法自带兜底，如果你改成定长缓冲，边界要自己补。三是 q 的块方向可以并行、k 的块方向是顺序依赖的（alpha 修正必须按顺序来），多线程要切就切外层，别切内层。

上手三步

1. 先算 N×N×4 字节，跟可用内存比一比。中间矩阵没超过可用内存的一半就别改，标准写法代码量少、读起来也省事。
2. 块大小从 128 起试。这台 2 核机器上 256 最快，32 因为循环调度太碎反而慢一倍多；块大小对精度没有可见影响。
3. 用一个 N≤1024 的输入跟标准写法对一次最大偏差。量级停在 1e-07 就说明 alpha 修正写对了，涨到 1e-01 就是漏了修正项。