Forget, then remember: checkpointing the adjoint sweep on a GPU

Here is a number that should not exist.

Take the 252-step arithmetic Asian call from the previous articles — 1,536 tape nodes, 252 normal draws per path — and run it on the Vulkan engine, one path per GPU lane. Price only, it does 125 million paths per second. Ask for the value and its five Greeks from the same tape, and it does 15. Eight times slower. On the CPU engines the same request costs 1.7×, and the bytecode article explained why it cannot reasonably cost less: the reverse sweep is a second walk over the tape, roughly as long as the first.

So where did the other 6× go? Not into arithmetic. Into remembering. And the fix — which is the whole of this article — is to remember less: keep 38 numbers per path instead of 760, throw everything else away the moment it is used, and compute it a second time when the reverse sweep comes back for it. That makes the Greeks pass 2.9× faster on its own and 6.8× faster once four markets share a dispatch, returns the same seventeen digits as the kernel it replaces, and switches itself on only for tapes long enough to need it.

Every number below was measured on referential machine — the same Ryzen 7 8845HS as the previous articles, this time on its integrated Radeon 780M through Vulkan (RADV), fp32 — as the median of three seven-point spot ladders, one fresh JVM per row.

Part one: what a Greeks pass has to remember

A tape is a receipt. Recording the Asian call writes down every arithmetic step the price took: node 15 is a random draw, node 16 multiplies it by vol·√dt, node 17 adds the drift, node 18 takes the exponential, node 19 multiplies by the running spot, node 20 adds it to the running sum — and then the same six lines again, 251 more times, and nine lines of payoff at the end. The forward sweep reads the receipt top to bottom and produces a price.

The reverse sweep reads it bottom to top and produces the Greeks. It carries one number per node — call it the node's blame: how much the final price would move if that node's value moved by a tiny amount. The last node's blame is 1, by definition. Every other node's blame is computed from the blame of the nodes that used it, by a rule that depends on the operation. And the rules need some of the forward values back:

forward node	to pass blame back to its inputs, the reverse sweep needs
c = a + b	nothing — the blame splits evenly
c = a × b	a and b — a's share is scaled by b, and vice versa
c = exp(a)	c itself — the slope of exp is exp
c = a / b	b and c
c = max(a, b)	a and b — to know which one won

Count what that means for one time step of the Asian. The multiplication by vol·√dt needs the draw. The exponential needs its own output. The multiplication by the running spot needs the running spot. Three numbers per step; 252 steps; about 760 forward values that the plain kernel has to keep from the moment they are computed until the reverse sweep, hundreds of nodes later, reads them.

On a CPU that is a non-event. The cpu-jit engine keeps the whole tape in a double[] of 12 KB, one path at a time per thread, and 12 KB is a fraction of the core's first-level cache. The GPU is a different shape.

Part two: a GPU lane's pockets

A GPU gets its speed from running thousands of paths at once, one per lane, and switching between them whenever one is waiting for something. Each lane has a small private store — its registers — and the arithmetic is fast precisely because it only ever reads and writes those. A lane gets at most a few hundred 32-bit registers, and the more a kernel uses, the fewer paths the scheduler can keep in flight to cover each other's waits.

Ask a lane to hold 760 values and it cannot. The compiler does the only thing it can: it spills — writes the overflow to "scratch", a private slice of main memory with a small cache in front, and reads it back when the reverse sweep asks. One path spilling 3 KB is nothing. Every lane on the chip spilling 3 KB at the same time, then all of them reading it back in reverse order, is a memory workload dressed up as an arithmetic one, and memory is what a GPU has least of. The 8× is that.

You can see the exact moment it happens by making the tape longer:

At 50 and 100 steps the plain kernel is fine: everything fits, and it is the fastest thing on the chart. At 150 steps it has lost two thirds of its speed; at 252, ninety percent. Nothing about the arithmetic changed. The tape just got too long to hold.

Part three: bookmarks, and re-reading

Here is the idea, without the maths first.

