diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 1e7841dcb1..eab08903c1 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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). diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index ccfb26ea74..9e3d26ce49 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -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