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

A deep-learning compiler finds fast GPU code for each layer by generating thousands of candidate programs and asking a learned cost model which to measure, but running that model on every candidate is itself nearly 40 percent of tuning time — Pruner first scores candidates with a cheap formula that counts how well each one fills the GPU's memory levels and parallel units, keeps only the best 512 of 8,000, and only then runs the learned model, cutting search time 2.6 to 4.8 times against Ansor; MoA-Pruner adds a momentum-updated copy of a cost model pre-trained on other hardware so tuning on a new GPU starts informed instead of from scratch.

The argument in one glance

Problem

Finding fast code for each layer means searching billions of schedules

To run a neural network efficiently on a GPU, each operator (a convolution, a matrix multiply) must be turned into a loop nest with the right tiling, thread mapping, caching and unrolling. Hand-written kernel libraries do this for common shapes; for everything else, search-based compilers such as TVM explore a schedule space that reaches billions of combinations. Measuring every candidate on the device is impossible, so a cost model predicts which few hundred are worth measuring.
What people did

Learned cost models: accurate, and applied to every candidate

Ansor, TenSet and TLP train neural networks to predict a schedule's latency from hand-designed features, 164 per statement in Ansor, thousands in TIRAMISU, and use those predictions to steer a genetic search. In online tuning, the model is retrained on each round's measurements; in offline tuning, it is pre-trained on a large dataset first.
Why it fell short

The model is slow to apply, hard to train, and stuck to one platform

On a Jetson Orin, running the learned model on every explored candidate is about 40 percent of Ansor's tuning time, and more for heavier models like TLP. TLP's Transformer needs a large offline dataset because its one-hot features barely differ between programs (1.4 percent of feature values for a GEMM), and it sometimes crashed when fine-tuned online. And a model trained on one GPU cannot be applied online to another: the paper calls this cross-platform online unawareness, and transfer tricks like TenSet's gap model or Moses' distillation cost as much as training from scratch.
This paper's idea

Draft with a formula, verify with the learned model, adapt with momentum

Pruner splits exploration into two stages. The Latent Schedule Explorer drafts: it scores each candidate with a Symbol-based Analyzer, an empirical formula built from eight hardware-aware symbols that count memory allocation, computation and parallelism at three memory levels, and uses that score as the fitness for a genetic search, keeping the best 512 of about 8,000 candidates. The Pattern-aware Cost Model verifies: a Transformer over 23-dimensional temporal dataflow features, extracted from the recurring multi-level tiling pattern, plus Ansor's statement features, ranks the 512 and picks ten to measure. MoA-Pruner adds momentum online adaptation: a cost model pre-trained on another GPU initializes the target model and is itself nudged by the target's gradients with momentum 0.99, so online tuning on new hardware starts well and stays stable. Against Ansor, Pruner and MoA-Pruner reach the same result 2.6 and 4.82 times faster on average; offline, Pruner is 4.75 times faster than TenSet and 4.05 times faster than TLP; on tensor cores it is 4.08 times faster than MetaSchedule.

Why this problem is worth solving

A team wants to deploy a vision model on an NVIDIA Jetson in a car. cuDNN covers the standard convolutions, but the model has a few odd shapes and a fused operator the library does not know. A

search-based deep-learning compiler

A compiler that, for each operator, generates a space of possible loop schedules and searches it, measuring candidates on the target device, to find the fastest program.

like TVM will find good code for those, but the paper’s Table 1 shows what it costs on a Jetson Orin: tuning ResNet-50 for 2,000 trials with Ansor takes about 85 minutes, of which 35 minutes is the exploration step, running the

learned cost model

A model, usually a neural network trained on measured programs, that predicts a candidate schedule's latency so the search can avoid measuring it.

on candidates, 5 minutes is training that model, and 44 minutes is measuring the chosen programs on the device. Inception-V3 and DeTR are similar. Across a full model suite that is hours; with a heavier cost model it is longer.

