Matrix Multiplication on Blackwell
6 days ago
- #GPU Programming
- #Blackwell Architecture
- #Matrix Multiplication
- 矩阵乘法是大型语言模型的核心运算，像Llama 8B这样的模型中超过83%的运行时间都花在各种矩阵乘法的变体上。
- GPU提供大规模并行性（例如B200最多可处理151,552个线程）和专用张量核心，通过MMA运算实现高效的矩阵乘法。
- Blackwell GPU引入了第五代张量核心，支持256x256x16 MMA运算，利用张量内存减少寄存器使用，并增强了流水线处理能力。
- GPU编程采用SIMT模型，线程并行计算元素，组织为warp、块（CTA）和集群以实现内存共享。
- 从Ampere的异步拷贝到Hopper的TMA和WGMMA，再到Blackwell的张量内存，通过重叠数据传输和计算提高了性能。
- 一个简单的Mojo GPU矩阵乘法内核实现了5 TFLOPs，但cuBLAS在B200上达到了1763 TFLOPs，这突显了利用Blackwell特性进行优化的潜力。