6 ms·
Hi, I'm the author. Thanks for sharing, was great to wake up to my blog post on the front page! Would love to hear any feedback or if I missed anything.
by reasonableklout 2y ago
Hi, I'm the author. Thanks for sharing, was great to wake up to my blog post on the front page! Would love to hear any feedback or if I missed anything.
- deleted 2y ago[deleted]
- jmorgan 2y agoThank you for writing this!
- freetonik 2y agoOff topic, but would you consider adding an RSS feed to your blog?
- upghost 2y agoI tried this (and backprop) in GPGPU compute shaders once. One issue that really tripped me up is matrix multiplication along dimensions that are not powers of 2. It will take me awhile to learn to read CUDA code -- is this just an easier thing to do in CUDA or did you figure out a good trick for it?
- ryao 2y agoI know a trick, although it relies on the closed source cuBLAS library. The matrix multiplication in non-batched forward passes is really matrix-vector multiplication, but you can shoehorn that into matrix-matrix multiplication by setting the right dimension to 1. cuBLAS uses column major ordering to match Fortran instead of the row major ordering of C/C++, although you can swap transa and transb, m and n, a and b, lda and ldb; and if you are using cublasGemmEx() to use tensor cores, atype and btype. This relies on the property of mathematics where A * B = C and B’ * A’ = C’ for matrices. You can experiment with this on a CPU using a BLAS library like OpenBLAS or the Intel MKL before trying to do this on a GPU. That said, to do matrix vector multiplication (GEMV), you just need to compute dot products on all of the rows of the matrix with the vector to get your output vector. You could probably just handle each dot product in power of 2 chunks on the GPU to compute partial sums. When you get to the last chunk that is not a power of 2, just have the threads that go pass the end of the dot product contribute 0 to the partial sums. Then you would have the threads handle the horizontal sum of the partial sums to finish the computation of an entry in the output vector. As for matrix-matrix multiplication (GEMM), that is much harder to do in a performant manner. If you do it by treating the multiplication as a series of matrix-vector multiplications, you will be bandwidth limited rather than compute limited. You would need to implement tiling to get good performance, but there are many tricks needed and it is very hard to outperform cuBLAS on Nvidia hardware. Here is one person’s attempts to do it: https://siboehm.com/articles/22/CUDA-MMM https://siboehm.com/articles/22/CUDA-MMM Here is a fairly cryptic description of how to get close to cuBLAS that mentions a few key techniques: https://accu.org/journals/overload/32/181/schuetze/ https://accu.org/journals/overload/32/181/schuetze/ Here is another guy’s attempts that actually managed to outperform cuBLAS on the H100: https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog https://cudaforfun.substack.com/p/outperforming-cublas-on-h1...
- upghost 2y agoWow, great answer. That deserves more upvotes than it's probably going to get. So what that means for compute shaders is I will need to partition the vectors being multiplied into power of 2 sized partitions and repeatedly call the kernels, but with offsets appropriate to the partition size, and then have the overrun contribute zero as you said. Obviously not as efficient as doing it all GPU side, still faster than doing it on the CPU. Thank you!
- ryao 2y agoYou don’t need to repeatedly call the kernels. You just need to loop inside the kernels to have each process every nth multiplication according to their thread id, adding it to the thread’s partial sum and have a branch in the loop skip the addition (or add 0) when you exceed the range. Afterward, you would do a horizontal sum of the partial sums to get the value for the corresponding field in the output vector, and you would have a “grid” do this for all members of the output vector. The original poster’s mathmul kernel does this. See the example under “we can perform a warp sum reduction to combine results from all threads”: https://andrewkchan.dev/posts/yalm.html#section-3.2 https://andrewkchan.dev/posts/yalm.html#section-3.2 It does “if (i >= d) return;“ to handle non-multiples in one dimension and skips the addition in its dot product, matmul_row(), to handle non-multiples in the other dimension. In CUDA, you have threads that are organized into 32 thread warps, that are organized into blocks of warps, that are organized into grids. Threads within a block can share memory, but blocks within a grid cannot. The hardware will effectively loop over blocks in the grid until it finishes. The dimension for the output vector does not require sharing memory, so we can set that to the grid size. The hardware might actually round the grid up, so “if (i >= d) return;” is a check for the case that we are executing past the end of the output vector in case the hardware schedules pass the end of the vector. As for the other dimension, we only have a limited number of threads (in this example 32, which is a single warp), so we need each to do partial sums of the multiplications in the dot product. We essentially break the dot product into thread count (32 here) chunks. Each thread will know which multiplication in the thread count chunk to compute based on its thread id, and we keep adding the thread count (32 here) to that to get the next one until we have finished. The for loop’s loop condition protects us from going past the end of the dot product dimension. Then we have thread count number of partial sums, so we need to do a horizontal sum across the threads in a block to get the final sum. This part is a bit magical and is done via a warp sum reduction. Since the threads share memory, you basically have them start summing their neighbors’ values into their own. Basically, all threads find pairs and one of the pair does the sum, cutting the threads doing the sum in half. Then the ones that did the sums pair up again. This again cuts the number in half. This continues until only 1 thread remains. Then only the remaining thread is allowed to write to “out[i] = rowSum” due to “if (threadIdx.x == 0)”. The warp itself is somewhat special since it can do a really fast horizontal sum via shuffles that cannot be done when using thread blocks ghat utilize multiple warps. That is why the author chose to have the block size be exactly 1 warp rather than a multiple of warps. It helps to think about the threads in the block as a spreadsheet. The columns are the threads and the rows represent instructions. Every thread executes the same instruction, and when a branch says not to do something for only some threads, the ones excluded do noops until execution reaches a point where they all do the same thing. An if else similarly has noops in threads not executing true, followed by noops in the threads not executing false. This is called branch divergence. There are diagrams showing a spreadsheet like view of this, such as the following, which shows a cutdown example: https://image1.slideserve.com/2427671/branch-divergence-l.jpg https://image1.slideserve.com/2427671/branch-divergence-l.jp... I have mentioned a horizontal sum multiple times. That just means summing across the columns inside rows of the “spreadsheet”. It is called horizontal since rows are horizontal. Coincidentally, the partial sum calculation is called a vertical sum since the columns are vertical. Now we need to consider the grid, which is yet another dimension and is the dimension along which we populate the output vector. Basically, you can view the grid as doing multiple spreadsheets in parallel, but not all at once. One comical way to imagine it is to imagine having printed the same spreadsheet hundreds of times an and then feeding them into a paper shredder head first. A normal paper shredder can only do a few at a time, and similarly, the GPU can only do so many at a time too. You can imagine the processing to proceed along the rows in all of the spreadsheets being shredded at the same time.
- ryao 2y agoIf you ever want to do this in pure C, I show how here: https://news.ycombinator.com/item?id=42432802 https://news.ycombinator.com/item?id=42432802 Technically, the device side code is in CUDA C, which is really C++, but no one has written a C compiler that can compile a lightly extended version of ANSI C to PTX, so the device code must be C style CUDA C there, CUDA FORTRAN or hand written PTX assembly. I went with what was easiest, which was C style CUDA C.
- ryao 2y agoYou can probably get this to run slightly faster using CUDA graphs, which should reduce the amount of communication between the GPU and CPU. I notice that llama.cpp uses them, but glancing at your blog post, it appears that you do not. You might want to consider switching to bfloat16. The model weights are published as bfloat16. LLMs are more sensitive to the exponent bits than the mantissa bits, so using fp16 over bfloat16 should slightly hurt the model’s intelligence. Using bfloat16 would require you to use the tensor cores instead of the CUDA cores. Interestingly, llama.cpp performs very poorly when asked to use bfloat16, so performance comparisons of a bfloat16 would be best done against llama.cpp using fp16, since there is no reason for performance of fp16 and bfloat16 to be different, other than a lack of optimization. Using the tensor cores should not make your forward pass any faster despite the additional computational power since the GEMV operations that your forward pass uses should be memory bandwidth bound, although if you ever decide to try doing multiple queries in parallel, you will see a performance boost from being able to use that additional compute via GEMM operations. If you write a prompt processing version of the forward pass, you can use GEMM operations to process all prompt tokens simultaneously even without doing multiple queries in parallel. This would give you fast prompt processing speeds and open another front on which you can try to outdo llama.cpp.
- reasonableklout 2y agoInteresting! CUDA graphs indeed look very promising (https://developer.nvidia.com/blog/cuda-graphs/ https://developer.nvidia.com/blog/cuda-graphs/). I'd expected that kernel dispatch is not a problem because it is done asynchronously by the host and profiles suggested that the device thread is rarely idle, but from the NVIDIA blog I see the effective time per kernel goes from 3.8usec to 3.4usec. I guess there is still dispatch overhead that occurs on the device? Re: bfloat16. Yes, this is a good suggestion and I'm aware that bfloat16 is better suited for ML applications for the reason you mentioned (the range being more important than the precision) while offering the same footprint as float16. TBH, the main reason I haven't done this is because of lack of native support on my CPU (EPYC 7702P). This makes it a bit more annoying to maintain the format parity for the CPU and GPU backends, which is nice because it allows for easy testing, although if I wanted to preserve that I could just write non-performant test-only CPU code for bfloat16. Prompt pre-filling optimizations using GEMMs are great and I plan to explore this when I have free time. I talk about this in the "What's next" section briefly. Optimizing GEMM is a rabbit hole, as you point out in some of the other blog posts you shared. Will be a good next challenge!