打破内存墙：FlatAttention 如何用"本地协作"重构 AI 推理的底层逻辑

参考对象：理查德·费曼（从第一性原理思考）、高德纳（算法美学）、林纳斯·托瓦兹（务实工程哲学）

一、问题的本质：当算力不再是瓶颈

想象你正在经营一家餐厅。

你的厨师（GPU 计算单元）动作飞快，切菜炒菜行云流水。但问题是——食材（数据）在仓库（HBM 显存）里，每次要用都得派人去取。仓库虽然大，但门很窄，一次只能过一个人。

这就是今天 AI 推理面临的困境。

1.1 注意力机制的"二次方诅咒"

Transformer 架构的核心是注意力机制（Attention）。简单来说，就是让每个词都能"看到"句子里的其他所有词，计算它们之间的关联度。

数学上，这相当于计算一个 \(N \times N\) 的矩阵，其中 \(N\) 是序列长度。

二次方增长——这是数学对工程的无情嘲讽。

1.2 DeepSeek-v3 的警示

DeepSeek-v3 是一个 6710 亿参数的 MoE（混合专家）模型。研究者发现：

在解码（decode）阶段，注意力机制占了 71% 的总计算量。

更可怕的是，这不是因为计算太复杂，而是因为数据搬运太慢。

就像你请了一百个厨师，但只有一个采购员在仓库和厨房之间来回跑。厨师们大部分时间都在等食材。

1.3 "内存墙"的物理本质

现代 GPU 的架构是层次化的：

寄存器（Register）    ← 最快，但容量极小（KB 级）
    ↓
共享内存（SRAM）     ← 很快，容量有限（100KB-几MB）
    ↓
HBM（高带宽显存）    ← 容量大，但慢（GB 级）
    ↓
主内存（DRAM）       ← 最慢，容量最大

FlashAttention 的天才之处在于：它通过精巧的分块（tiling）策略，把计算尽可能留在 SRAM 里，减少 HBM 访问。

但这里有个根本性的限制——单个计算单元（SM/Tile）的 SRAM 是有限的。

FlashAttention-2/3 的优化都在单个 SM 内部打转，SM 之间的通信必须通过 HBM，这就像是让厨师们通过仓库传纸条。

二、新架构的曙光：从单芯片到晶圆级

2.1 Tile-based 加速器：把厨房扩大 1000 倍

传统的 GPU 是一个"大厨房"里有很多"小灶台"（Streaming Multiprocessors）。

新一代的 AI 加速器采用了不同的思路：Tile-based 架构。

想象把 1024 个小厨房（Tiles）连成一片，每个小厨房都有自己的厨师（矩阵引擎）、配菜员（向量引擎）和储物柜（本地内存）。这些小厨房之间通过高速通道（片上网络，NoC）直接相连，不需要经过仓库。

这就是晶圆级（wafer-scale）计算的愿景。

2.2 片上网络（NoC）：厨房之间的高速公路

片上网络（Network-on-Chip）是 Tile-based 架构的灵魂。它允许 Tiles 之间直接通信，而不必经过 HBM。

关键在于：NoC 可以支持"集合通信"（Collective Communication）原语。

什么是集合通信？

• 广播（Multicast）：一个人喊话，所有人同时听到
• 归约（Reduction）：所有人同时报数，结果自动汇总

在传统 GPU 上，这些操作需要通过 HBM 中转。但在支持硬件级集合通信的 NoC 上，它们可以直接在片上完成。

三、FlatAttention：重新思考数据流

3.1 核心洞察：I/O 复杂度才是关键

FlatAttention 的作者们提出了一个关键问题：

如果我们有 1024 个 Tiles 通过高速 NoC 相连，为什么还要把数据搬回 HBM？

他们的解决方案是：让注意力计算"摊平"在整个 Tile 网格上。

3.2 算法直觉：分而治之的艺术

传统 FlashAttention 的伪代码大致是：

# FlashAttention-2 简化版
for block_q in Q:           # 遍历查询块
    acc = 0
    for block_k, block_v in zip(K, V):  # 遍历键值对
        s = block_q @ block_k.T         # 计算分数
        m = max(s)                      # 更新最大值
        p = exp(s - m)                  # Softmax
        acc = acc * exp(old_m - m) + p @ block_v  # 累加
    write_to_HBM(acc)       # 写回 HBM

每个 Tile 独立工作，中间结果必须通过 HBM 交换。

FlatAttention 的思路完全不同：

# FlatAttention 简化版
# 1. 对角线 Tiles 从 HBM 加载 K、V
# 2. 通过 NoC 广播到同组所有 Tiles
# 3. 各 Tile 并行计算局部 Softmax
# 4. 通过 NoC 归约全局 Softmax 统计量
# 5. 最终写回 HBM

关键区别：

• 数据分布：K、V 只需从 HBM 读一次，然后通过 NoC 在 Tiles 之间共享
• 计算并行：多个 Tiles 同时计算不同部分的注意力
• 通信本地：Softmax 的归约通过 NoC 完成，不经过 HBM

3.3 技术细节：组（Group）的概念

FlatAttention 引入了一个关键抽象：组（Group）。

整个 Tile 网格被划分为多个组（比如 \(8 \times 8\) 的 Tiles 组成一个组）。

• 组内的 Tiles 通过 NoC 进行高效的集合通信
• 每个组独立计算一部分注意力输出
• 最后把各组的结果合并

