NVIDIA CUDA Tile高性能矩阵乘法实现

用户576110558132
0 阅读7分钟

如何使用NVIDIA CUDA Tile编写高性能矩阵乘法

2026年1月14日 | 作者：Jinman Xie 和 Qiqi Xiao

本博文是一个系列文章的一部分，旨在帮助开发者学习NVIDIA CUDA Tile编程，以构建高性能GPU内核，并以矩阵乘法为核心示例。

在本文中，您将学习：

  • 如何使用NVIDIA cuTile实现高性能矩阵乘法：理解Tile加载、计算和存储的流程
  • 块级并行编程思维：从线程级思维转向块级思维
  • Tile编程最佳实践：从代码中学习性能优化策略

环境要求

开始之前，请确保您的环境满足以下要求：

  • CUDA 13.1 或更高版本
  • GPU架构：某机构 Blackwell（例如某机构 RTX 50系列）
  • Python：3.10 或更高版本

安装cuTile Python包：

pip install cuda-tile

注：cuTile是面向某机构的下一代GPU编程框架。目前仅支持Blackwell架构（计算能力10.x和12.x）的优化，后续CUDA Toolkit版本将提供对更多架构的支持。

什么是矩阵乘法？

矩阵乘法是现代技术计算中的基本运算。它是求解方程组的基础操作，支撑着图形学、仿真、优化以及大部分机器学习，并且非常适合映射到GPU等高性能硬件。

给定输入矩阵A（M×K）和B（K×N），结果矩阵C（M×N）中元素的计算公式如下：

Cᵢⱼ = Σₖ₌₁ᴷ Aᵢₖ × Bₖⱼ

从公式可以看出，C的一个元素是通过计算A的一行与B的一列的点积得到的。

Tile编程通过将输出矩阵划分为多个Tile来简化实现，同时获得卓越的性能。每个Block负责一个输出Tile的计算，cuTile自动处理内存访问和线程同步。具体来说：

  • 每个Block处理输出矩阵C的一个（tm × tn）Tile
  • 沿K维度循环，逐个加载A和B的对应Tile
  • 使用ct.mma()执行矩阵乘累加（自动调用Tensor Core）
  • 最后将累加结果存回全局内存

图1展示了计算过程，类似于逐元素算法，但这里用Tile取代了单个元素的位置。

（图1：矩阵乘法A×B=C分解为Tile的示意图）

GPU内核实现

介绍了核心思想后，来看完整的实现代码。代码分为两部分：GPU上运行的内核和CPU上的启动代码。

import cuda.tile as ct
from math import ceil
import torch

# 编译时常量类型别名
ConstInt = ct.Constant[int]

# 步骤1：定义内核
@ct.kernel
def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt):
    # 1.1 获取Block ID并映射到输出Tile位置
    bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M)
    
    # 1.2 计算K维度上的Tile数量
    num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))
    
    # 1.3 初始化累加器
    accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
    
    # 1.4 沿K维度循环
    for k in range(num_tiles_k):
        # 从A和B加载Tile
        a = ct.load(A, index=(bidx, k), shape=(tm, tk))
        b = ct.load(B, index=(k, bidy), shape=(tk, tn))
        
        # 矩阵乘累加
        accumulator = ct.mma(a, b, accumulator)
    
    # 1.5 存储结果
    ct.store(C, index=(bidx, bidy), tile=accumulator)

# 步骤2：启动内核
def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    # 选择Tile尺寸
    tm, tn, tk = 128, 256, 64  # 用于float16
    
    # 计算网格维度
    grid_x = ceil(m / tm)
    grid_y = ceil(n / tn)
    grid = (grid_x * grid_y, 1, 1)
    
    # 创建输出张量并启动
    C = torch.empty((m, n), device=A.device, dtype=A.dtype)
    ct.launch(stream, grid, matmul_kernel, (A, B, C, tm, tn, tk))
    return C

1. 定义GPU内核

在cuTile中，@ct.kernel装饰器用于将标准Python函数标记为GPU内核：

@ct.kernel
def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt):
    # 内核代码

该装饰器表示：

  • 此函数将在GPU上执行
  • 每个Block将运行此函数的一个独立实例
  • 不能直接调用，必须使用ct.launch()启动

2. 编译时优化：常量类型注解

参数tm、tn和tk使用特殊的类型注解ct.Constant[int]，表示它们是编译时常量。cuTile会为不同的Tile尺寸值生成专门的机器码，使编译器能够：

  • 执行循环展开
  • 优化内存访问模式
  • 生成最优的Tensor Core指令

