pyptx ra mắt: viết PTX kernel cho Hopper & Blackwell ngay trong Python

TL;DR

pyptx là một DSL Python mới cho phép viết PTX kernel cho NVIDIA Hopper (sm_90a) và Blackwell (sm_100a) theo nguyên tắc một PTX instruction = một Python call. Cài bằng pip install pyptx[torch] hoặc pip install pyptx[jax]. Đi kèm GEMM (77–78% cuBLAS), grouped GEMM (~10× PyTorch), RMSNorm (88% HBM3 peak), SwiGLU, Flash Attention, và một transpiler PTX→Python round-trip byte-identical. Đây là thư viện Python duy nhất hiện nói được tcgen05 trực tiếp.

What's new

Patrick C Toulme vừa công bố pyptx trên X kèm trang chủ pyptx.dev. Triết lý cốt lõi: hàm Python chính là PTX stream. Mỗi DSL call emit đúng một PTX instruction — không optimizer, không autotuner, không codegen ngầm. Bạn viết ptx.wgmma.mma_async(...), output là đúng một câu lệnh WGMMA, không hơn không kém.

Toàn bộ launch shim chỉ gói gọn ~150 dòng C++, phần còn lại (DSL, tracer, parser, emitter, transpiler) là pure Python. Base package zero dependency. Kernel có thể gọi từ torch.compile, PyTorch eager, hay jax.jit qua một typed FFI — launch overhead đo được ~14 µs eager, ~4 µs khi replay CUDA graph.

Why it matters

Triton và Pallas là hai DSL phổ biến nhất cho kernel writing trong Python, nhưng cả hai không expose các instruction quyết định performance trên hardware mới: WGMMA, TMA với multicast, mbarriers, cluster launch, stmatrix — và đặc biệt là toàn bộ họ tcgen05 trên Blackwell. Muốn chạm vào những instruction này, lựa chọn duy nhất trước nay là CUTLASS C++ templates hoặc inline PTX trong CUDA. Cả hai đều kéo bạn ra khỏi vòng lặp Python.

pyptx là cây cầu đầu tiên: bạn viết PTX thật, theo từng câu lệnh, nhưng vẫn ở trong môi trường Python — debug được, REPL được, gọi thẳng từ PyTorch / JAX được. Với Blackwell vừa lên kệ và wgmma đã bị deprecated thay bằng tcgen05.mma, đây là thời điểm cần một công cụ Python-native nói được ngôn ngữ mới của tensor core.

Technical facts

Instruction coverage:

• Hopper (sm_90a): WGMMA, TMA 2D/3D + multicast, mbarriers, cluster launch, stmatrix.
• Blackwell (sm_100a): tcgen05.mma, TMEM alloc/load/store, SMEM & instruction descriptors, warp specialization, 2SM MMA path.

Comparison

pyptx hiện là thư viện Python duy nhất nói được tcgen05 trực tiếp, và là công cụ duy nhất có transpiler hai chiều: nuốt PTX biên dịch từ nvcc, Triton hay Pallas, đẻ ra Python pyptx tương đương, round-trip byte-identical. Có chế độ --sugar để demangle tên và đơn giản hoá control flow cho dễ đọc.

Transpiler đảo chiều này mở ra một workflow lạ: lấy kernel Triton đã tune sẵn, chuyển về pyptx, sửa tay vài instruction tới hạn (thường là phần WGMMA / TMA pipeline), rồi chạy lại. Bạn không phải viết toàn bộ kernel từ đầu chỉ để chạm vào một micro-optimization. Đây là chiến thuật cũ trong thế giới C/asm (compile rồi sửa output), nhưng lần đầu khả thi ở mức Python cho GPU kernel.

Use cases

• Kernel researcher săn 20% performance còn lại so với cuBLAS trên Blackwell — chỗ Triton / Pallas chạm trần.
• Inference engineer viết fused MoE / attention / norm cho LLM serving production. Grouped GEMM nhanh ~10× PyTorch trực tiếp giúp các model MoE-heavy kiểu Mixtral / DeepSeek.
• Tooling author muốn xuất phát từ PTX có sẵn (nvcc, Triton, Pallas output) rồi sửa interactively trong Python qua transpiler.
• GPU education — per-instruction Python dễ đọc hơn raw .ptx hay CUTLASS template.
• Hardware bring-up & microbenchmark — emit instruction predictable, không có optimizer ngầm xen vào.

Limitations & pricing

Miễn phí, mã nguồn mở. Nhưng đây không phải Triton thay thế cho người dùng phổ thông:

• Không có autotuner, không có tensor compiler — bạn phải tự specialize kernel theo từng shape. Đó là triết lý, không phải feature thiếu.
• Chỉ hỗ trợ Hopper (sm_90a) và Blackwell (sm_100a). Ampere, Ada không phải target.
• Cần kiến thức PTX nghiêm túc — warpgroup semantics, TMEM, mbarriers, descriptor layout.
• Debug ở mức PTX. Tooling còn beta.

What's next

Bài toán lớn của hệ sinh thái Blackwell trong 2026 là làm sao cho engineer Python tận dụng được tcgen05 mà không phải nhảy sang C++. pyptx là phát đạn mở màn cho hướng đi này. Trong vài tháng tới, kỳ vọng thấy nhiều reference kernel Blackwell hơn (attention variant, mixed-precision GEMM), transpiler phủ rộng hơn output của Triton/Pallas, và contribution cộng đồng khi nguồn cung Blackwell mở rộng.

Cài thử trong một dòng:

pip install pyptx[torch]
# hoặc
pip install pyptx[jax]

Nguồn: pyptx.dev, @PatrickToulme trên X, SemiAnalysis: Dissecting Blackwell Tensor Cores.