[wip][core] propagate sage attention updates.

[sayakpaul]

What does this PR do?

Propagates the latest changes from SAGE and SAGE2 upstream through kernels.

Summary of the speedups (used black-forest-labs/FLUX.2-klein-9B DiT on an L4) for SAGE2:

Before we jump to any conclusions, here is a table benchmarking just the attention kernel:

So, as we can see that just the attention kernel is doing fine. But since the underlying model is itself dominates on MLP, results in the context of the full model become somewhat diluted.

The usage doesn't change: pipe.transformer.set_attention_backend("sage_hub").

Full code is below:

Unfold

import argparse
import time
from pathlib import Path

import numpy as np
import torch

from diffusers import Flux2KleinPipeline
from diffusers.models.attention_dispatch import AttentionBackendName, _HUB_KERNELS_REGISTRY


STAGING_REPO_ID = "kernels-staging/sage-attention"
STAGING_REVISION = "pr-1095"

MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
PROMPT = "A cat holding a sign that says hello world"


def use_staged_kernel(repo_id: str, revision: str) -> None:
    """Point the `sage_hub` backend at a staged build.

    The v3 kernel is not published to `kernels-community` yet, so without this the loader
    resolves the still-published version and the run does not test the new build.
    """
    config = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB]
    config.repo_id = repo_id
    config.revision = revision
    config.version = None  # `revision` pins the build; a version pin would fight it
    print(f"[setup] sage_hub -> {config.repo_id}@{config.revision}", flush=True)


def _infer(pipe, steps: int, size: int, seed: int):
    return pipe(
        prompt=PROMPT,
        height=size,
        width=size,
        guidance_scale=1.0,
        num_inference_steps=steps,
        generator=torch.Generator(device="cuda").manual_seed(seed),
    ).images[0]


def generate(pipe, tag: str, args, out_dir: Path):
    if args.warmup_steps > 0:
        start = time.perf_counter()
        _infer(pipe, args.warmup_steps, args.size, args.seed)
        torch.cuda.synchronize()
        print(f"[{tag}] warmup ({args.warmup_steps} steps) {time.perf_counter() - start:.1f}s", flush=True)

    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    image = _infer(pipe, args.steps, args.size, args.seed)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

    path = out_dir / f"flux-klein-{tag}.png"
    image.save(path)
    print(
        f"[{tag}] {elapsed:.1f}s ({args.steps} steps) "
        f"| peak GPU {torch.cuda.max_memory_allocated() / 1e9:.2f} GB "
        f"| saved {path}",
        flush=True,
    )
    return image


def compare(reference, candidate) -> None:
    a = np.asarray(reference, dtype=np.float32)
    b = np.asarray(candidate, dtype=np.float32)
    mae = float(np.abs(a - b).mean())
    flat_a, flat_b = a.ravel(), b.ravel()
    cosine = float(flat_a @ flat_b / (np.linalg.norm(flat_a) * np.linalg.norm(flat_b)))
    print(f"[compare] native vs sage: MAE={mae:.3f}/255  cosine={cosine:.5f}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--backend",
        choices=["both", "native", "sage_hub"],
        default="both",
        help="Which attention backend(s) to run. 'both' also reports the numeric difference.",
    )
    parser.add_argument("--steps", type=int, default=4)
    parser.add_argument(
        "--warmup-steps",
        type=int,
        default=1,
        help="Steps for the discarded warmup generation run before each timed run. 0 disables it.",
    )
    parser.add_argument("--size", type=int, default=1024)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--out-dir", type=Path, default=Path.home())
    parser.add_argument("--model-id", default=MODEL_ID)
    parser.add_argument(
        "--no-offload",
        action="store_true",
        help="Keep the pipeline on the GPU instead of using enable_model_cpu_offload().",
    )
    parser.add_argument("--repo-id", default=STAGING_REPO_ID)
    parser.add_argument("--revision", default=STAGING_REVISION)
    parser.add_argument(
        "--no-staged",
        action="store_true",
        help="Resolve the published kernels-community kernel instead of a staged build.",
    )
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("This script needs a CUDA device.")

    capability = torch.cuda.get_device_capability(0)
    print(
        f"[env] torch {torch.__version__} (cuda {torch.version.cuda}) "
        f"| {torch.cuda.get_device_name(0)} sm{capability[0]}{capability[1]}",
        flush=True,
    )

    if not args.no_staged:
        use_staged_kernel(args.repo_id, args.revision)

    print(f"[load] {args.model_id} ...", flush=True)
    start = time.perf_counter()
    pipe = Flux2KleinPipeline.from_pretrained(args.model_id, torch_dtype=torch.bfloat16)
    if args.no_offload:
        pipe.to("cuda")
        placement = "resident on GPU"
    else:
        pipe.enable_model_cpu_offload()  # save some VRAM by offloading the model to CPU
        placement = "model cpu offload"
    print(f"[load] done in {time.perf_counter() - start:.1f}s ({placement})", flush=True)

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    native_image = None

    if args.backend in ("both", "native"):
        print("[run] native baseline", flush=True)
        native_image = generate(pipe, "native", args, out_dir)

    if args.backend in ("both", "sage_hub"):
        print("[run] sage_hub", flush=True)
        # `set_attention_backend` lives on ModelMixin, so it is set on the transformer rather
        # than on the pipeline.
        pipe.transformer.set_attention_backend("sage_hub")
        sage_image = generate(pipe, "sage", args, out_dir)
        if native_image is not None:
            compare(native_image, sage_image)


