pyptx

pyptx 是一个用于在 Python 中编写 NVIDIA PTX 内核的领域特定语言，支持 Hopper 和 Blackwell 架构，并可在 JAX 和 PyTorch 中调用。项目的核心功能包括对显式寄存器、谓词、屏障及共享内存的支持，以及针对 Ampere、Hopper 和 Blackwell 系列 GPU 的特定指令集特性。它允许开发者直接从 Python 函数生成 PTX 代码，无需中间优化或自动调优步骤。适用于需要高性能计算且能够充分利用 NVIDIA 最新 GPU 特性的场景，如深度学习模型训练、大规模矩阵运算等。通过简洁的 API 设计和高效的执行路径，pyptx 在多个基准测试中展示了接近甚至超过 cuBLAS 的性能表现。