把 KV 长序列分成两段分别算 Attention再把两个输出平均通常是错的。两个块的 Softmax 分母不同它们的归一化输出没有相同权重。分块计算必须保留能够恢复全局归一化的信息。PyTorch 的 Flash-Decoding 介绍发布于 2023 年 10 月 13 日。它沿 KV 序列维度增加并行度再单独归约结果主要面向小 batch、长上下文场景。原文的 GPU 数字属于其配置下文仅验证合并关系。用 log-sum-exp 给块结果定权重每块输出 o 是该块归一化后的加权和同时保存lselog(sum(exp(scores)))。合并时把 lse 再做一次 Softmax所得权重用于组合各块 o。对于没有有效 key 的空块不应产生一个“零输出、正常权重”的假块。各 KV 块保留归一化统计量合并时再定权重。import math def partial(s,v): if not s or len(s)!len(v): raise ValueError(nonempty matched block required) mmax(s); w[math.exp(x-m) for x in s]; zsum(w) return sum(a*b for a,b in zip(w,v))/z, mmath.log(z) def reduce_blocks(parts): if not parts: raise ValueError(no valid blocks) mmax(lse for _,lse in parts) w[math.exp(lse-m) for _,lse in parts] return sum(a*p[0] for a,p in zip(w,parts))/sum(w) s[0.,0.,10.]; v[0.,0.,9.] p[partial(s[:2],v[:2]),partial(s[2:],v[2:])] gotreduce_blocks(p); refpartial(s,v)[0] assert math.isclose(got,ref,rel_tol1e-12) assert abs(sum(o for o,_ in p)/2-ref)1 print(weighted reduction passed; naive average rejected)分块数量是一笔成本账块数增加可以暴露更多独立工作也增加中间结果、启动和末级归约成本。短上下文下这些成本可能比并行收益大。CPU 如果每次生成 token 都新建线程测到的主要可能是线程管理不能借用 GPU 的利用率结论。掩码必须在各块内保持一致因果位置之外的 key 不参与分母分页后的逻辑顺序也不能由物理块编号替代。GQA 的 query 头和 KV 头映射需要单独校验这里的标量 value 例子未覆盖这些维度。为什么这个反例特别容易看出平均法错误前两个分数都为零最后一个分数为十。前一块的局部输出为零后一块只有一个 value因此局部输出为九。直接平均得到四点五却完全忽略最后一个 key 在全局 Softmax 中占据绝大部分权重。归一化后的两个局部输出看起来同样“完成了计算”其分母承载的信息却不同。log-sum-exp 正是为恢复这个权重服务。合并时再减去块间最大值可以避免对很大的 lse 直接求指数。它和保存最大值、指数和、加权和的做法属于两种等价的信息组织不能把一套公式里的归一化输出误当成另一套公式里的未归一化累加器。实际内核可以用一组互不相同的 value 作为诊断输入。若所有 value 都是同一个数不管归一化权重错成什么样输出都可能仍是那个数。分数也应该包含块间幅值差距很大的情景以及不等长尾块。这样才能让错误权重在结果中显现出来。最终归约的浮点顺序还会影响舍入。设备实现可能以树形方式合并参考则顺序处理不应因此要求逐位相同。合理做法是比较数值误差并进一步检查误差是否随块数异常放大。若出现 NaN 或权重和明显失常应先追空块、mask 与重标定而不是先把阈值放宽。还应区分分块计算的工作区和最终输出的生命周期。各块结果可以并行产生但末级归约必须等所需结果都可见归约结束前也不能回收这些统计量。主机提交了全部任务只说明任务进入执行系统不说明结果已经准备好。若采用异步设备需要用依赖或事件表达这一点。分块数改变时工作区大小与事件数量也可能改变应由调度层同步维护而不是只改内核里的循环次数。实际调度应按 batch、上下文、head_dim 与设备选路径。本文没有部署模型、GPU 或多线程 Attention。可以信任的是这一个反例不带归一化统计量的局部输出不能随意拼成全局答案。