mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU by ishovkun · Pull Request #3975 · flashinfer-ai/flashinfer

ishovkun · 2026-07-14T19:28:37Z

📌 Description

Adds two things to the checkpointing-SSU (Mamba-2 MTP-replay) kernel.

1. Two-kernel split. An alternative to the single monolithic kernel that wins
at larger batch. A precompute kernel produces the conv1d coefficients (the
C·B contraction and the cumulative-A·dt decays); a persistent, grid-stride
main kernel does the state replay + output projection, consuming the
precompute's outputs through caller-provided scratch
(cb_scaled / cumAdt_vec / cb_old / cumAdt_old — graph-safe, allocated like
out). The public API auto-dispatches on batch × nheads (monolith below
~1 work-unit/SM, split above); algorithm="monolith"|"two-kernel" forces either
path for benchmarking. The split covers 2- and 4-byte state (bf16/fp16/fp32);
8-bit state (int8/fp8) always runs the monolith (checkpointing_ssu_kernel_8bit)
— its quantize + per-block-scale flow doesn't factor cleanly into precompute→main.

2. Ring-buffer cache contract (breaking). The MTP-replay cache tensors change
from the double-buffered old_x / old_B / old_dt / old_cumAdt +
cache_buf_idx to single-buffered, head-major ring buffers:

tensor	shape	dtype
`x_cache`	`(state_cache_size, nheads, L, dim)`	input
`B_cache`	`(state_cache_size, ngroups, L, dstate)`	input
`dt_cache`	`(state_cache_size, nheads, L)`	f32
`ring_start`	`(state_cache_size,)`	i32

with L = max_window + npredicted. Live tokens are (ring_start + j) mod L for
j ∈ [0, pnat); appends land at (ring_start + pnat + i) mod L. The host owns
all ring bookkeeping (advancing ring_start on flush) — the kernels only read
ring_start / prev_num_accepted and append. No decay is cached: cumAdt is a
prefix sum and is not invariant under the ring-start advance a flush performs, so
it is recomputed from the dt ring; on the two-kernel path the precompute stages
it into the cumAdt_old scratch for the main. Matches the ReplaySSM contract.

🔍 Related Issues

Design follows the ReplaySSM cache contract from vLLM PR vllm-project/vllm#48018.

🚀 Pull Request Checklist

✅ Pre-commit Checks

I have installed pre-commit.
I have installed the hooks with pre-commit install.
The hooks (ruff / clang-format / mypy) pass on the changed files.

🧪 Tests

Tests added/updated — tests/mamba/test_checkpointing_ssu.py rewritten for
the ring contract: two-kernel-vs-monolith bit-exactness across write/no-write,
d_split, pipeline-stages, meta-ring refill, and per-state-dtype (f16/f32/mixed)
cases; ring wraparound; varlen; non-contiguous strides; determinism. All
validated against the Triton reference (replay_selective_state_update.py).
All tests passing — the two-kernel-vs-monolith parity battery + the broad
correctness subset pass on B200.

Reviewer Notes

Breaking change. The cache API changes shape and semantics (double-buffer →
ring). Callers must switch to x_cache / B_cache / dt_cache / ring_start and
take over ring_start advancement on flush. The custom-op mutates_args set
and the wrapper validation are updated to match.
8-bit scope. The two-kernel split is 2/4-byte-state only; 8-bit stays on the
monolith (see Description).
Perf (B200/B300, mixed-PNAT sweep with conv1d + PDL — the production shape):
the two-kernel split overtakes the monolith from ~256 work-units up. The ring
migration is roughly perf-neutral: the monolith is at parity (slightly faster
from the gather rewrite) and the two-kernel no-write path is at parity; the
large-batch two-kernel path carries a small floor (~1.5%) from recomputing the
decay the ring no longer caches, exposed to ~10% at one low-occupancy batch
(~32) — a precompute head-tiling tune can close that if needed.

Summary by CodeRabbit

New Features
- Added ring-buffer-based SSU checkpointing with monolithic and two-kernel execution modes.
- Added variable-length sequence support, optional causal convolution integration, and improved cache handling.
- Added Triton replay reference implementations and expanded CUDA/Triton compatibility options.
Benchmarks
- Added mixed-checkpointing analysis, PNAT sampling, CSV export, plotting, and benchmark collection tools.
Bug Fixes
- Improved GPU feature detection and stochastic-rounding support across supported hardware.
Tests
- Expanded coverage for replay modes, ring caches, varlen inputs, quantized states, determinism, and two-kernel execution.

