ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-09 06:13:18 +00:00
commit 390bfae9e2
17 changed files with 1172 additions and 609 deletions
Download patch file Download diff file Expand all files Collapse all files

300
scripts/fp8_layer_ablation.py
View file

@ -67,7 +67,6 @@ from core.inference.diffusion_transformer_quant import ( # noqa: E402
# ── cuda helpers ────────────────────────────────────────────────────────────────
def _empty() -> None:
import torch
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
@ -121,8 +120,9 @@ class _Stop(Exception):
pass
def _capture_forward_tuple(pipe, *, steps, width, height, num_frames, guidance, seed,
capture_call, gvg):
def _capture_forward_tuple(
pipe, *, steps, width, height, num_frames, guidance, seed, capture_call, gvg
):
"""Run a bf16 generation and grab the ``capture_call``-th transformer forward's
(args, kwargs) to CPU, then abort. This is a REAL mid-schedule DiT input."""
import torch
@ -138,11 +138,15 @@ def _capture_forward_tuple(pipe, *, steps, width, height, num_frames, guidance,
raise _Stop()
return None
handle = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs=True)
g = torch.Generator(device="cuda").manual_seed(seed)
handle = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs = True)
g = torch.Generator(device = "cuda").manual_seed(seed)
kwargs = dict(
prompt=PROMPT, width=width, height=height, num_frames=num_frames,
num_inference_steps=steps, generator=g,
prompt = PROMPT,
width = width,
height = height,
num_frames = num_frames,
num_inference_steps = steps,
generator = g,
)
if gvg:
guider = getattr(pipe, "guider", None)
@ -192,8 +196,12 @@ def _score(ref, cand) -> dict:
a = canf[mask]
b = reff[mask]
cos = float(torch.nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item())
rel = float((torch.linalg.vector_norm(a - b) / (torch.linalg.vector_norm(b) + 1e-12)).item())
norm_ratio = float((torch.linalg.vector_norm(a) / (torch.linalg.vector_norm(b) + 1e-12)).item())
rel = float(
(torch.linalg.vector_norm(a - b) / (torch.linalg.vector_norm(b) + 1e-12)).item()
)
norm_ratio = float(
(torch.linalg.vector_norm(a) / (torch.linalg.vector_norm(b) + 1e-12)).item()
)
else:
cos, rel, norm_ratio = 0.0, float("inf"), 0.0
return {
@ -201,7 +209,9 @@ def _score(ref, cand) -> dict:
"relL2": round(rel, 5),
"frac_nonfinite": round(frac_nonfinite, 6),
"norm_ratio": round(norm_ratio, 5),
"cand_amax": round(float(canf[fin_c].abs().max().item()) if fin_c.any() else float("inf"), 3),
"cand_amax": round(
float(canf[fin_c].abs().max().item()) if fin_c.any() else float("inf"), 3
),
}
@ -217,12 +227,19 @@ def _verdict(sc: dict) -> str:
# ── quant helpers ──────────────────────────────────────────────────────────────
def _fp8_filter(min_features: int, exclude_tokens=(), only_tokens=()):
def _fp8_filter(
min_features: int,
exclude_tokens = (),
only_tokens = (),
):
"""A production fp8 filter (require_bf16) optionally narrowed:
- exclude_tokens: skip linears whose fqn contains any token (keep them bf16)
- only_tokens : quantise ONLY linears whose fqn contains any token."""
base = make_filter_fn(min_features, exclude_name_tokens=exclude_tokens,
require_bf16=(TQ_FP8 in _REQUIRE_BF16_SCHEMES))
base = make_filter_fn(
min_features,
exclude_name_tokens = exclude_tokens,
require_bf16 = (TQ_FP8 in _REQUIRE_BF16_SCHEMES),
)
def filt(module, fqn: str = "") -> bool:
if not base(module, fqn):
@ -237,16 +254,24 @@ def _fp8_filter(min_features: int, exclude_tokens=(), only_tokens=()):
return filt
def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(), only_tokens=(),
instrument=False):
def _ablate(
dense_cpu,
tup,
*,
min_features,
fast_accum = None,
exclude_tokens = (),
only_tokens = (),
instrument = False,
):
"""Deepcopy the dense DiT -> GPU -> fp8-quantise the selected subset -> one forward.
Returns (cpu_output_tensor, finiteness_records|None)."""
import torch
from torchao.quantization import quantize_
m = copy.deepcopy(dense_cpu).to("cuda")
filt = _fp8_filter(min_features, exclude_tokens=exclude_tokens, only_tokens=only_tokens)
quantize_(m, _make_quant_config(TQ_FP8, fast_accum=fast_accum), filter_fn=filt)
filt = _fp8_filter(min_features, exclude_tokens = exclude_tokens, only_tokens = only_tokens)
quantize_(m, _make_quant_config(TQ_FP8, fast_accum = fast_accum), filter_fn = filt)
records = None
hooks = []
@ -255,7 +280,11 @@ def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(),
def mk(nm):
def hook(mod, inp, out):
o = out if torch.is_tensor(out) else (out[0] if isinstance(out, (list, tuple)) and out else None)
o = (
out
if torch.is_tensor(out)
else (out[0] if isinstance(out, (list, tuple)) and out else None)
)
if o is None:
return
of = o.detach().float()
@ -265,15 +294,20 @@ def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(),
xin = inp[0].detach().float()
in_amax = float(xin.abs().max().item())
if xin.dim() >= 2:
ra = xin.reshape(-1, xin.shape[-1]).abs().amax(dim=-1)
ra = xin.reshape(-1, xin.shape[-1]).abs().amax(dim = -1)
in_frac_zero_rows = round(float((ra < 1e-6).float().mean().item()), 4)
records.append({
"name": nm,
"out_finite": bool(torch.isfinite(of).all().item()),
"out_amax": float(of.abs().max().item()) if torch.isfinite(of).any() else float("inf"),
"in_amax": in_amax,
"in_frac_zero_rows": in_frac_zero_rows,
})
records.append(
{
"name": nm,
"out_finite": bool(torch.isfinite(of).all().item()),
"out_amax": float(of.abs().max().item())
if torch.isfinite(of).any()
else float("inf"),
"in_amax": in_amax,
"in_frac_zero_rows": in_frac_zero_rows,
}
)
return hook
for nm, mod in m.named_modules():
@ -332,7 +366,7 @@ def _zero_row_diag(tup: dict) -> dict:
for key, val in tup.get("kwargs", {}).items():
if torch.is_tensor(val) and val.dim() >= 2 and val.is_floating_point():
x = val.float().reshape(-1, val.shape[-1])
row_amax = x.abs().amax(dim=-1)
row_amax = x.abs().amax(dim = -1)
n = int(row_amax.numel())
n_zero = int((row_amax < 1e-6).sum().item())
diag[key] = {
@ -360,16 +394,19 @@ def _bf16_reference_and_stats(dense_gpu, tup, *, min_features):
return
x = inp[0].detach().float()
x2 = x.reshape(-1, x.shape[-1]).abs() # [tokens, channels]
amax_tok = x2.amax(dim=-1)
med_tok = x2.median(dim=-1).values
amax_tok = x2.amax(dim = -1)
med_tok = x2.median(dim = -1).values
spread = float((amax_tok / (med_tok + 1e-9)).mean().item())
stats.append({
"name": nm,
"in_amax": float(x2.max().item()),
"spread": round(spread, 2),
"in_features": int(mod.in_features),
"out_features": int(mod.out_features),
})
stats.append(
{
"name": nm,
"in_amax": float(x2.max().item()),
"spread": round(spread, 2),
"in_features": int(mod.in_features),
"out_features": int(mod.out_features),
}
)
return hook
for nm, mod in dense_gpu.named_modules():
@ -380,31 +417,35 @@ def _bf16_reference_and_stats(dense_gpu, tup, *, min_features):
ref = ref.detach().to("cpu")
for h in hooks:
h.remove()
stats.sort(key=lambda s: s["spread"], reverse=True)
stats.sort(key = lambda s: s["spread"], reverse = True)
return ref, stats
def main(argv=None) -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--family", default="wan2.2-ti2v-5b", choices=sorted(_FAMILIES))
ap.add_argument("--phase", type=int, default=0, choices=[0, 1, 2])
ap.add_argument("--only", default="", help="phase 2: comma tokens -- fp8 ONLY these linears")
ap.add_argument("--exclude", default="", help="phase 2: comma tokens -- fp8 all EXCEPT these")
ap.add_argument("--steps", type=int, default=20)
ap.add_argument("--num-frames", type=int, default=25)
ap.add_argument("--width", type=int, default=512)
ap.add_argument("--height", type=int, default=320)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--capture-call", type=int, default=8,
help="which transformer forward call to capture (~mid schedule)")
ap.add_argument("--min-features", type=int, default=DEFAULT_MIN_LINEAR_FEATURES)
ap.add_argument("--out", default="outputs/fp8_ablation")
def main(argv = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("--family", default = "wan2.2-ti2v-5b", choices = sorted(_FAMILIES))
ap.add_argument("--phase", type = int, default = 0, choices = [0, 1, 2])
ap.add_argument("--only", default = "", help = "phase 2: comma tokens -- fp8 ONLY these linears")
ap.add_argument("--exclude", default = "", help = "phase 2: comma tokens -- fp8 all EXCEPT these")
ap.add_argument("--steps", type = int, default = 20)
ap.add_argument("--num-frames", type = int, default = 25)
ap.add_argument("--width", type = int, default = 512)
ap.add_argument("--height", type = int, default = 320)
ap.add_argument("--seed", type = int, default = 42)
ap.add_argument(
"--capture-call",
type = int,
default = 8,
help = "which transformer forward call to capture (~mid schedule)",
)
ap.add_argument("--min-features", type = int, default = DEFAULT_MIN_LINEAR_FEATURES)
ap.add_argument("--out", default = "outputs/fp8_ablation")
args = ap.parse_args(argv)
import torch
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
out.mkdir(parents = True, exist_ok = True)
spec = _FAMILIES[args.family]
repo = spec["repo"]
@ -416,79 +457,111 @@ def main(argv=None) -> int:
fam_obj = detect_video_family(repo)
gvg = bool(getattr(fam_obj, "guidance_via_guider", False))
print(f"== fp8 layer ablation: family={args.family} repo={repo} phase={args.phase} ==", flush=True)
print(
f"== fp8 layer ablation: family={args.family} repo={repo} phase={args.phase} ==", flush = True
)
t0 = time.perf_counter()
pipe = _build_pipe(repo, force_fp32)
print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush=True)
print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush = True)
tup = _capture_forward_tuple(
pipe, steps=args.steps, width=args.width, height=args.height,
num_frames=args.num_frames, guidance=guidance, seed=args.seed,
capture_call=args.capture_call, gvg=gvg,
pipe,
steps = args.steps,
width = args.width,
height = args.height,
num_frames = args.num_frames,
guidance = guidance,
seed = args.seed,
capture_call = args.capture_call,
gvg = gvg,
)
ks = {k: (tuple(v.shape) if torch.is_tensor(v) else type(v).__name__)
for k, v in tup.get("kwargs", {}).items()}
print(f"[capture] call #{args.capture_call} kwargs={ks}", flush=True)
ks = {
k: (tuple(v.shape) if torch.is_tensor(v) else type(v).__name__)
for k, v in tup.get("kwargs", {}).items()
}
print(f"[capture] call #{args.capture_call} kwargs={ks}", flush = True)
# Keep the dense bf16 DiT (still on GPU) as reference source; copy to CPU as the ablation seed.
dense_gpu = pipe.transformer
ref, stats = _bf16_reference_and_stats(dense_gpu, tup, min_features=args.min_features)
print(f"[bf16 ref] out shape={tuple(ref.shape)} amax={float(ref.float().abs().max()):.3f} "
f"n_quantized_linears={len(stats)}", flush=True)
ref, stats = _bf16_reference_and_stats(dense_gpu, tup, min_features = args.min_features)
print(
f"[bf16 ref] out shape={tuple(ref.shape)} amax={float(ref.float().abs().max()):.3f} "
f"n_quantized_linears={len(stats)}",
flush = True,
)
dense_cpu = copy.deepcopy(dense_gpu).to("cpu")
# Free the pipeline (VAE/text encoders) to leave the GPU for one transformer at a time.
del pipe, dense_gpu
_empty()
report: dict[str, Any] = {"family": args.family, "repo": repo, "capture_call": args.capture_call,
"input_shapes": ks, "ref_amax": round(float(ref.float().abs().max()), 3),
"zero_row_diag": _zero_row_diag(tup)}
print(f"[zero-row diag] {report['zero_row_diag']}", flush=True)
report: dict[str, Any] = {
"family": args.family,
"repo": repo,
"capture_call": args.capture_call,
"input_shapes": ks,
"ref_amax": round(float(ref.float().abs().max()), 3),
"zero_row_diag": _zero_row_diag(tup),
}
print(f"[zero-row diag] {report['zero_row_diag']}", flush = True)
if args.phase == 0:
# ── Phase 0a: proxy validity -- fp8-ALL (production accum = None auto-detect) ──
out_all, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None)
out_all, _ = _ablate(dense_cpu, tup, min_features = args.min_features, fast_accum = None)
sc_all = _score(ref, out_all)
report["fp8_all_autoaccum"] = {**sc_all, "verdict": _verdict(sc_all)}
print(f"[0a fp8-ALL auto-accum] {sc_all} -> {_verdict(sc_all)}", flush=True)
print(f"[0a fp8-ALL auto-accum] {sc_all} -> {_verdict(sc_all)}", flush = True)
# ── Phase 0b-ii: accumulate flip ──
for fa in (False, True):
o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=fa)
o, _ = _ablate(dense_cpu, tup, min_features = args.min_features, fast_accum = fa)
sc = _score(ref, o)
report[f"fp8_all_fast_accum_{fa}"] = {**sc, "verdict": _verdict(sc)}
print(f"[0b accum={fa}] {sc} -> {_verdict(sc)}", flush=True)
print(f"[0b accum={fa}] {sc} -> {_verdict(sc)}", flush = True)
# ── Phase 0b-i: per-Linear finiteness (instrumented fp8-ALL forward) ──
_, records = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None,
instrument=True)
_, records = _ablate(
dense_cpu, tup, min_features = args.min_features, fast_accum = None, instrument = True
)
if records:
nonfinite = [r for r in records if not r["out_finite"]]
report["first_nonfinite_linears"] = nonfinite[:10]
report["n_nonfinite_linears"] = len(nonfinite)
finite_sorted = sorted([r for r in records if r["out_finite"]],
key=lambda r: r["out_amax"], reverse=True)
finite_sorted = sorted(
[r for r in records if r["out_finite"]], key = lambda r: r["out_amax"], reverse = True
)
report["top_out_amax_linears"] = finite_sorted[:10]
print(f"[0b finiteness] {len(nonfinite)}/{len(records)} linears non-finite; "
f"first={[r['name'] for r in nonfinite[:5]]}", flush=True)
print(f"[0b top out_amax] {[(r['name'], round(r['out_amax'],1)) for r in finite_sorted[:5]]}",
flush=True)
print(
f"[0b finiteness] {len(nonfinite)}/{len(records)} linears non-finite; "
f"first={[r['name'] for r in nonfinite[:5]]}",
flush = True,
)
print(
f"[0b top out_amax] {[(r['name'], round(r['out_amax'],1)) for r in finite_sorted[:5]]}",
flush = True,
)
report["top_spread_linears"] = stats[:20]
print(f"[stats] top-spread linears: "
f"{[(s['name'], s['spread']) for s in stats[:8]]}", flush=True)
print(
f"[stats] top-spread linears: " f"{[(s['name'], s['spread']) for s in stats[:8]]}",
flush = True,
)
va = report["fp8_all_autoaccum"]["verdict"]
if va.startswith("BROKEN"):
print("\n[GATE] fp8-ALL reproduces the black signal in one forward -> proxy VALID; "
"proceed to Phase 1 bucket ablation.", flush=True)
print(
"\n[GATE] fp8-ALL reproduces the black signal in one forward -> proxy VALID; "
"proceed to Phase 1 bucket ablation.",
flush = True,
)
elif va == "CLEAN":
print("\n[GATE] fp8-ALL is CLEAN in one forward but full gen is black -> failure is "
"multi-step/cache; single-forward bisection cannot localise -> STOP, keep int8.",
flush=True)
print(
"\n[GATE] fp8-ALL is CLEAN in one forward but full gen is black -> failure is "
"multi-step/cache; single-forward bisection cannot localise -> STOP, keep int8.",
flush = True,
)
else:
print(f"\n[GATE] fp8-ALL verdict={va}: borderline; inspect scores.", flush=True)
print(f"\n[GATE] fp8-ALL verdict={va}: borderline; inspect scores.", flush = True)
elif args.phase == 1: # ── Phase 1: bucket ablation (necessity + sufficiency) ──
kind = _dit_kind(args.family)
@ -496,53 +569,68 @@ def main(argv=None) -> int:
results: dict[str, Any] = {}
# baseline: fp8-ALL (should be BROKEN, matching Phase 0)
o, _ = _ablate(dense_cpu, tup, min_features=args.min_features)
o, _ = _ablate(dense_cpu, tup, min_features = args.min_features)
results["fp8_all"] = {**_score(ref, o)}
results["fp8_all"]["verdict"] = _verdict(results["fp8_all"])
print(f"[1 baseline fp8-ALL] {results['fp8_all']['verdict']} {results['fp8_all']}", flush=True)
print(
f"[1 baseline fp8-ALL] {results['fp8_all']['verdict']} {results['fp8_all']}", flush = True
)
for bname, toks in buckets.items():
# NECESSITY: fp8 everything EXCEPT this bucket -- if CLEAN, this bucket is the culprit.
o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, exclude_tokens=toks)
o, _ = _ablate(dense_cpu, tup, min_features = args.min_features, exclude_tokens = toks)
sc_ex = _score(ref, o)
# SUFFICIENCY: fp8 ONLY this bucket -- if BROKEN, this bucket alone reproduces damage.
o2, _ = _ablate(dense_cpu, tup, min_features=args.min_features, only_tokens=toks)
o2, _ = _ablate(dense_cpu, tup, min_features = args.min_features, only_tokens = toks)
sc_only = _score(ref, o2)
results[bname] = {
"exclude": {**sc_ex, "verdict": _verdict(sc_ex)},
"only": {**sc_only, "verdict": _verdict(sc_only)},
}
print(f"[1 {bname}] EXCLUDE->{_verdict(sc_ex)} (cos {sc_ex['cosine']}, "
f"nf {sc_ex['frac_nonfinite']}) | ONLY->{_verdict(sc_only)} "
f"(cos {sc_only['cosine']}, nf {sc_only['frac_nonfinite']})", flush=True)
print(
f"[1 {bname}] EXCLUDE->{_verdict(sc_ex)} (cos {sc_ex['cosine']}, "
f"nf {sc_ex['frac_nonfinite']}) | ONLY->{_verdict(sc_only)} "
f"(cos {sc_only['cosine']}, nf {sc_only['frac_nonfinite']})",
flush = True,
)
report["phase1"] = results
else: # ── Phase 2: instrument a specific only/exclude set, find first overflow ──
only = tuple(t.strip() for t in args.only.split(",") if t.strip())
exclude = tuple(t.strip() for t in args.exclude.split(",") if t.strip())
o, records = _ablate(dense_cpu, tup, min_features=args.min_features,
only_tokens=only, exclude_tokens=exclude, instrument=True)
o, records = _ablate(
dense_cpu,
tup,
min_features = args.min_features,
only_tokens = only,
exclude_tokens = exclude,
instrument = True,
)
sc = _score(ref, o)
report["phase2"] = {"only": only, "exclude": exclude, **sc, "verdict": _verdict(sc)}
print(f"[2 only={only} exclude={exclude}] {sc} -> {_verdict(sc)}", flush=True)
print(f"[2 only={only} exclude={exclude}] {sc} -> {_verdict(sc)}", flush = True)
if records:
# first non-finite linears in execution order, with their input amax: an in_amax of
# ~0 at the first-broken layer confirms a zero-amax (padding) input row -> scale 0 -> inf.
nonfinite = [r for r in records if not r["out_finite"]]
report["phase2_first_nonfinite"] = nonfinite[:15]
report["phase2_n_nonfinite"] = len(nonfinite)
print(f"[2 finiteness] {len(nonfinite)}/{len(records)} quantised-path linears non-finite",
flush=True)
print(
f"[2 finiteness] {len(nonfinite)}/{len(records)} quantised-path linears non-finite",
flush = True,
)
for r in nonfinite[:8]:
print(f" first-inf {r['name']} in_amax={r['in_amax']} "
f"in_frac_zero_rows={r['in_frac_zero_rows']} out_amax={r['out_amax']}",
flush=True)
print(
f" first-inf {r['name']} in_amax={r['in_amax']} "
f"in_frac_zero_rows={r['in_frac_zero_rows']} out_amax={r['out_amax']}",
flush = True,
)
dest = out / f"phase{args.phase}_{args.family}.json"
with open(dest, "w", encoding="utf-8") as fh:
json.dump(report, fh, indent=2)
print(f"wrote {dest}", flush=True)
with open(dest, "w", encoding = "utf-8") as fh:
json.dump(report, fh, indent = 2)
print(f"wrote {dest}", flush = True)
return 0

