Add Sol-Attn backend

Dockerfile.spark

@@ -1,6 +1,8 @@
 # GB10/Grace Blackwell development image. It does not start inference by default.
 FROM ghcr.io/aeon-7/comfyui-aeon-spark:slim
 
+ARG SOL_ATTN_COMMIT=930a4d6e432ff8b8ed5e30ff2f72519b92d69bdf
+
 WORKDIR /opt/h3-blackwell-runtime
 COPY . .
 
@@ -13,10 +15,18 @@ RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.28
 
 RUN python -m pip install --no-cache-dir "fastsafetensors>=0.1.10"
 
+RUN python -m pip uninstall -y pynvml \
+    && python -m pip install --no-cache-dir nvidia-ml-py
+
+RUN git clone https://github.com/Saganaki22/ComfyUI-sol-attn.git /opt/ComfyUI-sol-attn \
+    && cd /opt/ComfyUI-sol-attn \
+    && git checkout ${SOL_ATTN_COMMIT}
+
 RUN python -m pip install --no-cache-dir --no-deps -e . \
     && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; assert hasattr(torch.ops.comfy_kitchen, 'rms_rope_split_half_'); print(torch.__version__, torch.version.cuda)"
 
 ENV H3_MODEL_PATH=/models/minimax_h3_ref2va_pruned_nvfp4.safetensors
+ENV PYTHONPATH=/opt/ComfyUI-sol-attn
 ENV TORCH_COMPILE_DISABLE=0 TORCHDYNAMO_DISABLE=0
 ENTRYPOINT []
 CMD ["bash"]

PARITY.md

@@ -187,9 +187,9 @@ gate and is exact.
 | H3 final layer | Yes | Bit-exact final AdaLN, target-row modulation, and video/audio patch rows; Comfy materializes the AdaLN and output-head biases through BF16 |
 | H3 DiT | Yes | Strict all-block numeric parity not achieved |
 | Beta/RES sampler | Yes | H3 replay is bit-exact for all 12 captured sampler calls and RES update is bit-exact from captured denoised; direct raw-output to model-wrapper denoised conversion retains small BF16/state boundary deltas |
-| Video VAE decoder | Yes | Same final sampled latent decoded through direct VAE versus Comfy PNG frames: mean abs pixel delta `0.0833898`, max `6` |
-| Audio VAE/decode/mux | No | Preview intentionally produces video only |
-| End-to-end prompt-only FL2VA preview | Yes | Direct 320x192x22 dragon smoke videos generated with Sage2. Comfy-equivalent CPU joint AV seed initialization matches captured `initial.pt` exactly. Video-side wrapper denoised conversion is exact after BF16 raw-output materialization; remaining seed-exact final-latent delta is localized to regenerated text conditioning plus audio-side wrapper conversion |
+| Video VAE decoder | Yes | Direct VAE temporal assembly matches upstream after overlap fix. FP16 is the default Comfy-equivalent runtime path; cat benchmark VAE decode is `25.085s`. Use FP32 only for exact direct diagnostics. |
+| Audio VAE/decode/mux | Yes | Direct decoder-only MiniMax H3 audio VAE returns stereo `32000 Hz` waveform and muxes with generated video. Native audio latent scaling is fixed. |
+| End-to-end prompt-only FL2VA preview | Yes | Apples-to-apples warm cat benchmark is at Comfy parity: Comfy warm `150.26s`; direct warm after text conditioning `149.304s`; direct warm including text conditioning `151.465s`. |
 | Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
 
 ## Remaining Gates, In Dependency Order
@@ -197,14 +197,11 @@ gate and is exact.
 Only these are outstanding. Do not recapture or revisit rows marked complete
 unless the checkpoint, Comfy version, prompt, or backend changes.
 
