Efficient Attention & FlashAttention: 解释 IO 复杂度分析与 roofline 模型。

Core Concept: 用'访问字节数'衡量算法成本；roofline 用算术强度（FLOPs/字节）判断算力受限还是带宽受限。

数学机理：<strong>roofline 模型</strong>——一个算子的性能上界由两个约束的较小值决定：P ≤ min(P_peak, I × BW)，其中 I=FLOPs/访问字节数（<strong>算术强度</strong>）、P_peak 为峰值算力、BW 为显存带宽。当 I 小于'拐点强度'（P_peak/BW）时，算子受<strong>带宽限制</strong>（memory-bound）；大于时受<strong>算力限制</strong>（compute-bound）。<strong>注意力的算术强度</strong>——算 QKᵀ 需读 Q（L×d）与 K（L×d）、写 S（L×L），FLOPs 为 2L²d；若把 S 写入 HBM 再读回做 softmax，则访问字节数约 O(L²)，故 I≈O(d/L)——<strong>序列越长、算术强度越低</strong>，越 memory-bound。这正是 Flash Attention 的动机：通过分块把 L×L 矩阵保留在 SRAM 中，访问字节数降到 O(Ld + L²d²/M)，大幅提高有效算术强度。<strong>对比其他算子</strong>——(a) <strong>大矩阵乘</strong>（如 FFN 的 L×d×d）：I≈O(d)（较高），通常 compute-bound（可利用张量核心）；(b) <strong>逐元素算子</strong>（如 GELU）：I≈O(1)，严格 memory-bound；(c) <strong>归约算子</strong>（如 LayerNorm）：I≈O(1)，memory-bound；(d) <strong>decode 阶段的注意力</strong>：每步只生成 1 token，需读取整个 KV cache（∝S），I≈O(1/S)——<strong>极低</strong>，故 decode 是典型的 memory-bound，瓶颈在 KV 读取带宽（这解释了 MQA/GQA/MLA 与 KV 量化的价值）。<strong>应用</strong>——roofline 指导优化方向：memory-bound 的算子应减少访存（融合、分块、量化 KV）、compute-bound 的应减少 FLOPs（稀疏、低秩）或用更快的张量核心。

深度剖析与工程权衡：① <strong>'先定位瓶颈再优化'的方法论</strong>——盲目降低 FLOPs（如稀疏化）在 memory-bound 场景无效；应先算算术强度判断瓶颈。这是性能工程的核心素养。② <strong>prefill vs decode 的截然不同</strong>——prefill 是 compute-bound（大矩阵乘、可并行）、decode 是 memory-bound（读 KV、串行）；故两者需要<strong>不同的优化策略</strong>（prefill 用 Flash + 大 batch，decode 用 GQA/MLA/量化 KV + 投机解码）。这解释了为何推理引擎要区分两阶段（见 PD 分离题）。③ <strong>与硬件参数的关系</strong>——A100：算力 312 TFLOPS（BF16）、带宽 2 TB/s → 拐点强度约 156 FLOPs/字节；H100：约 989 TFLOPS、3.35 TB/s → 拐点约 295。可见新硬件的拐点提高，<strong>更多算子落入 memory-bound</strong>，故访存优化愈发重要。④ <strong>与量化/融合的关系</strong>——量化减少'字节数'（提高有效强度）、算子融合减少'中间张量的读写'（也提高强度）；两者都是 memory-bound 优化的手段。⑤ <strong>与 KV cache 压缩的关系</strong>——decode 的瓶颈是 KV 读取，故 MQA/GQA/MLA（减少 KV 元素数）与 KV 量化（减少每元素字节）直接提升 decode 吞吐；这是 roofline 分析的直接推论。⑥ <strong>面试要点</strong>——被问'如何优化一个算子'，应先算<strong>算术强度</strong>判断 memory-bound 还是 compute-bound，再选对应策略；能给出'注意力在长序列与 decode 下都是 memory-bound'与'A100 拐点约 156 FLOPs/字节'这类量化细节，会显著加分。