if __name__ == "__main__":
    main()

Native	Sage

This PR additionally adds a sage_blackwell_hub which is basically SAGE3 (the consumer blackwell variant of SAGE). Results:

Native	Sage Blackwell

Script

# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "torch==2.13.0", "kernels>=0.16", "transformers", "accelerate", "safetensors",
#   "huggingface_hub", "numpy", "Pillow", "sentencepiece", "protobuf",
# ]
# [tool.uv.sources]
# torch = { index = "pytorch-cu130" }
# pytorch-triton = { index = "pytorch-cu130" }
# [[tool.uv.index]]
# name = "pytorch-cu130"
# url = "https://download.pytorch.org/whl/cu130"
# explicit = true
# ///
"""PR #14584's benchmark script, with `sage_blackwell_hub` as the attention backend.

Needs an SM120 Blackwell GPU (RTX 50-series / RTX PRO 6000). On HF Jobs:

    hf jobs uv run flux2_klein_sage_blackwell.py \
        --flavor rtx-pro-6000 --secrets HF_TOKEN --timeout 60m \
        -v <diffusers-payload>:/diffusers:ro -d

where <diffusers-payload> holds the `src/` and `tests/` of the diffusers checkout under test.
Unlike the original there is no staging indirection: `kernels-community/sage-blackwell` v1 is
published, so the registry entry resolves it directly.

Set UPLOAD_REPO_ID to "" to keep the images in OUT_DIR instead of pushing them to the Hub.
"""

import os
import shutil
import sys
import time
from pathlib import Path

import numpy as np
import torch

DIFFUSERS_SRC = "/diffusers"
if Path(DIFFUSERS_SRC).exists():
    shutil.copytree(DIFFUSERS_SRC, "/work", dirs_exist_ok=True)
    sys.path.insert(0, "/work/src")

from diffusers import Flux2KleinPipeline  # noqa: E402

BACKEND = "sage_blackwell_hub"
MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
PROMPT = "A cat holding a sign that says hello world"
STEPS, WARMUP_STEPS, SIZE, SEED = 4, 1, 1024, 0
OUT_DIR = Path(os.environ.get("OUT_DIR", "/tmp/out"))
UPLOAD_REPO_ID = os.environ.get("UPLOAD_REPO_ID", "sayakpaul/sage-blackwell-flux2-outputs")


def _infer(pipe, steps):
    return pipe(
        prompt=PROMPT,
        height=SIZE,
        width=SIZE,
        guidance_scale=1.0,
        num_inference_steps=steps,
        generator=torch.Generator(device="cuda").manual_seed(SEED),
    ).images[0]


