flashinfer.gdn_prefill.chunk_gated_delta_rule

flashinfer.gdn_prefill.chunk_gated_delta_rule(q: Tensor, k: Tensor, v: Tensor, g: Tensor | None = None, beta: Tensor | None = None, scale: float | None = None, initial_state: Tensor | None = None, output_final_state: bool = False, cu_seqlens: Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, output: Tensor | None = None, output_state: Tensor | None = None, state_checkpoints: Tensor | None = None, checkpoint_cu_starts: Tensor | None = None, checkpoint_every_n_tokens: int = 0, use_cp: Literal['auto'] | bool = 'auto', state_indices: Tensor | None = None) → Tensor | Tuple[Tensor, Tensor]

Chunked Gated Delta Rule (GDN) attention for prefill.

Implements the gated delta rule linear attention mechanism for efficient training and inference. Supports both GQA (grouped query attention) and GVA (grouped value attention) configurations.

Parameters:

  • q (torch.Tensor) – Queries of shape [total_seq_len, num_q_heads, head_size]. Must be contiguous and on CUDA.

  • k (torch.Tensor) – Keys of shape [total_seq_len, num_k_heads, head_size]. Must be contiguous and on CUDA.

  • v (torch.Tensor) – Values of shape [total_seq_len, num_v_heads, head_size]. Must be contiguous and on CUDA.

  • g (torch.Tensor, optional) – Forget gate (alpha) of shape [total_seq_len, num_sab_heads] where num_sab_heads = max(num_q_heads, num_v_heads). Must be float32. Defaults to all ones when None.

  • beta (torch.Tensor, optional) – Update gate (beta) of shape [total_seq_len, num_sab_heads]. Must be float32. Defaults to all ones when None.

  • scale (float, optional) – Scale factor for the attention scores. Defaults to 1 / sqrt(head_size) when None.

  • initial_state (torch.Tensor, optional) – Initial KV state. Packed, sequence-ordered shape [num_seqs, num_sab_heads, head_size, head_size]. Must be float32 on SM90/SM120. The SM100 path also accepts bfloat16, float16, float8_e4m3fn, and float8_e5m2. Starts from zero state when None. When state_indices is given (SM100/SM103 only), this is instead the state pool [N_pool, num_sab_heads, head_size, head_size] and sequence i reads its initial state from row state_indices[i]; the pool may be non-compact (padded first-dimension stride, inner [H, V, K] block contiguous).

  • output_final_state (bool) – Whether to output the final state. Default: False.

  • cu_seqlens (torch.Tensor) – Cumulative sequence lengths of shape [num_seqs + 1], integer dtype on the same CUDA device as q. Required for variable-length sequences (varlen mode); must not be None (asserted at the top of the function body). Internally cast to int32 for the SM100/Blackwell CuTe-DSL kernel and to int64 for the SM90/Hopper C++ kernel, so the caller can pass either dtype.

  • use_qk_l2norm_in_kernel (bool) – Whether to use QK L2 normalization in kernel. Default: False.

  • output (torch.Tensor, optional) – Pre-allocated output tensor of shape [total_seq_len, num_o_heads, head_size] where num_o_heads = max(num_q_heads, num_v_heads). Allocated automatically when None.

  • output_state (torch.Tensor, optional) – Pre-allocated output state tensor. Packed, sequence-ordered shape [num_seqs, num_sab_heads, head_size, head_size]. Must be float32 on SM90/SM120. The SM100 path also accepts bfloat16, float16, float8_e4m3fn, and float8_e5m2. Required when output_final_state=True. When state_indices is given it is instead the output state pool [N_pool, ...] and sequence i’s final state is written to row state_indices[i] (in place when output_state is initial_state); it must be provided by the caller (auto-allocation is rejected, since a compact [num_seqs, ...] buffer would be indexed out of bounds by the pool slot ids).

  • state_checkpoints (torch.Tensor, optional) – Pre-allocated checkpoint tensor of shape [total_checkpoints, num_sab_heads, head_size, head_size]. Must be float32 on SM90/SM120. The SM100 path also accepts bfloat16, float16, float8_e4m3fn, and float8_e5m2. Required when checkpoint_every_n_tokens > 0.

  • checkpoint_cu_starts (torch.Tensor, optional) – Cumulative checkpoint counts of shape [num_seqs + 1], int64. checkpoint_cu_starts[i+1] - checkpoint_cu_starts[i] is the number of checkpoints for sequence i (= seq_len_i // checkpoint_every_n_tokens). Required when checkpoint_every_n_tokens > 0.

  • checkpoint_every_n_tokens (int) – Store intermediate state every N tokens. Must be a multiple of the chunk size (64). 0 disables checkpointing (default).

  • use_cp (Literal["auto"] | bool, optional:) – Whether to use the SM90/SM120 context-parallel DSL implementation when low-parallelism heuristics match. "auto" enables conservative routing, True requires CP support, and False disables CP. Default: "auto".

  • state_indices (torch.Tensor, optional) –

    Int32 tensor of shape [num_seqs] (SM100/SM103 only). When provided, initial_state and output_state are treated as a state pool whose first dimension is indexed by these slot ids rather than laid out in sequence order: sequence i reads its initial state from row state_indices[i] and writes its final state back to the same row (in place when output_state is initial_state). This lets callers that keep a paged/indexed state pool avoid gathering the active rows into a packed buffer and scattering the result back. The pool may be non-compact (padded first-dimension stride). None (default) keeps the packed, sequence-ordered layout.

    The ids must be unique: as with any indexed scatter, two sequences sharing a slot id would concurrently write the same pool row across work tiles, leaving that row’s final state nondeterministic. Uniqueness is a caller precondition (not checked at launch, to avoid a per-call host sync); the caller’s slot allocator is expected to guarantee it.

Returns:

When output_final_state=False, the output tensor of shape [total_seq_len, num_o_heads, head_size]. Otherwise a tuple (output, final_state) where final_state has shape [num_seqs, num_sab_heads, head_size, head_size] — or, when state_indices is given, the state pool [N_pool, ...] itself (i.e. output_state), whose rows named by state_indices now hold the updated final states.

Return type:

torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

Notes

  • Supports GQA (num_q_heads > num_k_heads = num_v_heads) and GVA (num_v_heads > num_q_heads = num_k_heads).

  • The final state layout is [N, H, V, K].

  • Requires SM90 (Hopper) or SM100 (Blackwell) architecture. The SM100 path requires head_size == 128 and nvidia-cutlass-dsl[cu13]>=4.4.2 (pip install flashinfer-python[cu13]).