Dataflow Computing for AI Inference with Kunle Olukotun - #751

0:00I'd like to thank our friends at Capital One for sponsoring today's episode. Capital One's tech team isn't just talking about multi-agenticentic AI. they already deployed one. It's called Chat Concierge and it's simplifying car shopping. Using self-reflection and layered reasoning with live API checks, it doesn't just help buyers find a car they love. It helps schedule a test drive, get pre-approved for financing, and estimate trade in value. Advanced, intuitive, and deployed. That's how they stack. That's technology at Capital One.

0:34Hey folks, Stephen Johnson here, co-founder of Notebook LM. As an author, I've always been obsessed with how software could help organize ideas and make connections. So we built Notebook LM as an AI-first tool for anyone trying to make sense of complex information. Upload your documents and Notebook LM instantly becomes your personal expert, uncovering insights and helping you brainstorm. Try it at notebooklm.google.com.

1:41all right everyone welcome to another episode of the twimloidai podcast i am your host sam sharrington today i'm joined by kunle olukotin kunle is a professor of electrical engineering and computer science at stanford university and co-founder and chief technologist at SambaNova Systems. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Kunle, welcome back to the podcast. It has been a minute. Good to be back and glad to be here. We were just chatting. It was, I think, at a NURPS when we first spoke back in 2018. clearly a very different world in AI since then and it is going to be really good to catch up with you and what's been new in your research and at SambaNova.

2:34Let's get started by having you share a little bit about your background and what you've been working on recently. Well thank you for having me on again and it's uh good to be uh back on on the uh the cast podcast and uh you know my background is is is i'm a computer architect uh well known uh for doing some of the uh pioneering multi-core work uh back in the in the uh mid-90s and uh done a lot of work on parallel programming environments using domain-specific languages. And this work has led to the work we're doing now in research and at Sandanova about how to take reconfigurable data flow architectures and efficiently execute these domain-specific language models to create environments that are better for doing AI computation, right?

3:35And so our focus is on how you do very efficient, fast inference on big models as energy efficiently as possible, right? And so the focus has been on very large models, you know, models with trillions of parameters that can fit in a single rack, on models that are good for agentic solutions, where you've got hundreds of agents that potentially need to work together and switch between these agents with very low latency, and then sort of very fast inference because you have interactive environments that you want to support multiple LLM model calls sequentially for a chain of thought kinds of reasoning or because you've got a workflow that requires lots of LLM calls.

4:35You mentioned in there the idea of reconfigurable data flow architecture. Let's kind of step back for a second and talk a little bit about what that means. What is a reconfigurable data flow architecture and where does that idea come from? Well, it comes basically from thinking about how you take the computational paradigm that comes from machine learning and implement it in hardware, right? So you think about designing a AI or ML algorithm, right? And you do it in your favorite domain-specific language framework. And these days, everybody's favorite is PyTorch, right? You know, back in the day, we did some early work on a language that you've probably never heard of called Optimal, but it predated PyTorch, right?

5:36But it was the same idea, right? And then, of course, you know, Google had TensorFlow and for a while there it was in contention, but I think, you know, it's fallen by the wayside. So today, it's all PyTorch. And what'd you get out of a PyTorch uh you know once you write your your algorithm in pytorch what you get is what computer scientists would call a data flow graph right and so the data flow graph is a graph which has nodes and edges right because all of us have nodes nodes and edges and the nodes in this case are kernels of computation, like matrix multiply or convolution or softmax or maybe your favorite kernel that you came up with to support your specialized model.

6:31And then the edges are basically tensor data, right? Matrices that flow between these nodes, right? Right. So there's a famous computer architect named Jim Smith, and he designed Cray computers back in the day. And his favorite quote that I often use is, you know, if you have a vector problem, and most of the problems were vector problems, then you should build a vector computer. Right. So. All right. So we have a data flow problem, right? So let's build a data flow computer, right? So let's build a computer that matches the data flows as represented by the algorithm in the hardware. In other words, these graphs.

