AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures

Modern ML models spend most of their GPU time not in matrix math but in memory-bound glue operations, which compilers either fuse badly (recomputing values hundreds of times) or refuse to fuse (launching kernel storms) — AStitch stitches whole subgraphs into single kernels by giving intermediate values a home at the right level of the memory hierarchy.

Why this problem is worth solving

Machine-learning graphs contain two species of operator. Compute-intensive operators — matrix multiplication, convolution — do heavy arithmetic per byte and have received a decade of optimization attention. Memory-intensive operators do trivial arithmetic per byte: element-wise operations (add, multiply, but also expensive ones like tanh, power, log), reductions (summing a tensor dimension), and broadcasts (replicating values across a dimension). They are bounded by memory bandwidth, not compute.

The paper’s opening measurement rearranges the agenda: across five production models on a V100 GPU — BERT and Transformer (language), DIEN (recommendation), ASR and CRNN (speech and character recognition) — memory-intensive operators average 63.2% of execution time and 89.6% of all launched kernels. And the trend is against them: GPU arithmetic throughput grows faster than memory bandwidth (a 5.6× ratio jump from V100 to A100), pushing that time share to a projected 76.7% on A100. Attention- style architectures did this: a Transformer contains 1,666 reduce operators alone, roughly 10% of its operator count, woven into arbitrary topologies.

The costs come in three currencies: off-chip memory traffic (every operator writes its output to DRAM and the next reads it back), kernel-launch and framework scheduling overhead (thousands of tiny kernels), and CPU-GPU context switching. Kernel fusion — compiling many operators into one GPU kernel — is the standard cure, and for simple chains it works. The paper is about why it stops working exactly where modern models live.

What people did before, and where it broke

XLA and TVM fuse memory-intensive operators with one code-generation strategy: inline the producer into the consumer, per element, through per-thread registers. Each thread computes its own output element and recomputes whatever producer values it needs. This handles one-to-one element dependencies (pure element-wise chains) perfectly. It shatters on one-to-many dependencies, which arise from precisely two patterns that saturate modern models: a reduction feeding consumers (the reduced value feeds many downstream elements) and an expensive element-wise operation feeding a broadcast (one tanh/power result replicated across a dimension).

The dilemma is exact:

• Fuse anyway → redundant computation. With register-only communication, every consumer thread independently recomputes the shared value. The paper’s example: fusing power → broadcast → add where 128 add-elements need each power-result makes every power value be computed 128 times in 128 different threads — and power is expensive. Fusing across a reduction is worse: each consumer thread would redo the whole reduction. TVM fuses the broadcast pattern and eats this cost.
• Skip fusion → kernel storm. XLA declines to fuse both patterns, which is why a Transformer generates roughly 3× more memory-intensive kernels than compute-intensive ones, each launch paying scheduling and context-switch overhead, each boundary spilling intermediates to DRAM. Operator-level one-to-many dependencies add insult: a producer with two consumers in different kernels gets inlined — recomputed — into both.

A human CUDA expert would simply buffer the shared value in shared memory (the fast on-chip scratchpad visible to a thread block). But an expert hand-codes one kernel; a compiler must decide automatically, for arbitrary just-in-time subgraphs of hundreds of operators, where every intermediate should live and how threads should be arranged — jointly, because locality choices constrain parallelism choices.

The second, independent failure is irregular tensor shapes. Production shapes are not model-zoo shapes. Row-reducing a 750000×32 tensor (a real DIEN case) makes XLA launch 750,000 thread blocks of 32 threads — blocks too small to fill the machine. Row-reducing 64×30000 (a real Transformer case) makes 64 blocks of 1024 on a GPU that can run 160 concurrently — too few blocks. Both are parallelism disasters from the same rigid mapping of one row-task to one block.

The idea

Four stitching schemes: a home for every dependency

AStitch’s core abstraction assigns each operator one of four stitching schemes, jointly determined by dependency type, memory level, and the locality-versus-parallelism trade:

• Independent: no dependency — no communication.
• Local: one-to-one element dependencies — values pass through per-thread registers (this is all XLA/TVM had).
• Regional: one-to-many dependencies — values are buffered in shared memory, with the thread mapping arranged so producers and consumers of a region land in the same thread block (locality first).
• Global: any dependency at all — values are buffered in global memory, visible to every thread, with a lightweight in-kernel global barrier replacing what would otherwise be a kernel boundary (parallelism first).

