企业官网建设流程全解析

把 KV 长序列分成两段，分别算 Attention，再把两个输出平均，通常是错的。两个块的 Softmax 分母不同，它们的归一化输出没有相同权重。分块计算必须保留能够恢复全局归一化的信息。

PyTorch 的 Flash-Decoding 介绍发布于 2023 年 10 月 13 日。它沿 KV 序列维度增加并行度，再单独归约结果，主要面向小 batch、长上下文场景。原文的 GPU 数字属于其配置；下文仅验证合并关系。

用 log-sum-exp 给块结果定权重

每块输出 o 是该块归一化后的加权和，同时保存lse=log(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") m=max(s); w=[math.exp(x-m) for x in s]; z=sum(w) return sum(a*b for a,b in zip(w,v))/z, m+math.log(z) def reduce_blocks(parts): if not parts: raise ValueError("no valid blocks") m=max(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:])] got=reduce_blocks(p); ref=partial(s,v)[0] assert math.isclose(got,ref,rel_tol=1e-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。可以信任的是这一个反例：不带归一化统计量的局部输出，不能随意拼成全局答案。