7:27These graphs, right? So why does it need to be reconfigurable? Well, because you don't have one graph, right? Right. The graph is your program. It's your model. It's how you change what the hardware does. You reconfigure it, right? Instead of thinking about executing instructions, you think about configuring your hardware to match the graph that you've developed in this domain-specific language. That is the fundamental idea of reconfigurable data flow. So compare and contrast this to kind of traditional computer architectures, right? You know, there are data flows that I guess by, you know, you can squint and also call them reconfigurable data flows.

8:18You've got registers and you've got memory buses and data flows in different directions. And you've got. Yeah. So I think the key difference is whether you have instructions, right? You know. So meaning instructions versus kernels or? Instructions versus configuration, right? Versus do you fetch an instruction every cycle and decide what to do? Or do you configure your hardware and leave it for that way for the whole program? Okay. That's the difference, right? That's one difference. The other difference is a kind of technical issue, but also is very important, which is how do you synchronize the computation, right?

9:09You have parallel computation, and the way that parallel computation typically works is the way that you communicate is through a single address space called shared memory, right? And whenever you do this, you need some sort of synchronization to make sure that, you know, the data that one thread put in the memory is ready for the next thread to go pick up, right? So locks, barriers, all this stuff. Reconfigurable data flow has none of that. You don't have any of that, right? So this is a huge benefit because it means... shared memory or because there's an alternative to locks for ensuring synchronization and accessing shared memory?

9:57Both. I thought those were mutually exclusive propositions.

10:07There's no shared memory. Okay. But you can access, there's no explicit shared memory. That's true. And the way that you do the synchronization is by using data flow tags, right? So you've got hardware mechanisms that essentially tokens, which determine when data is ready, right? So I send you a stream of data and I've got some tokens there, which you can look at and decide when the data is ready. Okay, so that is the way that you do communication, and it matches very nicely with this model of the graph where one kernel does a computation, and it generates data, and it sends that to another kernel, which is executing on a different piece of the chip.

11:10One of the pictures that I'm getting of this is kind of analogous to like stream processing versus traditional database, like centralized processing. Is that? Streaming is definitely a mechanism that we use. Like at the hardware, like chip level. At the chip level, you're doing streaming and you've got mechanisms that make sure that the streaming has a zero overhead. Okay. Yeah. That's the way to think about it. Interesting.

11:46And with regards to reconfigurability, I'm trying to get a picture for, on the one side, and this is not the right spectrum, I don't think, but on the one side, you're fetching instructions every cycle on another side like you've got like uh um like uh like a reconfigurable chip like a xilinx or like a uh i'm the blanking the words blanking for me but like fpga yeah exactly yeah so it has some similarities to fpga but it's it's it's much coarser grained right you're not thinking doing things at the at the logic gate level you're doing things at the coarser tensor It's a core level, right? And then you've got this very explicit mechanism for doing this synchronization based on tokens.

12:43But if it's FPGA-oriented, like, I would... So don't think FPGA because it's not, you know, you don't program it with a hardware description. That's what I'm thinking of, like, the timescale of utility of configuration is, like, what I'm trying to get at. Yeah, yeah, yeah. So it's much, much quicker, right? So you're on the order of microseconds to reconfigure. Oh, okay. Got it. So you can push a PyTorch graph and reconfigure it. Maybe you wouldn't want to do that per transaction or per inference, but you could? You might. You might. You might. Because, you know, so think of it, right? Eventually, you're going to get a graph that is bigger than the system that you're trying to map it on.

13:29So now you have to reuse that hardware for different parts of the graph. Okay. Interesting. So you have to be able to reconfigure. Taking a step back here, we're talking about this idea of reconfigurable data flow architectures. And at the point we were talking, this kind of was, I think, a little bit closer to the cusp of research to practice. Now this is something that you've been doing, you know, particularly since like the, I think the, you know, the chat GPT moment, you know, shortly after there we saw a shift. I remember seeing a shift from Salmonova from like this theoretical computer that is scalable that you go figure out what you want to do with it to like, we're going to get really good at this LLM thing and, you know, make it, you know, available to folks.

