关键词: Design in Tiles (DiT)、Tile-Based Many-PE Accelerators、GEMM、Automated Deployment、Network on Chip (NoC)、Collective Primitives

- Design in Tiles: Automating GEMM Deployment on Tile-Based Many-PE Accelerators
- https://arxiv.org/pdf/2512.13638
针对基于瓦片的多PE(处理单元)加速器在部署通用矩阵乘法(GEMM)时面临的编程复杂度高、手动优化难度大,以及现有方案未能有效利用片上网络(NoC)的集体通信原语等问题,本文提出了自动化框架“Design in Tiles(DiT)”,旨在实现此类加速器上GEMM的高效部署。
DiT框架由两大核心组件构成:
* SoftHier:一个基于GVSoC事件模拟器的建模仿真框架,支持_硬件级的NoC集体通信原语,可配置计算瓦片(包含PE、DMA、本地L1暂存器)与分布式高带宽内存(HBM),用于完成功能与性能评估。
* DaCe:一个基于数据中心流图(SDFG)中间表示(IR)的部署优化工具,通过扩展其后端以生成与SoftHier兼容的RISC-V代码,其IR能够精确建模数据迁移与计算依赖关系。_
框架工作流包含数据预加载、SDFG生成与优化、C代码转换、基准测试四个阶段,实现了从算法描述到硬件执行的端到端自动化。

图 4:DiT 工作流程。 流程分为四个阶段实现端到端自动化:(1) 预加载(处理原始数据与布局描述,生成预加载文件以定义HBM数据分布);(2) 生成与优化(根据数据形状与调度策略生成SDFG,并插入瓦片间通信等优化);(3) 转换为C(将优化后的SDFG编译为SoftHier兼容的RISC-V可执行文件);(4) 基准测试(初始化HBM、执行程序、收集性能数据并验证结果正确性)。
DiT引入了一套参数化的瓦片式部署调度抽象,涵盖三大核心模块:
* 分块映射:支持2D/3D分块与集群索引重映射。
* 数据布局:通过数据分割与放置方案优化HBM带宽利用率。
* 数据流:支持通信与计算重叠、SUMMA/systolic等计算原语、以及BSP(整体同步并行)抽象。
实验评估基于类NVIDIA GH200的配置(32×32瓦片阵列、峰值算力1979 TFLOPS@FP8、4TB/s带宽)。结果显示,DiT在多种矩阵形状下,其PE利用率均超越了GH200上专家优化的GEMM库,获得了1.2至2.0倍的性能加速,验证了框架的有效性。本文的贡献还包括对PE工作负载的显式建模IR以及对GEMM性能的详细研究。
本文目录
- 一、引言
- 二、设计概述
- 2.1 SoftHier
- 2.2 以数据为中心的并行编程(DaCe)
- 2.3 工作流概述
- 三、基于瓦片的部署调度抽象
- 3.1 瓦片划分与映射
- 3.2 数据布局
- 3.3 数据流
- 四、评估
- 4.1 性能分解分析
- 五、结论
- 参考文献

