Is this page helpful?

core.ssm.ops.gdp.chunk_h#

Inter-chunk state recurrence for the Gated Delta Product.

Carries the matrix-valued state forward across chunks, emitting the state at each chunk boundary (for the intra-chunk output kernel) and the corrected values v_new. The Householder expansion shows up here as the stride between state stores: the state is only checkpointed once every num_householder expanded chunks, because that is one chunk of the original token stream.

The key dimension is handled in fixed 64-wide registers, which is why K is capped at 256.

Module Contents#

Functions#

chunk_gated_delta_product_fwd_kernel_h_blockdim64

Sweep one sequence’s chunks, carrying the [K, V] state in registers.

chunk_gated_delta_product_fwd_h

Run the inter-chunk state recurrence.

API#

core.ssm.ops.gdp.chunk_h.chunk_gated_delta_product_fwd_kernel_h_blockdim64(
k,
v,
w,
v_new,
g,
h,
h0,
ht,
state_indices,
ht_slot_stride,
ht_head_stride,
cu_seqlens,
chunk_offsets,
T,
num_householder: core.ssm.ops.gdp.common.tl.constexpr,
H: core.ssm.ops.gdp.common.tl.constexpr,
K: core.ssm.ops.gdp.common.tl.constexpr,
V: core.ssm.ops.gdp.common.tl.constexpr,
BT: 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_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,
SAVE_NEW_VALUE: core.ssm.ops.gdp.common.tl.constexpr,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
)#

Sweep one sequence’s chunks, carrying the [K, V] state in registers.

core.ssm.ops.gdp.chunk_h.chunk_gated_delta_product_fwd_h(
k: torch.Tensor,
w: torch.Tensor,
u: torch.Tensor,
g: torch.Tensor | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
chunk_size: int = 64,
save_new_value: bool = True,
cu_seqlens: torch.Tensor | None = None,
num_householder: int = 1,
chunk_indices: torch.Tensor | None = None,
chunk_offsets: torch.Tensor | None = None,
state: torch.Tensor | None = None,
state_indices: torch.Tensor | None = None,
states_dtype: torch.dtype | None = None,
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Run the inter-chunk state recurrence.

Parameters:

  • k – Keys [B, T, H, K] on the Householder-expanded token stream.

  • w – WY key factor, shaped like k.

  • u – WY value factor [B, T, H, V].

  • g – Within-chunk cumulative log2 decays on the expanded stream, or None.

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

  • output_final_state – Whether to return the final state.

  • chunk_size – Chunk length.

  • save_new_value – Whether to emit the corrected values v_new.

  • cu_seqlens – Sequence boundaries [N+1] on the expanded stream.

  • num_householder – Number of Householder copies M.

  • chunk_indices – Chunk descriptors for the unexpanded stream. Derived from cu_seqlens // num_householder when omitted, which synchronizes on the device.

  • chunk_offsets – Per-sequence prefix sum of unexpanded chunk counts. Derived from cu_seqlens when omitted.

  • state – [S, H, K, V] per-request state cache for dynamic batching, written in place at state_indices instead of into a dense final_state. -1 slots are skipped.

  • state_indices – [N] cache slot per sequence, or None for a dense [N, H, K, V] final state.

  • states_dtype – dtype for the returned per-chunk states h. Defaults to the input dtype (bf16/fp16). Pass the state-cache dtype when the caller snapshots h for prefix caching: the recurrence accumulates in fp32 and only rounds on store, so an fp32 h preserves the full snapshot precision instead of the bf16 the working dtype would keep.

Returns (h, v_new, final_state). h holds the state at each unexpanded chunk boundary.