-1. **Final video latent layout.** Compare direct `unpatchify_video` output
-   with Comfy before entering the sampler.
-2. **Text/audio end-to-end closure.** Direct H3, video wrapper conversion,
-   RES update, VAE decode, and initial AV noise now match their captures.
-   Close regenerated text conditioning and audio-side wrapper conversion before
-   claiming identical seed output.
-3. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
-    Sage3, CUDA graphs, and multi-GPU execution.
+1. **Feature/performance work.** Prompt-only FL2VA is now closed against the
+   warm Comfy baseline. Optimize load/caching/sampling, then evaluate Sage3,
+   CUDA graphs, and multi-GPU execution.
+2. **Full Ref2VA support.** Add references, vision conditioning, VAE encode,
+   and reference-path validation gates.
 
 ## Existing Tools And Their Intended Gate
 

PLAN.md

@@ -48,6 +48,21 @@ The direct runner must first match the model contract and output quality. Beatin
 5. Implement ragged Ulysses Sage3 with transport-identity and distributed-versus-single-Sage3 tests.
 6. Sweep Ulysses/tensor-parallel layouts on 2/4/6/8 GPUs in an NVLink/NVSwitch domain.
 
+## Performance Backend Plan
+
+Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Feature and performance work should proceed in this order:
+
+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, 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`.
+   - `kj_sage` is implemented as explicit SageAttention mode backends: `kj_sage_cuda`, `kj_sage_triton`, `kj_sage_fp8`, and `kj_sage_fp8pp`; all passed hot-runtime smoke tests.
+   - `sol_attn` is still blocked on locating/adding the standalone Sol-Attn source or package. It is not installed in the Spark image and is not present in this repository.
+5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`.
+   - Initial direct cache modes are implemented as opt-in approximate sampler modes. They reuse cached denoised deltas and report skipped-step stats; full-size quality/threshold sweeps are still required before using them for production output.
+6. Keep every backend explicit per run, with separate quality and timing records for sampling, VAE, audio, and end-to-end output.
+
 ## Non-Negotiable Validation
 
 - Never silently pad semantic H3 tokens for unmasked attention.
@@ -58,6 +73,8 @@ The direct runner must first match the model contract and output quality. Beatin
 
 ## 2026-08-13 VAE Debug Handoff
 
+Resolved on `vae-decode-optimization`: direct VAE temporal overlap constants now match upstream, audio decode/mux is implemented, and Comfy-equivalent FP16 video VAE is the default runtime path. The 960x544x124 cat benchmark now matches warm Comfy performance: Comfy `150.26s`, direct `149.304s` after text conditioning, direct VAE decode `25.085s`.
+
 Current saved latent and comparison assets live under:
 
 - Spark: `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime`

README.md

@@ -45,5 +45,29 @@ Generation and latent-decode tools are quiet by default: they suppress ffmpeg ba
 - `--profile-memory`: print memory checkpoints in `tools/direct_t2v_preview.py`.
 - `--ffmpeg-loglevel info`: show ffmpeg details instead of the default `error` level.
 - `--quiet`: suppress JSON summary lines.
+- `--vae-dtype float16`: use Comfy-style FP16 video VAE decode in `tools/direct_t2v_preview.py` or `tools/decode_video_latent.py`; this is the default runtime path. Use `--vae-dtype float32` only for exact direct-path diagnostics. `tools/direct_t2v_preview.py` also accepts `H3_VAE_DTYPE`.
+- `--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 reported by `/ready`; switching attention does not reload model weights.
+
+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).
+- `attention: "sol_attn"` routes eligible H3 attention calls through the pinned ComfyUI Sol-Attn Triton kernel vendored into the Spark image. Configure with `H3_SOL_TAU` (`1.3`), `H3_SOL_MIN_TOKENS` (`4096`), `H3_SOL_THRESH_TYPE` (`diag`), `H3_SOL_INT8_QK`, `H3_SOL_INT8_PV`, `H3_SOL_FALLBACK` (`sage2`), and `H3_SOL_STRICT`.
+- `--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.