3. 确定工作范围：Block ID映射

每个Block计算输出矩阵的一个特定Tile。通过swizzle_2d()函数获取当前处理的Block索引：

bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M)

该代码的功能是确定当前Block应处理输出矩阵的哪个Tile。

4. 准备累加器：初始化输出Tile

在沿K维度循环之前，需要创建累加器来存储中间结果：

num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
  • num_tiles_k：计算K维度需要处理的Tile数量
  • accumulator：形状为(tm, tn)的零矩阵，用于累加结果
  • 使用float32确保数值精度，避免累加误差

5. 核心计算循环：遍历K维度

这是矩阵乘法的核心。循环遍历K维度的每个Tile并累加结果：

for k in range(num_tiles_k):
    # 加载Tile
    a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad)
    b = ct.load(B, index=(k, bidy), shape=(tk, tn), padding_mode=zero_pad)
    
    # 累加
    accumulator = ct.mma(a, b, accumulator)
  • ct.load()：从矩阵加载Tile，越界时用零填充
  • ct.mma()：执行矩阵乘累加，当a和b的形状满足Tensor Core要求时，cuTile自动调用GPU的Tensor Core加速

循环结束后，accumulator存储了输出Tile的完整结果。

6. 写回结果：存储到全局内存

accumulator = ct.astype(accumulator, C.dtype)
ct.store(C, index=(bidx, bidy), tile=accumulator)

首先将float32累加器转换为输出矩阵的数据类型，然后使用ct.store()将Tile写回全局内存的对应位置。

启动内核：主机端代码

def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    # 根据dtype确定Tile尺寸
    if A.dtype.itemsize == 2:  # float16/bfloat16
        tm, tn, tk = 128, 256, 64
    else:  # float32
        tm, tn, tk = 32, 32, 32
    
    m, k = A.shape
    _, n = B.shape
    
    # 计算网格维度
    grid_x = ceil(m / tm)
    grid_y = ceil(n / tn)
    grid_size = grid_x * grid_y
    grid = (grid_size, 1, 1)
    
    # 创建输出张量
    C = torch.empty((m, n), device=A.device, dtype=A.dtype)
    
    # 启动内核
    ct.launch(torch.cuda.current_stream(), grid, matmul_kernel,
              (A, B, C, tm, tn, tk))
    
    return C

启动内核需要三个关键步骤：

  1. 计算网格大小：根据输入矩阵维度和Tile尺寸计算需要的Block数量
  2. 设置Tile尺寸：根据数据类型选择适当的Tile维度（编译时常量）
  3. 调用ct.launch()：启动内核执行

性能优化：Swizzle

swizzle_2d_from_bid函数的代码如下：

def swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid):
    num_bid_m = ct.cdiv(M, tm)
    num_bid_n = ct.cdiv(N, tn)
    num_bid_in_group = GROUP_SIZE_M * num_bid_n
    group_id = bid // num_bid_in_group
    first_bid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_bid_m - first_bid_m, GROUP_SIZE_M)
    bid_m = first_bid_m + (bid % group_size_m)
    bid_n = (bid % num_bid_in_group) // group_size_m
    return bid_m, bid_n

在矩阵乘法中，使用GROUP_SIZE_M=8。

Swizzle如何提升性能？ 它通过分组和交错重新映射Block ID到Tile索引，从而更高效地使用缓存。

  • 线性行访问：需要读取20个数据块，右矩阵数据频繁加载且快速被替换，缓存命中率低
  • Swizzle/Tiled块访问：只需读取16个数据块（减少20%），更好的数据局部性带来更高的缓存命中率

性能基准测试

在某机构 GeForce RTX 5080（计算能力12.0）上进行了测试：

  • 数据类型：float16
  • 矩阵形状：标准方阵（N×N）
  • 测试尺寸：N = 1024, 2048, 4096, 8192, 16384

结果：在大矩阵规模下，cuTile实现能充分利用GPU计算能力。通过适当的Tile尺寸配置和Swizzle优化，cuTile实现达到了与SOTA实现（PyTorch调用cuBLAS）相比90%以上的性能。

总结

这个经典矩阵乘法示例展示了使用cuTile实现GPU内核的完整过程。虽然矩阵乘法本身很简单，但它包含了Tile编程的核心思想。掌握这些概念后，您将能够使用cuTile实现各种高性能GPU内核。FINISHED

avatar
文章
阅读
粉丝