示例#

cuBLASDx 库提供了多个块级 BLAS 示例，涵盖了各种精度和类型的基本 GEMM 运算，以及一些突显 cuBLASDx 性能优势的特殊示例。

示例

组

示例

描述

子组

入门示例

introduction_example

cuBLASDx API 入门示例

简单 GEMM 示例

基本示例

simple_gemm_fp32

执行 fp32 GEMM

simple_gemm_int8_int8_int32

使用 Tensor Cores 执行整数 GEMM

simple_gemm_cfp16

执行复数 fp16 GEMM

simple_gemm_fp8

执行 fp8 GEMM

额外示例

simple_gemm_leading_dimensions

使用非默认前导维度执行 GEMM

simple_gemm_fp32_decoupled

使用 16 位输入类型执行 fp32 GEMM，以节省存储和传输

simple_gemm_std_complex_fp32

使用 cuda::std::complex 作为数据类型执行 GEMM

simple_gemm_mixed_precision

为矩阵 A、B 和 C 使用不同数据类型执行 GEMM

simple_gemm_transform

使用变换运算符执行 GEMM

simple_gemm_custom_layout

使用自定义矩阵布局执行 GEMM

simple_gemm_aat

执行 C = A * A^T 的 GEMM

NVRTC 示例

nvrtc_gemm

执行 GEMM，内核使用 NVRTC 编译

GEMM 性能

single_gemm_performance

单 GEMM 的基准测试

fused_gemm_performance

融合到单个内核的 2 个 GEMM 的基准测试

device_gemm_performance

使用 cuBLASDx 对单 tile 进行完整设备 GEMM 的基准测试

高级示例

融合

gemm_fusion

在单个内核中执行 2 个 GEMM

gemm_fft

在单个内核中执行 GEMM 和 FFT

gemm_fft_fp16

在单个内核中执行 GEMM 和 FFT（半精度复数类型）

gemm_fft_performance

融合到单个内核的 GEMM 和 FFT 的基准测试

深度学习

scaled_dot_prod_attn

使用 cuBLASDx 的缩放点积注意力

scaled_dot_prod_attn_batched

使用 cuBLASDx 的多头注意力

其他

batched_gemm_fp64

单个 CUDA 块中的手动批处理

blockdim_gemm_fp16

具有不同块维度的 BLAS 执行

入门示例#

  • introduction_example

入门示例用于文档中解释 cuBLASDx 库及其 API 的基础知识。introduction_example 用于初学者指南中，以使用 cuBLASDx API 运行 GEMM：使用 cuBLASDx 的通用矩阵乘法。

简单 GEMM 示例#

  • simple_gemm_fp32

  • simple_gemm_fp32_decoupled

  • simple_gemm_int8_int8_int32

  • simple_gemm_cfp16

  • simple_gemm_fp8

  • simple_gemm_leading_dimensions

  • simple_gemm_std_complex_fp32

  • simple_gemm_mixed_precision

  • simple_gemm_transform

  • simple_gemm_custom_layout

  • simple_gemm_aat

在上面列出的每个示例中，都在 CUDA 块中执行通用矩阵乘法 (GEMM) 运算。这些示例展示了如何创建完整的 BLAS 描述、分配内存、设置块维度和必要的共享内存量。输入数据（矩阵 A、B 和 C）在主机上生成，复制到设备缓冲区，并加载到寄存器和共享内存中。execute() 计算 GEMM 并将结果存储在矩阵 C 或寄存器片段累加器中，然后将它们复制到全局内存，最后复制回主机。结果会针对 cuBLAS 进行验证。

simple_gemm_leading_dimensions 展示了如何通过 LeadingDimension 运算符为矩阵 A、B 和 C 设置静态（编译时）前导维度。出于性能原因，建议尝试使用 suggested_leading_dimension_of 建议的前导维度。

simple_gemm_fp32_decoupled 示例展示了如何将输入精度与计算精度分离。

simple_gemm_std_complex_fp32 示例证明 cuBLASDx 接受类型不同于 BLAS::a_value_type、BLAS::b_value_type 和 BLAS::c_value_type 的矩阵。在这种情况下，它是来自 CUDA C++ 标准库的复数类型 - cuda::std::complex<float>，但也可以是 CUDA 提供的 float2。

