checkpointing_ssu kernel: fused replay + conditional state-write for Mamba2 (#3324)

<!-- .github/pull_request_template.md -->

## 📌 Description

Adds `flashinfer.mamba.checkpointing_ssu` — a fused CUDA kernel for the
Mamba2 selective-state-update (SSU) operation with two distinct
features:

1. **Replay** — fast-forward the SSM state through previously-cached
tokens
without re-running the upstream model layers, then process the new
tokens
in the same launch. Used for speculative decoding (MTP) and re-attempt
   flows where the model needs to "rewind and replay" recent history.
2. **Checkpointing** — at runtime, decide per request whether the new
tokens
fit in the cache or not. If they fit, skip the state HBM write entirely
   (the cache still has enough headroom).  If they overflow, write the
post-replay state to HBM and reset the cache. Eliminates the state-write
   path on most steps under typical workloads.

Replaces an upstream two-kernel Triton reference (precompute + main)
with a
single Ampere-class `mma.sync.m16n8k16` kernel that does the entire
`replay(prev_k) → output(new_T)` recurrence in one launch.

### What it computes

Per `(batch, head)`, unrolling the SSD recurrence to closed form:

```
h_t[d,n] = h_0[d,n] * exp(cumAdt[t])
         + Σ_{j≤t} exp(cumAdt[t] - cumAdt[j]) * dt_proc[j] * B[j,n] * x[j,d]
y[t,d]   = decay[t] * Σ_n C[t,n] * h_0[d,n]      (init_out — matmul-3)
         + Σ_{j≤t} CB_scaled[t,j] * x[j,d]        (cb_out   — matmul-4)
         + D[d] * x[t,d]
y[t,d]  *= z[t,d] * sigmoid(z[t,d])
```

with `CB_scaled[t,j] = exp(cumAdt[t]-cumAdt[j]) * dt_proc[j] * <C[t,:],
B[j,:]>`
the lower-triangular `[T,T]` scaling matrix.

**Replay** fast-forwards `h_0` with `prev_k` previously-cached tokens
*before*
the new `T` tokens are processed:

```
h_0 ← h_0 * exp(total_old_cumAdt) + old_x^T @ (coeff * old_B)
```

This decomposes into four matmuls per call:

| # | Op | Shape |

|----|---------------------------------|-----------------------------------|
| 1 | C @ Bᵀ (precompute CB_scaled) | `[T,N] @ [N,T] → [T,T]` |
| 2 | old_xᵀ @ dB_scaled (replay) | `[D,K] @ [K,N] → [D,N]` |
| 3 | C @ stateᵀ (init_out) | `[T,N] @ [N,D] → [T,D]` |
| 4 | CB_scaled @ x (cb_out) | `[T,T] @ [T,D] → [T,D]` |

All four ride `mma.sync.m16n8k16.f32.bf16/f16` (Ampere;
forward-compatible to
Hopper/Blackwell).  Hopper `wgmma` (M ≥ 64) and Blackwell `tcgen05.mma`
(≥ 64×64) are too coarse for these tile sizes (M ≤ 16, K = 16).

### Checkpointing semantics

The cache holds up to `MAX_WINDOW` historical tokens per sequence
(double-buffered).  Each call to the kernel can:

* **Append** the `seq_len` new tokens to the active buffer at offset
  `prev_k` (no-checkpoint path: `must_checkpoint = False`).  State stays
  in registers, no HBM write.
* **Checkpoint** to the staging buffer at offset 0, flipping the
  double-buffer pointer (checkpoint path: `must_checkpoint = True`).
  Post-replay state gets written to HBM (the only path that touches
  state gmem on the write side).

The trigger is computed per CTA at runtime:

```cpp
must_checkpoint = (prev_k + seq_len > MAX_WINDOW)
```

Under typical serving load (long context, small new-token chunks), most
calls hit Branch B and never touch state HBM.  Branch A only fires when
the cache fills up.

The kernel is template-specialized on the dispatch — `if constexpr
(must_checkpoint)` is **runtime**, dispatched via two helper functions
(`ssu_checkpoint{,_8bit}` / `ssu_nocheckpoint{,_8bit}`) on the per-CTA
boolean, so warp divergence inside the dispatch is balanced.

### Architecture: two kernel headers

The kernel is split into **two source files** because the 8-bit state
path
needs a fundamentally different MMA layout from the 16/32-bit path:

#### `kernel_checkpointing_ssu.cuh` — bf16 / fp16 / fp32 state

* **N-shard replay** (`Layout<_1, _4>` warp tiling): each warp owns a
contiguous slice of the `dstate` axis. Loads state from smem into a B-
fragment, hits HMMA in fp32 accumulation, writes back as bf16/fp16/fp32.
* `D_SPLIT ∈ {1, 2}` supported.  Splits each head's DIM axis across
multiple CTAs to lift small-batch occupancy (`D_SPLIT=4` deferred — the
  output MMA's `_1×4` warp layout needs `D_PER_CTA ≥ 32`).
* Philox stochastic rounding via PTX `cvt.rs.f16x2.f32` is wired in for
  fp16 state (HW path on sm_100a / Blackwell B200+, software emulation
  elsewhere).

#### `kernel_checkpointing_ssu_8bit.cuh` — int8 / fp8_e4m3fn state
(1-byte)

* **M-shard replay chain** (`Layout<_4, _1>` warp tiling): each warp
owns
  16 D-rows of the post-replay state, the chain matmul-3 happens *in
  registers* (replay's fp32 C-frag is converted to bf16 A-frag in place
  via `convert_layout_acc_Aregs_sm80` — no smem.new_state staging).
* Followed by a transposed matmul-4 (`x` as A with M=D, CBᵀ as B) →
  `output^T(D, T)` in regs → smem transpose → cooperative STG.128 to
  `(T, D)` gmem.
* Per-(head, dim) channel decode scale computed from the post-replay
  amax; stored alongside the quantized state.
* `D_SPLIT=1` only (the M-shard replay needs `D_PER_CTA ≥ 64`).
* Philox stochastic rounding:
  * **int8** / **int16**: pure `floor(x + uniform_noise)` —
    runs anywhere.
* **fp8_e4m3fn**: PTX `cvt.rs.satfinite.e4m3x4.f32` packs 4 fp32 values
    + 32-bit random seed into one fp8 vector store (HW path on sm_100a,
    bit-exact SW emulation elsewhere).

#### `kernel_checkpointing_ssu_common.cuh` — shared infrastructure

`load_data` (cp.async load of `x`/`dt`/`B`/`C`/`z` + cache scalars),
`compute_CB_scaled_2warp` (the `[T, T]` causal-masked CB precompute),
`compute_CB_old_2warp` (the no-write path's `CB_old @ old_x` extension),
`store_old_x` / `store_old_B` (cache writebacks), `precompute_dB_coeff`,
swizzle layouts, MMA traits.

#### Public dispatcher: `launch_checkpointing_ssu.cuh`

Routes on `(d_split, varlen)`:

* `d_split` switch reads `params.d_split` (1 or 2 for the 16/32-bit
path,
  always 1 for the 8-bit path).
* `varlen` switch reads `params.cu_seqlens != nullptr`.

Both `D_SPLIT` and `VARLEN` are kernel template parameters; the same
`.so`
holds all four `(D_SPLIT, VARLEN)` specializations.

### Quantized state path (int8 / int16 / fp8_e4m3fn)

For 8-bit states the state HBM is `int8` or `fp8_e4m3fn`; bandwidth
drops
~4× vs fp32 for the same recurrence.  Quantization is **per-(cache_slot,
head, dim) channel** — one decode scale per channel, broadcast over
`dstate`.  The kernel computes scale + quantize on checkpoint steps:

```
amax        = max(|state[d, :]|)  over dstate
encode      = quant_max / amax     (quant_max = 127 for int8, 448 for fp8)
state_q     = round_or_sr(state * encode)
state_scale = 1 / encode           (stored, broadcast over dstate on read)
```

On read, state is `state_q * state_scale` (decode_scale factored out of
the matmul inner product — applied post-matmul as a single FMUL per
output column).

### Varlen support

When `cu_seqlens` is provided, inputs are **packed** as `(1,
total_tokens,
nheads, dim)` / `(1, total_tokens, ngroups, dstate)` — vLLM's "no
padding"
batch ABI.  Each `cu_seqlens[i]` gives the token-axis base of sequence
`i`; `seq_len_i = cu_seqlens[i+1] - cu_seqlens[i]`.

Implementation:

* `VARLEN` is a kernel template parameter; dispatched at launch time.
* The wrapper picks `*_stride_seq = x.stride(0)` (per-batch) in
non-varlen
  or `x.stride(1)` (per-token) in varlen.  Kernel uses one uniform
  formula `outer * *_stride_seq` regardless of mode.
* `NPREDICTED` stays compile-time (max `seq_len` the caller commits to);
  per-sequence variable `T` realized via masking, not by re-sizing smem
  or MMA tiles.
* `must_checkpoint` uses the tighter `prev_k + seq_len > MAX_WINDOW`
  (vs. the conservative `prev_k + NPREDICTED > MAX_WINDOW`) — avoids
  unnecessary checkpoint rotation when a sequence ends short.

### Files

#### CUDA

* `include/flashinfer/mamba/checkpointing_ssu.cuh` — params struct
* `include/flashinfer/mamba/kernel_checkpointing_ssu.cuh` — 16/32-bit
kernel
* `include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh` — 8-bit
kernel
* `include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh` —
shared helpers
* `include/flashinfer/mamba/launch_checkpointing_ssu.cuh` — public
dispatcher
* `csrc/checkpointing_ssu.cu` — TVM-FFI binding (input validation +
param population)
* `csrc/checkpointing_ssu_kernel_inst.cu` — explicit kernel template
instantiation
* `csrc/checkpointing_ssu_jit_binding.cu` — JIT FFI export
* `csrc/checkpointing_ssu_customize_config.jinja` — JIT-stamped
constants
  (`NPREDICTED`, `MAX_WINDOW`, `DIM`, `DSTATE`, `HEADS_PER_GROUP`,
  `PHILOX_ROUNDS`, state dtypes)

#### Python

* `flashinfer/jit/mamba/checkpointing_ssu.py` — JIT module generator
(URI + config render)
* `flashinfer/mamba/checkpointing_ssu.py` — `checkpointing_ssu()` user
API + dispatch
* `flashinfer/mamba/__init__.py` — exports

#### Triton reference

* `tests/mamba/triton_reference/checkpointing_state_update.py` —
independent
  reference implementation.  Supports the same dtype matrix (fp16, bf16,
  fp32, int8, int16, fp8_e4m3fn), the same Philox SR paths, and the same
  varlen (`cu_seqlens` + `max_seqlen`) interface.  Used as the
  cross-validation oracle for the CUDA kernel.

### API

```python
flashinfer.mamba.checkpointing_ssu(
    state,                         # (cache, nheads, dim, dstate) — updated in-place
    old_x,                         # (cache, T, nheads, dim) cache, single-buffered
    old_B,                         # (cache, 2, T, ngroups, dstate) double-buffered
    old_dt_proc,                   # (cache, 2, nheads, T) f32
    old_cumAdt,                    # (cache, 2, nheads, T) f32
    cache_buf_idx,                 # (cache,) int32 — which buffer is "active"
    prev_num_accepted_tokens,      # (cache,) int32 — how many old tokens in active
    x, dt, A, B, C, out,           # standard SSU inputs/outputs
    D=None, z=None, dt_bias=None,
    dt_softplus=False,
    state_batch_indices=None,      # (batch,) — optional paged-cache mapping
    pad_slot_id=-1,
    state_scale=None,              # (cache, nheads, dim) f32 — required for int8/fp8 state
    rand_seed=None,                # single-element int64 CUDA tensor — enables Philox SR
    philox_rounds=10,
    d_split=None,                  # {1, 2} for 16/32-bit; auto-heuristic
    cu_seqlens=None,               # (batch+1,) int32 — enables varlen mode
    max_seqlen=None,               # upper bound on max(cu_seqlens diff), required in varlen
)
```

## 🔍 Related Issues

Replaces the discarded [#3217 `ssu_incremental`
PR](https://github.com/flashinfer-ai/flashinfer/pull/3217).

### Differences from #3217

The original PR added an "always-write replay" path (every call wrote
post-replay state to HBM, no conditional checkpoint logic).  This branch
extends that with:

1. **Conditional checkpointing** — runtime `must_checkpoint` decision
   per CTA; Branch B skips state HBM write entirely.
2. **8-bit quantized state** (int8, fp8_e4m3fn) with per-channel decode
   scales and stochastic rounding (Philox-driven SR; PTX
   `cvt.rs.satfinite.e4m3x4.f32` on Blackwell + bit-exact SW emulation
   elsewhere).
3. **int16 state** support (via the same per-channel scale
infrastructure
   as int8).
4. **Varlen** support (cu_seqlens) for vLLM's packed-batch ABI.
5. **Independent Triton reference** with full feature parity (was a
   replay-only Triton ref in the original PR).
6. Renamed `ssu_incremental` → `checkpointing_ssu` throughout to match
   the v19 split (checkpoint / no-checkpoint paths).

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

`tests/mamba/test_checkpointing_ssu.py` — 138 cases covering the CUDA
kernel + the merged Triton reference.

### CUDA kernel correctness

* **`test_checkpointing_ssu_max_window_gt_npredicted`** (12 cases) —
the main correctness sweep. `(NPREDICTED, MAX_WINDOW) ∈ {(4,8),
(10,16)}`
  × `state_dtype ∈ {fp16, bf16, fp32}` × `paged_cache ∈ {True, False}`.
  Inner loop sweeps `prev_k ∈ [0, NPREDICTED]` crossing the
  `must_checkpoint` boundary — exercises Branch A and Branch B in the
  same test.  Compared against Triton reference.
* **`test_checkpointing_ssu_int8_rn_parity`** (8 cases) — int8 state
with
round-to-nearest quantization. Cross-validated against the fp32 Triton
  reference path with explicit per-channel dequantization on the state.
* **`test_checkpointing_ssu_fp8_rn_parity`** (8 cases) — same as int8
but
  fp8_e4m3fn.  Skipped on SM < 89 (Ada/Hopper/Blackwell required for the
  fp32↔fp8 cvt PTX).
* **`test_checkpointing_ssu_int8_philox`** / **`…_fp8_philox`** /
**`…_philox`** (8 + 8 + 8 cases) — Philox stochastic rounding for int8 /
fp8 / bf16 state.
* **`test_checkpointing_ssu_philox_no_checkpoint`** (4 cases) — fp16
  Philox in Branch B (no-write).  Validates that state HBM stays
  byte-identical when must_checkpoint=False.
* **`test_checkpointing_ssu_philox_with_checkpoint`** (4 cases) — fp16
  Philox in Branch A across different `(prev_k, NPREDICTED, MAX_WINDOW)`
  triples that cross the must_checkpoint boundary.
* **`test_checkpointing_ssu_mixed_checkpoint_batch`** (8 cases) — mixed
  `prev_k` across the batch so different CTAs hit different branches
  in the same launch.

### Statistical SR unbiasedness

* **`test_philox_rounding_unbiased`** /
**`test_checkpointing_ssu_int8_philox_unbiased`**
  / **`test_checkpointing_ssu_fp8_philox_unbiased`** — verify that
  Philox SR preserves the expected value (mean residual ≈ 0) over many
  samples.

### Varlen (vLLM packed-batch ABI)

* **`test_checkpointing_ssu_varlen_mixed_no_checkpoint`** (5 cases) —
  mixed `seq_lens=[3,1,4,2,4]` with `prev_ks=[0,5,10,12,8]` chosen so
  that every sequence stays in Branch B.  Compares CUDA varlen against
  per-batch padded CUDA non-varlen.  Sweeps all 5 dtypes.
* **`test_checkpointing_ssu_varlen_mixed_checkpoint`** (5 cases) — same
  pattern with `prev_ks` chosen to force every sequence into Branch A.
* **`test_checkpointing_ssu_varlen_cuda_vs_triton_no_checkpoint`** (5
cases)
  / **`...checkpoint`** (5 cases) — **independent reference** check:
  CUDA varlen vs Triton varlen.  Catches bugs that the
  CUDA-vs-CUDA-padded tests wouldn't (e.g. a math error present in both
  varlen and non-varlen CUDA branches).

### Stride / layout

* **`test_checkpointing_ssu_noncontig`** (4 cases) — every batch-side
  tensor (`x`, `dt`, `B`, `C`, `z`, `out`, plus cache-side `old_x`,
  `old_B`, `old_dt_proc`, `old_cumAdt`) gets a different outer-dim
  padding → distinct non-default strides.  Bit-exact comparison
  against the contiguous-clone reference.  Sweeps `{varlen, non-varlen}
  × {bf16, int8}`.
* **`test_checkpointing_ssu_d_split2`** (1 case) — smoke test for the
  `D_SPLIT=2` D-output-split path.
* **`test_checkpointing_ssu_heads_per_group`** (1 case) — smoke test
  with `HPG = nheads/ngroups = 8` to cover the multi-group routing
  (all other tests trivially hit `group_idx = 0` with HPG=16).
* **`test_checkpointing_ssu_contiguous`** (1 case) — sanity that fully
  contiguous inputs work end-to-end.

### Boundary / input validation

* **`test_checkpointing_ssu_rejects_large_T`** (3 cases) — wrapper
  rejects `T > MAX_WINDOW_CAP = 16` with a clear error.
* **`test_checkpointing_ssu_int8_smoke`** — minimal end-to-end smoke for
  int8 state.

### Triton-reference tests (no CUDA JIT)

* **`test_checkpointing_state_update`** (30 cases) — explicit-tuple
  parametrization (was a 288-element cartesian; trimmed to 30
  representatives covering each dtype × T-bucket × write_ckpt × paged
corner). Cross-validates the merged Triton `checkpointing_state_update`
  against the upstream `selective_state_update` reference.
* **`test_checkpointing_state_update_philox`** (12 cases) — Philox
  variant of the above.
* **`test_checkpointing_philox_rounding_unbiased`** (4 cases) —
statistical
  SR unbiasedness on the Triton reference.

### How to run

```bash
# Full checkpointing_ssu suite (138 cases)
uv run pytest tests/mamba/test_checkpointing_ssu.py -v

# Just the CUDA kernel paths (varlen + non-varlen)
uv run pytest tests/mamba/test_checkpointing_ssu.py -v \
    -k "not (state_update and not _ssu_)"
```

Cold-rebuild estimate: ~5-10 min for the unique JIT keys
(`state_dtype × NPREDICTED × MAX_WINDOW × philox_rounds ×
heads_per_group`).
Warm-cache run: ~3-5 min.

## Reviewer Notes

### Architecture support

The JIT module gen (`flashinfer/jit/mamba/checkpointing_ssu.py`) sets a
single `-gencode` arch list: SM 80+ (Ampere → Blackwell).  No
dtype-specific subset — every fp32-to-narrow cvt the kernel uses has
both a hardware PTX path on the relevant architecture and a software
fallback for older arches.

| Op | Hardware path | Software fallback |

|-----------------------------------|--------------------------------------------|--------------------------------------------------|
| fp32 → fp16 SR (`cvt_rs_f16_f32`) | `cvt.rs.f16x2.f32`, sm_100a+ |
`cvt_rs_f16_sw` (`conversion.cuh:148`, bit-exact) |
| fp32 → fp8 e4m3 SR (`cvt_rs_e4m3x4_f32`) |
`cvt.rs.satfinite.e4m3x4.f32`, sm_100a+ | `cvt_rs_e4m3_sw`
(`conversion.cuh:270`, bit-exact) |
| fp32 → fp8 e4m3 RN (`__nv_fp8_e4m3(x)`) |
`cvt.rn.satfinite.e4m3x2.f32`, sm_89+ | cuda_fp8 lib SW path
(`cuda_fp8.hpp:263+`) |
| fp32 → int8 RN (`cvt_rni_sat_s8`) | `cvt.rni.sat.s8.f32`, SM 80+ |
`__float2int_rn` + clamp |
| fp32 → int8 / int16 SR | (no HW PTX op) | uniform-noise + libdevice
floor (runs anywhere) |

Every dtype × rounding-mode the kernel supports runs on any SM80+ arch.
The HW PTX paths are performance optimizations on Blackwell / Ada, not
correctness gates.  Bit-exact SW fallbacks for fp16 / fp8 SR are
validated against hardware on sm_100a.

The kernel source uses `#ifdef __CUDA_ARCH__ >= 1000 &&
__CUDA_ARCH_FEAT_SM100_ALL` guards in `conversion.cuh` to pick between
HW and SW at compile time, so nvcc only emits the HW PTX where it's
legal.

Perf target: SM100 (B200).  Correctness target: SM80+.

### Other notes

* `HEADS_PER_GROUP` is JIT-stamped (one specialization per `.so`,
  vs. the 7-way `dispatchRatio` in the original design) — cuts compile
  time ~7× per `.so`.
* Kernel-launch overhead: 1 launch (was 2 in the Triton precompute+main
  reference) — saves ~1-2 µs out of a ~10-30 µs end-to-end budget.
* No gmem round-trip for `CB_scaled` / `decay_vec` (kept in smem).
* Branch B (no-checkpoint) saves the state HBM write entirely on most
steps under typical workload patterns (`prev_k + seq_len ≤ MAX_WINDOW`).
* Quantized state (int8 / fp8) saves ~4× state HBM bandwidth vs fp32.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Checkpointing SSU: JIT-backed kernels, Torch custom-op, runtime
dispatch for checkpoint/replay selective state updates;
varlen/fixed-length and 8-bit/FP8 paths.
* Benchmarks: new benchmark comparing incremental/checkpointing backends
and runtime GPU bandwidth detection.
* Quantization: expanded int8/FP8 e4m3 support with stochastic rounding
and extended Philox-offset handling.

* **Tests**
* New Philox-offset and FP8 stochastic-rounding reference tests
(hardware + SW fallback).

* **Documentation**
  * Updated CODEOWNERS for MAMBA-related kernel paths.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3324)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>