面向Tensix架构的大语言模型推理中的算子融合

# Tensix架构上大语言模型推理的算子融合
来源：https://arxiv.org/html/2606.09879
11institutetext:麒麟软件有限公司，天津，中国22institutetext:openKylin社区，天津，中国
22email:li\_ke@hl\-it\.cn33institutetext:国防科技大学计算机科学与技术学院，长沙，中国

###### 摘要

本研究针对Transformer模型在Tenstorrent Tensix架构上的设备端推理瓶颈，提出了一种增强数据局部性的算子融合策略。将RMSNorm与自注意力机制和前馈神经网络中的矩阵乘法进行融合，使内存受限算子和计算受限算子在片上SRAM中连续执行，显著减少中间结果的DRAM读写次数以及调度开销。为支持多核并行，利用基于片上网络（NoC）的组播机制，其中行/列主节点将输入和权重高效分发到核心网格中，缓解DRAM带宽竞争。在Wormhole平台上使用Qwen2.5-0.5B、Qwen3-0.6B和Qwen3-4B进行的实验表明，注意力机制延迟最多降低37.44%，MLP延迟最多降低15.89%，每个解码器层延迟最多降低7.91%，同时皮尔逊相关系数（PCC）保持在98.75%以上，证明了在数值一致性下的端到端效率显著提升。

## 1 引言

随着大型语言模型在代码生成、问答和多模态推理中的广泛应用，边缘推理的性能和能效变得日益关键[yuan2024llm]。在推理过程中，Transformer工作负载呈现出显著的“读取-计算-写入”模式：内存受限算子（归一化、激活）与计算受限算子（矩阵乘法）交替出现。这导致中间结果在片外DRAM和片上SRAM之间频繁移动，带来高内存延迟和调度开销。在资源受限的设备（如Tenstorrent Tensix）上，提高数据局部性、减少冗余数据移动以及实现计算与内存访问的重叠，是实现高效推理的关键。

我们提出了一种面向数据流和内存层次结构的算子融合策略，专门针对仅解码器Transformer，并受Tensix架构限制。具体而言，我们将RMSNorm与自注意力机制中的QKV投影矩阵乘法融合，并将RMSNorm与第二个残差块中第一个FFN矩阵乘法融合。这些操作在片上SRAM中连续执行，避免DRAM回写和后续读取。对于多核并行，我们利用基于NoC的行/列组播机制，其中主节点将输入和权重高效分发到核心网格中，缓解DRAM带宽竞争并提高内存效率。

我们的实验采用Tenstorrent的Tensix架构和Wormhole N300加速器作为评估平台[Wormhole18:online]。Tensix处理器采用2D网格互连，强调计算、控制和通信的解耦。每个核心集成了多个RISC-V控制核心、专用的FPU/SFPU计算单元以及大容量本地SRAM，支持片上数据流和算子级并行[tenstorr70:online]。Wormhole N300集成了两个Tensix芯片，并可在64/128核心配置之间灵活切换，为所提出的融合和组播机制提供了可重构的硬件基础和真实的评估环境。

本研究的主要贡献如下：

- • 提出一种针对Transformer的单核融合方案：在Tensix核心内以直通方式融合RMSNorm与QKV/FFN矩阵乘法，最大化局部性并减少调度开销。
- • 采用2D网格映射并利用NoC组播机制进行多核执行，利用行/列数据共享来减少Tensix核心网格的带宽压力。
- • 在Tenstorrent Wormhole N300上使用Qwen2.5-0.5B、Qwen3-0.6B和Qwen3-4B进行了系统评估。结果表明，关键模块和解码器层的延迟显著降低，同时保持数值一致性（PCC ≥ 98.75%）。

## 2 相关工作

### 2.1 大型语言模型

现代大型语言模型（LLM）主要采用仅解码器架构。单个解码器层通常包含两个残差模块：自注意力模块和前馈网络（FFN）。这两个模块都围绕矩阵乘法展开：自注意力执行线性投影以生成查询、键和值向量以及最终输出投影；FFN包含多个线性层（例如，Llama中的SwiGLU包括门控、上升和下降投影矩阵）以执行非线性变换和维度映射。在这些模块的输入处，通常使用RMSNorm作为归一化层以提高训练稳定性。

参见图注图 1：仅解码器架构概览
### 2.2 算子融合

算子融合将轻量级的内存受限算子（如归一化、激活）与计算受限算子（如矩阵乘法、卷积）合并在一起执行。这种方法提高了整体计算效率，并减少了深度模型中的内存占用和调度开销。核心目标是最大化片上数据局部性并最小化对片外DRAM的访问，从而提升边缘推理性能[11177548]。

从实现角度来看，算子融合通常将神经网络建模为有向无环图（DAG），并将其划分为可融合的子图。在这些子图中，上游算子的输出直接从片上缓存供下游算子消耗，避免了不必要的中间读写[alwani2016fused]。文献探索了多种策略：Optimus[cai2022optimus]使用自底向上的动态规划和基于查找表（LUT）的剪枝来减少片外流量搜索空间；TileFlow[zheng2023tileflow]在搜索最优融合方案时，结合遗传算法与蒙特卡洛树搜索（MCTS）来量化数据移动并优化资源分配。

