Selective-scan CUDA kernels

Five CUDA implementations of the recurrence at the heart of selective state-space models (Mamba-style), built from the textbook version up to the single-pass decoupled look-back scan, each one measured against the previous one and against a memory-traffic model that predicts what the speedup should be before the kernel is written.

The fastest kernel reaches 79.6-91.5 % of the theoretical memory bandwidth of an RTX 4050 Laptop and is 54x to 69x faster than the best pure-PyTorch baseline.

WRITEUP.md is the long version: why the recurrence is parallelisable at all, what each kernel fixed about the previous one, and the two hypotheses that turned out to be wrong.

Results

RTX 4050 Laptop (Ada, sm_89, 192 GB/s theoretical), memory pinned at 8001 MHz, 42-45 W. DRAM regime, 96 MB per input array so nothing fits in L2. Median of 15 interleaved samples, repetitions inside the timed window.

L	PyTorch Hillis-Steele	Level 2, 3 passes	Level 3 serial	Level 3b warp	Level 3c + items	best % of peak
2,048	92.4 ms	4.208	2.091	2.114	1.720	91.5 %
4,096	101.0	4.324	2.387	2.341	1.747	90.0 %
8,192	109.6	4.321	2.599	2.475	1.867	84.3 %
16,384	118.6	4.343	2.727	2.541	1.910	82.4 %
32,768	127.0	4.310	2.813	2.522	1.928	81.6 %
65,536	135.7	4.311	2.986	2.577	1.975	79.6 %

Reading the right-hand panel of the figure: level 2 sits above level 3c in efficiency at long lengths while being 2.2x slower. That is not a contradiction. Level 2 moves seven arrays worth of traffic where level 3 moves three, so it has more bytes to hide latency behind. Percentage of peak measures how well a kernel uses the bus; time measures how much work it asked the bus to do. Level 3 wins by not doing the work.

The problem

The selective-SSM forward pass is

h[t] = a[t] * h[t-1] + b[t]

which looks strictly sequential. It is not. Each step is an affine map, affine maps compose, and composition is associative, so the recurrence is a prefix scan over pairs:

combine((a1,b1), (a2,b2)) = (a1*a2, a2*b1 + b2)

Associativity is the entire licence for running this in parallel. Commutativity does not hold, and every operand swap in the code below is a silent wrong-answer bug rather than a crash. That is what the float64 CPU mirrors in sim_blelloch.py and sim_multiblock.py exist to catch: each one contains a deliberately broken variant, so the test suite is shown to detect an operand-order error before any CUDA is written.

The five kernels

	file	ceiling on L	idea
Level 1	csrc/scan_blelloch.cu	2,048	Blelloch up-sweep and down-sweep in shared memory, one block per row
Level 2	csrc/scan_multiblock.cu	4,194,304	block scan, spine scan, apply prefix: three kernel launches
Level 3	csrc/scan_lookback.cu	none	decoupled look-back, single pass, thread 0 walks predecessors
Level 3b	csrc/scan_lookback_warp.cu	none	32 lanes inspect 32 predecessors at once
Level 3c	csrc/scan_lookback_items.cu	none	plus ITEMS elements per thread, scanned sequentially in registers

Decoupled look-back follows Merrill and Garland (NVIDIA tech report NVR-2016-002), which is what CUB's DeviceScan does. Each tile publishes its local aggregate with flag A, looks backwards until it finds a predecessor with an inclusive prefix (flag P), and then publishes its own P. Tiles are claimed with an atomicAdd counter rather than by blockIdx, because the hardware does not promise that block i is resident before block i+1, and a tile waiting on a predecessor that has not been scheduled yet deadlocks.

Findings

Five predictions were written down before measuring. Three held, two did not.

1. Memory traffic explains level 1 vs level 2 - confirmed, with a caveat. Level 1 touches 3 arrays, level 2 touches 7, so level 2 should be 2.33x slower. Measured 1.82x. The shortfall has an identified cause: at C = 2048 level 1 needs 16.9 KB of shared memory per block, which caps occupancy, and it only reaches 68 % of peak where level 2 reaches 85-87 %.

2. The serial look-back is the level 3 bottleneck - confirmed. The warp version's advantage grows with chain length, 1.04x at G = 2 up to 1.80x at G = 256. The same signature is visible in the table: the serial version decays from 75 % to 53 % of peak as L grows while the warp version stays flat at 61-62 %.

3. Per-call allocations were the remaining cost - refuted. A variant with a pre-allocated workspace measures identically, 0 % difference, inside the IQR. PyTorch's caching allocator already made torch.empty() nearly free. The hypothesis was reasonable and it was wrong; the workspace parameter is kept in the source only as documented dead weight.

4. The registered expectation for Hillis-Steele - refuted. The PyTorch Hillis-Steele does not hold up against the single-pass kernel, 54x-69x slower across the sweep. The comparison is not fair to the algorithm and the number should not be read as one: it runs in PyTorch, not as a register-resident CUDA kernel, so part of what it measures is the framework rather than the scan. A CUDA Hillis-Steele is listed as future work for exactly that reason.

