From 40e3747d437ba099cd5b7beaced0b175585d643d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 06:46:40 +0000 Subject: [PATCH] Absorb the first-generation compile hitch with a post-load background prewarm vLLM and SGLang finish every compilation at server startup (dummy batches through each compiled shape) so no request ever pays a compile mid-serving. The video backend's compiled tier instead paid a first-generation extra after every restart: ~54 s cold and ~11.3 s even with a warm Mega-cache bundle (the residual is dynamo tracing plus cudnn.benchmark autotune, which the bundle cannot carry). After a compiled DEFAULT-tier resident load commits, a daemon thread now runs one tiny throwaway generation (192x128 snapped, 4k+1-lattice 9 frames, 2 steps) under the generate lock. The default tier compiles dynamic=True, so the small trace serves every later resolution. Measured through the real backend (HunyuanVideo-1.5-480p, B200, 480x288/17f/30 steps): warm bundle: first-generation extra 11.3 s -> 2.1 s (9.6 s background warmup) cold start: the full compile moves off the user's first request entirely (14.5 s background; first generation extra 2.1 s), and the warmup persists the Mega-cache bundle itself steady state: unchanged (2.4-2.5 s per 30-step clip in every phase) Exactly lossless by construction: the warmup only changes when compilation work happens. It resets its step-cache residuals, the real generation seeds its own generator, and no process-wide flag is touched. The warmup registers itself as the active cancellable job, so unload, a new load, or cancel_generate abort it at a step boundary (verified: unload 2 s into a running prewarm returns in 6.7 s with the warmup cancelled). It yields untouched when a real request arrived first and is token-scoped against superseded loads. Gated per family (supports_compile_prewarm), skipped for speed=max (static per-shape graphs a warmup shape cannot serve), offload (every warmup forward would stream the DiT over PCIe), and CFG parallel (its planner owns compile-sensitive runs); the UNSLOTH_DIFFUSION_COMPILE_PREWARM kill switch disables it. The decision and reason land in the resolved record. Tests: +4 hermetic (decision gates, engaged-load spawn with snapped tiny shape, skip-without-compile, yield to generations / stale tokens); related backend set 551 passed; ruff clean. --- studio/backend/core/inference/video.py | 182 ++++++++++++++++++ .../backend/core/inference/video_families.py | 5 + studio/backend/tests/test_video_backend.py | 99 ++++++++++ 3 files changed, 286 insertions(+) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index eab08903c1..f25c794634 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -236,6 +236,63 @@ def _scheduler_step_progress(pipe: Any, on_step: Any): scheduler.step = original +# ── post-load compile prewarm ───────────────────────────────────────────────── +# vLLM and SGLang guarantee that ALL compilation finishes at server startup +# (dummy batches through every compiled shape) so no request ever pays a compile +# mid-serving. The video backend's analogue: after a compiled-tier load commits, +# a background thread runs one tiny throwaway generation so the first REAL +# request starts at steady state. Measured through the real backend +# (HunyuanVideo-1.5-480p, B200, 480x288/17f/30 steps, temp/vs_warm_probe.py): +# with a warm Mega-cache bundle the first-generation extra drops 11.3 s -> 2.1 s +# (a 9.0 s background warmup absorbs the dynamo tracing + cudnn autotune the +# bundle cannot carry); on a cold start the full ~54 s compile moves off the +# user's first request entirely. Steady state is untouched (the warmup adds no +# lasting state: cache residuals are reset and the real generation re-seeds its +# own generator). The shape is deliberately tiny -- the default tier compiles +# dynamic=True, so one small trace serves every later resolution. +_PREWARM_ENV = "UNSLOTH_DIFFUSION_COMPILE_PREWARM" +_PREWARM_WIDTH = 192 +_PREWARM_HEIGHT = 128 +_PREWARM_FRAMES = 9 +_PREWARM_STEPS = 2 + + +def compile_prewarm_decision( + fam: VideoFamily, + *, + speed_mode: str, + speed_optims: tuple, + offload_policy: str, + cfg_parallel_active: bool, +) -> tuple[bool, str]: + """Whether the post-load background compile prewarm should run, with the + resolved-record reason. Pure so it unit-tests without the runtime.""" + if (os.environ.get(_PREWARM_ENV) or "").strip().lower() in ("0", "off", "false", "no"): + return False, f"disabled via {_PREWARM_ENV}" + if "compiled" not in speed_optims: + return False, "not applicable (no regional compile engaged; nothing to prewarm)" + if speed_mode != SPEED_DEFAULT: + # speed=max compiles dynamic=False: a graph is keyed to the exact shape, + # so a tiny warmup shape would compile a graph the user's request never + # runs and the real shape would still pay its own compile. + return False, "skipped: speed=max compiles static per-shape graphs a warmup shape cannot serve" + if not bool(getattr(fam, "supports_compile_prewarm", True)): + return False, "family opted out (supports_compile_prewarm=False)" + if offload_policy != "none": + # Offload wraps block forwards in disabled onload hooks (the compiled-inner + # arming skips them) and every warmup forward would stream the full DiT + # over PCIe -- all cost, none of the measured warmup benefit. + return False, "skipped: offload streams weights per forward; the warmup would only churn transfers" + if cfg_parallel_active: + # The CFG-parallel proxy serialises compile-sensitive runs through its own + # per-(shape, steps, cache) planner; an unplanned warmup would bypass it. + return False, "skipped: CFG parallel plans compile-sensitive runs through its own dispatcher" + return True, ( + "background warmup generation absorbs the first-generation compile hitch " + "(measured 11.3s -> 2.1s warm-cache, ~54s cold, HunyuanVideo-1.5-480p on B200)" + ) + + def _detect_load_family( repo_id: str, gguf_filename: Optional[str], family_override: Optional[str] ) -> Optional[VideoFamily]: @@ -421,6 +478,9 @@ class VideoBackend: # a second begin_generate() is refused while the first still runs (or is # about to run: generate() only sets _gen after taking its locks). self._generate_job_active = False + # The post-load background compile prewarm thread (None until a compiled + # load spawns one); kept for tests/diagnostics, never joined on the hot path. + self._prewarm_thread: Optional[threading.Thread] = None # ── validation ─────────────────────────────────────────────────────────── @@ -1603,6 +1663,13 @@ class VideoBackend: logger = logger, ) + prewarm_on, prewarm_reason = compile_prewarm_decision( + fam, + speed_mode = effective_speed, + speed_optims = speed_optims, + offload_policy = offload_policy, + cfg_parallel_active = cfg_parallel_proxy is not None, + ) resolved = build_resolved_record( { "memory_mode": ( @@ -1681,6 +1748,11 @@ class VideoBackend: if getattr(fam, "vae_force_fp32", False) else "not engaged (dense VAE loaded)", ), + "compile_prewarm": ( + None, + "on" if prewarm_on else "off", + prewarm_reason, + ), } ) @@ -1739,6 +1811,18 @@ class VideoBackend: effective_speed, transformer_quant_engaged or "off", ) + if prewarm_on: + # Post-commit so a real request is never blocked behind an uncommitted + # load; the thread re-checks the token and yields to any generation + # that arrived first (which then pays -- and absorbs -- the warmup + # itself, exactly the pre-prewarm behaviour). + self._prewarm_thread = threading.Thread( + target = self._compile_prewarm, + args = (_load_token,), + daemon = True, + name = "video-compile-prewarm", + ) + self._prewarm_thread.start() return self.status() @staticmethod @@ -1757,6 +1841,104 @@ class VideoBackend: return Path(hf_hub_download_with_xet_fallback(repo_id, gguf_filename or "", hf_token)) + # ── post-load compile prewarm ───────────────────────────────────────────── + + def _compile_prewarm(self, token: Optional[int]) -> None: + """Run one tiny throwaway generation so the first REAL request skips the + compile/trace warmup (see the module-level prewarm constants for the + measured numbers). Runs on a daemon thread under ``_generate_lock``, + registers itself as the active cancellable job (so unload / a new load / + cancel_generate can abort it at a step boundary, exactly like a real + generation), and never touches the user-visible ``_gen`` progress. A + failure or cancellation is logged and the load simply keeps the + pre-prewarm behaviour: the first real generation pays the warmup.""" + import torch + + cancel = threading.Event() + with self._generate_lock: + with self._lock: + state = self._state + if state is None or (token is not None and token != self._load_token): + return # superseded/unloaded before the warmup could start + if self._generate_job_active or self._gen.get("active"): + return # a real request beat us; it absorbs the warmup itself + self._active_generate_cancel = cancel + started = time.monotonic() + try: + pipe = state.pipe + fam = state.family + width, height = snap_video_size(fam, _PREWARM_WIDTH, _PREWARM_HEIGHT) + frames = snap_num_frames(fam, _PREWARM_FRAMES) + call_params = inspect.signature(pipe.__call__).parameters + kwargs: dict[str, Any] = { + "prompt": "warmup", + "num_inference_steps": _PREWARM_STEPS, + "width": width, + "height": height, + "num_frames": frames, + "generator": torch.Generator(device = state.device).manual_seed(0), + } + # Default guidance keeps the CFG branch structure of a real run + # (both guider branches trace), mirroring generate(). + if fam.guidance_via_guider: + pipe.guider.guidance_scale = float(fam.default_guidance) + else: + kwargs[fam.cfg_kwarg] = float(fam.default_guidance) + if "frame_rate" in call_params: + kwargs["frame_rate"] = float(fam.default_fps) + + def _on_step(p, step_index, timestep, callback_kwargs): + if cancel.is_set(): + p._interrupt = True + return callback_kwargs + + def _on_scheduler_step(done: int) -> None: + if cancel.is_set(): + raise _VideoGenerationCancelled() + + if "callback_on_step_end" in call_params: + kwargs["callback_on_step_end"] = _on_step + progress_ctx = contextlib.nullcontext() + else: + progress_ctx = _scheduler_step_progress(pipe, _on_scheduler_step) + with torch.inference_mode(), progress_ctx: + pipe(**kwargs) + if cancel.is_set(): + raise _VideoGenerationCancelled() + # Drop the warmup's step-cache residuals so the next real + # generation starts from the same state as an unwarmed load + # (generate() also resets per request; this is defence in depth). + if state.transformer_cache: + self._reset_step_cache(pipe) + # The warmup just paid the compile: persist the Mega-cache bundle + # now (env-gated, idempotent) instead of waiting for the first + # real generation, mirroring generate()'s save point. + try: + compile_cache.save(state.compile_cache_ctx, logger = logger) + except Exception: # noqa: BLE001 -- cache persistence is best-effort + pass + logger.info( + "video.compile_prewarm: warmup absorbed the compile hitch in %.1fs " + "(%dx%d, %d frames, %d steps)", + time.monotonic() - started, + width, + height, + frames, + _PREWARM_STEPS, + ) + except _VideoGenerationCancelled: + logger.info("video.compile_prewarm: cancelled (unload / new load / user cancel)") + except Exception as exc: # noqa: BLE001 -- warmup is best-effort, never fatal + logger.warning( + "video.compile_prewarm: failed (%s); the first generation pays the " + "compile warmup instead", + exc, + ) + finally: + with self._lock: + if self._active_generate_cancel is cancel: + self._active_generate_cancel = None + # ── generation ─────────────────────────────────────────────────────────── @staticmethod diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index 4254210e58..39752e1f5c 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -77,6 +77,11 @@ class VideoFamily: # True when the family's DiT compiles cleanly with regional torch.compile # (Wan/LTX-2 declare _repeated_blocks; set False until verified per family). supports_torch_compile: bool = True + # True when the post-load background compile prewarm (a tiny throwaway + # generation that absorbs the first-generation compile/trace hitch, the + # vLLM/SGLang "compile before serving" pattern) may run for this family. + # Only consulted when the load actually engaged the regional compile. + supports_compile_prewarm: bool = True # Families whose activations overflow float16 -> the loader promotes fp16 to # float32. Video DiTs are bf16-native, so this defaults True (fp16 is never # the right resolution for them; bf16 or float32 only). diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 4f9e31cc62..18b1d5fcca 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -1924,3 +1924,102 @@ def test_rollback_precommit_compile_cache_is_token_scoped(fake_runtime, monkeypa assert calls == [ctx] and backend._precommit_compile_cache is None backend._rollback_precommit_compile_cache(7) # idempotent assert len(calls) == 1 + + +def test_compile_prewarm_decision_gates(monkeypatch): + # The pure gate: on only for a compiled DEFAULT-tier resident load on a family + # that allows it, with the env kill switch and the cfg-parallel/offload/max + # exclusions each carrying their own resolved reason. + import dataclasses + + from core.inference import video as video_mod + from core.inference.video_families import detect_video_family + + fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + base = dict( + speed_mode = "default", + speed_optims = ("compiled",), + offload_policy = "none", + cfg_parallel_active = False, + ) + + on, reason = video_mod.compile_prewarm_decision(fam, **base) + assert on is True and "absorbs" in reason + + monkeypatch.setenv("UNSLOTH_DIFFUSION_COMPILE_PREWARM", "0") + on, reason = video_mod.compile_prewarm_decision(fam, **base) + assert on is False and "UNSLOTH_DIFFUSION_COMPILE_PREWARM" in reason + monkeypatch.delenv("UNSLOTH_DIFFUSION_COMPILE_PREWARM") + + on, reason = video_mod.compile_prewarm_decision( + fam, **{**base, "speed_optims": ("cudnn_benchmark",)} + ) + assert on is False and "no regional compile" in reason + + on, reason = video_mod.compile_prewarm_decision(fam, **{**base, "speed_mode": "max"}) + assert on is False and "static per-shape" in reason + + opted_out = dataclasses.replace(fam, supports_compile_prewarm = False) + on, reason = video_mod.compile_prewarm_decision(opted_out, **base) + assert on is False and "family opted out" in reason + + on, reason = video_mod.compile_prewarm_decision(fam, **{**base, "offload_policy": "group"}) + assert on is False and "offload" in reason + + on, reason = video_mod.compile_prewarm_decision(fam, **{**base, "cfg_parallel_active": True}) + assert on is False and "CFG parallel" in reason + + +def test_compile_prewarm_runs_after_compiled_load(fake_runtime, monkeypatch): + # A compiled default-tier load must spawn the background prewarm: one tiny + # snapped throwaway generation (192x128, 4k+1 frames, 2 steps) through the + # pipe, no user-visible progress, cancel slot cleared afterwards, and the + # resolved record says why it ran. + from core.inference import video as video_mod + + monkeypatch.setattr(video_mod, "apply_speed_optims", lambda *a, **k: {"compiled": True}) + backend = VideoBackend() + status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + assert status["resolved"]["compile_prewarm"]["value"] == "on" + + thread = backend._prewarm_thread + assert thread is not None + thread.join(timeout = 5) + assert not thread.is_alive() + + call = backend._state.pipe.last_kwargs + assert call is not None, "prewarm never reached the pipe" + assert call["prompt"] == "warmup" + assert call["num_inference_steps"] == 2 + assert (call["width"], call["height"]) == (192, 128) + assert call["num_frames"] == 9 # 4k+1 lattice for Wan's frame_step=4 + # The warmup is invisible: no generation progress, no leaked cancel event. + assert backend._gen.get("active") is False + assert backend._active_generate_cancel is None + backend.unload() + + +def test_compile_prewarm_skipped_without_compile(fake_runtime): + # The fake runtime engages no speed optims, so the load has nothing to warm: + # no thread, no pipe call, and the resolved record carries the reason. + backend = VideoBackend() + status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + assert status["resolved"]["compile_prewarm"]["value"] == "off" + assert "no regional compile" in status["resolved"]["compile_prewarm"]["reason"] + assert backend._prewarm_thread is None + assert backend._state.pipe.last_kwargs is None + + +def test_compile_prewarm_yields_to_generations_and_stale_tokens(fake_runtime): + # The worker must abort without touching the pipe when a real generation got + # in first (it absorbs the warmup itself) or when its load was superseded. + backend = VideoBackend() + backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + + backend._compile_prewarm(backend._load_token + 1) # superseded load + assert backend._state.pipe.last_kwargs is None + + backend._generate_job_active = True + backend._compile_prewarm(backend._load_token) # a request beat the warmup + assert backend._state.pipe.last_kwargs is None + backend._generate_job_active = False