Suppose you are going to be quizzed on a novel, chapter by chapter, in reverse order — chapter 16's questions first, then chapter 15's, and so on back to chapter 1. One way to prepare is to memorise the whole book. The other is to read it once, leave a bookmark at the start of each chapter, and forget the pages. When chapter 16's questions come, open the book at bookmark 16, re-read that one chapter, answer, close it. Then bookmark 15. At any moment you are holding sixteen bookmarks and one chapter — never the book.

The price is obvious: every chapter is read twice. The forward sweep runs once to produce the price and the bookmarks, and then once more, a segment at a time, inside the reverse sweep. On a CPU that would be a bad trade — you would pay a second forward sweep to save memory you were not short of. On a GPU lane that is spilling, it is the best trade available: arithmetic is the thing the lane has in abundance, and memory traffic is the thing that was costing 8×.

There is one condition, and it is the same property the draw cache was built on. A chapter has to read the same the second time. Every node in the tape is a function of the nodes above it, except one kind: the random draws. If the draws came from a generator with internal state — "give me the next number" — then re-reading a chapter would need that state saved at every bookmark, or it would produce different dice and a different, wrong, adjoint. NablaTensor's GPU kernels use Philox, a counter-based generator: draw k of path p is a pure function of (p, k), computed from scratch every time it is asked for, and the comment on the generator has said since it was written that this is so "the adjoint sweep can recompute a checkpointed segment's randoms without saving RNG state". That sentence sat there for a while before anything used it.

Part four: where to cut

The analogy hides one question: what, exactly, is a bookmark? In a novel it is a page number. In a tape it is the set of values computed before the cut that some node after the cut still reads. Those have to survive; everything else in the segment can be recomputed from them.

The interesting fact about the Asian tape — and about most Monte-Carlo tapes — is how small that set is. Look at three time steps' worth:

At the boundary between step 1 and step 2, exactly seven numbers cross: the strike, rate and maturity (inputs the payoff will need), the drift and vol·√dt (computed once, used by every step), and the two things a step actually carries forward — the running spot and the running sum. It is seven at the next boundary, and the one after that, and at every boundary of the 252-step tape. The forward computation is a narrow pipe. Only the reverse sweep's appetite is wide, and the whole trick is to feed that appetite segment by segment instead of all at once.

The code that finds the cuts is short. For every node it records the last node that reads it; then for each of the fifteen cut positions it looks in a window around the evenly spaced target and picks the spot where the fewest values are still alive:

int[] lastUse = new int[n];
for (int i = 0; i < n; i++) {
  if (refsA(tape.op(i))) lastUse[tape.argA(i)] = i;
  if (refsB(tape.op(i))) lastUse[tape.argB(i)] = i;
}
for (int s = 1; s < segments; s++) {
  int target = (int) ((long) s * n / segments);
  int window = Math.max(1, n / (4 * segments));
  // best cut in [target - window, target + window]: fewest j < c with lastUse[j] >= c
  ...
}

Everything that crosses any cut — plus the inputs and the output — is marked persistent and declared once, at the top of the path loop. Everything else is local to its segment. The generated GLSL then has this shape, with the real node numbers from the 252-step tape:

// 38 persistent values live for the whole kernel
float v0; float v1; ... float v11; float v13; ... float v1535;

// forward sweep: sixteen blocks; a block's locals die at its closing brace
{
  v0 = in0;  v1 = in1;  v2 = in2;  v3 = in3;  v4 = in4;
  float v5 = v2 * v2;
  ...
  v11 = v8 * v10;          // drift          — persistent
  v13 = v2 * v12;          // vol·√dt        — persistent
  float v15 = rng_normal(0u);
  float v16 = v13 * v15;
  float v17 = v11 + v16;
  float v18 = exp(v17);    // needed by the reverse sweep — and forgotten
  float v19 = v0 * v18;
  float v20 = v14 + v19;
  ...
}
{ ... }   // × 15 more

A block's local variables go out of scope at its closing brace, and the compiler frees their registers. That is the "forget". Six lines per step, and at the end of the forward sweep the lane is holding 38 numbers.

Then the reverse sweep, last segment first. Each block re-runs its segment's forward lines — rng_normal(238u) produces the identical draw it produced the first time, because (path, 238) is all it depends on — then runs the segment's adjoint rules against those freshly recomputed values:

