flashinfer.cake_gdp.cake_gated_delta_product_mtp

flashinfer.cake_gdp.cake_gated_delta_product_mtp(q: Tensor, k: Tensor, v: Tensor, initial_state: Tensor, initial_state_indices: Tensor, A_log: Tensor, a: Tensor, dt_bias: Tensor, b: Tensor, scale: float | None = None, output: Tensor | None = None, ssm_state_indices: Tensor | None = None, disable_state_update: bool | None = None, use_qk_l2norm: bool = True, output_state_indices: Tensor | None = None) → tuple[Tensor, Tensor]

Gated DeltaProduct decode / MTP verification over an FP32 state pool.

One fused kernel per call: for every batch row and head the state is read from its pool slot, decayed once per token, updated by the token’s num_householder Householder sub-steps, read out at the real token and, after the last token, written back to the row’s write slot. The gate is fused: alpha_t = exp(-exp(A_log) * softplus(a_t + dt_bias)) per token (softplus threshold 20) and beta_tj = sigmoid(b_tj) per sub-step. The parameters are those of flashinfer.gated_delta_rule_mtp(), with k, v and b carrying the Householder axis.

Parameters:

  • q (torch.Tensor) – Queries [B, T, H, 128], bfloat16 or float16, one per real token. q and k views are read in place when their last dimension is dense and the view is 8-byte aligned (base address and outer strides multiples of four elements); other layouts are refused. float16 inputs are converted to bfloat16 for the kernel.

  • k (torch.Tensor) – Keys [B, T, num_householder, H, 128] in q.dtype: the num_householder (1 to 4) sub-steps of token t are k[:, t, 0 .. n - 1].

  • v (torch.Tensor) – Values [B, T, num_householder, HV, 64] in q.dtype. HV must equal H (no GQA / GVA).

  • initial_state (torch.Tensor) – FP32 state pool [pool_size, HV, 64, 128] (K-last), read and written in place. Its rows must be dense, its strides multiples of four elements and its base address 16-byte aligned; a non-contiguous pool with these properties (padding between heads or slots) is used as is.

  • initial_state_indices (torch.Tensor) – Read slots [B] (int32 or int64). A negative entry marks a padding row: the kernel skips it entirely – no state read, no recurrence, no output store, no writeback and no scatter – so its output row keeps the value of the caller’s output (zero when output is None).

  • A_log (torch.Tensor) – Log decay parameter [HV], float32 (read in place).

  • a (torch.Tensor) – Input-dependent decay [B, T, HV] in q.dtype, one per real token.

  • dt_bias (torch.Tensor) – Decay bias [HV]. A float32 vector is read in place; a bfloat16 or float16 vector is converted once per call, which adds one small launch – keep the parameter in float32 for the single-launch path.

  • b (torch.Tensor) – Update-gate inputs [B, T, num_householder, HV] in q.dtype, one per sub-step.

  • scale (float, optional) – Query scale; 1 / sqrt(128) when None. Any other value, 0.0 included, is applied as given.

  • output (torch.Tensor, optional) – Pre-allocated output [B, T, HV, 64], bfloat16 or float16. A bfloat16 buffer whose last dimension is dense and whose base address and strides are 8-byte aligned is written directly; any other buffer receives a copy after the launch (its values in skipped rows are preserved through a bfloat16 round trip). Allocated as zeros in q.dtype when None.

  • ssm_state_indices (torch.Tensor, optional) – Per-token scatter slots [B, T], int32. When given, the state after every real token (its num_householder sub-steps completed) is written to initial_state[ssm_state_indices[b, t]] (a negative entry skips that write) and the end-of-sequence writeback to the write slot is skipped, as in flashinfer.gated_delta_rule_mtp(). Requires disable_state_update=False; T == 1 is accepted.

  • disable_state_update (bool, optional) – True leaves the pool unchanged (verification without commit). None (the default) means False: the state is written back.

  • use_qk_l2norm (bool) – Apply the L2 normalization (x * rsqrt(sum(x^2) + 1e-6)) to q and k inside the kernel. Default True.

  • output_state_indices (torch.Tensor, optional) – Write slots [B] (int32 or int64); initial_state_indices when None. A negative entry skips the writeback of that row only (its output is still computed and stored).

Returns:

(output, initial_state): the output [B, T, HV, 64] in the requested dtype and the pool with the writebacks / scatters applied (unchanged when disable_state_update=True).

Return type:

tuple[torch.Tensor, torch.Tensor]

Notes

  • Requires an SM90, SM100 or SM103 device (sm_90a / sm_100a / sm_103a programs).

  • T is a runtime value up to the exported maximum (16); T of 1, 4 and 8 with num_householder == 3 run programs compiled for that token count, selected by B x HV.

  • The call reads no tensor contents on the host and synchronizes nothing. With bfloat16 activations, float32 A_log / dt_bias, int32 indices and a bfloat16 output it launches the kernel alone; it allocates only the output (when output is None) and, when ssm_state_indices is None, an uninitialized [B, T] int32 scatter table the kernel never reads. It is CUDA-graph safe.

  • Every batch row is independent of the others and three launches on the same inputs are bitwise identical.