突破注意力瓶颈！FlatAttention数据流+片上集合通信协同优化，晶圆级系统实现2.9倍吞吐量提升

关键词：晶圆级系统、大语言模型推理、FlatAttention 数据流、片上集合通信、软硬协同设计

随着大语言模型（LLM）推理的序列长度急剧增长以及混合专家（MoE）架构的普及，注意力（Attention）算子已成为制约系统整体吞吐量的核心瓶颈。传统的优化手段多局限于算法改进或硬件带宽升级等单一维度。

图 1：(a) 大语言模型在预填充与解码阶段的浮点运算量占比。数据显示，在 MoE 模型中，注意力机制的运算占比随模型规模显著上升，例如 DeepSeek-v3-671B 的解码阶段占比高达 71%，印证了“注意力瓶颈”的存在。(b) 在英伟达 GH200 GPU 上，基于 FP16 精度评估的 FlashAttention-3（预填充）与 FlashMLA（解码）性能上限。结果表明，当前主流注意力实现的硬件利用率仅为 26% 至 64%，说明 GPU 架构上的注意力优化已接近瓶颈。这两点共同表明，突破 LLM 推理的注意力瓶颈需要全新的平铺（Tile）架构与专用数据流协同设计。

最新研究揭示了一个关键方向：通过在基于平铺（Tile）架构的多处理单元（PE）加速器上，深度协同优化数据流与片上互联的集合通信，可以突破单点优化的极限。

• FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Large Attention-Based Model Inference on Tile-Based Accelerators
• https://arxiv.org/pdf/2604.02110

该研究提出了一种名为 FlatAttention 的全新架构数据流范式。它在计算密集型和访存密集型场景下，分别实现了平均 86% 和 78% 的资源利用率。在基准测试中，其性能平均达到 1.9 倍加速比，全面超越了在英伟达顶级 GH200 GPU 上经过极致优化的 FlashAttention 和 FlashMLA 算子实现。

更进一步的成果体现在晶圆级系统上。在一个集成 64 个加速芯片（单芯片算力 1976 TFLOPS FP8）的晶圆级多裸片系统中，FlatAttention 助力 DeepSeek-v3-671B 模型进行端到端解码推理，实现了比现有最佳商用系统（由 96 张 H800 GPU 组成的集群）高出 2.9 倍 的系统吞吐量，同时将每个输出 Token 的时间（TPOT）降低了 1.4 倍。

本文将解析 FlatAttention 的核心创新机制，探讨其如何通过底层数据流与片上网络（NoC）集合通信的深度协同，重塑未来大模型推理的硬件基础设施。

一、核心痛点：为什么传统的 Attention 优化在多 Tile 架构上失效？

在探讨 FlatAttention 之前，必须理解当前架构面临的根本困境。

大量早期研究致力于优化 Attention 数据流，其核心目标是消除 Softmax 归一化作为独立阶段带来的高昂片外内存访问开销。业界普遍采用“算子融合”方案，将 Attention 的核心操作融合，使中间张量驻留在片上内存中。

然而，以往工作主要聚焦于单 Tile（或 NVIDIA GPU 架构中的单流多处理器 SM）加速器。以 FlashAttention-2 为例，它通过优化的切块（Tiling）和 SM 间工作划分，在 GPU 上实现了领先的融合 Attention 数据流。

算法 1：FlashAttention-2 算法。该算法是 GPU 上经典的注意力数据流实现，通过分块加载 K、V 矩阵至片上缓存，融合矩阵乘与 Softmax 计算以提升数据复用，减少高带宽内存（HBM）访问。但其优化仅限于单个计算单元（SM）内部，缺乏跨单元通信设计。当移植到多 Tile 架构时，各 Tile 独立处理数据，完全丧失了跨 Tile 的数据复用可能，且分块大小受单 Tile 缓存容量限制，无法进一步降低 I/O 复杂度。这一缺陷凸显了在 Tile 架构下设计跨 Tile 协作数据流的必要性，也是 FlatAttention 提出的核心动因之一。

