CuTe新功能：100倍加速吃透Python GPU编程红利！

各位跨境电商、游戏、支付和贸易领域的同行们，大家都在追求极致的效率和性能，尤其是在AI技术飞速发展的当下。今天，新媒网跨境获悉，我们有一个重磅消息要分享给大家——NVIDIA最新推出的CUTLASS 4中的CuTe DSL（领域特定语言），这无疑是为我们Python开发者量身打造的一款性能优化神器。

简单来说，CuTe DSL让我们这些习惯了Python便捷的开发者，也能不费力地编写出底层GPU核函数，而无需深陷C++模板元编程的复杂泥沼。更重要的是，它能在不同代次的NVIDIA GPU芯片（包括最新的Blackwell和之前的Ampere架构）上，保持与C++版本不相上下的Tensor Core计算效率。这意味着，无论是深度学习模型训练、推理，还是其他高性能计算场景，我们的Python代码都能跑出C++般的“飞快”速度。

长久以来，NVIDIA的CUTLASS 3.x及其核心组件CuTe，为开发者们在Tensor Cores上实现峰值性能提供了强大的抽象能力。但不得不说，C++模板的大量使用，也带来了编译时间过长的问题。随着Python在人工智能研究和生产工作流中越来越普及，以及即时编译（JIT）技术的兴起，CUTLASS 4的诞生就显得水到渠成了。它积极响应了时代需求，将强大的GPU编程能力带给更广泛的Python社区。

这篇文章，我们就来深入探讨CuTe DSL的优势。它不仅提供了与C++版本一致的编程接口，还能在不同GPU芯片上保持卓越的Tensor Core效率，最关键的是，它的编译成本比C++大大降低，这对我们快速迭代和部署应用至关重要。

如果大家想深入了解CuTe和CUTLASS 3.x的基础知识，可以查阅《CUTLASS：处理多维数据的原则性抽象》和《CUTLASS 3.x：GEMM核函数设计的正交、可重用和可组合抽象》这两份资料。

CuTe DSL：CUTLASS 4的核心基石

CUTLASS 4中的全新CuTe DSL（目前处于Beta测试阶段），将CuTe的强大功能带给了Python开发者。它让编写底层GPU核函数变得像搭积木一样简单，彻底摆脱了C++模板元编程的繁琐。为了降低学习门槛，CuTe DSL沿用了CuTe底层的核心概念，所以只要理解了CuTe的精髓，上手CuTe DSL就会非常快。

大家可以访问美国英伟达公司在GitHub上的CUTLASS项目，那里有一些CuTe DSL的实战示例，包括密集GEMM的持久化变体、分组GEMM以及融合多头注意力（FMHA）等。

CuTe DSL与CuTe C++的巅峰对决

CuTe凭借其强大的布局表示和代数运算能力，在过去十多年间为多代NVIDIA GPU架构提供了统一的GPU编程模型。CuTe DSL完美继承了CuTe C++的编程模型，同时融入了Python的易用性。这带来的好处显而易见：编译速度大幅提升、错误信息更友好、学习曲线更平缓，并且能近乎即时地融入到基于Python的深度学习框架中。

通过对比C++和DSL代码，我们会发现它们的编程模型和模式几乎一模一样，唯一的区别只在于C++和Python各自的语法。

TiledMMA：精细化GPU算力管理

cute::TiledMma是一个空间微核函数，它负责描述硬件MMA原子在“线程”和数据之间的平铺和置换。通过这种方式，它能够为任何硬件MMA编写规范的三重循环，无论是SIMT FP64还是最新的Blackwell架构NVFP4 Tensor Core指令，都能被高效驾驭。

auto tiled\_mma = make\_tiled\_mma(SM100\_MMA\_F16BF16\_SS<TA, TB, TC, 128, 128, UMMA::Major::MN, UMMA::Major::MN>{}, Layout<Shape<\_1,\_1>>{});

// Allocate "fragments" -- these are actually umma tmem and smem descriptors
Tensor tCrA = tiled\_mma.make\_fragment\_A(sA); // (MMA,MMA\_M,MMA\_K,PIPE)
Tensor tCrB = tiled\_mma.make\_fragment\_B(sB); // (MMA,MMA\_M,MMA\_K,PIPE)

// Allocate TMEM
Tensor tCtC = tiled\_mma.make\_fragment\_C(tCgC);// (MMA,MMA\_M,MMA\_N)

for (int k\_block = 0; k\_block < size<2>(tCrA); ++k\_block) {
  static\_assert(size<2>(tCrA) == size<2>(tCrB), "A and B contraction modes do not match!");
  gemm(mma, tCrA(\_,\_,k\_block), tCrB(\_,\_,k\_block), tCtC)
}

