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.

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,
cu_seqlens,
scale,
T,
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,
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,
) → 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.

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