Follow on X: https://x.com/PatrickToulme

Code reproduction: https://github.com/patrick-toulme/justabyte/blob/main/ai_compute_1_percent_post/compute_model.py

Subscribe now

What happens when a software engineer stops using one AI agent and starts commanding a swarm of fifty? When a quarter of adult Americans each have an AI that staffs a team of agents the way a manager staffs a project? In this scenario you don’t get one agent per person. You get a planet trying to run billions of always-on agents, with nowhere near the silicon to do it.

This article will argue that the US and world economies are changing rapidly, that AI silicon is the bedrock of that change, and that we have built very little of the compute this future requires.

Today the world has roughly 19 million H100 equivalent AI chips. That is every Hopper, every Blackwell, every Google TPU, AWS Trainium, Meta MTIA, AMD, Cerebras, Groq etc. ever shipped and still racked, added up and normalized to one unit. It is the core of the AI economy: about 38 zettaFLOP/s of FP8 compute, most of it built in the last eighteen months.

Here is the thing people get wrong about the demand side. They picture one agent per human and ask whether that is realistic. No human is going to manage a hundred agents by hand. You talk to one AI, and it runs the swarm of agents for you. It plans, spawns subagents, runs them in parallel, and hands back a result. The human stays in a single conversation; the machine fans out into hundreds. That is already how the best agentic coding tools work, and open-source models such as GLM-5.2 and Kimi K3 are now capable of orchestrating these massive swarms of agents.

So take a grounded case, and a deliberately partial one. Not all 8 billion humans. Not even all of the world’s roughly 1 billion knowledge workers, because not everyone will adopt this or afford it. Just half of them, about 500 million people, each with an AI that runs a modest 50-agent swarm. How many chips does that need?

About 830 million chips. Roughly 44 times every AI chip that exists.

This article shows the math behind that number, and the scale at which the world has to build compute to make this future real.

That 830 million is not the aggressive case. It is a conservative reading, on a cheap open model, for a deliberately narrow slice of the workforce. I picked one legible group on purpose, as a frame of reference: analysts, engineers, and lawyers using a team of agents at work, and only half of them. One ordinary slice already needs 44× every chip on Earth, and the slices I left out stack fast. The other half of knowledge workers: another 44×. Three billion consumers with even a five agent helper for daily life: another ~26×. The floor climbs past 100×, of which we’ve built about 1%. We are not oversupplied on AI compute. We are not even close.

I want to make this concrete and reproducible, so I built a small parametric model and sourced every input to SemiAnalysis numbers where they exist. The experiment is ~150 lines of Python, linked at the end. This code powers the graphs in this article, and you can change any assumption you disagree with and watch the conclusion move.

The conclusion does not move much.

The unit, and the question

Everything here is in H100 equivalents (H100e), FP8 throughput convention: a Blackwell counts as ~2.3 H100e because it does ~2.3× the FP8 work.

The thought experiment has one assumption worth stating upfront: what does “running AI” mean per person? I model always-on agents: not chats you prompt a few times a day, but workers that generate tokens continuously. And every frontier model people actually use now is a reasoning model, one that generates a long hidden chain of thought before it answers, roughly 8× the tokens of a non-reasoning model. So I count generated tokens, reasoning included, at a sustained 40 tokens/sec per agent, averaged over the full day, with a floor of 15 and an aggressive case of 100.

The unit of demand is the swarm, not the person. Here’s how that actually works. A user opens one session and states a goal. An orchestrator model reads it, breaks it into pieces, and spawns subagents (a planner, a handful of coders, reviewers, test writers, researchers) that run in parallel, each in its own context, and report back. Chain those steps and you have a workflow: the orchestrator fans work out, gathers the results, and decides what to do next. No human is monitoring the 50 agents. The human is managing one agent, and that agent is managing the rest. So the thing that scales isn’t how many people adopt AI. It’s how many agents each person’s AI is running, and that number is climbing from one toward the hundreds.

I sweep that swarm size across grounded populations: all software engineers (~28 million, per JetBrains/SlashData), a fraction of adult Americans, and half the world’s ~1 billion knowledge workers. I sweep one more axis underneath everything: model size, from 10B to 100T parameters because the entire argument bends on it.

What the world has now

The installed compute base is ~15–20M H100e (Epoch put global capacity at ~15M in early 2026, and it’s doubling roughly every seven months). NVIDIA’s own fleet, its millions of Hopper GPUs plus the fast-ramping Blackwell, comes to roughly 11M H100e, but it does not own the whole fleet. The custom accelerators are not a rounding error: Google is in fact the single largest owner of AI compute, holding about a quarter of the world’s total, mostly its own TPUs (Ironwood does 4.6 PFLOPS FP8 with ~4.3M TPUs shipping in 2026). AWS has 1.4 million Trainium chips deployed across three generations, over a million of them running Anthropic’s Claude; AMD’s Instinct is at a $10B-a-quarter run rate, and Microsoft’s Maia, Meta’s MTIA, and Huawei’s Ascend add more. Let’s call the world total ~19M H100e.

AI compute production is rapidly growing. NVIDIA went from ~3.5M datacenter GPUs in 2024 to ~6.5–7M in 2025, up ~55% in units and ~68% in revenue, and the fleet’s compute is roughly doubling every nine months. Across every vendor, that’s ~16M H100e added per year and rising.

Here is the part that matters: you cannot simply print these. The binding constraint on AI chip production is not logic dies; those are abundant, and AI silicon vendors just outbid everyone else for leading node wafers. The constraint is advanced packaging (CoWoS) and HBM. In 2025 the four largest AI chip designers consumed ~90% of global CoWoS capacity and HBM supply while using only ~12% of advanced logic production. CoWoS has roughly doubled every year (16k → 40k → 70k → 130k wafers/month from 2023 to 2026), and HBM has been sold out through 2026, with SK Hynix at 62% share supplying ~90% of NVIDIA’s stacks.

The supply of AI chips is a physical pipeline with significant resource constraints.

What one model costs to serve

To serve a model you need two things from your chips: enough memory to hold the weights, and enough memory bandwidth to stream them per token. For frontier models, bandwidth is the wall.

Three facts set the per-chip economics, and all three are well grounded in real 2025 deployments:

1. Frontier models are sparse, and getting sparser. No one serves a dense 10T model. Frontier models are Mixture of Experts, and the active fraction shrinks as total size grows: DeepSeek-V3 activates 37B of 671B (5.5%), the open GLM-5.2 activates ~40B of 744B (5.4%), Kimi K2 activates 32B of 1T (3.2%), and the largest open model yet, Kimi K3, activates ~104B of 2.8T (3.7%). The optimal sparsity curve bends down with scale, and deployed models are more aggressive than optimal because serving economics reward it. Active params grow sub-linearly: I model 10B→3B active, 100B→12B, 1T→40B, 10T→200B, 100T→1T.

2. Decode is memory-bandwidth bound. To generate one token you stream every active parameter out of HBM once. The roofline is simply tokens/sec ≈ HBM_bandwidth / (2 × active_param_bytes). On an H100 (3.35 TB/s) with 37B active params in FP8, that’s ~45 tokens/sec, single stream. This is not a FLOPs problem. It’s a “how fast can you read the weights” problem, and it’s why bandwidth, not compute, is the decode bottleneck for essentially every model chip pair.

3. Batching helps, but reasoning caps it. That 45 tok/s is one user at batch one. Batch many users together and you amortize each weight read across all of them, but only down to the latency your users will tolerate, and reasoning users tolerate very little, as a 25,000-token thinking trace served at low interactivity would take half an hour. So you serve at a usable speed, and SemiAnalysis’s InferenceMAX measures exactly what that costs. For DeepSeek-R1 (a reasoning MoE, 37B active) at ~42 tokens/sec/user, an H100 manages just 266 tok/s/GPU and a B200 4,792, about 2,080 per H100-equivalent. Blend a Hopper and Blackwell fleet and you land near 1,300 tok/s/H100e for a 37B-active model; I scale it inversely with active params.

(And 1,300 is still generous: it assumes a Blackwell-heavy fleet. An H100-only fleet serves ~5× worse, at 266 tok/s/H100e, and pushing interactivity higher, to 71 or 100 tok/s/user, collapses throughput further still. Every one of those moves makes the numbers below bigger, not smaller.)

One caveat the other way: I run all of this in FP8. The frontier is sliding to FP4 (DeepSeek quantizes FP8→NVFP4 with under 1% accuracy loss; Together and Fireworks already serve it, and Kimi K3 now ships its weights natively in MXFP4, quantization-aware trained from the SFT stage on), which roughly halves every chip count here.

Put it together and you get the per-chip serving rate for each model size. Now apply it to a realistic slice of the workforce.

A fraction of people, each running a swarm

All humans do not need to use AI agents for the world to be extremely compute constrained. We only need a slice of the workforce running swarms on a cheap, open model, and that is already happening.

Every line crosses the world’s entire chip fleet almost immediately. All ~28 million software engineers, each driving a 50-agent coding swarm, need ~47 million chips, 2.5× every AI chip on Earth, and they cross the fleet at about 20 agents apiece. A quarter of adult Americans (~65 million people) on the same 50-agent swarm is 5.7× the fleet, crossing at roughly 9 agents each. And half the world’s knowledge workers, ~500 million people whose work happens on a screen, reach the entire global fleet at barely one agent apiece; give each of them a 50-agent swarm and you are at ~830 million chips, 44× everything ever built.

None of this assumes a giant model, universal adoption, or anyone managing more than one conversation. It assumes a cheap open model, a fraction of the professional world, and the swarm sizes agentic coding tools already ship.

And notice what this slice leaves out. Half of one profession is a frame of reference, not a forecast. I chose it because it’s easy to picture and easy to argue with. It counts no one using AI for their health, their kids’ education, their finances, their research, or their small business: the things that will reach billions of people who never call themselves knowledge workers. Each of those is its own swarm story. I left them all out, and the number is still 44× every chip on Earth. Add any one of them back and it only moves one way.

The model is a free download. The silicon isn’t.

The reason a fifty-agent swarm is realistic and not a thought experiment is that the model to run it already shipped, and it’s open. GLM-5.2, from Zhipu / Z.ai, is a 744 billion parameter model with ~40 billion active, released in June 2026 under an unrestricted MIT license with a 1-million-token context. It posts frontier-class agentic coding scores: 62 on SWE-bench Pro and an MCP-Atlas tool-use score of 77, within a few points of Claude Opus 4.8 on the long-horizon coding marathons. It speaks the Anthropic and OpenAI APIs, so it drops into Claude Code, Cline, or Cursor with a base URL swap, and the hosted API runs about $2 per million input tokens and $6 per million output, roughly a sixth of the closed leaders, with $18-a-month coding plans. I’ve run it myself: open model, open harness, local serving, frontier coding at essentially zero marginal cost.

And GLM-5.2 is no longer alone. Weeks after it shipped, Moonshot AI open-sourced Kimi K3: 2.8 trillion total parameters, ~104 billion active, native vision, the same 1-million-token context, weights free to download under a permissive commercial license. It posts frontier agentic scores of its own, 42.0 on SWE Marathon against Claude Opus 4.8’s 40.0, 88.3 on Terminal-Bench 2.1, with a hosted API at $3 per million input tokens and $15 per million output. Two labs, weeks apart, each handing the world a frontier-class agent orchestrator for nothing but the cost of the silicon to run it. The open frontier isn’t an event. It’s a cadence.

Here’s the part that matters for compute, and it’s the opposite of what “open and free” sounds like. Open weights don’t remove the silicon. They remove the brake. To actually serve GLM-5.2 you still need the chips: its weights are ~744 GB in FP8, a box of roughly eight datacenter GPUs before it generates a single token, and Kimi K3’s weigh ~1.4 TB even packed in 4-bit. What open weights remove is the meter. There’s no per-token bill capping how many agents you run, no vendor rate limit, no one deciding you’ve had enough tokens for the day. The number of agents you can run stops being a function of your API budget and becomes a function of how many accelerators you can get.

That is the demand engine for this whole piece. A closed, metered model self-limits: you run the agents you’re willing to expense. An open model you host yourself has one limit left, and it’s silicon. Make the best coder a free download and the rational move isn’t to run one agent and bank the savings; it’s to run fifty, because the marginal token is nearly free and not running the swarm is leaving value on the table. Open weights are Jevons’ paradox with the last brake cut.

It also globalizes the demand. A model on Hugging Face runs in any datacenter, any country, any startup’s basement. When the best model is an open download, no single vendor can even meter the world’s demand for it, let alone supply it. The only thing standing between an open frontier model and an unbounded swarm is the chip.

So run every scenario in this piece on GLM-5.2 specifically, with its real 40 billion active params setting the per-chip rate. That’s a deliberately conservative anchor: Kimi K3 activates ~104 billion params per token to GLM’s 40 billion, so a world whose swarms run on the newest open flagship instead needs roughly 2.6× every chip count below. Here is the bill:

single profession on a fifty-agent swarm already needs more chips than every vendor has ever built. The model is free. The hundred million plus chips to run it are not.

To put the swarm in everyday units: half the world’s knowledge workers, 50 agents each, generate on the order of 80,000 trillion tokens a day. The entire AI industry today, every ChatGPT session and every Gemini call, runs about 50 trillion tokens/day. One slice of one profession is already on the order of 1,000× the world’s current token output.

What happens when some people want faster inference?

Everything above assumes standard reasoning serving: ~40 tokens a second per agent, the InferenceMAX operating point. That’s fine for a background worker. It is too slow for the thing people actually pay for.

Watch a developer drive an agentic coding tool and you’ll see the real demand. They want the answer now. Not 40 tokens a second but 100, 200, as fast as the model can think, because a human is sitting there waiting and their time is worth more than the GPU’s. More importantly, they will pay for it. A coding agent that’s 3× faster is worth a lot more than 3× the price, because for that workload latency is the product.

But speed and batching are opposites. To make one request faster you pull it out of the batch and hand it more of the chip. On InferenceMAX, pushing DeepSeek-R1 from ~42 to ~100 tokens/sec/user on an H100 cuts throughput from 266 to 23 tok/s/GPU, more than 10× the silicon per token. I’ll be conservative and price the premium tier at a flat 3×.

Put a slice of the world on that fast tier and the whole curve lifts:

If that knowledge worker swarm ran on the fast tier, the bill goes from ~830 million chips to ~2.5 billion, 131× the world’s entire fleet. Even a modest 20% premium segment adds ~330 million chips on top. And this is the high-margin demand: the customers who pay 2–3× are precisely the ones who make a chip most worth selling. The latency-sensitive tier doesn’t just add load. It adds the load people are most willing to fund. Those who are running fast-tier inference for agentic coding will be running a lot of it and paying virtually whatever it costs.

What about AI training? Pretraining is a rounding error. RL isn’t.

You might expect pretraining these large models to dominate the bill. It doesn’t, and the gap is very illuminating.

Training cost is 6 × active_params × tokens. Run the numbers (FP8, 35% MFU, calibrated to GPT-4’s 2.1e25 FLOP and the 100k-H100 cluster that retrains GPT-4 in four days): training and continuously refreshing a fleet of 20 frontier model lines at 10T params costs ~130k H100e as a standing fleet. At 100T, ~880k.

Against ~830 million chips for serving, training is under 0.03% of the bill. The two demand curves above are visually identical because the training contribution vanishes beneath them. This is the inference economy in one fact: once a model is good, the cost is in running it, forever, for everyone. The market still prices these chips as if training clusters were the demand. The demand is the billions of agents that haven’t been spun up yet.

That covers pretraining. But there’s a much more compute-hungry paradigm of training: RL post-training.

The scale up of reinforcement learning is the steepest curve in AI right now. Reinforcement learning was ~1% of pretraining compute before 2024; then under 10% for OpenAI’s o1, ~20% for DeepSeek-R1, and for xAI’s Grok 4, roughly parity, with RL run at pretraining scale on a 200,000-GPU cluster. Epoch estimates RL compute is growing ~10× every three to five months, against ~4× a year for pretraining. The two are converging now.

And here’s why it belongs in this article and not a training footnote: RL is inference-shaped. Its cost isn’t the gradient step; it’s rollout generation. To run RL you sample many long answers per prompt and grade them (DeepSeek-R1 drew 16 completions per question, up to 32,768 tokens each), and that generation, autoregressive memory-bandwidth-bound decode, is 70–90% of the wall-clock. As SemiAnalysis puts it, inference is “no longer just the end of the pipeline, but an integral part of training.” RL doesn’t dent the “serving dominates” claim. It is serving, pointed inward towards improving the model.

Pretraining plus RL across 20 frontier lines is ~260k chips, a rounding error against the ~830 million for serving. But at today’s scale it’s the whole story. Most of the current fleet isn’t serving the public at all: only ~30% of OpenAI’s 2024 compute went to inference, the rest to training, RL, and experiments, and ~40% of the world’s AI compute is training of one kind or another. That’s exactly why only ~60% of the fleet is free to serve, and why we’re so far from one agent per human. The compute to build the models is still eating the compute that would run them, and the fastest-growing part of that, RL, is itself inference.

You can’t ignore the gap in compute

The realistic scenarios need anywhere from a few times to a few hundred times the world’s AI silicon. How fast can we build it?

At ~16M H100e of net new production per year, closing this gap of ~830 million chips is ~50 years of the entire planet’s AI chip output, devoted to this one slice of the workforce, ignoring replacement of dying hardware. Production is growing ~1.6× a year in units, faster in H100e, so the real timeline is shorter. But the throat is CoWoS and HBM, and those double roughly once a year, not once a quarter. You cannot software your way out of a packaging constraint. You cannot fine-tune your way past an HBM shortage. The gap between “19 million chips” and “what a fraction of the workforce running swarms needs” is not a market sentiment. It’s a decade plus of the hardest manufacturing on Earth running flat out.

This is the whole argument, and I’ve kept it in units of silicon on purpose. I’m not going to model market caps or P/E multiples. The installed base is the floor of the AI economy’s capacity, every credible demand scenario sits one to two orders of magnitude above that floor, and the floor is bottlenecked on two of the most capital-intensive, slowest-to-scale processes in existence.

And then you have to power them

All in (GPU, networking, cooling, grid losses), a modern AI chip pulls on the order of a kilowatt from the wall; a GB200 NVL72 rack is 120 kW for 72 GPUs, and since perf per watt keeps improving I’ll call it ~1 kW per H100 equivalent and be generous. Serving half the world’s knowledge workers a 50-agent swarm each, ~830 million chips, is then about 830 gigawatts of continuous draw.

Every data center on Earth combined draws ~55 GW today (≈485 TWh/yr, ~1.5% of global electricity). This one realistic case is roughly 15 times the entire planet’s current data center load, nearly double all the electricity the United States consumes, for a single slice of one profession. Push the swarm to 500 agents and it’s ~8,300 GW, many times US generating capacity.

And power moves slower than silicon. New generation and grid interconnects take 24 to 72 months; fewer than 5% of existing data centers can even feed a 50 kW rack, never mind 120. CoWoS doubles in a year; a substation does not. The demand curve isn’t sitting above one slow physical wall; it’s sitting above two in series, packaging and power, and it clears neither for a decade. The chips are useless without the watts, and the watts are harder to get than the chips.

Hardware vendors aren’t competing. There is no pie.

Here’s the pervasive story I want to kill. Every quarter someone writes that NVIDIA, Google’s TPU, AWS Trainium, AMD, and Cerebras are locked in a fight for the AI accelerator market: a fixed pie, each new entrant taking a slice from the rest, winners and losers.

This is wrong. There is effectively infinite demand for AI silicon. One vendor cannot supply the entire market.

Putting the numbers on it:

Add up every AI chip from every vendor: ~11M H100 equivalents of NVIDIA, ~4.8M of Google TPU, ~1.2M of AWS Trainium, AMD’s ~1M, and the rest (Maia, MTIA, Huawei). You get about 19 million chips. That total, the entire industry’s combined output, doesn’t even cover all the world’s software engineers running a 50-agent coding swarm, which needs ~47 million chips, 2.5× every chip every vendor has ever made. A quarter of adult Americans on the same swarm needs ~109 million, ~6×; half the world’s knowledge workers, ~830 million chips, 44×.

When demand sits anywhere from a couple times to many hundreds of times above the combined supply of every manufacturer, those manufacturers are not competing for customers. There aren’t enough chips in existence to fight over one. NVIDIA at ~58% share isn’t winning a zero-sum game against Trainium; they are both selling everything they can build into a market that would swallow dozens of times the volume. Cerebras and Groq aren’t stealing NVIDIA’s inference business. They’re additional capacity aimed squarely at the low-latency premium tier from two sections ago, which is additive demand NVIDIA can’t fully serve anyway.

The real competition isn’t vendor versus vendor for customers. It’s all vendors versus physics for CoWoS and HBM, the shared, sold-out chokepoint upstream of all of them. They’re not fighting over the pie. Which is also why portability is a myth: when every vendor is supply limited, each one optimizes its own stack to the metal and sells out regardless. Nobody is losing deals over a missing CUDA shim; they’re losing them over packaging substrate.

How to break this

I held the whole model to one rule: every uncertain assumption breaks toward supply. So the real numbers are likelier to be worse, not better. The knobs, and which way they push:

• Throughput. I used ~1,300 tok/s/H100e: InferenceMAX at 42 tok/s/user, blended across a Blackwell-heavy fleet. An older H100-heavy fleet serves ~5× worse (266/H100e), and the faster interactivity coding agents demand cuts it further. Realistic serving pushes the numbers up, not down.

• Sparsity. If active fractions don’t keep shrinking (if a 10T model needs 500B active, not 200B), serving costs roughly 2.5× more. My schedule is optimistic for supply here too: Kimi K3 already runs above it, at ~104B active where my curve implies ~80B for a 2.8T model.

• Population intensity. Every number above uses the mid rate of 40 tok/s/agent. The floor is 15 (chips ÷ ~2.7); the aggressive case is 100 (chips × 2.5). Even the floor blows past the fleet at modest swarm sizes.

• Model size is the real variable. This is the one that helps supply. If a few billion parameter model turned out to be all anyone needed, today’s fleet could nearly cover one agent per human. The entire bull case for “we have enough chips” is “models get smaller and we stop wanting more of them.”

• What I don’t model. Fleet depreciation and failures (GPUs last ~3–5 years, so a slice of annual production just replaces dead silicon), node-level memory and networking limits, and downtime. Each of these shrinks the effective fleet, so leaving them out is again supply-favorable.

What I’m not claiming: that a fifty-agent swarm lands in every knowledge worker’s hands next year. Adoption is slow, latency is hard, and most of the workforce isn’t running agents yet. This is a demand curve made of physics, population, and the economics of cheap agents, not a forecast of next year’s bookings.

“But models will get more efficient”

This is the last objection standing, and the one people are most certain of. Models will get smaller. A 10 billion parameter model will eventually match today’s 10 trillion parameter frontier through distillation, better data, and architectures we haven’t found yet. Serving collapses 60×, and the shortage solves itself.

It doesn’t. It inverts. This is Jevons’ paradox, the most durable pattern in the history of resource use: make something cheaper to consume and total consumption goes up, not down. Cheaper steam engines burned more coal, not less. Cheaper compute didn’t empty the datacenters; it built hyperscale. Make an agent 60× cheaper and nobody runs one agent and banks the savings. They run sixty. Then six hundred. Then a thousand, because at that price the only expensive thing left is restraint. We are watching it happen in real time. GLM-5.2 already put a frontier coder at a sixth the price as an open download, and the response wasn’t fewer agents. It was swarms. And the labs are now selling efficiency itself as the headline: Moonshot pitches Kimi K3 as 2.5× the intelligence per unit of compute of its predecessor. That gain won’t be banked either. It will be spent, on more agents.

So watch what a breakthrough actually does. Say a 100B model matches today’s 10T quality, and people respond the obvious way, not with one agent each but with a swarm:

A 100B model is 16× cheaper per agent than the 10T. So nobody who pockets that saving runs one agent and banks it; they grow the swarm. Take the same half-billion knowledge workers and give each a hundred agents on the cheaper model and you need ~500 million chips; push to a thousand, the number cheap agents actually invite, and it’s ~5 billion chips, 263× the entire world fleet, six times the realistic headline we started from. Even the most efficient case, a 10B model at frontier quality with a thousand agents each, is ~1.25 billion chips, 66× the fleet.

The efficiency breakthrough didn’t shrink the problem. It detonated it. And it compounds with everything else: a cheaper model is also a cheaper fast model, so more of that swarm runs on the premium tier instead of the batched one. Every gain in efficiency gets spent on more agents, running faster, for more people, until it slams back into the one wall that doesn’t move: how fast the world can package silicon and stack memory.

This is the deep reason the chips are the floor and not the ceiling. You cannot make models efficient enough to escape it, because efficiency is the thing that creates the demand. The cheaper inference gets, the more of it the world buys.

Does this mean the stocks are cheap?

Regardless of what you believe about a future in which humans are running agents 24/7, the direction the world is taking is undeniable. Models keep getting more useful. Agents multiply rather than consolidate; every coding tool already spawns dozens at once. Token volume is growing severalfold a year; ChatGPT went from 400M to 900M weekly users in twelve months. You don’t need my exact scenario. You need any point on that curve a few times past today, and every one of them sits far above the installed base, behind two walls that take years to move.

That’s the valuation argument, and it’s why I kept the whole piece in chips. When demand is structurally multiples above supply, and supply is gated by fabs and power on multi year clocks, the gap clears exactly one way: the companies that make the packaging, the HBM, the accelerators, and the power have to get much bigger, in absolute terms, for a long time. “Undervalued” here isn’t a P/E call or an entry price. It means the market is implicitly pricing this fleet as if it were near its mature size, when the arithmetic says it’s a small multiple of the floor.

Conclusion: There is no bubble in AI.

The reflexive take on AI compute is that it’s a bubble: too many GPUs, too much capex, a glut waiting to clear. Run the demand side in units of silicon and the opposite falls out. Today’s entire global fleet can’t cover even a fraction of the workforce running modest agent swarms on a cheap open model. That is the most grounded scenario in this piece, and it’s already ~44× every chip on Earth, for one ordinary slice of working life. Every step up (bigger swarms, more people, faster tiers, larger models) multiplies the requirement further past everything that exists. Agents are multiplying from one toward the hundreds. The world has 19 million chips.

The chips aren’t overbuilt. They’re the floor of a new economy we’ve barely started building. The companies that make the packaging, the HBM, and the accelerators aren’t selling into a glut. They’re selling the first 1% of the substrate, against a demand curve set by the size of the workforce and the swarms it will command.

Every one of these companies is already among the most valuable on Earth. Every one of them is still too small for the floor it’s standing on. That’s the part the market hasn’t priced: not that the chips are overbuilt, but that the people who build them have barely begun.

A new economy is forming, one in which the foundations are in silicon.

Connect on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Subscribe now

Model: compute_model.py, ~150 lines, every figure and number reproducible. Chip counts in FP8 H100-equivalents; supply figures SemiAnalysis-derived via IFP and Epoch; serving economics calibrated to SemiAnalysis InferenceMAX and LMSYS; training to SemiAnalysis cluster figures and Epoch. Every source is linked inline above. Change any number you doubt and rerun; the gap survives it.

Follow on X: https://x.com/PatrickToulme

Motivation

Disclaimer: This blog is more opinionated than my prior blog posts.

The AI industry says it wants portable kernels. It keeps building hardware-specific ones.

TPUs have Pallas. Trainium has NKI. NVIDIA has CUDA C, CUTLASS, Triton, CuTile, CuTe — five DSLs and counting. AMD adopted Triton, the closest thing to a “portable” kernel DSL — then built FlyDSL anyway. Tenstorrent has tt-Metalium. Mojo targets NVIDIA and AMD.

If portability worked, we’d have one universal DSL. We have many — and the number keeps growing.

282 lines of Pallas Python on TPU. 4 million lines of generated CUDA C++ on Blackwell. Same operation — MoE grouped matrix multiplication. Zero shared code. Not because anyone failed at portability, but because the hardware is different enough that the optimal algorithms, tiling strategies, memory staging, and synchronization are all different. Portability at the performance layer was never possible.

And yet teams keep chasing it. Startups burn months building “hardware-agnostic” training stacks that run everywhere and run fast nowhere. Hardware vendors contort their silicon to fit someone else’s programming model. Engineers write portable kernels that achieve 30% MFU on two platforms instead of 90% MFU on one — and at frontier scale, that gap is millions of dollars in wasted compute.

This post argues that the best AI stacks of the future will not be portable — and that this is a good thing. It lays out why portability at the kernel level is a myth, why every hardware vendor will end up building their own DSL, and what those DSLs should look like to maximize both performance and usability.

Three Layers of the AI Stack

Every AI training or inference stack has three layers. Understanding which layers are portable — and which aren’t — is the key to understanding why the industry keeps rebuilding infrastructure from scratch.

Layer 1: The Math (Portable)

# This runs on TPU, GPU, CPU, Trainium. It's the same math everywhere.
y = jnp.dot(x, w)                    # matrix multiply
y = jax.nn.softmax(logits)            # softmax
loss = optax.softmax_cross_entropy(logits, labels)

PyTorch and JAX are portable at this layer. torch.matmul and jnp.dot express a mathematical operation — multiply two matrices. The framework doesn’t specify how. This is real portability. It works, and it matters.

But jnp.dot is not what runs on the hardware. What runs on the hardware is the output of a compiler that made thousands of hardware-specific decisions about tiling, memory placement, instruction selection, synchronization, and scheduling. The math is the same. Everything else is different.

Layer 2: The Compiler (The Portability Illusion)

XLA, torch.compile, Triton’s compiler — these sit between your portable math and the hardware. They take jnp.dot and produce TPU VLIW bundles, or CUDA PTX, or Trainium NeuronCore instructions. The API is portable. The output is not.

This is where the portability illusion lives. When someone says “JAX runs on TPU and GPU,” what they mean is: JAX’s tracing and HLO generation are shared, but XLA-TPU and XLA-GPU are different compilers that produce fundamentally different code for fundamentally different hardware. They share an IR format. They don’t share a backend.

The same is true for PyTorch. torch.compile with the Inductor backend targets NVIDIA GPUs. torch.compile with the XLA backend targets TPUs. The torch.matmul call is the same. Everything below it diverges.

Layer 3: The Hardware-Native Code (Not Portable)

This is where MFU is won or lost.

On TPU, a fused attention kernel is a Pallas program that manages VMEM explicitly, tiles across 8 sublanes × 128 lanes, issues DMA operations for HBM↔VMEM transfers, and produces VLIW bundles that pack MXU matmul, VPU elementwise, and XLU cross-lane ops into single cycles.

On Blackwell, the same logical operation is a CuTile program that allocates TMEM columns with atomic contention handling, elects a warp leader with elect.sync, issues tcgen05 MMA instructions, manages 20 mbarrier objects for async pipelining, and backs off with NANOSLEEP on TMEM allocation failure.

On Trainium, that operation is an NKI kernel that partitions work across NeuronCores, stages data through SBUF (on-chip SRAM) and PSUM (accumulation memory), and issues tensor engine instructions with explicit partition-dimension tiling.

These are not different implementations of the same algorithm. They are different algorithms driven by different hardware constraints. You cannot port one to another. You rewrite from scratch.

The Evidence: Same Math, Zero Shared Code

Here’s the most concrete version of this argument. Mixture-of-Experts grouped matrix multiplication — the same operation compiled for two different hardware targets.

On TPU: MaxText MegaBlox (Pallas)

From my last post: MaxText’s MoE layer compiles into 29 GMM Pallas kernel calls. Each one is a tpu_custom_call that:

• Runs on a 3D grid (tiles_n, num_active_tiles, tiles_k)

• DMA-fetches [512, 1024] tiles of sorted tokens from HBM to VMEM

• Selects the correct expert’s weight tile via group_ids[grid_id]

• Accumulates dot(lhs_tile, rhs_tile) in VMEM f32 scratch

• Applies group boundary masks and stores results back to HBM

The backward pass uses two separate kernels: gmm with transpose_rhs=True for input gradients, and tgmm — a structurally different kernel — for weight gradients. Tiling is (512, 1024, 1024), tuned for the TPU MXU’s 256×256 systolic arrays.

Total handwritten kernel code: 282 lines of Pallas Python targeting TPU memory hierarchy and MXU.

On NVIDIA: flashinfer MoE (CUDA)

flashinfer PR #2917: 300 files of generated CUDA kernels for the same logical operation — MoE batched GEMM. Each file targets SM100 (Blackwell) and encodes its optimization parameters directly in the filename:

batched_gemm_e2_sm100_s128x128x128_...bf16_tma_warpspecialized_cooperative_align.cu
batched_gemm_e2_sm100_s128x256x128_...fp8_tma_warpspecialized_pingpong_align.cu

Tile sizes (128×128×128, 128×256×128). Data types (bf16, fp8). Memory access patterns (TMA). Scheduling strategies (cooperative, pingpong). Architecture target (SM100). All baked in. The PR adds approximately 4 million lines of generated CUDA — not because the developers enjoy writing CUDA, but because peak performance on SM100 requires SM100-specific code.

Why a Portable DSL Can’t Fix This

The obvious response is: “What if we had one DSL that abstracts over these differences?”

This is the promise of Triton, and to some extent the MLIR ecosystem. Write your kernel once in a hardware-agnostic DSL, and the compiler lowers it to TPU, GPU, or Trainium.

