mirror of
https://github.com/vladmandic/automatic
synced 2026-09-06 13:00:44 +02:00
CalamitousFelicitousness
bb0cc5328e
The flex backend never called the sdpa it replaced, so any backend stacked before it was unreachable and every call it could not serve, cpu or 3d inputs included, failed inside flex_attention. It is now an ordinary entry gated on what flex_attention accepts: 4d tensors on one non-cpu device. The mask path drops the 2d special case, which indexed attn_mask.size and reshaped the mask onto the wrong axis; expanding to (batch, heads, q, kv) already follows sdpa broadcast semantics.
40 lines
1.9 KiB
Python
40 lines
1.9 KiB
Python
import torch
|
|
from modules.logger import log
|
|
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
|
|
|
|
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): # pylint: disable=unused-argument
|
|
score_mod = None
|
|
block_mask = None
|
|
if attn_mask is not None:
|
|
batch_size, num_heads = query.shape[:2]
|
|
seq_len_q = query.shape[-2]
|
|
seq_len_kv = key.shape[-2]
|
|
attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv) # sdpa masks broadcast over the trailing dims
|
|
if attn_mask.dtype == torch.bool:
|
|
def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
|
|
return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
|
|
block_mask = create_block_mask(mask_mod, batch_size, None, seq_len_q, seq_len_kv, device=query.device)
|
|
else:
|
|
def score_mod_fn(score, batch_idx, head_idx, q_idx, kv_idx):
|
|
return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
|
|
score_mod = score_mod_fn
|
|
elif is_causal:
|
|
block_mask = create_block_mask(causal_mask, query.shape[0], query.shape[1], query.shape[-2], key.shape[-2], device=query.device)
|
|
return flex_attention(query, key, value, score_mod=score_mod, block_mask=block_mask, scale=scale, enable_gqa=enable_gqa)
|
|
|
|
log.debug('Torch attention: type="Flex attention"')
|
|
return call
|
|
|
|
|
|
backend = AttentionBackend(
|
|
name='flex', label='Flex attention', priority=20, prepare=prepare,
|
|
constraints=Constraints(min_ndim=4, same_device=True), # flex_attention takes 4d tensors on one device and compiles on cpu
|
|
)
|