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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Sonnet 5
Claude Sonnet 5
parent
fe8d530a13
commit
79ca8cfb6b
@@ -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",
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user