flashinfer.msa_ops.MSASparseAttentionWorkspace

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

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

On SM107, this workspace is supported only by the packed-NVFP4 paged-KV prefill and decode routes.

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.

The multi-request long-prefill route and the persistent K-major route read the request boundaries from cu_seqlens_q and cu_seqlens_k / seqused_k to the host in the eager warm call (the uniform FP8 decode route reads nothing: its tier comes from the page-table width); the workspace keeps them for that binding: pass the same tensor objects to every later call (a rebuilt cu_seqlens tensor – a new object, even at the old address – is read again at the cost of one synchronization) and rewrite their contents only through torch operations, which advance the tensor version the binding is keyed on.

__init__(device: device | str) → None

Methods

__init__(device)