Pyptx：用Python手写英伟达PTX，绕过所有编译器

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值得放进评估清单。否则，保持关注但暂不迁移。