For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

kda_chunk

kda_chunk()

max.nn.state_space.kda_chunk(q, k, v, raw_gate, beta_logits, a_log, dt_bias, cu_seqlens, state_pool, state_indices, *, output_dtype, use_computebound=False, gate_mode='original', beta_mode='logits', state_layout='K_FIRST')

source

Runs the KDA chunk-parallel prefill recurrence, mutating state_pool.

This is the chunk-parallel prefill path: same recurrence as kda_decode but evaluated in parallel across 16-token chunks. It is intended for multi-token sequences; single-token decode should stay on kda_decode.

use_computebound=True selects the fused compute-bound kernel instead of the L1/L2/L3 pipeline on supported shapes.

On an AMD accelerator (gfx950 only) the op runs a separate pipeline over 64-token chunks with a narrower contract: key_head_dim == value_head_dim == 128, num_key_heads == num_value_heads, gate_mode="safe", bf16 q/k/v and output, fp32 gates and state, and use_computebound=False.

Parameters:

• q (TensorValue) – [total_tokens, num_key_heads, key_head_dim].
• k (TensorValue) – [total_tokens, num_key_heads, key_head_dim].
• v (TensorValue) – [total_tokens, num_value_heads, value_head_dim].
• raw_gate (TensorValue) – [total_tokens, num_value_heads, key_head_dim] forget-gate pre-activation, before dt_bias is added.
• beta_logits (TensorValue) – [total_tokens, num_value_heads].
• a_log (TensorValue) – [num_value_heads].
• dt_bias (TensorValue) – [num_value_heads, key_head_dim].
• cu_seqlens (TensorValue) – [batch_size + 1] int32 exclusive prefix offsets.
• state_pool (BufferValue) – [max_slots, num_value_heads, key_head_dim, value_head_dim] mutable pool, laid out per state_layout.
• state_indices (TensorValue) – [batch_size] int32 pool slot per sequence.
• output_dtype (DType) – Dtype of the returned tensor.
• use_computebound (bool) – Whether to dispatch the fused compute-bound kernel.
• gate_mode (Literal['original', 'safe']) – Forget-gate form; see KDA_GATE_LOWER_BOUND before selecting "safe".
• beta_mode (Literal['logits', 'probability']) – Whether the kernel applies the sigmoid to beta_logits.
• state_layout (Literal['K_FIRST', 'V_FIRST']) – Pool axis order; must match how the pool was allocated.

Returns:

[total_tokens, num_value_heads, value_head_dim].

Raises:

ValueError – If the dtype groupings bind no kernel, the ranks and extents disagree, or, on AMD, the arguments fall outside the AMD contract. The kernel’s own shape guards are debug_assert on NVIDIA, so they are absent from a production build.

Return type:

TensorValue