Operator Fusion for LLM Inference on the Tensix Architecture

2026-06-02

AI Summary

This work addresses the memory bandwidth bottleneck and scheduling overhead encountered when deploying Transformer models on edge devices with the Tensix architecture. The authors propose a deep fusion of RMSNorm with matrix multiplications in both self-attention and feed-forward network (FFN) modules, enabling contiguous execution of memory-intensive and compute-intensive operators within on-chip SRAM. A novel NoC-based multicast strategy leveraging row- and column-major master nodes is introduced to support efficient multi-core parallelism, significantly enhancing data locality and throughput. Evaluation on the Wormhole platform using Qwen-family models demonstrates up to 37.44% and 15.89% latency reduction in attention and MLP layers, respectively, with per-decoder-layer latency reduced by as much as 7.91%, while preserving numerical consistency (PCC > 98.75%).

Abstract

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.

Problem

• LLM Inference
• Operator Fusion
• Tensix Architecture
• On-device Inference
• Transformer Models

Innovation

• Operator Fusion
• Tensix Architecture
• On-device LLM Inference
• NoC-based Multicast
• Data Locality