在 Hopper 架构之前的经典 GPU 中，SM 之间的通信必须通过全局内存进行，这严重阻碍了跨 SM 的算子融合，使得片上 SM 间通信成本极高。因此，FlashAttention-2 被迫采用了一种“极度并行”的数据流策略，刻意避免这种高成本通信。

• 直至最新的 Hopper 架构，NVIDIA 才在 GPU 处理集群（GPC）内部引入了高效的 SM 到 SM 通信机制，这为 GPC 内的 SM 之间进行高效的片上集合通信操作创造了条件。
• 尽管近期有研究尝试利用 Hopper 的片上集合通信来扩大 LLM 工作负载的层融合范围，但其主要局限于 GEMM 链和 FFN 层，并未触及计算模式更为复杂的 Attention 算子。

图 2：(a) 基于处理单元阵列（Tile）的架构模板。(b) 依托片上硬件支持的集合通信实现的行式多播，与两种基于软件的集合通信实现（软件树型、软件串行）的对比。(c) 由多个 Tile 式处理单元加速器组成的晶圆级多芯系统，采用二维网状片间互连拓扑。
该图展示了研究的硬件基础。图 2a 的 Tile 架构整合了计算、存储与通信单元，是当前大模型推理加速器的主流设计范式。图 2b 表明，硬件原生支持的集合通信（行式多播）在性能上显著优于软件实现，这为采用硬件通信原语进行优化提供了依据。图 2c 的晶圆级多芯系统则旨在满足千亿参数模型对内存容量和计算规模的需求，以解决单芯片资源受限的问题。这些共同构成了后续优化所面向的硬件场景。

在多 Tile 架构及晶圆级系统中，若数据流设计与底层硬件集合通信缺乏深度协同，简单的算子融合将难以克服由通信延迟和带宽利用率构成的“互联墙”。

二、 核心创新：FlatAttention 的软硬协同设计

为应对上述挑战，本文提出 FlatAttention。这是一种与底层芯片级集合通信原语深度协同设计的数据流，专门用于在基于 Tile 的多处理单元（PE）加速器上，高效执行大语言模型中的多种注意力变体，包括多头注意力（MHA）、分组查询注意力（GQA）和多头潜在注意力（MLA）。

图 3：(a) 主流大语言模型架构概述，以及 (b) 预填充阶段的多头注意力、(c) 自回归解码阶段的多头注意力、(d) 自回归解码阶段的分组查询注意力原理示意图。
该图阐明了不同注意力变体的计算特性。MHA 在预填充和解码阶段的计算形态存在显著差异；GQA 则在解码阶段通过分组共享 KV 缓存，将矩阵-向量乘运算还原为更高效的矩阵乘。这些特性是 FlatAttention 能够泛化至各注意力变体的关键。设计需根据计算形态差异，适配不同的 Tile 分组与分块策略，例如 GQA 解码阶段可采用更长的查询序列分组。图 3a 所示的 MoE 与 MLA 混合架构，也解释了本文选择 DeepSeek-v3 作为端到端验证案例的原因。

图 4：(a) FlatAttention 的参数定义。(b) 详细的 FlatAttention 数据流。(c) 朴素的 FlatAttention 调度方式。(d) 优化后的异步 FlatAttention 调度方式。
该图是 FlatAttention 数据流的核心展示。图 4a 的组参数定义为多 Tile 协同处理奠定了基础。图 4b 拆解了其“加载-多播-计算-归约-存储”的核心流程，体现了跨 Tile 数据复用的设计思想。图 4c 与 4d 的对比则凸显了异步调度的价值：在朴素调度中，数据移动与 Softmax 计算会导致矩阵计算单元闲置；而异步调度通过双任务并发，使矩阵计算单元接近满负载运行。这一优化是 FlatAttention 实现高达 92.3% 硬件利用率的关键，也印证了数据流调度与硬件特性协同的重要性。