关键问题
问题一:小规模瓦片/非规则NoC拓扑下,DiT框架的适配能力与硬件参数依赖阈值
论文中DiT框架的性能优势(1.2-2.0倍加速、更高PE利用率)基于对标NVIDIA GH200的32×32瓦片、特定NoC拓扑配置。若迁移至更小规模瓦片加速器(如8×8/16×16)或非规则NoC拓扑(如环形、树形),其自动化分块映射与集体通信适配能力是否会显著衰减? 核心 优化逻辑对硬件参数的依赖阈值 是什么?
DiT框架的适配能力不会显著衰减。其核心设计通过参数化抽象与灵活的硬件适配机制对冲了硬件配置变化的影响,且对硬件参数的依赖存在明确的可调整阈值。
从框架设计来看:
* 首先,集群索引重映射机制支持将物理瓦片网格(如4×4)重新解释为任意逻辑网格(1×16、2×8等),无论瓦片规模是8×8、16×16还是32×32,均可通过该机制调整任务分块粒度,避免小规模瓦片导致的分块碎片化。
* 其次,NoC集体通信基于掩码寻址系统,支持自定义瓦片分组(行、列、矩形或非规则区域),即使是环形、树形等非规则NoC拓扑,也可通过配置掩码定义通信组,无需修改核心框架逻辑。
在硬件参数依赖阈值方面,核心依赖三大类可配置参数:
1. 瓦片规模(需支持软件管理的L1暂存器与DMA传输);
2. NoC特性(需支持可编程集体通信原语与掩码寻址);
3. 存储带宽(HBM通道数与分布需与数据布局策略匹配)。
只要硬件满足“支持NoC集体通信”与“具备软件管理的内存层级”这两个核心特征,框架即可通过调整分块大小、通信掩码、数据布局等参数进行适配,不存在刚性的依赖阈值。论文中虽以32×32瓦片为评估对象,但SoftHier建模工具的全配置特性与DaCe的多后端适配能力,已验证了其在不同瓦片规模与NoC拓扑下的可移植性。
问题二:极端不规则矩阵维度下,参数化配置的有效性与性能优势持续性
论文强调DiT无需手动开发内核,但部署调度抽象仍需用户指定分块策略、数据流模式等参数。在极端不规则矩阵维度(如M、N、K维度相差3个数量级)场景下,现有参数化配置是否会出现“最优解搜索失效”? 此时其性能能否仍超越针对性手动优化的瓦片加速器内核,而非仅对标GPU专家库?
在极端不规则矩阵维度(M、N、K相差3个数量级)下,DiT框架的3D分块与动态数据流选择机制可有效避免“最优解搜索失效”,且其性能仍有望超越针对性手动优化的瓦片加速器内核。
针对不规则维度的核心解决方案是三维瓦片划分(Split-K)策略:
- 作者明确指出,二维瓦片划分在处理小维度矩阵(如 N=2112)时,会生成低效的小尺寸瓦片(例如 66×…),导致处理单元利用率仅约 50%。
- 三维瓦片划分则通过拆分 K 维度,将多个处理单元绑定到同一输出瓦片的计算任务上,从而在小维度方向放大瓦片尺寸(例如 N 维度从 66 提升至 528)。这更好地适配了矩阵引擎的计算特性,显著提升了利用率。同时,框架的参数化数据布局(拆分策略与放置策略) 能够动态调整矩阵在高带宽内存通道中的分布,避免极端维度导致的内存通道争用。数据流模式(如 SUMMA、Systolic、分层调度等)可根据“计算密集型”或“存储密集型”场景自适应选择(例如计算密集型场景选用 SUMMA,存储密集型场景选用带有限流水线的 Systolic),从而进一步优化极端场景下的性能。
在性能优势方面,论文验证了 DiT 在多样化矩阵形状下,性能达到 GH200 专家手动优化库的 1.2 至 2.0 倍,而 GH200 的优化库本身已代表了行业顶尖的手动调优水平。即便矩阵维度 M、N、K 相差三个数量级,三维瓦片划分的维度拆分灵活性、数据布局的动态适配能力,以及片上网络集体通信的高效数据复用,仍能确保生成的代码在并行度和数据局部性上优于手动开发的内核——手动优化难以兼顾“极端维度分块”与“多硬件资源协同”的复杂性,而 DiT 的端到端自动化流程能够实现全局最优解的搜索,因此其性能优势得以持续保持。
一、引言
在大规模人工智能与高性能计算工作负载的驱动下,计算需求呈爆炸式增长。这一趋势推动了大规模并行加速器的构建,但传统图形处理器架构在扩展过程中暴露出利用率问题。

