flashinfer.mamba.checkpointing_ssu

flashinfer.mamba.checkpointing_ssu(state: Tensor, x_cache: Tensor, B_cache: Tensor, dt_cache: Tensor, ring_start: Tensor, prev_num_accepted_tokens: Tensor, x: Tensor, dt: Tensor, A: Tensor, B: Tensor, C: Tensor, out: Tensor, D: Tensor | None = None, z: Tensor | None = None, dt_bias: Tensor | None = None, dt_softplus: bool = False, state_batch_indices: Tensor | None = None, pad_slot_id: int = -1, state_scale: Tensor | None = None, rand_seed: Tensor | None = None, philox_rounds: int = 10, d_split: int | None = None, cu_seqlens: Tensor | None = None, max_seqlen: int | None = None, enable_pdl: bool = False, cb_scaled: Tensor | None = None, cumAdt_vec: Tensor | None = None, cb_old: Tensor | None = None, precompute_heads_per_cta: int = 0, algorithm: str = 'auto') → Tensor

Checkpointing SSU with MTP replay using matmul-based parallel token processing.

Parameters:

  • state (torch.Tensor) – SSM state, shape (state_cache_size, nheads, dim, dstate). Updated in-place.

  • x_cache (torch.Tensor) – Ring of cached x, shape (state_cache_size, nheads, RING_BUFFER_LEN, dim). RING_BUFFER_LEN is implicit (= size(2)); the LOGICAL replay window is max_window = RING_BUFFER_LEN - T (flush rule pnat + 2T > RING_BUFFER_LEN).

  • B_cache (torch.Tensor) – Ring of cached B, shape (state_cache_size, ngroups, RING_BUFFER_LEN, dstate).

  • dt_cache (torch.Tensor) – Ring of cached processed dt, shape (state_cache_size, nheads, RING_BUFFER_LEN), f32. Replay decays are recomputed from it (no cumAdt is cached — prefix sums are not ring-shift-invariant).

  • ring_start (torch.Tensor) – Ring head per slot (oldest live row), shape (state_cache_size,), int32. The HOST owns bookkeeping: advance by the replayed count on flush.

  • prev_num_accepted_tokens (torch.Tensor) – Number of old tokens to replay, shape (state_cache_size,), int32.

  • x (torch.Tensor) – New token inputs, shape (batch, T, nheads, dim).

  • dt (torch.Tensor) – Delta time, shape (batch, T, nheads, dim) with tie_hdim (stride[-1]=0). Accepted in native dtype (e.g. bf16) — converted to f32 internally.

  • A (torch.Tensor) – Decay rate, shape (nheads, dim, dstate) with tie_hdim.

  • B (torch.Tensor) – Input projection, shape (batch, T, ngroups, dstate).

  • C (torch.Tensor) – Output projection, shape (batch, T, ngroups, dstate).

  • out (torch.Tensor) – Preallocated output, shape (batch, T, nheads, dim).

  • D (Optional[torch.Tensor]) – Skip connection, shape (nheads, dim).

  • z (Optional[torch.Tensor]) – Gate, shape (batch, T, nheads, dim).

  • dt_bias (Optional[torch.Tensor]) – Bias added to dt, shape (nheads, dim) with tie_hdim.

  • dt_softplus (bool) – Whether to apply softplus to dt.

  • state_batch_indices (Optional[torch.Tensor]) – Maps batch index to cache slot, shape (batch,), int32 | int64.

  • pad_slot_id (int) – Sentinel value for padded entries.

  • state_scale (Optional[torch.Tensor]) – Block-scale decode factors for quantized state, shape (state_cache_size, nheads, dim), f32.

  • rand_seed (Optional[torch.Tensor]) – Single-element int64 CUDA tensor for stochastic rounding seed.

  • philox_rounds (int) – Philox PRNG rounds for stochastic rounding (default 10).

  • d_split (Optional[int]) – Per-head DIM split factor. This is only exposed for benchmarking. Do not use it cause it will make things slow.

  • cu_seqlens (Optional[torch.Tensor]) – Cumulative sequence boundaries for packed variable-length input, as a one-dimensional int32 CUDA tensor of shape (batch + 1,). It must start at 0, be monotonically nondecreasing, and end at total_tokens. When provided, x must have shape (1, total_tokens, nheads, dim) and max_seqlen is required.

  • max_seqlen (Optional[int]) – Upper bound on every packed sequence length. In variable-length mode this is the JIT-specialized predicted-token count used to derive the logical replay window from the ring-buffer size. Must be None when cu_seqlens is not provided.

  • precompute_heads_per_cta (int) – Two-kernel PRECOMPUTE head-tiling: heads per precompute CTA. 0 (default) uses the launcher’s co-residency heuristic; >0 overrides it (must divide nheads/ngroups, snapped to the HEADS_PER_GROUP>>k chain). Tuning knob — two-kernel path only.

  • algorithm (str) – Kernel selection: "auto" (default), "monolith", or "two-kernel". "auto" runs the two-kernel split iff the scratch quartet is provided AND batch * nheads >= sm_count. The crossover collapses in batch * nheads across nheads (measured at TP 8/4/2) and state widths: the monolith wins <= 128 work-units and the split wins from 256 on a 148-SM B200 (mixed-PNAT + conv1d/PDL bench; ties take the split). "two-kernel" forces the split (scratch quartet required); "monolith" forces the monolith (scratch ignored). Benches/tests that must pin the path should force it.

  • enable_pdl (bool) – When True the kernel is launched with cudaLaunchAttributeProgrammaticStreamSerialization, enabling the in-kernel griddepcontrol.{wait,launch_dependents} PTX to gate on the upstream (e.g. conv1d) and signal the downstream kernel. Caller’s responsibility: upstream/downstream kernels must also be PDL-paired for the wait/signal to have effect. Defaults to False.

  • cb_scaled (Optional[torch.Tensor]) – Pre-allocated input-dtype (same as x) scratch for the precomputed new-token CB matrix, fragment-native layout (batch, nheads, WARP_SIZE, MMA_FRAG_SIZE) — each (batch, head)’s CB is one m16n8k16 MMA A-fragment stored as [warp lane, register]. Providing it (together with cumAdt_vec / cb_old) makes the two-kernel (precompute + main) path available — algorithm decides whether it runs; leaving all four None always runs the monolithic kernel. Caller-allocated so the path is CUDA-graph-safe (no in-wrapper allocation, like out).

  • cumAdt_vec (Optional[torch.Tensor]) – Pre-allocated fp32 scratch for the per-head raw cumAdt vector, shape (batch, nheads, T_pad); the main kernel exponentiates it on the fly to get the decay/β factor. Must be provided iff cb_scaled is.

  • cb_old (Optional[torch.Tensor]) – Pre-allocated input-dtype (same as x) scratch for the precomputed old-token CB matrix, fragment-native layout (batch, nheads, WARP_SIZE, K_old // 2) where K_old = next_multiple_of_8(max_window) — the m16n8k{K_old} MMA A-fragment consumed on the no-write (replay) path, stored as [warp lane, register]. Must be provided iff cb_scaled is.

Returns:

out – Output tensor, shape (batch, T, nheads, dim).

Return type:

torch.Tensor