daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add KJ Sage attention backends

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 20:28:56 +07:00
parent 1c0883a54b
commit eaf9324145
6 changed files with 43 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -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.

27
src/h3_blackwell_runtime/attention.py
View file

@ -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}")

26
tools/compare_attention_backends.py
View file

@ -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}")

3
tools/compare_fl2va_free_run.py
View file

@ -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()

3
tools/direct_t2v_preview.py
View file

@ -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")

9
tools/serve_hot_runtime.py
View file

@ -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")