图 1:CUTLASS 3.9 版本下 NVIDIA A100 与 GH200 两款 GPU 的利用率对比。该图直观展示了两款不同代际 NVIDIA GPU 在运行通用矩阵乘法任务时的处理单元利用率差异。即便使用专家优化的 GEMM 库 CUTLASS 3.9,更新且规模更大的 GH200 在相同 GEMM 矩阵形状下的平均利用率仍低于老款 A100。这一现象揭示了传统 GPU 架构在硬件管理缓存层级上的缺陷,也为后续提出基于瓦片的多 PE 加速器方案提供了现实依据。
如图 1 所示,即使使用专家优化的通用矩阵乘法库 CUTLASS,较新且规模更大的 GH200 在相同 GEMM 形状下的平均利用率仍低于较旧的 A100。其核心问题在于 GPU 的硬件管理缓存层级结构。
尽管近年来多项研究聚焦于提高缓存局部性(指数据在缓存中被重复访问的概率,局部性高可减少内存访问延迟)和避免缓存颠簸(指数据频繁在缓存和内存间交换,导致缓存命中率骤降的现象)的技术,但实现高且可预测的利用率正变得越来越具挑战性。
为应对利用率难题,基于瓦片的多处理单元加速器逐渐兴起。
* 瓦片:指将加速器芯片划分为包含计算单元和本地存储的独立模块。
* 处理单元:执行计算任务的基本硬件单元。
这种架构模板摒弃了统一共享的硬件缓存,由大量计算瓦片构成——每个计算瓦片包含软件管理的一级暂存存储器(L1 SPM,一种高速本地存储器,由软件直接控制数据存取,区别于硬件自动管理的缓存)和处理单元。
- 高带宽可编程片上网络连接大规模处理单元集群,多个高带宽内存芯片则分布在芯片边缘。
- 这种基于软件管理存储层级的设计,让程序员能直接控制整个片上数据流:既可调度数据从高带宽内存传输到 L1 SPM 以实现本地数据复用,也可利用片上网络实现瓦片间高效的数据共享与移动。
该方案不仅消除了传统基于缓存的存储层级中存在的冲突和潜在颠簸问题,还能将更多芯片面积分配给计算集群。典型案例包括 Dojo 系统、SambaNova SN40L 和 Tenstorrent Blackhole。
此外,近期一项研究证实,通过利用基于硬件的片上网络集合通信(指多个计算单元通过片上网络协同完成的通信操作,如广播、归约等),基于瓦片的多 PE 加速器在多头注意力任务上的利用率可高于商用 GPU。
然而,基于瓦片的加速器需要更高程度的显式映射、调度与编排,这带来了可编程性难题:手动管理数千个处理单元的数据移动、同步与内存布局极为复杂,不仅导致编程工作繁琐且层级低下、优化难度大,还会延缓不同硬件配置下的部署效率——最优软件策略与硬件配置深度耦合,手动重写复杂代码以测试新硬件配置几乎不现实。
现有研究从多个角度尝试解决可编程性挑战:
* 一方面,已有工作探索了自上而下的方法,即利用瓦片表示进行调度搜索与优化,但这些研究要么缺乏对 GEMM 部署的详细讨论,要么依赖固定的片上通信模式。
* 另一方面,存在针对特定架构的 GEMM 自下而上分析,但该方案依赖手动编程,且其数据流设计受限于片外内存的缺失。
关键问题在于,当前所有研究均未纳入基于硬件的片上网络集合通信——因此,在基于瓦片的多 PE 加速器上,如何通过有效利用片上网络集合原语实现自动化 GEMM 部署,仍是现有文献中的空白领域。
为填补这一空白,我们构建了“瓦片化设计”框架——一种自动化 GEMM 部署框架。DiT 通过利用硬件集合原语,可在大范围参数化的基于瓦片的加速器上高效部署 GEMM;这种自动化方案既降低了编程难度,又能在不同硬件配置下进行详细的性能分析。本文的贡献如下:
* 一套端到端 GEMM 部署框架,涵盖代码生成、编译、周期精确分析与数值验证。
* 一种中间表示,可显式建模每个处理单元的工作负载,包括数据移动、工作负载映射及瓦片间通信。
* 一种参数化、可配置的高层部署调度抽象,包含瓦片划分与映射策略、分布式高带宽内存通道中的数据布局、集合通信模式,用于生成上述中间表示。
* 一项全面的 GEMM 性能研究,提供具体的 GEMM 部署见解与可移植性讨论。
在评估中,针对与 NVIDIA GH200 性能相当的大规模加速配置,DiT 实现了比 GH200 专家优化 GEMM 库更高的处理单元利用率,在多种矩阵形状下均能达到 1.2 至 2.0 倍的加速比。
二、设计概述
本文提出的框架由两部分组成:可执行模型与部署工具。
2.1 SoftHier

