[feat] Optional cache-dit step caching for the Wan DiT

docs/design/inference_schema_parity_inventory.yaml

@@ -101,6 +101,13 @@ surfaces:
       moba_config: "Derived runtime config loaded from moba_config_path."
       model_paths: "Runtime bookkeeping."
       model_loaded: "Runtime bookkeeping."
+      use_cachedit: "Opt-in cache-dit step-caching toggle (Wan DiT); runtime plumbing, not part of the typed API yet."
+      cachedit_fn_compute_blocks: "cache-dit leading-block count; runtime plumbing for use_cachedit."
+      cachedit_bn_compute_blocks: "cache-dit trailing-block count; runtime plumbing for use_cachedit."
+      cachedit_residual_threshold: "cache-dit residual-diff skip threshold; runtime plumbing for use_cachedit."
+      cachedit_max_warmup_steps: "cache-dit warmup-step count; runtime plumbing for use_cachedit."
+      cachedit_taylorseer: "cache-dit TaylorSeer calibrator toggle; runtime plumbing for use_cachedit."
+      cachedit_taylorseer_order: "cache-dit TaylorSeer expansion order; runtime plumbing for use_cachedit."
 
   pipeline_config_base:
     moved:

fastvideo/fastvideo_args.py

@@ -159,6 +159,28 @@ class FastVideoArgs:
     vae_parallel_encode: bool = False
     vae_parallel_decode_strategy: str | None = None
 
+    # Step caching via cache-dit (https://github.com/vipshop/cache-dit).
+    # LOSSY: skips DiT blocks on steps whose features barely change, so the
+    # output is NOT bit-identical (SSIM<1.0). Opt-in, default OFF, Wan DiT
+    # only for now. When on, the first ``cachedit_fn_compute_blocks`` blocks
+    # always run and produce an L1 "stable" residual; if its relative change
+    # from the previous step is below ``cachedit_residual_threshold`` the
+    # middle blocks are skipped and a cached residual reused; the last
+    # ``cachedit_bn_compute_blocks`` blocks always run to refine. No caching
+    # during the first ``cachedit_max_warmup_steps`` steps. ``cachedit_
+    # taylorseer`` swaps the constant-residual reuse for a Taylor-expansion
+    # extrapolation of the residual (higher fidelity at the same skip rate).
+    # Requires ``pip install cache-dit`` and is incompatible with DiT
+    # offloading (see DenoisingStage — caching skips blocks, offload assumes
+    # every block runs each step).
+    use_cachedit: bool = False
+    cachedit_fn_compute_blocks: int = 8
+    cachedit_bn_compute_blocks: int = 0
+    cachedit_residual_threshold: float = 0.08
+    cachedit_max_warmup_steps: int = 8
+    cachedit_taylorseer: bool = False
+    cachedit_taylorseer_order: int = 1
+
     # Compilation
     # ``enable_torch_compile`` covers the DiT path (transformer,
     # transformer_2, and the LTX-2 stage-2 transformer_refine).
@@ -711,6 +733,47 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             help="Disable autocast for denoising loop and vae decoding in pipeline sampling",
         )
 
+        # cache-dit step caching (lossy; Wan DiT). Requires `pip install
+        # cache-dit` and is incompatible with DiT offloading.
+        parser.add_argument(
+            "--use-cachedit",
+            action=StoreBoolean,
+            help="Enable cache-dit step caching for the Wan DiT (lossy; skips DiT blocks on steps whose features "
+            "barely change). Requires `pip install cache-dit`; incompatible with DiT offloading.",
+        )
+        parser.add_argument(
+            "--cachedit-fn-compute-blocks",
+            type=int,
+            help="cache-dit: number of leading DiT blocks always computed (default 8).",
+        )
+        parser.add_argument(
+            "--cachedit-bn-compute-blocks",
+            type=int,
+            help="cache-dit: number of trailing DiT blocks always computed to refine (default 0).",
+        )
+        parser.add_argument(
+            "--cachedit-residual-threshold",
+            type=float,
+            help="cache-dit: relative L1 residual-diff threshold below which middle blocks are skipped (default 0.08; "
+            "higher = faster, lower quality).",
+        )
+        parser.add_argument(
+            "--cachedit-max-warmup-steps",
+            type=int,
+            help="cache-dit: number of initial steps that always compute every block (default 8).",
+        )
+        parser.add_argument(
+            "--cachedit-taylorseer",
+            action=StoreBoolean,
+            help="cache-dit: use a TaylorSeer calibrator (extrapolates the cached residual instead of holding it "
+            "constant — higher fidelity at the same skip rate).",
+        )
+        parser.add_argument(
+            "--cachedit-taylorseer-order",
+            type=int,
+            help="cache-dit: TaylorSeer expansion order / number of derivatives (default 1; 2 = quadratic).",
+        )
+
         # VSA parameters
         parser.add_argument(
             "--VSA-sparsity",

fastvideo/pipelines/stages/denoising.py

@@ -76,6 +76,55 @@ def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, va
             requested=component_attention_backend(self.transformer),
         )
 
