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

@@ -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

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).

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