The problem is that a portable DSL only works if the compiler is sufficiently complex to bridge the gap between abstract operations and hardware-specific instructions. And that compiler would need to do something no compiler has ever done.

What the Portable Compiler Would Need to Do

Take tile_matmul(a, b) in some abstract DSL. The compiler must decide:

For TPU:

• Allocate VMEM buffers for both operands

• Issue DMA copy-start/copy-done pairs for HBM→VMEM transfers

• Tile to (512, 1024, 1024) for the 256×256 MXU systolic arrays

• Pack MXU matmul + VPU elementwise + XLU cross-lane shuffle into VLIW bundles

• Double-buffer DMA and compute across loop iterations

For Blackwell:

• Allocate TMEM columns with atomic UTCATOMSWS.FIND_AND_SET

• Elect a warp leader with elect.sync before every MMA

• Issue tcgen05 MMA instructions through the leader thread only

• Manage 20 mbarrier objects with parity-based reuse for async pipelining

• Back off with 100ns NANOSLEEP on TMEM allocation contention

• Await results with 10ms barrier wait timeouts

For Trainium:

• Stage data through SBUF (on-chip SRAM) with explicit partition-dimension tiling

• Issue tensor engine instructions across NeuronCore partitions

• Accumulate in PSUM memory space

• Manage NeuronCore-level parallelism

These aren’t different instruction selections for the same algorithm. They’re different algorithms. The double-buffered DMA approach on TPU, the leader-elected async MMA on Blackwell, and the partition-tiled tensor engine on Trainium are structurally different programs that solve the same mathematical problem through entirely different mechanisms.

The Union Problem

A portable DSL that exposes enough hardware detail to write high-performance kernels would need to contain the union of all hardware concepts:

• TPU: VMEM, HBM, DMA engines, MXU, VPU, XLU, sublanes/lanes, VLIW bundles

• Blackwell: SMEM, TMEM, L2 cache, tensor cores, tcgen05, mbarriers, TMA, warp specialization

• Trainium: SBUF, PSUM, HBM, tensor engine, NeuronCores, partition dimensions

That’s not a DSL — it’s three DSLs wearing a trench coat. Every kernel written in it would need if target == TPU: ... elif target == NVIDIA: ... elif target == TRAINIUM: ... escape hatches that defeat the purpose of a shared abstraction.

Alternatively, a portable DSL that hides these differences produces generic code that doesn’t exploit any of them — and you’re back at the compiler layer, hoping the compiler is smart enough to recover the performance you left on the table.

The DSL Proliferation Is the Proof

If portability at the kernel level were achievable, the market would have converged on one DSL by now. Instead, every hardware vendor has independently built their own:

NVIDIA alone has four DSLs for writing kernels — CUDA C, CUTLASS/CuTe DSL, Triton, and CuTile — spanning from thread-level C to tile-level Python. Four abstraction levels for one vendor’s hardware, because even within a single architecture family, one DSL isn’t enough.

AMD’s trajectory is particularly telling. They adopted Triton — the closest thing the industry has to a “portable” kernel DSL — and forked it for ROCm. But adopting Triton wasn’t enough. AMD still built FlyDSL: a Python DSL with explicit register-level layouts, lane-level operations, and CuTe-style tiling algebra, lowered through their own MLIR Fly dialect to ROCDL machine code. Even when you adopt a “portable” DSL, you still end up building a hardware native one to get the performance the portable one can’t deliver.

This isn’t a failure of engineering. It’s evidence that hardware-specific kernel programming is irreducibly hardware-specific.

For Highest MFU: Write the ISA in Python

Here’s the principle that unifies all the successful DSLs: the highest-performing kernel DSLs are thin Python wrappers around the hardware’s instruction set.

Pallas doesn’t abstract away the MXU — it gives you dot() that maps directly to MXU matmul instructions, load()/store() that map to DMA operations, and Ref types that correspond to VMEM/HBM/SMEM memory spaces.

CuTile doesn’t abstract away tcgen05 — it gives you tile_matmul() that maps to tcgen05 MMA, tile_load() that maps to TMA descriptors, and automatic TMEM allocation that generates UTCATOMSWS.FIND_AND_SET instructions.

NKI doesn’t abstract away the tensor engine — it gives you nisa.nc_matmul() that maps to tensor engine instructions and explicit SBUF/PSUM memory staging. When you write NKI, you’re writing the Trainium ISA in Python.

The value of these DSLs is not abstraction — it’s ergonomics. They let you write hardware-specific code in Python instead of C or assembly. The key is that they don’t hide the hardware. They expose it through a familiar syntax.

This is the right trade-off. C-level ISA access (CUDA C++, raw PTX, assembly) gives you maximum control but poor iteration speed and readability. An abstract portable DSL gives you ergonomics but takes away the hardware control that determines MFU. A Python DSL over the hardware ISA gives you both: hardware control with Python iteration speed.

The flashinfer PR is 4 million lines of generated CUDA C++. MaxText’s MegaBlox kernel is 282 lines of Pallas Python. Both achieve high MFU on their respective hardware. The difference in developer experience is staggering — but the reason Pallas works isn’t that it’s more abstract than CUDA. It’s that it maps Python directly to TPU hardware concepts, so the 282 lines express the same level of hardware-specific intent as the generated CUDA, just more concisely.

Recommendations for Hardware Vendors

If you’re building AI accelerator hardware and want developers to achieve peak performance on your silicon, here’s what the evidence from every successful AI hardware platform tells us:

1. Build Your Own Python DSL

Don’t wait for an open-source DSL to support your hardware. Google built Pallas. NVIDIA built CuTile. AWS built NKI. In every case, the vendor built DSL outperforms third-party alternatives because the DSL designers have deep knowledge of the hardware’s performance model.

Your DSL should let developers write kernels that map to your hardware’s execution model in Python. Not C. Not a custom language. Python — because the ML ecosystem lives in Python, and the fastest path from researcher insight to high-performance kernel is a Python file that directly expresses your hardware’s concepts.

Build the DSL for the silicon. Don’t build the silicon for the DSL.

2. Expose the ISA, Not Just High-Level Abstractions

Your DSL should operate at two levels:

Low level: virtual ISA access. Let developers express the hardware’s instruction set — your matmul unit’s tile operations, your memory hierarchy’s explicit staging, your synchronization primitives. This is where peak MFU comes from. Developers who understand your hardware should be able to write near ISA level code in Python and know exactly what instructions will be generated.

High level: compiler-lowered operations. Not all your users will want to write ISA — and that’s fine. For developers who don’t need peak performance or are prototyping, provide higher-level operations (tile_matmul, tile_load) that the compiler lowers to the ISA. This is where Pallas’s dot() and CuTile’s tile_matmul() live — high-level enough to be ergonomic, low-level enough that an expert can predict the generated code.

Critically, the two levels must be mixable in the same kernel. Pallas GPU gets this right: you can write a kernel that uses high level jnp operations for the parts where the compiler is good enough, and drop into explicit tcgen05 instructions for the parts where you need direct hardware control — in the same function, in the same Python file. The developer chooses the abstraction level per-operation, not per-kernel.

This is how you scale adoption. New users start with the high-level ops and get reasonable performance. Expert users drop into ISA-level control for the hot path. A DSL that only provides the high level interface leaves performance on the table for expert users. A DSL that only provides ISA access has too steep a learning curve for adoption. The DSLs that win will let developers slide between both levels seamlessly.

3. Do Not Constrain Your Hardware to an Open-Source DSL

This is the most counterintuitive recommendation, but the evidence is clear.

Triton is a good DSL for NVIDIA GPUs that are not Blackwell or Rubin. It is not a good target for designing your hardware around. If you build your chip to be “Triton-compatible,” you’ve constrained your hardware architecture to the abstractions Triton exposes — shared memory, thread blocks, warp-level operations. Those abstractions reflect NVIDIA’s hardware design choices. They may not reflect yours.

NVIDIA itself doesn’t constrain Blackwell to Triton. CuTile exposes tcgen05, TMEM, and mbarrier — hardware features that don’t exist in Triton’s programming model. Triton will eventually support them, but CuTile had day one access. If NVIDIA doesn’t limit itself to the open-source DSL it helped create, you shouldn’t limit yourself to it either.

The risk of designing hardware for an existing DSL: your chip becomes a less-efficient implementation of someone else’s hardware model. The opportunity of designing a DSL for your hardware: your chip can exploit architectural innovations that no existing DSL anticipated.

Build the DSL for the silicon. Don’t build the silicon for the DSL.

4. Control Your Own IR

When you own your DSL, you own the intermediate representation — and the IR is where hardware specific semantics live. When you adopt someone else’s DSL, you inherit their IR’s assumptions about how hardware works.

If your chip ships a new memory space or a novel synchronization primitive, and you own the IR, you add an op and it’s available to users on day one. The compiler understands it, can optimize around it, can fuse across it, can schedule it relative to other operations. Your new hardware feature is a first class citizen in the optimization pipeline.

If you’re lowering from someone else’s IR — say Triton’s TTIR — your options are worse. You either wait for upstream to add support, hack around it with extension intrinsics, or fork the IR and maintain it yourself.

Triton’s tl.extra.nvidia and tl.extra.cuda extensions illustrate the problem. They give you access to hardware-specific instructions from within Triton — but those calls are escape hatches, not integrated operations. The Triton compiler can’t reason about them, can’t fuse across them, can’t schedule them relative to surrounding Triton ops. You get the syntax of a high level DSL with the semantics of inline assembly.

Compare that to CuTile, where tcgen05 is a first-class IR concept. The compiler understands it, generates TMEM allocation and mbarrier synchronization automatically, and schedules it as part of the overall kernel optimization. The difference isn’t access to the instruction — it’s whether the compiler understands the instruction.

That’s the cost of building on someone else’s IR: your hardware-specific features will always be second-class citizens, bolted on through extensions rather than integrated into the optimization pipeline. This is why NVIDIA built CuTile instead of extending Triton for Blackwell — and it’s why you should own yours too.

5. Build Your Compiler in MLIR — Open Source the Frontend, Keep the Backend

MLIR is often cited as the solution to the portability problem. It’s not — it’s a framework for building compilers, and the compilers people build with it are hardware specific. But that’s exactly why you should use it.

MLIR gives you common infrastructure: parsing, verification, transformation utilities, and a pass pipeline framework that every compiler needs. Instead of building all of that from scratch, you define your own MLIR dialect with ops that match your hardware’s semantics, and plug into the existing ecosystem. Your DSL lowers to your dialect. Your dialect lowers to your backend. MLIR handles the plumbing.

The smart split is: open source your frontend compiler, keep your backend codegen closed. The frontend — the DSL to MLIR lowering, frontend optimization passes, the dialect definition — should be open. This is what developers interact with, debug against, and build tooling around. Transparency here builds trust and adoption.

The backend — the final lowering from your MLIR dialect to machine code, the instruction selection, the register allocation, the hardware specific scheduling — can stay closed. This is where your proprietary IP lives, and keeping it closed doesn’t hurt developer experience because users don’t need to read your codegen to write high performance kernels.

We’re seeing this pattern emerge across the industry. The Mosaic compiler in the JAX repository is open source — developers can read the TPU and GPU Pallas frontend passes, understand how their kernels are lowered, and debug compilation issues. The backend TPU codegen (libtpu) is closed. AWS’s NKI compiler is on the path to being open sourced at the frontend level, with the Trainium backend remaining proprietary.

This is the right model. MLIR makes it cheap to build a non portable compiler. Open-sourcing the frontend makes it usable. Keeping the backend closed protects your IP. Everyone wins.

The Implication for Frontier Labs

If you’re building a frontier training or inference stack, the evidence from every successful large scale deployment says the same thing:

Pick your hardware. Go deep. Don’t plan to port.

The “we’ll stay hardware agnostic and port later” strategy sounds prudent but costs you twice: you lose MFU today because you avoid hardware specific optimizations, and you spend engineering time later porting code that can’t actually be ported — it has to be rewritten.

MaxText achieves high performance on TPU v6e for MoE pretraining because it uses TPU native Pallas kernels, TPU native XLA fusion, and TPU native SPMD partitioning. The flashinfer MoE kernels will achieve peak MFU on SM100 because they’re 4 million lines of SM100 specific CUDA. Neither can run on the other’s hardware. Neither should.

The math is portable. Write your model in JAX or PyTorch — the jnp.dot and torch.matmul calls work everywhere. But the moment you need peak performance — the attention kernel, the MoE routing, the collective overlap, the memory scheduling — you’re writing for one hardware target. Accept it early, and you’ll ship faster than the team that’s still trying to write kernels that run everywhere and run fast nowhere.

Inference: Where This Matters Even More

The portability argument applies to training, but it’s even more acute for inference. In training, you amortize kernel performance over weeks or months of runs. In inference, kernel performance is your cost per token — and cost per token is the metric that determines whether your product is economically viable at scale. Every percentage point of MFU you leave on the table by using a “portable” kernel instead of a hardware native one translates directly into higher serving costs — multiplied by every token you generate, forever.

The Cost of Kernel Development Is Collapsing

One historical argument for portability was that writing hardware-specific kernels was too expensive. If it takes a team of experts six months to write a custom attention kernel, you can’t afford to do it for every hardware target. Better to write it once and port.

That calculus is changing. Agentic AI coding tools are dramatically reducing the cost of kernel development. An engineer with a Python DSL and an AI coding assistant can iterate on kernel implementations in hours instead of weeks. The DSLs themselves — Pallas, CuTile, NKI — already reduced the barrier from “CUDA C++ expert” to “Python developer who understands the hardware.” AI tooling reduces it further: generate a kernel skeleton, profile it, have the agent iterate on tiling and memory staging, converge on high MFU code in a fraction of the time.

When kernel development cost was high, portability was an economic argument: write once, amortize across targets. When kernel development cost is low, that argument collapses. It becomes cheaper to write two hardware native kernels that each achieve 90%+ MFU than to write one “portable” kernel that achieves 30% MFU everywhere and needs constant tuning to close the gap.

Cost Per Token Is All That Matters

At the end of the day, at frontier scale the economics are driven by one metric: cost per token. Frontier labs and organizations will use whatever silicon provides the best cost per token — and is usable. That second part is the bottleneck. The silicon that wins isn’t necessarily the silicon with the best theoretical FLOP/s. It’s the silicon where an engineer can write a high MFU kernel in a week instead of a quarter. A solid Python DSL that lets you write the ISA directly is a massive step in the usability direction — and usability is what determines whether your hardware gets adopted or sits in a rack collecting dust.

Conclusion

The frameworks are portable. jnp.dot runs on TPU, GPU, and Trainium. That’s a genuine achievement, and it matters — your model architecture, your training recipe, your evaluation harness all transfer across hardware.

The kernels are not. 282 lines of Pallas and 4 million lines of CUDA solve the same math with zero shared code, and that gap isn’t closing — it’s widening with every generation of silicon. Blackwell introduced tcgen05 and TMEM. TPU v6e has a different MXU layout than v5e. Trainium’s tensor engine has its own tiling constraints. Each generation adds hardware specific concepts that reward hardware-specific code.

The industry should stop treating this as a problem to solve and start treating it as a reality to build around. Hardware vendors: build your own Python DSL, control your own IR, open source your frontend compiler. Frontier labs: pick your hardware and go deep — the cost of writing native kernels is collapsing, and the cost of not writing them is millions in wasted compute per year.

The race isn’t to build the most portable stack. It’s to build the deepest one.

Repo: https://github.com/patrick-toulme/pyptx

Docs: https://pyptx.dev/

Connect on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Launching pyptx — a Python DSL for writing NVIDIA PTX kernels directly.

Today I’m open-sourcing a project I’ve been building on personal time: pyptx, a Python DSL where the function body is the PTX instruction stream. One PTX instruction = one Python call. No optimizer, no autotuner, no tile IR between you and the hardware.

Why? Because the newest GPU features — Hopper’s wgmma, TMA multicast, mbarrier-based pipelines, Blackwell’s tcgen05.mma + TMEM + cooperative 2-SM MMA often only exist at the PTX level. For developers chasing peak performance, that has historically meant writing inline PTX inside CUDA C++.

pyptx brings that whole path into Python. Callable from JAX (via typed XLA FFI) and PyTorch (eager, torch.compile, and a C++ extension fast path).

A few numbers from real silicon:

• H100 bf16 GEMM: 815 TFLOPS, competitive with cuBLAS at matrix sizes ≥ 6K

• B200 bf16 GEMM: 1240 TFLOPS on the 1SM kernel

• RMSNorm: 2.6 TB/s (88% of HBM3 peak, 3.9× PyTorch eager)

• SwiGLU: 2.8 TB/s (94% HBM3)

The other half of the project is a transpiler in the opposite direction:

python -m pyptx.codegen kernel.ptx --sugar

takes PTX from anywhere — nvcc, Triton, CUTLASS output, DeepGEMM kernels — and emits editable pyptx Python. The parser/emitter round-trips byte-identical on 218+ real-world kernels. So you can read someone else’s kernel as Python, modify it, and ship the result.

Built end-to-end: parser, IR, emitter, transpiler, JAX integration, PyTorch integration, full Hopper + Blackwell ISA coverage, multi-arch wheels published to PyPI. ~17K lines of Python total.

Ships with maintained GEMM, grouped GEMM, RMSNorm, LayerNorm, and SwiGLU kernels for both Hopper and Blackwell, plus the PTX → Python transpiler.

pip install pyptx[torch] # for PyTorch

pip install pyptx[jax] # for JAX

pip install pyptx[all] # both

Repo: https://github.com/patrick-toulme/pyptx

Docs: https://pyptx.dev

If you write GPU kernels — especially if you’ve ever wished Triton would let you express a specific wgmma pattern, or wanted to read a CUTLASS PTX dump as editable Python — try it. PRs welcome, especially Blackwell tuning and new ISA wrappers.

Full Code / IR Dumps: https://github.com/patrick-toulme/justabyte/tree/main/maxtext_pretraining

Connect on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Motivation

3.5 billion parameters. 16 MoE experts with top-2 routing. 2-way FSDP sharded across 2 chips, expert-parallel across the other 2 chips. 200 training steps at 22 TFLOP/s per device — every step generating 887 fused kernels, 104 async all-gathers, 12 ragged all-to-all collectives, 8 Pallas attention kernels, and 24 MegaBlox grouped matmuls. All compiled, scheduled, and overlapped automatically from a single python3 -m maxtext.trainers.pre_train.train command.

The cost of this entire experiment was under 10 dollars.

I keep watching startups, neo-frontier labs and major organizations spend months rebuilding training infrastructure from scratch in PyTorch on NVIDIA GPUs. FSDP sharding, kernel fusion, collective overlap, MoE routing. Team after team duplicating the same work, and it's not clear most of them even achieve as high an MFU as MaxText gets out of the box. Meanwhile, I rarely hear MaxText mentioned as even an option.

I’m not claiming a frontier lab can take MaxText off the shelf and ship. But I am claiming they could get to a production training job dramatically faster by building on top of it — and this post shows why. I understand some organizations have bespoke requirements or privacy reasons that prevent them from using open source — this post isn't aimed at them.

This experiment runs on 4 chips, but MaxText and JAX's SPMD model scale with no code changes. You change ici_fsdp_parallelism=2 to ici_fsdp_parallelism=128 and the same code runs on 128 chips. The compilation pipeline I traced here — the fusions, the async collectives, the VMEM scheduling — is the same compilation pipeline at any scale. The only things that change are the mesh dimensions in a config flag.

I ran a GPT-OSS MoE pretraining job on 4 TPU v6e chips, dumped the IR, and traced the full compilation pipeline from Python down to fused TPU kernels. What I found: the XLA compiler does a staggering amount of work that would take months of manual kernel engineering. The fusion, the async scheduling, the memory management, the SPMD partitioning — it’s all generated automatically. This is the infrastructure teams are rebuilding by hand.

Setup

Hardware: a single TPU v6e-4 node – 4 TPU v6e chips in a 2×2 mesh, each with 32 GB HBM. Total: 128 GB HBM across the node.

[TpuDevice(id=0, coords=(0,0,0)), TpuDevice(id=1, coords=(1,0,0)),
 TpuDevice(id=2, coords=(0,1,0)), TpuDevice(id=3, coords=(1,1,0))]

Software: MaxText cloned from AI-Hypercomputer/maxtext, installed with pip install -e ".[tpu]". JAX 0.9.0, libtpu 0.0.36. Nothing else.

The model: a scaled variant of GPT-OSS (OpenAI’s recently open-sourced MoE architecture). The released gpt-oss-20b uses 32 experts with 4 active experts per token; I adjusted width/depth and used 16 experts with top-2 routing to fit on 4 chips—the point here is tracing the compiler + SPMD stack, not exact parity with the released checkpoints.

The parallelism: 2-way FSDP × Expert Parallelism-2. Weights are sharded across 2 chips (FSDP), and the 16 experts are partitioned across 2 chips (EP). This requires real multi-axis communication – all-gather for FSDP weight reconstruction and ragged all-to-all for MoE token dispatch.

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  model_name="gpt-oss-20b" override_model_config=true \
  base_emb_dim=2048 base_num_decoder_layers=16 \
  num_experts=16 num_experts_per_tok=2 \
  megablox=true sparse_matmul=true capacity_factor=-1.0 \
  ici_fsdp_parallelism=2 ici_expert_parallelism=2 \
  attention='flash' remat_policy=full dtype=bfloat16 \
  dump_hlo=true dump_jaxpr=true ...

That ici_fsdp_parallelism=2 ici_expert_parallelism=2 line is the entire parallelism specification. The same codebase handles single chip debugging and multi thousand chip production training – the parallelism strategy is a configuration parameter, not code.

Note: 2-way FSDP × 2-way EP is not the most optimal parallelism configuration for this model and chip count. The point of this experiment is demonstrating the capability of the stack, not maximizing utilization or FLOP/s.

What is MaxText?

MaxText is Google’s open-source JAX/Flax training framework for large language models. It sits at AI-Hypercomputer/maxtext and is the reference implementation for training on TPUs.

The architecture has four layers:

Configuration. MaxText uses OmegaConf + Pydantic for model config. You pick a base architecture with model_name="gpt-oss-20b" (which selects the decoder layer class, attention pattern, MoE routing style) and then override any dimension with CLI flags: base_emb_dim=2048 base_num_decoder_layers=16 num_experts=16.

Model layers (Flax NNX). The model definition is pure Python using Flax’s NNX API. GptOssDecoderLayer in models/gpt_oss.py is 289 lines – it composes an Attention module, an RMSNorm, and a RoutedMoE with residual connections and dropout. The scannable block wraps multiple decoder layers with alternating attention patterns (local sliding-window, global) for nn.scan compilation.

Custom kernels (Pallas). Performance-critical operations – grouped matrix multiplication for MoE, flash attention – are hand-written Pallas kernels. Pallas is JAX’s kernel authoring language that compiles to tpu_custom_call ops. These are the only hand-written kernels in the stack; everything else is generated by XLA.

Parallelism (JAX SPMD). MaxText defines a 13-axis logical mesh (fsdp, expert, tensor, data, sequence, …) and annotates tensors with nn.with_logical_constraint. JAX’s Shardy partitioner reads these annotations and automatically inserts all-gathers, reduce-scatters, and ragged all-to-all collectives.

The key insight: MaxText is a thin layer. The model definition is high level Python. The kernels handle the MXU-level compute. Everything in between – fusion, scheduling, communication, memory management – is the compiler’s job.

Training Results

step:   0, seconds: 30.309, TFLOP/s/device:  0.133, loss: 10.871  # compilation
step:   1, seconds:  0.531, TFLOP/s/device:  7.598, loss: 10.871  # warmup
step:   3, seconds:  0.182, TFLOP/s/device: 22.194, loss: 10.730  # steady state
step:  50, seconds:  0.180, TFLOP/s/device: 22.441, loss:  0.008
step: 100, seconds:  0.180, TFLOP/s/device: 22.446, loss:  0.002
step: 199, seconds:  0.180, TFLOP/s/device: 22.368, loss:  0.001

Step 0 takes 30 seconds – XLA compilation. By step 3, we’re at steady state. Note: this run uses dataset_type=synthetic, so the rapid loss collapse is expected (fast memorization). The goal is validating throughput and end-to-end correctness of MoE + collectives + optimizer, not model quality. 182 ms per step, 22.2 TFLOP/s per device, 5,640 tokens/s/device. The loss drops from 10.87 to 0.001 – synthetic data memorization confirming the full pipeline works: forward through 16 MoE layers, loss computation, backward through MegaBlox kernels and routing gradients, gradient reduce-scatter across devices, and AdamW optimizer update.

One Binary Per Step

Here’s something that surprises people coming from PyTorch: the entire training step – forward pass through 16 MoE decoder layers, cross-entropy loss, backward pass through every layer including custom MegaBlox VJPs, gradient reduce-scatter across all 4 devices, and AdamW optimizer update – is compiled into a single XLA binary.

This is what @jax.jit does. When MaxText calls train_step(state, batch), JAX traces the full Python function into a Jaxpr (18,079 lines of functional IR). XLA then compiles that Jaxpr into a single HLO module (22,636 lines after optimization). That module is compiled once into a TPU binary. Every subsequent step just re-executes the same binary with new inputs.

This is why step 0 takes 30 seconds and step 1 takes 0.5 seconds. Step 0 is compilation. Steps 1-199 are execution.

The consequence: XLA sees everything at compile time. That turns the training step into one schedulable program instead of a chain of compiled islands. XLA can fuse across the forward/backward/optimizer boundary, and it can overlap communication with compute by issuing async all-gathers early and consuming them later—e.g., pulling layer 5’s shards while layer 3’s backward pass runs. With separate compilation units, those boundaries act like barriers: more launches, more forced materialization, and less automatic comm/compute overlap.

The Compilation Stack: Four IR Layers

Between python3 -m maxtext.trainers.pre_train.train and the TPU executing 182 ms training steps, the code passes through four IR layers. MaxText’s dump_hlo=true and dump_jaxpr=true flags capture all of them.

Layer 1: Jaxpr

JAX traces the Python training step into Jaxpr – a functional IR that captures the computation graph. Our train step produces an 18,079-line Jaxpr with 207 helper function definitions encoding the full MoE routing logic:

train_step → jit → TransformerLinenPure.apply → Decoder.__call__
  → scan (16 layers) → GptOssScannableBlock.__call__
    → GptOssDecoderLayer.__call__
      → shard_map → Attention (splash_mha Pallas kernel)
      → shard_map → RoutedMoE.__call__
        → permute → ragged_all_to_all → gmm (MegaBlox Pallas kernel) → unpermute
  → cross_entropy_with_logits
  → TrainState.apply_gradients → adam → reduce_scatter

The Jaxpr reveals the primitive inventory: 112 dot_general (matrix multiplies), 120 sharding_constraint (SPMD annotations), 32 custom_vjp_call (custom backward passes for GMM and attention), 16 shard_map (manual sharding regions), 12 pallas_call (custom TPU kernels), and 16 scan iterations.

The mesh configuration appears directly in the Jaxpr:

ctx_mesh=Mesh('diloco': 1, 'data': 1, 'stage': 1, 'fsdp': 2,
              'fsdp_transpose': 1, 'sequence': 1, 'context': 1,
              'context_autoregressive': 1, 'tensor': 1,
              'tensor_transpose': 1, 'tensor_sequence': 1,
              'expert': 2, 'autoregressive': 1)

Thirteen logical mesh axes, but only two are active: fsdp=2 and expert=2. XLA’s SPMD partitioner reads these annotations and generates all the necessary collectives.

Layer 2: HLO Before Optimization

Jaxpr lowers to HLO (High-Level Operations), XLA’s main IR. Before optimization, our train step is 14,541 lines of HLO with 11,207 individual instructions and zero fusion blocks.

This is the key: everything is separate. Every add, multiply, convert, broadcast, reduce is its own HLO instruction. Let me show you what this actually looks like.

RMSNorm – 10 separate instructions for one normalization:

%convert.105 = f32[4,1024,2048] convert(%input_bf16)           // bf16 → f32
%square.4   = f32[4,1024,2048] multiply(%convert.105, %convert.105)  // x²
%reduce.92  = f32[4,1024]      reduce(%square.4, zero), dims={2}     // Σx²
%reshape.1  = f32[4,1024,1]    reshape(%reduce.92)              // broadcast prep
%div.48     = f32[4,1024,1]    divide(%reshape.1, 2048.0)       // mean(x²)
%add.217    = f32[4,1024,1]    add(%div.48, 1e-5)               // + ε
%rsqrt.4    = f32[4,1024,1]    rsqrt(%add.217)                  // 1/√(mean(x²)+ε)
%bcast.1    = f32[4,1024,2048] broadcast(%rsqrt.4)              // expand
%mul.128    = f32[4,1024,2048] multiply(%convert.105, %bcast.1) // x * scale
%convert.106 = bf16[4,1024,2048] convert(%mul.128)              // f32 → bf16

Ten HLO instructions. Each one reads its input from HBM and writes its output back to HBM. For a f32[4,1024,2048] tensor, that’s 32 MB per read/write. This RMSNorm alone would generate 320 MB of memory traffic without fusion. There are 32 RMSNorm instances in this model (2 per layer × 16 layers).

Adam optimizer – 15 separate instructions per parameter update:

%mul.1073  = f32[2048,8] multiply(grad, 0.1)               // (1-β₁) * grad
%mul.1074  = f32[2048,8] multiply(mu_old, 0.9)             // β₁ * μ
%add.1253  = f32[2048,8] add(%mul.1073, %mul.1074)         // μ_new
%div.614   = f32[2048,8] divide(%add.1253, bias_correction) // μ̂
%pow.59    = f32[2048,8] multiply(grad, grad)               // grad²
%mul.1155  = f32[2048,8] multiply(%pow.59, 0.05)            // (1-β₂) * grad²
%mul.1156  = f32[2048,8] multiply(nu_old, 0.95)             // β₂ * ν
%add.1294  = f32[2048,8] add(%mul.1155, %mul.1156)          // ν_new
%div.696   = f32[2048,8] divide(%add.1294, bias_correction) // ν̂
%sqrt.64   = f32[2048,8] sqrt(%div.696)                     // √ν̂
%add.1336  = f32[2048,8] add(%sqrt.64, 1e-8)               // √ν̂ + ε
%div.759   = f32[2048,8] divide(%div.614, %add.1336)       // μ̂/(√ν̂+ε)
%mul.1219  = f32[2048,8] multiply(param, 0.1)               // weight decay
%add.1377  = f32[2048,8] add(%div.759, %mul.1219)          // update + WD
%add.1422  = f32[2048,8] add(param, lr * %add.1377)        // θ_new

This pattern repeats for 40+ parameter tensors. That’s 600+ element-wise HLO instructions just for the optimizer.

Layer 3: HLO After Optimization

After XLA’s optimization passes, the picture changes dramatically:

The optimized IR is longer because XLA outlines each fused computation as a separate block, but those 11,207 individual instructions have been compressed into hundreds of fused kernels. Each fusion reads its inputs from HBM once, executes all operations in VMEM, and writes outputs once.

Let me show you what these fusions look like.

Deep Dive: What XLA Actually Fuses

Before looking at individual fusions, here’s the full picture.

The backward pass dominates: 492 fusions vs. 274 forward. This isn’t surprising – the backward pass through MoE requires separate gmm (input gradient) and tgmm (weight gradient) kernels for each forward GMM, plus chain-rule fusions through every activation, normalization, and residual connection. The 22 backward Pallas kernels vs. 10 forward kernels reflect the same 2:1+ ratio at the custom kernel level.

The optimizer is compact – just 41 fusions – but each one is a monster. A single AdamW fusion handles gradient scaling, both moment EMAs, bias correction, weight decay, and the parameter update for an entire tensor. The L2 norm reduction for gradient clipping is fused in too.

Nearly a quarter of all fusions (188 out of 832) contain bf16↔f32 type conversions. The model computes in bf16 for throughput but accumulates in f32 for numerical stability. XLA fuses these conversions into the surrounding operations so they never hit HBM as separate ops.

Now let me show you what these fusions look like in practice.

Fusion 1: Logits Matmul + Softmax Prep (kOutput fusion)

This is where XLA goes beyond element-wise fusion. The forward logits computation – RMSNorm scaling, the [1024,2048] × [2048,32768] matmul into vocabulary space, and the reduce_max for softmax numerical stability – all in a single kOutput fusion:

%fused_computation.1802 (weights: bf16[2048,32768],
    norm_weight: bf16[2048], rsqrt_scale: f32[1024],
    activation: bf16[1,1024,2048])
    -> (bf16[1024], bf16[1024,32768]) {
  // RMSNorm: scale activation (nested kLoop fusion)
  %normed    = fusion(norm_weight, rsqrt_scale, activation)  // bf16[1024,2048]
  %w_reshaped = fusion(weights)                               // layout bitcast

  // THE MATMUL: [1024,2048] × [2048,32768] → logits
  %logits    = bf16[1024,32768] convolution(%normed, %w_reshaped),
                 dim_labels=bf_io->bf                         // logits_dense forward

  // FUSED: reduce_max over vocab dim (softmax numerics)
  %row_max   = bf16[1024] reduce(%logits, -inf), dimensions={1}

  ROOT tuple(%row_max, %logits)
}