这种设计有几个好处：

1. 可扩展性：组的大小可以根据问题规模调整
2. 局部性：组内通信走 NoC，延迟低带宽高
3. 灵活性：支持不同的注意力变体（MHA、GQA、MLA）

3.4 异步流水线：让厨师不停歇

FlatAttention 还采用了异步调度策略：

• 矩阵乘法（在矩阵引擎上执行）
• Softmax 计算（在向量引擎上执行）
• 数据移动（通过 NoC）

这三者可以重叠执行。就像一个厨师在炒菜的同时，配菜员在准备下一道菜，采购员在去仓库的路上。

四、惊人的实验结果

4.1 单芯片性能：碾压 GH200

研究者在 32×32 Tile 配置上测试（峰值算力对标 NVIDIA GH200）：

关键洞察：在 Tile-based 加速器上，FlashAttention 的 HBM 利用率高达 80%，说明它已经被内存带宽饱和了。而 FlatAttention 通过 NoC 通信，大幅降低了 HBM 压力。

4.2 不同注意力变体的表现

注意：在内存受限场景，FlatAttention 实现了 78% 的 HBM 带宽利用率。这不是浪费，而是说明系统已经不再被 HBM 带宽束缚——瓶颈转移到了其他地方。

4.3 端到端 DeepSeek-v3：晶圆级系统的胜利

这是最让人兴奋的部分。

研究者在一个晶圆级多芯片系统上部署了 DeepSeek-v3-671B：

• 64 个 Tile-based 加速器
• 8×8 网格拓扑
• D2D 互联：1 TB/s 带宽，256 ns 延迟
• FP8 精度
• 对比基线：96 张 NVIDIA H800 GPU 集群

结果：

关键洞察：尽管 FlatAttention 系统的峰值算力只有 H800 集群的 2/3，但实际的吞吐量却是 1.9 倍。这就是"内存墙"被打破后的威力。

五、深层启示：架构与算法的协同设计

5.1 软硬协同不是口号

FlatAttention 给我们的最大启示是：算法和架构必须一起设计。

传统的 AI 优化是这样的：

1. 硬件工程师设计芯片
2. 算法工程师在芯片上优化算法
3. 两者之间的鸿沟靠编译器和 runtime 填补

FlatAttention 展示了另一种可能：

1. 算法工程师提出"理想的计算模式"
2. 硬件工程师设计支持这种模式的架构
3. 两者在设计空间探索（DSE）中共同进化

5.2 从"数据搬运"到"计算迁移"

传统思路：数据在哪里，计算就去哪里。

FlatAttention 的思路：计算需要在哪里发生，数据就通过网络搬到哪里。

这就像是：

• 传统餐厅：厨师去仓库取食材
• FlatAttention 餐厅：食材通过传送带直接送到厨师面前

在晶圆级系统中，"传送带"（NoC）的速度和带宽使得这种"计算迁移"策略成为最优解。

5.3 对 AI 基础设施的影响

DeepSeek-v3 的成功（以及 FlatAttention 的优化）预示着几个趋势：

1. 从通用 GPU 到专用加速器：当模型架构相对固定（Transformer + MoE），专用芯片的效率优势会越来越明显

2. 从单芯片到晶圆级：1024 个 Tiles 的协同计算，将重新定义"单机"的概念

3. 从 HBM 为中心到 NoC 为中心：片上通信带宽将超过片外存储带宽，成为系统设计的关键

六、局限与未来

6.1 当前局限

1. 通用性：FlatAttention 针对 Tile-based 架构优化，在现有 GPU 上无法发挥优势

2. 编译器复杂性：异步调度和集合通信需要复杂的编译器支持

3. 问题规模限制：对于小规模序列（<1K），FlatAttention 的开销可能不划算

6.2 未来方向

1. 3D 堆叠内存：将 HBM 直接堆叠在计算芯片上，进一步缩短数据路径

2. 光学互联：芯片间通信使用光而非电，突破电气信号的带宽和延迟限制

3. 自适应数据流：根据输入规模动态选择 FlashAttention 或 FlatAttention 策略

七、结语：工程美学的胜利

FlatAttention 不是某个单一突破的产物，而是系统级优化的典范。

它结合了：

• 算法层面的 I/O 复杂度分析
• 架构层面的 NoC 集合通信
• 系统层面的异步流水线调度
• 应用层面的端到端优化

正如费曼所说：

"What I cannot create, I do not understand."

FlatAttention 的创造者们不仅理解了注意力机制的数学本质，更理解了它在硅片上的物理实现。他们用工程的手段，解决了理论的瓶颈。

当 DeepSeek-v3 在 64 个 Tiles 上流畅运行时，我们看到的不仅是更快的推理速度，更是一种新的计算范式：从单兵作战到群体协作，从数据搬运到计算迁移。

内存墙仍然存在，但 FlatAttention 告诉我们——墙不是用来撞的，是用来绕的。

参考

• Zhang et al., FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Large Attention-Based Model Inference on Tile-Based Accelerators, arXiv:2604.02110, 2026
• Dao et al., FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023
• DeepSeek-AI, DeepSeek-V3 Technical Report, 2024
• Benini & De Micheli, Networks on Chips: A New SoC Paradigm, IEEE Computer, 2002

"The first principle is that you must not fool yourself—and you are the easiest person to fool."
——理查德·费曼

#AI推理 #注意力机制 #DeepSeek #芯片架构 #小凯