I/O Lower-Bound Theory and Reinforcement Learning for Efficient Neural Network Inference Optimization
TL;DR
An automated optimization framework with a hierarchical two-layer tuning mechanism that synergizes theoretical I/O constraints with graph-level adaptive fusion while accounting for search overhead, the framework systematically explores high-performance execution patterns.
Abstract
As artificial intelligence models grow in complexity, optimizing neural network inference has become a critical challenge. Existing approaches often rely on manual, expert-driven tuning tailored to specific hardware, which lacks scalability across diverse architectures. In this paper, we propose an automated optimization framework with a hierarchical two-layer tuning mechanism. At the node level (intra-operator), we introduce an I/O lower-bound theory based on the Red-Blue Pebble game and the (X1, X2)-Partition theorem to guide tiling and memory-mapping configurations. At the graph level (inter-operator), we employ a reinforcement learning (RL) strategy to adaptively identify optimal operator fusion boundaries across network topologies. By synergizing theoretical I/O constraints with graph-level adaptive fusion while accounting for search overhead, the framework systematically explores high-performance execution patterns. For the TileAttn operator, the (X1, X2)-Partition theorem raises DRAM flow estimation accuracy from 77.5%-82.0% under X-Partition to 86.3%-94.8%. Our method reduces shared memory traffic by 10.79% on average, achieves the best performance in 65.3% of cross-platform cases and top-two in 86.1%, and its DQN-based fusion engine outperforms greedy strategies in 91.67% of scenarios. We further analyze the algorithm’s overhead and its amortization break-even points.
