Related Resources
AI 平台/部署

借助启发式算法和 CUTLASS 4.2 提高 NVIDIA GPU 上的 GEMM 内核自动调整效率

2025年 9月 2日
By , , , , and

为特定问题和硬件选择合适的通用矩阵乘法（GEMM）核函数是一项重大挑战。GEMM 核函数的性能由一系列编译时和运行时的元参数共同决定，包括 CTA、线程束和指令级的图块大小、内核调度策略、光栅化方式、集群维度以及 split-k 因子等。

传统的内核搜索方法通常需要生成数千种可能的内核配置，逐一进行编译，并运行详尽的自动调优流程，以找出性能最高的配置。整个流程可能耗时数小时，给开发者带来显著的效率瓶颈，也限制了离线编译库（如 CUTLASS）。对于即时编译库（如 Torch Inductor 或 OpenAI Triton）而言，这一问题尤为突出，因为模型的快速编译至关重要。这种复杂的调优过程带来的使用障碍，可能导致用户转而采用性能较差的内核，以规避漫长的优化周期。

本文将介绍 NVIDIA Matmul Heuristics（nvMatmulHeuristics），这是一个用于 GPU 内核元参数优化的模块，可为 GEMM 运算提供高效的启发式策略。该模块通过分析具体操作参数及目标硬件特性，筛选出一组性能表现优异的内核配置。本文示例基于 CUTLASS 4.2 实现。

基于启发式算法的 CUTLASS 配置和自动调整

将 nvMatmulHeuristics 集成到 CUTLASS 生态系统中，通过优化内核生成与调优流程，显著提升了用户体验。

该集成引入了一种更高效的新工作流程，取代了传统的暴力方法。对于给定的 GEMM 问题，用户现在可以利用 nvMatmulHeuristics 来预测少量具有高潜力的针对性内核配置。

工作流程包含以下步骤：

  1. 启发式预测: nvMatmulHeuristics 采用 GEMM 问题定义 (形状、数据类型等) 和目标后端 (本例中为 CUTLASS) ，检测硬件属性，并根据其内部模型生成前景良好的内核配置简短列表。这些配置包括最佳 CTA 形状、split-k 因子和其他元参数。
  2. 内核生成: 此预测配置列表将传递给 CUTLASS 库内核生成器，该生成器仅生成此类相关的小型内核集，从而大大缩短了编译数千个不必要的变体的耗时过程。
  3. 自动调整: CUTLASS 分析器采用相同的列表，仅自动调整运行时参数，仍然遵循为少量已编译内核选择的 nvMatmulHeuristics。然后，它会快速找到表现最好的候选项。

这种方法能显著缩短寻找高性能内核所需的端到端时间。

要在 CUTLASS 中使用此功能，首先需要准备一个 JSON 格式的 GEMM 问题列表。以下示例对应一个 FP16 GEMM 运算，其中 tnn 分别表示矩阵 A 的行数（若转置则为列数）、矩阵 B 的列数（若不转置），以及矩阵 C/D 的列数。

[
{
     "m" : 4096,
     "n" : 4096,
     "k" : 4096,
     "batch_count" : 1,
     "layout" : "tnn",
     "dtype_a" : "f16",
     "dtype_b" : "f16",
     "dtype_c" : "f16",
     "dtype_acc" : "f32",
     "dtype_d" : "f16"
}
]

接下来，在配备目标 GPU 的计算机上，正常构建 CUTLASS，并提供 -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE=<path_to_your_problem_list.json> 和 -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=N，其中 N 表示 nvMatmulHeuristics 为输入列表中每个 GEMM 所生成的配置数量。例如：

cmake ${SRC_DIR} \
    -DCUTLASS_NVCC_ARCHS=90a \
    -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE=<path_to_your_problem_list.json> \
    -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=8 \
    -DCUTLASS_LIBRARY_HEURISTICS_TESTLIST_FILE=my_testlist_file.csv
    -DCMAKE_BUILD_TYPE=Release
...
...

make cutlass_profiler -j

“cmake 步骤将生成一个 CSV 测试列表，其中包含为当前配置执行自动调优所需运行的全部测试用例。您可以将该列表与自定义的基准测试代码结合使用，以实现自主调优；也可以直接使用 cutlass_profiler，它能够读取此 CSV 文件并开箱即用地完成配置的自动调优。”

