Fusion

Picture two operations chained together, one feeding the next. Run them unfused and the first op writes its whole result to memory before the second op is allowed to start reading it back. Merge them into a single kernel instead and the picture changes completely: the first op's output for one element flows straight into the second op's computation on that element, inside one loop, and the intermediate value never leaves fast memory at all. That merge is what XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → calls fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →, and the instruction it produces is a kFusion, one op standing in for what used to be several.

It is easy to picture fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → as the compiler finding a smarter formula, some algebraic shortcut dressed up as an optimization pass. That picture is wrong in a way worth fixing early, because it will mislead you at every later chapter. A fused kernel computes the identical sequence of arithmetic the unfused version computed; nothing about the math is rewritten, simplified, or reordered at the level of results. The only thing fusion changes is whether an intermediate value round-trips through memory or stays resident where the next op can reach it immediately.

Fusion is about the intermediates, never the math.

Two questions gate every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → decision XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → makes, and the compiler runs the cheaper one first. Fusibility asks a structural question: can these two instructions even be merged, given their opcodes and how their operands connect? Profitability asks the harder one: would merging them actually help? XLA answers profitability with a cost model, backed by an is_fusible callback each backend gets to customize for its own hardware, since what counts as a win on a CPU and what counts as a win on a GPU are not the same arithmetic.

On GPU that cost model runs as a priority queue: every candidate producer gets ranked by the time it would save if fused, and the pass fuses in that order, highest expected benefit first. Nothing here is exhaustive search. It is a greedy ordering over real cost estimates, and you can watch it work directly: the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → x-ray at /gym/xla#fusion holds three real before-and-after pairs pulled from a TPU, each one a case where this exact ranking made a call you can check against your own intuition. The kernel path already measured what a fusion like this is worth in practice, a fused softmax beating the unfused chain outright, recorded with full provenance at /bench.

Every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → pass, however aggressive, is bounded by one hard rule: it can only merge instructions that already sit on a dataflow edge in the graph it was handed. It never invents a new algorithm and it never restructures the computation into a shape that does less work. If two values are related in your head but not connected by a direct producer-consumer edge in the HLO, fusion has nothing to offer them.

This is the wall the kernel path spends an entire stage teaching you to see, at /l/xla: naive softmax needs the row maximum before it can compute a single exponential, so the score matrix gets written to memory, read back, and only then reduced. No pass in this pipeline can fold that spill away, because folding it away would mean changing the algorithm from multi-pass to streaming, and fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → does not do algorithms. It does edges.

The dump below is real, pulled from the optimized module for the attention program this path has been tracking. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → looked at a broadcast feeding directly into a divide and merged the two into one %fused_computation: the broadcast materializes its input into the divide's shape, and the divide consumes it, all inside a single computation with its own parameters and one ROOT instruction. Nothing outside this fused block ever sees the broadcast's output on its own.

%fused_computation (param_0: f32[64,64], param_1.1: f32[64]) -> f32[64,64] {
  %param_0 = f32[64,64]{1,0} parameter(0)
  %param_1.1 = f32[64]{0} parameter(1)
  %broadcast.3 = f32[64,64]{1,0} broadcast(f32[64]{0} %param_1.1), dimensions={0}, metadata={op_name="jit(attend)/jit(main)/div"}
  ROOT %divide.0 = f32[64,64]{1,0} divide(f32[64,64]{1,0} %param_0, f32[64,64]{1,0} %broadcast.3), metadata={op_name="jit(attend)/jit(main)/div" source_file="attend.py" source_line=12}
}

This section used to publish a disagreement. LAB·X2 timed this chapter's attention on a Colab TPU v6e at 2048 by 128 in bf16, medians of twenty runs in isolated processes, and reported 187.5 microseconds with the pipeline untouched against 139.3 microseconds with --xla_disable_hlo_passes=fusion, a ratio of 0.74x. Disabling fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → appeared to make the program faster, which is the opposite of everything above, so it went up here as an open question.

The question is now closed, and not in the direction anyone expected. Compiling both versions and reading what each one emitted returned the same module twice: 66 instructions either way, seven fusionsSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → either way, no custom calls in either, and the same fusion at the root. The flag changed nothing. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → accepted a pass name, found nothing registered under it in this backend's pipeline, and carried on without a word. Both timings measured the same program, so the ratio was never about fusion at all.

A flag that names a pass the backend does not have is not an experiment. It is a typo the compiler agreed to.

Pass names are backend-specific, which chapter 4's dump shows plainly: the TPU pipeline files its work under names like Phase_1_pre_layout_assignment_passes and X64_elimination that a CPU compile never mentions. A name lifted from one backend's vocabulary can be silently absent from another's, and nothing in the output distinguishes a disabled pass from a misspelled one. The habit that costs ten seconds and would have caught this before it was published: compile both ways, diff the modules, and only time them once you have proof the flag moved something.

The follow-up run answered the measurement half. Six identical processes, no flags anywhere, timing the same program on the same chip came back 170.2, 147.6, 147.6, 182.7, 166.5, and 200.4 microseconds. The first process sat within two percent of the median of the rest, so position was never the issue, but the spread across identical work is roughly 35%. A single pair of process timings on this hardware cannot resolve a 26% difference, which is all the retracted measurement ever was.

If the same work twice spans 35 percent, a single pair of numbers is not a measurement.

The vocabulary half turned up something the flag hunt had missed. Grepping the dumped pass names for fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → returns not one stage but five: tpu_fusion, main_fusion, tpu_multi_output_fusion, fusion, and async-collective-fusion. Disabling each in turn and diffing the compiled module settles which of them do the work here: tpu_fusion and tpu_multi_output_fusion change it, while fusion, main_fusion, and async-collective-fusion leave it untouched. So the backend does register a pass called exactly fusion, and disabling it does nothing for this program, which is why the original flag looked accepted and behaved like a no-op at the same time.

The experiment then ran twice, with tpu_fusion confirmed to change the module first. Six alternating pairs, then six more: 191.9 against 168.9, 153.4 against 164.9, and so on through twelve. Fusion came out ahead in ten of the twelve, and the marginal ranges still overlap in both runs, which is what the lab reported at the time.

The overlap test was the wrong test, and this section owes that correction as much as it owed the first one. The design pairs the two configurations and alternates them precisely so that drift lands on both sides of a pair, which means the paired difference is the measurement and the marginal range is not. Read that way the twelve pairs give a median difference of 16.7 microseconds in fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →'s favour, ten of twelve in the same direction, which a sign test puts at about p = 0.02. Comparing the outer ranges threw the pairing away and hid a result that was sitting in the data.

Pair the runs, then compare the pairs. Comparing the envelopes discards the design.

So the answer, finally: on a v6e, disabling tpu_fusion for this attention program costs somewhere around ten percent, and the mechanism this chapter argues for holds. It took a retraction, a flag that did nothing, a stage hunt, twelve paired runs, and one more correction to the analysis to say that with a straight face, which is roughly the honest price of a performance claim.