Add selectable hot attention backends

.dockerignore

@@ -0,0 +1,5 @@
+.git
+.pytest_cache
+artifacts
+**/__pycache__
+**/*.pyc

Dockerfile.spark

@@ -11,7 +11,7 @@ COPY wheels/sageattn3-*.whl /tmp/wheels/
 RUN python -m pip install --no-cache-dir --no-deps /tmp/wheels/sageattn3-*.whl \
     && rm -rf /tmp/wheels
 
-RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.28
+RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.31
 
 RUN python -m pip install --no-cache-dir "fastsafetensors>=0.1.10"
 
@@ -23,7 +23,7 @@ RUN git clone https://github.com/Saganaki22/ComfyUI-sol-attn.git /opt/ComfyUI-so
     && 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)"
+    && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; assert hasattr(torch.ops.comfy_kitchen, 'rms_rope_split_half_'); assert hasattr(comfy_kitchen, 'int8_attention'); assert hasattr(comfy_kitchen, 'int8_attention_is_available'); print(torch.__version__, torch.version.cuda)"
 
 ENV H3_MODEL_PATH=/models/minimax_h3_ref2va_pruned_nvfp4.safetensors
 ENV PYTHONPATH=/opt/ComfyUI-sol-attn

README.md

@@ -58,11 +58,13 @@ Standalone `tools/compare_*`, `tools/trace_*`, `tools/inspect_*`, and `tools/pat
 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.
+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. The hot image includes Sage2, forced cuDNN SDPA, and Comfy Kitchen INT8 attention. Sage2 is the default based on the 960x544x124 GB10 benchmark and the existing parity baseline.
 
 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: "cudnn_sdpa"` forces cuDNN SDPA with no fallback to another PyTorch kernel.
+- `attention: "ck_int8"` uses Comfy Kitchen's approximate INT8 Q/K/V attention kernel.
 - `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).
 

compose.spark.yml

@@ -31,4 +31,4 @@ services:
       H3_NVFP4_SCALE_BACKEND: "vortex"
       H3_NVFP4_SCALE_VERSION: "1"
       H3_SOL_QKV_LAYOUT: "native"
-    command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sol_attn", "--warmup"]
+    command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"]

pyproject.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 description = "Direct MiniMax H3 Blackwell inference research runtime"
 requires-python = ">=3.12"
 dependencies = [
-  "comfy-kitchen==0.2.28",
+  "comfy-kitchen==0.2.31",
   "fastsafetensors>=0.1.10",
   "safetensors>=0.5.0",
   "torch==2.9.1+cu130",

src/h3_blackwell_runtime/attention.py

@@ -9,14 +9,16 @@ 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", "cudnn_sdpa", "ck_int8", "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")
-DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sol_attn")
+DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sage2")
 
 
 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({"cudnn_sdpa": "available: forced cuDNN SDPA with no backend fallback"})
+    status.update({"ck_int8": "available: approximate Comfy Kitchen INT8 Q/K/V attention"})
     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"})
@@ -101,6 +103,17 @@ def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend:
         from sageattention import sageattn
 
         return sageattn(q, k, v, is_causal=is_causal, tensor_layout="HND", smooth_k=False)
+    if backend == "cudnn_sdpa":
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
+            return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+    if backend == "ck_int8":
+        if is_causal:
+            raise ValueError("Comfy Kitchen INT8 attention does not support causal H3 attention")
+        import comfy_kitchen
+
+        return comfy_kitchen.int8_attention(q, k, v)
     if backend == "sage3":
         from sageattn3 import sageattn3_blackwell
 

tests/test_fl2va_contracts.py

@@ -1,4 +1,5 @@
 import math
+import sys
 from types import SimpleNamespace
 import unittest
 from unittest.mock import patch
@@ -8,6 +9,7 @@ from torch import nn
 from torch.nn import functional as F
 
 from h3_blackwell_runtime.packing import FRAME_RESCALE, H3PromptPacker, _video_t_spans
+from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, run_attention
 from h3_blackwell_runtime.qwen3vl_vision import (
     TEXT_HEAD_DIM,
     TEXT_ROPE_DIMS,
@@ -59,6 +61,38 @@ class Fl2vaVAEContracts(unittest.TestCase):
         self.assertEqual(converted.flatten().tolist(), [0, 127, 255])
 
 
+class AttentionBackendContracts(unittest.TestCase):
+    def test_hot_backends_include_benchmark_candidates(self):
+        self.assertTrue({"sage2", "cudnn_sdpa", "ck_int8"}.issubset(AVAILABLE_BACKENDS))
+
+    def test_cudnn_backend_is_forced_without_fallback(self):
+        q = torch.randn(1, 2, 3, 4)
+        expected = torch.randn_like(q)
+        context = unittest.mock.MagicMock()
+        with (
+            patch("torch.nn.attention.sdpa_kernel", return_value=context) as kernel,
+            patch("h3_blackwell_runtime.attention.functional.scaled_dot_product_attention", return_value=expected) as sdpa,
+        ):
+            actual = run_attention(q, q, q, backend="cudnn_sdpa", is_causal=False)
+        self.assertIs(actual, expected)
+        self.assertEqual(kernel.call_args.args[0], [torch.nn.attention.SDPBackend.CUDNN_ATTENTION])
+        sdpa.assert_called_once_with(q, q, q, is_causal=False)
+
+    def test_comfy_kitchen_int8_backend_dispatches_hnd_tensors(self):
+        q = torch.randn(1, 2, 3, 4)
+        expected = torch.randn_like(q)
+        kitchen = SimpleNamespace(int8_attention=unittest.mock.MagicMock(return_value=expected))
+        with patch.dict(sys.modules, {"comfy_kitchen": kitchen}):
+            actual = run_attention(q, q, q, backend="ck_int8", is_causal=False)
+        self.assertIs(actual, expected)
+        kitchen.int8_attention.assert_called_once_with(q, q, q)
+
+    def test_comfy_kitchen_int8_rejects_causal_attention(self):
+        q = torch.randn(1, 2, 3, 4)
+        with self.assertRaisesRegex(ValueError, "does not support causal"):
+            run_attention(q, q, q, backend="ck_int8", is_causal=True)
+
+
 class Fl2vaVisionContracts(unittest.TestCase):
     def test_visual_rotary_coordinates_are_block_major(self):
         class CoordinateTable(nn.Module):

tools/compare_generation_latents.py

@@ -0,0 +1,32 @@
+"""Compare saved hot-runtime video and audio latents against one reference run."""
+
+import argparse
+import json
+
+import torch
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("reference")
+parser.add_argument("candidates", nargs="+")
+args = parser.parse_args()
+
+reference = torch.load(args.reference, map_location="cpu", weights_only=False)
+results = {}
+for path in args.candidates:
+    candidate = torch.load(path, map_location="cpu", weights_only=False)
+    metrics = {}
+    for name in ("latent", "audio_latent"):
+        expected = reference[name].float()
+        actual = candidate[name].float()
+        delta = actual - expected
+        metrics[name] = {
+            "max_abs": delta.abs().max().item(),
+            "mean_abs": delta.abs().mean().item(),
+            "rmse": delta.square().mean().sqrt().item(),
+            "relative_rmse": (delta.square().mean().sqrt() / expected.square().mean().sqrt()).item(),
+            "cosine": torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item(),
+        }
+    results[path] = metrics
+
+print(json.dumps({"reference": args.reference, "results": results}, indent=2))

tools/serve_hot_runtime.py

@@ -13,7 +13,7 @@ from urllib.parse import urlparse
 import torch
 from PIL import Image
 
-from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND
+from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND, attention_backend_status
 from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
 
 
@@ -80,6 +80,7 @@ def service_status() -> dict:
     return {
         "ready": True,
         "attention_backends": list(AVAILABLE_BACKENDS),
+        "attention_backend_status": attention_backend_status(),
         "runtime": runtime.status(),
     }
 

tools/smoke_attention_backends.py

@@ -0,0 +1,58 @@
+"""Verify and time request-selectable H3 attention kernels on the active GPU."""
+
+import argparse
+import json
+import time
+
+import torch
+
+from h3_blackwell_runtime.attention import run_attention
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--backends", nargs="+", default=("sage2", "cudnn_sdpa", "ck_int8"))
+parser.add_argument("--sequence", type=int, default=512)
+parser.add_argument("--heads", type=int, default=56)
+parser.add_argument("--head-dim", type=int, default=128)
+parser.add_argument("--warmup", type=int, default=2)
+parser.add_argument("--iterations", type=int, default=5)
+parser.add_argument("--seed", type=int, default=440407)
+args = parser.parse_args()
+
+torch.manual_seed(args.seed)
+q = torch.randn(1, args.heads, args.sequence, args.head_dim, device="cuda", dtype=torch.bfloat16)
+results = {}
+reference = None
+
+with torch.inference_mode():
+    for backend in args.backends:
+        for _ in range(args.warmup):
+            output = run_attention(q, q, q, backend=backend, is_causal=False)
+        torch.cuda.synchronize()
+        elapsed = []
+        for _ in range(args.iterations):
+            started = time.perf_counter()
+            output = run_attention(q, q, q, backend=backend, is_causal=False)
+            torch.cuda.synchronize()
+            elapsed.append(time.perf_counter() - started)
+        if reference is None:
+            reference = output
+        delta = (output.float() - reference.float()).abs()
+        results[backend] = {
+            "mean_seconds": sum(elapsed) / len(elapsed),
+            "min_seconds": min(elapsed),
+            "finite": bool(torch.isfinite(output).all()),
+            "shape": list(output.shape),
+            "dtype": str(output.dtype),
+            "max_abs_vs_reference": delta.max().item(),
+            "mean_abs_vs_reference": delta.mean().item(),
+        }
+
+print(json.dumps({
+    "gpu": torch.cuda.get_device_name(),
+    "torch": torch.__version__,
+    "cuda": torch.version.cuda,
+    "shape": list(q.shape),
+    "reference": args.backends[0],
+    "results": results,
+}, indent=2))