CUTLASS Blackwell Blockwise/Groupwise GEMM 完全指南FP8 分块缩放与分组 GEMM 的配置、剖析与性能调优【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass导读本文围绕 CUTLASS 在 NVIDIA BlackwellSM100架构上提供的 Blockwise分块缩放与 Groupwise分组缩放GEMM 以及 Grouped GEMM分组批 GEMM展开讲解如何通过按累加器类型的软件缩放在可配置粒度上控制数值精度特别适用于量化神经网络中对张量不同区域采用不同缩放需求per-block/per-group 缩放的场景。读完本文你将掌握缩放因子张量 SFA/SFB 的数学语义与 CuTe 布局表示、Sm100BlockwiseScaleConfig的配置方法、跨框架如 PyTorch张量格式转换要点、利用 CUTLASS Profiler 自动调优选择最优 kernel 的完整命令行流程以及 kernel 命名规范和 MMA 维度相关的性能调优技巧。本指南的主体基于 examples/81_blackwell_gemm_blockwise/README.md并辅以 include/cutlass/detail/blockwise_scale_layout.hpp、include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp 等源码与 examples/81_blackwell_gemm_blockwise 目录下的四个示例程序进行纵深佐证。背景为什么需要 Blockwise/Groupwise 缩放标准的 GEMM 计算为 $D \alpha A B \beta C$。在量化推理场景中若将 A、B 量化到 FP8如e4m3一个全局缩放因子往往无法同时兼顾张量不同区域动态范围的差异。Blockwise/Groupwise GEMM 的核心思想是引入两个缩放因子张量SFA 与 SFB把计算改写为$$D \alpha \ (\text{SFA} * A) \ (\text{SFB} * B) \beta C$$其中*表示逐元素相乘。缩放因子以可配置的粒度作用于 A 和 B使得张量的不同区域可以拥有各自独立的缩放从而实现更细粒度的数值精度控制。CUTLASS 通过**软件方式按累加器类型**完成这一缩放缩放因子以f32累加器类型表示而 A、B 以e4m3等低精度类型参与矩阵乘缩放发生在计算过程中。从源码结构看这类 kernel 的实现分布在 include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp单 CTA 的 Warp Specialized 主循环与 include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp指针数组/Grouped 形式中对应 include/cutlass/gemm/dispatch_policy.hpp 中定义的调度策略KernelScheduleSm100BlockwiseBlockwise 调度的基类KernelTmaWarpSpecializedBlockwise1SmSm100/KernelTmaWarpSpecializedBlockwise2SmSm1001SM 与 2SM 的 TMA Warp Specialized 变体KernelScheduleSm100PtrArrayBlockwise与KernelPtrArrayTmaWarpSpecializedBlockwise1SmSm100服务于 Grouped GEMM 的指针数组PtrArray形式。缩放因子张量 SFA 与 SFB缩放因子张量由两个粒度参数完全确定SFA作用于 A。在由scale granularity m与scale granularity k定义的分块内广播同一个缩放值。这两个粒度参数也常被称为scale vector m与scale vector k。SFB作用于 B。在由scale granularity n与scale granularity k定义的分块内广播同一个缩放值。同理也称为scale vector n与scale vector k。也就是说SFA 的尺寸为 $(M / \text{scale granularity M}) \times (K / \text{scale granularity K})$SFB 的尺寸为 $(N / \text{scale granularity N}) \times (K / \text{scale granularity K})$均不计 batch 维。CuTe 布局表示这两个张量在 CuTe 中可以分别表示为如下 LayoutSFA Layout $((\text{scale granularity M},\ M / \text{scale granularity M}),\ (\text{scale granularity K},\ K / \text{scale granularity K})) : ((0,\ int),\ (0,\ int))$SFB Layout $((\text{scale granularity N},\ N / \text{scale granularity N}),\ (\text{scale granularity K},\ K / \text{scale granularity K})) : ((0,\ int),\ (0,\ int))$其中内层形如(scale granularity, count)的 mode 对配合 stride 中的0 元素步长(0, int)确保同一分块内的所有坐标都映射到缩放因子张量中的同一个元素——这正是块内广播的布局级实现。在 include/cutlass/detail/blockwise_scale_layout.hpp 中可以看到与之对应的类型定义using ShapeSFA ShapeShapeIntSFVecSizeM, int32_t, ShapeIntSFVecSizeK, int32_t, int32_t; using StrideSFA conditional_tmajorSFA UMMA::Major::MN, StrideStride_0,_1, Stride_0,int32_t, int32_t, StrideStride_0,int32_t, Stride_0,_1, int32_t;配置Sm100BlockwiseScaleConfig为方便使用Blockwise 与 Groupwise 的实现提供了配置类cutlass::detail::Sm100BlockwiseScaleConfigScaleGranularityM, ScaleGranularityN, ScaleGranularityK它负责推导 SFA/SFB 的 Layout并管理紧凑张量。默认情况下该配置使所有张量以 M/N mode 为主序major但可以通过模板参数覆盖。例如cutlass::detail::Sm100BlockwiseScaleConfigScaleGranularityM, ScaleGranularityN, ScaleGranularityK, UMMA::Major::K, UMMA::Major::MN表示 SFA 在 K 维度上为主序而 SFB 在 N 维度上为主序。模板定义的完整签名见 include/cutlass/detail/blockwise_scale_layout.hpp为templateint SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA UMMA::Major::MN, UMMA::Major majorSFB UMMA::Major::MN struct Sm1xxBlockwiseScaleConfig { ... }; templateint SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA UMMA::Major::MN, UMMA::Major majorSFB UMMA::Major::MN using Sm100BlockwiseScaleConfig Sm1xxBlockwiseScaleConfigSFVecSizeM, SFVecSizeN, SFVecSizeK, majorSFA, majorSFB;同一文件还提供了Sm120BlockwiseScaleConfig与Sm90BlockwiseScaleConfigSM90 目前仅支持 MN 主序的 SFA/SFB以及平凡trivial配置辅助函数sm100_trivial_blockwise_scale_config(MmaTileShape_MNK{})——它将缩放粒度直接取为 MMA tile 的各维大小。该配置类提供了两组关键接口deduce_layoutSFA()/deduce_layoutSFB()返回仅含静态零步长信息的原子 Layout步长只遍历标量、不含零smem_atom_layoutSFA()/smem_atom_layoutSFB()根据 CTA tile 形状推导共享内存中的原子布局tile_atom_to_shape_SFA(problem_shape)/tile_atom_to_shape_SFB(problem_shape)将原子布局平铺tile到实际的动态问题尺寸M、N、K、L上此时步长才包含零。示例程序 81_blackwell_gemm_blockwise.cu 中即采用了 trivial 配置using ScaleConfig decltype(cutlass::detail::sm100_trivial_blockwise_scale_config(MmaTileShape_MNK{})); using LayoutSFA decltype(ScaleConfig::deduce_layoutSFA()); using LayoutSFB decltype(ScaleConfig::deduce_layoutSFB());而 81_blackwell_gemm_groupwise.cu 则展示了显式指定粒度的 Groupwise 用法constexpr int ScaleGranularityM 1; constexpr int ScaleGranularityN 128; constexpr int ScaleGranularityK 128; using ScaleConfig cutlass::detail::Sm100BlockwiseScaleConfigScaleGranularityM, ScaleGranularityN, ScaleGranularityK;粒度约束编译期检查从 include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp 的实现可以看出缩放粒度并非任意取值主循环在编译期通过static_assert施加如下约束ScaleGranularityM 必须整除 TileShape 的 M 维且不大于 TileShape 的 M 维ScaleGranularityN 必须整除 TileShape 的 N 维且不大于 TileShape 的 N 维ScaleGranularityK 必须整除 TileShape 的 K 维且不大于 TileShape 的 K 维ScaleGranularityK 必须能被 MMA 指令的 K 维MMA_K整除——缩放粒度不能小于单个 MMA 原子指令在 K 方向覆盖的范围这保证了缩放可以按 MMA 指令为单位对齐执行。与其他框架如 PyTorch的集成从 Torch 等框架迁移时SFA 的形状通常为 $(M / \text{ScaleGranularityM},\ K / \text{ScaleGranularityK})$而 SFB 的形状为 $(K / \text{ScaleGranularityK},\ N / \text{ScaleGranularityN})$。接入 CUTLASS 时需要注意对 SFB 与 B 执行转置使其符合 CUTLASS 的规范 CuTe 布局形式保证K 始终是第二个 mode利用张量自身的 strides 判断每个张量是 MN 主序还是 K 主序据此直接构造布局或使用上述便捷包装类Sm100BlockwiseScaleConfig等完成布局推导。换句话说只要保证 SFA 的 mode 顺序为 (M, K)、SFB 的 mode 顺序为 (K, N)并把每个 mode 内的粒度信息正确映射到(granularity, count)的 CuTe mode 对中即可无缝对接。Kernel 选择与 Profiling要为自己的 workload 确定性能最优的 Blockwise/Groupwise GEMM 或 Grouped GEMM kernel官方推荐使用 CUTLASS Profiler。编译期裁剪 kernel 数量所有使用f32缩放、e4m3或运行时f8类型的 Blockwise/Groupwise GEMM 及 Group GEMM都可以通过在 CMake 配置时传入 kernel 子集来启用/裁剪-DCUTLASS_LIBRARY_KERNELScutlass3x*f32xe4m3_*f32xe4m3*,cutlass3x*f32xf8_*f32xf8*进一步地可以通过指定 SFA 与 SFB 的缩放粒度来减少生成的 kernel 数量例如-DCUTLASS_LIBRARY_KERNELScutlass3x*1x128f32xe4m3_*128x128f32xe4m3*使用 Profiler 自动调优使用 profiler 最简单的方式是传入m、n、k以及scale_vec_size_m、scale_vec_size_n、scale_vec_size_k。这三个scale_vec_size_*参数对应 profiler 源码 tools/profiler/src/blockwise_gemm_operation_profiler.cu 中注册的命令行参数Scale vector size in GEMM M/N/K dimension。加上enable-best-kernel-for-fixed-shape后profiler 会对每个 kernel 执行自动调优寻找最佳的 rasterization 顺序、swizzle 与 cluster 尺寸。通过operation标志传入blockwiseGemm或GroupedGemm可以决定剖析哪一组操作。例如下面的命令剖析所有支持scale granularity m 1、scale granularity n 128、scale granularity k 128的已编译 kernel在 8192x8192x8192 问题规模上的性能cutlass_profiler --operationblockwiseGemm \ --enable-best-kernel-for-fixed-shape \ --m8192 --n8192 --k8192 \ --scale_vec_size_m1 --scale_vec_size_n128 --scale_vec_size_k128 \ --verification-enabledfalseKernel 命名规范Blockwise 与 Groupwise kernel 的命名引入了新模式对于每对张量缩放组合采用scale_granularity_m 或 scale_granularity_nxscale_granularity_k累加器类型x被缩放张量类型。以cutlass3x_sm100_tensorop_gemm_64x128f32xe4m3_1x128f32xe4m3_f32_f16_f16_64x128x128_1x1x1_0_nnn_align16_1sm为例各段含义依次为命名片段含义cutlass3x_sm100_tensorop_gemmCUTLASS 3、面向 SM100、使用 Tensor Core 的 GEMM64x128f32xe4m3SFA 为f32scale granularity m 64、scale granularity k 128A 矩阵为e4m31x128f32xe4m3SFB 为f32scale granularity n 1、scale granularity k 128B 矩阵为e4m3f32累加器epilogue在f32中完成f16/f16C 矩阵为f16、D 矩阵为f1664x128x128MMA tile 形状MxNxK1x1x1cluster 形状0_nnnA、B、C、D 均按序为列主序n column-majoralign16A、B、C、D 主 mode 的对齐为 16 个元素1smMMA 变体为 1SM 指令另外值得注意如果不需要beta * C缩放C 可以为void即不提供 C 张量。性能技巧与调优建议MMA 维度选择在 Blackwell 与 Hopper 的 Tensor Core 上最小的MMA_M维是 64但某些指令的MMA_N维可以小到 8。因此对于 M 较小的 problem size应当考虑改为计算$$D^T \alpha B^T A^T \beta C^T$$交换 A、B 并转置后原本较小的 M 变成了 N 维配合小的MMA_N可以更高效地分块tiling避免无谓的多余计算。布局交换Layout Swapping使用 profiler 优化时可以交换m与n输入并相应调整布局以反映这种交换与转置。例如若原始布局为 row-major A、column-major B、row-major D则交换张量后可以运行这样的 kernel左手矩阵原 B 转置后为 row-major右手矩阵原 A 转置后为 column-major输出原 D 转置后为 column-major。使用 Blockwise/Groupwise GEMM 时做上述优化必须同步交换缩放向量的大小例如原本 scale granularity M 1、scale granularity N 128交换后应运行 scale granularity M 128、scale granularity N 1 的 kernel。该技巧同样体现在源码注释中81_blackwell_gemm_groupwise.cu 指出当一个 tile 内有多个缩放因子如 M 方向每个 tile 有 128 个 scale时实现会尽可能限制在 16B 对齐以内即 M 方向至少有 16B 的 scale此时可执行的最小 M 为 16。对于更小的 M可以通过交换 A、B 并转置 A、B、C 与 scale 来规避因为 $B^T A^T C^T$。参考示例与运行方式本目录examples/81_blackwell_gemm_blockwise共包含四个示例示例说明81_blackwell_gemm_blockwise.cu单问题 Blockwise 缩放 GEMMtrivial 缩放配置MMA tile 128x128x128、cluster 1x1x181_blackwell_gemm_groupwise.cu单问题 Groupwise 缩放 GEMM显式 ScaleGranularityM1/N128/K128MMA tile 256x128x128、cluster 2x1x181_blackwell_grouped_gemm_blockwise.cuGrouped GEMM Blockwise 缩放使用GroupProblemShape与指针数组 TMA 调度81_blackwell_grouped_gemm_groupwise.cuGrouped GEMM Groupwise 缩放这些示例统一以ElementA ElementB cutlass::float_e4m3_t、ElementAccumulator float为数据布局见 81_blackwell_gemm_blockwise.cu且要求编译架构包含100a、CUDA 12、GPU compute capability 为 100a示例代码在cudaGetDeviceProperties中检查props.major 10 props.minor 0。构建注册见 CMakeLists.txt其中CUTLASS_NVCC_ARCHS需匹配100a并通过cutlass_example_add_executable注册四个可执行目标。命令行可选项以 blockwise 示例为例--m、--n、--k、--lbatch 维、--alpha、--beta、--iterations、--skip-verification。例如./81_blackwell_gemm_blockwise --m1024 --n512 --k1024 --alpha2 --beta0.707示例在verify()阶段使用 include/cutlass/util/reference/host/gett.hpp 提供的Gemm3x参考实现含GettBlockScalingMainloopParams对 kernel 输出做逐元素校验并统计平均运行时间与 GFLOPS2 * m * n * k次浮点运算。小结Blockwise/Groupwise GEMM 是 CUTLASS 面向 Blackwell SM100 提供的高精度量化 GEMM 方案通过 SFA/SFB 两个缩放因子张量在 M/N/K 三个维度上以可配置粒度做软件缩放结合Sm100BlockwiseScaleConfig的布局推导、Profiler 的自动调优与命名规范的快速定位开发者可以为自己的量化模型快速找到最优 kernel 配置。性能调优时牢记两条主线小 M 时交换 A/B 并转置计算 $D^T$以及做布局交换时同步交换缩放粒度即可在 Blackwell Tensor Core 上充分发挥 tcgen05 MMA 的吞吐能力。如需深入可继续阅读media/docs/cpp/profiler.mdProfiler 完整使用手册、include/cutlass/detail/blockwise_scale_layout.hpp缩放布局推导源码以及 include/cutlass/gemm/dispatch_policy.hppBlockwise 调度策略定义。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
