CUTLASS 4.0：Python接口赋能CUDA开发者，CuTe重塑GPU高性能计算范式
NVIDIA CUTLASS库在生成式AI时代迎来重大升级，发布4.0版本并引入Python接口，旨在赋能CUDA开发者最大化GPU性能。
CUTLASS 3.x引入了核心组件CuTe（CUDA Tensors and Spatial Microkernels），通过将“布局”（Layout）提升为编程模型的一等公民，极大地简化了线程-数据组织。CuTe以分层、可组合的抽象模型，提供了对线程和数据布局的统一描述与操作代数，使得开发者能清晰、静态地检查索引逻辑，同时保持与CUTLASS 2.x同等级别的性能和Tensor Core覆盖率。新版本强调了可定制化层设计、编译时正确性检查、简化的API以及对NVIDIA Hopper H100和Blackwell B200等最新硬件特性的优化支持，如WGMMA、UMMA、TMA及线程块聚簇。
CuTe通过Layout和Tensor对象封装了数据类型、形状、内存空间及布局，并自动处理复杂的索引计算。其分层布局概念支持静态和动态信息，允许用户从简单布局构建复杂布局，或跨布局进行分区，使开发者能专注于算法逻辑，将线程-数据映射的复杂性交给CuTe处理。这一范式消除了GPU编程中线程到数据映射的难题，实现了线程布局与数据布局的解耦，并支持任意分区模式。
在原子（Atom）层面，CuTe将PTX指令与线程-数据形状元数据相结合，封装为执行硬件加速数学或拷贝操作的最小单元。通过“分块MMA”（Tiled MMA）和“分块拷贝”（Tiled Copy），CuTe能够将这些原子构建成可复用的组件，实现更复杂的GPU微内核，并提供统一的API来处理数据分区，例如将单个warp的MMA原子扩展到整个线程块，以支持更大的维度。这为构建高效的GEMM（General Matrix Multiply）主循环提供了基础，同时后续的CUTLASS层面还将处理计算与拷贝指令的时间重叠和流水线优化。
Cris Cecka
网友讨论