GPU 词汇表中新增关于 CuTe DSL、CUTLASS 和 CuTe 的文章——这些工具用于在当代数据中心 GPU 上编写一些性能最高的内核。https://modal.com/gpu-glossary/host-software/cute-dsl…

什么是 CuTe DSL？ | GPU 词汇表

来源：https://modal.com/gpu-glossary/host-software/cute-dsl CuTe DSL 是一种基于 Python 的领域特定语言 (DSL)，用于编写和动态编译内核，兼具高性能和高开发效率。

CuTe DSL 是 CUTLASS 的一部分，CUTLASS 是一组 CUDA C++ 模板和 DSL。与提供常用操作即用内核的 cuBLAS 或 cuDNN 不同，CUTLASS 堆栈提供了用于可组合地定义高性能内核的工具。

CuTe DSL 的核心抽象包括布局（layouts）、张量（tensors）、硬件原子（hardware atoms）和平铺操作（tiled operations）。布局描述数据在内存中和跨线程的组织方式。张量将数据指针或迭代器与布局元数据结合。原子表示基本的硬件操作，如矩阵乘累加（MMA）或内存拷贝。平铺操作描述原子如何在线程块（thread block）和线程束（warp）上应用。底层细节参见 CuTe。

从 Python 启动 CuTe DSL 内核时，Python 程序调用一个 @cute.jit 函数，该函数再启动一个 @cute.kernel 函数。

@cute.jit 装饰器声明一个 JIT 编译的函数，可以从 Python 或其他 CuTe DSL 函数调用。@cute.kernel 装饰器定义了一个 GPU 内核函数，可以从 @cute.jit 函数启动。Python 代码不能直接调用 @cute.kernel 函数。

例如，我们来看一个简单的（未优化）CuTe DSL 内核，用于两个一维张量的逐元素加法——这是 GPU 编程的“Hello World”，可追溯到先于 CUDA 并启发了 CUDA 的 Ian Buck 的 Brook 框架。你可以编辑此内核并使用此 Modal Notebook 在 B200 GPU 上执行它。

python

`` import cutlass.cute as cute import torch

Tensor = cute.Tensor | torch.Tensor

@cute.kernel def elem_add_kernel(a: cute.Tensor, b: cute.Tensor, out: cute.Tensor): block_x, _, _ = cute.arch.block_idx() block_dim_x, _, _ = cute.arch.block_dim() thread_x, _, _ = cute.arch.thread_idx()

i = block_x * block_dim_x + thread_x

if i < out.shape[0]:
    out[i] = a[i] + b[i]

@cute.jit def elem_add(a: Tensor, b: Tensor, out: Tensor): n = out.shape[0] threads_per_block = 128 blocks = (n + threads_per_block - 1) // threads_per_block

elem_add_kernel(a, b, out).launch(
    grid=(blocks, 1, 1),
    block=(threads_per_block, 1, 1),
)

``

elem_add_kernel 函数是内核。每个线程计算一个输出元素。全局元素索引 i 由线程块索引、块内线程数和块内线程索引计算得出：

python

i = block_x * block_dim_x + thread_x

elem_add 函数计算覆盖输出张量所需的线程块数，并用一维线程块网格启动内核。

这个例子是教学性的，未优化。即便如此，它展示了一个良好的基本访问模式：相邻线程读取 a 和 b 的相邻元素，然后写入 out 的相邻元素。这正是对全局内存进行合并访问所需的模式；参见内存合并。

布局问题是 CuTe DSL 对高性能内核有用的原因之一。性能工程很困难，因为内核必须紧密映射到硬件：哪些线程处理哪些数据、如何访问内存、如何平铺工作，以及生成的代码应使用哪些硬件操作。CuTe DSL 允许程序员显式表达这些映射，同时在各种形状和流式多处理器架构上重复使用大量相同的内核代码。

这可能会让来自其他领域的关注性能的工程师感到惊讶——用 Python 这样的解释型语言编写的程序如何能与用编译型语言编写的程序竞争？

答案是 CuTe DSL 内核是即时编译（JIT）的。Python 源代码被转换为抽象语法树（AST），用代理参数进行跟踪，然后编译。注意，JIT 编译代码仅支持 Python 语义的一个子集。

在撰写本文时，对于 CUTLASS 4.x，编译堆栈会依次经过多层中间表示（MLIR）、PTX IR，再到设备特定的 SASS，然后才被执行。

以 FlashAttention-4 内核为例。我们对开源代码的分析讲述了它如何使用流水线线程束特化、Tensor Core 操作以及张量内存（Tensor Memory）和张量内存加速器（Tensor Memory Accelerator）操作，直接从 CuTe DSL 实现最先进的性能。

有关 CuTe DSL 的更多详细信息，请参阅 NVIDIA 的 CuTe DSL 文档和 CuTe DSL 概述博客。