Is this page helpful?

core.ssm.ops.gdp.cumsum#

Within-chunk cumulative sum of the scalar log decays.

Only the scalar, non-reversed, variable-length path is provided – the one the Gated Delta Product prefill calls.

Module Contents#

Functions#

chunk_local_cumsum_scalar_kernel

Cumulative sum within each chunk of a [B, T, H] scalar sequence.

chunk_local_cumsum

Cumulative sum of g within each chunk of length chunk_size.

API#

core.ssm.ops.gdp.cumsum.chunk_local_cumsum_scalar_kernel(
s,
o,
scale,
cu_seqlens,
chunk_indices,
T,
B: core.ssm.ops.gdp.common.tl.constexpr,
H: core.ssm.ops.gdp.common.tl.constexpr,
BT: core.ssm.ops.gdp.common.tl.constexpr,
REVERSE: core.ssm.ops.gdp.common.tl.constexpr,
HAS_SCALE: core.ssm.ops.gdp.common.tl.constexpr,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
HEAD_FIRST: core.ssm.ops.gdp.common.tl.constexpr,
)#

Cumulative sum within each chunk of a [B, T, H] scalar sequence.

core.ssm.ops.gdp.cumsum.chunk_local_cumsum(
g: torch.Tensor,
chunk_size: int,
reverse: bool = False,
scale: float = None,
cu_seqlens: torch.Tensor | None = None,
head_first: bool = False,
output_dtype: torch.dtype | None = torch.float,
chunk_indices: torch.Tensor | None = None,
) → torch.Tensor#

Cumulative sum of g within each chunk of length chunk_size.

Parameters:

  • g – Scalar sequence [B, T, H] (or [B, H, T] when head_first).

  • chunk_size – Chunk length; must be a power of two.

  • reverse – Accumulate from the end of each chunk instead of the start.

  • scale – Optional scale applied to the result.

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

  • head_first – Whether g is laid out head-major.

  • output_dtype – Result dtype; None keeps g’s dtype.

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

Returns the cumulative sums, shaped like g.