d1535 = 1.0;                       // the price's blame is 1
{                                  // segment 16
  rng_init((pathLo ^ pc.seedLo) ^ pc.zero, (pathHi ^ pc.seedHi) ^ pc.zero);
  float v11 = opq(v11);            // the bookmark, read back
  float v13 = opq(v13);
  float v1436 = opq(v1436);        // running sum at the cut
  float v1441 = opq(v1441);        // running spot at the cut
  float v1443 = rng_normal(238u);  // the same dice as the first time
  float v1444 = v13 * v1443;
  float v1445 = v11 + v1444;
  float v1446 = exp(v1445);        // recomputed, not remembered
  ...
  float d1444 = 0.0;  float d1445 = 0.0;  ...
  d1531 += d1535 * v1534;          // the adjoint rules, highest node first
  d1534 += d1535 * v1531;
  ...
}
{ ... }                            // segment 15, then 14, down to 1

The blame of a persistent node — d11, d13, the inputs' — is declared at the top level so that it accumulates across segments; the blame of a local node is declared inside its block and dies with it. When the last block closes, d0 through d4 hold the five Greeks, and they are added into the per-invocation Kahan sums exactly as before.

Part five: why there is an XOR in the reverse sweep

Two lines in that listing look odd: the opq(...) wrappers on the bookmark values and the ^ pc.zero on the generator's key. They are the one piece of the mechanism that is about the compiler rather than the maths, and the kernel does not work without them.

A recomputed segment is, by construction, the same arithmetic on the same inputs as the first time round. That is precisely the pattern an optimising compiler is built to find: "these two expressions are identical — compute it once, keep the result, use it twice." Common-subexpression elimination is correct, it is on by default in every shader compiler, and applied here it puts every forward value straight back on the live list for the whole kernel. The generated source is perfectly checkpointed and the compiled kernel is the plain one again.

The way out is to make the recomputation look different while being the same. The pipeline's push-constant block carries one extra uint, zero, which the host sets to 0 on every dispatch. Every value read from a bookmark passes through

float opq(float x) { return uintBitsToFloat(floatBitsToUint(x) ^ pc.zero); }

and the generator is re-keyed with key ^ pc.zero at the top of each reverse block. At run time x ^ 0 is x, bit for bit — not "approximately", not "up to rounding"; the same 32 bits, including the sign of a zero. At compile time pc.zero is an unknown value that arrives with the dispatch, so v11 ^ pc.zero is not v11, the chain built on it is not the original chain, and nothing merges. One XOR per bookmark value per segment, and the recomputation stays a recomputation.

Part six: what it is worth

The methodology is the one from the draw-cache article: one kernel, warm on a different seed, then seven calls with only spot changed, three such ladders, median.

The 252-step Asian, 300,000 paths, value and five Greeks:

kernel	ms per seven-point ladder	Mpath/s	vs plain
plain	138	15	—
checkpointed, 16 segments	47	46	2.9×

The segment count barely matters. Two segments, four, eight, sixteen, thirty-two, sixty-four, a hundred and twenty-six: every one of them lands between 47 and 50 ms. That is the signature of a memory problem rather than an arithmetic one — the cost was never "how much do you recompute", it was "does the working set fit", and any cut at all makes it fit. Sixteen is the default because it is comfortably inside the flat region on this tape and keeps the compiled shader small; the checkpointed kernel takes about three seconds to compile against one for the plain one, paid once per process and cached by the driver after that.

Below the cliff, the same kernel costs: 0.85× at 50 and 100 steps, because the second forward sweep is pure overhead when the first one never spilled. So checkpointing is not always on. The engine switches it on when a tape has at least 768 nodes — the middle of the gap between the last length that fits and the first that does not, on this tape — and leaves shorter tapes on the plain kernel. That threshold, and the segment count, are the two knobs in the next section.

Part seven: it computes the same number

A recomputation that returned different bits would be a bug, not an optimisation, and it would be an easy one to write: a draw regenerated with the wrong index, a bookmark read one segment early. So, spot 100, seventeen significant digits, the plain kernel against the checkpointed one:

252-step Asian, 300,000 paths, Vulkan fp32
kernel           price               delta                vega
plain            5.3010181964111330  0.56163036491394050  22.385515967610676
checkpointed     5.3010181964111330  0.56163036491394050  22.385515967610676

Identical. Not to a tolerance — the same double, because every float partial that went into it was the same float. The forward values that the reverse sweep reads are recomputed by the same instructions from the same inputs, the dice are a pure function of (path, k), and the XOR that keeps the compiler honest is the identity on every bit pattern. The engine's tests can assert == here, and do.

Sharing the dice across a ladder

There is a second thing in this release that composes with the first, and it is the draw cache's idea in the form a GPU actually likes.

A revaluation ladder — seven spots, one seed — asks the generator the same question seven times. The CPU engines answer that with a memo table in RAM. On a GPU the memory is the scarce thing, so the Vulkan kernel does something else: VulkanAadKernel.replayMany takes several input sets at once and prices all of them per path inside one dispatch, generating each path's draws exactly once and holding them in registers while the tape runs for market 1, market 2, market 3. The draws are never written anywhere; there is nothing to cache because nothing is asked for twice. Two things fall out of it: the dice are paid for once per ladder rather than once per market, and the GPU gets several independent chains of arithmetic per lane to interleave, which it is very good at.

tape	pass	one call per market	fused	vs
European, 200M paths	price	101 ms	25 ms (7 markets)	4.0×
European, 200M	value + 5 Greeks	115 ms	41 ms (7 markets)	2.8×
Asian 252, 1M paths	price	58 ms	16 ms (7 markets)	3.6×
Asian 252, 300k, checkpointed	value + 5 Greeks	47 ms	20 ms (4 markets)	2.3×

That last row is the headline: 138 ms for the plain Greeks ladder, 20 ms for the checkpointed one with four markets per dispatch — 6.8×, and still the same seventeen digits. Register pressure grows with the number of markets, so on the long tape four per dispatch is the sweet spot; the one-draw European is happy with all seven. replayMany is an engine-level method today — the Nabla model API still issues one call per market, and a ladder request on it is the next step.

Part eight: three more engines, and the two bugs that made them look hopeless

vulkan is not the only engine with a checkpointed kernel. cuda, opencl and rocm all share one implementation — CudaAadCodegen.generateCheckpointed — compiled three different ways: NVRTC for cuda, HIPRTC for rocm unchanged, and a small dialect rewrite for opencl. Naively benchmarked on the same 252-step Asian tape, all three came back with the same answer, and it was not a good one:

engine	hardware	settled, plain	settled, checkpointed	vs plain
opencl	referential machine's Radeon 780M	69.9 ms	80.8 ms	0.87×
rocm	referential machine's Radeon 780M	73.5 ms	80.5 ms	0.91×
cuda	a Colab Tesla T4	25.2 ms	28.1 ms	0.90×

