FlashTensor: Optimizing Tensor Programs by Leveraging Fine-grained Tensor Property

Abstract

Deep neural networks (DNNs) have shown significant effectiveness in natural language processing and video applications. However, DNN models, especially for long-context tasks, introduce extremely large intermediate tensors, producing substantial memory overhead. Although considerable efforts have been made to optimize DNNs, insufficient awareness of tensor properties has hindered effective memory optimization and can lead to inefficient computations in a long-context scenario.

In this paper, we present FlashTensor1, a DNN optimization system that reduces memory overhead and improves inference performance by leveraging fine-grained tensor properties. We first extract and identify essential tensor properties from a computation graph, such as reduce dependency and broadcastability. Then, we apply various optimizations involving transformation and kernel mapping based on these properties. Experiments on seven models demonstrate that FlashTensor achieves speedups of 1.50× and 3.24× on average for end-to-end and core module performance, respectively, compared to eight state-of-the-art works, on H100 (1.86× and 3.70× on A100).