Claude
09b6c70db7
A scene picture a video-only pipeline cannot use is named in Generate. CHANGELOG Brownian wording and Veda GPU labels match the code. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjVsBg3ixjnT1Fgdy12
207 lines
12 KiB
Python
207 lines
12 KiB
Python
"""Veda learned sparse attention for MiniMax H3, as the loader's SLA method "veda".
|
|
|
|
Veda (Han, Yang, Hu, Mei, Jiang, Qi; MIT; arXiv 2605.30325) ported from Veda-on-ComfyUI: see veda_core.py for
|
|
the method and VEDA-NOTICE.txt for the licences. EXPERIMENTAL, tested on a rental GPU: the pure-torch half is
|
|
tested on CPU too; the Triton kernel (veda_kernel.py) needs CUDA and cannot run in the CPU suite.
|
|
|
|
It rides the same hook as SLA (`transformer_options["optimized_attention_override"]`, plus SLA's per-step wrapper
|
|
for the step counter, the dense steps and the rising share), so it needs no weight change and meets q/k/v only after
|
|
the LoRA-patched projections: LoRA, int8_convrot and MXFP8 (weights only) leave what it sees as bf16/fp16 q/k/v, which
|
|
is what the predictor was trained on. Calls it does not take (the text refiner, masked calls, short sequences, dense
|
|
steps, a head count the predictor was not made for) go to the chosen attention backend, and the log says how many
|
|
and why. When Veda cannot run at all, the loader says why and the backend is installed on its own.
|
|
"""
|
|
import collections
|
|
import logging
|
|
|
|
from . import sla_attention as sla, veda_core as core, veda_download
|
|
|
|
_NAMED = ("auto",)
|
|
|
|
|
|
def predictor_choices():
|
|
try:
|
|
return list(_NAMED) + veda_download.files()
|
|
except Exception: # noqa: BLE001 - no ComfyUI folders here (tests, dev server): a list of nothing is fine
|
|
return list(_NAMED)
|
|
|
|
|
|
def _pick(name):
|
|
"""The predictor file's path, or (None, why). "auto": the T2VA file when it is there, else the first."""
|
|
try:
|
|
have = veda_download.files()
|
|
except Exception: # noqa: BLE001
|
|
have = []
|
|
if not have:
|
|
return None, "no predictor in ComfyUI/models/veda (the Veda predictor row in the app offers the download)"
|
|
if name in have:
|
|
return veda_download.folder() / name, None
|
|
if name not in ("", *_NAMED):
|
|
return None, f"the predictor {name} is not in ComfyUI/models/veda"
|
|
t2va = veda_download.MODELS[0]["file"]
|
|
return veda_download.folder() / (t2va if t2va in have else have[0]), None
|
|
|
|
|
|
def _fold(veda):
|
|
"""The step that just ended: its work joins the run's totals (one read of the device per step, not per call), and
|
|
its own kept share is what the summary reports, as SLA reports its last step."""
|
|
stats = veda["stats"]
|
|
if stats.get("dense"):
|
|
veda["run"]["work"] += float(stats["work"])
|
|
veda["run"]["dense"] += stats["dense"]
|
|
veda["last"] = (float(stats["kept"]) / stats["pairs"] if stats.get("pairs") else 1.0, stats.get("global", 0))
|
|
stats.clear()
|
|
|
|
|
|
def make_override(state, cfg, bundle, attend, dense_fn=None, dense_label=None):
|
|
"""The `optimized_attention_override`. `attend` is the kernel callable, or None for veda_kernel.attend (checked once
|
|
per device, on the first call, against the fp32 reference). A kernel that fails costs speed, not the run."""
|
|
import torch
|
|
veda = state["veda"]
|
|
ok_dtypes = (torch.bfloat16, torch.float16)
|
|
|
|
def override(func, q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
def dense(why):
|
|
state["dense"] += 1
|
|
veda["why"][why] += 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 {}
|
|
layout, layer = to.get("minimax_h3_layout"), to.get("block_index")
|
|
if mask is not None:
|
|
return dense("masked call")
|
|
if not skip_reshape or q.ndim != 4 or q.shape[0] != 1 or q.shape != k.shape or q.shape != v.shape:
|
|
return dense("not H3's packed self-attention")
|
|
if q.shape[2] < cfg["min_seq_len"]:
|
|
return dense("shorter than min_seq_len")
|
|
if layout is None or getattr(layout, "seq_len", None) != q.shape[2]:
|
|
return dense("no layout for this call")
|
|
if q.dtype not in ok_dtypes:
|
|
return dense(f"{q.dtype} q/k/v")
|
|
if to.get("_funui_sla_dense", False):
|
|
return dense("dense step")
|
|
if not isinstance(layer, int):
|
|
layer = state["layer"]
|
|
state["layer"] += 1
|
|
if veda["step"] != state["step"]:
|
|
_fold(veda)
|
|
veda["step"] = state["step"]
|
|
if layer >= bundle.layers:
|
|
return dense("layer outside the predictor")
|
|
if q.shape[1] != bundle.heads or q.shape[3] != bundle.dim:
|
|
return dense(f"{q.shape[1]} heads of {q.shape[3]}, predictor has {bundle.heads} of {bundle.dim}")
|
|
if state["failed"] is not None:
|
|
return dense("kernel failed earlier")
|
|
try:
|
|
fn = attend
|
|
if fn is None:
|
|
from . import veda_kernel
|
|
fn = veda_kernel.attend
|
|
if str(q.device) not in veda["checked"]:
|
|
veda["checked"][str(q.device)] = veda_kernel.check(q.device)
|
|
if veda["checked"][str(q.device)]:
|
|
raise RuntimeError(veda["checked"][str(q.device)])
|
|
spec = veda["specs"].get(id(layout))
|
|
if spec is None or spec[0] is not layout:
|
|
spec = veda["specs"][id(layout)] = (layout, core.describe_layout(layout))
|
|
spec = spec[1]
|
|
plan = veda["plans"].get(spec[1][1])
|
|
if plan is None:
|
|
plan = veda["plans"][spec[1][1]] = core.choose_plan(bundle.plans, spec[1][1])
|
|
if cfg["timing"] and veda["timer"] is None:
|
|
veda["timer"] = core.PhaseTimer(q.device)
|
|
with sla.pause_malloc_graph(): # the cached tile layouts outlive the call: made outside ComfyUI's compiler graph
|
|
out = core.attention(q[0].transpose(0, 1), k[0].transpose(0, 1), v[0].transpose(0, 1), layer, spec, plan[0], bundle,
|
|
(bundle.generated, bundle.reference), fn, veda["tiles"], veda["stats"], veda["timer"] or core._NO_TIMER)
|
|
veda["geometry"][spec[1][1]] = plan[2] # every size this run met (Two stages meets two), and the plan each got
|
|
state["calls"] += 1
|
|
state["seq"] = q.shape[2]
|
|
return out.transpose(0, 1).unsqueeze(0) if skip_output_reshape else out.reshape(1, q.shape[2], -1)
|
|
except Exception as exc: # noqa: BLE001 - a bad kernel must not kill the run
|
|
state["failed"] = f"{exc.__class__.__name__}: {exc}"
|
|
logging.warning("[FunUI Veda] sparse attention failed, dense from here this run: %s", state["failed"], exc_info=True)
|
|
return dense("kernel failed")
|
|
|
|
return override
|
|
|
|
|
|
def summarise(state, cfg):
|
|
"""One line per sampling run (the wrapper calls it when the run ends), then the per-run caches are let go."""
|
|
veda, stats = state["veda"], state["veda"]["stats"]
|
|
why = ", ".join(f"{n} {w}" for w, n in veda["why"].most_common())
|
|
if state["calls"] == 0:
|
|
logging.warning("[FunUI Veda] installed but never invoked: attention was NOT sparsified (%d dense calls%s). Check that the "
|
|
"model reaching the sampler is the one this loader produced.", state["dense"], f": {why}" if why else "")
|
|
else:
|
|
# Same fields in the same order as the "[FunUI SLA]" line, so the two read side by side. Tile pairs are the video x
|
|
# video ones, counted over the whole run (SLA counts its last step); the global tiles (text, audio) are always exact.
|
|
_fold(veda)
|
|
sizes = veda["geometry"]
|
|
done = min(1.0, veda["run"]["work"] / veda["run"]["dense"]) if veda["run"]["dense"] else 0.0
|
|
kept, pinned = veda["last"]
|
|
logging.info("[FunUI Veda] %s | %d calls | S=%d | video tile pairs %.1f%% kept in the last step (%.1f%% sparse, asked %s) | %d global tiles pinned "
|
|
"| %d dense fall-throughs%s | displaced %s | computed %.1f%% of full attention | video %s | %s",
|
|
veda["kernel"], state["calls"], state["seq"], kept * 100.0, (1.0 - kept) * 100.0, cfg["asked"], pinned,
|
|
state["dense"], f" ({why})" if why else "", state["backend"] or "?", done * 100.0, "; ".join(f"{core.describe_grid(g)}, tile plan: {how}" for g, how in sizes.items()),
|
|
cfg["predictor_name"])
|
|
if veda["timer"] is not None:
|
|
phases = veda["timer"].totals() # a phase with a dot is a part of the one before it: not added again
|
|
top = [k for k in phases if "." not in k]
|
|
text = []
|
|
for k in top:
|
|
inner = ["%s %.0f" % (n.split(".")[1], v) for n, v in phases.items() if n.startswith(k + ".")]
|
|
text.append("%s %.0f%s" % (k, phases[k], " (%s)" % ", ".join(inner) if inner else ""))
|
|
logging.info("[FunUI Veda] time in each phase this run (ms): %s | total %.0f", ", ".join(text), sum(phases[k] for k in top))
|
|
logging.info("[FunUI Veda] full attention on %s of %d", f"steps {sla._runs(state['dense_at'])}" if state["dense_at"] else "no step", state["n"])
|
|
if state["failed"] is not None:
|
|
logging.warning("[FunUI Veda] the kernel fell back to dense at least once: %s", state["failed"])
|
|
stats.clear()
|
|
veda["run"].update(work=0.0, dense=0.0)
|
|
veda["step"] = None
|
|
veda["tiles"].clear()
|
|
veda["geometry"].clear()
|
|
veda["timer"] = None
|
|
veda["why"].clear()
|
|
veda["specs"].clear()
|
|
|
|
|
|
def install_veda(model, cfg, dense_fn=None, dense_label=None, attend=None):
|
|
"""Give `model` Veda attention. Returns (model, status, installed), the contract of sla.install_sla: `installed` is
|
|
False when Veda did not take, with the reason in the status, so the caller installs the chosen backend alone.
|
|
`attend` replaces the Triton kernel (tests only)."""
|
|
if attend is None and not sla.sla_available():
|
|
return model, "SLA: veda needs CUDA and Triton (its kernel is Triton's), not here", False
|
|
path, why = _pick(cfg.get("veda_predictor"))
|
|
if path is None:
|
|
return model, f"SLA: veda skipped, {why}", False
|
|
try:
|
|
bundle = core.load_bundle(str(path))
|
|
except (core.VedaError, OSError) as exc:
|
|
return model, f"SLA: veda skipped, {exc}", False
|
|
unused = [n for n, on in (("sla_sparsity", float(cfg["sparsity_ratio"]) != sla.SLA_DEFAULTS["sparsity_ratio"]),
|
|
("sla_sparsity_start", float(cfg["sparsity_start"]) > 0.0), ("sla_references", cfg["references"] != "off")) if on]
|
|
cfg = {**cfg, "sparsity_ratio": sla.SLA_DEFAULTS["sparsity_ratio"], "sparsity_start": 0.0, "references": "off",
|
|
"min_seq_len": int(cfg["min_seq_len"]), "dense_last_steps": int(cfg["dense_last_steps"]),
|
|
"predictor_name": path.name, "engine": "veda", "summarise": summarise, "timing": bool(cfg.get("veda_timing"))}
|
|
cfg["dense_steps"], bad = sla.parse_steps(cfg["dense_steps"])
|
|
notes = [f"dense_steps: ignored {', '.join(bad)} (steps count from 1: '1' is the first, '1-2' the first two)"] if bad else []
|
|
if unused:
|
|
notes.append(f"{', '.join(unused)}: not used by veda, which runs at the predictor's trained budget")
|
|
state = sla.new_state()
|
|
state["veda"] = {"stats": {}, "tiles": {}, "specs": {}, "plans": {}, "checked": {}, "why": collections.Counter(), "geometry": {}, "timer": None, "step": None, "run": {"work": 0.0, "dense": 0.0}, "last": (1.0, 0), "kernel": "triton-int8" if attend is None else "stand-in kernel"}
|
|
patched = sla.patching.clone(model)
|
|
to = patched.model_options.get("transformer_options", {}).copy()
|
|
to["optimized_attention_override"] = make_override(state, cfg, bundle, attend, dense_fn, dense_label)
|
|
patched.model_options["transformer_options"] = to
|
|
sla.add_wrappers(patched, state, cfg)
|
|
cfg["asked"] = f"trained ({core.describe_budget(bundle.generated)})"
|
|
return patched, "\n".join([
|
|
f"SLA on | veda (EXPERIMENTAL, tested on a GPU) | predictor {core.describe_bundle(bundle)} | sparsity {cfg['asked']} | "
|
|
f"references {core.describe_budget(bundle.reference)} | min_seq_len={cfg['min_seq_len']} | full attention on steps "
|
|
f"{sla._runs(i + 1 for i in cfg['dense_steps']) or 'none'} and the last {cfg['dense_last_steps']} | dense calls -> {dense_label or 'as launched'}",
|
|
*(f"SLA: {n}" for n in notes)]), True
|