(1,000,000 paths/dispatch, fp32, value + 5 Greeks. opencl and rocm share a tape and a machine with the Vulkan numbers above; cuda ran on a discrete datacenter GPU with sixteen gigabytes of its own dedicated memory — about as far from referential machine's memory-starved APU as GPU hardware gets.)

Checkpointing was not a small win on any of them. It was a consistent, repeatable loss, by roughly the same 9–13%, on three GPUs from two vendors, one of which has no shared memory bus to blame. That consistency is itself a clue: whatever is wrong lives in the shared code, not in any one driver.

The scratch buffer is laid out backwards

generateCheckpointed gives every GPU thread a private, contiguous slice of the scratch buffer:

const unsigned long long srow = tid * SLOTS_PER_PATH;
...
scratch[srow + slot] = v1441;          // write a checkpointed value
...
float v1441 = scratch[srow + slot];    // read it back, later, in reverse

That is an array of structures: thread 0's checkpoint sits at [0, SLOTS_PER_PATH), thread 1's at [SLOTS_PER_PATH, 2·SLOTS_PER_PATH), and so on. It reads naturally — "here is thread tid's little bundle of saved values" — and it is exactly the layout a GPU's memory system is built to punish. Thirty-two threads in a warp execute scratch[srow + slot] on the same instruction, at the same slot, so their addresses are tid · SLOTS_PER_PATH apart — for this tape, 90 floats, 360 bytes. A memory controller wants those thirty-two addresses inside one 128-byte line; instead they land thirty-two different lines apart, and the load that should be one transaction becomes up to thirty-two.

The fix is the standard one: lay the buffer out structure of arrays instead — every thread's slot-k value adjacent to the next thread's slot-k value — so a warp's access to one slot is one contiguous 128-byte read instead of thirty-two scattered ones:

- const unsigned long long srow = tid * SLOTS_PER_PATH;
- ...
- scratch[srow + slot] = v1441;
+ scratch[(unsigned long long) slot * invocations + tid] = v1441;

invocations — the total thread count across the launch — was already a kernel parameter, put there to size the buffer; it just was not being used to shape it. One line changes, in three places, in one file shared by all three engines.

Rebuilt and re-measured, right there: the loss shrinks from 8% to 3%. Real, and not nearly enough.

Every segment boundary has a fixed cost the plain kernel does not

Part six of this article measured Vulkan's segment count and found it did not matter: "two segments, four, eight, sixteen, thirty-two... every one of them lands between 47 and 50 ms." That is true for Vulkan, and it is why the number sixteen never got a second thought. It is not true here, because Vulkan's checkpointing has no scratch buffer to argue with. Every extra segment boundary in the register-recompute kernel costs a few more lines of ALU; every extra segment boundary in the scratch-buffer kernel costs a full round trip to memory, coalesced or not. AadCheckpointPlan's segment length defaults to max(48, √n) — 48 for this 1,536-node tape, which cuts it into 32 segments. That default reads like it was chosen for a cost model where segment count is nearly free. It is not free here; it is the second knob nobody had turned.

Sweeping -Dnablatensor.checkpoint.segLen on the now-coalesced kernel:

segLen	segments	settled	vs plain
24	64	98.4 ms	0.77×
48 (default)	32	77.3 ms	0.97×
100	16	60.6 ms	1.24×
200	8	58.1 ms	1.30×
300	6	69.9 ms	1.08×
400	4	77.1 ms	0.98×
600	3	77.7 ms	0.97×
800	2	70.8 ms	1.06×

Fewer, bigger segments amortise the fixed per-boundary cost over more useful recomputation, and the optimum for this tape sits at eight segments, not Vulkan's own default of sixteen for the same 1,536 nodes — the two kernels pay for a boundary in different currencies, so there is no reason to expect the same segment count to suit both. The shape past the optimum is not a clean curve: 800 (two segments) ticks back up slightly, most likely because at that point almost nothing crosses a cut at all — slotsPerPath drops to zero, meaning that configuration barely touches the scratch buffer in the first place, a different regime from the rest of the sweep. The unambiguous part is the left half: too many small segments is far worse than too few large ones.

Coalesced layout and a segment length chosen for this cost model, together, on the same tape, same machine, same harness that measured every other number in this section:

opencl, 252-step Asian, 1,536 nodes, 1,000,000 paths, fp32
kernel                              settled     Mpath/s   vs plain
plain                               75.4 ms     13.3      —
checkpointed, as shipped            81.8 ms     12.2      0.92×
+ coalesced scratch buffer          77.3 ms     12.9      0.97×
+ segLen 200 (8 segments)           58.1 ms     17.2      1.30×

Bit-exact against the plain kernel at every step, on every configuration — the fixes touch addressing and segment boundaries, not arithmetic.

So it was never "Vulkan-only"

It never was, mechanically. Checkpointing's saving does not care which GPU API compiled the kernel; what it needs is a working set that genuinely spills and a way to shrink it that costs less than the spill did. cuda, opencl and rocm had that too. What they did not have was a scratch buffer anyone had checked for coalescing, or a segment length tuned for a strategy that, unlike Vulkan's, pays a real price for every boundary it draws.

I verified the fix end to end — bit-exact prices and Greeks, twice, back to back — on opencl on referential machine, because that is what was asked. I have not re-run it on cuda or rocm. But the kernel source touched by both fixes is character-for-character the same file those two engines compile unchanged, and both showed the identical pre-fix loss (0.90–0.91×) that opencl showed before either fix. There is no honest reason to expect a different answer there — that is a prediction, though, not a measurement, and it is marked as one until someone runs it.

Neither change is in the engine as shipped. CudaAadCodegen.java's scratch indexing and AadCheckpointPlan's segment-length default are exactly as they were before this article was written; what exists right now is a diagnosis and a demonstrated fix, not a merged one.

What to actually use

This part is vulkan-specific; cuda, opencl and rocm checkpointing is off by default and stays off — Part eight's fixes are a diagnosis, not a config flag, until someone lands them.

Nothing to switch on for vulkan. A kernel built from a tape of 768 nodes or more gets the checkpointed adjoint sweep automatically; shorter tapes keep the plain one. The knobs, for when a tape does not look like this one:

-Dnablatensor.vulkan.ckpt=<segments>   # 1 switches it off; any n > 1 forces it on with n segments
-Dnablatensor.vulkan.ckpt.nodes=<n>    # the automatic threshold, default 768
-Dnablatensor.vulkan.dump=<dir>        # write the generated GLSL, to see the cuts

If your product's per-step state is much wider than the Asian's — a basket with twenty underlyings carries twenty running spots across every cut — the bookmark is bigger, the cliff comes sooner, and lowering the threshold is the first thing to try. If it is narrower, the default is already right.

The public API is unchanged:

try (var mc = MonteCarlo.of(Products.asianCall())
        .market(EquityMarket.atmOneYear()).steps(252)
        .greeks().fp32().on("vulkan").build()) {
  var r = mc.run(300_000, 42L);        // checkpointed: the tape has 1,536 nodes
  System.out.println(r.price() + "  delta " + r.greek(EquityMarket::spot));
}

The ranking of what mattered, in the order it mattered:

1. Noticing that the Greeks pass was memory-bound. An 8× gap where the CPU shows 1.7× is not an arithmetic gap, and no amount of tuning the arithmetic would have found it.
2. The draws being a pure function of (path, k). Without it, re-reading a segment would need the generator's state saved at every bookmark, and the memory saving would be a fraction of what it is.
3. Cutting where the pipe is narrow. Seven values cross a step boundary of this tape. Fifteen cuts of seven collapse to 38 persistent values once the shared ones are counted once; 38 fits.
4. One XOR. Without it the compiler merges the recomputation into the original and the kernel is bit-for-bit the plain one, at the plain one's speed. A technique that works by computing the same thing twice has to say so in a way the compiler cannot hear.
5. Checking the arithmetic before believing the excuse. "Vulkan recomputes into registers, the others round-trip memory, so of course they lose" sounds complete, and it predicted nothing that held up against a Tesla T4 with VRAM to spare. The real reasons — an uncoalesced buffer and a segment length tuned for the wrong kernel's cost model — were both sitting in the source the whole time.

Which is the moral, and it is a GPU-shaped one. Every previous article on this site made something faster by not computing it twice — compile the tape, cache the draws, run the adjoint once instead of bumping. This one made something three times faster by computing it twice on purpose, because the thing being saved was not arithmetic. Know which resource you are short of before you start conserving one.

And a second moral, arrived at by re-deriving the first one badly before getting it right: a tidy explanation is not the same as one that predicts a number correctly on hardware you have not measured yet. "It only works on Vulkan" was tidy, and wrong. The two actual bugs were untidy — one line of address arithmetic, one mistuned default — and the only way to tell the tidy story from the true one was to open the file and run it.

Source: VulkanAadCodegen.java (the cut selection in emitCheckpointed, the segment blocks, opq, and the Philox generator whose comment promised this), VulkanAadKernel.java (the automatic threshold, the zero push constant, and replayMany), and CudaAadCodegen.java (generateCheckpointed's scratch-buffer indexing, shared by cuda, opencl and rocm) and AadCheckpointPlan.java (the segment-length default) for Part eight.