Optimizing Dynamic-Shape Neural Networks on Accelerators via On-the-Fly Micro-Kernel Polymerization

An accelerator has a fixed number of processing units, so a large matrix multiplication runs in WAVES — and if the work does not divide evenly, the final wave is nearly empty while still costing a full wave, which is why the same vendor library reaches 84 percent of an A100's peak on one shape and 7 percent on another; the shapes that matter are now decided at runtime by sentence length or image resolution, far too late for the hours-long search that compilers use to pick tile sizes; this paper (Chinese Academy of Sciences and UNSW Sydney) tunes about forty fixed-size tiles offline, then at runtime CUTS the output into regions and gives each region a DIFFERENT tile size, so the ragged remainder is made of cheap work rather than expensive work — a choice made in roughly two microseconds that lands within 4 percent of an exhaustive search taking 1.6 seconds.

Why this problem is worth solving

The gap between what a shape should cost and what it does cost is far larger than most people expect.

The mechanism behind those collapses is worth understanding precisely, because it is what the paper attacks.

What people did before, and where it broke

• Hand-written vendor libraries. Every vendor ships highly tuned routines, and they are excellent on the shapes they were tuned for. But a routine holds a few implementations, and prior work established that no single carefully designed implementation suits all shapes.
• The measured consequence. The same library delivers 84 percent of an A100’s peak on one compute-bound shape and 7 percent on another. Being slow on unusual shapes is not a rare corner case once shapes are chosen by the input.
• Search for the tile sizes. Tensor compilers explore an enormous space of tiling choices for a given shape, measure candidates on real hardware, and emit an implementation tuned to it. This is the strongest known approach and routinely beats hand-written code.
• But the search costs about a third of a CPU hour. That is a fine price to pay once, offline, for a shape you will run millions of times. It is not a price you can pay while a user waits for a sentence to be processed.
• Pre-tune across a declared range. DietCode and Nimble ask the developer to state the range of each dynamic dimension, generate a family of programs offline each covering part of it, and pick one at runtime. This removes the online cost and is a sensible engineering answer.
• But a declared range is a promise about the future. Outside it the system fails rather than degrading: with a range of 1 to 128 and 8,192 shapes presented, only 254 runs completed, the others hitting out-of-bounds errors or unavailable resources. Widening the range fixes that and costs performance inside it, since one program covering more shapes fits each less well.
• The gap. Everyone was choosing between tuning offline for shapes you must know, or paying an unaffordable search online. Nobody had built a system where the expensive part is done offline WITHOUT knowing the shape, leaving only a cheap selection online.

How the work splits between offline and online

The idea

Move the expensive work to where the shape is not needed. The deadlock in this field was that tuning requires knowing the shape and the shape arrives too late. The escape is to notice that not every part of the tuning depends on the overall shape. How to arrange the innermost loops so a tile of data lands in local memory and is reused there is a question about the TILE and the HARDWARE, not about how many tiles there will eventually be. So that half can be searched exhaustively offline, exactly as a static compiler would, producing a set of tuned fixed-size micro-kernels. What genuinely depends on the shape is only the outer structure — how many instances, arranged how — and that turns out to be answerable by arithmetic rather than search. Splitting a problem so the costly analysis lands on the shape-independent part is the same separate-what-varies-from-what-does-not discipline covered elsewhere on this site where an expensive offline analysis is done once against stable structure and only a cheap lookup remains on the request path.

Make the leftover work cheap instead of trying to eliminate it. The mechanism is more surprising than “pick a better tile size”. Because processing units are finite, work runs in waves, and a shape that overflows a wave boundary by a little pays for a whole wave that is nearly empty. The obvious response is to choose tiles that divide the shape evenly, minimising waves. This paper does something better: it cuts the output into REGIONS and gives each a different micro-kernel, using a large one for the bulk and a small one for the remainder. In the worked example this produces THREE waves where the single-kernel version needed two — and runs 1.21 times faster, because the extra wave is made of short tasks rather than long ones. The tail was not removed; it was made cheap. That inversion, where the design deliberately does worse on the metric everyone watches because that metric had stopped predicting time, is the same lesson covered elsewhere on this site where a system accepts more of the quantity being minimised and wins decisively on the quantity that matters.

