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, 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 descriptor storage and small-BH packet-ring storage. Persistent M128 is an eager-only B200/GB200 route; explicit workspaces use direct M128 or M64 so graph capture never synchronizes sequence lengths to construct host task bins.

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)