fastvideo.v1.attention.backends.video_sparse_attn

`fastvideo.v1.attention.backends.video_sparse_attn`#

Module Contents#

Classes#

`VideoSparseAttentionBackend`
`VideoSparseAttentionImpl`
`VideoSparseAttentionMetadata`
`VideoSparseAttentionMetadataBuilder`

Data#

logger

API#

class fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionBackend[source]#

Bases: fastvideo.v1.attention.backends.abstract.AttentionBackend

accept_output_buffer: bool[source]#: True

static get_builder_cls() → Type[fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadataBuilder][source]#

static get_impl_cls() → Type[fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionImpl][source]#

static get_metadata_cls() → Type[fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata][source]#

static get_name() → str[source]#

static get_supported_head_sizes() → List[int][source]#

class fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionImpl(num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: Optional[int] = None, prefix: str = '', **extra_impl_args)[source]#

Bases: fastvideo.v1.attention.backends.abstract.AttentionImpl

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, gate_compress: torch.Tensor, attn_metadata: fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata) → torch.Tensor[source]#

postprocess_output(output: torch.Tensor, attn_metadata: fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata) → torch.Tensor[source]#

preprocess_qkv(qkv: torch.Tensor, attn_metadata: fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata) → torch.Tensor[source]#

tile(x: torch.Tensor) → torch.Tensor[source]#

untile(x: torch.Tensor) → torch.Tensor[source]#

class fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata[source]#

Bases: fastvideo.v1.attention.backends.abstract.AttentionMetadata

VSA_sparsity: float[source]#: None

current_timestep: int[source]#: None

dit_seq_shape: List[int][source]#: None

class fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadataBuilder[source]#

Bases: fastvideo.v1.attention.backends.abstract.AttentionMetadataBuilder

build(current_timestep: int, forward_batch: fastvideo.v1.pipelines.pipeline_batch_info.ForwardBatch, fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs) → fastvideo.v1.attention.backends.video_sparse_attn.VideoSparseAttentionMetadata[source]#

prepare()[source]#

fastvideo.v1.attention.backends.video_sparse_attn.logger[source]#: ‘init_logger(…)’