flashinfer.msa_ops.MSASparseAttentionWorkspace

class flashinfer.msa_ops.MSASparseAttentionWorkspace(device: device | str)

Caller-owned storage for SM100/SM103 MSA CUDA graph capture.

Construct one workspace per captured sparse-attention invocation. Warm it by calling the operation eagerly with the exact tensors, options, and CUDA stream that will be captured, then synchronize that stream before capture. The workspace owns output and temporary tensors whose addresses must stay stable for graph replay.

A workspace binds to its first stream. Once it participates in capture it cannot be passed through Python again; graph replay remains valid for the lifetime of the workspace.

__init__(device: device | str) → None

Methods

__init__(device)