Is this page helpful?

core.ssm.ops.gdp.wy_fast#

Recompute the w and u factors of the WY representation.

Given the inverted transition block A from solve_tril, this produces the two per-chunk factors the state recurrence consumes: u, the effective values, and w, the effective keys carrying the decay.

Module Contents#

Functions#

recompute_w_u_fwd_kernel

Apply the inverted transition block to the betas-scaled keys and values.

recompute_w_u_fwd

Recompute the WY factors w and u for every chunk.

API#

core.ssm.ops.gdp.wy_fast.recompute_w_u_fwd_kernel(
k,
v,
beta,
w,
u,
A,
g,
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,
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,
)#

Apply the inverted transition block to the betas-scaled keys and values.

core.ssm.ops.gdp.wy_fast.recompute_w_u_fwd(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
g: torch.Tensor | None = None,
cu_seqlens: torch.Tensor | None = None,
chunk_indices: torch.Tensor | None = None,
) → tuple[torch.Tensor, torch.Tensor]#

Recompute the WY factors w and u for every chunk.

Parameters:

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

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

  • beta – Betas [B, T, HV].

  • A – Inverted transition blocks [B, T, HV, BT] from solve_tril.

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

  • 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.

Returns (w, u), shaped [B, T, HV, K] and like v respectively.