refactor(loaders): SLA composes with the attention backend instead of replacing it

It is a toggle now, not a value in the list: SLA takes H3's packed self-attention
and the chosen backend takes the text refiner, masked calls and dense-last steps.
When SLA does not take, the chosen backend is still installed on its own.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

CHANGELOG.md

@@ -7,9 +7,12 @@
   Loader's `attention` list rather than a node to wire. ComfyUI ships no sparse-attention
   backend for H3, which is why lightx2v's SLA turbo LoRA gives no speedup on its own — the
   LoRA is the adaptation to sparsity, not the acceleration. Roughly 3.7x the attention
-  throughput at 768p/15s. Six settings, folded under Advanced and validated at their
-  defaults; offered only where CUDA and Triton can run it, and ignored on any model that is
-  not H3. Kernel and block map vendored from LightX2V (Apache-2.0) via
+  throughput at 768p/15s. It is a toggle beside the attention backend, not one of its
+  values: SLA takes H3's long packed self-attention and the chosen backend (sage3, int8,
+  flash) takes the text refiner, masked calls and any trailing dense steps, so the two
+  compose. Five settings folded under Advanced and validated at their defaults; skipped
+  with a stated reason on anything that is not H3 or without Triton, leaving the chosen
+  backend installed on its own. Kernel and block map vendored from LightX2V (Apache-2.0) via
   ComfyUI-H3-SLA-Attention (MIT).
 
 ### Fixed

loaders.py

@@ -61,8 +61,10 @@ def attention_choices():
     'default' leaves the model on whatever ComfyUI itself selected (the --use-sage-attention
     / --use-flash-attention launch flags), which is why it stays the default here.
 
