flashinfer.gdn_paged.chunk_gated_delta_rule_paged

flashinfer.gdn_paged.chunk_gated_delta_rule_paged(q: Tensor, k: Tensor, v: Tensor, g: Tensor | None = None, beta: Tensor | None = None, scale: float | None = None, *, cu_seqlens: Tensor, state_pool: Tensor, page_indices: Tensor, page_indptr: Tensor, page_size: int, last_page_len: Tensor, use_qk_l2norm_in_kernel: bool = False, output: Tensor | None = None, gate_domain: Literal['linear', 'log'] = 'linear', use_gate_in_kernel: bool = False, A_log: Tensor | None = None, dt_bias: Tensor | None = None, beta_is_logit: bool = False, slots: GDNPagedPrefillSlots | None = None) → Tensor

Chunked GDN prefill that reads and writes a paged recurrent-state pool.

For every sequence b the call processes its new tokens q[cu_seqlens[b]:cu_seqlens[b + 1]], which are the last tokens of the sequence’s (num_pages_b - 1) * page_size + last_page_len[b] tokens (num_pages_b = page_indptr[b + 1] - page_indptr[b]). The tokens before them form its prefix. The call

• reads the initial state from the page that holds the last prefix token (the zero state when the prefix is empty),

• writes the state at every page boundary it crosses to that page, and

• writes the final state to the last page.

Pages that end at or before the prefix are only read, so a full-page prefix shared with other sequences is never written. A prefix that ends inside a page continues that (private) page in place. Prefix lengths need not be multiples of the 64-token chunk.

Parameters:

• q – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• k – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• v – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• g – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• beta – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• scale – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• use_qk_l2norm_in_kernel – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• output – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• gate_domain – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• use_gate_in_kernel – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• A_log – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• dt_bias – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• beta_is_logit – As in flashinfer.gdn_prefill.chunk_gated_delta_rule().

• cu_seqlens (torch.Tensor) – [num_seqs + 1] contiguous int32/int64 offsets of the new tokens of each sequence in q/k/v. Zero-length sequences are allowed and leave the pool unchanged.

• state_pool (torch.Tensor) – [num_pages, num_sab_heads, head_size, head_size] state pool, updated in place. Each row must be one dense [H, V, K] block; the row stride may be padded. Its dtype is the state dtype.

• page_indices (torch.Tensor) – [num_indices] contiguous int32/int64 page ids. The ids may be in any order and need not be consecutive.

• page_indptr (torch.Tensor) – [num_seqs + 1] contiguous int32/int64 CSR offsets into page_indices. Sequence b owns page_indices[page_indptr[b]:page_indptr[b + 1]] and must own at least one page.

• page_size (int) – Tokens per page, a positive multiple of 64.

• last_page_len (torch.Tensor) – [num_seqs] contiguous int32/int64 number of tokens in the last page of each sequence after this call, in [1, page_size].

• slots (GDNPagedPrefillSlots, optional) – Rows from gdn_paged_prefill_slots() for the same cu_seqlens and page table, typically shared by every layer of a step. Without them the call computes them (one extra launch). They are not checked against the page table.

Returns:

The output [total_seq_len, num_o_heads, head_size].

Return type:

torch.Tensor

Notes

A page written by one sequence must not be read or written by another sequence of the same call. The page table and lengths are caller preconditions; they are not validated against each other on the device.