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

Wire the Mega-cache compile prewarm into the video backend

Browse source 
The pre-warmed torch.compile cache (diffusion_compile_cache.py: fingerprinted
bundles over torch.compiler.save/load_cache_artifacts plus a persistent per-key
TORCHINDUCTOR_CACHE_DIR) was wired into the image backend only, so every video
load re-paid the full first-generation compile after every process restart (the
stock inductor dir lives in /tmp).

video.py now mirrors the image backend exactly: compile_cache.begin runs after
the attention-backend set (the fingerprint keys on the engaged kernel) and
before apply_speed_optims on a compile-eligible default/max tier, keyed on the
same fullgraph decision as the compile itself (an engaged or still-toggleable
step cache and a planned offload both drop it); the context is committed to
_VideoLoadState and compile_cache.save persists the bundle after the first
successful generation (env-gated distributor / first-run-warm mode);
_teardown_state restores the inductor dir, and a token-scoped
_rollback_precommit_compile_cache covers loads that die before the state
commit, mirroring the globals and CFG-parallel rollbacks.

Measured on HunyuanVideo-1.5-480p through the real VideoBackend (B200,
480x288/17f/30 steps): the first-generation extra drops from 107.5 s cold to
13.8 s when the 12.8 MB bundle loads into a fresh inductor dir (0.10 s load)
and to 11.7 s from the persistent per-key dir alone; through the wired
production path a restart lands at 10.5-10.8 s (bundle-only included) vs
86.5 s cold. Steady state is unchanged (2.4-2.6 s), and the loaded artifacts
are the same bits a local compile would produce, so numerics are untouched.