Replace inline PTX `griddepcontrol.wait` / `griddepcontrol.launch_dependents` with `cudaGridDependencySynchronize()` / `cudaTriggerProgrammaticLaunchCompletion()` as required by CUDA PDL guidelines (see flashinfer#2558). Remove the now-unused `gdc_wait` / `gdc_launch_dependents` wrappers from `kernel_checkpointing_ssu_common.cuh`; call sites in the bf16 and 8-bit kernels use the CUDA device APIs directly inside their existing `if constexpr (ENABLE_PDL)` guards.

… reference - Add _persistent_main_impl, _persistent_rectangle_impl, and _persistent_main_kernel to the Triton checkpointing reference, matching the upstream TRT-LLM persistent architecture with double-buffered caches (old_x shape changes from 4D to 5D) and per-slot runtime dispatch - Add replay_selective_state_update wrapper that routes to the legacy checkpointing_state_update path for 4D old_x and to the new persistent IS_DYNAMIC path for 5D old_x - Fix continuous dA_cumsum semantics: on no-write steps with prev_k > 0, stored cumsum is now offset by the previous tail value so the buffer holds a monotonically increasing cumsum across back-to-back no-write steps - Update test_checkpointing_ssu.py to verify against analytically computed expected dA_cumsum values; mark CUDA comparisons as xfail for no-write + prev_k > 0 (CUDA does not yet store the continuous prefix)

The persistent path in tests/mamba/triton_reference/checkpointing_state_update.py dropped the IS_DYNAMIC runtime is_write resolution during its migration, so persistent_dynamic + rectangle_for_nowrite=False never wrote checkpoint state: replay still ran (output looked correct), only the HBM state store was compiled out, making the bug silent. Port the gating from the standalone reference: - add IS_DYNAMIC and WRITE_CHECKPOINT_IS_CONSTEXPR to _persistent_main_impl - resolve is_write per-slot at runtime when IS_DYNAMIC - gate write_offset / write_buf / state store / old_x write on is_write - thread the constexprs through both _persistent_main_kernel call sites Tests: add _old_x_to_5d helper and migrate _run_checkpointing_ssu_case to the 5D persistent path (CUDA kernel stays 4D). Also bundles the prior CUDA no-write continuous-dA_cumsum prefix fix (kernel_checkpointing_ssu.cuh) and its test updates (xfail removals + continuous-cumsum multistep test).

Port cu_seqlens varlen support into the persistent Triton path (_replay/_rectangle_precompute_impl, _dynamic_precompute_kernel, _persistent_main_impl, _persistent_rectangle_impl, _persistent_main_kernel, and the replay_selective_state_update wrapper). Varlen is batch-side packing only: bos = cu_seqlens[pid_b], outer offset switches to bos*stride_*_T, T-axis masks use the per-sequence seq_len; cache (state/old_*) indexing is unchanged. All varlen code is gated behind IS_VARLEN: tl.constexpr (derived by @triton.heuristics from cu_seqlens_ptr on the kernels, passed explicitly to the device-function impls), so the non-varlen path DCEs it entirely. The HAS_Z per-token output store is guarded with `t < seq_len` so packed neighbors aren't clobbered when seq_len < T. Tests: add _assert_old_x_5d_matches_4d helper; migrate the CUDA-vs-Triton tests (_run_checkpointing_ssu_case, max_window_gt_npredicted, philox_no/with_checkpoint, mixed_checkpoint_batch) and the varlen tests (_run_varlen_cuda_vs_triton) to the 5D persistent path (CUDA kernel stays 4D; old_x converted via _old_x_to_5d). Full tests/mamba/test_checkpointing_ssu.py suite: 156 passed.

…5D path Switch test_checkpointing_state_update, _philox, and _philox_rounding_unbiased from the old 4D checkpointing_state_update to the 5D replay_selective_state_update (old_x built via _old_x_to_5d), and drop the now-unused checkpointing_state_update import. test_checkpointing_state_update forced write_checkpoint as a constexpr, which the persistent-dynamic path cannot honor — it decides is_write = (prev_k+T) > max_window per slot at runtime. Drive the state/cache expectations from that runtime is_w instead of the forced flag, and extend the old_x cache postcondition to the 5D double-buffer layout (write buffer + untouched sibling buffer). Full tests/mamba/test_checkpointing_ssu.py suite: 156 passed.

…al/cuda-incr Add a `triton-replay` kernel that calls replay_selective_state_update on the 5D persistent path (single persistent_dynamic launch, per-slot runtime write decision), so it can be compared against the old `incremental` (4D checkpointing_state_update) and `cuda-incr` rows. Build a 5D double-buffered old_x once in build_kernel_inputs (active buffer = cache_buf_idx[slot]) so the 4D->5D conversion isn't timed. Gated by --triton-replay (default on). The old incremental row is left intact for comparison.

…he faithful TMA standalone The merged checkpointing_state_update copy dropped TMA + tuning during its migration (37 vs 277 TMA refs), so the persistent path was slow on Blackwell. Vendor a small import shim into the faithful standalone reference (replay_selective_state_update.py) — local get_sm_version + REPLAY_WORK_* column constants — so it is importable in the flashinfer test env without a tensorrt_llm dependency, and point the benchmarks at it: - bench_checkpointing_ssu.py: build a 5D double-buffered old_x + write-first replay_work_items once (untimed); add `triton-replay` (persistent_dynamic) and `triton-replay-pm` (persistent_main) rows, gated by --triton-replay[-pm]. - bench_ssu_checkpoint_mixed.py: add both persistent kernels to DEFAULT_KERNELS and a _PHILOX_KERNELS set so they honor stochastic rounding. Mixed f16-philox-5 sweep (CUPTI, median us): persistent_main scales smoothly and lands within ~1.1-1.5x of cuda-incr (2.4-3.2x faster than the old 4D Triton); persistent_dynamic is faster than the old path but has perf cliffs (e.g. batch 512). Standalone has TMA + tuning; merged copy does not.

…t test migration Phase A: mirror the merged-copy varlen port (commit 165f3a4) onto the TMA standalone replay_selective_state_update.py — cu_seqlens_ptr + IS_VARLEN through _replay/_rectangle_precompute_impl, _dynamic_precompute_kernel, _persistent_main_impl, _persistent_rectangle_impl, _persistent_main_kernel, and the wrapper (cu_seqlens/max_seqlen + per-launch cu_seqlens_ptr). Varlen is batch-side only (bos*stride_*_T on x/dt/B/C/out plain loads, seq_len masks); the TMA state paths, cache indexing, and cb_scaled/decay_vec stores are untouched. Tests (start of migration to the standalone): import it as `replay_persistent`, add _make_replay_work_items (write-first sorted n_writes/replay_work_items), and migrate _run_checkpointing_ssu_case (de-risk) + _run_varlen_cuda_vs_triton to it with mode="persistent_dynamic". Remaining tests still use the merged copy. Verified: 10/10 varlen + 3/3 de-risk on the standalone; full suite 156 passed.

…ma_pm replay configs Route every replay test through the _call_replay helper and the _TRITON_IMPLS list so each runs across non_tma (merged copy, persistent_dynamic), tma_pd (standalone TMA, persistent_dynamic), and tma_pm (standalone TMA, persistent_main). Validates the standalone TMA path alongside the merged kernel; _TRITON_IMPLS is the single point to drop non_tma when the merged kernel is removed.

- Capture per-kernel name + device start/end from the CUPTI Activity records we already pull (SSU_CUPTI_DEBUG dumps the conv1d->SSU timeline), enabling PDL-overlap analysis without nsys/sqlite. - Wire triton-replay (persistent_dynamic) and triton-replay-pm (persistent_main) rows to the standalone TMA kernel. - conv1d total-span timing (--with-conv1d) plus --no-cupti / --no-cuda-graph toggles for clean external-profiler (ncu) capture. - Add the standalone causal_conv1d_triton kernel and histogram_T6.csv, which the bench imports/reads (previously untracked).

… test checkpointing_ssu gains optional cb_scaled/decay_vec scratch tensors; their presence selects the (not-yet-implemented) two-kernel precompute+main path, else the monolithic kernel runs. Both-or-neither ValueError; the two-kernel launch raises NotImplementedError for now. Scratch is caller-provided (graph-safe, like out). Adds strict-xfail test_two_kernel_matches_monolithic (monolithic vs two-kernel equivalence on out/state/caches, bf16 write+nowrite).

- New kernel_checkpointing_ssu_precompute.cuh (WIP prototype, not yet wired into the JIT sources): per-group grid (batch, ngroups, head_tiles); raw C*B MMA once per group, reused across the head loop; stores CB_scaled in FRAGMENT-NATIVE layout [batch, nheads, lane, 8] = matmul-4 fragA via one STG.128/thread (no smem/swizzle either side; (t,j) computed on the fly from the register index). load/MMA/cache helpers marked TODO(iterate). - CheckpointingSsuParams: cb_scaled/decay_vec documented as bf16 fragA-native + f32 (b,h,NPREDICTED_PAD_MMA_M); appended at struct end (no layout shift). - T0 allocates the bf16 fragA-native cb_scaled. - .plans/ssu_split.md: fragA-native scratch contract throughout.

Replace the serial head loop with the warp-per-head structure: 2-warp raw C*B MMA -> swizzled smem (once per group), then a uniform NUM_ITER loop (= ceil(HEADS_PER_CTA/NUM_WARPS)) so the iter==0 __syncthreads() is reached by every warp exactly once (deadlock-proof, no assert), with the barrier deferred past the per-warp cumAdt compute to overlap the raw-CB smem store. CUTLASS helper extraction still TODO(iterate).

Precompute uses a lean CheckpointingSsuPrecomputeStorage (B/C swizzled like the monolithic + fp32 CB + per-warp cumAdt/dt_proc/decay), dropping the monolithic's state/old_x/old_B/x/z -> ~3x less smem -> more CTAs/SM. Raw C*B is kept FP32 in smem and scaled in fp32 (cast to bf16 only at the gmem STG.128), matching the monolithic's precision (storing raw as bf16 would round before scaling). scale_store_cb_gmem reads CB fp32 from smem (per-warp coeffs). CB smem layout index is a flagged placeholder until compute_cb_2warp fixes the CuTe store.

compute_cb_2warp = the MMA half of compute_CB_scaled_2warp (_1x1 m16n8k16 tiled-mma + DSTATE/16 K-loop over swizzled smem.C/smem.B) with the scale/mask epilogue dropped: stores the RAW fp32 accumulator to ROW-MAJOR smem.CB (per-head decay/dt scaling + causal mask deferred to scale_store_cb_gmem). The row-major store makes scale_store's t*S+j read address the identical (t,j), resolving the CB layout placeholder. Wired into the kernel (warps 0/1; warps 2/3 idle, then pick up heads). Remaining: load_group_BC, load_dt_proc + compute_cumAdt_pw, decay_vec + old_dt/old_cumAdt cache writes.

The raw fp32 CB was hand-rolled row-major; even non-LDSM (t,j) reads of a 16-wide fp32 smem buffer bank-conflict. Switch to SmemSwizzle<float> (the f32 specialization, ATOM_COLS=32) — exactly how the monolithic swizzles its bf16 CB — so the accumulator store (compute_cb_2warp) and the per-warp reads (scale_store_cb_gmem) are bank-conflict-free. Both sides go through the same make_swizzled_layout_rc<float, M, M, CB_STRIDE>, so they address identical (t,j). CB smem grows 1 KB -> 2 KB (pad to the f32 swizzle atom).

Wire the three input helpers into the precompute kernel: - load_group_BC: cp.async this group's C/B (conv1d outputs) into swizzled smem, warps 0/1 only (register pressure; mirrors load_post_pdl_wait_data minus x). - load_dt_pw (C1): scalar per-lane strided LDG of dt + softplus -> per-warp smem. - compute_cumAdt_pw (C2): per-warp Hillis-Steele scan -> cumAdt/decay. Loop: dt+cumAdt -> deferred barrier (iter0 __syncthreads / else __syncwarp) -> decay_vec store + scale_store_cb_gmem (C5). Trailing __syncwarp closes the WAR on the single-slot per-warp coeffs (double-buffer would remove it). Remaining: old_dt/old_cumAdt cache writes (C7), then S4 main + S5 launcher to test. Still a prototype (not yet in the JIT sources).

Rename load_dt_pw -> load_dt (no non-pw variant). Make dt_softplus a compile- time template bool (if constexpr) on load_dt + the precompute kernel, instead of a per-lane runtime if(params.dt_softplus): the launcher reads params.dt_softplus once and dispatches the matching instantiation, so the softplus branch folds away. (2 kernel instantiations, like the existing VARLEN bool.)

No-write path now uses all 4 warps: W0/1 -> C.B (C5, new) -> smem.CB, W2/3 -> C.old_B (C6, old) -> smem.CB_old, in parallel (mirrors the monolithic's compute_CB_scaled_2warp || compute_CB_old_2warp). The write path skips CB_old (W2/3 idle) since the main folds old tokens into state via the replay. must_checkpoint is uniform per CTA -> divergence-free branch. - SmemT: + CB_old (f32 swizzled) + old_B (swizzled) + MAX_WINDOW template arg. - load_group_BC: load C on ALL warps (it's the A-operand of both MMAs now), B only on W0/1. - load_old_B: cp.async the buffered B (cache) for W2/3, no-write only. - compute_cb_old_2warp: C.old_B MMA mirror of compute_cb_2warp (N-operand old_B, output CB_old). - params: + cb_old gmem hook (fragA-native, written only on no-write). Next: scale_store_cb_old (C6 decay/coeff scaling -> cb_old gmem) + old_dt/ old_cumAdt load for coeff. Still a prototype (not in the JIT sources).

The trailing WAR __syncwarp was standalone after the has_head block, so a headless warp ran it back-to-back with the (A) sync for nothing. has_head is warp-uniform (h = warp + iter*NUM_WARPS, no lane) and monotonic in iter, so a warp that drops out never writes again -> no WAR to guard. Put the syncwarp right after scale_store inside the block: headless warps skip it, head-carrying warps still order this iter's cross-lane reads before next iter's writes.

The WAR sync was after the scale_store_cb_gmem CALL (+ a comment) in the loop — not after the actual read. The cross-lane smem.cumAdt/dt reads live inside the function, so the syncwarp belongs there, right after the read loop and before the STG.128 (so the gmem store overlaps the next head's scan). Drop the standalone syncwarp from the kernel loop; note C7's cache writes must read smem.cumAdt/dt before scale_store (which now ends with that WAR sync).

Completes the precompute's no-write path and cache tape: - scale_store_cb_old: C6 scaling exp(cumAdt[t])*coeff[i] (coeff = exp(total_old - old_cumAdt[i])*old_dt[i], mask i<prev_k, no causal) -> cb_old, fragA-native STG.128 + internal WAR syncwarp. Validated vs monolithic compute_CB_old. - Wire into head loop: load_old_dt_cumAdt (write block) + scale_store_cb_old (read block), both !must_checkpoint. SmemT gains per-warp old_dt/old_cumAdt. - C7: persist this head's new dt/cumAdt into the old_dt/old_cumAdt tape at write_offset (cumAdt continuity offset on no-write), mirroring the monolithic kernel_*_ssu.cuh:1112. Runs before the stores so their WAR syncwarp also closes the cross-lane old_cumAdt[prev_k-1] read. - Dedup: extract load_old_dt_cumAdt into common.cuh; the monolithic's two inline copies (load_pre_pdl_wait_data, load_data) now call it (forceinline -> identical codegen). Verified byte-identical via diff; single test passes in isolation. Precompute now does C1/C2/C5/C6/C7 + decay_vec. Still a prototype (not in JIT).

The two-kernel cache-write split is forced by data locality: only the precompute holds smem.B (the main reads cb_scaled and never loads B), so the precompute owns the old_B writeback too — alongside old_dt/old_cumAdt (C7). The main owns old_x (it loads x) and state. - store_old_B on W0/1, hoisted before the CB MMA so the STG overlaps it (like the monolithic); self-gates on head%HEADS_PER_GROUP via first_head. - Drop the now-used (void)buf_write/write_offset casts. - Update the file header to reflect C1/C2/C5/C6 + the cache-write ownership. Precompute is now functionally complete (produces cb_scaled, cb_old, decay_vec, old_dt, old_cumAdt, old_B). Still a prototype (not in the JIT sources).

Store raw cumAdt[t] in the scratch instead of exp(cumAdt[t]). The main loads it into smem.cumAdt and the monolithic output path's existing exp(smem.cumAdt) epilogue + the no-write exp(total_old_cumAdt) beta-factor (from cache old_cumAdt) compute the decay on the fly — so when ssu_checkpoint/ssu_nocheckpoint are parameterized (READ_PRECOMPUTED_CB) only the two frag_CB loads need a branch; decay/beta/replay/matmul-3 stay byte-identical. This is the clean (a) realization. - Precompute stores smem.cumAdt[warp] (raw); drop the unused smem.decay slot + its write in compute_cumAdt_pw. - Rename decay_vec -> cumAdt_vec in the param struct, precompute, Python wrapper (arg + docstring + ValueError), and the T0 test. triton_reference untouched (its decay_vec is the upstream baseline's own scratch).

…cb_old ssu_checkpoint/nocheckpoint's output functions now take a defaulted READ_PRECOMPUTED_CB template + cb_gmem_head/cb_old_gmem_head args, so the monolithic call sites are unchanged (if constexpr(false) -> original LDSM, byte-identical codegen) and the two-kernel main can reuse them with the CB loaded from gmem instead of smem. - load_cb_fragA<REGS>: one vectorized LDG of a lane's REGS-element fragA chunk (REGS = M*K/32 = K/2, M=16: 8 @ K16 / 4 @ K8) — main's substitute for the LDSM. - compute_and_store_output / compute_no_write_output: branch only the frag_CB (+ frag_CB_old) load on READ_PRECOMPUTED_CB; decay/beta/replay/matmul-3 untouched (raw cumAdt in smem.cumAdt does it). - K-aware cb_old: matmul-4-old is m16n8k{MAX_WINDOW_PAD_MMA_K}, so scale_store_cb_old emits MAX_WINDOW_PAD_MMA_K/2 regs (the [0,8)-col prefix of the m16n8k16 mapping; STG.64 @ K8 / STG.128 @ K16); main LDGs the matching count. Still a prototype path (main kernel + launcher pending); monolithic unaffected.

…point The per-path dispatchers gain the defaulted READ_PRECOMPUTED_CB template + cb_gmem_head (ssu_checkpoint) / cb_gmem_head+cb_old_gmem_head (ssu_nocheckpoint), forwarded to compute_and_store_output / compute_no_write_output. Monolithic call sites in checkpointing_ssu_kernel are unchanged (defaults -> false/nullptr). The two-kernel main will call these with READ_PRECOMPUTED_CB=true + the gmem CB pointers. Completes the output-path reuse foundation; main kernel pending.

New checkpointing_ssu_main_kernel + load_main_data (new file). Grid (D_SPLIT, batch, nheads); reuses the monolithic ssu_checkpoint/ssu_nocheckpoint with READ_PRECOMPUTED_CB=true (matmul-4 CB LDG'd from gmem cb_scaled/cb_old; beta from cumAdt_vec -> smem.cumAdt, exp'd by the existing epilogue). - load_main_data: subset of load_data composed from the shared blocks (load_state_*, load_tile_async for C/x/old_x/old_B/z, load_old_dt_cumAdt) minus new-B + dt/C1/C2, plus cumAdt_vec -> smem.cumAdt. Composable for a future persistent variant. - cb_scaled/cb_old gmem offsets (REGS = K/2 per lane) match the precompute's scale_store_* writes; cumAdt_vec index matches the precompute's store. - Cache: writes old_x + state (replay C8); precompute owns old_B/old_dt/old_cumAdt. - PDL first cut: gdc_wait before the load (no-op w/o a programmatic parent, T0); replay-overlap (wait between matmul-3/4) is the S9 refinement. Remaining to make it run: S5 (PDL-chain launcher + cb_old/cumAdt_vec plumbing + drop the NotImplementedError stub) + S6 (JIT sources). Still a prototype.

When the caller provides scratch (cb_scaled/cumAdt_vec/cb_old), launch precompute -> main instead of the monolithic: - launch_checkpointing_ssu.cuh: two-kernel branch in launchCheckpointingSsuImpl (precompute grid (batch,ngroups,1) + main grid (D_SPLIT,batch,nheads), both with the programmatic-serialization attr; bf16/fp16 only). DT_SOFTPLUS dispatched on the runtime flag. First cut has no early-trigger -> the main's gdc_wait resolves at precompute completion (correct; overlap is S9). - csrc binding + launcher.cu: thread cb_scaled/cumAdt_vec/cb_old (Optional) into the signature + fill params. - Python: thread through _checkpointing_ssu (custom op + fake), add to mutates_args (precompute writes them), public wrapper gains cb_old + requires all three together, drop the NotImplementedError stub. New kernels compile transitively via the headers (no new JIT sources).

First compile + run of the precompute + main kernels via the two-kernel path. Bit-exact match to the monolithic on out / state / old_x / old_B / old_dt / old_cumAdt across a nowrite case (k=0) and a write case (k=T). - load_cb_fragA: read the gmem fragA bytes AS the fragment's element type (MMA operand type, e.g. cutlass bf16) instead of input_t — the assignment was ill-typed (nv_bfloat16 vs the operand type). Sole compile error from the first build. - T0 test: allocate cb_old (m16n8k{K_old} fragA, REGS = K_old/2), drop the strict-xfail marker — the two-kernel path now matches the monolithic. The entire two-kernel split (precompute + main + READ_PRECOMPUTED_CB reuse + launcher + binding/Python plumbing) is functional and validated.

…ithic (T0)

ishovkun · 2026-07-16T04:26:27Z

/bot run tests/mamba

flashinfer-bot · 2026-07-16T04:27:22Z

GitLab MR !979 has been created, and the CI pipeline #58224873 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-07-16T06:39:58Z

[FAILED] Pipeline #58224873: 10/20 passed

Replace the session-scoped autouse fixture with a `pytest_collection_modifyitems` hook that batch-precompiles only the missing `checkpointing_ssu` JIT variants, skipping the prewarm when fewer than 8 matching tests are collected. Also switch the variant table from abbreviated string tuples to explicit `torch.dtype` values matching the actual `_get_module` call signature, eliminating the local dtype-mapping dict.

ishovkun · 2026-07-16T17:25:30Z

/bot run tests/mamba

flashinfer-bot · 2026-07-16T17:26:13Z

GitLab MR !979 has been updated with latest changes, and the CI pipeline #58318552 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-07-17T06:41:44Z

[FAILED] Pipeline #58318552: 13/20 passed

The conftest prewarm expanded 14-element tuples into gen_checkpointing_ssu_module positionally, but the signature has 15 parameters: every value after heads_per_group landed one slot to the left (0 -> num_groups, False -> philox_rounds). The rendered config then contained 'constexpr int PHILOX_ROUNDS = False;' which fails to compile, so on a cold JIT cache the prewarm built nothing and every CI runner fell back to serial first-touch compiles. Call the generator with keyword arguments via a strict zip against an explicit field-name tuple so an arity mismatch skips the prewarm loudly instead of building garbage URIs, and regenerate the variant list by re-instrumenting _get_module over a full suite run (57 variants; the old hand-maintained list of 43 was also stale). AI-assisted.

ishovkun · 2026-07-17T22:03:42Z

/bot run tests/mamba

flashinfer-bot · 2026-07-17T22:04:03Z

GitLab MR !979 has been updated with latest changes, and the CI pipeline #58505240 is currently running. I'll report back once the pipeline job completes.

jimmyzho · 2026-07-17T22:55:02Z

Merge pending CI tests passing.

flashinfer-bot · 2026-07-17T23:39:20Z

[FAILED] Pipeline #58505240: 9/20 passed

ishovkun · 2026-07-18T04:11:41Z

/bot run tests/mamba

flashinfer-bot · 2026-07-18T04:12:32Z

GitLab MR !979 has been created, and the CI pipeline #58537967 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-07-18T04:48:43Z

[FAILED] Pipeline #58537967: 9/20 passed

The two-kernel split launches the persistent main with internal PDL always on, and the precompute fires its launch-completion trigger at the top — main CTAs can be resident while the precompute is still running. Everything the main touches before cudaGridDependencySynchronize() must therefore be previous-step data, but the tile-0 (prologue) prefetch read the cumAdt_old scratch — old-token decay rows on the checkpoint path and the beta tail on the no-write path — which the precompute writes THIS step, and the replay consumed it before the gdc. On an idle GPU the precompute owns all SMs so the race is always won; under SM contention the grids interleave and the main reads stale coefficients, corrupting replay decay and beta scaling. This is the flaky two-kernel out-mismatch on the shared RTX 5090 CI runner (reproducible locally in seconds by running any GPU load concurrently with the tests; monolith unaffected, k=0 cases unaffected — both have no scratch read). Move the scratch reads into load_old_coeff, issued from prefetch_async_post_gdc, and reorder the prologue so the gdc precedes the post-gdc prefetch and the replay. The pre-gdc cp.asyncs (state + ring tiles, previous-step rows only) still overlap the gdc wait, and the steady-state loop is unchanged (its prefetches already ran post-gdc; the tile-0-output wait_prior becomes a no-op backstop). Verified: burner-loop repro fails 8 tests before, 0 after (5/5 runs); full test_checkpointing_ssu.py 223 passed idle. AI-assisted.

ishovkun · 2026-07-21T04:28:08Z

/bot run tests/mamba

flashinfer-bot · 2026-07-21T04:28:40Z

GitLab MR !979 has been updated with latest changes, and the CI pipeline #58871658 is currently running. I'll report back once the pipeline job completes.

The main kernel now recomputes old-token decay in-registers via `warp_scan_old_cumAdt` instead of reading from a precompute-staged gmem scratch buffer. This eliminates the `cumAdt_old` allocation from the caller-visible API, the params struct, validation, and all benchmarks/tests.

ishovkun · 2026-07-21T05:21:43Z

/bot run tests/mamba

flashinfer-bot · 2026-07-21T05:22:07Z

GitLab MR !979 has been updated with latest changes, and the CI pipeline #58876328 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-07-21T10:48:08Z

[FAILED] Pipeline #58876328: 12/20 passed

Revert #3975 for api-breakage. #4129 is the associated doc change which is no longer relevant.  ## 📌 Description  ## 🔍 Related Issues  ## 🚀 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 - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] 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 - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes

Mirrors the v0.6.16 revert (#4246, commit 0a27ba3) so that 0.6.17 does not ship the API-breaking ring-buffer cache contract that 0.6.16 left out. Reverts: afd4754 mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU (#3975) f90e9c4 docs(mamba): document checkpointing varlen arguments (#4129) Unlike release-v0.6.16, this branch also carries #4129, so both are reverted here; 0.6.16 only needed #3975. Verified no collateral damage: the SM107 changes from #4280 in tests/mamba/conftest.py and the #4029 changes in flashinfer/utils.py are preserved. NOTE: like the 0.6.16 revert, this also reverts the is_cvt_rs_supported correctness fix that rode along in #3975 (back to `major in (10, 11)`). See the release notes discussion -- that hunk is a candidate to keep. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Mirrors the v0.6.16 revert (#4246, commit 0a27ba3) and its v0.6.17 counterpart (09fd5fc) so that 0.6.18 does not ship the API-breaking ring-buffer cache contract that neither 0.6.16 nor 0.6.17 shipped. Reverts: afd4754 mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU (#3975) f90e9c4 docs(mamba): document checkpointing varlen arguments (#4129) Like release-v0.6.17, this branch carries #4129 as well, so both are reverted; 0.6.16 only needed #3975. main still carries both. Verified no collateral damage: the SM107 change from #4280 in tests/mamba/conftest.py and the #4029/#4078 changes in flashinfer/utils.py are preserved, and the five core reverted files now byte-match release-v0.6.17. NOTE: like both earlier reverts, this also reverts the cvt_rs fix that rode along in #3975 -- is_cvt_rs_supported goes back to `major in (10, 11)` (wrong for SM110a) and the CUDA guard back to SM100_ALL only (B300/sm_103a falls to software emulation). That matches what 0.6.17 shipped, but the hunk remains a candidate to keep.

ishovkun added 30 commits June 8, 2026 10:41

docs(ssu): drop PROTOTYPE header — precompute/main validated vs monol…

fe90bb7

…ithic (T0)

jimmyzho approved these changes Jul 17, 2026

View reviewed changes

askliar mentioned this pull request Jul 20, 2026

[RFC]: ReplaySSM: cache SSM inputs instead of state for faster standard and speculative decode (Mamba2 + GDN) vllm-project/vllm#49232

Open

1 task

jimmyzho merged commit afd4754 into flashinfer-ai:main Jul 21, 2026
30 of 37 checks passed

kangbintNV mentioned this pull request Jul 28, 2026

docs(mamba): document checkpointing varlen arguments #4129

Merged

jimmyzho mentioned this pull request Jul 29, 2026

revert: #3975 and #4129 (mamba checkpointing ssu 2-split + doc) #4246

Merged

5 tasks

Uh oh!

Conversation

ishovkun commented Jul 14, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

ishovkun commented Jul 16, 2026

Uh oh!

flashinfer-bot commented Jul 16, 2026

Uh oh!

flashinfer-bot commented Jul 16, 2026

Uh oh!

ishovkun commented Jul 16, 2026

Uh oh!

flashinfer-bot commented Jul 16, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

ishovkun commented Jul 17, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

jimmyzho commented Jul 17, 2026

Uh oh!

flashinfer-bot commented Jul 17, 2026

Uh oh!

ishovkun commented Jul 18, 2026

Uh oh!

flashinfer-bot commented Jul 18, 2026

Uh oh!

flashinfer-bot commented Jul 18, 2026

Uh oh!

ishovkun commented Jul 21, 2026

Uh oh!

flashinfer-bot commented Jul 21, 2026

Uh oh!

ishovkun commented Jul 21, 2026

Uh oh!

flashinfer-bot commented Jul 21, 2026

Uh oh!

flashinfer-bot commented Jul 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ishovkun commented Jul 14, 2026 •

edited by coderabbitai Bot

Loading