QuACK GEMM内核深度剖析：Hopper与Blackwell架构性能对比

QuACK GEMM内核深度剖析：Hopper与Blackwell架构性能对比

【免费下载链接】quackA Quirky Assortment of CuTe Kernels项目地址: https://gitcode.com/gh_mirrors/quack15/quack

QuACK（Quirky Assortment of CuTe Kernels）是一个基于CuTe DSL的高性能CUDA内核库，专门针对现代GPU架构进行优化。本文将深入剖析QuACK中的GEMM（通用矩阵乘法）内核，对比NVIDIA Hopper（SM90）与Blackwell（SM100）架构的性能差异，帮助开发者理解在不同GPU架构上的优化策略。

架构演进：从Hopper到Blackwell

QuACK的GEMM实现针对不同GPU架构采用了专门的优化策略。Hopper架构（SM90）和Blackwell架构（SM100）在硬件特性上有着显著差异，QuACK通过不同的内核实现充分利用了各自的优势。

Hopper架构（SM90）的核心特性：

• 支持WGMMA（Warp Group MMA）指令，可直接从共享内存读取数据
• 支持TMA多播，实现跨CTA的数据广播
• 最大集群大小为4，共享内存为228KB
• 使用WGMMA进行矩阵乘加运算

Blackwell架构（SM100）的改进：

• 引入TMEM（Tensor Memory）和UMMA（Unified MMA）指令
• 支持更大的集群大小（最大16）
• 支持硬件随机舍入
• 引入CLC（Cluster Launch Control）进行持久化内核调度

性能对比分析

通过QuACK的基准测试，我们可以看到两个架构在不同工作负载下的性能表现：

GEMM内核的关键差异：

1. 内存访问模式

Hopper架构使用TMA（Tensor Memory Accelerator）多播技术，允许单个TMA加载操作将数据广播到集群中的多个CTA。而SM120（RTX 5090）虽然支持TMA加载和存储，但不支持多播功能，这影响了多CTA GEMM的数据共享效率。

2. 计算核心差异

• Hopper：使用WGMMA指令，直接从SMEM读取A/B矩阵
• Blackwell：使用tcgen05.mma指令，将累加器写入TMEM
• SM120：使用warp级MMA指令，需要通过ldmatrix从SMEM复制到RMEM

3. 集群支持能力

Hopper的最大集群大小为4，而Blackwell支持最大16的集群大小。在QuACK的实现中，这直接影响了几何内核的并行效率。

实际性能数据

根据QuACK的基准测试结果，我们可以看到：

在大型矩阵运算中（M=N=K=4096）：

• Hopper架构在fp16精度下达到约1500+ TFLOP/s
• Blackwell架构在相同精度下性能提升15-20%
• SM120由于缺少TMA多播，在多CTA场景下性能受限

内存带宽利用率：

• Hopper：充分利用228KB共享内存
• Blackwell：优化TMEM使用，减少寄存器压力
• SM120：受限于99KB共享内存，需要更紧凑的内存布局

优化策略对比

Hopper架构优化

在quack/gemm_sm90.py中，Hopper GEMM实现采用了以下策略：

1. TMA多播优化：减少L2内存流量
2. WGMMA直接访问：直接从SMEM读取数据
3. 多级流水线：重叠计算和内存访问

Blackwell架构优化

在quack/gemm_sm100.py中，Blackwell GEMM实现的关键特性：

1. 持久化瓦片调度：更好地重叠内存加载/存储与MMA操作
2. Warp专业化：避免主循环加载和MMA之间的显式流水线
3. TMEM利用：专用张量内存区域

SM120架构适配

在quack/gemm_sm120.py中，针对SM120的适配：

1. 显式SMEM→RMEM复制：通过ldmatrix指令
2. 集群限制：cluster_shape=(1,1,1)避免多播依赖
3. 共享内存优化：适应99KB限制

配置参数对比

QuACK的GEMM配置在quack/gemm_config.py中提供了针对不同架构的优化参数：

Hopper典型配置：

tile_shape_mnk = (128, 256, 64) cluster_shape_mnk = (2, 1, 1)

Blackwell典型配置：

mma_tiler_mnk = (128, 256, 64) cluster_shape_mnk = (4, 2, 1) use_2cta_instrs = True

性能调优建议

1. 数据类型选择

• fp16/bf16：两个架构都提供良好支持
• fp8：Blackwell提供更好的fp8支持
• int8/uint8：仅Blackwell原生支持

2. 瓦片大小优化

• Hopper：M必须是64/128，N必须是64/128/256
• Blackwell：M可以是64/128（use_2cta_instrs=False）或128/256（use_2cta_instrs=True）
• K维度通常设置为4个MMA指令的倍数

3. 集群配置

• 根据问题规模和GPU规格选择集群大小
• 考虑内存对齐要求（至少16字节对齐）
• 平衡计算密度和内存带宽

实际应用场景

大语言模型推理

在LLM推理中，QuACK的GEMM内核表现出色：

• 注意力机制：softmax和线性层的高效实现
• 前馈网络：矩阵乘法的优化执行
• 层归一化：rmsnorm的高性能计算

科学计算

• 大型稠密矩阵运算
• 批量矩阵乘法
• 高精度数值计算

未来展望

随着GPU架构的不断演进，QuACK团队持续优化内核实现：

1. 支持新数据类型：如MXFP4、NVFP4等低精度格式
2. 自适应调度：根据硬件特性动态选择最优内核
3. 跨架构兼容：确保代码在Hopper、Blackwell和未来架构上的可移植性

总结

QuACK的GEMM内核展示了在不同GPU架构上的精细优化策略。Hopper架构通过WGMMA和TMA多播实现了高效的计算，而Blackwell架构则通过TMEM和UMMA指令进一步提升了性能。对于开发者而言，理解这些架构差异有助于选择最佳的配置参数，从而在实际应用中获得最优性能。

无论是进行大规模AI模型训练还是高性能科学计算，QuACK都提供了经过充分优化的GEMM实现，让开发者能够充分利用现代GPU的硬件特性，实现极致的计算性能🚀。

【免费下载链接】quackA Quirky Assortment of CuTe Kernels项目地址: https://gitcode.com/gh_mirrors/quack15/quack