### 2.3 Tensix核心

Tensix架构[tenstorr70:online]由Tenstorrent提出，是一种异构、可扩展的AI处理器设计。Tenstorrent是一家专注于RISC-V CPU和AI加速器的半导体公司，提供以Tensix为核心的芯片以及Wormhole[Wormhole18:online]等加速器，配合TT-metal[tenstorr91:online]低级编程范式和算子库。与传统GPU不同，Tensix将计算、控制和通信解耦，并使用2D网格互连在芯片上形成一个可扩展的、面向数据流的计算网络。该网络包含专用节点：Tensix核心作为计算节点，辅以存储节点、芯片管理/控制节点以及高带宽通信节点，以提高并行性和数据流效率。

参见图注\(a\)Tensix核心架构[ttmetalM8:online]
参见图注\(b\)Tensix内核

图 2：Tensix架构概览Tensix核心是Tensix处理器的计算单元（图2(a) (https://arxiv.org/html/2606.09879#S2.F2.sf1)）。每个核心包含五个小型RISC-V控制核心[waterman2014risc]（“婴儿核心”）、用于矩阵运算的专用硬件单元（FPU）、用于向量运算的专用单元（SFPU）以及1.5 MB本地SRAM。典型的数据流：数据通过片上NoC传递到核心，解包，由专用计算单元处理，重新打包，然后通过NoC发送到DRAM或其他Tensix核心。

如图2(b) (https://arxiv.org/html/2606.09879#S2.F2.sf2)所示，婴儿核心并行处理指令调度和控制流，向FPU/SFPU发出命令以执行计算（而不是直接在RISC-V核心上运行计算）。芯片提供两个独立的NoC接口，沿相反方向环绕芯片，允许并行读写操作，这些操作与计算重叠，形成高效的生产者-消费者流水线。

### 2.4 TT-metalium

TT-metalium[tenstorr91:online]是Tenstorrent的低级内核编程范式，用于在Tensix处理器上实现自定义AI计算。其API类似于OpenCL，提供用于高级操作的C++接口和直接硬件控制，能够访问FPU、SFPU、NoC接口和SRAM，以编写高度优化的内核，实现细粒度控制。TT-metalium采用单程序多数据（SPMD）执行模型[Singlepr87:online]：多个处理单元在不同数据分片上运行相同程序。在Tensix上，核心网格（常见配置为64/128核心）使用SPMD协作处理同一计算的不同块。

如图2(b) (https://arxiv.org/html/2606.09879#S2.F2.sf2)所示，TT-metalium将核心计算流水线分为三种内核类型，并使用核心SRAM中的循环缓冲区协调它们，实现生产者-消费者并行。这三种内核在时间上完全重叠，形成高效的流水线：

- • 读取内核：由RISC-V CPU0控制，使用双缓冲异步从DRAM获取数据到SRAM；
- • 计算内核：由RISC-V CPU1-3控制，将数据从SRAM驱动到FPU/SFPU进行计算，并将结果写回SRAM；
- • 写入内核：由RISC-V CPU4控制，异步将结果从SRAM写回DRAM。

Tensix核心架构和TT-metalium范式已广泛应用于各种工作负载，包括求解拉普拉斯方程[brown2024accelerating]、快速傅里叶变换[10.1007/978-3-032-07612-0_46]、矩阵乘法[pizzini2025assessing]、注意力机制[thuning2024attention]以及高性能计算[almerol2025accelerating]。

## 3 方法

### 3.1 算子描述

#### 3.1.1 RMSNorm算子

RMSNorm是Transformer模型中常用的归一化层，通过对输入特征进行归一化来防止梯度消失或爆炸。其计算为

RMSNorm\(x;γ\)=xE\[x2\]\+ε⊙γ,\\mathrm\{RMSNorm\}\(x;\\gamma\)\\;=\\;\\frac\{x\}\{\\sqrt\{\\mathrm\{E\}\[x^\{2\}\]\+\\epsilon\}\}\\,\\odot\\,\\gamma,其中ε\\epsilon是一个小常数，γ\\gamma是可学习参数，E\[x2\]\\mathrm\{E\}\[x^\{2\}\]是输入特征的均方值。

在现代仅解码器LLM（如Llama、Qwen）中，RMSNorm作为预归一化层出现在每个残差块的开头，即在自注意力之前和FFN之前。它稳定后续线性投影和非线性映射的数值范围，提高了训练和推理中的收敛性和鲁棒性。

#### 3.1.2 MatMul算子

MatMul是Transformer中基本的线性变换原语；其标准形式是GEMM：

C←αAB\+βC,A∈RM×K,B∈RK×N,C∈RM×N\.C\\;\\leftarrow\\;\\alpha\\,AB\\;\+\\;\\beta\\,C,\\quad A\\in\\mathbb\{R\}^\{M\\times K\},\\;B\\in\\mathbb\{R\}^\{K\\times N\},\\;C\\in\\mathbb\{R\}^\{M\\times N\}.
在解码器模块中，输入X∈RT×dmodelX\\in\\mathbb\{R\}^\{T\\times d\_\{\\mathrm\{model\}\}\}被投影为

Q=XWQ,K=XWK,V=XWV,W\{Q,K,V\}∈Rdmodel×\(nhdh\),Q=XW\_\{Q\},\\quad K=XW\_\{K\},\\quad V=XW\_\{V\},\\quad W\_\{\\\{\\\!Q,K,V\\\!\\\}\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{model\}\}\\times\(n\_\{\\mathrm\{h\}\}d\_\{\\mathrm\{h\}\}\)\},并重排为Q,K,V∈Rnh×T×dhQ,K,V\\in\\mathbb\{R\}^\{n\_\{\\mathrm\{h\}\}\\times T\\times d\_\{\\mathrm\{h\}\}\}. FFN包含

U=XWup,G=XWgate,Y=φ\(G\)⊙U,Z=YWdown,U=XW\_\{\\mathrm\{up\}\},\\quad G=XW\_\{\\mathrm\{gate\}\},\\quad Y=\\phi\(G\)\\odot U,\\quad Z=YW\_\{\\mathrm\{down\}\},其中Wup∈Rdmodel×dffW\_\{\\mathrm\{up\}\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{model\}\}\\times d\_\{\\mathrm\{ff\}\}\}和Wdown∈Rdff×dmodelW\_\{\\mathrm\{down\}\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{ff\}\}\\times d\_\{\\mathrm\{model\}\}\}. 这些算子占据了大部分FLOPs并且是计算受限的，随时间与内存受限算子（如归一化和激活）交替出现。为了提高边缘吞吐量，我们将前面的RMSNorm与上述投影矩阵乘法连续融合。归一化后的输出在片上SRAM中原地消耗，减少了调度和片外访问，同时与NoC数据流形成了高效的生产者-消费者流水线。