The reduce_max consumes logits tiles as the convolution (matmul) produces them – the full bf16[1024,32768] tensor (64 MB) never needs a separate read pass. In the unfused version, XLA would write the logits to HBM, then read them back for the max. The kind=kOutput annotation means this fusion is anchored on the convolution – XLA built the fusion outward from the matmul, pulling in both its input preparation (RMSNorm) and its consumer (reduce_max).

Fusion 2: Backward Logits Matmul + RMSNorm Gradient (kOutput fusion)

The backward pass shows an even deeper kOutput fusion. The gradient flows backward through the logits projection and directly into the RMSNorm gradient – matmul and normalization backward fused into one kernel:

%fused_computation.1488 (activation: bf16[1,1024,2048],
    norm_weight: bf16[2048], weights: bf16[2048,32768],
    dLogits: bf16[1024,32768], row_max: f32[1024],
    softmax_denom: f32[1024], labels: bf16[1024],
    label_indices: s32[1024], loss_scale: f32[1024])
    -> (f32[1024], bf16[1024,2048]) {
  // Softmax gradient correction (nested kLoop fusion)
  %dLogits_corrected = fusion(dLogits, row_max, softmax_denom,
                              labels, label_indices, loss_scale)  // bf16[1024,32768]

  // THE MATMUL: dLogits × W^T → dActivation
  %dAct      = bf16[1024,2048] convolution(%dLogits_corrected, weights),
                 dim_labels=bf_oi->bf                     // backward logits_dense

  // FUSED: RMSNorm backward -- no HBM round-trip after matmul
  %scaled    = bf16[1,1024,2048] multiply(bitcast(%dAct), broadcast(norm_weight))
  %f32_grad  = f32[1,1024,2048] convert(%scaled)         // bf16 → f32
  %f32_act   = f32[1,1024,2048] convert(activation)      // bf16 → f32
  %chain     = f32[1,1024,2048] multiply(%f32_act, %f32_grad)  // chain rule
  %grad_sum  = f32[1024] reduce(%chain, 0.0), dimensions={0,2} // RMSNorm weight grad

  ROOT tuple(%grad_sum, %dAct)
}

The [1024,32768] × [32768,2048] backward matmul produces dAct, which is immediately consumed by the RMSNorm gradient chain: scale by norm weight, convert bf16→f32, elementwise multiply with saved activations, reduce-sum over hidden dim. Seven operations after the matmul, all executing on the convolution’s output tiles without an HBM round-trip. That’s 64 MB of avoided intermediate traffic for the matmul output alone.

Fusion 3: Grouped Attention Conv + Residual + RMSNorm Backward (kOutput fusion)

The deepest matmul fusion in the model. Inside the decoder backward pass, XLA fuses a 32-head grouped convolution with residual gradient accumulation and the full RMSNorm backward chain:

%fused_computation.372 (saved_activations: bf16[8,1,1024,2048],
    layer_idx: s32[], norm_weight: bf16[2048],
    O_weight: bf16[1,2048,32,64],
    dQ_heads: bf16[1,1024,32,32], dK_heads: bf16[1,1024,32,32],
    residual_grad_1: bf16[1024,2048,1],
    residual_grad_2: bf16[1024,2048,1])
    -> (f32[1024], bf16[1024,2048]) {
  // Extract this layer's activations from scan buffer
  %act_slice  = dynamic-slice(saved_activations, layer_idx, 0, 0, 0)  // bf16[1,1024,2048]
  %f32_act    = f32[1,1024,2048] convert(%act_slice)                   // bf16 → f32

  // Accumulate two residual gradient streams
  %residual   = bf16[1024,2048] add(residual_grad_1, residual_grad_2)

  // Prepare attention head gradients and weights
  %dHeads     = fusion(dQ_heads, dK_heads)         // bf16[1024,32,64]: merged Q+K grads
  %W_reshaped = fusion(O_weight)                   // bf16[2048,32,64]: reshaped O projection

  // THE MATMUL: 32-head grouped attention projection backward
  %dHidden    = bf16[1024,2048,1] convolution(%dHeads, %W_reshaped),
                  window={size=32}, dim_labels=b0f_o0i->bf0       // 32 attention heads

  // Add matmul result to residual chain
  %combined   = bf16[1024,2048] add(%residual, bitcast(%dHidden))

  // RMSNorm backward: scale, convert, chain-rule multiply, reduce
  %scaled     = bf16[1,1024,2048] multiply(bitcast(%combined), broadcast(norm_weight))
  %f32_grad   = f32[1,1024,2048] convert(%scaled)                 // bf16 → f32
  %chain      = f32[1,1024,2048] multiply(%f32_act, %f32_grad)    // chain rule
  %grad_sum   = f32[1024] reduce(%chain, 0.0), dimensions={0,2}   // RMSNorm weight grad

  ROOT tuple(%grad_sum, %combined)
}

Count the operations: dynamic-slice from the scan buffer, bf16→f32 convert, two residual adds, a 32-head grouped convolution (window={size=32} – this is how TPU HLO represents grouped matmuls), broadcast, multiply, bf16→f32 convert, multiply, reduce-sum. ~20 operations across three distinct algorithmic stages (residual accumulation, attention projection, normalization backward) in a single kernel. The dynamic-slice indexed by layer_idx is particularly notable – this fusion runs inside the scan while-loop, extracting per-layer activations from the recomputation buffer each iteration.

Fusion 4: Full AdamW Update (~40 ops → 1 kernel)

The entire AdamW optimizer step for a single MoE expert weight tensor – gradient clipping, both moment EMAs, bias correction, weight decay, parameter update, and L2 norm computation – all in one fusion:

%fused_computation.1373 (weight: f32[8,8,2048,1024],
    lr: f32[], beta1_correction: f32[], beta2_correction: f32[],
    nu_prev: f32[8,8,2048,1024], grad_scale: f32[],
    mu_prev: f32[8,8,2048,1024], is_finite: pred[],
    gradient: f32[8,8,2048,1024])
    -> (f32[], f32[...], f32[...], f32[...], f32[]) {
  // Gradient clipping: zero out if non-finite, else scale
  %scaled_grad = divide(gradient, grad_scale)
  %safe_grad   = select(is_finite, %scaled_grad, gradient)
  // First moment: μ_t = 0.9·μ_{t-1} + 0.1·g
  %mu_new      = add(multiply(%safe_grad, 0.1), multiply(mu_prev, 0.9))
  // Second moment: ν_t = 0.95·ν_{t-1} + 0.05·g²
  %grad_sq     = multiply(%safe_grad, %safe_grad)
  %nu_new      = add(multiply(%grad_sq, 0.05), multiply(nu_prev, 0.95))
  // Bias-corrected update: μ̂/(√ν̂ + ε)
  %nu_hat      = divide(%nu_new, beta2_correction)
  %update      = divide(%mu_new, multiply(beta1_correction, add(sqrt(%nu_hat), 1e-8)))
  // Weight decay + parameter update
  %new_w       = add(weight, multiply(lr, add(%update, multiply(weight, 0.1))))
  // L2 norms for gradient clipping and logging
  %w_norm      = reduce(multiply(%new_w, %new_w), dims={0,1,2,3})
  %g_norm      = reduce(%grad_sq, dims={0,1,2,3})

  ROOT tuple(%w_norm, %new_w, %nu_new, %mu_new, %g_norm)
}

The tensor shape f32[8,8,2048,1024] is an MoE expert weight: [num_experts=8, layers_per_scan=8, hidden=2048, ffn=1024]. That’s a 512 MiB tensor. Without fusion, the 15+ intermediate reads and writes would generate ~8 GB of HBM traffic per parameter update. With fusion, it’s three reads (weight, gradient, both moments) and three writes (new weight, new μ, new ν). Two full reductions (weight norm, gradient norm) are computed in the same pass. This pattern repeats for every parameter in the model.

Deep Dive: How MoE Compiles

This is the most complex part of the compilation pipeline. The MoE layer transforms from ~500 lines of Python into a choreographed sequence of sorts, collectives, and Pallas kernels that spans all 4 TPU chips simultaneously.

The Python: What the Developer Writes

The MoE forward pass in MaxText (layers/moe.py) follows this flow:

# 1. Gate: project to expert logits
gate_logits = self.gate(inputs)                    # [batch, seq, num_experts]
top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, k=2)
top_k_weights = jax.nn.softmax(top_k_weights)

# 2. Permute: sort tokens by assigned expert
sorted_inputs = _sort_activations(inputs, argsort(experts))
group_sizes = jnp.bincount(experts, length=num_experts)

# 3. Dispatch: send tokens to expert-owning devices
x = jax.lax.ragged_all_to_all(sorted_inputs, offsets, sizes, ...)

# 4. Expert FFN: three grouped matrix multiplies
gate_out = gmm(x, w0, group_sizes, tiling=(512,1024,1024))  # gate proj
up_out   = gmm(x, w1, group_sizes, tiling=(512,1024,1024))  # up proj
hidden   = silu(gate_out) * (up_out + 1)                     # activation
output   = gmm(hidden, wo, group_sizes, tiling=(512,1024,1024))  # down proj

# 5. Combine: send results back and weight-sum
x = jax.lax.ragged_all_to_all(output, ...)         # reverse dispatch
output = unpermute(x, argsort(experts))
result = einsum("BKE,BK->BE", output, weights)     # weighted combination

The gmm function wraps a Pallas kernel with a custom VJP: _gmm_fwd calls the kernel for the forward pass, _gmm_bwd calls gmm with transpose_rhs=True for the input gradient and tgmm (a separate kernel) for the weight gradient. Three different tiling configurations optimize each independently.

The shard_map wrapper around this code gives each device a local view, and ragged_all_to_all – a JAX primitive – handles the variable-length token shuffle: each device sends different numbers of tokens to each peer based on the routing decisions.

The HLO: What XLA Compiles

Here’s what the MoE pipeline looks like in the optimized HLO. I’ve annotated each step:

Step 1 – Token routing (sort by expert):

%sort.508 = (s32[2048], s32[2048]) sort(%experts, %iota), dimensions={0}, is_stable=true

Step 2 – Sort activations to match expert ordering:

%gather_custom_fusion.263 = bf16[2048,2048] fusion(%activations, %sort_indices)

Step 3 – Dispatch tokens via ragged all-to-all:

%ragged_all_to_all.295 = bf16[4096,2,1024] ragged-all-to-all(
    %sorted_tokens,       // bf16[2048,2,1024]: local tokens
    %output_buffer,       // bf16[4096,2,1024]: receive buffer (2x for worst case)
    %send_sizes,          // s32[2]: how many tokens to send to each peer
    %send_offsets,        // s32[2]: where they start in my buffer
    %recv_sizes,          // s32[2]: how many I'll receive from each peer
    %recv_offsets),       // s32[2]: where to put them
    channel_id=1,
    replica_groups={{0,1},{2,3}}

The replica groups {{0,1},{2,3}} show the expert-parallel communication pattern: chips 0↔1 and 2↔3 exchange tokens along the expert axis of the 2×2 mesh.

Step 4 – GMM forward (Pallas custom call):

%gmm.1 = bf16[4096,2048] custom-call(
    %group_index,           // s32[]: scalar loop state
    %group_sizes,           // s32[9]: tokens per expert (+padding)
    %group_offsets_lhs,     // s32[15]: tile boundary metadata
    %group_offsets_rhs,     // s32[15]: tile boundary metadata
    %num_actual_groups,     // s32[1]: 8 experts per shard
    %dispatched_tokens,     // bf16[4096,2048]: sorted tokens (LHS)
    %expert_weights),       // bf16[8,2048,2048]: expert matrices (RHS)
    custom_call_target="tpu_custom_call",
    kernel_metadata={"tiling": {"tile_m": 512, "tile_k": 1024, "tile_n": 1024},
                     "cost_estimate": {"flops": "34359738368"}}

Each GMM call does ~34.4 GFLOP – tiling [512, 1024, 1024] across the M (token), K (input), and N (output) dimensions. The group metadata tells the kernel where each expert’s tokens start and end in the sorted input.

Step 5 – SiLU activation (fused into a single kernel):

%fusion.1232 = bf16[4096,2048] fusion(%gmm.0, %gmm_shuffled, %gmm.1, %weights),
    kind=kLoop, calls=%fused_computation.341   // ← the SiLU fusion shown earlier

Step 6 – Combine via reverse ragged all-to-all back to original device mapping.

The Pallas Kernel: What Runs on the MXU

Inside the tpu_custom_call, the Pallas GMM kernel runs on a 3D grid (tiles_n, num_active_tiles, tiles_k):

grid=(tiles_n, num_active_tiles, tiles_k)
# tiles_n=2:      N-dimension tiles (parallel)
# num_active_tiles: M-tiles covering all groups (sequential)
# tiles_k=2:      K-dimension accumulation (sequential)

The first dimension is marked "parallel" – independent axis. The kernel body:

1. Fetch: DMA a [512, 1024] tile of sorted tokens from HBM to VMEM

2. Fetch: DMA the correct expert’s [1024, 1024] weight tile (selected by group_ids[grid_id])

3. Accumulate: acc_scratch += dot(lhs_tile, rhs_tile) in VMEM f32 scratch

4. On last k-tile: Apply group boundary mask and store result to HBM

The boundary mask handles the case where a single tile spans two experts: only rows belonging to the current expert are written.

The Backward Pass: Three Kernels, Not One

The backward pass through MoE requires custom VJPs at the ops.py level. For each of the three forward GMM calls:

• gmm with transpose_rhs=True computes ∂L/∂input = grad @ W^T – the input activation gradient

• tgmm computes ∂L/∂W = input^T @ grad – the weight gradient

The tgmm kernel has a fundamentally different structure: it produces [num_experts, K, N] output, detecting group boundaries to accumulate then store per-expert weight gradients. Its Pallas grid is (tiles_n, tiles_k, num_active_tiles) – note the reordered axes.

In the optimized HLO, you can see both:

// Input gradient: gmm with transposed weights
%gmm.10 = bf16[4096,2048] custom-call(..., %expert_weights), target="tpu_custom_call"

// Weight gradient: tgmm (separate kernel body)
%tgmm.2 = bf16[8,2048,2048] custom-call(
    ..., %saved_activations, %upstream_grad), target="tpu_custom_call",
    kernel_metadata={"tiling": {"tile_m": 512, "tile_k": 1024, "tile_n": 1024},
                     "num_actual_groups": 8}

Deep Dive: Collectives and Async Overlap

The 2×2 mesh creates three distinct communication topologies, all visible in the IR. In total: 104 async all-gathers, 12 ragged all-to-all, 10 plain all-to-all, 8 reduce-scatters, and 22 all-reduces – 156 collective operations in a single training step.

Pattern 1: Async Collective Fusions (Compute + Comms Overlap)

This is the most sophisticated pattern. XLA doesn’t just make all-gathers async – it fuses them with compute into a single pipelined operation. While the current layer’s backward matmul runs, the all-gather for the next layer’s FSDP-sharded weights is in flight:

%async_collective_fusion.1497 (
    shard: bf16[1,8,2048,1024],     // FSDP-sharded expert weight (half)
    full_weight: bf16[1,8,2048,2048],  // previous all-gather result
    semaphore: s32[2],              // collective sync state
    flags: u32[], u32[], ...,       // S(2) flag registers
    gate_bias: bf16[16],            // router bias
    gate_weight: bf16[1,2048,16],   // router projection
    activation: bf16[1024,2048])    // layer input
    -> (bf16[1,1024,16], ..., bf16[1,8,2048,2048], ...) {

  // ===== COMPUTE: router backward matmul =====
  %dGate   = bf16[1024,16] convolution(activation, gate_weight),
               dim_labels=bf_io->bf                    // [1024,2048] × [2048,16]
  %dGate   = add(bitcast(%dGate), broadcast(gate_bias))  // + bias

  // ===== COMMS: all-gather for NEXT layer (overlapped) =====
  %gathered = bf16[1,8,2048,2048] all-gather(shard),
               channel_id=38,
               replica_groups=[2,2]<=[2,2]T(1,0),      // FSDP: {0,2},{1,3}
               dimensions={3},                          // double dim 3: 1024 → 2048
               frontend_attributes={chain_id="4"},      // pipeline ordering
               backend_config={async_collective_fusion_config:
                 {flag_start:"2", flag_end:"8"}}        // semaphore window

  ROOT tuple(%dGate, shard, full_weight, %gathered, semaphore, flags...)
}

The chain_id="4" and flag_start/flag_end annotations reveal XLA’s double-buffered all-gather pipeline: chain IDs sequence the all-gathers across layers, and flag registers manage the handoff between pipeline stages. The semaphore state (s32[2] in memory space S(4)) and flag registers (u32[] in S(2)) are threaded through the tuple output as “continuation state” – passed from one async collective fusion to the next.

There are 39 async collective fusions in total, wrapping 104 individual all-gather operations. The replica_groups=[2,2]<=[2,2]T(1,0) pattern (39 all-gathers) reconstructs FSDP-sharded weights across chips {0,2} and {1,3}. The [1,4]<=[4] pattern (65 all-gathers) gathers across all 4 chips for activations and non-expert parameters.

Pattern 2: MoE Ragged All-to-All (Token Dispatch)

Expert parallelism uses ragged_all_to_all with variable length messages – each device sends different numbers of tokens to each peer depending on the routing decisions:

%ragged_all_to_all.295 = bf16[4096,2,1024]
    ragged-all-to-all(
        %sorted_tokens,       // bf16[2048,2,1024]: local tokens (input)
        %output_buffer,       // bf16[4096,2,1024]: receive buffer (zeros)
        %send_sizes,          // s32[2]: how many tokens to send each peer
        %recv_sizes,          // s32[2]: how many I'll receive
        %send_offsets,        // s32[2]: where they start in my buffer
        %recv_offsets),       // s32[2]: where to put them
    channel_id=1, replica_groups={{0,1},{2,3}},
    barrier_config={"barrier_type":"CUSTOM","id":"3"}

The replica groups {{0,1},{2,3}} show the expert-parallel topology: chips 0↔1 and 2↔3 exchange tokens along the horizontal axis of the 2×2 mesh. There are 12 ragged all-to-all operations (4 forward dispatch + 4 forward combine + 4 backward), plus 8 plain all-to-all exchanging small s32[2,1,1] metadata tensors for token count coordination.

The output lands in S(1) (VMEM) – dispatched tokens go straight to on-chip memory, avoiding an HBM round-trip before the GMM kernels consume them.

Pattern 3: Batched Gradient Reduce-Scatter

After the backward pass, gradients are reduce-scattered back to FSDP-sharded form. XLA batches 10 parameter gradients from 2 transformer layers into a single collective:

%all-reduce-scatter.17 (
    Q_grad:    bf16[2048,8,64],    // Q projection (layer A)
    K_grad:    bf16[32,64,2048],   // K projection (layer A)
    V_grad:    bf16[2048,32,64],   // V projection (layer A)
    O_grad:    bf16[2048,8,64],    // output projection (layer A)
    gate_grad: bf16[2048,16],      // router gate (layer A)
    Q_grad_B:  bf16[2048,8,64],    // Q projection (layer B)
    K_grad_B:  bf16[32,64,2048],   // K projection (layer B)
    V_grad_B:  bf16[2048,32,64],   // V projection (layer B)
    O_grad_B:  bf16[2048,8,64],    // output projection (layer B)
    gate_grad_B: bf16[2048,16])    // router gate (layer B)
    → (bf16[512,8,64], bf16[32,64,512], ...) {
  // All-reduce across all 4 devices, then scatter via partition-id indexing
  %all-reduce = all-reduce(inputs...), replica_groups={{0,1,2,3}}
  // Each device slices its 1/4 shard: dynamic-slice(..., partition_id * 512, ...)
}

Ten gradient tensors, one collective launch. Each output is scattered from 2048 down to 512 along the FSDP dimension via dynamic-slice indexed by partition-id. For MoE expert weights, separate reduce-scatters operate along the FSDP axis only (replica_groups={{0,2},{1,3}}).

The Scheduling Trade-Off: Coalescing vs. Overlap

These patterns reveal a tension in the compiler’s scheduling strategy. Coalescing wants to batch collectives together – the 10 input reduce-scatter amortizes launch overhead by waiting until all 10 gradient tensors from 2 layers are ready, then issuing one large collective instead of 10 small ones. Overlap wants to start collectives as early as possible – the async collective fusions issue all-gathers for layer N+1 while layer N’s compute is still running, hiding latency behind useful work.

XLA makes different choices for different collective types. All-gathers are overlap optimized: each weight gather launches individually inside its own async collective fusion, pipelined with chain_id ordering, because waiting to batch them would stall the compute pipeline. Reduce-scatters are coalescing optimized: batching 10 gradients into one call is worth the delay because the gradients aren’t needed until the optimizer runs, so there’s no compute to overlap them with anyway. Ragged all-to-alls are neither – they’re synchronous barriers because the MoE routing decisions create data-dependent communication patterns that can’t be predicted at compile time.

The scheduler is balancing three constraints: minimize collective launch overhead (coalesce), maximize compute-comms overlap (pipeline early), and respect data dependencies (barrier where required). The fact that it makes different choices for all-gather vs. reduce-scatter vs. ragged all-to-all shows this isn’t a one size fits all heuristic – it’s a per collective scheduling decision informed by the dependency graph.

What This Would Take to Hand-Fuse

This is the infrastructure teams are reimplementing when they rebuild training stacks from scratch. Let me quantify it.

Kernels you’d need to write: The hand-written kernel count alone is significant: 9 for MoE (gate projection, top-k selection, argsort, bincount, GMM forward, gated activation, GMM backward, TGMM for weight gradients, and unsort/combine) and 3 for attention (forward, dQ backward, dKV backward). Each needs its own tiling strategy, VMEM allocation, and DMA scheduling. But that's just the custom kernels — it doesn't include RMSNorm, softmax, embedding, cross-entropy, residual connections, or the AdamW optimizer. On TPU, XLA generates fused kernels for all of those automatically. On a from scratch PyTorch stack, you're writing or tuning those too.

Memory management decisions: ~8 critical choices. VMEM scratch allocation for GMM accumulators. HBM buffer sizing for ragged all-to-all (worst case capacity). Double-buffering strategy for weight prefetch. Activation checkpointing (MaxText uses remat_policy=full – all activations recomputed in backward). Padding strategies for tile boundaries. When to materialize vs. recompute intermediate results.

Communication patterns: ~6 distinct collective types. Ragged all-to-all for MoE dispatch (variable-length). Ragged all-to-all for MoE combine (reverse direction). FSDP all-gather for weight reconstruction. All-gather for routing metadata. Reduce-scatter for gradient sharding. All-reduce for loss aggregation. Each one needs to be overlapped with compute for efficiency.

Custom backward passes: The MoE routing sort requires a custom VJP (_sort_activations_custom_bwd) because JAX’s automatic backward for indexing is inefficient. The GMM needs separate forward/backward kernel implementations. The attention uses three different Pallas kernels for forward, dQ, and dKV.

Interacting parallelism dimensions: FSDP weight sharding, expert partitioning, and data parallelism interact in subtle ways. The weight_gather_axes logic in moe.py handles cases where weights are sharded across FSDP and must be all-gathered before GMM but reduce-scattered in backward.

The MaxText/JAX/XLA stack automates all of this. The developer writes ~500 lines of high-level Python (the RoutedMoE class) and the GMM kernel interface (282 lines). The compiler generates the remaining equivalent of ~50,000+ lines of kernel code, memory management, and communication scheduling.

The numbers tell the story: 11,207 individual HLO instructions compressed into 887 fused kernels. 39 async collective fusions containing 104 all-gathers with compute overlap. 12 ragged all-to-all collectives. 156 total collective operations. 8 Splash Attention calls and 24 GMM calls – the only handwritten code in the entire pipeline. XLA generated everything else.

What This Means

This entire pretraining run — MoE architecture, Pallas kernels, XLA compilation, multi-axis SPMD sharding, async collective overlap — is open source. Every file is at AI-Hypercomputer/maxtext. The model definition: 289 lines of Python. The MoE layer: layers/moe.py. The MegaBlox kernel: 282 lines.

Out of the 22,636 lines of optimized HLO, only the Splash Attention and MegaBlox GMM calls are handwritten kernels. Everything else — 887 fused kernels, 39 async collective fusions wrapping 104 all-gathers, ~960 prefetch operations across 4 memory spaces, the entire SPMD partitioning — is generated by XLA from high level JAX code.

This is what I keep watching teams spend months rebuilding by hand. FSDP overlap, kernel fusion, MoE routing, collective scheduling — the compiler already does it, and it’s already open source. I’m not claiming a frontier lab can take MaxText off the shelf and ship tomorrow. But the distance between this codebase and a production training job is months shorter than starting from scratch on a custom PyTorch stack — and the MFU you get out of the box is likely higher than what most teams achieve after those months of work.

MaxText isn’t a ceiling — it’s a floor, and the floor is already really good. If you’re a frontier lab, you could rent V7 pods, fork MaxText, and have an MoE training loop running at high MFU before your PyTorch team finishes writing their first custom kernel. And if XLA’s codegen isn’t enough for a specific op, you write a Pallas kernel for that op — the same way Splash Attention and MegaBlox already exist in the stack. You’re not replacing the compiler; you’re surgically overriding it where it matters.

The architecture, the kernels, the compiler flags, the SPMD partitioning — it’s all the same code. The only difference between this experiment and a production training run is the number of chips.

Reproduce This

# Clone and install (requires Python 3.12+)
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext && pip install -e ".[tpu]"

# Set dump flags
export DECOUPLE_GCLOUD=TRUE
export XLA_FLAGS="--xla_dump_to=/tmp/xla_dump \
  --xla_dump_hlo_module_re=jit_train_step \
  --xla_dump_hlo_as_text --xla_dump_hlo_as_proto"
export LIBTPU_INIT_ARGS=" \
  --xla_tpu_enable_async_collective_fusion=true \
  --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \
  --xla_tpu_enable_async_collective_fusion_multiple_steps=true \
  --xla_enable_async_all_gather=true \
  --xla_tpu_overlap_compute_collective_tc=true \
  --xla_tpu_scoped_vmem_limit_kib=98304 \
  --xla_tpu_enable_data_parallel_all_reduce_opt=true \
  --xla_tpu_data_parallel_opt_different_sized_ops=true"

# Train (adjust for your chip count)
python3 -m maxtext.trainers.pre_train.train \
  src/maxtext/configs/base.yml \
  model_name="gpt-oss-20b" override_model_config=true \
  dataset_type=synthetic steps=200 \
  base_emb_dim=2048 base_num_decoder_layers=16 \
  base_num_query_heads=32 base_num_kv_heads=8 head_dim=64 \
  base_mlp_dim=2048 base_moe_mlp_dim=2048 \
  num_experts=16 num_experts_per_tok=2 \
  megablox=true sparse_matmul=true capacity_factor=-1.0 \
  per_device_batch_size=1 max_target_length=1024 vocab_size=32768 \
  attention='flash' sa_block_q=512 \
  ici_fsdp_parallelism=2 ici_expert_parallelism=2 \
  remat_policy=full enable_checkpointing=false reuse_example_batch=1 \
  base_output_directory=/tmp/maxtext_output run_name=gpt_oss_3b \
  dump_hlo=true dump_hlo_local_dir=/tmp/hlo_dumps \
  dump_jaxpr=true dump_jaxpr_local_dir=/tmp/jaxpr_dumps \
  gcs_metrics=false

