core.ssm.ops.gdp.decode_prepare

Fused decode-step preparation for Gated Delta Product.

Between the short-conv update and the recurrent kernel, GDP decode does nothing but reshape and gate: it splits the post-conv activations into value, key and query, GQA-expands query/key, interleaves the Householder copies (query on the last copy, the decay on the first), 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, and a decode step carries one token per request, so their launch overhead outweighs the work they do.

gdp_decode_prepare does all of it in one Triton kernel, writing the five tensors the recurrent kernel consumes in exactly the layouts it wants. It is 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).

It is 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.

It also aims 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 asserts 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_decode_prepare_kernel	One program per (request, Householder copy, head).
gdp_decode_prepare	Split, reshape and gate one decode step’s post-conv activations.

API

core.ssm.ops.gdp.decode_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.decode_prepare.gdp_decode_prepare_kernel(

x,

x_n_stride,

ba,

ba_n_stride,

A_log,

dt_bias,

q,

k,

v,

beta,

g,

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,

BP: core.ssm.ops.gdp.common.tl.constexpr,

BN: core.ssm.ops.gdp.common.tl.constexpr,

)

One program per (request, Householder copy, head).

x is the post-conv [n, M*H*P + M*G*N + G*N] row – value, key, query concatenated – and ba is the [n, M*H + H] gating row. Each program emits one head-slice of every output.

core.ssm.ops.gdp.decode_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, 1, M*H*P + M*G*N + G*N], value/key/query concatenated along the last dimension.

• ba – The ba slice of the input projection, [n, 1, 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, M, H, N], [n, M, H, N], [n, M, H, P], [n, M, H] and [n, M, H] (fp32) – the exact layouts fused_recurrent_gated_delta_rule_update expects for an M-length sequence per request.