flashinfer.kda_prefill.RecurrentKDAPrefillWorkspace

class flashinfer.kda_prefill.RecurrentKDAPrefillWorkspace(device: device | str)

Caller-owned storage required for recurrent-KDA CUDA graph capture.

Construct one workspace per captured flashinfer.kda.recurrent_kda() invocation on the graph’s CUDA device. Warm it by invoking that function eagerly with the exact tensors and capture stream, then synchronize that stream before capture. The workspace owns optional final-state scratch for calls without an initial state, beta padding, and M64/M128 TMA descriptor storage for the lifetime of the graph.

A workspace binds to its first stream. Once it participates in capture it cannot be passed to Python again, either eagerly or in another capture. Graph replay does not invoke Python and remains valid for the lifetime of the workspace.

__init__(device: device | str) → None

Methods

__init__(device)