Predict instead of measure, and check how much you gave up. Selection online must be nearly free, so measurement is impossible and a cost model is required. Theirs is deliberately simple: for each region, multiply the number of waves needed by the predicted cost of one pipelined task, and add up the regions. Both halves matter, which the paper demonstrates by breaking them. A version using only the wave count picks large micro-kernels and reaches 81 percent of an exhaustive search; a version using only the pipeline cost picks small ones and reaches 72 percent; the combination reaches 96 percent. The comparison against exhaustive search is the honest way to report this, because it separates “our model is good” from “our search space is good”. Two microseconds for 96 percent of what 1.6 seconds achieves is a clean statement of the trade. Replacing an unaffordable measurement with a cheap estimate and then quantifying the loss is the same bound-what-you-cannot-count discipline covered elsewhere on this site where an approximation is chosen precisely because it keeps the decision tractable, and its accuracy is reported rather than assumed.

Cutting a matrix multiplication in two

occupancy is limited because each block holds a large tile of data in on-chip memory, which restricts how many blocks fit per multiprocessor
the operation is not finished until its last block is, so a wave costs full price however few blocks are in it

Why it works

• Tuning the innermost loops depends on the tile and the hardware, not on the operator’s overall shape → that half of the search can be done offline without knowing anything → so the 0.33 CPU hours are paid once, in advance, instead of on the request path.
• Micro-kernels are compiled offline into binaries whose addresses and trip counts are parameters → the online step selects and parameterises rather than generating code → costing about 2 microseconds per shape rather than hours.
• Finite processing units mean work runs in waves, and a nearly empty final wave costs nearly a full one → giving the remainder a smaller, faster micro-kernel makes that wave cheap → so three waves of mixed work beat two waves of uniform work by 1.21 times.
• The cost model multiplies waves needed by predicted cost per pipelined task → neither factor alone is sufficient, since one favours large tiles and one small → and combining them reaches 96 percent of exhaustive search against 81 and 72 percent for each half alone.
• Local padding guarantees a valid micro-kernel combination for every shape → no shape range has to be declared in advance → producing 8,192 valid runs out of 8,192 where a range-based system managed 254.

What it costs (trade-offs and scope)

• The benefit depends on the shape being awkward. Gains are largest for small shapes and for operators whose shapes vary frequently over a wide range, because that is when the wave imbalance is pronounced. On a large language model the operator gains are 1.08 to 1.24 times and end-to-end gains fall from 1.05 to 1.01 times as batch size grows from 1 to 8, since bigger batches make shapes rounder.
• Knowing the range would still help. The authors note that when shape ranges ARE known at compile time, a better-matched micro-kernel set and a refined cost model would do better. Giving up that knowledge is a real cost, paid for robustness.
• The online set is fixed and pruned. Forty micro-kernels are kept out of a much larger generated space, chosen by average performance on synthetic workloads. A shape poorly served by all forty has no recourse, and the hyperparameters were set empirically per platform.
• It optimises operators, not graphs. Operator fusion and other graph-level optimisations are listed as future work, so this composes with rather than replaces those gains.
• Convolution goes through matrix multiplication only. Alternative algorithms such as Winograd are not implemented, which the authors flag; for fairness they also force the baseline libraries to use the same route.

The evaluation

Testing covered two quite different accelerators — an NVIDIA A100 with Ampere Tensor Cores, and a Huawei Ascend 910 with Da Vinci Cores — using each vendor’s own library as the baseline. The shape sets were large: 1,433 matrix multiplication shapes from DeepBench and from real models, and 5,485 convolution shapes. End-to-end tests replaced the operators inside eight models, driven by 150 sentences of length 5 to 500 for the language models and by eight batch sizes and ten resolutions for the vision models. Importantly, the reported end-to-end times INCLUDE the cost of running the selection model, not only the kernels it picks.

Check yourself

Sources

• Optimizing Dynamic-Shape Neural Networks on Accelerators via On-the-Fly Micro-Kernel Polymerization (ACM Digital Library, DOI 10.1145/3620665.3640390) — ASPLOS 2024; the full 16-page paper, read in full for this explainer. Published open access under a Creative Commons licence.