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.