S-GPT: Shape-aware graph partitioning and tuning for sustainable inference of large-scale deep neural networks

Abstract

The rapid growth of large-scale deep neural networks has pushed inference workloads toward increasingly diverse shapes. In practical inference, these models are often invoked with tensors of varying shapes rather than a single fixed shape. Such shape variability reshapes operator loops and intermediate tensors, leading to shape-dependent operator costs, memory usage, fusion benefits, and backend tuning complexity. Consequently, graph partitions and kernel schedules optimized for one shape may become inefficient for another. Existing compilers either rely on fixed fusion patterns and static graph partitioning, or handle variable shapes mainly through symbolic representation and runtime specialization, leaving limited ability to jointly optimize graph partitioning and kernel tuning across representative shape distributions. This work proposes S-GPT, a shape-aware compiler framework that co-optimizes graph-level partitioning and backend kernel tuning. At the graph level, we introduce a shape-aware operator weighting method that models both the expected tuning complexity and the cross-shape variance of each operator. Then, we further perform weight-guided subgraph partitioning, enabling different shape intervals to obtain customized partitions while preserving fusion legality and backend schedulability. Finally, we develop a shape-aware kernel tuning mechanism, where structurally similar subgraphs and neighboring shape intervals can reuse historical schedules, measurement records, and cost-model states. Experimental results demonstrate that, compared to the state-of-the-art compilation framework, S-GPT reduces compilation time by 1 . 9× , improves average speedups of 1 . 72× and 1 . 63× on the NVIDIA Tesla A40 and NVIDIA GeForce RTX3090, respectively.