Forget, then remember: checkpointing on a GPU

The 252-step Asian tape from Module 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

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×.

Two checkpointing strategies, one engine

nablatensor-core's AadCheckpointPlan cuts a tape into fixed-length segments — the actual, current algorithm is plainer than it might sound:

int segLen = Integer.getInteger("nablatensor.checkpoint.segLen",
    Math.max(48, (int) Math.round(Math.sqrt(n))));
int segCount = (n + segLen - 1) / segLen;

For the 1,536-node Asian tape: segLen = max(48, round(√1536)) = max(48, 39) = 48, so segCount = 1536 / 48 = 32 — evenly spaced cuts, no attempt to find a cheaper one nearby.

Did you know?

AadCheckpointPlan.java's own class doc says it is "shared by the CUDA and Vulkan adjoint code generators." It isn't, anymore — VulkanAadCodegen .emitCheckpointed never references AadCheckpointPlan at all. It has its own, separate cut-selection loop that searches a window around each evenly spaced target for the position where the fewest values are still live, and only CudaAadCodegen.generateCheckpointed — shared by cuda, opencl, and rocm — actually takes an AadCheckpointPlan and uses its plainer fixed-length segments. Two genuinely different strategies live in one engine, and only one of them is named in the class comment that describes both.

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:

kernel	ms / 7-point ladder	Mpath/s	vs plain
plain (fully unrolled)	138	15	—
checkpointed, 16 segments	47	46	2.9×
checkpointed, 4 markets/dispatch	20	—	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

Using segLen = max(48, round(√n)), work out whether checkpointing would even 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. AadCheckpoint Plan.of refuses to build a plan at all below its minNodes threshold — every tape before this page, including the 1,536-node one on cpu-jit or cpu, never goes near this code path, because it's the CUDA/OpenCL/ROCm codegen only, and the tape has to run on one of those three engines for AadCheckpointPlan to matter at all.

▶️ Run it

This one needs a GPU this book's other pages didn't — nothing here runs on cpu-jit. The command, for 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 only covers vulkan's register-only checkpointing in depth. The cuda/opencl/rocm path this page's second sidenote points at has its own, separate story — a scratch-buffer layout bug and a mistuned segment length, diagnosed but, per the source article, not yet merged into the engine as shipped — left for the deeper link below. 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.