From 01b97dea88e108fccc325d7103a860dcf6934b49 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 23 Aug 2026 00:45:14 +0100 Subject: [PATCH] fix(attention): compile the flex backend flex_attention called eagerly materializes the whole score matrix. At a video sequence length that is tens of gigabytes, and it surfaces as a CUDA driver error rather than a clean allocation failure: MiniMax H3 at 9505 tokens over 56 heads asks for about 20 GB for the scores alone and the generation dies. The backend now calls the compiled entry point for every path, which is also the only form that reads block lists. --- modules/attention/backends/flex.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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: