From a4839c36f265bc0c5566cb08ec1f7f44ed65edf0 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:36:01 +0300 Subject: [PATCH] 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 --- CHANGELOG.md | 9 ++++-- loaders.py | 58 +++++++++++++++++++++---------------- sla_attention.py | 51 ++++++++++++++++++++------------ tests/test_loaders.py | 11 +++---- tests/test_sla_attention.py | 43 ++++++++++++++++++++++----- 5 files changed, 113 insertions(+), 59 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 363e0db8..631d9737 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/loaders.py b/loaders.py index 86fcbda7..24cf0716 100644 --- a/loaders.py +++ b/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)) diff --git a/sla_attention.py b/sla_attention.py index 079bfb0f..25b5ba2b 100644 --- a/sla_attention.py +++ b/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 diff --git a/tests/test_loaders.py b/tests/test_loaders.py index ec47af49..191667c1 100644 --- a/tests/test_loaders.py +++ b/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 diff --git a/tests/test_sla_attention.py b/tests/test_sla_attention.py index 2e2905e2..3ec0a4fa 100644 --- a/tests/test_sla_attention.py +++ b/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")