flashinfer.kda_prefill.RecurrentKDAPrefillWorkspace

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

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

Auto-selected cuDNN leaves this workspace unused and manages its own scratch.

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. Packed offsets must retain their warmed values through capture and replay. Inference tensors have no version counter, so their offsets are reread on every eager call.

The workspace owns optional final-state scratch, backend metadata, TMA descriptors, and schedule-specific scratch for the lifetime of the graph. On SM100-family devices this includes beta padding, M64/M128-N32/M128-N16 and BT64 descriptor storage, small-BH packet-ring storage, and BT16 prepare/chain metadata, factors, and independent descriptor storage. Cake’s persistent M128 route is eager-only; its explicit workspaces use non-persistent schedules to avoid reading sequence lengths on the host.

Native use binds the workspace to its first stream. After native capture, another native Python call cannot reuse it. Graph replay remains valid for the lifetime of the workspace.

__init__(device: device | str) → None

Methods

__init__(device)