-    sla_h3 is not in the registry — it is FunPack's own block-sparse kernel, listed only
-    where CUDA and Triton can run it, on the same principle as the rest of this list.
+    SLA is deliberately NOT in this list: it is not a rival backend but a layer above one.
+    It handles only H3's long packed self-attention and hands everything else — the text
+    refiner, masked calls, the trailing dense steps — to whichever backend is chosen here,
+    so the two compose instead of displacing each other. Its switch is `sla`.
     """
     names = []
     try:
@@ -70,8 +72,6 @@ def attention_choices():
         names = sorted(REGISTERED_ATTENTION_FUNCTIONS.keys())
     except Exception:  # noqa: BLE001 - older ComfyUI without the registry
         pass
-    if sla_attention.sla_available():
-        names.append(sla_attention.SLA_NAME)
     return ["default"] + names
 
 
@@ -141,17 +141,26 @@ class FunPackDiffusionModelLoader:
                     "default": "default",
                     "tooltip": "Attention backend for THIS model. 'default' uses whatever "
                                "ComfyUI was launched with. Only backends installed on this "
-                               "machine are listed. 'sla_h3' is block-sparse attention for "
-                               "MiniMax H3 — the inference path lightx2v's SLA turbo LoRA was "
-                               "distilled against, which is why that LoRA gives no speedup "
-                               "without it. Roughly 3.7x the attention throughput at 768p/15s; "
-                               "ignored on any other model family."}),
+                               "machine are listed. Composes with `sla`: sparse attention "
+                               "handles H3's long packed sequence, this backend handles "
+                               "everything else."}),
             },
             "optional": {
                 "fp16_accumulation": ("BOOLEAN", {
                     "default": False,
                     "tooltip": "Accumulate fp16 matmuls in fp16. Faster on recent NVIDIA cards; "
                                "needs a torch build that supports it, ignored otherwise."}),
+                "sla": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "Block-sparse attention for MiniMax H3 — the inference path "
+                               "lightx2v's SLA turbo LoRA was distilled against, which is why "
+                               "that LoRA gives no speedup on its own. Roughly 3.7x the "
+                               "attention throughput at 768p/15s. It runs ALONGSIDE the "
+                               "attention backend above rather than replacing it: SLA takes "
+                               "H3's long packed self-attention, the chosen backend (sage3, "
+                               "int8, flash…) takes the text refiner, masked calls and any "
+                               "trailing dense steps. Skipped, with the reason, on anything "
+                               "that is not MiniMax H3 or on a machine without Triton."}),
                 # SLA settings. Every one is validated at its default (see sla_attention);
                 # they are here because a knob whose value was measured is still a knob.
                 "sla_sparsity": ("FLOAT", {
@@ -187,13 +196,6 @@ class FunPackDiffusionModelLoader:
                     "tooltip": "Sequences shorter than this stay dense. Guards the short text "
                                "refiner, which must never be sparsified, and low-resolution "
                                "runs where block selection costs more than it saves."}),
-                "sla_enabled": ("BOOLEAN", {
-                    "default": sla_attention.SLA_DEFAULTS["enabled"],
-                    "advanced": True,
-                    "tooltip": "Off runs dense attention while leaving sla_h3 selected — a "
-                               "like-for-like speed and quality baseline that keeps the "
-                               "settings you were testing, instead of losing them to the "
-                               "backend dropdown."}),
                 "sla_dense_last_steps": ("INT", {
                     "default": sla_attention.SLA_DEFAULTS["dense_last_steps"],
                     "min": 0, "max": 8,
@@ -206,7 +208,7 @@ class FunPackDiffusionModelLoader:
     def load_model(self, model_name, weight_dtype, compute_dtype, attention,
                    fp16_accumulation=False, sla_sparsity=None, sla_block_size=None,
                    sla_protect_audio=None, sla_min_seq_len=None, sla_dense_last_steps=None,
-                   sla_enabled=None):
+                   sla=False):
         notes = [f"FunPack Diffusion Model Loader | {model_name}"]
 
         accum = set_fp16_accumulation(fp16_accumulation)
@@ -235,23 +237,29 @@ class FunPackDiffusionModelLoader:
         elif dtype is not None:
             notes.append("compute dtype: unsupported by this ComfyUI, ignored")
 
-        if attention == sla_attention.SLA_NAME:
-            # Its own install: SLA needs a per-step wrapper as well as the override, to
-            # read the packed layout (which never reaches the attention call site) and to
-            # report at the end of a run whether it actually fired.
-            model, note = sla_attention.install_sla(
+        override = attention_override(attention)
+
+        # SLA wraps the chosen backend rather than replacing it: there is one override
+        # slot, so the backend goes in as SLA's dense fall-through. When SLA does not
+        # take (not H3, no Triton, switched off) the backend is installed on its own —
+        # asking for sparse attention must never cost the backend you picked.
+        installed = False
+        if sla:
+            model, note, installed = sla_attention.install_sla(
                 model,
                 sparsity_ratio=sla_sparsity, block_size=sla_block_size,
                 min_seq_len=sla_min_seq_len, dense_last_steps=sla_dense_last_steps,
-                protect_audio=sla_protect_audio, enabled=sla_enabled)
+                protect_audio=sla_protect_audio,
+                dense_fn=override, dense_label=attention)
             notes.append(note)
-        else:
-            override = attention_override(attention)
+        if not installed:
             if override is not None:
                 model.model_options.setdefault("transformer_options", {})["optimized_attention_override"] = override
                 notes.append(f"attention: {attention}")
             else:
                 notes.append("attention: default (as launched)")
+        else:
+            notes.append(f"attention: {attention} (dense calls)")
 
         return (model, "\n".join(notes))
 

sla_attention.py

@@ -104,8 +104,16 @@ def _summarise(state, sparsity, blkq, blkk):
                         state["failed"])
 
 
-def make_override(state, sparsity_ratio, blkq, blkk, min_seq_len, protect_audio=True):
-    """The `optimized_attention_override` wrap_attn hands every attention call to."""
+def make_override(state, sparsity_ratio, blkq, blkk, min_seq_len, protect_audio=True,
+                  dense_fn=None, dense_label=None):
+    """The `optimized_attention_override` wrap_attn hands every attention call to.
+
+    `dense_fn` is another override (same signature) handling everything SLA does not:
+    the text refiner, masked calls, the trailing dense steps, any non-H3 model. There is
+    only one override slot, so without this, choosing SLA would silently discard the
+    backend the user picked and drop those calls onto whatever ComfyUI was launched with.
+    Sparse where sparsity applies, the chosen backend everywhere else.
+    """
     torch = _torch()
     ok_dtypes = (torch.bfloat16, torch.float16)
     topk_ratio = 1.0 - sparsity_ratio
@@ -114,12 +122,13 @@ def make_override(state, sparsity_ratio, blkq, blkk, min_seq_len, protect_audio=
                  skip_reshape=False, skip_output_reshape=False, **kwargs):
         def dense():
             state["dense"] += 1
-            return func(q, k, v, heads, mask=mask, attn_precision=attn_precision,
-                        skip_reshape=skip_reshape,
-                        skip_output_reshape=skip_output_reshape, **kwargs)
+            run = (lambda *a, **kw: dense_fn(func, *a, **kw)) if dense_fn is not None else func
+            return run(q, k, v, heads, mask=mask, attn_precision=attn_precision,
+                       skip_reshape=skip_reshape,
+                       skip_output_reshape=skip_output_reshape, **kwargs)
 
         if state["backend"] is None:
-            state["backend"] = getattr(func, "__name__", repr(func))
+            state["backend"] = dense_label or getattr(func, "__name__", repr(func))
 
         to = kwargs.get("transformer_options") or {}
 
@@ -239,12 +248,16 @@ def make_wrapper(state, sparsity_ratio, blkq, blkk, dense_last_steps):
 
 
 def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None,
-                dense_last_steps=None, protect_audio=None, enabled=None):
-    """Give `model` block-sparse H3 attention. Returns (model, status line).
+                dense_last_steps=None, protect_audio=None, enabled=None,
+                dense_fn=None, dense_label=None):
+    """Give `model` block-sparse H3 attention. Returns (model, status line, installed).
 
-    Weights are untouched: this installs an attention override and a per-step wrapper
-    on a clone. A model that is not MiniMax H3 comes back unchanged and says so —
-    the sparsity is only safe at the ratio H3's turbo LoRA was distilled to tolerate.
+    Weights are untouched: this installs an attention override and a per-step wrapper on
+    a clone. A model that is not MiniMax H3 comes back unchanged and says so — the
+    sparsity is only safe at the ratio H3's turbo LoRA was distilled to tolerate.
+
+    `installed` is False whenever SLA did not take, so the caller can fall back to
+    installing the chosen backend on its own rather than leaving the model with nothing.
     """
     def pick(value, key):
         return SLA_DEFAULTS[key] if value is None else value
