Skip to content

Fix FlashAttention3 optimal use when available #195

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
3manifold wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix FlashAttention3 optimal use when available #195

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6940d0e
Fix FlashAttention3 use when available
3manifold Apr 17, 2026
f1aeb4c
Add print statements for attention method selection
3manifold May 19, 2026
465db25
Merge branch 'Lightricks:main' into patch-1
3manifold May 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 19 additions & 10 deletions packages/ltx-core/src/ltx_core/model/transformer/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,12 @@
except ImportError:
flash_attn_interface = None

if flash_attn_interface is not None:
print("Using Flash Attention.")
elif memory_efficient_attention is not None:
print("Using xformers.")
else:
print("Neither Flash Attention nor xFormers xformers is available. Use PytorchAttention.")

class AttentionCallable(Protocol):
def __call__(
Expand Down Expand Up @@ -90,6 +96,7 @@ def __call__(
out = out.reshape(b, -1, heads * dim_head)
return out

_pytorch_attention = PytorchAttention()

class FlashAttention3(AttentionCallable):
def __call__(
Expand All @@ -103,18 +110,15 @@ def __call__(
if flash_attn_interface is None:
raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")

if mask is not None:
# FA3 does not support arbitrary attention masks — fall back to PyTorch SDPA
return _pytorch_attention(q, k, v, heads, mask)

b, _, dim_head = q.shape
dim_head //= heads

q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))

if mask is not None:
raise NotImplementedError("Mask is not supported for FlashAttention3")

out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
out = out.reshape(b, -1, heads * dim_head)
return out

return out.reshape(b, -1, heads * dim_head)

class AttentionFunction(Enum):
PYTORCH = "pytorch"
Expand All @@ -132,8 +136,13 @@ def to_callable(self) -> AttentionCallable:
elif self is AttentionFunction.FLASH_ATTENTION_3:
return FlashAttention3()
else:
# Default behavior: XFormers if installed else - PyTorch
return XFormersAttention() if memory_efficient_attention is not None else PytorchAttention()
# Default behavior: FA3 > XFormers > PyTorch
if flash_attn_interface is not None:
return FlashAttention3()
elif memory_efficient_attention is not None:
return XFormersAttention()
else:
return PytorchAttention()


class Attention(torch.nn.Module):
Expand Down