图 2:SoftHier 架构模板。该模板展示基于 GVSoC 事件模拟器的 SoftHier 框架核心结构,每个计算瓦片集成处理单元(标量、向量、矩阵引擎)、直接内存访问控制器和本地 L1 暂存器,片上网络连接所有瓦片,高带宽内存芯片分布在芯片边缘并通过内存控制器连接片上网络。模板还支持基于掩码的片上网络集体通信原语,可通过配置文件实例化特定加速器设计,满足功能与性能评估需求。
SoftHier 是面向基于瓦片的多 PE 加速器的建模与仿真框架,构建于 GVSoC 事件驱动模拟器之上,同时支持功能验证与性能评估。 其模型通过开源寄存器传输级的周期精确仿真进行校准,且可通过架构配置文件完全配置,允许用户实例化特定的加速器设计。
在 SoftHier 中,每个计算瓦片包含处理单元(支持标量、向量、矩阵引擎)、直接内存访问控制器与本地 L1 暂存器;片上网络连接所有计算瓦片;片外高带宽内存沿网格边缘分布,并通过内存控制器与片上网络相连。
SoftHier 的核心特性是支持硬件级集合通信原语[10],它通过灵活的基于掩码的寻址系统实现广播与归约操作。该系统允许单个计算瓦片向自定义的计算瓦片组(如行、列或矩形区域)广播数据,或向单个目标瓦片归约数据。瓦片组(Tile group)被定义为满足以下坐标匹配规则的瓦片集合:
其中:
* $T_{i,j}$ 代表位于第 $i$ 行、第 $j$ 列的瓦片;
* 数据包头部携带选择器坐标 $(s_r, s_c)$ 与掩码 $(m_r, m_c)$;
* 若某个瓦片 $T_{i,j}$ 同时满足 $(i & m_r) = s_r$(行坐标与行掩码按位与后等于行选择器)和 $(j & m_c) = s_c$(列坐标与列掩码按位与后等于列选择器),则该瓦片属于当前瓦片组。
2.2 以数据为中心的并行编程(DaCe)
DaCe[3]是 DiT 框架中的部署与优化工具,其核心是基于以数据为中心的图状结构——有状态数据流多图(SDFG)的中间表示(IR)。
我们利用 SDFG 来表示并转换面向 SoftHier 的应用数据流。该 IR 能够捕获瓦片内的所有数据移动,包括 HBM 与本地内存间的传输,以及计算瓦片本地内存间的通信。
* DaCe 能为多核 CPU、GPU、现场可编程门阵列(FPGA)生成高效代码[26],支持异构架构的性能可移植性;
* 我们对 DaCe 的后端进行了扩展,新增了对 SoftHier 的支持,可将 SDFG IR 下转为 C 代码,该代码可进一步编译为 RISC-V 指令集的可执行文件,在 SoftHier 上运行。

图 3:面向 SoftHier 的 SDFG 数据迁移与计算原语。图中呈现 DaCe 框架的 SDFG 中间表示(IR)核心操作:(a) HBM 到 L1 数据传输(瓦片按索引并行获取数据)、(b) MMAD 计算任务块(用 L1 数据执行矩阵乘加)、(c)(d) 瓦片间收发(通过流数组实现,类似异步 MPI)。圆圈表示数据访问节点,八边形表示计算任务块,边缘捕获数据依赖,完整建模了瓦片内/间的数据流动与计算逻辑。
图 3 展示了 SDFG IR 在 SoftHier 上的三类基本操作:瓦片级数据移动用数据访问节点(圆形节点)间的路径表示;计算用任务块(八边形节点,被视为“黑盒”计算单元,其内部实现细节对 IR 透明)表示;数据依赖关系则通过边(连接节点的线段)捕获。具体而言:
* 图 3a(HBM 到 L1 传输):每个计算瓦片根据自身的瓦片索引(tile_I、tile_J),将 HBM 中对应的数据分段传输到本地内存,这些传输操作在并行区域内同步执行。
* 图 3b(MMAD 任务块):计算瓦片利用 L1 SPM 中存储的数据,执行矩阵乘加(MMAD,GEMM 的基本计算操作,即矩阵乘法后叠加原有结果)任务块。
* 图 3c(PE 到 PE 发送)与图 3d(PE 到 PE 接收):瓦片间的数据移动用 SDFG 流数组(Stream Array,为每个计算瓦片实例化的专用数据传输结构,如对应 tile_I、tile_J 的流数组)的路径表示。
这些流数组共同构成了分布式 L1 SPM 中
local_A数据的全局视图:从数据访问节点到流节点(虚线边框圆形)的路径代表发送(send)操作,从流节点到数据访问节点的路径代表接收(recv)操作,其通信模式类似异步消息传递接口(MPI)通信[7](即发送方无需等待接收方确认即可继续执行后续操作)。
2.3 工作流概述