Then the team buys an A100 for the data centre. The cost model that learned the Orin’s behaviour is useless there: the same tile sizes perform differently, and the paper shows that trying to reuse it online is worse than starting over, because the first rounds of measurements on a new platform are few and biased and can wreck the model’s training. Building a proper offline dataset for the new GPU, as TLP requires, took the authors more than ten days for 1.5 million programs. The paper’s target is to make the search itself cheap and to make a model from one GPU useful on another without that expense.

What people did before, and where it broke

Ansor and TenSetMLP. Multi-layer perceptrons over 164 hand-extracted features per innermost statement, retrained online or pre-trained on TenSet’s 16 million programs. Where it broke. Every candidate in every search round passes through feature extraction on the CPU and inference on the GPU: nearly 40 percent of tuning time on Orin.
TLP. A Transformer over one-hot encodings of schedule primitives, cheaper to extract. Where it broke. The features hardly vary between programs, so the model needs a large offline dataset to train and sometimes fails outright when fine-tuned online, leaving no result at all for some workloads in the paper’s experiments.
Constraint-based and gradient-based search: Heron, Felix. Reduce the number of iterations with hardware constraints or differentiable spaces. Where it broke. They still run the expensive model on each candidate, and Felix fails on irregular operators.
Rule-based generation: Roller. Derive candidates from empirical formulas without a learned model. Where it broke. Fast, but its rules discard solutions that lie outside them, and it misses the optimum; on TITAN V it is slower than Ansor on ResNet-50 at batch 128.
Cross-platform transfer: TenSet’s gap model, Moses’ distillation, TLP’s multi-task heads. Where it broke. The gap model costs as much to train as a fresh model and doubles inference; distillation needs extra evaluation of every weight; multi-task learning needs a dataset for the target platform and does not adapt online.
The gap. No system both avoided running the heavy model on every candidate and let a model from one platform help online on another without extra machinery.

Where Pruner sits

One tuning round for one subgraph

Task scheduler

Ansor's gradient-based scheduler picks which subgraph of the network to tune this round.
↓ →
Draft: Latent Schedule Explorer new

Generate the schedule space, sample candidates, and run a genetic search whose fitness is the Symbol-based Analyzer formula. Keep the 512 fittest, plus some random ones.
↓ →
Verify: Pattern-aware Cost Model new

Extract temporal dataflow and statement features for the 512 and rank them with a multi-branch Transformer. Select 10 to measure.
↓ →
Measure

Compile and time the 10 on the device; record the results.
↓ →
Update new

Fine-tune the cost model on the new measurements. In MoA-Pruner, also push the pre-trained Siamese model toward the target model with momentum.

The two new boxes in the middle replace the single step where prior systems ran the learned model on all 8,000 explored candidates. Pruner is implemented inside TVM's Ansor and MetaSchedule.

The idea

Hardware-aware symbols. Pruner’s observation is that a tensor program’s performance follows how well it fits the accelerator’s hierarchy of parallel units and memories. It models a GPU as three levels: L0 for registers, L1 for shared memory and the warp, L2 for global memory and the streaming multiprocessors. From the schedule primitives that generate a candidate, the split factors, reorders and cache reads, it extracts eight

hardware-aware symbols

Quantities such as memory allocated, operations counted, or parallel blocks launched at one level of the memory hierarchy, computed symbolically from a schedule's split factors without compiling it.

: memory allocation and computation count at L0, allocation and parallelism at L1, footprint, parallelism, innermost transaction dimension and computation count at L2. Each is a product or sum of loop extents, so it costs almost nothing to evaluate.

Penalties and the Symbol-based Analyzer. The symbols become six penalty terms between 0 and 1. If register allocation exceeds the L0 limit, a memory penalty applies. At L1, the number of scheduling blocks divided by the smallest multiple of the hardware’s parallel units that covers them measures how much of a warp or SM is wasted: six blocks on four units gives 0.75. At L2, a similar term covers SM utilization and another covers memory transaction alignment. The

Symbol-based Analyzer

Pruner's draft model: an empirical formula that multiplies the hardware's peak compute and bandwidth by the penalty terms and sums per-statement compute and memory latencies.

