1,001 Ways to Write CUDA Kernels in Python

Leo Fang, CUDA Python Tech Lead | GTC 2025

目录

• CUDA 与 Python：以光速提升生产力的故事
• CUDA加速Python：贯穿整个技术栈的一流Python支持
• 使命宣言
• 案例研究：分段归约 (Segmented Reduction

• 编写CUDA内核的挑战

  • 编写CUDA内核曾经很困难
  • 编写CUDA内核的难点：内核实现
  • 编写CUDA内核困难的原因

• 简化CUDA编程的途径

  • 基础知识的可移植性
  • 方法一：使用现代C++
  • 方法二：考虑使用Python

• Python方式：即时编译 (JIT) 的优势

  • JIT vs AOT 编译流程
  • JIT编译的优势

• Pythonic CUDA 工具链与库

  • 介绍 cuda.core：以Pythonic方式访问CUDA特性
  • （重）介绍 numba-cuda：一个 Pythonic 的 CUDA DSL + JIT 编译器
  • 介绍 cuda.cooperative：在 Python 核函数中使用 CUB 设备算法
  • 分段归约作为线性代数问题
  • 介绍 nvmath-python：Python 设备数学库
  • 介绍 cuTile：CUDA 在 SIMT 之外新增了 Tile 编程模型
  • Warp：具备 Tile 支持的自动微分 Python Kernel
  • 在 Python 中引入 CUTLASS 编程

• Python Kernel DSLs/Compilers 大观

• 如何选择最适合编写 Kernel 的工具

• 编写 Kernel 的最佳方式是不编写任何 Kernel

  • 利用现有生态系统
  • 引入 cuda.parallel
  • nvmath-python 再次前来救场
  • Python 中的数组编程
  • CUB：C++ 的解决方案

• 总结

• 附录：CUDA Kernel 的生命周期与编译器架构

  • Numba 核函数的生命周期（更新版）
  • JIT-time linker：最后的缺失环节
  • “一个 Warp 内核的生命周期”
  • CUTLASS Python 架构

• 致谢

CUDA 与 Python：以光速提升生产力的故事

本演讲是关于CUDA和Python系列演讲的第二部分，重点深入探讨内核（Kernels）。

系列演讲日程安排如下：
- 第0步：聆听我们的Python故事
- 加速Python：社区与生态系统 (Accelerated Python: The Community and Ecosystem)
- 星期二, 3月18日, 3:00 PM - 3:40 PM PDT

• 第1步：学习CUDA Python工具

  • CUDA Python开发者工具箱 (The CUDA Python Developer's Toolbox)
  • 星期三, 3月19日, 10:00 AM - 10:40 AM PDT

• 第2步：深入内核

  • 1,001种用Python编写CUDA内核的方法 (1,001 Ways to Write CUDA Kernels in Python)
  • 星期三, 3月19日, 11:00 AM - 11:40 AM PDT

• 第3步：掌握Tensor Core

  • 使用CUTLASS 4.0在Python中启用Tensor Core编程 (Enable Tensor Core Programming in Python With CUTLASS 4.0)
  • 星期五, 3月21日, 11:00 AM - 11:40 AM PDT

CUDA加速Python：贯穿整个技术栈的一流Python支持

"一流"意味着CUDA现在将Python融入到平台的每一个API和每一个特性中。CUDA Python生态系统提供了跨越不同层级的互操作性，允许开发者混合和匹配各个层次的库。

技术栈各层级及其对应的Python库包括：

• 框架与领域特定语言 (Frameworks & DSLs): TensorRT, Omniverse, Warp
• 软件开发工具包 (SDKs): RAPIDS, CUDA-Q, Holoscan
• 领域特定库 (Domain-Specific Libraries): cuPy/Numeric, cuDNN-python
• 加速库 (Accelerated Libraries): cuda.parallel, nvmath-python
• 通信库 (Communication Libraries): mpi4py, nvshmem4py
• 设备库 (Device Libraries): cuda.cooperative, nvmath-python
• 内核编写 (Kernel Authoring): numba-cuda, cuTe, CUTLASS
• 编译器栈 (Compiler Stack)
• 主机运行时与工具 (Host Runtimes & Tools): cuda.bindings, cuda.core, Nsight Tools

使命宣言

本次演讲将带领我们踏上一段自下而上的旅程，旨在实现Python中高效的内核开发。我们将重点关注技术栈的底层部分。

案例研究：分段归约 (Segmented Reduction)

本次演讲将使用“分段归约”作为贯穿始终的案例。该操作将一个一维数组分成若干段，并对每一段进行求和（或其他归约操作）。

例如，将数组 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] 分成三段，每段4个元素，分别求和得到 [10, 26, 42]。

编写CUDA内核的挑战

编写CUDA内核曾经很困难

因为它是C++的语言扩展。以下是使用C++和Thrust库实现分段归约的主机端代码示例。核心问题是：如何编写这个内核（segmented_reduction_kernel）？

编写CUDA内核的难点：内核实现

以下是一个用C++编写的分段归约CUDA内核的示例代码。它涉及手动管理共享内存（shared memory）、计算线程索引、展开循环（#pragma unroll）、线程同步（__syncthreads()）以及在线程块（thread block）内进行并行归约。

从教学角度出发，我们假设段的长度等于线程块大小的整数倍。在实际应用中，我们还需要额外的逻辑来处理多个分片（tile）和可能剩余的“部分”分片，这使得段长度逻辑需要作为内核参数传入。

编写CUDA内核困难的原因

传统的CUDA C++编程之所以困难，是因为开发者需要掌握大量底层知识，并处理复杂的编译流程。

所需知识点：
- 理解CUDA线程层级: 线程（threads）、线程束（warps）、线程块（blocks）、线程块集群（thread block clusters）、网格（grids）等。
- 理解CUDA内存层级: 寄存器（registers）、本地内存（local memory）、共享内存（shared memory）、分布式共享内存（distributed shared memory）、全局内存（global memory）等。
- 理解CUDA编程模型: 异步执行、内存分配、内核启动、流排序（stream ordering）、同步等。
- 理解GPU架构: 流式多处理器（streaming multiprocessor）、占用率（occupancy）、合并加载/存储（coalesced load/store）、矢量化（vectorization）、内存带宽、张量核心（tensor core）等。
- 理解用于表达算法的C++机制: 模板（templates）、迭代器（iterators）、指针（pointers）、引用（references）、容器（containers）、类型萃取（type traits）等。
- 理解CUDA C++编译工作流: 离线编译、主机编译器、设备编译器、对象链接、可重定位设备代码等。

右侧图表展示了使用NVCC编译器将.cu文件（包含主机C++代码和设备C++代码）编译为最终可执行文件的过程。

简化CUDA编程的途径

基础知识的可移植性

尽管CUDA C++编程很复杂，但其中涉及的CUDA基础知识（如线程/内存层级、编程模型、GPU架构）在不同编程语言之间是通用的。

对于希望深入学习这些基础知识的开发者，推荐以下GTC演讲：
- S72897 - 如何编写CUDA程序：并行编程版 (S. Jones)
- S72683 - 最大化内存带宽和隐藏延迟的CUDA技术 (Hendriksen & Elafrou)
- S72685 - 最大化计算和指令吞吐量的CUDA技术 (Pinzone & Clark)
- S72686 - 最大化并发和系统利用率的CUDA技术 (Terentyev & Papadopoulou)
- S72687 - 从Grace Hopper获取最佳性能 (A. Subramaniam)

同时，可以查阅新的加速计算中心 (Accelerated Computing Hub)获取学习资源。

方法一：使用现代C++

一种方法是使用现代C++。
- 使用标准的CUDA C++库（如Thrust, CUB, libcu++, CUTLASS, MathDX等）以及用户定义的迭代器和/或lambda函数。
- 推荐以下关于现代CUDA C++的演讲：
- S72472: CUDA C++开发者工具箱 (B. Leilbach)
- S72575: 你应该如何编写一个CUDA C++内核 (G. Evtushenko)
- S72574: 以光速构建CUDA (R. Maynard)
- S72897: 如何编写CUDA程序：并行编程版 (S. Jones)

• 加速计算中心也提供了相关资源：https://github.com/NVIDIA/accelerated-computing-hub/tree/main/gpu-cpp-tutorial

方法二：考虑使用Python

我们可以用Python替代C++来简化开发流程，尤其是编译过程。当前的cuda.core流程依然需要用户编写C++内核。那么，我们能否在Python中直接编写内核，而无需接触C++呢？

当前流程回顾：
1. 用户提供一个C++内核 (.cuh 文件)。
2. Python代码通过cuda.core调用NVRTC（NVIDIA运行时编译器）。
3. NVRTC将C++代码编译成PTX或CUBIN。
4. Python代码通过cuLaunchKernel启动编译好的内核。

接下来的内容将探讨完全在Python中编写内核的方法。

Python方式：即时编译 (JIT) 的优势

JIT vs AOT 编译流程

• Python方式 (“在线”/即时编译 JIT):

  • 用户在Python中定义内核。
  • 通过Clang/LLVM、NVRTC/nvJitLink/NVVM、MLIR等工具链在运行时即时编译成.ptx/.cubin。
  • 通过cuLaunchKernel启动。
  • 优势: 无需C++主机编译器，所有操作都在进程内完成，提高了生产力和性能。

• C++方式 (“离线”/预先编译 AOT):

  • 开发者编写.cu文件。
  • 使用NVCC进行离线编译，分别处理主机代码和设备代码，最终链接成一个可执行文件。

JIT编译的优势

即时编译（JIT）允许我们仅编译当前需要的部分。

• Python (JIT):

  • segmented_reduction_kernel[...] (arr, ...)
  • 内核的具体类型由运行时的输入参数arr决定。

• C++ (AOT):

  • 需要编写复杂的模板元编程代码（如if constexpr (std::is_same_v<T, double>)）来为所有可能的数据类型预先生成不同版本的内核。

JIT避免了为所有可能类型预先编译内核的开销，使得代码更简洁，编译更高效。

Pythonic CUDA 工具链与库

介绍 cuda.core：以Pythonic方式访问CUDA特性

cuda.core库使得在Python中开发C++内核变得更加容易。它提供了对运行时编译器和链接器工具链的Pythonic封装。

核心特性：

• cuda.core.{Program, Linker, ObjectCode, Kernel1}：封装了运行时编译链接工具。
• cuda.core.launch：提供了Pythonic的内核启动方式。
• C++编译时模板实例化变为Python的运行时实例化，大大简化了开发流程。

代码示例：
如上图所示，可以直接在Python中读取.cuh文件内容，使用Program对象编译它，并通过指定模板参数（如<float, int, {block_size}>）来实例化和获取内核，最后在CUDA流上启动该内核。

cuda.core 的更多功能:
- cuda.core.{Device, Stream, Event}: CUDA运行时的入口点。
- cuda.core.{Buffer, MemoryResource}: 受RAPIDS内存管理器(RMM)启发的内存管理Pythonic封装。
- cuda.core.utils.StridedMemoryView 和 as_strided_memory: 对DLPack和CUDA Array Interface的Pythonic封装，支持稠密张量/数组。
- cuda.core v0.2.0 已于本周发布：
- 目前处于beta/实验阶段。
- pip/conda 可安装。
- 支持Linux/Windows。
- 支持最新的两个CUDA主版本(11/12)。

未来工作:

• 支持CUDA Graphs、内存管理、进程间通信、绿色上下文等。

（重）介绍 numba-cuda：一个 Pythonic 的 CUDA DSL + JIT 编译器

Numba 的 CUDA 目标正在从上游 Numba 项目中分离出来。
- 项目地址：https://github.com/NVIDIA/numba-cuda
- 由 NVIDIA 官方支持。

该幻灯片通过一个 square 函数示例，对比了使用 C++ 和 Numba 编写 CUDA 核函数的异同。Numba 提供了更 Pythonic 的语法，同时保留了对 CUDA 执行模型的控制。

使用 numba-cuda 进行 Python 化实现

尝试使用 numba-cuda 将分段归约的 C++ 代码 Python 化。代码结构与 C++ 版本非常相似，包括：
1. 分配共享内存。
2. 获取线程索引。
3. 从全局内存加载值并计算部分聚合。
4. 将部分聚合存储到共享内存。
5. 在共享内存内进行并行归约。
6. 将归约结果写回输出。

尽管代码变成了 Python，但其底层逻辑和复杂性与 C++ 版本相差无几。这引出了一个问题：我们确实摆脱了 C++，但代码真的更容易编写了吗？答案是“也许并非如此”。

许多操作，如从全局内存加载数据和对其求和，是常见的例程。开发者希望有人能预先写好这些例程，从而简化代码编写。

介绍 cuda.cooperative：在 Python 核函数中使用 CUB 设备算法

cuda.cooperative 库旨在解决手动编写复杂并行原语的问题。它提供了：
- 对基础 CUDA 设备协作算法的 Pythonic 访问。
- 利用对 CUB C++ 库的投入，并将其暴露给 Python。
- 填补了 Python GPU 库开发者的空白。
- 故意设计得比 Python 数组/张量库更底层。
- 支持自定义数据类型和用户提供的操作。
- 在不离开 Python 的情况下，提供与 C++ 相当的性能和控制力。
- 使用 LTO IR 是实现高性能的关键。
- JIT 编译生成小巧的二进制文件。
- 基于 Apache 2.0 协议开源。

代码示例展示了如何使用 coop.block.load 和 coop.block.reduce 来极大简化分段归约的实现。

即将推出:
- 首次 Beta 版本发布。
- 可通过 pip/conda 安装。
- 适用于所有 CUDA 工作环境。
- 更多算法（scan, histogram, shuffle, ...）。
- 支持 N-D 线程块。

协作设备 API 使 SIMT（单指令多线程）核函数的编写更加容易。下图直观地展示了协作求和的过程：首先并行求和，然后在块内进行串行求和，最后通过协作求和得到最终结果。

分段归约作为线性代数问题

用不同的方式表述同一个问题可以带来更多的可能性。分段归约问题可以被看作是一个矩阵-向量乘法问题。如下图所示，将输入数据视为一个矩阵，与一个全为1的向量相乘，即可得到每行的和。

介绍 nvmath-python：Python 设备数学库

nvmath-python 是一个涵盖所有 NVIDIA CPU 和 GPU 数学库（BLAS/LAPACK/FFTW/...）的 Python 库。
- 可通过 pip/conda 安装。
- Host API 支持多种内存和执行空间（CPU, GPU 等）。
- 设备 API 基于 MathDX，并以 Numba/Warp 为目标。
- 基于 Apache 2.0 协议开源。
- v0.3.0 新特性：支持 FP8/MAXFP8 的 matmul。
- 即将推出：分布式（MGMN）执行空间、稀疏线性代数、线性求解器设备 API。

代码示例展示了如何调用 nvmath.device 中的 matmul 编译器，在 CUDA 核函数内部直接执行矩阵乘法，从而实现分段归约。

介绍 cuTile：CUDA 在 SIMT 之外新增了 Tile 编程模型

cuTile 是一个新的编程模型，它将数据作为一个 tile（瓦片）进行处理，进一步抽象了底层的线程管理。
如下面的代码所示，使用 cuTile 的 segmented_reduce_kernel 变得极其简洁。开发者只需关注 load 一个 tile 的数据，对其进行 sum 操作，然后 store 结果即可，无需关心线程同步、共享内存等细节。

cuTile 编程模型 vs. CUDA SIMT 编程模型

下图详细对比了两种模型的差异：

特性	CUDA Tile 编程模型	CUDA SIMT 编程模型
执行单元	一个 tile 块（tile block）	一个线程（thread）
线程	隐式（implicit）	显式（explicit）
数据映射	系统自动将线程映射到数据	用户负责将线程映射到数据
协作 API	编程模型的内置概念	需要在 SIMT 结构之上手动实现
程序设计	程序需要指定 tile 大小	块大小决定了程序设计
编程风格	在 CUDA 核函数中编写类似 NumPy 的代码	-

编程模型的演进

通过一个简单的 square 函数，我们可以看到从 C++ 到 Numba SIMT 再到 Numba cuTile 的演进过程，抽象层次不断提高。

1. C++: 底层、冗长，需要手动管理线程和内存。
2. Numba (SIMT): Pythonic 语法，但仍需遵循 SIMT 模型，手动进行线程索引计算和循环。
3. Numba (cuTile): 更高层次的抽象，以数据 tile 为中心，代码更简洁，更接近于 NumPy 的表达方式。

Warp：具备 Tile 支持的自动微分 Python Kernel

Warp 是一个用于开发 GPU 加速和可微分核函数的 Python DSL，旨在扩展深度学习（DL）模拟流程。

核心特性:
* Warp 核函数是即时编译（JIT-compiled）的，并能自动合成反向模式（reverse-mode）的自动微分核函数。
* 内置了用于几何处理的数据结构和算法。
* 具备 PyTorch / JAX / CuPy 的互操作性。

通过新的 Tile 原语扩展 Warp:
* wp.tile() 函数由块（block）中的所有线程协同执行。
* 支持协同的 wp.tile_load() / wp.tile_store() 操作。
* 用户在启动时指定 block_dim，JIT 会为新的、未见过的块大小编译新版本的核函数。
* 基于 Tile 的矩阵乘法（matmul）、快速傅里叶变换（fft）和 Cholesky 分解由 MathDx 提供支持。
* Warp 现已采用 Apache 2.0 许可证。

GitHub: https://github.com/NVIDIA/warp

在 Python 中引入 CUTLASS 编程

现在可以直接在 Python 内部为 Tensor Core 编程，并获得与 C++ 相当的性能。

主要优势:
* 在 Python 中进行原生设备代码编写。
* 对 Tensor Core 的完全底层访问。
* 在功能和性能上与 C++ 版本保持一致。

相较于 C++ 的改进:
* 无需再编写复杂的 C++ 模板。
* 编译速度极快。
* 调试信息更清晰，自动调优（autotuning）更容易。
* 原型开发周期更快。
* 通过缓存实现 JIT 编译，减少了开销。

性能对比:
* 编译时间：对于 (8k, 8k, 8k) GEMM，CUTLASS C++ 编译耗时约 25842 毫秒，而 CUTLASS Python 仅需 180 毫秒，速度提升超过 100 倍。
* 峰值性能：Python 版本的性能与 C++ 版本持平，均达到理论峰值（SOL）。

相关 GTC 讲座:
* S74639: Python DSL: Enable Tensor Core Programming in Python With CuTe
* S72720: Blackwell and C++ API: Programming Blackwell Tensor Cores with CUTLASS

Python Kernel DSLs/Compilers 大观

本页展示了多种用于实现同一个目标——分段求和核函数（segmented sum kernel）——的 Python 内核 DSL 和编译器。殊途同归，条条大路通罗马。

如何选择最适合编写 Kernel 的工具？

简短的回答是：视情况而定。

在选择工具时，需要考虑以下一系列问题：
* 我当前的工作负载中已经使用了哪些包？
* 我的依赖项是否已经提供了编写 Python 内核的方式？
* 我的支持矩阵是怎样的？（Python 版本、CUDA 工具包版本、CUDA 驱动版本、操作系统、CPU/GPU 架构等）
* 我们计划如何构建/打包/发布/部署我们的 Python 项目？
* 我们期望的用户体验是怎样的？
* 我们需要的包是否在我公司/组织的白名单上？
* 我们团队的偏好/专业知识是什么？
* 维护和演进一个选定的解决方案有多容易？
* 它能否满足我们项目的性能要求？

最终的结论是：如果我们能避免编写任何 Kernel，那便是最佳选择！

编写 Kernel 的最佳方式是不编写任何 Kernel

利用现有生态系统

CUDA 生态系统已经足够成熟，解决方案可能早已存在于某个地方。下图展示了 CUDA 软件栈与 CUDA Python 生态系统中的各个库的对应关系，开发者可以利用这些现有的高级库来满足需求，而无需从头编写底层核函数。

引入 cuda.parallel

在许多情况下，您不需要编写任何核函数。cuda.parallel 库提供了 Pythonic 的接口来访问基础的 CUDA 主机并行算法。

主要特性:
* 利用了 NVIDIA 在 Thrust C++ 库上的投入，并将其暴露给 Python。
* 填补了 Python GPU 库开发者的空白。
* 其抽象层次有意低于 Python 数组/张量库。
* 性能：
* 在不离开 Python 的情况下，提供与 C++ 相同的性能和控制。
* 链接时优化（LTO）是性能的关键。

• 编译：

  • 即时编译（JIT）生成小体积的二进制文件！
  • LTO 在后台自动完成。

• 发布状态：

  • 即将推出：第一个 beta 版本。
  • 可通过 pip/conda 安装。
  • 在任何支持 CUDA 的地方都能工作。
  • 将提供更多算法（如 transform, segmented sort, reduce_by_key 等）。

GitHub: https://github.com/NVIDIA/cccL/tree/main/python/cuda_parallel

nvmath-python 再次前来救场

在许多情况下，您不需要编写任何核函数。例如，分段求和问题可以通过 nvmath-python 库中的矩阵乘法（matmul）巧妙地解决。通过将输入矩阵与一个全为 1 的列向量相乘，可以高效地计算出每行的和。

Python 中的数组编程

在许多情况下，您不需要编写任何核函数。对于之前的分段求和问题，其本质上是在一个 C-连续的二维（稠密）矩阵上进行按行归约。这可以通过高级的数组编程接口轻松实现。

• 使用 xp.sum(in_arr, axis=-1) 即可完成。
• 该方法符合 Python 数组 API 标准，这意味着它在许多 Python 数组库（如 CuPy, PyTorch, Jax, Dask 等）之间是可移植的。

CUB：C++ 的解决方案

在许多情况下，您不需要编写任何核函数，即使是在 C++ 中。对于分段求和问题，与其编写一个自定义的 CUDA C++ 核函数（如下方代码所示），不如直接使用 CUB 库中高度优化的原语。

下图中的自定义 segmented_reduction_kernel 调用可以被一行 cub::DeviceSegmentedReduce::Sum 调用所替代，从而获得更好的性能和可维护性。

总结

CUDA Python: 当性能遇上生产力

• 完全在 Python 中开发核函数，兼顾生产力和性能。

• 现代编译器工具链使代码编写更容易。

  • Python 设备库已成为现实。
  • 新的 CUDA Tile 编程模型（cuTile）将进一步提供帮助。

• 如果我们能避免编写任何核函数，情况会更好。

• 根据我们的需求、约束、性能目标和个人偏好，混合和匹配各种工具。
• 了解您的工具，随时向我们提问，我们随时提供帮助！
  • CWE72433: CUDA 开发者最佳实践 (今天下午)

附录：CUDA Kernel 的生命周期与编译器架构

以下是与 CUDA Kernel 相关的不同层次和库的列表：

• cuda.core
• numba-cuda
• nvmath.device & cuda.cooperative
• cuda.parallel
• cuTile
• Warp
• CUTLASS DSL

下图比较了传统的 C++ CUDA 编译流程和设想中的 Python 流程。C++ 流程清晰，但较为复杂；而如何从 Python 代码（.py）生成可执行文件，则充满了待解决的问题。

cuda.core 提供了对核心 CUDA 功能的 Pythonic 访问。它通过封装 NVRTC（NVIDIA 运行时编译库），使得在 Python 中动态编译和启动 CUDA C++ 内核成为可能。

下图展示了用户定义的 C++ 内核和用户定义的 numba-cuda Python 内核的编译流程对比。

Numba 核函数的生命周期（更新版）

为 Python 设备库提供动力

下图展示了一个 Numba 核函数的完整编译和链接流程：
1. 用户定义的 Python 核函数和设备函数首先被编译成字节码（bytecode）。
2. Numba 将字节码转换成 Numba IR。
3. Numba IR 被转换成 LLVM IR。
4. LLVM IR 最终转换成 NVVM IR。
5. NVVM IR 与 LTO-IR 设备库（LTO-IR device libraries）一起，通过 NVRTC（NVIDIA Runtime Compilation）和 nvJitLink 进行链接时优化（Link-Time Optimization, LTO）。
6. 最终生成可执行的 PTX 或 CUBIN 代码。
7. 通过 cuLaunchKernel 在 GPU 上启动核函数。

LTO IR 是实现高性能和链接外部库（如 C++ 设备库）的关键。

JIT-time linker：最后的缺失环节

即时（JIT）链接器是实现 C++ 和 Python 设备库混合编程的关键。它允许多个独立的编译单元在运行时链接成一个可执行的内核。

如下图所示，无论是用户定义的 C++ 内核、设备函数，还是 C++ 或 LTO-IR（链接时优化中间表示）设备库，都可以被编译成 LTO IR。然后，nvJitLink 工具将这些 LTO IR 模块链接在一起，生成最终的 .ptx 或 .cubin 文件，以供 cuLaunchKernel 调用。

该 JIT 链接功能通过 Python API 暴露给用户。开发者可以通过 Python 代码动态编译和链接 C++ 内核。右侧的流程图展示了其 Python API 的使用方式：多个 cuda.core.ObjectCode 对象可以被送入一个 cuda.core.Linker 实例，链接后生成一个新的 cuda.core.ObjectCode。这个对象随后被转换成一个 cuda.core.Kernel，最终通过 launch() 方法执行。

该机制为在 Python 中使用高性能的 NVIDIA C++ 设备库（如协作组和数学库 cuda.cooperative/nvmath.device）提供了支持。

如下图所示，用户定义的 Python 内核和设备函数通过 Numba 编译流程生成 NVVM IR。同时，NVIDIA C++ 设备库被预编译成 LTO IR 格式。nvJitLink 链接器在 JIT 阶段将来自 Python 的 NVVM IR 和来自 C++ 库的 LTO IR 结合起来，生成最终的可执行 .ptx/.cubin 文件。这使得 Python 内核可以直接调用高性能的 C++ 设备库函数。

cuda.parallel 模块进一步加强了 C++ 和 Python 之间的互操作性。它允许两种语言编写的设备代码相互调用和链接。编译流程图中的虚线箭头表示，从 Python 代码生成的 LLVM IR 可以与从 C++ 代码生成的 LTO IR 通过 nvJitLink 进行链接。这实现了在 C++ 和 Python 代码之间进行更深层次的集成，构建跨语言的复杂 GPU 应用。

cuTile 是一个新的 Pythonic CUDA DSL 和 JIT 编译器，它引入了一个新的中间表示——Tile IR。在 cuTile 的编译流程中，用户定义的 Python 内核被转换成 Tile IR，这是一种更高级、面向 Tile（数据块）计算的中间表示。

免责声明：该编译流程可能会发生变化。

“一个 Warp 内核的生命周期”

在 Warp 框架中，实现不同编译器之间的协同工作是核心。下图展示了其 GPU 后端的编译路径。

用户定义的 Python 代码可以同时通过两条路径进行处理：
1. AST 路径：Python 代码的抽象语法树（AST）可以直接与用户定义的 C++ 内核和设备函数进行交互，然后通过 NVRTC/nvJitLink 编译。
2. Bytecode 路径：Python 字节码通过标准的 Numba 流程（Numba IR -> LLVM IR -> NVVM IR）或 cuTile 流程（Numba IR -> Tile IR）进行编译。

这种双路径处理实现了 Python 和 C++ 之间前所未有的紧密集成，让编译器能够协同优化代码。

注：此图仅涉及 Warp 的 GPU 后端；其 CPU 后端采用不同的编译路径。

CUTLASS Python 架构

CUTLASS Python 接口采用基于 MLIR（多层中间表示）的编译器栈，提供了从高层 Python 应用到底层硬件指令的完整编译路径。

其架构分为以下几个层次：
1. Python-Native Application：用户编写的顶层 Python 应用。
2. CuTe DSL：应用通过 CuTe DSL 来表达计算。
3. MLIR Compiler Stack：CuTe DSL 被 lowering 到 MLIR 栈中的 CuTe IR。这一步生成包含 cute, nvgpu, scf, func, nvvm 等多种方言（dialect）的 IR 程序。
4. CUDA Compiler Stack：MLIR 的输出进一步被 CUDA 编译器栈处理，依次通过 NVVM/LLVM、PTX 和 SASS（着色器汇编）层级，生成最终的机器码。
5. JIT Executor & CUDA Driver：最终生成的 CUBIN 由 JIT 执行器加载，并通过 CUDA 驱动程序在 GPU 上运行。

致谢

感谢您的关注！

报告人：Leo Fang，CUDA Python 技术主管

欢迎通过以下链接分享您的反馈和问题：
- https://github.com/NVIDIA/accelerated-computing-hub
- https://github.com/NVIDIA/cuda-python