cutlass_profiler --operation=Gemm --testlist-file=my_testlist_file.csv 
--profiling-iterations=0 --profiling-duration=50 --verification-enabled=false 
--output=<path_to_outfile>

为获得一致的分析结果，建议使用锁定的时钟运行。

如需深入了解此功能，请参考 CUTLASS 文档。

nvMatmulHeuristics 现已开放抢先体验

nvMatmulHeuristics 作为 cuBLAS 启发式算法的核心组件，现已开放抢先体验，可供广泛使用，并已集成至 CUTLASS 库中。其主要特性包括：

  • 快速分析启发式算法和自动 GPU 配置检测 (例如，绿色上下文、MIG、图形中的 CUDA、时钟)
  • 支持 CUTLASS 和其他 GEMM 后端。后端支持空间也可以修改，例如，通过限制支持的 CTA 图块形状、分割 k 因子等
  • NVIDIA Ampere、Ada、Hopper 和 (初步) Blackwell GPU 架构支持所有基于 Tensor Core 的 GEMM 精度 ( FP4、FP8、FP16/ FB16、TF32、INT8 等)
  • Python 和 C++ API

首先，安装 nvMatmulHeuristics 工具：

pip install nvidia-matmul-heuristics

然后，您可以查询预测出的前 N 个 GEMM 配置。以 FP16 输入/输出和 FP32 计算（HSH）为例：

from nvMatmulHeuristics import *

# Load interface
nvmmh = NvMatmulHeuristicsInterface(NvMatmulHeuristicsTarget.CUTLASS3, 
                                    precision='HSH', 
                                    flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING)

# Create Hardware descriptor
# hw can be "None" to use the system's GPU instead
hw = nvmmh.createHardwareDescriptor()
nvmmh.setHardwarePredefinedGpu(hw, NvMatmulHeuristicsNvidiaGpu.H200_SXM)

# Select layout
layout = NvMatmulHeuristicsMatmulLayout.NN_ROW_MAJOR

# Load internal discovery set for improved accuracy
assert nvmmh.loadInternalDiscoverySet(layout, hw)

# Get best configurations
configs = nvmmh.get_with_mnk(4000, 16, 32768, layout, 8, hw)

# Print results
print(f"Found {len(configs)} configurations:\n")
for i, config in enumerate(sorted(configs, key=lambda d: d['runtime']), 1):
    print(f"Configuration {i}:")
    print(f"  Kernel: {config['kernel']}")
    print(f"  Estimated runtime: {config['runtime'] * 1000:.6f} ms")

输出：

Found 8 configurations:

Configuration 1:
  Kernel: layout(NN_ROW) stages(6) cta(128 16 128) warp(64 8 128) instr(64 8 16) splitK(4) swizz(1) ctaOrder(0) cluster(2 1)
  Estimated runtime: 0.083215 ms
...
Configuration 8:
  Kernel: layout(NN_ROW) stages(8) cta(64 16 64) warp(64 8 64) instr(64 8 16) splitK(1) swizz(1) ctaOrder(0) cluster(4 1)
  Estimated runtime: 0.102996 ms

在给定八个配置的情况下，现在可以编译一组包含八个 CUTLASS 内核的集合，并从中选出最适合当前问题的内核。

nvMatmulHeuristics 的表现如何？

通过将构建和分析工作聚焦于 nvMatmulHeuristics 推荐的少量候选方案，您仅需花费全面搜索所需的一小部分时间，即可获得接近最优的性能表现。

图 1 展示了在 NVIDIA H100 SXM GPU 上，Llama 3 405B 训练工作负载中一组 GEMM 的 geomean 性能得分。性能分数以详尽搜索所得性能最优的内核为基准（基准值为 1.0），表示相对加速比，并绘制了对应的总构建与分析时间。

如数据所示，详尽的搜索需要超过 700 分钟才能找到性能最佳的内核。相比之下，采用 nvMatmulHeuristics 仅筛选 16 个候选内核，便可在约 150 分钟内达到峰值性能的 96%。

通过增加候选项的数量，用户能够接近详尽搜索基准的性能，同时显著缩短构建与调优所需的时间。这使得在 PyTorch 等对即时编译时间敏感的环境中，高效寻找到高性能内核成为可能。

