Efficient Attention & FlashAttention: 解释 decode 阶段的注意力瓶颈与 Flash-Decoding。

Core Concept: decode 每步只算 1 个 query，但需读取全部 KV cache；并行度不足导致 GPU 空闲，Flash-Decoding 用'切分 KV + 两阶段归约'提升并行。

数学机理：<strong>decode 阶段的特性</strong>——每步只生成 <strong>1 个 token</strong>（1 个 query），但需要读取<strong>整个 KV cache</strong>（长度 S）来计算注意力。故：(a) <strong>算术强度极低</strong>——FLOPs 约 O(S·d)、访问字节数约 O(S·d)（读 KV），强度 O(1)，严格 memory-bound；(b) <strong>并行度不足</strong>——朴素实现的并行度只有 batch × heads（如 1×32=32 个并行单元），而 A100 有 108 个 SM，故长序列单请求时<strong>大量 SM 空闲</strong>，GPU 利用率极低。<strong>Flash-Decoding（Dao 等 2023）</strong> 的解法：<strong>沿 KV 序列维切分并行</strong>——把 KV cache 切成多个 chunk，每个 chunk 由一个'线程块'独立计算<strong>部分</strong>注意力（部分 softmax 的分子与分母：partial O 与 partial ℓ，以及该 chunk 的 max）；然后用一个<strong>第二阶段归约 kernel</strong> 把所有 chunk 的部分结果按在线 softmax 的规则合并（重缩放 + 相加），得到最终输出。<strong>为什么能提速</strong>——并行度从 batch×heads 提升到 batch×heads×chunks（可填满 GPU）；且每个 chunk 的读取是连续的（访存友好）。<strong>数值稳定性</strong>——两阶段归约仍用在线 softmax 的规则（每个 chunk 记录自己的 max，归约时按 e^{m_chunk−m_global} 重缩放），故与全量 softmax 等价。<strong>效果</strong>——在长序列（如 16k）单请求的 decode 场景可提速数倍（论文报告约 8 倍）。

深度剖析与工程权衡：① <strong>'并行度不足'是长序列 decode 的核心问题</strong>——即使单步总 FLOPs 很小，只要并行度不够就无法利用 GPU；故'增加并行度'（切分 KV）比'减少 FLOPs'更关键。这是 roofline 之外的另一维度（并行度/occupancy）。② <strong>split-K 的通用范式</strong>——'切分归约维 + 两阶段归约'是矩阵乘中的经典优化（split-K GEMM）；Flash-Decoding 把这一思想用到注意力的 KV 维，说明'注意力优化可借鉴 GEMM 优化'。③ <strong>chunk 数的选择</strong>——chunk 越多并行度越高但归约开销越大；实践中按 SM 数与序列长度动态选择。④ <strong>与投机解码/连续批处理的关系</strong>——这些技术都旨在<strong>提高 batch 内并行度</strong>（把多个请求/多个候选 token 凑成更大的矩阵乘）；Flash-Decoding 则解决'单请求长序列'的并行度问题；三者互补。⑤ <strong>与 prefill 的对比</strong>——prefill 的并行度天然高（L 个 query 并行），故不是瓶颈；decode 才是。这再次说明'两阶段需不同优化'。⑥ <strong>面试要点</strong>——被问'长序列推理为什么慢'，应指出'<strong>decode 是 memory-bound + 并行度只有 batch×heads → SM 空闲</strong>'，并给出'<strong>Flash-Decoding（切分 KV + 两阶段归约）</strong>'的解法；能提到'split-K GEMM 的类比'与'与其他提升并行度技术的互补'是明显加分。