Optimal Kernel Orchestration for Tensor Programs with Korch

Running a neural network on a GPU means deciding which operations get bundled into which GPU kernels, and every framework does it the same way — greedily merge neighbouring operators according to rules an expert wrote down; Korch (CMU, Illinois and Tsinghua) argues the operator boundary is the wrong unit entirely, since operators were named for their mathematical meaning rather than how they run, so it first breaks each one apart into primitives with uniform parallelism, then hands the packing problem to an integer programming solver — and lets the solver duplicate work, because arithmetic has got cheaper relative to memory much faster than anyone updated the rules for.

Why this problem is worth solving

The decision being made here is invisible in the source and dominant in the runtime.

What Korch changes is the unit those decisions are made about.

What people did before, and where it broke

• Apply expert-written fusion rules greedily. The mainstream approach across every major framework: enumerate which operator pairs may merge, and merge them whenever the pattern matches. Where it broke. Greedy application takes a locally good merge that can foreclose a better arrangement further along, and the rules cannot express anything but combining — so the case where an operator should be split is unreachable by construction.
• Classify operators by dependency type and write rules for the categories. A refinement that generalises from specific operator pairs to categories based on how output elements depend on input elements. Where it broke. It broadens which combinations a rule can cover without changing the granularity or the greediness. The unit is still the operator.
• Assume the operator boundary is a sensible execution boundary. Implicit in all of the above. Where it broke. Operators are named for their algebraic semantics. Ones with richer mathematics — normalisation, aggregation — bundle several kinds of computation with different parallelism, and running all of it in one kernel is often measurably worse than splitting it up.
• Partition the computation with no duplication. Every prior optimiser assigns each piece of work to exactly one kernel, which seems obviously right — why compute anything twice? Where it broke. It was right when arithmetic was the scarce resource. GPU floating-point throughput has grown much faster than memory bandwidth, so recomputing a cheap operation inside two kernels can beat computing it once and moving the result between them.
• Search the fusion space with heuristics. A reasonable middle ground. Where it broke. Nobody had a formulation of the problem clean enough to solve exactly. The contribution here is as much the formalisation — legal kernels are convex subgraphs, and the choice is a covering problem with dependency constraints — as it is any particular search.
• The gap. Two assumptions were never questioned together: that the operator is the right unit, and that each piece of work belongs to exactly one kernel. Dropping both turns an ad-hoc rule-matching exercise into a constrained optimisation with a solver.

How Korch builds a kernel schedule

The idea

Decompose before you compose, because the operator boundary was chosen for mathematics. This is the reframing and everything follows from it. A tensor operator is named for what it means — softmax, layer normalisation, attention — not for how it executes, and the richer ones bundle together computations with genuinely different characters. A softmax is elementwise work, then a reduction, then a broadcast, then more elementwise work, and a GPU treats those four things very differently. Insisting they share a kernel because they share a name is arbitrary, and it also blocks the interesting move: fusing part of one operator with part of the next. The paper’s evidence for this being the real lever is a neat ablation — feeding the decomposed graph to an unmodified commercial framework, changing nothing else, gives a 1.24 times speedup, so a quarter of a win comes from granularity alone. Discovering that the fusion space defined at operator granularity is too narrow for memory-bound work is the same finding as the memory-intensive fusion work covered elsewhere on this site, reached from a different direction.

Once the granularity is right, the choice is a covering problem, so stop guessing. With operators split into primitives, the question of which primitives share a kernel has clean structure. A set of primitives can legally share a kernel exactly when it forms a convex subgraph — no path may leave the set and come back, or the two kernels would deadlock waiting for each other. Korch enumerates all such sets by searching over reachable execution states, then does something unusual: it prices each candidate by compiling and timing it rather than by estimating. With measured costs in hand, choosing the set of kernels becomes a binary linear program with two constraint families, one saying every required output gets produced and one saying no kernel runs before its inputs exist, and an off-the-shelf solver finds the optimum. Turning a search over a large space of equivalent programs into a formal optimisation an existing solver can discharge is the same move as the differentiable extraction work covered elsewhere on this site, arriving at the opposite answer about whether to solve exactly.

Let the solver do the same work twice, because hardware balance changed. The assumption every prior optimiser makes is that the computation gets partitioned — each piece of work assigned to exactly one kernel, no duplication. That is correct when arithmetic is expensive. But GPU floating-point throughput has grown considerably faster than memory bandwidth, so the calculus has quietly inverted: recomputing a cheap layout change inside three separate kernels can beat computing it once and moving the result through memory twice. Korch simply drops the constraint and lets the solver decide, and the case study shows it exercising that freedom, executing a reshape-and-transpose sequence three times to eliminate kernel boundaries. The paper is careful to distinguish this from recomputation during training, which trades more memory traffic for less peak memory — here the duplication reduces traffic rather than increasing it. Recognising that recomputing can be cheaper than storing and reloading is the same economics as the adaptive recomputation work for pipeline parallelism covered elsewhere on this site.

Walking an attention block through Korch

The categories matter more than the count. Grouping primitives by parallelism and access pattern is what lets the later stages reason about which ones combine well — memory-bound primitives fuse freely, while a linear transformation generally wants a kernel to itself so the vendor library can be used.

Why it works

• Splitting first makes fusing across operator boundaries possible. An operator bundles computations with different parallelism because it was named for its mathematics → decomposing into uniform primitives lets part of one operator join part of the next → and the ablation shows a quarter of the total win comes from this granularity change alone, before any solver runs.
• Convexity is exactly the right legality condition. A kernel cannot contain two primitives with a path between them that leaves the set → because the two kernels would then wait on each other → so the legal candidates are precisely the convex subgraphs, which is enumerable by searching reachable execution states.
• Measured costs beat modelled ones for a discrete choice. Each candidate kernel is compiled and timed → so the solver optimises over real latencies rather than a cost model’s approximation of them → and candidates that cannot be generated at all simply drop out with infinite cost.
• The no-duplication assumption encodes obsolete hardware economics. It was right when arithmetic dominated → but floating-point throughput has outgrown memory bandwidth, so recomputing a cheap primitive can cost less than moving its result → dropping the constraint lets the solver find arrangements a partitioning formulation cannot represent.
• Real networks are deep and narrow, which tames an exponential. Candidate enumeration grows exponentially with graph width → but modern architectures are long chains rather than wide fans, and simple rejections remove most candidates → so a theoretically intractable enumeration finishes in practice.

What it costs (trade-offs and scope)

• Tuning takes hours. End-to-end tuning runs from around 9 to over 12 hours per model, dominated by compiling and timing every candidate kernel. One single kernel took 12 hours to tune by itself. The paper argues this is a one-time cost and parallelisable across GPUs, which is fair, and it is a very different proposition from a compiler that finishes in seconds.
• The enumeration is exponential in graph width, and the paper says so. It works because current networks are deep and narrow. The future work section names this as the barrier to applying the approach to general tensor programs, and proposes a cost model to discard candidates cheaply.
• Compute-heavy operators are not actually optimised. Matrix multiplications and convolutions are handed to vendor libraries on the reasoning that automatically generated kernels rarely beat them within reasonable tuning time. So the optimisation is really about the memory-bound work and the boundaries around the compute-bound work.
• The code generator is the weak link and the authors point at it. Their evaluation shows the generator underperforming a highly optimised commercial framework on the newer GPU, which they note leaves room for improvement — and it partly explains why the gain is smaller on the newer hardware.
• Averages are well below the headline. Up to 1.7 times and up to 1.6 times are the maxima; the averages are 1.39 and 1.30. On conventional convolutional networks the average is 1.31.
• The formulation has stated gaps. Only candidate subgraphs with a single output are considered, which the future work section identifies as an unnecessary restriction. Data layout is not a variable in the optimisation, even though the case study shows layout being worth 3.52 times on one kernel — it is currently exploited only incidentally through which primitives end up together.
• Five models, one batch size, two GPU generations. Three convolutional networks and two vision transformers, evaluated at batch size one except where noted.
• The optimality claim is bounded by the enumeration. The solver finds the best covering among candidates that survived enumeration and rejection heuristics, using measured costs. That is a meaningful guarantee and it is not optimality over all possible ways to execute the network.

The evaluation

Check yourself

Sources

The paper is “Optimal Kernel Orchestration for Tensor Programs with Korch” by Muyan Hu, Ashwin Venkatram, Shreyashri Biswas, Balamurugan Marimuthu, Bohan Hou, Gabriele Oliaro, Haojie Wang, Liyan Zheng, Xupeng Miao, Jidong Zhai and Zhihao Jia, published at ASPLOS 2024 Volume 3, pages 755-769, DOI 10.1145/3620666.3651383. The authors are at Carnegie Mellon University, the University of Illinois at Urbana-Champaign, Tsinghua University and industry. The implementation is open-source.

This explainer is written from the full text, read from the authors’ copy on arXiv (2406.09465). Everything above comes from the paper’s own sections 3 through 8: the argument that operator-level fusion is too coarse because operators are defined by algebraic semantics rather than implementation; the fission rules producing elementwise, reduce and broadcast, layout transformation and linear transformation primitives, with the softmax decomposition given as the worked example; the primitive-graph transformations converting a reduction into a matrix multiply against a vector of ones, moving an elementwise division past a subsequent multiply, and merging two multiplies with a shared input; the convex subgraph legality condition with its circular dependency justification, and the depth-first search over execution states used to enumerate candidates; the binary linear programming formulation with its output and dependency constraints, and the deliberate relaxation permitting a primitive to be executed multiple times, justified by floating-point throughput outgrowing memory bandwidth and distinguished from training rematerialisation; the implementation choices routing compute-intensive primitives to vendor libraries and memory-intensive ones through an auto-tuning code generator; the end-to-end results of up to 1.7 and 1.6 times with averages of 1.39 and 1.30 across two GPU generations; the adaptation study feeding the primitive graph to an unmodified commercial framework for 1.24 times; the attention block case study with its 3.29 times subgraph speedup, twelve kernels reduced to seven, three-fold duplication of a layout sequence, and 3.52 times from a layout change under the same vendor routine; the tuning time table showing 9.2 to 11.5 hours and a longest case of 12.2 hours; and the future work section naming exponential width scaling, the single-output restriction, and data layout as unaddressed.