[perf] Add Adaptive Guidance (CFG gating) for stale-uncond reuse (#1372)

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>

docs/inference/optimizations.md

@@ -11,6 +11,7 @@ This page describes the various options for speeding up generation times in Fast
   - [Sliding Tile Attention (Archived)](#sliding-tile-attention-archived)
   - [Sage Attention](#sage-attention)
   - [Sage Attention 3](#sage-attention-3)
+- [Adaptive Guidance (CFG gating)](#adaptive-guidance-cfg-gating)
 
 - [torch.compile](#torch-compile)
 
@@ -298,3 +299,37 @@ for backend in ["TORCH_SDPA", "FLASH_ATTN", "SAGE_ATTN"]:
 ```
 
 Note: reinstantiate `VideoGenerator` after changing `FASTVIDEO_ATTENTION_BACKEND`.
+
+## Adaptive Guidance (CFG gating)
+
+CFG gating accelerates classifier-free guidance by reusing the cached
+`noise_pred_cond - noise_pred_uncond` delta after a configurable fraction of
+the denoising schedule, skipping the unconditional model forward for the
+remaining steps. The technique is the LinearAG variant of Adaptive Guidance
+(Castillo et al. 2023, [arXiv:2312.12487](https://arxiv.org/abs/2312.12487)).
+
+### Enabling
+
+Set the `FASTVIDEO_CFG_GATE_STEP` environment variable to a float in `[0, 1]`:
+
+| Value | Behavior |
+|-------|----------|
+| `1.0` (default) | Disabled — legacy two-pass CFG every step. |
+| `0.5` | Cache the delta after `len(timesteps) * 0.5` steps; reuse for the rest. |
+| `0.0` | Cache from the very first step (most aggressive). |
+
+```bash
+export FASTVIDEO_CFG_GATE_STEP=0.5
+```
+
+### Trade-offs
+
+- **Memory**: one extra model-output-sized tensor per rank held during the
+  gating window.
+- **Quality**: VBench-measured quality is preserved within noise on 4 of 5
+  dimensions at `FASTVIDEO_CFG_GATE_STEP=0.5` for Wan T2V 1.3B per the PR's
+  reported numbers (see [#1372](https://github.com/hao-ai-lab/FastVideo/pull/1372)).
+- **Speed**: ~22% e2e on 4xL40S and ~24% on 1xH100 at the same settings.
+
+Default behavior is byte-for-byte equivalent to the legacy two-pass CFG path;
+the feature is fully opt-in.

fastvideo/envs.py

@@ -42,6 +42,7 @@ if TYPE_CHECKING:
     FASTVIDEO_TRACE_STEPS: str = ""
     FASTVIDEO_SERVER_DEV_MODE: bool = False
     FASTVIDEO_STAGE_LOGGING: bool = False
+    FASTVIDEO_CFG_GATE_STEP: float = 1.0
     FASTVIDEO_HOST_IP: str = ""
     FASTVIDEO_LOOPBACK_IP: str = ""
 
@@ -283,6 +284,32 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # taken for each stage
     "FASTVIDEO_STAGE_LOGGING":
     lambda: bool(int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0"))),
+
+    # CFG gating fraction for stale-uncond reuse (Adaptive Guidance / LinearAG
+    # variant — Castillo et al. 2023, arXiv:2312.12487).  Float in [0, 1].
+    # Interpretation: for step index `i < len(timesteps) * X`, run both
+    # cond and uncond forwards and refresh delta_cached = cond - uncond.
+    # Once `i >= len(timesteps) * X`, skip the uncond forward and reuse
+    # the cached delta:  noise_pred = cond + (guidance_scale - 1) * delta.
+    #
+    # Edge cases:
+    #   1.0 (default) : disables gating; identical to baseline two-pass CFG.
+    #   0.5           : run uncond for the first half of steps, reuse delta
+    #                    for the second half (~25% inference time saved on
+    #                    bandwidth-bound SP setups).
+    #   0.0           : step 0 still computes uncond fresh (cache is empty
+    #                    at start) — all subsequent steps reuse the step-0
+    #                    delta.  This is the most aggressive setting; does
+    #                    NOT mean "no uncond forward ever."
+    #
+    # Caveats:
+    #   - Algorithmically approximate; not bit-exact vs baseline CFG.
+    #     Validate per-pipeline with SSIM / VBench before lowering below 1.0.
+    #   - Interaction with `guidance_rescale > 0` is unvalidated; the
+    #     denoising stage logs a warning when both are active.
+    #   - Wan2.2 high/low-noise expert switch invalidates the cache.
+    "FASTVIDEO_CFG_GATE_STEP":
+    lambda: float(os.getenv("FASTVIDEO_CFG_GATE_STEP", "1.0")),
 }
 
 # end-env-vars-definition

fastvideo/pipelines/stages/denoising.py

@@ -11,6 +11,7 @@ from typing import Any
 import torch
 from tqdm.auto import tqdm
 
+import fastvideo.envs as envs
 from fastvideo.attention import get_attn_backend
 from fastvideo.distributed import (get_local_torch_device, get_world_group)
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -229,6 +230,49 @@ class DenoisingStage(PipelineStage):
         # written to, so we allocate once.
         v2v_zero_pad = torch.zeros_like(latents) if batch.video_latent is not None else None
 
+        # CFG gating / stale-uncond reuse setup (Adaptive Guidance LinearAG
+        # variant, Castillo et al. 2023).  When envs.FASTVIDEO_CFG_GATE_STEP
+        # < 1.0, the uncond forward is skipped after the gating step and the
+        # guidance delta (cond - uncond) is reused from the last fresh
+        # compute.  See envs.py for semantics.  delta_cached_model_id tracks
+        # which underlying transformer produced the cache so we invalidate on
+        # Wan2.2 expert switch.
+        _cfg_gate_fraction = envs.FASTVIDEO_CFG_GATE_STEP
+        if not 0.0 <= _cfg_gate_fraction <= 1.0:
+            raise ValueError(f"FASTVIDEO_CFG_GATE_STEP must be in [0.0, 1.0], got {_cfg_gate_fraction!r}. "
+                             "Use 1.0 (default) to disable; lower values trade quality for speed.")
+        _cfg_gate_active = _cfg_gate_fraction < 1.0 and batch.do_classifier_free_guidance
+        _is_rank0 = get_world_group().local_rank == 0
+        if _cfg_gate_active:
+            # Use len(timesteps), not num_inference_steps: the loop iterates
+            # over timesteps directly, and for schedulers with order > 1
+            # (e.g. DPM-Solver++ 2M, Heun) len(timesteps) is a multiple of
+            # num_inference_steps. Using num_inference_steps would cause the
+            # gate to fire at fraction/order of the loop instead of fraction.
+            _cfg_gate_step_idx = int(len(timesteps) * _cfg_gate_fraction)
+            if _is_rank0:
+                logger.info("CFG gating enabled: fraction=%.3f, gate_step=%d/%d", _cfg_gate_fraction,
+                            _cfg_gate_step_idx, len(timesteps))
+            if batch.guidance_rescale > 0.0 and _is_rank0:
+                # guidance_rescale rescales CFG output stats to match cond
+                # stats (Lin et al. §3.4).  When `delta_cached` goes stale,
+                # the rescaling still computes but is no longer guaranteed
+                # to preserve the original quality semantics.  Warn so the
+                # caller knows this combo is unvalidated; tighten or
+                # fallback once VBench data lands.
+                logger.warning(
+                    "CFG gating (fraction=%.3f) combined with guidance_rescale=%.3f is unvalidated; "
+                    "quality may degrade beyond CFG-gating-alone expectations.", _cfg_gate_fraction,
+                    batch.guidance_rescale)
+        else:
+            _cfg_gate_step_idx = len(timesteps) + 1  # never gates
+        delta_cached: torch.Tensor | None = None
+        delta_cached_model_id: int | None = None
+        # Telemetry — logged at end of denoising loop on rank 0.
+        _cfg_gate_fresh_uncond = 0
+        _cfg_gate_reused_delta = 0
+        _cfg_gate_invalidations = 0
+
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -362,26 +406,58 @@ class DenoisingStage(PipelineStage):
                         )
 
                     if batch.do_classifier_free_guidance:
-                        batch.is_cfg_negative = True
-                        with set_forward_context(
-                                current_timestep=i,
-                                attn_metadata=attn_metadata,
-                                forward_batch=batch,
-                        ):
-                            noise_pred_uncond = current_model(
-                                latent_model_input,
-                                neg_prompt_embeds,
-                                t_expand,
-                                guidance=guidance_expand,
-                                **image_kwargs,
-                                **neg_cond_kwargs,
-                                **action_kwargs,
-                                **camera_kwargs,
-                                **timesteps_r_kwarg,
-                            )
+                        # CFG gating: invalidate cached delta when the underlying
+                        # transformer changes (Wan2.2 high/low-noise expert
+                        # switch at `boundary_timestep`).  delta_cached is tied
+                        # to the model that produced it; reusing it across the
+                        # boundary is silently wrong.
+                        if delta_cached_model_id is not None and delta_cached_model_id != id(current_model):
+                            delta_cached = None
+                            delta_cached_model_id = None
+                            _cfg_gate_invalidations += 1
+
+                        _use_cached_delta = (i >= _cfg_gate_step_idx and delta_cached is not None)
 
                         noise_pred_text = noise_pred
-                        noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_text - noise_pred_uncond)
+                        if _use_cached_delta:
+                            # Reuse frozen delta = cond - uncond from the last
+                            # fresh compute.  Algebra:
+                            #   pred = uncond + s * (cond - uncond)
+                            #        = cond + (s - 1) * (cond - uncond)
+                            #        = cond + (s - 1) * delta_cached
+                            noise_pred = noise_pred_text + (current_guidance_scale - 1.0) * delta_cached
+                            _cfg_gate_reused_delta += 1
+                        else:
+                            batch.is_cfg_negative = True
+                            with set_forward_context(
+                                    current_timestep=i,
+                                    attn_metadata=attn_metadata,
+                                    forward_batch=batch,
+                            ):
+                                noise_pred_uncond = current_model(
+                                    latent_model_input,
+                                    neg_prompt_embeds,
+                                    t_expand,
+                                    guidance=guidance_expand,
+                                    **image_kwargs,
+                                    **neg_cond_kwargs,
+                                    **action_kwargs,
+                                    **camera_kwargs,
+                                    **timesteps_r_kwarg,
+                                )
+                            _cfg_gate_fresh_uncond += 1
+
+                            # Refresh cache only when gating is active; under the
+                            # default (FASTVIDEO_CFG_GATE_STEP=1.0, _cfg_gate_step_idx
+                            # > len(timesteps)) we never reuse, so skip the
+                            # tensor allocation.
+                            if _cfg_gate_step_idx <= len(timesteps):
+                                delta_cached = noise_pred_text - noise_pred_uncond
+                                delta_cached_model_id = id(current_model)
+                                noise_pred = noise_pred_uncond + current_guidance_scale * delta_cached
+                            else:
+                                noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_text -
+                                                                                           noise_pred_uncond)
 
                         # Apply guidance rescale if needed
                         if batch.guidance_rescale > 0.0:
@@ -408,6 +484,22 @@ class DenoisingStage(PipelineStage):
                                                (i + 1) % self.scheduler.order == 0 and progress_bar is not None):
                     progress_bar.update()
 
+        # CFG gating telemetry — log once on rank 0 after the loop ends.  When
+        # gating is disabled (or CFG itself is off) fresh_uncond equals the
+        # number of CFG-on steps and reused/invalidations are zero; we still
+        # emit the line so users can confirm the env var is wired through.
+        if _is_rank0 and batch.do_classifier_free_guidance:
+            logger.info(
+                "CFG gating summary: fraction=%.3f gate_step=%d/%d "
+                "fresh_uncond=%d reused=%d invalidations=%d",
+                _cfg_gate_fraction,
+                _cfg_gate_step_idx if _cfg_gate_active else -1,
+                len(timesteps),
+                _cfg_gate_fresh_uncond,
+                _cfg_gate_reused_delta,
+                _cfg_gate_invalidations,
+            )
+
         trajectory_tensor: torch.Tensor | None = None
         if trajectory_latents:
             trajectory_tensor = torch.stack(trajectory_latents, dim=1)

fastvideo/tests/inference/test_cfg_gating.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+class RecordingLogger:
+    def __init__(self):
+        self.infos = []
+        self.warnings = []
+
+    def info(self, msg, *args):
+        self.infos.append(msg % args if args else msg)
+
+    def warning(self, msg, *args):
+        self.warnings.append(msg % args if args else msg)
+
+
+class NullProgressBar:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+    def update(self):
+        pass
+
+
+class TinyScheduler:
+    order = 1
+    num_train_timesteps = 1000
+
+    def scale_model_input(self, latents, timestep):
+        return latents
+
+    def step(self, noise_pred, timestep, latents, return_dict=False):
+        return (latents.float() - 0.01 * noise_pred.float(), )
+
+
+class TinyDenoiser(torch.nn.Module):
+    hidden_size = 1
+    num_attention_heads = 1
+
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(use_meanflow=False)
+        self.calls = []
+
+    def forward(self, latent_model_input, prompt_embeds, timestep, guidance=None):
+        prompt_value = prompt_embeds[0].to(latent_model_input.device, torch.float32)
+        is_uncond = bool(prompt_value.item() < 0)
+        self.calls.append("uncond" if is_uncond else "cond")
+
+        prompt_term = prompt_value.reshape((1, ) * latent_model_input.ndim)
+        timestep_term = timestep.to(latent_model_input.device, torch.float32).reshape(
+            latent_model_input.shape[0], *([1] * (latent_model_input.ndim - 1)))
+        return latent_model_input.float() * 0.2 + prompt_term * 0.5 + timestep_term * 0.001
+
+
+def _tiny_args():
+    return SimpleNamespace(
+        disable_autocast=True,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        use_fsdp_inference=False,
+        moba_config={},
+        VSA_sparsity=0.0,
+        model_loaded={"transformer": True},
+        model_paths={"transformer": "unused"},
+        pipeline_config=SimpleNamespace(
+            embedded_cfg_scale=None,
+            ti2v_task=False,
+            dit_config=SimpleNamespace(boundary_ratio=None, patch_size=(1, 1, 1)),
+        ),
+    )
+
+
+def _tiny_batch():
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    return ForwardBatch(
+        data_type="video",
+        latents=torch.tensor([[[[[0.125]]]]], dtype=torch.float32),
+        prompt_embeds=[torch.tensor([1.0], dtype=torch.float32)],
+        negative_prompt_embeds=[torch.tensor([-1.0], dtype=torch.float32)],
+        timesteps=torch.tensor([4.0, 3.0, 2.0, 1.0], dtype=torch.float32),
+        num_inference_steps=4,
+        guidance_scale=2.0,
+        height=1,
+        width=1,
+        num_frames=1,
+        raw_latent_shape=(1, 1, 1, 1, 1),
+        save_video=False,
+    )
+
+
+def _patch_denoising_module(monkeypatch, cfg_gate_step):
+    if cfg_gate_step is None:
+        monkeypatch.delenv("FASTVIDEO_CFG_GATE_STEP", raising=False)
+        expected_gate_step = 1.0
+    else:
+        monkeypatch.setenv("FASTVIDEO_CFG_GATE_STEP", str(cfg_gate_step))
+        expected_gate_step = float(cfg_gate_step)
+
+    import fastvideo.pipelines.stages.denoising as denoising
+
+    # envs.py evaluates FASTVIDEO_CFG_GATE_STEP lazily via __getattr__, so the
+    # stage sees monkeypatched values without reloading the module.
+    assert denoising.envs.FASTVIDEO_CFG_GATE_STEP == expected_gate_step
+
+    logger = RecordingLogger()
+    monkeypatch.setattr(denoising, "logger", logger)
+    monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(denoising, "get_world_group", lambda: SimpleNamespace(local_rank=0))
+    monkeypatch.setattr(denoising, "get_attn_backend", lambda **kwargs: object())
+    monkeypatch.setattr(denoising, "set_forward_context", lambda **kwargs: nullcontext())
+    return denoising, logger
+
+
+def _run_stage(monkeypatch, cfg_gate_step):
+    denoising, logger = _patch_denoising_module(monkeypatch, cfg_gate_step)
+    model = TinyDenoiser()
+    stage = denoising.DenoisingStage(model, TinyScheduler())
+    stage.progress_bar = lambda iterable=None, total=None: NullProgressBar()
+
+    result = stage.forward(_tiny_batch(), _tiny_args())
+    return result.latents, model, logger
+
+
+def _run_legacy_two_pass():
+    batch = _tiny_batch()
+    model = TinyDenoiser()
+    scheduler = TinyScheduler()
+    assert batch.latents is not None
+    assert batch.timesteps is not None
+    latents = batch.latents.clone()
+
+    for timestep in batch.timesteps:
+        latent_model_input = scheduler.scale_model_input(latents.to(torch.bfloat16), timestep)
+        timestep_expand = timestep.repeat(latent_model_input.shape[0])
+        noise_pred_text = model(latent_model_input, batch.prompt_embeds, timestep_expand)
+        noise_pred_uncond = model(latent_model_input, batch.negative_prompt_embeds, timestep_expand)
+        noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred_text - noise_pred_uncond)
+        latents = scheduler.step(noise_pred, timestep, latents, return_dict=False)[0]
+
+    return latents
+
+
+@pytest.mark.parametrize("cfg_gate_step", [None, "1.0"])
+def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, cfg_gate_step):
+    out, model, logger = _run_stage(monkeypatch, cfg_gate_step)
+    legacy_out = _run_legacy_two_pass()
+
+    assert torch.equal(out, legacy_out)
+    assert model.calls == ["cond", "uncond"] * 4
+    assert not any("CFG gating enabled" in msg for msg in logger.infos)
+    assert any("gate_step=-1/4" in msg and "reused=0" in msg for msg in logger.infos)
+
+
+def test_cfg_gating_reuses_cached_delta_after_gate(monkeypatch):
+    out, model, logger = _run_stage(monkeypatch, "0.5")
+    legacy_out = _run_legacy_two_pass()
+
+    assert model.calls == ["cond", "uncond", "cond", "uncond", "cond", "cond"]
+    assert any("CFG gating enabled: fraction=0.500, gate_step=2/4" in msg for msg in logger.infos)
+    assert any("fresh_uncond=2 reused=2 invalidations=0" in msg for msg in logger.infos)
+    assert torch.allclose(out, legacy_out, atol=1e-3, rtol=0.0)