jax4j / NablaTensor
← quantChapter 4 · Picking an engine4 min read

Forget, then remember: checkpointing on a GPU

See how GPU checkpointing saves memory for Monte Carlo Greeks by keeping selected checkpoints and recomputing values during the adjoint sweep.

The 252-step Asian tape from Chapter 3 costs a CPU only 1.7× more to get all five Greeks instead of just the price. On a GPU, the same tape costs 8×. Same tape, same arithmetic, same reverse sweep. Where did the other 6× go?

The whole story

The plain kernel needs to hold about 760 forward values per path until the reverse sweep reads them — too many for a GPU lane's registers, so it spills to memory and throughput collapses. The checkpointed kernel keeps only a handful of bookmark values per cut, forgets everything else, and recomputes each segment from its bookmark when the reverse sweep needs it. On the 252-step Asian tape this makes the Greeks pass 2.9 times faster, or 6.8 times faster once four markets share one dispatch, and returns the same seventeen digits.

Did you know?

Recomputing a segment only works because NablaTensor's GPU kernels draw randoms from Philox, a counter-based generator: draw k of path p is a pure function of (p, k), not the next value pulled from some mutable stream state. Re-running a segment's forward arithmetic a second time asks for the same draws by index and gets bit-identical numbers back, with nothing saved anywhere for the random stream itself. A stateful generator would need that state checkpointed too, at every cut — and the memory saving would be a fraction of what it is.

Why a CPU doesn't notice and a GPU can't ignore it

The reverse sweep needs some of the forward sweep's intermediate values back — an exp node needs its own output; a multiplication needs both its inputs. Add that up over 252 time steps and it comes to roughly 760 forward values a plain kernel has to keep alive from the moment they're computed until the reverse sweep, hundreds of nodes later, reads them.

cpu-jit keeps the whole tape in a 12 KB double[], one path at a time per thread — a fraction of an L1 cache. A GPU lane gets, at best, a few hundred registers, shared across however many paths the scheduler is running at once. Ask it to hold 760 values and the compiler spills the overflow to scratch memory; every lane doing that at once turns the reverse sweep into a memory workload wearing an arithmetic costume. That's the 8×.

Shared checkpoint analysis, backend-specific kernels

AadCheckpointPlan in nablatensor-core provides shared checkpoint analysis across CUDA, ROCm, OpenCL, Vulkan, and WebGPU: it finds each node's last use and tracks which values cross segment boundaries. Vulkan and WebGPU use its balanced cut selection, which favors cuts with fewer crossing values. CUDA, ROCm, and OpenCL use the same last-use and carry analysis, while retaining their fixed segment-length policy. Each backend still emits its own kernel code and uses its own runtime/storage layout.

Did you know?

The plan is shared; the kernels are not. GLSL, WGSL, and CUDA-family backends each emit and replay their own code, but they can rely on the same core analysis of last uses and values carried across cuts.

What it's worth

This part is Vulkan's own numbers, from the engine's blog article on exactly this technique — the same 252-step Asian tape, 300,000 paths, value plus five Greeks, measured as the median of three seven-point spot ladders:

kernelms / 7-point ladderMpath/svs plain
plain (fully unrolled)13815—
checkpointed, 16 segments47462.9×
checkpointed, 4 markets/dispatch20—6.8×

Every row returns the same price and Greeks to all seventeen significant digits — not "close," bit-exact, because every recomputed value is the same instructions on the same inputs the first pass already produced. Below the spill point (50–100 time steps) the same trick costs 0.85×, which is why Vulkan only turns it on automatically above 768 nodes.

Try it yourself

For the CUDA/ROCm/OpenCL fixed-segment policy, use segLen = max(48, round(√n)). Would checkpointing apply to any tape this Learn section has built so far? 1.3's European call: 26 nodes. 3.1's Asian call at 252 steps: 1,536 nodes. The policy applies only above its minNodes threshold; Vulkan and WebGPU instead choose balanced cuts using the shared analysis.

▶️ Run it

This interactive example builds the same production Asian-call tape and asks NablaTensor's AadCheckpointPlan for its fixed-length CUDA-family segment layout. TeaVM shows the planner's checkpoint layout, but does not execute a GPU kernel. Vulkan and WebGPU use the shared last-use and carry analysis too, with balanced cuts instead.

Java · compile and run in this browser

To measure the checkpointed GPU kernel itself, use a machine with a Vulkan-capable device:

mvn -o -q install
mvn -o -q -pl nablatensor-examples exec:java \
  -Dexec.mainClass=com.nablatensor.bench.CheckpointBench \
  -Dengine=vulkan -Dsteps=252 -Dscenarios=300000 -Dmode=checkpointed

Nothing to switch on beyond the engine name — a tape of 768 nodes or more gets the checkpointed sweep automatically on vulkan.

⚠️ What this doesn't do

This page focuses on vulkan's register-only checkpointing in depth. The cuda/opencl/rocm path has a different scratch-buffer layout and its own segment-length policy; the deeper link below covers the historical layout bug and tuning work in more detail. It also doesn't cover sharing random draws across a whole revaluation ladder in one GPU dispatch, a related but different trick from the same article.

What's next

→ Deeper: Forget, then remember: checkpointing the adjoint sweep on a GPU has the full derivation, the register-spill cliff chart, the compiler- merging bug that needed an XOR to defeat, and the unmerged CUDA/OpenCL/ROCm fix this page only named. → Next: Four ways to talk to a GPU — CUDA, Vulkan, OpenCL, and ROCm, compared instead of just named.

Questions or corrections? open an issue