From 1c0883a54b2dd038c55f38c1fbb4cc800f4f8f95 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 20:13:48 +0700 Subject: [PATCH] Add persistent hot runtime service --- PLAN.md | 2 +- README.md | 10 ++ compose.spark.yml | 17 ++ src/h3_blackwell_runtime/runtime.py | 235 ++++++++++++++++++++++++++++ tools/serve_hot_runtime.py | 115 ++++++++++++++ 5 files changed, 378 insertions(+), 1 deletion(-) create mode 100644 src/h3_blackwell_runtime/runtime.py create mode 100644 tools/serve_hot_runtime.py diff --git a/PLAN.md b/PLAN.md index 8f6bb63..366f3cb 100644 --- a/PLAN.md +++ b/PLAN.md @@ -54,7 +54,7 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe 1. Validate and benchmark the existing `sage3` backend against the same cat prompt, seed, dimensions, and FP16 VAE runtime path used for Sage2 parity. - First cat benchmark result: Sage3 runs successfully but is slower than Sage2 in this direct path. Sampling was `123.675s` versus Sage2 `114.414s`; warm after text conditioning was `158.111s` versus Sage2 `149.304s`. Same-seed MP4 frame diff versus Sage2 was mean `46.563`, max `255`, so keep Sage3 experimental pending human visual review and stricter tensor gates. -2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory for the selected attention backend, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, backend selection, and timing fields that separate resident-model request latency from cold startup. +2. Build a persistent hot runtime service instead of measuring only process-per-run CLIs. A warm container must preload and retain Qwen, H3, video VAE, and audio VAE in GPU memory, then accept video jobs without model-load latency. Add explicit startup warmup, readiness reporting, request-level attention selection (`sage2`, `sdpa`, `sage3` initially), and timing fields that separate resident-model request latency from cold startup. 3. Add exact memory/lifetime optimizations next: `kj_head_sliced` and `kj_chunked_ffn`. These must preserve the validated direct outputs before being kept. 4. Evaluate prior H3-tested attention candidates as standalone adapters: `sol_attn` and `kj_sage`. 5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`. diff --git a/README.md b/README.md index 8c3f00b..50fcecb 100644 --- a/README.md +++ b/README.md @@ -49,3 +49,13 @@ Generation and latent-decode tools are quiet by default: they suppress ffmpeg ba - `--vae-tile-size 256`: set the direct video VAE spatial tile size. `tools/direct_t2v_preview.py` also accepts `H3_VAE_TILE_SIZE`. Standalone `tools/compare_*`, `tools/trace_*`, `tools/inspect_*`, and `tools/patch_comfy_*` scripts are debugging utilities and remain opt-in by being separate commands. + +## Hot Runtime Service + +`tools/serve_hot_runtime.py` keeps Qwen, H3, video VAE, and audio VAE resident in one process. Start the optional Spark service with: + +```bash +docker compose -f compose.spark.yml up -d h3-hot-runtime +``` + +Use `GET /ready` to confirm resident model readiness. Use `POST /generate` with JSON fields like `prompt`, `output`, `width`, `height`, `frames`, `steps`, `seed`, and optional `attention`. Supported request-level attention values are `sage2`, `sdpa`, and `sage3`; switching attention does not reload model weights. diff --git a/compose.spark.yml b/compose.spark.yml index a7a4934..f4d6017 100644 --- a/compose.spark.yml +++ b/compose.spark.yml @@ -12,3 +12,20 @@ services: - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro - /home/daniel/StoryStudioAssets/H3-output:/output command: ["sleep", "infinity"] + h3-hot-runtime: + build: + context: . + dockerfile: Dockerfile.spark + image: h3-blackwell-runtime:dev + gpus: all + volumes: + - /home/daniel/aeon-spark-test/h3/comfy-models/diffusion_models:/models:ro + - /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro + - /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro + - /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro + - /home/daniel/StoryStudioAssets/H3-output:/output + ports: + - "8001:8000" + environment: + H3_DISABLE_MMAP: "1" + command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"] diff --git a/src/h3_blackwell_runtime/runtime.py b/src/h3_blackwell_runtime/runtime.py new file mode 100644 index 0000000..c2756a0 --- /dev/null +++ b/src/h3_blackwell_runtime/runtime.py @@ -0,0 +1,235 @@ +"""Resident prompt-only H3 runtime used by the hot service and tests.""" + +from __future__ import annotations + +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path + +import torch + +from .audio_vae_decoder import MiniMaxH3AudioVAE +from .attention import AVAILABLE_BACKENDS +from .checkpoint import H3Checkpoint +from .denoiser import H3PackedDenoiser +from .packing import H3PromptPacker +from .qwen3vl_text import Qwen3VLPromptConditioner +from .sampler import sample_video_res_multistep +from .t2v import random_av_latents +from .token_refiner import H3TokenRefiner +from .vae_decoder import MiniMaxH3VideoVAE, dtype_from_name + + +@dataclass(frozen=True) +class RuntimeConfig: + model_path: str = "/models/minimax_h3_fl2va_pruned_nvfp4.safetensors" + text_encoder_path: str = "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors" + tokenizer_path: str = "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer" + video_vae_path: str = "/vae/minimax_h3_video_vae_fp16.safetensors" + audio_vae_path: str = "/vae/minimax_h3_audio_vae_fp32.safetensors" + attention: str = "sage2" + vae_dtype: str = "float16" + vae_tile_size: int = 256 + device: str = "cuda" + + +def _sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _ffmpeg_command(loglevel: str, *parts: str) -> list[str]: + return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts] + + +class H3HotRuntime: + """Keep all prompt-only H3 models resident for repeated requests.""" + + def __init__(self, config: RuntimeConfig): + self.config = config + self.attention = config.attention + self.loaded_at = time.time() + self.load_stages: list[dict] = [] + + self.checkpoint = H3Checkpoint(config.model_path, device=config.device) + self.conditioner = self._timed_load( + "qwen_loaded", + lambda: Qwen3VLPromptConditioner(config.text_encoder_path, config.tokenizer_path), + ) + self.model = self._timed_load( + "h3_loaded", + lambda: H3PackedDenoiser.from_checkpoint(self.checkpoint, attention_backend=config.attention).eval(), + ) + self.refiner = self._timed_load( + "token_refiner_loaded", + lambda: H3TokenRefiner(self.checkpoint, attention_backend=config.attention).eval(), + ) + self.packer = H3PromptPacker(self.checkpoint) + self.video_vae = self._timed_load( + "video_vae_loaded", + self._load_video_vae, + ) + self.audio_vae = self._timed_load( + "audio_vae_loaded", + lambda: MiniMaxH3AudioVAE.from_safetensors(config.audio_vae_path, device=config.device).eval(), + ) + + def _timed_load(self, stage: str, fn): + _sync() + start = time.perf_counter() + value = fn() + _sync() + self.load_stages.append({"stage": stage, "seconds": time.perf_counter() - start}) + return value + + def _load_video_vae(self) -> MiniMaxH3VideoVAE: + vae = MiniMaxH3VideoVAE.from_safetensors( + self.config.video_vae_path, + device=self.config.device, + dtype=dtype_from_name(self.config.vae_dtype), + ).eval() + vae.tile_size = self.config.vae_tile_size + return vae + + def status(self) -> dict: + return { + "ready": True, + "initial_attention": self.config.attention, + "current_attention": self.attention, + "vae_dtype": self.config.vae_dtype, + "vae_tile_size": self.config.vae_tile_size, + "loaded_at": self.loaded_at, + "load_stages": self.load_stages, + } + + @torch.inference_mode() + def set_attention(self, attention: str) -> None: + if attention not in AVAILABLE_BACKENDS: + raise ValueError(f"Unsupported attention backend: {attention}") + if attention == self.attention: + return + for module in self.model.modules(): + if hasattr(module, "backend"): + module.backend = attention + for block in self.refiner.blocks: + block.attention_backend = attention + self.attention = attention + + @torch.inference_mode() + def generate( + self, + *, + prompt: str, + output: str | Path, + width: int, + height: int, + frames: int, + steps: int, + seed: int, + attention: str | None = None, + mux_audio: bool = True, + ffmpeg_loglevel: str = "error", + save_latent: str | Path | None = None, + ) -> dict: + stages: list[dict] = [] + + def timed(stage: str, fn): + _sync() + start = time.perf_counter() + value = fn() + _sync() + stages.append({"stage": stage, "seconds": time.perf_counter() - start}) + return value + + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + if attention is not None: + self.set_attention(attention) + + video, audio, aligned_frames = timed( + "latents_initialized", + lambda: random_av_latents(width, height, frames, seed, device=self.config.device), + ) + text = timed("text_conditioned", lambda: self.refiner(self.conditioner(prompt))) + sampled = timed( + "sampled", + lambda: sample_video_res_multistep( + self.model, + self.packer, + text, + video, + audio, + steps=steps, + return_audio=mux_audio, + ), + ) + if mux_audio: + latent, audio_latent = sampled + else: + latent, audio_latent = sampled, None + if save_latent is not None: + latent_path = Path(save_latent) + latent_path.parent.mkdir(parents=True, exist_ok=True) + state = {"latent": latent.detach().cpu(), "frames": aligned_frames, "width": width, "height": height, "prompt": prompt, "seed": seed} + if audio_latent is not None: + state["audio_latent"] = audio_latent.detach().cpu() + torch.save(state, latent_path) + pixels = timed("vae_decoded", lambda: self.video_vae.decode(latent.to(next(self.video_vae.parameters()).dtype))[:, :, :aligned_frames]) + pixels = timed("pixels_cpu", lambda: ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()) + + raw = output.with_suffix(".rgb") + video_output = output.with_name(output.stem + ".video.mp4") if mux_audio else output + timed("raw_write", lambda: pixels.numpy().tofile(raw)) + timed( + "video_encode", + lambda: subprocess.run( + _ffmpeg_command( + ffmpeg_loglevel, + "-y", + "-f", + "rawvideo", + "-pixel_format", + "rgb24", + "-video_size", + f"{pixels.shape[2]}x{pixels.shape[1]}", + "-framerate", + "24", + "-i", + str(raw), + "-an", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + str(video_output), + ), + check=True, + ), + ) + raw.unlink() + + audio_output = None + if mux_audio: + audio_output = output.with_suffix(".wav") + waveform = timed("audio_decoded", lambda: self.audio_vae.decode(audio_latent.to(next(self.audio_vae.parameters()).dtype)).clamp(-1, 1).cpu()[0]) + audio_raw = audio_output.with_suffix(".f32le") + timed("audio_raw_write", lambda: waveform.transpose(0, 1).contiguous().numpy().tofile(audio_raw)) + timed("audio_encode", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_output)), check=True)) + audio_raw.unlink() + timed("mux", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-i", str(video_output), "-i", str(audio_output), "-c:v", "copy", "-c:a", "aac", "-shortest", str(output)), check=True)) + video_output.unlink() + + return { + "output": str(output), + "audio_output": str(audio_output) if audio_output is not None else None, + "frames": aligned_frames, + "width": width, + "height": height, + "seed": seed, + "attention": self.attention, + "vae_dtype": self.config.vae_dtype, + "vae_tile_size": self.config.vae_tile_size, + "stages": stages, + "request_seconds": sum(stage["seconds"] for stage in stages), + } diff --git a/tools/serve_hot_runtime.py b/tools/serve_hot_runtime.py new file mode 100644 index 0000000..df0a88c --- /dev/null +++ b/tools/serve_hot_runtime.py @@ -0,0 +1,115 @@ +"""Serve a resident prompt-only H3 runtime over a small JSON HTTP API.""" + +from __future__ import annotations + +import argparse +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import urlparse + +from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig + + +parser = argparse.ArgumentParser() +parser.add_argument("--host", default="0.0.0.0") +parser.add_argument("--port", type=int, default=8000) +parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2", help="Initial attention backend. Requests can switch with the JSON attention field.") +parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16") +parser.add_argument("--vae-tile-size", type=int, default=256) +parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.") +parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4")) +args = parser.parse_args() + +runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size)) +runtime_lock = threading.Lock() +warmup_result = None +if args.warmup: + warmup_result = runtime.generate( + prompt="A small warmup cat blinks in soft light.", + output=args.warmup_output, + width=320, + height=192, + frames=22, + steps=2, + seed=440501, + mux_audio=True, + ) + + +def service_status() -> dict: + return { + "ready": True, + "attention_backends": ["sage2", "sdpa", "sage3"], + "runtime": runtime.status(), + } + + +def write_json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None: + body = json.dumps(payload, indent=2).encode("utf-8") + handler.send_response(status) + handler.send_header("Content-Type", "application/json") + handler.send_header("Content-Length", str(len(body))) + handler.end_headers() + handler.wfile.write(body) + + +class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args) -> None: + return + + def do_GET(self) -> None: + path = urlparse(self.path).path + if path in {"/health", "/ready"}: + write_json(self, 200, {**service_status(), "warmup_result": warmup_result}) + return + write_json(self, 404, {"error": "not found"}) + + def do_POST(self) -> None: + path = urlparse(self.path).path + if path != "/generate": + write_json(self, 404, {"error": "not found"}) + return + try: + length = int(self.headers.get("Content-Length", "0")) + payload = json.loads(self.rfile.read(length).decode("utf-8")) if length else {} + prompt = payload["prompt"] + output = payload["output"] + width = int(payload.get("width", 960)) + height = int(payload.get("height", 544)) + frames = int(payload.get("frames", 124)) + steps = int(payload.get("steps", 12)) + seed = int(payload.get("seed", 440407)) + attention = payload.get("attention") + if attention is not None and attention not in {"sage2", "sdpa", "sage3"}: + write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": ["sage2", "sdpa", "sage3"]}) + return + mux_audio = bool(payload.get("mux_audio", True)) + ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error") + save_latent = payload.get("save_latent") + started = time.perf_counter() + with runtime_lock: + result = runtime.generate( + prompt=prompt, + output=output, + width=width, + height=height, + frames=frames, + steps=steps, + seed=seed, + attention=attention, + mux_audio=mux_audio, + ffmpeg_loglevel=ffmpeg_loglevel, + save_latent=save_latent, + ) + result["wall_seconds"] = time.perf_counter() - started + write_json(self, 200, result) + except Exception as exc: + write_json(self, 500, {"error": type(exc).__name__, "message": str(exc)}) + + +server = ThreadingHTTPServer((args.host, args.port), Handler) +print(json.dumps({"serving": True, "host": args.host, "port": args.port, **service_status(), "warmup_result": warmup_result}, indent=2), flush=True) +server.serve_forever()