# Examine the IR
ls /tmp/xla_dump/  # HLO before/after optimization
wc -l /tmp/xla_dump/*after_optimizations.txt
grep -c 'fusion' /tmp/xla_dump/*after_optimizations.txt

The IR dumps, training scripts, and full logs from this experiment are at https://github.com/patrick-toulme/justabyte.

If you found this useful, subscribe to Just a Byte for more deep dives into AI compilers, silicon, and systems.

Connect with me on LinkedIn.

Subscribe now

My fully custom TPU now runs a Llama model. If you haven’t seen the original post, I built a TPU from scratch — custom Verilog RTL, custom ISA, custom compiler — with a full PJRT backend so JAX treats it as a real device. The compiler takes the entire model and fuses it into a single megakernel binary. No op-by-op dispatch. No handwritten kernels. Pure compiler codegen.

Full Video of JAX Execution on X:

Getting here wasn’t straightforward. The biggest challenge was accuracy. When your logits don’t match the reference implementation, the bug could be anywhere — the RTL, the compiler, the register allocation, the instruction scheduling. Debugging this through waveform traces alone was brutal, so I built an ISA simulator in Python. It executes the same binary the hardware runs, instruction by instruction, which made it far easier to isolate where values started diverging.

Once accuracy was solid, I shifted to performance. The main optimizations were DMA/compute overlap so the machine isn’t stalling on memory fetches, and enhancements to the VLIW scheduler to pack instructions more tightly. The result is a pipeline that actually keeps the functional units busy instead of sitting idle between operations.

Same stack as before: JAX → HLO → MLIR → ASM → VLIW → Binary. Still no CUDA. Still running on my custom Verilog RTL.

Questions? Comments?

Reach out on X or Linkedin:

X: https://x.com/PatrickToulme

Linkedin: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Subscribe now

I built a TPU from scratch. Custom Verilog RTL. JAX→HLO→MLIR→ASM→VLIW→Binary. Real PJRT backend. Runs on Verilog simulator. No CUDA. No handwritten kernels. Pure compiler codegen. Below is a transformer layer compiling and executing from JAX.

Full Video of JAX Execution on X: https://x.com/PatrickToulme/status/2017706776288719194?s=20

Questions? Comments?

Reach out on X or Linkedin:

X: https://x.com/PatrickToulme

Linkedin: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Subscribe now

Connect on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Motivation

86 lines of Python. 180KB of shared memory. 46 tcgen05 instructions. A contention handling retry loop with nanosleep backoff.

NVIDIA is building a compiler moat, and making a play for the GPU kernel DSL market.

CuTile compiled a Mixture of Experts kernel for Blackwell, and the gap between what I wrote and what the compiler generated is the evidence. Last post, I traced Pallas through the TPU compiler stack. Pallas exists because XLA can’t rewrite your algorithm — it optimizes what you give it, but it can’t infer that you never need to materialize the full attention matrix. CuTile is NVIDIA’s answer: a tile level DSL that lets you express algorithms the compiler can’t find automatically, while hiding the thread level complexity that makes GPU programming painful.

The difference is what CuTile targets. Blackwell’s tcgen05 instruction family is publicly documented. Triton, Pallas, and CUTLASS all support it to some degree. But high performance requires patterns that are hard to discover — single thread MMA issuance via leader election, multi barrier pipelining, TMEM allocation with contention handling (retry/backoff). CuTile generates these automatically. Triton built Gluon as an escape hatch because the high level abstractions couldn’t find the right codegen. CuTile doesn’t need an escape hatch.

This post traces a real MoE kernel through the compilation stack: Python → CuTile IR → [MLIR black box] → PTX → SASS. The compiler is a black box, but the outputs tell you everything.

Setup

All experiments were performed on a B200 GPU rented from Lambda Labs. The B200 is NVIDIA’s Blackwell architecture (sm_100a), requiring:

• CUDA 13.1+ (CuTile is not available on earlier versions)

• NVIDIA Driver 580+ with open kernel modules (nvidia-driver-580-open)

• CuTile Python package from NVIDIA (github source)

Extracting Compilation Artifacts

For the CuTile IR, I used the library’s built in print_ir() functionality. For PTX and SASS, I extracted them from the compiled .cubin files using cuobjdump.

The MLIR intermediate passes are not publicly accessible — they require an internal extension (cuda.tile_internal._internal_cext) that isn’t available with the public release. I extracted the pass names by running strings on the tileiras compiler binary.

The Kernel: NVIDIA’s Fused MoE

The kernel comes from NVIDIA’s CuTile examples — a production quality Mixture of Experts implementation. Here’s the core computation:

@ct.kernel
def fused_moe_kernel(
    A,                     # Input tokens, shape (batch, K)
    B,                     # Expert weights, shape (num_experts, N, K)
    C,                     # Output tensor, shape (num_tokens * topk, N)
    topk_weights,          # Router weights for each token-expert pair
    sorted_token_ids,      # Token indices sorted by expert assignment
    sorted_expert_ids,     # Expert index for each TILE_M
    num_token_replicas: int,
    mul_routed_weight: ConstBool,
    TILE_M: ConstInt,
    TILE_N: ConstInt,
    TILE_K: ConstInt,
):
    M = sorted_token_ids.shape[0]
    N = B.shape[1]
    K = B.shape[2]

    GROUP_SIZE_M = 8
    bid_m, bid_n = swizzle_2d(M, N, TILE_M, TILE_N, GROUP_SIZE_M)

    # Gather token indices for this block
    token_id_indices = bid_m * TILE_M + ct.arange(TILE_M, dtype=ct.int32)
    token_ids = ct.gather(sorted_token_ids, token_id_indices)
    a_row_indices = token_ids // num_token_replicas

    # Each TILE_M block processes one expert
    expert_id = ct.load(sorted_expert_ids, index=bid_m, shape=())

    # Main matmul loop
    accumulator = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32)
    for k in range(0, ct.cdiv(K, TILE_K)):
        a_col_indices = k * TILE_K + ct.arange(TILE_K, dtype=ct.int32)
        a = ct.gather(A, (a_row_indices[:, None], a_col_indices[None, :]))

        b = ct.load(B, (expert_id, k, bid_n), shape=(1, TILE_K, TILE_N),
                    order=(0, 2, 1), padding_mode=ct.PaddingMode.ZERO)
        b = b.reshape((TILE_K, TILE_N))

        accumulator = ct.mma(a, b, accumulator)

    if mul_routed_weight:
        moe_weight = ct.gather(topk_weights, token_ids)
        accumulator = accumulator * moe_weight[:, None]

    # Scatter results back
    c_col_indices = bid_n * TILE_N + ct.arange(TILE_N, dtype=ct.int32)
    accumulator = ct.astype(accumulator, C.dtype)
    ct.scatter(C, (token_ids[:, None], c_col_indices[None, :]), accumulator)

What’s in the Kernel

• Tile level operations: ct.gather, ct.load, ct.mma, ct.scatter

• 2D block swizzling for L2 cache locality

• Irregular memory access patterns (gather/scatter for token routing)

• Compile time constants (ConstInt, ConstBool) for specialization

What’s NOT in the Kernel

• Thread indices (threadIdx.x)

• Shared memory allocation

• Synchronization barriers

• Async copy operations

• Register allocation hints

• The tcgen05 instruction family

The kernel is ~86 lines of tile level Python. The compiler generates ~1,900 lines of PTX.

CuTile IR: The First Layer

CuTile traces the kernel into a typed intermediate representation. Here’s the full IR for fused_moe_kernel:

CuTile IR

Tensor Views and Memory Layout

The IR creates explicit tensor views with stride information:

B{$177, $179, $181, $183, $185, $187, B_6}: Array[bfloat16,(?,?,?):(?,?,1)] =
    make_tensor_view(base_ptr=$177, shape=($179, $181, $183), dynamic_strides=($185, $187))

sorted_expert_ids{$201, $203, sorted_expert_ids_2}: Array[int32,(?):(1)] =
    make_tensor_view(base_ptr=$201, shape=($203), dynamic_strides=())

The (?,?,?):(?,?,1) notation encodes shape and strides — the last dimension has stride 1 (contiguous), while outer dimensions have dynamic strides.

The Swizzle Computation

The 2D swizzle becomes explicit integer arithmetic:

$261: int32 = tile_bid(axis=0)
$307: Tile[int32,()] = raw_binary_arith(lhs=$200, rhs=$357, fn="cdiv", ...)
$312: Tile[int32,()] = raw_binary_arith(lhs=$181, rhs=$371, fn="cdiv", ...)
$315: Tile[int32,()] = raw_binary_arith(lhs=$380, rhs=$312, fn="mul", ...)
$318: Tile[int32,()] = raw_binary_arith(lhs=$389, rhs=$315, fn="floordiv", ...)
...
$331: Tile[int32,()] = raw_where(cond=$431, x=$432, y=$425)
$332: Tile[int32,()] = raw_binary_arith(lhs=$321, rhs=$331, fn="add", ...)

The swizzle_2d() helper expands into ~20 IR operations for group-based block assignment.

The Token System

CuTile uses tokens for memory ordering — similar to Pallas’s implicit ordering through Refs:

$40: Tile[int32,(128)], $token.0: Token =
    load_pointer_tko(pointer=$546, mask=$543, padding_value=$548, token=$token, latency=None)

$578: Tile[int32,(1)], $token.2: Token =
    tile_load_token_ordered(array=sorted_expert_ids{...}, index=($332), token=$token,
                            order=(0,), padding_mode=PaddingMode.UNDETERMINED, latency=None)

The _tko suffix means “token-keyed ordering” — loads and stores carry tokens that establish memory dependencies.

The Main Loop

The k-loop becomes explicit with accumulator threading:

$792: Tile[float32,(128,128)] = for k in range($60, $65, $624)
    (with accumulator.0: Tile[float32,(128,128)] = $58)
do (k: int32, accumulator.0: Tile[float32,(128,128)]):
    $69: int32 = raw_binary_arith(lhs=k, rhs=$169, fn="mul", ...)
    ...
    $91: Tile[bfloat16,(128,64)], $token.8: Token =
        load_pointer_tko(pointer=$717, mask=$713, padding_value=$720, token=$token, latency=None)

    $108: Tile[bfloat16,(1,64,128)], $token.10: Token =
        tile_load_token_ordered(array=B{...}, index=($49, k, $337), token=$token,
                                order=(0, 2, 1), padding_mode=PaddingMode.ZERO, latency=None)

    $113: Tile[bfloat16,(64,128)] = tile_reshape(x=$108)
    $119: Tile[float32,(128,128)] = tile_mma(x=$91, y=$113, acc=accumulator.0)
    continue $119

The tile_mma operation is the matrix multiply-accumulate — this single IR operation will expand into Blackwell's tcgen05.mma instructions.

The Compiler Wall

Everything after CuTile IR is a black box. The MLIR passes, the layout transformations, the async scheduling etc. are not visible.

The MLIR Black Box: What I Found

I ran strings on the tileiras compiler binary /usr/local/cuda-13.1/bin/tileiras) and extracted the pass pipeline. Here are the 30+ MLIR passes that transform CuTile IR into PTX:

Layout Assignment

tileas-assign-dot-layouts         # Assign layouts for MMA operations
tileas-assign-load-store-layouts  # Assign layouts for memory operations
tileas-assign-pipeline-layouts    # Assign layouts for async pipeline

Kernel Planning

tileas-plan-cta                   # Plan CTA (thread block) structure
tileas-generate-schedule          # Generate execution schedule

Optimization Passes

tileas-optimize-alloc-tensor         # Optimize tensor allocations
tileas-optimize-dot-accumulation     # Optimize MMA accumulation
tileas-optimize-reduce               # Optimize reduction operations
tileas-refine-atom-by-resource       # Refine atomics based on resources

Pipelining

tileas-dynamic-persistent         # Dynamic persistent thread scheduling
tileas-materialize-async          # Materialize async operations
tileas-materialize-schedule       # Materialize the schedule
tileas-unspecialized-pipeline     # Create unspecialized pipeline
tileas-prepare-for-scheduling     # Prepare IR for scheduling

Loop Transformations

tileas-slice-and-fuse                  # Slice and fuse loops
tileas-slicing                         # Loop slicing
tileas-unroll-register-loops           # Unroll register-bound loops

Final Lowering

tileas-attach-tma-desc-args      # Attach TMA descriptor arguments
tileas-insert-OCG-knobs          # Insert optimization hints for backend
convert-nv-tileas-to-llvm        # Convert to LLVM IR

The Dialect Hierarchy

Based on the pass names, this is what I think the dialect progression looks like:

ccuda_tile (Python IR)
    ↓
nv_tileaa (NVIDIA address-annotated)
    ↓
nv_tileas (NVIDIA address-specialized)
    ↓
LLVM IR
    ↓
PTX
    ↓
SASS

The cuda_tile dialect has 100+ operations. The nv_tileaa and nv_tileas dialects are NVIDIA internal. We can see the inputs and outputs, but the transformations are proprietary.

What the Passes Reveal

From error messages in the binary:

"TileASRemoveBufferAliasPass failed to converge"
"failed to assign pipelineLayout in assignPipelineLayouts pass"
"Materialize the pipeline schedule to generate warp-specialized or unspecialized IR"

The compiler can generate warp-specialized IR — different code paths for different warps within a thread block. This is how it implements the leader election pattern we'll see in the PTX.

PTX: Where Blackwell Shows Up

The compiled PTX reveals what the MLIR passes generated. Let’s trace each kernel operation to its PTX implementation.

Kernel Header

.version 9.1
.target sm_100a
.address_size 64

.visible .entry fused_moe_kernel(
    .param .u64 .ptr .global .align 1 fused_moe_kernel_param_0,
    .param .u32 fused_moe_kernel_param_1,
    ...
)
.reqntid 256
.minnctapersm 1
.reg .pred  %p<4315>;
.reg .b16   %rs<5>;
.reg .b32   %r<1178>;
.reg .b64   %rd<1202>;
.shared .align 128 .b8 global_smem[180424];

Key details:

• PTX 9.1

• sm_100a — Blackwell architecture target

• 256 threads per block (.reqntid 256)

• 4,315 predicates

• 1,202 b64 PTX virtual registers declared

• 180KB shared memory

The tcgen05 Allocation System

The first thing the kernel does is allocate TMEM columns (and write the TMEM address into shared memory).:

// Get warp ID
mov.u32     %r1, %tid.x;
shr.u32     %r2, %r1, 5;                    // warp_id = tid / 32
shfl.sync.idx.b32  %r3, %r2, 0, 31, -1;     // Broadcast warp 0's ID

// Elect a single leader thread
elect.sync  %r205|%p3937, -1;
bar.sync    0;

// Only warp 0, thread 0 executes allocation
setp.ne.s32 %p3938, %r3, 0;
@%p3938 bra $L__BB0_2;

// Allocate 128 TMEM columns (write TMEM base address into shared memory)
mov.b32     %r206, 128;
tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [global_smem+180416], %r206;
tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;

What is tcgen05?

tcgen05 is Blackwell’s new instruction family that changes how threads cooperate on matrix operations. This is introduced on Blackwell (sm_100a) and represents NVIDIA’s new approach to tile level execution.

The tcgen05 instruction family includes:

The typical tcgen05 lifecycle follows this pattern:

// 1. Allocate TMEM columns (TMEM base address written into shared memory)

tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [smem+offset], %r;

// 2. Release the allocation permit
tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;

// 3. Perform matrix operations (repeated in main loop)
tcgen05.mma.cta_group::1.kind::f16 [col], A, B, desc, {...}, pred;

// 4. Commit results and signal barrier
tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [barrier];

// 5. Deallocate columns (once at kernel end)
tcgen05.dealloc.cta_group::1.sync.aligned.b32 col, count;

The .cta_group::1 suffix indicates thread block group scope, while .shared::cluster enables cluster-wide synchronization across multiple CTAs.

This is not in my kernel. The compiler determined that the MMA operations need this allocation and inserted it automatically.

Barrier Initialization

After allocation, the kernel initializes 20 mbarrier objects:

// Only thread 0 initializes barriers
@%p3939 bra $L__BB0_4;

mov.b32     %r207, 128;
mbarrier.init.shared.b64 [global_smem+180400], %r207;
mov.b32     %r208, 1;
mbarrier.init.shared.b64 [global_smem+180408], %r208;
mbarrier.init.shared.b64 [global_smem+180304], %r207;
mbarrier.init.shared.b64 [global_smem+180336], %r208;
mbarrier.init.shared.b64 [global_smem+180312], %r207;
mbarrier.init.shared.b64 [global_smem+180344], %r208;
mbarrier.init.shared.b64 [global_smem+180320], %r207;
mbarrier.init.shared.b64 [global_smem+180352], %r208;
mbarrier.init.shared.b64 [global_smem+180328], %r207;
mbarrier.init.shared.b64 [global_smem+180360], %r208;
mbarrier.init.shared.b64 [global_smem+180224], %r208;
mbarrier.init.shared.b64 [global_smem+180264], %r208;
mbarrier.init.shared.b64 [global_smem+180232], %r208;
mbarrier.init.shared.b64 [global_smem+180272], %r208;
mbarrier.init.shared.b64 [global_smem+180240], %r208;
mbarrier.init.shared.b64 [global_smem+180280], %r208;
mbarrier.init.shared.b64 [global_smem+180248], %r208;
mbarrier.init.shared.b64 [global_smem+180288], %r208;
mbarrier.init.shared.b64 [global_smem+180256], %r208;
mbarrier.init.shared.b64 [global_smem+180296], %r208;

$L__BB0_4:
fence.mbarrier_init.release.cluster;
bar.sync    0;

The barriers are initialized with different expected arrival counts:

• Some expect 128 arrivals

• Some expect 1 arrival (leader-only operations)

The fence.mbarrier_init.release.cluster ensures barrier visibility across the cluster — cluster-scoped synchronization.

Mapping: swizzle_2d() → Integer Arithmetic

Kernel:

bid_m, bid_n = swizzle_2d(M, N, TILE_M, TILE_N, GROUP_SIZE_M)

PTX (60+ instructions):

mov.u32     %r5, %clusterid.x;              // Get cluster ID (not just CTA ID!)
shr.s32     %r211, %r1112, 31;              // Sign extension for ceiling div
shr.u32     %r212, %r211, 25;
add.s32     %r213, %r1112, %r212;
shr.s32     %r214, %r213, 7;                // ÷ 128 (TILE_M)
and.b32     %r215, %r213, -128;
setp.ne.s32 %p3945, %r1112, %r215;
setp.gt.s32 %p3946, %r1112, -1;
and.pred    %p3947, %p3946, %p3945;
selp.b32    %r216, 1, 0, %p3947;
add.s32     %r7, %r214, %r216;              // num_bid_m = cdiv(M, TILE_M)
...
div.s32     %r219, %r5, %r16;               // group_id = bid / num_bid_in_group
mul.lo.s32  %r220, %r219, %r16;
...
rem.s32     %r231, %r5, %r18;               // bid % group_size_m
...
shl.b32     %r20, %r241, 7;                 // bid_m * TILE_M

The swizzle pattern interleaves blocks across the N dimension within groups of 8, improving L2 cache reuse. The compiler expands this into explicit integer operations.

Mapping: ct.gather() → Predicated Global Loads

Kernel:

token_ids = ct.gather(sorted_token_ids, token_id_indices)

PTX:

// Compute index and check bounds
setp.ge.u32   %p3959, %r22, %r1121;         // index >= array_length?
mov.b32       %r1152, 0;                     // Default value if out-of-bounds
@%p3959 bra   $L__BB0_8;                     // Skip load if OOB
ld.global.b32 %r1152, [%rd1];                // Actual global load

// Repeat for each element (8x unrolled)
add.s32       %r243, %r22, 16;
setp.ge.u32   %p3960, %r243, %r1122;
mov.b32       %r1153, 0;
@%p3960 bra   $L__BB0_10;
ld.global.b32 %r1153, [%rd2];

add.s32       %r245, %r22, 32;
setp.ge.u32   %p3961, %r245, %r1123;
mov.b32       %r1154, 0;
@%p3961 bra   $L__BB0_12;
ld.global.b32 %r1154, [%rd3];
// ... continues for 8 elements

The compiler generates predicated loads with bounds checking. Each gather element becomes a separate load with its own bounds check.

Mapping: K-Loop → Async Pipeline with Double Buffering

Kernel:

for k in range(0, ct.cdiv(K, TILE_K)):
    a = ct.gather(A, ...)
    b = ct.load(B, ...)
    accumulator = ct.mma(a, b, accumulator)

PTX (software-pipelined):

$L__BB0_33:
// Wait for previous iteration's data to arrive
mov.b32     %r386, 10000000;                // 10ms timeout
mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 %p4141, [%r385], %r1165, %r386;
not.pred    %p4142, %p4141;
@%p4142 bra $L__BB0_33;                     // Spin until barrier ready

// Issue async copies for NEXT iteration (double buffering!)
selp.b32    %r415, 16, 0, %p4298;           // Copy size or 0 if masked
cp.async.cg.shared.global [%r414], [%rd1170], 16, %r415;
cp.async.cg.shared.global [%r419], [%rd1171], 16, %r420;
cp.async.cg.shared.global [%r424], [%rd1172], 16, %r425;
cp.async.cg.shared.global [%r429], [%rd1173], 16, %r430;
cp.async.cg.shared.global [%r434], [%rd1174], 16, %r435;
cp.async.cg.shared.global [%r439], [%rd1175], 16, %r440;
cp.async.cg.shared.global [%r444], [%rd1176], 16, %r445;
cp.async.cg.shared.global [%r449], [%rd1177], 16, %r450;

// Signal that async copies are issued
cp.async.mbarrier.arrive.noinc.shared.b64 [%r452+180304];

My simple for loop became a software pipelined async copy with:

• mbarrier.try_wait for synchronization

• cp.async.cg.shared.global for async global→shared memory copies

• Double buffering (issue copies for iteration N+1 while computing iteration N)

• Parity based barrier reuse

Mapping: ct.mma() → tcgen05.mma

Kernel:

accumulator = ct.mma(a, b, accumulator)

PTX:

// Elect a leader thread for this MMA
elect.sync    %r1087|%p4269, -1;
not.pred      %p4270, %p4269;
@%p4270 bra   $L__BB0_133;                   // Non-leaders skip

// Encode shared memory addresses
or.b64        %rd1130, %rd516, 4611756662049538048;
or.b64        %rd1131, %rd515, 4611756662049538048;
mov.b32       %r1088, 0;
mov.b32       %r1089, 136316048;             // MMA descriptor

// Issue the tcgen05 MMA — only the leader executes this!
tcgen05.mma.cta_group::1.kind::f16 [%r4], %rd1131, %rd1130, %r1089,
                                   {%r1088, %r1088, %r1088, %r1088}, %p4314;

$L__BB0_133:
// Second MMA tile (unrolled)
elect.sync    %r1090|%p4271, -1;
not.pred      %p4272, %p4271;
@%p4272 bra   $L__BB0_135;
add.s64       %rd1132, %rd516, 4611756662049538050;
add.s64       %rd1133, %rd515, 4611756662049538050;
tcgen05.mma.cta_group::1.kind::f16 [%r4], %rd1133, %rd1132, %r1092, {...}, %p4273;

$L__BB0_135:
// Third MMA tile...

Key insight: The tcgen05.mma is issued by a single elected thread, not all threads. The hardware handles distributing the matrix multiply across the warp. This is fundamentally different from previous architectures where each thread contributed to the MMA.

The elect.sync instruction selects one thread per warp to issue the instruction. Other threads skip directly to the next synchronization point.

Mapping: Loop Commit → Barrier Arrival

After each iteration’s MMAs complete:

// Leader commits results and signals barrier
elect.sync    %r1099|%p4280, -1;
not.pred      %p4281, %p4280;
@%p4281 bra   $L__BB0_141;
tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%r179+40];

$L__BB0_141:
elect.sync    %r1100|%p4282, -1;
not.pred      %p4283, %p4282;
@%p4283 bra   $L__BB0_143;
tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%r178+32];

The tcgen05.commit instruction:

1. Commits the MMA results to shared memory

2. Signals an mbarrier arrival (.mbarrier::arrive::one)

3. Has cluster scope (.shared::cluster) — visible across CTAs in a cluster

Kernel Cleanup: Deallocation

At the end of the kernel:

// Only warp 0 deallocates
setp.ne.s32 %p4289, %r3, 0;
bar.sync    0;
@%p4289 bra $L__BB0_150;

mov.b32     %r1103, 128;
tcgen05.dealloc.cta_group::1.sync.aligned.b32 %r4, %r1103;

$L__BB0_150:
ret;

The tcgen05.dealloc releases the allocated TMEM columns allocated at kernel start.

SASS: The Final Layer

The SASS (Shader ASSembly) is the actual machine code that runs on the GPU. Here’s what the PTX becomes.

The TMEM allocation: contention-handling retry pattern

The PTX tcgen05.alloc lowers into a contention-handling atomic retry loop in SASS:

/*0100*/  ELECT P0, URZ, PT ;                    // Elect leader
/*0110*/  @!P0 BRA 0x2a0 ;                       // Non-leaders skip

/*0120*/  UMOV UR4, 0x4 ;
/*0130*/  IMAD.MOV.U32 R3, RZ, RZ, 0xf ;
/*0140*/  IMAD.MOV.U32 R5, RZ, RZ, 0x1 ;
/*0150*/  DEPBAR.LE SB0, 0x36 ;

// Atomic find-and-set with retry loop
/*0160*/  UTCATOMSWS.FIND_AND_SET.ALIGN UP0, UR4, UR4 ;
/*0170*/  IMAD.U32 R2, RZ, RZ, UR4 ;
/*0180*/  BRA.U UP0, 0x1f0 ;                     // If acquired, continue
/*0190*/  NANOSLEEP 0x64 ;                       // 100ns backoff!
/*01a0*/  UMOV UR4, 0x4 ;
/*01b0*/  DEPBAR.LE SB0, 0x36 ;
/*01c0*/  UTCATOMSWS.FIND_AND_SET.ALIGN UP0, UR4, UR4 ;
/*01d0*/  IMAD.U32 R2, RZ, RZ, UR4 ;
/*01e0*/  BRA.U !UP0, 0x190 ;                    // Retry if not acquired

New Blackwell instructions:

• UTCATOMSWS.FIND_AND_SET.ALIGN — Atomic find-and-set on the thread coalescing state

• NANOSLEEP 0x64 — Sleep for ~100ns (bounded [0, 2t]) (backoff for contention)

This is a contention handling retry loop with nanosleep backoff — the compiler generated a sophisticated synchronization primitive from my simple kernel.

Barrier Initialization in SASS

/*0420*/  FENCE.VIEW.ASYNC.S ;                   // System-scope fence
/*0430*/  SYNCS.EXCH.64 URZ, [UR14+0x2c0b0], UR4 ;
/*0440*/  SYNCS.EXCH.64 URZ, [UR14+0x2c0b8], UR6 ;
/*0450*/  SYNCS.EXCH.64 URZ, [UR14+0x2c050], UR4 ;
/*0460*/  SYNCS.EXCH.64 URZ, [UR14+0x2c070], UR6 ;
// ... 20 total barrier initializations

The SYNCS.EXCH.64 is Blackwell’s barrier exchange instruction — it atomically swaps a value into the barrier object.

Virtual Memory Accounting

/*02f0*/  UVIRTCOUNT.DEALLOC.SMPOOL 0x80 ;

UVIRTCOUNT.DEALLOC.SMPOOL deallocates from the shared memory pool’s virtual counter. This is part of Blackwell’s resource accounting system — tracking how much shared memory is “logically” allocated even when the physical allocation is different.

Optimized Barrier Waiting

// Try to acquire barrier (non-blocking)
SYNCS.PHASECHK.TRANS64.TRYWAIT P4, [UR10+0x2c070], R14 ;

// If not ready, sleep and retry
@!P4 NANOSLEEP.SYNCS 0x989680 ;               // Sleep 10ms while waiting
@!P4 SYNCS.PHASECHK.TRANS64 P4, [R15+URZ+0x2c070], R14 ;  // Blocking check

The NANOSLEEP.SYNCS variant sleeps while waiting for a synchronization operation — more power efficient than a busy-wait loop.

tcgen05 in SASS: Under the Hood

At the SASS level, the tcgen05 PTX instructions expand into sequences that reveal more about the hardware implementation. The allocation uses UTCATOMSWS (Unified Thread Coalescing Atomic on Shared Workspace):

// tcgen05.alloc expands to:
/*0160*/  UTCATOMSWS.FIND_AND_SET.ALIGN UP0, UR4, UR4 ;  // Find free column slot
/*0170*/  IMAD.U32 R2, RZ, RZ, UR4 ;                      // Extract result
/*0180*/  BRA.U UP0, 0x1f0 ;                              // Branch if acquired
/*0190*/  NANOSLEEP 0x64 ;                                // 100ns backoff on contention

The UTCATOMSWS.FIND_AND_SET.ALIGN atomically searches a bitmap for a free slot and sets it — this is how the hardware manages dynamic column allocation. The .ALIGN suffix ensures the operation respects alignment constraints for the tcgen05 column structure.

For deallocation, the compiler generates:

// tcgen05.dealloc expands to:
/*xxxx*/  UTCATOMSWS.AND URZ, UR4 ;                      // Clear allocation bits
/*xxxx*/  UVIRTCOUNT.DEALLOC.SMPOOL 0x80 ;              // Return 128 bytes to pool

The UVIRTCOUNT.DEALLOC.SMPOOL instruction is part of Blackwell’s virtual memory accounting system — it tracks “logical” allocations separately from physical shared memory.

What CuTile Generated That I Didn’t Write

Let’s be explicit about the gap between the kernel and the compiled code.

Memory Management

What I wrote → What the compiler generated:

• No shared memory allocation → 180KB global_smem[180424]

• No explicit barriers → 20 mbarrier objects

• Simple for-loop → Double-buffered async pipeline

• No register hints → 1,202 b64 PTX virtual registers declared

Synchronization

What I wrote → What the compiler generated:

• No thread coordination → elect.sync for leader election

• No barriers → mbarrier.try_wait.parity with 10ms timeout

• No atomics → UTCATOMSWS.FIND_AND_SET with NANOSLEEP backoff

• No cluster operations → fence.mbarrier_init.release.cluster

Instruction Selection

What I wrote → What the compiler generated:

• ct.mma(a, b, acc) → tcgen05.mma.cta_group::1.kind::f16

• ct.gather() → Predicated loads with bounds checks

• ct.scatter() → Predicated stores with bounds checks

• for k in range() → Software pipelined loop with async copies

Where This Is Heading

The technical evidence points to an uncomfortable conclusion: NVIDIA’s first party compiler advantage isn’t just real — it’s structural and widening.

The moat isn’t just documentation. Yes, tcgen05 is in PTX ISA 9.1. Yes, elect.sync and nanosleep are public. But NVIDIA has a history of undocumented PTX and SASS instructions that show up in cuBLAS and cuDNN before they appear in any ISA reference — sometimes they never do. I can’t claim CuTile is using undocumented instructions in this kernel, but it would be naive to assume the public PTX ISA is the complete picture of what NVIDIA’s own compilers can target. Even setting that aside, the orchestration problem remains: knowing which patterns, with what parameters, for which kernel shapes. My 86 line kernel became 20 barriers with specific arrival counts, 100ns backoff on TMEM allocation, 10ms timeouts on barrier waits, leader election before every MMA. You can find pieces in Flash Attention 4, in CUTLASS examples, in blog posts. The specific combination? That’s encoded in CuTile’s MLIR passes.

The gap compounds with each generation. Blackwell introduced tcgen05, cluster-scope barriers, TMEM, CTA pairs. Each feature multiplies the configuration space. The Triton team struggled enough with Blackwell attention that they built Gluon as an escape hatch — a lower level frontend that exposes more control because the high level abstractions couldn’t find the right codegen. CuTile doesn’t need an escape hatch. The compiler already knows.

Rubin will repeat this pattern. Whatever instruction family Rubin introduces — tcgen06, new cluster primitives, new memory hierarchies — NVIDIA’s compiler team has been generating code for it since before the silicon taped out. CuTile will ship Rubin support on day one with tuned codegen. Triton and Pallas will start reverse engineering patterns from publicly released kernels and CUTLASS examples. This isn’t a criticism of those projects; it’s the structural reality of targeting hardware you don’t design.

Frameworks will notice. If CuTile consistently generates faster code for complex kernels — MoE, attention, whatever the next bottleneck is — how long before PyTorch’s Inductor or JAX’s Pallas add CuTile as a backend option? Vendor lock-in concerns are real, but frameworks already have NVIDIA specific paths through cuDNN and cuBLAS. A CuTile backend is the same trade-off with better codegen for custom ops.

The black box is the product. CuTile’s closed MLIR passes let NVIDIA iterate on compiler internals without breaking user code. Your tile level kernel stays stable while the codegen improves underneath. This is the opposite of Triton’s model, where you sometimes need to understand the compiler to get good performance. For NVIDIA, hiding the complexity is the point — they want you writing tile level Python, not debugging PTX.

My predictions:

• Within two years, at least one major framework adds CuTile as an optional backend for NVIDIA GPUs.

• Rubin ships with CuTile support day one. Triton gets functional Rubin support within 6 months, competitive codegen within 18.

Conclusion

NVIDIA is building a compiler moat.

Not just through documentation — though the gap between public PTX and what their own tools can emit is its own question. The deeper moat is orchestration: which patterns to combine, which parameters to use, which configurations work for which workloads. My 86 line kernel became 1,900 lines of PTX with 20 barriers, contention handling TMEM allocation with nanosleep backoff, and software pipelined async copies. Every one of those primitives is discoverable. The specific combination that makes this MoE kernel fast is not.

Triton can emit tcgen05 instructions. Matching CuTile’s codegen quality — for this kernel, for attention, for whatever comes next — requires reproducing years of tuning encoded in proprietary MLIR passes. And that’s assuming the public ISA is the whole story, which history suggests it isn’t.

Each architecture adds complexity. Each generation widens the gap.

If you want to see this yourself, the artifacts are in the repo. Run strings on tileiras, extract the PTX from your .cubin files, trace your operations through the stack. The compiler is a black box, but the outputs tell you everything you need to know about where GPU programming is heading.

Full Code / IR Dumps: https://github.com/patrick-toulme/justabyte/tree/main/cutile_blackwell_post

Connect on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Connect On Linkedin: https://www.linkedin.com/in/patrick-toulme-150b041a5/

Follow on X: https://x.com/PatrickToulme

Motivation

Last post, the TPU compiler did everything: 8 lines of JAX became 250 VLIW bundles across 5 fused kernels, with automatic memory orchestration, dual-MXU scheduling, and async DMA overlap. The thesis was simple—write high-level code, let the compiler figure out the rest.

This post is about where that breaks down.

Same hardware, same compiler: naive attention compiles to 10,245 VLIW bundles and moves 526MB through HBM. Splash Attention (block_size=512) compiles to 1,651 bundles and moves 14MB. XLA optimized both aggressively. The 6× gap isn’t a compiler failure—it’s a compiler limitation. XLA can fuse ops and schedule instructions, but it can’t rewrite your algorithm. It can’t infer the streaming/online-softmax reformulation, that you never need to materialize the full attention matrix, that there’s a numerically-stable streaming formulation hiding inside your einsum.

In principle, a compiler could learn this. But no one has taught one yet. FlashAttention exists because a human figured out the trick and wrote a kernel. Until compilers get smarter, we need an escape hatch.

Pallas is that escape hatch. This post traces Splash Attention—JAX’s FlashAttention for TPUs—through the TPU Pallas Compiler, and compares it to naive attention at every layer of IR.

Setup:

All experiments were performed on a TPU V6e Trillium from Google Cloud. We use the same dump flags from the previous post with the addition of xla_mosaic flags which dumps the Mosaic compiler IR for Pallas.

os.environ["XLA_FLAGS"] = (
    f"--xla_dump_hlo_as_text "
    f"--xla_dump_to={HLO_PATH} "
    f"--xla_dump_hlo_pass_re=.* "
)
os.environ["LIBTPU_INIT_ARGS"] = (
    f"--xla_jf_dump_to={LLO_PATH} "
    f"--xla_jf_dump_hlo_text=true "
    f"--xla_jf_dump_llo_text=true "
    f"--xla_jf_dump_llo_html=false "
    f"--xla_jf_dump_llo_static_gaps=true "
    f"--xla_jf_emit_annotations=true "
    f"--xla_jf_debug_level=2 "
    f"--xla_mosaic_dump_to={MOSAIC_PATH} "
    f"--xla_mosaic_enable_dump_debug_info=true "
    f"--xla_mosaic_enable_llo_source_annotations=true"
)

Full Code / IR Dumps: https://github.com/patrick-toulme/justabyte/tree/main/tpu_pallas_post

What this post covers:

• The Compiler’s Limitation: Why XLA can’t fuse across the attention matrix

• Splash Attention Kernel: How Pallas expresses online softmax with grid, BlockSpec, and scratch

• Pallas Compiler Pipeline: The path from Pallas code → Mosaic → LLO → VLIW

• HLO/Mosaic/LLO Walkthrough: What each IR layer reveals

• Reference Attention Comparison: Same algorithm through XLA’s standard path

• Analysis: Bundle counts, HBM traffic, block size tradeoffs

The Compiler’s Limitation

Why can’t XLA automatically achieve Pallas-level performance for every case? Looking at what XLA does fuse for reference JAX attention:

fusion.5: Q @ K^T + mask + reduce_max  →  (max, scores)
fusion.2: exp(scores - max) + reduce_sum  →  sum
fusion:   normalize + S @ V  →  output

XLA fuses the matmul with the mask and max reduction. It fuses the exp with the sum reduction. It’s doing real work here. But the 128MB attention matrix still gets written to HBM between fusion.5 and fusion.2.

The problem is that standard softmax is inherently multi-pass: you need the max over all values before you can compute any exp. XLA can fuse along dataflow edges, but it can’t restructure the algorithm. The attention matrix has to exist somewhere because multiple operations need to read it.

Online softmax sidesteps this by maintaining running statistics — you compute a tile’s contribution to max/sum/output, then update your estimates as you process the next tile. This is an algorithmic transformation, not a fusion pattern. XLA would need to:

1. Recognize that the softmax + matmul pattern permits streaming

2. Prove the online reformulation is numerically equivalent

3. Restructure the computation to process tiles incrementally

Concretely: it would have to replace scores = QKᵀ; p = softmax(scores); o = pV with a loop that carries (m, l, o) and updates them per KV tile.

No production compiler does this today (to my knowledge). Pallas lets you write the streaming algorithm directly.

Splash Attention Kernel

Splash Attention is JAX’s FlashAttention for TPUs, built on Pallas. The algorithm itself is well-documented elsewhere — here we focus on how Pallas expresses it.

// Standard softmax (requires full matrix):
max = reduce_max(scores)        // Need ALL scores first
exp_scores = exp(scores - max)  // Then process all
sum = reduce_sum(exp_scores)    // Need ALL exp values
output = exp_scores / sum

// Online softmax (streaming):
for each KV_block:
    m_new = max(m_prev, local_max)           // Update running max
    correction = exp(m_prev - m_new)         // Rescale factor
    l_new = correction * l_prev + local_sum  // Update running sum
    o_new = correction * o_prev + local_out  // Update running output

Pallas Kernel Structure

