Dockerfile.spark

@@ -1,8 +1,6 @@
 # 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 . .
 
@@ -15,18 +13,10 @@ 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 | 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. |
+| 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 | Yes | Direct decoder-only MiniMax H3 audio VAE returns stereo `32000 Hz` waveform and muxes with generated video. Native audio latent scaling is fixed. |
+| Audio VAE/decode/mux | No | Preview intentionally produces video only |
-| 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`. |
+| 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 |
 | Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
 
 ## Remaining Gates, In Dependency Order
@@ -197,11 +197,14 @@ 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. **Feature/performance work.** Prompt-only FL2VA is now closed against the
+1. **Final video latent layout.** Compare direct `unpatchify_video` output
-   warm Comfy baseline. Optimize load/caching/sampling, then evaluate Sage3,
+   with Comfy before entering the sampler.
-   CUDA graphs, and multi-GPU execution.
+2. **Text/audio end-to-end closure.** Direct H3, video wrapper conversion,
-2. **Full Ref2VA support.** Add references, vision conditioning, VAE encode,
+   RES update, VAE decode, and initial AV noise now match their captures.
-   and reference-path validation gates.
+   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.
 
 ## Existing Tools And Their Intended Gate
 

PLAN.md

@@ -48,21 +48,6 @@ 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.
@@ -73,8 +58,6 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe
 
 ## 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,29 +45,5 @@ 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,35 +18,5 @@
   "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,20 +12,3 @@ 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,6 +1,5 @@
 """Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""
 
-import os
 import torch
 import torch.nn.functional as functional
 from torch import nn
@@ -9,18 +8,20 @@ from .checkpoint import H3Checkpoint
 from .nvfp4 import Nvfp4Linear
 
 
-AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
+AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3")
-PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn")
+PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "sol_attn", "kj_sage", "kj_chunked_ffn", "kj_head_sliced")
 
 
 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({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"})
+    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"})
     return status
 
 
@@ -30,45 +31,6 @@ 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
 
@@ -77,26 +39,6 @@ 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,8 +39,6 @@ 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):
@@ -50,25 +48,10 @@ 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

@@ -1,266 +0,0 @@
-"""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,14 +57,6 @@ 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,
@@ -77,12 +69,6 @@ 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)
@@ -91,68 +77,25 @@ 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
+        sigma_audio = _audio_sigma(sigma)
-        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
+        carry = sigma_audio / sigma
-        skipped = False
+        native_audio = audio_carried.to(torch.bfloat16) * carry
-        if can_cache:
+        step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
-            video_now = _cache_sample(video, cache_subsample_factor)
+        hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps)
-            audio_now = _cache_sample(audio_carried, cache_subsample_factor)
+        raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
-            input_change = (video_now - cache["video_prev"]).flatten().abs().mean() + (audio_now - cache["audio_prev"]).flatten().abs().mean()
+        raw_video = raw_video.to(torch.bfloat16).float()
-            input_norm = cache["prev_norm"].clamp_min(1e-8)
+        raw_audio = raw_audio.to(torch.bfloat16)
-            rate = (input_change / input_norm).item()
+        velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
-            if cache_mode == "easycache":
+        velocity_audio = (
-                cache["cumulative_rate"] += rate
+            (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
-                skipped = cache["cumulative_rate"] < cache_threshold
+            + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
-                if not skipped:
+        ).float()
-                    cache["cumulative_rate"] = 0.0
+        video_denoised = video - sigma * velocity_video
-            else:
+        audio_denoised = audio_carried - sigma * velocity_audio
-                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,14 +23,6 @@ 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:
@@ -206,7 +198,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.float16) -> "MiniMaxH3VideoVAE":
+    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float32) -> "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,8 +10,6 @@ from pathlib import Path
 
 import torch
 
-from h3_blackwell_runtime.vae_decoder import dtype_from_name
-
 
 def sync() -> None:
     if torch.cuda.is_available():
@@ -28,6 +26,14 @@ 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]
 
@@ -38,7 +44,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="float16")
+    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
     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,14 +21,20 @@ 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 AVAILABLE_BACKENDS:
+for name in ("sdpa", "sage2", "sage3"):
     torch.cuda.synchronize()
     start = time.perf_counter()
-    try:
+    if name == "sdpa":
-        output = run_attention(q, k, v, backend=name, is_causal=False)
+        output = functional.scaled_dot_product_attention(q, k, v, is_causal=False)
-        torch.cuda.synchronize()
+    elif name == "sage2":
-        output = out_proj(output.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1).reshape(q.shape[2], -1).contiguous())
+        from sageattention import sageattn
-        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}")
+        output = sageattn(q, k, v, is_causal=False, tensor_layout="HND", smooth_k=False)
-    except Exception as exc:
+    else:
-        print(f"{name} error={type(exc).__name__}: {exc}")
+        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}")

tools/compare_fl2va_free_run.py

@@ -5,7 +5,6 @@ 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
@@ -19,7 +18,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=AVAILABLE_BACKENDS, default="sage2")
+parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), 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, dtype_from_name
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
 
 
 parser = argparse.ArgumentParser()
@@ -17,8 +17,6 @@ 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)
@@ -30,9 +28,7 @@ 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
 
@@ -40,10 +36,7 @@ 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():
@@ -68,4 +61,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, "vae_dtype": args.vae_dtype, "vae_tile_size": vae.tile_size})
+    print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation})

tools/direct_t2v_preview.py

@@ -14,15 +14,13 @@ 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, dtype_from_name
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
 
 
 parser = argparse.ArgumentParser()
@@ -33,7 +31,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=AVAILABLE_BACKENDS, default="sage2")
+parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), 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")
@@ -44,15 +42,6 @@ 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
@@ -95,7 +84,6 @@ 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")
@@ -106,26 +94,7 @@ 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
-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)
-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:
@@ -146,8 +115,7 @@ 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", dtype=dtype_from_name(args.vae_dtype)).eval()
+vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").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

@@ -1,128 +0,0 @@
-"""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()