feat(attention): exclude known bad models from sparse attention

sparse_attention_exclude is a comma separated denylist matched case insensitively against the architecture, the pipeline class and the denoiser class, so one entry works whichever name is to hand. It resolves once per model rather than per call and declines with a log line. Seeded with CosmosTransformer3DModel, the transformer Anima runs, which returns banded noise at every budget tested against a sound dense baseline; listing the class rather than the architecture covers the other models built on it, none of which have been checked.

modules/attention/context.py

@@ -11,7 +11,7 @@ class GenerationContext:
     step: int = 0 # index of the denoiser forward about to run
     steps: int = 0 # forwards in the current pass
     forwards: int = 0
-    model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, telemetry only
+    model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, for telemetry and the sparse exclusion list
     step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph
     layout: object | None = None # TokenLayout published by whoever knows the packing, None until something does
 

modules/attention/sparse/stage.py

@@ -18,7 +18,7 @@ pattern = os.environ.get('SD_SPARSE_PATTERN', 'adaptive').strip().lower()
 DEFAULT_MIN_TOKENS = 8192
 
 # settings the stage reads, so a change to any of them rebuilds the chain
-OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared')
+OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared', 'sparse_attention_exclude')
 
 
 @dataclass(frozen=True)
@@ -29,6 +29,7 @@ class StageOptions:
     schedule_steps: int = 0
     schedule_bump: float = 0.0
     head_shared: bool = False
+    exclude: tuple = () # architectures, pipeline classes or denoiser classes that stay dense
 
 
 def read_options() -> StageOptions:
@@ -41,9 +42,21 @@ def read_options() -> StageOptions:
         schedule_steps=int(getattr(opts, 'sparse_attention_schedule_steps', 0)),
         schedule_bump=float(getattr(opts, 'sparse_attention_schedule_bump', 0)) / 100.0,
         head_shared=bool(getattr(opts, 'sparse_attention_head_shared', False)),
+        exclude=parse_exclusions(getattr(opts, 'sparse_attention_exclude', '')),
     )
 
 
+def parse_exclusions(raw) -> tuple:
+    """The exclusion list as lowercase entries, each naming an architecture, a pipeline class or a denoiser class."""
+    return tuple(entry.strip().lower() for entry in str(raw or '').split(',') if len(entry.strip()) > 0)
+
+
+def match_exclusion(model_key, arch, exclude: tuple) -> str:
+    """The entry the loaded model matches, empty when none of them do."""
+    names = {str(name).lower() for name in (*(model_key or ()), arch) if name}
+    return next((entry for entry in exclude if entry in names), '')
+
+
 def resolve_layout(seq: int, reported: set) -> layout_mod.TokenLayout:
     """The published layout when there is one, otherwise sparsify the whole sequence and say so once."""
     published = context.current.layout
@@ -65,6 +78,7 @@ def make_stage(options: StageOptions):
     notified: set = set()
     cache: dict = {}
     static: dict = {}
+    excluded: dict = {}
 
     def static_selection(query, key, spec, pins, drops, cache_key):
         """A density matched band, built once per geometry, honoring the same layout pins so the control differs from the selector only in how it chooses video tiles."""
@@ -100,6 +114,19 @@ def make_stage(options: StageOptions):
             cache[key] = table
         return table[min(state.step, len(table) - 1)] if table else options.budget
 
+    def on_exclusion_list(state) -> bool:
+        """Whether the loaded model is excluded, resolved once per model since the answer cannot change within one."""
+        if not options.exclude:
+            return False
+        hit = excluded.get(state.model_key)
+        if hit is None:
+            from modules import shared
+            hit = match_exclusion(state.model_key, getattr(shared, 'sd_model_type', None), options.exclude)
+            excluded[state.model_key] = hit
+            if hit: # an enabled setting that cannot act says so rather than doing nothing quietly
+                log.info(f'Sparse attention: "{hit}" is on the exclusion list; attention stays dense')
+        return len(hit) > 0
+
     def decline(reason: str):
         stage.last_skip = reason
         return None
@@ -108,6 +135,8 @@ def make_stage(options: StageOptions):
         state = context.current
         if state.role != 'transformer' or not state.active:
             return decline('not the denoiser')
+        if on_exclusion_list(state):
+            return decline('excluded')
         if is_causal: # the selection keeps the diagonal but encodes no causality
             return decline('causal')
         if attn_mask is not None and 'masked_block' not in caps: # flex would need a mask_mod to combine the two

modules/ui_definitions.py

