Run the GPU kernels as AoT CUDA written for their layouts, and ship them as blochsim[cu12]/[cu13]

[mcencini]

The GPU kernels run ahead-of-time CUDA, at or under Triton's time on every case measured, and ship to PyPI as blochsim[cu12] / blochsim[cu13]. This supersedes #31 (AoT CUDA in place of Triton) and #32 (tiled PERK), whose commits it contains.

What changed

Kernels written for their layouts (_layout*.hpp). A layout is what is compiled: pools, how a pulse is formed, per-voxel maps, problems per thread, one train or several. Every other switch is read at run time and steers whole blocks once per event, so one compile serves every combination of them. The loops are written once over a number type: at float they are the forward simulation and the adjoint, at num::Dual the JVP and the second-order adjoint.

• EPG kernels. Complex and real forward, JVP, VJP and VJP-JVP, for 0–3 pools.
• Many-pool kernels. Forward and adjoint for 2–8 pools.
• Checkpointing. The adjoints keep every fourth state and replay each stretch, so they need no recording launch and no full trajectory.
• Tile kernels remain for what a layout does not take (rows wider than a warp) and for the host build that the suite checks them on.

Other changes.

• The per-combination specialisations from Compile the GPU kernels ahead of time as CUDA and drop Triton #31 are removed; the layouts cover them.
• PERK splits its features across a second grid axis when a launch is too small to fill the card.
• The package's own structural derivatives (packing binding, transition tables) are taken in reverse mode. A simulation and its gradients therefore no longer make PyTorch register its forward-mode decompositions through TorchScript.

Packaging

The base wheel is CPU-only. The card's module is a package per CUDA major version (src/cuda/12, src/cuda/13), the same CMake project with BLOCHSIM_CUDA_PACKAGE set. Each one:

• links libcudart from torch's nvidia wheel through an rpath, rather than carrying a runtime;
• carries machine code for 7.5, 8.0 and 9.0 plus PTX for 9.0, in compressed fatbinaries;
• pins the blochsim it was built with.

_gpu_launch loads the build for torch's CUDA major version. It refuses one for another major or release with CudaBuildMismatch, which names the extra to install.

pip install torch --index-url https://download.pytorch.org/whl/cu126
pip install 'blochsim[cu12]'        # or blochsim[cu13] beside a CUDA 13 torch

Wheel (75;80;90 + PTX 90), as CI builds it	Size
blochsim-cuda12, CUDA 12.6	61.7 MB
blochsim-cuda13, CUDA 13.0	48.8 MB

wheels.yml builds both in a manylinux_2_28 container with the oldest toolkit of each major that torch ships. It installs each beside torch cu126/cu130 on Python 3.10 and 3.14 and checks that libcudart resolves into torch's nvidia wheel. It publishes each project from its own environment (pypi, pypi-cuda12, pypi-cuda13).

Speed against Triton

Device time per case, the median of three interleaved runs on an RTX 4060 Laptop (8.9), sm_89 build.

Case	n	Triton µs	CUDA µs	Ratio
epg_real	2000	863	932	1.08*
epg_real	100000	13598	9427	0.69
epg_complex	2000	1231	1192	0.97
epg_complex	100000	24645	20083	0.81
epg_real_grad	2000	4105	2598	0.63
epg_real_grad	100000	115635	69236	0.60
epg_complex_grad	20000	61689	42196	0.68
epg_real_hvp	2000	13975	8005	0.57
epg_real_hvp	20000	94107	52168	0.55
epg_complex_hvp	5000	60277	46209	0.77
pooled (3 pools)	2000	2004	1352	0.67
pooled (3 pools)	50000	21183	8279	0.39
pooled_grad	20000	75813	44555	0.59
pooled_hvp	5000	150924	66449	0.44
pooledb_hvp (2 pools)	20000	356946	98822	0.28
pooled5	500	21786	11056	0.51
pooled5	5000	134057	75933	0.57
pooled5_grad	2000	171391	68360	0.40
perk	2000	245	93	0.38
perk	200000	15232	7724	0.51
perk_grad	2000	361	278	0.77
perk_grad	200000	21451	14882	0.69

* The kernel itself is 0.71× (99 µs against 140). The total at this size is dominated by torch's own host-to-device copies, which take 78 µs or 114 µs per process on either build and in either order.

Per kernel, against Triton's on the same captured launches:

Kernels	Ratio
Complex forward	0.16–0.98
Real forward	0.53–0.85
JVPs	0.5–1.0
Real VJP	0.52–0.86
Complex VJP	0.05–0.83
VJP-JVP	0.11–0.78
Many-pool forward	0.04–0.22
Many-pool adjoint	0.13–0.30
Three-pool table kernels	0.10–0.16

The first call takes 0.01–1.1 s, against 0.6–2.5 s for Triton.

Memory against Triton

• Peak torch allocation: identical in every case except pooled5_grad, which drops from 265 to 81 MiB. The many-pool adjoint keeps checkpoints, not Triton's recording.
• Device memory beyond the allocator: equal or lower in every case but one. The exception is pooled_hvp (three pools, second order), which leaves 84 MiB against Triton's 52.
  • That is the driver's local-memory reservation for the three-pool dual adjoint's 3.6 KB stack. Down from 7.0 KB, since this PR contracts wide intervals one direction at a time, which is also faster.
  • Getting under Triton's 2.5 KB needs that adjoint's contraction redesigned. It is left as a follow-up.

Verified

• Full suite: 1614 passed, 3 skipped on the card.
• cu13 wheel: installed beside torch 2.13+cu130, the full suite passes with CUDA_DISABLE_PTX_JIT=1, so an 8.9 card runs the wheel's sm_80 code. scripts/check_wheel.py loads the module bare and finds libcudart.so.13 in torch's nvidia wheel through the rpath.
• cu12 wheel: loads bare against a CUDA 12.6 runtime placed where its nvidia wheel would put it. Not run against a CUDA 12 torch locally; the CI job does that.
• Layout against tile paths: they agree to float32 round-off through autograd: values, first and second derivatives, CPU kernels as reference.
• Not run: on a T4, A40 or H100.

Publishing

blochsim-cuda12 and blochsim-cuda13 have trusted publishers for wheels.yml with environments pypi-cuda12 and pypi-cuda13, so a v*.*.* tag publishes all three projects.

benchmarks/README.md still reports measurements taken on the Triton build and should be re-measured.

🤖 Generated with Claude Code