129
scripts/hunyuan_attn_diag.py
View file

@ -10,6 +10,7 @@ so we know whether the win is the N-reduction (trim) or the mask-elimination (nu
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_attn_diag.py [--repo ...] [--frames 121]
"""
from __future__ import annotations
import argparse
@ -62,8 +63,8 @@ def _block_pre_hook(module, args, kwargs):
CAP["dtype"] = hs.dtype
if amask is not None:
m = amask.bool()
CAP["text_valid_per_batch"] = m.sum(dim=1).tolist()
CAP["text_cols_valid_any"] = int(m.any(dim=0).sum()) # what our global-trim would keep
CAP["text_valid_per_batch"] = m.sum(dim = 1).tolist()
CAP["text_cols_valid_any"] = int(m.any(dim = 0).sum()) # what our global-trim would keep
raise _StopCapture
@ -72,14 +73,16 @@ def _model_pre_hook(module, args, kwargs):
def g(name):
return kwargs.get(name)
for key, mkey in (("encoder_hidden_states", "encoder_attention_mask"),
("encoder_hidden_states_2", "encoder_attention_mask_2")):
for key, mkey in (
("encoder_hidden_states", "encoder_attention_mask"),
("encoder_hidden_states_2", "encoder_attention_mask_2"),
):
s = g(key)
m = g(mkey)
if s is not None:
CAP.setdefault("streams", {})[key] = {
"len": int(s.shape[1]),
"valid": (m.bool().sum(dim=1).tolist() if m is not None else None),
"valid": (m.bool().sum(dim = 1).tolist() if m is not None else None),
}
ie = g("image_embeds")
if ie is not None:
@ -88,44 +91,52 @@ def _model_pre_hook(module, args, kwargs):
return None
def _time_sdpa(q, k, v, mask, iters=30):
def _time_sdpa(
q,
k,
v,
mask,
iters = 30,
):
# q,k,v: [B, H, N, D]
torch.cuda.synchronize()
for _ in range(3): # warmup
F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
F.scaled_dot_product_attention(q, k, v, attn_mask = mask)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
F.scaled_dot_product_attention(q, k, v, attn_mask = mask)
torch.cuda.synchronize()
return (time.perf_counter() - t0) / iters * 1e3 # ms
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type=int, default=121)
ap.add_argument("--width", type=int, default=832)
ap.add_argument("--height", type=int, default=480)
ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type = int, default = 121)
ap.add_argument("--width", type = int, default = 832)
ap.add_argument("--height", type = int, default = 480)
args = ap.parse_args()
diffusers = _import_diffusers()
dev = "cuda:0"
print(f"loading {args.repo} ...", flush=True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16)
print(f"loading {args.repo} ...", flush = True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16)
pipe = pipe.to(dev)
pipe.transformer.register_forward_pre_hook(_model_pre_hook, with_kwargs=True)
pipe.transformer.transformer_blocks[0].register_forward_pre_hook(_block_pre_hook, with_kwargs=True)
pipe.transformer.register_forward_pre_hook(_model_pre_hook, with_kwargs = True)
pipe.transformer.transformer_blocks[0].register_forward_pre_hook(
_block_pre_hook, with_kwargs = True
)
print("running 1 capture step ...", flush=True)
print("running 1 capture step ...", flush = True)
try:
pipe(
prompt="a cat playing piano",
num_frames=args.frames,
width=args.width,
height=args.height,
num_inference_steps=1,
prompt = "a cat playing piano",
num_frames = args.frames,
width = args.width,
height = args.height,
num_inference_steps = 1,
)
except _StopCapture:
pass
@ -133,14 +144,24 @@ def main():
# the StopCapture may surface wrapped; if we captured, continue
if "n_video" not in CAP:
raise
print(f"(generation aborted after capture: {type(exc).__name__})", flush=True)
print(f"(generation aborted after capture: {type(exc).__name__})", flush = True)
print("\n===== CAPTURED SHAPES =====", flush=True)
for kk in ("batch", "n_video", "n_text", "heads", "dim_head", "dtype",
"text_valid_per_batch", "text_cols_valid_any", "image_embeds_len",
"image_is_t2v", "streams"):
print("\n===== CAPTURED SHAPES =====", flush = True)
for kk in (
"batch",
"n_video",
"n_text",
"heads",
"dim_head",
"dtype",
"text_valid_per_batch",
"text_cols_valid_any",
"image_embeds_len",
"image_is_t2v",
"streams",
):
if kk in CAP:
print(f" {kk}: {CAP[kk]}", flush=True)
print(f" {kk}: {CAP[kk]}", flush = True)
B = CAP["batch"]
H = CAP["heads"]
@ -152,42 +173,62 @@ def main():
keep_text = CAP.get("text_cols_valid_any", n_text)
N_trim = n_video + keep_text
dtype = CAP["dtype"]
print(f"\n joint N = {N} (video {n_video} + text {n_text}); "
f"trimmed N = {N_trim} (text kept {keep_text})", flush=True)
print(
f"\n joint N = {N} (video {n_video} + text {n_text}); "
f"trimmed N = {N_trim} (text kept {keep_text})",
flush = True,
)
def mk(n):
return torch.randn(B, H, n, D, device=dev, dtype=dtype)
return torch.randn(B, H, n, D, device = dev, dtype = dtype)
# (a) dense mask over full N (current). Build [B,1,N,N] bool (mostly True).
print("\n===== SDPA TIMING (ms/call, real shapes) =====", flush=True)
print("\n===== SDPA TIMING (ms/call, real shapes) =====", flush = True)
q, k, v = mk(N), mk(N), mk(N)
dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev)
dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev)
# emulate text padding: last (n_text - keep_text) columns invalid
if n_text - keep_text > 0:
dense[:, :, :, n_video + keep_text:] = False
dense[:, :, n_video + keep_text:, :] = False
dense[:, :, :, n_video + keep_text :] = False
dense[:, :, n_video + keep_text :, :] = False
t_dense = _time_sdpa(q, k, v, dense)
mask_gb = dense.numel() / 1e9
print(f" (a) dense [B,1,N,N] mask N={N:>6} : {t_dense:7.3f} ms (mask {mask_gb:.2f} GB)", flush=True)
print(
f" (a) dense [B,1,N,N] mask N={N:>6} : {t_dense:7.3f} ms (mask {mask_gb:.2f} GB)",
flush = True,
)
# (b) no mask over full N (upper bound of flash path if all valid)
t_none = _time_sdpa(q, k, v, None)
print(f" (b) attn_mask=None N={N:>6} : {t_none:7.3f} ms ({t_dense/t_none:.2f}x vs a)", flush=True)
print(
f" (b) attn_mask=None N={N:>6} : {t_none:7.3f} ms ({t_dense/t_none:.2f}x vs a)",
flush = True,
)
# (c) trimmed N, dense all-True mask (text padding removed but mask still built)
qt, kt, vt = mk(N_trim), mk(N_trim), mk(N_trim)
dense_t = torch.ones(B, 1, N_trim, N_trim, dtype=torch.bool, device=dev)
dense_t = torch.ones(B, 1, N_trim, N_trim, dtype = torch.bool, device = dev)
t_dense_trim = _time_sdpa(qt, kt, vt, dense_t)
print(f" (c) dense mask @trimmed N={N_trim:>6} : {t_dense_trim:7.3f} ms ({t_dense/t_dense_trim:.2f}x vs a)", flush=True)
print(
f" (c) dense mask @trimmed N={N_trim:>6} : {t_dense_trim:7.3f} ms ({t_dense/t_dense_trim:.2f}x vs a)",
flush = True,
)
# (d) trimmed N, no mask (trim + null: the full proposed fast path)
t_none_trim = _time_sdpa(qt, kt, vt, None)
print(f" (d) no mask @trimmed N={N_trim:>6} : {t_none_trim:7.3f} ms ({t_dense/t_none_trim:.2f}x vs a)", flush=True)
print(
f" (d) no mask @trimmed N={N_trim:>6} : {t_none_trim:7.3f} ms ({t_dense/t_none_trim:.2f}x vs a)",
flush = True,
)
print("\n Interpretation:", flush=True)
print(f" trim-only ceiling (a->c): {(1-t_dense_trim/t_dense)*100:5.1f}% attn saving", flush=True)
print(f" null-only ceiling (a->b): {(1-t_none/t_dense)*100:5.1f}% attn saving", flush=True)
print(f" trim+null (a->d): {(1-t_none_trim/t_dense)*100:5.1f}% attn saving", flush=True)
print("\n Interpretation:", flush = True)
print(
f" trim-only ceiling (a->c): {(1-t_dense_trim/t_dense)*100:5.1f}% attn saving",
flush = True,
)
print(f" null-only ceiling (a->b): {(1-t_none/t_dense)*100:5.1f}% attn saving", flush = True)
print(
f" trim+null (a->d): {(1-t_none_trim/t_dense)*100:5.1f}% attn saving", flush = True
)
if __name__ == "__main__":

97
scripts/hunyuan_int8_profile.py
View file

@ -38,7 +38,12 @@ def _dynamo_counts():
"""(recompiles, graph_breaks, unique_compiles) from dynamo counters, best-effort."""
try:
from torch._dynamo.utils import counters
rc = sum(v for k, v in counters.get("recompiles", {}).items()) if "recompiles" in counters else 0
rc = (
sum(v for k, v in counters.get("recompiles", {}).items())
if "recompiles" in counters
else 0
)
gb = sum(counters.get("graph_break", {}).values())
# total frames compiled
stats = counters.get("stats", {})
@ -48,8 +53,9 @@ def _dynamo_counts():
return -1, -1, -1
def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames: int, guidance: float,
seed: int):
def _profile_mode(
mode: str, *, steps: int, width: int, height: int, num_frames: int, guidance: float, seed: int
):
import torch
from core.inference.video_families import detect_video_family
from core.inference.diffusion_cache import maybe_toggle_step_cache # noqa: F401
@ -67,39 +73,59 @@ def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames:
default_steps = getattr(fam_obj, "default_steps", 50)
pipe = _build_pipe(_REPO, False)
cfg = dict(te="none", vae="none", dit=("int8" if mode == "int8" else "none"),
speed="default", attn="native", cache="off")
engaged = _apply_levers(pipe, cfg, fam_name="hunyuanvideo-1.5", fam_obj=fam_obj,
force_fp32_vae=False, default_steps=default_steps)
print(f"[{mode}] dit_scheme={engaged['dit'] or 'dense'} speed={engaged.get('_effective_speed')}",
flush=True)
cfg = dict(
te = "none",
vae = "none",
dit = ("int8" if mode == "int8" else "none"),
speed = "default",
attn = "native",
cache = "off",
)
engaged = _apply_levers(
pipe,
cfg,
fam_name = "hunyuanvideo-1.5",
fam_obj = fam_obj,
force_fp32_vae = False,
default_steps = default_steps,
)
print(
f"[{mode}] dit_scheme={engaged['dit'] or 'dense'} speed={engaged.get('_effective_speed')}",
flush = True,
)
# per-forward GPU timing via cuda events on the transformer.
fwd_ms: list[float] = []
starts: list = []
def _pre(mod, args, kwargs):
ev = torch.cuda.Event(enable_timing=True)
ev = torch.cuda.Event(enable_timing = True)
ev.record()
starts.append(ev)
return None
def _post(mod, args, output):
end = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing = True)
end.record()
torch.cuda.synchronize()
if starts:
fwd_ms.append(starts[-1].elapsed_time(end))
h1 = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs=True)
h1 = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs = True)
h2 = pipe.transformer.register_forward_hook(_post)
def _gen(tag):
fwd_ms.clear()
starts.clear()
g = torch.Generator(device="cuda").manual_seed(seed)
kwargs = dict(prompt=PROMPT, width=width, height=height, num_frames=num_frames,
num_inference_steps=steps, generator=g)
g = torch.Generator(device = "cuda").manual_seed(seed)
kwargs = dict(
prompt = PROMPT,
width = width,
height = height,
num_frames = num_frames,
num_inference_steps = steps,
generator = g,
)
if gvg:
guider = getattr(pipe, "guider", None)
if guider is not None and hasattr(guider, "guidance_scale"):
@ -118,36 +144,47 @@ def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames:
p10 = srt[max(0, n // 10)] if n else 0.0
p90 = srt[min(n - 1, (9 * n) // 10)] if n else 0.0
slow = sum(1 for x in fwd_ms if x > 2.0 * med) if med else 0
print(f"[{mode}:{tag}] total={total:.2f}s n_fwd={n} med={med:.1f}ms "
f"p10={p10:.1f} p90={p90:.1f} max={max(fwd_ms) if fwd_ms else 0:.1f} "
f"slow(>2x med)={slow} | recompiles={rc} graph_breaks={gb} unique_graphs={uc}",
flush=True)
print(
f"[{mode}:{tag}] total={total:.2f}s n_fwd={n} med={med:.1f}ms "
f"p10={p10:.1f} p90={p90:.1f} max={max(fwd_ms) if fwd_ms else 0:.1f} "
f"slow(>2x med)={slow} | recompiles={rc} graph_breaks={gb} unique_graphs={uc}",
flush = True,
)
return total
_gen("warmup") # pays compile
_gen("timed")
_gen("timed2")
h1.remove(); h2.remove()
h1.remove()
h2.remove()
del pipe
import gc
gc.collect()
torch.cuda.empty_cache()
def main(argv=None):
def main(argv = None):
ap = argparse.ArgumentParser()
ap.add_argument("--modes", default="dense,int8")
ap.add_argument("--steps", type=int, default=30)
ap.add_argument("--num-frames", type=int, default=25)
ap.add_argument("--width", type=int, default=512)
ap.add_argument("--height", type=int, default=320)
ap.add_argument("--guidance", type=float, default=6.0)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--modes", default = "dense,int8")
ap.add_argument("--steps", type = int, default = 30)
ap.add_argument("--num-frames", type = int, default = 25)
ap.add_argument("--width", type = int, default = 512)
ap.add_argument("--height", type = int, default = 320)
ap.add_argument("--guidance", type = float, default = 6.0)
ap.add_argument("--seed", type = int, default = 42)
args = ap.parse_args(argv)
for mode in [m.strip() for m in args.modes.split(",") if m.strip()]:
_profile_mode(mode, steps=args.steps, width=args.width, height=args.height,
num_frames=args.num_frames, guidance=args.guidance, seed=args.seed)
_profile_mode(
mode,
steps = args.steps,
width = args.width,
height = args.height,
num_frames = args.num_frames,
guidance = args.guidance,
seed = args.seed,
)
return 0

90
scripts/hunyuan_trim_e2e.py
View file

@ -7,6 +7,7 @@ wall-clock. Stock is only run at MODEST settings (a full 121-frame stock gen is
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_e2e.py
"""
from __future__ import annotations
import argparse
@ -23,7 +24,7 @@ import torch
_REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend"))
OUT = _REPO_ROOT / "outputs" / "video_speedmem"
OUT.mkdir(parents=True, exist_ok=True)
OUT.mkdir(parents = True, exist_ok = True)
def _import_diffusers():
@ -36,12 +37,18 @@ def _import_diffusers():
def _gen(pipe, seed, frames, steps, w, h):
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
torch.cuda.synchronize()
t0 = time.perf_counter()
out = pipe(prompt="a cat playing piano on a stage, cinematic",
num_frames=frames, width=w, height=h,
num_inference_steps=steps, generator=g, output_type="np")
out = pipe(
prompt = "a cat playing piano on a stage, cinematic",
num_frames = frames,
width = w,
height = h,
num_inference_steps = steps,
generator = g,
output_type = "np",
)
torch.cuda.synchronize()
wall = (time.perf_counter() - t0) * 1e3
frames_np = out.frames[0] # [F,H,W,C] in [0,1]
@ -54,7 +61,12 @@ def _luma(frames):
return float((0.299 * r + 0.587 * gg + 0.114 * b).mean())
def _lpips_mean(loss_fn, a, b, stride=4):
def _lpips_mean(
loss_fn,
a,
b,
stride = 4,
):
vals = []
for i in range(0, len(a), stride):
ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
@ -66,14 +78,14 @@ def _lpips_mean(loss_fn, a, b, stride=4):
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--cmp-frames", type=int, default=25)
ap.add_argument("--cmp-steps", type=int, default=50)
ap.add_argument("--cmp-w", type=int, default=832)
ap.add_argument("--cmp-h", type=int, default=480)
ap.add_argument("--full-frames", type=int, default=121)
ap.add_argument("--full-steps", type=int, default=50)
ap.add_argument("--seed", type=int, default=1234)
ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--cmp-frames", type = int, default = 25)
ap.add_argument("--cmp-steps", type = int, default = 50)
ap.add_argument("--cmp-w", type = int, default = 832)
ap.add_argument("--cmp-h", type = int, default = 480)
ap.add_argument("--full-frames", type = int, default = 121)
ap.add_argument("--full-steps", type = int, default = 50)
ap.add_argument("--seed", type = int, default = 1234)
args = ap.parse_args()
diffusers = _import_diffusers()
@ -81,43 +93,61 @@ def main():
from core.inference.video_families import detect_video_family
import lpips
print(f"loading {args.repo} ...", flush=True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda")
print(f"loading {args.repo} ...", flush = True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to(
"cuda"
)
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
loss_fn = lpips.LPIPS(net="alex").cuda()
loss_fn = lpips.LPIPS(net = "alex").cuda()
fr, st, w, hh = args.cmp_frames, args.cmp_steps, args.cmp_w, args.cmp_h
# ---- STOCK x2 (nondeterminism floor) ----
print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush=True)
print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush = True)
stock1, w_stock = _gen(pipe, args.seed, fr, st, w, hh)
print(f"[stock#1] wall={w_stock:.0f} ms luma={_luma(stock1):.4f} frames={stock1.shape}", flush=True)
print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush=True)
print(
f"[stock#1] wall={w_stock:.0f} ms luma={_luma(stock1):.4f} frames={stock1.shape}",
flush = True,
)
print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush = True)
stock2, _ = _gen(pipe, args.seed, fr, st, w, hh)
# ---- TRIM (same seed) ----
engaged = install_hunyuan_attention_trim(pipe, fam, logger=None)
print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush=True)
engaged = install_hunyuan_attention_trim(pipe, fam, logger = None)
print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush = True)
trim, w_trim = _gen(pipe, args.seed, fr, st, w, hh)
print(f"[trim ] wall={w_trim:.0f} ms luma={_luma(trim):.4f} ({w_stock/w_trim:.2f}x vs stock)", flush=True)
print(
f"[trim ] wall={w_trim:.0f} ms luma={_luma(trim):.4f} ({w_stock/w_trim:.2f}x vs stock)",
flush = True,
)
floor = _lpips_mean(loss_fn, stock1, stock2)
lp = _lpips_mean(loss_fn, stock1, trim)
print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f} <- nondeterminism floor", flush=True)
print(f"LPIPS(stock#1, trim ) = {lp:.5f} <- trim vs stock", flush=True)
print(f" => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor", flush=True)
print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f} <- nondeterminism floor", flush = True)
print(f"LPIPS(stock#1, trim ) = {lp:.5f} <- trim vs stock", flush = True)
print(
f" => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor",
flush = True,
)
# ---- FULL-RES TRIM (wall-clock + black-frame guard; stock at full-res is ~19 min, skipped) ----
print(f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...", flush=True)
print(
f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...",
flush = True,
)
full, w_full = _gen(pipe, args.seed, args.full_frames, args.full_steps, args.cmp_w, args.cmp_h)
print(f"[trim-full] wall={w_full/1000:.1f} s luma={_luma(full):.4f} frames={full.shape}", flush=True)
print(
f"[trim-full] wall={w_full/1000:.1f} s luma={_luma(full):.4f} frames={full.shape}",
flush = True,
)
try:
from PIL import Image
Image.fromarray((trim[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_cmp.png")
Image.fromarray((full[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_full.png")
print(f"\nsaved sample frames to {OUT}", flush=True)
print(f"\nsaved sample frames to {OUT}", flush = True)
except Exception as exc: # noqa: BLE001
print(f"(png save skipped: {exc})", flush=True)
print(f"(png save skipped: {exc})", flush = True)
if __name__ == "__main__":

74
scripts/hunyuan_trim_fp32ref.py
View file

@ -8,6 +8,7 @@ quality loss. If trim is clearly farther from fp32 than stock, it is a real regr
Run: CUDA_VISIBLE_DEVICES=1 python scripts/hunyuan_trim_fp32ref.py
"""
from __future__ import annotations
import argparse
@ -35,14 +36,25 @@ def _import_diffusers():
def _gen(pipe, seed, fr, st, w, h):
g = torch.Generator(device="cuda").manual_seed(seed)
out = pipe(prompt="a cat playing piano on a stage, cinematic",
num_frames=fr, width=w, height=h, num_inference_steps=st,
generator=g, output_type="np")
g = torch.Generator(device = "cuda").manual_seed(seed)
out = pipe(
prompt = "a cat playing piano on a stage, cinematic",
num_frames = fr,
width = w,
height = h,
num_inference_steps = st,
generator = g,
output_type = "np",
)
return np.asarray(out.frames[0])
def _lpips_mean(loss_fn, a, b, stride=3):
def _lpips_mean(
loss_fn,
a,
b,
stride = 3,
):
vals = []
for i in range(0, len(a), stride):
ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float().cuda() * 2 - 1
@ -54,12 +66,12 @@ def _lpips_mean(loss_fn, a, b, stride=3):
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type=int, default=25)
ap.add_argument("--steps", type=int, default=30)
ap.add_argument("--w", type=int, default=832)
ap.add_argument("--h", type=int, default=480)
ap.add_argument("--seed", type=int, default=1234)
ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type = int, default = 25)
ap.add_argument("--steps", type = int, default = 30)
ap.add_argument("--w", type = int, default = 832)
ap.add_argument("--h", type = int, default = 480)
ap.add_argument("--seed", type = int, default = 1234)
args = ap.parse_args()
diffusers = _import_diffusers()
@ -68,37 +80,45 @@ def main():
import lpips
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
loss_fn = lpips.LPIPS(net="alex").cuda()
loss_fn = lpips.LPIPS(net = "alex").cuda()
fr, st, w, h, seed = args.frames, args.steps, args.w, args.h, args.seed
# ---- bf16 stock + bf16 trim on one pipe ----
print(f"loading bf16 {args.repo} ...", flush=True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda")
print(f"[bf16 stock] gen {fr}f/{st}steps ...", flush=True)
print(f"loading bf16 {args.repo} ...", flush = True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to(
"cuda"
)
print(f"[bf16 stock] gen {fr}f/{st}steps ...", flush = True)
stock = _gen(pipe, seed, fr, st, w, h)
install_hunyuan_attention_trim(pipe, fam, logger=None)
print("[bf16 trim ] gen ...", flush=True)
install_hunyuan_attention_trim(pipe, fam, logger = None)
print("[bf16 trim ] gen ...", flush = True)
trim = _gen(pipe, seed, fr, st, w, h)
del pipe
gc.collect(); torch.cuda.empty_cache()
gc.collect()
torch.cuda.empty_cache()
# ---- fp32 reference (stock masked attention, upcast) ----
print("loading fp32 reference ...", flush=True)
pipe32 = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.float32).to("cuda")
print(f"[fp32 gold ] gen {fr}f/{st}steps ...", flush=True)
print("loading fp32 reference ...", flush = True)
pipe32 = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.float32).to(
"cuda"
)
print(f"[fp32 gold ] gen {fr}f/{st}steps ...", flush = True)
gold = _gen(pipe32, seed, fr, st, w, h)
d_stock = _lpips_mean(loss_fn, gold, stock)
d_trim = _lpips_mean(loss_fn, gold, trim)
d_st = _lpips_mean(loss_fn, stock, trim)
print("\n===== ACCURACY vs fp32 reference =====", flush=True)
print(f" LPIPS(fp32, bf16-stock) = {d_stock:.5f}", flush=True)
print(f" LPIPS(fp32, bf16-trim ) = {d_trim:.5f}", flush=True)
print(f" LPIPS(bf16-stock, trim) = {d_st:.5f}", flush=True)
print("\n===== ACCURACY vs fp32 reference =====", flush = True)
print(f" LPIPS(fp32, bf16-stock) = {d_stock:.5f}", flush = True)
print(f" LPIPS(fp32, bf16-trim ) = {d_trim:.5f}", flush = True)
print(f" LPIPS(bf16-stock, trim) = {d_st:.5f}", flush = True)
if d_stock is not None and d_trim is not None:
verdict = "NOT less accurate (trim ~= stock vs fp32)" if d_trim <= d_stock * 1.25 + 0.01 \
verdict = (
"NOT less accurate (trim ~= stock vs fp32)"
if d_trim <= d_stock * 1.25 + 0.01
else "LESS accurate (trim farther from fp32 than stock)"
print(f"\n VERDICT: {verdict}", flush=True)
)
print(f"\n VERDICT: {verdict}", flush = True)
if __name__ == "__main__":

91
scripts/hunyuan_trim_validate.py
View file

@ -7,6 +7,7 @@ real production shape (default 121 frames / 480p).
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_validate.py
"""
from __future__ import annotations
import argparse
@ -45,14 +46,18 @@ def _capture_hook(module, args, kwargs):
raise _Stop
def _forward(transformer, no_grad=True):
def _forward(transformer, no_grad = True):
ctx = torch.no_grad() if no_grad else torch.enable_grad()
with ctx:
out = transformer(*CAP["args"], **CAP["kwargs"])
return out[0] if isinstance(out, tuple) else out.sample
def _median_ms(fn, iters=8, warmup=2):
def _median_ms(
fn,
iters = 8,
warmup = 2,
):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
@ -70,7 +75,7 @@ def _median_ms(fn, iters=8, warmup=2):
def _compare(a, b):
a = a.float().flatten()
b = b.float().flatten()
cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
cos = torch.nn.functional.cosine_similarity(a, b, dim = 0).item()
max_abs = (a - b).abs().max().item()
denom = a.abs().max().item() or 1.0
return cos, max_abs, max_abs / denom
@ -78,11 +83,11 @@ def _compare(a, b):
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type=int, default=121)
ap.add_argument("--width", type=int, default=832)
ap.add_argument("--height", type=int, default=480)
ap.add_argument("--compile", action="store_true", help="also test regional compile of blocks")
ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
ap.add_argument("--frames", type = int, default = 121)
ap.add_argument("--width", type = int, default = 832)
ap.add_argument("--height", type = int, default = 480)
ap.add_argument("--compile", action = "store_true", help = "also test regional compile of blocks")
args = ap.parse_args()
diffusers = _import_diffusers()
@ -90,16 +95,26 @@ def main():
from core.inference.video_families import detect_video_family
dev = "cuda:0"
print(f"loading {args.repo} ...", flush=True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to(dev)
print(f"loading {args.repo} ...", flush = True)
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to(
dev
)
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
print(f"family: {getattr(fam, 'name', None)} / transformer_class={getattr(fam, 'transformer_class', None)}", flush=True)
print(
f"family: {getattr(fam, 'name', None)} / transformer_class={getattr(fam, 'transformer_class', None)}",
flush = True,
)
h = pipe.transformer.register_forward_pre_hook(_capture_hook, with_kwargs=True)
print("capturing one real forward input ...", flush=True)
h = pipe.transformer.register_forward_pre_hook(_capture_hook, with_kwargs = True)
print("capturing one real forward input ...", flush = True)
try:
pipe(prompt="a cat playing piano", num_frames=args.frames,
width=args.width, height=args.height, num_inference_steps=1)
pipe(
prompt = "a cat playing piano",
num_frames = args.frames,
width = args.width,
height = args.height,
num_inference_steps = 1,
)
except _Stop:
pass
h.remove()
@ -108,46 +123,58 @@ def main():
ehs = k.get("encoder_hidden_states")
m = k.get("encoder_attention_mask")
ie = k.get("image_embeds")
print(f"\ncaptured: encoder_hidden_states={tuple(ehs.shape)} mask_valid={m.bool().sum(1).tolist()}"
f" image_embeds={tuple(ie.shape) if ie is not None else None}"
f" image_all_zero={bool(torch.all(ie==0).item()) if ie is not None else None}", flush=True)
print(
f"\ncaptured: encoder_hidden_states={tuple(ehs.shape)} mask_valid={m.bool().sum(1).tolist()}"
f" image_embeds={tuple(ie.shape) if ie is not None else None}"
f" image_all_zero={bool(torch.all(ie==0).item()) if ie is not None else None}",
flush = True,
)
# ---- STOCK forward (reference) + timing ----
transformer = pipe.transformer
out_stock = _forward(transformer).detach().clone()
t_stock = _median_ms(lambda: _forward(transformer))
print(f"\nSTOCK forward: {t_stock:8.2f} ms out={tuple(out_stock.shape)}", flush=True)
print(f"\nSTOCK forward: {t_stock:8.2f} ms out={tuple(out_stock.shape)}", flush = True)
# ---- install trim, re-run same inputs ----
engaged = install_hunyuan_attention_trim(pipe, fam, logger=None)
print(f"install_hunyuan_attention_trim engaged = {engaged}", flush=True)
engaged = install_hunyuan_attention_trim(pipe, fam, logger = None)
print(f"install_hunyuan_attention_trim engaged = {engaged}", flush = True)
out_trim = _forward(transformer).detach().clone()
t_trim = _median_ms(lambda: _forward(transformer))
cos, max_abs, rel = _compare(out_stock, out_trim)
print(f"TRIM forward: {t_trim:8.2f} ms ({t_stock/t_trim:.2f}x faster)", flush=True)
print(f"\nACCURACY stock-vs-trim: cosine={cos:.8f} max_abs={max_abs:.4e} rel_max={rel:.4e}", flush=True)
print(f"TRIM forward: {t_trim:8.2f} ms ({t_stock/t_trim:.2f}x faster)", flush = True)
print(
f"\nACCURACY stock-vs-trim: cosine={cos:.8f} max_abs={max_abs:.4e} rel_max={rel:.4e}",
flush = True,
)
finite = bool(torch.isfinite(out_trim).all().item())
print(f"trim output finite: {finite}", flush=True)
print(f"trim output finite: {finite}", flush = True)
if args.compile:
print("\ncompiling blocks (compile_repeated_blocks, mode=default, dynamic=True) ...", flush=True)
print(
"\ncompiling blocks (compile_repeated_blocks, mode=default, dynamic=True) ...",
flush = True,
)
try:
for _a in ("recompile_limit", "cache_size_limit"):
if hasattr(torch._dynamo.config, _a):
setattr(torch._dynamo.config, _a, 64)
transformer.compile_repeated_blocks(fullgraph=False, dynamic=True)
transformer.compile_repeated_blocks(fullgraph = False, dynamic = True)
out_c = _forward(transformer).detach().clone() # triggers compile
t_c = _median_ms(lambda: _forward(transformer), iters=5, warmup=1)
t_c = _median_ms(lambda: _forward(transformer), iters = 5, warmup = 1)
cos_c, ma_c, rel_c = _compare(out_stock, out_c)
cnt = torch._dynamo.utils.counters
print(f"TRIM+COMPILE forward: {t_c:8.2f} ms ({t_stock/t_c:.2f}x vs stock)", flush=True)
print(f" accuracy vs stock: cosine={cos_c:.8f} max_abs={ma_c:.4e}", flush=True)
print(f" dynamo recompiles={sum(cnt['recompiles'].values()) if 'recompiles' in cnt else '?'}"
f" graph_breaks={sum(cnt['graph_break'].values()) if 'graph_break' in cnt else 0}", flush=True)
print(f"TRIM+COMPILE forward: {t_c:8.2f} ms ({t_stock/t_c:.2f}x vs stock)", flush = True)
print(f" accuracy vs stock: cosine={cos_c:.8f} max_abs={ma_c:.4e}", flush = True)
print(
f" dynamo recompiles={sum(cnt['recompiles'].values()) if 'recompiles' in cnt else '?'}"
f" graph_breaks={sum(cnt['graph_break'].values()) if 'graph_break' in cnt else 0}",
flush = True,
)
except Exception as exc: # noqa: BLE001
import traceback
print(f"COMPILE FAILED: {type(exc).__name__}: {exc}", flush=True)
print(f"COMPILE FAILED: {type(exc).__name__}: {exc}", flush = True)
traceback.print_exc()

224
scripts/quant_accuracy_sweep.py
View file

@ -71,14 +71,17 @@ _VAE_FAMILIES: dict[str, dict[str, Any]] = {
def _check_deps() -> None:
import importlib.util as ilu
missing = [m for m in ("torch", "torchao", "diffusers", "lpips", "numpy", "PIL") if not ilu.find_spec(m)]
missing = [
m
for m in ("torch", "torchao", "diffusers", "lpips", "numpy", "PIL")
if not ilu.find_spec(m)
]
if missing:
print(
"missing deps: " + ", ".join(missing) + "\n"
" uv pip install torch torchao diffusers lpips numpy pillow",
file=sys.stderr,
flush=True,
file = sys.stderr,
flush = True,
)
raise SystemExit(2)
@ -102,7 +105,7 @@ def _load_vae(repo: str, subfolder: str, device: str):
import torch
diffusers = _import_diffusers()
vae = diffusers.AutoModel.from_pretrained(repo, subfolder=subfolder, torch_dtype=torch.bfloat16)
vae = diffusers.AutoModel.from_pretrained(repo, subfolder = subfolder, torch_dtype = torch.bfloat16)
vae = vae.to(device).eval()
return vae
@ -148,11 +151,13 @@ def _ref_images(args: argparse.Namespace, size: int) -> list:
from PIL import Image
ref_dir = Path(args.ref_image_dir)
files = sorted(glob.glob(str(ref_dir / "*.jpg")) + glob.glob(str(ref_dir / "*.png")))[: args.num_samples]
files = sorted(glob.glob(str(ref_dir / "*.jpg")) + glob.glob(str(ref_dir / "*.png")))[
: args.num_samples
]
imgs = []
for f in files:
im = Image.open(f).convert("RGB").resize((size, size), Image.BICUBIC)
imgs.append(np.asarray(im, dtype=np.uint8))
imgs.append(np.asarray(im, dtype = np.uint8))
return imgs
@ -189,11 +194,11 @@ def _make_latents(vae: Any, args: argparse.Namespace, device: str):
x = torch.from_numpy(arr).float().permute(2, 0, 1).unsqueeze(0).div(127.5).sub(1.0)
if is_3d:
x = x.unsqueeze(2).repeat(1, 1, args.enc_frames, 1, 1) # static clip [1,3,T,H,W]
x = x.to(device=device, dtype=torch.bfloat16)
x = x.to(device = device, dtype = torch.bfloat16)
lat.append(_encode_latent(vae, x))
if lat:
return lat, is_3d
print(" (no ref images found; falling back to random latents)", flush=True)
print(" (no ref images found; falling back to random latents)", flush = True)
for seed in range(args.num_samples):
g = torch.Generator().manual_seed(1000 + seed)
shape = (
@ -201,8 +206,8 @@ def _make_latents(vae: Any, args: argparse.Namespace, device: str):
if is_3d
else (1, channels, args.latent_hw, args.latent_hw)
)
z = torch.randn(shape, generator=g, dtype=torch.float32)
lat.append(z.to(device=device, dtype=torch.bfloat16))
z = torch.randn(shape, generator = g, dtype = torch.float32)
lat.append(z.to(device = device, dtype = torch.bfloat16))
return lat, is_3d
@ -215,7 +220,7 @@ def _decode(vae: Any, z: Any):
try:
out = vae.decode(z)
except TypeError:
out = vae.decode(z, return_dict=True)
out = vae.decode(z, return_dict = True)
sample = out.sample if hasattr(out, "sample") else out[0]
sample = sample.float().clamp(-1, 1)
# [B,C,H,W] (image) or [B,C,T,H,W] (video). Emit one frame per temporal slot.
@ -228,9 +233,11 @@ def _decode(vae: Any, z: Any):
frames.append(sample[0])
imgs = []
for f in frames:
arr = ((f.permute(1, 2, 0).cpu().numpy() + 1.0) * 127.5).round().clip(0, 255).astype(np.uint8)
arr = (
((f.permute(1, 2, 0).cpu().numpy() + 1.0) * 127.5).round().clip(0, 255).astype(np.uint8)
)
if arr.shape[2] == 1:
arr = np.repeat(arr, 3, axis=2)
arr = np.repeat(arr, 3, axis = 2)
imgs.append(arr)
return imgs # list of HxWx3 uint8
@ -247,13 +254,21 @@ class _Lpips:
self.torch = torch
self.device = device
self.fn = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
self.fn = lpips.LPIPS(net = "alex", verbose = False).to(device).eval()
def __call__(self, a: Any, b: Any) -> float:
t = self.torch
def to_t(x):
return t.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0).div(127.5).sub(1.0).to(self.device)
return (
t.from_numpy(x)
.float()
.permute(2, 0, 1)
.unsqueeze(0)
.div(127.5)
.sub(1.0)
.to(self.device)
)
with t.no_grad():
return float(self.fn(to_t(a), to_t(b)).item())
@ -299,12 +314,18 @@ def _apply_fp8_dynamic_no1x1(vae_q: Any) -> None:
if w is None or w.dim() < 2 or w.shape[0] % 16 or w.shape[1] % 16:
return False
ks = getattr(module, "kernel_size", None)
if isinstance(ks, tuple) and all(k == 1 for k in ks): # pointwise conv -> torchao kernel fails
if isinstance(ks, tuple) and all(
k == 1 for k in ks
): # pointwise conv -> torchao kernel fails
return False
name = fqn.lower() if fqn else ""
return not any(tok in name for tok in _VAE_KEEP_DENSE_TOKENS)
quantize_(vae_q, Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()), filter_fn=filter_fn)
quantize_(
vae_q,
Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()),
filter_fn = filter_fn,
)
def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[dict]:
@ -318,7 +339,6 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
class _Target:
def __init__(self):
import torch
self.device = "cuda"
self.dtype = torch.bfloat16
@ -326,30 +346,33 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
rows: list[dict] = []
for family in args.family:
repo = _VAE_FAMILIES[family]["repo"]
print(f"\n=== VAE {family} ({repo}) ===", flush=True)
print(f"\n=== VAE {family} ({repo}) ===", flush = True)
t0 = time.time()
try:
vae = _load_vae(repo, "vae", "cuda")
except Exception as exc: # noqa: BLE001
print(f" load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush=True)
print(f" load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush = True)
rows.append({"family": family, "scheme": "-", "error": f"load: {exc}"})
continue
ch, is_3d = _latent_spec(vae)
cls = type(vae).__name__
print(f" {cls} latent_ch={ch} {'3D' if is_3d else '2D'} loaded {time.time()-t0:.0f}s", flush=True)
print(
f" {cls} latent_ch={ch} {'3D' if is_3d else '2D'} loaded {time.time()-t0:.0f}s",
flush = True,
)
latents, _ = _make_latents(vae, args, "cuda")
ref_by_sample = [_decode(vae, z) for z in latents]
fam_dir = out_dir / family
fam_dir.mkdir(parents=True, exist_ok=True)
fam_dir.mkdir(parents = True, exist_ok = True)
Image.fromarray(ref_by_sample[0][0]).save(fam_dir / "dense_s0.png")
for scheme in args.scheme:
try:
vae_q = copy.deepcopy(vae)
except Exception as exc: # noqa: BLE001
print(f" [{scheme}] deepcopy FAILED: {exc}", flush=True)
print(f" [{scheme}] deepcopy FAILED: {exc}", flush = True)
continue
pipe = type("P", (), {"vae": vae_q})()
# "fp8_dynamic_no1x1" is a diagnostic that bypasses quantize_vae to apply the
@ -359,17 +382,19 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
_apply_fp8_dynamic_no1x1(vae_q)
engaged = scheme
except Exception as exc: # noqa: BLE001
print(f" [{scheme}] apply FAILED: {str(exc)[:120]}", flush=True)
print(f" [{scheme}] apply FAILED: {str(exc)[:120]}", flush = True)
del vae_q
_empty_cache()
continue
else:
engaged = vq.quantize_vae(
pipe, target, mode=scheme, family=family, offload_active=False, force_fp32=False
pipe, target, mode = scheme, family = family, offload_active = False, force_fp32 = False
)
if engaged != scheme:
print(f" [{scheme}] NOT engaged (returned {engaged}); skipping", flush=True)
rows.append({"family": family, "vae_class": cls, "scheme": scheme, "verdict": "NOT_ENGAGED"})
print(f" [{scheme}] NOT engaged (returned {engaged}); skipping", flush = True)
rows.append(
{"family": family, "vae_class": cls, "scheme": scheme, "verdict": "NOT_ENGAGED"}
)
del vae_q
_empty_cache()
continue
@ -377,9 +402,15 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
q_by_sample = [_decode(vae_q, z) for z in latents]
except Exception as exc: # noqa: BLE001 — the shipped caster produced a VAE that crashes at decode
emsg = f"{type(exc).__name__}: {str(exc)[:100]}"
print(f" [{scheme}] DECODE CRASH: {emsg}", flush=True)
print(f" [{scheme}] DECODE CRASH: {emsg}", flush = True)
rows.append(
{"family": family, "vae_class": cls, "scheme": scheme, "verdict": "CRASH", "error": emsg}
{
"family": family,
"vae_class": cls,
"scheme": scheme,
"verdict": "CRASH",
"error": emsg,
}
)
del vae_q
_empty_cache()
@ -402,7 +433,7 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
rows.append(row)
print(
f" [{scheme}] LPIPS={m['lpips']} PSNR={m['psnr']} SSIM={m['ssim']} -> {verdict}",
flush=True,
flush = True,
)
del vae_q
_empty_cache()
@ -414,7 +445,6 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
def _empty_cache() -> None:
try:
import torch
torch.cuda.empty_cache()
except Exception:
pass
@ -448,14 +478,16 @@ def _apply_auto(pipe: Any, family: str, components: list[str]) -> dict[str, Opti
tgt = _Target()
engaged: dict[str, Optional[str]] = {}
if "transformer" in components:
engaged["transformer"] = tq.quantize_transformer(pipe, tgt, mode="auto", family=family)
engaged["transformer"] = tq.quantize_transformer(pipe, tgt, mode = "auto", family = family)
if "text_encoder" in components:
try:
engaged["text_encoder"] = dp.quantize_text_encoders(pipe, tgt, mode="auto", family=family)
engaged["text_encoder"] = dp.quantize_text_encoders(
pipe, tgt, mode = "auto", family = family
)
except Exception as exc: # noqa: BLE001
engaged["text_encoder"] = f"err:{type(exc).__name__}"
if "vae" in components:
engaged["vae"] = vq.quantize_vae(pipe, tgt, mode="auto", family=family)
engaged["vae"] = vq.quantize_vae(pipe, tgt, mode = "auto", family = family)
return engaged
@ -466,7 +498,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
rows: list[dict] = []
for family in args.family:
model = args.e2e_model or _VAE_FAMILIES.get(family, {}).get("repo")
print(f"\n=== E2E {family} ({model}) ===", flush=True)
print(f"\n=== E2E {family} ({model}) ===", flush = True)
prompts = args.prompts or _E2E_PROMPTS
seeds = args.seeds
@ -474,13 +506,13 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
imgs = []
for pi, prompt in enumerate(prompts):
for seed in seeds:
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
kw = dict(
prompt=prompt,
num_inference_steps=args.steps,
generator=g,
height=args.height,
width=args.width,
prompt = prompt,
num_inference_steps = args.steps,
generator = g,
height = args.height,
width = args.width,
)
if args.guidance is not None:
kw["guidance_scale"] = args.guidance
@ -490,10 +522,10 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
try:
pipe = diffusers.AutoPipelineForText2Image.from_pretrained(
model, torch_dtype=torch.bfloat16
model, torch_dtype = torch.bfloat16
).to("cuda")
except Exception as exc: # noqa: BLE001
print(f" pipe load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush=True)
print(f" pipe load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush = True)
rows.append({"family": family, "error": f"load: {exc}"})
continue
ref = _gen(pipe)
@ -501,10 +533,10 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
_empty_cache()
pipe2 = diffusers.AutoPipelineForText2Image.from_pretrained(
model, torch_dtype=torch.bfloat16
model, torch_dtype = torch.bfloat16
).to("cuda")
engaged = _apply_auto(pipe2, family, args.e2e_components)
print(f" engaged: {engaged}", flush=True)
print(f" engaged: {engaged}", flush = True)
q = _gen(pipe2)
del pipe2
_empty_cache()
@ -513,7 +545,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
ls = []
fam_dir = out_dir / f"e2e_{family}"
fam_dir.mkdir(parents=True, exist_ok=True)
fam_dir.mkdir(parents = True, exist_ok = True)
for (pi, seed, a), (_, _, b) in zip(ref, q):
aa, bb = np.asarray(a.convert("RGB")), np.asarray(b.convert("RGB"))
ls.append(lp(aa, bb))
@ -524,7 +556,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
rows.append(
{"family": family, "engaged": engaged, "mean_lpips": mean_l, "verdict": verdict}
)
print(f" mean LPIPS={mean_l} -> {verdict}", flush=True)
print(f" mean LPIPS={mean_l} -> {verdict}", flush = True)
return rows
@ -532,28 +564,33 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di
def _write(out_dir: Path, mode: str, rows: list[dict]) -> None:
(out_dir / f"{mode}_results.json").write_text(json.dumps(rows, indent=2))
print(f"\nwrote {out_dir / f'{mode}_results.json'}", flush=True)
print(f"\n=== {mode.upper()} RESULTS ===", flush=True)
(out_dir / f"{mode}_results.json").write_text(json.dumps(rows, indent = 2))
print(f"\nwrote {out_dir / f'{mode}_results.json'}", flush = True)
print(f"\n=== {mode.upper()} RESULTS ===", flush = True)
if mode == "vae":
print(f" bars: LPIPS <= {LPIPS_BAR}, SSIM >= {SSIM_BAR}", flush=True)
print(f" {'family':<20}{'scheme':<14}{'LPIPS':>9}{'PSNR':>9}{'SSIM':>9} verdict", flush=True)
print(f" bars: LPIPS <= {LPIPS_BAR}, SSIM >= {SSIM_BAR}", flush = True)
print(
f" {'family':<20}{'scheme':<14}{'LPIPS':>9}{'PSNR':>9}{'SSIM':>9} verdict", flush = True
)
for r in rows:
if "error" in r:
print(f" {r['family']:<20}{'(error)':<14} {r['error'][:60]}", flush=True)
print(f" {r['family']:<20}{'(error)':<14} {r['error'][:60]}", flush = True)
continue
print(
f" {r['family']:<20}{r['scheme']:<14}{_f(r.get('lpips')):>9}"
f"{_f(r.get('psnr')):>9}{_f(r.get('ssim')):>9} {r.get('verdict')}",
flush=True,
flush = True,
)
else:
print(f" bar: mean LPIPS <= {E2E_LPIPS_BAR}", flush=True)
print(f" bar: mean LPIPS <= {E2E_LPIPS_BAR}", flush = True)
for r in rows:
if "error" in r:
print(f" {r['family']}: (error) {r['error'][:80]}", flush=True)
print(f" {r['family']}: (error) {r['error'][:80]}", flush = True)
continue
print(f" {r['family']:<20} mean_lpips={r.get('mean_lpips')} {r.get('verdict')} {r.get('engaged')}", flush=True)
print(
f" {r['family']:<20} mean_lpips={r.get('mean_lpips')} {r.get('verdict')} {r.get('engaged')}",
flush = True,
)
def _f(v: Any) -> str:
@ -565,36 +602,49 @@ def _f(v: Any) -> str:
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="Decoded-image accuracy sweep for the auto VAE / end-to-end quantisation.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
description = "Decoded-image accuracy sweep for the auto VAE / end-to-end quantisation.",
formatter_class = argparse.ArgumentDefaultsHelpFormatter,
)
p.add_argument("--mode", choices=["vae", "e2e"], default="vae")
p.add_argument("--family", nargs="+", default=list(_VAE_FAMILIES.keys()))
p.add_argument("--scheme", nargs="+", default=["fp8_dynamic", "fp8_dynamic_no1x1", "fp8"])
p.add_argument("--num-samples", type=int, default=5, help="latents to average over")
p.add_argument("--latents", choices=["encode", "random"], default="encode",
help="encode natural photos (in-distribution) or seeded N(0,1) latents")
p.add_argument("--ref-image-dir", default="outputs/quant_accuracy/_refs",
help="natural photos to encode for the round-trip")
p.add_argument("--enc-hw", type=int, default=512, help="2D encode pixel H=W")
p.add_argument("--enc-hw-3d", type=int, default=256, help="3D encode pixel H=W")
p.add_argument("--enc-frames", type=int, default=9, help="3D encode pixel frame count")
p.add_argument("--latent-hw", type=int, default=64, help="2D random-latent H=W (x8 -> 512px)")
p.add_argument("--latent-hw-3d", type=int, default=32, help="3D random-latent H=W")
p.add_argument("--latent-t", type=int, default=3, help="3D random-latent temporal length")
p.add_argument("--lpips-device", default="cpu", help="device for the LPIPS net (keep off the measured GPU)")
p.add_argument("--out-dir", default="outputs/quant_accuracy")
p.add_argument("--mode", choices = ["vae", "e2e"], default = "vae")
p.add_argument("--family", nargs = "+", default = list(_VAE_FAMILIES.keys()))
p.add_argument("--scheme", nargs = "+", default = ["fp8_dynamic", "fp8_dynamic_no1x1", "fp8"])
p.add_argument("--num-samples", type = int, default = 5, help = "latents to average over")
p.add_argument(
"--latents",
choices = ["encode", "random"],
default = "encode",
help = "encode natural photos (in-distribution) or seeded N(0,1) latents",
)
p.add_argument(
"--ref-image-dir",
default = "outputs/quant_accuracy/_refs",
help = "natural photos to encode for the round-trip",
)
p.add_argument("--enc-hw", type = int, default = 512, help = "2D encode pixel H=W")
p.add_argument("--enc-hw-3d", type = int, default = 256, help = "3D encode pixel H=W")
p.add_argument("--enc-frames", type = int, default = 9, help = "3D encode pixel frame count")
p.add_argument("--latent-hw", type = int, default = 64, help = "2D random-latent H=W (x8 -> 512px)")
p.add_argument("--latent-hw-3d", type = int, default = 32, help = "3D random-latent H=W")
p.add_argument("--latent-t", type = int, default = 3, help = "3D random-latent temporal length")
p.add_argument(
"--lpips-device", default = "cpu", help = "device for the LPIPS net (keep off the measured GPU)"
)
p.add_argument("--out-dir", default = "outputs/quant_accuracy")
# e2e-only
p.add_argument("--e2e-model", default=None, help="full model repo for --mode e2e")
p.add_argument("--e2e-components", nargs="+", default=["transformer", "text_encoder", "vae"],
choices=["transformer", "text_encoder", "vae"],
help="which components to auto-quantise for the e2e (isolate the VAE with: --e2e-components vae)")
p.add_argument("--prompts", nargs="*", default=None)
p.add_argument("--seeds", nargs="*", type=int, default=[12345])
p.add_argument("--steps", type=int, default=8)
p.add_argument("--guidance", type=float, default=None)
p.add_argument("--height", type=int, default=1024)
p.add_argument("--width", type=int, default=1024)
p.add_argument("--e2e-model", default = None, help = "full model repo for --mode e2e")
p.add_argument(
"--e2e-components",
nargs = "+",
default = ["transformer", "text_encoder", "vae"],
choices = ["transformer", "text_encoder", "vae"],
help = "which components to auto-quantise for the e2e (isolate the VAE with: --e2e-components vae)",
)
p.add_argument("--prompts", nargs = "*", default = None)
p.add_argument("--seeds", nargs = "*", type = int, default = [12345])
p.add_argument("--steps", type = int, default = 8)
p.add_argument("--guidance", type = float, default = None)
p.add_argument("--height", type = int, default = 1024)
p.add_argument("--width", type = int, default = 1024)
return p
@ -602,7 +652,7 @@ def main(argv: Optional[list[str]] = None) -> int:
args = _build_parser().parse_args(argv)
_check_deps()
out_dir = Path(args.out_dir).resolve()
out_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents = True, exist_ok = True)
lp = _Lpips(args.lpips_device)
if args.mode == "vae":
rows = _sweep_vae(args, lp, out_dir)

324
scripts/quant_speedmem_bench.py
View file

@ -69,33 +69,28 @@ _TE_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3")
# ── cuda memory / timing helpers (lifted from diffusion_bench.py) ──────────────
def _sync() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
def _reset_peak() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
def _alloc_gb() -> float:
import torch
return torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _peak_gb() -> float:
import torch
return torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _empty() -> None:
import torch
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
@ -116,12 +111,11 @@ def _lpips_alex(ref_arr, arr):
fn = _LP.get("fn")
if fn is None:
fn = lpips.LPIPS(net="alex", verbose=False).eval()
fn = lpips.LPIPS(net = "alex", verbose = False).eval()
_LP["fn"] = fn
def _t(a):
import torch as _torch
return _torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
@ -137,7 +131,7 @@ def _timed_generate(pipe, *, steps, res, seed):
import torch
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
step_ts: list[float] = []
last = [0.0]
@ -152,8 +146,12 @@ def _timed_generate(pipe, *, steps, res, seed):
_sync()
t0 = _time.perf_counter()
img = pipe(
prompt=PROMPT, width=res, height=res, num_inference_steps=steps, generator=g,
callback_on_step_end=_cb,
prompt = PROMPT,
width = res,
height = res,
num_inference_steps = steps,
generator = g,
callback_on_step_end = _cb,
).images[0]
_sync()
return img, (_time.perf_counter() - t0), step_ts
@ -177,15 +175,14 @@ def _target():
import torch
return types.SimpleNamespace(device="cuda", dtype=torch.bfloat16)
return types.SimpleNamespace(device = "cuda", dtype = torch.bfloat16)
# ── model_index component class resolution ────────────────────────────────────
def _model_index(repo: str) -> dict:
from huggingface_hub import hf_hub_download
path = hf_hub_download(repo, "model_index.json")
with open(path, "r", encoding="utf-8") as fh:
with open(path, "r", encoding = "utf-8") as fh:
return json.load(fh)
@ -203,7 +200,7 @@ def _load_named(repo: str, name: str):
lib, cls_name = spec
module = importlib.import_module(lib)
klass = getattr(module, cls_name)
return klass.from_pretrained(repo, subfolder=name, torch_dtype=torch.bfloat16)
return klass.from_pretrained(repo, subfolder = name, torch_dtype = torch.bfloat16)
def _load_text_encoders(repo: str, device: str):
@ -244,7 +241,7 @@ def _load_tokenizers(repo: str):
("text_encoder_3", "tokenizer_3"),
):
try:
toks[te_attr] = AutoTokenizer.from_pretrained(repo, subfolder=tk_sub)
toks[te_attr] = AutoTokenizer.from_pretrained(repo, subfolder = tk_sub)
except Exception:
toks[te_attr] = None
return toks
@ -256,15 +253,15 @@ def _encoder_hidden(te, ids, mask):
with torch.inference_mode():
try:
out = te(input_ids=ids, attention_mask=mask, output_hidden_states=True)
out = te(input_ids = ids, attention_mask = mask, output_hidden_states = True)
except TypeError:
out = te(ids, output_hidden_states=True)
out = te(ids, output_hidden_states = True)
hs = getattr(out, "last_hidden_state", None)
if hs is None:
hidden = getattr(out, "hidden_states", None)
hs = hidden[-1] if hidden else (out[0] if isinstance(out, (tuple, list)) else out)
m = mask.unsqueeze(-1).to(hs.dtype)
v = (hs * m).sum(1) / m.sum(1).clamp(min=1)
v = (hs * m).sum(1) / m.sum(1).clamp(min = 1)
return v.float().flatten()
@ -280,7 +277,7 @@ def _te_hidden_refs(bag, toks, device):
continue
vecs = []
for p in _PROMPT_SUITE:
enc = tok(p, return_tensors="pt", padding="max_length", truncation=True, max_length=64)
enc = tok(p, return_tensors = "pt", padding = "max_length", truncation = True, max_length = 64)
ids = enc["input_ids"].to(device)
mask = enc.get("attention_mask")
mask = mask.to(device) if mask is not None else torch.ones_like(ids)
@ -289,7 +286,12 @@ def _te_hidden_refs(bag, toks, device):
return refs
def measure_te_accuracy(family: str, *, schemes=("fp8", "fp8_dynamic"), logger=None) -> list[dict]:
def measure_te_accuracy(
family: str,
*,
schemes = ("fp8", "fp8_dynamic"),
logger = None,
) -> list[dict]:
"""Hidden-state cosine / relL2 of each TE scheme vs the dense bf16 encoder (per encoder),
on real prompts. Bar (PR#150): cosine >= 0.99 and min_cosine >= 0.98."""
import torch
@ -310,14 +312,16 @@ def measure_te_accuracy(family: str, *, schemes=("fp8", "fp8_dynamic"), logger=N
rows: list[dict] = []
for scheme in schemes:
bag = _load_text_encoders(repo, device)
engaged = quantize_text_encoders(bag, _target(), mode=scheme, family=family, logger=logger)
engaged = quantize_text_encoders(bag, _target(), mode = scheme, family = family, logger = logger)
cur = _te_hidden_refs(bag, toks, device)
for attr, ref_vecs in refs.items():
q_vecs = cur.get(attr, [])
if not q_vecs:
continue
cosines = [F.cosine_similarity(r, q, dim=0).item() for r, q in zip(ref_vecs, q_vecs)]
rell2 = [((q - r).norm() / r.norm().clamp(min=1e-8)).item() for r, q in zip(ref_vecs, q_vecs)]
cosines = [F.cosine_similarity(r, q, dim = 0).item() for r, q in zip(ref_vecs, q_vecs)]
rell2 = [
((q - r).norm() / r.norm().clamp(min = 1e-8)).item() for r, q in zip(ref_vecs, q_vecs)
]
mean_cos = sum(cosines) / len(cosines)
min_cos = min(cosines)
rows.append(
@ -340,7 +344,6 @@ def _encode_once(bag) -> None:
"""One forward through every present text encoder on a fixed short token batch. Uses a length
within each encoder's max positions (CLIP caps at 77) so position embeddings never overflow."""
import torch
with torch.inference_mode():
for attr in _TE_ATTRS:
te = getattr(bag, attr, None)
@ -350,10 +353,10 @@ def _encode_once(bag) -> None:
vocab = int(getattr(cfg, "vocab_size", 30000) or 30000)
maxpos = int(getattr(cfg, "max_position_embeddings", 64) or 64)
length = max(8, min(64, maxpos))
ids = torch.randint(1, min(vocab, 30000), (1, length), device="cuda")
ids = torch.randint(1, min(vocab, 30000), (1, length), device = "cuda")
mask = torch.ones_like(ids)
try:
te(input_ids=ids, attention_mask=mask)
te(input_ids = ids, attention_mask = mask)
except TypeError:
te(ids)
@ -363,7 +366,7 @@ def _load_vae(repo: str, device: str):
import torch
diffusers = _import_diffusers()
vae = diffusers.AutoModel.from_pretrained(repo, subfolder="vae", torch_dtype=torch.bfloat16)
vae = diffusers.AutoModel.from_pretrained(repo, subfolder = "vae", torch_dtype = torch.bfloat16)
return vae.to(device).eval()
@ -395,12 +398,11 @@ def _latent_spec(vae) -> tuple[int, bool]:
def _decode_once(vae, z) -> None:
import torch
with torch.inference_mode():
try:
out = vae.decode(z)
except TypeError:
out = vae.decode(z, return_dict=True)
out = vae.decode(z, return_dict = True)
_ = out.sample if hasattr(out, "sample") else out[0]
@ -411,8 +413,8 @@ def _make_latent(vae, device: str):
channels, is_3d = _latent_spec(vae)
g = torch.Generator().manual_seed(1234)
shape = (1, channels, 3, 32, 32) if is_3d else (1, channels, 64, 64)
z = torch.randn(shape, generator=g, dtype=torch.float32)
return z.to(device=device, dtype=torch.bfloat16)
z = torch.randn(shape, generator = g, dtype = torch.float32)
return z.to(device = device, dtype = torch.bfloat16)
# ── measurement primitives ────────────────────────────────────────────────────
@ -431,7 +433,14 @@ def _time_median(fn, *, warmup: int, iters: int) -> float:
# ── mode: te ──────────────────────────────────────────────────────────────────
def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", logger=None) -> list[dict]:
def measure_te(
family: str,
*,
warmup: int,
iters: int,
scheme: str = "auto",
logger = None,
) -> list[dict]:
from core.inference.diffusion_precision import quantize_text_encoders
repo = _FAMILIES[family]["repo"]
@ -439,20 +448,22 @@ def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", lo
rows: list[dict] = []
# dense
_empty(); _reset_peak()
_empty()
_reset_peak()
bag = _load_text_encoders(repo, device)
_sync()
mem_dense = _alloc_gb()
_reset_peak()
lat_dense = _time_median(lambda: _encode_once(bag), warmup=warmup, iters=iters)
lat_dense = _time_median(lambda: _encode_once(bag), warmup = warmup, iters = iters)
peak_dense = _peak_gb()
# quant in place (scheme="auto" resolves the ladder; else force an explicit scheme to compare)
engaged = quantize_text_encoders(bag, _target(), mode=scheme, family=family, logger=logger)
_empty(); _sync()
engaged = quantize_text_encoders(bag, _target(), mode = scheme, family = family, logger = logger)
_empty()
_sync()
mem_quant = _alloc_gb()
_reset_peak()
lat_quant = _time_median(lambda: _encode_once(bag), warmup=warmup, iters=iters)
lat_quant = _time_median(lambda: _encode_once(bag), warmup = warmup, iters = iters)
peak_quant = _peak_gb()
del bag
@ -470,38 +481,55 @@ def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", lo
"peak_quant_gb": round(peak_quant, 3),
"lat_dense_ms": round(lat_dense, 2),
"lat_quant_ms": round(lat_quant, 2),
"lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) if lat_dense else None,
"lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1)
if lat_dense
else None,
}
)
return rows
# ── mode: vae ───────────────────────────────────────────────────────────────
def _measure_vae_scheme(family: str, repo: str, mode: str, *, warmup: int, iters: int, logger=None) -> dict:
def _measure_vae_scheme(
family: str,
repo: str,
mode: str,
*,
warmup: int,
iters: int,
logger = None,
) -> dict:
import types
from core.inference.diffusion_vae_quant import quantize_vae
device = "cuda"
_empty(); _reset_peak()
_empty()
_reset_peak()
vae = _load_vae(repo, device)
z = _make_latent(vae, device)
_sync()
mem_dense = _alloc_gb()
_reset_peak()
lat_dense = _time_median(lambda: _decode_once(vae, z), warmup=warmup, iters=iters)
lat_dense = _time_median(lambda: _decode_once(vae, z), warmup = warmup, iters = iters)
peak_dense = _peak_gb()
# quantize_vae reads pipe.vae, so hand it a bag exposing .vae (it mutates that module in place).
# Pass the family's force_fp32 (Wan) so the real dense-only behaviour is reflected.
force_fp32 = bool(_FAMILIES.get(family, {}).get("vae_force_fp32", False))
engaged = quantize_vae(
types.SimpleNamespace(vae=vae), _target(), mode=mode, family=family, force_fp32=force_fp32, logger=logger
types.SimpleNamespace(vae = vae),
_target(),
mode = mode,
family = family,
force_fp32 = force_fp32,
logger = logger,
)
_empty(); _sync()
_empty()
_sync()
mem_quant = _alloc_gb()
_reset_peak()
lat_quant = _time_median(lambda: _decode_once(vae, z), warmup=warmup, iters=iters)
lat_quant = _time_median(lambda: _decode_once(vae, z), warmup = warmup, iters = iters)
peak_quant = _peak_gb()
del vae, z
@ -519,40 +547,65 @@ def _measure_vae_scheme(family: str, repo: str, mode: str, *, warmup: int, iters
"peak_quant_gb": round(peak_quant, 3),
"lat_dense_ms": round(lat_dense, 2),
"lat_quant_ms": round(lat_quant, 2),
"lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) if lat_dense else None,
"lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1)
if lat_dense
else None,
}
def measure_vae(family: str, *, warmup: int, iters: int, logger=None) -> list[dict]:
def measure_vae(
family: str,
*,
warmup: int,
iters: int,
logger = None,
) -> list[dict]:
repo = _FAMILIES[family]["repo"]
rows = [_measure_vae_scheme(family, repo, "auto", warmup=warmup, iters=iters, logger=logger)]
rows = [_measure_vae_scheme(family, repo, "auto", warmup = warmup, iters = iters, logger = logger)]
if family == "flux.2":
# the one image family where the explicit fp8_dynamic conv opt-in is measured in-bar.
rows.append(_measure_vae_scheme(family, repo, "fp8_dynamic", warmup=warmup, iters=iters, logger=logger))
rows.append(
_measure_vae_scheme(
family, repo, "fp8_dynamic", warmup = warmup, iters = iters, logger = logger
)
)
return rows
# ── mode: e2e (qwen-image) ────────────────────────────────────────────────────
def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters: int, family: str, logger=None):
def _e2e_run(
repo: str,
*,
quant: bool,
steps: int,
res: int,
seed: int,
iters: int,
family: str,
logger = None,
):
import torch
from core.inference.diffusion_precision import quantize_text_encoders
from core.inference.diffusion_vae_quant import quantize_vae
diffusers = _import_diffusers()
_empty(); _reset_peak()
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
_empty()
_reset_peak()
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype = torch.bfloat16)
pipe = pipe.to("cuda")
load_peak = _peak_gb()
te_scheme = vae_scheme = None
if quant:
te_scheme = quantize_text_encoders(pipe, _target(), mode="auto", family=family, logger=logger)
vae_scheme = quantize_vae(pipe, _target(), mode="auto", family=family, logger=logger)
te_scheme = quantize_text_encoders(
pipe, _target(), mode = "auto", family = family, logger = logger
)
vae_scheme = quantize_vae(pipe, _target(), mode = "auto", family = family, logger = logger)
_empty()
weights_gb = _alloc_gb()
def _gen():
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
step_ts: list[float] = []
last = [0.0]
@ -567,12 +620,12 @@ def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters:
_sync()
t0 = time.perf_counter()
img = pipe(
prompt=PROMPT,
width=res,
height=res,
num_inference_steps=steps,
generator=g,
callback_on_step_end=_cb,
prompt = PROMPT,
width = res,
height = res,
num_inference_steps = steps,
generator = g,
callback_on_step_end = _cb,
).images[0]
_sync()
return img, (time.perf_counter() - t0), step_ts
@ -600,7 +653,15 @@ def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters:
def measure_e2e(
family: str, *, steps: int, res: int, seed: int, iters: int, out: Path, variant: str = "both", logger=None
family: str,
*,
steps: int,
res: int,
seed: int,
iters: int,
out: Path,
variant: str = "both",
logger = None,
) -> list[dict]:
repo = _FAMILIES[family]["repo"]
# "both" runs dense then auto in one process (fast, but the 2nd run is on a hotter GPU / a
@ -610,14 +671,21 @@ def measure_e2e(
rows = []
for quant in variants:
row, img = _e2e_run(
repo, quant=quant, steps=steps, res=res, seed=seed, iters=iters, family=family, logger=logger
repo,
quant = quant,
steps = steps,
res = res,
seed = seed,
iters = iters,
family = family,
logger = logger,
)
try:
img.save(out / f"e2e_{family}_{row['variant']}.png")
except Exception:
pass
rows.append(row)
print(f" e2e {row['variant']:5s}: {json.dumps(row)}", flush=True)
print(f" e2e {row['variant']:5s}: {json.dumps(row)}", flush = True)
return rows
@ -638,8 +706,16 @@ def _compile_blocks(transformer) -> bool:
def _dit_run(
repo: str, family: str, *, dit_quant: str, steps: int, res: int, seed: int, iters: int,
compile_blocks: bool = True, logger=None,
repo: str,
family: str,
*,
dit_quant: str,
steps: int,
res: int,
seed: int,
iters: int,
compile_blocks: bool = True,
logger = None,
):
"""Load the full pipeline dense, quantise ONLY the transformer (TE + VAE stay dense to isolate
the DiT), regional-compile it (the real feature path), then measure per-step + total latency and
@ -649,21 +725,24 @@ def _dit_run(
from core.inference.diffusion_transformer_quant import quantize_transformer
diffusers = _import_diffusers()
_empty(); _reset_peak()
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16).to("cuda")
_empty()
_reset_peak()
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype = torch.bfloat16).to("cuda")
load_peak = _peak_gb()
engaged = None
if dit_quant and dit_quant != "none":
engaged = quantize_transformer(pipe, _target(), mode=dit_quant, family=family, logger=logger)
engaged = quantize_transformer(
pipe, _target(), mode = dit_quant, family = family, logger = logger
)
_empty()
weights_gb = _alloc_gb()
compiled = _compile_blocks(getattr(pipe, "transformer", None)) if compile_blocks else False
img, _, _ = _timed_generate(pipe, steps=steps, res=res, seed=seed) # warmup (triggers compile)
img, _, _ = _timed_generate(pipe, steps = steps, res = res, seed = seed) # warmup (triggers compile)
_reset_peak()
dts, steps_ms, last_img = [], [], img
for _ in range(iters):
last_img, dt, st = _timed_generate(pipe, steps=steps, res=res, seed=seed)
last_img, dt, st = _timed_generate(pipe, steps = steps, res = res, seed = seed)
dts.append(dt)
steps_ms.append(_median(st) if st else 0.0)
gen_peak = _peak_gb()
@ -682,14 +761,24 @@ def _dit_run(
}, last_img
def measure_dit(family: str, *, schemes, steps: int, res: int, seed: int, iters: int, out: Path, logger=None):
def measure_dit(
family: str,
*,
schemes,
steps: int,
res: int,
seed: int,
iters: int,
out: Path,
logger = None,
):
"""Dense reference + each DiT scheme (auto/fp8/int8/mxfp8), reporting speedup, peak-memory drop,
and LPIPS(AlexNet) vs the dense render (the whole-image accuracy metric)."""
import numpy as np
repo = _FAMILIES[family]["repo"]
dense_row, dense_img = _dit_run(
repo, family, dit_quant="none", steps=steps, res=res, seed=seed, iters=iters, logger=logger
repo, family, dit_quant = "none", steps = steps, res = res, seed = seed, iters = iters, logger = logger
)
try:
dense_img.save(out / f"dit_{family}_dense.png")
@ -699,78 +788,105 @@ def measure_dit(family: str, *, schemes, steps: int, res: int, seed: int, iters:
dense_row["lpips_vs_dense"] = 0.0
dense_row["speedup_vs_dense"] = 1.0
rows = [dense_row]
print(f" dit dense: {json.dumps(dense_row)}", flush=True)
print(f" dit dense: {json.dumps(dense_row)}", flush = True)
base_lat = dense_row["gen_latency_s"] or 1.0
for scheme in schemes:
row, img = _dit_run(
repo, family, dit_quant=scheme, steps=steps, res=res, seed=seed, iters=iters, logger=logger
repo,
family,
dit_quant = scheme,
steps = steps,
res = res,
seed = seed,
iters = iters,
logger = logger,
)
row["lpips_vs_dense"] = _lpips_alex(ref_arr, np.array(img))
row["speedup_vs_dense"] = round(base_lat / row["gen_latency_s"], 3) if row["gen_latency_s"] else None
row["speedup_vs_dense"] = (
round(base_lat / row["gen_latency_s"], 3) if row["gen_latency_s"] else None
)
try:
img.save(out / f"dit_{family}_{scheme}.png")
except Exception:
pass
rows.append(row)
print(f" dit {scheme:5s}: {json.dumps(row)}", flush=True)
print(f" dit {scheme:5s}: {json.dumps(row)}", flush = True)
return rows
# ── main ──────────────────────────────────────────────────────────────────────
def main(argv=None) -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--family", required=True, choices=sorted(_FAMILIES))
ap.add_argument("--mode", required=True, choices=("te", "vae", "e2e", "teacc", "dit"))
ap.add_argument("--dit-schemes", default="auto", help="dit mode: comma list e.g. auto,fp8,int8,mxfp8")
ap.add_argument("--warmup", type=int, default=2)
ap.add_argument("--iters", type=int, default=5)
ap.add_argument("--steps", type=int, default=20, help="e2e denoise steps")
ap.add_argument("--res", type=int, default=1024, help="e2e image size")
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--e2e-iters", type=int, default=3)
ap.add_argument("--variant", choices=("both", "dense", "auto"), default="both", help="e2e variant(s)")
ap.add_argument("--te-scheme", default="auto", help="te mode: auto | fp8_dynamic | fp8 | int8")
ap.add_argument("--out", default="outputs/quant_speedmem")
def main(argv = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("--family", required = True, choices = sorted(_FAMILIES))
ap.add_argument("--mode", required = True, choices = ("te", "vae", "e2e", "teacc", "dit"))
ap.add_argument(
"--dit-schemes", default = "auto", help = "dit mode: comma list e.g. auto,fp8,int8,mxfp8"
)
ap.add_argument("--warmup", type = int, default = 2)
ap.add_argument("--iters", type = int, default = 5)
ap.add_argument("--steps", type = int, default = 20, help = "e2e denoise steps")
ap.add_argument("--res", type = int, default = 1024, help = "e2e image size")
ap.add_argument("--seed", type = int, default = 42)
ap.add_argument("--e2e-iters", type = int, default = 3)
ap.add_argument(
"--variant", choices = ("both", "dense", "auto"), default = "both", help = "e2e variant(s)"
)
ap.add_argument("--te-scheme", default = "auto", help = "te mode: auto | fp8_dynamic | fp8 | int8")
ap.add_argument("--out", default = "outputs/quant_speedmem")
args = ap.parse_args(argv)
import logging
logging.basicConfig(level=logging.INFO, format="%(message)s")
logging.basicConfig(level = logging.INFO, format = "%(message)s")
logger = logging.getLogger("speedmem")
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
out.mkdir(parents = True, exist_ok = True)
print(f"== speed+mem bench: family={args.family} mode={args.mode} ==", flush=True)
print(f"== speed+mem bench: family={args.family} mode={args.mode} ==", flush = True)
if args.mode == "dit":
schemes = [s.strip() for s in args.dit_schemes.split(",") if s.strip()]
rows = measure_dit(
args.family, schemes=schemes, steps=args.steps, res=args.res, seed=args.seed,
iters=args.e2e_iters, out=out, logger=logger
args.family,
schemes = schemes,
steps = args.steps,
res = args.res,
seed = args.seed,
iters = args.e2e_iters,
out = out,
logger = logger,
)
elif args.mode == "teacc":
rows = measure_te_accuracy(args.family, logger=logger)
rows = measure_te_accuracy(args.family, logger = logger)
elif args.mode == "te":
rows = measure_te(args.family, warmup=args.warmup, iters=args.iters, scheme=args.te_scheme, logger=logger)
rows = measure_te(
args.family, warmup = args.warmup, iters = args.iters, scheme = args.te_scheme, logger = logger
)
elif args.mode == "vae":
rows = measure_vae(args.family, warmup=args.warmup, iters=args.iters, logger=logger)
rows = measure_vae(args.family, warmup = args.warmup, iters = args.iters, logger = logger)
else:
rows = measure_e2e(
args.family, steps=args.steps, res=args.res, seed=args.seed, iters=args.e2e_iters,
out=out, variant=args.variant, logger=logger
args.family,
steps = args.steps,
res = args.res,
seed = args.seed,
iters = args.e2e_iters,
out = out,
variant = args.variant,
logger = logger,
)
for r in rows:
print(" " + json.dumps(r), flush=True)
print(" " + json.dumps(r), flush = True)
suffix = ""
if args.mode == "e2e" and args.variant != "both":
suffix = f"_{args.variant}"
elif args.mode == "te" and args.te_scheme != "auto":
suffix = f"_{args.te_scheme}"
dest = out / f"{args.mode}_{args.family}{suffix}.json"
with open(dest, "w", encoding="utf-8") as fh:
json.dump(rows, fh, indent=2)
print(f"wrote {dest}", flush=True)
with open(dest, "w", encoding = "utf-8") as fh:
json.dump(rows, fh, indent = 2)
print(f"wrote {dest}", flush = True)
return 0

16
scripts/sdpa_mask_backend_probe.py
View file

@ -4,6 +4,7 @@
"""Which SDPA backends tolerate a dense bool attn_mask, and at what cost, at Hunyuan's real
joint shape (B=1, H=16, N=50345, D=128, bf16)? Decides whether nulling the all-True mask is
the real win on the PRODUCTION cuDNN path (not just the native math fallback)."""
import time
import torch
import torch.nn.functional as F
@ -14,10 +15,10 @@ dev, dt = "cuda:0", torch.bfloat16
def mk():
return torch.randn(B, H, N, D, device=dev, dtype=dt)
return torch.randn(B, H, N, D, device = dev, dtype = dt)
def timed(fn, iters=20):
def timed(fn, iters = 20):
torch.cuda.synchronize()
for _ in range(3):
try:
@ -33,7 +34,7 @@ def timed(fn, iters=20):
q, k, v = mk(), mk(), mk()
dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev)
dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev)
backends = {
"default(dispatch)": None,
@ -46,17 +47,18 @@ backends = {
print(f"shape B={B} H={H} N={N} D={D} {dt}\n")
print(f"{'backend':<20}{'mask=dense(ms)':>18}{'mask=None(ms)':>18}")
for name, bk in backends.items():
def run_dense():
if bk is None:
return F.scaled_dot_product_attention(q, k, v, attn_mask=dense)
return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
with sdpa_kernel(bk):
return F.scaled_dot_product_attention(q, k, v, attn_mask=dense)
return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
def run_none():
if bk is None:
return F.scaled_dot_product_attention(q, k, v, attn_mask=None)
return F.scaled_dot_product_attention(q, k, v, attn_mask = None)
with sdpa_kernel(bk):
return F.scaled_dot_product_attention(q, k, v, attn_mask=None)
return F.scaled_dot_product_attention(q, k, v, attn_mask = None)
dms = timed(run_dense)
nms = timed(run_none)

292
scripts/video_speedmem_bench.py
View file

@ -84,33 +84,28 @@ _FAMILIES: dict[str, dict[str, Any]] = {
# ── cuda memory / timing helpers ───────────────────────────────────────────────
def _sync() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
def _reset_peak() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
def _alloc_gb() -> float:
import torch
return torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _peak_gb() -> float:
import torch
return torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _empty() -> None:
import torch
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
@ -131,12 +126,11 @@ def _lpips_alex(ref_arr, arr):
fn = _LP.get("fn")
if fn is None:
fn = lpips.LPIPS(net="alex", verbose=False).eval()
fn = lpips.LPIPS(net = "alex", verbose = False).eval()
_LP["fn"] = fn
def _t(a):
import torch as _torch
return _torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
@ -210,11 +204,10 @@ def _target():
"""The object the real casters/optimisers read: a stand-in for DiffusionDeviceTarget.
supports_default_torch_compile must be True or compile_eligible() bails."""
import torch
return types.SimpleNamespace(
device="cuda",
dtype=torch.bfloat16,
supports_default_torch_compile=True,
device = "cuda",
dtype = torch.bfloat16,
supports_default_torch_compile = True,
)
@ -224,33 +217,59 @@ def _target():
# `shipped` is the current default; `speedmax`/`flash4` probe untapped headroom.
_CONFIGS: dict[str, dict[str, Any]] = {
# te vae dit speed attn cache
"reference": dict(te="none", vae="none", dit="none", speed="off", attn="native", cache="off"),
"compile": dict(te="none", vae="none", dit="none", speed="default", attn="native", cache="off"),
"cudnn": dict(te="none", vae="none", dit="none", speed="default", attn="auto", cache="off"),
"fbcache": dict(te="none", vae="none", dit="none", speed="default", attn="auto", cache="auto"),
"ditquant": dict(te="none", vae="none", dit="auto", speed="default", attn="auto", cache="auto"),
"shipped": dict(te="auto", vae="auto", dit="auto", speed="default", attn="auto", cache="auto"),
"speedmax": dict(te="auto", vae="auto", dit="auto", speed="max", attn="auto", cache="auto"),
"flash4": dict(te="auto", vae="auto", dit="auto", speed="default", attn="flash4", cache="auto"),
"reference": dict(te = "none", vae = "none", dit = "none", speed = "off", attn = "native", cache = "off"),
"compile": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "native", cache = "off"),
"cudnn": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off"),
"fbcache": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "auto"),
"ditquant": dict(te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "auto"),
"shipped": dict(te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "auto", cache = "auto"),
"speedmax": dict(te = "auto", vae = "auto", dit = "auto", speed = "max", attn = "auto", cache = "auto"),
"flash4": dict(te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "flash4", cache = "auto"),
# diagnostics: isolate whether the DiT-quant + compile crash needs FBCache.
"diag_ditq_default_nocache": dict(te="none", vae="none", dit="auto", speed="default", attn="native", cache="off"),
"diag_ditq_max_nocache": dict(te="none", vae="none", dit="auto", speed="max", attn="native", cache="off"),
"diag_ditq_nocompile": dict(te="none", vae="none", dit="auto", speed="eager", attn="native", cache="off"),
"diag_ditq_default_nocache": dict(
te = "none", vae = "none", dit = "auto", speed = "default", attn = "native", cache = "off"
),
"diag_ditq_max_nocache": dict(
te = "none", vae = "none", dit = "auto", speed = "max", attn = "native", cache = "off"
),
"diag_ditq_nocompile": dict(
te = "none", vae = "none", dit = "auto", speed = "eager", attn = "native", cache = "off"
),
# which quant scheme survives torch.compile? (fp8/mslk fails "fake tensors"; test int8/mxfp8)
"diag_ditint8_compile": dict(te="none", vae="none", dit="int8", speed="default", attn="native", cache="off"),
"diag_ditmxfp8_compile": dict(te="none", vae="none", dit="mxfp8", speed="default", attn="native", cache="off"),
"diag_ditint8_fbcache": dict(te="none", vae="none", dit="int8", speed="default", attn="native", cache="auto"),
"diag_ditint8_compile": dict(
te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "off"
),
"diag_ditmxfp8_compile": dict(
te = "none", vae = "none", dit = "mxfp8", speed = "default", attn = "native", cache = "off"
),
"diag_ditint8_fbcache": dict(
te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "auto"
),
# isolate the quant x FBCache over-caching interaction at production size:
"ditfp8_nocache": dict(te="none", vae="none", dit="auto", speed="default", attn="auto", cache="off"),
"te_fbcache": dict(te="auto", vae="none", dit="none", speed="default", attn="auto", cache="auto"),
"ditfp8_fbcache": dict(te="none", vae="none", dit="auto", speed="default", attn="auto", cache="auto"),
"ditint8_fbcache_prod": dict(te="none", vae="none", dit="int8", speed="default", attn="auto", cache="auto"),
"ditfp8_nocache": dict(
te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "off"
),
"te_fbcache": dict(
te = "auto", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "auto"
),
"ditfp8_fbcache": dict(
te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "auto"
),
"ditint8_fbcache_prod": dict(
te = "none", vae = "none", dit = "int8", speed = "default", attn = "auto", cache = "auto"
),
# mixed-fp8 vs int8 head-to-head (Phase 3): the DiT-quant accuracy comparison on Wan/Hunyuan.
# fp8 here goes through the production quantize_transformer family exclude (input embedders
# kept bf16), so it is only non-black if the mixed-fp8 wiring is live. cache on AND off.
"ditfp8mixed_nocache": dict(te="none", vae="none", dit="fp8", speed="default", attn="native", cache="off"),
"ditfp8mixed_fbcache": dict(te="none", vae="none", dit="fp8", speed="default", attn="auto", cache="auto"),
"ditint8_nocache": dict(te="none", vae="none", dit="int8", speed="default", attn="native", cache="off"),
"ditfp8mixed_nocache": dict(
te = "none", vae = "none", dit = "fp8", speed = "default", attn = "native", cache = "off"
),
"ditfp8mixed_fbcache": dict(
te = "none", vae = "none", dit = "fp8", speed = "default", attn = "auto", cache = "auto"
),
"ditint8_nocache": dict(
te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "off"
),
}
@ -258,7 +277,7 @@ def _build_pipe(repo: str, force_fp32_vae: bool):
import torch
diffusers = _import_diffusers()
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype = torch.bfloat16)
pipe = pipe.to("cuda")
# Wan-style VAEs decode in fp32 for numerical stability (the loader pins this via
# vae_force_fp32); loading bf16 bands every clip, so mirror the loader.
@ -287,8 +306,16 @@ class _SecondExpertView:
setattr(self._pipe, "transformer_2" if name == "transformer" else name, value)
def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bool, default_steps: int,
logger=None) -> dict:
def _apply_levers(
pipe,
cfg: dict,
*,
fam_name: str,
fam_obj,
force_fp32_vae: bool,
default_steps: int,
logger = None,
) -> dict:
"""Apply the configured levers with the loader's own argument values, in the loader's order:
quant (dit -> te -> vae) THEN optimisation layers (cache -> attention -> speed). For a
dual-expert MoE (pipe.transformer_2 present) every DiT-touching lever is applied to BOTH experts
@ -305,7 +332,14 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo
)
tgt = _target()
engaged = {"dit": None, "te": None, "vae": None, "attn": None, "cache": None, "speed_optims": {}}
engaged = {
"dit": None,
"te": None,
"vae": None,
"attn": None,
"cache": None,
"speed_optims": {},
}
# DiT-touching levers run per expert: [pipe] for a single-DiT family, plus a second-expert view
# for a dual-expert MoE. Each view exposes the expert as ``.transformer``.
@ -316,7 +350,7 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo
# DiT quant (pipeline kind, resident): mutates each expert's transformer in place.
if cfg["dit"] not in ("none", "off"):
schemes = [
quantize_transformer(v, tgt, mode=cfg["dit"], family=fam_name, logger=logger)
quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger)
for v in views
]
engaged["dit"] = schemes[0]
@ -327,15 +361,20 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo
# TE quant (once; text encoders are shared, not per-expert).
if cfg["te"] not in ("none", "off"):
engaged["te"] = quantize_text_encoders(
pipe, tgt, mode=cfg["te"], family=fam_name, offload_active=False, logger=logger
pipe, tgt, mode = cfg["te"], family = fam_name, offload_active = False, logger = logger
)
_empty()
# VAE quant (once; Wan force_fp32 pins dense inside quantize_vae regardless).
if cfg["vae"] not in ("none", "off"):
engaged["vae"] = quantize_vae(
pipe, tgt, mode=cfg["vae"], family=fam_name, offload_active=False,
force_fp32=force_fp32_vae, logger=logger,
pipe,
tgt,
mode = cfg["vae"],
family = fam_name,
offload_active = False,
force_fp32 = force_fp32_vae,
logger = logger,
)
_empty()
@ -356,29 +395,51 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo
if cache_request is not None:
for v in views:
engaged["cache"] = apply_step_cache(
v, mode=cache_request, threshold=None,
quant_active=dit_quant_active, logger=logger,
v,
mode = cache_request,
threshold = None,
quant_active = dit_quant_active,
logger = logger,
)
cache_active = engaged["cache"] not in (None, "off")
# Attention (per expert).
backend = select_attention_backend(tgt, cfg["attn"], speed_active=speed_active)
backend = select_attention_backend(tgt, cfg["attn"], speed_active = speed_active)
for v in views:
engaged["attn"] = apply_attention_backend(v, backend, logger=logger)
engaged["attn"] = apply_attention_backend(v, backend, logger = logger)
# Speed profile (per expert; compiles each denoiser).
if speed != "off":
for v in views:
engaged["speed_optims"] = apply_speed_optims(
v, tgt, is_gguf=False, family=fam_obj, speed_mode=speed,
cache_active=cache_active, offload_active=False, logger=logger,
v,
tgt,
is_gguf = False,
family = fam_obj,
speed_mode = speed,
cache_active = cache_active,
offload_active = False,
logger = logger,
)
engaged["_effective_speed"] = speed
return engaged
def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cache_mode,
dit_quant_active, default_steps, guidance_via_guider=False, logger=None):
def _timed_video(
pipe,
*,
steps,
width,
height,
num_frames,
guidance,
seed,
cache_mode,
dit_quant_active,
default_steps,
guidance_via_guider = False,
logger = None,
):
"""One clip generation. Re-checks FBCache per generation (maybe_toggle_step_cache) exactly
like the loader, then times total + per-step. Returns (output, total_s, [per_step_ms])."""
import torch
@ -388,12 +449,12 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach
if cache_mode == "auto":
try:
maybe_toggle_step_cache(
pipe, steps=steps, quant_active=dit_quant_active, threshold=None, logger=logger
pipe, steps = steps, quant_active = dit_quant_active, threshold = None, logger = logger
)
except Exception:
pass
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
step_ts: list[float] = []
last = [0.0]
@ -406,8 +467,12 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach
return kw
kwargs = dict(
prompt=PROMPT, width=width, height=height, num_frames=num_frames,
num_inference_steps=steps, generator=g,
prompt = PROMPT,
width = width,
height = height,
num_frames = num_frames,
num_inference_steps = steps,
generator = g,
)
if guidance_via_guider:
# HunyuanVideo-1.5: CFG lives on a guider component and __call__ takes no
@ -429,8 +494,20 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach
return out, (time.perf_counter() - t0), step_ts
def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, height: int,
num_frames: int, seed: int, iters: int, out: Path, logger=None):
def _run_config(
name: str,
cfg: dict,
*,
family: str,
steps: int,
width: int,
height: int,
num_frames: int,
seed: int,
iters: int,
out: Path,
logger = None,
):
import numpy as np
from core.inference.video_families import detect_video_family
@ -443,21 +520,26 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he
default_steps = getattr(fam_obj, "default_steps", 50)
gvg = bool(getattr(fam_obj, "guidance_via_guider", False))
_empty(); _reset_peak()
_empty()
_reset_peak()
# Fresh dynamo state per config so a prior config's compiled graphs cannot leak into this one
# (each config builds a fresh pipe; without this, later configs in a multi-config run can be
# measured against a dirty compile cache).
try:
import torch
torch._dynamo.reset()
except Exception:
pass
pipe = _build_pipe(repo, force_fp32)
load_peak = _peak_gb()
engaged = _apply_levers(
pipe, cfg, fam_name=family, fam_obj=fam_obj, force_fp32_vae=force_fp32,
default_steps=default_steps, logger=logger,
pipe,
cfg,
fam_name = family,
fam_obj = fam_obj,
force_fp32_vae = force_fp32,
default_steps = default_steps,
logger = logger,
)
_empty()
weights_gb = _alloc_gb()
@ -466,18 +548,36 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he
# warmup (pays the one-time compile / autotune)
_timed_video(
pipe, steps=steps, width=width, height=height, num_frames=num_frames, guidance=guidance,
seed=seed, cache_mode=cache_mode, dit_quant_active=dit_active, default_steps=default_steps,
guidance_via_guider=gvg, logger=logger,
pipe,
steps = steps,
width = width,
height = height,
num_frames = num_frames,
guidance = guidance,
seed = seed,
cache_mode = cache_mode,
dit_quant_active = dit_active,
default_steps = default_steps,
guidance_via_guider = gvg,
logger = logger,
)
_reset_peak()
dts, steps_ms = [], []
last_out = None
for _ in range(iters):
last_out, dt, st = _timed_video(
pipe, steps=steps, width=width, height=height, num_frames=num_frames, guidance=guidance,
seed=seed, cache_mode=cache_mode, dit_quant_active=dit_active, default_steps=default_steps,
guidance_via_guider=gvg, logger=logger,
pipe,
steps = steps,
width = width,
height = height,
num_frames = num_frames,
guidance = guidance,
seed = seed,
cache_mode = cache_mode,
dit_quant_active = dit_active,
default_steps = default_steps,
guidance_via_guider = gvg,
logger = logger,
)
dts.append(dt)
steps_ms.append(_median(st) if st else 0.0)
@ -487,7 +587,6 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he
try:
if arrs:
from PIL import Image
Image.fromarray(arrs[len(arrs) // 2]).save(out / f"vid_{family}_{name}.png")
except Exception:
pass
@ -496,7 +595,6 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he
if name == "reference" and arrs:
try:
import numpy as _np
_np.savez_compressed(out / "ref_frames.npz", *arrs)
except Exception:
pass
@ -526,35 +624,39 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he
return row, arrs
def main(argv=None) -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--family", default="wan2.2-ti2v-5b", choices=sorted(_FAMILIES))
ap.add_argument("--configs", default=",".join(_CONFIGS),
help="comma list from: " + ",".join(_CONFIGS))
ap.add_argument("--steps", type=int, default=30)
ap.add_argument("--num-frames", type=int, default=25)
ap.add_argument("--width", type=int, default=512)
ap.add_argument("--height", type=int, default=320)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--iters", type=int, default=3)
ap.add_argument("--out", default="outputs/video_speedmem")
def main(argv = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("--family", default = "wan2.2-ti2v-5b", choices = sorted(_FAMILIES))
ap.add_argument(
"--configs", default = ",".join(_CONFIGS), help = "comma list from: " + ",".join(_CONFIGS)
)
ap.add_argument("--steps", type = int, default = 30)
ap.add_argument("--num-frames", type = int, default = 25)
ap.add_argument("--width", type = int, default = 512)
ap.add_argument("--height", type = int, default = 320)
ap.add_argument("--seed", type = int, default = 42)
ap.add_argument("--iters", type = int, default = 3)
ap.add_argument("--out", default = "outputs/video_speedmem")
args = ap.parse_args(argv)
import logging
logging.basicConfig(level=logging.INFO, format="%(message)s")
logging.basicConfig(level = logging.INFO, format = "%(message)s")
logger = logging.getLogger("videobench")
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
out.mkdir(parents = True, exist_ok = True)
names = [c.strip() for c in args.configs.split(",") if c.strip()]
for n in names:
if n not in _CONFIGS:
raise SystemExit(f"unknown config '{n}'; choose from {list(_CONFIGS)}")
print(f"== video speed+mem bench: family={args.family} configs={names} "
f"steps={args.steps} frames={args.num_frames} {args.width}x{args.height} ==", flush=True)
print(
f"== video speed+mem bench: family={args.family} configs={names} "
f"steps={args.steps} frames={args.num_frames} {args.width}x{args.height} ==",
flush = True,
)
rows = []
ref_arrs = None
@ -564,22 +666,32 @@ def main(argv=None) -> int:
if ref_npz.exists():
try:
import numpy as _np
with _np.load(ref_npz) as z:
ref_arrs = [z[k] for k in z.files]
except Exception:
ref_arrs = None
for n in names:
row, arrs = _run_config(
n, _CONFIGS[n], family=args.family, steps=args.steps, width=args.width,
height=args.height, num_frames=args.num_frames, seed=args.seed, iters=args.iters,
out=out, logger=logger,
n,
_CONFIGS[n],
family = args.family,
steps = args.steps,
width = args.width,
height = args.height,
num_frames = args.num_frames,
seed = args.seed,
iters = args.iters,
out = out,
logger = logger,
)
if n == "reference":
ref_arrs = arrs
row["lpips_vs_reference"] = _mean_lpips(ref_arrs, arrs) if ref_arrs is not None else None
rows.append(row)
print(f" [{n}] {json.dumps({k: row[k] for k in ('dit_scheme','te_scheme','attn','cache','effective_speed','weights_gb','gen_peak_gb','gen_latency_s','per_step_ms','lpips_vs_reference')})}", flush=True)
print(
f" [{n}] {json.dumps({k: row[k] for k in ('dit_scheme','te_scheme','attn','cache','effective_speed','weights_gb','gen_peak_gb','gen_latency_s','per_step_ms','lpips_vs_reference')})}",
flush = True,
)
# speedups relative to reference (if present)
ref_lat = next((r["gen_latency_s"] for r in rows if r["config"] == "reference"), None)
@ -589,9 +701,9 @@ def main(argv=None) -> int:
)
dest = out / f"video_{args.family}_{'-'.join(names)}.json"
with open(dest, "w", encoding="utf-8") as fh:
json.dump(rows, fh, indent=2)
print(f"wrote {dest}", flush=True)
with open(dest, "w", encoding = "utf-8") as fh:
json.dump(rows, fh, indent = 2)
print(f"wrote {dest}", flush = True)
return 0

44
studio/backend/core/inference/diffusion_attention.py
View file

@ -437,9 +437,9 @@ def _null_mask_processor_cls():
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
image_rotary_emb=None,
encoder_hidden_states = None,
attention_mask = None,
image_rotary_emb = None,
):
# Fast path only when the pre-hook removed all padding (attn_mask redundant); a
# constant python bool so torch.compile const-folds the branch (no graph break).
@ -447,9 +447,9 @@ def _null_mask_processor_cls():
return super().__call__(
attn,
hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
image_rotary_emb=image_rotary_emb,
encoder_hidden_states = encoder_hidden_states,
attention_mask = attention_mask,
image_rotary_emb = image_rotary_emb,
)
# Null path = the stock body with the mask block removed and attn_mask=None.
@ -466,9 +466,8 @@ def _null_mask_processor_cls():
if image_rotary_emb is not None:
from diffusers.models.embeddings import apply_rotary_emb
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim = 1)
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim = 1)
if encoder_hidden_states is not None:
encoder_query = attn.add_q_proj(encoder_hidden_states)
@ -484,19 +483,19 @@ def _null_mask_processor_cls():
if attn.norm_added_k is not None:
encoder_key = attn.norm_added_k(encoder_key)
query = torch.cat([query, encoder_query], dim=1)
key = torch.cat([key, encoder_key], dim=1)
value = torch.cat([value, encoder_value], dim=1)
query = torch.cat([query, encoder_query], dim = 1)
key = torch.cat([key, encoder_key], dim = 1)
value = torch.cat([value, encoder_value], dim = 1)
hidden_states = dispatch_attention_fn(
query,
key,
value,
attn_mask=None,
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
attn_mask = None,
dropout_p = 0.0,
is_causal = False,
backend = self._attention_backend,
parallel_config = self._parallel_config,
)
hidden_states = hidden_states.flatten(2, 3)
@ -530,7 +529,7 @@ def _trim_stream(states, mask):
if states is None or mask is None or mask.dim() != 2:
return states, mask, True # nothing to mask -> treat as no-padding
mb = mask.bool()
keep = mb.any(dim=0) # column valid for at least one batch element
keep = mb.any(dim = 0) # column valid for at least one batch element
if not bool(keep.all()):
states = states[:, keep]
mask = mask[:, keep]
@ -642,7 +641,12 @@ def _install_null_processors(dit: Any, logger: Any) -> bool:
return installed > 0
def install_hunyuan_attention_trim(pipe: Any, family: Any, *, logger: Any = None) -> bool:
def install_hunyuan_attention_trim(
pipe: Any,
family: Any,
*,
logger: Any = None,
) -> bool:
"""HunyuanVideo-1.5 only: make the joint attention skip padded text tokens (see module note).
Installs a null-mask processor on every denoiser DiT block plus an eager pre-hook that trims
@ -662,7 +666,7 @@ def install_hunyuan_attention_trim(pipe: Any, family: Any, *, logger: Any = None
continue
if getattr(dit, "_unsloth_trim_hook", None) is None:
try:
handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs=True)
handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs = True)
dit._unsloth_trim_hook = handle
except Exception as exc: # noqa: BLE001 — optimisation only
_warn(logger, "hunyuan_attn_trim", exc)

1
studio/backend/core/inference/diffusion_prequant.py
View file

@ -319,6 +319,7 @@ def _validate_checkpoint(
ckpt_excludes = meta.get("exclude_name_tokens")
if ckpt_excludes is not None:
from .diffusion_transformer_quant import exclude_tokens_for_scheme
# The exclude set is derived from scheme AND family (fp8 keeps a family's padded-conditioning
# embedder bf16); use the checkpoint's recorded family so an artifact baked under an older
# token list (e.g. a Wan fp8 checkpoint built before the condition_embedder exclude) is

17
studio/backend/core/inference/diffusion_vae_quant.py
View file

@ -73,12 +73,12 @@ _VAE_AUTO_LADDER = (VAE_QUANT_FP8,)
# ltx-2 SSIM 0.942. FLUX.2 (klein/dev) and Hunyuan-1.5 pass fp8_dynamic, so they are not denied.
# The vae_force_fp32 video families (Wan) are gated separately at the loader (never quantise).
_VAE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
"sdxl": frozenset({VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC}),
"flux.1": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"flux.1-kontext": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"qwen-image": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"qwen-image-edit": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"ltx-2": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"sdxl": frozenset({VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC}),
"flux.1": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"flux.1-kontext": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"qwen-image": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"qwen-image-edit": frozenset({VAE_QUANT_FP8_DYNAMIC}),
"ltx-2": frozenset({VAE_QUANT_FP8_DYNAMIC}),
}
# Cache of device -> bool for the fp8_dynamic conv smoke probe (run once per device).
@ -255,7 +255,10 @@ def quantize_vae(
# leaves it dense; the big video Conv3d VAEs clear the floor and quantise. Explicit
# requests below skip this gate (the user asked for it directly).
if _vae_param_bytes(vae) < _VAE_AUTO_MIN_BYTES:
_note(logger, "vae auto: VAE under the ~1GB size floor; staying dense (quant saves ~nothing)")
_note(
logger,
"vae auto: VAE under the ~1GB size floor; staying dense (quant saves ~nothing)",
)
return None
mode = select_vae_quant_scheme(
target,

2
studio/backend/core/inference/video.py
View file

@ -1353,7 +1353,7 @@ class VideoBackend:
# attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask
# fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set
# so the requested kernel pins onto the new processors. No-op for every other family.
trim = install_hunyuan_attention_trim(view, fam, logger=logger)
trim = install_hunyuan_attention_trim(view, fam, logger = logger)
engaged = apply_attention_backend(
view,
select_attention_backend(

22
studio/backend/tests/test_diffusion_attention.py
View file

@ -462,7 +462,7 @@ def test_trim_stream_layout_agnostic_drops_only_global_padding():
def test_trim_stream_full_mask_is_noop():
states = torch.ones(1, 4, 2)
mask = torch.ones(1, 4, dtype=torch.long)
mask = torch.ones(1, 4, dtype = torch.long)
out_s, out_m, all_valid = att._trim_stream(states, mask)
assert out_s.shape == (1, 4, 2) and all_valid is True
@ -483,9 +483,9 @@ def test_trim_stream_mixed_batch_not_all_valid():
assert all_valid is False
def _fake_dit(n_blocks=2):
blocks = [types.SimpleNamespace(attn=types.SimpleNamespace()) for _ in range(n_blocks)]
return types.SimpleNamespace(transformer_blocks=blocks)
def _fake_dit(n_blocks = 2):
blocks = [types.SimpleNamespace(attn = types.SimpleNamespace()) for _ in range(n_blocks)]
return types.SimpleNamespace(transformer_blocks = blocks)
def test_trim_pre_hook_empties_t2v_image_and_trims_and_flags():
@ -508,7 +508,7 @@ def test_trim_stream_all_invalid_yields_empty_but_valid():
# A fully-padded secondary stream (e.g. unused byt5 in t2v) trims to 0 length and reports
# all_valid True (vacuous) so it does NOT drop the fast path -- it just contributes no tokens.
states = torch.ones(1, 5, 2)
mask = torch.zeros(1, 5, dtype=torch.long)
mask = torch.zeros(1, 5, dtype = torch.long)
out_s, out_m, all_valid = att._trim_stream(states, mask)
assert out_s.shape == (1, 0, 2) and all_valid is True
@ -522,7 +522,7 @@ def test_trim_pre_hook_byt5_all_invalid_keeps_fast_path():
"encoder_hidden_states": torch.arange(4.0).reshape(1, 4, 1),
"encoder_attention_mask": torch.tensor([[1, 1, 1, 0]]),
"encoder_hidden_states_2": torch.ones(1, 6, 1),
"encoder_attention_mask_2": torch.zeros(1, 6, dtype=torch.long), # all padding
"encoder_attention_mask_2": torch.zeros(1, 6, dtype = torch.long), # all padding
}
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
assert out["encoder_hidden_states"].shape == (1, 3, 1)
@ -538,7 +538,7 @@ def test_trim_pre_hook_empty_primary_reverts_and_disables():
kwargs = {
"image_embeds": torch.zeros(1, 5, 3),
"encoder_hidden_states": mllm,
"encoder_attention_mask": torch.zeros(1, 4, dtype=torch.long), # 0 valid
"encoder_attention_mask": torch.zeros(1, 4, dtype = torch.long), # 0 valid
}
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
assert out["encoder_hidden_states"] is mllm # reverted (not emptied)
@ -590,14 +590,14 @@ def test_trim_pre_hook_absent_stream_not_written_back():
def test_install_trim_noop_for_non_hunyuan_family():
fam = types.SimpleNamespace(transformer_class="WanTransformer3DModel")
pipe = types.SimpleNamespace(transformer=types.SimpleNamespace())
fam = types.SimpleNamespace(transformer_class = "WanTransformer3DModel")
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
assert att.install_hunyuan_attention_trim(pipe, fam) is False
def test_install_trim_noop_when_transformer_class_mismatch():
# Family claims Hunyuan but the loaded module isn't -> no processors touched, no diffusers
# import; returns False rather than swapping an unknown attention processor.
fam = types.SimpleNamespace(transformer_class="HunyuanVideo15Transformer3DModel")
pipe = types.SimpleNamespace(transformer=types.SimpleNamespace()) # class name mismatch
fam = types.SimpleNamespace(transformer_class = "HunyuanVideo15Transformer3DModel")
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) # class name mismatch
assert att.install_hunyuan_attention_trim(pipe, fam) is False

16
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -608,7 +608,10 @@ def test_family_deny_auto_skips_fp8_for_hunyuan(monkeypatch):
_stub_torch(monkeypatch, cc = (10, 0))
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5") == TQ_INT8
assert select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5-720p") == TQ_INT8
assert (
select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5-720p")
== TQ_INT8
)
# ltx-2 keeps the ladder head (fp8) -- it is not a black-frame family.
assert select_transformer_quant_scheme(_target(), "auto", family = "ltx-2") == TQ_FP8
@ -690,7 +693,12 @@ def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch):
torch.bfloat16 = "bfloat16"
class _Linear:
def __init__(self, inf, outf, dtype = "bfloat16"):
def __init__(
self,
inf,
outf,
dtype = "bfloat16",
):
self.in_features, self.out_features = inf, outf
self.weight = types.SimpleNamespace(dtype = dtype)
@ -711,4 +719,6 @@ def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch):
assert filt(big, "condition_embedder.time_embedder.linear_1") is False
assert filt(big, "blocks.0.attn1.to_q") is True
assert filt(big, "blocks.0.ffn.net.0.proj") is True
assert filt(big, "blocks.0.attn2.to_k") is True # cross-attn K/V stay fp8 (embedder bias rescues rows)
assert (
filt(big, "blocks.0.attn2.to_k") is True
) # cross-attn K/V stay fp8 (embedder bias rescues rows)

38
studio/backend/tests/test_diffusion_vae_quant.py
View file

@ -250,7 +250,11 @@ def test_fp8_dynamic_conv_filter(monkeypatch):
ff = captured["filter_fn"]
nn = torch.nn
def _mod(cls, shape, kernel_size = None):
def _mod(
cls,
shape,
kernel_size = None,
):
m = cls()
m.weight = _Weight(shape)
if kernel_size is not None:
@ -264,7 +268,10 @@ def test_fp8_dynamic_conv_filter(monkeypatch):
assert ff(_mod(nn.Linear, (512, 512)), "decoder.mid_block.attentions.0.to_q") is True
# POINTWISE (1x1 / 1x1x1) convs excluded even at %16 channels: torchao 0.17's fp8 conv
# kernel rejects them ("Activation and filter channels must match") -> crash at decode.
assert ff(_mod(nn.Conv2d, (128, 128, 1, 1), (1, 1)), "decoder.mid_block.attentions.0.proj_conv") is False
assert (
ff(_mod(nn.Conv2d, (128, 128, 1, 1), (1, 1)), "decoder.mid_block.attentions.0.proj_conv")
is False
)
assert ff(_mod(nn.Conv3d, (64, 64, 1, 1, 1), (1, 1, 1)), "decoder.time_mix.conv") is False
# Channels not a multiple of 16 excluded (torchao would skip them regardless): the RGB
# in/out head (C=3) and any off-16 dim.
@ -388,7 +395,9 @@ def test_quantize_vae_auto_resolves_and_applies(monkeypatch):
_stub_torch(monkeypatch, cc = (10, 0))
_stub_capability(monkeypatch, (10, 0))
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("auto must not use fp8_dynamic"))
monkeypatch.setattr(
vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("auto must not use fp8_dynamic")
)
calls: list = []
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
vae = object()
@ -406,7 +415,9 @@ def test_quantize_vae_auto_size_gate_skips_small(monkeypatch):
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 200_000_000) # ~0.2 GB image VAE
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("small VAE must stay dense"))
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("small VAE must stay dense"))
monkeypatch.setattr(
vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("small VAE must stay dense")
)
pipe = types.SimpleNamespace(vae = object())
assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") is None
@ -442,13 +453,18 @@ def test_quantize_vae_explicit_fp8_dynamic_probe_gates(monkeypatch):
_stub_torch(monkeypatch, cc = (10, 0))
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: False)
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("probe failed: must not cast"))
monkeypatch.setattr(
vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("probe failed: must not cast")
)
pipe = types.SimpleNamespace(vae = object())
assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") is None
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True)
calls: list = []
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v))
assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") == VAE_QUANT_FP8_DYNAMIC
assert (
quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein")
== VAE_QUANT_FP8_DYNAMIC
)
assert len(calls) == 1
@ -467,5 +483,11 @@ def test_real_family_deny_list_policy(monkeypatch):
assert select_vae_quant_scheme(_target(), "auto", family = "ltx-2") == VAE_QUANT_FP8
assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.1") is None
# FLUX.2 / Hunyuan keep fp8_dynamic available as an explicit opt-in (measured in-bar).
assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.2-klein") == VAE_QUANT_FP8_DYNAMIC
assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5") == VAE_QUANT_FP8_DYNAMIC
assert (
select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.2-klein")
== VAE_QUANT_FP8_DYNAMIC
)
assert (
select_vae_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5")
== VAE_QUANT_FP8_DYNAMIC
)