5. More items per thread closes the remaining gap - confirmed. Predicted 2.33x over level 2 from the traffic model; measured 2.18x-2.48x. Efficiency rises from level 3b's 61-74 % to 79.6-91.5 %, which is the real claim here: the single-pass kernel stopped waiting on the prefix frontier and went back to being limited by memory, which is the only thing a scan should be limited by.

Supporting detail for the mechanism: items=8 only beats items=4 at long lengths and they tie at short ones, exactly what you expect if the benefit comes from shortening the tile chain, since with few tiles there is no chain to shorten.

Correctness

Every kernel is checked against a float64 CPU reference over ragged tails, degenerate shapes (L = 1), decay near 1.0, very long rows (L = 2^22), and chains up to G = 32,768 tiles.

check.py                   18/18
check_multiblock.py        25/25
check_lookback.py          19/19  + 240 determinism runs
check_lookback_warp.py     23/23  + 240 determinism runs, bit-identical to serial
check_lookback_items.py    22/22  + 150 determinism runs

The determinism runs matter more than the accuracy ones. A look-back scan is a hand-rolled inter-block synchronisation protocol, and a missing __threadfence() or a non-volatile load produces a kernel that is correct almost every time.

One test expectation had to be corrected rather than a kernel. Level 3c with items=1 is not bit-identical to level 3b, and should not be: it groups combine(combine(block, thread), elem) where 3b groups combine(block, combine(thread, elem)). fp32 honours associativity only approximately, so the last bits differ by about 2 ULP. Demanding bit equality there was a wrong expectation, not a bug. See the docstring in check_lookback_items.py.

Two real bugs the tests caught

Rounding the thread count up to a full warp, threads = max(C/2, 32), writes past the end of the shared tile at C = 32. Level 1 had the same bug and its tests passed anyway, because the overflow happened to land in the bank-conflict padding. It was caught for real by the level 2 suite. A test that passes for the wrong reason is worse than one that fails.

The benchmark harness also exposed three faults in itself before it produced a usable number: kernels too short for the GPU to ever boost (14 W, never clocking up), a working set that fit in L2 and so reported over 100 % of peak, and a measured-copy denominator that flattered the kernels because cudaMemcpy moves 4-byte scalars while these kernels move 8-byte pairs.

Building and running

Needs a CUDA toolkit with nvcc. The pip CUDA wheels do not ship one on Windows, only ptxas.

Everything, in one command:

build_env.bat run_all.py

That runs the five correctness suites, each in its own process so a deadlock in one look-back kernel cannot take the other four down with it. Add --bench to measure and regenerate the figure afterwards. The benchmark is opt-in on purpose: it is slow, and on a laptop roughly every second sweep is unusable because the memory clock oscillates, so running it by accident and quoting whatever came out is the exact failure this project spent its time avoiding.

Individual stages still work on their own:

build_env.bat check_lookback_items.py

Build flags live in one place, kernels.py, which also reads the target architecture off the device that is present rather than hardcoding sm_89. There is no setup.py: the five .cu files each declare their own PYBIND11_MODULE, so they are five separate extensions, and an ahead-of-time build would have to pin an architecture at install time. JIT compiles once per machine and caches.

Three non-obvious flags, all environment rather than code:

• -allow-unsupported-compiler, because MSVC 14.50 is newer than CUDA 12.1 knows about
• -Xcompiler /D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH, because Microsoft's STL independently asserts CUDA 12.4 or newer
• vcvars64.bat needs vswhere.exe on the PATH

CUDA 12.6 or 12.8 would need neither of the first two. Picking 12.1 to match the torch build exactly was a mistake that cost two flags.

Limitations

• One GPU, one run. A laptop shares its power budget between CPU and GPU, so roughly every other sweep comes back with the memory clock oscillating between 5501 and 8001 MHz and is unusable. These numbers come from a sweep where the clock stayed pinned, verified by level 2 returning its known stable value, but they have not been reproduced on a second architecture.
• The L = 512 and L = 1024 rows in bench.json are polluted; the clocks were still settling. They are excluded from the figure and the table.
• No vectorised (float4) loads yet. Deliberately deferred to keep one variable changing at a time.
• The Hillis-Steele comparison is not fair to Hillis-Steele: it runs in PyTorch, not as a register-resident CUDA kernel.
• fp32 only, forward pass only, no backward.

Repository layout

csrc/                     the five CUDA kernels
sim_blelloch.py           float64 CPU mirrors of the index arithmetic,
sim_multiblock.py         each with a deliberately broken variant
check*.py                 correctness and determinism suites
kernels.py                build flags and architecture detection, one copy
run_all.py                one command that runs everything
bench.py                  measurement harness, writes results/<gpu>/bench.json
plot_results.py           turns bench.json into the figure above
build_env.bat             the compilation recipe for this machine
PLAN.md                   the original plan
STATE.md                  working notes (Spanish)
WRITEUP.md                the long-form writeup