本文全面介绍 NVIDIA CuTe（CUDA Tensor）领域特定语言的基础知识。我们将从第一个 GPU 核函数开始，逐步学习动态打印、数据类型、张量构建与切片，最终掌握使用向量化核函数、线程-值布局（TV布局）的高级技术。CuTe DSL 提供了一套强大的抽象，用于在 NVIDIA GPU 上表达和优化张量计算，其核心思想是将数据布局与计算解耦。通过理解形状、步幅、平铺和组合，开发者可以构建高效、可移植的 GPU 内核。
目录
- 摘要与概述
- 术语表：核心概念详解
- 第一个核函数与启动（Hello World）
- 打印与数值类型：静态 vs 动态
- 张量与布局：构建、索引、切片
- 寄存器张量与归约操作
- 使用 zipped_divide 的向量化加法
- TV 布局与每线程片段
- 布局代数：合并、组合、划分、乘积
- 完整代码示例：矩阵乘法实现
1. 摘要与概述
CuTe DSL 是 NVIDIA 开发的专门用于 GPU 张量计算的领域特定语言，它通过一系列优雅的抽象将数据布局与计算逻辑解耦，让开发者能够高效地构建复杂、高性能的 GPU 内核。
本指南的系统化学习路径如下：
- 基础入门：从最简单的核函数启动开始，掌握 CuTe 的基本语法和执行模型。
- 调试与类型系统：学习如何查看运行时值，理解静态与动态类型的区别及其应用。
- 核心数据结构：深入掌握张量和布局的概念，这是理解 CuTe 内存访问模式的基础。
- 寄存器级优化：学习如何使用 TensorSSA 进行向量化操作和高效归约。
- 并行化策略：掌握使用
zipped_divide 进行数据分块和线程映射的技术。
- 高级布局模式：了解 TV 布局如何将线程组织与数据平铺解耦。
- 布局代数：学习组合、合并、划分和乘积等布局操作，构建复杂访问模式。
通过本指南，您将不仅学会如何使用 CuTe DSL 编写 GPU 核函数，更能深入理解其设计哲学，从而能够设计出适合特定硬件特性的高效内存访问模式。
2. 术语表：核心概念详解
理解 CuTe DSL 需要掌握以下核心术语，这些概念构成了整个框架的基础：
- 布局 (Layout)：CuTe 中最核心的概念，定义了逻辑坐标到物理内存偏移的映射关系。每个布局由形状（shape）和步幅（stride）组成，采用
(形状):(步幅) 的表示法。布局可以是静态的（编译时已知）或动态的（运行时确定），支持复杂的嵌套和组合操作。
- 引擎 (Engine)：负责实际内存访问的组件，类似于智能指针。引擎知道如何根据给定的偏移进行解引用操作。当与布局组合时，引擎和布局共同构成完整的张量。
- TensorSSA：静态单赋值形式的寄存器级张量值。TensorSSA 代表驻留在 GPU 寄存器中的张量数据，支持向量化的逐元素操作和归约。它是实现计算密集型操作的关键抽象。
- TV 布局 (Thread/Value Layout)：将线程标识符和值索引映射到逻辑张量坐标的特殊布局。TV 布局解耦了线程组织（如何安排线程）和数据平铺（每个线程处理多少数据），是构建高效、可扩展核函数的重要工具。
- 平铺器 (Tiler)：用于指导布局划分操作的布局或布局元组。平铺器定义了如何将大张量划分为较小的块（tile），这些块可以分配给不同的线程或线程块进行处理。
- Constexpr：编译时常量，与运行时数值类型相区别。Constexpr 值在编译时就已经确定，可以用于静态布局的形状和步幅定义，使得编译器能够进行更积极的优化。
- DLPack 互操作：支持与不同深度学习框架进行零拷贝数据交换的标准化接口。CuTe 通过
from_dlpack 函数支持 DLPack 协议，使得 PyTorch、NumPy 等框架的张量可以无缝转换为 CuTe 张量，无需数据复制。
3. 第一个核函数与启动（Hello World）
让我们从如何在 CuTe 中定义和启动最简单的 GPU 核函数开始。CuTe 使用装饰器语法来清晰地区分设备端核函数和主机端启动代码。
要定义 GPU 核函数，需要使用 @cute.kernel 装饰器。在核函数内部，可以通过 cute.arch.thread_idx() 获取当前线程的索引。核函数的主机端启动逻辑则由 @cute.jit 装饰的函数负责，它使用 .launch(grid=..., block=...) 方法指定执行配置。
下面的代码示例展示了完整的 “Hello World” 核函数：
import cutlass
import cutlass.cute as cute
@cute.kernel
def hello_kernel():
tidx, _, _ = cute.arch.thread_idx()
if tidx == 0:
cute.printf("Hello from GPU")
@cute.jit
def hello_world():
hello_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1))
hello_world()
关键点总结：
@cute.kernel 用于定义设备端执行的核函数代码。@cute.jit 用于定义主机端启动逻辑和管理执行配置。
4. 打印与数值类型：静态 vs 动态
在 CuTe 中，理解静态和动态值的区别至关重要。Python 的内置 print 函数在 CuTe JIT 编译上下文中是静态的，这意味着它只能在编译时打印已知的值。对于运行时确定的动态值，静态打印会显示为问号(?)。
要查看运行时的实际值，必须使用 cute.printf 函数。这个函数专为设备端设计，支持在 @cute.kernel 和 @cute.jit 内部打印动态值。cutlass.Constexpr 用于表示编译时常量。
CuTe 要求在使用数值时显式声明类型，如 cutlass.Int32、cutlass.Float32 等。可以使用 .to() 方法在不同类型间进行转换。
以下示例展示了静态与动态打印的区别以及数值类型的使用：
import cutlass
import cutlass.cute as cute
@cute.jit
def print_demo(a: cutlass.Int32, b: cutlass.Constexpr[int]):
print("static a:", a) # 输出: ? (动态值)
print("static b:", b) # 输出: 2 (静态常量)
cute.printf("dynamic a: {}", a) # 运行时输出实际值
cute.printf("dynamic b: {}", b) # 运行时输出2
print_demo(cutlass.Int32(8), 2)
数值类型的转换和运算示例：
import cutlass
import cutlass.cute as cute
@cute.jit
def dtypes():
a = cutlass.Int32(42) # 32位整数
b = a.to(cutlass.Float32) # 转换为32位浮点数
c = b + 0.5 # 浮点运算
d = c.to(cutlass.Int32) # 转换回整数
cute.printf("a={}, b={}, c={}, d={}", a, b, c, d)
dtypes()
关键点总结：
- 静态打印 vs 动态打印：Python 的
print 是编译时的静态打印，cute.printf 是运行时的动态打印。
- 类型显式声明：在 JIT 代码中必须显式使用 CuTe 数值类型。
- 类型转换：使用
.to() 方法在不同数值类型间进行转换。
- Constexpr：用于表示编译时常量，优化编译时决策。
5. 张量与布局：构建、索引、切片
张量和布局是 CuTe 最核心的概念。在 CuTe 中，张量被定义为”引擎 ∘ 布局”，其中引擎负责底层存储访问，布局定义逻辑索引到物理偏移的映射关系。使用 cute.make_tensor(ptr_or_engine, layout) 可以构建张量，而 cute.make_layout(shape, stride=...) 用于创建布局。
CuTe 使用独特的 (形状):(步幅) 表示法来显示布局。通过 DLPack 协议，CuTe 可以与其他深度学习框架（如 PyTorch）进行零拷贝互操作。cute.print_tensor 函数专门用于以可读格式打印张量内容。
理解布局的步幅表示法对于优化内存访问模式至关重要。例如，形状为 (M, N) 的二维张量：
- 行优先布局表示为
(M, N):(N, 1)，表示沿行的连续元素在内存中相邻。
- 列优先布局表示为
(M, N):(1, M)，表示沿列的连续元素在内存中相邻。
以下示例展示了如何创建和操作张量与布局：
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.jit
def tensor_demo(t: cute.Tensor):
cute.printf("t[0,0] = {}", t[0, 0]) # 标量访问
sub = t[(None, 0)] # 第一行的视图
frag = cute.make_fragment(sub.layout, sub.element_type)
frag.store(sub.load()) # 加载数据到寄存器片段
cute.print_tensor(frag) # 打印片段
# 创建 PyTorch 张量并转换为 CuTe 张量
arr = torch.arange(0, 12, dtype=torch.float32).reshape(3, 4)
tensor_demo(from_dlpack(arr))
布局的步幅表示法示例：
import cutlass
import cutlass.cute as cute
@cute.jit
def layout_stride_demo(m: cutlass.Int32, n: cutlass.Int32):
row_major = cute.make_layout((m, n), stride=(n, cutlass.Int32(1)))
col_major = cute.make_layout((m, n), stride=(cutlass.Int32(1), m))
print("static row-major:", row_major) # 静态打印可能显示 ?
print("static col-major:", col_major)
cute.printf("dynamic row-major: {}", row_major) # 动态打印实际值
cute.printf("dynamic col-major: {}", col_major)
layout_stride_demo(cutlass.Int32(4), cutlass.Int32(3))
关键点总结：
- 张量结构：张量 = 引擎 ∘ 布局，分别负责存储访问和索引映射。
- 布局表示法：
(形状):(步幅) 格式直观显示索引到偏移的映射。
- 内存访问模式：步幅决定了数据在内存中的排列方式，影响访问效率。
- 切片操作：使用
None 保留维度，创建张量视图而非复制数据。
- 数据移动：
.load() 和 .store() 在内存和寄存器间传输数据。
6. 寄存器张量与归约操作
TensorSSA（静态单赋值形式的张量）是 CuTe 中用于表示寄存器级张量值的核心抽象。与内存中的张量不同，TensorSSA 驻留在 GPU 寄存器中，支持高效的向量化逐元素操作。在内存张量和 TensorSSA 之间，通过 vec = tensor.load() 和 tensor.store(vec) 进行数据加载和存储。
归约操作是张量计算中的常见模式，CuTe 通过 .reduce(op, init, reduction_profile=...) 方法支持各种归约操作。归约参数包括：
- op：归约操作类型，如
cute.ReductionOp.ADD、cute.ReductionOp.MUL、cute.ReductionOp.MAX 等。
- init：初始累加器值，同时也确定了累加器的数据类型。
- reduction_profile：指定要归约的轴，可以是标量
0（归约所有轴），或元组形式指定每个轴的归约行为。
以下示例展示了如何使用 TensorSSA 进行逐元素加法：
import numpy as np
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.jit
def ssa_add(dst: cute.Tensor, x: cute.Tensor, y: cute.Tensor):
xv = x.load() # 加载到 TensorSSA
yv = y.load()
dst.store(xv + yv) # 计算并存回内存
cute.print_tensor(dst)
x_array = np.ones((2, 3), dtype=np.float32)
y_array = np.full((2, 3), 2.0, dtype=np.float32)
z_array = np.zeros((2, 3), dtype=np.float32)
ssa_add(from_dlpack(z_array), from_dlpack(x_array), from_dlpack(y_array))
归约操作示例：
import numpy as np
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.jit
def ssa_reduce(a: cute.Tensor):
v = a.load()
# 所有元素的总和
total = v.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=0)
cute.printf("total sum = {}", total)
# 行方向求和 -> 形状 (行数,)
row_sum = v.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=(None, 1))
row_frag = cute.make_fragment(row_sum.shape, cutlass.Float32)
row_frag.store(row_sum)
print("Row-wise sum:")
cute.print_tensor(row_frag)
a_array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
ssa_reduce(from_dlpack(a_array))
关键点总结：
- TensorSSA：寄存器级张量表示，支持高效向量化操作。
- 数据加载/存储：
.load() 将内存数据加载到寄存器，.store() 将结果存回内存。
- 归约操作：
.reduce() 方法支持各种归约模式，通过 reduction_profile 指定归约轴。
- 归约参数：
1 表示归约该轴，None 表示保留该轴，0 表示归约所有轴。
7. 使用 zipped_divide 的向量化加法
在 GPU 编程中，简单地将每个线程映射到单个元素的方法虽然直观，但效率不高。CuTe 提供了 cute.zipped_divide 函数，用于将张量划分为每线程块，从而实现向量化操作。这种方法允许每个线程处理多个元素，提高了内存访问效率和计算吞吐量。
cute.zipped_divide 生成的平铺张量具有双模式结构：模式-0 表示每个线程处理的块，模式-1 索引原始张量中的块。通过切片操作 (None, (mi, ni))，可以提取特定线程处理的块，然后使用 .load() 加载数据到寄存器进行计算。
数学形式化：给定一个形状为 S 的张量和一个平铺器 T，zipped_divide 创建一个新的双模式布局，其中模式-0（局部模式）形状为 T，模式-1（平铺模式）形状为 S/T。访问原始张量中的元素 (m, n) 现在需要通过双索引访问：(local_idx, tile_idx)，这可以表示为逐元素运算。
以下示例展示了如何使用 zipped_divide 实现向量化加法：
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.kernel
def vadd_kernel(g_a: cute.Tensor, g_b: cute.Tensor, g_c: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
idx = bidx * bdim + tidx
m, n = g_a.shape[1] # 线程域索引
mi = idx // n
ni = idx % n
# 每个线程加载一个 (1,4) 的向量块，相加，然后存储
g_c[(None, (mi, ni))] = g_a[(None, (mi, ni))].load() + g_b[(None, (mi, ni))].load()
@cute.jit
def vadd(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
# 将张量划分为每线程块 (1,4)
g_a = cute.zipped_divide(a, (1, 4))
g_b = cute.zipped_divide(b, (1, 4))
g_c = cute.zipped_divide(c, (1, 4))
threads = 256
vadd_kernel(g_a, g_b, g_c).launch(
grid=(cute.size(g_c, mode=[1]) // threads, 1, 1),
block=(threads, 1, 1),
)
m_size, n_size = 1024, 1024
a_tensor = torch.randn(m_size, n_size, device="cuda", dtype=torch.float16)
b_tensor = torch.randn(m_size, n_size, device="cuda", dtype=torch.float16)
c_tensor = torch.zeros(m_size, n_size, device="cuda", dtype=torch.float16)
vadd_compiled = cute.compile(vadd, from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
vadd_compiled(from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
关键点总结：
- 向量化优势：每个线程处理多个元素，提高内存访问效率和计算吞吐量。
- 平铺张量结构：
zipped_divide 创建双模式张量，分离线程块索引和每线程数据块。
- 切片操作：
(None, (mi, ni)) 高效提取每个线程的工作集。
- 执行配置：网格大小基于平铺块数量，块大小基于线程数。
8. TV 布局与每线程片段
TV 布局（Thread/Value Layout）是 CuTe 中一个强大的抽象，它将线程组织与数据平铺解耦，使得开发者可以独立优化线程映射和向量加载大小。通过 cute.make_layout_tv(thread_layout, value_layout) 可以创建 TV 布局，它将 (tid, vid) 对映射到 (TileM, TileN) 块中的位置。
这种分离提供了极大的灵活性，允许开发者根据硬件特性和算法需求独立调整线程组织和数据分块。
数学表示：TV 布局 L_tv 将线程索引 tid 和值索引 vid 映射到逻辑坐标 (m, n)。如果线程布局 L_thr 的形状为 S_thr，值布局 L_val 的形状为 S_val，那么线程块平铺块大小为 S_thr × S_val，每块线程数为 size(S_thr)，每线程值数为 size(S_val)。
以下示例展示了如何使用 TV 布局实现高效的矩阵加法：
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@cute.kernel
def tv_add_kernel(g_a: cute.Tensor, g_b: cute.Tensor, g_c: cute.Tensor, tv_layout: cute.Layout):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
# 选择线程块对应的平铺块
blk_coord = ((None, None), bidx)
blk_a = g_a[blk_coord]
blk_b = g_b[blk_coord]
blk_c = g_c[blk_coord]
# 组合 TV 布局，将 (tid, vid) 映射到物理地址
tidfrg_a = cute.composition(blk_a, tv_layout)
tidfrg_b = cute.composition(blk_b, tv_layout)
tidfrg_c = cute.composition(blk_c, tv_layout)
# 切片每线程向量
thr_coord = (tidx, None)
thr_a = tidfrg_a[thr_coord]
thr_b = tidfrg_b[thr_coord]
thr_c = tidfrg_c[thr_coord]
thr_c[None] = thr_a.load() + thr_b.load()
@cute.jit
def tv_add(m_a: cute.Tensor, m_b: cute.Tensor, m_c: cute.Tensor):
# 线程布局 (4,32)：沿 M（行）4组，沿 N（列）32个连续线程
# 值布局 (4,8)：每个线程处理 4行 x 8个连续值
thr_layout = cute.make_layout((4, 32), stride=(32, 1))
val_layout = cute.make_layout((4, 8), stride=(8, 1))
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
# 将张量平铺成线程块块
g_a = cute.zipped_divide(m_a, tiler_mn)
g_b = cute.zipped_divide(m_b, tiler_mn)
g_c = cute.zipped_divide(m_c, tiler_mn)
# 每个模式-1中的平铺块对应一个线程块；每块线程数 = TV 线程数
tv_add_kernel(g_a, g_b, g_c, tv_layout).launch(
grid=[cute.size(g_c, mode=[1]), 1, 1],
block=[cute.size(tv_layout, mode=[0]), 1, 1],
)
m_size, n_size = 2048, 2048
a_tensor = torch.randn(m_size, n_size, device="cuda", dtype=torch.float16)
b_tensor = torch.randn(m_size, n_size, device="cuda", dtype=torch.float16)
c_tensor = torch.zeros(m_size, n_size, device="cuda", dtype=torch.float16)
tv_add_compiled = cute.compile(tv_add, from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
tv_add_compiled(from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
关键点总结：
- 解耦设计：TV 布局将线程组织与数据平铺分离，提供更大的优化灵活性。
- 硬件适配：可以根据 GPU 架构特性（如 warp 大小、内存总线宽度）优化线程和值布局。
- 内存访问优化：合理的线程布局可以确保合并内存访问，提高带宽利用率。
- 组合操作：通过
cute.composition 将 TV 布局与块局部张量组合，得到每线程片段。
9. 布局代数：合并、组合、划分、乘积
布局代数是 CuTe 的数学基础，提供了一组操作符用于构建和变换复杂的布局结构。这些操作允许开发者以声明式的方式构建高效的内存访问模式。
组合操作 (Composition)
组合操作将一个布局应用于另一个布局的结果，实现嵌套映射关系。数学上，组合布局 C = A ∘ B 满足 C(i) = A(B(i))。这意味着我们首先应用布局 B 将逻辑索引 i 映射到中间索引 j，然后应用布局 A 将 j 映射到物理偏移。
合并操作 (Coalesce)
合并操作将兼容的布局模式合并为更简单的形式，减少层次结构。例如，将两个连续的步幅为1的模式合并为一个模式。
划分操作 (Divide)
划分操作将布局按给定的平铺器划分为多个子布局。划分是数据平铺的基础，允许将大张量分解为适合线程块或线程处理的小块。
乘积操作 (Product)
乘积操作创建布局的副本，用于实现广播和重复模式。乘积布局常用于实现广播模式，其中某些维度被重复多次。
以下示例展示了布局代数的基本操作：
import cutlass
import cutlass.cute as cute
@cute.jit
def layout_algebra_demo():
# 创建基础布局
a_layout = cute.make_layout((6, 2), stride=(cutlass.Int32(8), 2))
b_layout = cute.make_layout((4, 3), stride=(3, 1))
# 组合操作：A∘B
result = cute.composition(a_layout, b_layout)
print("Composition R = A ∘ B:")
print(f" A = {a_layout}")
print(f" B = {b_layout}")
print(f" R = {result}")
# 合并操作：简化布局
coalesced = cute.coalesce(result)
print("Coalesced layout C:")
print(f" C = {coalesced}")
# 划分操作：使用平铺器划分布局
layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1)))
tiler = (cute.make_layout(3, stride=3),
cute.make_layout((2, 4), stride=(1, 8)))
divided = cute.logical_divide(layout, tiler=tiler)
print("Logical divide D:")
print(f" L = {layout}")
print(f" T = {tiler}")
print(f" D = {divided}")
# 乘积操作：创建重复模式
product = cute.logical_product(
cute.make_layout((2, 2), stride=(4, 1)),
cute.make_layout(6, stride=1),
)
print("Logical product P:")
print(f" P = {product}")
layout_algebra_demo()
关键点总结：
- 数学基础：布局代数提供了构建复杂访问模式的数学基础。
- 声明式编程：通过代数操作声明性定义布局变换，而非命令式实现。
- 优化机会：合并操作可以简化布局，提高后续操作的效率。
- 模式构建：组合和乘积操作可以构建复杂的嵌套和重复模式。
- 平铺基础：划分操作是实现数据平铺和线程映射的基础。
10. 完整代码示例：矩阵乘法实现
下面我们将之前的所有概念整合到一个完整的示例中，展示如何使用 CuTe DSL 实现一个高效的矩阵乘法核函数。
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
# 定义矩阵乘法核函数
@cute.kernel
def gemm_kernel(
g_a: cute.Tensor, # 全局内存中的矩阵A
g_b: cute.Tensor, # 全局内存中的矩阵B
g_c: cute.Tensor, # 全局内存中的矩阵C
tv_layout: cute.Layout, # TV布局，定义线程组织与数据平铺
k_tile: cutlass.Int32 # K维度的平铺大小
):
# 获取线程和块索引
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
# 选择线程块对应的平铺块
blk_coord = ((None, None), bidx)
blk_a = g_a[blk_coord]
blk_b = g_b[blk_coord]
blk_c = g_c[blk_coord]
# 组合TV布局，将(tid, vid)映射到物理地址
tidfrg_a = cute.composition(blk_a, tv_layout)
tidfrg_b = cute.composition(blk_b, tv_layout)
tidfrg_c = cute.composition(blk_c, tv_layout)
# 切片每线程向量
thr_coord = (tidx, None)
thr_a = tidfrg_a[thr_coord]
thr_b = tidfrg_b[thr_coord]
thr_c = tidfrg_c[thr_coord]
# 初始化累加器
accum = thr_c.load() * 0.0
# 循环遍历K维度
for k in range(0, k_tile, 1):
# 加载当前K切片的数据
a = thr_a[..., k].load()
b = thr_b[..., k].load()
# 累加矩阵乘法的结果
accum += a * b
# 将结果存回全局内存
thr_c[None] = accum
@cute.jit
def gemm(
m_a: cute.Tensor, # 矩阵A，形状(M, K)
m_b: cute.Tensor, # 矩阵B，形状(K, N)
m_c: cute.Tensor # 矩阵C，形状(M, N)
):
m, k = m_a.shape
k, n = m_b.shape
# 定义平铺大小
tile_m = 128
tile_n = 128
tile_k = 32
# 定义线程布局和值布局
thr_layout = cute.make_layout((4, 32), stride=(32, 1)) # 128个线程，组织为4x32
val_layout = cute.make_layout((4, 8), stride=(8, 1)) # 每个线程处理4x8个元素
# 创建TV布局
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
# 将矩阵A和B平铺
g_a = cute.zipped_divide(m_a, (tile_m, tile_k))
g_b = cute.zipped_divide(m_b, (tile_k, tile_n))
g_c = cute.zipped_divide(m_c, (tile_m, tile_n))
# 启动核函数
gemm_kernel(g_a, g_b, g_c, tv_layout, tile_k).launch(
grid=[cute.size(g_c, mode=[1]), 1, 1],
block=[cute.size(tv_layout, mode=[0]), 1, 1],
)
# 测试矩阵乘法
m_size, n_size, k_size = 1024, 1024, 1024
a_tensor = torch.randn(m_size, k_size, device="cuda", dtype=torch.float16)
b_tensor = torch.randn(k_size, n_size, device="cuda", dtype=torch.float16)
c_tensor = torch.zeros(m_size, n_size, device="cuda", dtype=torch.float16)
# 编译并执行
gemm_compiled = cute.compile(gemm, from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
gemm_compiled(from_dlpack(a_tensor), from_dlpack(b_tensor), from_dlpack(c_tensor))
# 验证结果
c_cpu = c_tensor.cpu()
a_cpu = a_tensor.cpu()
b_cpu = b_tensor.cpu()
expected = a_cpu @ b_cpu
print("Max error:", torch.max(torch.abs(c_cpu - expected)).item())
这个矩阵乘法实现的关键要点：
- TV布局设计：线程布局(4,32)将128个线程组织成4组，每组32个连续线程，优化了内存访问的局部性。值布局(4,8)确保每个线程处理32个元素，充分利用了向量化加载。
- 平铺策略：将大型矩阵划分为128×128的平铺块，每个平铺块由一个线程块处理。K维度使用32的平铺大小，通过循环累加完成计算。
- 内存访问优化：使用
zipped_divide 将全局内存张量划分为平铺块，然后通过TV布局将每个平铺块进一步划分为每线程片段。这种分层划分确保了合并内存访问。
- 执行配置：网格大小等于平铺块数量，块大小等于TV布局中的线程数（128个线程）。
- 循环累加：在K维度上循环累加矩阵乘积结果，每个线程独立计算自己负责的片段。
- 数值验证：将GPU计算结果与CPU参考实现进行比较，验证了计算的正确性。
这个实现展示了CuTe DSL如何通过布局代数、TV布局和向量化操作来构建高效的GPU核函数。通过将数据布局与计算逻辑解耦，CuTe使得开发者能够专注于算法设计，同时自动生成优化的内存访问模式。
总结
CuTe DSL 提供了一套强大的抽象，用于在 NVIDIA GPU 上表达和优化GPU张量计算。通过将数据布局与计算解耦，CuTe 允许开发者专注于算法逻辑，同时自动生成高效的内存访问模式。本指南从基础概念开始，逐步介绍了核函数启动、打印与类型系统、张量与布局、寄存器张量与归约、向量化操作、TV布局以及布局代数。
掌握这些概念后，您将能够使用 CuTe DSL 编写高效、可维护的 GPU 核函数，并能够根据具体硬件特性进行优化，尤其是在C++层面深入理解内存布局与性能的关系。CuTe 的学习曲线可能较陡，但一旦掌握，它将极大地提高您在 GPU 编程中的生产力和性能。希望本指南能为您打开 CuTe 世界的大门，也欢迎在云栈社区与其他开发者交流心得，共同进步。
发表于 5 天前
|
查看: 21|
回复: 0