2.1 极致的资源利用与分块策略

FlatAttention 在资源调度上追求极致精度。其首要目标是在将 L1 内存使用量严格控制在 384 KiB 预算内的同时，最大化底层矩阵乘法引擎（RedMulE）的利用率，目标设定为超过 95%。

图 11：(a) RedMulE 引擎利用率与分块大小的关系。(b) 异步 FlatAttention 数据流的 L1 本地内存占用与分块大小的关系。
该图为确定最优单 Tile 分块大小提供了量化依据。图 11a 显示，当分块大小为 128 时，矩阵引擎利用率超过 95%，继续增大分块对利用率提升有限。图 11b 显示，该分块大小下的 L1 内存占用未超出 384 KiB 的硬件预算。因此，分块大小 128 是兼顾计算利用率与存储资源的最优解，为后续组规模的扩展提供了基准。

研究团队通过详尽的设计空间探索发现，对于所评估的基于 Tile 的多 PE 配置，最优的分块与组规模缩放策略配置，使得系统在不超出 L1 容量限制的前提下，将 RedMulE 的利用率推至 98% 的峰值。

图 10：FlatAttention 的通用分块和组规模缩放策略。
该图概括了 FlatAttention 的核心调优方法：首先保证单 Tile 内矩阵引擎的高利用率，再根据注意力分数矩阵的形态和 Tile 网络拓扑扩展组规模。这一策略避免了“过度扁平化”问题，为 FlatAttention 泛化到 MHA、GQA、MLA 等不同变体提供了统一的调优框架，使其能灵活调整 Tile 组的行列规模，在最大化跨 Tile 数据复用的同时，避免计算单元利用率下降。

算法 2：基于 Tile 式多处理单元架构的 FlatAttention 算法。
该算法是 FlatAttention 数据流的核心实现。它依托 Tile 架构的片上网络集合通信原语，将多个 Tile 划分为协同工作组。算法中仅对角 Tile 从高带宽内存（HBM）加载数据，并通过组内多播进行分发，结合归约操作完成全局 Softmax 计算，从而大幅降低了 HBM 访问量。该设计不仅突破了单 Tile 的资源限制，其异步调度机制还能让矩阵计算与数据移动、Softmax 计算重叠执行，完美适配 Tile 架构的通信特性。

2.2 全面超越单芯片系统极限

为验证实际效能，研究在多种输入形状下，对预填充阶段的 MHA、解码阶段的 MHA、GQA 以及 MLA 进行了详尽评估。

表 1：系统规格参数。
该表格列出了性能对比的硬件基准。其 32×32 Tile 架构的峰值 FP16 算力与对比基线 NVIDIA GH200 相当，确保了对比的公平性。Tile 内整合了矩阵与向量双引擎，以兼顾注意力计算中的密集矩阵乘和 Softmax 向量操作。384 KiB 的本地内存为片上数据缓冲提供支持，而 2 TB/s 的 HBM 带宽则匹配了计算单元的数据吞吐需求。此配置为 FlatAttention 发挥集合通信原语优势提供了硬件基础。

在与 GH200 上最先进的注意力实现（针对 MHA 和 GQA 的 FlashAttention，以及针对 MLA 的 FlashMLA）的对比中，FlatAttention 展现出显著优势：
* 计算受限场景：实现了平均 86% 的计算资源利用率。
* 访存受限场景：实现了平均 78% 的 HBM 带宽利用率。

如评估数据所示，在绝大多数场景下，FlatAttention 在基于 Tile 的多 PE 加速器上的性能全面超越 GH200 上的优化实现，平均加速比达到 1.9 倍。

图 12：Tile 式多处理单元加速器上 FlatAttention 与英伟达 GH200 上优化注意力数据流的基准对比。对比覆盖预填充阶段（改变头维度和序列长度）与解码阶段（改变推测长度和键值缓存长度）。图中标注 C:x% 表示计算受限核（计算利用率 x%），M:y% 表示内存受限核（HBM 带宽利用率 y%）。

