Skip to content

vllm_omni.diffusion.attention.backends.cudnn_attn

CuDNNAttentionBackend

Bases: AttentionBackend

accept_output_buffer class-attribute instance-attribute

accept_output_buffer: bool = True

supports_prefix_kv_slicing class-attribute instance-attribute

supports_prefix_kv_slicing: bool = True

get_impl_cls staticmethod

get_impl_cls() -> type[CuDNNAttentionImpl]

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

supports_head_size classmethod

supports_head_size(head_size: int) -> bool

CuDNNAttentionImpl

Bases: AttentionImpl

backend_explicit instance-attribute

backend_explicit = bool(
    extra_impl_args.get("backend_explicit", False)
)

causal instance-attribute

causal = causal

softmax_scale instance-attribute

softmax_scale = softmax_scale

forward_cuda

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