A Pallas kernel is a function that operates on Refs — mutable views into memory regions. The runtime calls your kernel once per grid point, with Refs pointing to the appropriate tiles:

def flash_attention_kernel(
    # Prefetched scalar inputs (live in SMEM)
    data_next_ref,      # Which KV block to load
    block_mask_ref,     # Skip mask
    mask_next_ref,      # Partial mask index
    # Tiled inputs (sliced per grid point)
    q_ref,              # [block_q, head_dim] slice of Q
    k_ref,              # [block_kv, head_dim] slice of K  
    v_ref,              # [block_kv, head_dim] slice of V
    ...
    # Scratch space (persists across grid iterations)
    m_scratch_ref,      # [block_q, 128] running max
    l_scratch_ref,      # [block_q, 128] running sum
    o_scratch_ref,      # [block_q, head_dim] running output
    # Output tile
    o_ref,
    *,
    # Static config (compiled into kernel)
    bq: int,
    bkv: int,
    ...
):

Inside the kernel, you read/write Refs with ref[...] syntax. Pallas traces these operations and compiles them to Mosaic IR.

BlockSpec: Mapping Grid to Memory

BlockSpec defines how each grid point maps to a tile of the input/output arrays:

pl.BlockSpec(
    block_shape=(None, bq, head_dim),  # None = full dim, bq = tiled
    index_map=lambda h, i, j, *_: (h, i, 0)  # grid coords → tile origin
)

For Q, the index map is simple — grid point (h, i, j) reads Q block (h, i, :). But K and V need indirection for sparse attention:

def k_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref):
    # Look up which KV block to actually load (enables block skipping)
    next_j, *_ = _next_nonzero(h, i, j, data_next_ref, block_mask_ref, mask_next_ref)
    return (h // q_heads_per_kv_head, next_j, 0)

The data_next_ref indirection is how Splash skips masked blocks — if block j is fully masked, data_next[h,i,j] points to the next valid block instead.

Grid and Dimension Semantics

grid = (num_q_heads, q_seq_len // block_q, grid_width)

compiler_params = pltpu.CompilerParams(
    dimension_semantics=("parallel", "arbitrary", "arbitrary")
)

• Dimension 0 (heads): "parallel" — no cross-iteration dependencies, can execute in any order

• Dimension 1 (Q blocks): "arbitrary" — independent because scratch is per-(head, q_block)

• Dimension 2 (KV blocks): "arbitrary" — dependent because scratch accumulates across KV tiles

The "parallel" hint lets the compiler vectorize across heads. "arbitrary" means “don’t assume anything” — safe but conservative.

PrefetchScalarGridSpec

Splash uses PrefetchScalarGridSpec to overlap data loading with compute:

pltpu.PrefetchScalarGridSpec(
    num_scalar_prefetch=3,  # First 3 inputs go to SMEM, prefetched ahead
    in_specs=[...],
    out_specs=[...],
    grid=grid,
)

The first 3 arguments (data_next, block_mask, mask_next) are small index arrays. Putting them in SMEM with prefetch means the kernel can look up the next block address while the current block is still computing.

Memory Spaces

Pallas exposes TPU memory hierarchy:

# VMEM (on-chip SRAM) — default for BlockSpec refs
pl.BlockSpec((bq, head_dim), index_map)

# SMEM (scalar memory) — for small indexing data
pl.BlockSpec((num_heads,), lambda *_: (0,), memory_space=pltpu.SMEM)

# Scratch (persists across grid iterations)
jax.ShapeDtypeStruct((bq, NUM_LANES), jnp.float32)  # m_scratch, l_scratch

The scratch refs (m_scratch_ref, l_scratch_ref, o_scratch_ref) are the key to online softmax — they accumulate across the KV dimension without spilling to HBM.

Pallas Compiler Pipeline

When building a DSL compiler, you want to reuse as much of the existing compiler infrastructure as possible rather than writing everything from scratch. Pallas shows up in HLO as an opaque tpu_custom_call. XLA still does its usual whole-graph work (layout, scheduling, surrounding fusions), but it can’t see or rewrite the kernel body. The payload is serialized MLIR that goes straight to Mosaic → LLO → VLIW. The key insight is that both XLA's fusion region codegen and Pallas share the same LLO backend. Whether the TPU is running a compiler-generated fusion or a hand-written Pallas kernel, the final stages—MXU scheduling, DMA overlap, VLIW packing—are identical. Pallas just enters the pipeline at a different point.

What XLA Sees

The whole Pallas kernel becomes a single opaque HLO op:

%custom-call = bf16[8,2048,128] custom-call(
    bf16[8,2048,128] %q, bf16[8,2048,128] %k, bf16[8,2048,128] %v, ...
), custom_call_target="tpu_custom_call", backend_config="..."

XLA’s fusion passes can’t look inside — the kernel goes directly to Mosaic, which compiles it to LLO → VLIW bundles. The backend_config encodes the full Pallas program as serialized MLIR IR.

HLO

Here’s what the Pallas kernel looks like in HLO:

HloModule jit_splash_attention_kernel, entry_computation_layout={(bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)})->bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.28 (Arg_0.1: bf16[8,2048,128], Arg_1.2: bf16[8,2048,128], Arg_2.3: bf16[8,2048,128]) -> bf16[8,2048,128] {
  %constant.4 = s8[1,2,2]{2,1,0} constant({ { { 0, 0 }, { 0, 1 } } })
  %constant.5 = s8[1,2,2]{2,1,0} constant({ { { 1, 0 }, { 2, 1 } } })
  %Arg_0.1 = bf16[8,2048,128]{2,1,0} parameter(0), metadata={op_name="q"}
  %Arg_1.2 = bf16[8,2048,128]{2,1,0} parameter(1), metadata={op_name="k"}
  %Arg_2.3 = bf16[8,2048,128]{2,1,0} parameter(2), metadata={op_name="v"}
  %constant.6 = s32[2048]{0} constant({...})
  %broadcast.0 = s32[2048,128]{1,0} broadcast(%constant.6), dimensions={0}, metadata={op_name="jit(splash_attention_kernel)/jit(main)/splash_attention_kernel/splash_kernel_b1024_h8_s2048/jit(_splash_attention)/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=1045}
  %custom-call.0 = (f32[1024,128]{1,0}, f32[1024,128]{1,0}, f32[1024,128]{1,0}, bf16[8,2048,128]{2,1,0}) custom-call(%constant.4, %constant.5, %Arg_0.1, %Arg_1.2, %Arg_2.3, /*index=5*/%broadcast.0), custom_call_target="tpu_custom_call", operand_layout_constraints={s8[1,2,2]{2,1,0}, s8[1,2,2]{2,1,0}, bf16[8,2048,128]{2,1,0}, bf16[8,2048,128]{2,1,0}, bf16[8,2048,128]{2,1,0}, s32[2048,128]{1,0}}, metadata={op_name="jit(splash_attention_kernel)/jit(main)/splash_attention_kernel/splash_kernel_b1024_h8_s2048/jit(_splash_attention)/splash_mha_fwd_c602b16d/pallas_call" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=1100}, backend_config={"custom_call_config": {"body": "TUzvUgFNTElSMjEuMC4wZ2l0AAFRCwEDBQcJAQMLAzsNDxETFRcZGx0fISMlJykrLS8xMzU3OTs9P0FDRQMKCWYILwH7BwsXCxMXCxMLCwsbCwsLCysLCxMTExcLCwsLCw8PDw8XGxMTFxNzUwsLExcXFxcXCxMLExMTExMTCxMTExcPCxMTExMTEwsTDxcTCxMTExMPCxcPDxMXCxMXExMTExMPpYUPCwsLCwsLCwsLFxMPExcPExMLExMPExMTGwsFD2GBkY1hKgIqAgFGBgsLExMTExcTFwsLCwsTExMTHxcXCwszFwsLFxMXExMTIwsLawsbCwtzCw8LC0sfCx8LHwsfCycLJwsnCx8LGxsbGxsbGxsPDw8fCycLJxMLJxMLJxMLJxMLGxMLGxMfFwsfExMLHw8fExMTEw8PHw8fExMPHw8LJxMTDw8PExMPExMfDxMfDw8nExMPDw8TEw8PHxMTHw8THxMTHxcfCxMTHwsTExMfExMfExMTHxMTEx8TEx8rUw8TEx8TEx9TExMfExMfExMfExMfCw8LHxMTEx8PCxMTHxMTHxMTHxMTHxcfExMfMxMTHxMTHxMTEx8TFwsTEycfExMTExMTDxMTJxMTHw8TExMnFw8XCxcXCyMLFxMfFxMTExMTDx8TCxMfDxMLFx8LExMfKxMTEx8LExMfExMTHxcTEx8XExMfFx8LExMfKxMTHxcTEx8XExMfExMTHxMTEx8TEx8TEx8TEx8TExMfFx8TEx8zExMTHxMTHxMTEx8TExMfExMfExMfExMfExMfExMfExMfExMfBwVZWQkFXUkBLw8HHycHCw8rHwsnFx8jHwcnG0MvKx8fAlovHwVHAwMTLgMFSRV7XgUDAxMqAwVLFW4EJwVNBU8FUQMDmgRWCAVTBVUFVwVZAwUSAgoFFgIaAgVbBV0VV0oDFSIEJx0b2gQdogOmAwVfBWEFYwVlBWcdoykdo6kRAQUdow8DAxPuBAMDYgVaCB0NfgUdDYoFFU4C8gUdDZIHIw0HMQEAAAAAAAAAAAQAAAAAAACAAAAAAAAAACMNBSEABAAAAAAAAIAAAAAAAAAADScNKR0NMgMdQgNGAx1OA1IDHVoDXgMdZgNqAx1yA3YDBWkVV7IDBWsdDc4DHQ3mAx0NAgQdDS4EHQ1GBB0vWgQFbR0b/gQdGxoFHRsyBR1GBUoFEQsABW8dG3oHHRuGBx0bAggdGxoIHRsyCB0bSggFcR03NgMV2ScdfgOCAx03rgMFcx03ygMDA2W/HTfSAx034gMV4ScFdQMDZSoEFe0nFXEnHQ12BB2KBI4EBXcdDa4EAwMTTgUdDWYFHQ1yBR0NGgYVkgZJHQ3WBxENEWFmZmluZV9tYXA8KGQwLCBkMSwgZDIpIC0+IChkMCwgZDEsIGQyKT4AYWZmaW5lX21hcDwoZDAsIGQxKSAtPiAoZDAsIGQxKT4AEQ0BBXkFewV9BX8FgQWDBYUFhwWJHToDPgMdx4oDFdljHZfGAx3qA+4DFeFjHZf+Ax3pEgQFiwMDZcUdLzIEFe1jHZdCBAMDEz0d6V4EAwWvPfk9BY0jdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPgAjdHB1LnBpcGVsaW5lX21vZGU8c3luY2hyb25vdXM+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4AI3RwdS5tZW1vcnlfc3BhY2U8c21lbT4AI3RwdS5kb3RfZGltZW5zaW9uX251bWJlcnM8WzFdLCBbMV0sIFswXSwgWzBdLCBbMCwgMCwgMSwgMF0sIFtdLCBbXT4AI3RwdS5kb3RfZGltZW5zaW9uX251bWJlcnM8WzFdLCBbMF0sIFswXSwgWzFdLCBbMCwgMCwgMSwgMV0sIFtdLCBbXT4ABY8FkSMBAQEdH/IEHR8OBR0fJgUdPgVCBR0xUgUdNgKeBQWTBZUFlwWZHTGuBR0xugUdB+4FHX8OBgMFrz35JgYDAxOGBh2uBrIGBZsFnSMNAxEBAAAAAAAAAB3SBtYGBZ8FoR0aBx4HHR8uBx02ArYHHR8OCB0fJggdHz4IAwWWAr8ZmgIFowWlAw+iAqYCHaoCrgKyArYCugK+AsUZx8ICxgIFpwEHAgL//w0lBakjDQcxCAAAAAAAAAACAAAAAAAAAAIAAAAAAAAABasRDQkFrQWvARHKAtIC2gLiAuoC8gL6AgIDAwUjzgIlTQnJAwUj1gIlTQnLAwUj3gIlTQnNAwUj5gIlTwnPAweN/SPuAiVPCdEDB439I/YCJU8J0wMHjf0j/gIlTwnVAwUjBgMlTQnXAwUdURnJAwUdURnLAwUdURnNAwUdUxnPAwUdUxnRAwUdUxnTAwUdUxnVAwUdURnXEQEBEQMBFY+RLQMH5ggtWwWxLQMJRg8jTg8LBbMtAwkyER3KEQsVWVYDBbUtAwmiEhPaEgcVW2IDBbctAwkiJBNiJAcVXW4DBbktAwmuJBfKJAsVX3oDBbstYQeBJ08Vk4YDBb0tYQeNI0EV244DLWEHZgJBvR2SA5YDBb8tYQf+AgkVHRWeAxUtqgMFwS0DB94IK1EVld0tAwfmCB9dFVm2AxVbugMVXb4DFV/CAxWT2xWZkS0DB+YIH2UVnZEtAwcCCSVRHRXaAxUt3gMVn90tAwcCCRdTFY+hBcMtAwkmDyMuDwsdFfYDFS36AxWV4xWZoRWdoR0VCgQVLQ4EFZ/jFRYEJx0vGgQtAwdKCxcjHRUpHS8mBC0DB0oLByURDQUVj6ctAwliC2l+CwcdFToEFS0+BBWV7xWZpxWdpx0VTgQVLVIEFZ/vHRWpLQMHmgwHLRViBCcdL2YELQMHugwXPR0VDx0vcgQtAwe6DAc/FXoEDx0RfgQtAwfCDBE1AwMThgQTCRAAAOAPBcUVkgQPHRGWBC0DB8YMM0EFxx1zogQVpgQPHRGqBC0DB8YMGXsVsgQPHRG2BC0DB8oMJUkdMb4EFcIEDx0RxgQtAwfKDCVZHRXOBBXSBA8dEdYELQMHygwjgxXeBA8dEeIELQMHygwJHQMFEgLqBBYCGgIjAQkhAQAAAAEAAAADAAAAAAAAABMJARX2BA8dEfoELQMH5gwzbRUCBQ8dEQYFLQMH5gwJLSMBCSEBAAAAAQAAAAIAAAAAAAAAFRIFDx0RFgUtAwfqDDNtFR4FDx0RIgUtAwfqDAktFSoFDx0RLgUtAwfuDDNtFTYFDx0ROgUtAwfuDAktBckVe6kFyy0DB7IMCXERAQIgFVYFCR0HWgUtAweKCyllFXFjBc0VagUJHQduBS0DB44LK08VdgUJHQd6BS0DB44LU3cVggUJHQeGBS0DB54LESUVjgUJHQeSBS0DB6oLFTcDAxOaBSURCQAAAAAVogUJHQemBS0DB7YLE48DBzoCCgI+An1CAn0VsgUJHQe2BS0DB/4LI00VvgUJHQfCBS0DB/4LU48df8oFFc4FCR0H0gUtAwf+CyOPAwOvPR3eBeIFBc8V5gVJHTPqBS0DCaoJPbIJDy0DBxIME0UVe/YFFXH6BRVX/gUVWQIGFVsGBhVdCgYVX5MVEgZJHTMWBi0DCaoJJ7IJDxUeBkkdMyIGLQMHxgkVOxEBIR1zLgYVMgZJHTM2Bi0DCcIJJ8oJDwMDZT4GEQ0VHUYGSgYF0RVOBl4GHVIGVgYF0y1aBgcCBR89BdUVYgZqBh0zZgYtAwf2CSlzFU4CbgYVe3IGFXF2BhVXegYVWX4GFVuCBhVdXxMJkMzMzD8djga7BdcdM5YGLQMHhgoTUR0fux2iBrsF2QMDE6oGJRcJAACA/wXbFbYGCR0HugYtAwcaDBs5AwViAl4IZgJqAh0fxgYVygYJHQfOBi0DBxoMG0sF3RXaBgkdB94GLQMHIgwbUR1z5gYV6gYJHQfuBi0DB0YMNYcdcgL2BhX6BgkdB/4GLQMHRgwrhx12AgYHFQoHCR0HDgctAwdGDBuJAwMTFgclFwkAAAAABd8VIgcJHQcmBy0DB1IMTXMDBWICYghmAmoCFTIHCR0HNgctAwdSDBudHXICPgcVQgcJHQdGBy0DB14MKUcddgJOBxVSBwkdB1YHLQMHXgwZSR0xXgcVYgcJHQdmBy0DB2IMLUkdf24HFXIHCR0HdgctAwdiDBtJFX4HCR0HggctAwdmDAktFYoHCR0HjgctAwdmDDFVFZYHCR0HmgctAwd2DBU3HRWiBxWmBwkdB6oHLQMHggwRMwMDE7IHJQUJAAAAABW6BwkdB74HLQMHhgwbYQMHOgIOAj4CfUICfR1zygcVzgcJHQfSBy0DB44MHXsV2gcJHQfeBy0DB5IMQ2MdMeYHFeoHCR0H7gctAweSDC9jHX/2BxX6BwkdB/4HLQMHkgwvdRUGCAkdBwoILQMHkgwJKRUSCCkdNRYILQMHUgszbRUeCCkdNSIILQMHUgsJLRUqCCkdNS4ILQMHVgszgxU2CCkdNToILQMHVgsJLRVCCCkdNUYILQMHWgszbRVOCCkdNVIILQMHWgsJLSNhcml0aC5mYXN0bWF0aDxub25lPgAjYXJpdGgub3ZlcmZsb3c8bm9uZT4AI3ZlY3Rvci5raW5kPG1heGltdW1mPgAjdmVjdG9yLmtpbmQ8YWRkPgABAgIDJwUCIAIECRcGAgcFCQkTwQsBCQECBBf7BwUCIAIEH8EnBQIgAiAJAUEX+wUCIAIECcMnAwIgCScFAiACBB8nBwUCIAIEHycFAiACIAEHF/sFAiACBAHDJwUCIAUJBRsBAQEHBw8PDyEVFRUPAQULAQEBBwcHAQEBBQsBAQEHBwUBAScFAiACBAEnBQIgAiALBEIdBQERAZICBwMBJQ0RAZ4CBwNNdxsBAQEBAQEHAQcBDwEPAQ8BIQEVARUBFQEPAQMD5wsDAREH5+sDCwUFGx0GHgQDAQMdAwM5CwMBEQc5pQMLBR8hHxQ5AyMJAx9NAwOGAkEDCQkGhgIDBQNNAwOHBQMDAwOHBQMDBQaHAwUHF1FTCwWHIQlPF1FTAwOKAloCAwkJBooCAwUDVwMDiQUDAwMDiQUDAwUGiQMFBxNbXQsFiSEJWRNbXQMDjgJBAwkJBo4CAwUDYQMDiwUDAwMDiwUDAwUGiwMFBxVlZwsFiyEJYxVlZxkAOQMBBRkAOQMDbQUDAwcGbQMDAwMHBm0DAwMFFQZtAxMJCSUnKRcGNgQDAQMrAwPxCwMBEQfxmwMLBS0vAwNvBQMDBwZvAwMDAwcGbwMDAwUVBm8DEwkHMzU3FwZKBAMBAzkdBlYEAwEDMQMDOwsDAREHO6UDCwU9Px8UOwNBCQObigIDAyoCCwMBAwMuArMDASMHLgJDAwEFTU8DA7UFAwMDA7UFAwMFBrUDBQcTU1UDA7cFAwMDA7cFAwMFBrcDBQcVWVsDA0UFAwMDA0UFAwMDA0UFAwMFBkUDGwkLX2FjEwZFAxkDZQMDRwUDAwcGRwMDA1EDA0cFAwMFBkcDGwkNaWttEwZHAxkDbwMDMgKWBQMRJQcyAqoFAxEHZ3FzAwNGArMDASMHRgJDAwEFO3cDA0oCswMBIwdKAkMDAQVNeycHxgVDAwEFeX01A9oF1gUDHQkGUgIDHQN/JwdSAkMDHQWDgQMDuQUDAwMDuQUDAwUGuQMrBxGHiRsHKgZWAgMdA4sRB0IGOgYDLQWNhQMDigZaAgMJCQaaBgMRA5E3Bp4GAxEHj3WTAwNeAqYGAxcpB14CvgYDFwWVlxMGwgYDIwOZCQZuAgMFA5s5B24CFwMFBVedGwfiBlYCAxEDnysH8gYXAxEFlaEtBwIHFwMRA6MDA3oCEgcDFykHegIqBwMXBaWnEwZ+AgMjA6kJBn4CAwUDqysHOgcXAwUFV58tB0oHFwMFA68hB1oHFwMFBbFdLwdqBxcDBQWtswMDgQUDAwMDgQUDAwUGgQMFBxO3uQsFgSEJnxO3uQMDgwUDAwMDgwUDAwUGgwMFBxW9vwsFgyEJtRW9vwMDSwUDAwcGSwMDA1EDA0sFAwMFBksDGwkPw8XHEwZLAxkDyTsGngcDBQPLAwOCAq4HAwUlB4ICwgcDBQelzc8bB8YH9wMFA7EDA70FAwMDA70FAwMFBr0DBQcX1dchB+IHFwMFBdPZLwfyBxcDBQXb0QMDhQUDAwMDhQUDAwUGhQMFBxff4QsFhSEJ3Rff4QMDKgLzAwEZADsDAQUZADsDA/XzAwERB/XrAwsFBUMdBmoEAwEDRQMDPwsDAREHP6UDCwVHSR8UPwNLCQNDmQMDqwUDAwMDqwUDAwUGqwMFBxVNTwMDrYIEAwkJBq0DBQNTMQetFwMFBVVRGweeBPcDBQNXAwOxBQMDAwOxBQMDBQaxAwUHF1tdIQe6BBcDBQVfWTMGygQDGQNhAwMrBQMDAwMrBQMDAwMrBQMDBQYrAxsJGWVnaRMGKwMZA2sTBisDGwNjCwUr5gQLbxllZ2kDAx4CQQMJCQYeAgMFA3EDA3UFAwMDA3UFAwMFBnUDBQcTdXcLBXUhCXMTdXcDAyICQQMJCQYiAgMFA3sDA3cFAwMDA3cFAwMFBncDBQcVf4ELBXchCX0Vf4EDAyYCQQMJCQYmAgMFA4UDA3kFAwMDA3kFAwMFBnkDBQcXiYsLBXkhCYcXiYsZAD8DAQUZAD8PAAENEQEKAwcDDw8LAQEBAQEBBwEHAQMDAQsDAQMDAQsDAQ8EAQcBAwsNEQEOAwcDJz8LAQEBAQEBBwEHAQMDaQUDAwcGaQMDAwMHBmkDAwMFFQZpAxMJCQsNDxcG8gMDAQMRAwPlCwMBEQflmwMLBRMVAwNrBQMDBwZrAwMDAwcGawMDAwUVBmsDEwkHGRsdFwYGBAMBAx8DAwELAwEDAwELAwEPBAEHASEjDREBEgMHAyc/CwEBAQEBAQcBBwEDA1UFAwMHBlUDAwMDBwZVAwMDBRUGVQMTCQkLDQ8XBpoDAwEDEQMD3wsDAREH35sDCwUTFQMDZwUDAwcGZwMDAwMHBmcDAwMFFQZnAxMJBxkbHRcG1gMDAQMfAwMBCwMBAwMBCwMBDwQBBwEhIw0RARYDBwMPDwsBAQEBAQEHAQcBAwMBCwMBAwMBCwMBDwQBBQMLDREBGgMHAxETCwEBAQEBAQcBBwEDAwELAwEDAwELAwEDAwELAwEPBAEFCw0NEQEeAwcDERMLAQEBAQEBBwEHAQMDAQsDAQMDAQsDAQMDAQsDAQ8EAQULDQ0RASIDBwMREwsBAQEBAQEHAQcBAwMBCwMBAwMBCwMBAwMBCwMBDwQBBQsNDREBJgMHAw8PCwEBAQEBAQcBBwEDAwELAwEDAwELAwEPBAEHAQMLBgMBBQEA6h/hGQsZFQ0qAmUJDR1JDRMLX0ETI2U/JTM1Xx0jISMpMS0LCx8LHR0lGxEpDQkZGRkZGRkZGQsVDQkdCxEVdx1LMwsvHSUlHQ0TLQ1JC0syAhcfGxMbFxcTFy8XFxcXDxkXFRkZJRcZFSMjIxkfDw8NCR0RYnVpbHRpbgBzdGFibGVfbW9zYWljAHRwdQBhcml0aAB2ZWN0b3IAbW9kdWxlAGFyaXRoLmNvbnN0YW50AHZlY3Rvci5sb2FkAGFyaXRoLmluZGV4X2Nhc3QAdmVjdG9yLmJyb2FkY2FzdAB0cHUudmVjdG9yX3N0b3JlAGZ1bmMuZnVuYwBmdW5jLnJldHVybgBhcml0aC5jbXBpAHZlY3Rvci5zaGFwZV9jYXN0AG1lbXJlZi5sb2FkAGFyaXRoLmV4dHNpAHNjZi55aWVsZAB0cHUucmVwZWF0AGFyaXRoLmV4dHVpAHNjZi5pZgBhcml0aC5tdWxmAGFyaXRoLm11bGkAdHB1Lm1hdG11bABhcml0aC5hZGRpAHZlY3Rvci5tdWx0aV9yZWR1Y3Rpb24AYXJpdGguc3ViZgBtYXRoLmV4cABhcml0aC5hZGRmAGFyaXRoLmRpdmYAYXJpdGgudHJ1bmNmAHRwdS5pb3RhAGFyaXRoLnNlbGVjdABhcml0aC5tYXhpbXVtZgBhcml0aC5leHRmAC9ob21lL3B0b3VsbWUvbWluaWNvbmRhMy9lbnZzL3ZsbG0vbGliL3B5dGhvbjMuMTIvc2l0ZS1wYWNrYWdlcy9qYXgvZXhwZXJpbWVudGFsL3BhbGxhcy9vcHMvdHB1L3NwbGFzaF9hdHRlbnRpb24vc3BsYXNoX2F0dGVudGlvbl9rZXJuZWwucHkAZmxhc2hfYXR0ZW50aW9uX2tlcm5lbC48bG9jYWxzPi5ib2R5AC9nZXQAZmxhc2hfYXR0ZW50aW9uX2tlcm5lbC48bG9jYWxzPi5lbmQAdmFsdWUAL2NvbnZlcnRfZWxlbWVudF90eXBlAHN5bV9uYW1lAC9zd2FwAGZ1bmN0aW9uX3R5cGUAL2Jyb2FkY2FzdF9pbl9kaW0AdHJhbnNmb3JtX2luZGljZXMAd2luZG93X2JvdW5kcwBmbGFzaF9hdHRlbnRpb25fa2VybmVsAC9tdWwAX2FwcGx5X21hc2tfYW5kX3NvZnRfY2FwAGZsYXNoX2F0dGVudGlvbl9rZXJuZWwuPGxvY2Fscz4uaW5pdABfbmV4dF9ub256ZXJvAC9ob21lL3B0b3VsbWUvanVzdGFieXRlL3RwdV9wYWxsYXNfcG9zdC8uL3BhbGxhc19rZXJuZWwucHkAcHJlZGljYXRlAC9yZXBlYXQAL2FkZABwaXBlbGluZV9tb2RlAC9ndAAvY29uZABkaW1lbnNpb24AbWFpbgB0cmFuc2Zvcm1fMAB0cmFuc2Zvcm1fMQB0cmFuc2Zvcm1fMgB0cmFuc2Zvcm1fMwB0cmFuc2Zvcm1fNAB0cmFuc2Zvcm1fNQB0cmFuc2Zvcm1fNgB0cmFuc2Zvcm1fNwAvZXEAdGltZXMAb3BlcmFuZFNlZ21lbnRTaXplcwBzdHJpZGVzAC9kb3RfZ2VuZXJhbABkaW1lbnNpb25fbnVtYmVycwB0cmFuc3Bvc2VfbGhzAHRyYW5zcG9zZV9yaHMAa2luZAByZWR1Y3Rpb25fZGltcwAvc3ViAC9leHAAc3RhYmxlX21vc2FpYy52ZXJzaW9uAHNwbGFzaF9taGFfZndkX2M2MDJiMTZkAGRpbWVuc2lvbl9zZW1hbnRpY3MAaXRlcmF0aW9uX2JvdW5kcwBzY2FsYXJfcHJlZmV0Y2gAc2NyYXRjaF9vcGVyYW5kcwB3aW5kb3dfcGFyYW1zAF9zcGxhc2hfYXR0ZW50aW9uX2ZvcndhcmQuPGxvY2Fscz4udl9pbmRleF9tYXAAX3NwbGFzaF9hdHRlbnRpb25fZm9yd2FyZABfc3BsYXNoX2F0dGVudGlvbl9jdXN0b20AX3NwbGFzaF9hdHRlbnRpb24AU3BsYXNoQXR0ZW50aW9uS2VybmVsLl9fY2FsbF9fAGJlbmNobWFya19rZXJuZWwuPGxvY2Fscz4uc3BsYXNoX2F0dGVudGlvbl9rZXJuZWwAYmVuY2htYXJrX2tlcm5lbAA8bW9kdWxlPgBfbmV4dF9ub256ZXJvLjxsb2NhbHM+LjxsYW1iZGE+AF9zcGxhc2hfYXR0ZW50aW9uX2ZvcndhcmQuPGxvY2Fscz4ua19pbmRleF9tYXAAL2RpdgBmYXN0bWF0aAAvc2NhbgBmbGFzaF9hdHRlbnRpb25fa2VybmVsLjxsb2NhbHM+LnJ1bgBvdmVyZmxvd0ZsYWdzAC9pb3RhAC9nZQBDYXVzYWxNYXNrLl9faW5pdF9fLjxsb2NhbHM+LmNhdXNhbF9tYXNrX2Z1bmN0aW9uAC9ob21lL3B0b3VsbWUvbWluaWNvbmRhMy9lbnZzL3ZsbG0vbGliL3B5dGhvbjMuMTIvc2l0ZS1wYWNrYWdlcy9qYXgvZXhwZXJpbWVudGFsL3BhbGxhcy9vcHMvdHB1L3NwbGFzaF9hdHRlbnRpb24vc3BsYXNoX2F0dGVudGlvbl9tYXNrLnB5AC9waml0AC9zZWxlY3RfbgAvcmVkdWNlX21heAAvbWF4AC9yZWR1Y2Vfc3VtAA==", "serialization_format": 1, "needs_layout_passes": true}}
  ROOT %get-tuple-element.0 = bf16[8,2048,128]{2,1,0} get-tuple-element(%custom-call.0), index=3, metadata={op_name="jit(splash_attention_kernel)/jit(main)/splash_attention_kernel/splash_kernel_b1024_h8_s2048/jit(_splash_attention)/splash_mha_fwd_c602b16d/pallas_call" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=1100}
}

The key insight: XLA sees the entire Splash Attention kernel as one opaque custom-call. The backend_config contains the entire kernel body as base64-encoded MLIR (that "TUzvUg..." blob). XLA can’t optimize inside it — it just hands the blob to Mosaic for TPU-specific compilation.

Notice the scratch buffers in the return tuple: f32[1024,128] appears three times. These are the online softmax accumulators (m_scratch, l_scratch, o_scratch). The fourth element — bf16[8,2048,128] — is the actual output.

The s8[1,2,2] block masks encode which blocks to process. With seq_len=2048 and block_size=1024, we have 2048/1024 = 2 blocks per dimension, hence the 2×2 shape.

Mosaic Compilation Pipeline

When the HLO custom-call with custom_call_target="tpu_custom_call" reaches the TPU backend, the base64-encoded kernel body gets deserialized and fed to Mosaic. The compilation proceeds through 12 passes:

Stage 1: Original MLIR

The kernel arrives as high-level MLIR with explicit memory spaces and TPU-specific operations:

func.func @main(
  %arg0: i32,  // grid index: head
  %arg1: i32,  // grid index: q_block
  %arg2: i32,  // grid index: kv_block
  %arg3: memref<1x2x2xi8, #tpu.memory_space<smem>>,   // block_mask
  %arg4: memref<1x2x2xi8, #tpu.memory_space<smem>>,   // data_next
  %arg5: memref<1x1024x128xbf16, #tpu.memory_space<vmem>>,  // Q block
  %arg6: memref<1x1024x128xbf16, #tpu.memory_space<vmem>>,  // K block
  %arg7: memref<1x1024x128xbf16, #tpu.memory_space<vmem>>,  // V block
  %arg8: memref<1024x128xi32, #tpu.memory_space<vmem>>,     // q_sequence (for causal)
  %arg9:  memref<1024x128xf32, #tpu.memory_space<vmem>>,    // m_scratch
  %arg10: memref<1024x128xf32, #tpu.memory_space<vmem>>,    // l_scratch
  %arg11: memref<1024x128xf32, #tpu.memory_space<vmem>>,    // o_scratch
  %arg12: memref<1x1024x128xbf16, #tpu.memory_space<vmem>>  // output
) attributes {
  dimension_semantics = [#tpu.dimension_semantics<parallel>,   // heads
                         #tpu.dimension_semantics<arbitrary>,  // q_blocks
                         #tpu.dimension_semantics<arbitrary>], // kv_blocks
  iteration_bounds = array<i64: 8, 2, 2>,  // 8 heads, 2 q_blocks, 2 kv_blocks
  scalar_prefetch = 2 : i64,  // prefetch first 2 args (mask info)
  window_params = [...]
}

The function attributes encode Pallas’s grid specification. iteration_bounds = [8, 2, 2] means the kernel runs 8 × 2 × 2 = 32 iterations. The first dimension is parallel (heads can run independently), while Q and KV blocks are arbitrary (have data dependencies through scratch buffers).

The kernel body shows the attention computation:

// Load Q block, K block
%q = vector.load %arg5[%c0, %c0, %c0] : vector<1x1024x128xbf16>
%k = vector.load %arg6[%c0, %kv_idx, %c0] : vector<1x1024x128xbf16>

// QK^T matmul -> [1024, 1024] attention scores
%qk = tpu.matmul %21, %24, %cst {
  dimension_numbers = #tpu.dot_dimension_numbers<[1], [1], [0], [0], [0, 0, 1, 0], [], []>
} : vector<1024x128xbf16>, vector<1024x128xbf16>, vector<1024x1024xf32>

// Apply causal mask using iota comparison
%29 = tpu.iota {dimension = 1 : i32} : vector<1024x1024xi32>
%34 = arith.cmpi sge, %33, %31 : vector<1024x1024xi32>
%36 = arith.select %34, %25, %35 : vector<1024x1024xi1>, vector<1024x1024xf32>

// Online softmax: max reduction
%37 = vector.multi_reduction <maximumf>, %36, %cst_20 [1] : vector<1024x1024xf32> to vector<1024xf32>
%40 = arith.maximumf %18, %39 : vector<1024x128xf32>  // update running max

// exp(scores - max) and sum reduction
%42 = arith.subf %36, %41 : vector<1024x1024xf32>
%43 = math.exp %42 : vector<1024x1024xf32>
%44 = vector.multi_reduction <add>, %43, %cst_21 [1] : vector<1024x1024xf32> to vector<1024xf32>

// Rescale previous accumulator: alpha = exp(m_prev - m_next)
%47 = arith.subf %18, %40 : vector<1024x128xf32>
%48 = math.exp %47 : vector<1024x128xf32>
%49 = arith.mulf %48, %19 : vector<1024x128xf32>  // alpha * l_prev
%50 = arith.addf %46, %49 : vector<1024x128xf32>  // l_next = l_curr + alpha * l_prev

// S @ V matmul for current output contribution
%57 = tpu.matmul %43, %56, %cst_28 {
  dimension_numbers = #tpu.dot_dimension_numbers<[1], [0], [0], [1], [0, 0, 1, 1], [], []>
} : vector<1024x1024xf32>, vector<1024x128xf32>, vector<1024x128xf32>

// Update output accumulator: o_next = alpha * o_prev + o_curr
%60 = arith.mulf %58, %59 : vector<1024x128xf32>  // alpha * o_prev
%61 = arith.addf %60, %57 : vector<1024x128xf32>  // o_next

// Store updated accumulators
tpu.vector_store %arg9[...], %40   // m_scratch
tpu.vector_store %arg10[...], %50  // l_scratch
tpu.vector_store %arg11[...], %61  // o_scratch

Stage 2-3: Layout and Tiling

The infer-vector-layout and tiling passes map abstract vectors onto the VPU’s physical structure (8 sublanes × 128 lanes):

// Before: abstract vector and memref
%qk = tpu.matmul %q, %k, %zeros : vector<1024x1024xf32>
%arg5: memref<1x1024x128xbf16, #tpu.memory_space<vmem>>

// After: explicit layouts matching hardware
%qk = tpu.matmul %q, %k, %zeros {
  in_layout = [#tpu.vpad<"16,{0,0},(16,128)">,   // bf16: 16 sublanes
               #tpu.vpad<"16,{0,0},(16,128)">],
  out_layout = [#tpu.vpad<"32,{0,0},(8,128)">]   // f32: 8 sublanes
} : ...

%arg5: memref<1x1024x128xbf16, 
  #tpu.tiled<(8,128)(2,1)>,  // 8×128 tiles, 2 bf16 per 32-bit word
  #tpu.memory_space<vmem>>

The (8,128) tile shape matches the VPU dimensions. When layouts don’t match between operations, the compiler inserts explicit relayouts.

Stage 4: Lower to LLO

The lower-to-llo pass transforms everything to Low-Level Operations — the final MLIR dialect before machine code:

// High-level vector store
tpu.vector_store %arg11[%c0, %c0], %result : memref<1024x128xf32>, vector<1024x128xf32>

// Becomes LLO with explicit addressing
%addr = "llo.saddr_scaled"(%arg11, %offset) <{multiplier_in_bytes = 512}>
llo.vector_store %result into %addr : vector<8x128xf32> into i32
llo.vector_store %result into %addr + %8 : vector<8x128xf32> into i32
// ... (128 stores for 1024 rows at 8 rows per tile)

The 1024×128 vector gets broken into 128 tiles of 8×128, each stored with computed VMEM addresses.

Pallas LLO Compilation: From Mosaic to Machine Code

Next, the Mosaic IR is converted to native TPU LLO IR. We will now trace the computation through the LLO compiler.

For this LLO walkthrough, we use block_size=512 (grid 8×4×4) to show more loop iterations and better demonstrate software pipelining. The HLO and Mosaic sections above use block_size=1024.

Kernel Structure (Pass 02: Original)

The initial LLO reveals the kernel's memory layout and loop structure:

$region0: #{splash_mha_fwd_e9701070.1}
  // Memory allocations
  #allocation1 [shape = 'u32[144,128]', space=vmem, size = 0x12000, tag = 'internal scratch']
  #allocation3 [shape = 'u8[512]', space=smem, size = 0x200, tag = 'prefetched SMEM operand 0']
  #allocation4 [shape = 'u8[512]', space=smem, size = 0x200, tag = 'prefetched SMEM operand 1']
  
  // Kernel inputs/outputs
  %s0 = inlined_call_operand.hbm [shape: s8[1,4,4], index: 0]   // block_mask
  %s1 = inlined_call_operand.hbm [shape: s8[1,4,4], index: 1]   // data_next
  %s2 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 2]  // Q
  %s3 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 3]  // K
  %s4 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 4]  // V
  %s5 = inlined_call_operand.vmem [shape: s32[2048,128], index: 5]    // q_sequence (causal)
  
  // Scratch buffers for online softmax
  %s6 = inlined_call_operand.hbm [shape: f32[512,128], index: 6]  // m_scratch output
  %s7 = inlined_call_operand.hbm [shape: f32[512,128], index: 7]  // l_scratch output
  %s8 = inlined_call_operand.hbm [shape: f32[512,128], index: 8]  // o_scratch output
  %s9 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 9]  // final output
  
  // Prefetch mask data to SMEM before main loop
  %16 = dma.hbm_to_smem /*hbm=*/%s0, /*size=*/16, /*smem=*/[#allocation3]
  %18 = dma.hbm_to_smem /*hbm=*/%s1, /*size=*/16, /*smem=*/[#allocation4]
  %19 = dma.done [#allocation2], 32

The kernel operand layout reflects Pallas’s PrefetchScalarGridSpec:

• Operands 0-1: Block mask indices prefetched to SMEM

• Operands 2-4: Q, K, V matrices streaming from HBM

• Operand 5: Causal sequence indices (pre-loaded to VMEM)

• Operands 6-8: Online softmax scratch buffers

Loop Structure with Phi Nodes

With block_size=512 and seq_len=2048, the grid is (8 heads, 4 q_blocks, 4 kv_blocks). Software pipelining adds 2 iterations for warmup/cooldown:

loop: start=0, step=1, limit=130   // 8×4×4 = 128 + 2 for pipeline

$region107: #{splash_mha_fwd_e9701070.1} parent=1 // loop_body
  // Iteration tracking across 3 dimensions
  %s6473 = sphi 0, %s37    /* iteration index, stage = 0 */
  %s6475 = sphi 0, %s59    /* iter bound = 0 (head index) */
  %s6477 = sphi 0, %s55    /* iter bound = 1 (q_block index) */
  %s6479 = sphi 0, %s51    /* iter bound = 2 (kv_block index) */
  
  // Pipelined stage tracking (stage 1 lags stage 0)
  %s6481 = sphi 0, %s6475  /* stage = 1 iter bound = 0 */
  %s6483 = sphi 0, %s6477  /* stage = 1 iter bound = 1 */
  %s6485 = sphi 0, %s6479  /* stage = 1 iter bound = 2 */
  
  // Online softmax state (carries across KV blocks)
  %s6487 = sphi 0, %s66    /* running max state */
  %s6493 = sphi 0, %s128   /* running sum state */
  %s6499 = sphi 0, %s190   /* running output state */
  
  // Dimension wrap-around logic
  %s49 = sadd.s32 1, %s6479               /* kv_block + 1 */
  %p50 = scmp.ge.s32.totalorder %s49, 4   /* kv_block >= 4? */
  %s51 = scalar_select %p50, 0, %s49      /* wrap to 0 */
  
  %s52 = sadd.s32 1, %s6477               /* q_block + 1 (if kv wrapped) */
  %s53 = scalar_select %p50, %s52, %s6477 /* conditional increment */
  %p54 = scmp.ge.s32.totalorder %s53, 4   /* q_block >= 4? */
  %s55 = scalar_select %p54, 0, %s53      /* wrap to 0 */
  
  %s56 = sadd.s32 1, %s6475               /* head + 1 (if q wrapped) */
  %p58 = scmp.ge.s32.totalorder %s57, 8   /* head >= 8? */
  %s59 = scalar_select %p58, 0, %s57      /* wrap to 0 */

The phi nodes track iteration state across the flattened 3D grid. The stage = 1 variables implement software pipelining—while stage 0 loads the next tile, stage 1 processes the current tile.

DMA Operations with Double Buffering

The kernel uses double buffering for async DMA overlap:

// Compute buffer slot for double buffering
%s14824_s0 = sand.u32 1, %s8466_s28   /* slot = iteration % 2 */
%s344_s11 = scalar_lea.sflag [#allocation6], %s15337_s19  /* sync flag */

// Compute HBM source address
%s6544_s15 = sshll.u32 %s353_s16, 6   /* address offset */
%s355_s3 = scalar_lea.hbm %s15333_s2, %s6544_s15

// Compute VMEM destination
%s347_s24 = scalar_lea.vmem [#allocation5], %s6541_s30
%s358_s1 = int_to_ptr.vmem %s357_s1

// Issue async DMA: HBM → VMEM
%7188 = dma.hbm_to_vmem [thread:$0] (!%p8805_p8), 
        /*hbm=*/%s355_s3, 
        /*size_in_granules=*/4096,   /* 256KB for 512×128 bf16 */
        /*vmem=*/%s358_s1, 
        /*dst_syncflagno=*/%s344_s11
/* metadata:
   window_bounds: (1, 64, 1)    // tile shape
   iteration_bounds: (8, 4, 4)  // grid dimensions
   element_size: 2048 bytes */

// Wait for previous DMA before consuming buffer
%8381 = dma.done.wait (%p15353_p5), %s454_s2, 4096

The (!%p8805_p8) predicate enables block skipping—if the current tile is fully masked, the DMA is skipped entirely.

Matmul Operations (Pass 13: Post-MXU-Assigner)

After MXU assignment, all matmuls are initially on mxu0:

// Q @ K^T for attention scores (line 749 in splash_attention_kernel.py)
%1237 = vmatprep.subr.mxu0 0
%1238 = vmatpush1.bf16.xpose.msra.mxu0 %v1164
%1269 = vmatprep.mubr.bf16.mxu0 0
%1270 = vmatmul.mubr.bf16.gmra.mxu0 %v1047

// S @ V for output accumulation (line 801)
%1279 = vmatprep.mubr.bf16.mxu0 0
%1280 = vmatmul.mubr.bf16.gmra.mxu0 %v1050

The operations decode as:

• vmatprep.subr: Prepare RHS with subtraction (accumulator init)

• vmatpush1.bf16.xpose: Push bf16 tile to MXU, transposed

• vmatmul.mubr.bf16.gmra: Multiply-accumulate with bf16 inputs

Online Softmax in LLO

The online softmax pattern appears in three stages:

1. Max Reduction Tree:

// Tree reduction: max across 4 tiles
%v2530 = vmax.f32 %v9322, %v9326   /* max(tile0, tile1) */
%v2531 = vmax.f32 %v2530, %v9324   /* max(result, tile2) */
%v2532 = vmax.f32 %v2531, %v9343   /* max(result, tile3) */

// Cross-lane reduction using XLU
%2533 = vmax.xlane.f32.xlu0 %v2532   /* reduce across 128 lanes */
%v2534 = vpop.xlane.xlu0 %2533       /* pop result */

2. Exponential via Base Conversion:

// exp(x) = 2^(x × log₂(e)) where log₂(e) = 1.442695
%v3170 = vmul.f32 1.442695, %v2914   /* x × log₂(e) */
%7348 = vpow2.f32 %v3170             /* 2^result = exp(x) */
%v10295 = vpop.eup %7348             /* pop from EUP */

3. Running Max Update (the key online softmax insight):

// m_next = max(m_prev, m_curr)
%v10156 = vmax.f32 %v10130, %v2539   /* running max update */

// scores - max (for numerical stability)
%v2918 = vsub.f32 %v9345, %v10156
%v2919 = vsub.f32 %v9347, %v10156

Final VLIW Bundles (Pass 79)

The final bundles show aggressive instruction-level parallelism. Here’s the entry bundle:

0x0 : { 
  %s8468_s30 = smov [#allocation3]  ;;  // Load SMEM constant addresses
  %s8469_s12 = smov [#allocation4]  ;;
  %s14794_s0 = inlined_call_operand.hbm [shape: s8[1,4,4], index: 0]  ;;  // block_mask
  %s14795_s2 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 2]  ;;  // Q
  %s14796_s3 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 3]  ;;  // K
  %s14797_s4 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 4]  ;;  // V
  ...
} /* entry bundle */

The Online Softmax Interleaving

The most impressive bundles show matmul, softmax, and mask operations executing in parallel:

0x20b : > { 
  %v2535_v40 = vmax.f32 %v9345, %v9347  ;;        // Tree reduction: max
  %v1281_v41 = vpop.f32.mrf.mxu0  ;;               // Pop Q@K^T result from MXU0
  %7068 = vmatmul.mubr.bf16.gmra.mxu0 %v6597  ;;   // Start next matmul on MXU0
  %v2532_v42 = vmax.f32 %v2531, %v9343  ;;        // Continue max reduction
  %v9365_v48 = vsel %vm2025, %v1630, -2.38e+38  ;; // Apply causal mask
  %v1960_v34 = vld [vmem:[%s8897_s21 + $0x30]]  ;; // Load next Q sequence index
  %vm2041 = vcmp.ge.s32 %v1959, %v9314  ;;        // Compute next mask predicate
}

In a single VLIW bundle:

1. Pop matmul result from MXU

2. Start next matmul

3. Tree reduction for running max

4. Apply causal mask via vsel

5. Load next tile

6. Compute next mask predicate

This is the power of Pallas—the online softmax operations interleave with matmul, eliminating the HBM round-trips that plague reference attention.

Dual-MXU Distribution (Final Bundles)

The VLIW scheduler distributes work across both MXUs:

0x20c : > { 
  %v1634_v45 = vpop.f32.mrf.mxu1  ;;              // Pop from MXU1
  %7148 = vmatmul.mubr.bf16.gmra.mxu1 %v6597  ;;  // Matmul on MXU1
  %7069 = vmatprep.mubr.bf16.mxu0 %v9334  ;;      // Prepare MXU0
  %v9367_v49 = vsel %vm2026, %v1281, -2.38e+38  ;; // Mask application
  %vm2042 = vcmp.ge.s32 %v1960, %v9305  ;;
}

What Makes This Possible

The key insight from the LLO is that everything happens in one loop:

1. DMA prefetch loads next Q/K/V tiles while current tiles compute

2. Matmul computes attention scores (Q @ K^T)

3. Online softmax updates running max/sum/output in the same iteration

4. S @ V accumulates output using rescaled softmax weights

5. Scratch buffers in VMEM carry state across KV blocks

Reference attention requires three separate fusions with HBM materialization between them. Pallas keeps everything in VMEM—the 128MB attention matrix never exists.

The Compiler’s Role

The TPU Pallas compiler handled:

• MXU assignment (pass 13): Initially all on mxu0

• VLIW bundle scheduling (pass 29-33): Distributes across both MXUs

• DMA scheduling (pass 35-37): Overlaps memory access with compute

• Register allocation (pass 45-55): Manages VMEM pressure

• Final bundle packing (pass 73-79): Maximizes ILP

The Pallas programmer writes the algorithm; the compiler handles the hardware mapping. But the algorithm itself—online softmax with tiled accumulation—is something no compiler can currently invent. That’s why Pallas exists.

Reference Attention Jax

The reference implementation lives in JAX’s Splash Attention library. It’s the naive O(n²) memory baseline that materializes the full attention matrix:

def _attention_reference_default(mask, q, k, v, segment_ids, mask_value, ...):
    # Q @ K^T -> full [seq, seq] attention matrix
    logits = jnp.einsum("sd,td->st", q.astype(jnp.float32), k.astype(jnp.float32))
    
    # Apply causal mask - masked positions get a large negative value
    logits = jnp.where(mask, logits, mask_value)  # mask_value = -2.38e+38
    
    # Numerically stable softmax
    m = logits.max(axis=-1)              # reduce_max over key dimension
    s = jnp.exp(logits - m[..., None])      # subtract max for stability
    l = s.sum(axis=-1)                    # reduce_sum for normalization
    s = s / l[..., None]                   # normalize to probabilities
    
    # Output projection: S @ V
    o = jnp.einsum("st,td->sd", s, v.astype(jnp.float32))
    return o

Reference Attention HLO

Initial HLO (Before Optimization)

XLA traces the reference implementation into standard HLO ops. Here’s the complete structure:

HloModule jit_reference_attention, entry_computation_layout={(bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)})->f32[8,2048,128]{2,1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

%region_0.25 (Arg_0.22: f32[], Arg_1.23: f32[]) -> f32[] {
  %Arg_0.22 = f32[] parameter(0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max"}
  %Arg_1.23 = f32[] parameter(1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max"}
  ROOT %maximum.24 = f32[] maximum(%Arg_0.22, %Arg_1.23), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
}

%region_1.36 (Arg_0.33: f32[], Arg_1.34: f32[]) -> f32[] {
  %Arg_0.33 = f32[] parameter(0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum"}
  %Arg_1.34 = f32[] parameter(1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum"}
  ROOT %add.35 = f32[] add(%Arg_0.33, %Arg_1.34), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=198}
}

ENTRY %main.47 (Arg_0.1: bf16[8,2048,128], Arg_1.2: bf16[8,2048,128], Arg_2.3: bf16[8,2048,128]) -> f32[8,2048,128] {
  %constant.4 = pred[8,2048,2048]{2,1,0} constant({...})
  %Arg_0.1 = bf16[8,2048,128]{2,1,0} parameter(0), metadata={op_name="q"}
  %convert.0 = f32[8,2048,128]{2,1,0} convert(%Arg_0.1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/convert_element_type" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=184}
  %Arg_1.2 = bf16[8,2048,128]{2,1,0} parameter(1), metadata={op_name="k"}
  %convert.1 = f32[8,2048,128]{2,1,0} convert(%Arg_1.2), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/convert_element_type" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=184}
  %dot.0 = f32[8,2048,2048]{2,1,0} dot(%convert.0, %convert.1), lhs_batch_dims={0}, lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={2}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(sd,td->st)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=184}
  %constant.0 = f32[] constant(-2.38197633e+38)
  %broadcast.1 = f32[8,2048,2048]{2,1,0} broadcast(%constant.0), dimensions={}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}
  %select.1 = f32[8,2048,2048]{2,1,0} select(%constant.4, %dot.0, %broadcast.1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/select_n" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}
  %constant.1 = f32[] constant(-inf)
  %reduce.0 = f32[8,2048]{1,0} reduce(%select.1, %constant.1), dimensions={2}, to_apply=%region_0.25, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
  %reshape.0 = f32[8,2048,1]{2,1,0} reshape(%reduce.0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %broadcast.2 = f32[8,2048,1]{2,1,0} broadcast(%reshape.0), dimensions={0,1,2}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %reshape.1 = f32[8,2048]{1,0} reshape(%broadcast.2), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %broadcast.3 = f32[8,2048,2048]{2,1,0} broadcast(%reshape.1), dimensions={0,1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %subtract.0 = f32[8,2048,2048]{2,1,0} subtract(%select.1, %broadcast.3), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %exponential.0 = f32[8,2048,2048]{2,1,0} exponential(%subtract.0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/exp" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %constant.2 = f32[] constant(0)
  %reduce.1 = f32[8,2048]{1,0} reduce(%exponential.0, %constant.2), dimensions={2}, to_apply=%region_1.36, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=198}
  %reshape.2 = f32[8,2048,1]{2,1,0} reshape(%reduce.1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %broadcast.4 = f32[8,2048,1]{2,1,0} broadcast(%reshape.2), dimensions={0,1,2}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %reshape.3 = f32[8,2048]{1,0} reshape(%broadcast.4), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %broadcast.5 = f32[8,2048,2048]{2,1,0} broadcast(%reshape.3), dimensions={0,1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %divide.0 = f32[8,2048,2048]{2,1,0} divide(%exponential.0, %broadcast.5), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %Arg_2.3 = bf16[8,2048,128]{2,1,0} parameter(2), metadata={op_name="v"}
  %convert.2 = f32[8,2048,128]{2,1,0} convert(%Arg_2.3), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/convert_element_type" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=201}
  ROOT %dot.1 = f32[8,2048,128]{2,1,0} dot(%divide.0, %convert.2), lhs_batch_dims={0}, lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(st,td->sd)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=201}
}

Notice: XLA sees every operation—two dots, a select, two reductions, elementwise ops. This is fundamentally different from Pallas, where XLA sees one opaque custom-call.

Optimized HLO (After Fusion)

After optimization passes, XLA fuses operations into three main kernels plus helper fusions. Here’s the complete optimized HLO:

HloModule jit_reference_attention, is_scheduled=true, entry_computation_layout={(bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)})->f32[8,2048,128]{2,1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

%copy_fusion.2 (input.2: pred[8,2048,2048]) -> pred[8,2048,2048] {
  %input.2 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} parameter(0)
  ROOT %copy.4 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} copy(%input.2)
}

%fused_computation.1 (param_0.19: f32[8,2048], param_1.20: f32[8,2048], param_2.14: pred[8,2048,2048], param_3.4: f32[8,2048,2048]) -> f32[8,2048,2048] {
  %param_2.14 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} parameter(2)
  %fusion.12 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} fusion(%param_2.14), kind=kLoop, output_to_operand_aliasing={{}: (0, {})}, calls=%copy_fusion.2
  %param_3.4 = f32[8,2048,2048]{1,2,0:T(8,128)} parameter(3)
  %constant.21 = f32[]{:T(128)} constant(-2.38197633e+38)
  %broadcast.21 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%constant.21), dimensions={}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":[]}}
  %select.5 = f32[8,2048,2048]{1,2,0:T(8,128)} select(%fusion.12, %param_3.4, %broadcast.21), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/select_n" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}
  %param_1.20 = f32[8,2048]{1,0:T(8,128)S(1)} parameter(1)
  %broadcast.15 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%param_1.20), dimensions={0,1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["1","128"]}}
  %subtract.4 = f32[8,2048,2048]{1,2,0:T(8,128)} subtract(%select.5, %broadcast.15), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %exponential.4 = f32[8,2048,2048]{1,2,0:T(8,128)} exponential(%subtract.4), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/exp" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %param_0.19 = f32[8,2048]{1,0:T(8,128)S(1)} parameter(0)
  %broadcast.11 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%param_0.19), dimensions={0,1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["1","128"]}}
  ROOT %divide.2 = f32[8,2048,2048]{1,2,0:T(8,128)} divide(%exponential.4, %broadcast.11), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
}

%bitcast_fusion.1 (bitcast_input.1: bf16[8,2048,128]) -> bf16[8,2048,128] {
  %bitcast_input.1 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0)
  ROOT %bitcast.1 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.1)
}