14:23Talk a little bit about how you got from this idea of reconfigurable data flow architectures to we've got these LLMs, this is the right thing for it, and what we need to do to either accommodate the LLMs on the architecture or accommodate the architecture for the LLMs. Yeah. So it turns out the architecture was very well suited for LLMs. is fundamentally, especially for inference, right? And so inference was something that we saw was really going to be a bottleneck. And if you think about inference, if you think about doing inference, it's fundamentally constrained by the memory bandwidth, right?

15:17So I should back up and say, okay, You know, when I talked to you about 2018, in 2018, this idea was pretty theoretical. Now we have five generations of chip. The chip that we currently have in the market is called the SN40L, which is, you know, 100 billion transistors, five nanometer chip. And it has three tiers of memory, right? So it's got half a gigabyte of memory on chip. It's got HBM, 64 gigabytes of HBM connected to the accelerator. And then it also has a high capacity memory, 1.5 terabytes of DDR, right? So this is unique in that most accelerators do not have that high capacity memory.

16:10And so when you think about doing inference on a large language model, It's fundamentally constrained by the bandwidth to fetch the parameters and to fetch the KV cache from the HBM. That's the limit. So what you want to do then is make sure that you don't send any extraneous data that should should not be necessary over that highly constrained, that critical resource, that HBM bandwidth, right? And so reconfigurable data flow does this by, you know, you can go back to your streaming idea, right? So if you take a LLM model, right? So let's take a simple one like LAMA 3.1 8B, right? 8 billion parameter model, it's got these decoders, and it's got 32 decoders, right, which are essentially the same operation.

17:23What you can do is that you can take that whole decoder and map it onto a set of 16 RDU chips in space. So that whole... An RDU chip is... Is a reconfigurable data unit. The SN40, yeah. Yeah, yeah. So we call reconfigurable data flow architecture and the chip we call a reconfigurable data flow unit, RDU. So take 16 of these RDUs and you take the one decoder and you map it in space across all those chips. Okay, so now it's resonant on your system. And then to generate the tokens, then you just loop through that 16 times. And no data ever goes between the different components of the decoder ever crosses the HBM accelerator boundary.

18:31And the only data that comes... so you don't need like InfiniBand or, you know, some super low latency interconnects? Well, you do have interconnect between the chips that is lower latency, you know, some sort of custom chip implementation, custom interchip protocol. The 16 RDUs, like, should we think of these as cores that are on one die, or are they like chips that are on one board, or are they systems that are in one rack? So typically you've got a couple of chips per board, and so you can think of these as across eight boards. So you think of an RDU as an Argus to a GPU in some sense. So now, of course, this idea in the compiler land is known as fusion, right?

19:34So you essentially have created a single kernel, a single fused kernel that encompasses the whole decoder, right? So, you know, you probably have heard of this idea of flash attention, right? In fact, I was just going to ask you, you talked about the memory bandwidth limitation being a significant issue and the way we usually do inference, like with GPUs. And we've got all these tricks like flash attention and speculative decoding and other things that are trying to help us overcome that limitation. and the question that I was going to ask was like, do you not have to do these kinds of things or do you do them differently?

20:22You can do them on steroids, right? Oh, okay. The same tricks work, except you get to do them much better. Okay, okay. Talk more about that. So if you think about this idea of flash attention, right? What you're doing is you're taking some component of the decoder that does the attention mechanism, and you're fusing things together such that you can reduce the memory bandwidth requirements. We can go even further and fuse the whole decoder, right? So you get flash attention across benefits across the whole decoder, right? Which is a huge reduction in memory bandwidth requirements, okay? Okay? So then that's one thing.

