fastvideo.v1.attention.backends.sage_attn
#
Module Contents#
Classes#
Data#
API#
- class fastvideo.v1.attention.backends.sage_attn.SageAttentionBackend[source]#
Bases:
fastvideo.v1.attention.backends.abstract.AttentionBackend
- static get_impl_cls() Type[fastvideo.v1.attention.backends.sage_attn.SageAttentionImpl] [source]#
- class fastvideo.v1.attention.backends.sage_attn.SageAttentionImpl(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, attn_metadata: fastvideo.v1.attention.backends.abstract.AttentionMetadata) torch.Tensor [source]#