+    def _enable_or_refresh_cachedit(self, model, fastvideo_args, num_inference_steps) -> None:
+        """Wire ``model`` into cache-dit via a transformer-only BlockAdapter on
+        first use, then refresh the cache context each generation. cache-dit is
+        lazy-imported so it stays an optional dependency.
+
+        Wan runs cond + uncond as separate forwards (``enable_separate_cfg=True``),
+        cond first (``cfg_compute_first=False``). ``num_inference_steps`` lets
+        cache-dit auto-refresh at the generation boundary; we also call
+        ``refresh_context`` explicitly per generation so cache state never leaks
+        across prompts (the bare-transformer path has no pipeline call to reset
+        it otherwise).
+        """
+        try:
+            import cache_dit
+            from cache_dit import (BlockAdapter, DBCacheConfig, ForwardPattern, TaylorSeerCalibratorConfig)
+        except ImportError as e:
+            raise ImportError("use_cachedit requires the cache-dit package, which is not installed. Install it with "
+                              "`pip install \"fastvideo[cache]\"` (or `pip install cache-dit`).") from e
+
+        cache_config = DBCacheConfig(
+            Fn_compute_blocks=fastvideo_args.cachedit_fn_compute_blocks,
+            Bn_compute_blocks=fastvideo_args.cachedit_bn_compute_blocks,
+            residual_diff_threshold=fastvideo_args.cachedit_residual_threshold,
+            max_warmup_steps=fastvideo_args.cachedit_max_warmup_steps,
+            enable_separate_cfg=True,
+            cfg_compute_first=False,
+            num_inference_steps=num_inference_steps,
+        )
+        calibrator_config = None
+        if fastvideo_args.cachedit_taylorseer:
+            calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=fastvideo_args.cachedit_taylorseer_order)
+
+        if not getattr(model, "_cachedit_enabled", False):
+            adapter = BlockAdapter(
+                transformer=model,
+                blocks=model.blocks,
+                forward_pattern=ForwardPattern.Pattern_2,
+                has_separate_cfg=True,
+                check_forward_pattern=False,
+            )
+            cache_dit.enable_cache(adapter, cache_config=cache_config, calibrator_config=calibrator_config)
+            model._cachedit_enabled = True
+            logger.info("cache-dit enabled: Fn=%d Bn=%d threshold=%s warmup=%d taylorseer=%s",
+                        fastvideo_args.cachedit_fn_compute_blocks, fastvideo_args.cachedit_bn_compute_blocks,
+                        fastvideo_args.cachedit_residual_threshold, fastvideo_args.cachedit_max_warmup_steps,
+                        fastvideo_args.cachedit_taylorseer)
+        else:
+            cache_dit.refresh_context(model, num_inference_steps=num_inference_steps)
+
     def forward(
         self,
         batch: ForwardBatch,
@@ -367,6 +416,37 @@ def forward(
         _cfg_gate_reused_delta = 0
         _cfg_gate_invalidations = 0
 
+        # cache-dit step caching skips DiT blocks, which is incompatible with
+        # layerwise / CPU offload: the offload hook prefetches each block's
+        # params on the prior block's forward and releases them on its own,
+        # assuming every block runs exactly once per step. A skipped block
+        # leaves its params prefetched-but-never-released, desyncing the chain.
+        if fastvideo_args.use_cachedit and (fastvideo_args.dit_layerwise_offload or fastvideo_args.dit_cpu_offload):
+            raise ValueError("use_cachedit is incompatible with DiT offloading: caching skips "
+                             "blocks, but the layerwise/CPU offload hook assumes every block "
+                             "runs each step. Set dit_layerwise_offload=False and "
+                             "dit_cpu_offload=False (the model must fit in GPU memory).")
+        # cache-dit skips blocks via data-dependent control flow (a per-step
+        # residual-diff decision), which torch.compile cannot trace without
+        # graph breaks/recompiles. Disallow the combination for now (eager
+        # only); compile support is a follow-up. ``inference_torch_compile``
+        # (regional fullgraph compile of each DiT block) is even stricter —
+        # the loader compiles the blocks with fullgraph=True before this stage
+        # ever runs, so the skip decision cannot graph-break out.
+        if fastvideo_args.use_cachedit and (fastvideo_args.enable_torch_compile
+                                            or fastvideo_args.inference_torch_compile):
+            _flag = ("enable_torch_compile" if fastvideo_args.enable_torch_compile else "inference_torch_compile")
+            raise ValueError(f"use_cachedit is currently incompatible with {_flag}: cache-dit "
+                             "introduces data-dependent control flow that torch.compile cannot trace cleanly. "
+                             f"Set {_flag}=False (eager); compile support is a follow-up.")
+        # Enable cache-dit on the transformer(s) once, then refresh the cache
+        # context each generation so state never leaks across prompts.
+        if fastvideo_args.use_cachedit:
+            for _tf in (self.transformer, self.transformer_2):
+                _model = getattr(_tf, "module", _tf)
+                if _model is not None and hasattr(_model, "blocks"):
+                    self._enable_or_refresh_cachedit(_model, fastvideo_args, num_inference_steps)
+
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):

fastvideo/tests/modal/_cachedit_ab_inner.py

@@ -0,0 +1,183 @@
+"""
+Inner-script half of the cache-dit A/B harness.
+
+Runs one denoising pass (caching off or on) over the harness's prompt + seed
+set, records per-prompt walls, and writes a JSON results blob to the path
+given on argv. Invoked by ``cachedit_ab.py`` via ``/opt/venv/bin/python`` so
+FastVideo runs inside the image's venv (Modal's main process Python lacks
+torch).
+
+cache-dit step caching is LOSSY: the patched output is not bit-identical, so
+the SSIM column is a quality measurement (target >= ~0.95), not a 1.0 gate.
+"""
+import argparse
+import json
+import os
+import sys
+import time
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=("run_pass", "compute_ssim"), default="run_pass")
+    parser.add_argument("--config-json")
+    parser.add_argument("--results-json")
+    parser.add_argument("--baseline-results-json")
+    parser.add_argument("--patched-results-json")
+    parser.add_argument("--ssim-output-json")
+    args = parser.parse_args()
+
+    if args.mode == "compute_ssim":
+        return _compute_ssim_main(args)
+
+    if not args.config_json or not args.results_json:
+        parser.error("--config-json and --results-json are required for mode=run_pass")
+    with open(args.config_json) as f:
+        cfg = json.load(f)
+
+    model_id = cfg["model_id"]
+    num_gpus = cfg["num_gpus"]
+    use_cachedit = cfg["use_cachedit"]
+    fn = cfg.get("cachedit_fn_compute_blocks", 8)
+    bn = cfg.get("cachedit_bn_compute_blocks", 0)
+    threshold = cfg.get("cachedit_residual_threshold", 0.08)
+    warmup = cfg.get("cachedit_max_warmup_steps", 8)
+    taylorseer = cfg.get("cachedit_taylorseer", False)
+    taylorseer_order = cfg.get("cachedit_taylorseer_order", 1)
+    enable_compile = cfg.get("enable_compile", False)
+    height, width = cfg["height"], cfg["width"]
+    num_frames = cfg["num_frames"]
+    num_inference_steps = cfg["num_inference_steps"]
+    output_dir = cfg["output_dir"]
+    prompts = cfg["prompts"]
+    seed_base = cfg["seed_base"]
+
+    from fastvideo.attention.backends.flash_attn import fa_version
+    print(f"[inner] resolved flash-attn version: FA{fa_version}", flush=True)
+
+    from fastvideo import VideoGenerator
+
+    # cache-dit skips blocks, incompatible with layerwise/CPU offload (the
+    # offload prefetch chain assumes every block runs each step). Disable it on
+    # BOTH passes so the A/B isolates the cache effect. Wan 1.3B fits a single
+    # L40S (48GB) without offload.
+    generator = VideoGenerator.from_pretrained(
+        model_id,
+        num_gpus=num_gpus,
+        use_cachedit=use_cachedit,
+        cachedit_fn_compute_blocks=fn,
+        cachedit_bn_compute_blocks=bn,
+        cachedit_residual_threshold=threshold,
+        cachedit_max_warmup_steps=warmup,
+        cachedit_taylorseer=taylorseer,
+        cachedit_taylorseer_order=taylorseer_order,
+        dit_layerwise_offload=False,
+        dit_cpu_offload=False,
+        enable_torch_compile=enable_compile,
+    )
+    resolved = getattr(generator.fastvideo_args, "use_cachedit", None)
+    if resolved is not use_cachedit:
+        raise RuntimeError(f"use_cachedit did not propagate: requested {use_cachedit}, resolved {resolved}")
+    if use_cachedit:
+        print(f"[inner] cache-dit ON: Fn={fn} Bn={bn} threshold={threshold} warmup={warmup} "
+              f"taylorseer={taylorseer}(order={taylorseer_order})", flush=True)
+    else:
+        print("[inner] caching OFF (baseline)", flush=True)
+
+    os.makedirs(output_dir, exist_ok=True)
+    records = []
+    for i, prompt in enumerate(prompts):
+        prompt_out = os.path.join(output_dir, f"prompt_{i:02d}")
+        t0 = time.perf_counter()
+        generator.generate_video(
+            prompt,
+            output_path=prompt_out,
+            save_video=True,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            seed=seed_base + i,
+        )
+        wall = time.perf_counter() - t0
+        # generate_video appends _1, _2, ... rather than overwriting; pick the
+        # mp4 we just wrote by mtime (newest).
+        mp4s = sorted([os.path.join(prompt_out, f) for f in os.listdir(prompt_out) if f.endswith(".mp4")],
+                      key=os.path.getmtime)
+        if not mp4s:
+            raise RuntimeError(f"no .mp4 produced for prompt {i} at {prompt_out}")
+        records.append({"i": i, "wall_s": wall, "mp4": mp4s[-1]})
+        print(f"[inner] [{'cachedit' if use_cachedit else 'baseline'}] prompt {i}: {wall:.3f}s -> {mp4s[-1]}",
+              flush=True)
+
+    with open(args.results_json, "w") as f:
+        json.dump(records, f)
+    print(f"[inner] wrote {len(records)} records to {args.results_json}", flush=True)
+    return 0
+
+
+def _compute_ssim_main(args) -> int:
+    """Pairwise SSIM between baseline and patched output mp4s. Avoids importing
+    fastvideo (its top-level import pulls triton, which needs a CUDA driver);
+    runs on CPU with pytorch_msssim + torchvision/av directly."""
+    if not args.ssim_output_json or not (args.baseline_results_json and args.patched_results_json):
+        raise SystemExit("compute_ssim needs --baseline-results-json, --patched-results-json, --ssim-output-json")
+
+    import torch
+    from pytorch_msssim import ssim as pm_ssim
+
+    def _read_video_frames(path):
+        try:
+            from torchvision.io import read_video
+            frames, _, _ = read_video(path, pts_unit="sec", output_format="TCHW")
+            if frames.shape[0] > 0:
+                return frames
+        except Exception:
+            # torchvision's backend (FFmpeg/PyAV) can raise more than
+            # ImportError/AttributeError; fall back to the PyAV path below.
+            pass
+        import av
+        container = av.open(path)
+        frames = []
+        for frame in container.decode(video=0):
+            frames.append(torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1))
+        container.close()
+        if not frames:
+            raise RuntimeError(f"No video frames decoded from {path}")
+        return torch.stack(frames)
+
+    def _ssim(p1, p2):
+        f1, f2 = _read_video_frames(p1), _read_video_frames(p2)
+        n = min(f1.shape[0], f2.shape[0])
+        if n == 0:
+            raise RuntimeError(f"no decodable frames to compare: {p1} ({f1.shape[0]}) vs {p2} ({f2.shape[0]})")
+        f1 = (f1[:n].float() / 255.0).contiguous()
+        f2 = (f2[:n].float() / 255.0).contiguous()
+        return [pm_ssim(f1[i:i + 1], f2[i:i + 1], data_range=1.0).item() for i in range(n)]
+
+    with open(args.baseline_results_json) as f:
+        baseline = json.load(f)
+    with open(args.patched_results_json) as f:
+        patched = json.load(f)
+
+    rows = []
+    for b, p in zip(baseline, patched, strict=True):
+        if b["i"] != p["i"]:
+            raise ValueError(f"baseline/patched prompt index mismatch: {b['i']} != {p['i']}")
+        vals = _ssim(b["mp4"], p["mp4"])
+        rows.append({
+            "i": b["i"],
+            "baseline_wall_s": b["wall_s"],
+            "patched_wall_s": p["wall_s"],
+            "ssim_mean": float(sum(vals) / len(vals)),
+            "ssim_worst": float(min(vals)),
+        })
+        print(f"[inner]   prompt {b['i']} SSIM mean={rows[-1]['ssim_mean']:.6f} worst={rows[-1]['ssim_worst']:.6f}",
+              flush=True)
+    with open(args.ssim_output_json, "w") as f:
+        json.dump(rows, f)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())