pyptx

pyptx 是一个用于在 Python 中直接编写 NVIDIA PTX 汇编代码的领域专用语言（DSL），支持 Ampere、Ada、Hopper 和 Blackwell 架构（sm_80/sm_89/sm_90a/sm_100a/sm_120）。它提供细粒度控制：显式寄存器、谓词、屏障、共享内存及硬件特性（如 WGMMA、TMA、tcgen05、mbarrier、wgmma）；生成的 PTX 与手写汇编严格等价，无中间优化层。支持从 JAX（via jax.jit）、PyTorch eager 及 torch.compile 直接调用，并自动适配当前 GPU 架构。适用于需要极致性能调控的场景，如自定义算子开发、GPU 内核原型验证、AI 系统底层加速（如 MoE、RMSNorm、高吞吐 GEMM）等。