该结果验证了 FlatAttention 在不同注意力变体（MHA、GQA、MLA）及计算场景下的通用优势。无论是在计算受限还是内存受限场景，FlatAttention 均实现了更高的硬件利用率：计算受限场景平均利用率为 86%，内存受限场景 HBM 带宽平均利用率为 78%。相较于 GH200 上的优化实现，FlatAttention 平均带来 1.9 倍的性能提升，且该优势覆盖预填充、解码两个阶段及全部注意力变体。这证明 FlatAttention 的设计能有效适配不同注意力变体的计算特性，具备良好的通用性。

更重要的是，FlatAttention 不仅是一种优化方案，还可作为一种强大的反馈机制，用于指导和反向优化基于 Tile 的多 PE 加速器的架构设计参数，从而推动真正的“软硬协同设计”。

三、 晶圆级系统上的端到端性能：DeepSeek-v3 解码推理实战

单芯片的优异表现是起点，大模型的真正挑战在于分布式集群性能。研究团队进一步在晶圆级多裸片系统上，对 DeepSeek-v3 进行了端到端解码推理性能分析。

3.1 晶圆级系统架构配置

评估系统通过晶圆级中介层，以 8×8 的 Mesh 拓扑结构互联了 64 个基于 Tile 的多 PE 加速器。节点间通过 D2D 链路连接，提供 1 TB/s 的超高带宽与仅 256 ns 的极低延迟。

图 5：(a) Tile 式加速器上用于通用矩阵乘的 SUMMA 数据流。DeepSeek-v3 在 (b) 流水线并行、(c) 全专家并行、(d) 专家-流水线混合并行下的工作负载分布。(e) 晶圆级多芯片系统的执行模式。

图 5 为大模型在晶圆级系统的并行部署提供了设计依据：图 5a 的 SUMMA 数据流为 FlatAttention 之外的 GEMM 计算提供了高效实现，保障了解码阶段其他计算核的性能；图 5b-d 对比了三种并行策略的负载分布特性，为选择 EP32-PP2 混合并行策略提供了参考；图 5e 的执行模式明确了核执行与片间通信的同步方式。基于这些设计，研究实现了 DeepSeek-v3-671B 在 64-Tile 系统上的高效部署，证明了并行策略与数据流优化的协同性。

由于 DeepSeek-v3 推理完全在 FP8 精度下运行，所有加速器均配置为相同的峰值 FP8 吞吐量（未开启稀疏性时为 1976 TFLOPS）以及 4 TB/s 的 HBM 带宽。在 RedMulE 矩阵引擎中，FP8 的峰值吞吐量与 FP16 相匹配。

为达成此算力目标，加速器工作频率提升至 1.9 GHz，与 NVIDIA GH200 处于同一梯队。每个加速器配备 128 GiB 的 HBM 容量，使得包含 MTP 特性的 DeepSeek-v3-671B 庞大解码推理任务能够完整部署于此 64 芯片的晶圆级系统上。

3.2 FlatAttention 对吞吐量与延迟的双重优化

在系统级并行策略上，采用了 EP32–PP2（专家并行度 32，流水线并行度 2）的配置。我们重点观察不同注意力数据流对解码性能的影响。

通过调整每个加速器批处理的用户 Token 请求数，得到了系统级吞吐量与 TPOT（每个输出 Token 的时间）的关系曲线。

图 13：(a) 64 加速器系统上 DeepSeek-v3-671B 解码性能对比，分别采用 FlatAttention 和 FlashMLA 数据流。(b) 单芯片批处理量为 256 时，单个解码层的运行时间拆解。专家并行的影响：(c) 晶圆级系统的性能表现；(d) 批处理量为 256 时的片间通信开销。

