218个真实PTX文件字节级还原,4微秒启动开销——一个开发者单枪匹马搞出了绕过Triton和CUDA C++的第三条路。

为什么需要第四种写法

打开网易新闻 查看精彩图片

写GPU内核的选项已经够多了。Triton把硬件细节藏进编译器,CUDA C++让你用宏和内联汇编硬刚,Pallas走JAX路线。但总有人卡在中间:Triton的优化器做了太多假设,CUDA C++的编译流程太重,而纯PTX又像是回到汇编时代。

Patrick Toulme的Pyptx瞄准了这个缝隙。它让你用Python语法写PTX指令——每个函数调用严格对应一条PTX指令,中间没有优化器、没有自动调参、没有张量中间表示(IR)。

这不是又一个Python绑定。Pyptx是一个完整的领域特定语言(DSL),包含真实的PTX解析器、发射器和转译器。测试集里218个来自CUTLASS、Triton、fast.cu、DeepGEMM、ThunderKittens和LLVM测试套件的PTX文件,它能做到字节级还原。

核心设计很直白:保留Python的语法糖,去掉所有可能改变语义的中间层。你想要控制寄存器分配?直接操作。想精确调度共享内存?自己写。这种"零抽象"的定位,让它和现有工具形成鲜明对照。

Hopper和Blackwell的新指令怎么接

Pyptx对英伟达两代新架构的支持不是敷衍的。Hopper(sm_90a)这边,WGMMA(线程组矩阵乘累加)、TMA 2D/3D多播、mbarrier、集群启动全部暴露为Python API。Blackwell(sm_100a)更激进:tcgen05.mma、tcgen05.ld、TMEM(张量内存)、SMEM描述符、线程束特化(warp specialization)都进来了。

这些不是简单的封装。以WGMMA为例,Pyptx让你直接构造wgmma_tile,指定数据类型、维度和主存储方向。TMA的多播模式、mbarrier的等待/到达语义,都是显式调用而非编译器推断。

Blackwell的tcgen05指令尤其值得关注。这是英伟达新一代张量核心接口,Pyptx选择直接暴露而非等待上游编译器支持。对于需要抢发论文或产品上线的团队,这种"指令级抢先"可能是关键优势。

寄存器、谓词、屏障、共享内存的管理全部开放。没有隐藏的寄存器分配器替你决定,也没有隐式的同步点。好处是预测性极强,代价是写错就崩。

三种调用路径的实测开销

内核写完了怎么跑?Pyptx给了三条路,延迟差异很大。

CUDA图(CUDA graph)回放最快,约4微秒。这是把内核预录进执行图后的开销,适合推理阶段的固定计算图。缓存的C++扩展路径约14微秒,torch.compile路径14-22微秒。

这个对比很有意思。torch.compile的"编译一次到处运行"承诺,在这里变成了实实在在的延迟惩罚。如果你的场景是交互式训练或动态形状,4微秒和22微秒的差距会累积成显著瓶颈。

技术实现上,PTX通过cuModuleLoadData即时编译,再用约150行C++的启动垫片(launch shim)注册。这个垫片足够薄,没有额外的抽象层摊销。

Pyptx还埋了一个逆向工具:PTX转Python。运行python -m pyptx.codegen kernel.ptx --sugar --name my_kernel > my_kernel.py,可以把nvcc、Triton或Pallas生成的PTX还原成Python代码。--sugar标志会做符号还原、把自旋循环提升为ptx.loop(...)、折叠mbarrier等待块、分组表达式链。

这对调试和学习很有价值。你可以拿Triton生成的PTX,看它被"去糖"后长什么样,再决定哪些部分值得用手写Pyptx替换。

和Triton、CUDA C++的三角关系

Pyptx的定位需要仔细理解。Triton用编译器抽象硬件细节,你写tile-level算法,它帮你决定寄存器、调度、内存布局。Pyptx相反:你写的就是最终运行的指令,没有编译器做假设。

CUDA C++加内联PTX是另一条路,但开发体验割裂。你得在.cu文件里写宏,用nvcc编译,再和Python框架对接。Pyptx把整个流程留在Python里,和JAX、PyTorch eager、torch.compile的集成是原生的。

这不是说Pyptx要取代谁。它的适用场景很明确:当你发现Triton的优化器做了错误假设,或者需要用到编译器尚未支持的新指令,又或者想对关键内核做最后一轮指令级调优。

单开发者项目的事实也限制了它的当前适用范围。没有企业级支持,文档在pyptx.dev但明显还在建设中,社区生态为零。这决定了它现阶段更适合有PTX阅读能力的进阶用户,而非广谱采用。

一个最小GEMM内核长什么样

看代码比描述更直接。Pyptx的GEMM示例展示了它的语法风格:

装饰器@kernel定义了输入输出规格、网格划分、线程块配置和目标架构。Tile抽象处理张量布局,smem.wgmma_tile构造共享内存中的WGMMA兼容分块,reg.array显式分配寄存器数组。

函数体里的ptx.wgmma.mma_async(...)就是直接的PTX指令调用。没有隐式的同步,没有自动的流水线展开。每个操作的成本和时序都是可见的。

这种写法的心智负担明显高于Triton的@triton.jit装饰器。但作为交换,你获得了对Hopper张量核心行为的完全控制——包括那些Triton编译器可能尚未优化或错误优化的边缘情况。

谁会真的用它

Pyptx的受众画像很清晰:已经在用Triton但遇到编译器瓶颈的人,需要抢先支持Blackwell新指令的团队,以及想把关键内核从CUDA C++迁移到Python生态的开发者。

对于大多数人,Triton仍是更务实的选择。它的自动调参和优化器覆盖了90%的场景,社区支持和文档成熟度也远胜。Pyptx的价值在于那10%——当你需要打开 hood 检查时,它提供了一种不离开Python的方式。

PTX转Python的功能可能被低估。它降低了阅读和学习现有高性能内核的门槛,你可以把CUTLASS或DeepGEMM的输出当成教材,逐行理解工业级优化是怎么做的。

单开发者项目的风险也需要正视。如果Patrick Toulme的精力转移,或者英伟达PTX格式发生重大变更,维护连续性没有保障。这在生产环境选型时是硬约束。

这件事的真正意义

Pyptx的出现印证了一个趋势:Python正在吞噬GPU编程的每一层抽象。从Numba到Triton到Pyptx,开发者不断用Python语法包裹底层硬件,同时保留对关键细节的控制权。

它的独特之处在于"零中间层"的承诺。不是用Python生成CUDA C++再编译,不是用Python配置Triton的编译器选项,而是Python语法直接映射到PTX指令。这种激进简化对特定人群有不可替代的价值。

对于在Hopper或Blackwell上榨取最后10%性能的团队,Pyptx提供了一条绕过编译器假设的捷径。对于想理解GPU指令级行为的学习者,它是比反汇编更友好的入口。对于Python优先的ML基础设施团队,它填补了Triton和原始PTX之间的工具空白。

判断标准很简单:如果你最近三个月内因为Triton的优化器行为而重写内核,或者正在等待某个CUDA版本支持Blackwell新指令——Pyptx值得放进评估清单。否则,保持关注但暂不迁移。