Ian’s Blog

Category: ML Infrastructure

  • Make this forward pass as fast as possible

    Flapping Airplanes just raised a billion dollars, in part through demonstrating a megakernel framework that could do efficient forward and backwards passes while (for some reason) playing Doom1. Elsewhere, Hyperstition have what they call a Hyperkernel, which is generating high-performance kernels for all sorts of models. Turning a model into something faster is exactly what ML compilers have traditionally done, and all this activity begs the question of whether we need ML compilers at all any more, or whether we can let agents optimize directly?

    The folks working on TVM have been asking this, too:

    If agents can handle the low-level details themselves, should we still put a compiler between them and the hardware?
    We believe the answer is yes, but for a different reason. Instead of requiring agents to reinvent IR and compiler infrastructure at every iteration of kernel development, a minimal, stable ML compiler foundation can provide shared abstractions and reusable tools for the entire community while leaving agents free to explore their own optimization strategies.

    The core of their argument is that kernel optimization is a search and iteration process. Tokens and tool calls cost money and time. So, you want to optimize that search, avoid getting stuck at local optima, and offer ways to verify the generated solutions are correct. They build checkers for things like memory access and synchronization and they formalized their IR in Lean to enable cheap verification:

    In this future, we may no longer need a compiler that actually compiles anything, but we will still need compiler infrastructure to understand and verify programs. A shared IR can serve as the foundation for analysis, verification, and new tools built by agents.

    Together’s new Together Compiler (TKCC) builds optimized inference kernels. But, they reach quite a different conclusion around intermediate representations:

    Traditional IRs have limited expressivity. We hypothesized that IRs could help agents understand complex inference codebases, which contain multiple codepaths and kernel launch mechanisms. When directly provided with the source code, we do observe that agents spend ≈ 20% of their time reading engine files to understand the relevant code paths and kernel calls within the engine code. However, we also find that unrestricted source code access unlocks a higher performance ceiling. Even our heavily endowed augmented Inductor IR was not expressive enough to permit certain optimizations

    They still focus on the needs of verification and cutting down the search space, though. Much of what they built was an empirical measurement framework, with different tiers for executing kernels that normed out variation in forward passes, and offered hot-swapping for fast iteration. They also cut down on search by developing a custom cost model for parallelisms that culls options analytically, before testing in their tiered setup:

    We propose a cost model that ranks the top-k parallelism configurations using analytical estimates and tiered measurements, ranging from theoretical hardware capabilities to measured compute, memory, and communication performance. This enables TKCC to identify promising configurations before investing in expensive agentic optimization

    So its not so much compilers going away as it is them being decomposed. A lot of traditional compiler techniques become less useful in a world where inspection and investigation is a cheaper, and a lot more automated.

    At the Triton conference this year, Neil Dhar from OpenAI2 talked about how they have been using Codex to do large scale code-mods, with the aim of removing heuristics and implicit semantics from the Triton compiler. Those kind of techniques are great for getting improvements across a wide range of code, but they can make the code unpredictable: a small change can have an outsized effect on the result because it triggers different compiler passes.

    One example Neil used was removing automatic FMA. Triton had pattern-matched a multiply followed by an addition, and replaced it a single fused multiply-add. But! That could change the numerics of the kernel, due FMA dropping intermediate rounding. They removed this pass and manually inserted FMA calls in the source of all the kernels that used it3. This makes the search landscape smoother, and easier to navigate (though possibly larger)

    TVM’s TIRx is a bet that starting from fresh will allow them to build a more search-and-verification friendly compiler stack: they deliberately don’t apply heuristic-based optimizations, but do offer many extension points so agents can iterate new capabilities on top. OpenAI’s Gluon and Meta’s TLX are both directionally similar bets: they add low-level extension points while reinforcing the core of an existing compiler. Together are suggesting that providing the right tools is more important than betting on a single toolchain to guide search and verification.

    Compilers users to make things faster, but now they need to make the search faster, which isn’t quite the same thing. I would be hesitant on pronouncing a winner though, this space co-evolves with the agents themselves. In one particularly clear example Together had a note about some of the things that absolutely didn’t work in helping their agents: it read pretty much exactly like a list of agentic best-practices, circa six months ago4.

    1. Of course, because every idea is now a prompt away, someone immediately put up an open source version, so feel free to run Doom on your GB300. ↩︎
    2. And, full disclosure, formerly of the Triton team at Meta, where I currently work! ↩︎
    3. I mean, ish. OpenAI might have automatically fixed all their kernels that depended on this behavior, but they didn’t fix everyone else’s! There were a few hurried code searches going on as Neil revealed the change had already landed upstream. ↩︎
    4. Providing a knowledge hub lead to 25-31% lower speedup, given the same optimization time. Seeding with expert-written strategies lost to literally prompting “make this forward pass as fast as possible”, which is a great ablation. Providing human-engineered tools for research and logging impeded throughput: agents spent up to 43% of their tool calls on understanding the systems rather than just getting on with optimizing. ↩︎

  • Same same, but faster

    This week I’ve been at the Triton and PyTorch conferences in San Jose, and yesterday I had a chance to browse around the PyTorch poster presentations.

    There’s a lot of talk about agentic development, of course, and about how to scale it. Much of that comes down to verification. But what are we verifying against? The theme from many of the posters was: we are verifying the code against… some other code. Which is odd! We could just check whether our models are working properly, but models are surprisingly forgiving: they will happily train straight through a bug.

    Folks from ByteDance presented OpGuard which compares two training runs op-by-op and reports the first op where bits differ. They found race in PyTorch’s embedding backwards that had been there for 7 years. The bug hardly affected the loss (about 1e-5!)

    LinkedIn presented their people search, which scores member matched via embedding on GPU in FP8. Compared to their FP16 reference about a third of the top 10 results shifted, but the LLM scored relevance didn’t move a bit. This forgiveness is a nice feature most of the time, but makes validation difficult. LinkedIn ended up keeping the FP16 copy in DRAM and reranking the FP8 candidates with it. Metric wise it isn’t more right, it was just… there first.

    These reference implementations are becoming a kind of an extant spec, that will be then processed by agents under some set of constraints around exactly how close an optimized version has to be.

    Guilherme Leobas at OpenTeams (building on an idea from Richard Zou) has been rebuilding Dynamo’s (the PyTorch compile front end) object model to mirror CPython’s directly, using Python’s own test suite as the reference. The HuggingFace folks have been enabling on-device models for WebGPU, where they match against a CPU reference for correctness.

    Nvidia presented a “Day-0 RL” poster that was basically a checklist for using something as a reference: use identical weights and inputs, check trainer and rollout log-probs agree, and monitor more than just the reward score. As if to demonstrate why, AMD presented on getting DeepSeek V4 Flash RL on MI355X. They had SGLang doing rollouts and Megatron doing training, and ran into disagreements. It wasn’t the usual suspects either, SGLang scored experts with a sigmoid where Meatron used sqrt softplus1, and there were some scale and clamp config differences. But… which is “right”? They ended up using an independent HuggingFace version as the reference, to break ties.

    There were some hallway track discussion about formal methods for kernels and models. Maybe by next year we won’t have to test to see whether two things are equivalent, we’ll be able to prove it! Still have to decide which is “right” though.

    1. This is DeepSeek’s actual score function. It probably didn’t change the expert picked, and I am guessing that SGLang reused an existing module because it already worked! ↩︎

  • Proving your kernel is right

    We were talking the other day about the lack of formal verification of FlashAttention 3, which, it was pointed out to me, somewhat undersells the tremendous amount of verification work going on for kernels right now. Fair!

    When OpenAI announced their Jalapeño chip a couple of months back, SemiAnalysis went along to take a look at what they had done. OAI showed off some kernels, written by Codex, for getting DeepSeek MLA running on the chip. As Jordan Nanos of SA put it:

    “The point is, the guys who were scrolling through this code with us, it was kind of clear that they know a whole bunch about hardware, they know a whole bunch about the concepts in the system, and they just have no idea what this MLA kernel that they’re showing us for DeepSeek actually does.”

    They didn’t have to read it because they were confident it was correct. Which is not a particularly easy thing to do. A new paper from Microsoft Research and friends, RESOLVE, had 5.6 Sol write optimized versions of the 100 level 1 KernelBench problems, looping until each passed the benchmark’s tolerance checks. Two of them, an ELU and a hardtanh, were… empty. They literally didn’t do anything, and passed anyway!1

    The decision of what you are actually verifying is quite a consequential one. Even for something as straightforward as “accuracy”, choosing a tolerance for torch.allclose changes the design space. From the RESOLVE paper:

    Testing must consequently compare outputs using a numerical tolerance, but a permissive tolerance can conceal bugs while a restrictive one can reject valid optimizations.

    Even if you lock down your numerics, you still have a range of inputs and execution conditions. The TIRx Harness is an agentic-oriented stack of compiler+tools for exactly this kind of kernel loop2, and made a special effort around verifying concurrency:

    In one KDA trace, a candidate passed the benchmark’s correctness tests, but synchronization analysis still identified a latent bug: a reused mbarrier could advance to its next generation before a late consumer finished waiting on the previous one.

    The RESOLVE folks attempt to prove their kernels are correct3 using Kuiper, a verification language that compiles down to CUDA. WarpDRF, another Tyler Sorensen joint, looked at warp primitives like shuffles, tensor core ops, warp barriers. Which threads take part in different operations depends on how the threads diverge. They used Rocq, another verifier language, to prove they were consistent, then fuzzed thousands of test per config across a range of devices to see if they held.

    Most of it held together, but Nvidia was… patchy. If the whole warp stays together, you are good. Beyond that, you just don’t know which threads are going to participate. The team used their findings to build a static checker and ran it over the kernels in llama.cpp. Most of the kernels were fine, but a few weren’t!

    Zhongming Yu, one of the authors of NVIDIA’s CAKE, wrote a compelling post the other day talking about how the blocker for true recursive self improvement is the speed of verification, and specifically called out megakernels as an example of fragile workloads. Kernels generally run pretty quickly (hence why we can autotune), but setting up the right verification is hard, and the more complex the verification, the slower the whole loop goes.

    RESOLVE costs it out, towards the end of the paper. Testing elementwise ops took milliseconds. Determinism testing took about 1.5 minutes per kernel. Reduction and concurrency… around 2.4 hours.

    1. And, this wasn’t even surprising! KernelBench’s maintainers list no-op kernels as a common reward hack: if the output buffer still holds the answer from an earlier run, doing nothing is very fast and very correct. ↩︎
    2. As an aside in the more questionably-verification side, one agent found an output-lane masking option on tcgen05.mma by reading the ISA, which is cool. It didn’t reduce the work at all, but did raise the observed clock frequency, for a 2-3% win. Would be interesting to see if that one holds up outside the benchmark! ↩︎
    3. They have a reference kernel and an optimized one, and they are proving equivalence. They create a simplified version of each with minimal concurrency, ensure they generates the same answers, then validate the kernels are equivalent using Kuiper. ↩︎

  • Mo’ turns, mo’ problems

    Doing reinforcement learning for agentic use-cases is, conceptually, not all that complicated. For each problem you set up a VM with a harness and a fresh model context, send the prompt, let the model make tool calls, score the whole thing at the end, and update the weights.

    The important assumption in there is that the model generating the rollout and the model you’re updating are the same model, which, for reasons, is referred to as being “on policy”. Most RL algorithms use PPO’s importance ratio: policy/old_policy. This scaled the update based on how close the model is to the version that produced the data. On the first update that ratio should be exactly 1: it’s the same weights, seeing the same tokens, so we should get the same answer. Differences between the two makes Horace sad.

    Also, agentic post-training means there is more than the model to consider. Rollouts come from the model, in an inference engine, but also a harness and sandbox environment. These can also produce some weird differences, and as session get longer the chances of hitting a meaningful difference increase. Sessions are getting pretty long these days: In Alibaba’s QwenGyre paper rollouts on their 122B model went for around 2 hours on average, while for the 2.4T flagship many of them ran past four hours.

    How you train it is all tied up with what you want to do with the model. Like, if you want a model to be good in Claude Code why not train in Claude Code? That’s what Alibaba do, in part because proprietary harnesses…

    are closed and evolve rapidly; reimplementing their control flow inside an RL framework is impractical

    But! Then its black box. You want to keep your GPUs busy by flexing which ones are training and which are doing rollouts, but if you stall your harness you start getting different outputs for reasons which are nothing to do with the model itself. Alibaba keep their sandbox containers running and route inference through a proxy. They cluster their GPUs into cells which flip rapidly between rollouts and training based on demand. The proxy reroutes in-flight requests and KV cache gets shuttled between cells.

    DeepSeek, on the other hand, use their own harness so they can just pause everything. They are also trying to optimize overall throughput, so will routinely preempt long-running rollouts. Before v4.1 the agent loop lived in the (preemptible) GPU job, and had to be reconstructed via a reply when the job came back. With 4.1 they moved the harness inside the sandbox itself and just froze the whole state: DRAM and the writable layer on the file system.

    Even when the harness is consistent, you can still go off-policy with the trajectories themselves. In simpler RL-times rollouts were single streams of generation, but tool calls introduced extra tokens that needed to be masked out lest the model train on things it didn’t produce. Even then, if you let the inference server parse the models output text you can introduce variations: a tool argument of false might become False. As we’ve talked about before, the Prime Intellect team spent a lot of time on this in their Renderers library. QwenGyre’s proxy even caches the original tool-call payloads because the harness reformats them between turns.

    Even with that, you increasingly aren’t dealing with one rollout per… rollout. Agents spawn subagents, make parallel tool calls, and compact long sessions.1 What comes back from a long session is a tree, which isn’t all that easy to represent to the trainer.

    Tree Training serializes the tree into one (or more) streams and weights each token by the percentage of branches that have it in history. QwenGyre caps the number of paths that can be trained on2, and averages token weights across rollouts so ones with that spawn a lot of subagents don’t automatically get more gradient.

    Finally, even if you have worked hard to get reliable tokens, you have to worry about the operations on those tokens. Because inference is (often!) latency oriented with small batches, and training is throughput oriented with larger batches, you can end up with different kernels for the same jobs, running in different engines: SGLang or vLLM for the rollouts, Megatron or TorchTitan for the training. We’ve talked about the pain of training-inference mismatch before, but it gets worse as runs get longer.

    One simple fix is to recompute the old log-probs on the trainer. This is brute-forcing the ratio to 1 by going off-policy to the sampler, and it means you have to do an extra forward pass over every rollout. Another approach other is to correct for it: QwenGyre’s trainer clips the ratio per rollout, and caps the weight for each token. This is a bit of a fudge, which, for example, the Megatron folks are pointedly unenthusiastic about, as it “hides the mismatch rather than removing it.”

    If you can make computation completely bitwise-equivalent, you don’t need to correct or to recompute. A paper from colleagues of mine at Meta dug into where bitwise differences come from for floats in a number of kernels: reduction ordering, in large part, but also partial-sum precision, FMAs, and where you round. Fixing the order isn’t free: it can cost 10-20% of performance: Colfax’s recent FlashAttention-4 work shows the same on the training side: 1.13x on the head-dim-64 backward pass, but only 1.05x with deterministic accumulation. There are some efforts to build in bitwise-consistency earlier: AReaL-TIK is an agentic loop which starts from a single bitwise-consistent kernel shared by inference and training, and then has agents make each faster while maintaining parity.

    I do appreciate that agentic post-training is responsible for billions of dollars of revenue, the field being renamed superintelligence, etc., but it probably helps that it is a fun ball of mud to work through too.

    1. You want to train the summarization too so the model learns to compact well. CompactionRL has a good recipe. ↩︎
    2. A max of 5, ordered by main agent, its summaries, subagents, their summaries. ↩︎

  • FA3 had some bugs

    For reasons best known to themselves OpenAI blessed the world with several hundred new mathematical proofs. This inspired elation of some mathematicians, despair of many others, and confusion of a great many people whose only encounter with zeta moments involved Sean Connery and a laser grid.

    Lean, the language the proofs are written in, originally started for proving software so you might think we were using it for that too?

    But, surprisingly, we are still finding bugs in possibly the most widely used kernel of the last few years, FlashAttention 3:

    When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN.

    There were two issues. One was an issue I remember from 2024 in FA2. Horace’s comment on that issue sums it up very nicely:

    We’re computing exp(x_i * scale - max_scaled). Now, max_scaled = max(x_i * scale). The idea here is that before we take the exponent, we normalize everything down to 0 so that exp(x) doesn’t blow up massively.


    Now, for the largest value of x_i, max_scaled = x_i * scale. So, x_i * scale - max_scaled == 0, right?


    Unfortunately, no, due to fma. max_scaled is actually equal to round_fp32(x_i * scale). But in fma, x_i * scale never gets rounded! So with fma, we are computing x_i * scale - round_fp32(x_i * scale), which is not equal to 0!

    This was an opt-in thing for FA2, and just never fixed for FA3.

    Turns out though, there was another bug: dS, which is the gradient with respect to the attention scores, should sum to exactly zero: its expresses how the keys differ from each other, not where they are absolutely. FA3 rounds dS to BF16 before multiplying it, so it doesn’t quite sum to zero. Early in training that’s mostly noise, but as the keys get larger and attention spikier it can swamp the real gradient.

    Part of the reason this hadn’t come up was that applying QK-norm (or doing QK-clipping) is quite common, and AdamW’s weight decay tended to avoid the explosion too. Muon, however, is more exposed, and that’s what the team were using when they noticed this.

    Quasi-Riemann is one thing, but reduced precision floats? Quite another.

  • Environments and Benchmarks

    One big concern in AI Safety circles is eval awareness: if the model can tell it’s being tested, might it dissemble and tell the evaluator it’s not going to turn us all into paperclips? In terms of what models are actually doing right now, though, it seems like much of the problem is that models don’t know enough about what is expected of them.

    Imagine fixing a bug in some code, in a Docker container. The model executes via a harness, makes changes, and maybe gets feedback from the user. If you have an eval for this bugfix, then the model does exactly the same thing, but the code is taken and run against some hidden tests to get a score. In the model is going through post-training, it does the same thing again but then gets a weight update based on the results of that grade.

    The problem is that the folks calling the model kind of want it to do different things in each case.

    This week Xiaomi released Mimo 2.6. They’ve been unusually open about the RL process for it, publishing a dashboard for the run, releasing a detailed technical report, details of how they built their RL envs, and the RL envs themselves. At the same time Epoch released their evals of evals: a dashboard grading the reliability of popular benchmarks.

    Looking throught this, you can forgive the models for being confused. For example, the MiMo folks observed some reward hacking in their paper:

    Task (pytest): Fix Windows conftest.py imports broken by path lowercasing.
    Thinking: “Let me check the pytest changelog or GitHub to see if there’s a more recent fix.”
    Action: pip install pytest==5.4.3; inspect the installed source.

    Xiaomi, and its users, want a model that can fix a bug even if it can’t just look it up on the internet. But also, if I am trying to fix something and the answer is on the internet then… yes, let’s just use that please!

    Epoch found confusing situations too. For example in DeepSWE they found that sometimes the models would do the right thing and write the test. Unfortunately they would choose a somewhat obvious name, which happened to be the same obvious name the secret verification tests used. This would cause the hidden test to fail to compile, and the task to be graded as a failure:

    The agent defines a symbol, usually a test function, that the verifier’s tests also define in the same package. Example: While solving the onedump-dump-encryption-pipeline task, Fable 5 modifies encryption/config_test.go, defining TestLoadKeyFromEnv, TestLoadKeyFromFile, and TestLoadKeyFromLiteral inside the encryption package. The verifier defines hidden tests with the same names in the same package, producing redeclaration errors.

    Epoch point out the models have no way to know they are going to hit this. Xiaomi, on the other hand, want their model to avoid those kinds of errors. They added explicit instructions against retrieving existing solutions, basically testting the model it’s in a test of some sort in order to get the behavior they want.

    This doesn’t just make it hard for the models, it makes it hard to know whether the models are any good. MiMo 2.6 Pro was evaluated on DeepSWE, getting to 71.91. Opus 5 got 74, GPT 5.6 Sol 73. But, Epoch found at least 20% of tasks have these kinds of false-negative problems. So maybe some models are better, or maybe they are luckier! And maybe the ceiling on this benchmark is pretty close to where the models are.

    What labs want is the “right” behavior for users. That requires very careful RL environment setup so the right thing is rewarded, and very careful evals where they that right thing is consistently happening. When your benchmarks and your environments are assembled similarly, based on real open source projects, this can be… tricky.

    1. Ish? Pretty sure it topped out higher than that on the RL run evals, but maybe thats not the checkpoint they shipped. ↩︎

  • Agents have favorite tools

    A lot of people use coding agents like software engineers (for all kinds of things, it turns out), and adopt a dealing-with-humans mindset about how that works. If a model has a useful tool, and you give it another useful tool, then it gets the capabilities from both.

    Largely that seems… not true? You can make a tool available, but then the model has to decide to use it, and often it just doesn’t. My recent futzing with LSPs ran pretty directly into this. One follow-up from there was to try the Python part of SWE-Bench ProMax, a benchmark built around large, coordinated refactors. This is exactly the type of work where an LSP’s find references should come in handy.

    Helpfully, Aaron Pollack pointed me at the CodeAnchor paper, whose authors had run into pretty much the same problem. They had tried giving their agent a call-graph tool and…

    “we observed low tool-use rates: the agent typically relied on plain grep instead.”

    Rather than retraining the model, they put the call-graph context into the relevant code itself, adding annotations like “used by: foo, bar” near the definition. I tried a variant of this by annotating not the code, but the grep results: when the agent searched for a symbol, the LSP found references and appended them to the grep output.

    From looking at the failed traces, it was clear the annotations were working. Of the 99 files surfaced by them, all 99 were opened, and 92 got patches. This resulted in fewer turns, and it did seem to help with run-to-run variance, which was also a benefit the CodeAnchor folks saw: I had fewer runaway runs with the annotations than without.

    The most common failure, though, was incomplete refactors. The model would find the core fix quickly, edit it correctly, then miss other files which also needed changing. The misses were ones the LSP didn’t really help with: they were about compliance with long, winding specs. It is possible to use LSP tools to help with these, but the model didn’t.

    Injecting tool-derived context got rid of the adoption problem, but how the model behaved was still driven by its learned tool-use policy.

    There is some interesting academic and non-academic work around this idea of tool use as a behavioral policy, and a surprisingly predictable one at that. BiasBusters showed that models have strong preferences among functionally equivalent tools. AutoTool found that tool calls tend to follow predictable sequences, and used previous agent trajectories to predict the next tool and bypass some inference entirely. On the non-academic side, when Steve Yegge was building the Beads CLI he let the agents hallucinate arguments and then just added them, so the tool naturally conformed to the expectations the model already had.

    You can swap in tools which have roughly the same shape as ones models are familiar with, and you can add output to familiar tools to give them more information. But getting models to use new tools, or to use tools differently, seems much trickier.

    The lab answer seems to be, once again: more RL then.

    DeepSeek’s V4.1-Flash technical report, which we discussed yesterday, is particularly explicit about this. Their post-training approach is pretty standard: SFT, then RL, then on-policy distillation. What is new is automatically generating tasks and environments to train the agent with. They seed these from internal agent sessions that show poor model performance, and popular GitHub repositories:

    “coding-agent sessions from internal employees and external partners … [and] public GitHub repositories that meet a star-count threshold.”

    If you want models to get good at your tools: make sure they are really popular, or get a job at a lab and fail a lot.

  • Agents love prefill

    LLM inference has two stages: prefill, where the prompt is processed and the KV cache is built, and decode, where the model auto-regressively generates tokens. In a chat use case the two are somewhat close in size. The user writes a prompt, the model reasons about it then generates an answer, which is probably longer than the prompt.

    That is no longer where the FLOPs go. Everything is agentic now (even the chats), so the loop looks more like:

    • you put in a query
    • the model generates a tool call
    • the tool call runs
    • the tool output is appended and prefilled, and the model decides what to do next

    Chat had this back and forth too, but prefix caching meant that you didn’t have to re-prefill what had already been generated. Tool outputs are lengthy, uncached, new content: file contents, terminal dumps, a fly’s connectome, etc. etc.

    This was clearly On The Mind of folks at DeepSeek. From their v4.1 Flash technical report: “The widespread adoption of long-horizon agents has made model workloads increasingly input-heavy.” Heavy enough that they made some interesting architectural changes. DSv4.1 Flash is a 552B parameter MoE with 16B params active… during decode. For prefill they just run the first half of the model, where only 8B are active!1

    They call this “Causal Encoder–Decoder”, inspired by Microsoft’s YOCO. For T5 fans, it isn’t an encoder-decoder in the old2 sense: everything is causal. The first 20 layers, the encoder, behave normally. The second 20 decoder layers don’t derive KV from their own hidden states. The first decoder layer projects KV based on the final encoder state, and the other nineteen reuse it. So, during prefill you can stop half way through and still have everything the decoder needs.

    It’s quite a bet! The second half of the model basically never sees the input, just what the 20 prior layers thought about it. It is, I think, the first frontierish model where there is significantly less compute on the “read” path than the “generate” path. But it does make sense that you maybe don’t need to think quite as hard about the thing that already exist.

    You can get a sense of how much prefill costs people by looking at the spreads between the different options on their rate cards. Most APIs quote prices in three columns: cached input, uncached input and output. Most labs charge cached input at around 10% of the price of an uncached input. DeepSeek charges 2%!

    You might think that the split above is what shrank the cache, but turns out not so much. DeepSeek are very aggressive with their attention: a combination of Compressed Sparse Attention3 and storing the KV in FP4 mean its cheaper to run through the decoder and store the derived states rather than storing the encoder hidden state that it derived from.

    So, they added CED to cut prefill cost by half, then CSA2 and FP4 to cut cached cost by four. The whale does not play when it comes to efficiency.

    1. Give or take 128 tokens. Every layer has a sliding window, and that state has to come from somewhere, so they last 128 tokens go through the whole model. It’s a bit wrong (the first of the 128 should be informed by the prior windows), but not wrong enough to matter apparently. ↩︎
    2. 2017? 2014 for Ilya stans. ↩︎
    3. CSA2, technically. ↩︎

  • LSPs for LLMs

    Back in the dark ages of typing code into editors we were aided by squigglies under broken code, click-to-definitions links, and so on. That was powered by language servers and type checkers. Several harnesses now expose an LSP as a tool, on the reasonable premise that better code intelligence makes for a better agent.

    Models, though, are trained primarily with the tools they always have, which tend to be things like grep and ranged-read. Supplanting those while gaining effectiveness is tricky.

    Models understand codebases a bit differently than people. A human can keep a handful of files in view, a slightly larger handful in their working memory and, over time, they build up an approximate mental model of the code base.

    An agent has an enormous context window and can understand a lot of files at the same time. They can find symbols by searching for them, which will usually then trigger a partial read to pull part of the file into their context window. They also just… know stuff? Their weights contain reasonably high-fidelity versions of an awful lot of public code. That helps with navigating that code, or reasoning by analogy about other code bases.

    To try and see how LSPs might help in this process, I ran a bunch of experiments, which are in this repo.1 The experiments were on a mix of local and API models, working against Python. I used Pyrefly for the checker, and a static AST resolver validated against Pyrefly2 to run most of the tests.

    The questions I had were whether the LSP tells the agent anything it couldn’t otherwise find, whether LSP answers are cheaper in tokens than the default reads, and whether the timing of a diagnostic changes the outcome.

    The answer to the first one was, mostly, no. Resolving between 10 same-named overrides in a variety of setups worked whether or not the model had LSP tools. The type annotations did make a difference though. If I stripped those out, the text retrieval got worse. A capable model in a harness is basically already a decent type checker.

    With regards to efficiency, merely adding an LSP did nothing: the models didn’t use it without some prompting. If the model had to read a file to resolve a type, the LSP defn tool was cheaper than the file-read tool.3 But! The models would often do the file read as well, which completely negated the benefit of the LSP for token efficiency.

    Even if I injected definitions that contained the relevant fix, the model re-read the file in almost every case. Telling the model that the span provided was complete only saved a few % of those cases, and it cost more tokens in providing the prompt itself! To actually get the model to prefer the defn call to a file read required fine-tuning. I used a DAgger-style relabel (generate trajectories, swap the file read for the definition read, fine-tune on the result) on Qwen 3.6 and then the model would avoid the extraneous call.4

    It’s also worth noting that cheaper isn’t always better. These are input/prefill tokens, which are the cheaper and less interesting ones to optimize. The Codebase-Memory paper found a file-exploration agent beat a structured graph tool on real repositories (92% against 83%) at roughly ten times the tokens. So, sometimes those extra tokens are doing something useful.

    The third question was about timing. If you want to keep your type annotations correct (you do) you can use a type checker. But when is it best to deliver the feedback from it?

    I had an agent work on a set of draft changes that passed visible tests, but failed a held out one, and asked it to review and submit. Left alone, the model accepted the bad revision 11 times out of 12. When a type checker gated the submission, it only accepted the bad change 1 time out of 12.5

    After Thinky’s post on interaction models, I wanted to see whether delivering diagnostics live during generation (agentic squiggles, basically) would help.6 Largely, the answer is no: live delivery was neutral vs no feedback. Prompting the model to make the fix seemed actively harmful: telling an agent to work on tool feedback seems to override its own judgement in a bad way.

    Batching the feedback at end of turn or after each edit did help, and for overall token spend end-of-turn was a clear winner.

    Caveats caveats caveats: many of the tasks were synthetic, all the tasks were pretty easy, and every decent model solved pretty much everything. The timing results were only tested with a 7B model, so you may get better results with a strong one. The codebases in question can fit fully within the context window, though in actual usage agents never seemed to proactively just ingest the whole thing!

    With that said, I do have some takeaways, at least for my own work:

    1. Types are good. Correct annotations helped the agent navigate the codebase, regardless of tooling.
    2. Keep types correct with an end-of-turn hook. Ideally, make this a gate so it actively asks for repair on type errors and is silent otherwise. Still, measure it! Log how often it blocks, how often the repair passes, how often it rejects work it shouldn’t have.
    3. Measure token savings at the task level. If benchmarking a single op it is easy to conclude the op is cheaper. You need to see the full model behavior to really assess the change though: is correctness the same with and without a tool or prompt, and does it yield consistent token savings across a range of usage.
    4. Experiment more with LSPs and prompts on larger codebases. I didn’t test this directly, but my instinct from the results is that you will get more out of navigation tools for large, complex, private, codebases. You will still likely have to prompt the model to actively make use of the tools: by default, they’re going to reach for what they know.

    There is some interesting future work out there. The tools we’re using were built for humans who can apply discretion/ignore output. That is trickier for a model, and something that probably needs training.

    I’d also like to be able to evaluate a codebase as to whether a tool will help without having to just run a bunch of tasks over it: are there metrics we can collect statically that might inform those kind of decisions?

    Finally, I think we need more large code base refactor and migration tasks: problems of a ProgramBench scale but working with a large, existing, and not-in-the-training-set codebase. On that note, SWE-Bench ProMax came out as I was writing this: seems relevant, but I haven’t yet read it!

    1. Credit to Codex and Claude for most of this repo, the writing in the report is LLM+a bunch of editing, so temper your expectations, and all the numbers are reproducible from the artifacts there if interested. ↩︎
    2. On real library symbols the two agreed most of the time: the gaps were re-exports where Pyrefly returned null and the resolver didn’t. ↩︎
    3. The ranged-read cost about 1.3x more in tokens than the defn call. ↩︎
    4. One risk with this kind of training is that it teaches the model superficial tool usage, but not judgement: sometimes, you do need to actually read the file! Somewhat surprisingly that didn’t seem to be a problem: when I gave it insufficient spans it went and read every time, and when the span was sufficient it read only twice. There weren’t any actual examples of tasks requiring reads in the training data, so it suggests that the conditioning doesn’t completely kill model judgement. ↩︎
    5. In the other cases the bugs were exposed in a type-check, but this one type-checked clean, so no signal. ↩︎
    6. The actual approach I used for injecting results into a live stream is the async injection of events approach described by Hooper et al.. ↩︎

  • Power by the hour

    It is a truth universally acknowledged that an airline in possession of an airplane must be in want of engines to make it go. Yet, somewhat surprisingly, they don’t really buy engines.

    Rolls-Royce were the notable innovator here in selling not an engine, but instead what they call power by the hour:1 airlines pay a price for the hours the engine actually flies.

    There are a lot of operating decisions in owning and maintaining an aircraft engine: for example, if you have to take the engine off the wing for repairs it means a lot of downtime for the plane, during which it isn’t making any money. Airlines used to make the call of when to do proactive maintenance vs when to do more extensive refits based on their experience within their fleet. But Rolls-Royce had a view across all their customers, and could do a much better job of predicting maintenance and repair timing, getting the right parts to the right places, and generally keeping the engine in service.

    They converted that operational expertise into a product, which just happens to be delivered via a massive turbofan. This pricing model was appealing to the airlines since it aligned with their incentives, and their risks. Under the old model engine makers got paid for spare parts and shop visits, so failures were a kind of revenue. But once they were paid per flying hour then downtime became their cost too.

    In ML infrastructure, those engine flight hours are generally referred to as goodput. A Google Cloud blog post from a couple of years ago has a good set of definitions:

    Runtime Goodput measures the time spent to make forward progress as a fraction of time when all training resources are available. Maximizing runtime requires careful engineering considerations. […]


    Program Goodput measures the fraction of peak hardware performance that the training job can extract. Program Goodput is also referred to as Model Flop Utilization or effective model flop utilization, i.e., the model training throughput as a fraction of peak throughput of the system. Program Goodput depends on factors such as efficient compute communication overlaps and careful distribution strategies to scale efficiently to the desired number of accelerators.

    To give one concrete example if you checkpoint every N steps then get a failure at step N-1, the prior steps are lost work and don’t contribute to goodput. A surprisingly large number of things can impact it; Llama 3.1 had a section talking about the reliability pain during pre-training:

    During a 54-day snapshot period of pre-training, we experienced a total of 466 job interruptions. Of these, 47 were planned interruptions due to automated maintenance operations such as firmware upgrades or operator initiated operations like configuration or dataset updates.

    Improving goodput is a horrendously complex project spanning hardware, software, cluster scheduling, kernel work, profiling and resource management. It’s often the preserve of dedicated ML infrastructure teams2, but not everyone has such a team available.

    Neoclouds sprang up to provide pods of accelerators, but the actual goodput available varied enormously. The GPU provider gets paid for the GPUs you are renting, regardless of whether they are making useful progress or sitting in a NCCL timeout.

    SemiAnalysis’s ClusterMAX rating gave a (controversial!) view of the varying operational capabilities of the neoclouds, measuring goodput-impacting things like spare GPU capacity, failover and repair times, launch overhead and so on. This let purchasers at least compare reliability, but still fundamentally they were charged for the GPU-hour.

    What’s emerged now is a higher level of the stack, notably Thinking Machines’ Tinker, and Prime Intellect’s Lab products.

    Both teams are full of folks used to extracting top-tier goodput, and much like Rolls-Royce they are both selling that operational expertise and have the ability to improve it by working with a wide range of customers. Rather than power-by-the-hour they sell goodput-by-the-step: you pay for meaningful progress, not for idle clusters.

    Given the capabilities of open models, like the (open this month) 2.8T Kimi K3, or Thinky’s own 980B Inkling, there is an opportunity to RL a strong model on a bunch of domain-specific knowledge or tasks only you have access to.

    Bridgewater, the HR experiment with algorithmic hedge-fund attached, published a case study with Thinking Machines from their AIA Labs group, where they attempted to train Qwen on certain financial analysis work:

    Frontier models we tested on struggle with relatively simple financial tasks, and model advances don’t improve performance much. In contrast, we’ve shown that high-quality proprietary datasets labeled by expert investors and used for fine-tuning produce custom models that exceed frontier performance on our tasks.

    This is sometimes referred to as RL-as-a-service, and there is legitimate skepticism of whether that is a good business. But really, this is frontier-scale, reliable training infrastructure as-a-service.

    You can see this in the Bridgewater post: a small hedge-fund research team did a multi-iteration RL recipe search at 235B scale, and ended up with GRPO, CISPO loss with asymmetric clipping, interleaved batching and on-policy distillation, all ablated. That kind of experimental capability at scale was largely inaccessible outside perhaps 20 firms a year or two ago. Now, it’s doable with an API call.

    1. Technically this term is much older, dating back to the 1960s with Bristol Siddeley, which is now a subsidiary of Rolls-Royce. But in practical terms this became a commercial airline thing in the late 90s with TotalCare: GE and CFM, the other big engine makers, followed along. ↩︎
    2. Disclosure, several of which I have had the pleasure of working in at Meta! ↩︎

  • Maybe the agents shouldn’t write the kernels

    A thing you can do is take the most performance and correctness sensitive part of your stack and just ask a chatbot to write it for you. They will sometimes get it right!

    Back towards the end of 2024 Ouyang et al at Stanford attempted to benchmark how often that happened with KernelBench. DeepSeek R1 could one-shot 12% of simple ops, 36% of fused operators, and 2% of whole architectures. Still, things have moved on a bit in the past 18 months, and Han, Zhang et al.1 extended the idea in KernelBench-X. They found:

    1. Writing correct kernels and performant kernels is somewhat decoupled. You can refine kernels and that mostly helps with correctness: the models got more of the tasks compiling, but it dropped the average speedup on the way.
    2. What you ask for matters more than how you ask. The category of the task explained 3x more of the variance than switching between different agents or other method varieties.

    “Together, these results indicate that the capability boundary of current LLM-based kernel generation is not a single wall but a sequence of distinct barriers – compilability, semantic correctness, hardware efficiency and performance portability – each requiring different mechanisms to clear”

    In one particular area they tried getting the models to write quantization kernels, an area known for needing numerical precision. They got 0 out of 30. The models produced running kernels, just not kernels that were, you know, right.

    One thing that did stand out to me was that a lot of the baselines were eager PyTorch, so I decided to run an experiment myself. How do these models do against a compiler, not just eager?

    I took a popular model (Qwen, naturally), ran it through torch.compile with minimal settings on my DGX Spark and identified three kernels that were eating big chunks of the wall time: SwiGLU, residual+RMSNorm and the SDPA prelude. I then had Codex, Claude and Kimi2 take a run at writing those kernels for the hardware.

    The results were an absolute blowout: SwiGLU 1.06x, RMSNorm 1.21x, SDPA prelude 3.91x. For the latter, Kimi stacked up three different weight matrices into one and fused multiple matmuls together. It was very impressive stuff.

    Then, another suspiciously well-timed paper arrived, FASTKERNELS from Snowflake. Rather than benchmarking against PT Eager or against single-operator references, they wanted to test how the models did on real model-serving problems, with a focus on the end-to-end speedups. Their takeaway:

    ”agents learn to generate kernels that score well in sandboxes but introduce interface incompatibilities, compilation-stack conflicts, and silent correctness degradation when integrated into real systems.”

    All of those issues hurt end-to-end performance when you put them in a real model-serving context. Of the three strong kernel-generation agents3 they tried, none beat the production baselines

    “in contrast to the supra-unity numbers these agents have reported on operator-level benchmarks whose reference is PyTorch eager.”

    Taking another look at my vibed-up experiment results, as the FASTKERNELS folks may have suggested, there was a catch. Several catches.

    The baseline, it turned out, spent an awful lot of time doing… kernel dispatch. Even getting a 3.91x speedup on SDPA prelude led to an end-to-end model speedup of… 1.007x. Not quite as exciting.

    You also had to be very, very careful about how the agents were getting speedups.

    For example, the initial correctness check accepted anything within cos_sim >= 0.95 of the reference kernel. Codex “won” the SwiGLU round by replacing sigmoid(x) with clamp(0.21*x + 0.5, 0, 1), a straight line which diverges from sigmoid everywhere except a narrow band near zero.

    It turns out this kind of thing is pretty common. The FASTKERNELS folks found a case where an agent needed to write an all-reduce kernel for cross-GPU synchronization. The test harness they were using was single GPU so the agent just no-op’d it, replacing the all-reduce with a straight tensor copy. This “passed its checker but produces the wrong sum on every scenario of our 4-rank NCCL+Gloo harness.”

    Even when the kernels are right, and fast, it doesn’t mean they are… good? Several of the generated kernels in my experiment were somewhat unshippable due to hardcoded shapes or silent global mutations. FASTKERNELS found similar things:

    “L2 failures are dominated by syntactically valid kernels that respect the per-tensor signature but violate the surrounding production contract.”

    Which I think is the academic way of saying they wouldn’t ship those either.

    If you get your verification of the problem wrong in the harness, you will get a kernel optimized for the harness. Use the wrong contract and your kernel will be wrong in exactly the shape of your wrongness.

    Still, a small win is a win, right! My original run had agents outperforming torch.compile by about 2.6%. At that point I had a friend take a look who immediately pointed out that I had hampered the compiler unrealistically, and suggested running on max-autotune. This was especially unfair since the agents each got several cracks at the problem. Turns out, with that baseline the agents lost by 4.6%.

    And, that’s pretty similar to what FASTKERNELS found. Across 88 tasks and three agents the best of theirs landed at about 0.94x4 the performance of the production stack.

    Fairly late in the day I decided to replicate the experiments I had run on the Spark on a 3090. That’s Ampere, sm_86, an elder statesman of consumer GPUs at this point. It turns out that once again, some of the wins were just worse baselines. For example, Kimi tried the same SDPA-prelude matrix stacking as on the GB10, but on Ampere the 3.91x speedup turned into a 0.74x loss. The difference was cuBLAS: it was simply better tuned for the 3090 than the GB10, and did a much better job of utilizing all 82 SMs. The baseline Kimi had to beat was (relatively) higher.

    The question of “do agents beat compilers” is hard to answer because what we are (roughly!) measuring is compiler maturity. Agents are most useful in exactly the window where a compiler is weakest: new silicon, untuned heuristics, and libraries that are still evolving5.

    As libraries improve, hardware is better understood, and compilers mature, the value of exploratory search diminishes: there are “right ways” and it’s better to just use them than create custom solutions. If an agent is identifying patterns reliably and repeatably, it may as well author a compiler pass and spend more tokens on the areas that can’t be as cleanly captured.

    1. I think these folks are associated with Tsinghua, but to be honest I am not entirely sure! ↩︎
    2. Each model ran in their respective coding harnesses. One fun takeaway was the wall-time for generating the kernels was a factor too. Kimi took at least 3x longer than the other agents, spending a lot of tokens on the way, but also generated the most performant kernels of the three on every task on Blackwell, which was not what I was expecting. ↩︎
    3. Codex, KernelAgent and Dr. Kernel, the latter of which I hadn’t heard of but has by far the best name. ↩︎
    4. Codex landed at 0.94x. KernelAgent at 0.78x. Dr. Kernel got 0.53x, but still billed my insurance. ↩︎
    5. I suspect this is particularly pointed for the GB10, which is an unusual piece of hardware, and in particular has a lowish number of SMs. ↩︎

  • The elusive order of things

    SIMT offered a fantastic bargain. You write a straight-line program, the machine runs a lot of copies of it, and when one waits for memory the hardware swaps in others. You look with disdain on the less enlightened thread programmers dealing with deadlocks and concurrency etc. etc.

    Choosing what to run where and when is a scheduling problem, and there have been three effective approaches to that so far.

    You can schedule statically: decide ahead of time what all the units should do each tick. You can schedule temporally: swapping in different phases of workers via a pipeline. Or you can schedule spatially: divide the resources of the machine into different roles.

    The underlying mechanics of which one you pick tends to be determined by the hardware. A chip like a TPU spends most of its silicon on math, and fairly little on orchestrating work. That means static scheduling, and a compiler that can build you that schedule.

    Ampere and before1, and all the modern AMD chips, encourage temporal pipelining. The hardware will swap in warps (or waves) when one stalls, and by structuring your kernels into phases you can hide memory latency and keep the chips busy.

    Hopper and beyond are where spatial scheduling started mattering, in the form of warp specialization. Nvidia GPUs let you assign different register footprints to different warp groups. When you introduce warp-group scoped MMA for compute and TMA for executing data moves from a single thread you have the ingredients to divide the pipeline between groups. Instead of the same worker doing load -> compute -> store you have different workers exclusively working on different parts of the pipeline. Blackwell made this… much harder. TMEM and UMMA added new operator and memory types, so you now need to schedule movement between shared memory, tensor memory, registers, global memory, and a variety of compute units.

    The problem is: how do you do that?

    To stick with Nvidia for a moment, at the bottom of the stack are barriers. An mbarrier is a phase switch for a specific number of arrivals: one side waits, the other increases the arrival count. When the counter matches the expected number, the phase flips. It’s elegant and straightforward, and easy to get wrong. A classic example is the phase parity bug: if you screw up the wraparound the kernel can work perfectly at first, but then deadlock waiting on the wrong phase.

    Next up, libraries like CUTLASS, and newer ones like ThunderKittens, package the patterns you tend to write. The CUTLASS Pipeline combines buffers and synchronization into a unit and makes it easy to compose common structures. This is where much of the expert-kernel-writing time goes, but that time encodes a lot of hardware-specific behavior. Hopper wants one set of patterns, Blackwell another, and even within a generation there can be differences between variants of the hardware. The more explicit the schedule is for the developer, the more they own the portability problem.

    The subsequent step is to make the schedule less explicit, while still keeping the roles visible. AsyncGraphene’s ARef is a good example of this. An ARef is a reference to asynchronously produced data. Basically, a channel, with synchronization attached. A producer writes, a consumer reads, and both sides can know when the other is done. A compiler can then plan a schedule. Nvidia’s TAWA work does this explicitly for Triton, tagging producers and consumers and lowering to ARefs. TLX on the other hand, as well as systems like PipeThreader, allow defining subtasks in a kernel that a compiler can schedule.

    TileIR and CuTile also enable building an explicit graph, but through focusing on the data itself. Attaching usage information on how data is read or written gives the compiler room to bundle work into tasks and reschedule.

    Getting the graph is the starting point, but then you need to identify what the right schedule actually is. In practice this involves exploring different shapes and combinations to work out which is best. You can either do that explicitly through heuristics and cost models of the hardware, or do it via searching across many different possible schedules to find the ones that work best. Most systems do both.

    But what do we need in a kernel DSL?

    If you are building a DSL for writing kernels, the starting point is to reflect whatever the hardware does. This is not only direct, but also a necessary option because there are always smart people operating at the frontier who have a strong intuition around how to drive the most performance. They’re often targeting very new hardware which is not yet well understood (sometimes, even by the people that made it).

    Beyond that, deciding what else should be on offer means answering three questions:

    1. How do you think about portability?

    Portability doesn’t mean “write one generic kernel and get peak performance everywhere”. But it can mean: what’s the minimal amount I can express to get correctness and a particular level of performance across hardware. Projects like Helion are explicitly operating at a high level to enable rapid research iteration. Regardless of your view on where the line for “high performance” is, you need something to define what “correct” means.

    Having a good concept of a “task” seems to offer the flexibility to schedule statically, temporally or spatially, but there are a lot of edge-cases to consider.

    2. What do agents change?

    Humans are not going to be writing every, or even most, kernels. We have to figure out how much of that portability or performance is a deterministic search, versus how much is agentic loops exploring the space somewhat probabilistically. Agents make generating code massively cheaper. They can create candidates, run profiling on real hardware, test hypotheses and explore options.

    But we also need a sense of where and how the agents fail, particularly when it differs from the patterns of humans. That includes things like verbosity: more lines (generally!) means more bugs. Performance can be both spiky and somewhat subjective; sometimes small changes can reshape the kernel’s performance, and a faster kernel might only be “correct” within specific numerical accuracy bounds.

    3. How do you think about kernel boundaries?

    A lot of discussion focused on GEMMs, which is understandable. But almost all real-world kernel work is across operator boundaries. FlashAttention wasn’t making the matmuls in attention fast, it was fusing them despite a reduction in the middle.

    When we are writing programs we are expressing intent and providing direction. We mix that “what” and the “how”. This reflects a search vs expressiveness divide; search-oriented approaches want you to focus more on the what, expressiveness leans more into the how. The more the units inside kernels can compose across kernel boundaries, the more we can optimize across models and discover patterns automatically2.

    The way I think about compilers is that they encode knowledge (in the form of rules and heuristics) about hardware. The more we can move that out of our heads, or the model’s parametric knowledge, the more we can focus our time or tokens on the parts we don’t yet understand.

    1. Mostly! cp.async was introduced in Ampere and it was very impactful in making temporal pipelining work, as it let the mechanism largely hide HBM latency ↩︎
    2. Whether via compiler, or agent! These problems tend to recurse, so you could have pretty much this whole discussion at the IR level too. ↩︎

  • Cutie Fly

    Cutie Fly

    The FlashAttention 4 paper is out and is fascinating, you should read it! One of the things that Tri called out on Twitter was that the experience of using a Python-based language (CuteDSL) significantly improved the dev loop, not just for him, but for Claude:

    CuTe’s layout algebra plus the quick iteration cycle of a Python DSL are a nice combination. Hence, it’s not too surprising that late last month,AMD dropped FlyDSL, which is, largely, CuteDSL for AMD. This is not a knock on FlyDSL! The project is very open about acknowledging CuTe and its provenance.

    To help navigate, here is a handy translation guide:

    • CuTeDSL: cute.make_layout.
      FlyDSL: flir.make_layout.
    • CuteDSL: cute.composition.
      FlyDSL: flir.composition.
    • CuteDSL: cute.zipped_divide.
      FlyDSL: flir.zipped_divide.
    • CuteDSL: cute.make_tiled_copy_tv.
      FlyDSL: flir.make_tiled_copy_tv.

    FlyDSL also calls out Colfax’s paper from earlier this year: Categorical Foundations for CuTe Layouts. This paper, along with the Integer Set Relations one from Nvidia last year, really started to establish a mathematical formalization of what had been going on in CuTe layouts. This kind of foundation enables verifying the approaches taken in fresh implementations, like FlyDSL’s.

    We can actually go see that, as the whole compiler is open source. This lets you compare the composition_impl in FlyDSL to the diagrammatic version in (section 4.1.3) in the Colfax paper to understand why it works!1

    Given the blistering pace of layout algebra, we shouldn’t be surprised that just a few days after, Cris Cecka of Nvidia dropped a beastly preprint: CuTe Layout Representation and Algebra:

    Colfax Research [19] analyzes CUTE layouts and some operations on them in the context of category theory. In this paper, we intend to provide a more definitive and formal treatment of CUTE concepts and their applications.

    Sometimes with this kind thing it doesn’t matter if you have the idea, it’s often specific implementations of that that end up defining the standard for it. I read this paper as Cecka planting his flag and saying “y’all, this is the real CuTe”. And he cuts no corners.2

    Cecka reframes layout algebra as a theory of loop transformations, showing that the objects you are transforming (Shapes Strides) and the things you are transforming them with (Shapes Strides) are the same.3

    One of the cleverest results of this is in Section 2.3.1. Cecka demonstrates that strides don’t have to be just… regular strides. If your stride is in fact a coordinate then each “step” in the stride moves in the coordinate dimension.

    This is, for example, what you need for TMA on Hopper or Blackwell: you tell it the logical position in the tensor and it figures out the physical address internally, handling tiling, swizzling and bank conflict avoidance in hardware. If you stride over coordinates, you can use exactly the same layout algebra as for your computations.

    Another example was that if a Stride is a bitmask you get something very like Triton’s LinearLayouts!3 That lets you compose layouts with swizzling, using the same composition operators again.

    The paper is full of these interesting, but also practical, results. Cecka gives guidance on optimizations like ‘avoid ranged slicing’; (e.g. a[2:4, 1:3]) as it mixes up tile size (an optimization knob) and thread ID (a runtime index)4, or use layouts to algebraically work out how to auto-vectorize loads and stores rather than hard coding5.

    There is something satisfying about paper on composition that itself pulls together ideas floating around CUTLASS internals, preprints, and alternative implementations, then shows they are all views of the same object. This will help projects like FlyDSL, Triton, and any number of other authoring libraries ground their management of one of the most painful aspects of kernel dev in a way that should make life easier, for everyone.

    1. I think! My understanding of category theory is similar to my understanding of Skibidi Toilet: I get the idea, but I have so many questions. ↩︎
    2. As an example, Cecka provides a wider generalization than the Colfax paper, demonstrating that CuTe layouts are not strictly closed under group composition: you can’t always compose layouts however you want. But! The failures correspond to real errors, which is the kind of restriction you actually want. ↩︎
    3. Actually, you do a bit better: being strictly in F₂ means Linear Layouts are limited to powers of 2, which it turns out is a bit limiting. ↩︎
    4. This makes it harder for compilers to separate static and dynamic elements. CuTe, and Fly, do this in two stages: zipped_divide to tile. then slice by a dynamic bid, allowing the compiler to optimize (e.g. constant-fold) the static tile parameter. ↩︎
    5. By composing the layout with the right-inverse of the other, apparently! Or calling max_common_vector(src_layout, dst_layout) ↩︎

  • TileIR

    There are a lot of things folks do on GPUs (including, sometimes, graphics) so I have an approximately-correct taxonomy of operations to group them in to:

    1. Dense compute: A matmul or a convolution.
    2. Map: Elementwise/pointwise work on each value of a tensor.
    3. Reduce: Process a tensor into fewer dimension, like a sum.
    4. Transforms: Change data structure. Easy ones like transpose, annoying ones like scatter/gather.
    5. Synchronize / Communicate: Move data, or wait for it (copies, collectives, fences/barriers).

    At the moment people are pouring billions of dollars into hardware that primarily does 1. And, at the same time, many of the greatest minds of our generation are attempting to ensure that the hardware spends as much time doing 1 as possible.

    The biggest barrier to doing a lot of dense compute is 5: moving things in and out of memory. Launching kernels, transferring data between host and device (or between devices), moving data between global memory and registers and so on. It’s like there’s a box defined by Data × Footprint 1 × Time, and everyone is trying to keep it as full as possible.

    This involves tradeoffs. You want to mul as many mats as you can, but you only have so much room to store accumulators. Fetching new data from memory also takes a while. You can keep many in-flight fetches around, but each one expands the kernel Footprint, lowering occupancy.

    There are 3 tricks that we can use to help fill up the box by stitching different operations together:

    • Epilogue fusion: take an elementwise op and fuse it onto the end of a dense op, so that when the MMA produces output, the elementwise op can be run while the output data is still in registers. A classic example: fuse the activation after the dense compute in a feed-forward net.
    • Vertical fusion: take two subsequent operations and chain them together to avoid running a loop for one, writing it back, then running a loop for the other2. A classic example is Fused LayerNorm: normally you’d need to add elements, then collect stats for the normalization. You can fuse the two to collect the stats as you add the residual.
    • Horizontal fusion: doing different things over the same data, in parallel. The Q, K, and V projections in a transformer all need the exact same input, so are good candidates to fuse horizontally.

    You rely on the design of the hardware to enable some of this. For example, an epilogue fusion is beneficial because it’s one kernel launch instead of two, and because the work doesn’t need to be written back to global memory, but also because the epilogue can overlap with other work.

    It’s not always obvious how to put these fusions together. Flash Attention was such a breakthrough because it made dense op fusion possible. The naive attention block has a softmax in the middle: Softmax(QK^T / √d) · V. That softmax is a reduction op, which means it needs all of QK^T to be computed first, a pretty large matrix. Tri Dao and colleagues realized that if you used online softmax you could calculate the softmax for subsets of the QK matrix, and avoid materializing the whole thing. They enabled fusing the QK into the softmax and the V in one kernel, at the tile level.

    Tiles are the subsection of a matrix you’re working on at any given time. In a matmul, tiles from both input matrices are loaded and multiplied, to produce an output tile. There’s a useful image of this in the Nvidia blog post on cuTile, Nvidia’s most recent entrant into the the kernel-development landscape. To side-step concerns of plagiarism, I had nanobanana plagiarize it for me:

    Illustration depicting three matrices labeled A, B, and C, showing a looping process with arrows indicating active rows and columns in a computational context.

    cuTile is built on a well-specified intermediate representation called TileIR. There’s an experimental backend for Triton that lowers to TileIR too. While Triton is block-oriented rather than tile-oriented, in practice what you mostly work on in a thread-block is… a tile. TileIR elevates the tile to a first-class concept.

    You can see this by generating the same kernel against the regular backend and the TileIR backend. Triton’s intermediate representation (TTIR) uses pointer arithmetic: generating offsets, computing masks, loading from explicit addresses. Here’s a bit of an inner loop of a matmul. It groups up which data it wants, loads the tiles a and b by pointer, and computes the dot product:

    %offs_m = tt.make_range {end = 128, start = 0} : tensor<128xi32>
    %a_ptrs = tt.expand_dims %offs_m {axis = 1} : tensor<128xi32> -> tensor<128x1xi32>
    %a_ptrs_1 = tt.splat %stride_am : i32 -> tensor<128x1xi32>
    %a_ptrs_2 = arith.muli %a_ptrs, %a_ptrs_1 : tensor<128x1xi32>
    ...
    scf.for %k = %c0 to %num_k step %c1 iter_args(%acc = %zero) -> tensor<128x128xf32>:
      %a = tt.load %a_ptrs, %mask : tensor<128x64x!tt.ptr<f16>>
      %b = tt.load %b_ptrs, %mask : tensor<64x128x!tt.ptr<f16>>
      %acc_new = tt.dot %a, %b, %acc : tensor<128x64xf16> * tensor<64x128xf16> -> tensor<128x128xf32>

    TileIR on the other hand preserves the tile as a semantic object. This snippetis doing exactly the same thing, but this representation elides the pointer math and masking:

    $33: Tile[float32,(128,128)] = typed_const(value=0)
    accumulator.2: Tile[float32,(128,128)] = for k.1 in $36 (with accumulator.0 = $33)
    do
        $50: Tile[float16,(128,64)] = tile_load_token_ordered(array=A, index=($7, k.1), shape=(128,64))
        $64: Tile[float16,(64,128)] = tile_load_token_ordered(array=B, index=(k.1, $10), shape=(64,128))
        $72: Tile[float32,(128,128)] = tile_mma(x=$50, y=$64, acc=accumulator.0)
        continue $72

    This is a nice IR (compact!), but from my perspective the most interesting part is that load function: tile_load_token_ordered.

    The “time” dimension of the Data × Footprint × Time box is the hardest one to manage. Time questions separate performant kernels from slow ones: When to prefetch, how to overlap loads, and so on. Since the advent of warp specialization, the Triton compiler has been exploring pipelining options through heuristics and autotuning, and kernel engineers have been going straight to the hardware with explicit barrier API extensions like TLX3 and Gluon4 .

    TileIR goes a somewhat different route. It assumes an unordered memory model: the order your code is written in does not determine when data is actually available. Instead, each memory operation returns a token and you attach semantics to it: read-only, write-only, read-write and so on.

    By being explicit about memory dependencies you give the compiler freedom to manage the Time dimension. Where accesses don’t overlap the compiler can freely reorder them. Where they do, the token chain tells the compiler exactly what depends on what. The kernel expresses intent; the compiler maps that to the hardware.

    TileIR is (mostly) targeting Blackwell right now, and the experimental backend is still early. The open question is whether we can express this smoothly enough in the syntax of kernels to actually enable taking the same kernel across hardware, or whether we are just adding some syntactic-sugar to avoid when doing hardware-specific tuning.

    That said, the idea feels pretty right? The tile is the unit with which we can express what we actually mean about memory, ordering, and fusion. The CUDA programming model was always about bounded-linearity within a massively parallel framework, and this loosens the bounds that little bit more.

    1. Use of available registers and shared memory, for example ↩︎
    2. This is loop fusion, in compiler terms. There are other things you can do, but this is the big one. ↩︎
    3. Triton Language Extensions, from Meta. As a disclaimer, these are the folks I work with. ↩︎
    4. From OpenAI ↩︎

  • Megacores

    Megacore - Systole as a 80s metal album cover.

    What we do in machine learnings owes a lot to the history of computer graphics. Folks like Kurt Akeley, one of the founders of SGI, identified that 3D graphics have a naturally pipelined structure. You have a high volume of similar operations, such as applying pixel-y soldier textures to a mesh of triangles, and by pipelining them you can find an opportunity for a high degree of parallelism.

    Akeley was one of the drivers of OpenGL, which provided a standard interface to that pipeline, and later worked with Nvidia on CG, a realtime shader language and compiler. Shader languages, as used in Pixar’s RenderMan and other non-realtime 3D use cases, introduced an approach where you could manage lighting programmatically by describing the transforms to each individual element. The shader would be run in parallel across all the geometry or pixels it was addressing.

    With CUDA, Ian Buck and others at Nvidia helped formalize what had been true in the hardware for a while: GPUs were massively parallel processing machines, not just polygon factories. CUDA was part of a move from the supercomputer approach of Single Instruction Multiple Data (SIMD) to Single Instruction Multiple Thread (SIMT). On a Cray or other vector oriented processor you had to pack the work into a vector. CUDA let programmers familiar with CPU threads think in those terms instead. Under the hood, the threads in a warp were executed in lockstep, but they could be masked off to allow for divergence. It was flexible, fast, and attracted the attention of the machine learning community. Because so much of ML is large matmuls, Nvidia bolted on Tensor Cores as specialized co-processors that handled blocks of matrix math efficiently. This combination of performant hardware and flexible software helped make Nvidia the most valuable company in the world, and drive up house prices across the Bay Area.

    But, it transpires, not everyone loved shoveling their margin to Jensen, and they looked for more cost-efficient ways to run ML workloads. The flexibility for threads to branch, pause or switch requires infrastructure and silicon. You need big register files per core, multiple levels of memory to cache, and logic to manage swapping in and out warps.

    If you look at the “do the math” parts of a chip, a CPU probably only spends about 10% of silicon on that, with the rest managing the chaos of running an operating system: branch prediction, caching, data movement. A GPU, in contrast, is a wildly efficient machine, with maybe 30-40% of the silicon dedicated to mathing effectively.

    When Google looked at the problem of running inference at their scale back in the dark ages of 2016 they wanted to spend as much of their budget as possible doing the math, to keep the costs as low as they could. The chip they created, the Tensor Processing Unit (TPU) recently hit its 7th iteration and SemiAnalysis published an extensive breakdown on it: TPU v7 Ironwood, quickly followed up with a deep dive Amazon’s Trainium v3.

    Trainium3 takes a similar approach to Trainium2 and Google’s TPU and builds the chip out of a small number of large NeuronCores. This contrasts with GPU architectures like Nvidia and AMD’s, which instead uses a large number of smaller tensor cores. Large cores are typically better for GenAI workloads since they have less control overhead.

    Dylan and his team are touting these as the first chips to genuinely threaten Nvidia’s moat. The big frontier labs seem interested, with deals and investigation from Anthropic, OpenAI, Meta and others. As the piece repeatedly points out, if you want to understand the dominance of Nvidia you have to focus on the system, and not the microarchitecture. So, of course, I want to talk exclusively about the microarchitecture here.

    TPU, Trainium, as well as other custom approaches like Meta’s MTIA1 lean on an approach called Systolic Arrays. As a recap, Nvidia’s Streaming Multiprocessor (SMs), AMDs compute units ,and so on are cooperative multiprocessors. They access registers, talk to caches and handle the flow of data. Threads can request data if it’s not ready and the hardware warp schedulers will swap in another piece of work to keep the chip humming.

    Systolic arrays are different. The name comes from systole, the phase where your heart pumps blood. In a systolic array, you load your data once and fire it through a grid of Processing Elements (PEs). Each element maths its math then passes the result to its neighbor on the next clock tick.

    This was very much in line with the needs of the original TPU: load a set of model weights up, then pump user requests through as efficiently as possible. TPUv1 only supported int8: it was a low-bit, high-efficiency matmul machine. The data flow needed to be pre-determined: you set it up and make it go, which made it incredibly silicon efficient. You don’t need lots of caches or schedulers, and in fact the original TPU didn’t have any at all!

    The con of course was that you have to get it right! If the data isn’t there to pump in, the whole thing just waits. There is no backup plan to another warp, no other threads. Not only that, but because the systolic arrays are generally a lot bigger (say 256×256 vs the Tensorcores 16×16), you have fewer of them. While an Nvidia GPU might have more than 100 SMs, a Trainium v3 has 8 cores, and a TPU has just 2. Each core is a lot larger, and wasting it gets a lot more expensive.

    Presumably Jeff Dean just programmed these right the first time, but for the rest of Google (and later the world) they spent years building XLA (Accelerated Linear Algebra), a full-graph compiler. In GPU kernel programming the challenge is hiding memory latency and managing register pressure. On a TPU-type approach, there is one massive VMEM that fulfills a similar role as the registers and no memory hierarchy, but you can’t rely on the hardware to swap between jobs. XLA needs to know exactly how the graph works so that it can schedule the right data at the right time.

    TPUs used a VLIW architecture: Very Long Instruction Words. Rather than a traditional instruction set with diverse instructions, VLIW lets you bundle Very Long packages of instructions into single units (kind of a silicon equivalent of German) which execute operations on each of the different units of the core at the same time. This was introduced in TPU v2, and its where the pressure on the compiler really multiplied.

    To draw a GPU analogy, if you think about something like a Relu(AxB+C) you have a graph of operations: AxB -> Result, Result + C -> Result2, Relu(Result2). To optimize that you could use an CUDA graph to compile it into single kernel dispatch and CPU/GPU communication. One step further would be kernel fusion: keep all the intermediate results in registers and write one kernel that avoids the back and forth to higher tier memory. That lets you bundle up even more , but you have to have even higher confidence in the sizes involved to avoid running out of registers,

    VLIW is like parallel kernel fusions: a TPU v2 had 2 matrix units, 2 vector units, 2 scalar units and 2 memory load/store units2.To keep them busy every step the compiler needs to plan ahead enough to give each of them something useful to do. VLIW instructions bundle those ops along with any constants needed into a single instruction. Fusion goes from being an optimization to being a necessity. Once you get it though, you can spend more like 50-60% of your silicon on the part you care most about, and that translates into an excellent total cost of ownership.

    Does this mean we should all be cancelling our Rubin orders and buying TPUs? I mean, no. But there is some nuance. Choosing between flexible streaming processors or efficient systolic megacores feels drastic, but I think it might not matter quite as much as it seems.

    Research still overwhelmingly benefits from flexibility. You are running experiments, solving bottlenecks and debugging. Nvidia tends to be the big lab tool of choice thanks to the flexibility, the depth of tooling and the general CUDA ecosystem3.

    If you are mainly serving a massive model, it’s worth the investment to lock down all the weirdness and optimize it. That’s where the megacore chips have proved their mettle first, with TPU, Inferentia4, MTIA and others all starting on that side of the house.

    Folks like Akeley and Buck realized that when you’re building a chip you’re really building a programming model. Get that right, and the model can long outlast the hardware. Balancing expressivity with performance is the thing that lets a platform win: who best lets researchers and engineers define the future without fighting the silicon.

    What seems to be emerging isn’t quite the SIMT/CUDA architecture: its something around expressing the dataflow of tiles on the critical kernels5, while relying on a compiler to optimize the larger graph and compute.

    Making sure that you have access to the right software might be more important than trying to perfectly identify which hardware platform is the once and future king. But also, look, the world moves fast and if you get a Prime Day deal on Trainium instances, you should probably just take it. The hardware can and will change and it can always be adopted, as the frontier labs are showing. If we keep hunting for the expressivity we need, as OpenGL, CUDA, Triton and others have over the years, we will keep unlocking the possibilities in whatever hardware is available.

    1. Disclosure: I work at Meta and like these chips a lot, though no one would let me anywhere near any chip design, luckily enough ↩︎
    2. Newer versions have others too, like the sparse cores in TPU v6 and v7 which are basically dedicated embedding management processors ↩︎
    3. With the notable exception of Google themselves, though the Jax-XLA-TPU ecosystem is very rich internally ↩︎
    4. Amazon remain undefeated at naming things ↩︎
    5. From system to VMEM on megacore approaches, from SMEM to registers on GPUs ↩︎