通过算子融合提升大模型在Tensix架构上的推理效率
Operator Fusion for LLM Inference on the Tensix Architecture
- 将RMSNorm与矩阵乘法融合，减少片上内存访问次数
- 在Qwen系列模型上实现最高37.44%注意力层延迟降低
- 适合关注边缘设备高效推理的开发者与研究人员
本研究针对Tenstorrent Tensix架构上Transformer模型的本地推理瓶颈，提出一种算子融合策略以增强数据局部性。将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%以上，验证了数值一致性下的显著端到端性能提升。
原文摘要 · Abstract (English)
This study addresses on-device inference bottlenecks of Transformer models on Tenstorrent's Tensix architecture and proposes an operator fusion strategy that enhances data locality. RMSNorm is fused with matrix multiplication in self-attention and in the FFN, enabling back-to-back execution of memory-bound and compute-bound operators in on-chip SRAM to significantly reduce DRAM reads/writes of intermediate results and scheduling overhead. To support multi-core parallelism, a NoC-based multicast mechanism is leveraged in which row/column master nodes efficiently distribute inputs and weights across the core mesh, alleviating DRAM bandwidth contention. Experiments on the Wormhole platform with Qwen2.5-0.5B, Qwen3-0.6B, and Qwen3-4B show up to 37.44% latency reduction for attention and 15.89% for MLP, with up to 7.91% reduction per decoder layer, while Pearson Correlation Coefficient (PCC) remains above 98.75%, confirming significant end-to-end efficiency gains under numerical consistency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。