# Construct a tiled\_mma item
atom = tcgen05.MmaF16BF16Op(
    io\_dtype,
    acc\_dtype,
    mma\_inst\_shape\_mnk, #(128, 128, 64)
    tcgen05.CtaGroup.ONE,
    tcgen05.OperandSource.SMEM,
    tcgen05.OperandMajorMode.K,
    tcgen05.OperandMajorMode.K,
)

tiled\_mma = cute.make\_tiled\_mma(atom)

tCrA = tiled\_mma.make\_fragment\_A(sA) # (MMA, MMA\_M, MMA\_K,PIPE)
tCrB = tiled\_mma.make\_fragment\_B(sB) # (MMA, MMA\_N, MMA\_K,PIPE)
tCtC = tiled\_mma.make\_fragemnt\_C(tCgC) # (MMA, MMA\_M, MMA\_N)

for k\_block\_idx in cute.size(tCrA, mode = 2):
    assert(cute.size(tCrA, mode = 2) == cute.size(tCrB, mode = 2), "A and B contraction modes do not match!");
    cute.gemm(
    tiled\_mma,
    tCtC,
    tCrA[None, None, k\_block\_idx],
    tCrB[None, None, k\_block\_idx],
    tCtC)

TiledCopy：GPU上的数据搬运艺术

经典的cute::copy是一个单循环操作，负责发出数据移动指令，将一个张量复制到另一个。它利用张量的布局来处理可能发生的任何转置或置换。cute::TiledCopy则是一种类型，用于表示和验证优化后的数据传输在任意两个张量之间（例如，在不同内存空间如全局内存到共享内存，或在内存内部，无论是否包含布局转换）的可行性。它可以使用任何硬件加速的复制原子操作来完成。

using TMEM\_LOAD = typename std::conditional<sizeof(TC) == 4, SM100\_TMEM\_LOAD\_16dp256b1x, SM100\_TMEM\_LOAD\_16dp256b1x\_16b>::type;
// tCtC are accumuator layout
auto tiled\_ldtm = make\_tmem\_copy(TMEM\_LOAD{}, tCtC);
auto thr\_ldtm = tiled\_ldtm.get\_slice(threadIdx.x);
Tensor tDtC = thr\_ldtm.partition\_S(tCtC); // ((TMEM\_LOAD,#TMEM\_LOAD),MMA\_M,MMA\_N)
Tensor tDgC = thr\_ldtm.partition\_D(tCgC); // ((TMEM\_LOAD,#TMEM\_LOAD),MMA\_M,MMA\_N)
Tensor tDrC = make\_tensor<TC>(shape(tDgC));// ((TMEM\_LOAD,#TMEM\_LOAD),MMA\_M,MMA\_N)

// TMEM\_LOAD
copy(tiled\_ldtm, tDtC, tDrC);

# Construct a tensor memory to register memory (T2R) tiled\_copy item
# tCtACC are accumulator tensor, layout as (MMA, MMA\_M, MMA\_N)
# tCgC is the partitioned results (MMA, MMA\_M, MMA\_N, RestM, RestN, RestL) of global tensor C (M, N)
copy\_atom = cute.make\_copy\_atom(
    tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE),
    cutlass.Float32)

tiled\_copy\_t2r = tcgen05.make\_tmem\_copy(copy\_atom, tCtACC)
thr\_copy\_t2r = tiled\_copy\_t2r.get\_slice(tidx)

# This is tensor memory layout (T2R\_M, T2R\_N, EPI\_M, EPI\_N)
tT2R\_tAcc = thr\_copy\_t2r.partition\_S(tCtACC)
# (T2R\_M, T2R\_N, EPI\_M, EPI\_N, RestM, RestN, RestL)
tT2R\_gC = thr\_copy\_t2r.partition\_D(tCgC)

# Construct register memory layout from the partitioned global tensor
tT2R\_rAcc = cute.make\_fragment(
    tT2R\_gC[None, None, None, None, 0, 0, 0].shape,
    cutlass.Float32)

cute.copy(tiled\_copy\_t2r, tT2R\_tAcc, tT2R\_rAcc)

CuTe DSL在多代GPU上的性能表现

CUTLASS C++之所以能被训练和推理框架广泛采用，其核心原因之一就是它能提供“快如闪电”的性能。新媒网跨境了解到，CuTe DSL在性能上几乎能与CUTLASS C++持平，并且还有更多的优化正在路上。更棒的是，CUTLASS 3及其底层的CuTe已经在过去几代GPU硬件上的研发和生产环境中得到了广泛应用，而这些GPU在生产环境中往往有很长的生命周期，有时甚至在异构设置中。CuTe DSL一经推出，就支持从NVIDIA Ampere到Blackwell的全系列GPU，完美兼容这些现有部署。