21:11The other thing that we do, it comes back to this idea of synchronization, right? And making sure that the synchronization is very low overhead. As you know, in any computer system, you have to make sure that the critical resource is used 100 % of the time. Because that's going to limit your performance. And the performance limiter for inference is, as we've just discussed, HBM bandwidth for the parameters and the KV cache. And so what you want to do is you want to make sure that you never have that interface idle and you overlap the computation happening in the CPU, with the memory access on the HBM and the communication across chip to the other RDUs.

22:12You keep all of this, you know, it's like everything everywhere all at once, right? You know, you want everything to be busy all at the same time. And you want to make sure that you use that interface to the HBM all the time. And what we do is we can achieve the fraction of busy time that we can achieve on the RDU is two to three times higher than GPUs can achieve. because we have this data flow idea of asynchronous execution of all the components, right? So there's no sequential instruction access, right? So one of the other, you alluding to your earlier question of why is this different? Well, the problem with sequential instruction execution is, you know, has a benefit is that humans can understand what's going on.

23:23The big downside is it keeps things sequential, right? Sequential, right? So what you want is you want asynchronous, right? So GPUs are starting to get into this idea, but we've taken it to the extreme in data flow. As I said, it's everything all at once, all the time, right? So it's all happening at the same time, much like you would a hardware system. And so there is nothing keeping us from keeping things busy all the time. And so is that HBM utilization, is that like just a simple following from the parallelization? Or are you also layering tricks like imagining like speculative prefetching or, you know, various things you might want to...

24:23It comes naturally. It comes naturally from once you kind of design things in this data flow manner, you don't have to prefetch, right? Because you naturally have prefetching going on because that unit accessing the HBM is working independently at the rate at which the HBM can go, right? You don't have to prefetch. I'm imagining then that like the, you're thinking a lot about, you're either thinking about InDesign or you're limited by like some ratio of RDUs to, to like interconnect or in order to keep that HBM utilized. Like you can't go beyond a certain scale of compute because that determines like that utilization level.

25:13Yeah, so part of the trick is sort of how you're going to map the model to your architecture, but you're going to map it in such a way that you make sure that the HBM bandwidth is the limiter, right, for the data you need to move. And so the idea is you don't move any data that you shouldn't have to across to using the HBM bandwidth. So you only use the HBM bandwidth for the absolutely necessary data, right? So a GPU puts intermediate data between the kernels that flows back and forth the GPU. That's wasted bandwidth, right? So you don't do any of that because you've got this data flow mechanisms that move the data between the kernels on chip and between chips using chip to chip interconnect.

26:09And then you also make sure that that critical resource is used at the highest utilization that you can muster, which is upwards of 90%, which GPUs cannot touch for multiple reasons. So talk a little bit about that process of mapping a model to the architecture. You talked about one of the key constraints there is utilizing that HPM bandwidth, but what goes into it when Meta comes out with a new model or Quinn or whoever comes out with a new model? How do you then ensure that that's optimized for this approach? Yeah, yeah, that's a good question. So, you know, we start with PyTorch models, right?

27:00And so PyTorch gives us a graph of operators, right? And then we have a implementation of these operators in our environment. And so we've got a way of, you know, so if you think about these operators, you can describe them as they're either gems or there may be some custom operation, right? And this custom operation could be described using a map reduce paradigm. So we've got a Python environment that allows you to describe specialized operators using MapReduce. And so now you've got a set of kernels. And now the question is, how should you fuse these kernels together to create a representation that can be mapped?

28:08And the way that you do things is that you need to paralyze within the chip by splitting up the tensors. You might want to reduce the size of things by tiling, just by taking the tile of computation, thinking back to what you had to do manually in something like flash attention. And then you also want to think about how you shard the computation across multiple chips. And so our compiler has mechanisms for you to specify what you want based on the tensors. This is all about the tensors. You take a tensor, you say, am I going to tile the tensor? Am I going to paralyze the tensor? Am I going to shard the tensor?

