How does XLA decide which GPU emitter compiles a fusion?

In two places, and the first one is not a decision at all. Every fusion instruction can carry a backend config with a string field named kind, and an earlier pass may have written a name into it: __triton_gemm from the GEMM fusion pass, __triton or __triton_nested_gemm_fusion or __triton_collective from other Triton claims, __cudnn$fusion from a cuDNN pattern matcher, or __custom_fusion for a fusion backed by a pre-compiled kernel. GetEmitterFusionKind checks those strings first and returns immediately, so a stamped fusion is never inspected. Only when there is no stamp does it look at the fusion itself, and then it asks six questions in a fixed order: is there a real reduction hero whose fellow roots have a compatible element count, does the analysis carry a tiled transpose, are there more than one root, is the single root a scatter, is it a sort, does the concatenate emitter accept it. If none of those fires, the last line returns the loop kind unconditionally. The resulting kind is then handed to GetFusionEmitter, a switch that constructs a concrete emitter class for it.

Why is my XLA matmul a Triton kernel instead of a cuBLAS call, or the other way round?

Because of the order two rewrite passes are added to the pipeline, seven lines apart. GemmFusion is added first, inside a condition that requires xla_gpu_enable_triton_gemm, which defaults on, and a compute capability of at least Ampere on NVIDIA or any ROCm card. It rewrites the dots it accepts into fusions stamped __triton_gemm. The GemmRewriter passes are added after it, and they rewrite whatever dots remain into __cublas$gemm or __cublas$lt$matmul custom calls. A pass can only rewrite what is still there, so Triton gets first refusal, cuBLAS takes the rest, and a dot that neither wanted stays an HLO dot and eventually gets a generic loop kernel. That is only the first claim, though: the autotuner later compiles a Triton kernel and a set of cuBLASLt algorithms for the same instruction and keeps whichever is faster, so the pass order decides who gets asked first, not who wins.

What does the XLA GPU autotuner actually do, and why is my first compile so slow?

It is the only pass in the pipeline that compiles real kernels and runs them on the device as part of compiling. It sits in its own pipeline at the very end of GPU optimization, named autotuner, and it does four things in order for each instruction: check the cache, try a cost-model estimate if one is enabled, take the first config that compiles if tuning is disabled, and only then tune. Tuning means enumerating every config every registered backend supports, compiling all of them, running all of them, and picking a winner. The default Triton GEMM config sets in the tree are 25 configs for Hopper, 30 for Ampere and 40 for Blackwell, and cuBLASLt separately enumerates up to 128 algorithms unless xla_gpu_blas_max_algorithms caps it lower. So a model with a dozen distinct matmul shapes is hundreds of full kernel compiles plus timed runs on a cold cache. On a warm cache it is a lookup, which is why the second run of the same program is dramatically faster.

Why did the XLA autotuner pick a config that is not the fastest one it measured?

Because the selector optimises for scratch memory among configs that are statistically tied on speed. PickBestConfig runs two loops. The first finds the minimum duration among configs that neither failed nor sit on an excluded backend. The second computes a limit of that minimum plus a window, defaulting to 2 microseconds, and among everything inside the window keeps the config with the fewest scratch bytes, breaking ties on duration. Scratch bytes are device memory a kernel needs as workspace beyond its inputs and outputs, and a cuBLASLt algorithm can want tens of megabytes of it for the life of the program. So a config 1.5 microseconds slower that needs no workspace will be chosen over the fastest one that needs 32 MiB. When the selection differs from the fastest, a log line at verbosity 2 names both configs and the tolerance, which is the only signal that the tiebreak fired.

How does the XLA autotuner know a kernel produced the wrong answer?

It does not compare against a reference implementation. It clusters the outputs and picks the biggest cluster that contains a trusted member. Every codegen backend declares whether it can produce wrong results, and the candidate list is stably partitioned so the ones that cannot go first. As each candidate runs, its output is compared with the representative of every existing cluster at a relative tolerance, defaulting to 0.1 for GEMMs; a match joins the cluster, a mismatch founds a new one. Once the trusted candidates have run and formed at least one cluster, the untrusted ones are run with new clusters disallowed, so an untrusted kernel that agrees with nothing trusted gets no cluster at all. Everything outside the winning cluster is then marked failed with kind kWrongResults and never reaches the selector. Correctness here is therefore agreement rather than truth, and it only runs at all when xla_gpu_autotune_level is 4 or higher, which is the default. A separate check surrounds every buffer with red-zone padding, 8 MiB by default, and fails any kernel that writes outside its own allocation.

