I am broadly interested in High Performance Computing, GPU Programming, and Deep Learning. You can find all of my code on my Github <...
TL;DR: I have spent the last few months learning about and implementing a 2:4 bfloat16 Sparse GEMM kernel in CUDA. On a range of (nice) problem sizes it achieves between 94% and 98% the throughput of cuSPARSE. There are many excellent posts covering dense GEMM optimization in CUDA, building up the final kernel piece-by-piece, this post is not that. Instead, we will provide an overview of how…
Memory coalescing occurs when all threads in a warp access a contiguous chunk of data from global memory. The CUDA C++ Best Practices guide states this in a more precise way: ... the concurrent accesses of the threads of a warp will coalesce into a number of transactions equal to the number of 32-byte transactions necessary to service all of the threads of the warp. For example, if all threads in…
TL;DR: I develop a IO-aware kernel for the RG-LRU in Triton. When training on sequences greater than 4K tokens, the kernel is faster than PyTorch's Flash Attention 2 implementation. Introduction Hawk is an RNN proposed by Google Deepmind in “ Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models”. 1 At its core is the Real-Gated Linear Recurrent Unit (RG-LRU)…
TL;DR I improve upon my earlier codebase by implementing ZeRO-1 optimizer sharding using a combination of xmap and pjit . The resultant method is more performant and scales better across multiple TPU hosts achieving 67% MFU on a TPU v3-32. I use this method to train a 1.3B parameter decoder-only transformer on 200B tokens. The codebase I wrote can be found here . Background For the past 4 months,…