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

refactor(attention): backend registry and a single sdpa router

Browse Source 
Replace the six closure hijacks stacked in devices.set_sdpa_params with
a registry of declarative backends and one router installed in their
place. Each backend declares the constraints its closure carried as a
predicate, a priority matching its old stacking position, and a prepare
step that imports and configures the implementation; the router walks
the prepared entries by priority and hands declined calls to the
terminal backend (dynamic, flex) or the original sdpa, so fallback is
the router's job rather than each closure's.

- parity held: gates transcribed literally, the same kernel kwargs,
  enable_gqa passed to the original only when set, torch_info keeps the
  last prepared backend, the dynamic pin still set
- a backend enabled on a platform without it warns instead of silently
  doing nothing
- the legacy set_* entry points are gone; devices.py installs the router
- test/test-attention-router.py checks every override subset against the
  old stacking order, gate parity over 16,000 shape cases, dispatch,
  terminal handoff and prepare isolation, offline
This commit is contained in:
CalamitousFelicitousness
2026-08-22 21:51:29 +01:00
parent 6ed1b99aaa
commit 3302e78af6
13 changed files with 748 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/attention/__init__.py
+7 -3
View File
@@ -1,8 +1,12 @@
"""Attention backends: the SDPA hijacks stacked by devices.set_sdpa_params, plus the diffusers-side processor and dispatcher setup."""
from modules.attention.hijacks import set_dynamic_attention, set_sdnq_attention, set_triton_flash_attention, set_flex_attention, set_ck_flash_attention, set_sage_attention
"""Attention backends: one scaled_dot_product_attention router over the registered backends, plus the diffusers-side processor and dispatcher setup."""
from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry
from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router
from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
from modules.attention import backends
__all__ = [
'set_dynamic_attention', 'set_sdnq_attention', 'set_triton_flash_attention', 'set_flex_attention', 'set_ck_flash_attention', 'set_sage_attention',
'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry',
'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router',
'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
'backends',
]
modules/attention/backends/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""Built-in backends, registered in ascending priority."""
from modules.attention.registry import registry
from modules.attention.backends import dynamic, flex, triton_amd, flash_ck, sage, sdnq
registry.register(dynamic.backend)
registry.register(flex.backend)
registry.register(triton_amd.backend)
registry.register(flash_ck.backend)
registry.register(sage.backend)
registry.register(sdnq.backend)
modules/attention/backends/dynamic.py
+11
View File
@@ -0,0 +1,11 @@
from modules.attention.registry import AttentionBackend, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from modules import devices
devices.sdpa_pre_dyanmic_atten = original # the sliced path calls this pin for every slice
from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
return dynamic_scaled_dot_product_attention
backend = AttentionBackend(name='dynamic', label='Dynamic attention', priority=10, prepare=prepare, terminal=True)
modules/attention/backends/flash_ck.py
+44
View File
@@ -0,0 +1,44 @@
from installer import install, installed
from modules import rocm
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
if platform.backend == 'rocm':
if not installed('flash-attn'):
log.info('Torch attention: type="Flash attention" building...')
agent = rocm.Agent(platform.device)
install(rocm.get_flash_attention_command(agent), reinstall=True)
else:
install('flash-attn')
from flash_attn import flash_attn_func
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
is_unsqueezed = False
if query.dim() == 3:
query = query.unsqueeze(0)
is_unsqueezed = True
if key.dim() == 3:
key = key.unsqueeze(0)
if value.dim() == 3:
value = value.unsqueeze(0)
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
if is_unsqueezed:
attn_output = attn_output.squeeze(0)
return attn_output
log.debug('Torch attention: type="Flash attention"')
return call
backend = AttentionBackend(
name='flash', label='Flash attention', priority=40, prepare=prepare,
constraints=Constraints(max_head_dim=128, allow_mask=False, allow_float32=False, same_device=True),
)
modules/attention/backends/flex.py
+38
View File
@@ -0,0 +1,38 @@
import torch
from modules.logger import log
from modules.attention.registry import AttentionBackend, 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=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs): # 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]
if attn_mask.ndim == 2:
attn_mask = attn_mask.view(attn_mask.shape[0], 1, attn_mask.size[1], 1)
attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv)
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, terminal=True)
modules/attention/backends/sage.py
+53
View File
@@ -0,0 +1,53 @@
import torch
from installer import install
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
install('sageattention')
use_cuda_backend = False
if platform.backend == 'cuda' and torch.cuda.get_device_capability(platform.device) == (8, 6):
use_cuda_backend = True # sm86 needs the cuda backend, sage attention over triton produces NaNs there
try:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
except Exception:
use_cuda_backend = False
if use_cuda_backend:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn_qk_int8_pv_fp16_cuda(
q=query, k=key, v=value,
tensor_layout="HND",
is_causal=is_causal,
sm_scale=scale,
return_lse=False,
pv_accum_dtype="fp32",
)
else:
from sageattention import sageattn
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn(
q=query, k=key, v=value,
attn_mask=None,
dropout_p=0.0,
is_causal=is_causal,
scale=scale,
)
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
return sage_attn_impl(query, key, value, is_causal, scale)
log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
return call
backend = AttentionBackend(
name='sage', label='Sage attention', priority=50, prepare=prepare,
constraints=Constraints(head_dims=frozenset({64, 96, 128}), allow_mask=False, same_device=True),
)
modules/attention/backends/sdnq.py
+28
View File
@@ -0,0 +1,28 @@
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from modules import shared
from sdnq.kernels.triton_atten import sdnq_triton_atten
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
return sdnq_triton_atten(
query=query, key=key, value=value, attn_mask=attn_mask,
is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
matmul_dtype=shared.opts.sdnq_attention_matmul_type,
pv_matmul_dtype=shared.opts.sdnq_attention_pv_matmul_type,
smooth_k=shared.opts.sdnq_attention_smooth_k,
use_hadamard=shared.opts.sdnq_attention_use_hadamard,
hadamard_group_size=shared.opts.sdnq_attention_hadamard_group_size,
use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum,
)
log.debug(f'Torch attention: type="SDNQ attention" matmul={shared.opts.sdnq_attention_matmul_type}:{shared.opts.sdnq_attention_pv_matmul_type} smooth={shared.opts.sdnq_attention_smooth_k} hadamard={shared.opts.sdnq_attention_use_hadamard} fp16_accum={shared.opts.sdnq_attention_use_fp16_accum}')
return call
backend = AttentionBackend(
name='sdnq', label='SDNQ attention', priority=60, prepare=prepare,
constraints=Constraints(min_tokens=32, min_long_side=512, min_heads=2), # sequences of 512 or fewer are text encoders, single-head calls the vae
)
modules/attention/backends/triton_amd.py
+32
View File
@@ -0,0 +1,32 @@
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 modules.flash_attn_triton_amd import interface_fa
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
if scale is None:
scale = query.shape[-1] ** (-0.5)
head_size_og = query.size(3)
if head_size_og % 8 != 0:
query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
out_padded = torch.zeros_like(query)
interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
return out_padded[..., :head_size_og].transpose(1, 2)
log.debug('Torch attention: type="Triton Flash attention"')
return call
backend = AttentionBackend(
name='triton', label='Triton Flash attention', priority=30, prepare=prepare,
constraints=Constraints(max_head_dim=128, allow_mask=False, same_device=True),
platforms=frozenset({'rocm', 'zluda'}),
)
modules/attention/hijacks.py
-242
View File
@@ -1,242 +0,0 @@
from functools import wraps
import torch
from modules import rocm
from modules.logger import log
from installer import install, installed, torch_info
def set_dynamic_attention():
try:
sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention
from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
torch.nn.functional.scaled_dot_product_attention = dynamic_scaled_dot_product_attention
torch_info.set(attention='dynamic')
return sdpa_pre_dyanmic_atten
except Exception as err:
log.error(f'Torch attention: type="dynamic attention" {err}')
return None
def set_sdnq_attention():
try:
from modules import shared
from sdnq.kernels.triton_atten import sdnq_triton_atten
sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sdnq_atten)
def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
if (
query.device.type != "cpu"
and (query.shape[-2] >= 32 and key.shape[-2] >= 32)
and (query.shape[-2] > 512 or key.shape[-2] > 512) # Skip TE
and query.shape[-3] > 1 # Skip VAE
):
return sdnq_triton_atten(
query=query, key=key, value=value, attn_mask=attn_mask,
is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
matmul_dtype=shared.opts.sdnq_attention_matmul_type,
pv_matmul_dtype=shared.opts.sdnq_attention_pv_matmul_type,
smooth_k=shared.opts.sdnq_attention_smooth_k,
use_hadamard=shared.opts.sdnq_attention_use_hadamard,
hadamard_group_size=shared.opts.sdnq_attention_hadamard_group_size,
use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum,
)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_sdnq_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_sdnq_atten
torch_info.set(attention='sdnq')
log.debug(f'Torch attention: type="SDNQ attention" matmul={shared.opts.sdnq_attention_matmul_type}:{shared.opts.sdnq_attention_pv_matmul_type} smooth={shared.opts.sdnq_attention_smooth_k} hadamard={shared.opts.sdnq_attention_use_hadamard} fp16_accum={shared.opts.sdnq_attention_use_fp16_accum}')
except Exception as err:
log.error(f'Torch attention: type="SDNQ attention" {err}')
def set_triton_flash_attention(backend: str):
try:
if backend in {"rocm", "zluda"}: # flash_attn_triton_amd only works with AMD
from modules.flash_attn_triton_amd import interface_fa
sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_triton_flash_atten)
def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_triton = (
query.shape[-1] <= 128
and attn_mask is None
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_triton:
if scale is None:
scale = query.shape[-1] ** (-0.5)
head_size_og = query.size(3)
if head_size_og % 8 != 0:
query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
out_padded = torch.zeros_like(query)
interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
return out_padded[..., :head_size_og].transpose(1, 2)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten
torch_info.set(attention='triton')
log.debug('Torch attention: type="Triton Flash attention"')
except Exception as err:
log.error(f'Torch attention: type="Triton Flash attention" {err}')
def set_flex_attention():
try:
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
def flex_attention_causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
return q_idx >= kv_idx
sdpa_pre_flex_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_flex_atten)
def sdpa_flex_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor: # 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]
if attn_mask.ndim == 2:
attn_mask = attn_mask.view(attn_mask.shape[0], 1, attn_mask.size[1], 1)
attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv)
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(flex_attention_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)
torch.nn.functional.scaled_dot_product_attention = sdpa_flex_atten
torch_info.set(attention="flex")
log.debug('Torch attention: type="Flex attention"')
except Exception as err:
log.error(f'Torch attention: type="Flex attention" {err}')
def set_ck_flash_attention(backend: str, device: torch.device):
try:
if backend == "rocm":
if not installed('flash-attn'):
log.info('Torch attention: type="Flash attention" building...')
agent = rocm.Agent(device)
install(rocm.get_flash_attention_command(agent), reinstall=True)
else:
install('flash-attn')
from flash_attn import flash_attn_func
sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_flash_atten)
def sdpa_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_flash = (
query.shape[-1] <= 128
and attn_mask is None
and query.dtype != torch.float32
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_flash:
is_unsqueezed = False
if query.dim() == 3:
query = query.unsqueeze(0)
is_unsqueezed = True
if key.dim() == 3:
key = key.unsqueeze(0)
if value.dim() == 3:
value = value.unsqueeze(0)
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
if is_unsqueezed:
attn_output = attn_output.squeeze(0)
return attn_output
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
torch_info.set(attention="flash")
log.debug('Torch attention: type="Flash attention"')
except Exception as err:
log.error(f'Torch attention: type="Flash attention" {err}')
def set_sage_attention(backend: str, device: torch.device):
try:
install('sageattention')
use_cuda_backend = False
if (backend == "cuda") and (torch.cuda.get_device_capability(device) == (8, 6)):
use_cuda_backend = True # Detect GPU architecture - sm86 confirmed to need CUDA backend workaround as Sage Attention + Triton causes NaNs
try:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
except Exception:
use_cuda_backend = False
if use_cuda_backend:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn_qk_int8_pv_fp16_cuda(
q=query, k=key, v=value,
tensor_layout="HND",
is_causal=is_causal,
sm_scale=scale,
return_lse=False,
pv_accum_dtype="fp32",
)
else:
from sageattention import sageattn
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn(
q=query, k=key, v=value,
attn_mask=None,
dropout_p=0.0,
is_causal=is_causal,
scale=scale,
)
sdpa_pre_sage_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sage_atten)
def sdpa_sage_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_sage = (
query.shape[-1] in {128, 96, 64}
and attn_mask is None
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_sage:
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
# Call preselected sage attention implementation
return sage_attn_impl(query, key, value, is_causal, scale)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_sage_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_sage_atten
torch_info.set(attention="sage")
log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
except Exception as err:
log.error(f'Torch attention: type="Sage attention" {err}')
modules/attention/registry.py
+91
View File
@@ -0,0 +1,91 @@
"""Declarative backend registry behind the scaled_dot_product_attention router."""
from dataclasses import dataclass, field
from typing import Callable
import torch
AttentionCall = Callable[..., torch.Tensor]
@dataclass(frozen=True)
class Platform:
"""Where the router runs: the devices backend name and the selected device."""
backend: str
device: torch.device | None = None
@dataclass(frozen=True)
class Constraints:
"""Shape, dtype and device conditions a backend serves; a call failing any of them moves on to the next entry."""
allow_cpu: bool = False
allow_mask: bool = True
allow_float32: bool = True
same_device: bool = False
head_dims: frozenset[int] | None = None
max_head_dim: int | None = None
min_tokens: int = 0 # query and key sequences both at least this long
min_long_side: int = 0 # query or key sequence longer than this
min_heads: int = 0
def accepts(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None) -> bool:
if not self.allow_cpu and query.device.type == 'cpu':
return False
if not self.allow_mask and attn_mask is not None:
return False
if not self.allow_float32 and query.dtype == torch.float32:
return False
if self.same_device and (key.device != query.device or value.device != query.device):
return False
head_dim = query.shape[-1]
if self.head_dims is not None and head_dim not in self.head_dims:
return False
if self.max_head_dim is not None and head_dim > self.max_head_dim:
return False
if self.min_tokens and (query.shape[-2] < self.min_tokens or key.shape[-2] < self.min_tokens):
return False
if self.min_long_side and query.shape[-2] <= self.min_long_side and key.shape[-2] <= self.min_long_side:
return False
if self.min_heads and query.shape[-3] < self.min_heads:
return False
return True
@dataclass(frozen=True)
class AttentionBackend:
"""One attention implementation: how to prepare it once and which calls it serves."""
name: str
label: str # the sdp_overrides choice that enables it
priority: int # higher priority entries are tried first
prepare: Callable[[Platform, AttentionCall], AttentionCall | None] # imports and configures the implementation, returns its call or None
constraints: Constraints = field(default_factory=Constraints)
terminal: bool = False # serves every call the entries decline, in place of the original sdpa
platforms: frozenset[str] | None = None # devices backends the implementation exists for, None for all
def available_on(self, platform: Platform) -> bool:
return self.platforms is None or platform.backend in self.platforms
class Registry:
def __init__(self):
self.backends: dict[str, AttentionBackend] = {}
def register(self, backend: AttentionBackend) -> AttentionBackend:
if backend.name in self.backends:
raise ValueError(f'attention backend registered twice: name={backend.name}')
if self.by_label(backend.label) is not None:
raise ValueError(f'attention backend label registered twice: label="{backend.label}"')
self.backends[backend.name] = backend
return backend
def by_label(self, label: str) -> AttentionBackend | None:
return next((backend for backend in self.backends.values() if backend.label == label), None)
def ordered(self) -> list[AttentionBackend]:
"""Backends by ascending priority, the order they are prepared in."""
return sorted(self.backends.values(), key=lambda backend: backend.priority)
def labels(self) -> list[str]:
return [backend.label for backend in self.ordered()]
registry = Registry()
modules/attention/router.py
+90
View File
@@ -0,0 +1,90 @@
"""The single scaled_dot_product_attention entry point over the prepared backends."""
from dataclasses import dataclass
from functools import wraps
import torch
from installer import torch_info
from modules.logger import log
from modules.attention.registry import AttentionBackend, AttentionCall, Platform, Registry, registry as default_registry
@dataclass(frozen=True)
class PlanEntry:
backend: AttentionBackend
call: AttentionCall
@dataclass(frozen=True)
class Plan:
"""The prepared chain for one set of overrides: entries by descending priority, then the terminal or the original sdpa."""
entries: tuple[PlanEntry, ...]
terminal: PlanEntry | None
original: AttentionCall
platform: Platform
labels: tuple[str, ...]
def chain(self) -> list[str]:
names = [entry.backend.name for entry in self.entries]
names.append(self.terminal.backend.name if self.terminal is not None else 'sdpa')
return names
current_plan: Plan | None = None
def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
reg = reg if reg is not None else default_registry
entries: list[PlanEntry] = []
terminal: PlanEntry | None = None
for backend in reg.ordered(): # ascending priority: the last prepared backend is tried first and owns the torch_info record
if backend.label not in labels:
continue
if not backend.available_on(platform):
log.warning(f'Torch attention: type="{backend.label}" not available on backend={platform.backend}')
continue
try:
call = backend.prepare(platform, original)
except Exception as err:
log.error(f'Torch attention: type="{backend.label}" {err}')
continue
if call is None:
continue
entry = PlanEntry(backend=backend, call=call)
if backend.terminal:
terminal = entry
else:
entries.append(entry)
torch_info.set(attention=backend.name)
entries.reverse()
return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
def make_router(plan: Plan) -> AttentionCall:
entries = plan.entries
terminal = plan.terminal.call if plan.terminal is not None else None
original = plan.original
@wraps(original)
def sdpa_router(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs):
for entry in entries:
if entry.backend.constraints.accepts(query, key, value, attn_mask):
return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
if terminal is not None:
return terminal(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **kwargs)
if enable_gqa: # older sdpa signatures and platform wrappers reject the keyword, so it only travels when set
kwargs['enable_gqa'] = enable_gqa
return original(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
return sdpa_router
def install_router(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
"""Prepare the enabled backends and install the router; an empty plan leaves the original sdpa in place."""
global current_plan # pylint: disable=global-statement
plan = build_plan(labels, platform, original, reg)
torch.nn.functional.scaled_dot_product_attention = make_router(plan) if (plan.entries or plan.terminal is not None) else original
current_plan = plan
return plan
def get_plan() -> Plan | None:
return current_plan
modules/devices.py
+1 -21
View File
@@ -515,27 +515,7 @@ def set_sdpa_params():
except Exception as err:
log.warning(f'Torch attention: type="sdpa" {err}')
# Stack hijcaks in reverse order. This gives priority to the last added hijack.
# If the last hijack is not compatible, it will use the one before it and so on.
if 'Dynamic attention' in opts.sdp_overrides:
global sdpa_pre_dyanmic_atten # pylint: disable=global-statement
sdpa_pre_dyanmic_atten = attention.set_dynamic_attention()
if 'Flex attention' in opts.sdp_overrides:
attention.set_flex_attention()
if 'Triton Flash attention' in opts.sdp_overrides:
attention.set_triton_flash_attention(backend)
if 'Flash attention' in opts.sdp_overrides:
attention.set_ck_flash_attention(backend, device)
if 'Sage attention' in opts.sdp_overrides:
attention.set_sage_attention(backend, device)
if 'SDNQ attention' in opts.sdp_overrides:
attention.set_sdnq_attention()
attention.install_router(opts.sdp_overrides, attention.Platform(backend=backend, device=device), sdpa_original)
from importlib.metadata import version
try:
test/test-attention-router.py
+343
View File
@@ -0,0 +1,343 @@
#!/usr/bin/env python
"""
Offline unit tests for the attention router in modules.attention.
Covers:
- plan construction over every subset of the sdp_overrides choices on cuda, rocm, zluda and cpu
against an oracle of the stacking order the closure hijacks used: priority, terminal selection,
platform gating
- gate parity: every backend's declared constraints against a literal transcription of the
predicate its closure carried, over a grid of shapes, dtypes, devices and masks
- every sdp_overrides choice maps to a registered backend and every backend to a choice
- router dispatch: the first accepting entry wins, the terminal receives declined calls, the
original sdpa only receives enable_gqa when it is set
- a backend whose prepare raises is skipped without disturbing the rest
- install_router leaves the original sdpa in place for an empty plan
- the dynamic backend pins the pre-dynamic sdpa the sliced path reads
No running server required. Nothing is moved to the accelerator.
Usage:
python test/test-attention-router.py
"""
import itertools
import logging
import os
import sys
from dataclasses import replace
import torch
script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, script_dir)
os.chdir(script_dir)
os.environ['SD_INSTALL_QUIET'] = '1'
# Bootstrap cmd_args before any module that pulls in shared.py.
import modules.cmd_args # pylint: disable=wrong-import-position
import installer # pylint: disable=wrong-import-position
orig_argv = sys.argv
sys.argv = [sys.argv[0]]
try:
modules.cmd_args.parse_args()
finally:
sys.argv = orig_argv
installer.add_args(modules.cmd_args.parser)
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
stock_sdpa = torch.nn.functional.scaled_dot_product_attention # importing shared installs the configured hijacks in-process
from modules.errors import log # pylint: disable=wrong-import-position
from modules import attention # pylint: disable=wrong-import-position
from modules.attention import router as attention_router # pylint: disable=wrong-import-position
# ============================================================
# Test infrastructure
# ============================================================
results: dict[str, dict] = {}
def category(name: str):
if name not in results:
results[name] = {'passed': 0, 'failed': 0, 'tests': []}
return name
def record(cat: str, passed: bool, name: str, detail: str = ''):
status = 'PASS' if passed else 'FAIL'
results[cat]['passed' if passed else 'failed'] += 1
results[cat]['tests'].append((status, name))
msg = f' {status}: {name}'
if detail:
msg += f' ({detail})'
if passed:
log.info(msg)
else:
log.error(msg)
def run_test(cat: str, fn):
name = fn.__name__
try:
ok = fn()
if ok is False:
record(cat, False, name)
else:
record(cat, True, name)
except AssertionError as e:
record(cat, False, name, str(e))
except Exception as e: # pylint: disable=broad-except
record(cat, False, name, f'exception: {e}')
import traceback
traceback.print_exc()
# ============================================================
# The closure hijacks this router replaces, transcribed
# ============================================================
# devices.set_sdpa_params applied the hijacks in this order; each wrapped the previous, so the
# last applied was tried first. Dynamic and flex replaced the chain end instead of wrapping it.
OLD_ORDER = ['Dynamic attention', 'Flex attention', 'Triton Flash attention', 'Flash attention', 'Sage attention', 'SDNQ attention']
OLD_TERMINALS = {'Dynamic attention', 'Flex attention'}
OLD_NAMES = {
'Dynamic attention': 'dynamic',
'Flex attention': 'flex',
'Triton Flash attention': 'triton',
'Flash attention': 'flash',
'Sage attention': 'sage',
'SDNQ attention': 'sdnq',
}
# mirrors shared_defaults.get_default_modes: five choices everywhere, Triton Flash attention added on rocm and zluda
CHOICES = OLD_ORDER
TRITON_PLATFORMS = {'rocm', 'zluda'}
OLD_GATES = {
'sdnq': lambda q, k, v, m: q.device.type != "cpu" and (q.shape[-2] >= 32 and k.shape[-2] >= 32) and (q.shape[-2] > 512 or k.shape[-2] > 512) and q.shape[-3] > 1,
'triton': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
'flash': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.dtype != torch.float32 and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
'sage': lambda q, k, v, m: q.shape[-1] in {128, 96, 64} and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
}
def oracle_chain(labels, platform_backend):
enabled = [label for label in OLD_ORDER if label in labels]
if platform_backend not in TRITON_PLATFORMS:
enabled = [label for label in enabled if label != 'Triton Flash attention']
terminal = None
entries = []
for label in enabled:
if label in OLD_TERMINALS:
terminal = label
else:
entries.append(label)
entries.reverse()
return [OLD_NAMES[label] for label in entries], (OLD_NAMES[terminal] if terminal else None)
def stub_registry(failing=()):
"""The registered backends with prepares that return a tagged call instead of importing anything."""
reg = attention.Registry()
for backend in attention.registry.ordered():
def prepare(platform, original, name=backend.name): # pylint: disable=unused-argument
if name in failing:
raise RuntimeError(f'{name} unavailable')
def call(*args, **kwargs): # pylint: disable=unused-argument
return name
return call
reg.register(replace(backend, prepare=prepare))
return reg
def shaped(shape, dtype=torch.float16, device='meta'):
"""A tensor of the given shape without allocating it."""
return torch.empty(1, dtype=dtype, device=device).expand(*shape)
def sdpa_stub(**kwargs): # pylint: disable=unused-argument
return 'sdpa'
# ============================================================
# Tests
# ============================================================
def test_plan_matches_stacking_oracle():
level = log.level
log.setLevel(logging.ERROR) # platform gating warns per plan
try:
plans = 0
for platform_backend in ('cuda', 'rocm', 'zluda', 'cpu'):
reg = stub_registry()
platform = attention.Platform(backend=platform_backend)
for count in range(len(OLD_ORDER) + 1):
for labels in itertools.combinations(OLD_ORDER, count):
plan = attention.build_plan(list(labels), platform, sdpa_stub, reg)
expected_entries, expected_terminal = oracle_chain(labels, platform_backend)
got_entries = [entry.backend.name for entry in plan.entries]
got_terminal = plan.terminal.backend.name if plan.terminal is not None else None
assert got_entries == expected_entries, f'{platform_backend} {labels}: entries {got_entries} != {expected_entries}'
assert got_terminal == expected_terminal, f'{platform_backend} {labels}: terminal {got_terminal} != {expected_terminal}'
assert plan.chain() == got_entries + [got_terminal or 'sdpa']
plans += 1
finally:
log.setLevel(level)
log.info(f' {plans} plans match the stacking oracle')
return True
def test_gates_match_transcribed_predicates():
cases = 0
lengths = (16, 32, 512, 513, 4096)
for q_device, kv_device, dtype, heads, q_len, k_len, head_dim, masked in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True)):
q = shaped((1, heads, q_len, head_dim), dtype, q_device)
k = shaped((1, heads, k_len, head_dim), dtype, kv_device)
v = shaped((1, heads, k_len, head_dim), dtype, kv_device)
m = shaped((1, 1, q_len, k_len), torch.bool, q_device) if masked else None
for name, gate in OLD_GATES.items():
expected = bool(gate(q, k, v, m))
got = attention.registry.backends[name].constraints.accepts(q, k, v, m)
assert got == expected, f'{name}: q={tuple(q.shape)} k={tuple(k.shape)} dtype={dtype} devices={q_device}/{kv_device} mask={masked} got={got} expected={expected}'
cases += 1
log.info(f' {cases} gate cases match the transcribed predicates')
return True
def test_terminals_carry_no_gate():
for name in ('dynamic', 'flex'):
backend = attention.registry.backends[name]
assert backend.terminal, name
assert backend.constraints == attention.Constraints(), name
return True
def test_choices_match_backends():
labels = attention.registry.labels()
assert sorted(labels) == sorted(CHOICES), f'registered={labels} choices={CHOICES}'
for label in CHOICES:
assert attention.registry.by_label(label) is not None, label
triton = attention.registry.backends['triton']
assert triton.platforms == frozenset(TRITON_PLATFORMS), triton.platforms
for name, backend in attention.registry.backends.items():
if name != 'triton':
assert backend.platforms is None, name
return True
def test_router_dispatch_prefers_priority_then_terminal_then_original():
calls = []
def original(**kwargs):
calls.append(('sdpa', kwargs))
return 'sdpa'
reg = attention.Registry()
def add(name, constraints, priority, terminal=False):
def prepare(platform, orig): # pylint: disable=unused-argument
def call(*args, **kwargs): # pylint: disable=unused-argument
calls.append((name, kwargs))
return name
return call
reg.register(attention.AttentionBackend(name=name, label=f'{name} attention', priority=priority, prepare=prepare, constraints=constraints, terminal=terminal))
add('narrow', attention.Constraints(head_dims=frozenset({64})), priority=20)
add('wide', attention.Constraints(), priority=10)
platform = attention.Platform(backend='cuda')
router = attention_router.make_router(attention.build_plan(['narrow attention', 'wide attention'], platform, original, reg))
q64 = shaped((1, 8, 128, 64))
q128 = shaped((1, 8, 128, 128))
cpu = shaped((1, 8, 128, 64), device='cpu')
assert router(q64, q64, q64) == 'narrow'
assert router(q128, q128, q128) == 'wide'
assert router(cpu, cpu, cpu) == 'sdpa'
assert 'enable_gqa' not in calls[-1][1], calls[-1]
assert router(cpu, cpu, cpu, enable_gqa=True) == 'sdpa'
assert calls[-1][1].get('enable_gqa') is True, calls[-1]
add('term', attention.Constraints(), priority=5, terminal=True)
router = attention_router.make_router(attention.build_plan(['narrow attention', 'term attention'], platform, original, reg))
assert router(q64, q64, q64) == 'narrow'
assert router(cpu, cpu, cpu, extra=1) == 'term'
assert calls[-1][1].get('extra') == 1 and calls[-1][1].get('enable_gqa') is False, calls[-1]
return True
def test_prepare_failure_skips_backend():
reg = stub_registry(failing=('sage',))
plan = attention.build_plan(['Sage attention', 'SDNQ attention', 'Flash attention'], attention.Platform(backend='cuda'), sdpa_stub, reg)
assert [entry.backend.name for entry in plan.entries] == ['sdnq', 'flash'], plan.chain()
return True
def test_install_router_keeps_original_for_empty_plan():
saved = torch.nn.functional.scaled_dot_product_attention
saved_plan = attention_router.current_plan
try:
platform = attention.Platform(backend='cuda')
plan = attention.install_router([], platform, sdpa_stub, stub_registry())
assert torch.nn.functional.scaled_dot_product_attention is sdpa_stub
assert plan.chain() == ['sdpa'], plan.chain()
plan = attention.install_router(['SDNQ attention', 'Sage attention'], platform, sdpa_stub, stub_registry())
assert torch.nn.functional.scaled_dot_product_attention is not sdpa_stub
assert plan.chain() == ['sdnq', 'sage', 'sdpa'], plan.chain()
assert attention.get_plan() is plan
finally:
torch.nn.functional.scaled_dot_product_attention = saved
attention_router.current_plan = saved_plan
return True
def test_dynamic_backend_pins_pre_dynamic_sdpa():
from modules import devices
saved = devices.sdpa_pre_dyanmic_atten
try:
call = attention.registry.backends['dynamic'].prepare(attention.Platform(backend='cuda'), sdpa_stub)
from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
assert call is dynamic_scaled_dot_product_attention
assert devices.sdpa_pre_dyanmic_atten is sdpa_stub
finally:
devices.sdpa_pre_dyanmic_atten = saved
return True
def run_all():
log.warning('=== attention router ===')
cat = category('router')
for fn in [
test_plan_matches_stacking_oracle,
test_gates_match_transcribed_predicates,
test_terminals_carry_no_gate,
test_choices_match_backends,
test_router_dispatch_prefers_priority_then_terminal_then_original,
test_prepare_failure_skips_backend,
test_install_router_keeps_original_for_empty_plan,
test_dynamic_backend_pins_pre_dynamic_sdpa,
]:
run_test(cat, fn)
log.warning('=== Results ===')
total_passed = 0
total_failed = 0
for cat_name, info in results.items():
ok = info['failed'] == 0
status = 'PASS' if ok else 'FAIL'
log.info(f" {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]")
total_passed += info['passed']
total_failed += info['failed']
log.warning(f'Total: {total_passed} passed, {total_failed} failed')
return total_failed == 0
if __name__ == '__main__':
import time
t0 = time.time()
ok = run_all()
torch.nn.functional.scaled_dot_product_attention = stock_sdpa
log.warning(f'Total time: {time.time() - t0:.2f}s')
sys.exit(0 if ok else 1)