This enables hierarchical data reuse at two levels: element-level (a value needed by many threads is computed once and read many times) and operator-level (a producer with many consumer operators is computed once, not inlined into each). The dilemma dissolves: fusion no longer implies recomputation, so it never needs to be skipped. A subgraph that XLA splits into four kernels — with intermediates round-tripping through DRAM and shared parameters loaded twice — becomes one kernel with three barriers, shared-memory handoffs, and single loads. The global barrier’s price, measured: at most 2.72 microseconds, cheaper than the kernel launches it replaces, but valid only if all blocks are resident simultaneously — a constraint the thread mapping must honor.

Adaptive thread mapping: fit the threads to the shape

Irregular shapes get a task packing and splitting toolkit for reductions. Horizontal packing fuses many small row-tasks into one big thread block: the 750000×32 case packs 32 rows into each 1024-thread block, fixing the small-block-size disease. Vertical packing stacks multiple blocks’ work into one block, shrinking block count to fit within one resident wave — exactly what the global barrier requires. Task splitting goes the other way: the 64×30000 case splits each row across several blocks with cross-block atomics, curing the small-block-count disease. The mapping is chosen just-in-time from the actual shapes.

Making it automatic: dominants and propagation

Choosing a scheme and mapping per operator would explode combinatorially. Two observations collapse the space. First, an operator with one-to-one dependencies (Local scheme) can inherit its thread mapping from its consumer by index propagation — so mappings need deciding only at a few dominant operators (reductions, expensive-ops-feeding-broadcasts, outputs), then flow outward through element-wise territory. Second, dominants connected only through element-wise paths can merge, keeping one schedule-setter (preferring the reduction, whose cost most deserves the best schedule). What remains per dominant is Regional-versus-Global: a passive block-locality check computes whether producer and consumer blocks align on contiguous regions — shared memory if yes, global if no — while pure element-wise groups proactively reshape their mapping to match their producer’s locality. Reductions prioritize parallelism; element-wise groups prioritize locality; the compiler encodes the expert’s instincts as rules.

kernels: 4 | intermediates: via DRAM
parameter.2, broadcast.2: loaded twice

Why it works

Because the fusion dilemma was never about fusion decisions but about code generation — registers being the only communication channel — adding shared-memory and global-memory channels with matching barriers removes the dilemma at its root: one-to-many values are computed once and read many times. → Because the in-kernel global barrier costs microseconds while kernel boundaries cost launches, scheduling, and DRAM round-trips, even the “worst” stitching scheme beats splitting kernels. → Because thread mappings propagate through one-to-one dependencies, deciding schedules only at dominant operators makes whole-subgraph optimization tractable just-in-time — no autotuning, which production customization cannot wait for. → Because task packing and splitting decouple logical row-tasks from physical thread blocks, irregular production shapes get full-machine parallelism where rigid one-row-one-block mappings starved or flooded it. → Therefore whole memory-intensive subgraphs become single, well-shaped kernels: 1.84× average over XLA (to 2.73×), 2.47× average over TensorRT on inference, kernel counts halved against Ansor — and the gains compound as hardware’s compute-to-bandwidth ratio keeps rising.

What it costs

• Bigger kernels manage scarcer resources. Stitching hundreds of operators into one kernel pressures registers and shared memory; the resource-aware launch configuration bounds register usage and block counts, trading some per-thread headroom for the barrier’s validity.
• Global barriers constrain occupancy. All blocks must be resident in one wave, capping block counts — vertical packing exists to pay this cost, and it serializes work that separate waves would pipeline.
• The heuristics encode today’s patterns. Dominant selection, the locality checks, and reduce-first preferences are rules distilled from the two problematic patterns; new operator species would need new rules (the authors frame the schemes as general, the rules as extensible).
• Scope is memory-intensive subgraphs. Compute-intensive operators keep their vendor kernels; AStitch optimizes the glue, so models dominated by GEMM see proportionally less benefit (training speedups, at 1.30× over XLA, are smaller than inference’s 1.84×).
• Framework coupling. The implementation rides TensorFlow/XLA’s infrastructure; the ideas port, the artifact is TensorFlow-specific.

The evaluation

Five production models (CRNN, ASR, BERT, Transformer, DIEN) with production batch sizes, on a V100 GPU (T4 and mixed-precision results additionally reported), against TensorFlow 1.15, XLA, and TensorRT 7.0 for inference; plus a BERT case study against TVM’s Ansor auto-scheduler (2,000 tuning trials) and a 10-month production deployment.