Is this page helpful?

core.ssm.ops.gdp.chunk#

Chunked Gated Delta Product prefill.

The stage sequence:

  1. Interleave the log decays across the Householder copies, then take the within-chunk cumulative sum on both the unexpanded and the expanded stream, converting to base 2 on the way (the kernels exponentiate with exp2).

  2. Build the WY representation: beta * K K^T, then invert I + A, then recompute the w and u factors.

  3. Sweep the inter-chunk state recurrence, emitting the chunk-boundary states.

  4. Combine those with the within-chunk attention to get the outputs.

Two chunkings are in play throughout. The queries and the outputs live on the token stream as written; the keys, values and betas live on the Householder-expanded stream, whose sequences are M times longer. The second chunking is not a rescaling of the first, because ceil(L*M/64) is not M*ceil(L/64) unless L is a multiple of the chunk size.

Module Contents#

Functions#

chunk_gated_delta_product_varlen

Variable-length chunked Gated Delta Product forward pass.

API#

core.ssm.ops.gdp.chunk.chunk_gated_delta_product_varlen(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
num_householder: int,
cu_seqlens: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
chunk_indices: torch.Tensor | None = None,
chunk_indices_dp: torch.Tensor | None = None,
chunk_offsets: torch.Tensor | None = None,
state: torch.Tensor | None = None,
state_indices: torch.Tensor | None = None,
return_chunk_states: bool = False,
) → tuple[torch.Tensor, ...]#

Variable-length chunked Gated Delta Product forward pass.

Parameters:

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

  • k – Keys [1, T*M, H, K] (Householder-expanded).

  • v – Values [1, T*M, H, V] (Householder-expanded).

  • g – Log decays [1, T, H].

  • beta – Betas [1, T*M, H].

  • num_householder – Number of Householder copies M.

  • cu_seqlens – Sequence boundaries over the unexpanded stream, [N+1].

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

  • initial_state – Starting state [N, H, 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 first.

  • chunk_indices – Chunk descriptors for the token stream as written.

  • chunk_indices_dp – The same for the Householder-expanded stream, whose sequences are M times longer. Not a rescaling of chunk_indices: ceil(L*M/64) != M*ceil(L/64) in general.

  • chunk_offsets – Per-sequence prefix sum of unexpanded chunk counts.

  • state – [S, H, K, V] per-request state cache for dynamic batching, written in place at state_indices rather than returned densely.

  • state_indices – [N] cache slot per sequence; -1 marks padding.

  • return_chunk_states – Also return the per-chunk states the scan passes through, [NT, H, K, V]. Row chunk_offsets[i] + c is sequence i’s state entering its chunk c, i.e. after its first 64 * c tokens – which is the mid-sequence state prefix caching snapshots. Note this differs from the Mamba2 chunk scan, whose raw states are indexed by the chunk they come out of.

Returns (o, final_state) with o shaped [1, T, H, V], or (o, final_state, chunk_states) when return_chunk_states is set.

Passing the three descriptor arguments is what makes this capturable in a CUDA graph: deriving them here reads a device tensor on the host and yields a data-dependent length, which also sizes every launch grid below. Built once per step and padded to a fixed length by metadata, they keep the grids constant for a captured batch shape.