multiplies peak compute throughput by the compute penalties to get utilized throughput, does the same for memory bandwidth at L2, and estimates each innermost statement’s latency as operations divided by utilized compute plus bytes divided by utilized bandwidth. The sum is the draft score. No features are extracted, no neural network runs.

The draft stage. The Latent Schedule Explorer generates the schedule sketch as Ansor does, samples an initial population, and runs a genetic algorithm over tiling factors for a fixed number of steps with the analyzer as fitness, filtering the accumulated candidates to a set of 512 with the best hardware-fitness scores. A few random schedules are added for diversity. In the ablation on TenSet’s T4 data, this draft set almost always contains the true best schedule of each subgraph (Best-1 score 0.995 at size 512), while a random genetic search does not; removing the compute penalties drops the score to 0.88.

The verify stage. The Pattern-aware Cost Model runs only on the draft set. Its new feature is the

temporal dataflow feature

A 23-dimensional vector per data-block movement (load, compute, store) between memory levels in the multi-level tiling pattern, recording reuse, stride, access type, allocation size, direction and compute density.

: because every tiled program repeats the same pattern of moving blocks from global to shared to registers, computing, and writing back, the sequence of those movements can be extracted from the low-level IR and encoded consistently across programs. Element-wise operators, which are not tiled and make up under 3 percent of TenSet, get zero padding. A self-attention branch models the sequence; linear layers encode Ansor’s statement features; the branches are concatenated and trained with a ranking loss. On TenSet, PaCM reaches top-1 accuracy of 0.892 on T4 and 0.897 on K80, against 0.862 and 0.880 for TLP, and it converges with far less training data.

Momentum online adaptation. For online tuning on a new GPU, MoA-Pruner keeps two copies of the cost model with the same architecture: a

Siamese model

A second copy of the cost model, initialized from a model pre-trained on other hardware, whose weights are updated slowly toward the target model rather than by direct training.

pre-trained on TenSet’s K80 data and the target model. Each round, the target is initialized from the Siamese weights, fine-tuned on the round’s measurements, and then the Siamese weights are moved toward the target by a momentum update with factor 0.99, the same rule MoCo uses in self-supervised learning. The Siamese model never trains directly on the noisy early data, so it stays a good initializer, while it drifts toward the target platform over time. The strategy adds no transfer model, no distillation pass, and works with any learned cost model.

Tracing one tuning round on a GEMM

A 128 by 128 by 128 matrix multiply with ReLU, one round on an A100

Loops i, j, k are each split into tiles (I0..I4, J0..J4, K0..K2) as in the paper's Figure 3. Numbers of candidates match the paper's settings; scores and latencies are illustrative.

Step 1 of 6

Generate and sample

Ansor's rules produce a schedule sketch: split i and j five ways and k three ways, reorder, cache A and B in shared memory, compute C in registers, fuse the ReLU. The explorer samples an initial population of concrete tile sizes.

space: all (I0..I4, J0..J4, K0..K2) factorizations of 128 | initial population sampled

Why it works

A cheap filter before an expensive judge. The analyzer runs in microseconds and, on TenSet, its top 512 contain the optimum with probability 0.995. Because the learned model sees only those 512, its inference cost per round drops by roughly 16 times, and exploration stops dominating tuning time.
Hardware fitness is a good proxy for latency ordering. Tile choices that waste warps, spill registers or misalign memory transactions are slow regardless of the fine details, so a formula over those quantities separates good regions of the space from bad ones. The ablation confirms the compute-utilization penalties matter most.
Dataflow features vary where one-hot features do not. Every tiled program has the same movement pattern with different sizes, strides and reuse, so the 23-dimensional vectors differ meaningfully between candidates and a Transformer can learn from small data, where TLP’s one-hot vectors changed in only 1.4 percent of positions.
Momentum shields the initializer from bad early data. The target model absorbs the noisy first rounds; the Siamese model changes by only one percent per round, so it keeps its cross-platform knowledge while slowly specializing. The ablation shows MoA beats direct online fine-tuning on four of five workloads.
The two stages compound. Removing the draft stage hurts most in the ablation because the learned model alone cannot find good candidates quickly; removing the dataflow features hurts more than removing statement features. Together they make the tuning curve drop early and keep dropping.