图 4:DiT 工作流程。流程分四阶段实现端到端自动化:(1) 预加载(处理原始数据与布局描述生成预加载文件,定义 HBM 数据分布);(2) 生成优化(根据数据形状与调度生成 SDFG,插入瓦片间通信等优化);(3) 转换为 C(将优化 SDFG 编译为 SoftHier 兼容的 RISC-V 可执行文件);(4) 基准测试(初始化 HBM、执行程序、收集性能数据并验证结果正确性)。
本框架将 DaCe 与 SoftHier 整合为一套端到端工作流(图 4),支持从高层调度与数据布局描述出发,自动完成代码生成、执行与性能验证。该工作流包含以下四个阶段:
1. 预加载(Preload):将原始数据与数据布局描述处理为预加载文件。该文件定义了初始输入张量(Tensor,AI 与数值计算中常用的多维数组结构)及其在 SoftHier 的 HBM 通道上的分布方式。
2. 生成与优化(Generate and Optimize):结合原始数据的形状与部署调度,生成并优化 SDFG。若部署调度中指定了优化选项(如插入瓦片间通信),则对 SDFG 执行图转换优化。
3. 下转为 C 代码(Lower to C):基于优化后的 SDFG,后端将其下转为 C 代码,并生成与 SoftHier 兼容的可执行二进制文件。
4. 基准测试(Benchmark):SoftHier 从预加载文件初始化 HBM,执行编译后的二进制文件,收集性能指标(如计算吞吐量、延迟),并将计算结果与参考输出对比以验证正确性。
三、基于瓦片的部署调度抽象
本框架采用基于瓦片的部署调度抽象,生成面向 SoftHier 的 SDFG。该抽象以参数化方式描述工作负载如何分解并映射到 SoftHier 的硬件资源,包含三个核心组件:
1. 将工作负载分区并分配给 PE 计算瓦片(瓦片划分与映射);
2. 在片外内存中组织数据(数据布局);
3. 描述执行过程中的数据移动(数据流)。
通过显式定义这些组件,可从高层描述自动生成高性能代码。
3.1 瓦片划分与映射
框架将计算任务分解为多个子任务块(Chunk),并将其映射到各个计算瓦片。瓦片划分与映射共同定义了每个计算瓦片的计算任务,包括其处理的输入数据区域与生成的输出数据部分。
3.1.1 瓦片划分与映射策略
框架支持多种遵循“输出固定”( output-stationary,GEMM 调度的经典策略,指输出数据在计算过程中固定存储于本地内存,避免频繁移动以提高复用率 )的瓦片划分与映射选项:
* 选项 1:单个计算瓦片独立计算完整的输出瓦片(如 2D GEMM 场景)。
* 选项 2:多个计算瓦片协同计算一个输出瓦片(如 3D 拆分 K 维度(Split-K)GEMM 场景)——此时需对各瓦片生成的部分结果执行归约操作以得到最终输出。当需要归约时,框架还提供可配置策略,用于确定哪些计算瓦片负责执行最终归约,并将结果写入 HBM。
3.1.2 集群索引重映射
计算瓦片网格的物理布局(如 4×4 网格)是固定的,但最优映射方式依赖于工作负载与 GEMM 的维度(如 M、N、K 维度大小)。为此,框架引入“集群索引重映射”机制,将物理网格重新解释为逻辑网格(例如,将 4×4 物理网格重映射为 1×16 或 2×8 逻辑网格)。
该机制与基于掩码的集合通信无缝集成:当用户在逻辑拓扑上指定集合操作时,框架会自动在底层 4×4 物理网格上生成对应的掩码(用于选择参与操作的瓦片)。
3.2 数据布局
SoftHier 采用软件管理的、分布式多通道 HBM 系统,每个通道拥有独立的地址空间。
因此,框架需显式控制数据在不同 HBM 通道上的物理分布——选择合理的数据布局对避免内存通道竞争(多个访问请求同时占用同一 HBM 通道导致延迟增加)与 NoC 拥塞(NoC 中数据传输量过大导致通信延迟升高)至关重要,最终可最大化 HBM 带宽利用率。
本框架通过两个关键参数定义数据布局:
3.2.1 拆分方案
拆分方案定义将 M×N 矩阵逻辑划分为块网格的方式。

