Fuse the GPU frame chain: up to 2.6x on small frames, bit-identical output

[jacotay7]

Small GPU frames were launch-bound. On the Arm bench (Neoverse-N1), an 80x80 Pyramid frame spent about 430 µs of host time issuing roughly twenty CuPy calls, three of them uploads of a scalar input, while the GPU sat idle. This PR cuts that to about 160 µs. Seeded output stays bit-identical on both CPU and GPU.

What changed

• Fused expectation kernel (_cuda.fused_expectation): (rate + background) * exposure·QE * PRNU, then + dark + extra, in one ElementwiseKernel. It writes the photo map and the total.
• Fused readout kernel (_cuda.fused_readout): full-well clip, defects, the reset/avalanche/correlated/read noise terms, the gain divide, bias pedestal, amplifier offset, bias structure, common mode, rint, the ADC clip and the uint32 cast, in one kernel. It replaces about 10 launches.
• No per-frame scalar uploads. noise._working_operand keeps scalar background/extra_electrons on the host as working-dtype NumPy scalars. The _CuPyGenerator wrapper reuses a cached device copy of scalar gamma/poisson parameters (EM-gain scale, CIC rate).
• Poisson draws go straight into the working dtype (dtype= on CuPy's sampler). This is the same conversion as .astype but without the int64 intermediate.
• The truth reuses the chain's total: photo + dark + extra, same operands and order. This applies with no workspace, no ROI and binning 1.
• Identity arithmetic is skipped on both paths: + 0, / 1 and * 1 by Python scalars.

Why the output is unchanged

• Noise is drawn by the same calls in the same stream order.
• The kernels repeat the same float32/float64 operations in the same order: CuPy's own clip expression, Python scalars rounded to the working dtype exactly as CuPy does, and --fmad=false.
• Any operand the kernels cannot take exactly (a dtype mismatch, a NumPy-scalar config value, an odd shape, or a value the unfused path rejects) falls back to the separate operations.

Evidence

• A golden-hash harness covered 1,028 cases per device: every preset plus four kitchen-sink configs, single and double precision, and the paths expose (array/scalar/background/extra), binning digital/on-chip, workspace with and without truth, out, ROI, dark/bias/flat, unseeded sequences, series, NDR, CDS and simulate_frame. It hashed every ADU and truth array. 0 differences vs main on CPU and on the RTX 4060.
• New test_gpu_fused_kernels_are_bit_identical_to_separate_operations (10 cases) compares the fused path with the forced-unfused path bit for bit, truth included. A mutation (v / gain → v * (1 / gain)) makes it fail.
• On the bench: full suite 595 passed; -m gpu 23 passed. Locally: ruff 0.16.10, mypy --strict and mkdocs build --strict are clean.

Numbers

Method: bench_devices.py --seconds 2 --warmup 10, Neoverse-N1 pinned to 12 cores, main and this branch interleaved, median of 3 runs (A400: 2 runs). Values are frames/s.

Workflow	RTX 4060 main	RTX 4060 PR	RTX A400 main	RTX A400 PR	CPU (12 cores) main	CPU PR
Pyramid 80x80	2,422	6,418 (2.65x)	2,403	6,315 (2.63x)	2,146	2,158
Shack-Hartmann 160x160	2,424	6,340 (2.62x)	1,698	1,794 (1.06x)	576	584
OCAM2K 240x240	1,650	2,733 (1.66x)	517	531 (1.03x)	143	145
SAPHIRA 256x320	1,809	2,072 (1.15x)	384	393 (1.02x)	114	116
1024x1024	240	247 (1.03x)	42	45 (1.06x)	13.4	14.1

What didn't help, or can't be done without changing seeded output

• Large GPU frames: CuPy's Poisson kernel takes about 3.9 of the 4.0 ms of device time for a 1024x1024 frame. It runs PTRS in double precision, which the RTX 4060 executes at 1/64 rate. A float32 sampler would change the values.
• CUDA graphs: CuPy advances the Poisson/Gamma seed and the cuRAND offset on the host at launch. A captured graph would replay the same noise.
• CPU: NumPy's Poisson sampler takes about 80% of a 1024x1024 frame, and Gamma dominates the EMCCD/eAPD cases. Both draw from one sequential stream, so they cannot be threaded without changing seeded frames. Reusing buffers made no difference (no allocation penalty measured). The remaining gain on the NumPy path is 1–5%.

🤖 Generated with Claude Code