Skip to content

SD2 inference pass mesh for flahattention with spmd #12

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

Merged
merged 4 commits into from
Jan 28, 2025
Merged

SD2 inference pass mesh for flahattention with spmd #12

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
25d1431
Update attention_processor.py with inference spmd mesh for flahattent…
ManfeiBai Jan 28, 2025
2ff4275
Update attention_processor.py
ManfeiBai Jan 28, 2025
c141e0a
Update attention_processor.py
ManfeiBai Jan 28, 2025
8acfe07
format change
ManfeiBai Jan 28, 2025
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
15 changes: 10 additions & 5 deletions src/diffusers/models/attention_processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,14 @@
from torch import nn

from ..image_processor import IPAdapterMaskProcessor
from ..utils import deprecate, logging, is_torch_xla_available
from ..utils import deprecate, is_torch_xla_available, logging
from ..utils.import_utils import is_torch_npu_available, is_xformers_available
from ..utils.torch_utils import is_torch_version, maybe_allow_in_graph


if is_torch_xla_available():
import torch_xla.distributed.spmd as xs
import torch_xla.runtime as xr
from torch_xla.experimental.custom_kernel import flash_attention
XLA_AVAILABLE = True
else:
Expand Down Expand Up @@ -2369,17 +2372,19 @@ def __call__(

# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1


mesh = xs.get_global_mesh()
if XLA_AVAILABLE and all(tensor.shape[2] >= 4096 for tensor in [query, key, value]):
if attention_mask is not None:
attention_mask = attention_mask.view(batch_size, 1, 1, attention_mask.shape[-1])
# Convert mask to float and replace 0s with -inf and 1s with 0
attention_mask = attention_mask.float().masked_fill(attention_mask == 0, float('-inf')).masked_fill(attention_mask == 1, float(0.0))

# Apply attention mask to key
key = key + attention_mask
query /= math.sqrt(query.shape[3])
hidden_states = flash_attention(query, key, value, causal=False, partition_spec=('data', None, None, None))
partition_spec = ('data', None, None, None) if xr.use_spmd() else None
hidden_states = flash_attention(query, key, value, causal=False, partition_spec=partition_spec, mesh=mesh)
else:
hidden_states = F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
Expand Down Expand Up @@ -4317,4 +4322,4 @@ def __init__(self):
PAGIdentitySelfAttnProcessor2_0,
PAGCFGHunyuanAttnProcessor2_0,
PAGHunyuanAttnProcessor2_0,
]
]
Loading