Is this page helpful?

core.ssm.ops.gdp.prepare#

Fused decode- and prefill-step preparation for Gated Delta Product.

Between the short conv and the recurrence, GDP dynamic inference does nothing but reshape and gate: it splits the post-conv activations into value, key and query, GQA-expands query/key, lays value/key out with the Householder copies folded into the sequence dimension, and turns the ba slice into beta = sigmoid(b) and g = -exp(A_log) * softplus(a + dt_bias). Each of those steps is a separate elementwise or copy kernel in PyTorch – ~14 per layer in prefill, three of them strided copies of most of the post-conv tensor – and a decode step carries one token per request, so their launch overhead outweighs the work they do.

gdp_decode_prepare and gdp_prefill_prepare each do all of it in one launch of the same Triton kernel. They differ only in what the downstream kernel wants:

  • Decode feeds fused_recurrent_gated_delta_rule_update, which walks the Householder copies as a sequence, so query and g are interleaved onto the copies (query on the last, the decay on the first, zeros elsewhere). One program covers one head, which keeps a small decode batch spread over the GPU.

  • Prefill feeds chunk_gated_delta_product_varlen, which places query and g on the copies itself, so they stay one per token. One program covers every head of a (token, copy), because a prefill step carries thousands of tokens.

Both are CUDA-graph safe for the same reasons as the rest of this package: static shapes, no host synchronization, and padding handled by the conv kernel upstream (a -1 slot zeroes its conv output, which propagates here). The prefill token count is a runtime argument, so prefill lengths do not each compile a kernel; decode’s (1 plus the draft length) is a compile-time constant.

Both are deterministic: a pure elementwise map where each program owns a disjoint slice of every output, with no atomics, no reductions and no autotuning, so nothing depends on scheduling order.

Both also aim to be bitwise identical to the eager path, which the inference functional tests rely on. query/key/value are plain copies and match for free; beta and g match only because the transcendentals and the division go through libdevice rather than Triton’s fp32 defaults. test_decode_prepare.py and test_prefill_prepare.py assert that with torch.equal, so a Triton or torch upgrade that breaks the agreement fails there rather than in a functional test.

Module Contents#

Functions#

softplus

log1p(exp(x)) in fp32, saturating to the identity above 20 like torch’s.

gdp_prepare_kernel

One program per (request-token, Householder copy, block of BH heads).

_check_prepare_args

_launch_prepare

gdp_decode_prepare

Split, reshape and gate one decode step’s post-conv activations.

gdp_prefill_prepare

Split, reshape and gate a prefill step’s post-conv activations.

API#

core.ssm.ops.gdp.prepare.softplus(x)#

log1p(exp(x)) in fp32, saturating to the identity above 20 like torch’s.

log1p rather than log(1 + ...) so the far-negative tail, where exp(x) falls below fp32’s epsilon and 1 + exp(x) rounds to exactly 1, keeps its significant digits instead of flushing to zero.

Mirrors F.softplus’s (x * beta) > threshold ? x : log1p(exp(x * beta)) / beta at the default beta == 1.0, where the multiply and divide are exact.

core.ssm.ops.gdp.prepare.gdp_prepare_kernel(
x,
x_n_stride,
x_t_stride,
ba,
ba_n_stride,
ba_t_stride,
A_log,
dt_bias,
q,
k,
v,
beta,
g,
T,
T_CONST: core.ssm.ops.gdp.common.tl.constexpr,
H: core.ssm.ops.gdp.common.tl.constexpr,
G: core.ssm.ops.gdp.common.tl.constexpr,
P: core.ssm.ops.gdp.common.tl.constexpr,
N: core.ssm.ops.gdp.common.tl.constexpr,
M: core.ssm.ops.gdp.common.tl.constexpr,
HEADS_PER_GROUP: core.ssm.ops.gdp.common.tl.constexpr,
K_BASE: core.ssm.ops.gdp.common.tl.constexpr,
Q_BASE: core.ssm.ops.gdp.common.tl.constexpr,
BH: core.ssm.ops.gdp.common.tl.constexpr,
BP: core.ssm.ops.gdp.common.tl.constexpr,
BN: core.ssm.ops.gdp.common.tl.constexpr,
INTERLEAVE_QG: core.ssm.ops.gdp.common.tl.constexpr,
)#

One program per (request-token, Householder copy, block of BH heads).

x is the post-conv [n, T, M*H*P + M*G*N + G*N] tensor – value, key, query concatenated – and ba is the [n, T, M*H + H] gating tensor. For decode T is 1 plus the speculative draft length, passed as T_CONST so the per-program split into (request, token) is a constant division; prefill passes T_CONST == 0 and its packed token count at runtime, so step sizes do not each compile a kernel.

core.ssm.ops.gdp.prepare._check_prepare_args(name, x, ba, M, H, G, P, N)#
core.ssm.ops.gdp.prepare._launch_prepare(
x,
ba,
A_log,
dt_bias,
outputs,
M,
H,
G,
P,
N,
interleave_qg,
)#
core.ssm.ops.gdp.prepare.gdp_decode_prepare(
x: torch.Tensor,
ba: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
num_householder: int,
num_heads: int,
num_groups: int,
head_dim: int,
state_dim: int,
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Split, reshape and gate one decode step’s post-conv activations.

Parameters:

  • x – Post-conv activations [n, S, M*H*P + M*G*N + G*N], value/key/query concatenated along the last dimension. S is 1 plus the number of speculative draft tokens.

  • ba – The ba slice of the input projection, [n, S, M*H + H]. May be a non-contiguous view; only the last dimension must be contiguous.

  • A_log – Log decay rates [H].

  • dt_bias – Softplus bias [H].

  • num_householder – Number of Householder copies M.

  • num_heads – TP/CP-local head count H.

  • num_groups – TP/CP-local group count G; H must be a multiple of it.

  • head_dim – Value head dimension P.

  • state_dim – Key/query head dimension N.

Returns (query, key, value, beta, g) shaped [n, S*M, H, N], [n, S*M, H, N], [n, S*M, H, P], [n, S*M, H] and [n, S*M, H] (fp32) – the exact layouts fused_recurrent_gated_delta_rule_update expects for an S*M-length sequence per request.

core.ssm.ops.gdp.prepare.gdp_prefill_prepare(
x: torch.Tensor,
ba: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
num_householder: int,
num_heads: int,
num_groups: int,
head_dim: int,
state_dim: int,
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Split, reshape and gate a prefill step’s post-conv activations.

Parameters:

  • x – Post-conv activations [n, T, M*H*P + M*G*N + G*N], value/key/query concatenated along the last dimension (dynamic-batching prefill packs every request into n == 1). Rows may be strided (e.g. a view of the conv output); only the last dimension must be contiguous.

  • ba – The ba slice of the input projection, [n, T, M*H + H]. May be a non-contiguous view; only the last dimension must be contiguous.

  • A_log – Log decay rates [H].

  • dt_bias – Softplus bias [H].

  • num_householder – Number of Householder copies M.

  • num_heads – TP/CP-local head count H.

  • num_groups – TP/CP-local group count G; H must be a multiple of it.

  • head_dim – Value head dimension P.

  • state_dim – Key/query head dimension N.

Returns (query, key, value, beta, g) shaped [n, T, H, N], [n, T*M, H, N], [n, T*M, H, P], [n, T*M, H] and [n, T, H] (fp32) – what chunk_gated_delta_product_varlen consumes, with the L2 norm left to it.