flashinfer.kda_backward.RecurrentKDABackwardWorkspace

class flashinfer.kda_backward.RecurrentKDABackwardWorkspace(device: device | str)

Caller-owned scratch for recurrent_kda_backward().

Construct one workspace per invocation that will be captured by a CUDA Graph. Warm it by calling recurrent_kda_backward() eagerly with the exact input and out tensors on the intended capture stream, then synchronize that stream before capture. The warm call allocates every route-specific intermediate, validates packed offsets, and prepares the high-head TMA descriptors. Capture performs no allocation or descriptor preparation.

A workspace binds to its first stream. Once it participates in capture it cannot be passed through Python again; graph replay remains valid while the workspace and all warmed tensors stay alive.

__init__(device: device | str) → None

Methods

__init__(device)