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

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

Author: rich7420

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 @@
     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 @@ def maybe_convert_int(value: str | None) -> int | None:
     # 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 @@
 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 @@ def forward(
         # 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 @@ def forward(
                         )
 
                     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 @@ def forward(
                                                (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)