flashinfer.gdn_decode.gated_delta_rule_replayssm_commit

flashinfer.gdn_decode.gated_delta_rule_replayssm_commit(checkpoint_state: Tensor, rawv_cache: Tensor, rawk_cache: Tensor, g_cache: Tensor, beta_cache: Tensor, state_indices: Tensor, accept_lens: Tensor, *, track_state_indices: Tensor | None = None, track_steps: Tensor | None = None, use_qk_l2norm: bool = True, null_block_id: int = -1, backend: Literal['auto', 'simt', 'tcgen05'] = 'auto') → None

Commit an accepted ReplaySSM prefix into all layers’ checkpoints.

Call gated_delta_rule_mtp() with cache_replayssm=True for each layer first. Verify leaves checkpoints frozen. After acceptance is known, call this function once, then overwrite the windows during the next verify. The cache is not shifted or cleared; unaccepted tokens are discarded.

Parameters:

  • checkpoint_state (torch.Tensor) – Contiguous FP32 [layers, slots, HV, 128, 128] state, K-last. Updated in place. All tensors must be on the same SM100/SM103 device.

  • rawv_cache (torch.Tensor) – Contiguous BF16 [layers, slots, HV, T, 128] raw V windows.

  • rawk_cache (torch.Tensor) – Contiguous BF16 [layers, slots, H, T, 128] pre-normalization K.

  • g_cache (torch.Tensor) – Contiguous FP32 [layers, slots, HV, T] log-decay windows.

  • beta_cache (torch.Tensor) – Contiguous FP32 [layers, slots, HV, T] sigmoid-beta windows.

  • state_indices (torch.Tensor) – Contiguous int32 [B] pool slots. Live slots must be unique and in range. Slots <= null_block_id are ignored.

  • accept_lens (torch.Tensor) – Contiguous int32 [B] accepted token counts in [0, T]. Zero is a no-op, including for optional track writes.

  • track_state_indices (torch.Tensor, optional) – Contiguous int32 [B] destinations for an additional checkpoint. Live track slots must be in range, mutually distinct, and disjoint from every live state_indices slot. Paired with track_steps.

  • track_steps (torch.Tensor, optional) – Contiguous int32 [B] zero-based token indices to save. A step outside the accepted prefix leaves the destination unchanged.

  • use_qk_l2norm (bool) – Normalize cached K with epsilon 1e-6. Must match verify. Default True.

  • null_block_id (int) – Largest reserved/invalid slot; must be >= -1. Default -1.

  • backend ({"auto", "simt", "tcgen05"}) – Without tracking and with normalized K, auto selects TF32 tcgen05 for T=8, or T=6/7 when B>1 and layers * B * HV >= 256. Otherwise it uses FP32 SIMT. Shorter accepted prefixes may favor explicit simt; auto does not read accepted lengths back to the host. Explicit tcgen05 supports normalized T=4..8 without tracking. Use simt to retain FP32 contractions.

Notes

Index values, accepted lengths and non-aliasing are caller contracts; they are not read back to the host. All tensor base addresses must be 16-byte aligned for vectorized loads and TMA. Contiguous storage-offset views must also satisfy this alignment. Warm up before CUDA graph capture.