[wip][core] propagate sage attention updates.

[sayakpaul]

What does this PR do?

Refer to huggingface/kernels-community#1095. This PR will be marked for review after that PR is merged and the builds are successfully updated on the Hub repo.

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

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

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