Add approximate H3 cache modes

README.md

@@ -64,3 +64,9 @@ Exact memory/lifetime options:
 
 - `attention: "kj_head_sliced"` slices attention heads and runs the slice backend from `H3_HEAD_SLICE_BACKEND` (`sage2` by default) with `H3_HEAD_SLICE_SIZE` heads per slice (`8` by default).
 - `--mlp-chunks N` on `tools/serve_hot_runtime.py` or `tools/direct_t2v_preview.py` chunks H3 SwiGLU rows exactly to reduce peak activation memory. Default is `1` (disabled).
+
+Approximate cache options are opt-in and must be quality-gated per prompt:
+
+- `cache_mode: "easycache"` reuses cached denoised deltas while cumulative latent input change stays below `cache_threshold`.
+- `cache_mode: "h3_cache"` reuses cached denoised deltas when the current per-step latent input change is below `cache_threshold`.
+- Both modes accept `cache_start_percent`, `cache_end_percent`, and `cache_subsample_factor` in `POST /generate`; the CLI exposes equivalent `--cache-*` flags.

src/h3_blackwell_runtime/runtime.py

@@ -141,8 +141,14 @@ class H3HotRuntime:
         mux_audio: bool = True,
         ffmpeg_loglevel: str = "error",
         save_latent: str | Path | None = None,
+        cache_mode: str | None = None,
+        cache_threshold: float = 0.0,
+        cache_start_percent: float = 0.0,
+        cache_end_percent: float = 1.0,
+        cache_subsample_factor: int = 2,
     ) -> dict:
         stages: list[dict] = []
+        cache_stats: dict = {}
 
         def timed(stage: str, fn):
             _sync()
@@ -172,6 +178,12 @@ class H3HotRuntime:
                 audio,
                 steps=steps,
                 return_audio=mux_audio,
+                cache_mode=cache_mode,
+                cache_threshold=cache_threshold,
+                cache_start_percent=cache_start_percent,
+                cache_end_percent=cache_end_percent,
+                cache_subsample_factor=cache_subsample_factor,
+                cache_stats=cache_stats,
             ),
         )
         if mux_audio:
@@ -241,5 +253,6 @@ class H3HotRuntime:
             "vae_dtype": self.config.vae_dtype,
             "vae_tile_size": self.config.vae_tile_size,
             "stages": stages,
+            "cache": cache_stats,
             "request_seconds": sum(stage["seconds"] for stage in stages),
         }

src/h3_blackwell_runtime/sampler.py