Scatter plot showing the build time and achieved performance of nvMatmulHeuristics compared with an exhaustive baseline.
图 1. nvMatmulHeuristics 在 NVIDIA H100 SXM 上的 Llama 3 405B 性能

图2展示了NVIDIA B200 GPU上DeepSeek-R1 671B训练工作负载的类似结果。通过详尽搜索，性能可达99%，构建与自动调优速度提升五倍以上。由于仅需构建少量核函数，因此能够静态地在核函数中集成更多行为，从而获得类似“JIT”的优势；而预编译的核函数集合通常会在运行时动态迁移更多行为，以减少编译后核函数的数量。

在这种情况下，基准测试采用动态集群大小，这在 NVIDIA Blackwell 预编译的 CUTLASS 内核中较为常见；而 nvMatmulHeuristics 推荐的内核则是基于编译时已知的静态集群大小构建的。因此，这些 GEMM 的性能达到了基准的 104%。

Scatter plot showing build and achieved geomean speedup performance for nvMatmulHeuristics against a CUTLASS exhaustive baseline search. The experiment is conducted on a B200 with a fixed clock at 1530 MHz. The problem set used represents GEMMs from DeepSeek-R1 671B.
图 2. nvMatmulHeuristics 在 NVIDIA B200 上 DeepSeek-R1 675B 上的性能

开始改进 GEMM 自动调整工作流程

这些结果表明，优秀的启发式算法能够在无需手动干预或 exhaustive 调整的情况下，帮助工作负载达到峰值性能。nvMatmulHeuristics 有望突破性能瓶颈，显著提升深度学习框架、编译器和内核库中的性能表现与开发效率。

下载 nvMatmulHeuristics 以开始使用。如有疑问，欢迎前往 NVIDIA 开发者论坛 参与讨论，或发送邮件至 Math-Libs-Feedback@nvidia.com 联系我们。

致谢

我们谨向所有 CUTLASS 开源项目贡献者。正是得益于他们的基础性贡献，CUTLASS 4 的实现才成为可能。

 

标签

AI 平台/部署 | 开发与优化 | 模型/库/框架 | AI Platforms / Deployment | Blackwell | CUDA | CUDA-X | CUTLASS | Development & Optimization | General | Hopper | Intermediate Technical | Models / Libraries / Frameworks | 功能 | 最佳实践

关于作者

Harrison Barclay 是 NVIDIA 的高级性能架构师，致力于确保 DL 内核性能达到光速。他毕业于德克萨斯大学奥斯汀分校，拥有电气和计算机工程硕士学位。
Ian 是 NVIDIA TensorRT 团队的高级系统软件工程师，他专注于 MLPerf 推理。在加入 TensorRT 团队之前，他曾为 NVIDIA 自动驾驶软件开发实时调度系统。伊恩毕业于多伦多大学工程科学专业，主修电气和计算机工程。
Micha Kuku a 拥有华沙大学计算机科学硕士学位和数学学士学位。他目前担任 NVIDIA cuBLAS 的高级数学库工程师，专门开发启发式策略，以优化 NVIDIA 最新架构的性能。他的主要工作重点是设计和实施新的启发式方法，以确保 NVIDIA 在先进的 NVIDIA 硬件平台上提供最佳性能。
Michel Migdal 拥有巴黎大学 Saclay 分校和 ENSIIE 工程与计算机科学双硕士学位。Michel 目前在 NVIDIA 的 cuBLAS 团队担任高级 GPU 软件工程师，在该团队中，他启动了一项针对内核生成器的新的分析性 GEMM 启发式算法计划。同时，他还致力于为当前和即将推出的架构优化和开发新的内核。此前，Michel 曾在 nvJpeg 团队工作，负责开发新的编解码器。
Nicholai Tukanov 是 NVIDIA 的软件工程师，从事 cuBLAS 和 NVPL 的研究。Nicholai 毕业于德克萨斯大学奥斯汀分校，拥有计算机科学学士学位和卡内基梅隆大学电气与计算机工程硕士学位。
Leopold Cambier 于 2021 在斯坦福大学获得了计算和数学工程博士学位。在博士期间，他专注于从理论和并行计算的角度研究大型稀疏线性系统的快速求解器。 2016 年和 2017 年，他还在NVIDIA cuDNN 团队实习。自从他在 2021 年 1 月加入全职以来，利奥波德一直在cuFFT。