Is this page helpful?

core.ssm.ops.gdp.fused_recurrent#

Fused recurrent Gated Delta Rule step, used by the decode path.

Gated Delta Product decode reaches this kernel by folding the M Householder copies into the sequence dimension, so a single decode token becomes an M-length sequence with the query placed on the last copy and the decay on the first; the caller slices the answer back out. Under speculative decoding a step carries S draft tokens per request, so the folded sequence is S * M long.

Speculative decoding also needs the state between draft tokens, because verification may accept only a prefix of them and the recurrence must roll back to the last accepted token. intermediate_states asks the kernel to snapshot b_h every steps_per_token recurrence steps – once per draft token, after its M Householder updates have all been applied – which is the direct analog of the intermediate_ssm_states dump in the Mamba2 selective_state_update kernel. The snapshot is taken inside the step loop, so it costs one extra store per draft token and no re-read of the running state.

Module Contents#

Functions#

fused_recurrent_gated_delta_rule_fwd_kernel

Walk one sequence token by token, carrying the [K, V] state.

fused_recurrent_gated_delta_rule_update

Run the recurrent Gated Delta Rule forward pass.

API#

core.ssm.ops.gdp.fused_recurrent.fused_recurrent_gated_delta_rule_fwd_kernel(
q,
k,
v,
g,
beta,
o,
h0,
ht,
state_indices,
state_slot_stride,
state_head_stride,
intermediate_states,
int_slot_stride,
int_token_stride,
int_head_stride,
cu_seqlens,
scale,
T,
STEPS_PER_TOKEN: core.ssm.ops.gdp.common.tl.constexpr,
H: core.ssm.ops.gdp.common.tl.constexpr,
HV: core.ssm.ops.gdp.common.tl.constexpr,
K: core.ssm.ops.gdp.common.tl.constexpr,
V: core.ssm.ops.gdp.common.tl.constexpr,
BK: core.ssm.ops.gdp.common.tl.constexpr,
BV: core.ssm.ops.gdp.common.tl.constexpr,
USE_G: core.ssm.ops.gdp.common.tl.constexpr,
USE_QK_L2NORM_IN_KERNEL: core.ssm.ops.gdp.common.tl.constexpr,
IS_BETA_HEADWISE: core.ssm.ops.gdp.common.tl.constexpr,
USE_INITIAL_STATE: core.ssm.ops.gdp.common.tl.constexpr,
STORE_FINAL_STATE: core.ssm.ops.gdp.common.tl.constexpr,
HAS_STATE_INDICES: core.ssm.ops.gdp.common.tl.constexpr,
STORE_INTERMEDIATE: core.ssm.ops.gdp.common.tl.constexpr,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
)#

Walk one sequence token by token, carrying the [K, V] state.

core.ssm.ops.gdp.fused_recurrent.fused_recurrent_gated_delta_rule_update(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
cu_seqlens: torch.Tensor | None = None,
state: torch.Tensor | None = None,
state_indices: torch.Tensor | None = None,
intermediate_states: torch.Tensor | None = None,
steps_per_token: int = 1,
) → tuple[torch.Tensor, torch.Tensor]#

Run the recurrent Gated Delta Rule forward pass.

Parameters:

  • q – Queries [B, T, H, K].

  • k – Keys [B, T, H, K].

  • v – Values [B, T, HV, V].

  • g – Log decays [B, T, HV], or None.

  • beta – Betas [B, T, HV] (head-wise) or [B, T, HV, V].

  • scale – Score scale; defaults to K ** -0.5.

  • initial_state – Starting state [N, HV, K, V], or None for zeros.

  • output_final_state – Whether to return the final state.

  • use_qk_l2norm_in_kernel – Whether to L2-normalize q and k in-kernel.

  • cu_seqlens – Sequence boundaries [N+1] for variable-length input.

  • state – [S, HV, K, V] per-request state cache for dynamic batching, read and written in place at state_indices. Supersedes initial_state / output_final_state, which gather and scatter a dense state instead.

  • state_indices – [N] cache slot per sequence; -1 marks a padding request, whose output is zeroed and whose state is untouched.

  • intermediate_states – [S, T // steps_per_token, HV, K, V] buffer for speculative decoding, written at the same cache slots as state: entry [slot, i] receives the state after draft token i, so verification can roll the recurrence back to any accepted prefix. Padding requests write nothing. None disables the snapshots.

  • steps_per_token – Recurrence steps that make up one draft token, i.e. the Householder count M folded into the sequence dimension. Only meaningful together with intermediate_states.

Returns (o, final_state) with o shaped like v. When state is given, final_state is that same cache tensor, updated in place.