From 79ca8cfb6b1c68bc8c1f55258ba04417c8633d46 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:13:05 +0300 Subject: [PATCH] feat: SLA block-sparse attention for MiniMax H3, ported from v4 Loader-level attention option: mean-pool Q and a smoothed K into blocks, score them with one small matmul, attend only the top 1-sparsity fraction of key blocks per query block, via a vendored Triton kernel. Runs alongside the chosen attention backend rather than replacing it (one override slot; the backend becomes SLA's dense fall-through for the text refiner, masked calls, and trailing dense steps), and refuses cleanly -- naming why -- on anything that is not MiniMax H3 or a machine without CUDA+Triton. modules/loaders/sla_block_map.py and sla_kernel.py are near-verbatim ports of v4's vendored files (no FunPack-specific imports, so nothing to adapt); sla_attention.py is the loader-facing glue (H3 gate, override composition, per-step wrapper), same content as v4's with a "not a module" note added to match gguf_support.py's convention. Wired into FunPackDiffusionModelLoader as six new optional widgets, following v4's loaders.py wiring but expressed in v5's io.Schema style. The Triton kernel itself is unexercisable on this Mac (no CUDA) -- matches how real generation is verified in this project (rented GPU, not local). Ported the CPU-testable half of v4's test suite (the override contract, the H3 gate, the per-step wrapper); the three kernel-correctness tests carry the same CUDA+Triton skip guard v4 used. Co-Authored-By: Claude Sonnet 5 --- core/tests/test_shipped_nodes.py | 2 + modules/loaders/diffusion_model/nodes.py | 88 +++++- modules/loaders/sla_attention.py | 302 ++++++++++++++++++ modules/loaders/sla_block_map.py | 130 ++++++++ modules/loaders/sla_kernel.py | 165 ++++++++++ modules/loaders/tests/test_sla_attention.py | 324 ++++++++++++++++++++ 6 files changed, 1003 insertions(+), 8 deletions(-) create mode 100644 modules/loaders/sla_attention.py create mode 100644 modules/loaders/sla_block_map.py create mode 100644 modules/loaders/sla_kernel.py create mode 100644 modules/loaders/tests/test_sla_attention.py diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 756b7618..cca04df1 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -57,6 +57,8 @@ FROZEN = { "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size"], "FunPackDiffusionModelLoader": [ "model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation", + "sla", "sla_sparsity", "sla_block_size", "sla_protect_audio", "sla_min_seq_len", + "sla_dense_last_steps", ], "FunPackLoraLoader": [ "model", "lora_name", "strength_model", "clip", "strength_clip", diff --git a/modules/loaders/diffusion_model/nodes.py b/modules/loaders/diffusion_model/nodes.py index 75ac9461..3c155174 100644 --- a/modules/loaders/diffusion_model/nodes.py +++ b/modules/loaders/diffusion_model/nodes.py @@ -11,7 +11,7 @@ import folder_paths from comfy_api.latest import io from ..._core import log -from .. import gguf_support +from .. import gguf_support, sla_attention from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices, attention_override, dtype_of, set_fp16_accumulation, weight_model_options) @@ -49,6 +49,55 @@ class FunPackDiffusionModelLoader(io.ComfyNode): "listed; 'default' keeps ComfyUI's launch choice."), io.Boolean.Input("fp16_accumulation", default=False, optional=True, tooltip="Faster fp16 matmuls where the torch build supports it."), + io.Boolean.Input("sla", default=False, optional=True, + 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. Runs ALONGSIDE the attention " + "backend above rather than replacing it: SLA takes " + "H3's long packed self-attention, the chosen backend " + "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. + io.Float.Input("sla_sparsity", default=sla_attention.SLA_DEFAULTS["sparsity_ratio"], + min=0.0, max=0.95, step=0.05, optional=True, + tooltip="Fraction of key blocks skipped, when attention is " + "sla_h3. 0.90 is validated; 0.85 is lightx2v's own value " + "and ~15% slower. Break-even is around 0.60 -- below that " + "the kernel is SLOWER than dense, so a low value is a " + "loss, not a safe fallback."), + io.Combo.Input("sla_block_size", + options=["64", "128"], + default=str(sla_attention.SLA_DEFAULTS["block_size"]), + optional=True, + tooltip="How many sequence tokens share one key selection. H3 " + "packs audio at 80 rows per second, so a 128-row block " + "forces 1.6s of speech down one attention pattern while " + "the same rows are 3% of a video frame. Use 128 only " + "when the audio does not matter."), + io.Boolean.Input("sla_protect_audio", + default=sla_attention.SLA_DEFAULTS["protect_audio"], optional=True, + tooltip="Always attend the [text | cond | audio] prefix, " + "whatever top-k picks. Audio is ~1% of the packed " + "sequence, so plain top-k regularly drops all of it " + "and the soundtrack degrades while the video still " + "looks fine. Costs about 7%."), + io.Int.Input("sla_min_seq_len", + default=sla_attention.SLA_DEFAULTS["min_seq_len"], + min=0, max=1000000, step=1024, optional=True, + 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."), + io.Int.Input("sla_dense_last_steps", + default=sla_attention.SLA_DEFAULTS["dense_last_steps"], + min=0, max=8, optional=True, + tooltip="Run the last N sampling steps at full attention. 0 " + "matches lightx2v; 1 was tested and did not help, for " + "+20% time."), ], outputs=[ io.Model.Output(display_name="model"), @@ -58,7 +107,10 @@ class FunPackDiffusionModelLoader(io.ComfyNode): @classmethod def execute(cls, model_name: str, weight_dtype: str, compute_dtype: str, - attention: str, fp16_accumulation: bool = False) -> io.NodeOutput: + attention: str, fp16_accumulation: bool = False, sla: bool = False, + sla_sparsity: float = None, sla_block_size: str = None, + sla_protect_audio: bool = None, sla_min_seq_len: int = None, + sla_dense_last_steps: int = None) -> io.NodeOutput: notes = [f"FunPack Diffusion Model Loader | {model_name}"] accumulation = set_fp16_accumulation(fp16_accumulation) @@ -120,18 +172,38 @@ class FunPackDiffusionModelLoader(io.ComfyNode): notes.append("compute dtype: unsupported by this ComfyUI, ignored") override = attention_override(attention) - if override is not None: - model.model_options.setdefault("transformer_options", {})[ - "optimized_attention_override"] = override - notes.append(f"attention: {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 that was picked. + installed = False + if sla: + model, sla_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, + dense_fn=override, dense_label=attention) + notes.append(sla_note) + 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("attention: default (as launched)") + notes.append(f"attention: {attention} (dense calls)") # Never reach into the model's shape to describe it: a log line that # assumes structure can fail the load it was only meant to narrate. kind = type(getattr(model, "model", model)).__name__ + applied_attention = "sla_h3" if installed else ( + attention if override is not None else "default (as launched)") log.info("FunPack Diffusion Model Loader", f"{model_name} loaded as {kind}, weights " f"{weight_dtype}, compute {compute_dtype}, attention " - f"{attention if override is not None else 'default (as launched)'}") + f"{applied_attention}") return io.NodeOutput(model, "\n".join(notes)) diff --git a/modules/loaders/sla_attention.py b/modules/loaders/sla_attention.py new file mode 100644 index 00000000..143b863d --- /dev/null +++ b/modules/loaders/sla_attention.py @@ -0,0 +1,302 @@ +"""SLA: block-sparse attention for MiniMax H3, as a loader choice. + +Not a module: see common.py's own note -- a plain file at the domain level is a +helper the loaders import, and ships no node of its own. + +ComfyUI has 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. This module supplies the missing half — mean-pool Q and a smoothed K +into blocks, score them with one small matmul, and attend only the top +``1 - sparsity`` fraction of key blocks per query block. Nothing is trained and +nothing is loaded; the published SLA files contain only ordinary LoRA tensors. + +Ported from ComfyUI-H3-SLA-Attention (MIT), whose kernel and block map are vendored +from LightX2V (Apache-2.0) — see sla_kernel.py / sla_block_map.py for the change +lists. What is FunPack's here: this file, the H3 model check, and the loader wiring, +so turning it on is picking a value in the diffusion model loader rather than adding +and wiring a node. + +The hook is ``transformer_options["optimized_attention_override"]``, which +``wrap_attn`` consults and H3's one attention call site reaches. The legacy +``set_model_attn1_patch`` is the SD-UNet path a DiT never consults: a patch installed +there reports success and silently does nothing, which is why the invocation counter +below exists and why a run that never sparsified says so in the log. +""" +import logging + +SLA_NAME = "sla_h3" + +_H3_HEAD_DIM = 128 +_H3_MODEL_NAMES = ("MiniMaxH3Model",) + +# Validated on an RTX 5090 at 768p/15s (ComfyUI-H3-SLA-Attention's own measurements): +# 3.7x the attention throughput of stock ComfyUI at sparsity 0.90 / block 64. +SLA_DEFAULTS = { + "sparsity_ratio": 0.90, + "block_size": 64, + "min_seq_len": 8192, + "dense_last_steps": 0, + "protect_audio": True, + "enabled": True, +} + + +def _torch(): + import torch + return torch + + +def sla_available(): + """True when this machine could actually run the kernel. + + Triton and CUDA, nothing else — the ladder in sla_kernel.py handles differing + shared-memory limits per architecture. Offered as a choice only where it can run, + like every other backend in the loader's list. + """ + try: + import torch + import triton # noqa: F401 + except Exception: # noqa: BLE001 + return False + return bool(torch.cuda.is_available()) + + +def is_h3_model(model): + """Whether this patcher holds a MiniMax H3 diffusion model. + + SLA is only sound on H3: the sparsity ratio it runs at is the one the turbo LoRA + was distilled to tolerate, and H3's packed [text | cond | audio | video] sequence + is what the prefix pinning is for. Head shape alone would also match LTX, and + silently sparsifying LTX attention is a quality loss with no LoRA compensating it. + """ + diffusion = getattr(getattr(model, "model", None), "diffusion_model", None) + return type(diffusion).__name__ in _H3_MODEL_NAMES + + +def new_state(): + return { + "calls": 0, # sparse invocations this run + "dense": 0, # fall-throughs this run + "step": 0, + "seq": 0, + "kept": 0, + "blocks": 0, + "pinned": 0, + "backend": None, # what we displaced + "failed": None, # first kernel failure, if any + } + + +def _summarise(state, sparsity, blkq, blkk): + """One line per sampling run. Never one per block — there are 50 of those.""" + if state["calls"] == 0: + logging.warning( + "[FunPack SLA] installed but never invoked — attention was NOT sparsified " + "(%d dense fall-throughs). Check that the MODEL reaching the sampler is the " + "one this loader produced.", state["dense"]) + return + real = 1.0 - (state["kept"] / state["blocks"]) if state["blocks"] else 0.0 + logging.info( + "[FunPack SLA] %d calls | S=%d | blocks %d/%d kept (%.1f%% sparse, asked %.0f%%) " + "| %d pinned | BLK=%dx%d | %d dense fall-throughs | displaced %s", + state["calls"], state["seq"], state["kept"], state["blocks"], real * 100.0, + sparsity * 100.0, state["pinned"], blkq, blkk, state["dense"], + state["backend"] or "?") + if state["failed"] is not None: + logging.warning("[FunPack SLA] kernel fell back to dense at least once: %s", + state["failed"]) + + +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 + + def override(func, q, k, v, heads, mask=None, attn_precision=None, + skip_reshape=False, skip_output_reshape=False, **kwargs): + def dense(): + state["dense"] += 1 + 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"] = dense_label or getattr(func, "__name__", repr(func)) + + to = kwargs.get("transformer_options") or {} + + # Anything that is not the packed H3 self-attention goes straight through. The + # min_seq_len guard is what keeps the 2-block token refiner (S = text length) + # and low-resolution runs dense, where selection costs more than it saves. + if ( + not skip_reshape + or mask is not None + or q.ndim != 4 + or q.shape[-1] != _H3_HEAD_DIM + or q.dtype not in ok_dtypes + or q.shape[2] < min_seq_len + or to.get("_funpack_sla_dense", False) + ): + return dense() + + try: + # Imported here, not at module scope: triton is absent on plenty of machines + # and importing it eagerly would take the whole module — and the loader — + # down with it. + try: + from .sla_block_map import get_block_map + from .sla_kernel import block_sparse_attention + except ImportError: + from sla_block_map import get_block_map + from sla_kernel import block_sparse_attention + + B, H, S, D = q.shape + + # [1, H, S, D] -> [1, S, H, D]. H3 builds q/k/v as [S, H, D] and transposes + # for the call, so this transposes back onto the original memory and the + # copy is a no-op. A BHSD kernel would cost a real ~1.3 GB copy per tensor. + qb, kb, vb = (t.transpose(1, 2) for t in (q, k, v)) + if not qb.is_contiguous(): + qb, kb, vb = qb.contiguous(), kb.contiguous(), vb.contiguous() + + # Pin the [text | cond | audio] prefix into every query's selection. Audio is + # ~1% of the packed sequence, so plain top-k routinely drops all of it and the + # soundtrack degrades while the video still looks fine. 0 when the layout is + # unavailable, which disables the protection rather than guessing at it. + prefix = int(to.get("_funpack_sla_prefix", 0) or 0) if protect_audio else 0 + if prefix >= S: + prefix = 0 + + lut, topk = get_block_map(qb, kb, topk_ratio, blkq, blkk, protect_upto=prefix) + out = block_sparse_attention(qb, kb, vb, lut, topk, blkq, blkk) + + state["calls"] += 1 + state["seq"] = S + state["kept"] = topk + state["blocks"] = (S + blkk - 1) // blkk + state["pinned"] = (prefix + blkk - 1) // blkk + + if skip_output_reshape: + return out.transpose(1, 2) + return out.reshape(B, S, H * D) + + except Exception as exc: # noqa: BLE001 - a bad kernel must not kill the run + if state["failed"] is None: + state["failed"] = "%s: %s" % (exc.__class__.__name__, exc) + logging.debug("[FunPack SLA] kernel failed", exc_info=True) + return dense() + + return override + + +def make_wrapper(state, sparsity_ratio, blkq, blkk, dense_last_steps): + """DIFFUSION_MODEL wrapper: per-step state, and the end-of-run summary. + + Registered once and then reused — ComfyUI caches node outputs, so this closure + outlives a single sampling run. The step counter therefore has to reset itself, or + every run after the first drifts permanently into the trailing-dense window and + silently stops sparsifying. + """ + + def wrapper(executor, x, timestep, context, transformer_options={}, + minimax_payload=None, **kwargs): + to = transformer_options + n_steps = max(1, len(to.get("sample_sigmas", [])) - 1) + + if state["step"] >= n_steps: # new run + state["step"] = 0 + state["calls"] = 0 + state["dense"] = 0 + state["failed"] = None + state["step"] += 1 + + # PackedLayout.segments is [(start, stop, kind), …] over + # [text | cond/ref | audio | video]; the video start is therefore the length of + # everything that must stay exactly attended. It lives on the payload, which + # never reaches the attention call site, so the wrapper is the only place it can + # be picked up. + prefix = 0 + layout = minimax_payload.get("layout") if minimax_payload else None + for seg in getattr(layout, "segments", ()) or (): + if len(seg) == 3 and seg[2] == "video": + prefix = int(seg[0]) + break + to["_funpack_sla_prefix"] = prefix + to["_funpack_sla_dense"] = bool( + dense_last_steps > 0 and state["step"] > n_steps - dense_last_steps) + + # Forward minimax_payload only when H3 actually supplied one: every other + # diffusion model would raise TypeError on the unexpected kwarg, turning a + # graceful no-op into a crash mid-sampling. + if minimax_payload is not None: + kwargs["minimax_payload"] = minimax_payload + out = executor.original(x, timestep, context, + transformer_options=transformer_options, **kwargs) + + if state["step"] >= n_steps: + _summarise(state, sparsity_ratio, blkq, blkk) + return out + + return wrapper + + +def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None, + 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. + + `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 + + sparsity_ratio = float(pick(sparsity_ratio, "sparsity_ratio")) + blkq = int(pick(block_size, "block_size")) + min_seq_len = int(pick(min_seq_len, "min_seq_len")) + dense_last_steps = int(pick(dense_last_steps, "dense_last_steps")) + protect_audio = bool(pick(protect_audio, "protect_audio")) + + 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, "SLA: off (dense baseline)", False + if not is_h3_model(model): + return model, ("SLA: skipped — not a MiniMax H3 model"), False + if not sla_available(): + 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 + # fastest. LightX2V picks the same split for its sage2 path off sm90. + blkk = 64 if blkq == 128 else blkq + + state = new_state() + 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, + 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"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} | dense calls -> " + f"{dense_label or 'as launched'}"), True diff --git a/modules/loaders/sla_block_map.py b/modules/loaders/sla_block_map.py new file mode 100644 index 00000000..b6755d5f --- /dev/null +++ b/modules/loaders/sla_block_map.py @@ -0,0 +1,130 @@ +"""Block selection: which key blocks each query block is allowed to see. + +Vendored into FunPack from ComfyUI-H3-SLA-Attention, which vendored it from +LightX2V (Apache-2.0). The provenance note and change list below are the +original author's and still describe this file exactly. + +Vendored from LightX2V (Apache-2.0): + lightx2v/common/ops/attn/utils/sla_util_blhd.py + https://github.com/ModelTC/LightX2V + +This is the whole of what "SLA" does at inference time. Mean-pool Q into blocks, +mean-pool a smoothed K into blocks, score the two against each other with one +small matmul, and keep the top ``topk_ratio`` fraction of key blocks per query +block. No weights, nothing trained, nothing to load -- the published SLA LoRA +adapts the *model* to tolerate the resulting sparsity, it does not parameterise +this step. + +Three changes against upstream, marked FIX: + +1. ``other=0.0`` on the masked load. The masked lanes feed ``tl.sum`` two lines + later, and Triton leaves them undefined, so the final (partial) block of the + sequence pooled whatever was in memory. Upstream fixed exactly this in the + BHSD twin of this file and did not port the fix here. +2. ``max(1, ...)`` on the top-k count, so a short sequence keeps one key block + rather than zero. Also upstream's BHSD-only fix. +3. Smooth-k is folded into the pooled result instead of being materialised. + Pooling is a mean over L and the correction is constant along L, so + ``pool(k - mu) == pool(k) - mu`` exactly. Upstream builds the whole smoothed + copy of K, which at 768p/15s is a needless ~1.3 GB allocation. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _compress_kernel( + X, + XM, + L: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_bh = tl.program_id(1) + + idx_b = idx_bh // H + idx_h = idx_bh - idx_b * H + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + + x_offset = idx_b * L * H * D + idx_h * D + xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D + # FIX vs upstream: other=0.0 -- these lanes are summed below. + x = tl.load( + X + x_offset + offs_l[:, None] * (H * D) + offs_d[None, :], + mask=offs_l[:, None] < L, + other=0.0, + ) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx + tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) + + +def mean_pool(x, BLK): + """``(B, L, H, D)`` -> ``(B, H, ceil(L/BLK), D)``, mean over each L block.""" + assert x.is_contiguous() + B, L, H, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + # fp32, not x.dtype. The Triton reduction already accumulates in fp32; only + # the store rounds. Upstream stores bf16, which quantises the block *scores* + # and makes top-k pick different blocks than exact arithmetic would -- and + # at 85-90% sparsity only 3-4 blocks per query survive, so one wrong pick is + # a large error. Keeping fp32 here costs a few MB and no measurable time. + x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=torch.float32) + + grid = (L_BLOCKS, B * H) + _compress_kernel[grid](x, x_mean, L, H, D, BLK) + return x_mean + + +def get_block_map(q, k, topk_ratio, BLKQ=128, BLKK=128, protect_upto=0): + """Return ``(lut, topk)``: the key blocks each query block should attend to. + + ``q``/``k`` are ``(B, L, H, D)`` contiguous. ``lut`` comes back as + ``(B, H, ceil(LQ/BLKQ), topk)`` int32, contiguous, ready for the kernel. + + ``protect_upto`` pins the first N tokens into every query block's selection. + For H3 that is the ``[text | cond | audio]`` prefix, and it exists because + plain top-k starves audio: at 768p/15s the audio is ~1% of the packed + sequence (19 key blocks of 1794), so nothing makes a query keep any of it, + and the smooth-k mean it is scored against is 99% video. The pinned blocks + are added on top of the top-k budget rather than displacing video, so video + coverage is unchanged and the extra cost is the prefix itself (~7%). + """ + pooled_q = mean_pool(q, BLKQ) + # Smooth-k (SageAttention's trick), folded in rather than materialised. + mu = k.mean(dim=1, dtype=torch.float32) # (B, H, D) + pooled_k = mean_pool(k, BLKK) - mu[:, :, None, :] + + # GQA, for completeness -- H3 is MHA so this is a no-op there. + num_q_heads, num_kv_heads = pooled_q.shape[1], pooled_k.shape[1] + if num_q_heads != num_kv_heads: + assert num_q_heads % num_kv_heads == 0 + pooled_k = pooled_k.repeat_interleave(num_q_heads // num_kv_heads, dim=1) + + pooled_score = pooled_q @ pooled_k.transpose(-1, -2) # (B, H, NQ, NK) + + NK = pooled_score.shape[-1] + # FIX vs upstream: keep at least one key block. + topk = max(1, min(NK, int(topk_ratio * NK))) + + n_pinned = 0 + if protect_upto > 0: + n_pinned = min((int(protect_upto) + BLKK - 1) // BLKK, NK) + # Ranking them above everything else is what pins them; widening topk by + # the same amount is what stops them evicting the blocks top-k chose. + if n_pinned > 0: + pooled_score[..., :n_pinned] = float("inf") + topk = min(NK, topk + n_pinned) + + lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices + + return lut.to(torch.int32).contiguous(), topk diff --git a/modules/loaders/sla_kernel.py b/modules/loaders/sla_kernel.py new file mode 100644 index 00000000..b4328f62 --- /dev/null +++ b/modules/loaders/sla_kernel.py @@ -0,0 +1,165 @@ +"""Block-sparse attention forward kernel for MiniMax-H3. + +Vendored into FunPack from ComfyUI-H3-SLA-Attention, which vendored it from +LightX2V (Apache-2.0). The provenance note and change list below are the +original author's and still describe this file exactly. + +Vendored and reduced from LightX2V (Apache-2.0): + lightx2v/common/ops/attn/kernels/sla_kernel_ar.py -- ``_attn_fwd`` + https://github.com/ModelTC/LightX2V + +Only the forward pass survives the port; sampling runs under ``no_grad`` so the +backward half and the ``autograd.Function`` wrapper are dead weight here. + +The ``_ar`` name upstream refers to the models it was written for, not to any +causal structure: the kernel walks whatever key blocks the lookup table names +and applies no triangular mask. That is what makes it usable for H3, whose +packed ``[text | cond/ref | audio | video]`` sequence is fully bidirectional. + +Layout is BLHD -- ``(B, L, H, D)`` -- which is deliberate. H3 materialises q/k/v +as ``[S, H, D]`` and only then transposes to ``[1, H, S, D]`` for the attention +call, so transposing back is free, while a BHSD kernel would force a real +``.contiguous()`` copy costing ~1.3 GB per tensor at 768p/15s. + +Two changes against upstream, both marked FIX below: masked loads now pass +``other=0.0``, because Triton leaves masked lanes undefined and ``0 * NaN`` is +NaN, not zero -- the sequence-tail block can poison a whole row otherwise. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _attn_fwd( + Q, + K, + V, + qk_scale: tl.constexpr, + topk: tl.constexpr, + LUT, + OS, + H: tl.constexpr, + LQ: tl.constexpr, + LK: tl.constexpr, + M_BLOCKS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + idx_b = idx_bh // H + idx_h = idx_bh % H + + HD: tl.constexpr = H * D + + # Q/K/V/O: (B, L, H, D) contiguous. + q_offset = idx_b * LQ * HD + idx_h * D + kv_offset = idx_b * LK * HD + idx_h * D + lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk + + offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + Q_ptrs = Q + q_offset + offs_m[:, None] * HD + offs_d[None, :] + OS_ptrs = OS + q_offset + offs_m[:, None] * HD + offs_d[None, :] + LUT_ptr = LUT + lut_offset + + m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32) + + # FIX vs upstream: other=0.0 on every masked load. + q = tl.load(Q_ptrs, mask=offs_m[:, None] < LQ, other=0.0) + for block_idx in tl.range(topk): + idx_n = tl.load(LUT_ptr + block_idx).to(tl.int64) + k_start = idx_n * BLOCK_N + k_mask = (k_start + offs_n) < LK + + K_ptrs = K + kv_offset + (k_start + offs_n)[None, :] * HD + offs_d[:, None] + V_ptrs = V + kv_offset + (k_start + offs_n)[:, None] * HD + offs_d[None, :] + + k = tl.load(K_ptrs, mask=k_mask[None, :], other=0.0) + qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # 1/ln(2), for exp2 + qk = tl.where(k_mask[None, :], qk, float("-inf")) + + v = tl.load(V_ptrs, mask=k_mask[:, None], other=0.0) + local_m = tl.max(qk, 1) + new_m = tl.maximum(m_i, local_m) + qk = qk - new_m[:, None] + + p = tl.math.exp2(qk) + l_ij = tl.sum(p, 1) + alpha = tl.math.exp2(m_i - new_m) + o_s = o_s * alpha[:, None] + o_s += tl.dot(p.to(v.dtype), v) + + l_i = l_i * alpha + l_ij + m_i = new_m + + o_s = o_s / l_i[:, None] + tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < LQ) + + +# (BLOCK_M, BLOCK_N) -> ladder of (num_warps, num_stages) to try, best first. +# +# These are not upstream's numbers. LightX2V hardcodes num_warps=4, num_stages=3, +# which on sm_120 needs 160 KB of shared memory against a 99 KB limit -- it does +# not merely run slowly, it fails to launch. Consumer Blackwell has less shared +# memory than the datacentre parts these kernels were tuned on, so we probe and +# fall back. Measured on a 5090 at S=32768, H=56, D=128 (ms, lower better): +# 128x64 w8 s3 -> 22.4 128x64 w4 s3 -> 23.9 128x128 w8 s2 -> 25.4 +# 64x64 w4 s1 -> 26.8 64x128 w4 s2 -> 30.5 128x128 w4 s1 -> 40.5 +_LADDER = { + (128, 64): ((8, 3), (4, 3), (8, 2), (4, 1)), + (128, 128): ((8, 2), (4, 2), (8, 1), (4, 1)), + (64, 128): ((4, 2), (8, 2), (4, 1)), + (64, 64): ((4, 1), (4, 3), (8, 3), (8, 1)), +} +_CHOSEN: dict = {} + + +def block_sparse_attention(q, k, v, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None): + """Attend each query block to only the key blocks named in ``lut``. + + ``q``/``k``/``v`` are ``(B, L, H, D)`` contiguous; ``lut`` is + ``(B, H, M_BLOCKS, topk)`` int32, contiguous. Returns ``(B, L, H, D)``. + """ + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert lut.is_contiguous() + assert BLOCK_M in (64, 128) and BLOCK_N in (64, 128) + + B, LQ, H, D = q.shape + LK = k.shape[1] + if qk_scale is None: + qk_scale = D**-0.5 + + M_BLOCKS = triton.cdiv(LQ, BLOCK_M) + o_s = torch.empty_like(q) + grid = (M_BLOCKS, B * H) + + key = (BLOCK_M, BLOCK_N, D) + ladder = (_CHOSEN[key],) if key in _CHOSEN else _LADDER[(BLOCK_M, BLOCK_N)] + + last = None + for cfg in ladder: + num_warps, num_stages = cfg + try: + _attn_fwd[grid]( + q, k, v, qk_scale, topk, lut, o_s, + H, LQ, LK, M_BLOCKS, D, BLOCK_M, BLOCK_N, + num_warps=num_warps, num_stages=num_stages, + ) + except triton.runtime.errors.OutOfResources as exc: + last = exc + continue + _CHOSEN[key] = cfg + return o_s + + raise last if last is not None else RuntimeError("no viable launch config") diff --git a/modules/loaders/tests/test_sla_attention.py b/modules/loaders/tests/test_sla_attention.py new file mode 100644 index 00000000..22aa6e97 --- /dev/null +++ b/modules/loaders/tests/test_sla_attention.py @@ -0,0 +1,324 @@ +"""SLA block-sparse attention: the ComfyUI override contract, and the H3 gate. + +The contract half matters more than it looks. The failure this guards against is +not a crash -- it is a patch that installs cleanly, logs success and never runs, +because it hooked an API the model does not consult. So these drive the override +exactly the way `wrap_attn` does, with H3's real argument shape, and assert which +path fired. + +Ported from FunPack v4's own suite, itself ported from ComfyUI-H3-SLA-Attention +(MIT). The CUDA half is skipped without a GPU; everything else runs anywhere +torch imports. +""" +import types + +import pytest +import torch + +from modules.loaders import sla_attention as sla + +H, D = 56, 128 # MiniMax H3: 56 heads, head_dim 128 + +CUDA = torch.cuda.is_available() +try: + import triton # noqa: F401 + HAS_TRITON = True +except Exception: # noqa: BLE001 + HAS_TRITON = False + + +def _backend(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, + skip_output_reshape=False, **kwargs): + """Stand-in for ComfyUI's undecorated attention backend.""" + if not skip_reshape: + b, s, _ = q.shape + q, k, v = (t.view(b, s, heads, -1).transpose(1, 2) for t in (q, k, v)) + o = torch.nn.functional.scaled_dot_product_attention(q, k, v) + if skip_output_reshape: + return o + b, _, s, _ = q.shape + return o.transpose(1, 2).reshape(b, s, -1) + + +def _call(override, q, k, v, **kw): + # wrap_attn hands the override the UNDECORATED backend as arg 0, then q/k/v/heads + # with H3's kwargs: mask is always None, skip_reshape True, and skip_output_reshape + # is not passed at all. + opts = dict(mask=None, skip_reshape=True, transformer_options={}, + _inside_attn_wrapper=True) + opts.update(kw) + return override(_backend, q, k, v, H, **opts) + + +# ── the override contract ───────────────────────────────────────────────────── + +def test_a_short_sequence_stays_dense(): + """H3's text refiner is a few hundred tokens and must never be sparsified.""" + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=8192) + q = torch.randn(1, H, 512, D) + out = _call(ov, q, q.clone(), q.clone()) + assert (state["calls"], state["dense"]) == (0, 1) + assert out.shape == (1, 512, H * D) + + +def test_masked_attention_stays_dense(): + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=0) + q = torch.randn(1, H, 256, D) + _call(ov, q, q.clone(), q.clone(), mask=torch.zeros(1, 1)) + assert state["calls"] == 0 + + +def test_float32_never_reaches_the_kernel(): + """The dtype guard catches it first, so nothing is recorded as a kernel failure.""" + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=0) + q = torch.randn(1, H, 256, D) + out = _call(ov, q, q.clone(), q.clone()) + assert (state["calls"], state["dense"]) == (0, 1) + assert state["failed"] is None + assert out.shape == (1, 256, H * D) + + +def test_the_run_records_which_backend_it_displaced(): + """So the log can say what dense fall-throughs will actually use.""" + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=8192) + _call(ov, *(torch.randn(1, H, 256, D),) * 3) + assert state["backend"] == "_backend" + + +def test_a_kernel_failure_costs_speed_not_the_run(): + """bf16 on CPU passes every guard and then cannot launch a CUDA kernel -- the + closest stand-in for a driver or Triton mismatch on a real machine.""" + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=0) + q = torch.randn(1, H, 256, D, dtype=torch.bfloat16) + out = _call(ov, q, q.clone(), q.clone()) + assert out.shape == (1, 256, H * D) + assert state["failed"] is not None + assert (state["calls"], state["dense"]) == (0, 1) + + +# ── the per-step wrapper ────────────────────────────────────────────────────── + +def _run(wrapper, n_steps, payload=None): + """One sampling run: n_steps forwards through the wrapper. Returns the per-step + value of the force-dense flag.""" + seen = [] + + class Ex: + @staticmethod + def original(*a, **kw): + seen.append(kw["transformer_options"]["_funpack_sla_dense"]) + return None + + to = {"sample_sigmas": [0.0] * (n_steps + 1)} + for _ in range(n_steps): + wrapper(Ex, None, None, None, transformer_options=to, minimax_payload=payload) + return seen + + +def test_the_step_counter_resets_between_runs(): + """ComfyUI caches node outputs, so this closure outlives one run. Without the reset + every later run sits permanently inside the trailing-dense window and silently stops + sparsifying.""" + state = sla.new_state() + w = sla.make_wrapper(state, 0.90, 64, 64, dense_last_steps=1) + first = _run(w, 4) + assert first == [False, False, False, True] + assert _run(w, 4) == first + + +def test_dense_last_steps_zero_never_forces_dense(): + state = sla.new_state() + w = sla.make_wrapper(state, 0.90, 64, 64, dense_last_steps=0) + assert _run(w, 4) == [False] * 4 + + +def test_the_protected_prefix_is_read_from_the_packed_layout(): + """The video segment's start is the length of what must stay exactly attended. It + lives on minimax_payload, which never reaches the attention call site, so the + wrapper is the only place it can be picked up.""" + layout = types.SimpleNamespace(segments=[ + (0, 512, "text"), (512, 800, "cond"), (800, 2000, "audio"), (2000, 114785, "video")]) + state = sla.new_state() + w = sla.make_wrapper(state, 0.90, 64, 64, dense_last_steps=0) + to = {"sample_sigmas": [0.0] * 5} + + class Ex: + @staticmethod + def original(*a, **kw): + return None + + w(Ex, None, None, None, transformer_options=to, minimax_payload={"layout": layout}) + assert to["_funpack_sla_prefix"] == 2000 + + +def test_no_layout_disables_the_protection_rather_than_guessing(): + state = sla.new_state() + w = sla.make_wrapper(state, 0.90, 64, 64, dense_last_steps=0) + to = {"sample_sigmas": [0.0] * 5} + + class Ex: + @staticmethod + def original(*a, **kw): + return None + + w(Ex, None, None, None, transformer_options=to, minimax_payload=None) + assert to["_funpack_sla_prefix"] == 0 + + +def test_a_non_h3_model_is_not_handed_a_minimax_kwarg(): + """Every other diffusion model would raise TypeError on the unexpected kwarg -- a + crash mid-sampling instead of the graceful no-op it should be.""" + seen = {} + + class Ex: + @staticmethod + def original(x, timestep, context, transformer_options=None, **kw): + seen.update(kw) + return None + + state = sla.new_state() + w = sla.make_wrapper(state, 0.90, 64, 64, dense_last_steps=0) + w(Ex, None, None, None, transformer_options={"sample_sigmas": [0.0] * 2}) + assert "minimax_payload" not in seen + + +# ── the H3 gate ─────────────────────────────────────────────────────────────── + +class _Patcher: + def __init__(self, class_name): + inner = type(class_name, (), {})() + self.model = types.SimpleNamespace(diffusion_model=inner) + self.model_options = {} + self.wrappers = [] + + def clone(self): + out = _Patcher("x") + out.model, out.model_options = self.model, dict(self.model_options) + return out + + def add_wrapper_with_key(self, *a): + self.wrappers.append(a) + + +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, 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, installed = sla.install_sla(model) + if sla.sla_available(): + assert installed is True and out is not model and "sparsity=0.90" in note + else: + assert installed is False and out is model and "CUDA+Triton" in note + + +def test_the_defaults_are_the_validated_ones(): + """0.90 / 64 / protect on: the settings the port was measured at.""" + assert sla.SLA_DEFAULTS == {"sparsity_ratio": 0.90, "block_size": 64, + "min_seq_len": 8192, "dense_last_steps": 0, + "protect_audio": True, "enabled": True} + + +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, 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") + + +@pytestmark_cuda +def test_keeping_every_block_is_just_attention(): + from modules.loaders.sla_block_map import get_block_map + from modules.loaders.sla_kernel import block_sparse_attention + + S = 4096 + torch.manual_seed(0) + q, k, v = (torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + lut, topk = get_block_map(q, k, 1.0, 64, 64) + got = block_sparse_attention(q, k, v, lut, topk, 64, 64) + ref = torch.nn.functional.scaled_dot_product_attention( + *(t.transpose(1, 2) for t in (q, k, v))).transpose(1, 2) + assert not torch.isnan(got).any() + rel = ((got.float() - ref.float()).abs().max() / ref.float().abs().max()).item() + assert rel < 1e-2 + + +@pytestmark_cuda +def test_every_query_block_keeps_the_whole_protected_prefix(): + """This is the audio fix: audio is ~1% of the packed sequence, so plain top-k + routinely drops all of it and the soundtrack degrades while the video looks fine.""" + from modules.loaders.sla_block_map import get_block_map + + S, prefix = 16384, 2048 + torch.manual_seed(0) + q = torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, S, H, D, device="cuda", dtype=torch.bfloat16) + plain_lut, plain_topk = get_block_map(q, k, 0.10, 64, 64) + lut, topk = get_block_map(q, k, 0.10, 64, 64, protect_upto=prefix) + + n_pinned = prefix // 64 + assert topk == plain_topk + n_pinned # widened, so it displaces no video + got = lut.long().sort(dim=-1).values[..., :n_pinned] + want = torch.arange(n_pinned, device=lut.device) + assert torch.equal(got, want.expand_as(got)) + # and the failure it exists to prevent + covered = (plain_lut.long() < n_pinned).sum(-1).float().mean().item() + assert covered < n_pinned + + +@pytestmark_cuda +def test_the_override_fires_and_returns_h3s_shape(): + state = sla.new_state() + ov = sla.make_override(state, 0.90, 64, 64, min_seq_len=1024) + S = 8192 + q, k, v = (torch.randn(1, H, S, D, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + out = _call(ov, q, k, v) + assert (state["calls"], state["dense"]) == (1, 0) + assert out.shape == (1, S, H * D) + assert not torch.isnan(out).any()