图 5:矩阵分块与瓦片划分方案。方案包含两部分:分割方案将 M×N 矩阵按例 4×4 划分为 BM×BN 块,块为数据分布最小单位,默认轮询分配给 HBM 通道;放置方案将每块拆为 8×2 的 TM×TN 瓦片,按行优先连续存储,TM/TN 由工作负载瓦片划分指定。该方案可避免 HBM 通道冲突与 NoC 拥堵,最大化带宽利用率。
3.2.2 放置方案
放置方案定义了每个数据块内部的瓦片在 HBM 通道一维地址空间中的具体排列方式。
如图 5 所示,每个数据块进一步被划分为一个 8×2 的瓦片网格。这些瓦片按照“行优先”的顺序连续存储,即先存储第一行的所有瓦片,再存储第二行,依此类推。其中,瓦片的尺寸参数 TM 和 TN 由 3.1 节中定义的工作负载瓦片划分参数指定。

图 5:矩阵分块与瓦片划分方案。该方案包含两部分:分割方案将 M×N 矩阵(例如按 4×4)划分为 BM×BN 个数据块,块是数据分布的最小单位,默认以轮询方式分配给 HBM 通道;放置方案将每个数据块进一步拆分为 8×2 的 TM×TN 瓦片网格,并按行优先顺序连续存储。此方案旨在避免 HBM 通道冲突与片上网络拥堵,最大化内存带宽利用率。
3.3 数据流
数据流定义了 SoftHier 存储层次结构内的数据移动规则:
* 所有计算瓦片的 L1 SPM 通过可编程片上网络互连,形成一个分布式、软件管理的片上存储系统。
* 数据流不仅定义了片外 HBM 与片上 L1 SPM 之间的传输调度,也定义了计算瓦片间通过片上网络的通信调度,同时显式确定了数据移动的方式(如广播、瓦片间点对点通信或 HBM 读取)。
有效的数据流策略是最大化片上数据复用、最小化资源竞争的关键。数据复用指同一数据在计算中被多次使用,从而减少从 HBM 重复读取的次数,降低访存延迟与带宽消耗。
3.3.1 通信/计算重叠
该技术通过将数据传输与计算任务并行执行,以缓解内存访问延迟。内存延迟是指从发起内存访问请求到实际获取数据所需的时间间隔。
一个典型的例子是本地内存中的“双缓冲”技术:系统使用两个独立的缓冲区,一个用于存储当前计算所需的数据,另一个则同时预加载下一次计算所需的瓦片数据,从而实现计算与数据加载的并行执行。
3.3.2 数据流模式原语
框架利用基于掩码寻址的集合通信,支持任意的归约数据流模式。