@@ -57,6 +57,14 @@ def _decode_audio_latent(audio_carried: torch.Tensor, *, shift_video: float = 12
     return audio_carried * (shift_audio / shift_video)
 
 
+def _cache_sample(x: torch.Tensor, factor: int) -> torch.Tensor:
+    if factor <= 1:
+        return x
+    if x.ndim == 5:
+        return x[..., ::factor, ::factor]
+    return x[..., ::factor]
+
+
 @torch.inference_mode()
 def sample_video_res_multistep(
     model,
@@ -69,6 +77,12 @@ def sample_video_res_multistep(
     model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
     return_audio: bool = False,
     progress: bool = False,
+    cache_mode: str | None = None,
+    cache_threshold: float = 0.0,
+    cache_start_percent: float = 0.0,
+    cache_end_percent: float = 1.0,
+    cache_subsample_factor: int = 2,
+    cache_stats: dict | None = None,
 ) -> torch.Tensor:
     """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
     sigmas = beta_sigmas(steps, device=video.device)
@@ -77,10 +91,47 @@ def sample_video_res_multistep(
     video_history_sigma = audio_history_sigma = None
     total_steps = len(sigmas) - 1
     started = time.perf_counter()
+    cache_mode = None if cache_mode in {None, "", "disabled", "none"} else cache_mode
+    if cache_mode not in {None, "easycache", "h3_cache"}:
+        raise ValueError(f"Unsupported cache mode: {cache_mode}")
+    if cache_stats is not None:
+        cache_stats.update({"mode": cache_mode, "threshold": cache_threshold, "skipped_steps": 0, "rates": []})
+    cache = {
+        "video_diff": None,
+        "audio_diff": None,
+        "video_prev": None,
+        "audio_prev": None,
+        "prev_norm": None,
+        "cumulative_rate": 0.0,
+    }
     for index, sigma in enumerate(sigmas[:-1], start=1):
         step_started = time.perf_counter()
         previous_index = index - 1
         sigma_down = sigmas[index]
+        current_percent = previous_index / total_steps
+        can_cache = cache_mode is not None and cache_threshold > 0 and cache_start_percent <= current_percent <= cache_end_percent and cache["video_diff"] is not None
+        skipped = False
+        if can_cache:
+            video_now = _cache_sample(video, cache_subsample_factor)
+            audio_now = _cache_sample(audio_carried, cache_subsample_factor)
+            input_change = (video_now - cache["video_prev"]).flatten().abs().mean() + (audio_now - cache["audio_prev"]).flatten().abs().mean()
+            input_norm = cache["prev_norm"].clamp_min(1e-8)
+            rate = (input_change / input_norm).item()
+            if cache_mode == "easycache":
+                cache["cumulative_rate"] += rate
+                skipped = cache["cumulative_rate"] < cache_threshold
+                if not skipped:
+                    cache["cumulative_rate"] = 0.0
+            else:
+                skipped = rate < cache_threshold
+            if cache_stats is not None:
+                cache_stats["rates"].append({"step": previous_index, "rate": rate, "skipped": skipped})
+        if skipped:
+            video_denoised = video + cache["video_diff"]
+            audio_denoised = audio_carried + cache["audio_diff"]
+            if cache_stats is not None:
+                cache_stats["skipped_steps"] += 1
+        else:
             sigma_audio = _audio_sigma(sigma)
             carry = sigma_audio / sigma
             native_audio = audio_carried.to(torch.bfloat16) * carry
@@ -96,6 +147,12 @@ def sample_video_res_multistep(
             ).float()
             video_denoised = video - sigma * velocity_video
             audio_denoised = audio_carried - sigma * velocity_audio
+            if cache_mode is not None:
+                cache["video_diff"] = (video_denoised - video).detach()
+                cache["audio_diff"] = (audio_denoised - audio_carried).detach()
+                cache["video_prev"] = _cache_sample(video, cache_subsample_factor).detach().clone()
+                cache["audio_prev"] = _cache_sample(audio_carried, cache_subsample_factor).detach().clone()
+                cache["prev_norm"] = video.flatten().abs().mean() + audio_carried.flatten().abs().mean()
         previous_sigma = sigmas[previous_index - 1] if previous_index else None
         video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
         audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)

tools/direct_t2v_preview.py

@@ -48,6 +48,11 @@ parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), d
 parser.add_argument("--vae-tile-size", type=int, default=int(os.getenv("H3_VAE_TILE_SIZE", "256")))
 parser.add_argument("--mlp-chunks", type=int, default=int(os.getenv("H3_MLP_CHUNKS", "1")))
 parser.add_argument("--mlp-chunk-threshold", type=int, default=int(os.getenv("H3_MLP_CHUNK_THRESHOLD", "4096")))
+parser.add_argument("--cache-mode", choices=("disabled", "easycache", "h3_cache"), default="disabled")
+parser.add_argument("--cache-threshold", type=float, default=0.0)
+parser.add_argument("--cache-start-percent", type=float, default=0.0)
+parser.add_argument("--cache-end-percent", type=float, default=1.0)
+parser.add_argument("--cache-subsample-factor", type=int, default=2)
 args = parser.parse_args()
 started = time.perf_counter()
 last_report = started
@@ -101,7 +106,26 @@ if args.model_timesteps_capture is not None:
         for index in range(args.steps)
     ]
 want_audio = args.save_audio_latent is not None or args.audio_output is not None or args.mux_audio
-sampled = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps, return_audio=want_audio, progress=args.progress)
+cache_stats = {}
+sampled = sample_video_res_multistep(
+    model,
+    H3PromptPacker(checkpoint),
+    text,
+    video,
+    audio,
+    steps=args.steps,
+    model_timesteps=model_timesteps,
+    return_audio=want_audio,
+    progress=args.progress,
+    cache_mode=args.cache_mode,
+    cache_threshold=args.cache_threshold,
+    cache_start_percent=args.cache_start_percent,
+    cache_end_percent=args.cache_end_percent,
+    cache_subsample_factor=args.cache_subsample_factor,
+    cache_stats=cache_stats,
+)
+if cache_stats:
+    report({"cache": cache_stats})
 if want_audio:
     latent, audio_latent = sampled
 else:

tools/serve_hot_runtime.py

@@ -92,6 +92,11 @@ class Handler(BaseHTTPRequestHandler):
             mux_audio = bool(payload.get("mux_audio", True))
             ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
             save_latent = payload.get("save_latent")
+            cache_mode = payload.get("cache_mode")
+            cache_threshold = float(payload.get("cache_threshold", 0.0))
+            cache_start_percent = float(payload.get("cache_start_percent", 0.0))
+            cache_end_percent = float(payload.get("cache_end_percent", 1.0))
+            cache_subsample_factor = int(payload.get("cache_subsample_factor", 2))
             started = time.perf_counter()
             with runtime_lock:
                 result = runtime.generate(
@@ -106,6 +111,11 @@ class Handler(BaseHTTPRequestHandler):
                     mux_audio=mux_audio,
                     ffmpeg_loglevel=ffmpeg_loglevel,
                     save_latent=save_latent,
+                    cache_mode=cache_mode,
+                    cache_threshold=cache_threshold,
+                    cache_start_percent=cache_start_percent,
+                    cache_end_percent=cache_end_percent,
+                    cache_subsample_factor=cache_subsample_factor,
                 )
             result["wall_seconds"] = time.perf_counter() - started
             write_json(self, 200, result)