Skip to content

vllm_omni.diffusion.attention.backends.flash_attn_hub

logger module-attribute 

logger = init_logger(__name__)

FlashAttention3HubBackend

Bases: AttentionBackend

accept_output_buffer class-attribute instance-attribute 

accept_output_buffer: bool = True

supports_piecewise_spans class-attribute instance-attribute 

supports_piecewise_spans: bool = True

get_impl_cls staticmethod 

get_impl_cls() -> type[FlashAttention3HubImpl]

get_name staticmethod 

get_name() -> str

get_supported_head_sizes staticmethod 

get_supported_head_sizes() -> list[int]

supports_attention_mask classmethod 

supports_attention_mask(
    attention_spec: object | None = None,
) -> bool

FlashAttention3HubImpl

Bases: AttentionImpl[AttentionMetadata]

causal instance-attribute 

causal = causal

flash_attn_func instance-attribute 

flash_attn_func: Callable[..., Any] | None = getattr(
    hub_module, "flash_attn_func", None
)

flash_attn_varlen_func instance-attribute 

flash_attn_varlen_func: Callable[..., Any] | None = getattr(
    hub_module, "flash_attn_varlen_func", None
)

num_heads instance-attribute 

num_heads = num_heads

qkv_layout instance-attribute 

qkv_layout = qkv_layout

softmax_scale instance-attribute 

softmax_scale = softmax_scale

forward_cuda 

forward_cuda(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> Tensor

FlashAttentionHubBackend

Bases: AttentionBackend

accept_output_buffer class-attribute instance-attribute 

accept_output_buffer: bool = True

supports_piecewise_spans class-attribute instance-attribute 

supports_piecewise_spans: bool = True

get_impl_cls staticmethod 

get_impl_cls() -> type[FlashAttentionHubImpl]

get_name staticmethod 

get_name() -> str

get_supported_head_sizes staticmethod 

get_supported_head_sizes() -> list[int]

supports_attention_mask classmethod 

supports_attention_mask(
    attention_spec: object | None = None,
) -> bool

FlashAttentionHubImpl

Bases: AttentionImpl[AttentionMetadata]

causal instance-attribute 

causal = causal

flash_attn_func instance-attribute 

flash_attn_func: Callable[..., Any] | None = getattr(
    hub_module, "flash_attn_func", None
)

flash_attn_varlen_func instance-attribute 

flash_attn_varlen_func: Callable[..., Any] | None = getattr(
    hub_module, "flash_attn_varlen_func", None
)

num_heads instance-attribute 

num_heads = num_heads

qkv_layout instance-attribute 

qkv_layout = qkv_layout

softmax_scale instance-attribute 

softmax_scale = softmax_scale

forward_cuda 

forward_cuda(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> Tensor