SuperAPI AI 工具集

kernel-cute-writing

编程与开发

查看项目

更新于 11-11 02:25

技能介绍

本技能专注于使用 NVIDIA CuTe DSL(CUTLASS 4.x 的 Python API)编写 GPU 内核。CuTe DSL 是一种基于 Python 的领域特定语言,通过 MLIR 和 ptxas 将 CUTLASS C++ 模板抽象为 Python 接口,并即时编译为优化的 CUDA 内核。

主要能力涵盖:编写逐元素、GEMM 和自定义操作内核;优化 GPU 内存访问模式(向量化加载、TMA、共享内存);为 Ampere/Hopper/Blackwell 架构构建张量核心(MMA)内核;集成自定义 GPU 内核到 PyTorch 或 JAX;无需 C++ 元编程即可原型化高性能内核。

使用方式上,技能提供多个工作流:推荐从官方示例出发,先阅读参考文档理解模式,再根据实际负载修改而非直接复制;逐元素内核遵循四条不变原则(使用 from_dlpack 对齐加载、根据元素宽度推导向量大小、使用 zipped_divide 合并访问、使用 dynamic_expr 进行边界检查),并通过 cute.compile() 预编译避免重复编译开销;GEMM 内核采用三级分区(CTA 分块、TiledCopy 复制分区、TiledMMA 计算分区),使用 swizzled 共享内存布局和异步流水线;框架集成支持通过 TVM FFI 实现零开销张量传递;调试与性能分析提供环境变量(CUTE_DSL_PRINT_IR、CUTE_DSL_KEEP_PTX 等)、cute.printf() 和 Nsight 工具。

适用场景包括:需要在 Python 中编写高性能 CUDA 内核的开发者、需要张量核心 GEMM 或注意力机制内核的研究人员、需要将自定义内核集成到 PyTorch/JAX 的工程师。限制条件:仅支持 Linux x86_64、Python 3.10–3.13、NVIDIA Ampere 及以上 GPU(SM80/SM90/SM100)、CUDA 驱动 ≥ 575.51.03;对于 Hopper/Blackwell 上带额外张量(如偏置、激活输入)的 GEMM 内核,因 TMA 描述符修改过于复杂而不支持,建议采用两内核方案。安装命令为 pip install nvidia-cutlass-dsl。