“CUDA Python”是一套生态，不是单一语言

Python 用户通常处在四个层次之一：

业务表达
├─ PyTorch / JAX / RAPIDS / CuPy：张量、数组、数据处理
├─ cuBLAS / cuFFT / cuSPARSE / cuSOLVER 的 Python 封装：成熟算子
├─ cuda.tile / cuda.lang / Numba-CUDA：用 Python 写 GPU kernel
└─ cuda.core / cuda.bindings：编译、内存、stream、launch、Driver/Runtime API

选择原则和 C++ 一致：能用成熟库就先用库；需要融合、自定义数据布局或特殊算法时再写 kernel；需要控制上下文、模块加载、JIT 和资源生命周期时才下沉到 core/bindings。

组件选型表

组件发展很快。安装前应查各项目的 support policy，固定兼容的 Python、Driver、Toolkit 和包版本，而不是只写无上限的 pip install。

三种最常见工作流

1. 只组合数组与库

import cupy as cp

x = cp.arange(1_000_000, dtype=cp.float32)
y = cp.sin(x) * 0.5 + 2.0
result = cp.sum(y)

# 只有真正需要 host 标量时才同步和传回。
print(float(result.get()))

这条路线开发最快，但每个数组表达式都可能形成独立 kernel 和临时数组。先用 profiler 检查 launch 数量与分配开销，再决定是否 fuse。

2. 用 Numba 写 SIMT kernel

from numba import cuda

@cuda.jit
def saxpy(a, x, y):
    i = cuda.grid(1)
    stride = cuda.gridsize(1)
    for j in range(i, x.size, stride):
        y[j] = a * x[j] + y[j]