图 6:数据流模式示例。图中展示了五种依托片上网络集体通信实现的数据流原语:(a) SUMMA(水平广播 A 瓦片、垂直广播 B 瓦片);(b) 脉动阵列(A 瓦片向右传递、B 瓦片向下传递,由近邻通信驱动计算);(c)(d) 两级分层调度(脉动套 SUMMA、SUMMA 套脉动);(e) Split-K(分割 K 维度,通过掩码广播处理子块后进行归约)。这些原语支持灵活的数据流动,以适应不同的 GEMM 计算需求。
目前已实现并评估了多种数据流模式原语(图 6),具体如下:
- SUMMA:如图 6a 所示,遵循经典的 SUMMA 算法方案。在每次迭代中,A 瓦片沿水平方向广播,B 瓦片沿垂直方向广播。
- 脉动阵列:如图 6b 所示,采用脉动阵列的 GEMM 调度方式:A 瓦片向右传播,B 瓦片向下传播。计算过程由完全基于“最近邻通信”的“空间波前”驱动。最近邻通信指每个计算瓦片仅与相邻瓦片通信,以减少通信距离与延迟;空间波前指计算像波浪一样在空间上逐步推进。
- 分层调度:SUMMA 之上的脉动:采用分层分解策略。如图 6c 所示,一个 4×4 的物理计算瓦片集群被逻辑划分为 2×2 的瓦片组网格。在每个内部的 2×2 瓦片组内,系统利用 SUMMA 调度对缩减维度执行 GEMM;同时,外部的瓦片组之间通过全局的脉动调度来协调各子问题的数据移动。
- 分层调度:脉动之上的 SUMMA:与上一种相反,如图 6d 所示,每个 2×2 的内部瓦片组可对其缩减子问题执行本地的脉动 GEMM,而外部的 2×2 瓦片组之间则对完整的 GEMM 维度执行 SUMMA 传播。
- 拆分 K 维度:框架还支持将 K 维度拆分为子计算的数据流原语(图 6e)。利用基于掩码多寻址支持的“跨步广播”,不同的计算瓦片子集处理互不相交的 K 切片,随后通过本地归约操作累积部分结果。跨步广播指按固定间隔向指定瓦片广播数据,而非全量广播;K 切片指将 GEMM 的 K 维度拆分成的片段。
3.3.3 整体同步并行抽象
我们采用 BSP 抽象来定义和描述数据流调度候选方案。BSP“超步”是数据流描述的最小单位,每个超步包含以下三个阶段:
* 计算:每个计算瓦片对其瓦片分布式共享内存中的数据执行本地操作。
* 通信:瓦片通过片上网络交换数据,或从 HBM 读取数据。
* 屏障同步:通过“屏障”同步机制确保所有通信操作完成后,再进入下一个超步。所有参与超步的瓦片必须等待其他瓦片完成当前阶段操作,才能共同推进到下一阶段。
每个超步由 Python 抽象语法树指定,可以清晰、模块化地描述对应阶段的计算、通信与同步逻辑。AST 表示通过显式指定每个超步中“用于计算的缓冲区”与“同时用于通信的缓冲区”,编码了双缓冲机制。借助上述数据流模式原语,框架能够以最小开发工作量引入新的数据流组合。
四、评估
4.1 性能分解分析
4.1.1 基准测试设置
硬件配置
本节实验基于 SoftHier 配置开展,该配置的性能规模与 NVIDIA GH200 的峰值性能相匹配,具体参数如表 1 所示。

表 1:系统规格。该表格是 DiT 框架评估实验的核心硬件配置说明,专门针对模拟 NVIDIA GH200 的性能规格设计。32×32 瓦片布局、4096 位片上网络链路等参数,均对标 GH200 的硬件规模;矩阵引擎的 64×16 计算单元阵列与 1.93 TFLOPS@FP8 性能,确保单瓦片计算能力匹配实际高端加速器;4 TB/s 峰值 HBM 带宽则为高吞吐数据传输提供支撑,为后续对比 DiT 与 GH200 的 GEMM 性能奠定统一硬件基准。
每个计算瓦片包含一个由 64×16 个计算单元构成的矩阵引擎,所有计算瓦片组成一个 32×32 的集群。参考 GH200 的规格,该系统的理论峰值性能为 1979 TFLOPS,内存带宽上限为 4096 GB/s。
部署调度设置
为区分不同优化技术对性能的贡献,我们对所实现的调度方案进行了评估——每种调度方案均由不同的调度原语组合而成。除非另有说明,否则我们始终报告通过选择“性能最佳的数据布局候选方案”所获得的性能;若未选择该候选方案,则报告“基础布局”与“优化布局”对应的性能。
其中,基础布局以行优先方式存储矩阵,且不跨 HBM 通道进行数据分布;基准方案则是未采用专门数据布局或片上通信优化的参考方案,用于性能对比。
本次评估涵盖的调度方案包括:
* SUMMA 调度
* 脉动阵列调度
* 分层调度
4.1.2 数据布局与数据流对基准方案的性能提升
我们首先展示通过优化的数据布局与数据流,基准方案可实现的逐步性能提升。

图 7:GEMM 案例研究评估图表。包含三类核心数据:屋顶线图对比 SUMMA 与基准调度;2D 瓦片 GEMM 性能图;2D 与 3D SUMMA 对比图。
图 7a 在 Roofline 模型上,展示了 SUMMA 调度与基准调度的性能,分别包含“优化数据布局”和“未优化数据布局”两种情况。

