Optimizing Matrix Multiplication on NVIDIA H100 GPUs: A Deep Dive

• 00:00:42 Introduction: Bronal (Pel) discusses his blog post on outperforming cuBLAS in matrix multiplication on NVIDIA's H100 GPU, inspired by Seb Wlasny's A100 article.

• 00:01:21 Focus: The talk centers on optimizing matrix multiplication kernels on the H100 GPU, primarily using bf16 data types with fp32 accumulation for enhanced accuracy.

• 00:01:29 Learning Resources: Recommended resources include "Programming Massively Parallel Processors," Simon's Anthropic blog, Andrej Karpathy's "LMC", and NVIDIA's C++ programming guide.

• 00:02:54 Challenge: Limited educational content exists on H100 architecture and maximizing its features for state-of-the-art performance.

• 00:03:49 GPU Architecture Overview: H100 GPU features high bandwidth memory, 132 processors, shared L2 cache, and tensor cores.

• 00:04:37 Tensor Cores: Tensor cores are hardware units that accelerate matrix multiplication. On H100, they support specific matrix sizes, like 64x16 for A and B, with N varying from 8 to 256, requiring PTX assembly for programming.

• 00:06:39 Traditional Matrix Multiplication Parallelization: Dividing the C matrix into tiles, assigning each tile to a thread block, loading portions of A and B into shared memory (BM cross BK), and accumulating results in registers before writing back to global memory.

• 00:09:14 Benchmarking Gap: cuBLAS achieves 76 teraflops, while a basic algorithm reaches only 31.8, highlighting a need for H100-specific optimizations.

• 00:10:11 Tensor Core Instruction: Single hardware instruction can multiply matrices residing in shared memory and store the results in registers, making use of 128 threads.

• 00:11:12 PTX Assembly Requirement: Implementation requires PTX assembly, as Cuda doesn't provide the necessary abstraction. cuBLAS is optimized for PTX, not CUDA.

• 00:12:36 CuTLASS Abstraction: Cutlass is the mode of abstraction for this new GPU stuff.

• 00:13:50 First Generation ISA (SM90a): There is no forward compatibility, and this specific instruction is not support in Blackwell.

• 00:15:11 Memory Considerations: A and B storage consumes around 2.5 KB of shared memory, well within the H100's capabilities; C's register requirements are distributed over 128 threads.

• 00:17:15 Tensor Core Memory Layout Requirements: Tensor cores demand specific memory layouts for A and B to avoid shared bank conflicts, with NVIDIA providing a complex (and reportedly incorrect) diagram for guidance.

• 00:17:15 Stanford: spent a few days trying to reproduce this, but this indexing the diagram is wrong so it's just super hard to load these things

• 00:20:23 Tensor Memory Accelerator (TMA): TMA helps load rectangular tiles from global memory into continuous shared memory asynchronously.

• 00:20:30 Tensor Memory Accelerator (TMA): separate Hardware unit that lets us do this asynchronously so we can just tell the TMA that hey I need this style load into this shared memory address it will do to me at some point and just notify me when it's done.

• 00:21:09 TMA-based Implementation: Code structure involves loading chunks of A and B into shared memory using TMA, performing tensor core multiplication, and storing results.

• 00:23:37 Asynchronous Load Performance Boost: Separating loading and computation through asynchronous loads yields a performance improvement, reaching 317 teraflops, equivalent to approximately 10x

• 00:23:55 Hyperparameter Tuning: Adjusting tile sizes and using larger tensor core instructions boosts performance further, attaining 423 teraflops.

• 00:25:50 Asynchronous Processing: Separate load/compute tasks to decouple.

• 00:25:50 First thought: CPU mindset -> one guy just loads and this guy is Computing so that they can happen in a pipeline

• 00:26:29 Interruption & Comments by Chris from cuTLASS Core Team: Highlighted that the public PTX documentation regarding shared memory layouts is incorrect, emphasizing the use of unit-tested cuTLASS layouts.

• 00:28:55 Thread Assignment Strategy: 128 threads load and the other 128 do the Matrix multiplication

• 00:29:18 Nvidia recommends: that we separate these things in multiples of 128 threads

• 00:29:41 Shared Queue Implementation: Utilizes a barrier for synchronization.

• 00:30:48 Add a Buffer: If you load more things and just store then the tensor guy can just keep Computing it

• 00:31:50 Add a Buffer: load as many of them as we want till we have enough memory

• 00:32:32 Try to exhaust: the shared memory.

• 00:32:32 bandwidth: usually saturated.

• 00:32:53 set the value: of this shared Q so that we ex exhaust all of the shared memory which is around 256 kb per SM

• 00:34:04 GPU Historically: managed all of that for registers, automatically. Using asynchronous load all is developer managed now!

• 00:34:28 Performance Boost from Parallelism: Asynchronous loads lead to a performance jump from 59% to 70% of cuBLAS.

• 00:35:02 Pushing Tile Size Limits: Employing the largest available tensor core instruction (64x256x16) proves advantageous.

