fastvideo.v1.attention.backends.sliding_tile_attn#

Module Contents#

Classes#

Functions#

Data#

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

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.dict_to_3d_list(mask_strategy) List[List[List[Optional[torch.Tensor]]]][source]#
fastvideo.v1.attention.backends.sliding_tile_attn.logger[source]#

β€˜init_logger(…)’