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

Mojo function

gated_delta_recurrence_fwd_gpu

def gated_delta_recurrence_fwd_gpu[work_dtype: DType, state_dtype: DType, KEY_HEAD_DIM: Int, VALUE_HEAD_DIM: Int, recurrence_output_LT: TensorLayout, qkv_conv_output_LT: TensorLayout, decay_per_token_LT: TensorLayout, beta_per_token_LT: TensorLayout, recurrent_state_LT: TensorLayout, slot_idx_LT: TensorLayout, input_row_offsets_LT: TensorLayout](batch_size: Int, num_value_heads: Int, num_key_heads: Int, key_dim: Int, recurrence_output: TileTensor[work_dtype, recurrence_output_LT, MutUntrackedOrigin], recurrent_state: TileTensor[state_dtype, recurrent_state_LT, MutUntrackedOrigin], slot_idx: TileTensor[DType.uint32, slot_idx_LT, MutUntrackedOrigin], qkv_conv_output: TileTensor[work_dtype, qkv_conv_output_LT, MutUntrackedOrigin], decay_per_token: TileTensor[work_dtype, decay_per_token_LT, MutUntrackedOrigin], beta_per_token: TileTensor[work_dtype, beta_per_token_LT, MutUntrackedOrigin], input_row_offsets: TileTensor[DType.uint32, input_row_offsets_LT, MutUntrackedOrigin], qkv_conv_output_seqlen_stride: UInt32, qkv_conv_output_channel_stride: UInt32, per_token_seqlen_stride: UInt32, per_token_head_stride: UInt32, recurrent_state_slot_stride: UInt32, recurrent_state_value_head_stride: UInt32, recurrent_state_key_dim_stride: UInt32, recurrent_state_value_dim_stride: UInt32, recurrence_output_seqlen_stride: UInt32, recurrence_output_valuedim_stride: UInt32)

GPU kernel: slot-indexed gated delta rule recurrence, one CTA per head.

One CTA owns one (batch_item, value_head); thread tid == vd_element owns the KD-element state column recurrent_state[slot, value_head, :, tid] in registers for the whole sequence. The per-token raw Q/K for this value head's key head are staged once per block in shared memory (one element per thread, coalesced) so the KD reductions read them from shared memory rather than every vd-thread re-reading the same KD elements from global memory; L2 normalisation and the 1/sqrt(KD) query scale are folded in as scalars factored out of the reductions.

Parameters:

work_dtype (DType): DType for the per-token input and output tensors (qkv_conv_output, decay_per_token, beta_per_token, recurrence_output), float32.
state_dtype (DType): DType for the recurrent_state pool (bfloat16).
KEY_HEAD_DIM (Int): Compile-time key head dimension (e.g. 128 for Qwen3.5).
VALUE_HEAD_DIM (Int): Compile-time value head dimension; must equal KEY_HEAD_DIM.
recurrence_output_LT (TensorLayout): TensorLayout for recurrence_output.
qkv_conv_output_LT (TensorLayout): TensorLayout for qkv_conv_output.
decay_per_token_LT (TensorLayout): TensorLayout for decay_per_token.
beta_per_token_LT (TensorLayout): TensorLayout for beta_per_token.
recurrent_state_LT (TensorLayout): TensorLayout for recurrent_state.
slot_idx_LT (TensorLayout): TensorLayout for slot_idx.
input_row_offsets_LT (TensorLayout): TensorLayout for input_row_offsets.

Args:

batch_size (Int): Number of sequences in the ragged batch.
num_value_heads (Int): Number of value heads (nv).
num_key_heads (Int): Number of key heads (nk); the GQA expansion ratio is num_value_heads / num_key_heads.
key_dim (Int): Total key dimension, equal to num_key_heads * KEY_HEAD_DIM.
recurrence_output (TileTensor[work_dtype, recurrence_output_LT, MutUntrackedOrigin]): Flat output of shape [total_seq_len, value_dim] holding the recurrence result for every token.
recurrent_state (TileTensor[state_dtype, recurrent_state_LT, MutUntrackedOrigin]): Mutable state pool of shape [max_slots, num_value_heads, KEY_HEAD_DIM, VALUE_HEAD_DIM]; the kernel reads and writes slot slot_idx[batch_item] in place.
slot_idx (TileTensor[DType.uint32, slot_idx_LT, MutUntrackedOrigin]): Pool slot index for each batch item, shape [batch_size], uint32.
qkv_conv_output (TileTensor[work_dtype, qkv_conv_output_LT, MutUntrackedOrigin]): Conv1d output from Pass 1, shape [total_seq_len, conv_dim], with Q in channels [0, key_dim), K in [key_dim, 2*key_dim), V in [2*key_dim, 2*key_dim + value_dim).
decay_per_token (TileTensor[work_dtype, decay_per_token_LT, MutUntrackedOrigin]): Per-token per-head scalar decay factor, shape [total_seq_len, num_value_heads].
beta_per_token (TileTensor[work_dtype, beta_per_token_LT, MutUntrackedOrigin]): Per-token per-head beta gate, shape [total_seq_len, num_value_heads].
input_row_offsets (TileTensor[DType.uint32, input_row_offsets_LT, MutUntrackedOrigin]): Ragged offsets of shape [batch_size + 1]; sequence b spans flat indices [input_row_offsets[b], input_row_offsets[b+1]).
qkv_conv_output_seqlen_stride (UInt32): Stride between consecutive sequence positions in qkv_conv_output.
qkv_conv_output_channel_stride (UInt32): Stride between consecutive channels in qkv_conv_output.
per_token_seqlen_stride (UInt32): Stride between consecutive sequence positions in decay_per_token and beta_per_token.
per_token_head_stride (UInt32): Stride between consecutive heads in decay_per_token and beta_per_token.
recurrent_state_slot_stride (UInt32): Stride between consecutive slots in recurrent_state.
recurrent_state_value_head_stride (UInt32): Stride between consecutive value heads in recurrent_state.
recurrent_state_key_dim_stride (UInt32): Stride between consecutive key-dim elements in recurrent_state.
recurrent_state_value_dim_stride (UInt32): Stride between consecutive value-dim elements in recurrent_state.
recurrence_output_seqlen_stride (UInt32): Stride between consecutive sequence positions in recurrence_output.
recurrence_output_valuedim_stride (UInt32): Stride between consecutive value-dim elements in recurrence_output.