图 13 展示了 FlatAttention 在端到端 LLM 推理中的性能提升：
* 图 13a 和 13b 表明，FlatAttention 将注意力计算在解码层中的时间占比从 71% 降至 42%，实现了注意力核 4.5 倍的加速，并最终带来端到端解码 2.1 倍的加速。
* 图 13c 和 13d 揭示了专家并行度的调优规律：在低至中等批处理量下，增加专家并行度能提升吞吐量；但在高批处理量下，片间通信开销会成为瓶颈。

主要结论如下：
* 低批处理规模下：系统的 TPOT 达到最小，但整体吞吐量较低。此时 FlatAttention 与 FlashMLA 性能相当。
* 中高批处理规模下：随着每芯片批处理量的增加，FlatAttention 展现出优异的扩展性，其系统吞吐量最高达到 FlashMLA 的 2.1 倍，并且始终保持更低的 TPOT 延迟。

对运行时间进行微观分解：在满足 50 ms TPOT 约束、且每芯片批处理量为 256 时，注意力计算环节在整个解码时间中占据主导地位。
* 使用 FlashMLA 时，注意力计算耗时占比高达 71%，表明加速器利用率极低。
* 使用 FlatAttention 时，加速器利用率提升至 83%，注意力部分实现 4.5 倍加速，其整体占比降至 42%，最终转化为端到端 2.1 倍的加速比。

3.3 直面通信墙：专家并行与拓扑瓶颈

研究进一步探讨了专家并行度的影响。在完全流水线并行的架构下，若处于低批处理区域，增加批处理量无法提升整体吞吐量，反而会劣化 TPOT，其根本原因在于并非所有专家节点都被激活。仅当批处理量超过 16 时，所有专家被全面唤醒，系统吞吐量才开始爬升，尽管 TPOT 也会相应增加。

数据表明，在低至中等批处理规模范围内，将专家并行度从零提升至满载，能够同时改善系统吞吐量和 TPOT。 然而，在极高的批处理规模下，芯片到芯片的互联通信开销成为了主导解码延迟的关键瓶颈。

特别是考虑到 2D Mesh 互联拓扑中不可避免的多跳传输，盲目增加专家并行度会急剧放大通信开销。 这直接指出了未来的研究方向：探索更优且物理上可实现的晶圆级系统互联拓扑结构。

四、 对标行业标杆：DeepSeek-v3 解码性能的全面胜利

为确立 FlatAttention 晶圆级系统的业界地位，研究人员将其与由商用 GPU 和 NPU 构建的顶尖服务系统进行了对标。下表提炼了原论文中的核心对比数据：

表 2：DeepSeek-v3-671B 解码性能与主流 GPU/NPU 方案对比

该表格直观地验证了 FlatAttention 数据流在晶圆级多芯系统上的实际部署优势。即使本文采用的 64 Tile 加速器在峰值系统算力上低于由 96 张 H800 GPU 组成的对比方案，其最终仍实现了更高的单芯片吞吐量和更低的 TPOT（每次输出令牌时间）。

值得注意的是，即使将晶圆级系统内部（D2D）的链路带宽降低至与 NVLink 相当的水平，得益于优化的 Tiling 策略和组缩放机制，本文方案仍能实现 1.6 倍 的解码吞吐量提升。这一结果的核心驱动力在于 FlatAttention 数据流显著提升了注意力计算核心的硬件利用率。这证明，数据流与硬件架构的协同优化，比单纯追求峰值算力更能有效提升大语言模型推理的实际性能。

注：以上对比数据基于满足 50ms TPOT 约束的操作点得出。表中 Ours1 与 Ours2 均代表搭载 FlatAttention 的晶圆级系统。

深度数据对比分析：
我们与当前性能强劲的商用系统 DS-Prof（由 96 张 H800 组成）进行对比。