How do I reuse XLA autotuning results instead of tuning on every run?

Set xla_gpu_per_fusion_autotune_cache_dir. The cache is two tiers: an in-memory store for the life of the process, and that directory as an optional second tier. Lookups try the primary first and then the secondary, hits in the secondary are promoted, and inserts go to every writable tier. The directory obeys a mode from xla_gpu_experimental_autotune_cache_mode, so you can tune with it writable on a build machine and deploy with it read-only. Add xla_gpu_require_complete_aot_autotune_results and a cache miss becomes a hard error rather than a silent tuning run on the serving host. The constraint is what a key holds: the device model, the canonicalised HLO of the instruction, and a version. Canonicalisation is why one tuning run covers every identically shaped layer of a transformer, and the device model is why a cache tuned on one accelerator does not transfer to another. One caveat is stated in the source: the key used to exchange results between sharded compilations hashes the module fingerprint and the backend names but not the module config, so compiling the same module twice with different debug options can recover the earlier, possibly inferior, results.

The GPU Codegen Path: From a Fused HLO to a Kernel

When the optimizer stops, your model is a few hundred fusions and a handful of library calls, and none of them is machine code yet. This reads the two decisions that finish the job, at a pinned commit of openxla/xla. First the kind: a string an upstream pass stamped on the backend config, or, when there is no stamp, ten questions asked about the fusion roots in a fixed order until one answers yes, with a tenth that always does. Then the class the switch builds from that kind, including the two extra questions the loop arm asks before it settles, one of which turns a copy into a memcpy and skips the kernel entirely. Then the three families behind those classes: six emitters that write a module in an MLIR dialect XLA defined for the purpose and lower it through forty-one passes to LLVM; the Triton path, which builds one module and hands it to another compiler; and the library paths, where cuBLAS arrives as a custom call and cuDNN arrives as a fusion, for reasons that follow from what each library accepts. Then the autotuner: the only pass that compiles and runs real kernels while it is still compiling, the four ways it can get a config, the sentences it uses to refuse, the two-microsecond window that lets a thriftier kernel win over a faster one, the clustering that treats correctness as agreement rather than truth, and the two-tier cache with a key caveat the source states out loud. The signature exhibit runs the real dispatch function over seven candidate fusions and shows which question fires and which class gets built.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of The GPU Codegen Path: From a Fused HLO to a Kernel — built to be understood, not skimmed.

Questions