• 00:36:33 Splitting Tile Size (Multiple Consumers): Splitting the tile across the M Dimension, splitting the tile acros the M Dimension into 2 64 cross 12 256 tiles. Both of these can use one tensor Coe instruction.

• 00:38:12 Scheduling Considerations (Thread Block Overlap): Instead of focusing on one thread, a wider look sees optimizations that the SM will load some things and to tensor course then do work.

• 00:40:37 Knife Scheduling Strategy: We have a good distribution of SM for the first time and these 132 this 132 things will be running in parallel.

• 00:40:37 contiguous: the blue portion of the Matrix goes to First SM and the green portion of the Matrix goes to 14th. so for all the 132 SMs will receive some tiles looks right like they access some contiguous tiles just assign them hope it all works out we should get some performance boost right sadly not.

• 00:42:17 Contiguous Memory Accesses (L2 Cache Usage): Arranging the initial tiles loaded by each SM to be contiguous in memory improves L2 cache utilization, resulting in a performance boost to 660 teraflops. The other ones are contiguous in the Rectangle.

• 00:44:17 Fast barriers: 92% of cuBLAS using ptX code to fully use them! (Not fully covered due to complexity).

• 00:44:47 Optimization: Implement this.

• 00:45:02 TMA Multicast: *Instead of a TMA like one TMA call from each of the two SMS make only one TMA call from one of the SMS. Use multicasting to reduce loading (block clusters).

• 00:45:02 TMA Multicast: Instead of doing one call between the two SMS we do it together and get a 102% performance boost

• 00:45:02 TMA Multicast and Performance Boost: Using TMA multicast with thread block clusters improves performance, slightly exceeding cuBLAS.

• 00:45:23 Nvidia implements this as clusters: size two.

• 00:45:23 TMA Multicast and Performance Boost: You need to make this for the shared memory across the clusters

• 00:47:52 Results Achieved: Surpassed cuBLAS for a specific sizes/setups

• 00:47:52 In theory: support clusters up to size eight Nvidia

• 00:47:52 communication: override becomes too large

• 00:49:04 Limited SMS: cluster size can be limited due to communication overhead.

• 00:49:13 A very quick feedback: There needs to be an even faster to optimize the

• 00:49:39 General Take Away: The code optimizations from this point on are for a very specific setup. It is about having fun with it.

• 00:51:05 More Micro Optimization: More optimizations on the PTX assembly and the barriers.

• 00:51:05 *Tensor Core *There is a flag to not accumulate as needed! *So now you get a performance boost!

• **00:51:18 **Another cute trick *Skip L2 Cache! *The registers go to global memory, but we need to reduce stress. *So the memory is saved for A and B, and the data being written and not occupying the data!

• 00:51:42 All the sudden You get a performance boost!

• 00:53:49 Next Optimization Let's Do The ASYNC Store!

• 00:53:49 Memory stores in the registers are: A memory bottleneck!

• 00:54:19 Async Memory Store: Write a small pointer. Let some async operator push at some time

• 00:54:41 Now you get: More parrallel stores!

• 00:55:03 The Big Guns Scheduling the memory is super efficient!

• 00:55:46 Shifted 3 and four to reuse things from the old part!

• 00:57:25 This is using the big big Hilberd curve (recursive thing)

• 00:57:25 The hbert curve has been applied.

• 00:58:06 Some of this code optimization depends on a few tricks but the PTX is still reliable.

• 00:59:13 I'm not really sure about the instruction scheduling

• 01:00:30 The code for memory is now super tight.

• 01:00:53 All that code that we went over to write this out the compiler does it by default now. So a modern compiler will do this.

• 01:01:16 You need some padding and the code that makes it good is now standard so we can remove some code.

• 01:01:21 This one is just a bit More of the details!

• 01:01:50 St Matrix the barrier between it is reduced!

• 01:02:51 Cuda still needs to maintain some complexity.

• 01:04:53 Eric working as well

• 01:07:18 About engergy concumption You get to where the power limited.

• 01:07:18 *Energy consumption is important.

• 01:08:00 In this case it can be done in software to lower the code.

• 01:08:29 HORUS has done a very good test case.

• 01:08:29 HORUS Blog Post 333w out of 330W this is power bottle necked

• 01:10:51 Some of the results will show low clocks cause power use matters

• 01:11:18 H200 can only run 1450MHZ while the GH200 can run 2MHZ

• 01:11:18 so to get some test case you might want to be at the edge of things.

• **01:12:43 For the last few years NVIDIA did release a new compiler update

• 01:12:43 The new Cuda now has a new compiler For small things they are about the same as before. The big stuff is a still the same improvement of about 1%

• 01:12:50 There is hope

• 01:13:16 The speed here takes about two years!

• 01:13:16 In short It takes time for the compiler to catch up!

• 01:13:34 You need to always test The results.

• 01:13:34 Who do I need to stop sharing!

• 01:13:34 Last Take Away Points Do people want to see code?*

• 01:14:03 This is not production-ready There a lot of test and see as the final step!

• 01:14:15 It doesn't matter ** what you think! *You still need to do testing