1. 全速配置（Ours1）： 配备 1 TB/s D2D 链路的晶圆级系统，实现了高达 2.9 倍 的系统吞吐量提升，同时将 TPOT 延迟从 50.2 ms 显著降低至 35.8 ms。
2. 带宽受限配置（Ours2）： 即使将 D2D 链路带宽从 1 TB/s 降低至 160 GB/s（与 DS-Prof 中的 NVLink 带宽持平），该晶圆级多裸片系统依然在性能上超越了 DS-Prof。
3. 核心优势归因： 需要指出的是，本文晶圆级系统采用的 2D Mesh 互联拓扑网络直径更大，且未对计算与通信重叠（Kernel-communication overlap）进行专门优化。在面临这些不利的物理条件下，系统仍能达成 1.6 倍的解码吞吐量加速，其核心驱动力完全归功于 FlatAttention 数据流带来的注意力计算核心利用率的显著提升。

从整体系统性能来看，尽管本文的 64 节点系统在峰值理论算力上比 96 节点的 DS-Prof 集群低了约 1.5 倍，但凭借 FlatAttention 的深度优化，其最终的系统级吞吐量实现了高达 1.9 倍的超越。

五、 相关工作：架构演进中的坐标系

表 3：相关研究工作对比

该表格通过与现有研究的多个维度对比，凸显了 FlatAttention 方案的创新性与全面性。现有研究要么仅支持单 Tile 优化，要么虽支持多 Tile 但未采用硬件级的集合通信原语（多播/归约）。WaferLLM 虽利用了硬件多播，但其假设模型参数可完全放入片上内存，从而放弃了层融合优化，使其无法适配参数量更大的模型。FlatAttention 是首个同时实现层融合、注意力全变体优化、多 Tile 支持，并深度结合硬件级集合通信原语的方案，填补了大模型推理中数据流与 Tile 架构协同优化的空白。

注：Wafer-scale 假设指模型完全适配在片上内存，无需进行层融合，仅适用于较小模型。

• GPU 架构上的优化尝试： ClusterFusion 利用 Hopper GPU 的集合通信实现了优化的全融合 Attention 数据流，与本文工作有相似之处，但其灵活性欠佳（仅沿单一维度“展平”），且受限于单个 GPU 计算簇的架构规模。
• Tile/NoC 架构的初步探索：
  • Zen-Attention 在 AMD NPU 上开发了利用片上空间规约的平铺框架，但缺乏对层融合方案的详细阐述。
  • COMET 提出了评估片上网络（NoC）集合通信成本的框架，但仅将其视为调度层的开销，未深入到架构与数据流协同设计的层面。
• 硬件加速集合通信的缺失： 绝大多数前沿工作未评估硬件加速集合通信的使用。唯一的例外是 WaferLLM，它利用了 Cerebras WSE-2 的硬件多播支持。然而，WaferLLM 极度依赖模型能完全存入片上内存的假设，一旦模型过大导致参数溢出，其数据流效率将变得低下。

在此之前，业界尚无工作全面应对并解决为 LLM 协同设计优化、包含集合通信的多 Tile 加速器这一挑战，而 FlatAttention 填补了这一空白。

六、 结语与展望

FlatAttention 不仅仅是一项算法优化，它代表了 AI 芯片基础软件栈从“适应硬件”到“定义硬件通信”的范式转变。

通过将数据流与芯片级（Fabric）集合通信原语深度协同，FlatAttention 在 GH200 级别的单节点上证明了 1.9 倍 的性能优势，更在 64 节点的晶圆级系统上，以更少的计算单元，超越了业界顶尖的大规模商用 GPU 集群。

展望未来，随着模型参数规模持续增长，算力瓶颈将愈发下沉至通信与互联层。我们期待，基于 FlatAttention 所提供的数据流反馈机制，下一代晶圆级加速器能够在传统 Mesh 拓扑之外，探索并实现更具物理可行性的终极互联拓扑结构，从而彻底突破阻碍 AGI 发展的“互联墙”。

关注“鲸栖”小程序，掌握最新AI资讯