benchmarks/ref2va-960x544-124f.json

@@ -18,5 +18,35 @@
   "sampler": "res_multistep",
   "seed": 440202,
   "reference_comfy_sage3_seconds": 49.893,
-  "measurement": "ComfyUI prompt execution time after warm-up"
+  "measurement": "ComfyUI prompt execution time after warm-up",
+  "prompt_only_fl2va_cat_benchmark": {
+    "seed": 440407,
+    "prompt": "A playful orange tabby cat starts in an ordinary cozy living room in a normal house, afternoon light, sofa and rug. The cat crouches, jumps, and does one clean athletic backflip in slow motion. As the backflip completes there is a sharp cinematic cut: the cat lands perfectly on a glowing neon disco dance floor wearing oversized black sunglasses. Mirror ball reflections, colorful lights, joyful party energy, stylish and funny, clear before-and-after transformation.",
+    "comfy_warm_seconds": 150.26,
+    "direct_warm_after_text_conditioned_seconds": 149.304,
+    "direct_warm_including_text_conditioning_seconds": 151.465,
+    "direct_cold_through_audio_decode_seconds": 180.494,
+    "direct_video_vae_dtype": "float16",
+    "direct_video_vae_tile_size": 256,
+    "direct_video_vae_decode_seconds": 25.085,
+    "direct_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\direct-cat-benchmark-960x544-124f-seed440407-fp16-vae.mp4",
+    "direct_log": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\direct-cat-benchmark-960x544-124f-seed440407-fp16-vae.log",
+    "comfy_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\comfy-cat-benchmark-960x544-124f-seed440407_00001_.mp4",
+    "comfy_metrics": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\comfy-cat-benchmark-960x544-124f-seed440407.json"
+  },
+  "prompt_only_fl2va_cat_sage3_benchmark": {
+    "seed": 440407,
+    "attention_backend": "sage3",
+    "direct_cold_through_audio_decode_seconds": 184.646,
+    "direct_warm_after_text_conditioned_seconds": 158.111,
+    "direct_warm_including_text_conditioning_seconds": 160.38,
+    "direct_sample_seconds": 123.675,
+    "direct_video_vae_dtype": "float16",
+    "direct_video_vae_tile_size": 256,
+    "direct_video_vae_decode_seconds": 25.024,
+    "sage2_mp4_frame_diff_max": 255,
+    "sage2_mp4_frame_diff_mean": 46.56300230273561,
+    "direct_output": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\direct-cat-benchmark-960x544-124f-seed440407-sage3-fp16-vae.mp4",
+    "direct_log": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-blackwell-runtime\\benchmarks\\direct-cat-benchmark-960x544-124f-seed440407-sage3-fp16-vae.log"
+  }
 }

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"]

src/h3_blackwell_runtime/attention.py