Tests: begin/save/restore lifecycle with the fullgraph keying, the Speed=off
and compile-ineligible skips, and the token-scoped pre-commit rollback.
This commit is contained in:
Daniel Han 2026-07-11 06:06:26 +00:00
commit b24a94d1c2
2 changed files with 176 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -75,10 +75,13 @@ from .diffusion_memory import (
plan_diffusion_memory,
snapshot_device_memory,
)
from . import diffusion_compile_cache as compile_cache
from .diffusion_speed import (
SPEED_DEFAULT,
SPEED_MAX,
SPEED_OFF,
apply_speed_optims,
compile_eligible,
resolve_speed_mode,
restore_backend_flags,
snapshot_backend_flags,
@ -318,6 +321,10 @@ class _VideoLoadState:
# dispatch on it and _teardown_state frees the replica through it.
cfg_parallel: Optional[str] = None
cfg_parallel_handle: Any = None
# Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a
# compiled tier ran begin(); generate() persists the bundle after the first compiled
# generation (when saving is enabled) and _teardown_state restores the inductor dir.
compile_cache_ctx: Any = None
resolved: Optional[dict] = None
@ -712,6 +719,7 @@ class VideoBackend:
# the globals a newer in-flight load now owns).
self._rollback_precommit_globals(token)
self._rollback_precommit_cfg_parallel(token)
self._rollback_precommit_compile_cache(token)
if self._load_token != token:
return
logger.error("video.load_failed: %s", exc)
@ -764,6 +772,19 @@ class VideoBackend:
self._precommit_cfg_parallel = None
teardown_cfg_parallel(pipe, proxy, logger = logger)
def _rollback_precommit_compile_cache(self, token: Optional[int]) -> None:
"""Restore TORCHINDUCTOR_CACHE_DIR for a load that ran ``compile_cache.begin``
but died BEFORE committing _VideoLoadState (the committed path restores via
_teardown_state instead). Token-scoped exactly like _rollback_precommit_globals."""
stored = getattr(self, "_precommit_compile_cache", None)
if stored is None:
return
stored_token, ctx = stored
if token is not None and stored_token is not None and stored_token != token:
return
self._precommit_compile_cache = None
compile_cache.restore(ctx)
# Base-repo subfolders an LTX-2.3 assembly reads: the checkpoint (plus the GGUF
# repo's extras files) supplies the DiT, connectors, both VAEs and the vocoder,
# so only the 2.0 base's scheduler / text encoder / tokenizer are pulled.
@ -1427,13 +1448,13 @@ class VideoBackend:
attention_engaged = None
attention_trim_engaged = False
speed_optims: tuple = ()
# A dense torchao transformer on the pipeline path is not a GGUF one, so is_gguf
# keys off the load kind (gguf) AND no quant having engaged.
gguf_transformer = kind == "gguf" and transformer_quant_engaged is None
for view in views:
# apply_attention_backend / apply_speed_optims both act on ``view.transformer``;
# calling them once per view sets the kernel and compiles each expert. The
# engaged values match across experts (same device/family/mode), so record the
# first pass; a dense torchao transformer on the pipeline path is not a GGUF one,
# so is_gguf keys off the load kind (gguf) AND no quant having engaged.
gguf_transformer = kind == "gguf" and transformer_quant_engaged is None
# apply_attention_backend acts on ``view.transformer``; calling it once per
# view sets the kernel on each expert. The engaged values match across
# experts (same device/family/mode), so record the first pass.
# HunyuanVideo-1.5 only: drop the ~99% zero-padded text tokens from the joint
# 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
@ -1452,6 +1473,47 @@ class VideoBackend:
),
logger = logger,
)
if view is pipe:
attention_engaged = engaged
attention_trim_engaged = trim
# Pre-warmed torch.compile cache (Mega-cache), mirroring the image backend: when a
# compiled tier will run, point inductor at a per-fingerprint dir and load a matching
# bundle BEFORE the first compiled forward. Measured on HunyuanVideo-1.5-480p (B200):
# the first-generation compile extra drops 107.5 s -> 13.8 s from a 12.8 MB bundle
# (fresh inductor dir), and the persistent per-key dir alone recovers a restart to
# 11.7 s -- with the stock /tmp inductor dir that ~100 s is repaid after every reboot.
# A miss is silent -> local compile, exactly as before. Must run AFTER the attention
# backend set (the fingerprint keys on the engaged kernel) and BEFORE
# apply_speed_optims (whose compile the loaded artifacts serve).
compile_ctx = None
if effective_speed in (SPEED_DEFAULT, SPEED_MAX) and compile_eligible(
target, is_gguf = gguf_transformer, family = fam
):
compile_ctx = compile_cache.begin(
family = fam.name,
transformer = getattr(pipe, "transformer", None),
dtype = getattr(target, "dtype", None),
quant = transformer_quant_engaged,
attention_backend = attention_engaged,
compile_kwargs = {
# Mirrors apply_speed_optims' fullgraph decision: an active step cache
# (or one that may still toggle on) OR a planned offload graph-breaks,
# so the cached bundle must be keyed on the same fullgraph setting.
"fullgraph": cache_engaged is None
and not cache_may_toggle
and plan.offload_policy == "none",
"dynamic": effective_speed != SPEED_MAX,
"mode": "max-autotune-no-cudagraphs"
if effective_speed == SPEED_MAX
else "default",
},
logger = logger,
)
# Until the state commit below transfers ownership to _teardown_state, a
# failed or cancelled load must restore TORCHINDUCTOR_CACHE_DIR itself
# (_run_load's error handler, token-scoped like the globals).
self._precommit_compile_cache = (_load_token, compile_ctx)
for view in views:
applied = apply_speed_optims(
view,
target,
@ -1465,10 +1527,8 @@ class VideoBackend:
offload_active = plan.offload_policy != "none",
)
if view is pipe:
attention_engaged = engaged
attention_trim_engaged = trim
speed_optims = tuple(k for k, v in applied.items() if v) + (
("hunyuan_attn_trim",) if trim else ()
("hunyuan_attn_trim",) if attention_trim_engaged else ()
)
with self._generate_lock:
# A cancelled/superseded load must not place weights on the GPU the arbiter
@ -1662,12 +1722,14 @@ class VideoBackend:
vae_quant = vae_quant_engaged,
cfg_parallel = "on" if cfg_parallel_proxy is not None else None,
cfg_parallel_handle = cfg_parallel_proxy,
compile_cache_ctx = compile_ctx,
resolved = resolved,
)
# Ownership of the globals and the CFG-parallel proxy transferred
# to _state / _teardown_state.
# Ownership of the globals, the CFG-parallel proxy, and the compile
# cache context transferred to _state / _teardown_state.
self._precommit_globals = None
self._precommit_cfg_parallel = None
self._precommit_compile_cache = None
logger.info(
"video.loaded: %s (%s, %s, offload=%s, speed=%s, quant=%s)",
repo_id,
@ -2113,6 +2175,13 @@ class VideoBackend:
cfg_proxy.note_generation_done()
except Exception: # noqa: BLE001
pass
# The first compiled generation just paid the compile cost; persist the
# warm torch.compile cache bundle when saving is enabled (distributor /
# first-run warm). Idempotent + best-effort -- never fails a generation.
try:
compile_cache.save(state.compile_cache_ctx, logger = logger)
except Exception: # noqa: BLE001 -- cache persistence is best-effort
pass
self._gen.update(phase = "export", eta_seconds = None)
video_frames = output.frames[0]
@ -2206,6 +2275,9 @@ class VideoBackend:
state, self._state = self._state, None
if state is not None:
restore_backend_flags(state.backend_flags)
# Restore TORCHINDUCTOR_CACHE_DIR so a later load (or the image backend)
# does not inherit this load's per-fingerprint inductor dir. Idempotent.
compile_cache.restore(state.compile_cache_ctx)
# A GGUF video load may have installed the process-wide compiled GGUF
# dequantizer; restore the stock kernels so a later load that asked for
# speed_mode=off gets the bit-identical path (mirrors the image unload).

93
studio/backend/tests/test_video_backend.py
View file

@ -1833,3 +1833,96 @@ def test_detect_load_family_arch_fallback_for_local_gguf(tmp_path, monkeypatch):
lambda p: {"general.architecture": "ltxv"},
)
assert vid._detect_load_family(str(d), "model.gguf", "ltx-2").name == "ltx-2"
# ── pre-warmed torch.compile cache (Mega-cache) wiring ───────────────────────────
def _stub_compile_cache(monkeypatch, ctx = None):
"""Record begin/save/restore calls on the compile-cache module video.py imports."""
from core.inference import video as video_mod
calls = {"begin": [], "save": [], "restore": []}
monkeypatch.setattr(
video_mod.compile_cache,
"begin",
lambda **kwargs: calls["begin"].append(kwargs) or ctx,
)
monkeypatch.setattr(
video_mod.compile_cache,
"save",
lambda c, logger = None: calls["save"].append(c) or False,
)
monkeypatch.setattr(
video_mod.compile_cache, "restore", lambda c: calls["restore"].append(c)
)
return calls
def test_video_compile_cache_begin_save_restore_lifecycle(fake_runtime, monkeypatch):
# A compiled-tier load must run compile_cache.begin BEFORE the speed profile
# (mega-cache load precedes the first compiled forward), commit the context to
# _VideoLoadState, persist the bundle after the first successful generation, and
# restore TORCHINDUCTOR_CACHE_DIR on unload -- mirroring the image backend.
from core.inference import video as video_mod
monkeypatch.setattr(video_mod, "compile_eligible", lambda *a, **k: True)
ctx = object()
calls = _stub_compile_cache(monkeypatch, ctx = ctx)
backend = VideoBackend()
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert len(calls["begin"]) == 1
kwargs = calls["begin"][0]
assert kwargs["family"] == "wan2.2-ti2v-5b"
# The wan5b auto step cache engages (or may toggle) on the default 50-step
# schedule, so the cached bundle must be keyed fullgraph=False like the compile.
assert kwargs["compile_kwargs"]["fullgraph"] is False
assert kwargs["compile_kwargs"]["dynamic"] is True
assert backend._state.compile_cache_ctx is ctx
backend.generate(prompt = "a sloth")
assert calls["save"] == [ctx]
backend.unload()
assert calls["restore"] == [ctx]
def test_video_compile_cache_skipped_on_speed_off_and_ineligible(fake_runtime, monkeypatch):
# Speed=off (bit-exact reference) or a compile-ineligible target must never touch
# the compile cache: no begin, no context, no restore side effects to leak.
from core.inference import video as video_mod
monkeypatch.setattr(video_mod, "compile_eligible", lambda *a, **k: True)
calls = _stub_compile_cache(monkeypatch, ctx = object())
backend = VideoBackend()
backend.load_pipeline(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", speed_mode = "off"
)
assert calls["begin"] == []
assert backend._state.compile_cache_ctx is None
backend.unload()
monkeypatch.setattr(video_mod, "compile_eligible", lambda *a, **k: False)
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert calls["begin"] == []
assert backend._state.compile_cache_ctx is None
def test_rollback_precommit_compile_cache_is_token_scoped(fake_runtime, monkeypatch):
# A load that ran compile_cache.begin and then died before committing
# _VideoLoadState must restore TORCHINDUCTOR_CACHE_DIR itself -- but only for its
# own token, so a superseded worker cannot clobber the redirect a newer in-flight
# load now owns. Mirrors _rollback_precommit_cfg_parallel.
from core.inference import video as video_mod
calls = []
monkeypatch.setattr(video_mod.compile_cache, "restore", lambda ctx: calls.append(ctx))
backend = VideoBackend()
ctx = object()
backend._precommit_compile_cache = (7, ctx)
backend._rollback_precommit_compile_cache(8) # stale worker: leave the stash alone
assert calls == [] and backend._precommit_compile_cache is not None
backend._rollback_precommit_compile_cache(7) # owning worker: restored + cleared
assert calls == [ctx] and backend._precommit_compile_cache is None
backend._rollback_precommit_compile_cache(7) # idempotent
assert len(calls) == 1