%fused_computation (param_0.1: bf16[8,2048,128], param_1.18: f32[8,2048], param_2.12: f32[8,2048], param_3.3: pred[8,2048,2048], param_4: f32[8,2048,2048]) -> f32[8,2048,128] {
  %param_1.18 = f32[8,2048]{1,0:T(8,128)S(1)} parameter(1)
  %param_2.12 = f32[8,2048]{1,0:T(8,128)S(1)} parameter(2)
  %param_3.3 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} parameter(3)
  %param_4 = f32[8,2048,2048]{1,2,0:T(8,128)} parameter(4)
  %fusion.1 = f32[8,2048,2048]{1,2,0:T(8,128)} fusion(%param_1.18, %param_2.12, %param_3.3, %param_4), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/div" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=199}
  %param_0.1 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0)
  %fusion.8 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1), kind=kLoop, calls=%bitcast_fusion.1
  ROOT %convolution-base-dilated.2 = f32[8,2048,128]{2,1,0:T(8,128)} convolution(%fusion.1, %fusion.8), window={size=8 stride=7 lhs_dilate=8}, dim_labels=0bf_0io->0bf, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(st,td->sd)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=201}
}

%region_1.36 (Arg_0.33: f32[], Arg_1.34: f32[]) -> f32[] {
  %Arg_1.34 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum"}
  %Arg_0.33 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum"}
  ROOT %add.35 = f32[]{:T(128)} add(%Arg_0.33, %Arg_1.34), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=198}
}

%copy_fusion.1 (input.1: pred[8,2048,2048]) -> pred[8,2048,2048] {
  %input.1 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} parameter(0)
  ROOT %copy.3 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} copy(%input.1)
}

%fused_computation.3 (param_0.24: f32[8,2048], param_1.26: pred[8,2048,2048], param_2.19: f32[8,2048,2048]) -> f32[8,2048] {
  %param_1.26 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} parameter(1)
  %fusion.11 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} fusion(%param_1.26), kind=kLoop, output_to_operand_aliasing={{}: (0, {})}, calls=%copy_fusion.1
  %param_2.19 = f32[8,2048,2048]{1,2,0:T(8,128)} parameter(2)
  %constant.16 = f32[]{:T(128)} constant(-2.38197633e+38)
  %broadcast.23 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%constant.16), dimensions={}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":[]}}
  %select.7 = f32[8,2048,2048]{1,2,0:T(8,128)} select(%fusion.11, %param_2.19, %broadcast.23), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/select_n" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}
  %param_0.24 = f32[8,2048]{1,0:T(8,128)S(1)} parameter(0)
  %broadcast.16 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%param_0.24), dimensions={0,1}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["1","128"]}}
  %subtract.6 = f32[8,2048,2048]{1,2,0:T(8,128)} subtract(%select.7, %broadcast.16), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/sub" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %exponential.6 = f32[8,2048,2048]{1,2,0:T(8,128)} exponential(%subtract.6), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/exp" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=197}
  %constant.14 = f32[]{:T(128)} constant(0)
  ROOT %reduce.2 = f32[8,2048]{1,0:T(8,128)S(1)} reduce(%exponential.6, %constant.14), dimensions={2}, to_apply=%region_1.36, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=198}
}

%region_0.25 (Arg_0.22: f32[], Arg_1.23: f32[]) -> f32[] {
  %Arg_1.23 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max"}
  %Arg_0.22 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max"}
  ROOT %maximum.24 = f32[]{:T(128)} maximum(%Arg_0.22, %Arg_1.23), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
}

%bitcast_fusion (bitcast_input: bf16[8,2048,128]) -> bf16[8,2048,128] {
  %bitcast_input = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0)
  ROOT %bitcast = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input)
}

%bitcast_fusion.2 (bitcast_input.2: bf16[8,2048,128]) -> bf16[8,2048,128] {
  %bitcast_input.2 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} parameter(0)
  ROOT %bitcast.2 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.2)
}

%copy_fusion (input: pred[8,2048,2048]) -> pred[8,2048,2048] {
  %input = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)} parameter(0)
  ROOT %copy.2 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} copy(%input)
}

%fused_computation.7 (param_0.25: pred[8,2048,2048], param_1.28: bf16[8,2048,128], param_2.21: bf16[8,2048,128]) -> (f32[8,2048], f32[8,2048,2048]) {
  %param_0.25 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)} parameter(0)
  %fusion.10 = pred[8,2048,2048]{1,2,0:T(8,128)(4,1)} fusion(%param_0.25), kind=kLoop, output_to_operand_aliasing={{}: (0, {})}, calls=%copy_fusion
  %param_1.28 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(1)
  %fusion.7 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} fusion(%param_1.28), kind=kLoop, calls=%bitcast_fusion
  %param_2.21 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} parameter(2)
  %fusion.9 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} fusion(%param_2.21), kind=kLoop, calls=%bitcast_fusion.2
  %convolution-base-dilated.3 = f32[8,2048,2048]{1,2,0:T(8,128)} convolution(%fusion.7, %fusion.9), window={size=8 stride=7 lhs_dilate=8}, dim_labels=0bf_0oi->0bf, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(sd,td->st)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=184}
  %constant.22 = f32[]{:T(128)} constant(-2.38197633e+38)
  %broadcast.25 = f32[8,2048,2048]{1,2,0:T(8,128)} broadcast(%constant.22), dimensions={}, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/broadcast_in_dim" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":[]}}
  %select.9 = f32[8,2048,2048]{1,2,0:T(8,128)} select(%fusion.10, %convolution-base-dilated.3, %broadcast.25), metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(jit(_where))/select_n" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=195}
  %constant.15 = f32[]{:T(128)} constant(-inf)
  %reduce.3 = f32[8,2048]{1,0:T(8,128)S(1)} reduce(%select.9, %constant.15), dimensions={2}, to_apply=%region_0.25, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
  ROOT %tuple = (f32[8,2048]{1,0:T(8,128)S(1)}, f32[8,2048,2048]{1,2,0:T(8,128)}) tuple(%reduce.3, %convolution-base-dilated.3)
}

ENTRY %main.47 (Arg_0.1: bf16[8,2048,128], Arg_1.2: bf16[8,2048,128], Arg_2.3: bf16[8,2048,128]) -> f32[8,2048,128] {
  %Arg_0.1 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} parameter(0), metadata={op_name="q"}
  %copy-start = (bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%Arg_0.1), cross_program_prefetch_index=0
  %constant.4 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)} constant({...})
  %Arg_2.3 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} parameter(2), metadata={op_name="v"}
  %Arg_1.2 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)} parameter(1), metadata={op_name="k"}
  %copy-start.1 = (pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)}, pred[8,2048,2048]{1,2,0:T(32,128)(4,1)}, u32[]{:S(2)}) copy-start(%constant.4)
  %copy-done = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start)
  %fusion.5 = (f32[8,2048]{1,0:T(8,128)S(1)}, f32[8,2048,2048]{1,2,0:T(8,128)}) fusion(%constant.4, %copy-done, %Arg_1.2), kind=kOutput, calls=%fused_computation.7, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(sd,td->st)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=184}
  %get-tuple-element.1 = f32[8,2048,2048]{1,2,0:T(8,128)} get-tuple-element(%fusion.5), index=1, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
  %get-tuple-element = f32[8,2048]{1,0:T(8,128)S(1)} get-tuple-element(%fusion.5), index=0, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_max" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=196}
  %copy-start.2 = (bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[8,2048,128]{2,1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%Arg_2.3)
  %copy-done.1 = pred[8,2048,2048]{1,2,0:T(32,128)(4,1)S(1)} copy-done(%copy-start.1)
  %fusion.2 = f32[8,2048]{1,0:T(8,128)S(1)} fusion(%get-tuple-element, %copy-done.1, %get-tuple-element.1), kind=kLoop, calls=%fused_computation.3, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/reduce_sum" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=198}
  %copy-done.2 = bf16[8,2048,128]{2,1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start.2)
  ROOT %fusion = f32[8,2048,128]{2,1,0:T(8,128)} fusion(%copy-done.2, %fusion.2, %get-tuple-element, %copy-done.1, %get-tuple-element.1), kind=kOutput, calls=%fused_computation, metadata={op_name="jit(reference_attention)/jit(main)/reference_attention/reference_attention_h8_s2048/jit(_wrapped)/vmap(st,td->sd)/dot_general" source_file="/home/ptoulme/miniconda3/envs/vllm/lib/python3.12/site-packages/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py" source_line=201}
}

Key Optimization Patterns

1. Multi-Output Fusion (fusion.5)

The first matmul fuses with mask application and reduce_max. It returns a tuple with both outputs:

• f32[8,2048] — the row-wise maximum values

• f32[8,2048,2048] — the full attention scores

2. Recomputation vs Storage Trade-off

Notice that exp(scores - max) is computed three times:

• In fusion.2 to compute the sum

• In fused_computation.1 (inside fusion) to compute normalized weights

• The mask is also re-applied multiple times

This is intentional: recomputing is cheaper than storing and reloading the full [8,2048,2048] tensor.

The O(n²) Problem is Visible:

%get-tuple-element.1 = f32[8,2048,2048]{1,2,0:T(8,128)} get-tuple-element(%fusion.5), index=1

This 128MB tensor (8 × 2048 × 2048 × 4 bytes) lives in HBM between fusion.5 and downstream fusions. Pallas avoids this entirely with online softmax.

