diff --git a/README.md b/README.md index 50fcecb..5e3141e 100644 --- a/README.md +++ b/README.md @@ -58,4 +58,4 @@ 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 `sage2`, `sdpa`, and `sage3`; 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. diff --git a/src/h3_blackwell_runtime/attention.py b/src/h3_blackwell_runtime/attention.py index 84de387..bf570db 100644 --- a/src/h3_blackwell_runtime/attention.py +++ b/src/h3_blackwell_runtime/attention.py @@ -8,8 +8,8 @@ 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") +PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "sol_attn", "kj_chunked_ffn", "kj_head_sliced") def attention_backend_status() -> dict[str, str]: @@ -18,8 +18,7 @@ def attention_backend_status() -> dict[str, str]: 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({"sol_attn": "experimental: prior H3-tested sparse Triton attention; source/package pending"}) status.update({"kj_chunked_ffn": "planned: exact memory-lifetime adapter"}) status.update({"kj_head_sliced": "planned: exact memory-lifetime adapter"}) return status @@ -39,6 +38,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}") diff --git a/tools/compare_attention_backends.py b/tools/compare_attention_backends.py index 7872201..3378150 100644 --- a/tools/compare_attention_backends.py +++ b/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}") diff --git a/tools/compare_fl2va_free_run.py b/tools/compare_fl2va_free_run.py index ef39a5f..4e7403f 100644 --- a/tools/compare_fl2va_free_run.py +++ b/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() diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index 7800bd0..1c031f6 100644 --- a/tools/direct_t2v_preview.py +++ b/tools/direct_t2v_preview.py @@ -14,6 +14,7 @@ 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.denoiser import H3PackedDenoiser from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner @@ -31,7 +32,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") diff --git a/tools/serve_hot_runtime.py b/tools/serve_hot_runtime.py index df0a88c..a620809 100644 --- a/tools/serve_hot_runtime.py +++ b/tools/serve_hot_runtime.py @@ -10,13 +10,14 @@ 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=("sage2", "sdpa", "sage3"), default="sage2", help="Initial attention backend. Requests can switch with the JSON attention field.") +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("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.") @@ -42,7 +43,7 @@ if args.warmup: def service_status() -> dict: return { "ready": True, - "attention_backends": ["sage2", "sdpa", "sage3"], + "attention_backends": list(AVAILABLE_BACKENDS), "runtime": runtime.status(), } @@ -83,8 +84,8 @@ class Handler(BaseHTTPRequestHandler): steps = int(payload.get("steps", 12)) seed = int(payload.get("seed", 440407)) attention = payload.get("attention") - if attention is not None and attention not in {"sage2", "sdpa", "sage3"}: - write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": ["sage2", "sdpa", "sage3"]}) + 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")