图 9:计算密集型 GEMM 性能比较
- “未采用优化布局的基准方案”(Baseline w/o Optimal Layout)受内存带宽限制——由于片外数据传输量过大,其计算强度较低。
- 采用优化布局后,“带优化布局的基准方案”(Baseline w Optimal Layout)性能曲线向上移动,体现了带宽利用率的提升,但仍远未达到峰值性能。
- SUMMA 调度通过广播增强片上数据复用,从而实现了更高的计算强度。
- 而“带优化布局的 SUMMA 调度”(SUMMA w Optimal Layout)则接近计算性能上限。

图 10:Flat GEMM 性能对比

图 11:平面 GEMM 带宽对比
洞察 1:优化的数据布局可提升 HBM 带宽利用率,优化的数据流可提高计算强度。
4.1.3 数据流模式的影响
DiT 支持多种用户可配置的数据流模式。本节将对比这些模式的性能,并从基准测试结果中提炼关键洞察。
2D 分块(2D Tiling)
图 7b 针对不同形状的 2D 分块通用矩阵乘法(GEMM),对比了不同数据流模式的性能。这些模式的关键差异在于“计算 tile 是否同时启动”——这一差异直接导致了观察到的性能变化。

图 8:数据流模式性能差异详细分析图。图中细分计算密集型与存储密集型两类场景。
图 8 详细解释了上述性能差异:
* 在计算密集型场景中(如图 8a,矩阵规模为 4096×2112×7168),流水线会增加不必要的数据等待时间,因此不采用流水线的数据流(如 SUMMA)性能最佳。
* 在存储密集型场景中(如图 8b,矩阵规模为 16384×32768×512),引入流水线可减少 HBM 存储竞争,但流水线级数过多会导致性能下降。
洞察 2:
* 应尽可能使用基于硬件的集合组播。
* 在存储密集型场景中,需限制流水线级数。
3D 分块(3D Tiling)
- 图 7a 和图 7b 显示:当采用 2D 分块时,规模为 4096×2112×7168 的 GEMM 性能仍远低于峰值性能(1979 TFLOPS)。
- 图 7c 对比了 2D SUMMA 与 3D SUMMA 的性能,清晰体现了“额外分块维度”带来的优势。
3D 分块的主要优势在于灵活性更高:它支持非规则维度,且能在 M、N 维度上实现大得多的分块大小。
* 在 2D 分块中,32×32 的计算集群需将较小的 N 维度拆分为 32 个“小型且低效”的切片——这些小型非规则分块在 SoftHier 矩阵引擎上的利用率仅约 50%。
* 相比之下,3D 分块可将相同的 N 维度仅拆分为少数几个大得多的分块:例如,若 8 个计算 tile 共享一个 N 维度分块,但各自处理不同的 K 维度拆分,则 N 维度的分块大小可显著提高,从而显著提高引擎利用率与整体性能。

图 12:模拟 A100 与 GH200 硬件配置下的性能利用率对比图。结果显示,DiT 在类 GH200 配置上,比 CUTLASS 高 1.2-2.0 倍加速。
洞察 3:对于非规则形状的 GEMM,应采用 3D 分块选择“适合矩阵引擎的分块大小”,并利用基于 NoC 的归约操作,以最大化归约效率与 GEMM 性能。
五、结论
本文提出了 DiT(Design in Tiles)——一个专为基于 Tile 的多处理单元(PE)加速器设计的 GEMM 自动化部署框架。该框架通过提供可配置的高层部署调度抽象,彻底消除了人工开发内核的需求,并能在参数范围广泛的基于 Tile 的加速器上,实现可移植、高利用率的 GEMM 部署。
通过 DiT,我们阐明了“最优 GEMM 配置如何随矩阵形状变化”,并证明了:片上网络(NoC)上“基于硬件支持的集合通信”可进一步提升性能。
在与现有基于加速器的 GEMM 部署方案的对比实验中,我们将 SoftHier 配置为“匹配 GH200 峰值性能”的规格,并在其上评估 GEMM 性能——结果显示,DiT 的硬件利用率比“在 GH200 上运行的专家优化 GEMM 内核”高出 1.2–2.0 倍。
参考文献


关注“鲸栖”小程序,掌握最新AI资讯
本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:http://www.itsolotime.com/archives/16090
