flashinfer.gdn_paged

Gated Delta-Rule prefill and decode on a paged recurrent-state cache. The caller owns one state pool state_pool[num_pages, H, V, K] and gives every sequence a list of page ids in CSR form, as for a paged KV cache. Page j of a sequence holds the recurrent state after the last token of that page, so a new request that shares a full-page prefix with an earlier one can start from the last shared page instead of recomputing the prefix.

chunk_gated_delta_rule_paged(q, k, v[, g, ...])

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

gated_delta_rule_decode_paged(q, k, v, ...)

Single-token GDN decode on a paged recurrent-state pool.

gdn_paged_prefill_slots(cu_seqlens, ...)

Pool rows that chunk_gated_delta_rule_paged() derives from a page table, computed in one launch without a host synchronization.

gdn_paged_decode_slots(page_indices, ...)

Pool rows that gated_delta_rule_decode_paged() derives from a page table, computed in one launch without a host synchronization.

GDNPagedPrefillSlots(state_indices, ...)

Per-sequence pool rows of a paged prefill batch, all [num_seqs] int32.

GDNPagedDecodeSlots(initial_state_indices, ...)

Per-sequence pool rows of a paged decode batch, both [B] int32.

Page table

A sequence b owns pages page_indices[page_indptr[b]:page_indptr[b + 1]] (any page ids, in any order) and has (num_pages_b - 1) * page_size + last_page_len[b] tokens after the call. Its new tokens in q/k/v are the last ones; the tokens before them are its prefix. With page_size = 1024, a sequence of 2500 tokens owns three pages and last_page_len = 452:

tokens:   [0, 1024)   [1024, 2048)   [2048, 2500)
page:     p0          p1             p2  (partial, last_page_len = 452)
state:    after 1024  after 2048     after 2500

chunk_gated_delta_rule_paged derives every pool row on the device in one small launch, without a host synchronization. The rows depend only on the batch, so a model can compute them once per step with gdn_paged_prefill_slots (gdn_paged_decode_slots for decode) and pass them as slots to every layer’s call, which then skips that launch. The slots are contiguous int32 [num_seqs] tensors on the device of q; they are not checked against the page table. The rows come from a Triton kernel; on SM107 its compilation needs a ptxas from CUDA 13.4 or newer on PATH (or in TRITON_PTXAS_BLACKWELL_PATH), since older ones do not know sm_107a. The call:

  • It reads the initial state from the page that holds the last prefix token, or starts from the zero state when the prefix is empty.

  • It writes the state at every page boundary the new tokens cross to the page that ends there, and the final state to the last page.

  • It never writes a page that ends at or before the prefix, so full pages shared with other sequences stay unchanged. A prefix that ends inside a page continues that page in place.

Prefix lengths need not be multiples of the 64-token chunk, so chunked prefill can split a prompt at any token. A prefix that is not a multiple of 64 costs one more chunk of work, because the kernels pad the first chunk on the left. page_size must be a positive multiple of 64. Zero-length sequences leave the pool unchanged.

gated_delta_rule_decode_paged decodes one token per sequence (T = 1). When the token opens a new page (last_page_len[b] == 1) it reads the previous page and writes the new one; otherwise it updates the last page in place. Every sequence must hold at least one token before the decoded one, so a sequence whose token opens a new page owns at least two pages.

A pool that both entry points use must have the 4D layout [num_pages, H, V, K] (not the 5D pool of flashinfer.gdn_decode.ugpu_state_pool()) and a state dtype both support. The page tables hold real page ids only: -1 is not a padding value. To pad a batch, for example to a CUDA graph’s batch size, give each padded sequence its own scratch page. A padded prefill sequence has no new tokens and one scratch page. A padded decode sequence uses a scratch page with last_page_len >= 2.

# Prefill a 2500-token prompt whose first 2048 tokens are cached in pages
# [7, 3] (shared with another request); page 12 is new.
out = flashinfer.chunk_gated_delta_rule_paged(
    q, k, v, g, beta,                      # the 452 new tokens
    cu_seqlens=torch.tensor([0, 452], dtype=torch.int32, device="cuda"),
    state_pool=pool,                       # [num_pages, H, 128, 128]
    page_indices=torch.tensor([7, 3, 12], dtype=torch.int32, device="cuda"),
    page_indptr=torch.tensor([0, 3], dtype=torch.int32, device="cuda"),
    page_size=1024,
    last_page_len=torch.tensor([452], dtype=torch.int32, device="cuda"),
)
# Reads page 3, writes page 12; pages 7 and 3 are unchanged.

Slot arguments of chunk_gated_delta_rule

The paged prefill is a thin wrapper over three keyword arguments of flashinfer.gdn_prefill.chunk_gated_delta_rule(), which callers with other cache layouts can use directly:

output_state_indices

Pool row of each sequence’s final state, separate from the read row state_indices. A read row of -1 starts from the zero state; a write row of -1 skips the final-state store. A zero-length sequence writes nothing.

checkpoint_state_indices

Pool rows of the checkpoints: the j-th checkpoint of sequence i goes to row checkpoint_state_indices[checkpoint_cu_starts[i] + j] of state_checkpoints, and a row of -1 skips it. This also serves one selected checkpoint per sequence: give every other entry -1.

prefix_lens

Number of tokens each sequence has processed before this call. Checkpoints are taken at absolute positions: after new token n when (prefix_lens[i] + n) % checkpoint_every_n_tokens == 0. Sequence i then has (prefix_lens[i] + seq_len_i) // N - prefix_lens[i] // N checkpoints (N = checkpoint_every_n_tokens).

These arguments need state_indices, an output_state buffer and output_final_state=True for output_state_indices, and checkpoint_every_n_tokens > 0 for checkpoint_state_indices and prefix_lens. Each is a contiguous int32 or int64 tensor on the device of q.

initial_state, output_state and state_checkpoints may be the same pool. With checkpoint_state_indices every pool row must be one dense [H, V, K] block; the row stride may be padded. Rows written by one sequence must not be read or written by another sequence of the same call. The native non-context-parallel kernels (SM8x, SM90, SM100/SM103/SM107 and SM120) serve these arguments; use_cp=True, backend="cake_gdn" and backend="cudnn" reject them.