def generate(pipe, tag):
    if WARMUP_STEPS > 0:
        start = time.perf_counter()
        _infer(pipe, WARMUP_STEPS)
        torch.cuda.synchronize()
        print(f"[{tag}] warmup ({WARMUP_STEPS} steps) {time.perf_counter() - start:.1f}s", flush=True)

    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    image = _infer(pipe, STEPS)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

    path = OUT_DIR / f"flux-klein-{tag}.png"
    image.save(path)
    peak = torch.cuda.max_memory_allocated() / 1e9
    print(f"[{tag}] {elapsed:.1f}s ({STEPS} steps) | peak GPU {peak:.2f} GB | saved {path}", flush=True)
    return image, elapsed, peak


def compare(reference, candidate):
    a = np.asarray(reference, dtype=np.float32)
    b = np.asarray(candidate, dtype=np.float32)
    mae = float(np.abs(a - b).mean())
    fa, fb = a.ravel(), b.ravel()
    cosine = float(fa @ fb / (np.linalg.norm(fa) * np.linalg.norm(fb)))
    print(f"[compare] native vs {BACKEND}: MAE={mae:.3f}/255  cosine={cosine:.5f}", flush=True)
    return mae, cosine


capability = torch.cuda.get_device_capability(0)
device_name = torch.cuda.get_device_name(0)
print(
    f"[env] torch {torch.__version__} (cuda {torch.version.cuda}) "
    f"| {device_name} sm{capability[0]}{capability[1]}",
    flush=True,
)

OUT_DIR.mkdir(parents=True, exist_ok=True)
print(f"[load] {MODEL_ID} ...", flush=True)
start = time.perf_counter()
pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
pipe.to("cuda")
print(f"[load] done in {time.perf_counter() - start:.1f}s (resident on GPU)", flush=True)

print("[run] native baseline", flush=True)
native_image, native_s, native_gb = generate(pipe, "native")

head_dim = pipe.transformer.config.attention_head_dim
print(f"[run] {BACKEND} (transformer head dim: {head_dim})", flush=True)
pipe.transformer.set_attention_backend(BACKEND)
sage_image, sage_s, sage_gb = generate(pipe, "sage_blackwell")
mae, cosine = compare(native_image, sage_image)

summary = f"""# FLUX.2-klein-4B: native vs `{BACKEND}`

Script: `flux2_klein_sage_blackwell.py` (adapted from
[diffusers#14584](https://github.com/huggingface/diffusers/pull/14584)).

| | device | steps | time | peak GPU |
|---|---|---|---|---|
| native | {device_name} sm{capability[0]}{capability[1]} | {STEPS} | {native_s:.1f}s | {native_gb:.2f} GB |
| `{BACKEND}` | {device_name} sm{capability[0]}{capability[1]} | {STEPS} | {sage_s:.1f}s | {sage_gb:.2f} GB |

- prompt: `{PROMPT}`
- {SIZE}x{SIZE}, guidance_scale 1.0, seed {SEED}, warmup {WARMUP_STEPS} step(s)
- transformer `attention_head_dim` = {head_dim} (the kernel accepts 64 or 128 only)
- difference vs native: **MAE {mae:.3f}/255, cosine {cosine:.5f}**

At {STEPS} steps the timings are too short to resolve a speed difference; treat them as a
smoke test, not a benchmark. The image difference is what FP4 attention costs.
"""
(OUT_DIR / "README.md").write_text(summary)

if UPLOAD_REPO_ID:
    from huggingface_hub import HfApi

    api = HfApi()
    api.create_repo(UPLOAD_REPO_ID, repo_type="dataset", private=True, exist_ok=True)
    api.upload_folder(folder_path=str(OUT_DIR), repo_id=UPLOAD_REPO_ID, repo_type="dataset")
    print(f"[upload] https://huggingface.co/datasets/{UPLOAD_REPO_ID}", flush=True)

[sayakpaul]

Cc: @asomoza. I will do a separate one for Sage Blackwell (Sage Attention 3).

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.