Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning

2024-02-04

Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2

AI Summary

Tensor program tuning faces two key challenges: low search efficiency—due to reliance on high-accuracy but slow-to-train cost models requiring hours to days—and poor cross-platform online adaptability. This paper proposes a *Draft-then-Verify* two-stage exploration framework: (1) a lightweight symbolic analysis rapidly prunes the scheduling space to generate promising candidates; and (2) a high-accuracy learned cost model rigorously evaluates and refines the top candidates. We further introduce Momentum-based Online Adaptation (MoA), the first method enabling real-time, zero-shot transfer to unseen hardware across platforms. Our approach is deeply integrated into TVM. Experiments show that online tuning achieves 2.6×–4.82× speedup over Ansor; offline tuning attains 4.75× and 4.05× speedup over TenSet and TLP, respectively; and on TensorCore, it outperforms MetaSchedule by 4.08×.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: User modeling for targeted and personalized online advertising
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Tensor program tuning is essential for the efficient deployment of deep neural networks. Search-based approaches have demonstrated scalability and effectiveness in automatically finding high-performance programs for specific hardware. However, the search process is often inefficient, taking hours or even days to discover optimal programs due to the exploration mechanisms guided by an accurate but slow-learned cost model. Meanwhile, the learned cost model trained on one platform cannot seamlessly adapt online to another, which we call cross-platform online unawareness. In this work, we propose Pruner and MoA-Pruner. Pruner is a"Draft-then-Verify"exploration mechanism that accelerates the schedule search process. Instead of applying the complex learned cost model to all explored candidates, Pruner drafts small-scale potential candidates by introducing a naive Symbol-based Analyzer (draft model), then identifies the best candidates by the learned cost model. MoA-Pruner introduces a Momentum online Adaptation strategy to address the cross-platform online unawareness. We incorporate Pruner into the TVM and conduct extensive experiments on three GPU-based platforms. Results show considerable speedup in schedule search time. In online tuning scenarios, Pruner and MoA-Pruner achieve an average speedup of $2.6 imes$ and $4.82 imes$ compared to Ansor. In offline tuning scenarios, Pruner achieves an average speedup of $4.75 imes$ and $4.05 imes$ compared to TenSet and TLP, respectively. Furthermore, Pruner achieves an average speedup of $4.08 imes$ compared to MetaSchedule on TensorCore.

Problem

• Accelerate tensor program tuning for deep neural networks
• Improve search efficiency in finding high-performance programs
• Address cross-platform online unawareness in cost models

Innovation

• Draft-then-Verify mechanism accelerates tensor tuning
• Symbol-based Analyzer drafts potential candidates efficiently
• Momentum online Adaptation handles cross-platform unawareness