# CS336 Lecture 6: Kernels, Triton, XLA（算子与 Triton 实战） > [Stanford CS336 · Language Modeling from Scratch · Spring 2026](https://cs336.stanford.edu/) · [视频](https://www.youtube.com/watch?v=xnDHaNUvHBg) · [官方讲义 (Trace Viewer)](https://cs336.stanford.edu/lectures/?trace=lecture_06) > 讲师：Percy Liang · 时长 87 分钟 · 本文档 = 官方讲义截图 + 全程文稿（whisper large-v3-turbo 转录，已去除语气词） ## 目录 1. 课程引入：从 GPU 概览到写 kernel — [▶ 0:00](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=0) 2. GPU 速览：硬件规格与编程模型 — [▶ 0:24](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=24) 3. 方法论：Recipe for Success — [▶ 21:59](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1319) 4. Benchmarking：测端到端耗时 — [▶ 22:34](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1354) 5. Profiling：看时间花在哪 — [▶ 26:16](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1576) 6. GeLU 三种实现对比：朴素 / 内置 / torch.compile — [▶ 30:07](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1807) 7. Triton 入门：面向线程块的编程模型 — [▶ 35:21](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=2121) 8. Triton 实战 1：GeLU（逐元素操作） — [▶ 39:39](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=2379) 9. Triton 实战 2：softmax（行内归约） — [▶ 57:27](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=3447) 10. Triton 实战 3：行求和（分块归约） — [▶ 64:24](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=3864) 11. Triton 实战 4：矩阵乘 + ReLU（tiling 与共享内存） — [▶ 71:51](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=4311) 12. 总结与下一讲预告 — [▶ 83:54](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=5034) ## 1. 课程引入：从 GPU 概览到写 kernel *[0:00 – 0:24] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=0)* `0:00` . `0:04` So, welcome back everyone. ![Slide 0](https://cs336.sg.aidenovak.com/media/slides/l6/000.webp) `0:08` So, on Monday, Tatsu did an excellent job of describing a high-level overview of GPUs and how to think about performance, and all the quirks that come with GPUs. This lecture is going to be a continuation of that, where we're going to dive more deeply into the code, ## 2. GPU 速览：硬件规格与编程模型 *[0:24 – 21:59] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=24)* `0:24` write some tri-chain kernels, and as well as do some benchmarking and profiling. So, just to start, just to refresh what's going on with GPUs, here is a simplified diagram of what a typical GPU, you know, looks like. `0:44` There's memory, and then the actual GPU chip. And just to keep in your head, what are the characteristics of a GPU? So, of course, we're talking about NVIDIA GPUs, but Tatsu also talked about TPUs, and there's AMD, and so on. But focusing on NVIDIA GPUs, we've had multiple generations from A100s to H100s, to B200s. And in each generation, you have a GPU. It has a number of streaming multiprocessors, or SMs. And the number of SMs is, you know, about 100, between 100 and 200. So, that hasn't really changed that much. And then, within individual SM, there's a set of registers. ![Slide 3](https://cs336.sg.aidenovak.com/media/slides/l6/003.webp) `1:37` So, B200s have 65,000 registers, for a total of 256k per SM. And this also hasn't changed, you know, that much. In addition, there is a L1 cache, as well as shared memory. So, remember, these are the same memory, shared memory, you can control L1, you can't. This is per SM. And this is, like, in the same order, ballpark of size. And then there's a L2 cache. And this is not per SM. This is for the whole chip. `2:15` And this is a bit larger. And then finally, you have a high bandwidth memory, or HPM, which is large. And you see that this is the number that's actually going up quite a bit. Okay. And then you can, in addition to size, you can think about what's happening with the bandwidth. And essentially, it's, you know, kind of inversely correlated. So, registers are very fast. L1 is slightly less fast. L2 is less fast. ![Slide 4](https://cs336.sg.aidenovak.com/media/slides/l6/004.webp) `2:48` And high bandwidth memory is the slowest. Although, 8 terabytes a second is still, you know, not that slow in the grand scheme of things. So, this is the kind of main hierarchy you should have in your head. Memory, large memory is slow and far, but big. And, fast memory, like registers and L1, resides on SM. It's local, and it's fast but small. Okay. That's what, the environment we're dealing with. Okay. ![Slide 5](https://cs336.sg.aidenovak.com/media/slides/l6/005.webp) `3:18` So, then, how do you program a GPU? So, the, the program model is as follows. So, there's notions of threads, which each thread executes a piece of code on a small part of the data. So, each one of these arrows, you can think about as a thread. These threads are organized into thread blocks, also known as concurrent thread arrays or CTAs. And this is a group of threads. And then finally, you have a collection of thread blocks, which forms the grid. So, when you launch a kernel, you're basically launching a grid of threads and thread blocks, to all parallel, you know, simultaneously do some, you know, computation. ![Slide 6](https://cs336.sg.aidenovak.com/media/slides/l6/006.webp) `3:59` Okay. So, this is simplified. There's some other things that are happening. For example, there's H100s and B200s also have, thread block clusters, which are clusters of thread blocks, that enable some amount of, you know, distribution of memory. Also, B200s have tensor memory now for tensor cores, which is somewhere between registers and shared memory. some of these things are kind of invisible to the programmer, but they are in, in, in the hardware. but we won't worry about them, for this, intro lecture. Okay. So, one thing you might ask about is why do we have thread blocks? `4:37` Why can't we just have a grid of threads and each thread just takes a piece of data and does something? And this would be normally fine if all you want to do is element-wise operations, which we'll see. For example, Kalu is an activation function that applies element-wise. And threads are pretty natural. Each thread processes one element. So, it's like FI for I equals, you know, ranging over your, dataset. But for operations, that are, involve, communicating with, you know, threads such as Softmax or matrix multiplication, this view isn't really enough. And, and the reason is that, well, it, it could be enough if you were willing to pay the cost of writing and reading HPM, because then you can still have, for example, in a matrix, we'll see later, you can just have every element in a cell. It goes and computes that element of the matrix, and it can just read and write HPM. ![Slide 7](https://cs336.sg.aidenovak.com/media/slides/l6/007.webp) `5:40` But, as we've, noticed, HPM is very slow, so this would not be a good, strategy. So, instead, what we should do is use shared memory, which is local to, SM. And what the thread block allows you to think about is a collection of threads that are all going to access the shared, memory. So, so, consequently, you can think about a thread block. One of these, thread blocks is being scheduled on, one of these, semantic, not semantic, streaming multiprocessors, and it does its thing. And what it's going to do is it's going to read a bunch of data from HPM and then, process it where the processing might involve communication between, you know, the, the threads via the, shared memory, and then writes it back out. Okay? So, this is kind of a critical piece that comes up. Later we'll talk about tiling, and that's the whole, kind of game here. `6:45` Okay, in, in Triton, in fact, which we'll see as a primary way that we're going to write kernels, these days, we're going to think natively in thread blocks, and, you know, I, I think once you get the hang of, hang of thinking in thread blocks, this makes, a lot of sense, and makes your life a lot easier, as we'll see. ![Slide 8](https://cs336.sg.aidenovak.com/media/slides/l6/008.webp) `7:04` Okay, so the programming model is, is fairly, you know, simple, I think. There's threads, thread blocks, and grids. Okay, so, where it, I think things get a bit more, you know, complicated is the interaction between the programming model and the hardware. So, the programming model is actually very nice. It provides abstraction of the hardware. You, when you, you can write a kernel, you just, all you have to know is that there's a bunch of thread blocks. You define them, and then define what computations, each thread, within the block or the thread block has to do, and those are just like writing, it's like writing Python. so that part isn't hard. And, in fact, if you don't care about, if you just care about correctness, that's all you need to know. ![Slide 9](https://cs336.sg.aidenovak.com/media/slides/l6/009.webp) `7:55` But, in practice, the performance is very sensitive to the hardware, and so you need to really deeply understand the hardware to obtain high performance. And, in fact, the whole reason we're talking about, GPUs and kernels is that you're trying to squeeze out performance. So, there's sort of these two levels here where you're trying to understand the computation, and that is only, a part of the programming model. But, you know, how fast it runs is going to be strongly dependent on the hardware. Okay, so I'm going to give you some examples of why this matters, and some of this will be a review from what Tatsu, you know, talked about, but hopefully it'll reinforce some of these, concepts. So, I think there will be around five different examples, and, you know, this is just to give you a flavor of, the considerations. So, there's something called warps, which I didn't talk about. so in the very simplified view, warps are not really part of the, this clean picture of the program model. You have threads and thread blocks or, grids. ![Slide 10](https://cs336.sg.aidenovak.com/media/slides/l6/010.webp) `9:00` technically you can also exceed the warps, but you don't, you know, have to when you're just, you know, programming. so what is a warp? So, each thread block, remember, is a, a collection of threads. But these threads are actually grouped into warps. And, you know, there's basically 32 threads, per warp. So, for example, if you have 64 threads in a thread block, there's two warps. The first 32 and the second 32. And, as Tatsu mentioned last time, all threads within a warp, all these Ts, must execute the same instruction in lockstep on SM. So, every cycle, they have to execute exactly the same instruction. `9:44` And control divergence, is when, different threads within a warp need to actually execute different instructions. For example, if you have branching, if A, if something, then A, else B. And then what happens is that you can only do the As in your threads in your warp, and then you do the Bs. So, this gets kind of sequentialized. So, this is bad and inefficient, and that's why branching is, is something you want to generally, avoid. One thing that's, I think really, cool about warps is that, SM actually is running multiple warps. There's a warp scheduler, and there's a bunch of, you know, resident warps that are, about to, to run. Each warp has some, you know, and threads have, registers. and you, the SM can switch between them with, you know, zero cost. `10:40` Right? This is not true in general, like on a, on a CPU, but this way that it's designed is to hide, latency. So, this is important because one, one of the warps is, for example, reading from HBM, which remember is, is, very expensive, can take, like, a, you know, 100 cycles or something. You don't want to just sit around waiting for that warp to do nothing. You switch immediately to another warp where it can actually do some tensor core operations. Okay. So, so that's one thing to note about the warp. And another reason why warp comes, up is this idea of occupancy, more specifically warp occupancy. So each thread, just, you know, the hardware constraint says that each thread can use, at most 255 registers. `11:31` Okay? And so what happens is that the SM has a fixed number of registers, so the more registers each, thread is using, the fewer threads you can have. That's just, you know, math. and that is called, that can reduce your occupancy. but that's not necessarily bad because if you have fewer threads, but each thread is doing more work, that can actually be good. So, you know, occupancy is, is something that you can measure, but it's not necessarily the larger, the better because there are some other trade-offs here. And in fact, an example where, you might want to, you know, have fewer threads is this idea of thread coarsening, where let's say you have, you know, element-wise operation and you're going to, you can have one thread just perform, work on one element. Okay? So then you have a lot of threads. ![Slide 11](https://cs336.sg.aidenovak.com/media/slides/l6/011.webp) `12:29` But you can also say each thread processes multiple elements, like a constant, maybe eight. so that gives you fewer threads, which makes scheduling and these things easier, but each thread is doing more work. So if your threads are very light, maybe you want to fatten them up a little bit. Okay. So, just as an example of, what might happen here. So suppose you have, a block, a thread block with 128 threads in it. And each thread is using 160 registers. Okay? So there's a hardware constraint, so the, the B200s have a maximum of registers of 65,000 per, SM. `13:13` And it also has a constraint that says you can't run more than 65, 64 warps, at a time. So then you can, you can go and do some simple math to compute how many, what is the, you know, occupancy here. So, you know, number of registers per thread, this has to be less than 255, that's fine. the number of registers that you're using per block is the number of threads per block times the number of registers per thread. That's about 20,000. That means you can run at most three blocks on your SM concurrently. And then that corresponds to, 12 warps. And because the number of maximum warps is 64, that means you have, occupancy of, you know, 18%. So you're only using, running 18%, the total number of, you know, warps you have. `14:04` And this is because you have a lot of register use per thread. Okay? So this is an example of where, I guess this is sort of memory, is kind of constraining you in terms of how much compute you can, do. ![Slide 12](https://cs336.sg.aidenovak.com/media/slides/l6/012.webp) `14:21` Okay. Here's, another example, something called bank conflicts. This applies to shared memory. so remember, shared memory is, like L1. It's on a SM. And just the way the hardware works is that shared memory is divided into 32 banks, each one, four bytes wide. Okay? So, these are my 32 banks, and this is my memory, here. each bank has, many, elements here. `14:52` so there's a constraint, that says every, you know, clock cycle, each bank can only be accessed by one thread. At most one thread. assuming that it's not a location. So, for example, if, you can't access this location and this location, you know, at the same time. which means that if you have multiple threads that are trying to access the same bank, the accesses have to be serialized. and this leads, this is called a plane conflict. In the worst case example, imagine this were a matrix that was kind of laid out like this, and you had some operation where 32 threads, you think, wow, this 32 is so great, I can massively parallelize this, and you try to all access this first column. They're just going to just wait in line. And this is a 32 way thing, conflict, which is the worst possible, setting you can be in. `15:49` and now you could say, well, okay, why do you do that? Just like access rows, of course. But, you know, this is unavoidable because, for example, if you're doing a mat mall, you have to access rows of one matrix and columns of the other matrix. And sometimes you do transposes, so you can't always just get away with, choosing the order in which you, go down. you know, you know, for element-wise operations, it's fine. You can go in any kind of order. But if you're doing mat malls, you can't control, like, for every matrix, which, row major or column major you're doing. There's some solutions here which I won't get into. Something called, swizzling, which, ranges your shared memory so that when you're, going through, you can kind of avoid bank conflicts. ![Slide 13](https://cs336.sg.aidenovak.com/media/slides/l6/013.webp) `16:37` Okay. So that's, you know, another, you know, consideration. And when you're, you know, profiling, you can look at the bank conflicts. You can look at occupancy. And you can kind of see, you know, what's, what's happening. a final note which, actually two more things. So memory coalescing, which Tatsu talked about, so I'll just really quickly remind people, what this is about. So when you have, 32 threads in a warp and they try to access HPM, the memory accesses actually get combined into a transaction, of 20 to 128 bytes, which are called, cache lines. And it goes and fetches it all at, at once. `17:16` So, imagine you have memory laid out like this. this is, 32, wide. and, you know, in the, in the best case, which is called, full coalescing, all the threads are accessing the same cache line. So you have thread one accessing M00, thread two accessing M01, and so on. which means that, you know, all at once you're going to grab this entire cache line. Right? And whereas if you, go down the columns, then, then, you know, you're going to fetch a lot of memory here, which you're not going to use. same for the second row and, and so on. and so this is similar to, it feels similar to bank conflicts, but it's a very different constraint. `18:10` And this is about shared memory and this is about, HPM. Okay. Okay. So finally thing, block occupancy. so thread blocks are, remember, scheduled onto SMs. But, you know, we live in a, you know, finite, you know, logically you can define as many thread blocks as you want. But physically, on chip, you only have a certain number of SMs. For example, 148. so if you launch 160 thread blocks, then you can only schedule 148 of them. ![Slide 14](https://cs336.sg.aidenovak.com/media/slides/l6/014.webp) `18:41` And then you have to wait until they're done and then you schedule the 12. But what happens if you schedule 12 is that, there's a bunch of the SMs that are just not doing anything. so that's called, you know, low occupancy, when the last wave of, thread blocks gets, fewer, than the total maximum number of, thread blocks. Okay. So in general, you know, maybe it's a good idea to make the number of thread blocks divide the, the SMs. so, so maybe it's just to summarize here. You know, there is a very elegant, programming model where you have a grid of, thread blocks. And the thread blocks have individual threads in them. And, you know, at, and in terms of the memory, HPM is global to everyone. ![Slide 15](https://cs336.sg.aidenovak.com/media/slides/l6/015.webp) `19:31` Shared memory is local to a thread block, and registers are, local to a thread. But again, all the details of the hardware, warps, bank conflicts, memory coalescing occupancy, really determine, you know, performance. So, so I think, you know, and, and many of these details are sort of hard to know because, you don't, I mean, the profiling tells you, a bunch of information. But, you know, you have to know exactly how many, SMs there are and exactly the sizing of everything. And that gets, you know, sometimes a scheduler, does something and you don't really have control over what it's, it's doing. So it's a lot messier than the, the, programming model. `20:19` Okay. I will stop there for questions. Yeah. Yeah, so quick question. Can, is there a scenario in which a thread block can share an SM? Or, like, for example, like, in the problem that we see, like, 148 SMs, we have 165 blocks, right? And we have to, like, launch in two different ways. Is there no way in which a block can, like, share an SM? Yeah, so the question is, can a block, you know, share an SM? `20:48` I think the, the issue is, is if you're doing things right, then the, the, you know, the, the, I guess it depends on the block. This is, if you have a block that's, for example, using most of the tensor cores on SM, then putting another, you know, block there isn't really going to, you know, speed, things up. I think, and fundamentally there is this sort of jagged, you know, problem here of unevenness. The, the, because the blocks have to stay together. So you can't, like, take this and, like, spread it out over here. You can define your, you can, I think the thing, the thing you would do is to change your block size. So you should change the number of blocks. So you don't get this kind of, tail here. `21:35` Any other questions? `21:41` Any other questions? `21:44` Okay. So let's, move on. So that was, hopefully you guys are, you know, getting more comfortable with your GPUs now. so now I'm going to talk about benchmarking and, you know, profiling. ## 3. 方法论：Recipe for Success *[21:59 – 22:34] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1319)* ![Slide 16](https://cs336.sg.aidenovak.com/media/slides/l6/016.webp) ![Slide 17](https://cs336.sg.aidenovak.com/media/slides/l6/017.webp) `21:59` I'm not going to actually maybe say that much in terms of content. but I do want to emphasize the, the philosophy here. Which is, here's a recipe for success. you benchmark and profile your code. You make changes. And you benchmark your profile, your code again. and, and the reason I'm doing benchmarking and profiling as opposed to teaching you trying earlier is because, you know, you should always just measure what's going on, in your code and find, figure out what the bottlenecks are before you start, you know, ## 4. Benchmarking：测端到端耗时 *[22:34 – 26:16] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1354)* ![Slide 18](https://cs336.sg.aidenovak.com/media/slides/l6/018.webp) `22:34` writing kernels. Okay. So, benchmarking is basically how long things take. and it gives you just this end to end time. It doesn't tell you where things are, you know, spent. but nonetheless it's pretty useful because ultimately that's the thing that you, you care about. How long, you know, things are running for. and it also gives you, because it distills things into one number, you can see how things are scaling. Let's say with, dimension. ![Slide 19](https://cs336.sg.aidenovak.com/media/slides/l6/019.webp) `23:05` So, there's a, you know, nice, tool for benchmarking. but because this class is language models from scratch, I'm going to do it from scratch. but, you know, I, I think I'm doing it just to highlight a few, gotchas with, you know, benchmarking. So, let's say we have this, operation, matrix multiplication. so run operation two is this, this wrapper that basically instantiates two random square matrices of size dimension by dimension. And then it returns a function to perform the operation. Okay. So matmul, if you call this, basically does the matmul of those two random, matrices. Okay. ![Slide 20](https://cs336.sg.aidenovak.com/media/slides/l6/020.webp) `23:47` The matrix, random matrices are already generated. It's just like, multiplying the two. Okay. So, how do you benchmark? So, the naive thing to do is just start time, run it, and then stop time. but, you know, there's a few things. One is that always remember to do, your warmup. I think I mentioned this, you know, on the first, on the second lecture, but I think it's worth emphasizing. and this is because if you have some, some things are lazily compiled and you want to just make sure that, that, thing is, that is, that time doesn't, factor in too. ![Slide 21](https://cs336.sg.aidenovak.com/media/slides/l6/021.webp) `24:27` Because most of the time you care about, how fast something is because you're going to run it over and over again so the initial conditions don't really matter. and then you often want to time it multiple times because there is some variance. the, kind of the proper way to time things, is to use these CUDA events that, a start event and an end event, which you can kind of, call record on, actually do the computation and then hit this, the end event record. and then remember to, you know, synchronize to wait for the CUDA threads to, to finish because, everything on the GPU is happening, can't async and this is a synchronization barrier. and then you, record the time. Okay. And then you can do this, again. and then again. And then, you know, here we're just taking the average. `25:21` I think you might want to do, you know, if you are being very particular, you probably want to maybe look at the whole distribution, the P95 or, or whatever. But, we'll just do the average here. and then one thing that you can do with benchmarking is you can, let's say, scale up your, your matrices and see how the time, you know, changes. And you can see that in this case, matrix multiplication, you know, as expected, it should grow, you know, cubically. but, you know, notice that there is this floor where up until you get up to, like, almost 2,000, dimensional matrices, things are basically kind of a, you know, constant. Right? And this is because, as we've kind of discussed, the shapes of the, you know, these, you know, GPUs are built for fairly large matrix multiplications. ## 5. Profiling：看时间花在哪 *[26:16 – 30:07] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1576)* `26:16` And if you have, like, a two by two matrix, it's going to be very kind of inefficient. ![Slide 22](https://cs336.sg.aidenovak.com/media/slides/l6/022.webp) `26:22` Okay. So, really quickly on, profiling. profiling tells you where time is actually spent. Hopefully all of you are familiar with profiling and are, are, doing it. maybe less obviously, profiling, even if you don't care about the time, helps you, figure out what's actually happening under the hood. Because especially with these high level languages, you write some code and then it runs and you get some result. And sometimes it can be just, like, good to understand what's actually going on. So PyTorch has a built-in profiler in your assignment. You'll use insight, which gives you, more details. ![Slide 23](https://cs336.sg.aidenovak.com/media/slides/l6/023.webp) `27:00` but we're going to, skip that for the, in the interest of time. so let's start with just, you know, the, if you add two numbers and, or sorry, two tensors in, and High Torch. again, run operation creates two random, matrices and, and, applies the operation. So, in profiling, I'm warming up and then I am just, you know, putting this context and, of the profiling context and doing the, the run. and then, s- let's look at what, the profile looks like for just a plus b in, in PyTorch. So, if you're not, you know, normally doing PyTorch, you probably don't think about it. It's like, okay, well, these two tensors just get added. So what's actually going underneath the hood? Right? ![Slide 24](https://cs336.sg.aidenovak.com/media/slides/l6/024.webp) ![Slide 25](https://cs336.sg.aidenovak.com/media/slides/l6/025.webp) ![Slide 26](https://cs336.sg.aidenovak.com/media/slides/l6/026.webp) ![Slide 27](https://cs336.sg.aidenovak.com/media/slides/l6/027.webp) ![Slide 28](https://cs336.sg.aidenovak.com/media/slides/l6/028.webp) ![Slide 29](https://cs336.sg.aidenovak.com/media/slides/l6/029.webp) `27:54` So if you look at, the things that it's, it's calling, there is this, you know, long name, kernel at, scuda functor add. so this is basically, you know, a kernel that adds, to, to, you know, tensors. And, you know, the, the times aren't going to be that interesting because I'm only adding, so that's going to take 100% of the time. But this tells you that, well, underneath the hood, there is this thing called, add. what about matmul? so I'm going to, you know, in PyTorch, I'm doing a at b. similarly, there is this, you know, long name that describes this, particular matmul, you know, kernel. cutlet, F32, F32, 64, 64, 16, and so on. notice that if you change the, the dimensions, so now I'm doing a 128 by 128 matmul, you get, actually a different one. ![Slide 30](https://cs336.sg.aidenovak.com/media/slides/l6/030.webp) `28:56` if you look closely that this is 64, 64, 16, this is 32, 32, 16. Okay, so, underneath the hood, you know, PyTorch, it looks like you're doing just add, but underneath the hood, there could be all sorts of things that are, you know, happening. so observations, here you can see which CUDA kernels are actually being called. These are generally the ones that are, have the long names, and different CUDA kernels are invoked depending on, the, you know, tensor, or dimensions. the name actually tells you something about the implementation as well. So, this name, so cutlet is the NVIDIA CUDA library for, you know, linear algebra. SM100 corresponds to the Blackwell, you know, architecture. So this is a kernel that's specifically designed for, Blackwell. This is FP32. ![Slide 31](https://cs336.sg.aidenovak.com/media/slides/l6/031.webp) `29:47` And then 64, 64, 16 is, the shape of the tile, which we're going to talk about later when we talk about tiling in the context of mammals. `30:01` Okay. So, that's benchmarking and profiling. Just remember to, to do it. ## 6. GeLU 三种实现对比：朴素 / 内置 / torch.compile *[30:07 – 35:21] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=1807)* ![Slide 32](https://cs336.sg.aidenovak.com/media/slides/l6/032.webp) `30:07` I think we make you do it on the assignment, so you have no choice. okay. So, let's apply this, to another example here, the GALU. So, so remember the GALU activation, you know, function is, this, this function, which is a typical non-linearity that's used. Often it's approximated with this, you know, tanH approximation, which is, more computer-friendly. Okay. So, so naively, if you implement the, the GALU, you can do it in PyTorch, just like this. So you take a tensor, you just, you know, take this equation and you just put it into PyTorch, okay? that's fine. `31:01` And you get some number. PyTorch also has a built-in, you know, if you call it, an n-functional GALU, you can get that version as well. So, and you can check that they're, they're the same. You can run, the, the two, you know, GALU versions and make sure that the answer is the, the same for a random input. there's also something that, you know, people probably, some have discovered. It's, it's a, if you, you know, haven't encountered this, this is an important thing to know. Is that you can take any, PyTorch function, you call Torch compile it, on it, and it generates another, you know, function. and it does the same thing. okay. ![Slide 33](https://cs336.sg.aidenovak.com/media/slides/l6/033.webp) `31:48` So, we have three, horses in this race. We have the Naive, implementation. We have the built-in implementation and we have the compiled implementation. So, let's, we can benchmark them. Naive takes, you know, three, I guess this is, 3.75. Built-in is much faster and, compiled is, is much faster as well, but not maybe as quite as, fast as built-in. Okay? So, so what's, what's happening here? ![Slide 34](https://cs336.sg.aidenovak.com/media/slides/l6/034.webp) `32:22` Like, why are, why are these things, you know, different? They're, you know, all compute the same answer, but they have wildly different, performance, you know, characteristics. And so, here we can pull up the profiler and see what's actually going on underneath the hood. So, if you look at the Naive, Galeu here, and just do a profile, this, something's wrong with this view, so I'm going to go to, this. So, you see that, the profile shows, you know, how time is being spent. There's a bunch of different kernels, this binary functor, unary, add, tanh is a, as a, you know, a kernel here. And this corresponds to the fact that in PyTorch, when we write the PyTorch expression, you look at the computation graph, and each, you know, primitive in the computation graph is actually realized in a kernel. Right? And the reason this is slow is that, you know, when you launch a kernel, you know, the, the kernel has to read from HBM, pull it all the way over to your SM, do the computation, and write it back. `33:37` And then the next kernel picks it up from HBM, you know, and then writes it back and so on and so forth. So, you're doing a lot of reads, back, and writes back and forth, because between kernel invocations, you know, things have to go back to, you know, HBM. `33:54` So, and if you look at what the built-in is doing, this is actually not that interesting. There is this, Galu CUDA kernel implementation, which is a single kernel that just implements the Galu. And so why does this exist? Well, because, you know, people use Galu, so someone wrote a kernel for it and put it in the standard library. I mean, this is how, you know, there's no, nothing magical. ![Slide 35](https://cs336.sg.aidenovak.com/media/slides/l6/035.webp) `34:21` So, comp, compilation is really interesting here. I'm not going to say too much about, how it works, but the, it is really, I think, a fascinating, you know, topic where you can take a, a naive implementation, which, you know, remembrance in PyTorch, it has a computation graph, `34:42` which basically, if you look at what's underneath the hood, it's just a single kernel. And this is because it's, it's figured out to, you know, look at the computation graph and essentially write that kernel in, in Triton. So you can see that this is actually a Triton kernel. ![Slide 36](https://cs336.sg.aidenovak.com/media/slides/l6/036.webp) ![Slide 37](https://cs336.sg.aidenovak.com/media/slides/l6/037.webp) ![Slide 38](https://cs336.sg.aidenovak.com/media/slides/l6/038.webp) ![Slide 39](https://cs336.sg.aidenovak.com/media/slides/l6/039.webp) ![Slide 40](https://cs336.sg.aidenovak.com/media/slides/l6/040.webp) ![Slide 41](https://cs336.sg.aidenovak.com/media/slides/l6/041.webp) `35:03` Okay, so, naive implementation, multiple kernels requires multiple reads and writes to and from HBM. There's no, you know, kernel fusion here. this is, slow. The built-in and compiled version, there's one kernel. The, basically, all the operations in the Galu have been fused together into one kernel. ## 7. Triton 入门：面向线程块的编程模型 *[35:21 – 39:38] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=2121)* `35:21` So you read from HBM once, you write to HBM once per, you know, element. and you see that the compiled kernel is a, you know, is a Triton kernel. ![Slide 42](https://cs336.sg.aidenovak.com/media/slides/l6/042.webp) `35:34` Okay, so, that maybe is a good segue to talk about, what this Triton thing is all about. Yeah? what, you know what a kernel is, is it also a, do they just go straight up in CUDA? So I don't, let's see, I mean, I guess this says CUDA kernel implementation, so I imagine someone wrote it in CUDA. `36:03` Okay, any other questions? `36:07` Yeah? Why is the, waste, try, try to learn it faster than this one, this CUDA version? so why is the Triton kernel faster? So the Triton kernel is actually not faster in this, case. Oh, sorry. The, the, the compiled kernel is the, is one Triton kernel, and this is slower than the built-in. Yeah. I think last year when I did this, it was actually closer, or, but, you know, these things change and it's very hardware dependent, and, you know, I think none of this is, like, terribly optimized. This is just giving you the general idea here. ![Slide 43](https://cs336.sg.aidenovak.com/media/slides/l6/043.webp) `36:51` Okay, so let's, write some, Triton kernels. so remember our, programming model. You have a bunch of threads organized by thread groups, thread blocks, and there's a grid of thread blocks. Okay, and, if you were to, you know, write in CUDA, which was originally developed by NVIDIA, and that's been for years the thing that you do when you write kernels. the mental model is what does each thread do? So you write a piece of code which essentially has some ID that identifies which thread you're talking about, and it just, like, executes the code. So the nice thing about this is that it's very closely related to what is actually, happening underneath the hood. and it's, you know, fine-grained, gives you fine-grained control. but, you know, there's cons here, which is that, remember all these threads, are in a thread block. ![Slide 44](https://cs336.sg.aidenovak.com/media/slides/l6/044.webp) `37:55` And, you know, some operations require them to communicate. So what has to happen is that they have to synchronize and, like, they read from HBM all at once, and they have to synchronize and do the computation and then, or, and you have to basically do that, you know, bookkeeping. So if you were doing all element-wise operations, you know, this CUDA is just fine. It doesn't really, matter. But as you get more complex operations, then Triton provides, some value, of the abstraction. So Triton development bio was developed by OpenAI. I think by now it's been pretty, you know, standard. You basically specify what each thread block, you know, does. generally it's powerful enough, especially for, you know, this class and you're getting started. `38:46` If you really go and want to exploit every single new feature of the latest hardware, you know, it might not give you the full flexibility, but, you know, you know, let's not worry about that. And the conceptual framework to think about in Triton is that you're thinking about what does a block do? A block is going to load data into shared memory, operate it on, and write it back to, you know, global memory. So in some sense, these blocks are intermediate point between, thinking about what individual elements are doing, as well as thinking about the general operation. And so in PyTorch, you basically define these huge matrices and you say multiply them together, and that's the sort of atomic operation, and a lot of what you're thinking about is how do I get things into big mammals. Right? And at the, at some level Triton is sort of, you know, ## 8. Triton 实战 1：GeLU（逐元素操作） *[39:39 – 57:27] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=2379)* ![Slide 45](https://cs336.sg.aidenovak.com/media/slides/l6/045.webp) `39:39` a hybrid between that and, you know, the individual elements, as we'll see. Okay. So let's start with, the, the Gal-U example here. so let's define an 8,000, dimensional vector, and start writing some, Triton. Okay. So, so Triton is basically going to be, you know, you write, you know, Python. How many of you have written, Triton before? Just as a sharp hands. Okay. `40:08` Okay. So hopefully this, I'll try to, not that many, so, which is, good because then you won't be bored. okay. So, you know, so this is normal, this is normal PyTorch, right? There's no Triton here, but I'm just preparing. and so what I'm going to do is, I take this tensor, and I'm going to allocate an output, you know, tensor. Okay? Because in Triton, we're not thinking functionally anymore. We're just thinking about moving, you have to explicitly read and write. `40:44` so there's no, like, returning, value. So I'm going to allocate the output tensor which I'm going to, the kernel's going to write to. and then, so this tensor is, can be, you know, arbitrarily big, right? And I can't generally fit this all into one SM because, you know, it's just, it's just too big. so I need to break it up into blocks. So what I'm going to do is you can think about, you know, this, this, x is being this array. I'm going to chop it up into blocks. So the total number of elements is 8,000. I'm just going to set the block size to 1,024, you know, for now. `41:27` and, and then I have 8 blocks. Okay. `41:33` so then I'm going to, use this sort of weird syntax to call the, the kernel. so triton, galeu kernel, this basically in, in brackets tells me, the, essentially the shape of the grid. So basically this says the grid has num blocks, you know, blocks. and I'm going to, for every one of these blocks, invoke, this function, triton, galeu kernel, which I'll talk about in a bit, I'm going to pass in an x, y, num elements, and the block size. Okay? All right. So, unfortunately I'm not going to be able to trace through this, so I'll just show you the, what the code looks like, here. So, let me get rid of that. Okay. `42:31` So this is, probably the, kind of the simplest kernel, you can imagine. So, when you are looking at the triton, galeu kernel, now this is, you know, before we had x and y. So now these are pointers, right? You can think about these as just like, integers. They're addresses. so we have to get comfortable with that. and then we have the number of elements and the block size, which are basically passed, in, actually passed in from here. Okay. So the way to think about it is that for every block, we're going to have this function being, you know, called. `43:20` And, and what happens is the first thing I'm going to do is I'm going to wake up, this, this, you know, block wakes up and say, who am I? well, the program ID is, is basically, the PID is basically identifying the block. So PID would be zero for this block one, two, and three here. Okay. So now I have to figure out what data I'm going to operate on. So that's PID times block size. So this is the, this is the offset into x pointer. So, if PID is zero, then start is zero. If PID is one, then start is block size. `43:59` If PID is two, then it's two times block size and so on. `44:03` Okay. So now I'm going to figure out the, the span I'm operating on. So offsets is start plus, this is the, you know, the Triton, you know, library. a range zero block size, which gives, conceptually gives you the, the integer zero through block size minus one. So offsets is going to be, essentially, you know, for, let's say, block one, it's going to be block size, block size plus one, all the way to blocks two times block size minus, one. `44:39` Okay. So, you know, in this case, num elements divides by block size, but in general, that's not the, the case. So you often see in Triton code, there's this masking, which, says, well, sometimes if the, let's say the tensor only goes up to here, then I'm going to form a, a mask which is, going to be true up until that point and false after that point. And for blocks that are don't, are not the final block, it's going to be just, all ones. `45:11` Okay. So now I've done the, the setup. What I do is I read. And this is basically pointer arithmetic. so x pointer, remember, is the integer that I specialize in memory, a location of where x is. I'm going to add offsets to that, which gives me the, the first, block size number of elements there. you know, according to the mask. If I'm masked out, then I don't, you know, I don't, I don't read it. and then, now, I, I think, you know, you can just think about this as, you know, a, a vector. and then you do your normal computation, `45:47` and, and then you get, you know, y. And y is the same size as x, and then you do tl store, y pointer plus offsets, a y and, and the mask. Okay? So this load reloads from HBM, does some stuff, and then it writes back to HBM. Okay. I'm going to stop there, and take any questions about this first, kind of Triton kernel. yeah. So, like, what would the difference be between this and CUDA? For CUDA, it would just be like the load in the image, and the individual graph is the same, or the type of, and just go back. `46:33` Yeah. So the difference is, what is the difference between CUDA? And the, it's, for this element-wise, it looks pretty much the same, right? And, and, and, in fact, CUDA is even, I think, simpler, because it really is element-wise. You wake up, and you, I, identify with the thread, and then you just operate on that element. now the only thing here is that it's sort of like the vectorized version, where you, you have a block, and you operate on that block. later we'll see that, operating on blocks is, is doing more than, if you do more, something more than element-wise, then CUDA is going to be a lot more annoying to work with. I saw a bunch of questions. Yeah, back there. `47:16` How is this related to if you wanted to use the tensor units? So how is this related to if you want to use the tensor units? so, I'll later show you what this code actually compiles, down into, and, maybe we'll, I'll get, back to that. But the short answer is that you don't control that. Like, the hardware figures out what, where to put things. Yeah? Can you actually walk through what's happening at tl.load from HBM to share to register? What's, like, the step-by-step work date is actually smaller? `47:54` Yeah, so the question is what is actually happening when this, the statement is executed? you know, so the, the short version is that this is in some sense, a lie, right? It's not like the GPU actually, calls the tenant, you know, the Triton library on the GPU and it's like actually executing this code. This is basically for our consumption, to, you know, specify the computation. The compiler takes this, as we'll see later, and writes it into something called PTX and then it will actually do the, you know, the work. But conceptually, like, conceptually, so, so that's mechanically what's happening, which I'll get to later. if you're asking about the conceptual question, the way to think about this is that, you know, this, H, X pointer is, you know, a memory location in, in HBM. And this basically specifies a range of memory locations. And Lo takes those memory locations and returns the data associated with them. And it gets, you know, here I have a local variable called `49:06` X. in practice, this is going to be generally, you know, a, a register or shared memory. Actually, you, you know, you don't actually, you know, Triton kind of figures out what to, what to do there. `49:21` .. `49:33` . `49:55` . Yeah, yeah. So the question is, like, when does this actually get executed? And doesn't this, like, block? let me try to come back to this question. I'll show you the PTX and maybe hopefully it'll provide a bit more context on what's happening. Okay. Any other questions? Okay. So this is a, all the kernels are going to look something like this. So I, ![Slide 46](https://cs336.sg.aidenovak.com/media/slides/l6/046.webp) ![Slide 47](https://cs336.sg.aidenovak.com/media/slides/l6/047.webp) `50:19` I do want to make sure people understand, you know, the, just the general form. Which is that you have generally your inputs, your outputs, you wake up, you figure out, which index you're going to look at, you read, you do some stuff, and then you write to HVM. Okay. so, okay. So let's talk about PTX briefly. so let's see, does this work? Does this link work? Okay. so, when you, you know, write Triton, the compiler generates a PTX code, which is this intermediate assembly language for, you know, GPUs. you know, I am obviously not going to go through all this, but just to `51:11` give you the, the flavor of what this, you know, you know, looks like. let me actually start with observe, observation. So a few things. So, you know, this is now what a thread is actually doing. not a thread block because that's been kind of compiled, you know, away. and so, you know, if you look at, a few, a few notes here. So this LD global is, is basically saying load from, HVM into some registers, and the registers are denoted as, as like the R, R are integer registers, F are, you know, floating point, you know, registers. and then you have statements like, you know, move zero into R5, move zero into R6, and so on. and then you have multiplication. You're going to multiply this register by this, you know, constant, and then, put it in `52:16` this. So this code gets executed, and at the bottom you should see, you know, global, store. So this is just writing back into the HVM. So this gets executed. This is actually the code that, gets executed on, you know, a thread. you know, and Trine is basically, you know, a layer up, you know, above that. one other thing to notice is that, you know, you see all these, you know, kind of blocks, and what's going on there. And what's going on there is, you know, I alluded to thread coarsening, which is that this is one thread, but rather than processing a single element, it's actually processing eight, elements. So the compiler decided that, well, actually this thread is pretty lightweight. It doesn't do that much, so let's just try to `53:05` thicken it up a little bit. so, you know, looking at PTX can give you some sort of appreciation for, you know, what's, what's going on underneath the hood. Yeah? So it makes this for each individual thread? ![Slide 48](https://cs336.sg.aidenovak.com/media/slides/l6/048.webp) `53:24` so the question is, does it make this for each thread? So this is compiled, you know, once, and it's the same piece of code that each thread runs. And the way that the thread, distinguishes itself is that this, this piece of code gets passed, basically the, you know, the thread, you know, ID. so here the cta.x is the block index, and tid.x is the thread index. So, this basically, if this piece of code is running, this tells you which block I'm, in. And, tid.x tells you which thread inside that block I'm in. Okay. `54:20` Okay. any other questions? And I'm trying to figure out how to, okay. `54:31` Okay. Okay. So that's, generally a flavor of, you know, your first trying to kernel, kernel, load from HBM, compute right back to HBM. And, and, and then we saw the ptx, which is like kind of the, the grungy, what actually happens underneath the, the hood. There's still a lot of things that are not specified in a ptx. Like for example, which sms things are operating on, and, you know, the warps and everything. And that's, a lot of those are kind of hardware controlled. So, you don't even, you know, see. Yeah. `55:09` So ptx coded generated by the compiler. So not something that you would normally code. It looks like it's assembly for it. Yeah. So the question is to p, is ptx just, generated by the compiler. So there are people who do write ptx, if you really, you know, think you're better than, the compiler. and I, I think, you know, the NVIDIA compilers are generally pretty mature, but some other, you know, accelerators that are less developed, I think, you know, sometimes you just have to reach in and actually handhold a bit more. But generally, yeah, you shouldn't, need to do that. Yeah. `55:52` So, I, we're doing one of pgx. Like, so I, I, my warp will get scheduled into SM. I'll get to that tf.load and it's almost like a CPU trap call where, like, I'm just waiting for, like, you know, some, you know, something to happen. And, like, now I'm just, I am tuning my thumbs up, so some other warp will get scheduled over me. And then it will do operations. Now that gf.load is done, then I'll reschedule that warp, and then I'll continue on. Is that kind of- Yeah, that's right. So just to repeat the, the, the question, or the comment. `56:26` so if you look at this, this string code, these statements, like the, this load, is going to block for some number of cycles. And then, so this is running on some thread, on some warp, on some SM. And remember the SM is running multiple warps at the same time. So, you know, when you get to that point, you can just, like, you know, find another warp to run. And then when this is done, the warp schedule comes back and, takes over. Why are we taking four warp schedules? `56:58` Sorry. sorry. Why are we coming from last vector? We work for almost four warp schedules. Yeah, so the question is why four warp schedules? I don't know exactly the reason behind that. `57:15` Okay, so now let's go through some other examples. Maybe just as a, kind of a quick preview. We're going to do three more examples. Gel, Gel uses kind of the simplest, you know, form. ## 9. Triton 实战 2：softmax（行内归约） *[57:27 – 64:24] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=3447)* ![Slide 49](https://cs336.sg.aidenovak.com/media/slides/l6/049.webp) `57:27` Even though the, the computation is kind of like, you know, has a lot of messiness, it's just element-wise. So conceptually, in the context of, this, this lecture is actually very simple. now we're going to look at softmax where now, you're going to have to do reduction. but in this case, we're going to think about the case where row fits on a block. And then we're going to, consider the case where row doesn't fit on a block and then we're going to do, go up to MAML. And then hopefully by that point, you'll have all the ingredients that you need to, do the assignment and implement flash attention. Okay. So, so far we looked at element-wise operations. now let's think about other operations that aggregate over multiple, values. ![Slide 50](https://cs336.sg.aidenovak.com/media/slides/l6/050.webp) ![Slide 51](https://cs336.sg.aidenovak.com/media/slides/l6/051.webp) ![Slide 52](https://cs336.sg.aidenovak.com/media/slides/l6/052.webp) `58:15` So, remember what the softmax does, it, it takes a, let's just think of, as a matrix. You exponentiate and normalize each row of a matrix. Okay. So this is used in tension, it's used in, you know, generating probability- output, probabilities. generally a good, good thing to do. so let's just start with a naive implementation, and just keep track of what we're doing. So what's, what's happening. So here I'm defining this tensor and, here's the naive implementation. actually I think this is, it's not an assignment one. `58:58` okay. Anyway. So, here I have an M by N matrix. And so I need to, for every row, you know, compute the, you know, I'm going to compute the max of each row. And this is for numerical stability. I'm going to subtract off the max, and then I'm going to, you know, exponentiate element-wise. And I'm going to, sum and compute a normalization constant for every, every row. And I'm going to divide. and then that's it. ![Slide 53](https://cs336.sg.aidenovak.com/media/slides/l6/053.webp) `59:38` Okay? So, so if you count the number of reads and writes, remember this is just plain PyTorch. So this is a different kernel, this is a different kernel. And unless you call torch.compile, these are going to be different, you know, operations. And each operation is going to read and write, read and write from HPM. Okay? So you have, you know, five, MN, you know, reads, three MN writes. and in principle, you should really only have, you know, much fewer. `60:14` Okay. So this is, you know, piece of code. Kind of makes sense. Everyone should be familiar with what the softmax is doing. Okay. `60:23` So let's, let's write, you know, the, the Triton kernel. and, you know, in some sense the, the, the form is going to be very similar. Just like in the Galu, right? Once you do the scaffolding, the core computation looks very much like the naive version. so what we're going to do is we're going to, say each row is a block. Okay? And why do I make each row a block? Well, because now remember each row I have to normalize and sum. So it's not element wise. `60:56` Softmax is not element wise. But it is sort of, sort of row wise. So the blocks don't interact. so they don't need, and the blocks don't have, you know, shared memory. So that's, that's fine. there's no shared memory across blocks. so then within each row, I'm just going to do some stuff. Okay, so let's see what happens. so again, I'm going to, in Python, I'm going to allocate my output tensor. `61:24` I have this end by end, you know, matrix. I'm going to define the block size as, you know, basically the number of columns. you know, go to an x power of two for good luck. and then the number of blocks is just, the number of rows. Okay, and I'm going to call this, kernel. So how many blocks are there? `61:52` M. right, one for each row. And each block, I'm going to pass the, the input pointer, the output pointer. and then I'm going to also, pass these, you know, you know, strides which tell me, how, how far to move down. okay, so, okay, let me actually go and show you the softmax, kind of, okay. `62:23` okay, so what does this, look like here? so I wake up, I'm on a particular row. Okay? And this is going to give me the, all the columns from zero to the number of block size, which, you know, I'm assuming that's all the, the columns. I'm going to, you know, read. So, so basically I need to figure out which, where to read from memory, and that's going to be the start of my data, plus the, you know, which row. And the row stride basically gives me, you know, every row is, is, basically this is a number of columns, essentially. It's going to tell me where, how far to go down. and then this is x pointers, is basically the addresses of all the data I need to load. `63:15` I load them up. and here I do this thing where if it's, you know, if it's masked out, then I, you know, put minus infinity because I know those are, that's going to be kind of the, the, the equivalent of, a, a zero for the soft max operation. And then this part is, is then surely the same as the naive, soft max. I'm just going to compute the max, subtract it off, exponentiate, sum, and divide. And then I write it back. `63:48` Okay. Yeah. `63:58` Yeah. So this is the version where each block can just, span the entire row. So, yeah, this is sort of the, and you, you can see that, trying makes this very easy, right? Because this is as if you were just normal writing normal PyTorch code. `64:20` So, if everything fits in a block, you basically, the, the, the, the thing is, ## 10. Triton 实战 3：行求和（分块归约） *[64:24 – 71:48] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=3864)* ![Slide 54](https://cs336.sg.aidenovak.com/media/slides/l6/054.webp) `64:24` if you can fit in the entire block, you can just, like, write normal PyTorch almost. Yeah. what if my, number of columns and number of rows are, bigger than block size? Yeah. So, we'll get to that. Now, if the number of columns, rows, in general, it's going to be much larger than that block size. So, we'll come back to that. Yeah. So, if you wanted to do, like, a softmax, like, per column, would you have to do, like, a larger, the, the offset speed? `64:52` so the question is, if you want to do softmax by column, I think that should be fine because here we're tracking these pointers, right? So, the pointers are, can be anything. And here I've, basically, I think all you would have to do is change the stride, here, to basically access the, the, the columns. Actually, it would be here. The column offsets, you would just make, like, the, multiply this by, like, the row stride, I think. Okay. Let's, let's move on. All right. So, now, you know, warming up to the MAML, suppose that your row doesn't fit into a block. ![Slide 55](https://cs336.sg.aidenovak.com/media/slides/l6/055.webp) `65:33` Okay. So, what do you do in this case? So, for example, if you have, let's say, a, a row that's, 400, 4,000 columns, but the block size is only 1,024. So, you can't, you know, you have to do something here. So, here's the strategy. We're going to break up the row into tiles. So, in this case, there's going to be four tiles. And each thread is going to iterate over the tiles and accumulate a sum. And then, finally, at the end, we're going to do the, the, kind of, the, reduction. ![Slide 56](https://cs336.sg.aidenovak.com/media/slides/l6/056.webp) `66:11` by, you know, summing everything that each thread produced. So, I'll, I'll show you an example of this. And now, I'm, I'm sort of switching from softmax to row sum, because it's just easier to think about. okay. So, so here, okay, well, this is not very interesting. The built-in row sum does, what you would expect. so the row sum operation just basically takes, a matrix and then computes, the sum of each row. Okay? So, just no surprises there. `66:47` okay. So, conceptually, what we're going to do here is as follows. Okay, so each block is still in charge of one row. So, that part hasn't changed. So, suppose I'm block one, row one. I wake up and what do I do? I'm going to, remember now there's tiles. And suppose tile zero is columns zero through three, tile one is columns four through seven, col, and then tile two is columns eight through eleven. Okay, so what I'm going to do here is that I'm going to iterate now. `67:20` So, first I basically process, this, all the threads. Basically, each thread keeps a kind of accumulator and, and processes the first tile. And it's going to move on to the next tile. And it's going to add the, the current element to the accumulator. So, here I'm, putting three, one, four, one. And then the, the second loop, iteration, I'm going to add five to this. I get eight. I add nine to this. I get ten, two to four, and six to one, right? `67:54` So, each of these four threads is going to be kind of accumulating, its own thing. And then at the end, you know, I do another, for tile two, I add five and three to their respective accumulators. So, at the end of the day, I have this, you know, a vector of accumulators and then that is actually, you know, that I can just, like, sum up. ![Slide 57](https://cs336.sg.aidenovak.com/media/slides/l6/057.webp) `68:21` Okay, so let's see what the code looks like here. so I'm just going to call the, the row sum kernel and, let's see. `68:33` Okay, so here's what it looks like. So, wake up, I am on a particular row. And, and w, this is what one row looks like. There is one tile, a second tile, a third tile, and so on. So, n is the number of, elements of that row. And block size is this, size of this, this tile. Okay? So, remember, block size is the number of threads. I'm processing more data than the, larger than the block size. `69:04` I just have to do it iteratively now. So, I'm going to loop over all of the tiles. So, start is going to go from zero to block size to two times block size and so on and so forth. So, each, time I'm, you know, jumping across, I, I am doing, you know, calls, basically get the, the, the offsets at the particular tile. and then I'm going to load the data from HBM and I'm going to add it to the accumulation. So, this accumulation is going to be, either in registers or, shared memory. Okay? And then finally, after I loop over all the tiles, I process a role. I get basically, for every thread, I have a accumulator of what that thread has picked up. `69:58` And I can do just through a sum to get a scalar and I write it out. `70:05` Okay? So, this is a little bit more complicated than before because now we have a for loop within a thread. And this is, necessary when your data doesn't fit within a block. `70:18` Yeah? `70:20` . `70:28` Yeah. So, the question is, I think, if you can control where the accumulator resides. At least in this Triton program, you don't explicitly say, and this is up to the, Triton compiler to, figure out where to put it. `70:47` But in general, if the block size is large enough, it has to go and share, share memory. `70:55` Okay. All right. Does this, kind of make sense? So, you know, you, it might, you know, just, just to, I think, make sure people are on the same page. Remember in Aguilu, we also split a row into a bunch of, pieces. But those were blocks, right? And so, each of those pieces was a block that was processed independently. These are not blocks. These are tiles. `71:20` The block corresponds to this whole, row. And, and the block has to basically process all the tiles. And so, this is where it starts to not look like PyTorch because, you're not able to process all your data in one nice kind of, not everything fits into shared memory. Yeah. Is the correct that the accumulator is also stored in shared memory? Yeah. I, let's say the accumulator is stored in shared memory. Yeah. ## 11. Triton 实战 4：矩阵乘 + ReLU（tiling 与共享内存） *[71:51 – 83:54] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=4311)* ![Slide 58](https://cs336.sg.aidenovak.com/media/slides/l6/058.webp) `71:51` Okay. All right. So, let's now maybe go on to our finale, which is MatMul. Okay. So, MatMul, of course, is the bread and butter of deep learning. It's been optimized to death. and, and it's, you know, in some sense, very, you know, fundamental operation. So, you take two matrices, you multiply them. I'm going to add a little bit of a twist here. ![Slide 59](https://cs336.sg.aidenovak.com/media/slides/l6/059.webp) `72:18` I'm going to do a MatMul followed by a ReLU just because, just for kicks. Okay. and this is, well, this happens, right? Because if you have, you know, one linear layer, it's a MatMul, and then you apply a ReLU activation. So, this is not, like, completely out of nowhere. but I'll show you later why I, I, I did this. Okay. So, how do you build a MatMul kernel? All right. ![Slide 60](https://cs336.sg.aidenovak.com/media/slides/l6/060.webp) ![Slide 61](https://cs336.sg.aidenovak.com/media/slides/l6/061.webp) `72:44` So, here's a naive approach. Okay. So, here's my, let's say, matrix. I'm multiplying A times, B. And I'm trying to write out, the matrix C. And A is M by K. B is K by N. So, C is M by N. Okay. `73:04` So, what I'm going to do is I'm going to fix an, one of these elements. Let's say M equals 1, N equals 2. Okay. So, I'm processing, actually, let's do M equals 1, N equals 1. So, I'm doing C5. And then, basically, for every K, I'm going, so, I'm going to, yeah, I'm going to iterate over, over the, the K, the rows of A, and the columns of B. I'm going to read from HBM. I'm going to multiply them, accumulate that, and at the end, I'm going to write out to, this, this element. So, that is a valid MAML kernel. `73:47` Okay. So, what's wrong with it? ![Slide 62](https://cs336.sg.aidenovak.com/media/slides/l6/062.webp) `73:57` It's correct. But, you know, if you look at how many reads and writes it's doing, this is not good. Right? So, basically, for every M and N and K, I have to read from HBM. So, it's on the order of M times K times N reads, number of writes, I mean, it doesn't matter, but this is a bottleneck. Right? And, if you remember from the, the second lecture, if you look at the number of, operations you're doing, divided by the number of bytes that were transferred, that's the arithmetic intensity, which you want to be high. So, the number of operations is M times K times N order, and then the number of reads is also the same, so arithmetic intensity is a constant, which is, not good. ![Slide 63](https://cs336.sg.aidenovak.com/media/slides/l6/063.webp) `74:43` Okay. So, what, if you look carefully, you know, you notice that there is a lot of redundant reads. So, imagine computing C4. you needed to read A4, A5, and A6. And if you compute C5, you're going to have to read those over again as well. So, so if you can just read those once, then you really save on reads and that's, that's great. So, let's try to use, shared memory to do that. So, here's the idealized approach. I'm going to load all of A and B into shared memory, and then I'm going to just compute C. `75:24` So, if I can do that, then I get, now I don't have this cubic number of reads. I only get a quadratic number of reads, which means I get the error of intensity of order N, which in the second lecture I sort of, you know, said was a kind of ideal thing you could hope for. `75:45` Okay? So, this, if you can do that, that's, that's great because you basically, there's no redundant reads. You read everything once into shared memory, you do the computation, and you write it back. Okay, but what's the problem with this idealized approach? ![Slide 64](https://cs336.sg.aidenovak.com/media/slides/l6/064.webp) `76:02` the problem is that A and B are usually too large to fit into shared memory. ![Slide 65](https://cs336.sg.aidenovak.com/media/slides/l6/065.webp) `76:08` Okay. Right? So, so then, what do you have to do now? Okay, so the idea here is this, you know, very classic idea of tiling. And the idea basically is to, well, fit as much you can to share memory as you can, essentially. So, in some sense, it's, it's going to be, look like the naive, it's going to globally look like the naive approach, but locally look like the idealized approach is the way to, you know, think about it. so here's the, you know, the picture. I think Tatsu showed this as well. So, what we're going to do is we're going to take this matrix C, and instead of, remember the naive approach just said for every element, I'm going to compute it. `76:51` But now I'm just going to say for every tile, I'm going to do it. So, I break it up into tiles. and each of these tiles is going to be a thread block. And we're going to have a bunch of threads that is responsible for computing this. Now, a different tile is going to be computed, you know, completely separately by another, thread block. Okay, so I'm going to, you know, imagine I'm in Triton, I'm off, I wake up, I, I'm looking at this thread block. So, what do I have to do? Well, in some sense, it's the same as the, you know, naive approach where for every, you know, row tile of A, I'm going to go scan across the rows, and for every column tile of B, I'm going to load the, the corresponding tile, A, and the corresponding, so say I'm here, and the corresponding tile from B into shared memory. I'm going to multiply these two together, and, and that's going to be, you know, you know, like the, kind of the idealized approach. ![Slide 66](https://cs336.sg.aidenovak.com/media/slides/l6/066.webp) `77:58` And I'm going to accumulate that, in the partial sum, and this is all sitting in shared memory. And then after I finish all this, the sweep of the row and the column here, then I can finally write this output tile to HBM. ![Slide 67](https://cs336.sg.aidenovak.com/media/slides/l6/067.webp) ![Slide 68](https://cs336.sg.aidenovak.com/media/slides/l6/068.webp) ![Slide 69](https://cs336.sg.aidenovak.com/media/slides/l6/069.webp) `78:15` Okay? So that's, you know, the conceptually what's happening. and then the earth mega intensity here now goes up to order tile size. So you can generally not reach order n because that would require you to fit everything into shared memory, but, you know, if your tiles are big, then that's, you know, still not too bad. Okay. So just as kind of a bonus, you know, while you're doing all, writing a kernel for this anyway, sometimes if you want to apply an element-wise activation function, it's very easy to just, like, put it on at the end. And this is, you know, kernel fusion. Okay. So very quickly, the implementation here, so, oops. `79:03` Okay. So, just a kind of reminder about strides since that's just going to show up. So a tensor is, you know, multidimensional array, but in memory it's linearized. And strides of a tensor tells you basically how to map from a multidimensional index, such as a row, a, a comma column, into an actual index. And basically what you do is multiply the row by the stride, the, plus the column by the stride of the column. Okay. So in this case, every time you go advance to an X row, you go four, positions in your memory. And every time you advance a column, you go one. And this, if it were to transpose, it would be flipped. `79:46` Okay. So what does this, kernel, look like here? so there's, okay. So the launch is not interesting. so you wake up and you are on tile M and M. So it's like, okay, I'm responsible for computing, the C, matrix, but for the M slash, comma N tile. there's a bunch of, like, index manipulation which, you can kind of, I'll just, gloss through, but it's, it's, it's straightforward but a little bit, you know, you just have to, track the, the indices. So this basically tells you which rows of A, the matrix A I'm looking at, which rows of, sorry, which columns of B I'm looking at. And then this is just the numbers one through, one through K. ![Slide 70](https://cs336.sg.aidenovak.com/media/slides/l6/070.webp) `80:38` then I'm going to get the, the pointers into A and B at my tile location. and then I'm going to set up this accumulator matrix. This is going to be in shared memory. there's M by N. And then this is going to look kind of like the, row reduction, right? You have this, sum over all the, the tiles. But instead of just going across the, the tiles, I'm now going across the row tiles and also simultaneously down the column tiles of B. I load A, the small matrix, I load B, the small, B tile, and then I perform this dot. So remember, whenever things are in shared memory, things look like PyTorch and I can just, like, say, MAML it and, you know, it'll do the, do the thing. `81:23` And then, and then, and then I advance, to the next row tile of A and the next column tile of B. and this is like, you know, just, the bonus is that if I wanted to apply, LMM-wise non-linearity, I might as well, you know, do that here. before I write it out to HBM, I can do some, any sort of operation on it. And then finally, I just write it out. `81:51` Okay. So there's some sort of indices that you have to pay attention to, but hopefully the, form of the algorithm is, is clear. `82:06` Okay. Any questions about that? `82:12` All right. Maybe I'll just, summarize and you can ask me later. So today we talked about the, there's a programming model, which is, you know, you're talking about, either PyTorch or Triton or PTX even. This is, the, you know, what the kind of programmer can control. Right? And even PTX, you can write PTX and you can control every, you know, specialize however you want. but this is not the full picture because the reality is that your code has to run on hardware and there's only a finite number of SMs, a finite number of banks and, and, you know, sizing of memory and registers all are finite. So you come in with your big matrix and transformer and you want to, you have to kind of fit in with the constraints of, the, the hardware. so that's why, you know, benchmarking and profiling are really important to understand how the messiness of the hardware, translates to, you know, performance. `83:15` we talked about Triton, which is, I think, a pretty nice, neat language to think about, thread blocks. Hopefully, by now, you can appreciate that, things are easier to think about thread blocks than individual threads because you don't have to think about explicitly synchronizing threads or doing shared memory. And the way to think about it is that you figure out, you have your computation, you break it down into these thread blocks where you just need to read from shared memory, do some stuff, and write it back into HBM. And then we saw some examples of increasing, difficulty, element-wise is the easiest. ## 12. 总结与下一讲预告 *[83:54 – 86:40] · [▶ 从这里看视频](https://www.youtube.com/watch?v=xnDHaNUvHBg&t=5034)* ![Slide 1](https://cs336.sg.aidenovak.com/media/slides/l6/001.webp) `83:54` And then reduction over a row, reduction where the, it doesn't, fit into a row, and that's sort of, we introduced kind of the, baby tiling. And then MAML is kind of the, the, the canonical example where you actually do tiling. Okay, so that's all I'm going to say about how to program a single GPU next time. We're going to go to more GPU and talk about, multi GPU programming. yeah, question. So if I want to write valid kernels, what are alternatives that I have to Triton? And for each of my choices that I have, like how close could I get to the optimal level? so the question is, what are alternatives to, to Triton? so, you know, there's a trade-off between, every language has sort of inductive bias, right? `84:46` That makes certain things easier and certain things harder. So most of what we do, Triton was built, you know, by people who, you know, train transformers. So anything involving, transformers, I think it's going to be relatively, you know, easy there. of course in the extreme you can always go to PTX and, write that, but, you know, wouldn't advise that as a first step. There are a bunch of other, language, sort of libraries. There's, you know, Thunder Kittens, there's, you know, cute, you know, various DSLs that allow you to give you, you know, they're not necessarily, necessarily comparable, either up or down the stack. They just give you different, you know, characteristics. Yeah. You know, high-dimensional tensor, you know, multiplications or high-dimensional tensor, processing. `85:39` what is the best approach? Is that, is that, like, what I'm thinking is, can I load, unload the whole tensor, both of the tensors on, you know, a bunch of threads or thread blocks on my, GPU at the same time and compute them? Or would it be better to, sort of, you know, take each element, just what would it be alright, particularly, particularly component of, of my tensor, then process it and then write it back in the end? Yeah, so the, briefly the question, and we should probably wrap up, is, is it better to kind of read all out once or maybe process individual elements, at time? Yeah. I think it's hard to answer this in the abstract. It kind of depends on the nature of the computation. maybe I, we can talk offline about that. Okay. ![Slide 2](https://cs336.sg.aidenovak.com/media/slides/l6/002.webp) `86:34` Alright. See you next time. Okay. Okay. `86:40` Okay. Okay.