From 2ffe6fbeaba07a443e236e91717de08f53cad943 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:38:53 -0700 Subject: [PATCH] Studio diffusion (Phase 10): attention-backend selection (#6701) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/perf_levers_probe.py | 238 +++++++++++++++++ studio/backend/core/inference/diffusion.py | 25 ++ .../core/inference/diffusion_attention.py | 245 ++++++++++++++++++ studio/backend/models/inference.py | 29 +++ studio/backend/routes/inference.py | 1 + .../backend/tests/test_diffusion_attention.py | 225 ++++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 23 ++ 7 files changed, 786 insertions(+) create mode 100644 scripts/perf_levers_probe.py create mode 100644 studio/backend/core/inference/diffusion_attention.py create mode 100644 studio/backend/tests/test_diffusion_attention.py diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py new file mode 100644 index 0000000000..3964ab46ad --- /dev/null +++ b/scripts/perf_levers_probe.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Measure the next-phase diffusion levers on the real model, vs today's compiled baseline. + +Variants (Z-Image dense bf16, regional compile = the shipped "default" speed profile): + baseline -- channels_last + compile_repeated_blocks (reference image) + inductor_flags -- + the lossless inductor autotune flags (conv_1x1_as_mm, + coordinate_descent_tuning(+all_dirs), epilogue_fusion=False) + attn_cudnn -- + set_attention_backend("_native_cudnn") (exact) + attn_flash4 -- + set_attention_backend("flash_4_hub") (exact, SM100) + attn_sage -- + set_attention_backend("sage") (INT8 QK, quantized) + fbcache -- + First-Block-Cache (threshold 0.12) (few-step headroom test) + +Reports median latency, vs-baseline speedup, peak VRAM, and LPIPS vs baseline. One CUDA GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "perf_levers_images" + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + + # Keep the metric model on CPU: caching it on CUDA leaves it resident across + # variants, and each run resets peak-memory stats, so its VRAM would be charged + # to (and reduce headroom for) every later variant's measurement. + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval() + + def t(x): + return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _set_inductor_flags(): + import torch._inductor.config as ic + + ic.conv_1x1_as_mm = True + ic.coordinate_descent_tuning = True + ic.coordinate_descent_check_all_directions = True + ic.epilogue_fusion = False + try: + ic.force_fuse_int_mm_with_mul = True + except Exception: # noqa: BLE001 + pass + + +def _reset_inductor_flags(): + import torch._inductor.config as ic + + ic.conv_1x1_as_mm = False + ic.coordinate_descent_tuning = False + ic.coordinate_descent_check_all_directions = False + ic.epilogue_fusion = True + # Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into + # every later compiled row and the attention/fbcache measurements stop being isolated. + try: + ic.force_fuse_int_mm_with_mul = False + except Exception: # noqa: BLE001 + pass + + +def _load(): + import diffusers + import torch + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + try: + pipe.vae.to(memory_format = torch.channels_last) + except Exception: # noqa: BLE001 + pass + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def run( + tag, + steps, + seed, + res, + iters, + *, + attn = None, + fbcache = None, + inductor = False, +): + import torch + + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + _reset_inductor_flags() + if inductor: + _set_inductor_flags() + pipe = _load() + note = "" + if attn is not None: + try: + pipe.transformer.set_attention_backend(attn) + except Exception as exc: # noqa: BLE001 + note = f"attn({attn})={type(exc).__name__}:{str(exc)[:60]}" + print(f" [{tag}] {note}", flush = True) + del pipe # free the resident pipe so a skipped variant doesn't leak VRAM + torch.cuda.empty_cache() + return None + if fbcache is not None: + try: + from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbcache)) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] fbcache={type(exc).__name__}:{str(exc)[:60]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] compile={type(exc).__name__}:{str(exc)[:60]}", flush = True) + try: + _gen(pipe, steps, seed, res) # warmup / compile + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" [{tag}] FAILED first gen: {type(exc).__name__}:{str(exc)[:80]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, peak + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + args = p.parse_args(argv) + s, r, seed, it = args.steps, args.res, args.seed, args.iters + + print(f"== perf levers (Z-Image dense, {r}px, {s} steps) ==", flush = True) + base = run("baseline", s, seed, r, it) + if base is None: + print("baseline FAILED", flush = True) + return 1 + bmed, ref, bpeak = base + print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush = True) + rows = [("baseline", bmed, bpeak, 0.0)] + + variants = [ + ("inductor_flags", dict(inductor = True)), + ("attn_cudnn", dict(attn = "_native_cudnn")), + ("attn_flash4", dict(attn = "flash_4_hub")), + ("attn_sage", dict(attn = "sage")), + ("attn_sage_inductor", dict(attn = "sage", inductor = True)), + ("fbcache_0p12", dict(fbcache = 0.12)), + ] + for tag, kw in variants: + out = run(tag, s, seed, r, it, **kw) + if out is None: + rows.append((tag, None, None, None)) + continue + med, arr, peak = out + lp = _lpips(ref, arr) + rows.append((tag, med, peak, lp)) + spd = f"{bmed/med:.2f}x" if med else "-" + print(f" {tag:20s} {med:.3f}s ({spd} vs base) peak={peak:.1f}G LPIPS={lp}", flush = True) + + print("\n==== SUMMARY (ref = baseline compile) ====", flush = True) + for tag, med, peak, lp in rows: + if med is None: + print(f" {tag:20s} FAILED") + continue + spd = f"{bmed/med:.2f}x" if med else "-" + lpv = "ref" if (tag == "baseline") else (f"{lp:.3f}" if lp is not None else "n/a") + print(f" {tag:20s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush = True) + print("PERF-LEVERS-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index aac4069833..b61badb50b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -52,6 +52,10 @@ from .diffusion_speed import ( restore_backend_flags, snapshot_backend_flags, ) +from .diffusion_attention import ( + apply_attention_backend, + select_attention_backend, +) from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( load_prequantized_transformer, @@ -96,6 +100,9 @@ class _LoadState: # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. transformer_quant: Optional[str] = None + # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or + # None for the default SDPA. Set before compile; orthogonal to the weight quant. + attention_backend: Optional[str] = None @dataclass @@ -306,6 +313,7 @@ class DiffusionBackend: transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -340,6 +348,7 @@ class DiffusionBackend: transformer_quant = transformer_quant, transformer_quant_fast_accum = transformer_quant_fast_accum, transformer_prequant_path = transformer_prequant_path, + attention_backend = attention_backend, _load_token = token, ), daemon = True, @@ -476,6 +485,7 @@ class DiffusionBackend: transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -596,6 +606,18 @@ class DiffusionBackend: # first so unload can restore them: TF32 / cudnn.benchmark are global, # and a later `off` load must not inherit this load's settings. backend_flags_before = snapshot_backend_flags() + # Pick the attention kernel BEFORE compile (compile traces attention). auto + # upgrades to cuDNN fused attention on NVIDIA when a speed profile is active + # (~1.18x, near-lossless); an explicit backend is honored, falling back to + # the diffusers default if its kernel is unavailable. Orthogonal to the + # weight quant -- it speeds the QK/PV matmuls torchao does not touch. + attention_engaged = apply_attention_backend( + pipe, + select_attention_backend( + target, attention_backend, speed_active = effective_speed != SPEED_OFF + ), + logger = logger, + ) speed_applied = apply_speed_optims( pipe, target, @@ -649,6 +671,7 @@ class DiffusionBackend: backend_flags_before = backend_flags_before, text_encoder_quant = te_quant, transformer_quant = transformer_quant_engaged, + attention_backend = attention_engaged, ) logger.info( @@ -953,6 +976,7 @@ class DiffusionBackend: "speed_optims": [], "text_encoder_quant": None, "transformer_quant": None, + "attention_backend": None, } return { "loaded": True, @@ -969,6 +993,7 @@ class DiffusionBackend: "speed_optims": list(state.speed_optims), "text_encoder_quant": state.text_encoder_quant, "transformer_quant": state.transformer_quant, + "attention_backend": state.attention_backend, } diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py new file mode 100644 index 0000000000..e652b0068c --- /dev/null +++ b/studio/backend/core/inference/diffusion_attention.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Select the diffusion transformer's attention backend. + +diffusers exposes a unified ``transformer.set_attention_backend(name)`` dispatcher that +swaps the scaled-dot-product-attention kernel, validating hardware/package requirements at +set time and otherwise leaving the default (``native`` = ``F.scaled_dot_product_attention``). +Attention is memory-bandwidth bound, so a better kernel is a real end-to-end win that is +orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never +touches) and composes with torch.compile. + + auto - the best *exact* (non-quantized) backend for the device. On NVIDIA CUDA that is + cuDNN's fused attention (``_native_cudnn``), measured ~1.18x end-to-end on a B200 + with LPIPS ~0.004 vs the default (below the compile/quant noise floor). On + AMD/Intel/Apple/CPU it stays ``native`` (the dispatcher already routes those). + ``auto`` only upgrades when a speed profile is active, so ``speed_mode=off`` stays + bit-identical. + native - force the default SDPA (bit-identical reference). + cudnn - cuDNN fused attention (exact; NVIDIA). + flash / flash3 / flash4 - FlashAttention 2 / 3 (Hopper) / 4 (SM100); exact, kernel-gated. + sage - SageAttention (INT8 QK); quantized, a small quality cost, consumer-friendly. + xformers / aiter - memory-efficient (NVIDIA) / AITER (AMD ROCm). + +Best-effort: an unavailable backend (missing kernel / wrong arch) is caught and the load +falls back to the diffusers default rather than failing. torch/diffusers imported lazily. +""" + +from __future__ import annotations + +from typing import Any, Optional + +ATTN_AUTO = "auto" +ATTN_NATIVE = "native" + +# User-facing alias -> the diffusers dispatcher backend name. +_ALIASES: dict[str, str] = { + "native": "native", + "sdpa": "native", + "cudnn": "_native_cudnn", + "flash": "flash", + "flash2": "flash", + "flash3": "_flash_3_hub", + "flash4": "flash_4_hub", + "sage": "sage", + "xformers": "xformers", + "aiter": "aiter", +} +ATTN_ALIASES = (ATTN_AUTO,) + tuple(dict.fromkeys(_ALIASES)) + + +def normalize_attention_backend(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested attention backend; None / "" / "auto" -> "auto". + + Raises ValueError for an unsupported alias so a bad request is rejected cheaply.""" + if value is None: + return ATTN_AUTO + normalized = str(value).strip().lower() + if not normalized: + return ATTN_AUTO + if normalized not in ATTN_ALIASES: + raise ValueError( + f"Unsupported attention_backend '{value}'. Use one of: {', '.join(ATTN_ALIASES)}." + ) + return normalized + + +# Backends diffusers validates only by *package* at set time (``_check_attention_backend_ +# requirements`` checks the ``kernels`` install, not the GPU), but whose kernels need a +# specific CUDA arch at run time -- so an explicit request on the wrong card loads/sets fine +# and then crashes mid-generation. Gate them up front by a (min, max-exclusive) compute +# capability range. FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so it +# needs an upper bound: an explicit flash3 on a B200 (SM100) must drop to native instead of +# setting fine then crashing at generation. FlashAttention 4 is Blackwell+ (no upper bound). +_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = { + "_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only + "flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+ +} + + +def _cuda_capability() -> Optional[tuple[int, int]]: + """(major, minor) compute capability of the active CUDA device, or None if unknown.""" + try: + import torch + if not torch.cuda.is_available(): + return None + return tuple(torch.cuda.get_device_capability()) # type: ignore[return-value] + except Exception: # noqa: BLE001 + return None + + +def _backend_arch_supported(backend: str) -> bool: + """False only when ``backend`` needs a CUDA arch outside this device's supported range. + + Unknown capability (no CUDA / detection failure) returns True so we never block on a + guess -- diffusers' own set-time check still guards the package, and a genuine run-time + failure falls back to native.""" + bounds = _ARCH_CAPABILITY.get(backend) + if bounds is None: + return True + have = _cuda_capability() + if have is None: + return True + low, high = bounds + return have >= low and (high is None or have < high) + + +def _is_cuda_nvidia(target: Any) -> bool: + """CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies.""" + if getattr(target, "device", None) != "cuda": + return False + try: + import torch + return getattr(torch.version, "hip", None) is None + except Exception: # noqa: BLE001 + return False + + +def select_attention_backend( + target: Any, requested: Optional[str], *, speed_active: bool +) -> Optional[str]: + """The dispatcher backend name to apply, or None to leave the diffusers default. + + An explicit alias is honored verbatim (apply falls back if its kernel is unavailable). + ``auto`` upgrades to cuDNN on NVIDIA CUDA only when a speed profile is active (so + ``off`` stays bit-identical); everywhere else it returns None (native default).""" + alias = normalize_attention_backend(requested) + if alias != ATTN_AUTO: + backend = _ALIASES[alias] + if backend == "native": + return None + # An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine + # then crash mid-generation, so drop it to the native default up front. + if not _backend_arch_supported(backend): + return None + # cuDNN fused SDPA needs Ampere+ (SM80); diffusers accepts it on pre-SM80 cards + # (T4/V100) then fails at the first generation, so apply the same gate to an + # explicit cuDNN request as the auto path already does. + if backend == "_native_cudnn" and not _cudnn_attention_supported(): + return None + return backend + # auto + if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported(): + return "_native_cudnn" + return None + + +def _cudnn_attention_supported() -> bool: + """cuDNN fused SDPA needs Ampere+ (SM80). On pre-SM80 NVIDIA cards (T4 SM75 / + V100 SM70) diffusers accepts ``_native_cudnn`` at set time but the kernel fails at + the first generation, so gate the auto-cuDNN upgrade on capability. Unknown + capability allows it (diffusers' set-time check + the run-time fallback still guard).""" + have = _cuda_capability() + return have is None or have >= (8, 0) + + +def apply_attention_backend( + pipe: Any, + backend: Optional[str], + *, + logger: Any = None, +) -> Optional[str]: + """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher. + + Returns the backend actually engaged, or None when left at the native default (either + because ``backend`` was None or because the requested kernel was unavailable -> graceful + fallback, never a load failure). + + diffusers keeps a *process-wide* active attention backend that ``set_attention_backend`` + also updates, and a fresh transformer's processors follow it (their ``_attention_backend`` + defaults to None). So a load that wants native must restore it explicitly: otherwise it + silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), + breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" + transformer = getattr(pipe, "transformer", None) + fn = getattr(transformer, "set_attention_backend", None) + if not callable(fn): + return None + if backend is not None: + try: + fn(backend) + # set_attention_backend also pins the backend in diffusers' process-wide + # registry. This transformer's own processors keep it locally (their + # _attention_backend is now explicit), so reset the global default back to + # native -- otherwise a later component whose processors are unconfigured + # (backend None) silently inherits this kernel. + _reset_global_backend_to_native(logger) + if logger is not None: + logger.info("diffusion.attention: backend=%s", backend) + return backend + except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below + _warn(logger, backend, exc) + # No backend requested, or the requested one failed: pin the native default so a stale + # process-wide backend from a previous load can't leak into this one. + _restore_native_backend(fn, logger) + return None + + +def _active_attention_backend() -> Optional[str]: + """The diffusers process-wide active attention backend name, or None if undeterminable.""" + try: + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + # get_active_backend() returns a (AttentionBackendName, fn) tuple (or None), so + # take element 0 and read its .value (e.g. "native"); reading .value off the + # tuple itself would yield a junk string that never compares equal to a name. + active = _AttentionBackendRegistry.get_active_backend() + if active is None: + return None + name = active[0] if isinstance(active, tuple) else active + return getattr(name, "value", str(name)) + except Exception: # noqa: BLE001 + return None + + +def _reset_global_backend_to_native(logger: Any) -> None: + """Reset diffusers' process-wide active attention backend to native after a + successful per-transformer set, so a later component whose processors are + unconfigured (backend None) does not inherit this transformer's kernel. The + transformer's own processors keep the backend just set. Best-effort and silent: + if the diffusers internals move, the prior (leaking) behavior is unchanged.""" + if _active_attention_backend() == ATTN_NATIVE: + return + try: + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + except Exception: # noqa: BLE001 — best-effort; leave the global as-is on any change + pass + + +def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None: + """Force the native default when the global active backend isn't already native.""" + if _active_attention_backend() == ATTN_NATIVE: + return # already native -> avoid redundant work and an extra dispatcher warning + try: + set_backend_fn(ATTN_NATIVE) + except Exception as exc: # noqa: BLE001 — best-effort restore + _warn(logger, ATTN_NATIVE, exc) + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 08d309f257..402c931da1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1750,6 +1750,30 @@ class DiffusionLoadRequest(BaseModel): "the OS path separator). A bare on/off value such as '1' is deliberately not " "accepted -- it must name an allowed directory.", ) + attention_backend: Optional[ + Literal[ + "auto", + "native", + "sdpa", + "cudnn", + "flash", + "flash2", + "flash3", + "flash4", + "sage", + "xformers", + "aiter", + ] + ] = Field( + None, + description = "Attention kernel via the diffusers dispatcher. auto picks the best " + "exact backend for the device (cuDNN fused attention on NVIDIA, ~1.18x and " + "near-lossless, when a speed profile is active; native SDPA elsewhere and when " + "speed=off). native (alias sdpa) forces default SDPA; cudnn/flash/flash3/flash4 are exact " + "(kernel/arch-gated); sage is INT8 attention (a small quality cost, consumer " + "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An " + "unavailable kernel falls back to the default.", + ) class DiffusionGenerateRequest(BaseModel): @@ -1865,3 +1889,8 @@ class DiffusionStatusResponse(BaseModel): description = "Transformer quant engaged on the dense fast path: int8 | fp8 | " "nvfp4 | mxfp8 | null (null = the GGUF transformer was loaded)", ) + attention_backend: Optional[str] = Field( + None, + description = "Attention backend engaged via the diffusers dispatcher (e.g. " + "_native_cudnn), or null for the default SDPA", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ec5382e190..c89a0cc7eb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10334,6 +10334,7 @@ async def load_diffusion_model( transformer_quant = request.transformer_quant, transformer_quant_fast_accum = request.transformer_quant_fast_accum, transformer_prequant_path = request.transformer_prequant_path, + attention_backend = request.attention_backend, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py new file mode 100644 index 0000000000..3e43aab8a9 --- /dev/null +++ b/studio/backend/tests/test_diffusion_attention.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for attention-backend selection. No torch/diffusers needed: +``_is_cuda_nvidia`` is monkeypatched for the policy tests, and the apply path uses a fake +transformer that records / raises on ``set_attention_backend``. +""" + +from __future__ import annotations + +import types + +import pytest + +import core.inference.diffusion_attention as att +from core.inference.diffusion_attention import ( + ATTN_AUTO, + apply_attention_backend, + normalize_attention_backend, + select_attention_backend, +) + + +def _target(device = "cuda"): + return types.SimpleNamespace(device = device) + + +# ── normalize ──────────────────────────────────────────────────────────────────── +def test_normalize_defaults_and_aliases(): + assert normalize_attention_backend(None) == ATTN_AUTO + assert normalize_attention_backend("") == ATTN_AUTO + assert normalize_attention_backend("auto") == ATTN_AUTO + assert normalize_attention_backend("CuDNN") == "cudnn" + assert normalize_attention_backend("FLASH3") == "flash3" + assert normalize_attention_backend("sdpa") == "sdpa" + + +def test_normalize_rejects_unknown(): + with pytest.raises(ValueError): + normalize_attention_backend("bogus") + # dashes are no longer silently rewritten to underscores -> a dashed alias is rejected. + with pytest.raises(ValueError): + normalize_attention_backend("flash-3") + + +def test_sdpa_alias_maps_to_native(): + # sdpa is an alias for native -> nothing to set on the dispatcher. + assert select_attention_backend(_target(), "sdpa", speed_active = True) is None + + +# ── select policy ───────────────────────────────────────────────────────────────── +def test_auto_upgrades_to_cudnn_on_nvidia_when_speed_active(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) # Ampere+: cuDNN ok + assert select_attention_backend(_target(), "auto", speed_active = True) == "_native_cudnn" + + +def test_auto_does_not_pin_cudnn_below_sm80(monkeypatch): + # cuDNN fused SDPA fails at run time on pre-SM80 (T4 SM75 / V100 SM70); auto must stay + # on the native default there rather than pin a backend that crashes on first generation. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) # Turing T4 + assert select_attention_backend(_target(), "auto", speed_active = True) is None + + +def test_auto_stays_native_when_speed_off(monkeypatch): + # off must stay bit-identical -> no backend change even on NVIDIA. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + assert select_attention_backend(_target(), "auto", speed_active = False) is None + + +def test_auto_stays_native_off_nvidia(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) + assert select_attention_backend(_target(device = "mps"), "auto", speed_active = True) is None + + +def test_explicit_backend_honored_regardless_of_speed(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) + # Pin a high capability so the arch-gated flash4 isn't dropped by the runtime check. + monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0)) + assert select_attention_backend(_target(), "sage", speed_active = False) == "sage" + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" + + +def test_explicit_native_returns_none(): + # native is the default -> nothing to set. + assert select_attention_backend(_target(), "native", speed_active = True) is None + + +# ── arch gating (flash3/flash4 need a specific CUDA capability) ───────────────────── +def test_flash3_dropped_below_hopper(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 9)) # Ada / consumer + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + + +def test_flash4_dropped_below_blackwell(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) # Hopper, but FA4 needs SM100 + assert select_attention_backend(_target(), "flash4", speed_active = False) is None + # flash3 still allowed on Hopper. + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +def test_arch_gate_does_not_block_when_capability_unknown(monkeypatch): + # Unknown capability (e.g. no CUDA) must not block -> diffusers' set-time check still guards. + monkeypatch.setattr(att, "_cuda_capability", lambda: None) + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + + +def test_flash3_dropped_on_blackwell(monkeypatch): + # FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel: an explicit + # flash3 on a B200 (SM100) must drop to native rather than set fine then crash. + monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + # FA4 is still honored on Blackwell. + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + # flash3 is allowed exactly on Hopper SM90. + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +def test_explicit_cudnn_dropped_below_sm80(monkeypatch): + # An explicit cuDNN request on pre-Ampere (T4 SM75 / V100 SM70) must drop to native, + # not set fine and crash at first generation -- the same gate the auto path applies. + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) is None + # Ampere+ still honors it. + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" + + +# ── apply ───────────────────────────────────────────────────────────────────────── +class _FakeTransformer: + def __init__(self, *, fail = False): + self.fail = fail + self.set_to = None + + def set_attention_backend(self, name): + if self.fail: + raise RuntimeError(f"{name} kernel unavailable") + self.set_to = name + + +def _pipe(transformer): + return types.SimpleNamespace(transformer = transformer) + + +def test_apply_none_leaves_native_when_global_already_native(monkeypatch): + # Global already native -> no redundant set call, returns None. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to is None + + +def test_apply_none_restores_native_when_global_polluted(monkeypatch): + # A previous load pinned cuDNN process-wide; a native load must reset it so it can't + # silently inherit cuDNN (the bit-identical/off guarantee). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to == "native" + + +def test_apply_sets_backend(): + t = _FakeTransformer() + engaged = apply_attention_backend(_pipe(t), "_native_cudnn") + assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" + + +def test_apply_falls_back_on_unavailable_kernel(monkeypatch): + # an unavailable kernel must not fail the load -> returns None (diffusers default). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer(fail = True) + assert apply_attention_backend(_pipe(t), "sage") is None + + +def test_apply_failed_kernel_restores_native_when_polluted(monkeypatch): + # Requested kernel fails AND the global is polluted: restore native before returning. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + + class _FailOnceTransformer: + def __init__(self): + self.calls = [] + + def set_attention_backend(self, name): + self.calls.append(name) + if name != "native": + raise RuntimeError(f"{name} kernel unavailable") + + t = _FailOnceTransformer() + assert apply_attention_backend(_pipe(t), "sage") is None + assert t.calls == ["sage", "native"] + + +def test_apply_handles_missing_method(): + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert apply_attention_backend(pipe, "_native_cudnn") is None + + +def test_apply_resets_global_registry_after_success(monkeypatch): + # After a successful per-transformer set, the process-wide registry must be reset to + # native so a later component (unconfigured processors) can't inherit this kernel -- + # while the transformer's own backend stays the engaged one. + called = {"reset": False} + monkeypatch.setattr( + att, "_reset_global_backend_to_native", lambda logger: called.__setitem__("reset", True) + ) + t = _FakeTransformer() + engaged = apply_attention_backend(_pipe(t), "_native_cudnn") + assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" + assert called["reset"] is True + + +def test_active_attention_backend_reads_tuple_return(): + # get_active_backend() returns a (AttentionBackendName, fn) tuple; the helper must read + # the name's .value, not stringify the tuple (which never compares equal to a name). + pytest.importorskip("diffusers") + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + assert att._active_attention_backend() == "native" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 839d325443..dc34e60cef 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -424,6 +424,29 @@ def test_transformer_prequant_path_threads_through(client, monkeypatch): assert backend.last_load_kwargs.get("transformer_prequant_path") == "/data/zimage_fp8.pt" +def test_attention_backend_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "attention_backend": "cudnn", + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("attention_backend") == "cudnn" + + +def test_invalid_attention_backend_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "attention_backend": "bogus"}, + ) + assert resp.status_code == 422 + + def test_prequant_path_doc_describes_allowlist_not_toggle(): # The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a # directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots