Auto-research with codex: How I achieved a 232x Faster Kernel
sankalp.bearblog.devPersonally, I believe that LLMs should be treated like an advanced version of Prolog or linear programming: you give the constraints, you have a way of verifying correctness, and you give it a clear goal. If the LLM can verify itself and course-correct you can basically leave it on autopilot
Doing the same, re-implementing a lot of LLM/diffusion models in Rust+CUDA for my own usage, usually the initial implementation takes 1-2 days (of 100% autonomous work) then I put an agent to optimize the implementation which tends to get close to SOTA performance within another day or two.
As long as you can point the agent at "This is the correct baseline, make sure any optimizations still pass this", seemingly you can leave them and they come back after N hours with a faster program that just works.
with opus 5, I'm finding that asking itself to write such a loop skill that does benchmark - profile - verify - research - improve, keep logs for each iteration, etc. (just like autoresearch), and then using the skill, works very nicely.
I then asked claude to compare the C#/.NET implementation in the library with the C++ version, and it quickly identified that the C# library was missing a couple of fairly cheap optimisations that were present in the C++ version.
If I can help get a PR merged, then it'll be by far the biggest impact of any work I've ever done.
I also compared the Rust version, it had this specific optimisation. The far more popular Tokio/Prost library did not.
Given appropriate guardrails, LLMs are impossibly fast at iterating to find root causes and specific performance bottlenecks.
Could you have manually profiled and compared the execution paths? Sure. Could you have translated the C++ optimizations to C#? Sure. But in such an obvious case, the LLM managed itself.
I’ve been using fable recently not to write code, but to write briefing notes on projects. “Read and compare these three libraries. What are the unique differences in their implementations? Why were those ideas chosen? What are the pros and cons of those designs? What ideas would you copy into a new project? Write a report on your findings.”
Even the cheap LLMs are great in doing the awful crud work in the beginning: finding offsets, firmware update file structures, brute forcing checksums, etc.
It still produces a lot of crap in the later steps (understanding the implementation itself) but I'm happy doing this stuff myself then.
I've had success here by adding a phase called "grounding" that attempts to verify its "understanding" by creating tests that modify the running executable to ensure its made the right inference.
Is this variable really MARIO_X? Change it and see if Mario moves. Etc.
As an example in Donkey Kong - the system had trouble deciding if an array controlled barrels or fireballs. There was conflicting evidence.
After many trips through the loop - it realized it does BOTH, depending on which level you're on.
So the "understanding" grows with each iteration.
In Time Pilot - there are three routines that are called constantly from inside the main loop. Each routine computes the checksum of the other routine's code to see if it's been modified. If so it jumps into random junk data.
There are other less exotic routines that make sure the copyright string hasn't been modified, etc.
https://github.com/qarl/arcade-js/blob/main/games/timeplt/id...
Fascinating.
Watching claude and codex play winquake and age of empires, and debug support for Firefox 52 has been wild.
I do wonder if auto-research would have reached something similar, it did take a significant amount of steering from me to get it to the point where it was working realtime.
For anyone interested the ffmpeg is at https://github.com/poizan42/jellyfin-rpi-ffmpeg and a shim for using it with stock jellyfin at https://github.com/poizan42/jellyfin-rpi-ffmpeg-shim
This!
People keep measuring how good AI is by one-shotting a problem, but I e found that a back and fourth via the Superpowers SKILL.md is how the model builds these constraints… then when it has all the holes, it gets to work
Which I believe was the word intended.
yes, it gives labs edge and leads to self-recursive improvement loops.
also i was myself able to finish 7th in a later competition with 2-3 other approaches which are variants of the method discussed in this blog.
in general, having a harness as thin as possible with some problem specific instructions while controlling for context rot is the key.
point i am trying to make is there are a lot of optimisation surface areas possible.
You can't use Claude for this sort of thing if the goal is to make better AI systems. Anthropic finetunes Claude to dissuade people and the agent from using research that actually works. Anything that they use internally in their own models is poisoned, to protect their moat.
By proxy, that also means any openweights model that was distilled from Claude is equally useless for this purpose.
Thankfully, I don't believe OpenAI does this - they are far more honest and seem to care about their reputation. Anthropic is evil though.
Bro. Sam Altman?
I wonder if there is extra effort put into this because they are useful for the researchers working on the models or just a sub-domain that language models are a great fit for and humans have trouble with?
Off-topic, but imagine us collectively being okay with (or powerless to do anything about) this sentiment about any other software service provided like two years ago.
How the times have changed…
Obviously it's still hard, but the point is that, by construction (cause like, NVIDIA literally releases primitives like this, and/or people like TK build slightly higher-level primitives over the base hardware primitives), if you learn the complicated language, you can get really good results, and on some level you "know" you're right by construction even before you go to the actual empirical tests (since you're operating over a higher-level "language", and not arbitrary byte accesses).
Honestly a lot of interfaces and frameworks you could argue are like that, so it's not really a point for GPU kernels relative to other things. But maybe a hint as to what I personally think is important in the AI era - finding the right cuts, the right high leverage abstractions, as otherwise AI is going to produce spaghetti nonsense.
Because pre-LLMs humans partially "autogenerated" kernels through hyperparameter search and in some sense eating the code complexity in return for performance, and thus built tools for the same automatic verifiability that is useful for LLMs.
In some other tasks, we never built the same level of automatic verifiability since the level of automation in creation being much lower meant it's not giving you as much of a marginal benefit. We prefer code readability and simplicity and such in say, web services, because, say, the database IO time is going to dominate. Here getting an LLM to write a cromulent C# web service is more difficult since it's not easy to automatically verify whether code is cromulent or not. So if you put up LLMs to it, you end up with slop (which works).
OTOH, in kernel design, you give it access to every perf counter, every observable possible and have it optimise all of them. And all are verifiable/hill-climbable and you generally don't give a crap if the code is readable or reusable.
The existing models are surprisingly bad at it.
Or guessing colors rather than sampling from the image or pulling from figma is another stupid thing they do constantly.
1) it's the only way to get something out of models (or people for that matter) that they don't know yet.
2) it's harder to do with an LLM than without. Not easier.
3) and when you fuck it up, half the time the LLM (or other ML technique) makes a fool out of you and you spent $1000 to find the quickest way to get a robot leg on the ground is just to crash it into the ground.
IMO, LLMs will be a dead end to anything close to AGI because of this and hallucinations.
We're missing something in the mix, which I suspect is some kind of advanced JEPA model.
Humans do it ignorantly.
The LLMs will improve while average human IQ in the west dips closer and closer to the 80s on the global scale.
the code the LLM writes will be read and maintained and developed further by LLMs. so it doesn't really matter what it produces as long as all the tests are green and it achieves exactly what you want it to achieve.
Also QR is a primitive for operations like finding eigenvalues, and I don't think Cholesky can be used there.
welcome! check out my featured section
- accelerated launch of our new backends like polars, including a new lazy mode & planner, which are fundamentally new paths
- while we initially aimed for top GPU benchmark scores, we now also maintain top CPU scores too!
Long-term, more interesting to me is this opens rethinking what it means to be a query engine. Right now we are making it the fastest in general, especially on workloads from our own use, major industry benchmarks, and our users. At the same time, similar to jit and multistage computing, we're looking at new ahead-of-time optimization techniques users can do that are more interesting than plugging in custom indexes. Essentially, if our agents can do fast specializations, there should be safe hooks that we can expose to our user's agents too!
The only solutions that did not break when tested with OOD shapes were made by experts who know a lot about GPU programming and that did not create 25k lines of CUDA but followed and adjusted their solution in reasonable bounds.
The takeaway from this is that these approaches will always solve for specificity, but it's a much harder task to steer the model into making general solutions. So if you're an inference provider for some specific model shape, fantastic, go for it. If you are a maintainer of a open-source library, this is not useful.
The goal is not to create good, general or maintainable code. The only goal is to produce the fastest code.
i found out that as i learnt more domain wise, i was (obviously) able to steer better. doing a re-write can also remove lots of slop and context rot (and subsequently make it easier for both human and LLM to make solution more numerically stable, less reward hackish)
What I feel is lacking with the solutions that have for example agent-generated and tuned GPU kernels is that the use-cases for them are unclear. If you are a researcher on second-order optimizers, you probably want to be able to handle variable input shapes to experiment, you also might want something readable to understand intermediate steps and perhaps build on that. If you are a neolab running massive training runs for 80% of your VC funding, you need to know that every line in your training code is bitwise identical to the theory/reference because a divergent run from some LLM-generated numerical bug will set you back or bankrupt you, so you can't just plop in a random kernel even if it promises good performance.
So who is the agentically-looped end result for? Except for Openai and Anthropic of course who sold the tools.
This has nothing to do with LLMs. There have always been plenty of solutions that are more capable, but untrusted.
Your real question seems to be whether you can prioritize better. What are your project goals? If you have no say in or insight into those goals, you have an even bigger problem. What are you even working on?
For sure, you don't want to be maintaining your dependencies. LLMs make it trivial to rack up insane amounts of technical debt. Why is that appealing to anyone? How is it meaningfully different from the idiots wanting to fork everything on github and copypasta their way to startup success over a decade ago?
Also "will always" is way overconfident. A year ago nothing close to this existed at all. Next year it will fill yet a different role.
1) Realize that your problem isn't as general as it was
2) Realize that your problem has a better memory model
3) Realize that your problem can be parallelized further
Hyperoptimizing usually falsely leans on 1
Since most problems see fairly little optimization, that's still a big win most of the time.
- Have an agent generate unit tests until it gets to 100% path (not just statement) coverage, with every numerical test asserting checks against golden values to prevent regressions
- Let it rip on a performance improvement loop, for the widest E2E representative test case you have. Have it generate flamegraphs along the way so you can check in and steer it as necessary.
- Optionally allow for 1 ULP changes in output values so that it doesn't kill itself getting bit-exact results.
- Have it flag correctness errors as it goes, since your code probably isn't bug free.
This is also how I've done language ports from python to rust, and having the ironclad test coverage protects you from drifting.