@@ -258,13 +271,11 @@ def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None,
     if not bool(pick(enabled, "enabled")):
         # A dense baseline without touching the backend choice, so an A/B is one click
         # and the settings you were testing are still on the node when you switch back.
-        return model, "attention: sla_h3 off (dense baseline)"
+        return model, "SLA: off (dense baseline)", False
     if not is_h3_model(model):
-        return model, ("attention: sla_h3 requested but this is not a MiniMax H3 model — "
-                       "left on the launched backend")
+        return model, ("SLA: skipped — not a MiniMax H3 model"), False
     if not sla_available():
-        return model, ("attention: sla_h3 requested but this machine has no CUDA+Triton — "
-                       "left on the launched backend")
+        return model, ("SLA: skipped — this machine has no CUDA+Triton"), False
 
     # BLKK=64 is not a typo. On sm_120 the 128x128 tile needs 160 KB of shared memory
     # against a ~99 KB limit and cannot launch at all; 128x64 both fits and measured
@@ -275,12 +286,14 @@ def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None,
     patched = model.clone()
     to = patched.model_options.get("transformer_options", {}).copy()
     to["optimized_attention_override"] = make_override(
-        state, sparsity_ratio, blkq, blkk, min_seq_len, protect_audio)
+        state, sparsity_ratio, blkq, blkk, min_seq_len, protect_audio,
+        dense_fn=dense_fn, dense_label=dense_label)
     patched.model_options["transformer_options"] = to
     patched.add_wrapper_with_key(
         "diffusion_model", "funpack_sla_state",
         make_wrapper(state, sparsity_ratio, blkq, blkk, dense_last_steps))
 
-    return patched, (f"attention: sla_h3 | sparsity={sparsity_ratio:.2f} BLK={blkq}x{blkk} "
+    return patched, (f"SLA on | sparsity={sparsity_ratio:.2f} BLK={blkq}x{blkk} "
                      f"min_seq_len={min_seq_len} dense_last_steps={dense_last_steps} "
-                     f"protect_audio={protect_audio}")
+                     f"protect_audio={protect_audio} | dense calls -> "
+                     f"{dense_label or 'as launched'}"), True

