Compile the GPU kernels ahead of time as CUDA and drop Triton

[mcencini]

Requested by Matteo · project thread

What this changes

Before: the GPU kernels were Triton and were JIT-compiled on the first call. The suite checked them through Triton's CPU interpreter, behind the interpreted marker.

After: the EPG, pooled and PERK kernels are C++, written over a small tile runtime (_tile.hpp). nvcc compiles them ahead of time into blochsim._gpu, which links the CUDA runtime statically. Nothing compiles at the first call, and there is no Triton fallback. The public API is unchanged and the dispatch is the same, so consumers need no update.

How:

• The kernels were translated mechanically from the Triton source into _epg_kernels.hpp, _pools_kernels.hpp and _perk_kernels.hpp. _kernels.hpp is the launch table (parameters, block axes).
• _gpu_launch.Kernel(name)[grid](...) keeps Triton's call shape. CUDA tensors go to _gpu on torch's current stream; CPU tensors go to _gpu_host, the same source compiled for the host, one program at a time.
• CMake builds _gpu wherever it finds nvcc; BLOCHSIM_CUDA=ON/OFF overrides that.
  • Each kernel is compiled in its own file, generated from _gpu_kernel.cu.in, so the build runs in parallel.
  • Each kernel is compiled twice, with __launch_bounds__(256) and (1024). The launcher picks the 256 variant whenever the block fits.
  • The default architectures are 75-real;80-real;86-real;89-real;90, where 90 carries PTX for newer cards.
• cibuildwheel builds the x86-64 manylinux wheel with CUDA 12.8 (before-all installs cuda-nvcc and cuda-cudart-devel). Wheels leave out _gpu_host.
• Removed: _epg_triton.py, _pools_triton.py, _perk_triton.py, the interpreted marker, the Triton CI step and prune_triton_cache.py. The docs, skills and CLAUDE.md are updated to match.

Known limit: a tile is one element per thread, so a launch is at most 1024 threads. An EPG run with more than 1024 state orders on a card, or a pooled run with orders × pools above 1024, is refused with a clear error.

How it was checked

There is no GPU in the container, so nothing here has run on a card.

• Host build vs the C++ kernels: the host build of the GPU kernels matches the C++ kernels on all 13 former interpreter cases. These cover narrow, wide, chunked, streamed, washed, shimmed, profiled, one/two pools, real, real shimmed, spoiled and narrowed (tests/sequence/test_host_kernels.py).
• Pooled kernels: they match on all 7 cases, to at most 2.5e-6 (test_many_pools_host.py).
• PERK: forward and adjoint match Torch autograd (new test in test_perk_kernel.py).
• Full suite: pytest tests/ -n 6 gave 1269 passed, 315 skipped, 13 failed and 12 errors. Every failure and error comes from this container's broken torchvision install (operator torchvision::nms does not exist, the deepinv import). The one exception is test_an_image_quality_design_fits_in_its_budget, a timing budget that passes when run alone.
• nvcc 12.x for sm_80: every kernel compiles. The heaviest, _epg_vjp_jvp_kernel, takes about 4.5 min per architecture and uses 255 registers with some spills at the 256 bound. At the 1024 bound it spills heavily, and that variant only runs for blocks wider than 256.
• pre-commit run --all-files is clean.

Needs a card (Matteo):

• Build: pip install -e . --config-settings=cmake.define.BLOCHSIM_CUDA=ON, optionally adding --config-settings=cmake.define.CMAKE_CUDA_ARCHITECTURES=native.
• Then run pytest tests/ -n auto: test_cuda_parity.py, test_both_pools.py, test_perk_kernel.py, test_dynamic_transmit.py and test_subspace_streams.py exercise the device path.

Checklist

• pytest tests/ passes locally (apart from the environment failures above), and new tests cover the host build.
• pre-commit run --all-files is clean.
• Public functions and classes carry a numpydoc docstring.
• Comments and docstrings describe the code as it is now.
• A change to what the kernels compute is in both of them. Here the physics is unchanged.
• The documentation is updated where visible to a caller; the docs build was not run here.

🤖 Generated with Claude Code

https://claude.ai/code/session_014ND2A7uRuhWiay6B4rWF1H

---

Generated by Claude Code

[mcencini]

Superseded by #33, merged as 83a8e16: it contains this PR's commits (through 8f508bf). On top of them the EPG and many-pool kernels are written for their layouts, at or under Triton's device time on every case measured, and the card's module ships as blochsim-cuda12 / blochsim-cuda13.