What it costs

The analyzer is a model of the hardware. Its three-level abstraction, penalty definitions and limits must be written for each accelerator class; the authors added a tensor-core symbol and a shared-to-fragment dataflow to support MetaSchedule on tensor cores.
Slightly more GPU memory. PaCM uses 1,694 MB at inference batch 4,096 against 1,546 MB for TenSetMLP and Ansor, though far less than TLP’s 4,812 MB.
Hand-tuned libraries still win some shapes. When the reduction axis is long and the parallel axis short, cuBLAS uses split-K and beats Pruner: 36.7 versus 41.3 microseconds on a ViT linear layer, and on one GPT-2 projection 18.96 versus 23.46 microseconds. TensorRT beats it on several networks for the same reason.
The draft set size is a knob. At 256 the Best-k scores barely move; smaller sets start to lose the optimum. The paper fixes 512.
MoA needs a pre-trained model from somewhere. The Siamese model comes from TenSet’s K80 data. Without any pre-trained model there is nothing to adapt, and plain Pruner applies.
Some baselines could not run everywhere. TLP failed on some workloads after fine-tuning, Felix on irregular shapes, TLM on unseen subgraphs, Adatune on ConvTranspose2d; comparisons are reported only on the cases each baseline completed.

The evaluation

Fourteen networks (ResNet, WideResNet, Inception-V3, DenseNet-121, MobileNet-V2, DCGAN, DeepLab-V3, BERT tiny and base, GPT-2, Llama, OPT, Mistral, ViT, DeTR) on three GPUs: NVIDIA A100, TITAN V and Jetson Orin-AGX. Every search-based method gets 200 rounds of 10 measurements, 2,000 trials. Offline models are pre-trained on TenSet and fine-tuned on about 500,000 programs per platform; MoA’s Siamese model is pre-trained on TenSet K80.

Search time to reach each baseline's final result, and tuning quality

Property	Baseline	Pruner	MoA-Pruner
Online tuning versus Ansor, average over A100, TITAN V, Orin	1.0
Offline tuning versus TenSetMLP and TLP	1.0		—
Tensor cores versus MetaSchedule, A100, six language models	1.0		—
Ansor with more trials, A100 (latency ms, cost min)		—
Compilation time, 2,000 trials, TITAN V, five networks	Ansor: 99 to 125 min
Cost-model top-1 on TenSet T4 / K80	TenSetMLP 0.859 / 0.878; TLP 0.862 / 0.880		—
Versus inference frameworks on A100, ten networks	PyTorch 2.2, Triton, TensorRT	—

From Figures 6 to 12 and Tables 5, 7, 9 and 11 of the paper. Search-time speedups measure how long Pruner takes to reach the best result the baseline achieved over its whole search.Underlined values explain themselves when clicked.

In an actual implementation

Reading the ablation

Table 12 removes one component at a time from MoA-Pruner in online mode on TITAN V. On ResNet-50, Ansor reaches 2.019 ms; MoA-Pruner reaches 1.782. Removing the draft stage gives 1.995, almost back to Ansor, because the learned model alone cannot find good candidates quickly and its inference overhead returns. Removing temporal dataflow features gives 1.930, worse than removing statement features (1.863), so the new feature branch carries more of the accuracy. Replacing momentum adaptation with plain online fine-tuning gives 1.812, and dropping adaptation entirely gives 1.828. In offline mode, where a well-trained model already exists, keeping the draft stage still cuts compile time from 111 to 89 minutes on ResNet-50 and improves latency, because the learned model’s feature extraction and GPU inference are expensive even when accurate.

Where this stops working

