lloca.backbone.attention_backends.get_attention_backend

lloca.backbone.attention_backends.get_attention_backend(**kwargs)

Resolve the attention backend based on the extra keyword arguments.

Implemented backends:

• PyTorch native attention: torch.nn.functional.scaled_dot_product_attention

• PyTorch varlen attention: torch.nn.attention.varlen.varlen_attn

• xformers attention: xformers.ops.memory_efficient_attention

• PyTorch flex_attention: torch.nn.attention.flex_attention.flex_attention

• Flash attention (variable sequence length): flash_attn.flash_attn_varlen_func

The backend is selected explicitly via backend=... if provided, otherwise inferred from backend-specific kwargs (e.g. cu_seqlens_* triggers flash). Falls back to the native backend. Backends are imported lazily on first use.

Return type:

Callable