[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

importargparseimporttimefrompathlibimportPathimportnumpyasnpimporttorchfromdiffusersimportFlux2KleinPipelinefromdiffusers.models.attention_dispatchimportAttentionBackendName, _HUB_KERNELS_REGISTRYSTAGING_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"defuse_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_idconfig.revision=revisionconfig.version=None# `revision` pins the build; a version pin would fight itprint(f"[setup] sage_hub -> {config.repo_id}@{config.revision}", flush=True)
def_infer(pipe, steps: int, size: int, seed: int):
returnpipe(
prompt=PROMPT,
height=size,
width=size,
guidance_scale=1.0,
num_inference_steps=steps,
generator=torch.Generator(device="cuda").manual_seed(seed),
).images[0]
defgenerate(pipe, tag: str, args, out_dir: Path):
ifargs.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() -startpath=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,
)
returnimagedefcompare(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)
defmain() ->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()
ifnottorch.cuda.is_available():
raiseSystemExit("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,
)
ifnotargs.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)
ifargs.no_offload:
pipe.to("cuda")
placement="resident on GPU"else:
pipe.enable_model_cpu_offload() # save some VRAM by offloading the model to CPUplacement="model cpu offload"print(f"[load] done in {time.perf_counter() -start:.1f}s ({placement})", flush=True)
out_dir=args.out_dirout_dir.mkdir(parents=True, exist_ok=True)
native_image=Noneifargs.backendin ("both", "native"):
print("[run] native baseline", flush=True)
native_image=generate(pipe, "native", args, out_dir)
ifargs.backendin ("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)
ifnative_imageisnotNone:
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 -dwhere <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 ispublished, 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."""importosimportshutilimportsysimporttimefrompathlibimportPathimportnumpyasnpimporttorchDIFFUSERS_SRC="/diffusers"ifPath(DIFFUSERS_SRC).exists():
shutil.copytree(DIFFUSERS_SRC, "/work", dirs_exist_ok=True)
sys.path.insert(0, "/work/src")
fromdiffusersimportFlux2KleinPipeline# noqa: E402BACKEND="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, 0OUT_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):
returnpipe(
prompt=PROMPT,
height=SIZE,
width=SIZE,
guidance_scale=1.0,
num_inference_steps=steps,
generator=torch.Generator(device="cuda").manual_seed(SEED),
).images[0]
defgenerate(pipe, tag):
ifWARMUP_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() -startpath=OUT_DIR/f"flux-klein-{tag}.png"image.save(path)
peak=torch.cuda.max_memory_allocated() /1e9print(f"[{tag}] {elapsed:.1f}s ({STEPS} steps) | peak GPU {peak:.2f} GB | saved {path}", flush=True)
returnimage, elapsed, peakdefcompare(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)
returnmae, cosinecapability=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_dimprint(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 asmoke test, not a benchmark. The image difference is what FP4 attention costs."""
(OUT_DIR/"README.md").write_text(summary)
ifUPLOAD_REPO_ID:
fromhuggingface_hubimportHfApiapi=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.

[asomoza]

thanks! I did a test with a 5090 and Minimax-H3, the speedup is around ~2X, but I see some quality degradation specially on text, it is what I expected from last time I tried it. Since it is with a 5090 I used the SDNQ 8-Bit quantized model and with group offload.

Native

t2va_sdnq_8bit_pruned_42.1.mp4

Sage

t2va_sdnq_8bit_pruned_sage_42.mp4

What is unexpected to me is that it uses quite a bit more of VRAM.

	denoise, 19 steps	s/step	board peak	torch peak
native	10:06	31.8	20.21 GiB	11.96 GiB
sage, cold	6:31	16.0	~26.6 GiB	19.98 GiB
sage, warm	4:53	15.5	27.07 GiB	19.98 GiB

[sayakpaul]

Thanks @asomoza. What about a bit less complex prompts? Do we see equal degradations?

@DN6 if I could have your 👀 as well :)

[asomoza]

What about a bit less complex prompts? Do we see equal degradations?

yeah @sayakpaul, I can see it on all generations, but the more evident ones are the text and ones that needs more fine details, even with simple prompts, so that's not really the factor that determines it. Probably a simple one with a static subject and no details or a cartoon one will be more imperceptible.

Also I can see it in this video for example, but not sure if all people can see it too, by no means is bad, most of the times it will be a good trade off for speed so this is still really good IMO.

t2va_capybara_sage_42.mp4