Academia is for Ambition — Alex Zhang, MIT

Latent Space: The AI Engineer Podcast

Alex Zhang discusses his journey from GPU kernel optimization at GPU Mode to pioneering Recursive Language Models (RLMs) as a PhD student at MIT. He argues that

Key takeaways

  • RLMs are defined as harness designs where the only tool is code, allowing models to offload context to a file system and use subagents for recursive problem-solving.
  • GPU Mode evolved from a CUDA learning Discord into a competitive programming ecosystem that now heavily features AI-generated kernel solutions, though human intuition remains critical for stability.

Main topics

  • Recursive Language Models (RLMs)
  • GPU Kernel Optimization & Benchmarking

Notable quotes

"An RLM is basically just a harness design where the only tool in the harness is code, which is this programmatic subagent calling thing where it has the option to call itself as a tool."

Conclusion

The episode underscores a critical shift in AI development: maximizing model capability requires rethinking

Transcript preview

Speaker 2 (0:03) All right, we're here in the studio with Alex Zhang, I guess most famously of RLMs, but you have a few other affiliations. Welcome to the show. Yeah, thank you for having me. Yeah, I guess GPU mode as well? Yes, GPU mode as well. You were shepherded in by Mark Serafin, not everyone gets that kind of welcome. Yeah, I'm very, very close to all the people in Speaker 3 (0:22) GPU mode. Speaker 3 (0:24) often end up working together in various capacities, like even beyond just GPU mode itself. Yeah. Speaker 2 (0:29) Can we explain that? So people who are not that close don't know about this. It's a Discord. It used to be focused on, I guess, CUDA mode and then generalized a little bit. It was started by Mark. Yep. It was basically like, to me, it's like the hiring pipeline of the PyTorch team. And then he left PyTorch. Yep, Speaker 3 (0:46) yep, yep, yeah. Speaker 2 (0:49) So it Speaker 3 (0:49) used to be, I think, It started actually around when I was in college in like 2023. I think it was started by Mark Andreas and Jeremy Howard. The original premise was just like it was a GPU or is a discord dedicated to learning how to write GPU kernels. And they had like lectures. That was basically the extent of it. And I got interested in it because I was writing GPU kernels. It was actually out of like pure chance. I was interning at Snapchat at the time. And... Yeah, I was very bored with Rexus. So they had a project where, like, they were interested in writing... It was this paper called Infinite Attention. It was like a Google paper. Yes, we've covered it on Paper Club. Yes, yeah. So I was interested in whether or not you could write specialized kernels for it at Snapchat. It didn't... Nothing really came of it, but I joined GPU Mode at the time. It was called CUDA Mode. I think for like legal reasons or something, they changed the name. But I met Mark. I met Matei. I met a bunch of other people that were very involved in the community. And then Mark had pitched this idea called Popcorn, which was now what you see as the leaderboard today. But the general idea was like, I think all of us had this like intuition that GPU programming is like very similar to... if you guys have done like competitive programming, it's not, I don't mean to say like they're transferable skills. Speaker 2 (2:17) You have constraints, Speaker 3 (2:17) you Speaker 2 (2:17) code golf a little bit. Yeah, Speaker 3 (2:19) yeah, yeah. And there's like, there's actually a surprisingly small space of optimizations that people do. And there's actually not that many kernels per se that people are interested in optimizing. And so we kind of had this thought that like, if you had enough data, like in the same way that code forces, there's millions of problems. If you could do this with GPU code, like you could scale and automate. kind of GPU kernel development, which for researchers is a huge deal. Because I think one of the bigger bottlenecks, like if you look like Mamba, for example, like they released the paper with kernels because otherwise, like you can't really use it in any meaningful way, you know, and not everyone has like a tree down on their team. So we're very interested in this kernel bench kind of spawn from that too, of like, can we get LLMs to automate GPU kernel code? And I think that was like, it was a very, very fun time. It was like between college and my PhD. And yeah, I had a really pleasant time doing stuff with GPU mode. Now I kind of just help with the lectures sometimes. I'm not as involved. And I think in general, like we don't have as many competitions as we used to. But yeah, I still keep in touch a lot with everyone there. Speaker 2 (3:27) Is there a friendly rivalry? Because I think the previous community that used to do this was like MLSYS, MLPerf. Is there a friendly rivalry? Is this like just new generation MLPerf or what's going on? Speaker 3 (3:37) The nice thing about GPU mode is that it is also a community in the sense that like a lot of the lectures are very easy enough for a beginner to follow and ask questions and things like that. And like the competitions are like somewhat not secondary, but like you can participate in them to learn. I think with a lot of like MLSS, MLPerf kind of benchmarks like. for the most part, like only serious labs and companies participate, like seriously in them, at least. That was my understanding of it. I could be wrong. But I think also beyond GPU mode now, one thing that has been really exciting is there's a lot more websites and like people that work on hosting competitions. Like I think there's this... I think there's this website called like LeetGPU or something, and it's like LeetCode for GPU problems. There's like other ones too that I've seen. Like we've seen, like there's many that have kind of spawned and like talked on GPU mode. And like it's very exciting in the sense that I think GPU programming used to be super niche. Like when I was interested in it, and the only reason I got interested in it was Tridal gave a talk at Princeton because he was applying for faculty. I mean, he is faculty there now. But I listened to his talk on Flash Tension in like 2023, and I was like, wow, this is like the coolest thing ever. And I was like, this is like, this is what everyone should be working on. I guess like VLM and stuff had come out too, and it was like, oh, you know, we should be writing kernels. But now it's like, you know, everyone writes kernels. Like everyone, it's, I think it's actually almost saturated in some sense as a field. Speaker 1 (5:11) Any interesting takes for people that want to get into it? So I think one of the biggest news is GPT 5.6. wrote more efficient kernels. So Terra and Luna could be 80 % cheaper. And then we've seen other competitions where people are like setting records and they're like, we're doing some auto research loop. And these are people that don't have a background in any kernel writing, right? Yeah. Speaker 3 (5:36) So even on the GPU mode leaderboard, if you look at like a lot of the recent problems. Almost all the solutions are AI generated. However, you'll notice on the leaderboard, so there's this guy named Gaunerst, who's like a very, very like regular member of GPU mode. We've always known for a long time that he's like a super, super cracked like GPU kernel writer. One thing we discovered on this leaderboard is like almost every, he also used AI to help him with these solutions, but for the most part, like he helped prompt and move it in certain directions. We found that like his kernel was like, basically the only one in like the top 10 that was actually stable in like actual like end-to-end systems. And it does bring into question like, it's not like, I mean, GPU kernels have a verification problem. Like we've kind of known this. It's been a problem since kernel bench was released. Like there's a lot of reward hacking that goes on. But you also notice like the lines of code is a lot smaller. Yeah, Speaker 1 (6:33) I was going to ask, Speaker 3 (6:33) is Speaker 1 (6:33) that noticeable or is it just... Yeah, no, it's definitely Speaker 3 (6:36) like very important. And I think like... It's really interesting that still there's a lot of alpha in being good at brain GPU kernels. There definitely is. Yeah. I think like, and this applies to a lot of AI systems as well. Like I think, you know, even with the most recent like math proofs and stuff, like it doesn't necessarily mean mathematicians are obsolete. I mean, these companies still hire mathematicians like to do, whether it be like data labeling work or even just like steering the models to solve problems. Like there is still a lot of. alpha and being like knowledgeable in these things. Is Speaker 2 (7:11) it just knowledge or is it also there's just more planning and is there an emergent style of planning that works better? I think Speaker 3 (7:22) it's a mix of maybe this is what you mean like intuition for how to solve the problems. Like for Speaker 2 (7:29) example, I always diagram my code. Yeah. Right. And then like if there's a part of the diagram I don't understand. I work until I understand it. Otherwise, it's not allowed. Yeah, Speaker 3 (7:37) yeah. So I think it's like, it's a mix of those things of like the people who work, like the people who know how to look at these problems and how to solve them, like also know how to use AI to do them. Because like you're acting as a very strong verifier. Like if you are knowledge or if you know what to do and you're also, like I think the thing that we've kind of discovered with all these agent swarms and things like this is like when you throw enough compute at a problem, you can sufficiently explore solutions to that problem. But oftentimes, maybe you can burn 100 billion or a trillion tokens on something. But if you bring in someone who knows something about the problem, they can uncover something for the model that would erase that one trillion token spend. I mean, it's not super clear what exactly the trends are here. But I think like... There are so many problems in the wild still right now that we want to solve. And like we can't afford to just always, you know, throw as much compute as possible at it. Like there is still an efficiency aspect of all of these things that Speaker 2 (8:40) is super, super important. Is there like a theoretical right answer that you just calculate based on physics and then you just get close to the physics? Yes. Speaker 3 (8:48) So for GPU kernels, you can compute. It's actually not that easy to compute sometimes, like, depending on how complex the problem is. Like, for matrix multiplication, it's very easy to compute. This, like, speed of light kind of estimate of what the fastest kernel can be. And, like, this is also assuming, like, you know, maybe all of your data starts on the CPU or maybe it starts in DRAM on the GPU, et cetera. Like, this changes. these numbers slightly. The transfers and Speaker 2 (9:18) all these things. Yeah. Speaker 3 (9:19) But I will say like, it's not clear though, like in a lot of cases, if it's even possible to hit this theoretical number, if that makes sense. Like this is assuming like perfect overlapping and transfer of data. And like there's maybe some bottleneck that you can't get around. But often the kernels are not even close, like that we write are not nearly close enough to this number to be like meaningful at all. Speaker 2 (9:43) Yeah. And is it speed that matters? Do you also care about, obviously, memory, which feeds into speed? Do you care about power consumption? So one of our top pods of the year was Jeff Dean, who was like, actually, I just tracked the microjoules or the nanojoules, picojoules. Yeah, Speaker 3 (9:59) it's often picojoules. Picojoules. Do you care about that? So I don't. I guess maybe I'm not as... But Speaker 2 (10:06) everything here is speed, right? Everything here is speed. Nobody's counting picojoules. Speaker 3 (10:08) But there's a caveat here, which is, I think, like... There is speed in the context of a single kernel and there's speed in the context of a larger problem, like maybe the N10 model. Because like one thing to consider, and this is why it's important to talk about what speed of light is referring to. Because in these cases for the kernels, like we always start with everything in like HBM, for example, right? But you can imagine that like... an end-to-end kernel, like an end-to-end model, what you might want to do between two layers is like you might sacrifice the speed of the first operation to keep things in the cache for the second operation. And like these are things that like you can't really get out of isolation with like these kinds of kernels. And, you know, people call this like the fusion or like the kernel fusion problem. Yeah, or like mega kernel stuff. And it generally only applies like when you are like memory bound in most cases. But this is something that like also there is this question of like, as these models get better, like should we just be generating like megacernals? Is that like what we want? Speaker 1 (11:19) What's your take? I think that Speaker 3 (11:22) this is really difficult because you need the data to do this. And I think like, I have yet to see an example in the wild of like, we... bootstrap the ability to solve a very difficult class of problems without any examples. And I think like the other reason why I think maybe this isn't that interesting is that at the level of an individual kernel, A, like they're not that, they're not as complex, but B, you're somewhat confident that there's not as much structure in a single kernel. But like in a mega kernel, like I would be more inclined to believe that like a compiler would be better. here. Like, some compiler over, like, higher level ops makes sense. Because in general, like, actually, I think mega kernels are very, like, the pieces are very composable of, like, the individual kernels. There's some areas where you might want to do, like, weird fusions and everything. But in general, I think these are cases that, like, a compiler can probably handle. And there is a company that's working on this, from what I understand, that has given some talks on GPU mode as well. I want to basically cluster all the Speaker 2 (12:29) GPU mode discussions here because obviously there's other parts that we need to move on to. I think there Speaker 1 (12:33) is something to plug. You guys do host a lot of really good lectures. They're all on YouTube. People can follow. And you lead quite a bit of it. I Speaker 3 (12:40) used to. Sometimes I still do. I think they're mostly Mark. Mark is the one who usually does them. Matei does sometimes as well. But yeah, I highly recommend them. They are extremely good resources. Like I think it's kind of crazy how much people share. on there. So, yeah. Because, like, if you're Speaker 2 (13:00) there, like, you're very, very, like, you're exactly the right audience, you know? Yes, exactly. Like, this isn't going to reach the mainstream. And there's a lot Speaker 3 (13:05) of, like, introductory material as well that we've put on that I think is useful for people. Speaker 2 (13:10) Colonel Bench was kind of influential. I just want to see, like, you know, that was last year. Yep, yep. What other ongoing work? Do you want to shout out that people should pay attention to because obviously you're involved in this field? Yeah, I Speaker 3 (13:21) will give maybe the background story of like I am actually involved in a lot of benchmarks where I used to be maybe prior to my PhD. It started because I was at Princeton. I worked with the SWE bench team Speaker 2 (13:32) there. John, Speaker 3 (13:33) Carlos, Ophir. They're all great. I love them. Speaker 2 (13:37) There's basically this, I think people don't understand how much benchmarks come from the same group at Princeton. Speaker 3 (13:44) It is crazy. Do you know Shunyoo? Speaker 2 (13:45) Yes. We had him on the pod before. Now he's like running 10 seconds. Yeah, now he's like a superstar. When Speaker 3 (13:51) I met him, so he was advising my friend Michael. Michael Tang, who is now at Anthropic, but they work together a lot. We were like the two undergrads in Karthik's lab. I, and then some others joined later as well. But yeah, Shen Yu is great. I did not know he was like such a superstar until like later on, like after I left, but. Speaker 2 (14:12) Yeah. I mean like, like, okay. So there are very few PhD students. Like yours is, yours is like the next one. Like once a year we feature someone like who is like basically entire PhD has been like on target. There's not that many of them. Shen Yu was like clearly one of them. And, you know, Jack Morris is another one. And like, you know, we talked before the show, we talked about research tastes, right? Like somehow some grad students just have a very blessed career where like, yeah, mostly like, yep, yep, yep, yep, yep. This is like going to stick around relevant. Everyone should know this. And then others just nothing. I Speaker 3 (14:44) think this is also true of like even people within like industry labs as well. I think it's just like grad students are a lot more visible. So you just see, like, you see, like, you know, there are some people who really, like, get lucky in, like, or, I mean, it's a mix of being lucky and also being very smart and things like that. I think, like, with research tastes as well, like, I think it gets developed through opportunities, at least in my case. Like, I was very fortunate to have, like, taken the path that I took, like, working at Princeton and then, like, later, like, finding my advisor, like, Omar at MIT. Like, he's a fantastic advisor. I will say, though, I find that the most successful research from grad students or like in academia comes when people care about problems that maybe like most people in industry are not looking at. I think this is the issue that like a lot of grad students work on things that benefit like that look good to an industry lab. Like, for example, they'll work on some like some benchmark. It's really popular now. I mean, I think benchmarks in 2023 were a very different story than benchmarks now. There are a lot of people that work on like harnesses and like meta harnesses and like specific harnesses for XYZ task. And when you really think about it, the reason someone would work on this is like maybe there's like a clear goal shaped around the models that we have today of like, this is what I want to see. But, like, I'll give the, like, RLM, like, the recursive language model paper as an example because I think, like, it's a super, super simple idea. I think when it came out as well, like, there were a lot of people that were, like, when they see something like that, they're, like, what is even the purpose of this? This Speaker 2 (16:27) is Speaker 3 (16:27) too cool. Yeah, like, why, like, this is just subagents or something, right? And I think it's, like, when you get a reaction like that, it's almost like a good sign in the sense that, like, it's clear that people aren't thinking about. what the purpose of this is. And I will give another example of like Sweebench. When Sweebench came out, I mean, Ophir loves to tell this story. When it came out, like nobody cared. Like everybody was like, this is an impossible task. Like why would we ever even consider this as a benchmark? And it wasn't until Devin came out that everyone was like, whoa, like this is something we want to hill climb. And I think this is, this rings true for, you tend to see that a lot of ideas, I think like my favorite, I guess, example of this is Eric Seligman's work with like Star and like Quiet Star. Like I think when you read the paper, at least when I first read the paper, I was like, is this not like an obvious idea? Or maybe not. I don't know. I mean, I was like, oh, this seems really simple or like Chain of Thought. Same thing. Or like True News React. It's like, OK, like, yeah, I mean, sure. But then like when you really think about it, it's like why? What is the value of the paper? I think it comes from like, it tells a bit of a story as to like what you want the field to look like. And that is something that it's very hard to do this in academia because if you look at all these papers, Quiet Star, React, RLMs, Sweebench, none of these papers are. It's not like a GPT-6 Astra release. You know, it's not like everyone's like, oh, my gosh, like I'm going to use this now. This is the best thing in the world. Like academia just can't afford to do this, at least right now. I mean, there's a whole slew of reasons why I think that should change. But I think it's like if you don't have as a PhD student, I think you're in such a unique position where you can work on literally whatever you want for the most part. If you're not taking advantage of that and working on things that like nobody cares about or like people see as. some trivial thing like, oh, I thought about this, but like, I don't use it. I just think like in the end, the research is just never going to be that interesting because you kind of need to take big bets if you're going to be in academia. Because otherwise I think like, just go to an industry lab, like, you know, they have tons of resources, tons of talent. Why constrain yourself in an area where you don't have a lot of resources and like, there's not even that many people around. And I think it's just, it literally just comes down to like big bets. You just have to take big bets and like, a lot of them will fail, you know? Like, that's just, it's natural. But I think... that is, as a PhD student, like that's the biggest advantage you have over any single person at another lab because you don't have to deal with bureaucracy and all these other things. Fair Speaker 2 (19:07) enough. Speaker 3 (19:07) Yeah. I ask a lot of people this question and usually Speaker 2 (19:09) they hand wave away. So I think I appreciate that you're actually giving a thoughtful response on like, no, like this is your unfair advantage because everything else is biased against you, basically. Yeah, exactly. And so like, it's honestly, I will bring up Jev Speaker 3 (19:25) as an example because it's not an academic project. I want to bring this up because this also happened with RLMs and it happens with many other works. Like things get overhyped, right? To an extent, like something gets overhyped and then people are like, why is this overhyped? Like this is trivial. This is stupid. And I saw the same thing with Jev because I think the release was like, I mean, there is this whole thing about like. academics or they're not an academic group, but like people have to do branding and they have to like kind of market their research. And so like, I understand. But I think there is a lot of discourse about Jev just being like something we've known for years. And I think it's kind of missing the point of like, why is such a system so interesting? It's why is it not just some stupid MLP classifier that like we've we've been doing, you know, back in our intro ML classes or something. I think what's really interesting about Jev is that it kind of opens up this question of, are language models correct? Like, in the form that they're in, can we consider a different design space other than text-to-text? What they're doing is they're basically saying, like, I will take advantage. of this language model backbone. Like I know it captures a lot of information about language, but I'm going to change the output space of the model to give you a trade-off, which is I will do very, very fast inference over like if you have some prior about this problem, like let's say I only need to make a binary classification. Am I going to ask my language model to do this and pay like a 400x cost? No, it's like silly, right? And I think for the longest time, because the labs are the only places that control, you're never going to use something other than like GPT-6 or Fable because they're the best models. But because of that, like people have gotten kind of accustomed to this idea that a language model is just a autoregressive decoder. Like we have accepted this. And I think when RLMs came out, it was the same thing. Like one of the comments, like a very frequent. criticism I got was like, this is not a language model. Or like, when I look at this, like, I thought it was a new architecture, but it's actually not. And my response to that is like, well, a language model is just modeling language. It doesn't have to be this transformer decoder, you know? And Jeb is really, really interesting in that, like, we now have a new thing to tune, which is like, what is the output space and how does this affect inference latency? And I think we can actually start asking this about various parts of the language model itself. We are seeing this too with like loop transformers. It's a similar idea of a lot of the attention around it was like, this is a silly idea. Like why, who cares about this? But it's like, it is a simple idea, but it's actually, it opens up a whole new set of questions that I think like, especially if you're a PhD student. These are the things that you want to answer because I think it's like we don't know for Jev, for example, we don't know how far we can take this. For loop transformers, we also don't know how far we can take this. What if you loop only a subset of the model? What if you route to like only like you have some router to different parts of the model? Like, can you mimic what you do in a harness inside of the model? What can you bridge between a harness choice and a model architecture choice? These are all questions I think that open up with works like this. And that's like really exciting. Jeb in particular, when I saw it, I was like, this is actually really useful for RLMs. Like I think it makes sense because the biggest bottleneck in RLMs or swarms or systems like these is they're slow. When you do multiple language model calls all the time. you're not distributing your compute correctly because like maybe there's something trivial that you just want a simple model to do, but you can't do it because your language model is just this bulky thing, you know? So I'm very excited. I think we will start to see new types of models emerge beyond just the bog standard frontier model. And that is like, there's so many things that you can do with these Speaker 2 (23:32) like new trade-offs. Also shout out Thinky with their interaction model. Yes. Speaker 3 (23:36) Yeah. Yeah. Another great example. Speaker 2 (23:38) Yeah. So like basically try to break the paradigm from sequence to sequence, decoder only, and just literally do anything else. Yeah. There's Speaker 1 (23:47) a level of, if you're trying to compete, you're not going to compete with a frontier lab doing an autoregressive decoder on like. The amount of compute scaling resources they have, even Thinky will not, you know, okay, there may be one of the handful that can, but you're not really going to do much in that. I don't Speaker 3 (24:06) know Speaker 1 (24:06) too much about Thinky Speaker 3 (24:07) or I don't Speaker 1 (24:08) want Speaker 3 (24:08) to say anything either, but it's like, if their strategy is just to... replicate OpenAI or Anthropic, like, that's a horrible strategy. Because, like, Speaker 1 (24:16) I mean, they just don't Speaker 3 (24:17) have, like, you kind of just have to think of it in terms of, like, what advantage do you have? And if you're going to use the same setup, I'm sure they're not, but it's like, if you're going to do the same setup, like, you're basically competing on the things you can control, which is data and compute, and obviously they cannot compete with the Frontier Labs on that. So, yeah, it makes sense that, like, if you are a Neolab, like... Actually, I don't know if you would consider them to be a neolab, but I guess they're kind of a weird one. They're in that list. They should Speaker 2 (24:44) inkling. They can't. Speaker 3 (24:45) Anything other than OpenAR, Anthropic, maybe like Meta and GDM, you got to do something else. It's the sad reality, but I actually think it's a good thing. I'm very happy that scale works and these companies will just keep doing it because it opens up. potentially new players, like