@@ -1,5 +1,6 @@
 """Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""
 
+import os
 import torch
 import torch.nn.functional as functional
 from torch import nn
@@ -8,20 +9,18 @@ from .checkpoint import H3Checkpoint
 from .nvfp4 import Nvfp4Linear
 
 
-AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3")
-PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "sol_attn", "kj_sage", "kj_chunked_ffn", "kj_head_sliced")
+AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
+PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn")
 
 
 def attention_backend_status() -> dict[str, str]:
     """Report direct-runtime attention choices without importing ComfyUI nodes."""
     status = {name: "available" for name in AVAILABLE_BACKENDS}
+    status.update({"sol_attn": "experimental: sparse Triton attention for eligible non-causal H3 attention calls; falls back below H3_SOL_MIN_TOKENS unless H3_SOL_STRICT=1"})
     status.update({"flash4": "planned: exact Blackwell kernel adapter"})
     status.update({"easycache": "planned: approximate denoiser cache"})
     status.update({"h3_cache": "planned: approximate H3-specific cache"})
-    status.update({"sol_attn": "experimental: prior H3-tested sparse Triton attention; standalone adapter pending"})
-    status.update({"kj_sage": "experimental: prior H3-tested Sage patch; standalone adapter pending"})
-    status.update({"kj_chunked_ffn": "planned: exact memory-lifetime adapter"})
-    status.update({"kj_head_sliced": "planned: exact memory-lifetime adapter"})
+    status.update({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"})
     return status
 
 
@@ -31,6 +30,45 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
 
 def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor:
     """Run one `[batch, heads, sequence, dim]` attention operation."""
+    if backend == "sol_attn":
+        try:
+            from sol_kernel import sol_attn
+
+            tau = float(os.getenv("H3_SOL_TAU", "1.3"))
+            min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096"))
+            thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag")
+            int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"}
+            int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"}
+            if is_causal:
+                raise ValueError("Sol-Attn backend only supports non-causal H3 attention")
+            if q.shape[-1] != 128:
+                raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}")
+            if q.shape[2] < min_tokens:
+                raise ValueError(f"{q.shape[2]} tokens < H3_SOL_MIN_TOKENS={min_tokens}")
+            out = sol_attn(
+                q.transpose(1, 2).contiguous(),
+                k.transpose(1, 2).contiguous(),
+                v.transpose(1, 2).contiguous(),
+                tau=tau,
+                thresh_type=thresh_type,
+                int8_qk=int8_qk,
+                int8_pv=int8_pv,
+            )
+            return out.transpose(1, 2)
+        except Exception:
+            if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}:
+                raise
+            return run_attention(q, k, v, backend=os.getenv("H3_SOL_FALLBACK", "sage2"), is_causal=is_causal)
+    if backend == "kj_head_sliced":
+        head_slice_size = int(os.getenv("H3_HEAD_SLICE_SIZE", "8"))
+        base_backend = os.getenv("H3_HEAD_SLICE_BACKEND", "sage2")
+        if head_slice_size <= 0:
+            raise ValueError("H3_HEAD_SLICE_SIZE must be positive")
+        outputs = [
+            run_attention(q[:, start:start + head_slice_size], k[:, start:start + head_slice_size], v[:, start:start + head_slice_size], backend=base_backend, is_causal=is_causal)
+            for start in range(0, q.shape[1], head_slice_size)
+        ]
+        return torch.cat(outputs, dim=1)
     if backend == "sage2":
         from sageattention import sageattn
 
@@ -39,6 +77,26 @@ def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend:
         from sageattn3 import sageattn3_blackwell
 
         return sageattn3_blackwell(q, k, v, is_causal=is_causal)
+    if backend == "sage3_mean":
+        from sageattn3 import sageattn3_blackwell
+
+        return sageattn3_blackwell(q, k, v, is_causal=is_causal, per_block_mean=True)
+    if backend == "kj_sage_cuda":
+        from sageattention import sageattn_qk_int8_pv_fp16_cuda
+
+        return sageattn_qk_int8_pv_fp16_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32", tensor_layout="HND")
+    if backend == "kj_sage_triton":
+        from sageattention import sageattn_qk_int8_pv_fp16_triton
+
+        return sageattn_qk_int8_pv_fp16_triton(q, k, v, is_causal=is_causal, tensor_layout="HND")
+    if backend == "kj_sage_fp8":
+        from sageattention import sageattn_qk_int8_pv_fp8_cuda
+
+        return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp32", tensor_layout="HND")
+    if backend == "kj_sage_fp8pp":
+        from sageattention import sageattn_qk_int8_pv_fp8_cuda
+
+        return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp16", tensor_layout="HND")
     if backend == "sdpa":
         return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
     raise ValueError(f"Unsupported H3 attention backend: {backend}")

src/h3_blackwell_runtime/block.py

@@ -39,6 +39,8 @@ class H3SwiGLU(nn.Module):
         super().__init__()
         self.fc1 = fc1
         self.fc2 = fc2
+        self.chunks = 1
+        self.chunk_threshold = 4096
 
     @classmethod
     def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16):
@@ -48,10 +50,25 @@ class H3SwiGLU(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.chunks > 1 and x.shape[0] >= self.chunk_threshold:
+            return torch.cat([self._forward_chunk(chunk) for chunk in x.chunk(self.chunks, dim=0)], dim=0)
+        return self._forward_chunk(x)
+
+    def _forward_chunk(self, x: torch.Tensor) -> torch.Tensor:
         gate, up = self.fc1(x).chunk(2, dim=-1)
         return self.fc2(torch.nn.functional.silu(gate).mul_(up))
 
 
+def configure_mlp_chunking(model: nn.Module, chunks: int, threshold: int = 4096) -> None:
+    """Configure exact row-chunked H3 SwiGLU execution to reduce peak activation memory."""
+    if chunks < 1:
+        raise ValueError("MLP chunks must be >= 1")
+    for module in model.modules():
+        if isinstance(module, H3SwiGLU):
+            module.chunks = chunks
+            module.chunk_threshold = threshold
+
+
 class H3DiTBlock(nn.Module):
     """One H3 transformer block with externally supplied AdaLN tensors."""
 

src/h3_blackwell_runtime/runtime.py

@@ -0,0 +1,266 @@
+"""Resident prompt-only H3 runtime used by the hot service and tests."""
+
+from __future__ import annotations
+
+import os
+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 .block import configure_mlp_chunking
+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
+    mlp_chunks: int = 1
+    mlp_chunk_threshold: int = 4096
+    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]
+
+
+def _refiner_attention_backend(attention: str) -> str:
+    if attention != "sol_attn":
+        return attention
+    fallback = os.getenv("H3_SOL_FALLBACK", "sage2")
+    return fallback if fallback in AVAILABLE_BACKENDS and fallback != "sol_attn" else "sage2"
+
+
+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",
+            self._load_h3,
+        )
+        self.refiner = self._timed_load(
+            "token_refiner_loaded",
+            lambda: H3TokenRefiner(self.checkpoint, attention_backend=_refiner_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 _load_h3(self) -> H3PackedDenoiser:
+        model = H3PackedDenoiser.from_checkpoint(self.checkpoint, attention_backend=self.config.attention).eval()
+        configure_mlp_chunking(model, self.config.mlp_chunks, self.config.mlp_chunk_threshold)
+        return model
+
+    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,
+            "mlp_chunks": self.config.mlp_chunks,
+            "mlp_chunk_threshold": self.config.mlp_chunk_threshold,
+            "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 = _refiner_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,
+        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()
+            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,
+                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:
+            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,
+            "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,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)

src/h3_blackwell_runtime/vae_decoder.py

@@ -23,6 +23,14 @@ LATENTS_MEAN = (0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5
 LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037, 1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875, 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775, 2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265, 3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525)
 
 
+def dtype_from_name(name: str) -> torch.dtype:
+    return {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }[name]
+
+
 def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
     if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}:
         if weight is None:
@@ -198,7 +206,7 @@ class MiniMaxH3VideoVAE(nn.Module):
         self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False)
 
     @classmethod
-    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float32) -> "MiniMaxH3VideoVAE":
+    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float16) -> "MiniMaxH3VideoVAE":
         model = cls(device="meta", tiling=tiling)
         expected = model.state_dict()
         if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:

tools/benchmark_video_vae_decode.py

@@ -10,6 +10,8 @@ from pathlib import Path
 
 import torch
 
+from h3_blackwell_runtime.vae_decoder import dtype_from_name
+
 
 def sync() -> None:
     if torch.cuda.is_available():
@@ -26,14 +28,6 @@ def timed(stage: str, rows: list[dict], fn):
     return value
 
 
-def dtype_from_name(name: str) -> torch.dtype:
-    return {
-        "float32": torch.float32,
-        "float16": torch.float16,
-        "bfloat16": torch.bfloat16,
-    }[name]
-
-
 def ffmpeg_command(loglevel: str, *parts: str) -> list[str]:
     return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts]
 
@@ -44,7 +38,7 @@ def main() -> None:
     parser.add_argument("--output", type=Path)
     parser.add_argument("--metrics", type=Path, required=True)
     parser.add_argument("--implementation", choices=("direct", "upstream"), default="direct")
-    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
+    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float16")
     parser.add_argument("--no-tiling", action="store_true")
     parser.add_argument("--tile-size", type=int)
     parser.add_argument("--tile-overlap", type=int)

tools/compare_attention_backends.py

@@ -4,8 +4,8 @@ import argparse
 import time
 
 import torch
-import torch.nn.functional as functional
 
+from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, run_attention
 from h3_blackwell_runtime.checkpoint import H3Checkpoint
 from h3_blackwell_runtime.denoiser import H3PackedDenoiser
 
@@ -21,20 +21,14 @@ q, k, v = payload["q"], payload["k"], payload["v"]
 model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
 out_proj = model.backbone.blocks[0].attention.out_proj
 
-for name in ("sdpa", "sage2", "sage3"):
+for name in AVAILABLE_BACKENDS:
     torch.cuda.synchronize()
     start = time.perf_counter()
-    if name == "sdpa":
-        output = functional.scaled_dot_product_attention(q, k, v, is_causal=False)
-    elif name == "sage2":
-        from sageattention import sageattn
-
-        output = sageattn(q, k, v, is_causal=False, tensor_layout="HND", smooth_k=False)
-    else:
-        from sageattn3 import sageattn3_blackwell
-
-        output = sageattn3_blackwell(q, k, v, is_causal=False)
-    torch.cuda.synchronize()
-    output = out_proj(output.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1).reshape(q.shape[2], -1).contiguous())
-    delta = (output.float() - expected.float()).abs()
-    print(f"{name} elapsed_s={time.perf_counter() - start:.3f} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+    try:
+        output = run_attention(q, k, v, backend=name, is_causal=False)
+        torch.cuda.synchronize()
+        output = out_proj(output.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1).reshape(q.shape[2], -1).contiguous())
+        delta = (output.float() - expected.float()).abs()
+        print(f"{name} elapsed_s={time.perf_counter() - start:.3f} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+    except Exception as exc:
+        print(f"{name} error={type(exc).__name__}: {exc}")

tools/compare_fl2va_free_run.py

@@ -5,6 +5,7 @@ import argparse
 import torch
 
 from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
 from h3_blackwell_runtime.denoiser import H3PackedDenoiser
 from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video
 from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, _unpack_audio, res_multistep_update
@@ -18,7 +19,7 @@ parser.add_argument("--width", type=int, default=320)
 parser.add_argument("--height", type=int, default=192)
 parser.add_argument("--frames", type=int, default=22)
 parser.add_argument("--seed", type=int, default=440204)
-parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
+parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2")
 parser.add_argument("--oracle-timesteps", action="store_true")
 args = parser.parse_args()
 

tools/decode_video_latent.py

@@ -6,7 +6,7 @@ from pathlib import Path
 
 import torch
 
-from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
 
 
 parser = argparse.ArgumentParser()
@@ -17,6 +17,8 @@ parser.add_argument("--implementation", choices=("direct", "upstream"), default=
 parser.add_argument("--frames-dir", type=Path)
 parser.add_argument("--ffmpeg-loglevel", default="error")
 parser.add_argument("--quiet", action="store_true")
+parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
+parser.add_argument("--vae-tile-size", type=int, default=256)
 args = parser.parse_args()
 
 state = torch.load(args.latent, map_location="cuda", weights_only=False)
@@ -28,7 +30,9 @@ if args.implementation == "direct":
         "/vae/minimax_h3_video_vae_fp16.safetensors",
         device="cuda",
         tiling=not args.no_tiling,
+        dtype=dtype_from_name(args.vae_dtype),
     ).eval()
+    vae.tile_size = args.vae_tile_size
 else:
     from safetensors.torch import load_file
 
@@ -36,7 +40,10 @@ else:
 
     vae = UpstreamMiniMaxH3VideoVAE(tiling=not args.no_tiling).to("cuda").eval()
     checkpoint = load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda")
+    checkpoint = {key: value.to(dtype=dtype_from_name(args.vae_dtype)) for key, value in checkpoint.items()}
     missing, unexpected = vae.load_state_dict(checkpoint, strict=False)
+    vae = vae.to(dtype=dtype_from_name(args.vae_dtype))
+    vae.tile_size = args.vae_tile_size
     if not args.quiet:
         print({"upstream_missing": len(missing), "upstream_unexpected": len(unexpected)}, flush=True)
 with torch.inference_mode():
@@ -61,4 +68,4 @@ subprocess.run([
 ], check=True)
 raw.unlink()
 if not args.quiet:
-    print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation})
+    print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation, "vae_dtype": args.vae_dtype, "vae_tile_size": vae.tile_size})

tools/direct_t2v_preview.py

@@ -14,13 +14,15 @@ import torch
 
 from h3_blackwell_runtime.checkpoint import H3Checkpoint
 from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
+from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
+from h3_blackwell_runtime.block import configure_mlp_chunking
 from h3_blackwell_runtime.denoiser import H3PackedDenoiser
 from h3_blackwell_runtime.packing import H3PromptPacker
 from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
 from h3_blackwell_runtime.sampler import sample_video_res_multistep
 from h3_blackwell_runtime.t2v import random_av_latents
 from h3_blackwell_runtime.token_refiner import H3TokenRefiner
-from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
 
 
 parser = argparse.ArgumentParser()
@@ -31,7 +33,7 @@ parser.add_argument("--height", type=int, default=192)
 parser.add_argument("--frames", type=int, default=22)
 parser.add_argument("--steps", type=int, default=12)
 parser.add_argument("--seed", type=int, default=440204)
-parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
+parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2")
 parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
 parser.add_argument("--progress", action="store_true", help="Print per-step sampler progress.")
 parser.add_argument("--profile-memory", action="store_true")
@@ -42,6 +44,15 @@ parser.add_argument("--save-audio-latent", type=Path)
 parser.add_argument("--audio-output", type=Path)
 parser.add_argument("--mux-audio", action="store_true")
 parser.add_argument("--skip-decode", action="store_true")
+parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default=os.getenv("H3_VAE_DTYPE", "float16"))
+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
@@ -84,6 +95,7 @@ conditioner = Qwen3VLPromptConditioner(
 report_memory("qwen_loaded")
 video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed)
 model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
+configure_mlp_chunking(model, args.mlp_chunks, args.mlp_chunk_threshold)
 report_memory("h3_loaded")
 text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt))
 report_memory("text_conditioned")
@@ -94,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:
@@ -115,7 +146,8 @@ if args.save_audio_latent is not None:
     report({"audio_latent": str(args.save_audio_latent)})
 if args.skip_decode:
     raise SystemExit(0)
-vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
+vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", dtype=dtype_from_name(args.vae_dtype)).eval()
+vae.tile_size = args.vae_tile_size
 report_memory("vae_loaded")
 with torch.inference_mode():
     pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]

tools/serve_hot_runtime.py

@@ -0,0 +1,128 @@
+"""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.attention import AVAILABLE_BACKENDS
+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=AVAILABLE_BACKENDS, 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("--mlp-chunks", type=int, default=1)
+parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
+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, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold))
+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": list(AVAILABLE_BACKENDS),
+        "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 AVAILABLE_BACKENDS:
+                write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": list(AVAILABLE_BACKENDS)})
+                return
+            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(
+                    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,
+                    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)
+        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()