simple_gemm_mixed_precision 示例展示了如何计算混合精度 GEMM，其中矩阵 A、B 和 C 具有不同精度的数据。请注意，标量 alpha 和 beta 预计具有与矩阵 C 元素相同的精度和类型。

simple_gemm_transform 示例展示了如何使用变换运算符 a_load_op、b_load_op、c_load_op 和 c_store_opt 计算 GEMM，这些运算符在加载矩阵 A、B 和 C 以及存储矩阵 C 时按元素方式应用。

simple_gemm_custom_layout 示例展示了如何计算 GEMM，其中共享内存中的矩阵 A、B 和 C 使用自定义 CuTe 布局。

simple_gemm_aat 示例展示了如何执行 C = A * A^T，其中 A 和 A^T 都占用相同的共享内存，从而允许用户增加内核占用率或操作更大的矩阵。

NVRTC 示例#

  • nvrtc_gemm

NVRTC 示例介绍了如何将 cuBLASDx 与 NVRTC 运行时编译结合使用来执行 GEMM。使用 cuBLASDx 运算符创建的 BLAS 描述仅在设备代码中定义。头文件 cublasdx.hpp 也仅包含在传递给 NVRTC 的设备代码中。

注意

自 0.0.1 版本起，cuBLASDx 实验性地支持使用 NVRTC 进行编译。请参阅要求和功能 部分。

GEMM 性能#

  • single_gemm_performance

  • fused_gemm_performance

上面列出的示例说明了 cuBLASDx 的性能。

single_gemm_performance 程序展示了 cuBLASDx 设备函数执行通用矩阵乘法 (GEMM) 函数的性能。用户可以轻松修改此示例来测试他们需要的特定 GEMM 的性能。

fused_gemm_performance 示例展示了两个 GEMM 运算融合到单个内核中的性能。内核执行时间与 cuBLAS 执行相同计算所需的时间进行比较。在这两种情况下，测量的运算都运行多次，并报告平均速度。

高级示例#

  • gemm_fusion

  • gemm_fft

  • gemm_fft_fp16

  • gemm_fft_performance

  • scaled_dot_prod_attn

  • scaled_dot_prod_attn_batched

  • batched_gemm_fp64

  • blockdim_gemm_fp16

高级 cuBLASDx 示例旨在展示如何利用 cuBLASDx 通过将多个计算融合到单个内核中来提高性能，这最终意味着更少的全局内存访问。

示例 gemm_fusion、gemm_fft、gemm_fft_fp16 和 gemm_fft_performance 展示了如何将多个 GEMM 或 GEMM 和 FFT 融合到一个内核中。对于具有大量小输入矩阵的流水线来说，这可能特别有用，因为 Dx 库可以通过在网格中启动多个 CUDA 块来轻松适应批量执行。

scaled_dot_prod_attn 和 scaled_dot_prod_attn_batched 示例探索了深度学习和自然语言处理领域，展示了缩放点积注意力和多头注意力 (MHA) 算法的实现。在 H100 PCIe 80GB 上，将 scaled_dot_prod_attn_batched 示例中半精度 MHA 的性能与 PyTorch 的 scaled_dot_product_attention 函数进行了比较，结果在 图 1 中给出。

Multi-Head Attention (half precision) performance with PyTorch and cuBLASDx on H100 PCIe 80GB with maximum clocks set.

图 1 在 H100 PCIe 80GB 上，在 PyTorch（浅蓝色）和 cuBLASDx（绿色）之间执行了多头注意力算法的比较，并设置了最大时钟频率。该图表显示了对于不同长度的序列和设置为 64 的批大小，cuBLASDx 相对于 PyTorch 的 scaled_dot_product_attention 函数的加速比。NVIDIA PyTorch 容器镜像 23.04 用于 PyTorch 性能评估。#

示例 batched_gemm_fp64 和 blockdim_gemm_fp16 演示了 BlockDim 运算符的用法。在 batched_gemm_fp64 中，将 1D BlockDim 添加到 BLAS 描述并使用 2D 块维度启动内核，允许在单个 CUDA 块中手动批处理 GEMM（与通过在网格中启动多个块进行批处理相反）。blockdim_gemm_fp16 包括多种场景，这些场景展示了当使用与 BlockDim 中指定的布局和线程数不同的块维度启动内核时，如何安全正确地执行 BLAS 运算。