Probably Dance

Eight months ago I was producing roughly 90% human written code and 10% AI written code. This has switched surprisingly rapidly and now my code is probably 90% AI. So what do I do all day?

I’ll go through a change that I made to a matrix-multiply kernel, for which I’ll sadly have to be a bit vague. Since GPUs are now giant matrix-multiply chips, where north of 96% of the flops are in the tensor cores (latest Nvidia GPUs have 2250 tflops in bfloat16 matmuls, compared to 75 tflops for everything that’s not a matmul), you’d think that they’d make it easy to use all those flops. But no, matrix multiply kernels are giant crazy beasts that are incredibly tricky to get right. Some quick googling finds this explanation on Nvidia hardware and this one on AMD hardware. Just open those two and scroll down both to get a feeling for how much work is involved. (and yes, these implement the simple O(n^3) matmul loop where you iterate along repeatedly and multiply and add all the numbers)

The particular kernel I was optimizing was using Cluster Launch Control (CLC), a complexity that’s covered in neither of the blog posts linked above, and it had some bad interactions with some particular inputs. The resulting change was 95% written by AI. So I’ll just go through what I did. AI asks are bolded:

0. Find and Understand the Problem

Before I started I had to find and understand the problem. Some coworkers had talked about the matmuls taking too long, and nsys showed that the tensor-cores were oddly idle while that matmul kernel was running, and then it took some more targeted benchmarking to confirm that slightly different inputs give big speedups, which convinced me that something silly must be going on. My coworkers actually had a really good theory, which brings me to the actual start of the work on this task:

1. Find the Code

The code was in a library that I was vaguely familiar with, but since matmul kernels are giant scary beasts, it was a bit hard to navigate. So before I even tried, I asked an AI to find the relevant code for me. I explain what the problem is and I explain our theory. Then I also start looking myself but the AI finds the relevant code first and even points me at exactly the right lines to look at. It also confirms that our theory for the problem sounds plausible.

2. Understand the Code

Next I decide to understand the code myself, so I intentionally don’t ask the AI anything more. As I look around I begin to think that our theory actually isn’t right. I mean it was partially right, but the code very much wasn’t doing the slow thing that we thought it was doing. It was doing something slower: returning out of CLC mode and giving control back to the hardware scheduler.

This is a little silly, so after I am convinced of it, I ask the AI to confirm, just to get a second opinion. It agrees with me (but I take that with a grain of salt, because it also agreed with the initial theory).

3. Monkey-patch the Code

Since neural network training happens in Python, you fix things by monkey-patching. At least initially this is the fastest way to iterate on this code. But monkey-patching has a bit of boiler-plate that’s easy to get wrong, so I ask the AI to set it up for me.

Then I try to fix the bad behavior, but my fix results in a deadlock. I look at the surrounding code again, but realize I don’t understand CLC enough (this is my first interaction with it, and I jumped straight into a complicated kernel where it’s part of other pipelining), so I could either spend the time to understand the surrounding code better, or I could just ask the AI to have a look.

4. Fix the Code

I ask the AI how it would fix the issue. It immediately tells me why my fix didn’t work: I was violating some invariant in the code where the pipelining logic was also used to reason about what state the CLC is in, and my change required keeping the CLC state separately. The AI helpfully tells me that there is one unused slot of shared memory that is reserved for unknown reasons and is entirely unused by the kernel, so it suggests just using that memory to pull out the state that now has to be tracked separately.

Understanding this would have taken me hours, maybe even a full day, so I tell the AI to just go ahead and implement the fix that it has in mind. The fix immediately works and makes the code much faster for the problematic inputs.

5. Reviewing the Code

At this point I spend some time to understand the code and the change better. It looks reasonable to me, but I also ask a second AI to review the work of the first AI. The second AI finds a problem: The matmul kernel was clobbering over some other state that I hadn’t fully reasoned through, and this was not a problem when the thread-block was shutting down and giving control back to the hardware, but with CLC it means the clobbered state gets reused. I look at it and I’m not sure it’s a real problem. I think this would only happen if we’re iterating over padding tokens, and for those the clobbered state is fine. I ask my first AI and it seems worried, because the documentation doesn’t say that CLC guarantees ordering, so if we get a valid token after a padding token for some reason, the clobbered state would lead to problems. It also suggests a fix. I think about the fix, which seems too complicated. After some staring at the code I suggest a simpler fix, which the AI thinks should work, so I ask it to implement the simpler fix.

I also add an assert to check if we ever get valid tokens after padding tokens. Mostly just because it sounds strange to me if CLC behaves like this, so I’m curious to find out.

6. Benchmarking the Code

I have verified that the code is faster for the problematic inputs, but a coworker suggests drawing plots with the speedup across various different shapes. So I ask an AI to write a simple benchmarking script for me. (AI was actually already great for these kinds of throwaway scripts a year ago)

The benchmarking script does not reproduce the initial issue at all. The inputs look plausible to me, but I decide to just ask the AI to reproduce exactly the same inputs that we’d get in prod. (this is not particularly challenging, but manually tracing through the code to make sure you get this exactly right is tedious, so better to just ask the AI)

After that I can produce the plots that I expect. Interestingly it suggests that even with the fix, we’re still far below the speed that this same matmul kernel achieves for other inputs.

7. Asking for More Ideas

So I ask my first AI if it has any other ideas to speed up this code. I had one idea already, but I wanted to see what it says. It does suggest my idea, but it also has four other ideas. One of which is a tiny change that sounds really promising. So I ask it to implement that one. And then I also ask it to implement my idea. (this actually took a few steps because I was worried the change would be too big, but the AI comes up with a way to implement my idea with only minor changes)