Reference Attention LLO

After HLO optimization, each fusion compiles through the LLO (Low-Level Operations) pipeline. The fusion.5 kernel alone goes through 79 optimization passes before producing final VLIW bundles.

Original LLO: fusion.5 (Pass 02)

The initial LLO reveals the kernel’s structure before optimization:

$region0: #{fusion.5}
  // VMEM allocations for intermediate results
  #allocation0 [shape = 'f32[262144]{0}', space=vmem, size = 0x100000]     // 1MB scratch
  #allocation7 [shape = 'f32[131072]{0}', space=vmem, size = 0x80000]      // 512KB scratch
  #allocation8 [shape = 's32[1]{0}', space=sflag, size = 0x4]              // sync flag
  
  // Kernel inputs/outputs (from HLO fusion interface)
  %s0 = inlined_call_operand.hbm [shape: pred[8,2048,2048], index: 0]  // mask from HBM
  %s1 = inlined_call_operand.vmem [shape: bf16[8,2048,128], index: 1]  // Q in VMEM
  %s2 = inlined_call_operand.hbm [shape: bf16[8,2048,128], index: 2]   // K from HBM
  %s3 = inlined_call_operand.vmem [shape: f32[8,2048], index: 3]       // max output (VMEM)
  %s4 = inlined_call_operand.hbm [shape: f32[8,2048,2048], index: 4]   // scores output (HBM!)
  
  // Initialize max output to -inf (16 vector stores for 8×2048 shape)
  %7 = vst [vmem:[%s3] sm:$0xff] /*vst_source=*/-inf
  %s8 = scalar_lea.vmem %s3, 8
  %9 = vst [vmem:[%s8] sm:$0xff] /*vst_source=*/-inf
  // ... 14 more stores to initialize max buffer

Key observations:

• Q lives in VMEM but K streams from HBM via DMA

• The full attention scores write to HBM — 128MB output!

• The causal mask also streams from HBM — 32MB of predicates

