fastvideo.v1.attention.backends.sliding_tile_attn

`fastvideo.v1.attention.backends.sliding_tile_attn`#

Module Contents#

Classes#

`RangeDict`
`SlidingTileAttentionBackend`
`SlidingTileAttentionImpl`
`SlidingTileAttentionMetadata`
`SlidingTileAttentionMetadataBuilder`

Data#

logger

API#

class fastvideo.v1.attention.backends.sliding_tile_attn.RangeDict[source]#: Bases: dict

class fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend[source]#

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

accept_output_buffer: bool[source]#: True

static get_builder_cls() → Type[fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionMetadataBuilder][source]#

static get_impl_cls() → Type[fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionImpl][source]#

static get_metadata_cls() → Type[fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionMetadata][source]#

static get_name() → str[source]#

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

class fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionImpl(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(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_metadata: fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionMetadata) → torch.Tensor[source]#

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

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

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

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

class fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionMetadata[source]#

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

STA_param: List[List[Any]][source]#: None

current_timestep: int[source]#: None

class fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionMetadataBuilder[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.sliding_tile_attn.SlidingTileAttentionMetadata[source]#

prepare()[source]#

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