Is this page helpful?

core.ssm.ops.gdp.scaled_dot_kkt#

Strictly-lower-triangular beta * K K^T per chunk.

This is the first half of building the WY representation of the chunk’s transition matrix; solve_tril inverts I + A afterwards.

Module Contents#

Functions#

chunk_scaled_dot_kkt_fwd_kernel

Compute one chunk’s beta * K K^T, masked to below the diagonal.

chunk_scaled_dot_kkt_fwd

Compute beta * K K^T per chunk, decayed by g and masked below the diagonal.

API#

core.ssm.ops.gdp.scaled_dot_kkt.chunk_scaled_dot_kkt_fwd_kernel(
k,
g,
beta,
A,
cu_seqlens,
chunk_indices,
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,
BT: core.ssm.ops.gdp.common.tl.constexpr,
BK: core.ssm.ops.gdp.common.tl.constexpr,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
USE_G: core.ssm.ops.gdp.common.tl.constexpr,
)#

Compute one chunk’s beta * K K^T, masked to below the diagonal.

core.ssm.ops.gdp.scaled_dot_kkt.chunk_scaled_dot_kkt_fwd(
k: torch.Tensor,
g: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
cu_seqlens: torch.Tensor | None = None,
chunk_size: int = 64,
output_dtype: torch.dtype = torch.float32,
chunk_indices: torch.Tensor | None = None,
) → torch.Tensor#

Compute beta * K K^T per chunk, decayed by g and masked below the diagonal.

Parameters:

  • k – Keys [B, T, H, K], where H is the number of query/key heads.

  • beta – Betas [B, T, HV], where HV is the number of value/output heads. For GVA, H < HV and HV % H == 0; otherwise H == HV.

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

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

  • chunk_size – Chunk length.

  • output_dtype – Result dtype.

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

Returns [B, T, HV, BT], the per-chunk lower-triangular block.