美国英伟达公司Blackwell架构上的性能验证

我们测量了密集GEMM、分组GEMM和FMHA这三个关键操作在CUTLASS C++和CuTe DSL上的性能。总的来说，CuTe DSL的性能与CUTLASS C++旗鼓相当。

密集GEMM

我们分别在float16和float8 e4m3两种精度设置下测量了密集GEMM的性能，两者都使用float32作为累积精度。图1展示了在美国英伟达公司DGX B200服务器上，CuTe DSL密集GEMM与CUTLASS 3密集GEMM的对比基准测试结果。其中，x轴代表测试的问题规模，y轴则表示通过NVIDIA Compute Nsight捕获到的Tensor Core数学吞吐量效率。值得注意的是，对于小规模的GEMM-K问题（K=512），DSL核函数目前的性能略低于C++。这主要是由于进入核函数数学计算之前的同步成本较高，不过团队正在积极优化这一部分，相信很快就能迎头赶上。
Bar chart titled ‘B200 Dense GEMM Math Throughput Efficiency %’ showing that DSL performance is on par with C++ except for the small K cases.

图1：在美国英伟达公司DGX B200上，CuTe DSL密集GEMM与CUTLASS 3密集GEMM的对比基准测试。

分组GEMM

分组GEMM的对比基准测试同样使用了CuTe DSL分组GEMM和CUTLASS 3分组GEMM，两者表现一致。
Bar chart titled ‘B200 Float16 I/O Group GEMM Math Throughput Efficiency %’ showing that DSL performance is on par with C++.

图2：在美国英伟达公司DGX B200上，Float16 I/O的CuTe DSL分组GEMM与CUTLASS 3分组GEMM的对比基准测试。

融合多头注意力（FMHA）

融合多头注意力（FMHA）的对比基准测试采用了CuTe DSL FMHA和CUTLASS 3 FMHA，性能表现同样保持了高水准。
Bar chart titled ‘B200 Float16 I/O Flash Attention Math Throughput Efficiency %’ showing that DSL performance is on par with C++.

图3：在美国英伟达公司DGX B200上，Float16 I/O的CuTe DSL闪存注意力与CUTLASS 3闪存注意力的对比基准测试。

Ampere架构上的性能：密集GEMM

对比基准测试使用了CuTe DSL密集GEMM（Ampere）和CUTLASS 3密集GEMM（Ampere）。
Bar chart titled ‘A100 FP16 I/O Dense GEMM Math Throughput Efficiency %’ showing that DSL performance is slightly slower than C++ and perf gaps to be investigated.

图4：在美国英伟达公司A100上，Float16 I/O的CuTe DSL密集GEMM与CUTLASS 3密集GEMM的对比基准测试。在这里，DSL性能略低于C++，这些性能差距正在进一步调查中。

编译时间的效率革命

CuTe DSL为核函数开发者提供了利用CuTe抽象进行即时编译（JIT）的能力，这彻底解决了C++模板编译时间过长的问题。正如图5所示，编译时间的缩短令人惊叹，平均而言，可以达到两个数量级的提升。这不仅能让核函数开发者快速尝试更多的平铺尺寸和布局形状，迅速找到最佳配置，从而获得极高的性能，还能显著缩短PyTorch Inductor等框架的自动调优总时间。例如，在Blackwell架构上，GEMM的编译速度比C++快了约100倍，而闪存注意力（flash attention）的编译速度也提升了30-50倍。这样的效率提升，对我们快速迭代和抓住市场机遇的跨境人来说，无疑是如虎添翼。
Bar chart showing compilation time on B200 is far faster than C++ with speedups of 33 -116x.

图5：美国英伟达公司Blackwell架构上的编译时间比C++快得多。

轻松融入深度学习框架

CuTe DSL通过支持DLPack协议，能够直接接收主流深度学习框架的张量数据作为输入，并将其转换为cute.Tensor，而无需复制底层内存。这种“零拷贝”的机制极大地提升了效率。同时，CuTe DSL的原生Python接口，使得深度学习框架能够直接嵌入定制的核函数，不再需要编写繁琐的“胶水代码”，也无需深厚的CUDA C++专业知识。这无疑加速了开发周期，让研究人员和工程师能够在其现有的模型流水线中，快速原型设计和部署定制的线性代数核函数。CuTe DSL可组合的布局抽象简化了复杂内存和线程映射的表达，这对于高效利用NVIDIA Ampere、Hopper和Blackwell架构上的Tensor Core硬件至关重要。