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

Bench fidelity: explicit MagCache step re-sizing, generation-time cache reporting, per-config flag restore

Browse source 
Mirror production's explicit-MagCache re-engagement in the video bench: a
magcache row installed at the family default step count now re-interpolates
its curve, retention window, and skip budget when --steps differs, instead of
timing a stale schedule users never run.

Report the generation-time cache state: the row's cache field is derived from
the post-toggle transformer marker rather than the load-time engagement, so an
auto cache toggled off below the step threshold (or a re-sized explicit
magcache) is published as it actually ran; the load-time value stays available
as cache_at_load.

Restore the process-wide backend flags (cudnn.benchmark, TF32 and fp16
accumulation, emulate_precision_casts) after each config, like the production
unload path, so a speed-enabled row cannot poison a later reference or off row
in the same --configs run.
This commit is contained in:
Daniel Han 2026-07-10 18:56:39 +00:00
commit bd1c2e0774
1 changed files with 67 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

77
scripts/video_speedmem_bench.py
View file

@ -484,7 +484,11 @@ def _apply_levers(
_empty()
# ── optimisation layers ──
snapshot_backend_flags() # process-wide flags; benchmark process is short-lived so no restore.
# Process-wide flags (cudnn.benchmark, TF32/fp16 accumulation,
# emulate_precision_casts): snapshotted here and restored by _run_config after
# the row completes, like the production unload path. Without the restore a
# speed-enabled row would poison a later reference/off row in the same process.
engaged["_flags_snapshot"] = snapshot_backend_flags()
speed = cfg["speed"]
speed_active = speed != "off"
@ -591,20 +595,22 @@ def _timed_video(
import torch
from core.inference.diffusion_cache import (
_disengage_step_cache,
apply_step_cache,
auto_cache_mode,
auto_cache_quality,
maybe_toggle_step_cache,
normalize_cache_quality,
)
# Toggle on EVERY expert view, exactly like the loader's per-view recheck
# (video.py iterates _views_for): toggling only the primary pipe would disable
# FBCache on pipe.transformer while transformer_2 stays cached, measuring a
# mixed cache state production never runs on a dual-expert MoE.
views = [pipe]
if getattr(pipe, "transformer_2", None) is not None:
views.append(_SecondExpertView(pipe))
if cache_mode == "auto":
# Toggle on EVERY expert view, exactly like the loader's per-view recheck
# (video.py iterates _views_for): toggling only the primary pipe would disable
# FBCache on pipe.transformer while transformer_2 stays cached, measuring a
# mixed cache state production never runs on a dual-expert MoE.
views = [pipe]
if getattr(pipe, "transformer_2", None) is not None:
views.append(_SecondExpertView(pipe))
for v, expert in zip(views, ("transformer", "transformer_2")):
try:
maybe_toggle_step_cache(
@ -620,6 +626,37 @@ def _timed_video(
)
except Exception:
pass
elif cache_mode == "magcache":
# An explicit magcache row never toggles off, but production re-engages it
# when the actual step count differs from the configured one (video.py: the
# marker carries "#s{steps}"; endswith, not substring, so "#s5" cannot match
# inside "#s50"). _apply_levers installed the cache at default_steps, so a
# --steps override would otherwise time a stale curve/skip budget users
# never run.
for v, expert in zip(views, ("transformer", "transformer_2")):
transformer = getattr(v, "transformer", None)
marker = getattr(transformer, "_unsloth_step_cache", None)
if not marker or str(marker).endswith(f"#s{int(steps)}"):
continue
try:
_disengage_step_cache(
transformer,
reason = f"explicit magcache re-interpolating for {steps} steps",
logger = logger,
)
apply_step_cache(
v,
mode = "magcache",
threshold = cache_threshold,
quant_active = dit_quant_active,
family = family,
steps = steps,
quality = normalize_cache_quality(cache_quality) or auto_cache_quality(family),
expert = expert,
logger = logger,
)
except Exception:
pass
# Clear step-cache residuals before EVERY generation, mirroring production
# (VideoBackend.generate calls _reset_step_cache before each run): diffusers keys
@ -841,6 +878,12 @@ def _run_config(
except Exception:
pass
# Report the GENERATION-time cache state, not the load-time one: the per-generation
# recheck in _timed_video can toggle an auto cache off (steps below the threshold) or
# re-size an explicit magcache, and production status reports the post-toggle state.
# The marker is "{mode}@{threshold}[#s{steps}]", written by apply_step_cache and
# cleared by _disengage_step_cache, so its mode prefix IS the live state.
cache_marker = getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None)
row = {
"config": name,
"family": family,
@ -850,13 +893,14 @@ def _run_config(
"te_scheme": engaged["te"] or "dense",
"vae_scheme": engaged["vae"] or "dense",
"attn": engaged["attn"] or "native",
"cache": engaged["cache"] or "off",
"cache": str(cache_marker).split("@")[0] if cache_marker else "off",
"cache_at_load": engaged["cache"] or "off",
"effective_speed": engaged.get("_effective_speed"),
"speed_optims": engaged["speed_optims"],
"attn_trim": engaged.get("attn_trim", False),
"cache_threshold": cache_threshold,
"cache_quality": cache_quality,
"cache_marker": getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None),
"cache_marker": cache_marker,
"load_peak_gb": round(load_peak, 2),
"weights_gb": round(weights_gb, 2),
"gen_peak_gb": round(gen_peak, 2),
@ -868,6 +912,19 @@ def _run_config(
}
del pipe
_empty()
# Restore the process-wide backend flags this config's speed layer mutated
# (snapshot captured in _apply_levers before any mutation), mirroring the
# production unload path, so a later config in the same --configs run starts
# from clean flags instead of inheriting cudnn.benchmark / TF32 /
# emulate_precision_casts from this row. An exception skips this, but the main
# loop does not continue past a failed config, so no row can run with leaked
# flags.
try:
from core.inference.diffusion_speed import restore_backend_flags
restore_backend_flags(engaged.get("_flags_snapshot"))
except Exception:
pass
return row, arrs