flashinfer.kda.RecurrentKDAPrefillWrapper

class flashinfer.kda.RecurrentKDAPrefillWrapper(device: device | str)

Plan-and-run wrapper for packed recurrent-KDA prefill.

plan runs outside CUDA Graph capture. It reads cu_seqlens on the host, builds a stable descending-length sequence order and cumulative chunk prefix, and copies them into fixed-address device buffers. run consumes those buffers through recurrent_kda() as an explicit host plan.

The number of sequences, total token count, and total BT=16 chunk count are fixed by the first plan call so device buffer addresses, workspace capacity, and captured launch geometry remain valid across CUDA Graph replays. Call plan again before replay to update individual lengths, order, and chunk metadata in place when those totals remain unchanged.

This wrapper is specific to the CuTe DSL backend and intentionally uses its non-persistent schedule. One wrapper instance is a single-writer resource: do not call plan concurrently with run or while a kernel launched by run may still be reading the wrapper’s planned buffers.

__init__(device: device | str) → None

Methods

__init__(device)
plan(cu_seqlens, *[, non_blocking])	Plan a packed prefill sequence order outside CUDA Graph capture.
run(q, k, v, g, beta[, A_log, dt_bias, ...])	Run packed recurrent-KDA prefill using the most recent plan.