@@ -269,6 +269,7 @@ def create_settings(cmd_opts):
         "sparse_attention_schedule_steps": OptionInfo(0, "Sparse Attention dense steps", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
         "sparse_attention_schedule_bump": OptionInfo(30, "Sparse Attention dense step bonus", gr.Slider, {"minimum": 0, "maximum": 70, "step": 5}),
         "sparse_attention_head_shared": OptionInfo(False, "Sparse Attention share selection across heads", gr.Checkbox),
+        "sparse_attention_exclude": OptionInfo("CosmosTransformer3DModel", "Sparse Attention excluded models", gr.Textbox),
 
         "hf_attention_sep": OptionInfo("<h2>Attention Dispatcher</h2>", "", gr.HTML),
         "hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox),

test/test-attention-sparse.py

@@ -464,6 +464,44 @@ def test_minimum_sequence_of_zero_sparsifies_everything():
     return with_context(checks)
 
 
+def test_exclusion_list_parsing():
+    from modules.attention.sparse import stage as stage_mod
+    assert stage_mod.parse_exclusions(' Anima , CosmosTransformer3DModel ,, ') == ('anima', 'cosmostransformer3dmodel')
+    assert stage_mod.parse_exclusions('') == ()
+    assert stage_mod.parse_exclusions(None) == ()
+    return True
+
+
+def test_an_exclusion_entry_matches_any_of_the_three_names():
+    """One entry, matched against whichever of the architecture, pipeline class or denoiser class the user knew."""
+    from modules.attention.sparse import stage as stage_mod
+    key = ('AnimaTextToImagePipeline', 'CosmosTransformer3DModel')
+    for entry in ('cosmostransformer3dmodel', 'animatexttoimagepipeline', 'anima'):
+        assert stage_mod.match_exclusion(key, 'anima', (entry,)) == entry, entry
+    assert stage_mod.match_exclusion(key, 'anima', ('krea2',)) == '', 'an unlisted model matches nothing'
+    assert stage_mod.match_exclusion(None, 'anima', ('anima',)) == 'anima', 'the architecture stands in when no model key is published'
+    assert stage_mod.match_exclusion(('Pipe', None), None, ('pipe',)) == 'pipe', 'an absent denoiser class is skipped, not matched'
+    return True
+
+
+def test_an_excluded_model_stays_dense():
+    from modules.attention import context as ctx
+    from modules.attention.sparse import stage as stage_mod
+    listed = stage_mod.make_stage(stage_options(exclude=('cosmostransformer3dmodel',)))
+    unlisted = stage_mod.make_stage(stage_options(exclude=('somethingelse',)))
+    q, k, v = qkv(heads=2, seq=2048)
+
+    def checks():
+        ctx.current.model_key = ('AnimaTextToImagePipeline', 'CosmosTransformer3DModel')
+        assert listed(q, k, v, None, False) is None, 'a listed denoiser class stays dense'
+        assert listed.last_skip == 'excluded', listed.last_skip
+        assert unlisted(q, k, v, None, False) is not None, 'a list that matches nothing changes nothing'
+        ctx.current.model_key = ('Krea2Pipeline', 'Krea2Transformer2DModel')
+        assert listed(q, k, v, None, False) is not None, 'the exclusion applies to the listed model, not to every model'
+        return True
+    return with_context(checks)
+
+
 def test_published_segments_reach_the_stage():
     from modules.attention.sparse import layout as layout_mod
     from modules.attention.sparse import stage as stage_mod
@@ -566,6 +604,9 @@ def run_all():
     for fn in [
         test_stage_is_none_when_disabled_or_at_full_budget,
         test_stage_gates,
+        test_exclusion_list_parsing,
+        test_an_exclusion_entry_matches_any_of_the_three_names,
+        test_an_excluded_model_stays_dense,
         test_minimum_sequence_of_zero_sparsifies_everything,
         test_stage_follows_the_step_schedule,
         test_published_segments_reach_the_stage,

ui/locale/locale_en.json

@@ -1471,12 +1471,13 @@
     {"id":"","label":"SDNQ Attention PV MatMul type","localized":"","hint":"Precision the probability-value matmul is computed in, the second of the two matmuls in attention. Choices match <b><i>SDNQ Attention MatMul type</i></b>.<br><br>Quantizing this one as well takes out the floating point work the first setting leaves behind, and it is the more delicate of the two: its inputs are already normalized probabilities, and the small ones among them carry the fine detail.<br><b>disabled</b> keeps this matmul in the model's own precision.<br><br>Default disabled.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention Hadamard Group Size","localized":"","hint":"Width of the Hadamard rotation in channels. Wider groups mix more channels together and spread outliers further.<br><br>Clamped to the head dimension of the running model, rounded down to a power of two that divides it. On a model with 64 or 128 channels per head the upper part of this range resolves to that head dimension rather than to the number shown. Rotation is skipped below 4.<br>Applies while <b><i>SDNQ Attention use Hadamard</i></b> is enabled.<br><br>Default 256.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"},
-    {"id":"","label":"Sparse Attention","localized":"","hint":"Computes attention over a subset of the key tiles instead of all of them. Attention cost grows with the square of the sequence length, so on long sequences it dominates generation time, and skipping the tiles that contribute least buys much of it back.<br>The saving grows with sequence length: negligible on a short sequence, useful at high resolution, largest on video.<br><br>Requires an attention backend that accepts a block mask, <b>SDNQ attention</b> or <b>Flex attention</b> in <b><i>SDP overrides</i></b>; with neither enabled a warning is logged and attention stays dense. It also stays dense below <b><i>Sparse Attention minimum sequence</i></b>.<br>Not every architecture tolerates a reduced key set. If output breaks up rather than merely softens, leave this off for that model.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention","localized":"","hint":"Computes attention over a subset of the key tiles instead of all of them. Attention cost grows with the square of the sequence length, so on long sequences it dominates generation time, and skipping the tiles that contribute least buys much of it back.<br>The saving grows with sequence length: negligible on a short sequence, useful at high resolution, largest on video.<br><br>Requires an attention backend that accepts a block mask, <b>SDNQ attention</b> or <b>Flex attention</b> in <b><i>SDP overrides</i></b>; with neither enabled a warning is logged and attention stays dense. It also stays dense below <b><i>Sparse Attention minimum sequence</i></b>.<br>Not every architecture tolerates a reduced key set. The ones known to break are listed in <b><i>Sparse Attention excluded models</i></b> and stay dense; a model that breaks up rather than merely softening belongs on that list.<br><br>Disabled by default.","ui":"settings_cuda"},
     {"id":"","label":"Sparse Attention KV budget","localized":"","hint":"Percentage of the eligible key tiles each query tile keeps. Lower budgets are faster and coarser, and the cost shows first in fine detail and in consistency across the image.<br>Text, conditioning and audio tokens are always kept, as are the tiles on the diagonal, so the budget applies only to the bulk image or video tokens.<br><br><b>100</b> keeps every tile, which is ordinary dense attention and a convenient comparison run.<br>Default 30.","ui":"settings_cuda"},
     {"id":"","label":"Sparse Attention minimum sequence","localized":"","hint":"Shortest sequence that is sparsified. Below it attention stays dense, because choosing the tiles costs more than skipping them saves.<br>Sequence length is not resolution: a diffusion transformer sees roughly (width/16) x (height/16) tokens for an image, so 1024x1024 is about 4k tokens and 2048x2048 about 16k, and video multiplies that by the frame count.<br><br><b>0</b> sparsifies every sequence that reaches the stage.<br>Default 8192, around 1450x1450 for an image.","ui":"settings_cuda"},
     {"id":"","label":"Sparse Attention dense steps","localized":"","hint":"Number of steps at the start and end of sampling that receive a larger budget, where composition and fine detail are set. Counted at each end, and capped at half the run.<br><br>Only takes effect when <b><i>Sparse Attention dense step bonus</i></b> is above 0.<br><b>0</b> applies one budget to every step.<br>Default 0.","ui":"settings_cuda"},
     {"id":"","label":"Sparse Attention dense step bonus","localized":"","hint":"Percentage points added to the budget on the dense steps, capped at 100.<br><br>Only takes effect when <b><i>Sparse Attention dense steps</i></b> is above 0.<br>Default 30, so a budget of 30 rises to 60 on those steps.","ui":"settings_cuda"},
-    {"id":"","label":"Sparse Attention share selection across heads","localized":"","hint":"Computes one selection for all attention heads rather than one per head, by averaging the heads before scoring. Cheaper to select and coarser in what it keeps, since heads that attend to different regions are served by a single compromise.<br>Worth trying when selection is itself a visible share of the cost, which happens on models with many heads.<br><br>Disabled by default.","ui":"settings_cuda"}
+    {"id":"","label":"Sparse Attention share selection across heads","localized":"","hint":"Computes one selection for all attention heads rather than one per head, by averaging the heads before scoring. Cheaper to select and coarser in what it keeps, since heads that attend to different regions are served by a single compromise.<br>Worth trying when selection is itself a visible share of the cost, which happens on models with many heads.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention excluded models","localized":"","hint":"Models that stay dense no matter how the rest of this section is set. Comma separated, matched case insensitively against the architecture, the pipeline class and the denoiser class, so whichever of those names is to hand works as an entry.<br>The class names appear in the model load log; the architecture is the short name used elsewhere in the settings, such as <b>f1</b> or <b>anima</b>.<br><br>Listed by default is <b>CosmosTransformer3DModel</b>, the transformer Anima runs, which returns banded noise at every budget tested against a sound dense baseline. The class is listed rather than the architecture because the other models built on it have not been checked.<br>A model whose output breaks up rather than merely softening belongs here.<br><br>Default <b>CosmosTransformer3DModel</b>.","ui":"settings_cuda"}
   ],
   "t": [
     {"id":"txt2img_nav","label":"T2I","localized":"","hint":"Create image from text<br>Legacy interface that mimics original text-to-image interface and behavior"},