daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add approximate H3 cache modes

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 20:38:48 +07:00
parent fb257ef982
commit c3cc04e98d
5 changed files with 126 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
README.md
View file

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

13
src/h3_blackwell_runtime/runtime.py
View file

@ -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),
}

87
src/h3_blackwell_runtime/sampler.py
View file

@ -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,25 +91,68 @@ 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]
sigma_audio = _audio_sigma(sigma)
carry = sigma_audio / sigma
native_audio = audio_carried.to(torch.bfloat16) * carry
step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps)
raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
raw_video = raw_video.to(torch.bfloat16).float()
raw_audio = raw_audio.to(torch.bfloat16)
velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
velocity_audio = (
(1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
+ (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
).float()
video_denoised = video - sigma * velocity_video
audio_denoised = audio_carried - sigma * velocity_audio
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
step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps)
raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
raw_video = raw_video.to(torch.bfloat16).float()
raw_audio = raw_audio.to(torch.bfloat16)
velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
velocity_audio = (
(1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
+ (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
).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)

26
tools/direct_t2v_preview.py
View file

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

10
tools/serve_hot_runtime.py
View file

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