Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(attention): compile the flex backend

Browse Source 
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.
This commit is contained in:
CalamitousFelicitousness
2026-08-23 00:45:14 +01:00
parent 3fc5861445
commit 01b97dea88
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/attention/backends/flex.py
+5 -2
View File
@@ -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: