Tang, a compiler that derives GPU kernels and bounds how fast they can run

A compiler called Tang derives GPU kernels from a tensor IR built on affine views and monoid comprehensions, generating code for WGSL, CUDA tensor cores, Metal simdgroup matrices, OpenCL, GLSL and C without hand-written per-operator kernels. On an idle RTX 3090, a derived transformer block's forward pass at sequence length 512 ran 1.2× faster than PyTorch captured in a CUDA graph with TF32, while PyTorch was 1.2× faster at length 128; Qwen3-8B decoded as a single IR graph at 148.7 tok/s, tied with llama.cpp at short context and at 92% of it at 4k. The authors also prove weak duality in Lean for a linear-program bound on memory traffic, and a time floor built on that bound never exceeded measured time on the 24 GEMV-class and epilogue kernels validated on two GPUs.

Abstract Every fast model today runs on kernels someone wrote by hand, one per operator, per chip and per fusion. Each one freezes choices about what to fuse, where data lives and how threads share work, and those choices are right for one shape on one device. We think kernels should be derived instead. tang’s tensor IR is built from two constructs. Views are affine maps from an iteration space to tensor coordinates. Comprehensions are reductions under a monoid whose body is a scalar expression graph. Everything else is either a rewrite or a schedule. Fusion is one substitution rule. We do not write a flash attention kernel. The compiler derives its schedule by recognizing online softmax as a monoid. Gradients of comprehensions are comprehensions, so the same passes derive and optimize the backward pass. That includes a flash-style attention backward that writes no intermediates on the three GPU backends we measured on CUDA through a fused kernel that computes in bf16 and so misses an f32 tolerance, and on WGSL and Metal at a cost we report . The recent superoptimizers we discuss derive only the forward. Precision and thread mapping are schedule fields that cannot change the math precision changes only rounding, which an emulator budgets with known failure modes , chosen per kernel by measurement on the device. One graph lowers to WGSL, CUDA tensor cores, Metal simdgroup matrices, OpenCL, GLSL and C, from a datacenter GPU and a desktop GPU and a laptop down to a humanoid robot’s phone-class SoC. On an idle RTX 3090 a derived transformer block’s forward pass at sequence length 512 is 1.2× faster than PyTorch captured in a CUDA graph with TF32; at length 128 PyTorch is 1.2× faster, and forward plus backward ties. Qwen3-8B decodes as a single IR graph at 148.7 tok/s, tied with llama.cpp at short context and at 92% of it at 4k. We wrote no kernel for any one operator. A short list of emitters is keyed on the shape of a graph rather than on an operator: CUDA’s flash kernel the online-softmax monoid , the fused Hopper attention backward, the narrow GEMV matrix–vector product stage, the integer tensor-core stage and the chunked scan plan. Derived attention is not uniformly fast. At the forward trails FlashAttention-2 at every size we tested, and on the RTX 3090 no tensor-core forward is generated at , which makes it about 150× slower than FlashAttention-2 there. We then derive the kernels of a hand-tuned engine that serves an 80B mixture-of-experts model on one H200. With the work made equal, the derived MoE expert layer is 18% slower than the engine’s hand kernels at batch 1 and 7% slower at batch 4, ties at 8 and is 5% faster at 64. Inside the model, where the derived path runs only from 16 window rows up, a decode step takes 5–12% less time at 16 to 48 concurrent streams over five interleaved rounds on one H200, though the derived path misses the engine’s 48-stream agreement gate, 98.93% against a 99% bar. The derived head-dim-256 attention and gated-delta scan cut the engine’s own time to first token at 44k tokens from 8.33 s to 5.54 s, but vLLM on the FP8 checkpoint reaches first token about 12× faster than the engine, so this is a gain over our own baseline, not over the state of the art. At the derived attention runs at 0.43–0.54 of FlashAttention-3′s speed forward and 0.22–0.40 backward. Finally, we ask how far a kernel is from the fastest any implementation could be. For a comprehension over affine views, the least memory traffic and the tile that attains it are the primal and dual of one linear program, and we prove weak duality in Lean for arrays indexed by subsets of the loop coordinates, along with the soundness of the core rewrites. A time floor built on that bound, with launch cost measured on both GPUs and DRAM latency measured on the H200 and assumed on the RTX 3090, never exceeds the measured time on the 24 GEMV-class and epilogue kernels we validated on two GPUs. On the RTX 3090 it reaches 0.76–0.93 of measured time on the large kernels. On the H200 it reaches only 0.05–0.66 median 0.38 , and we have no validated term for the gap. Inside the schedule space, a two-level version of the program picks the tile the M4 Max autotuner picks, but only for register capacities between 4096 and 8192 words, and the traffic-optimal tile does not win on time. In a timed race of every candidate tile on the M4 Max and the RTX 3090, no candidate ran faster than the floor. A per-tile estimate built on it picks a tile within 2.5% of the fastest on seven of eight kernels and loses the eighth by 1.27×. 1 Introduction A modern model has a few dozen distinct operators. A modern framework has thousands of kernels. Every operator, backend pair needs one. So does every fusion that pays, such as attention with its softmax, a matmul with its bias and activation, or a norm with the projection after it. FlashAttention <sup 1</sup loc-46 is the best-known case. Its large speedup came from rewriting one operator by hand so its score matrix never reaches memory. Its backward pass took a second hand-written kernel. Hopper took a third <sup 2</sup loc-47 , and Blackwell’s FP4 tensor cores are taking a fourth. Each hand kernel is a separate copy of the same arithmetic, and each one freezes a schedule that was right for one shape on one device. We measured how often that schedule is wrong: the fastest schedule changes with the size, the chip and even between runs on the same machine §7.1 . The winning schedule is a property of the graph, the chip and the size, not of the operator. Our thesis is that a kernel is a theorem, not a program . Every fast kernel we know of is a claim of the form “this comprehension, reassociated under this monoid and placed in this memory, equals the naive one over the reals.” Such a claim can be derived, checked and searched. It does not need to be written. It can also be bounded: the least traffic any box-tiled schedule of a comprehension can move, and the tile that moves it, are dual solutions of one linear program. A theorem fixes what is computed. How fast it runs is still a schedule, and §11 lists the schedules tang does not yet have. If that is right, the kernel library should be replaced by an algebra small enough to reason about, together with a search over its schedules. Figure 1 fig-one shows one graph under the schedules that the field knows by name. The traffic column is not the time column. Checkpointing moves a third of the naive schedule’s traffic and is 2.2× slower. The score contraction accounts for more than the whole difference: 7.12 ms against 1.51 ms, with the same arithmetic and the same output. The naive schedule writes in its own small transpose kernel, so a warp’s 32 lanes then read 32 adjacent floats. Checkpointing folds the transpose into the reader’s view, so each lane reads its own row of . Nsight shows the same 67.6 M L1 requests carrying 5.4 GB in one case and 35.5 GB in the other. The output is exact to , so the lowering is correct but loses coalescing, and a traffic model that counts DRAM bytes once per buffer cannot see it. Flash, which also has the least traffic, wins. The backward tells a less flattering story. tang’s fused attention backward kernel computes in bf16 tensor cores even when the graph is f32. On this graph its gradients miss the experiment’s tolerance of , measured as max error over max reference against the f64 interpreter error in , so the search rejects it. Without the fused kernel the searched placement passes at in 1.91 ms, against 0.74 ms for the failing fused version. The search also rejected one placement whose output was simply wrong, at error 7.0; that bug is caught by the check but not yet explained. We propose writing the model once, as math, and letting the compiler derive every kernel. That works only if the IR is small enough to reason about and expressive enough that nothing important falls outside it. tang is a no std Rust library that started as a geometry and physics math crate. Its tang-expr crate traces Rust closures into scalar expression graphs that it can differentiate, simplify and compile, and its tensor IR uses those graphs as its body language. The IR has two constructs. - Views : affine index maps that express broadcast, transpose, slicing, striding, padding and per-head column selection without touching memory. - Comprehensions : , where body is a scalar expression graph and is a monoid. Contributions: 1. Kernels from rewrites. Fusion, flash attention, logsumexp and Welford variance are instances of a small set of rewrites §2 . 2. Derived backward passes. Reverse-mode differentiation maps comprehensions to comprehensions, including the online-softmax monoid. A proof procedure removes the stabilizing max’s gradient, and with nested contractions this yields an attention backward with no tensors §3 . 3. Precision as a schedule. An emulating interpreter checks a format choice against an error budget before it runs, with failure modes we report §4 . 4. One graph, six targets. Lowerings from shared emitters, measured against hand kernels, PyTorch and the bandwidth limit on four chips §5, §7 . 5. A production engine, derived. Top- routing as one monoid, integer tensor-core and narrow GEMV schedules, a segmented-sum gather backward, chunked scans, and Hopper attention schedules. With them, derived kernels replace hand kernels in an 80B MoE engine on an H200, faster in some regimes and slower in others, which we report §8, §11 . A whole derived decode step matches the hand engine on Qwen3-8B and trails it by 13–35% on the 80B MoE. 6. The speed of light. The traffic lower bound and the optimal tile are dual solutions of one linear program read off a kernel’s views, with weak duality in Lean. A time floor built on it never exceeds measured time on 24 GEMV-class and epilogue kernels on two GPUs; it is close on the large kernels of the RTX 3090 and loose on the H200 §9 . 2 The IR 2.1 Views A view maps a point of an iteration space to a source tensor’s coordinates with one affine form per source dimension, . Broadcast is a zero form, transpose a permutation, slice an offset, stride a coefficient. A form with two terms expresses “head reads columns “. Views compose by substitution. They are the only way a comprehension reads its inputs, so every layout decision is a view and every view can be analyzed. A view may be guarded , so that an out-of-range read returns the monoid’s identity zero for sum, for max rather than raising a shape error. Padded, dilated and transposed polyphase convolutions and overlap-add are guarded reads, with no padding operator. The adjoint of a guarded read is itself a guarded read, not a scatter. For , reads through the same guard, so out-of-range padding terms never reach . Grouped-query attention’s is not affine. The IR expresses heads as two dimensions, , and the views do the rest. Dimensions can be fixed or symbolic Dim::Sym "seq" . Construction, rewriting and differentiation work with symbols, and bind makes them concrete for scheduling and lowering. A view may also gather . One source coordinate is then data, read from an index tensor that is a sibling argument of the same comprehension, . MoE routing , paged KV caches and embedding lookups become ordinary reads rather than special operators. The gradient of a gathered read is a scatter-add, written as a dense gated sum, which costs times the needed work unless the keys are declared sorted §8 . An affine map read through a gather still fuses into it, and only nested gathers materialize. 2.2 Comprehensions and monoids The only compute operation is A map is a comprehension with no reduce dimensions. Sum and max are monoids. Matmul is . Attention scores, convolution and gather are the same operation with different views. The monoid’s combine is itself a scalar graph, and a comprehension may carry tuple-valued state. Online softmax is the monoid on , the running max and the rescaled sum, and attention carries the weighted values with it as ; arg-max is the monoid on . 2.3 Rewrites Fusion is one rule. A map is a pure function of its inputs, so substitute maps splices its body into every consumer and composes the views. Map→map, map→reduce and transpose→anything are the same case. The rule is cost-aware, so a map is recomputed inside a reduction only when it is cheap. Reductions are never substituted. They are inlined by a schedule decision, which separates algebra from scheduling. A row-major reshape that only merges dimensions is a view, not a copy. Epilogues run on the reducing thread. A map that reads a reduction point for point a bias, an activation, the that finishes softmax runs on the thread that holds the final accumulator. Sibling reductions merge. Two reductions over the same space and monoid with the same readers become one reduction under the product monoid merge siblings . The gate and up projections of a SwiGLU MLP become one tuple-valued matmul, the activation becomes its epilogue, and the normalized input is read once. Flash attention is a rewrite. stabilize softmax finds and every over the same canonical , with masks and scales allowed, and replaces them with one reduction under the online-softmax monoid. The value column of becomes a point dimension an output dimension, not reduced over of the fused comprehension. Nothing is materialized. Logsumexp, Welford variance and top-k are the same rewrite with different monoids. Figure 2 fig-rewrite shows the attention case. The rewrite is not specific to plain attention. Anything that changes the score or the mask leaves the monoid structure alone, so it derives the same way. The lab below runs tang’s compiler, built to WebAssembly from the same crates as this paper, in your browser. Type a variant: it builds the naive graph, runs the real pass pipeline, then runs both programs on your GPU through WebGPU and checks them against tang’s f64 interpreter. 2.4 Rediscoveries The rewrites above are general, so known algorithms fall out of them. Table 1 tab-rediscover lists the ones tang derives today. Each row was a paper or a hand kernel. In tang each is a monoid, a view or a schedule field, and none adds an operator. | algorithm | originally | in tang | | online softmax | Milakov & Gimelshein 2018 <sup 3</sup loc-48 | monoid on | | FlashAttention forward | Dao et al. 2022 <sup 1</sup loc-46 | stabilize softmax + staging rule which operands a tile copies into shared memory or registers | | FlashAttention backward | hand kernel <sup 1</sup loc-46 | derived gradient, max elimination, nested contractions §3 | | Welford variance | Welford 1962 <sup 4</sup loc-49 | monoid on | | logsumexp, arg-max | folklore | monoids on and | | top- routing | hand kernel | Monoid::topk , merge of sorted lists | | split-K GEMM | hand kernel | split on any associative monoid | | SwiGLU gate\|up fusion | hand kernel | merge siblings under a product monoid | | implicit-GEMM convolution | hand kernel | guarded views, flattened reduce dims | | paged KV cache | vLLM | gather view | | scatter-add backward, segmented | hand kernel | gated sum + set sorted segment §8 | | chunked SSD / gated delta rule | Dao & Gu 2024 <sup 5</sup loc-50 ; Yang et al. 2025 <sup 6</sup loc-51 | affine scan monoid + chunked plan §8 | 2.5 Schedules A comprehension says what to compute. Schedule fields say how , and they cannot change the math precision changes rounding, §4 : - inline : recompute inside the consumer instead of writing a buffer. - split : one workgroup per point, with threads striding the reduce space and tree-combining. This is valid for any monoid because combine is associative, including online softmax. - tile , tile 2d : blocks of threads, each owning a register tile of points. The emitter tracks every address as an affine form over loop variables. Anything that does not depend on the lane axis is staged in workgroup memory once per reduce chunk. That covers buffer reads K and V tiles, matmul operand tiles and inlined sub-computations the scores . - precision : storage and accumulation formats §4 . schedule auto encodes rules of thumb distilled from measurement. autotune kernels times each candidate mapping per materialized node on the device, using the graph’s real inputs, and caches the winner by device and kernel source. With substitution and the online-softmax rewrite, the staging rule produces the structure of flash attention: rows × lanes per workgroup, cooperative score tiles, and online softmax in registers. 3 Gradients are comprehensions For , the gradient with respect to through a pick view one whose forms are each a single coordinate is found by renaming coordinates. The coordinates that picks become ‘s output dimensions, the rest become reduce dimensions, and every other view is pushed through the same renaming. The body is differentiated by tang-expr ’s symbolic diff , which covers its nine arithmetic operations add, multiply, negate, reciprocal, square root, sine, atan2, , plus a select. A dimension read through any other form an offset, stride, diagonal or broadcast stays an output dimension, and iota , a node that turns coordinates into values, gates the body by . Max gates its term to the arg-max. The backward graph is an ordinary graph, so the same passes optimize it. Rotary embedding and GQA grouped-query attention gradients come from the gated scatter with no special handling. The online-softmax monoid has its own rule. With the forward’s and accumulators read back, , where is key ‘s value, so the flash-fused graph differentiates and its gradients equal the naive graph’s. A stabilizing max has no gradient. That the max can be stopped is standard practice in logsumexp implementations. What is new here is a proof the compiler runs per graph. Softmax is invariant to shifting its scores, so the terms sum to zero, but only in exact arithmetic. While those terms exist, the max must be materialized and masked tiles cannot be skipped. backward drops them when it can prove the loss does not depend on . The proof tracks how each node downstream of scales with . In softmax, has degree 1 in , so does , and so does ; their ratio has degree 0, a function of the scores alone. In general a node is either or with free of , combined by exact identities. <sup 1</sup loc-91 If every path from ends at a degree-0 node, and the gradient is not built. The f64 evaluation of a training step is unchanged largest absolute difference . The flash backward recomputes inside tiles Figure 3 fig-backward . hoist invariant factors factors row-invariant terms out of sums . schedule nested contractions then inlines a short contraction one reduce dimension of at most 128 into every tile that reads it. The attention backward reads and only through and , which the dQ, dK and dV tiles recompute per staged block on tensor cores. On CUDA this removed S, , dP, the max over S and three correction kernels from a transformer training step, reducing 63 kernels to 49 and the summed best-per-kernel time from 1.80 to 1.22 ms at sequence length 512 RTX 3090 . WGSL and Metal now nest the same contractions with scalar dot products in workgroup memory, so no GPU backend we measured CUDA, WGSL, Metal writes anything . For now this costs speed. Forward plus backward runs at 0.68–0.93× the materialized schedule on Metal and 0.25–0.34× under Vulkan on a 3090, until the nested dots are register-blocked the way the mma path blocks them. 4 Precision is a schedule f16 is a different way to run the same program, not a different program. Each node has a storage format, applied where a value is written to memory, and an accumulation format, applied to a reduction’s running state after each combine. The graph stays exact math. The interpreter gives an error budget, with known failure modes. eval is pure f64. eval emulated rounds every stored value and every accumulator step to its scheduled format round to nearest even, subnormals, overflow , with body arithmetic in f64. An ErrorBudget accepts or rejects a schedule before it runs. Emulation matches the GPU’s f16 kernels to 1 ulp when the reduction order is the same. Otherwise the order is part of the error: AdaIN’s variance over with an f16 accumulator has error under a 256-thread split and with one thread per row, while the sequential emulation predicts . Emulation models one sequential fold; a tree-combining split beats it, and the per-row GPU kernel, though also sequential, comes out somewhat worse, which we attribute to its different accumulation order and fused arithmetic rather than having verified it. The prediction can be far off when the reduction length or order differs: a conv1d sum of 3,584 terms measures against a predicted , a ConvNeXt LayerNorm statistic against , and a pointwise matmul against . Those predictions were made on smaller instances, whose shorter sums round less. At equal reduction length matmul, AdaIN prediction and measurement agree within about 1.7×. tune precision lowers buffers to half precision one move at a time and keeps a move only if the program stays within budget and gets faster. 5 Lowering Shared emitters produce the point, split and tiled schedules, register tiles, guarded views, epilogue fusion and tuple monoids once. A dialect switch then handles buffer access and surface syntax. - WGSL portable default, through wgpu on Metal, Vulkan and DX12 . f16 storage is packed two per u32 , because wgpu’s shader compiler naga rejected enable f16 on our targets. - CUDA. Matmul-shaped tiles run on mma.sync TF32, the tensor cores’ 19-bit float mode, m16n8k8, or f16/bf16 m16n8k16 when the operands are half precision with cp.async staging rings. Block shapes are tuning candidates. cuda flash kernel recognizes the online-softmax comprehension and keeps , the row max and the partial sums in mma registers. The score accumulator is reused as the A operand of after permuting keys within each 8-wide step, which is legal because the reduction is commutative and associative. - Metal MSL . Native half and bfloat . Contractions with f16 operands run on simdgroup half8x8 fragments. Reduce dimensions are flattened, so convolution becomes implicit GEMM, with dilated inputs staged once per chunk as a zero-padded, phase-de-interleaved window. - OpenCL and GLSL for Adreno GPUs that have neither CUDA nor Metal. On one Adreno driver, WGSL is routed through GLSL because the direct SPIR-V path was miscompiled. - C , for batch-1 control loops where a GPU round trip costs more than the arithmetic. A comprehension already is a loop nest, so to c emits it directly with a vectorizable innermost point dimension, and tang run performs no allocation. Program::tuned times each kernel’s native lowering against the portable one and keeps the faster. Kernel sources are canonicalized, so the same kernel at another layer is the same string and compiles once. 6 Verification Speed claims depend on every kernel being correct. tang-verify runs each GPU kernel against its CPU oracle on the same input, starting from the device’s own state. Before each kernel runs, it fills the kernel’s outputs with NaN, so unwritten memory cannot pass as a plausible 0. When a check fails, it names the kernel at fault by provenance: which earlier kernel’s correct output the failing buffer contains. We built it after four structural bugs in a hand-written engine slipped past every unit test. One was a byte offset added where a float offset was meant, which wrote k’s values over q. Derived kernels are less exposed to that class of bug, because addresses come from affine forms rather than pointer arithmetic. On the IR path there is a second check, because the f64 interpreter runs the same graph as the GPU. The rewrites themselves are mechanized in Lean 4 with Mathlib, with no sorry proofs/ , listed in docs/proofs.md . Over the reals, view composition, fusion, split and tiled reductions, online softmax as a commutative semigroup the proof has no identity element, and the compiler supplies , its equality with two-pass softmax, shift invariance, gradients of sum comprehensions, the zero gradient of a stabilizing max, and the flash backward identity are theorems. One combine exported from the compiler is proved equal to the verified online-softmax combine. Floating point is not modelled, and guarded and gathered views are outside the fusion theorem. 7 Evaluation Setup. Apple M4 Max Metal , NVIDIA RTX 3090 Vulkan and CUDA 12 , a Qualcomm QCS8550 in a Booster K1 humanoid robot Adreno 740, Cortex-X3 , and, for §8, an NVIDIA H200. PyTorch 2.10 for the shared-GPU runs and 2.8 for the idle-GPU runs, both cu128, timed with CUDA events. Unless a caption says otherwise Table 2 tab-hand includes submit and readback , tang times are kernel-only. Several runs shared the GPU with other tenants. We mark those, and report minima rather than means, because contention only adds time. Derived results are checked against an f64 interpreter, except where a caption says otherwise. The tolerance differs per experiment and is stated with it: to