The draft formula is only as good as its hardware abstraction; it was written for GPUs with a three-level memory hierarchy and needed extensions for tensor cores, and a different accelerator class needs new symbols and penalties. Its filtering also assumes the optimum has good hardware fitness, which the TenSet study supports for tiled operators but which is not proven for every kind of program. Pruner searches inside TVM’s tiling-rule space, so it inherits that space’s limits: on operators with long reduction axes and short parallel axes, cuBLAS and TensorRT’s split-K kernels remain faster. The reported speedups are search-time-to-match figures against baselines that were themselves given 2,000 trials, and several baselines could not complete every workload, so the comparison set is narrower than the full suite. Finally, MoA needs a pre-trained model to adapt from; it improves online tuning but does not create a cost model where none exists.

Background

Why this is speculative decoding for compilers

Draft-then-verify is the shape of speculative decoding in language-model inference: a small cheap model proposes several tokens, and the large model checks them in one pass, so the expensive model is called far less often while the output stays what it would have chosen. Pruner applies the same division to schedule search. The Symbol-based Analyzer is the small draft model, fast and approximately right; the Pattern-aware Cost Model is the large verifier, accurate and expensive; and the quality guarantee comes from measuring that the draft set almost always contains what the verifier would have picked from the full set. The idea transfers because the bottleneck is the same: an accurate model whose per-call cost dominates when it is applied to everything.

If you remember only three things

1. Search-based tensor compilers pick which of thousands of candidate schedules to measure by asking a learned cost model, and applying that model to every explored candidate is close to 40 percent of tuning time on a Jetson Orin. A model trained on one GPU cannot be used online on another, and the first noisy rounds of measurements on a new platform can wreck training from scratch.
2. Pruner drafts with a Symbol-based Analyzer, an empirical formula over eight hardware-aware symbols (allocation, computation and parallelism at register, shared-memory and global-memory levels) turned into utilization penalties, using it as the fitness of a genetic search that keeps 512 of about 8,000 candidates. It verifies with a Pattern-aware Cost Model, a multi-branch Transformer over 23-dimensional temporal dataflow features extracted from the multi-level tiling pattern plus Ansor's statement features, which ranks the 512 and selects 10 to measure. MoA-Pruner initializes the target model from a Siamese model pre-trained on other hardware and updates the Siamese weights toward the target with momentum 0.99.
3. Across A100, TITAN V and Orin with 2,000 trials, Pruner and MoA-Pruner reach Ansor's final result 2.6x and 4.82x faster on average; offline, Pruner is 4.75x faster than TenSetMLP and 4.05x faster than TLP; on tensor cores it is 4.08x faster than MetaSchedule. PaCM reaches top-1 of 0.892 on TenSet T4 against 0.862 for TLP, compile time falls to 75 to 84 percent of Ansor's, and MoA-Pruner at 2,000 trials beats Ansor at 10,000 on ResNet-50 in one eighth of the time. Hand-tuned split-K kernels still win on long-reduction shapes, and the draft formula must be rewritten for new accelerator classes.

Check yourself

Check your understanding

1. Why does Pruner run its learned cost model on only 512 candidates per round instead of all 8,000 explored?
2. What is a hardware-aware symbol, and why is it cheap to compute?
3. Predict what happens to search speed and result quality if the draft stage is removed but everything else is kept.
4. Why does MoA update the Siamese model with momentum instead of simply fine-tuning the pre-trained model online?
5. On which kinds of operators do hand-tuned libraries still beat Pruner, and why?

Sources

The paper is “Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning” by Liang Qiao, Jun Shi, Xiaoyu Hao, Xi Fang, Sen Zhang, Minfan Zhao, Ziqi Zhu, Junshi Chen, Hong An, Xulong Tang, Bing Li, Honghui Yuan and Xinyang Wang, University of Science and Technology of China, University of Pittsburgh and NIO, published at ASPLOS 2025, Volume 2, DOI 10.1145/3676641.3716269. This explainer was written from the full text of the arXiv preprint (2402.02361, version 3, April 2025). Numbers are taken from Tables 1, 5 and 7 to 13 and Figures 6 to 16. The code is open source at github.com/qiaolian9/Pruner.