threads = 256
blocks = min((x.size + threads - 1) // threads, 4096)
saxpy[blocks, threads](2.0, x, y)

Numba kernel 不返回普通 Python 值，结果写回数组；支持的是受限 Python/NumPy 子集。当前 NVIDIA 文档已将 Numba-CUDA 标记为维护模式，新特性目标转向 Numba-CUDA-MLIR，因此新项目要先核对路线图。

3. 用 cuda.core 编译并启动 CUDA C++

下面展示对象关系，具体导入名以所安装版本的 API 文档为准：

import cupy as cp
from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch

source = r'''
extern "C" __global__
void add(const float* a, const float* b, float* c, unsigned long long n) {
    auto i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}
'''

dev = Device()
dev.set_current()
stream = dev.create_stream()
opts = ProgramOptions(std="c++17", arch=f"sm_{dev.arch}")
module = Program(source, code_type="c++", options=opts).compile("cubin")
kernel = module.get_kernel("add")

n = 1 << 20
a = cp.ones(n, dtype=cp.float32)
b = cp.ones(n, dtype=cp.float32)
c = cp.empty_like(a)
cfg = LaunchConfig(grid=(n + 255) // 256, block=256)
launch(stream, cfg, kernel, a.data.ptr, b.data.ptr, c.data.ptr, cp.uint64(n))
stream.sync()

cuda.core 管理 CUDA 核心对象，但数据可继续由 CuPy、PyTorch 或自定义 memory resource 提供。

互操作真正难的是三件事

共享 GPU 数据时，需要同时传递：

存储：地址 + 大小 + dtype + shape + stride + device
顺序：哪个 stream 产生，哪个 stream 消费，依赖如何建立
生命周期：底层 allocation 在最后一次异步使用完成前必须存活

只有一个裸指针远远不够。最危险的 bug 往往不是立刻越界，而是生产者还没写完、消费者就在另一个 stream 读取，或 Python 对象被垃圾回收后 GPU 仍在使用它。

数据协议

CUDA Array Interface

许多 Python 数组库用 __cuda_array_interface__ 描述 device array，常见字段包括 shape、strides、typestr、data pointer、version，有些版本还带 stream 信息。它适合 Python 对象间零拷贝视图，但接收方必须遵循协议版本与同步要求。

meta = array.__cuda_array_interface__
print(meta["shape"], meta["typestr"], meta["data"][0])

不要手工伪造元数据来绕过所有权。让提供者对象持续存活，并用目标库公开的转换函数建立视图。

DLPack

DLPack 用标准 capsule 交换张量元数据与所有权，适合框架之间传递 dense tensor。现代接口允许消费者表达 stream，从而由生产者建立必要依赖。

# API 名称会因框架版本略有区别；优先使用其公开 from_dlpack。
y = target_framework.from_dlpack(source_tensor)

DLPack 通常是一次性所有权转移协议，不应重复消费同一个旧式 capsule。零拷贝只表示不复制存储，不表示无需同步或没有布局转换。

Array API 与框架桥接

对只需要逐元素/归约/线性代数的库，优先接受符合 Array API 的对象，而不是绑定某个实现。对 CUDA 特有 stream、event、pointer 行为，则显式定义扩展接口。

Stream 互操作

假设库 A 在 stream A 生成数据，库 B 要在 stream B 消费：

stream A: [produce]──record(event)────────────────
                                │
stream B: ─────────────────wait(event)──[consume]

全设备同步虽然简单，却破坏重叠并把隐藏错误变成性能瓶颈。正确做法通常是：

1. 获取或包装生产者的 stream；
2. 在生产者 stream 记录 event；
3. 让消费者 stream wait event；
4. 保持数据和 stream 所属 context 一致；
5. 在最后一次使用完成后再释放对象。

cuda.core.Stream.from_handle() 等“借用句柄”接口不会接管外部 stream 生命周期；原拥有者必须活得更久。跨库时还要区分 legacy default stream 与 per-thread default stream，它们的隐式同步语义不同。

Python 生命周期与异步错误

Python 代码返回并不代表 GPU 工作完成：

Python 发起 launch → 很快返回 → 后续某个 sync/copy 才观察到异步错误

因此：

• 调试时在关键阶段显式同步，以定位第一个失败操作；
• 生产路径用 event/stream wait 表达依赖，不要到处 device synchronize；
• 对 buffer、module、kernel、stream 建立明确所有权；
• 不依赖解释器退出顺序完成 CUDA 清理；
• 捕获异常时仍需考虑 context 中是否留下 sticky error。

Context 陷阱

Driver API context 是 device 资源的命名空间。多个框架各自创建 context 时，即使指向同一块 GPU，pointer/module/stream 也未必能跨 context 使用。

进程
├─ primary context（多数 Runtime API/主流框架共享）
└─ 自建 Driver context
   └─ 这里创建的 stream/module 不自动属于 primary context

通用 Python 组件应尽量与调用方的 current/primary context 协作。若必须 push/pop context，应在异常路径也恢复原状态，并避免在多个 OS 线程间无规则切换。

自定义扩展的四条路线

若 kernel 参与自动微分，需要定义 backward 或组合可微原语；“forward 数值正确”不等于整个训练图正确。

包装 C++/CUDA 扩展的接口契约

一个可靠的扩展入口至少检查：

• device 与当前 context；
• dtype、rank、shape、stride、contiguity；
• 对齐与空张量；
• 当前 stream，而不是擅自使用默认 stream；
• 输出 allocation 的 allocator 与生命周期；
• 错误通过目标语言的异常机制返回。

Python tensor
   │ validate metadata, capture current stream
   ▼
C++ binding
   │ dispatch dtype/layout/architecture
   ▼
CUDA kernel or CUDA-X library
   │ asynchronous on caller stream
   ▼
return tensor immediately; caller owns later synchronization

性能诊断

Python GPU 程序慢，不一定是 Python 本身：

可复现实验模板

# 伪代码：预热、event 计时、校验与同步位置必须明确。
for _ in range(warmup):
    fn(x)

start.record(stream)
for _ in range(repeat):
    y = fn(x)
stop.record(stream)
stop.synchronize()

check_against_reference(y)

报告中写明包版本、GPU、Driver、首次编译是否计入、是否包含数据传输，以及采用哪个 stream。只用 time.perf_counter() 包围异步 launch 会严重低估 GPU 时间。

练习

1. 分别用 CuPy 表达式、RawKernel 和 Numba 实现 SAXPY，比较冷启动、稳态时间和 launch 数量。
2. 在两个非默认 stream 间用 event 传递一块 CuPy 数据，移除 event 后构造并观察竞态。
3. 将 PyTorch tensor 零拷贝交给自定义 kernel，验证非连续 tensor、空 tensor 和不同 dtype 的错误路径。
4. 用 cuda.core 运行时编译两个架构目标，设计磁盘缓存 key：源码、选项、Toolkit/NVRTC、目标架构都必须进入 key。

官方参考

• CUDA Programming Guide: Intro to CUDA Python
• CUDA Python / cuda.core
• CUDA Python / cuda.bindings
• Numba-CUDA User Guide
• CuPy Documentation
• DLPack Specification