tests/test_loaders.py

@@ -224,9 +224,10 @@ def test_the_lora_list_is_a_picker_not_a_table_of_dropdowns():
     assert spec["optional"]["lora_stack"][1]["advanced"] is True
 
 
-def test_the_sla_backend_is_offered_only_where_it_can_run():
-    """Same rule as the rest of the list: a backend this machine cannot run is not
-    offered, rather than listed and failing at generation."""
+def test_sla_is_a_toggle_not_a_rival_backend():
+    """It layers ABOVE a backend — sparse for H3's packed sequence, the chosen backend for
+    everything else — so listing it as a choice would force an either/or that is not real."""
     import sla_attention
-    choices = loaders.attention_choices()
-    assert (sla_attention.SLA_NAME in choices) is sla_attention.sla_available()
+    assert sla_attention.SLA_NAME not in loaders.attention_choices()
+    spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES()
+    assert spec["optional"]["sla"][1]["default"] is False

tests/test_sla_attention.py

@@ -214,18 +214,18 @@ def test_sla_refuses_a_model_that_is_not_h3_and_says_so():
     """Head shape alone matches LTX too, and sparsifying LTX attention is a quality loss
     with no LoRA compensating for it. Silence would read as "it is on"."""
     model = _Patcher("LTXVModel")
-    out, note = sla.install_sla(model)
-    assert out is model
+    out, note, installed = sla.install_sla(model)
+    assert out is model and installed is False
     assert "not a MiniMax H3 model" in note
 
 
 def test_sla_reports_when_the_machine_cannot_run_it():
     model = _Patcher("MiniMaxH3Model")
-    out, note = sla.install_sla(model)
+    out, note, installed = sla.install_sla(model)
     if sla.sla_available():
-        assert out is not model and "sparsity=0.90" in note
+        assert installed is True and out is not model and "sparsity=0.90" in note
     else:
-        assert out is model and "CUDA+Triton" in note
+        assert installed is False and out is model and "CUDA+Triton" in note
 
 
 def test_the_defaults_are_the_validated_ones():
@@ -238,11 +238,40 @@ def test_the_defaults_are_the_validated_ones():
 def test_turning_sla_off_leaves_the_settings_in_place():
     """A dense A/B baseline must not cost the settings being tested."""
     model = _Patcher("MiniMaxH3Model")
-    out, note = sla.install_sla(model, sparsity_ratio=0.85, enabled=False)
-    assert out is model
+    out, note, installed = sla.install_sla(model, sparsity_ratio=0.85, enabled=False)
+    assert out is model and installed is False
     assert "off (dense baseline)" in note
 
 
+# ── composing with a chosen backend ───────────────────────────────────────────
+
+def test_dense_calls_go_to_the_chosen_backend_not_the_launched_one():
+    """There is one override slot. Without this, choosing SLA would silently discard the
+    backend the user picked and drop the text refiner onto whatever ComfyUI launched with."""
+    seen = []
+
+    def chosen(func, q, k, v, heads, **kw):
+        seen.append("chosen")
+        return func(q, k, v, heads, **kw)
+
+    state = sla.new_state()
+    ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=8192,
+                           dense_fn=chosen, dense_label="sage3")
+    q = torch.randn(1, H, 512, D)              # short -> dense fall-through
+    out = _call(ov, q, q.clone(), q.clone())
+    assert seen == ["chosen"]
+    assert state["backend"] == "sage3"         # what the log will name
+    assert out.shape == (1, 512, H * D)
+
+
+def test_without_a_chosen_backend_dense_still_runs():
+    state = sla.new_state()
+    ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=8192)
+    q = torch.randn(1, H, 512, D)
+    assert _call(ov, q, q.clone(), q.clone()).shape == (1, 512, H * D)
+    assert state["dense"] == 1
+
+
 # ── the kernel itself ─────────────────────────────────────────────────────────
 
 pytestmark_cuda = pytest.mark.skipif(not (CUDA and HAS_TRITON), reason="needs CUDA + Triton")