How does XLA decide which GPU emitter compiles a fusion?: In two places, and the first one is not a decision at all. Every fusion instruction can carry a backend config with a string field named kind, and an earlier pass may have written a name into it: __triton_gemm from the GEMM fusion pass, __triton or __triton_nested_gemm_fusion or __triton_collective from other Triton claims, __cudnn$fusion from a cuDNN pattern matcher, or __custom_fusion for a fusion backed by a pre-compiled kernel. GetEmitterFusionKind checks those strings first and returns immediately, so a stamped fusion is never inspected. Only when there is no stamp does it look at the fusion itself, and then it asks six questions in a fixed order: is there a real reduction hero whose fellow roots have a compatible element count, does the analysis carry a tiled transpose, are there more than one root, is the single root a scatter, is it a sort, does the concatenate emitter accept it. If none of those fires, the last line returns the loop kind unconditionally. The resulting kind is then handed to GetFusionEmitter, a switch that constructs a concrete emitter class for it.
Why is my XLA matmul a Triton kernel instead of a cuBLAS call, or the other way round?: Because of the order two rewrite passes are added to the pipeline, seven lines apart. GemmFusion is added first, inside a condition that requires xla_gpu_enable_triton_gemm, which defaults on, and a compute capability of at least Ampere on NVIDIA or any ROCm card. It rewrites the dots it accepts into fusions stamped __triton_gemm. The GemmRewriter passes are added after it, and they rewrite whatever dots remain into __cublas$gemm or __cublas$lt$matmul custom calls. A pass can only rewrite what is still there, so Triton gets first refusal, cuBLAS takes the rest, and a dot that neither wanted stays an HLO dot and eventually gets a generic loop kernel. That is only the first claim, though: the autotuner later compiles a Triton kernel and a set of cuBLASLt algorithms for the same instruction and keeps whichever is faster, so the pass order decides who gets asked first, not who wins.
What does the XLA GPU autotuner actually do, and why is my first compile so slow?: It is the only pass in the pipeline that compiles real kernels and runs them on the device as part of compiling. It sits in its own pipeline at the very end of GPU optimization, named autotuner, and it does four things in order for each instruction: check the cache, try a cost-model estimate if one is enabled, take the first config that compiles if tuning is disabled, and only then tune. Tuning means enumerating every config every registered backend supports, compiling all of them, running all of them, and picking a winner. The default Triton GEMM config sets in the tree are 25 configs for Hopper, 30 for Ampere and 40 for Blackwell, and cuBLASLt separately enumerates up to 128 algorithms unless xla_gpu_blas_max_algorithms caps it lower. So a model with a dozen distinct matmul shapes is hundreds of full kernel compiles plus timed runs on a cold cache. On a warm cache it is a lookup, which is why the second run of the same program is dramatically faster.
Why did the XLA autotuner pick a config that is not the fastest one it measured?: Because the selector optimises for scratch memory among configs that are statistically tied on speed. PickBestConfig runs two loops. The first finds the minimum duration among configs that neither failed nor sit on an excluded backend. The second computes a limit of that minimum plus a window, defaulting to 2 microseconds, and among everything inside the window keeps the config with the fewest scratch bytes, breaking ties on duration. Scratch bytes are device memory a kernel needs as workspace beyond its inputs and outputs, and a cuBLASLt algorithm can want tens of megabytes of it for the life of the program. So a config 1.5 microseconds slower that needs no workspace will be chosen over the fastest one that needs 32 MiB. When the selection differs from the fastest, a log line at verbosity 2 names both configs and the tolerance, which is the only signal that the tiebreak fired.
How does the XLA autotuner know a kernel produced the wrong answer?: It does not compare against a reference implementation. It clusters the outputs and picks the biggest cluster that contains a trusted member. Every codegen backend declares whether it can produce wrong results, and the candidate list is stably partitioned so the ones that cannot go first. As each candidate runs, its output is compared with the representative of every existing cluster at a relative tolerance, defaulting to 0.1 for GEMMs; a match joins the cluster, a mismatch founds a new one. Once the trusted candidates have run and formed at least one cluster, the untrusted ones are run with new clusters disallowed, so an untrusted kernel that agrees with nothing trusted gets no cluster at all. Everything outside the winning cluster is then marked failed with kind kWrongResults and never reaches the selector. Correctness here is therefore agreement rather than truth, and it only runs at all when xla_gpu_autotune_level is 4 or higher, which is the default. A separate check surrounds every buffer with red-zone padding, 8 MiB by default, and fails any kernel that writes outside its own allocation.
How do I reuse XLA autotuning results instead of tuning on every run?: Set xla_gpu_per_fusion_autotune_cache_dir. The cache is two tiers: an in-memory store for the life of the process, and that directory as an optional second tier. Lookups try the primary first and then the secondary, hits in the secondary are promoted, and inserts go to every writable tier. The directory obeys a mode from xla_gpu_experimental_autotune_cache_mode, so you can tune with it writable on a build machine and deploy with it read-only. Add xla_gpu_require_complete_aot_autotune_results and a cache miss becomes a hard error rather than a silent tuning run on the serving host. The constraint is what a key holds: the device model, the canonicalised HLO of the instruction, and a version. Canonicalisation is why one tuning run covers every identically shaped layer of a transformer, and the device model is why a cache tuned on one accelerator does not transfer to another. One caveat is stated in the source: the key used to exchange results between sharded compilations hashes the module fingerprint and the backend names but not the module config, so compiling the same module twice with different debug options can recover the earlier, possibly inferior, results.

The GPU Codegen Path: From a Fused HLO to a Kernel

Questions

Related explainers