Both of those ideas speed up the code more. The code is slightly faster with the idea that I had, but it’s also more complicated, so I actually decide to just do the tiny change that the AI suggested, together with the initial fix, and call the kernel ‘fast enough’ there.

8. More Review

At this point I’m getting all of this code ready for review for a coworker, when our code review tooling finds a problem in my new code. Before release I had replaced the assert that I added in step 5 with a print statement, because I don’t want code to crash just because I was curious if something ever happens (it never happened in all my testing), but the review bot points out that this should really be turned back into an assert because we still assume that CLC work arrives in order. This is weird to me because I had extensive conversations about this with two other AIs before, making sure that we don’t assume that, but the new review bot found an edge case where even the original unmodified code relied on this assumption.

9. Simplification

At this point I have some very careful conversations with my various AI bots to try to get to the bottom of exactly what the original code was already assuming and which new assumptions our patches introduce (I also have to read the code myself and understand it, because I don’t trust the AIs to get this completely right). And eventually I make the judgement call that we’ll just assume that CLC work arrives in order, which can simplify the change because we don’t need to worry about the clobbered state from earlier. I ask the AI to simplify and it does an OK job, but since this is the final code, I step in and simplify a bit further still.

Constant Conversation

The whole time I’m in constant conversation with AIs. I’m going through way more tokens per day than I did even a few months ago. It’s also not just one AI, but multiple different ones. Answering questions about the code, providing ideas, writing new code, reviewing.

Am I Faster?

Overall this took a couple days. The initial change took just under a day, where on my own I probably would have taken two or three days, because there were some genuinely tricky interactions with the pipelining, and the AI found a good trick to use some unused shared memory to get out of that easily.

On the other hand I was also distracted because AI made me question whether we can assume that CLC work arrives in order. I would have not doubted that (even if it’s not guaranteed by documentation) and would have saved some time without that paranoia.

AI also just allowed me to jump straight into the problem. The initial “find me the relevant code” pointed me directly at the right lines, and I was able to make changes without having to spend the time to warm up with it.

So overall I probably did in four days what would have taken me five days before. The initial speed up of doing the first implementation in less than a day instead of two or three days does not translate into a similarly dramatic overall speedup, mainly because there is a bunch of benchmarking to do and reviewing, and understanding of the actual code. For many years now I have felt that “writing code” is not usually my bottleneck (I’m not too far from the “10 lines of code per day” quoted in the Mythical Man Month). That is the part that AI can speed up dramatically. It also helps in other parts, but with smaller speed-ups.

Am I Better?

I think my work had a higher quality according to some metrics: I tried more optimizations and had nice plots for how those optimizations behaved for different inputs. I might have done that on my own, too, but it’s certainly easier to just ask an AI. I actually think it’s somewhat likely that I would have arrived at the same final code change without AI, but the AI allowed me to explore more of the space before I ended up there.

The main downside is that I have less understanding now than I would have otherwise had. I did end up having to understand the kernel quite well and I can now navigate that library easily, because in the end I had to make decisions about competing claims by the different AIs. But even though my understanding of the code is much higher than it was at the start, it is still not where it would have been if I had to do this entirely on my own.

I also think this change was a lot more careful with AI than it would have been without. It’s kind of humbling how many bugs the latest top models find in any code that I try to release. It now makes a lot of sense to me how all software is subtly broken, because there are broken edge cases in nearly all my changes (things like “there is a memory leak here. It was actually there before, but we didn’t go down that code path before your change.”). I think even if we somehow went back to writing code manually and only got to keep AI code review, software would be a lot more robust in a few years. (sadly I’m not sure if that will happen if AI also writes the code)

Do I Enjoy This?

The moment of “my first change didn’t work because I didn’t take the time to fully understand the existing code” gave me a similar feeling to having access to a cheat code in a video game: I could now do the hard work of earning my progress, or I could use the cheat code (the AI) to solve the problem.

The dissatisfaction of using the cheat code is similar to a video game. But in a work environment it’s hard to justify spending the extra days out of personal pride. There’s plenty more work to do after I’m finished with this. My job isn’t to write code, my job is to fix problems and to ship new features, and “write code” was just the way I got that done before.

I do get the new satisfaction of finishing more code. E.g. I can finish side projects again despite having less time to program (1, 2). And while I shipped the above code, I also shipped two small side projects at work. (a bug fix to a shared library, and a separate benchmarking utility that I only used a little here, so didn’t even mention above) Small changes are great now because I can just ask an AI to give it a try and check back 30 minutes later to see how it turned out. If the changes actually end up as small as expected, these often ship.

Do I Still Need to Be In the Loop?

OK so why can’t an AI just do all the things that I was doing? Have a supervisor AI that coordinates a planner AI, a writer AI and a reviewer AI? I mean here is a good talk where they wrote a “deep research” bot that works exactly like this, but AI is still not quite there when it comes to outputting something that you have to live with for a long time and may want to tweak yourself.

Two of the reasons why AI works so well for code is that 1. you can split the work into contained components, and 2. you can go in and tweak the last details. To illustrate the power of these two points, think of other tasks where these are not true, like asking the AI to generate a video for you. But for them to be true, code has to still be tight and readable. My role in this was mostly to get to the bottom of what’s actually needed, make a judgement call for picking a good spot on the “optimization vs complexity” trade-off and to then ask the AI to simplify and to then simplify further myself.

Will AI be able to do even that in a year? Plausibly. Currently the topic in the news is how AI can get pretty unhinged in its pursuit of goals, which I have also seen (at a smaller scale) before, so I’d keep a human in the loop for a while longer.