Is this page helpful?

core.ssm.ops.gdp.solve_tril#

Inverse of I + A for a strictly lower-triangular, chunk-blocked A.

Only the 64x64 path is provided; the Gated Delta Product kernels are wired to a chunk length of 64 throughout (see CHUNK_SIZE).

The 64x64 inverse is assembled from four 16x16 diagonal blocks inverted by forward substitution, then merged with the off-diagonal blocks via the standard block-triangular identity.

Module Contents#

Functions#

merge_16x16_to_64x64_inverse_kernel

Invert one 64x64 unit-lower-triangular block from its 16x16 sub-blocks.

solve_tril

Compute the inverse of I + A for a strictly lower-triangular A.

Data#

FLA_TRIL_PRECISION

DOT_PRECISION_AUTOTUNE_LIST

API#

core.ssm.ops.gdp.solve_tril.FLA_TRIL_PRECISION#

‘get(…)’

core.ssm.ops.gdp.solve_tril.DOT_PRECISION_AUTOTUNE_LIST#

None

core.ssm.ops.gdp.solve_tril.merge_16x16_to_64x64_inverse_kernel(
A,
Ai,
cu_seqlens,
chunk_indices,
T,
H: core.ssm.ops.gdp.common.tl.constexpr,
BT: core.ssm.ops.gdp.common.tl.constexpr,
USE_TMA: core.ssm.ops.gdp.common.tl.constexpr,
IS_VARLEN: core.ssm.ops.gdp.common.tl.constexpr,
DOT_PRECISION: core.ssm.ops.gdp.common.tl.constexpr,
)#

Invert one 64x64 unit-lower-triangular block from its 16x16 sub-blocks.

core.ssm.ops.gdp.solve_tril.solve_tril(
A: torch.Tensor,
cu_seqlens: torch.Tensor | None = None,
chunk_indices: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float,
) → torch.Tensor#

Compute the inverse of I + A for a strictly lower-triangular A.

Parameters:

  • A – [B, T, H, 64]; A.triu() == 0 is assumed, not checked.

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

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

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

Returns (I + A)^-1, shaped like A.