Is this page helpful?

core.ssm.ops.gdp.chunk_o#

Intra-chunk output for the Gated Delta Product.

Combines the inter-chunk state from chunk_h with the causal within-chunk attention. The queries live on the unexpanded token stream while the keys and values live on the Householder-expanded one, so the within-chunk term is accumulated over the num_householder copies, each strided H*K (or H*V) apart.

Module Contents#

Functions#

chunk_fwd_kernel_o

Emit one chunk of output from the chunk-boundary state plus local attention.

chunk_gated_delta_product_fwd_o

Compute the Gated Delta Product outputs.

Data#

BKV_LIST

NUM_WARPS

API#

core.ssm.ops.gdp.chunk_o.BKV_LIST#

None

core.ssm.ops.gdp.chunk_o.NUM_WARPS#

None

core.ssm.ops.gdp.chunk_o.chunk_fwd_kernel_o(
q,
k,
v,
h,
g,
o,
cu_seqlens,
chunk_indices,
scale,
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,
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,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
)#

Emit one chunk of output from the chunk-boundary state plus local attention.

core.ssm.ops.gdp.chunk_o.chunk_gated_delta_product_fwd_o(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
h: torch.Tensor,
g: torch.Tensor | None = None,
scale: float | None = None,
cu_seqlens: torch.Tensor | None = None,
chunk_size: int = 64,
num_householder: int = 1,
chunk_indices: torch.Tensor | None = None,
) → torch.Tensor#

Compute the Gated Delta Product outputs.

Parameters:

  • q – Queries [B, T, H, K] on the unexpanded token stream.

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

  • v – Corrected values [B, T*M, H, V] from chunk_h.

  • h – State at each chunk boundary, from chunk_h.

  • g – Within-chunk cumulative log2 decays [B, T, H], or None.

  • scale – Score scale.

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

  • chunk_size – Chunk length.

  • num_householder – Number of Householder copies M.

  • chunk_indices – Precomputed chunk descriptors for the unexpanded stream. Derived from cu_seqlens when omitted, which synchronizes on the device.

Returns the outputs [B, T, H, V].