### 3.2 单Tensix算子融合

参见图注\(a\)融合前
参见图注\(b\)融合后

图 3：单Tensix算子融合示意图图3(a) (https://arxiv.org/html/2606.09879#S3.F3.sf1)展示了Tensix上的标准编程范式：每个独立算子从共享DRAM读取数据到片上SRAM，在SFPU/FPU上执行计算，并将结果写回DRAM（见第2.2节 (https://arxiv.org/html/2606.09879#S2.SS2)）。这种“读取-计算-写入”模式导致中间数据的频繁片外移动，造成显著的内存延迟和调度开销，限制了整体效率。

为了解决这个问题，我们提出了一种单核算子融合策略，通过片上直通消除中间结果的DRAM读写：取消前一个算子（RMSNorm）的回写和下一个算子（矩阵乘法）的后续读取。融合后的计算为

C=RMSNorm\(x;γ\)⋅B,C\\;=\\;\\mathrm\{RMSNorm\}\(x;\\gamma\)\\,\\cdot\\,B,其中xx是激活块，γ\\gamma是归一化缩放因子，BB是权重矩阵。为确保高效的硬件实现，我们遵循以下原则：

- • 避免数据移动：将RMSNorm输出保留在片上SRAM中，并直接输入矩阵乘法，避免不必要的DRAM传输和格式转换，减少算子间的调度和同步开销。
- • 核心网格对齐：确保RMSNorm和矩阵乘法使用相同的核心网格分区，使得数据自然地位于正确的核心上，无需跨核心迁移。
- • 权重重用：在批量场景中，权重从DRAM读取一次，并在多个激活请求中重用，以提高片外带宽效率。

融合后的工作流（图3(b) (https://arxiv.org/html/2606.09879#S3.F3.sf2)）通过片上直通减少了片外往返，并通过网格对齐和权重重用提高了带宽效率和流水线并行性。

### 3.3 多Tensix算子融合

由于单个Tensix核心的计算能力有限，大规模深度学习任务通常需要多个核心并行协作。图4 (https://arxiv.org/html/2606.09879#S3.F4)展示了多核环境下的融合策略：将计算分区并映射到2D核心网格上。每个核心独立为其分配的输入块执行RMSNorm和矩阵乘法，并将中间结果存储在本地SRAM中，以最小化核间迁移。

多核并行的关键在于正确处理数据依赖关系。对于矩阵乘法，图4 (https://arxiv.org/html/2606.09879#S3.F4)采用了输出分块策略，将CC均匀映射到核心网格。代数上，同一行的核心共享AA的行块，而同一列的核心共享BB的列块。这一性质支撑了高效的数据分发和并行计算，使得核心能够独立并行地执行本地工作。

对于RMSNorm，由于同一行的核心后续共享AA的行块，行内的所有核心必须应用相同的归一化。为减少行内通信同时保持SPMD一致性，我们使用计算冗余：行内每个核心执行相同的归一化。