28:55And then that's the way you get the mapping. And then now you have something that can map to your machine, which is a collection of RDUs. And the whole goal is there will be some performance tweaking. But at the end of the day, what you're trying to do is to make sure that you are using the critical resource. If it's inference, it's HBM bandwidth to the maximum extent. extent. Yeah, I remember thinking that in the early days of Salmonova, like having to have, you know, enterprise engineers that have built their own custom models, like manage that, you know, mapping process, you know, sounded like a difficult problem just from like a go-to-market perspective.

29:47But it's a different world when you've got, even at the pace that we're seeing new models now, you know, that's work that your engineers can do. They know how to do it. And, um, right. And we've gotten very adept at doing this, this sort of thing. Right. And, and fundamentally they're all transformer based models, even variants. In fact, you know, once, once deep seek came out, which has a slightly different, uh, multi-headed latent attention, which is slightly different. It took us about a week to implement that and get it, get that running on our system. So, you know, as long as they're basically transformer-based models, then we can do this fairly quickly.

30:32But, you know, even if you have to create a custom kernel, you can do it at the Python level. You don't ever have to write kernel, I have the right kudo, I should say. And so we can make this transition fairly easily. You know, in thinking about kind of optimizing the system for inference, like are you primarily optimizing around tokens per second or are there, how do you think about the different metrics that are important for inference and how you tune your system to deliver those? We are optimized for both tokens per second because we have this kind of very low latency capability when you're doing a matrix vector style computation, right?

31:31So we do that very quickly and we're, you know, sort of 10x or faster than GPUs at that point. But even when you increase the batch size in order to get higher throughput, right, as you know, there's this latency throughput tradeoff that you see in inference where as you increase the batch size, you get higher throughput. And so from the point of view of the server, you know, this looks better because it costs you less power. You know, you can get a higher throughput in terms of the requests that come in. But the latency, of course, goes up, right? And if you look at GPUs, the latency really goes up once you get to sort of, you know, 1 ,024 or above.

32:26And what you see on a SN40L-based system is, yes, our latency increases, but it doesn't increase nearly as much. We're still kind of 5x better in terms of latency, even at these very, very high throughputs. So that's what cause you're not batching? or? We are batching, but fundamentally we still retain the advantages that we had, uh, you know, with the, the lower batch, uh, uh, implementations, right? Yeah. Why is that? Like I, I, when I think about the, you know, that throughput latency trade-off, I think that, you know, the reason why the latency grows up is because you're waiting to start your computation because of the batch, right?

33:14And so how do you not have to deal with that? Well, you still have to deal with the latency, the queuing latency, but the fundamental computation that we can do is more efficient because we can do tensor parallelism in ways that GPUs cannot, right? GPUs aren't efficient with tensor parallelism because they can't overlap communication very efficiently. And we can use tensor parallelism to decrease the time that the overall computation takes. And because we can overlap and hide the latency of communication, this is still effective for us, whereas it might not be, and typically it's not effective if you try and do that on GPUs.

34:05so yeah you you don't you don't get a you don't you the queuing latency you're going to pay but the question so there's two components of the latency there's one is the queuing latency uh and then in the gpu you know is it fair to think about as like a serialization latency um because the batch is processed serially and then in your case the batch is processed the batch is is processed in parallel but the thing is that that fundamentally yeah yeah fundamentally you can't take advantage of the kinds of things that we can do on the RDU, right? So yeah, you can think of it as some serialization that happens on a GPU that we can paralyze away on the RDU.

34:51That's one way to think about it. I assume you mentioned one other thing about inference while we're on that subject, which is one of the things that we have because we've got this high-capacity memory is we can hold multiple models at the same time, up to 5 trillion parameters in total. And then because you've got, you know, it's not nearly HBM bandwidth, but you've got much higher bandwidth to the accelerator than you would have if you were trying to communicate over PCIe to the host, Now you can switch models in about a millisecond, right? And so you can think about if you've got multiple custom models potentially, right?