Loop Structure (Pass 02)

  $region2: #{fusion.5} parent=0
    // Allocations for double-buffered DMA
    #allocation1 [shape = 'u8[524288]{0}', space=vmem, tag = 'operand span for K']  // 512KB
    #allocation2 [shape = 's32[2]{0}', space=sflag]                                  // sync flags
    #allocation4 [shape = 'u8[4194304]{0}', space=vmem, tag = 'operand span for mask'] // 4MB
    #allocation6 [shape = 'u8[16777216]{0}', space=vmem, tag = 'operand span for output'] // 16MB
    
    // Initialize sync flags for double buffering
    %38 = vsyncpa [#allocation2], 0
    %40 = vsyncpa [#allocation2 + $0x1], 0
    
    loop: start=0, step=1, limit=18
    
    $region4: #{fusion.5} parent=2 // loop_header
      // Phi nodes track iteration state across loop iterations
      %s48 = sphi 0, %s52 /* iteration index, stage = 0 */
      %p49 = scmp.ge.s32.totalorder %s48, 18 /* loop exit test */
      
      // Multi-dimensional iteration bounds (8 heads × 2 K-blocks)
      %s55 = sphi 0, %s88 /* iter bound = 0 (head index) */
      %s56 = sphi 0, %s84 /* iter bound = 1 (K block index) */
      // ... more phi nodes for software pipelining stages

Loop dimensions: The kernel tiles K into 2 blocks of 1024 rows each:

• 8 heads × 2 K-blocks = 16 main iterations

• Plus 2 for software pipelining = 18 total iterations

Matmul Operations (Pass 02 - Original)

Before optimization, each matmul tile uses explicit load-store patterns:

// Load Q slice, unpack bf16, push to MXU
%v322 = vld [vmem:[%s321] sm:$0xf]           // Load Q tile (bf16)
%v323 = vunpack.c.l.bf16 %v322               // Unpack low bf16 to f32
%325 = vst [vmem:[%s320] sm:$0xff] /*vst_source=*/%v323  // Store unpacked
%v326 = vld [vmem:[%s320] sm:$0xff]          // Reload (wasteful!)
%327 = vmatpush1.xpose.msra.mxu0 %v326       // Push to MXU systolic array

// Prepare RHS with zeros (accumulator init)
%319 = vmatprep.subr.mxu0 0.0                // Prepare zero for subtraction

// Repeat for all 16 tiles of the 128-dim K dimension...
%328 = vmatprep.subr.mxu0 0.0
%v331 = vld [vmem:[%s330] sm:$0xf]
%v332 = vunpack.c.l.bf16 %v331
%334 = vst [vmem:[%s329] sm:$0xff] /*vst_source=*/%v332
%v335 = vld [vmem:[%s329] sm:$0xff]
%336 = vmatpush1.xpose.msra.mxu0 %v335
// ... continues for all K tiles

This is verbose and inefficient — later passes eliminate redundant loads/stores.

VLIW Bundle Packing (Pass 29)

The bundle packer groups independent operations into VLIW bundles:

0x1   :  { %7 = vst [vmem:[%s3] sm:$0xff] /*vst_source=*/%v57179  ;;  
           %51239 = vst [vmem:[%s3 + $0x8] sm:$0xff] /*vst_source=*/%v57179  ;;  
           %51240 = vst [vmem:[%s3 + $0x10] sm:$0xff] /*vst_source=*/%v57179  ;; 
           // ... 16 parallel stores in one bundle!
         }

// Later: DMA + address computation in parallel
0x1e  : > { %57143 = dma.hbm_to_vmem [thread:$0] (!%p57141), /*hbm=*/%s192, 
                     /*size_in_granules=*/8192, /*vmem=*/%s195, /*dst_syncflagno=*/%s180 }
0x1f  : > { %p222 = pnand %p51265, %p221  ;;  
            %s202 = scalar_lea.vmem [#allocation4], %s51260 }

Final VLIW Bundles (Pass 79) - Matmul Section

After all optimizations, both MXUs execute in parallel:

// MXU operations now distributed across both units
0xa   : > { %55856 = vmatprep.subr.mxu0 %v465  ;;  
            %56016 = vmatprep.subr.mxu1 %v7203 }

0xb   : > { %55857 = vmatpush3.xpose.msra.mxu0 %v323  ;;  
            %56017 = vmatpush3.xpose.msra.mxu1 %v7059  ;;  
            %v332 = vunpack.c.h.bf16 %v51273  ;;  
            %v7068 = vunpack.c.h.bf16 %v52005 }

0xc   : > { %55858 = vmatprep.subr.mxu0 %v474  ;;  
            %56018 = vmatprep.subr.mxu1 %v7212  ;;  
            %v483 = vunpack.c.l.bf16 %v51288  ;;  
            %v7221 = vunpack.c.h.bf16 %v52022 }

Key optimization: The VLIW scheduler distributes work across both MXUs:

• mxu0: 1,396 operations (Q @ K^T lower tiles)

• mxu1: 1,383 operations (Q @ K^T upper tiles)

This effectively doubles throughput compared to single-MXU execution.

Final VLIW Bundles - Reduce_Max + Mask Application

The most complex bundles overlap matmul results with mask lookup and max reduction:

0xf4  : > { %v618_v36 = vpop.f32.mrf.mxu0  ;;                    // Pop matmul result
            %55937 = vmatmul.mubr.bf16.gmra.mxu0 %v57834_v53 }   // Start next matmul

0xf5  : > { %52037 = vst [vmem:[%s57903_s21 + $0x10] sm:$0xff] /*vst_source=*/%v7347_v37  ;;  
            %56098 = vmatprep.mubr.bf16.mxu1 %v57863_v16  ;;     // Prepare mxu1
            %v631_v54 = vpop.f32.mrf.mxu0  ;;                    // Pop another result
            %v7358_v6 = vsel %vm, %v7347, -2.38e+38 }            // Apply mask!

0xf9  : > { %v7392_v33 = vsel %vm, %v7381, -2.38e+38  ;;         // Mask application
            %52041 = vst [vmem:[%s57903_s21 + $0x90] sm:$0xff]  ;;  // Store to HBM buffer
            %v665_v37 = vmax.f32 %v627_v12, %v659_v20  ;;        // Tree reduction max
            %v669_v44 = vpop.f32.mrf.mxu0  ;;                    // Pop matmul result
            %vm58016_vm2 = vcmp.ne.s32.totalorder %v854, 0 }     // Prepare next mask

0xfa  : > { %v7399_v43 = vmax.f32 %v7358_v6, %v7392_v33  ;;      // Continue max tree
            %v7403_v53 = vpop.f32.mrf.mxu1  ;;                   // Pop from mxu1
            %v680_v59 = vsel %vm, %v669, -2.38e+38  ;;           // Mask more tiles
            %51310 = vst ... }                                    // Store result

In a single VLIW bundle (0xf9):

• Pop matmul result from MXU

• Apply mask via vsel

• Store to HBM buffer

• Compute tree reduction vmax.f32

• Prepare next mask predicate

This is 5 independent operations in parallel — the power of VLIW scheduling.

Comparison: Pallas vs Reference Jax

Now we can directly compare the two approaches using our test configuration: (8 heads, 2048 seq_len, 128 head_dim).

Algorithm: Online vs Naive Softmax

The fundamental difference lies in the softmax computation strategy:

The reference implementation computes the full [8, 2048, 2048] attention matrix (128MB in f32), stores it to HBM, then reads it back multiple times for softmax operations. Pallas processes small tiles (e.g., [512, 512] or [1024, 1024]) that fit in VMEM, computing and consuming each tile before moving to the next.

Kernel Structure

Reference: 3 Separate Fusions

fusion.5: Q @ K^T + mask + reduce_max  →  2,559 bundles
fusion.2: exp + reduce_sum             →  1,925 bundles  
fusion:   normalize + S @ V            →  5,761 bundles
─────────────────────────────────────────────────────────
Total:                                   ~10,245 bundles

Each fusion requires:

1. Reading inputs from HBM

2. Computing results

3. Writing outputs back to HBM

4. Synchronization before next fusion

Pallas: 1 Custom Kernel (varies by block size)

splash_mha_fwd (block_size=512)   →  1,651 bundles
splash_mha_fwd (block_size=1024)  →  4,302 bundles
splash_mha_fwd (block_size=2048)  → 15,153 bundles

Everything happens in a single kernel with explicit VMEM management. No intermediate HBM round-trips.

VLIW Bundle Comparison

The block size tradeoff is fascinating:

• block_size=512: More iterations (4×4=16 per head), but each iteration is simple. Fewest total bundles.

• block_size=1024: Balanced (2×2=4 per head). Still significantly fewer bundles than reference.

• block_size=2048: One iteration covers entire sequence (1×1=1 per head). More bundles than reference because the large tile operations can’t amortize overhead.

HBM Bandwidth Analysis

For our test case (8 heads, 2048 seq_len, 128 head_dim):

Reference HBM Traffic:

Inputs:
  Q, K, V: 3 × 8 × 2048 × 128 × 2 bytes = 6 MB

Intermediates (written then read):
  S matrix after fusion.5: 8 × 2048 × 2048 × 4 = 128 MB (write)
  S matrix for fusion.2:   8 × 2048 × 2048 × 4 = 128 MB (read)
  exp(S) after fusion.2:   8 × 2048 × 2048 × 4 = 128 MB (write)
  exp(S) for fusion:       8 × 2048 × 2048 × 4 = 128 MB (read)
  
Output:
  O: 8 × 2048 × 128 × 4 = 8 MB

Total: ~6 + 512 + 8 = ~526 MB HBM traffic

Pallas HBM Traffic:

Inputs:
  Q, K, V: 3 × 8 × 2048 × 128 × 2 bytes = 6 MB

Intermediates:
  None! Everything stays in VMEM

Output:
  O: 8 × 2048 × 128 × 4 = 8 MB

Total: ~14 MB HBM traffic

That’s a 37× reduction in HBM bandwidth. Since TPU performance is often memory-bound, this directly translates to speedup.

The Block Size Sweet Spot

Why does block_size=512 produce fewer bundles than block_size=1024, and why does block_size=2048 produce more bundles than reference attention?

At block_size=2048 (equal to seq_len), there’s no tiling—you compute the full [2048, 2048] attention matrix in one iteration. You’ve lost the memory benefit of online softmax but kept all its bookkeeping. It’s reference attention with extra overhead.

At block_size=512 vs 1024, the total math is the same—only the schedule changes. Smaller tiles reduce VMEM footprint, which improves double-buffering and DMA/compute overlap. Larger tiles reduce loop count but increase register pressure, making it harder for the backend to keep the MXUs busy while hiding memory latency.

Why Pallas Matters

Last post, the takeaway was: trust the compiler. This post’s takeaway is: know when not to.

XLA is doing real work on reference attention — it fuses the matmul with the mask and max reduction, it overlaps DMA with compute, it balances both MXUs. The 10,245 VLIW bundles represent a well-optimized implementation of the algorithm you wrote. The problem is the algorithm itself.

Pallas doesn’t replace XLA. For most operations, the automatic path is the right path. But when you’re hitting a fundamental algorithmic limit — when you know there’s a streaming formulation the compiler can’t discover — Pallas is how you escape. You write the tiled, memory-conscious kernel; Mosaic and LLO handle the rest.

The practical insights:

• Block size matters, but depends on your shapes. For our test case (8 heads, seq_len=2048, head_dim=128), block_size=512 produced 6× fewer bundles than reference attention. block_size=2048 produced more bundles. Smaller tiles pipeline better and reduce register pressure — but the optimal block size depends on your sequence length, head count, and head dimension. Profile your actual workload.

• HBM bandwidth is often the bottleneck. Reference attention moved 526MB through HBM; Splash moved 14MB. The 37× reduction in memory traffic is why online softmax wins, not the bundle count.

• The compiler still does heavy lifting. Pallas gives you control over algorithm-level tiling; Mosaic handles vector layouts, MXU scheduling, and VLIW packing automatically. You’re not writing assembly.

The dump flags for tracing Pallas kernels:

os.environ["XLA_FLAGS"] = (
    f"--xla_dump_hlo_as_text "
    f"--xla_dump_to={HLO_PATH} "
    f"--xla_dump_hlo_pass_re=.* "
)
os.environ["LIBTPU_INIT_ARGS"] = (
    f"--xla_jf_dump_to={LLO_PATH} "
    f"--xla_jf_dump_hlo_text=true "
    f"--xla_jf_dump_llo_text=true "
    f"--xla_jf_dump_llo_html=false "
    f"--xla_jf_dump_llo_static_gaps=true "
    f"--xla_jf_emit_annotations=true "
    f"--xla_jf_debug_level=2 "
    f"--xla_mosaic_dump_to={MOSAIC_PATH} "
    f"--xla_mosaic_enable_dump_debug_info=true "
    f"--xla_mosaic_enable_llo_source_annotations=true"
)

The HLO shows an opaque custom-call — the interesting stuff is in the Mosaic MLIR dumps (0001-original.txt through 0012-post-finalize-llo.txt) and the final LLO bundles.

In principle, a compiler could learn online softmax. Someone could teach XLA to recognize the attention pattern, prove the streaming reformulation is valid, and generate the tiled kernel automatically. Until then, we write Pallas.

Questions? Message me on LinkedIn: https://www.linkedin.com/in/patrick-toulme-150b041a5/ or follow me on X: https://x.com/PatrickToulme

Motivation

Eight lines of JAX code in Python become 250 VLIW bundles across 5 fused kernels. A matrix multiply, RMS normalization, softmax, and another matrix multiply — the kind of operation that runs billions of times inside every transformer. Here’s what happens between jax.jit(f)(x) and electrons moving through a TPU.

There’s surprisingly little public information about this — Google’s TPU compiler is closed-source, and the internal IRs are undocumented. I rented a TPU v6e for under a dollar and traced a small computation through four layers of compiler IR.

The key insight: TPUs reward experimentation. GPUs have automatic codegen too — Inductor, XLA, Triton — but in practice, peak performance often still requires hand-tuned kernels. FlashAttention exists because no compiler found that optimization automatically. The TPU compiler gets you closer to the ceiling without manual intervention: it fuses operations, schedules hardware, and orchestrates memory in ways that would require expert kernel engineering on GPU. The performance ceiling might be lower than a perfect custom kernel, but the floor is much higher. This post shows exactly how that works.

Whether you’re debugging why a custom op is slow, deciding between TPUs and GPUs for a new workload, or just curious what jax.jit actually does — this post traces the full compilation path with real IR dumps.

Setup

I have rented a TPU V6e Trillium on Google Cloud for these experiments. The cost is pretty minimal. I think this entire experiment was under one dollar.

Jax Code

I have written some simple Jax code that we can use to trace the TPU’s compilation.

import os

# Create dump directories
DUMP_ROOT = "compiler_dump/"
HLO_DUMP_PATH = os.path.join(DUMP_ROOT, "hlo")
LLO_DUMP_PATH = os.path.join(DUMP_ROOT, "llo")

os.makedirs(HLO_DUMP_PATH, exist_ok=True)
os.makedirs(LLO_DUMP_PATH, exist_ok=True)

os.environ["XLA_FLAGS"] = (
    f"--xla_dump_hlo_as_text "
    f"--xla_dump_to={HLO_DUMP_PATH} "
    f"--xla_dump_hlo_pass_re=.* "
)

os.environ["LIBTPU_INIT_ARGS"] = (
    f"--xla_jf_dump_to={LLO_DUMP_PATH} "
    f"--xla_jf_dump_hlo_text=true "
    f"--xla_jf_dump_llo_text=true "
    f"--xla_jf_dump_llo_html=false "
    f"--xla_jf_dump_llo_static_gaps=true "
    f"--xla_jf_emit_annotations=true "
    f"--xla_jf_debug_level=2"
)

# Import JAX after setting env vars
import jax
import jax.numpy as jnp


@jax.named_call
def matmul_1(x, w1):
    """Stage 1: Linear projection (like Q @ K^T)"""
    return x @ w1


@jax.named_call
def rms_norm(h):
    """Stage 2: RMS Normalization"""
    rms = jnp.sqrt(jnp.mean(h ** 2, axis=-1, keepdims=True) + 1e-6)
    return h / rms


@jax.named_call
def softmax(h):
    """Stage 3: Softmax (row-wise, numerically stable)"""
    h_max = jnp.max(h, axis=-1, keepdims=True)
    exp_h = jnp.exp(h - h_max)
    return exp_h / jnp.sum(exp_h, axis=-1, keepdims=True)


@jax.named_call
def matmul_2(h, w2):
    """Stage 4: Output projection (like attention @ V)"""
    return h @ w2


def mini_attention(x, w1, w2):
    """
    A minimal attention-like block:
    matmul → rms_norm → softmax → matmul
    
    """
    h = matmul_1(x, w1)
    h = rms_norm(h)
    h = softmax(h)
    out = matmul_2(h, w2)
    return out


def main():
    # Small shapes to keep IR readable
    batch, d_in, d_mid, d_out = 16, 64, 64, 32
    
    # Create inputs
    key = jax.random.PRNGKey(42)
    k1, k2, k3 = jax.random.split(key, 3)
    
    x = jax.random.normal(k1, (batch, d_in))
    w1 = jax.random.normal(k2, (d_in, d_mid)) * 0.02
    w2 = jax.random.normal(k3, (d_mid, d_out)) * 0.02
    
    # JIT compile and run
    jitted_fn = jax.jit(mini_attention)
    
    # First call triggers compilation (and IR dump)
    result = jitted_fn(x, w1, w2)
    
    # Block until computation is done
    result.block_until_ready()
    
    print(f"Input shape:  {x.shape}")
    print(f"Output shape: {result.shape}")
    print(f"Output sample: {result[0, :5]}")
    print(f"\nDumps written to:")
    print(f"  HLO: {HLO_DUMP_PATH}")
    print(f"  LLO: {LLO_DUMP_PATH}")


if __name__ == "__main__":
    main()
    

The above Jax code is really just a compiled matmul + rms_norm + softmax + matmul. We jit the computation with jax.jit. Jax then traces this computation into HLO (High Level Operations) IR. HLO is heavily open source - HLO.

TPU Compiler Top Level View

At a high level the TPU compiler pipeline is Jax→HLO→LLO→VLIW bundles. Technically there are other stages, such as StableHLO, Jaxpr, TPU TLP etc., but for our purposes we can follow this diagram.

HLO IR

The frontend TPU compiler uses XLA as it’s primary infrastructure. XLA works off an IR called HLO (High Level Operations), which is an SSA graph IR.

There is some information publicly available on the frontend TPU HLO compiler and some open sourcing at OpenXLA.

This is similar to PyTorch FX IR.

HloModule jit_mini_attention, entry_computation_layout={(f32[16,64]{1,0:T(8,128)}, f32[64,64]{1,0:T(8,128)}, f32[64,32]{0,1:T(8,128)})->f32[16,32]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

%region_0.15 (Arg_0.12: f32[], Arg_1.13: f32[]) -> f32[] {
  %Arg_0.12 = f32[] parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum"}
  %Arg_1.13 = f32[] parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum"}
  ROOT %add.14 = f32[] add(%Arg_0.12, %Arg_1.13), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
}

%region_1.28 (Arg_0.25: f32[], Arg_1.26: f32[]) -> f32[] {
  %Arg_0.25 = f32[] parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max"}
  %Arg_1.26 = f32[] parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max"}
  ROOT %maximum.27 = f32[] maximum(%Arg_0.25, %Arg_1.26), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max" source_file="/home/ptoulme/tpu.py" source_line=48}
}

%region_2.39 (Arg_0.36: f32[], Arg_1.37: f32[]) -> f32[] {
  %Arg_0.36 = f32[] parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum"}
  %Arg_1.37 = f32[] parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum"}
  ROOT %add.38 = f32[] add(%Arg_0.36, %Arg_1.37), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=50}
}

ENTRY %main.47 (Arg_0.1: f32[16,64], Arg_1.2: f32[64,64], Arg_2.3: f32[64,32]) -> f32[16,32] {
  %Arg_0.1 = f32[16,64]{1,0} parameter(0), metadata={op_name="x"}
  %Arg_1.2 = f32[64,64]{1,0} parameter(1), metadata={op_name="w1"}
  %dot.10 = f32[16,64]{1,0} dot(%Arg_0.1, %Arg_1.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mini_attention)/jit(main)/matmul_1/dot_general" source_file="/home/ptoulme/tpu.py" source_line=35}
  %multiply.11 = f32[16,64]{1,0} multiply(%dot.10, %dot.10), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/integer_pow" source_file="/home/ptoulme/tpu.py" source_line=41}
  %constant.9 = f32[] constant(0)
  %reduce.16 = f32[16]{0} reduce(%multiply.11, %constant.9), dimensions={1}, to_apply=%region_0.15, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
  %reshape.17 = f32[16,1]{1,0} reshape(%reduce.16), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/broadcast_in_dim" source_file="/home/ptoulme/tpu.py" source_line=41}
  %constant.6 = f32[] constant(64)
  %broadcast.7 = f32[16,1]{1,0} broadcast(%constant.6), dimensions={}
  %divide.18 = f32[16,1]{1,0} divide(%reshape.17, %broadcast.7), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=41}
  %constant.4 = f32[] constant(1e-06)
  %broadcast.5 = f32[16,1]{1,0} broadcast(%constant.4), dimensions={}
  %add.19 = f32[16,1]{1,0} add(%divide.18, %broadcast.5), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/add" source_file="/home/ptoulme/tpu.py" source_line=41}
  %sqrt.20 = f32[16,1]{1,0} sqrt(%add.19), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/sqrt" source_file="/home/ptoulme/tpu.py" source_line=41}
  %broadcast.21 = f32[16,1]{1,0} broadcast(%sqrt.20), dimensions={0,1}, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %reshape.22 = f32[16]{0} reshape(%broadcast.21), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %broadcast.23 = f32[16,64]{1,0} broadcast(%reshape.22), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %divide.24 = f32[16,64]{1,0} divide(%dot.10, %broadcast.23), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %constant.8 = f32[] constant(-inf)
  %reduce.29 = f32[16]{0} reduce(%divide.24, %constant.8), dimensions={1}, to_apply=%region_1.28, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max" source_file="/home/ptoulme/tpu.py" source_line=48}
  %reshape.30 = f32[16,1]{1,0} reshape(%reduce.29), metadata={op_name="jit(mini_attention)/jit(main)/softmax/broadcast_in_dim" source_file="/home/ptoulme/tpu.py" source_line=48}
  %broadcast.31 = f32[16,1]{1,0} broadcast(%reshape.30), dimensions={0,1}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %reshape.32 = f32[16]{0} reshape(%broadcast.31), metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %broadcast.33 = f32[16,64]{1,0} broadcast(%reshape.32), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %subtract.34 = f32[16,64]{1,0} subtract(%divide.24, %broadcast.33), metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %exponential.35 = f32[16,64]{1,0} exponential(%subtract.34), metadata={op_name="jit(mini_attention)/jit(main)/softmax/exp" source_file="/home/ptoulme/tpu.py" source_line=49}
  %reduce.40 = f32[16]{0} reduce(%exponential.35, %constant.9), dimensions={1}, to_apply=%region_2.39, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=50}
  %reshape.41 = f32[16,1]{1,0} reshape(%reduce.40), metadata={op_name="jit(mini_attention)/jit(main)/softmax/broadcast_in_dim" source_file="/home/ptoulme/tpu.py" source_line=50}
  %broadcast.42 = f32[16,1]{1,0} broadcast(%reshape.41), dimensions={0,1}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
  %reshape.43 = f32[16]{0} reshape(%broadcast.42), metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
  %broadcast.44 = f32[16,64]{1,0} broadcast(%reshape.43), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
  %divide.45 = f32[16,64]{1,0} divide(%exponential.35, %broadcast.44), metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
  %Arg_2.3 = f32[64,32]{1,0} parameter(2), metadata={op_name="w2"}
  ROOT %dot.46 = f32[16,32]{1,0} dot(%divide.45, %Arg_2.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mini_attention)/jit(main)/matmul_2/dot_general" source_file="/home/ptoulme/tpu.py" source_line=56}
}

The above HLO is the output of the Jax tracer after converting from StableHLO to HLO. No optimization passes have been performed yet.

We can clearly see - instructions, shapes, dtypes and metadata.

HLO Optimization Passes

The XLA compiler ran 71 optimization passes on this module. Here are the key transformations for our small toy program.

Algebraic Simplifier

The algebraic simplifier runs early in the pipeline (pass #6). The Algebraic Simplifier is open source at OpenXLA - algebraic_simplifier.cc The pass is nearly 10k lines of CPP code. One optimization is converting division by constants to multiplication:

// Before (pass #5)
%constant.6 = f32[] constant(64)
%broadcast.7 = f32[16,1]{1,0} broadcast(%constant.6), dimensions={}
%divide.18 = f32[16,1]{1,0} divide(%reshape.17, %broadcast.7)

// After (pass #6)
%constant = f32[] constant(0.015625)
%broadcast = f32[16,1]{1,0} broadcast(%constant), dimensions={}
%multiply = f32[16,1]{1,0} multiply(%reshape.17, %broadcast)

Division is expensive on most hardware. The compiler precomputes 1/64 = 0.015625 and converts the divide to a multiply. We see this in our toy program.

Layout Assignment

Layout assignment and tiling (pass #25 → #29) adds TPU tile annotations to all shapes.

Layout assignment is partially open source at OpenXLA - layout_assignment.cc

// Before
%dot.10 = f32[16,64]{1,0} dot(...)
%reduce.16 = f32[16]{0} reduce(...)

// After
%dot.10 = f32[16,64]{1,0:T(8,128)} dot(...)
%reduce.16 = f32[16]{0:T(128)} reduce(...)

The {1,0} specifies row-major layout (the last dimension is contiguous in memory). The :T(8,128) is TPU-specific tiling — the tensor is divided into tiles of 8 rows × 128 columns, matching the VPU's 8 sublanes × 128 lanes.

Fusion

The fusion pass (pass #38-39) combines independent HLO operations into fused kernels.

The TPU HLO fusion pass is also open source at OpenXLA - instruction_fusion.cc

ENTRY %main.47 (Arg_0.1: f32[16,64], Arg_1.2: f32[64,64], Arg_2.3: f32[64,32]) -> f32[16,32] {
  %Arg_0.1 = f32[16,64]{1,0:T(8,128)} parameter(0), metadata={op_name="x"}
  %Arg_1.2 = f32[64,64]{1,0:T(8,128)} parameter(1), metadata={op_name="w1"}
  %Arg_2.3 = f32[64,32]{0,1:T(8,128)} parameter(2), metadata={op_name="w2"}
  %convolution = f32[16,64]{1,0:T(8,128)} convolution(%Arg_0.1, %Arg_1.2), dim_labels=bf_io->bf
  %multiply_reduce_fusion = f32[16]{0:T(128)} fusion(%convolution), kind=kLoop, calls=%fused_computation.3
  %add_sqrt_fusion = f32[16]{0:T(128)} fusion(%multiply_reduce_fusion), kind=kLoop, calls=%fused_computation.9
  %fusion.5 = f32[16]{0:T(128)} fusion(%convolution, %add_sqrt_fusion), kind=kLoop, calls=%fused_computation.8
  %fusion.2 = f32[16]{0:T(128)} fusion(%fusion.5, %convolution, %add_sqrt_fusion), kind=kLoop, calls=%fused_computation.4
  ROOT %fusion = f32[16,32]{1,0:T(8,128)} fusion(%Arg_2.3, %fusion.2, %fusion.5, %convolution, %add_sqrt_fusion), kind=kOutput
}

The original ~25 HLO ops collapse into 6 operations. Each fusion wraps multiple operations into a single kernel. The kind=kLoop indicates elementwise fusion; kind=kOutput indicates producer-consumer fusion (the producer’s output is immediately consumed).

Note that dot became convolution — the TPU compiler canonicalizes matrix multiplications this way, likely because both operations map to the same MXU hardware.

Each fusion has a body. The add_sqrt_fusion contains:

%fused_computation.9 (param_0.29: f32[16]) -> f32[16] {
  %param_0.29 = f32[16]{0:T(128)} parameter(0)
  %constant.12 = f32[] constant(0.015625)
  %broadcast.37 = f32[16]{0:T(128)} broadcast(%constant.12), dimensions={}
  %multiply.7 = f32[16]{0:T(128)} multiply(%param_0.29, %broadcast.37)
  %constant.16 = f32[] constant(1e-06)
  %broadcast.36 = f32[16]{0:T(128)} broadcast(%constant.16), dimensions={}
  %add.5 = f32[16]{0:T(128)} add(%multiply.7, %broadcast.36)
  ROOT %sqrt.5 = f32[16]{0:T(128)} sqrt(%add.5)
}

This is the RMS norm denominator: sqrt(mean + epsilon).

Multi-Output Fusion (pass #43)

When multiple fusions share a common operand, they can be merged. Multi-output fusion (pass #43) identifies sibling fusions with shared inputs and combines them, returning a tuple:

%fused_computation.3 (param_0.33: f32[16,64], param_1.35: f32[64,64]) -> (f32[16], f32[16,64]) {
  %convolution.3 = f32[16,64]{1,0:T(8,128)} convolution(%param_0.33, %param_1.35), dim_labels=bf_io->bf
  %multiply.6 = f32[16,64]{1,0:T(8,128)} multiply(%convolution.3, %convolution.3)
  %constant.14 = f32[] constant(0)
  %reduce.0 = f32[16]{0:T(128)} reduce(%multiply.6, %constant.14), dimensions={1}, to_apply=%region_0.15
  ROOT %tuple = (f32[16]{0:T(128)}, f32[16,64]{1,0:T(8,128)}) tuple(%reduce.0, %convolution.3)
}

Before this pass, the matmul feeding into both the normalization path and the RMS reduction path would have been computed in separate fusions. Multi-output fusion merges them — the matmul happens once, and both the normalized result (convolution.3) and the squared sum (reduce.0) are returned together. The alternative would be either recomputing the matmul or spilling it to HBM between fusions.

The caller extracts both values:

%multiply_reduce_fusion = (f32[16]{0:T(128)}, f32[16,64]{1,0:T(8,128)}) fusion(%Arg_0.1, %Arg_1.2), kind=kOutput
%get-tuple-element = f32[16]{0:T(128)} get-tuple-element(%multiply_reduce_fusion), index=0
%get-tuple-element.1 = f32[16,64]{1,0:T(8,128)} get-tuple-element(%multiply_reduce_fusion), index=1

Memory Space Assignment & Async Scheduling

The final passes assign tensors to specific memory spaces and insert async memory operations. In HLO, memory spaces are annotated as S(n):

• S(0) (often omitted) — HBM (high bandwidth memory, off-chip)

• S(1) — VMEM (on-chip SRAM)

• S(2), S(3), etc. — additional device-specific memory spaces

The after_codegen.txt shows the scheduled program:

ENTRY %main.47 (Arg_0.1: f32[16,64], Arg_1.2: f32[64,64], Arg_2.3: f32[64,32]) -> f32[16,32] {
  %Arg_1.2 = f32[64,64]{1,0:T(8,128)} parameter(1), metadata={op_name="w1"}
  %copy-start = (f32[64,64]{1,0:T(8,128)S(1)}, f32[64,64]{1,0:T(8,128)}, u32[]{:S(2)}) copy-start(%Arg_1.2), cross_program_prefetch_index=0
  %Arg_2.3 = f32[64,32]{0,1:T(8,128)} parameter(2), metadata={op_name="w2"}
  %Arg_0.1 = f32[16,64]{1,0:T(8,128)} parameter(0), metadata={op_name="x"}
  %copy-done = f32[64,64]{1,0:T(8,128)S(1)} copy-done(%copy-start)
  %multiply_reduce_fusion = (f32[16]{0:T(128)S(1)}, f32[16,64]{1,0:T(8,128)S(1)}) fusion(%Arg_0.1, %copy-done), kind=kOutput, ...
  %copy-start.1 = (f32[64,32]{0,1:T(8,128)S(1)}, f32[64,32]{0,1:T(8,128)}, u32[]{:S(2)}) copy-start(%Arg_2.3)
  %get-tuple-element.1 = f32[16,64]{1,0:T(8,128)S(1)} get-tuple-element(%multiply_reduce_fusion), index=1
  %get-tuple-element = f32[16]{0:T(128)S(1)} get-tuple-element(%multiply_reduce_fusion), index=0
  %add_sqrt_fusion = f32[16]{0:T(128)S(1)} fusion(%get-tuple-element), kind=kLoop, ...
  %fusion.5 = f32[16]{0:T(128)S(1)} fusion(%get-tuple-element.1, %add_sqrt_fusion), kind=kLoop, ...
  %fusion.2 = f32[16]{0:T(128)S(1)} fusion(%fusion.5, %get-tuple-element.1, %add_sqrt_fusion), kind=kLoop, ...
  %copy-done.1 = f32[64,32]{0,1:T(8,128)S(1)} copy-done(%copy-start.1)
  ROOT %fusion = f32[16,32]{1,0:T(8,128)} fusion(%copy-done.1, %fusion.2, %fusion.5, %get-tuple-element.1, %add_sqrt_fusion), kind=kOutput, ...
}

Notice the parameters start without S(1) — they live in HBM. The copy-start/copy-done pairs are async DMA operations that move data to VMEM. The tuple returned by copy-start contains the destination buffer (in S(1)), the source reference, and a sync token (in S(2)).

The compiler overlaps memory transfers with computation:

1. copy-start(w1) — initiate DMA of w1 from HBM → VMEM

2. copy-done(w1) — wait for transfer, then use in matmul₁

3. copy-start(w2) — initiate DMA of w2 while RMS norm and softmax execute

4. copy-done(w2) — wait for transfer, then use in matmul₂

The backend_config also shows estimated cycle counts per fusion:

"estimated_cycles":"2248"  // multiply_reduce_fusion
"estimated_cycles":"2120"  // add_sqrt_fusion  
"estimated_cycles":"2143"  // fusion.5 (reduce_max)
"estimated_cycles":"2162"  // fusion.2 (reduce_sum)
"estimated_cycles":"3140"  // final fusion (matmul_2)

Below is the final HLO that the frontend TPU compiler emits.

HloModule jit_mini_attention, is_scheduled=true, entry_computation_layout={(f32[16,64]{1,0:T(8,128)}, f32[64,64]{1,0:T(8,128)}, f32[64,32]{0,1:T(8,128)})->f32[16,32]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation.1 (param_0.21: f32[16], param_1.23: f32[16], param_2.10: f32[16,64], param_3.5: f32[16]) -> f32[16,64] {
  %param_2.10 = f32[16,64]{1,0:T(8,128)S(1)} parameter(2)
  %param_3.5 = f32[16]{0:T(128)S(1)} parameter(3)
  %broadcast.29 = f32[16,64]{1,0:T(8,128)} broadcast(%param_3.5), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  %divide.5 = f32[16,64]{1,0:T(8,128)} divide(%param_2.10, %broadcast.29), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %param_1.23 = f32[16]{0:T(128)S(1)} parameter(1)
  %broadcast.24 = f32[16,64]{1,0:T(8,128)} broadcast(%param_1.23), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  %subtract.3 = f32[16,64]{1,0:T(8,128)} subtract(%divide.5, %broadcast.24), metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %exponential.3 = f32[16,64]{1,0:T(8,128)} exponential(%subtract.3), metadata={op_name="jit(mini_attention)/jit(main)/softmax/exp" source_file="/home/ptoulme/tpu.py" source_line=49}
  %param_0.21 = f32[16]{0:T(128)S(1)} parameter(0)
  %broadcast.18 = f32[16,64]{1,0:T(8,128)} broadcast(%param_0.21), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  ROOT %divide.1 = f32[16,64]{1,0:T(8,128)} divide(%exponential.3, %broadcast.18), metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
}

%bitcast_fusion.1 (bitcast_input.1: f32[64,32]) -> f32[64,32] {
  %bitcast_input.1 = f32[64,32]{0,1:T(8,128)S(1)} parameter(0)
  ROOT %bitcast.1 = f32[64,32]{0,1:T(8,128)} bitcast(%bitcast_input.1)
}

%fused_computation (param_0.1: f32[64,32], param_1.21: f32[16], param_2.9: f32[16], param_3.3: f32[16,64], param_4: f32[16]) -> f32[16,32] {
  %param_1.21 = f32[16]{0:T(128)S(1)} parameter(1)
  %param_2.9 = f32[16]{0:T(128)S(1)} parameter(2)
  %param_3.3 = f32[16,64]{1,0:T(8,128)S(1)} parameter(3)
  %param_4 = f32[16]{0:T(128)S(1)} parameter(4)
  %fusion.1 = f32[16,64]{1,0:T(8,128)} fusion(%param_1.21, %param_2.9, %param_3.3, %param_4), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(mini_attention)/jit(main)/softmax/div" source_file="/home/ptoulme/tpu.py" source_line=50}
  %param_0.1 = f32[64,32]{0,1:T(8,128)S(1)} parameter(0)
  %fusion.7 = f32[64,32]{0,1:T(8,128)} fusion(%param_0.1), kind=kLoop, calls=%bitcast_fusion.1
  ROOT %convolution.2 = f32[16,32]{1,0:T(8,128)} convolution(%fusion.1, %fusion.7), dim_labels=bf_io->bf, metadata={op_name="jit(mini_attention)/jit(main)/matmul_2/dot_general" source_file="/home/ptoulme/tpu.py" source_line=56}
}

%region_0.15 (Arg_0.12: f32[], Arg_1.13: f32[]) -> f32[] {
  %Arg_1.13 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum"}
  %Arg_0.12 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum"}
  ROOT %add.14 = f32[]{:T(128)} add(%Arg_0.12, %Arg_1.13), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
}

%bitcast_fusion (bitcast_input: f32[16,64]) -> f32[16,64] {
  %bitcast_input = f32[16,64]{1,0:T(8,128)} parameter(0)
  ROOT %bitcast = f32[16,64]{1,0:T(8,128)} bitcast(%bitcast_input)
}

%bitcast_fusion.2 (bitcast_input.2: f32[64,64]) -> f32[64,64] {
  %bitcast_input.2 = f32[64,64]{1,0:T(8,128)S(1)} parameter(0)
  ROOT %bitcast.2 = f32[64,64]{1,0:T(8,128)} bitcast(%bitcast_input.2)
}

%fused_computation.3 (param_0.33: f32[16,64], param_1.35: f32[64,64]) -> (f32[16], f32[16,64]) {
  %param_0.33 = f32[16,64]{1,0:T(8,128)} parameter(0)
  %fusion.6 = f32[16,64]{1,0:T(8,128)} fusion(%param_0.33), kind=kLoop, calls=%bitcast_fusion
  %param_1.35 = f32[64,64]{1,0:T(8,128)S(1)} parameter(1)
  %fusion.8 = f32[64,64]{1,0:T(8,128)} fusion(%param_1.35), kind=kLoop, calls=%bitcast_fusion.2
  %convolution.3 = f32[16,64]{1,0:T(8,128)S(1)} convolution(%fusion.6, %fusion.8), dim_labels=bf_io->bf, metadata={op_name="jit(mini_attention)/jit(main)/matmul_1/dot_general" source_file="/home/ptoulme/tpu.py" source_line=35}
  %multiply.6 = f32[16,64]{1,0:T(8,128)} multiply(%convolution.3, %convolution.3), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/integer_pow" source_file="/home/ptoulme/tpu.py" source_line=41}
  %constant.14 = f32[]{:T(128)} constant(0)
  %reduce.0 = f32[16]{0:T(128)S(1)} reduce(%multiply.6, %constant.14), dimensions={1}, to_apply=%region_0.15, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
  ROOT %tuple = (f32[16]{0:T(128)S(1)}, f32[16,64]{1,0:T(8,128)S(1)}) tuple(%reduce.0, %convolution.3)
}

%region_2.39 (Arg_0.36: f32[], Arg_1.37: f32[]) -> f32[] {
  %Arg_1.37 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum"}
  %Arg_0.36 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum"}
  ROOT %add.38 = f32[]{:T(128)} add(%Arg_0.36, %Arg_1.37), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=50}
}

%fused_computation.4 (param_0.30: f32[16], param_1.33: f32[16,64], param_2.17: f32[16]) -> f32[16] {
  %param_1.33 = f32[16,64]{1,0:T(8,128)S(1)} parameter(1)
  %param_2.17 = f32[16]{0:T(128)S(1)} parameter(2)
  %broadcast.32 = f32[16,64]{1,0:T(8,128)} broadcast(%param_2.17), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  %divide.7 = f32[16,64]{1,0:T(8,128)} divide(%param_1.33, %broadcast.32), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %param_0.30 = f32[16]{0:T(128)S(1)} parameter(0)
  %broadcast.25 = f32[16,64]{1,0:T(8,128)} broadcast(%param_0.30), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  %subtract.5 = f32[16,64]{1,0:T(8,128)} subtract(%divide.7, %broadcast.25), metadata={op_name="jit(mini_attention)/jit(main)/softmax/sub" source_file="/home/ptoulme/tpu.py" source_line=49}
  %exponential.5 = f32[16,64]{1,0:T(8,128)} exponential(%subtract.5), metadata={op_name="jit(mini_attention)/jit(main)/softmax/exp" source_file="/home/ptoulme/tpu.py" source_line=49}
  %constant.13 = f32[]{:T(128)} constant(0)
  ROOT %reduce.1 = f32[16]{0:T(128)S(1)} reduce(%exponential.5, %constant.13), dimensions={1}, to_apply=%region_2.39, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=50}
}

%region_1.28 (Arg_0.25: f32[], Arg_1.26: f32[]) -> f32[] {
  %Arg_1.26 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max"}
  %Arg_0.25 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max"}
  ROOT %maximum.27 = f32[]{:T(128)} maximum(%Arg_0.25, %Arg_1.26), metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max" source_file="/home/ptoulme/tpu.py" source_line=48}
}

%fused_computation.8 (param_0.32: f32[16,64], param_1.34: f32[16]) -> f32[16] {
  %param_0.32 = f32[16,64]{1,0:T(8,128)S(1)} parameter(0)
  %param_1.34 = f32[16]{0:T(128)S(1)} parameter(1)
  %broadcast.35 = f32[16,64]{1,0:T(8,128)} broadcast(%param_1.34), dimensions={0}, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]}}
  %divide.9 = f32[16,64]{1,0:T(8,128)} divide(%param_0.32, %broadcast.35), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=42}
  %constant.15 = f32[]{:T(128)} constant(-inf)
  ROOT %reduce.2 = f32[16]{0:T(128)S(1)} reduce(%divide.9, %constant.15), dimensions={1}, to_apply=%region_1.28, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max" source_file="/home/ptoulme/tpu.py" source_line=48}
}

%fused_computation.9 (param_0.29: f32[16]) -> f32[16] {
  %param_0.29 = f32[16]{0:T(128)S(1)} parameter(0)
  %constant.12 = f32[]{:T(128)} constant(0.015625)
  %broadcast.37 = f32[16]{0:T(128)} broadcast(%constant.12), dimensions={}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":[]}}
  %multiply.7 = f32[16]{0:T(128)} multiply(%param_0.29, %broadcast.37), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/div" source_file="/home/ptoulme/tpu.py" source_line=41}
  %constant.16 = f32[]{:T(128)} constant(1e-06)
  %broadcast.36 = f32[16]{0:T(128)} broadcast(%constant.16), dimensions={}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":[]}}
  %add.5 = f32[16]{0:T(128)} add(%multiply.7, %broadcast.36), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/add" source_file="/home/ptoulme/tpu.py" source_line=41}
  ROOT %sqrt.5 = f32[16]{0:T(128)S(1)} sqrt(%add.5), metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/sqrt" source_file="/home/ptoulme/tpu.py" source_line=41}
}

ENTRY %main.47 (Arg_0.1: f32[16,64], Arg_1.2: f32[64,64], Arg_2.3: f32[64,32]) -> f32[16,32] {
  %Arg_1.2 = f32[64,64]{1,0:T(8,128)} parameter(1), metadata={op_name="w1"}
  %copy-start = (f32[64,64]{1,0:T(8,128)S(1)}, f32[64,64]{1,0:T(8,128)}, u32[]{:S(2)}) copy-start(%Arg_1.2), cross_program_prefetch_index=0
  %Arg_2.3 = f32[64,32]{0,1:T(8,128)} parameter(2), metadata={op_name="w2"}
  %Arg_0.1 = f32[16,64]{1,0:T(8,128)} parameter(0), metadata={op_name="x"}
  %copy-done = f32[64,64]{1,0:T(8,128)S(1)} copy-done(%copy-start)
  %multiply_reduce_fusion = (f32[16]{0:T(128)S(1)}, f32[16,64]{1,0:T(8,128)S(1)}) fusion(%Arg_0.1, %copy-done), kind=kOutput, calls=%fused_computation.3, metadata={op_name="jit(mini_attention)/jit(main)/matmul_1/dot_general" source_file="/home/ptoulme/tpu.py" source_line=35}
  %copy-start.1 = (f32[64,32]{0,1:T(8,128)S(1)}, f32[64,32]{0,1:T(8,128)}, u32[]{:S(2)}) copy-start(%Arg_2.3)
  %get-tuple-element.1 = f32[16,64]{1,0:T(8,128)S(1)} get-tuple-element(%multiply_reduce_fusion), index=1, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
  %get-tuple-element = f32[16]{0:T(128)S(1)} get-tuple-element(%multiply_reduce_fusion), index=0, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=41}
  %add_sqrt_fusion = f32[16]{0:T(128)S(1)} fusion(%get-tuple-element), kind=kLoop, calls=%fused_computation.9, metadata={op_name="jit(mini_attention)/jit(main)/rms_norm/sqrt" source_file="/home/ptoulme/tpu.py" source_line=41}
  %fusion.5 = f32[16]{0:T(128)S(1)} fusion(%get-tuple-element.1, %add_sqrt_fusion), kind=kLoop, calls=%fused_computation.8, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_max" source_file="/home/ptoulme/tpu.py" source_line=48}
  %fusion.2 = f32[16]{0:T(128)S(1)} fusion(%fusion.5, %get-tuple-element.1, %add_sqrt_fusion), kind=kLoop, calls=%fused_computation.4, metadata={op_name="jit(mini_attention)/jit(main)/softmax/reduce_sum" source_file="/home/ptoulme/tpu.py" source_line=50}
  %copy-done.1 = f32[64,32]{0,1:T(8,128)S(1)} copy-done(%copy-start.1)
  ROOT %fusion = f32[16,32]{1,0:T(8,128)} fusion(%copy-done.1, %fusion.2, %fusion.5, %get-tuple-element.1, %add_sqrt_fusion), kind=kOutput, calls=%fused_computation, metadata={op_name="jit(mini_attention)/jit(main)/matmul_2/dot_general" source_file="/home/ptoulme/tpu.py" source_line=56}
}

That covers the HLO optimization passes. At this point we have fused, scheduled operations with memory annotations — but it's still relatively hardware-agnostic. Next, the TPU backend translates each fusion into LLO (Low Level Operators), a representation that maps directly to physical hardware units: the MXU for matmuls, VPU for elementwise ops, XLU for transposes, and DMA engines for memory movement.

LLO IR

After HLO optimizations, XLA’s TPU backend translates each fusion into LLO (Low Level Operators)—a TPU-specific intermediate representation that maps directly to hardware units. Each fusion becomes a separate LLO program that goes through in my case 78 optimization passes before producing final VLIW bundles.

TPU Hardware Units

Before diving into LLO, here’s a quick orientation. Each TPU Trillium V6e TensorCore contains:

• MXU (Matrix Unit): Two 256×256 systolic arrays for matrix multiplies. This is where the FLOPS come from.

• VPU (Vector Processing Unit): Handles elementwise ops (add, mul, exp, etc.) across 8 sublanes × 128 lanes.

• XLU (Transpose Unit): Cross-lane shuffles, transposes, and permutations.

• Scalar Unit: Scalar operations, address calculation, and control flow.

• DMA engines: Async memory transfers between HBM and VMEM.

The compiler’s job is to keep all of these busy simultaneously. A well-scheduled TPU program overlaps MXU matmuls with VPU elementwise ops with DMA transfers — all in the same VLIW bundle.

Deep Dive: Compiling multiply_reduce_fusion

Let’s trace how multiply_reduce_fusion—which computes our first matrix multiply and the squared values for RMS normalization—transforms from initial LLO to final machine code.

Initial LLO (Pass 02)

The compiler first translates the fusion HLO region into a verbose, unscheduled LLO representation. Here’s how the matrix multiply setup begins:

$region0: #{multiply_reduce_fusion}
  #allocation0 [shape = 'f32[1024]{0}', space=vmem, size = 0x1000, tag = 'scoped memory']
  #allocation1 [shape = 'f32[16]{0:T(1024)S(1)}', space=vmem, size = 0x1000, tag = 'reduce buffer']
  
  %s0 = inlined_call_operand.hbm [shape: f32[16,64], index: 0, kind: input]
  %s1 = inlined_call_operand.vmem [shape: f32[64,64], index: 1, kind: input]
  %s2 = inlined_call_operand.vmem [shape: f32[16], index: 2, kind: output]
  %s3 = inlined_call_operand.vmem [shape: f32[16,64], index: 3, kind: output]

The compiler allocates VMEM buffers and establishes that operand 0 comes from HBM (High Bandwidth Memory) while operand 1 is already in VMEM. This is a multi-output fusion—it produces both the matmul result (%s3) and the sum of squares (%s2) for RMS norm.

The initial MXU operations use vmatpush to load weight columns into the systolic array:

%v63 = vld [vmem:[%s55] sm:$0xff]
%64 = vmatpush.bf16.msra.mxu0 %v63
...
%v93 = vld [vmem:[#allocation0] sm:$0xff]
%94 = vmatmul.bf16.gmra.mxu0 %v93
%v95 = vpop.f32.mrf.mxu0

This sequence streams weight tiles through the MXU in bf16 precision, then pops the accumulated f32 results.

The Reduction Pattern

After the matmul, we need to sum squared values across columns. The compiler generates a cross-lane reduction using the XLU (transpose unit):

%141 = vxpose.xlu0.b32.start [1/2] (short) /*vx=*/%v138, /*width=*/128
%142 = vxpose.xlu0.b32.end [2/2] (short) /*vx=*/%v140, /*width=*/128
%v143 = vpop.trf.xlu0
%v144 = vpop.trf.xlu0
...
%v158 = vpop.trf.xlu0  // 16 pops total

After transposing, a tree reduction sums the 16 lanes:

%v161 = vadd.f32 0.0, %v143
%v165 = vadd.f32 %v161, %v144
%v169 = vadd.f32 %v165, %v145
...
%v221 = vadd.f32 %v217, %v158

Then a sublane rotation pattern completes the reduction:

%v223 = vrot.slane %v221, 4   // rotate by 4
%v226 = vadd.f32 %v221, %v223
%v228 = vrot.slane %v226, 2   // rotate by 2
%v231 = vadd.f32 %v226, %v228
%v233 = vrot.slane %v231, 1   // rotate by 1
%v236 = vadd.f32 %v231, %v233

This classic parallel reduction pattern uses log₂(n) steps—rotate by 4, 2, then 1—to sum 8 sublanes into a single scalar.

Final VLIW Bundles (Pass 78)

After all optimization passes, the compiler produces 71 tightly-packed VLIW bundles. Each bundle groups independent operations that execute in parallel across hardware units:

This is VLIW (Very Long Instruction Word) execution — the compiler statically packs independent operations into fixed-width bundles, and the hardware executes everything in a bundle simultaneously without runtime dependency checking.

0x9 : { %22 = dma.hbm_to_vmem [thread:$0] /*hbm=*/%s359_s0, ...
       %v28_v1 = vlaneseq  
       %v301_v2 = vmov 0.0
       %vm302_vm0 = vmmov 0
       %v247_v4 = vld [vmem:[%s360_s1 + $0x30] sm:$0xff]
       %v248_v5 = vld [vmem:[%s360_s1 + $0x38] sm:$0xff]
       %v249_v6 = vld [vmem:[%s360_s1 + $0x20] sm:$0xff] }

This single bundle executes 7 operations simultaneously: one DMA transfer from HBM, one lane sequence generation, two constant materializations, and three VMEM loads. The compiler has carefully scheduled these to avoid resource conflicts.

The MXU operations use the optimized vmatpush3 instruction:

0xc : { %261 = vmatpush3.bf16.msra.mxu0 %v61_v7
        %v89_v15 = vpack.c.bf16 %v253_v14, %v85_v13
        %297 = dma.done.wait [#allocation4], 256 }

Three operations in parallel: push a weight tile, pack bf16 values for the next push, and synchronize the DMA. The bf16 suffix indicates this matmul uses bfloat16 precision on the MXU, which provides 2x throughput compared to f32.

The matmul result extraction uses a masked variant:

0x14 : { %269 = vmatmul.mubr.msk.bf16.vlgmr.msra.gmra.mxu0 %vm257_vm2, %v258_v18 }
0x15 : { %v95_v19 = vpop.f32.mrf.mxu0 }
0x16 : { %v98_v20 = vmul.f32 %v95_v19, %v95_v19
         %105 = vst [vmem:[%s362_s3] sm:$0xff] /*vst_source=*/%v95_v19 }

The squaring happens immediately in bundle 0x16, overlapped with storing the matmul result—the multi-output fusion lets us reuse %v95_v19 for both outputs.

In the above video, we can visualize this fusion on a TPU V6e Trillium. Note we are using one MXU in this visualization, while a Trillium has 2 256x256 MXUs.

LLO Compilation Overview: All Fusions

Our mini_attention function compiles into five fusions, each with its own LLO compilation producing different bundle counts and hardware utilization patterns.

multiply_reduce_fusion (71 bundles)

This is our largest fusion, implementing both matmul_1 (the first matrix multiply) and the sum-of-squares reduction for RMS normalization. The fusion demonstrates sophisticated hardware utilization:

• MXU0 handles the 16×64 by 64×64 bf16 matmul

• VPU computes the squared values immediately after MXU results

• XLU0 performs the cross-lane transpose for reduction

• DMA streams the input activations from HBM

The 71 bundles break down into: DMA setup, weight loading, matmul execution, squaring, transpose-based reduction, and sublane rotation reduction.

add_sqrt_fusion (10 bundles)

The smallest and simplest fusion computes sqrt(mean + epsilon) for RMS normalization:

0x1 : { %v2_v0 = vld [vmem:[%s37_s0] sm:$0x1] }
0x2 : { %v5_v1 = vmul.f32 0.015625, %v2_v0 }
0x3 : { %v9_v2 = vadd.f32 1e-06, %v5_v1 }
0x4 : { %19 = vrsqrt.f32 %v9_v2
        %vm13_vm0 = vcmp.eq.f32.partialorder %v9_v2, inf
        %v16_v4 = vand.u32 2147483648, %v9_v2
        %vm15_vm1 = vcmp.eq.f32.partialorder %v9_v2, 0.0 }
0x5 : { %v20_v3 = vpop.eup %19 }
0x6 : { %v12_v5 = vmul.f32 %v20_v3, %v9_v2 }
0x7 : { %v14_v6 = vsel /*vm=*/%vm13_vm0, /*on_true_vy=*/%v9_v2, /*on_false_vx=*/%v12_v5 }
0x8 : { %v17_v7 = vsel /*vm=*/%vm15_vm1, /*on_true_vy=*/%v16_v4, /*on_false_vx=*/%v14_v6 }
0x9 : { %18 = vst [vmem:[%s38_s1] sm:$0x1] /*vst_source=*/%v17_v7 }

The constant 0.015625 is 1/64—the mean computation. The vrsqrt.f32 instruction computes reciprocal square root. Special case handling for infinity and zero ensures numerical correctness.

fusion.5 (56 bundles)

This fusion computes reduce_max across rows for the softmax numerically-stable computation. It follows a similar pattern to the sum reduction in multiply_reduce_fusion: XLU transpose followed by tree reduction, but using max operations instead of add.

fusion.2 (65 bundles)

Implements exp(x - max) + reduce_sum—the softmax numerator and denominator. The vpow2 instruction computes the exponential followed by another transpose-and-reduce pattern for the sum.

fusion (48 bundles)

The final matmul (matmul_2) multiplying the softmax output by the value matrix. This fusion is smaller than multiply_reduce_fusion because it’s a pure matmul without the additional reduction operations.

TLP: The Top Level Program

The TLP (Top Level Program) orchestrates the entire mini_attention execution. Looking at the correct TLP for our function, we can see how it coordinates all five fusions and manages data movement.

Program Structure Overview

The mini_attention TLP is 174 bundles (0x00-0xad) and follows this execution flow:

0x5f-0x65 : Program entry and mode checking
0x66-0x7d : DMA: Load weight matrix from HBM (copy-start)
0x7e-0x82 : Call multiply_reduce_fusion (matmul_1 + squares)
0x83-0x8d : DMA: Start loading value matrix (copy-start.1)
0x8e-0x91 : Call add_sqrt_fusion (sqrt for RMS norm)
0x92-0x96 : Call fusion.5 (reduce_max for softmax)
0x97-0x9b : Call fusion.2 (exp + reduce_sum for softmax)
0x9c-0xa0 : DMA: Wait for value matrix
0xa1-0xa6 : Call fusion (matmul_2)
0xa7-0xad : Program finalization

DMA Orchestration and Overlapping

The TLP carefully overlaps memory transfers with computation. First, it loads the weight matrix for the first matmul:

0x77 : { %26 = dma.hbm_to_vmem [thread:$1] /*hbm=*/%s10_s14, 
             /*size_in_granules=*/%s22_s17, /*vmem=*/%s24_s19, 
             /*dst_syncflagno=*/[#allocation8] }

The DMA uses thread:$1—the TPU has multiple DMA engines, allowing concurrent transfers. After initiating this transfer, the TLP immediately waits for completion and then calls the first fusion:

0x7a : { %293 = dma.done.wait [#allocation8], 1024 }
0x7b : { %294 = vsyncadd [#allocation8], 4294966272 }
0x7c : { %39 = vsyncpa [#allocation8], 1 }
0x81 : { %45 = inlined_call %s9_s0, %s299_s1, %s300_s2, %s301_s3 
              /* %multiply_reduce_fusion = fusion(%Arg_0.1, %copy-done) */ }

Overlapping DMA with Compute

While RMS normalization fusions run, the TLP starts loading the value matrix for the second matmul:

0x8c : { %56 = dma.hbm_to_vmem [thread:$1] /*hbm=*/%s354_s25, 
             /*size_in_granules=*/512, /*vmem=*/%s54_s23, 
             /*dst_syncflagno=*/[#allocation13] }

This DMA runs in the background while add_sqrt_fusion, fusion.5, and fusion.2 execute:

0x90 : { %66 = inlined_call ... /* %add_sqrt_fusion */ }
0x95 : { %71 = inlined_call ... /* %fusion.5 */ }
0x9a : { %76 = inlined_call ... /* %fusion.2 */ }

Only before the final matmul does the TLP wait for the value matrix:

0x9d : { %295 = dma.done.wait [#allocation13], 512 }
0xa5 : { %87 = inlined_call ... /* %fusion = final matmul */ }

This is double-buffering in action—the softmax computation hides the latency of loading the value matrix.

The above video visualizes this DMA and compute overlap on TPU.

Fusion Calling Convention

Each fusion call passes VMEM buffer addresses via inlined_call. For example, the final matmul:

0xa2 : { %s312_s0 = smov [#allocation12] /* materialized constant */
        %s355_s5 = sld [smem:[#allocation25_spill]]
        %s356_s5 = int_to_ptr.hbm [resolvable:$false] %s355_s5 }
0xa3 : { %s313_s1 = smov [#allocation16] /* materialized constant */
        %s314_s2 = smov [#allocation15] /* materialized constant */ }
0xa4 : { %s315_s3 = smov [#allocation11] /* materialized constant */
        %s316_s4 = smov [#allocation14] /* materialized constant */ }
0xa5 : { %87 = inlined_call %s312_s0, %s313_s1, %s314_s2, %s315_s3, %s316_s4, %s356_s5
              /* %fusion = fusion(%copy-done.1, %fusion.2, %fusion.5, 
                                   %get-tuple-element.1, %add_sqrt_fusion) */ }

The fusion receives six operands: the value matrix from HBM (copy-done.1), plus intermediate results from all previous fusions. Note the #allocation25_spill—this is a register spill to scalar memory, indicating the compiler ran out of scalar registers and had to temporarily store a value.

Trace Points for Profiling

Throughout execution, the TLP emits trace markers:

0x78 : { %30 = vtrace 2415919104 }  // After copy-start
0x82 : { %49 = vtrace 2415919106 }  // After multiply_reduce_fusion
0x91 : { %69 = vtrace 2415919108 }  // After add_sqrt_fusion
0x96 : { %74 = vtrace 2415919109 }  // After fusion.5
0x9b : { %79 = vtrace 2415919110 }  // After fusion.2
0xa6 : { %90 = vtrace 2415919112 }  // After final fusion

These trace IDs let TPU profiling tools measure the time spent in each fusion, helping identify performance bottlenecks.

Program Finalization

The TLP concludes with synchronization:

0xa9 : { %95 = vsettm %s317_s26 }        // Set timer mode
0xaa : { %96 = vdelay 1 }                 // Delay for pipeline drain
0xab : { %97 = sfence }                   // Memory fence
0xac : { %s318_s27 = smov 0 }
0xad : { %98 = sst [smem:[#allocation17]] %s318_s27 }  // Signal completion

The sfence ensures all memory operations have completed before the final store signals to the host that the TPU program has finished.

Conclusion/Takeaways

Eight lines of JAX code became 250 VLIW bundles across 5 fused kernels. Here’s what the compiler did to get there:

The TPU compiler is a sophisticated codegen compiler. GPUs have codegen too (Triton, XLA/GPU, torch.compile), but the TPU compiler is doing something impressive: it takes your entire computation graph, fuses across operation boundaries, schedules VLIW bundles that saturate 5+ hardware units simultaneously, and orchestrates async DMA — all automatically. No manual tiling, no explicit shared memory management, no @triton.autotune. You write JAX and the compiler figures out the rest.

It generalizes to novel workloads. This matters more than it might seem. On GPUs, peak performance often requires hand-tuned kernels — FlashAttention exists because the compiler couldn’t find that optimization automatically. The TPU compiler’s approach is different: rather than relying on a library of pre-optimized patterns, it reasons about your specific computation from first principles. A weird custom attention variant, a novel normalization scheme, some exotic activation function — the compiler will fuse them, schedule them, and generate reasonable code without anyone having to write a custom kernel. The performance ceiling might be lower than a hand-tuned GPU kernel, but the floor is much higher.

Fusion is everything. The original ~25 HLO ops collapsed into 5 fusions. Each fusion keeps intermediates in VMEM, avoiding HBM round-trips. RMS norm and softmax never touch HBM — they execute entirely in fast on-chip memory.

Async DMA hides memory latency. While the VPU computes softmax, the DMA engine loads the value matrix in the background. By the time softmax finishes, the weights are already in VMEM — zero stall before the final matmul.

VLIW packs independent ops. A single bundle can execute a DMA transfer, three VMEM loads, and two vector ops simultaneously. The compiler statically schedules everything; the hardware just executes.

What does this mean for you?

If you’re training or running inference on TPUs, you don’t need to understand any of this to get good performance — that’s the point. The compiler handles fusion, memory management, and scheduling automatically.

But if you’re debugging why a particular computation is slow, or curious whether a custom operation will perform well, you now know how to look under the hood. Dump the IR, find your fusion, count the bundles, check if DMA is overlapping with compute. The tools exist; they’re just undocumented.

And if you’re deciding between TPUs and GPUs for a new workload: TPUs reward experimentation. You can try unconventional architectures without writing custom kernels. The compiler will generate reasonable code for whatever you throw at it. That’s a different value proposition than “fastest possible matmul” — it’s “fast enough matmul for any shape, any fusion pattern, automatically.”

If you want to explore this yourself, the dump flags are straightforward:

XLA_FLAGS="--xla_dump_to=./hlo --xla_dump_hlo_pass_re=.*"
LIBTPU_INIT_ARGS="--xla_jf_dump_to=./llo --xla_jf_dump_llo_text=true"

The HLO is readable. The LLO takes some squinting, but the patterns emerge — look for vmatpush/vmatmul pairs for matmuls, vxpose/vpop.trf for transposes, and vrot.slane for reductions.

Most of this compiler is closed-source, but the IRs tell the story.

Questions? Message me on Linkedin: https://www.linkedin.com/in/patrick-toulme-150b041a5/