diff --git a/modules/attention/backends/flex.py b/modules/attention/backends/flex.py index e9b203786..c5e6a9d25 100644 --- a/modules/attention/backends/flex.py +++ b/modules/attention/backends/flex.py @@ -4,15 +4,18 @@ from modules.attention.registry import AttentionBackend, Constraints, Platform def prepare(platform: Platform, original): # pylint: disable=unused-argument - from torch.nn.attention.flex_attention import flex_attention, create_block_mask + from torch.nn.attention.flex_attention import create_block_mask + from modules.attention.sparse import flex as sparse_flex def causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument return q_idx >= kv_idx def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument if selection is not None: - from modules.attention.sparse import flex as sparse_flex return sparse_flex.attend(query, key, value, selection, scale=scale, enable_gqa=enable_gqa) + # compiled, always: eager flex_attention materializes the whole score matrix, which is + # tens of gigabytes at video sequence lengths and fails in the driver rather than cleanly + flex_attention = sparse_flex.flex_call() score_mod = None block_mask = None if attn_mask is not None: