From 59b84d263f2d8a3f183ed75587cb2cd0368f2453 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 12:31:43 -0500 Subject: [PATCH 01/17] first --- modules/launch_utils.py | 54 +++++--- modules/mps_flash_attention.py | 179 +++++++++++++++++++++++++ modules/mps_utils.py | 37 +++++ modules/sd_hijack_optimizations.py | 119 +++++++++++++++- modules/sub_quadratic_attention.py | 34 +++-- requirements_macos.txt | 3 + requirements_versions.txt | 7 +- scripts/benchmark_mps_attention.py | 97 ++++++++++++++ scripts/install_mps_flash_attention.py | 105 +++++++++++++++ test/test_mps_flash_attention.py | 36 +++++ test/test_mps_utils.py | 45 +++++++ test/test_sub_quadratic_attention.py | 46 +++++++ webui-macos-env.sh | 9 +- 13 files changed, 730 insertions(+), 41 deletions(-) create mode 100644 modules/mps_flash_attention.py create mode 100644 modules/mps_utils.py create mode 100644 requirements_macos.txt create mode 100644 scripts/benchmark_mps_attention.py create mode 100644 scripts/install_mps_flash_attention.py create mode 100644 test/test_mps_flash_attention.py create mode 100644 test/test_mps_utils.py create mode 100644 test/test_sub_quadratic_attention.py diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 804b802057a..9c9cd2de4d3 100644 --- a/modules/launch_utils.py +++ b/modules/launch_utils.py @@ -275,9 +275,6 @@ def run_extensions_installers(settings_file): startup_timer.record(dirname_extension) -re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*") - - def requirements_met(requirements_file): """ Does a simple parse of a requirements.txt file to determine if all rerqirements in it @@ -285,29 +282,28 @@ def requirements_met(requirements_file): """ import importlib.metadata - import packaging.version + import packaging.requirements with open(requirements_file, "r", encoding="utf8") as file: for line in file: - if line.strip() == "": + line = line.split("#", 1)[0].strip() + if not line: continue - m = re.match(re_requirement, line) - if m is None: + try: + requirement = packaging.requirements.Requirement(line) + except packaging.requirements.InvalidRequirement: return False - package = m.group(1).strip() - version_required = (m.group(2) or "").strip() - - if version_required == "": + if requirement.marker is not None and not requirement.marker.evaluate(): continue try: - version_installed = importlib.metadata.version(package) - except Exception: + version_installed = importlib.metadata.version(requirement.name) + except importlib.metadata.PackageNotFoundError: return False - if packaging.version.parse(version_required) != packaging.version.parse(version_installed): + if requirement.specifier and version_installed not in requirement.specifier: return False return True @@ -427,17 +423,37 @@ def ensure_build_dependencies(): """Ensure essential build tools are available""" if not is_installed("wheel"): run_pip("install wheel", "wheel") - # Check setuptools version compatibility + # Keep the historical A1111 pin for compatibility with old packages. try: - setuptools_version = run(f'"{python}" -c "import setuptools; print(setuptools.__version__)"', None, None).strip() - if setuptools_version >= "70": + setuptools_version = importlib.metadata.version("setuptools") + setuptools_match = re.match(r"(\d+)\.(\d+)", setuptools_version) + setuptools_tuple = tuple(map(int, setuptools_match.groups())) if setuptools_match else (0, 0) + if setuptools_tuple >= (70, 0): run_pip("install setuptools==69.5.1", "setuptools") except Exception: - # If setuptools check fails, install compatible version - run_pip("install setuptools==69.5.1", "setuptools") + print("Could not validate setuptools compatibility; leaving the installed version unchanged.") # Install build dependencies early ensure_build_dependencies() + if platform.system() == "Darwin" and platform.machine() == "arm64": + mps_flash_installer = os.environ.get("MPS_FLASH_ATTENTION_INSTALLER", "") + torch_version_match = re.match(r"(\d+)\.(\d+)", importlib.metadata.version("torch")) + torch_version = tuple(map(int, torch_version_match.groups())) if torch_version_match else (0, 0) + has_stream_safe_mps_flash = check_run_python( + "import metal_flash_sdpa; assert getattr(metal_flash_sdpa, 'A1111_MPS_STREAM_FIX', False)" + ) + if mps_flash_installer and torch_version >= (2, 3) and not has_stream_safe_mps_flash: + try: + run( + f'"{python}" "{mps_flash_installer}"', + "Installing stream-safe Metal Flash Attention", + "Couldn't install Metal Flash Attention", + live=True, + ) + startup_timer.record("install Metal Flash Attention") + except RuntimeError as exc: + print(f"Metal Flash Attention installation failed; continuing with native MPS attention: {exc}") + if not is_installed("clip"): run_pip(f"install --no-build-isolation {clip_package}", "clip") startup_timer.record("install clip") diff --git a/modules/mps_flash_attention.py b/modules/mps_flash_attention.py new file mode 100644 index 00000000000..0456a0fb0e0 --- /dev/null +++ b/modules/mps_flash_attention.py @@ -0,0 +1,179 @@ +"""Selective Draw Things-style Metal Flash Attention for Apple Silicon.""" + +from __future__ import annotations + +import importlib +import os +import platform +import subprocess +import sys + +import torch +import torch.nn.functional as F + + +MIN_TORCH_VERSION = (2, 3) + +_extension = None +_availability = None +_availability_error = None +_dispatch_count = 0 +_fallback_count = 0 +_runtime_failure_warned = False +_first_dispatch_logged = False +_native_sdpa_supports_gqa = "enable_gqa" in (F.scaled_dot_product_attention.__doc__ or "") + + +def _version_tuple(version): + components = version.split("+", 1)[0].split(".") + try: + return tuple(int(component) for component in components[:2]) + except ValueError: + return (0, 0) + + +def should_use_mfa_shape(query_tokens, key_tokens, head_dim): + """Return whether the measured M1 routing table favors MFA for this shape.""" + if query_tokens < 256: + return False + + # Dimension-40 attention wins by 2-3x and amortizes the Metal command + # buffer boundary. Smaller gains at dimensions 80/160 regress the full + # UNet because they introduce too many additional command buffers. + if head_dim == 40: + return True + return False + + +def _run_isolated_self_test(): + code = """ +import torch +import torch.nn.functional as F +from metal_flash_sdpa import MetalFlashAttentionForward + +q = torch.randn((1, 8, 256, 40), device='mps', dtype=torch.float16) +k = torch.randn_like(q) +v = torch.randn_like(q) +actual = MetalFlashAttentionForward.apply(q, k, v, 40 ** -0.5, False) +expected = F.scaled_dot_product_attention(q, k, v) +torch.mps.synchronize() +assert torch.isfinite(actual).all().item() +assert (actual.float() - expected.float()).abs().max().item() < 0.01 +""" + environment = os.environ.copy() + environment["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=90, + env=environment, + ) + if result.returncode != 0: + detail = result.stderr.strip().splitlines() + raise RuntimeError(detail[-1] if detail else f"native self-test exited {result.returncode}") + + +def is_available(): + """Load and crash-test the optional native extension once.""" + global _extension, _availability, _availability_error + if _availability is not None: + return _availability + + if platform.system() != "Darwin" or platform.machine() != "arm64": + _availability = False + return False + if not torch.backends.mps.is_available(): + _availability = False + return False + if _version_tuple(torch.__version__) < MIN_TORCH_VERSION: + _availability_error = f"requires PyTorch {MIN_TORCH_VERSION[0]}.{MIN_TORCH_VERSION[1]} or newer" + _availability = False + return False + + try: + _extension = importlib.import_module("metal_flash_sdpa") + if not hasattr(_extension, "MetalFlashAttentionForward"): + raise RuntimeError("extension does not expose MetalFlashAttentionForward") + if _version_tuple(torch.__version__) < (2, 11) and not getattr(_extension, "A1111_MPS_STREAM_FIX", False): + raise RuntimeError("native extension is missing the A1111 MPS stream safety patch") + _run_isolated_self_test() + except (ImportError, OSError, RuntimeError, subprocess.SubprocessError) as exc: + _availability_error = str(exc) + _availability = False + print(f"Metal Flash Attention unavailable: {_availability_error}") + return False + + _availability = True + print("Metal Flash Attention native self-test passed; selective MFA routing enabled.") + return True + + +def _can_dispatch(query, key, value, attn_mask, dropout_p, enable_gqa, training): + if not is_available() or training or enable_gqa or dropout_p != 0.0 or attn_mask is not None: + return False + if query.device.type != "mps" or key.device != query.device or value.device != query.device: + return False + if query.dtype != torch.float16 or key.dtype != query.dtype or value.dtype != query.dtype: + return False + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + return False + if query.shape[0] != key.shape[0] or key.shape != value.shape: + return False + if query.shape[1] != key.shape[1] or query.shape[-1] != key.shape[-1]: + return False + if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad): + return False + return should_use_mfa_shape(query.shape[-2], key.shape[-2], query.shape[-1]) + + +def scaled_dot_product_attention( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + enable_gqa=False, + *, + training=False, +): + """Use native MFA for measured winning shapes, otherwise use PyTorch SDPA.""" + global _dispatch_count, _fallback_count, _runtime_failure_warned, _first_dispatch_logged + if _can_dispatch(query, key, value, attn_mask, dropout_p, enable_gqa, training): + try: + _dispatch_count += 1 + if not _first_dispatch_logged: + print(f"Metal Flash Attention first dispatch: Q={tuple(query.shape)}, K={tuple(key.shape)}") + _first_dispatch_logged = True + attention_scale = scale if scale is not None else query.shape[-1] ** -0.5 + return _extension.MetalFlashAttentionForward.apply(query, key, value, attention_scale, is_causal) + except RuntimeError as exc: + if not _runtime_failure_warned: + print(f"Metal Flash Attention failed; using PyTorch SDPA: {exc}") + _runtime_failure_warned = True + torch.mps.empty_cache() + + _fallback_count += 1 + native_kwargs = { + "attn_mask": attn_mask, + "dropout_p": dropout_p, + "is_causal": is_causal, + "scale": scale, + } + if enable_gqa: + if not _native_sdpa_supports_gqa: + raise RuntimeError("grouped-query attention requires a newer PyTorch SDPA runtime") + native_kwargs["enable_gqa"] = True + + return F.scaled_dot_product_attention(query, key, value, **native_kwargs) + + +def diagnostics(): + return { + "available": bool(_availability), + "error": _availability_error, + "dispatches": _dispatch_count, + "fallbacks": _fallback_count, + } diff --git a/modules/mps_utils.py b/modules/mps_utils.py new file mode 100644 index 00000000000..5e930043ca8 --- /dev/null +++ b/modules/mps_utils.py @@ -0,0 +1,37 @@ +import psutil + + +def should_use_sdp(batch, heads, query_tokens, key_tokens, element_size, total_memory=None, available_memory=None): + """Use native SDPA only while its intermediates fit unified memory safely.""" + if total_memory is None or available_memory is None: + memory = psutil.virtual_memory() + total_memory = memory.total + available_memory = memory.available + + attention_bytes = batch * heads * query_tokens * key_tokens * element_size + estimated_peak = int(attention_bytes * 2.5) + budget = min( + int(total_memory * 0.10), + int(available_memory * 0.20), + 1536 * 1024 * 1024, + ) + return estimated_peak <= budget + + +def attention_query_chunk_size(requested_size, batch_heads, key_tokens, element_size, total_memory=None, available_memory=None): + """Bound an MPS query tile so attention intermediates stay memory-safe.""" + if total_memory is None or available_memory is None: + memory = psutil.virtual_memory() + total_memory = memory.total + available_memory = memory.available + + peak_budget = min( + int(total_memory * 0.025), + int(available_memory * 0.10), + 384 * 1024 * 1024, + ) + bytes_per_query = max(int(batch_heads * key_tokens * element_size * 2.5), 1) + maximum_size = max(peak_budget // bytes_per_query, 1) + if maximum_size >= 64: + maximum_size = (maximum_size // 64) * 64 + return min(requested_size, maximum_size) diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index d9af5a0d4d6..c36e7f53c5f 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -9,7 +9,7 @@ from ldm.util import default from einops import rearrange -from modules import shared, errors, devices, sub_quadratic_attention +from modules import shared, errors, devices, sub_quadratic_attention, mps_flash_attention, mps_utils from modules.hypernetworks import hypernetwork import ldm.modules.attention @@ -92,6 +92,36 @@ def apply(self): sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward +class SdOptimizationMpsAdaptive(SdOptimizationSdpNoMem): + name = "mps-adaptive" + label = "native Metal attention with a memory-safe fallback" + priority = 1100 + + def is_available(self): + return shared.device.type == 'mps' and super().is_available() + + def apply(self): + ldm.modules.attention.CrossAttention.forward = mps_adaptive_attention_forward + ldm.modules.diffusionmodules.model.AttnBlock.forward = mps_adaptive_attnblock_forward + sgm.modules.attention.CrossAttention.forward = mps_adaptive_attention_forward + sgm.modules.diffusionmodules.model.AttnBlock.forward = mps_adaptive_attnblock_forward + + +class SdOptimizationMpsFlash(SdOptimizationSdpNoMem): + name = "mps-flash" + label = "Draw Things-style Metal Flash Attention with native fallback" + priority = 1200 + + def is_available(self): + return shared.device.type == 'mps' and mps_flash_attention.is_available() + + def apply(self): + ldm.modules.attention.CrossAttention.forward = mps_flash_attention_forward + ldm.modules.diffusionmodules.model.AttnBlock.forward = mps_flash_attnblock_forward + sgm.modules.attention.CrossAttention.forward = mps_flash_attention_forward + sgm.modules.diffusionmodules.model.AttnBlock.forward = mps_flash_attnblock_forward + + class SdOptimizationSubQuad(SdOptimization): name = "sub-quadratic" cmd_opt = "opt_sub_quad_attention" @@ -148,6 +178,8 @@ def list_optimizers(res): SdOptimizationXformers(), SdOptimizationSdpNoMem(), SdOptimizationSdp(), + SdOptimizationMpsFlash(), + SdOptimizationMpsAdaptive(), SdOptimizationSubQuad(), SdOptimizationV1(), SdOptimizationInvokeAI(), @@ -430,6 +462,9 @@ def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_ _, k_tokens, _ = k.shape qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens + if q.device.type == 'mps': + q_chunk_size = mps_utils.attention_query_chunk_size(q_chunk_size, batch_x_heads, k_tokens, bytes_per_token) + if chunk_threshold is None: if q.device.type == 'mps': chunk_threshold_bytes = 268435456 * (2 if platform.processor() == 'i386' else bytes_per_token) @@ -503,9 +538,50 @@ def xformers_attention_forward(self, x, context=None, mask=None, **kwargs): return self.to_out(out) +_mps_sdp_fallback_warned = False + + +def mps_flash_attention_forward(self, x, context=None, mask=None, **kwargs): + def attention_function(query, key, value, **attention_kwargs): + return mps_flash_attention.scaled_dot_product_attention( + query, + key, + value, + training=self.training, + **attention_kwargs, + ) + + return scaled_dot_product_attention_forward( + self, + x, + context, + mask, + _attention_function=attention_function, + **kwargs, + ) + + +def mps_adaptive_attention_forward(self, x, context=None, mask=None, **kwargs): + key_tokens = context.shape[1] if context is not None else x.shape[1] + element_size = 4 if shared.opts.upcast_attn else x.element_size() + use_sdp = mps_utils.should_use_sdp(x.shape[0], self.heads, x.shape[1], key_tokens, element_size) + if not use_sdp: + return sub_quad_attention_forward(self, x, context, mask, **kwargs) + + try: + return scaled_dot_product_attention_forward(self, x, context, mask, **kwargs) + except RuntimeError as exc: + global _mps_sdp_fallback_warned + if not _mps_sdp_fallback_warned: + print(f"MPS scaled dot product attention failed; using sub-quadratic fallback: {exc}") + _mps_sdp_fallback_warned = True + torch.mps.empty_cache() + return sub_quad_attention_forward(self, x, context, mask, **kwargs) + + # Based on Diffusers usage of scaled dot product attention from https://github.com/huggingface/diffusers/blob/c7da8fd23359a22d0df2741688b5b4f33c26df21/src/diffusers/models/cross_attention.py # The scaled_dot_product_attention_forward function contains parts of code under Apache-2.0 license listed under Scaled Dot Product Attention in the Licenses section of the web UI interface -def scaled_dot_product_attention_forward(self, x, context=None, mask=None, **kwargs): +def scaled_dot_product_attention_forward(self, x, context=None, mask=None, _attention_function=None, **kwargs): batch_size, sequence_length, inner_dim = x.shape if mask is not None: @@ -532,7 +608,8 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None, **kwa q, k, v = q.float(), k.float(), v.float() # the output of sdp = (batch, num_heads, seq_len, head_dim) - hidden_states = torch.nn.functional.scaled_dot_product_attention( + attention_function = _attention_function or torch.nn.functional.scaled_dot_product_attention + hidden_states = attention_function( q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False ) @@ -634,7 +711,7 @@ def xformers_attnblock_forward(self, x): return cross_attention_attnblock_forward(self, x) -def sdp_attnblock_forward(self, x): +def sdp_attnblock_forward(self, x, _attention_function=None): h_ = x h_ = self.norm(h_) q = self.q(h_) @@ -648,7 +725,8 @@ def sdp_attnblock_forward(self, x): q = q.contiguous() k = k.contiguous() v = v.contiguous() - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) + attention_function = _attention_function or torch.nn.functional.scaled_dot_product_attention + out = attention_function(q, k, v, dropout_p=0.0, is_causal=False) out = out.to(dtype) out = rearrange(out, 'b (h w) c -> b c h w', h=h) out = self.proj_out(out) @@ -660,6 +738,37 @@ def sdp_no_mem_attnblock_forward(self, x): return sdp_attnblock_forward(self, x) +def mps_flash_attnblock_forward(self, x): + def attention_function(query, key, value, **attention_kwargs): + return mps_flash_attention.scaled_dot_product_attention( + query, + key, + value, + training=self.training, + **attention_kwargs, + ) + + return sdp_attnblock_forward(self, x, _attention_function=attention_function) + + +def mps_adaptive_attnblock_forward(self, x): + batch, _channels, height, width = x.shape + tokens = height * width + element_size = 4 if shared.opts.upcast_attn else x.element_size() + if not mps_utils.should_use_sdp(batch, 1, tokens, tokens, element_size): + return sub_quad_attnblock_forward(self, x) + + try: + return sdp_attnblock_forward(self, x) + except RuntimeError as exc: + global _mps_sdp_fallback_warned + if not _mps_sdp_fallback_warned: + print(f"MPS scaled dot product attention failed; using sub-quadratic fallback: {exc}") + _mps_sdp_fallback_warned = True + torch.mps.empty_cache() + return sub_quad_attnblock_forward(self, x) + + def sub_quad_attnblock_forward(self, x): h_ = x h_ = self.norm(h_) diff --git a/modules/sub_quadratic_attention.py b/modules/sub_quadratic_attention.py index 4cb561ef207..e811345a946 100644 --- a/modules/sub_quadratic_attention.py +++ b/modules/sub_quadratic_attention.py @@ -97,20 +97,26 @@ def chunk_scanner(chunk_idx: int) -> AttnChunk: ) return summarize_chunk(query, key_chunk, value_chunk) - chunks: list[AttnChunk] = [ - chunk_scanner(chunk) for chunk in torch.arange(0, k_tokens, kv_chunk_size) - ] - acc_chunk = AttnChunk(*map(torch.stack, zip(*chunks))) - chunk_values, chunk_weights, chunk_max = acc_chunk - - global_max, _ = torch.max(chunk_max, 0, keepdim=True) - max_diffs = torch.exp(chunk_max - global_max) - chunk_values *= torch.unsqueeze(max_diffs, -1) - chunk_weights *= max_diffs - - all_values = chunk_values.sum(dim=0) - all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0) - return all_values / all_weights + # FlashAttention-style online softmax. Merge one K/V tile into a running + # maximum, normalization sum, and output numerator, then discard the tile. + # This avoids stacking every partial result in unified memory. + accumulated = None + for chunk_idx in range(0, k_tokens, kv_chunk_size): + current = chunk_scanner(chunk_idx) + if accumulated is None: + accumulated = current + continue + + global_max = torch.maximum(accumulated.max_score, current.max_score) + accumulated_scale = torch.exp(accumulated.max_score - global_max) + current_scale = torch.exp(current.max_score - global_max) + accumulated = AttnChunk( + accumulated.exp_values * accumulated_scale.unsqueeze(-1) + current.exp_values * current_scale.unsqueeze(-1), + accumulated.exp_weights_sum * accumulated_scale + current.exp_weights_sum * current_scale, + global_max, + ) + + return accumulated.exp_values / accumulated.exp_weights_sum.unsqueeze(-1) # TODO: refactor CrossAttention#get_attention_scores to share code with this diff --git a/requirements_macos.txt b/requirements_macos.txt new file mode 100644 index 00000000000..585c01e4c07 --- /dev/null +++ b/requirements_macos.txt @@ -0,0 +1,3 @@ +# Constraints for Apple Silicon's Python 3.10 runtime. +# SciPy 1.15's arm64 wheel is rejected by macOS 27's stricter Mach-O loader. +scipy==1.13.1 diff --git a/requirements_versions.txt b/requirements_versions.txt index 389e2af33b0..13d15839dc5 100644 --- a/requirements_versions.txt +++ b/requirements_versions.txt @@ -1,4 +1,4 @@ -setuptools==69.5.1 # temp fix for compatibility with some old packages +setuptools==69.5.1 GitPython==3.1.32 Pillow==9.5.0 accelerate==0.21.0 @@ -24,10 +24,13 @@ pytorch_lightning==1.9.4 resize-right==0.0.2 safetensors==0.4.5 scikit-image==0.21.0 +scipy==1.13.1; platform_system == "Darwin" and platform_machine == "arm64" spandrel==0.3.4 spandrel-extra-arches==0.1.1 tomesd==0.1.3 -torch +torch; platform_system != "Darwin" or platform_machine != "arm64" +torch==2.3.1; platform_system == "Darwin" and platform_machine == "arm64" +torchvision==0.18.1; platform_system == "Darwin" and platform_machine == "arm64" torchdiffeq==0.2.3 torchsde==0.2.6 transformers==4.30.2 diff --git a/scripts/benchmark_mps_attention.py b/scripts/benchmark_mps_attention.py new file mode 100644 index 00000000000..e490135a621 --- /dev/null +++ b/scripts/benchmark_mps_attention.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""Compare native MPS SDPA with memory-bounded sliced attention.""" + +from __future__ import annotations + +import argparse +import statistics +import time + +import torch +import torch.nn.functional as F + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, default=4096, help="Query/key tokens (4096 = SD 512px latent)") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=40) + parser.add_argument("--chunk", type=int, default=1024, help="Query tokens per sliced-attention chunk") + parser.add_argument("--repeats", type=int, default=5) + return parser.parse_args() + + +def synchronize(): + torch.mps.synchronize() + + +def sliced_attention(query, key, value, chunk_size): + scale = query.shape[-1] ** -0.5 + output = torch.empty_like(query) + key_transposed = key.transpose(-1, -2) + for start in range(0, query.shape[-2], chunk_size): + end = min(start + chunk_size, query.shape[-2]) + scores = torch.matmul(query[:, :, start:end], key_transposed) * scale + probabilities = scores.softmax(dim=-1) + output[:, :, start:end] = torch.matmul(probabilities, value) + return output + + +def measure(operation, query, key, value, repeats): + for _ in range(2): + result = operation(query, key, value) + synchronize() + del result + + timings = [] + result = None + for _ in range(repeats): + synchronize() + started = time.perf_counter() + result = operation(query, key, value) + synchronize() + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings), result + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + shape = (args.batch, args.heads, args.tokens, args.head_dim) + torch.manual_seed(1) + query = torch.randn(shape, device="mps", dtype=torch.float16) + key = torch.randn_like(query) + value = torch.randn_like(query) + + native_ms, native_result = measure( + lambda q, k, v: F.scaled_dot_product_attention(q, k, v, dropout_p=0.0), + query, + key, + value, + args.repeats, + ) + sliced_ms, sliced_result = measure( + lambda q, k, v: sliced_attention(q, k, v, args.chunk), + query, + key, + value, + args.repeats, + ) + + difference = (native_result.float() - sliced_result.float()).abs() + fastest = "native SDPA" if native_ms <= sliced_ms else "sliced" + speedup = max(native_ms, sliced_ms) / min(native_ms, sliced_ms) + + print(f"PyTorch: {torch.__version__}") + print(f"Shape: {shape}; dtype: float16") + print(f"Native MPS SDPA: {native_ms:.2f} ms") + print(f"Sliced attention: {sliced_ms:.2f} ms") + print(f"Fastest: {fastest} ({speedup:.2f}x)") + print(f"Difference: max={difference.max().item():.6f}, mean={difference.mean().item():.6f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py new file mode 100644 index 00000000000..6a65d09bdd6 --- /dev/null +++ b/scripts/install_mps_flash_attention.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Install Metal Flash SDPA with the PyTorch 2.3 MPS stream safety backport.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import subprocess +import sys +import tarfile +import tempfile +import urllib.request + + +PACKAGE = "mps-flash-sdpa" +VERSION = "0.1.0" + + +def replace_exact(path, old, new, expected_count=1): + text = path.read_text() + count = text.count(old) + if count != expected_count: + raise RuntimeError(f"Expected {expected_count} patch sites in {path}, found {count}") + path.write_text(text.replace(old, new)) + + +def download_sdist(destination): + metadata_url = f"https://pypi.org/pypi/{PACKAGE}/{VERSION}/json" + with urllib.request.urlopen(metadata_url, timeout=30) as response: + metadata = json.load(response) + + source = next(item for item in metadata["urls"] if item["packagetype"] == "sdist") + digest = hashlib.sha256() + with urllib.request.urlopen(source["url"], timeout=120) as response, destination.open("wb") as output: + while chunk := response.read(1024 * 1024): + digest.update(chunk) + output.write(chunk) + + if digest.hexdigest() != source["digests"]["sha256"]: + raise RuntimeError("Downloaded Metal Flash SDPA source failed its SHA-256 check") + + +def extract_safely(archive, destination): + destination_resolved = destination.resolve() + with tarfile.open(archive, "r:gz") as source: + for member in source.getmembers(): + member_path = (destination / member.name).resolve() + if destination_resolved not in member_path.parents and member_path != destination_resolved: + raise RuntimeError(f"Unsafe path in source archive: {member.name}") + source.extractall(destination) + + +def patch_source(source): + bridge = source / "csrc" / "mfa_bridge.mm" + replace_exact( + bridge, + "#include \n", + "#include \n#include \n", + ) + replace_exact( + bridge, + " @autoreleasepool {\n id cmdBuf = torch::mps::get_command_buffer();", + " @autoreleasepool {\n at::mps::getCurrentMPSStream()->endKernelCoalescing();\n id cmdBuf = torch::mps::get_command_buffer();", + expected_count=2, + ) + + setup = source / "setup.py" + replace_exact( + setup, + "'cxx': ['-std=c++17', '-O2'],", + "'cxx': ['-std=c++17', '-O2', '-Wno-invalid-specialization'],", + ) + + package_init = source / "metal_flash_sdpa" / "__init__.py" + replace_exact( + package_init, + f'__version__ = "{VERSION}"\n', + f'__version__ = "{VERSION}"\nA1111_MPS_STREAM_FIX = True\n', + ) + + +def main(): + with tempfile.TemporaryDirectory(prefix="a1111-mps-flash-") as temporary: + temporary_path = Path(temporary) + archive = temporary_path / f"{PACKAGE}-{VERSION}.tar.gz" + download_sdist(archive) + extract_safely(archive, temporary_path) + source = temporary_path / f"mps_flash_sdpa-{VERSION}" + patch_source(source) + subprocess.check_call([ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-cache-dir", + "--force-reinstall", + "--no-deps", + str(source), + ]) + + +if __name__ == "__main__": + main() diff --git a/test/test_mps_flash_attention.py b/test/test_mps_flash_attention.py new file mode 100644 index 00000000000..e1fc9c569f9 --- /dev/null +++ b/test/test_mps_flash_attention.py @@ -0,0 +1,36 @@ +import torch + +from modules import mps_flash_attention +from modules.mps_flash_attention import should_use_mfa_shape + + +def test_dimension_40_routes_cross_attention_to_mfa(): + assert should_use_mfa_shape(4096, 77, 40) + + +def test_fallback_omits_unsupported_enable_gqa_keyword(): + previous_availability = mps_flash_attention._availability + mps_flash_attention._availability = False + try: + query = torch.randn(1, 2, 8, 4) + actual = mps_flash_attention.scaled_dot_product_attention(query, query, query) + expected = torch.nn.functional.scaled_dot_product_attention(query, query, query) + finally: + mps_flash_attention._availability = previous_availability + + assert torch.allclose(actual, expected) + + +def test_dimension_40_routes_self_attention_to_mfa(): + assert should_use_mfa_shape(4096, 4096, 40) + + +def test_regressing_dimensions_stay_on_pytorch_sdpa(): + assert not should_use_mfa_shape(4096, 4096, 64) + assert not should_use_mfa_shape(1024, 1024, 80) + assert not should_use_mfa_shape(1024, 77, 80) + assert not should_use_mfa_shape(256, 256, 160) + + +def test_short_attention_stays_on_pytorch_sdpa(): + assert not should_use_mfa_shape(128, 128, 40) diff --git a/test/test_mps_utils.py b/test/test_mps_utils.py new file mode 100644 index 00000000000..32245e15f37 --- /dev/null +++ b/test/test_mps_utils.py @@ -0,0 +1,45 @@ +from modules import mps_utils + + +GIB = 1024**3 + + +def test_sdp_is_used_for_normal_sd_attention_on_16gb_mac(): + assert mps_utils.should_use_sdp(2, 8, 4096, 4096, 2, total_memory=16 * GIB, available_memory=12 * GIB) + + +def test_sdp_is_avoided_for_high_resolution_self_attention(): + assert not mps_utils.should_use_sdp(1, 8, 9216, 9216, 2, total_memory=16 * GIB, available_memory=12 * GIB) + + +def test_sdp_budget_scales_down_on_8gb_mac(): + assert not mps_utils.should_use_sdp(2, 8, 4096, 4096, 2, total_memory=8 * GIB, available_memory=6 * GIB) + + +def test_cross_attention_remains_on_fast_path_at_high_resolution(): + assert mps_utils.should_use_sdp(2, 8, 9216, 77, 2, total_memory=8 * GIB, available_memory=4 * GIB) + + +def test_query_chunk_is_reduced_for_large_self_attention(): + chunk_size = mps_utils.attention_query_chunk_size( + 1024, + 16, + 9216, + 2, + total_memory=16 * GIB, + available_memory=12 * GIB, + ) + assert 1 <= chunk_size < 1024 + assert chunk_size % 64 == 0 + + +def test_query_chunk_is_unchanged_for_cross_attention(): + chunk_size = mps_utils.attention_query_chunk_size( + 1024, + 16, + 77, + 2, + total_memory=8 * GIB, + available_memory=4 * GIB, + ) + assert chunk_size == 1024 diff --git a/test/test_sub_quadratic_attention.py b/test/test_sub_quadratic_attention.py new file mode 100644 index 00000000000..503559db52b --- /dev/null +++ b/test/test_sub_quadratic_attention.py @@ -0,0 +1,46 @@ +import torch +import torch.nn.functional as F + +from modules.sub_quadratic_attention import efficient_dot_product_attention + + +def test_streaming_online_softmax_matches_sdpa(): + torch.manual_seed(123) + query = torch.randn(4, 37, 16) + key = torch.randn(4, 53, 16) + value = torch.randn(4, 53, 16) + + expected = F.scaled_dot_product_attention(query, key, value) + actual = efficient_dot_product_attention( + query, + key, + value, + query_chunk_size=13, + kv_chunk_size=11, + use_checkpoint=False, + ) + + assert torch.allclose(actual, expected, atol=2e-5, rtol=2e-5) + + +def test_streaming_online_softmax_gradients_match_sdpa(): + torch.manual_seed(321) + query = torch.randn(2, 17, 8, dtype=torch.float64, requires_grad=True) + key = torch.randn(2, 23, 8, dtype=torch.float64, requires_grad=True) + value = torch.randn(2, 23, 8, dtype=torch.float64, requires_grad=True) + gradient = torch.randn_like(query) + + expected = F.scaled_dot_product_attention(query, key, value) + actual = efficient_dot_product_attention( + query, + key, + value, + query_chunk_size=7, + kv_chunk_size=5, + use_checkpoint=False, + ) + + expected_gradients = torch.autograd.grad(expected, (query, key, value), gradient, retain_graph=True) + actual_gradients = torch.autograd.grad(actual, (query, key, value), gradient) + for actual_gradient, expected_gradient in zip(actual_gradients, expected_gradients): + assert torch.allclose(actual_gradient, expected_gradient, atol=1e-10, rtol=1e-8) diff --git a/webui-macos-env.sh b/webui-macos-env.sh index 00a36e1770e..8e4a7711b80 100644 --- a/webui-macos-env.sh +++ b/webui-macos-env.sh @@ -5,13 +5,20 @@ #################################################################### export install_dir="$HOME" -export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate" +export COMMANDLINE_ARGS="--skip-torch-cuda-test --no-half-vae --use-cpu interrogate" export PYTORCH_ENABLE_MPS_FALLBACK=1 if [[ "$(sysctl -n machdep.cpu.brand_string)" =~ ^.*"Intel".*$ ]]; then export TORCH_COMMAND="pip install torch==2.1.2 torchvision==0.16.2" else + export PIP_CONSTRAINT="${SCRIPT_DIR}/requirements_macos.txt" + # Direct Metal matrix multiplication is faster than MPSGraph for the + # projection sizes used by Stable Diffusion 1.x on M1. + export PYTORCH_MPS_PREFER_METAL=1 + # PyTorch 2.3 is the newest runtime verified to render correctly on the + # current macOS beta. The local MFA installer backports stream safety. export TORCH_COMMAND="pip install torch==2.3.1 torchvision==0.18.1" + export MPS_FLASH_ATTENTION_INSTALLER="${SCRIPT_DIR}/scripts/install_mps_flash_attention.py" fi #################################################################### From c5d897a990b54fc4462a20d32aa00a545f4d09cc Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 13:59:38 -0500 Subject: [PATCH 02/17] Optimize modern MPS inference operations --- modules/mac_specific.py | 16 +++++- scripts/benchmark_mps_unet_ops.py | 96 +++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 scripts/benchmark_mps_unet_ops.py diff --git a/modules/mac_specific.py b/modules/mac_specific.py index 039689f32e1..557765bcd7b 100644 --- a/modules/mac_specific.py +++ b/modules/mac_specific.py @@ -1,4 +1,5 @@ import logging +import os import torch from torch import Tensor @@ -30,6 +31,12 @@ def check_for_mps() -> bool: has_mps = check_for_mps() +def legacy_mps_workaround_required(fixed_in: str) -> bool: + """Keep old MPS safety copies off runtimes where the underlying bug is fixed.""" + force_legacy = os.environ.get("A1111_MPS_FORCE_LEGACY_OPS") == "1" + return force_legacy or version.parse(torch.__version__) < version.parse(fixed_in) + + def torch_mps_gc() -> None: try: if shared.state.current_latent is not None: @@ -84,10 +91,15 @@ def interpolate_with_fp32_fallback(orig_func, *args, **kwargs) -> Tensor: cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs) CondFunc('torch.cumsum', cumsum_fix_func, None) CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None) - CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None) + + # Early MPS builds could crash when a large narrow() view was consumed. + # Current runtimes handle the view correctly, so avoid cloning every slice. + if legacy_mps_workaround_required("2.1"): + CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None) # MPS workaround for https://github.com/pytorch/pytorch/issues/96113 - CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps') + if legacy_mps_workaround_required("2.0.1"): + CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps') # MPS workaround for https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14046 CondFunc('torch.nn.functional.interpolate', interpolate_with_fp32_fallback, None) diff --git a/scripts/benchmark_mps_unet_ops.py b/scripts/benchmark_mps_unet_ops.py new file mode 100644 index 00000000000..4221a66c3e5 --- /dev/null +++ b/scripts/benchmark_mps_unet_ops.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Benchmark the main SD 1.x UNet operation shapes on Apple MPS.""" + +from __future__ import annotations + +import argparse +import statistics +import time + +import torch +import torch.nn.functional as F + + +SD1_SHAPES = ( + (4096, 320), + (1024, 640), + (256, 1280), +) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2, help="UNet batch size (2 includes CFG)") + parser.add_argument("--warmup", type=int, default=4) + parser.add_argument("--repeats", type=int, default=12) + return parser.parse_args() + + +def measure(operation, warmup, repeats): + for _ in range(warmup): + operation() + torch.mps.synchronize() + + timings = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + torch.mps.synchronize() + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings) + + +def benchmark_shape(batch, tokens, channels, warmup, repeats): + side = int(tokens**0.5) + image = torch.randn((batch, channels, side, side), device="mps", dtype=torch.float16) + convolution_weight = torch.randn((channels, channels, 3, 3), device="mps", dtype=torch.float16) + convolution_bias = torch.randn((channels,), device="mps", dtype=torch.float16) + sequence = image.flatten(2).transpose(1, 2) + projection_weight = torch.randn((channels, channels), device="mps", dtype=torch.float16) + + heads = 8 + query = sequence.view(batch, tokens, heads, channels // heads).transpose(1, 2) + + results = { + "conv3x3": measure( + lambda: F.conv2d(image, convolution_weight, convolution_bias, padding=1), + warmup, + repeats, + ), + "groupnorm+silu": measure( + lambda: F.silu(F.group_norm(image, 32)), + warmup, + repeats, + ), + "linear": measure( + lambda: F.linear(sequence, projection_weight), + warmup, + repeats, + ), + "sdpa": measure( + lambda: F.scaled_dot_product_attention(query, query, query, dropout_p=0.0), + warmup, + repeats, + ), + } + + del image, convolution_weight, convolution_bias, sequence, projection_weight, query + torch.mps.empty_cache() + return results + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + torch.manual_seed(1) + print(f"PyTorch {torch.__version__}; batch={args.batch}; float16; MPS") + for tokens, channels in SD1_SHAPES: + results = benchmark_shape(args.batch, tokens, channels, args.warmup, args.repeats) + measurements = " ".join(f"{name}={milliseconds:.3f}ms" for name, milliseconds in results.items()) + print(f"tokens={tokens:4d} channels={channels:4d} {measurements}") + + +if __name__ == "__main__": + main() From 38ac556a9661f8762b6ff2e5885a1a76abe251a1 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 14:14:46 -0500 Subject: [PATCH 03/17] Coalesce Metal Flash Attention submissions --- modules/launch_utils.py | 2 +- modules/mps_flash_attention.py | 33 ++++++++++++++------------ scripts/install_mps_flash_attention.py | 7 +++++- test/test_mps_flash_attention.py | 8 +++---- 4 files changed, 29 insertions(+), 21 deletions(-) diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 9c9cd2de4d3..9dc0f3d66b5 100644 --- a/modules/launch_utils.py +++ b/modules/launch_utils.py @@ -440,7 +440,7 @@ def ensure_build_dependencies(): torch_version_match = re.match(r"(\d+)\.(\d+)", importlib.metadata.version("torch")) torch_version = tuple(map(int, torch_version_match.groups())) if torch_version_match else (0, 0) has_stream_safe_mps_flash = check_run_python( - "import metal_flash_sdpa; assert getattr(metal_flash_sdpa, 'A1111_MPS_STREAM_FIX', False)" + "import metal_flash_sdpa; assert getattr(metal_flash_sdpa, 'A1111_MPS_DEFERRED_COMMIT', False)" ) if mps_flash_installer and torch_version >= (2, 3) and not has_stream_safe_mps_flash: try: diff --git a/modules/mps_flash_attention.py b/modules/mps_flash_attention.py index 0456a0fb0e0..4a98426f929 100644 --- a/modules/mps_flash_attention.py +++ b/modules/mps_flash_attention.py @@ -34,15 +34,13 @@ def _version_tuple(version): def should_use_mfa_shape(query_tokens, key_tokens, head_dim): """Return whether the measured M1 routing table favors MFA for this shape.""" - if query_tokens < 256: + if query_tokens < 192: return False - # Dimension-40 attention wins by 2-3x and amortizes the Metal command - # buffer boundary. Smaller gains at dimensions 80/160 regress the full - # UNet because they introduce too many additional command buffers. - if head_dim == 40: - return True - return False + # These are the SD 1.x UNet head dimensions. With MFA encoded on the + # current command buffer, measured self- and cross-attention shapes all + # beat PyTorch SDPA without introducing a submission per attention call. + return head_dim in (40, 80, 160) def _run_isolated_self_test(): @@ -51,14 +49,19 @@ def _run_isolated_self_test(): import torch.nn.functional as F from metal_flash_sdpa import MetalFlashAttentionForward -q = torch.randn((1, 8, 256, 40), device='mps', dtype=torch.float16) -k = torch.randn_like(q) -v = torch.randn_like(q) -actual = MetalFlashAttentionForward.apply(q, k, v, 40 ** -0.5, False) +source = torch.randn((1, 256, 320), device='mps', dtype=torch.float16) +q = source.view(1, 256, 8, 40).transpose(1, 2) +k = q.clone() +v = q.clone() +projection = torch.randn((320, 320), device='mps', dtype=torch.float16) expected = F.scaled_dot_product_attention(q, k, v) +expected = F.linear(expected.transpose(1, 2).reshape(1, 256, 320), projection) +torch.mps.synchronize() +actual = MetalFlashAttentionForward.apply(q, k, v, 40 ** -0.5, False) +actual = F.linear(actual.transpose(1, 2).reshape(1, 256, 320), projection) torch.mps.synchronize() assert torch.isfinite(actual).all().item() -assert (actual.float() - expected.float()).abs().max().item() < 0.01 +assert (actual.float() - expected.float()).abs().max().item() < 0.05 """ environment = os.environ.copy() environment["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" @@ -95,8 +98,8 @@ def is_available(): _extension = importlib.import_module("metal_flash_sdpa") if not hasattr(_extension, "MetalFlashAttentionForward"): raise RuntimeError("extension does not expose MetalFlashAttentionForward") - if _version_tuple(torch.__version__) < (2, 11) and not getattr(_extension, "A1111_MPS_STREAM_FIX", False): - raise RuntimeError("native extension is missing the A1111 MPS stream safety patch") + if _version_tuple(torch.__version__) < (2, 11) and not getattr(_extension, "A1111_MPS_DEFERRED_COMMIT", False): + raise RuntimeError("native extension is missing the A1111 deferred MPS commit patch") _run_isolated_self_test() except (ImportError, OSError, RuntimeError, subprocess.SubprocessError) as exc: _availability_error = str(exc) @@ -105,7 +108,7 @@ def is_available(): return False _availability = True - print("Metal Flash Attention native self-test passed; selective MFA routing enabled.") + print("Metal Flash Attention native self-test passed; deferred-commit MFA routing enabled.") return True diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py index 6a65d09bdd6..5c625e1df55 100644 --- a/scripts/install_mps_flash_attention.py +++ b/scripts/install_mps_flash_attention.py @@ -64,6 +64,9 @@ def patch_source(source): " @autoreleasepool {\n at::mps::getCurrentMPSStream()->endKernelCoalescing();\n id cmdBuf = torch::mps::get_command_buffer();", expected_count=2, ) + # Keep MFA and the following PyTorch MPSGraph work on the same Metal + # command buffer. PyTorch submits it when the downstream graph is encoded. + replace_exact(bridge, "\n\n torch::mps::commit();", "", expected_count=2) setup = source / "setup.py" replace_exact( @@ -76,7 +79,9 @@ def patch_source(source): replace_exact( package_init, f'__version__ = "{VERSION}"\n', - f'__version__ = "{VERSION}"\nA1111_MPS_STREAM_FIX = True\n', + f'__version__ = "{VERSION}"\n' + 'A1111_MPS_STREAM_FIX = True\n' + 'A1111_MPS_DEFERRED_COMMIT = True\n', ) diff --git a/test/test_mps_flash_attention.py b/test/test_mps_flash_attention.py index e1fc9c569f9..7e70e225ccd 100644 --- a/test/test_mps_flash_attention.py +++ b/test/test_mps_flash_attention.py @@ -25,11 +25,11 @@ def test_dimension_40_routes_self_attention_to_mfa(): assert should_use_mfa_shape(4096, 4096, 40) -def test_regressing_dimensions_stay_on_pytorch_sdpa(): +def test_measured_sd1_dimensions_route_to_mfa(): assert not should_use_mfa_shape(4096, 4096, 64) - assert not should_use_mfa_shape(1024, 1024, 80) - assert not should_use_mfa_shape(1024, 77, 80) - assert not should_use_mfa_shape(256, 256, 160) + assert should_use_mfa_shape(1024, 1024, 80) + assert should_use_mfa_shape(1024, 77, 80) + assert should_use_mfa_shape(256, 256, 160) def test_short_attention_stays_on_pytorch_sdpa(): From 78c3fc988011add8f75dc66af259215d7fc56d2c Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 15:26:19 -0500 Subject: [PATCH 04/17] Fuse Metal GroupNorm and SiLU --- modules/launch_utils.py | 10 +- modules/mps_flash_attention.py | 19 ++- modules/mps_fused_ops.py | 82 +++++++++++ modules/sd_hijack_unet.py | 74 +++++++++- modules/shared_options.py | 5 +- scripts/install_mps_flash_attention.py | 33 ++++- scripts/mps_fused_group_norm.mm | 193 +++++++++++++++++++++++++ test/test_mps_fused_ops.py | 33 +++++ 8 files changed, 439 insertions(+), 10 deletions(-) create mode 100644 modules/mps_fused_ops.py create mode 100644 scripts/mps_fused_group_norm.mm create mode 100644 test/test_mps_fused_ops.py diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 9dc0f3d66b5..af3ee590495 100644 --- a/modules/launch_utils.py +++ b/modules/launch_utils.py @@ -439,10 +439,12 @@ def ensure_build_dependencies(): mps_flash_installer = os.environ.get("MPS_FLASH_ATTENTION_INSTALLER", "") torch_version_match = re.match(r"(\d+)\.(\d+)", importlib.metadata.version("torch")) torch_version = tuple(map(int, torch_version_match.groups())) if torch_version_match else (0, 0) - has_stream_safe_mps_flash = check_run_python( - "import metal_flash_sdpa; assert getattr(metal_flash_sdpa, 'A1111_MPS_DEFERRED_COMMIT', False)" - ) - if mps_flash_installer and torch_version >= (2, 3) and not has_stream_safe_mps_flash: + has_current_mps_flash = check_run_python( + "import metal_flash_sdpa; " + "assert getattr(metal_flash_sdpa, 'A1111_MPS_DEFERRED_COMMIT', False); " + "assert getattr(metal_flash_sdpa, 'A1111_MPS_FUSED_GROUP_NORM_SILU', False)" + ) + if mps_flash_installer and torch_version >= (2, 3) and not has_current_mps_flash: try: run( f'"{python}" "{mps_flash_installer}"', diff --git a/modules/mps_flash_attention.py b/modules/mps_flash_attention.py index 4a98426f929..edb8bd196e7 100644 --- a/modules/mps_flash_attention.py +++ b/modules/mps_flash_attention.py @@ -47,8 +47,9 @@ def _run_isolated_self_test(): code = """ import torch import torch.nn.functional as F -from metal_flash_sdpa import MetalFlashAttentionForward +from metal_flash_sdpa import MetalFlashAttentionForward, fused_group_norm_silu_forward +torch.manual_seed(1) source = torch.randn((1, 256, 320), device='mps', dtype=torch.float16) q = source.view(1, 256, 8, 40).transpose(1, 2) k = q.clone() @@ -62,6 +63,18 @@ def _run_isolated_self_test(): torch.mps.synchronize() assert torch.isfinite(actual).all().item() assert (actual.float() - expected.float()).abs().max().item() < 0.05 + +norm_source = torch.randn((1, 320, 48, 80), device='mps', dtype=torch.float16) +norm_weight = torch.randn((320,), device='mps', dtype=torch.float16) +norm_bias = torch.randn((320,), device='mps', dtype=torch.float16) +expected_norm = F.silu(F.group_norm(norm_source, 32, norm_weight, norm_bias, 1e-5)) +actual_norm = fused_group_norm_silu_forward(norm_source, norm_weight, norm_bias, 32, 1e-5) +actual_norm = actual_norm + 0 +torch.mps.synchronize() +norm_difference = (actual_norm.float() - expected_norm.float()).abs() +assert torch.isfinite(actual_norm).all().item() +assert norm_difference.max().item() < 0.02 +assert norm_difference.mean().item() < 0.001 """ environment = os.environ.copy() environment["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" @@ -100,6 +113,8 @@ def is_available(): raise RuntimeError("extension does not expose MetalFlashAttentionForward") if _version_tuple(torch.__version__) < (2, 11) and not getattr(_extension, "A1111_MPS_DEFERRED_COMMIT", False): raise RuntimeError("native extension is missing the A1111 deferred MPS commit patch") + if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU", False): + raise RuntimeError("native extension is missing fused GroupNorm+SiLU") _run_isolated_self_test() except (ImportError, OSError, RuntimeError, subprocess.SubprocessError) as exc: _availability_error = str(exc) @@ -108,7 +123,7 @@ def is_available(): return False _availability = True - print("Metal Flash Attention native self-test passed; deferred-commit MFA routing enabled.") + print("Metal self-test passed; deferred MFA and fused GroupNorm+SiLU routing enabled.") return True diff --git a/modules/mps_fused_ops.py b/modules/mps_fused_ops.py new file mode 100644 index 00000000000..68452aebde3 --- /dev/null +++ b/modules/mps_fused_ops.py @@ -0,0 +1,82 @@ +"""Measured native Metal fusions for Apple Silicon inference.""" + +from __future__ import annotations + +import os + +import torch +import torch.nn.functional as F + +from modules import mps_flash_attention + + +_dispatch_count = 0 +_fallback_count = 0 +_runtime_failure_warned = False +_first_dispatch_logged = False +_runtime_disabled = False + + +def _can_dispatch(input_tensor, norm): + if _runtime_disabled: + return False + if os.environ.get("A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU") == "1": + return False + from modules import shared + + if not getattr(shared.opts, "mps_fused_group_norm_silu", True): + return False + if input_tensor.device.type != "mps" or input_tensor.dtype != torch.float16: + return False + if input_tensor.ndim != 4 or not input_tensor.is_contiguous(): + return False + if norm.weight is None or norm.bias is None: + return False + if norm.weight.device != input_tensor.device or norm.bias.device != input_tensor.device: + return False + if norm.weight.dtype != input_tensor.dtype or norm.bias.dtype != input_tensor.dtype: + return False + if not norm.weight.is_contiguous() or not norm.bias.is_contiguous(): + return False + if input_tensor.shape[1] % norm.num_groups != 0: + return False + if torch.is_grad_enabled() and ( + input_tensor.requires_grad or norm.weight.requires_grad or norm.bias.requires_grad + ): + return False + return mps_flash_attention.is_available() + + +def group_norm_silu(input_tensor, norm): + """Fuse GroupNorm and SiLU when the measured MPS inference path supports it.""" + global _dispatch_count, _fallback_count, _runtime_failure_warned + global _first_dispatch_logged, _runtime_disabled + if _can_dispatch(input_tensor, norm): + try: + result = mps_flash_attention._extension.fused_group_norm_silu_forward( + input_tensor, + norm.weight, + norm.bias, + norm.num_groups, + norm.eps, + ) + _dispatch_count += 1 + if not _first_dispatch_logged: + print(f"Fused Metal GroupNorm+SiLU first dispatch: {tuple(input_tensor.shape)}") + _first_dispatch_logged = True + return result + except RuntimeError as exc: + _runtime_disabled = True + if not _runtime_failure_warned: + print(f"Fused Metal GroupNorm+SiLU failed; using PyTorch: {exc}") + _runtime_failure_warned = True + + _fallback_count += 1 + return F.silu(norm(input_tensor)) + + +def diagnostics(): + return { + "dispatches": _dispatch_count, + "fallbacks": _fallback_count, + } diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index b4f03b138a4..148351025dd 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -3,7 +3,7 @@ from einops import repeat import math -from modules import devices +from modules import devices, mps_fused_ops from modules.sd_hijack_utils import CondFunc @@ -36,6 +36,74 @@ def cat(self, tensors, *args, **kwargs): th = TorchHijackForUnet() +def fused_resblock_condition(_, self, x, emb): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 4 + and emb is not None + and not self.training + and not self.use_scale_shift_norm + and not getattr(self, "skip_t_emb", False) + and not getattr(self, "exchange_temb_dims", False) + and len(self.in_layers) == 3 + and len(self.out_layers) == 4 + and isinstance(self.in_layers[0], torch.nn.GroupNorm) + and isinstance(self.in_layers[1], torch.nn.SiLU) + and isinstance(self.out_layers[0], torch.nn.GroupNorm) + and isinstance(self.out_layers[1], torch.nn.SiLU) + ) + + +def fused_resblock_forward(_, self, x, emb): + if self.updown: + h = mps_fused_ops.group_norm_silu(x, self.in_layers[0]) + h = self.h_upd(h) + x = self.x_upd(x) + h = self.in_layers[2](h) + else: + h = self.in_layers[2](mps_fused_ops.group_norm_silu(x, self.in_layers[0])) + + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + h = h + emb_out + h = mps_fused_ops.group_norm_silu(h, self.out_layers[0]) + h = self.out_layers[2](h) + h = self.out_layers[3](h) + return self.skip_connection(x) + h + + +def fused_vae_resnet_condition(_, self, x, temb): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 4 + and not self.training + and isinstance(self.norm1, torch.nn.GroupNorm) + and isinstance(self.norm2, torch.nn.GroupNorm) + ) + + +def fused_vae_resnet_forward(_, self, x, temb): + h = self.conv1(mps_fused_ops.group_norm_silu(x, self.norm1)) + + if temb is not None: + h = h + self.temb_proj(torch.nn.functional.silu(temb))[:, :, None, None] + + h = self.dropout(mps_fused_ops.group_norm_silu(h, self.norm2)) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x + h + + # Below are monkey patches to enable upcasting a float16 UNet for float32 sampling def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): """Always make sure inputs to unet are in correct dtype.""" @@ -125,6 +193,10 @@ def hijack_ddpm_edit(): CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding) CondFunc('ldm.modules.attention.SpatialTransformer.forward', spatial_transformer_forward) +CondFunc('ldm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) +CondFunc('sgm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) +CondFunc('ldm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) +CondFunc('sgm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): diff --git a/modules/shared_options.py b/modules/shared_options.py index 03632ecc050..23e6f295b6b 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -232,8 +232,9 @@ options_templates.update(options_section(('optimizations', "Optimizations", "sd"), { "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), - "s_min_uncond": OptionInfo(0.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), - "s_min_uncond_all": OptionInfo(False, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), + "mps_fused_group_norm_silu": OptionInfo(True, "Fuse GroupNorm and SiLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), + "s_min_uncond": OptionInfo(1.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), + "s_min_uncond_all": OptionInfo(True, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), "token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"), "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"), "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for high-res pass", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio hr').info("only applies if non-zero and overrides above"), diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py index 5c625e1df55..a5927d6de86 100644 --- a/scripts/install_mps_flash_attention.py +++ b/scripts/install_mps_flash_attention.py @@ -6,6 +6,7 @@ import hashlib import json from pathlib import Path +import shutil import subprocess import sys import tarfile @@ -67,6 +68,20 @@ def patch_source(source): # Keep MFA and the following PyTorch MPSGraph work on the same Metal # command buffer. PyTorch submits it when the downstream graph is encoded. replace_exact(bridge, "\n\n torch::mps::commit();", "", expected_count=2) + replace_exact( + bridge, + '#include "mfa/ccv_nnc_mfa_attention.hpp"\n', + '#include "mfa/ccv_nnc_mfa_attention.hpp"\n\n' + 'void register_fused_ops(pybind11::module_& module);\n', + ) + replace_exact( + bridge, + "PYBIND11_MODULE(_C, m) {\n", + "PYBIND11_MODULE(_C, m) {\n register_fused_ops(m);\n", + ) + + fused_source = Path(__file__).with_name("mps_fused_group_norm.mm") + shutil.copyfile(fused_source, source / "csrc" / fused_source.name) setup = source / "setup.py" replace_exact( @@ -74,6 +89,12 @@ def patch_source(source): "'cxx': ['-std=c++17', '-O2'],", "'cxx': ['-std=c++17', '-O2', '-Wno-invalid-specialization'],", ) + replace_exact( + setup, + "mm_sources = [\n 'csrc/mfa_bridge.mm',\n]", + "mm_sources = [\n 'csrc/mfa_bridge.mm',\n" + " 'csrc/mps_fused_group_norm.mm',\n]", + ) package_init = source / "metal_flash_sdpa" / "__init__.py" replace_exact( @@ -81,7 +102,17 @@ def patch_source(source): f'__version__ = "{VERSION}"\n', f'__version__ = "{VERSION}"\n' 'A1111_MPS_STREAM_FIX = True\n' - 'A1111_MPS_DEFERRED_COMMIT = True\n', + 'A1111_MPS_DEFERRED_COMMIT = True\n' + 'A1111_MPS_FUSED_GROUP_NORM_SILU = True\n', + ) + replace_exact( + package_init, + "from metal_flash_sdpa._C import mfa_attention_forward, mfa_attention_backward\n", + "from metal_flash_sdpa._C import (\n" + " fused_group_norm_silu_forward,\n" + " mfa_attention_backward,\n" + " mfa_attention_forward,\n" + ")\n", ) diff --git a/scripts/mps_fused_group_norm.mm b/scripts/mps_fused_group_norm.mm new file mode 100644 index 00000000000..be76a5925a0 --- /dev/null +++ b/scripts/mps_fused_group_norm.mm @@ -0,0 +1,193 @@ +// Fused inference-only GroupNorm + SiLU for contiguous float16 MPS tensors. +#include +#include +#include +#include + +#import +#import + +namespace { + +struct FusedGroupNormParams { + uint32_t batch; + uint32_t channels; + uint32_t spatial; + uint32_t groups; + float epsilon; +}; + +static inline id getMTLBufferStorage(const at::Tensor& tensor) { + return __builtin_bit_cast(id, tensor.storage().data()); +} + +static inline size_t getMTLBufferOffset(const at::Tensor& tensor) { + return tensor.storage_offset() * tensor.element_size(); +} + +static id getFusedGroupNormSiLUPipeline() { + static id pipeline = nil; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +using namespace metal; + +struct FusedGroupNormParams { + uint batch; + uint channels; + uint spatial; + uint groups; + float epsilon; +}; + +kernel void fused_group_norm_silu_half( + device const half* input [[buffer(0)]], + device const half* weight [[buffer(1)]], + device const half* bias [[buffer(2)]], + device half* output [[buffer(3)]], + constant FusedGroupNormParams& params [[buffer(4)]], + uint tid [[thread_index_in_threadgroup]], + uint group_index [[threadgroup_position_in_grid]], + uint threads [[threads_per_threadgroup]]) { + threadgroup float partial_sum[256]; + threadgroup float partial_square_sum[256]; + + const uint channels_per_group = params.channels / params.groups; + const uint group_elements = channels_per_group * params.spatial; + const uint batch_index = group_index / params.groups; + const uint channel_group = group_index - batch_index * params.groups; + const uint base = + (batch_index * params.channels + channel_group * channels_per_group) * params.spatial; + + float sum = 0.0f; + float square_sum = 0.0f; + for (uint index = tid; index < group_elements; index += threads) { + const float value = float(input[base + index]); + sum += value; + square_sum += value * value; + } + partial_sum[tid] = sum; + partial_square_sum[tid] = square_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = threads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + partial_sum[tid] += partial_sum[tid + stride]; + partial_square_sum[tid] += partial_square_sum[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float mean = partial_sum[0] / float(group_elements); + const float variance = + max(partial_square_sum[0] / float(group_elements) - mean * mean, 0.0f); + const float inverse_stddev = rsqrt(variance + params.epsilon); + + for (uint index = tid; index < group_elements; index += threads) { + const uint local_channel = index / params.spatial; + const uint channel = channel_group * channels_per_group + local_channel; + float value = (float(input[base + index]) - mean) * inverse_stddev; + value = value * float(weight[channel]) + float(bias[channel]); + value = value / (1.0f + exp(-value)); + output[base + index] = half(value); + } +} +)METAL"; + + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK( + library != nil, + "Failed to compile fused GroupNorm+SiLU Metal library: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + id function = [library newFunctionWithName:@"fused_group_norm_silu_half"]; + TORCH_CHECK(function != nil, "Fused GroupNorm+SiLU Metal function was not found"); + pipeline = [device newComputePipelineStateWithFunction:function error:&error]; + TORCH_CHECK( + pipeline != nil, + "Failed to create fused GroupNorm+SiLU pipeline: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + }); + return pipeline; +} + +torch::Tensor fused_group_norm_silu_forward( + const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& bias, + int64_t groups, + double epsilon) { + TORCH_CHECK(input.device().is_mps(), "input must be an MPS tensor"); + TORCH_CHECK( + weight.device().is_mps() && bias.device().is_mps(), + "weight and bias must be MPS tensors"); + TORCH_CHECK(input.scalar_type() == at::kHalf, "input must be float16"); + TORCH_CHECK( + weight.scalar_type() == at::kHalf && bias.scalar_type() == at::kHalf, + "weight and bias must be float16"); + TORCH_CHECK(input.dim() == 4, "input must be NCHW"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK( + weight.is_contiguous() && bias.is_contiguous(), + "weight and bias must be contiguous"); + TORCH_CHECK( + groups > 0 && input.size(1) % groups == 0, + "channels must be divisible by groups"); + TORCH_CHECK( + weight.numel() == input.size(1) && bias.numel() == input.size(1), + "weight and bias must match the channel count"); + + auto output = torch::empty_like(input); + FusedGroupNormParams params = { + static_cast(input.size(0)), + static_cast(input.size(1)), + static_cast(input.size(2) * input.size(3)), + static_cast(groups), + static_cast(epsilon), + }; + auto pipeline = getFusedGroupNormSiLUPipeline(); + + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id encoder = [command_buffer computeCommandEncoder]; + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:getMTLBufferStorage(input) + offset:getMTLBufferOffset(input) + atIndex:0]; + [encoder setBuffer:getMTLBufferStorage(weight) + offset:getMTLBufferOffset(weight) + atIndex:1]; + [encoder setBuffer:getMTLBufferStorage(bias) + offset:getMTLBufferOffset(bias) + atIndex:2]; + [encoder setBuffer:getMTLBufferStorage(output) + offset:getMTLBufferOffset(output) + atIndex:3]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:4]; + [encoder dispatchThreadgroups:MTLSizeMake(params.batch * params.groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [encoder endEncoding]; + } + }); + } + return output; +} + +} // namespace + +void register_fused_ops(pybind11::module_& module) { + module.def( + "fused_group_norm_silu_forward", + &fused_group_norm_silu_forward, + "Fused Metal GroupNorm and SiLU forward pass", + pybind11::arg("input"), + pybind11::arg("weight"), + pybind11::arg("bias"), + pybind11::arg("groups"), + pybind11::arg("epsilon")); +} diff --git a/test/test_mps_fused_ops.py b/test/test_mps_fused_ops.py new file mode 100644 index 00000000000..f175a73709f --- /dev/null +++ b/test/test_mps_fused_ops.py @@ -0,0 +1,33 @@ +import torch +import torch.nn.functional as F + +from modules import mps_fused_ops + + +def test_cpu_fallback_matches_pytorch(): + torch.manual_seed(1) + norm = torch.nn.GroupNorm(4, 8) + source = torch.randn(1, 8, 6, 10) + + actual = mps_fused_ops.group_norm_silu(source, norm) + expected = F.silu(norm(source)) + + assert torch.equal(actual, expected) + + +def test_native_fusion_matches_pytorch(): + if not torch.backends.mps.is_available(): + return + torch.manual_seed(1) + norm = torch.nn.GroupNorm(32, 320).eval().half().to("mps") + source = torch.randn(1, 320, 48, 80, device="mps", dtype=torch.float16) + + with torch.no_grad(): + expected = F.silu(norm(source)) + actual = mps_fused_ops.group_norm_silu(source, norm) + 0 + torch.mps.synchronize() + + difference = (actual.float() - expected.float()).abs() + assert torch.isfinite(actual).all().item() + assert difference.max().item() < 0.02 + assert difference.mean().item() < 0.001 From 771259243a5a6e9a938dcedab80999512b78f5fb Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 16:59:31 -0500 Subject: [PATCH 05/17] Update README and webui-user.sh for Metal fork details and model sharing instructions Signed-off-by: Derek Anderson --- README.md | 566 ++++++++++++++++++++++++++++++++------------------ webui-user.sh | 8 + 2 files changed, 368 insertions(+), 206 deletions(-) diff --git a/README.md b/README.md index a93079fd19b..252837e6422 100644 --- a/README.md +++ b/README.md @@ -1,206 +1,360 @@ -# Stable Diffusion web UI -A web interface for Stable Diffusion, implemented using Gradio library. - -![](screenshot.png) - -## Features -[Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features): -- Original txt2img and img2img modes -- One click install and run script (but you still must install python and git) -- Outpainting -- Inpainting -- Color Sketch -- Prompt Matrix -- Stable Diffusion Upscale -- Attention, specify parts of text that the model should pay more attention to - - a man in a `((tuxedo))` - will pay more attention to tuxedo - - a man in a `(tuxedo:1.21)` - alternative syntax - - select text and press `Ctrl+Up` or `Ctrl+Down` (or `Command+Up` or `Command+Down` if you're on a MacOS) to automatically adjust attention to selected text (code contributed by anonymous user) -- Loopback, run img2img processing multiple times -- X/Y/Z plot, a way to draw a 3 dimensional plot of images with different parameters -- Textual Inversion - - have as many embeddings as you want and use any names you like for them - - use multiple embeddings with different numbers of vectors per token - - works with half precision floating point numbers - - train embeddings on 8GB (also reports of 6GB working) -- Extras tab with: - - GFPGAN, neural network that fixes faces - - CodeFormer, face restoration tool as an alternative to GFPGAN - - RealESRGAN, neural network upscaler - - ESRGAN, neural network upscaler with a lot of third party models - - SwinIR and Swin2SR ([see here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/2092)), neural network upscalers - - LDSR, Latent diffusion super resolution upscaling -- Resizing aspect ratio options -- Sampling method selection - - Adjust sampler eta values (noise multiplier) - - More advanced noise setting options -- Interrupt processing at any time -- 4GB video card support (also reports of 2GB working) -- Correct seeds for batches -- Live prompt token length validation -- Generation parameters - - parameters you used to generate images are saved with that image - - in PNG chunks for PNG, in EXIF for JPEG - - can drag the image to PNG info tab to restore generation parameters and automatically copy them into UI - - can be disabled in settings - - drag and drop an image/text-parameters to promptbox -- Read Generation Parameters Button, loads parameters in promptbox to UI -- Settings page -- Running arbitrary python code from UI (must run with `--allow-code` to enable) -- Mouseover hints for most UI elements -- Possible to change defaults/mix/max/step values for UI elements via text config -- Tiling support, a checkbox to create images that can be tiled like textures -- Progress bar and live image generation preview - - Can use a separate neural network to produce previews with almost none VRAM or compute requirement -- Negative prompt, an extra text field that allows you to list what you don't want to see in generated image -- Styles, a way to save part of prompt and easily apply them via dropdown later -- Variations, a way to generate same image but with tiny differences -- Seed resizing, a way to generate same image but at slightly different resolution -- CLIP interrogator, a button that tries to guess prompt from an image -- Prompt Editing, a way to change prompt mid-generation, say to start making a watermelon and switch to anime girl midway -- Batch Processing, process a group of files using img2img -- Img2img Alternative, reverse Euler method of cross attention control -- Highres Fix, a convenience option to produce high resolution pictures in one click without usual distortions -- Reloading checkpoints on the fly -- Checkpoint Merger, a tab that allows you to merge up to 3 checkpoints into one -- [Custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) with many extensions from community -- [Composable-Diffusion](https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/), a way to use multiple prompts at once - - separate prompts using uppercase `AND` - - also supports weights for prompts: `a cat :1.2 AND a dog AND a penguin :2.2` -- No token limit for prompts (original stable diffusion lets you use up to 75 tokens) -- DeepDanbooru integration, creates danbooru style tags for anime prompts -- [xformers](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Xformers), major speed increase for select cards: (add `--xformers` to commandline args) -- via extension: [History tab](https://github.com/yfszzx/stable-diffusion-webui-images-browser): view, direct and delete images conveniently within the UI -- Generate forever option -- Training tab - - hypernetworks and embeddings options - - Preprocessing images: cropping, mirroring, autotagging using BLIP or deepdanbooru (for anime) -- Clip skip -- Hypernetworks -- Loras (same as Hypernetworks but more pretty) -- A separate UI where you can choose, with preview, which embeddings, hypernetworks or Loras to add to your prompt -- Can select to load a different VAE from settings screen -- Estimated completion time in progress bar -- API -- Support for dedicated [inpainting model](https://github.com/runwayml/stable-diffusion#inpainting-with-stable-diffusion) by RunwayML -- via extension: [Aesthetic Gradients](https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients), a way to generate images with a specific aesthetic by using clip images embeds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients)) -- [Stable Diffusion 2.0](https://github.com/Stability-AI/stablediffusion) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20) for instructions -- [Alt-Diffusion](https://arxiv.org/abs/2211.06679) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#alt-diffusion) for instructions -- Now without any bad letters! -- Load checkpoints in safetensors format -- Eased resolution restriction: generated image's dimensions must be a multiple of 8 rather than 64 -- Now with a license! -- Reorder elements in the UI from settings screen -- [Segmind Stable Diffusion](https://huggingface.co/segmind/SSD-1B) support - -## Installation and Running -Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for: -- [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) -- [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs. -- [Intel CPUs, Intel GPUs (both integrated and discrete)](https://github.com/openvinotoolkit/stable-diffusion-webui/wiki/Installation-on-Intel-Silicon) (external wiki page) -- [Ascend NPUs](https://github.com/wangshuai09/stable-diffusion-webui/wiki/Install-and-run-on-Ascend-NPUs) (external wiki page) - -Alternatively, use online services (like Google Colab): - -- [List of Online Services](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Online-Services) - -### Installation on Windows 10/11 with NVidia-GPUs using release package -1. Download `sd.webui.zip` from [v1.0.0-pre](https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/tag/v1.0.0-pre) and extract its contents. -2. Run `update.bat`. -3. Run `run.bat`. -> For more details see [Install-and-Run-on-NVidia-GPUs](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) - -### Automatic Installation on Windows -1. Install [Python 3.10.6](https://www.python.org/downloads/release/python-3106/) (Newer version of Python does not support torch), checking "Add Python to PATH". -2. Install [git](https://git-scm.com/download/win). -3. Download the stable-diffusion-webui repository, for example by running `git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git`. -4. Run `webui-user.bat` from Windows Explorer as normal, non-administrator, user. - -### Automatic Installation on Linux -1. Install the dependencies: -```bash -# Debian-based: -sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0 -# Red Hat-based: -sudo dnf install wget git python3 gperftools-libs libglvnd-glx -# openSUSE-based: -sudo zypper install wget git python3 libtcmalloc4 libglvnd -# Arch-based: -sudo pacman -S wget git python3 -``` -If your system is very new, you need to install python3.11 or python3.10: -```bash -# Ubuntu 24.04 -sudo add-apt-repository ppa:deadsnakes/ppa -sudo apt update -sudo apt install python3.11 python3.11-venv - -# Manjaro/Arch -sudo pacman -S yay -yay -S python311 # do not confuse with python3.11 package - -# Only for 3.11 -# Then set up env variable in launch script -export python_cmd="python3.11" -# or in webui-user.sh -python_cmd="python3.11" -``` -2. Navigate to the directory you would like the webui to be installed and execute the following command: -```bash -wget -q https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh -chmod +x webui.sh -``` -Or just clone the repo wherever you want: -```bash -git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui -``` - -3. Run `webui.sh`. -4. Check `webui-user.sh` for options. -### Installation on Apple Silicon - -Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon). - -## Contributing -Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) - -## Documentation - -The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki). - -For the purposes of getting Google and other search engines to crawl the wiki, here's a link to the (not for humans) [crawlable wiki](https://github-wiki-see.page/m/AUTOMATIC1111/stable-diffusion-webui/wiki). - -## Credits -Licenses for borrowed code can be found in `Settings -> Licenses` screen, and also in `html/licenses.html` file. - -- Stable Diffusion - https://github.com/Stability-AI/stablediffusion, https://github.com/CompVis/taming-transformers, https://github.com/mcmonkey4eva/sd3-ref -- k-diffusion - https://github.com/crowsonkb/k-diffusion.git -- Spandrel - https://github.com/chaiNNer-org/spandrel implementing - - GFPGAN - https://github.com/TencentARC/GFPGAN.git - - CodeFormer - https://github.com/sczhou/CodeFormer - - ESRGAN - https://github.com/xinntao/ESRGAN - - SwinIR - https://github.com/JingyunLiang/SwinIR - - Swin2SR - https://github.com/mv-lab/swin2sr -- LDSR - https://github.com/Hafiidz/latent-diffusion -- MiDaS - https://github.com/isl-org/MiDaS -- Ideas for optimizations - https://github.com/basujindal/stable-diffusion -- Cross Attention layer optimization - Doggettx - https://github.com/Doggettx/stable-diffusion, original idea for prompt editing. -- Cross Attention layer optimization - InvokeAI, lstein - https://github.com/invoke-ai/InvokeAI (originally http://github.com/lstein/stable-diffusion) -- Sub-quadratic Cross Attention layer optimization - Alex Birch (https://github.com/Birch-san/diffusers/pull/1), Amin Rezaei (https://github.com/AminRezaei0x443/memory-efficient-attention) -- Textual Inversion - Rinon Gal - https://github.com/rinongal/textual_inversion (we're not using his code, but we are using his ideas). -- Idea for SD upscale - https://github.com/jquesnelle/txt2imghd -- Noise generation for outpainting mk2 - https://github.com/parlance-zz/g-diffuser-bot -- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator -- Idea for Composable Diffusion - https://github.com/energy-based-model/Compositional-Visual-Generation-with-Composable-Diffusion-Models-PyTorch -- xformers - https://github.com/facebookresearch/xformers -- DeepDanbooru - interrogator for anime diffusers https://github.com/KichangKim/DeepDanbooru -- Sampling in float32 precision from a float16 UNet - marunine for the idea, Birch-san for the example Diffusers implementation (https://github.com/Birch-san/diffusers-play/tree/92feee6) -- Instruct pix2pix - Tim Brooks (star), Aleksander Holynski (star), Alexei A. Efros (no star) - https://github.com/timothybrooks/instruct-pix2pix -- Security advice - RyotaK -- UniPC sampler - Wenliang Zhao - https://github.com/wl-zhao/UniPC -- TAESD - Ollin Boer Bohan - https://github.com/madebyollin/taesd -- LyCORIS - KohakuBlueleaf -- Restart sampling - lambertae - https://github.com/Newbeeer/diffusion_restart_sampling -- Hypertile - tfernd - https://github.com/tfernd/HyperTile -- Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user. -- (You) +# Stable Diffusion WebUI Metal + +An Apple Silicon performance fork of [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), focused on faster and more memory-aware inference through PyTorch MPS and native Metal kernels. + +The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, a fused GroupNorm + SiLU kernel, unified-memory-aware attention fallback, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. + +> [!IMPORTANT] +> This is an experimental performance fork, not a new Stable Diffusion engine. It favors measured M1 inference performance and safe fallback behavior over broad hardware tuning. If a native Metal path is unavailable or fails its startup test, the WebUI falls back to the corresponding PyTorch implementation. + +## How far is this from Automatic1111? + +The Metal implementation at commit [`78c3fc98`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/78c3fc988011add8f75dc66af259215d7fc56d2c) has an intentionally small, auditable delta from the official Automatic1111 `dev` branch. Documentation-only changes to this README are excluded from the implementation counts below. + +| Measure | Value | +| --- | ---: | +| Automatic1111 base | [`1937682a`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/1937682a20f7f0442311a1ede68f9f0cb480163b) | +| Base version | `v1.10.1-96-g1937682a` | +| Metal implementation | [`78c3fc98`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/78c3fc988011add8f75dc66af259215d7fc56d2c) | +| Implementation version | `v1.10.1-100-g78c3fc98` | +| Code relationship | 4 implementation commits ahead, 0 upstream commits behind official `dev` | +| Changed implementation paths | 20 of 329 tracked repository paths (6.1%) | +| New implementation files | 12 | +| Modified upstream implementation files | 8 | +| Implementation delta | 1,280 insertions, 46 deletions | + +The four fork commits are: + +1. Apple Silicon dependency, attention, memory, and benchmark foundation. +2. Removal of obsolete MPS safety copies on modern PyTorch. +3. Metal Flash Attention command-buffer coalescing. +4. Native fused GroupNorm + SiLU for compatible inference blocks. + +Most of the added lines are isolated Metal code, benchmark utilities, and tests. The fork does **not** change checkpoint formats, prompt syntax, the REST API contract, or the core Gradio workflow. + +
+Complete 20-path change surface + +| Area | Added | Modified | +| --- | --- | --- | +| Metal runtime | `modules/mps_flash_attention.py`
`modules/mps_fused_ops.py`
`modules/mps_utils.py` | `modules/mac_specific.py`
`modules/sd_hijack_optimizations.py`
`modules/sd_hijack_unet.py`
`modules/sub_quadratic_attention.py` | +| Startup and defaults | `requirements_macos.txt` | `modules/launch_utils.py`
`modules/shared_options.py`
`requirements_versions.txt`
`webui-macos-env.sh` | +| Native build and benchmarks | `scripts/install_mps_flash_attention.py`
`scripts/mps_fused_group_norm.mm`
`scripts/benchmark_mps_attention.py`
`scripts/benchmark_mps_unet_ops.py` | — | +| Tests | `test/test_mps_flash_attention.py`
`test/test_mps_fused_ops.py`
`test/test_mps_utils.py`
`test/test_sub_quadratic_attention.py` | — | + +
+ +You can reproduce the comparison locally: + +```bash +git rev-list --left-right --count 1937682a...78c3fc98 +git diff --shortstat 1937682a..78c3fc98 +git diff --name-status 1937682a..78c3fc98 +``` + +## What is different? + +### Selective Metal Flash Attention + +On supported Apple Silicon inference shapes, Automatic mode prefers a native Metal Flash Attention implementation derived from the `mps-flash-sdpa` package. + +- Routes measured SD 1.x head dimensions (`40`, `80`, and `160`) with at least 192 query tokens to the native kernel. +- Supports both self-attention and cross-attention on the measured path. +- Encodes work on PyTorch's current MPS command buffer instead of forcing a submission after every attention call. +- Uses PyTorch scaled dot product attention for unsupported shapes, training, masks, grouped-query attention, non-FP16 tensors, or runtime failure. +- Builds from `mps-flash-sdpa==0.1.0` source on first launch, checks the downloaded artifact against PyPI's published SHA-256 metadata, and installs only into the local virtual environment. +- Runs an isolated GPU self-test before enabling the native route, so a native crash cannot take down the main WebUI process during capability detection. + +The attention choices appear under **Settings → Optimizations → Cross attention optimization**: + +- `Automatic`: use Metal Flash Attention when its startup test succeeds. +- `mps-flash`: explicitly select the native Metal route with PyTorch fallback. +- `mps-adaptive`: use native PyTorch attention while it fits a unified-memory budget, then fall back to sub-quadratic attention. +- `sub-quadratic`, `sdp`, and the other upstream implementations remain available for comparison and compatibility. + +### Fused GroupNorm + SiLU + +A native inference-only Metal kernel combines GroupNorm and SiLU in one dispatch for compatible contiguous FP16 tensors. + +- Used in compatible SD/SGM UNet residual blocks. +- Used in compatible VAE residual blocks when the VAE is running in FP16. +- Preserves the normal PyTorch path for CPU, non-FP16 tensors, training/autograd, incompatible layouts, missing affine parameters, or runtime errors. +- Enabled by default through **Settings → Optimizations → Fuse GroupNorm and SiLU on Apple Silicon**. + +This is a focused fusion; convolutions and residual additions still use PyTorch MPS. A larger block-level MPSGraph prototype was tested and deliberately rejected because it was about 1% slower end to end and produced a larger numerical delta without a speed benefit. + +### Unified-memory-aware attention + +The fork treats system RAM and GPU memory as the same constrained resource instead of relying on a fixed attention threshold. + +- Estimates scaled dot product attention's temporary memory from batch, heads, token counts, and element size. +- Limits native attention to a fraction of total and currently available unified memory. +- Dynamically reduces sub-quadratic query tiles for large self-attention workloads. +- Uses streaming online softmax when K/V attention is chunked, merging one tile at a time instead of stacking all partial outputs in memory. +- Keeps cross-attention on the fast path when its short key sequence remains inexpensive. + +The adaptive path is especially useful for high resolutions and lower-memory Macs. The default Metal Flash Attention path remains the measured choice for normal SD 1.x shapes. + +### Modern MPS runtime cleanup + +Several workarounds needed by early PyTorch MPS releases are now gated by runtime version: + +- Avoids cloning every `narrow()` result on PyTorch versions where the underlying MPS bug is fixed. +- Avoids unconditional FP32 LayerNorm conversion on modern runtimes. +- Keeps an environment switch for diagnosing regressions with legacy behavior. +- Prefers direct Metal matrix multiplication for the SD 1.x projection sizes measured on M1. +- Removes Automatic1111's default `--upcast-sampling` flag on Apple Silicon; it can be restored locally when exact upstream behavior is more important than speed. + +### Apple Silicon dependency profile + +The default Apple Silicon environment is pinned to the combination verified for this fork: + +| Dependency | Version | +| --- | --- | +| Python | 3.10 recommended; 3.10.20 used during development | +| PyTorch | 2.3.1 | +| torchvision | 0.18.1 | +| SciPy | 1.13.1 | +| Native extension | `mps-flash-sdpa` 0.1.0 with local stream-safety and fusion patches | + +SciPy is constrained on Apple Silicon because newer wheels encountered loader problems on the macOS beta used during development. Requirements parsing was also updated to understand platform markers and normal Python package specifiers correctly. + +### Changed defaults + +The following defaults intentionally differ from the upstream `dev` branch: + +| Setting | Upstream | This fork | Effect | +| --- | --- | --- | --- | +| Negative Guidance minimum sigma (NGMS) | `0.0` | `1.0` | May skip unconditional guidance late in sampling | +| NGMS all steps | Off | On | Applies the configured NGMS rule on every eligible step | +| `--upcast-sampling` on macOS | On | Off | Keeps more sampling work in FP16 for speed | +| Cross-attention Automatic choice on MPS | Sub-quadratic | Metal Flash Attention | Uses the measured native route when available | +| Fused GroupNorm + SiLU | Not present | On | Reduces compatible normalization/activation dispatches | + +NGMS is the largest user-visible behavioral change. It is recorded in PNG generation metadata when active. Set NGMS to `0` and disable **NGMS all steps** if a workflow expects upstream guidance behavior. + +## Measured performance + +One recorded Apple M1 Mac mini comparison during development used the same checkpoint hash and compute shape: + +| Build | Workload | Time | +| --- | --- | ---: | +| Automatic1111 `v1.10.1-96-g1937682a` | 5 steps, DPM++ SDE, Karras, CFG 1.15, 384×640, SD 1.x checkpoint `8ecad70a19`, Clip skip 2, NGMS 1/all steps | 12.8 s | +| This fork `v1.10.1-99-g38ac556a` | Same sampler, schedule, dimensions, checkpoint hash, Clip skip, and NGMS settings | 8.7 s | + +That observed run was approximately **32% lower latency**, or **1.47× as fast**. The current head adds the fused GroupNorm + SiLU path after that recorded comparison. + +The two recorded generations used different seeds. This makes the table a throughput comparison at matching tensor shapes, not an image-parity A/B. + +Treat these numbers as a development result, not a universal guarantee. Timing varies with: + +- Apple Silicon generation and GPU core count +- Unified-memory capacity and pressure from other applications +- Model architecture and attention dimensions +- Resolution, batch size, sampler, and step count +- First-run shader compilation and warm-up +- VAE, LoRA, ControlNet, extensions, and live preview configuration + +For a meaningful comparison, use the same model hash, prompt, negative prompt, seed, sampler, schedule, steps, CFG, dimensions, Clip skip, VAE, and optimization settings. Run at least two warm-ups, then compare the median of several generations. + +## Installation + +### Requirements + +- An Apple Silicon Mac for the optimized native path +- macOS with Metal Performance Shaders support +- Python 3.10 +- Git +- Xcode Command Line Tools, required to compile the Objective-C++/Metal extension + +Install the command-line tools if needed: + +```bash +xcode-select --install +``` + +### Fresh installation + +```bash +git clone --branch dev https://github.com/dmikey/stable-diffusion-webui-metal.git +cd stable-diffusion-webui-metal +./webui.sh +``` + +The first launch creates the virtual environment, installs the pinned Apple Silicon dependencies, downloads and builds the native extension, runs its isolated self-test, and starts the normal Automatic1111 interface. Native extension compilation can make the first launch noticeably longer than later launches. + +Put checkpoints in: + +```text +models/Stable-diffusion/ +``` + +LoRAs, VAEs, embeddings, extensions, and outputs use the usual Automatic1111 directories. + +### Updating + +The repository's default branch is `dev`: + +```bash +git switch dev +git pull --ff-only origin dev +./webui.sh +``` + +Do not commit generated `config.json`, `ui-config.json`, `params.txt`, models, outputs, the virtual environment, or extension installations. They are local runtime state and are ignored by Git. + +### Local launch options + +`webui-macos-env.sh` contains the fork's tracked defaults. Put personal overrides in `webui-user.sh`, which is loaded afterward. For example, to restore sampling upcast while retaining the rest of the Metal work: + +```bash +export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate" +``` + +The default Apple Silicon launch options retain `--no-half-vae` to avoid FP16 VAE instability and run the interrogator on CPU. + +## Startup messages and fallback behavior + +A healthy optimized startup prints messages similar to: + +```text +Metal self-test passed; deferred MFA and fused GroupNorm+SiLU routing enabled. +Applying attention optimization: mps-flash... done. +``` + +The first compatible generation also reports the first native attention and GroupNorm dispatch. These messages are informational and print only once per process. + +If the extension cannot build or fails its isolated self-test, startup continues with native PyTorch MPS attention. If the fused GroupNorm kernel fails at runtime, that fusion is disabled for the process and PyTorch handles subsequent operations. + +## Compatibility and output parity + +### What should remain compatible + +- Automatic1111's txt2img, img2img, inpainting, high-resolution pass, API, and metadata workflow +- Existing `.safetensors` and `.ckpt` checkpoints +- Standard LoRA, embedding, VAE, and extension directory layouts +- Existing sampler names and generation parameter syntax +- CPU and non-MPS fallback implementations + +The optimized target is SD 1.x inference. Other architectures supported by this Automatic1111 base may run, but unsupported attention shapes fall back to PyTorch and may receive little or no speed benefit. Test model-specific extensions individually. + +### Why the same seed may differ from upstream + +Pixel-identical output is not guaranteed. Differences can come from: + +- NGMS being enabled by default +- Sampling no longer being upcast by default +- Native Flash Attention and fused GroupNorm changing FP16 reduction/rounding order +- A different selected attention implementation + +Small FP16 numerical differences can grow over multiple denoising evaluations even when both paths are deterministic. + +### Closest upstream behavior + +For an upstream-style comparison: + +1. Set **Negative Guidance minimum sigma** to `0`. +2. Disable **Negative Guidance minimum sigma all steps**. +3. Disable **Fuse GroupNorm and SiLU on Apple Silicon**. +4. Select `sub-quadratic` under **Cross attention optimization**. +5. Add `--upcast-sampling` to `COMMANDLINE_ARGS` in `webui-user.sh`. +6. Restart the WebUI after changing launch arguments. + +For diagnostics only, `A1111_MPS_FORCE_LEGACY_OPS=1` restores version-gated MPS safety copies, and `A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU=1` disables the native normalization fusion before startup. + +## Troubleshooting + +### Native extension does not build + +Confirm that Xcode Command Line Tools and the local virtual environment are available: + +```bash +xcode-select -p +./venv/bin/python scripts/install_mps_flash_attention.py +``` + +Then restart with `./webui.sh`. The installer intentionally rebuilds the package from source for the active Python and PyTorch environment. + +### Metal self-test fails + +The WebUI should continue on PyTorch MPS. Keep the final `Metal Flash Attention unavailable:` message when reporting the issue. Also include: + +- Mac model and memory capacity +- macOS version +- `./venv/bin/python -c "import torch; print(torch.__version__)"` +- The selected cross-attention optimization +- Model family, resolution, and batch size + +### Green, black, or corrupted output + +Keep `--no-half-vae` enabled first. Also compare with NGMS disabled, sampling upcast restored, `sub-quadratic` attention selected, and the fused GroupNorm option disabled. That separates model/VAE precision issues from the native Metal paths. + +### High-resolution out-of-memory errors + +Select `mps-adaptive - native Metal attention with a memory-safe fallback` or `sub-quadratic` in the optimization settings. Reduce batch size before reducing attention chunk limits manually. + +## Benchmarks and tests + +Two standalone benchmark scripts are included: + +```bash +./venv/bin/python scripts/benchmark_mps_attention.py +./venv/bin/python scripts/benchmark_mps_unet_ops.py --batch 2 +``` + +The first compares PyTorch MPS scaled dot product attention with sliced attention. The second measures representative SD 1.x convolution, GroupNorm + SiLU, linear projection, and attention shapes. + +Focused tests cover: + +- Metal Flash Attention routing and PyTorch fallback +- Native fused GroupNorm + SiLU correctness +- Unified-memory attention budgeting and dynamic query tiles +- Streaming online-softmax forward results and gradients + +With `pytest` installed in the virtual environment: + +```bash +./venv/bin/python -m pytest -q \ + test/test_mps_flash_attention.py \ + test/test_mps_fused_ops.py \ + test/test_mps_utils.py \ + test/test_sub_quadratic_attention.py +``` + +## Deliberately not included + +- Block-level MPSGraph execution: implemented and benchmarked, but rejected after a small regression. +- FP8 acceleration on M1: there is no matching M1 hardware fast path, so conversion would primarily add unpacking overhead. +- Core ML/ANE conversion: this would introduce a separate static execution engine and materially reduce Automatic1111 compatibility. +- Whole-UNet static graphs: potentially higher upside, but a much larger project with difficult LoRA, ControlNet, model-switching, and dynamic-resolution tradeoffs. +- Model-format changes or required quantization: existing Automatic1111 checkpoints are used directly. + +## Upstream features and documentation + +This README focuses on the fork. For the complete WebUI feature set, usage documentation, and extension ecosystem, see: + +- [Automatic1111 feature overview](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features) +- [Automatic1111 wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki) +- [Automatic1111 API documentation](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API) +- [Automatic1111 troubleshooting](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Troubleshooting) + +## Contributing + +Keep changes narrow, measurable, and safe to fall back from. + +For performance work: + +1. Record the exact model hash and generation settings. +2. Warm up both paths. +3. Compare multiple alternating runs rather than a single best time. +4. Verify deterministic behavior within each path. +5. Measure output deviation as well as latency and memory. +6. Retain the upstream PyTorch path for unsupported inputs and runtime failure. + +Changes that improve an isolated operator but do not improve an end-to-end generation should not be enabled by default. + +## License and credits + +This fork retains Automatic1111's license and third-party notices. Licenses for bundled and borrowed components are available under **Settings → Licenses** and in `html/licenses.html`. + +Primary credit remains with the [Automatic1111 project](https://github.com/AUTOMATIC1111/stable-diffusion-webui) and its contributors. The native attention work builds on [`mps-flash-sdpa`](https://pypi.org/project/mps-flash-sdpa/) and ideas explored by Draw Things, adapted here for Automatic1111's PyTorch MPS execution path. diff --git a/webui-user.sh b/webui-user.sh index 70306c60d5b..95db1c3d342 100644 --- a/webui-user.sh +++ b/webui-user.sh @@ -12,6 +12,13 @@ # Commandline arguments for webui.py, for example: export COMMANDLINE_ARGS="--medvram --opt-split-attention" #export COMMANDLINE_ARGS="" +# Share Models with your current install +# COMMANDLINE_ARGS=--ckpt-dir /path/to/A1111/models/Stable-diffusion \ +# --lora-dir /path/to/A1111/models/Lora \ +# --vae-dir /path/to/A1111/models/VAE \ +# --controlnet-dir /path/to/A1111/extensions/sd-webui-controlnet/models \ +# --embeddings-dir /path/to/A1111/embeddings + # python3 executable #python_cmd="python3" @@ -45,4 +52,5 @@ # Uncomment to disable TCMalloc #export NO_TCMALLOC="True" + ########################################### From 58e63e9f19bb0c6203f6a4ecececb3b1372b6917 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 18:52:04 -0500 Subject: [PATCH 06/17] Profile and accelerate M1 VAE decode --- README.md | 35 ++++- modules/mps_stage_profile.py | 227 ++++++++++++++++++++++++++++ modules/processing.py | 28 ++-- modules/sd_samplers_cfg_denoiser.py | 10 +- state-of-things-next.md | 91 +++++++++++ test/test_macos_launch_defaults.py | 35 +++++ test/test_mps_stage_profile.py | 90 +++++++++++ webui-macos-env.sh | 14 +- 8 files changed, 508 insertions(+), 22 deletions(-) create mode 100644 modules/mps_stage_profile.py create mode 100644 state-of-things-next.md create mode 100644 test/test_macos_launch_defaults.py create mode 100644 test/test_mps_stage_profile.py diff --git a/README.md b/README.md index 252837e6422..1d743c84baf 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ An Apple Silicon performance fork of [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), focused on faster and more memory-aware inference through PyTorch MPS and native Metal kernels. -The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, a fused GroupNorm + SiLU kernel, unified-memory-aware attention fallback, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. +The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, a fused GroupNorm + SiLU kernel, unified-memory-aware attention fallback, an M1-validated FP16 VAE path, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. > [!IMPORTANT] > This is an experimental performance fork, not a new Stable Diffusion engine. It favors measured M1 inference performance and safe fallback behavior over broad hardware tuning. If a native Metal path is unavailable or fails its startup test, the WebUI falls back to the corresponding PyTorch implementation. @@ -128,6 +128,7 @@ The following defaults intentionally differ from the upstream `dev` branch: | Negative Guidance minimum sigma (NGMS) | `0.0` | `1.0` | May skip unconditional guidance late in sampling | | NGMS all steps | Off | On | Applies the configured NGMS rule on every eligible step | | `--upcast-sampling` on macOS | On | Off | Keeps more sampling work in FP16 for speed | +| `--no-half-vae` on M1-family Macs | On | Off | Runs VAE encode/decode in FP16; Automatic1111 still retries in FP32 if VAE decode produces NaNs | | Cross-attention Automatic choice on MPS | Sub-quadratic | Metal Flash Attention | Uses the measured native route when available | | Fused GroupNorm + SiLU | Not present | On | Reduces compatible normalization/activation dispatches | @@ -146,6 +147,19 @@ That observed run was approximately **32% lower latency**, or **1.47× as fast** The two recorded generations used different seeds. This makes the table a throughput comparison at matching tensor shapes, not an image-parity A/B. +### M1 FP16 VAE validation + +A later controlled A/B isolated VAE precision on a 16 GB Apple M1 Mac mini. Both paths used checkpoint `8ecad70a19`, prompt `a dog`, seed `3163229250`, 5-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 384×640 output. Each result below is the median of five warm runs with coarse MPS stage profiling enabled. + +| VAE path | End-to-end client time | Sampler stage | VAE decode + transfer | +| --- | ---: | ---: | ---: | +| FP32 (`--no-half-vae`) | 8.450 s | 6.715 s | 1.536 s | +| FP16 | 7.795 s | 6.666 s | 0.972 s | + +FP16 reduced the measured VAE stage by about **37%** and end-to-end latency by about **7.8%**. The sampler time remained effectively unchanged, which is the expected result when only decode precision changes. + +Output quality was checked across three fixed-seed generations at 384×640 and 512×512. Compared with FP32 VAE output, every changed 8-bit RGB channel differed by at most 1 value, PSNR was 64.0–64.6 dB, and 97.4–97.7% of channels were byte-identical. All FP16 runs were deterministic and free of NaN, green, or corrupted output. The default is therefore enabled only on the tested M1 family; other Apple Silicon generations retain FP32 VAE until separately validated. + Treat these numbers as a development result, not a universal guarantee. Timing varies with: - Apple Silicon generation and GPU core count @@ -211,7 +225,7 @@ Do not commit generated `config.json`, `ui-config.json`, `params.txt`, models, o export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate" ``` -The default Apple Silicon launch options retain `--no-half-vae` to avoid FP16 VAE instability and run the interrogator on CPU. +M1-family Macs use the validated FP16 VAE path by default. Intel and other Apple Silicon generations retain `--no-half-vae`. Add `--no-half-vae` to a local `COMMANDLINE_ARGS` override at any time to force the conservative FP32 VAE path. Automatic1111's enabled-by-default VAE precision recovery also converts the VAE to FP32 and retries if an FP16 decode produces NaNs. ## Startup messages and fallback behavior @@ -258,7 +272,8 @@ For an upstream-style comparison: 3. Disable **Fuse GroupNorm and SiLU on Apple Silicon**. 4. Select `sub-quadratic` under **Cross attention optimization**. 5. Add `--upcast-sampling` to `COMMANDLINE_ARGS` in `webui-user.sh`. -6. Restart the WebUI after changing launch arguments. +6. On M1, also add `--no-half-vae`. +7. Restart the WebUI after changing launch arguments. For diagnostics only, `A1111_MPS_FORCE_LEGACY_OPS=1` restores version-gated MPS safety copies, and `A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU=1` disables the native normalization fusion before startup. @@ -287,7 +302,7 @@ The WebUI should continue on PyTorch MPS. Keep the final `Metal Flash Attention ### Green, black, or corrupted output -Keep `--no-half-vae` enabled first. Also compare with NGMS disabled, sampling upcast restored, `sub-quadratic` attention selected, and the fused GroupNorm option disabled. That separates model/VAE precision issues from the native Metal paths. +Add `--no-half-vae` to the local launch options and restart first. Also compare with NGMS disabled, sampling upcast restored, `sub-quadratic` attention selected, and the fused GroupNorm option disabled. That separates model/VAE precision issues from the native Metal paths. ### High-resolution out-of-memory errors @@ -304,10 +319,20 @@ Two standalone benchmark scripts are included: The first compares PyTorch MPS scaled dot product attention with sliced attention. The second measures representative SD 1.x convolution, GroupNorm + SiLU, linear projection, and attention shapes. +For an end-to-end stage breakdown, launch with the opt-in profiler: + +```bash +A1111_MPS_PROFILE=1 ./webui.sh +``` + +Each generation reports synchronized wall time for conditioning, sampling, VAE decode/transfer, and image processing; it also records UNet call shapes and MPS allocation snapshots in a machine-readable `MPS_PROFILE_JSON` line. Profiling is intentionally coarse because PyTorch 2.3 MPS timing events are unreliable on the tested runtime. When the environment variable is absent, the profiler adds no MPS synchronization points. + Focused tests cover: - Metal Flash Attention routing and PyTorch fallback +- M1-specific FP16 VAE launch defaults with conservative Intel and newer-chip behavior - Native fused GroupNorm + SiLU correctness +- Opt-in MPS stage profiling and its zero-synchronization disabled path - Unified-memory attention budgeting and dynamic query tiles - Streaming online-softmax forward results and gradients @@ -315,8 +340,10 @@ With `pytest` installed in the virtual environment: ```bash ./venv/bin/python -m pytest -q \ + test/test_macos_launch_defaults.py \ test/test_mps_flash_attention.py \ test/test_mps_fused_ops.py \ + test/test_mps_stage_profile.py \ test/test_mps_utils.py \ test/test_sub_quadratic_attention.py ``` diff --git a/modules/mps_stage_profile.py b/modules/mps_stage_profile.py new file mode 100644 index 00000000000..4a94ec69870 --- /dev/null +++ b/modules/mps_stage_profile.py @@ -0,0 +1,227 @@ +"""Opt-in coarse stage profiling for Apple Silicon inference. + +Set A1111_MPS_PROFILE=1 before launch to enable. The normal path never calls +torch.mps.synchronize(); profiling synchronizes only at request/stage boundaries. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from contextlib import contextmanager +from contextvars import ContextVar +import json +import os +import platform +import time + +import psutil +import torch + + +ENVIRONMENT_VARIABLE = "A1111_MPS_PROFILE" +OUTPUT_PREFIX = "MPS_PROFILE_JSON " + +_active_session: ContextVar[ProfileSession | None] = ContextVar("mps_profile_session", default=None) + + +def requested(): + return os.environ.get(ENVIRONMENT_VARIABLE) == "1" + + +def active(): + return _active_session.get() is not None + + +def _mps_available(): + return ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and torch.backends.mps.is_available() + ) + + +def _synchronize(): + torch.mps.synchronize() + + +def _memory_snapshot(): + process_rss = psutil.Process().memory_info().rss + system_memory = psutil.virtual_memory() + snapshot = { + "process_rss": process_rss, + "system_available": system_memory.available, + } + + if _mps_available(): + for name in ("current_allocated_memory", "driver_allocated_memory"): + function = getattr(torch.mps, name, None) + if function is not None: + try: + snapshot[f"mps_{name}"] = function() + except Exception: + pass + + return snapshot + + +class ProfileSession: + def __init__(self, metadata): + self.metadata = metadata + self.has_mps = _mps_available() + self.started = None + self.stages = defaultdict(lambda: {"calls": 0, "wall_ms": 0.0}) + self.unet_calls = 0 + self.unet_batches = 0 + self.unet_shapes = Counter() + self.memory_start = None + self.memory_end = None + self.memory_peaks = {} + self.synchronization_error = None + + def synchronize(self): + if not self.has_mps or self.synchronization_error is not None: + return + + try: + _synchronize() + except Exception as exc: + self.synchronization_error = str(exc) + + def observe_memory(self, snapshot=None): + snapshot = snapshot or _memory_snapshot() + for name, value in snapshot.items(): + self.memory_peaks[name] = max(self.memory_peaks.get(name, 0), value) + return snapshot + + def start(self): + self.synchronize() + self.memory_start = self.observe_memory() + self.started = time.perf_counter() + + def record_stage(self, name, elapsed_ms): + record = self.stages[name] + record["calls"] += 1 + record["wall_ms"] += elapsed_ms + self.observe_memory() + + def record_unet(self, input_tensor): + self.unet_calls += 1 + shape = tuple(input_tensor.shape) + if shape: + self.unet_batches += shape[0] + self.unet_shapes["x".join(map(str, shape))] += 1 + + def finish(self, error=None): + self.synchronize() + total_ms = (time.perf_counter() - self.started) * 1000 + self.memory_end = self.observe_memory() + stages = { + name: { + "calls": record["calls"], + "wall_ms": round(record["wall_ms"], 3), + } + for name, record in sorted(self.stages.items()) + } + accounted_ms = sum(record["wall_ms"] for record in self.stages.values()) + report = { + "metadata": self.metadata, + "total_wall_ms": round(total_ms, 3), + "accounted_stage_ms": round(accounted_ms, 3), + "unaccounted_wall_ms": round(max(total_ms - accounted_ms, 0.0), 3), + "stages": stages, + "unet": { + "calls": self.unet_calls, + "total_batch_elements": self.unet_batches, + "input_shapes": dict(sorted(self.unet_shapes.items())), + }, + "memory_bytes": { + "start": self.memory_start, + "end": self.memory_end, + "sampled_peaks": self.memory_peaks, + }, + } + if error is not None: + report["error"] = type(error).__name__ + if self.synchronization_error is not None: + report["synchronization_error"] = self.synchronization_error + + stage_summary = ", ".join( + f"{name}={record['wall_ms']:.1f}ms/{record['calls']}" + for name, record in sorted(self.stages.items()) + ) + print( + "MPS stage profile: " + f"total={total_ms:.1f}ms, {stage_summary or 'no stages'}, " + f"unet_calls={self.unet_calls}, unaccounted={max(total_ms - accounted_ms, 0.0):.1f}ms" + ) + print(OUTPUT_PREFIX + json.dumps(report, sort_keys=True, separators=(",", ":"), default=str)) + + +def metadata_for_processing(processing): + model = getattr(processing, "sd_model", None) + return { + "batch_size": getattr(processing, "batch_size", None), + "cfg_scale": getattr(processing, "cfg_scale", None), + "height": getattr(processing, "height", None), + "model_hash": getattr(processing, "sd_model_hash", None) or getattr(model, "sd_model_hash", None), + "n_iter": getattr(processing, "n_iter", None), + "sampler": getattr(processing, "sampler_name", None), + "scheduler": getattr(processing, "scheduler", None), + "steps": getattr(processing, "steps", None), + "width": getattr(processing, "width", None), + } + + +@contextmanager +def request(processing): + if not requested() or active(): + yield + return + + session = ProfileSession(metadata_for_processing(processing)) + token = _active_session.set(session) + try: + session.start() + except Exception as exc: + _active_session.reset(token) + print(f"MPS stage profiling unavailable; continuing without it: {exc}") + yield + return + + error = None + try: + yield + except BaseException as exc: + error = exc + raise + finally: + try: + try: + session.finish(error) + except Exception as exc: + print(f"MPS stage profiling report failed; generation result is unchanged: {exc}") + finally: + _active_session.reset(token) + + +@contextmanager +def stage(name): + session = _active_session.get() + if session is None: + yield + return + + session.synchronize() + started = time.perf_counter() + try: + yield + finally: + session.synchronize() + session.record_stage(name, (time.perf_counter() - started) * 1000) + + +def unet_call(function, input_tensor, *args, **kwargs): + session = _active_session.get() + if session is not None: + session.record_unet(input_tensor) + return function(input_tensor, *args, **kwargs) diff --git a/modules/processing.py b/modules/processing.py index 92c3582cc66..462eaedb8cc 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -16,7 +16,7 @@ from typing import Any import modules.sd_hijack -from modules import devices, prompt_parser, masking, sd_samplers, lowvram, infotext_utils, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors, rng, profiling +from modules import devices, prompt_parser, masking, sd_samplers, lowvram, infotext_utils, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors, rng, profiling, mps_stage_profile from modules.rng import slerp # noqa: F401 from modules.sd_hijack import model_hijack from modules.sd_samplers_common import images_tensor_to_samples, decode_first_stage, approximation_indexes @@ -843,8 +843,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed: # backwards compatibility, fix sampler and scheduler if invalid sd_samplers.fix_p_invalid_sampler_and_scheduler(p) - with profiling.Profiler(): - res = process_images_inner(p) + with mps_stage_profile.request(p): + with profiling.Profiler(): + res = process_images_inner(p) finally: sd_models.apply_token_merging(p.sd_model, 0) @@ -868,7 +869,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: else: assert p.prompt is not None - devices.torch_gc() + with mps_stage_profile.stage("initial_gc"): + devices.torch_gc() seed = get_fixed_seed(p.seed) subseed = get_fixed_seed(p.subseed) @@ -917,7 +919,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: infotexts = [] output_images = [] with torch.no_grad(), p.sd_model.ema_scope(): - with devices.autocast(): + with devices.autocast(), mps_stage_profile.stage("initialization"): p.init(p.all_prompts, p.all_seeds, p.all_subseeds) # for OSX, loading the model during sampling changes the generated picture, so it is loaded here @@ -963,7 +965,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if p.scripts is not None: p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds) - p.setup_conds() + with mps_stage_profile.stage("conditioning"): + p.setup_conds() p.extra_generation_params.update(model_hijack.extra_generation_params) @@ -984,7 +987,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: sd_models.apply_alpha_schedule_override(p.sd_model, p) - with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(): + with mps_stage_profile.stage("sampler"), devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(): samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts) if p.scripts is not None: @@ -999,17 +1002,20 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if opts.sd_vae_decode_method != 'Full': p.extra_generation_params['VAE Decoder'] = opts.sd_vae_decode_method - x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True) + with mps_stage_profile.stage("vae_decode_and_transfer"): + x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True) - x_samples_ddim = torch.stack(x_samples_ddim).float() - x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) + with mps_stage_profile.stage("image_tensor_processing"): + x_samples_ddim = torch.stack(x_samples_ddim).float() + x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) del samples_ddim if lowvram.is_enabled(shared.sd_model): lowvram.send_everything_to_cpu() - devices.torch_gc() + with mps_stage_profile.stage("post_decode_gc"): + devices.torch_gc() state.nextjob() diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py index b6fbf337243..687b7f3cf56 100644 --- a/modules/sd_samplers_cfg_denoiser.py +++ b/modules/sd_samplers_cfg_denoiser.py @@ -1,5 +1,5 @@ import torch -from modules import prompt_parser, sd_samplers_common +from modules import prompt_parser, sd_samplers_common, mps_stage_profile from modules.shared import opts, state import modules.shared as shared @@ -246,13 +246,13 @@ def apply_blend(current_latent): cond_in = catenate_conds([tensor, uncond]) if shared.opts.batch_cond_uncond: - x_out = self.inner_model(x_in, sigma_in, cond=make_condition_dict(cond_in, image_cond_in)) + x_out = mps_stage_profile.unet_call(self.inner_model, x_in, sigma_in, cond=make_condition_dict(cond_in, image_cond_in)) else: x_out = torch.zeros_like(x_in) for batch_offset in range(0, x_out.shape[0], batch_size): a = batch_offset b = a + batch_size - x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond=make_condition_dict(subscript_cond(cond_in, a, b), image_cond_in[a:b])) + x_out[a:b] = mps_stage_profile.unet_call(self.inner_model, x_in[a:b], sigma_in[a:b], cond=make_condition_dict(subscript_cond(cond_in, a, b), image_cond_in[a:b])) else: x_out = torch.zeros_like(x_in) batch_size = batch_size*2 if shared.opts.batch_cond_uncond else batch_size @@ -265,10 +265,10 @@ def apply_blend(current_latent): else: c_crossattn = torch.cat([tensor[a:b]], uncond) - x_out[a:b] = self.inner_model(x_in[a:b], sigma_in[a:b], cond=make_condition_dict(c_crossattn, image_cond_in[a:b])) + x_out[a:b] = mps_stage_profile.unet_call(self.inner_model, x_in[a:b], sigma_in[a:b], cond=make_condition_dict(c_crossattn, image_cond_in[a:b])) if not skip_uncond: - x_out[-uncond.shape[0]:] = self.inner_model(x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond=make_condition_dict(uncond, image_cond_in[-uncond.shape[0]:])) + x_out[-uncond.shape[0]:] = mps_stage_profile.unet_call(self.inner_model, x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond=make_condition_dict(uncond, image_cond_in[-uncond.shape[0]:])) denoised_image_indexes = [x[0][0] for x in conds_list] if skip_uncond: diff --git a/state-of-things-next.md b/state-of-things-next.md new file mode 100644 index 00000000000..849eef6e0d4 --- /dev/null +++ b/state-of-things-next.md @@ -0,0 +1,91 @@ +# State of Things and Next Work + +Last updated: 2026-08-11 + +Target machine: 16 GB Apple M1 Mac mini + +Branch: `dev` +Last committed head before this sprint: `771259243a5a6e9a938dcedab80999512b78f5fb` + +## Current result + +This fork remains Automatic1111 with targeted MPS and Metal acceleration rather than a separate inference engine. The working tree now adds two measured changes: + +1. An opt-in, coarse MPS stage profiler enabled with `A1111_MPS_PROFILE=1`. +2. FP16 VAE as the tracked launch default only on M1-family Macs. + +The normal path adds no profiler synchronization. Intel and non-M1 Apple Silicon retain `--no-half-vae` until separately validated. Automatic1111's existing NaN recovery remains enabled and retries VAE decode in FP32 if necessary. + +## Reference workload + +- Prompt: `a dog` +- Negative prompt: empty +- Checkpoint hash: `8ecad70a19` +- Steps: 5 +- Sampler: DPM++ SDE +- Schedule: Karras +- CFG: 1.15 +- Seed: `3163229250` +- Size: 384×640 +- Clip skip: 2 +- NGMS: 1.0, all steps +- Batch: 1 + +Five warm profiled runs produced these medians: + +| VAE precision | Client wall time | Sampler | VAE decode + transfer | +| --- | ---: | ---: | ---: | +| FP32 | 8.450 s | 6.715 s | 1.536 s | +| FP16 | 7.795 s | 6.666 s | 0.972 s | + +FP16 VAE saved about 0.65 seconds end to end and reduced VAE time by about 37%. The nine UNet calls were unchanged: five calls with input `2×4×80×48` and four with `1×4×80×48`, for 14 total batch elements. NGMS is responsible for the four batch-one calls. + +## Output validation + +FP32 and FP16 VAE output was compared for three fixed-seed generations at 384×640 and 512×512: + +| Case | Mean absolute RGB delta | PSNR | Largest 8-bit channel delta | Byte-identical channels | +| --- | ---: | ---: | ---: | ---: | +| `a dog` | 0.0257 | 64.02 dB | 1 | 97.43% | +| `a dog and a cat` | 0.0229 | 64.54 dB | 1 | 97.71% | +| `1man, batman, looking out across the city` | 0.0227 | 64.58 dB | 1 | 97.73% | + +All repeated FP16 runs were deterministic. No NaN, green, black, or corrupted images were observed. A normal `./webui.sh` launch reproduced the validated FP16 output hash. + +## Working-tree changes + +- `modules/mps_stage_profile.py`: request/stage timing, memory snapshots, UNet call accounting, JSON report. +- `modules/processing.py`: coarse generation-stage boundaries. +- `modules/sd_samplers_cfg_denoiser.py`: profiler-only UNet call/shape accounting. +- `webui-macos-env.sh`: M1-family FP16 VAE default; conservative fallback elsewhere. +- `test/test_mps_stage_profile.py`: profiler behavior and disabled-path checks. +- `test/test_macos_launch_defaults.py`: M1, M1 Max, M3, and Intel launch behavior. +- `README.md`: launch, quality, performance, profiling, and troubleshooting documentation. + +The focused suite currently passes 21 tests. Python compilation, shell syntax, `git diff --check`, native Metal self-tests, an API generation, and a normal WebUI launch also pass. + +## What the profile says next + +With FP16 VAE enabled, the sampler/UNet is now roughly 87% of measured generation time. Image conversion and orchestration are negligible. Another material gain cannot come from unified-memory copies, PNG conversion, conditioning, or more VAE tuning; it must reduce UNet work or execute the UNet more efficiently. + +Do not revisit these rejected directions without new evidence: + +- DPM++ 2M substitution: it changes the desired LCM result. +- FP8 on M1: there is no matching M1 hardware acceleration path. +- Per-operator MPS timing events on PyTorch 2.3: isolated event synchronization hung on the tested system. +- The previous block-level MPSGraph prototype: it was about 1% slower end to end and had a larger numerical delta. + +## Recommended next sprint + +The next useful experiment is an opt-in static UNet executor, not another broad rewrite. Keep the normal Automatic1111 model and sampler interfaces, and cache a compiled path by checkpoint, latent shape, conditional batch shape, and active network state. Start with the exact SD 1.x reference workload and refuse unsupported inputs rather than silently changing behavior. + +Suggested gates: + +1. Support only txt2img, SD 1.x, batch 1, 384×640 and 512×512, no ControlNet, and no live model mutation in the prototype. +2. Preserve all nine DPM++ SDE evaluations and both batch-one and batch-two call shapes. +3. Compare five warm alternating runs against the current PyTorch MPS path. +4. Require at least a 5% end-to-end improvement before expanding compatibility. +5. Require deterministic output and report image deviation; fall back to PyTorch for every unsupported or failed graph. +6. Add LoRA-aware cache invalidation before considering a default-on route. + +This is significant engine work. If the static executor cannot clear the 5% end-to-end gate on this M1, the current approximately 7.8-second profiled result is the practical stopping point for compatibility-preserving changes on the pinned PyTorch 2.3 runtime. diff --git a/test/test_macos_launch_defaults.py b/test/test_macos_launch_defaults.py new file mode 100644 index 00000000000..a567123d3be --- /dev/null +++ b/test/test_macos_launch_defaults.py @@ -0,0 +1,35 @@ +import os +from pathlib import Path +import subprocess + + +SCRIPT = Path(__file__).resolve().parents[1] / "webui-macos-env.sh" + + +def command_line_args(cpu_brand): + environment = os.environ.copy() + environment.update({ + "TEST_CPU_BRAND": cpu_brand, + "TEST_MACOS_ENV_SCRIPT": str(SCRIPT), + }) + command = """ +sysctl() { printf '%s\n' "$TEST_CPU_BRAND"; } +SCRIPT_DIR="$(dirname "$TEST_MACOS_ENV_SCRIPT")" +source "$TEST_MACOS_ENV_SCRIPT" +printf '%s' "$COMMANDLINE_ARGS" +""" + result = subprocess.run(["bash", "-c", command], check=True, capture_output=True, text=True, env=environment) + return result.stdout.split() + + +def test_m1_family_uses_fp16_vae_default(): + assert "--no-half-vae" not in command_line_args("Apple M1") + assert "--no-half-vae" not in command_line_args("Apple M1 Max") + + +def test_newer_apple_silicon_retains_fp32_vae_default(): + assert "--no-half-vae" in command_line_args("Apple M3 Pro") + + +def test_intel_retains_fp32_vae_default(): + assert "--no-half-vae" in command_line_args("Intel(R) Core(TM) i9") diff --git a/test/test_mps_stage_profile.py b/test/test_mps_stage_profile.py new file mode 100644 index 00000000000..14190c3891c --- /dev/null +++ b/test/test_mps_stage_profile.py @@ -0,0 +1,90 @@ +from contextlib import redirect_stdout +import io +import json +import os +from types import SimpleNamespace +from unittest import mock + +import torch + +from modules import mps_stage_profile + + +def processing_stub(): + return SimpleNamespace( + batch_size=1, + cfg_scale=1.15, + height=640, + n_iter=1, + sampler_name="DPM++ SDE", + scheduler="Karras", + sd_model_hash="8ecad70a19", + steps=5, + width=384, + ) + + +def report_from_output(output): + line = next(line for line in output.splitlines() if line.startswith(mps_stage_profile.OUTPUT_PREFIX)) + return json.loads(line.removeprefix(mps_stage_profile.OUTPUT_PREFIX)) + + +def test_disabled_profile_does_not_synchronize_or_print(): + output = io.StringIO() + source = torch.zeros((2, 4, 8, 8)) + environment = {key: value for key, value in os.environ.items() if key != mps_stage_profile.ENVIRONMENT_VARIABLE} + + with mock.patch.dict(os.environ, environment, clear=True), mock.patch.object(mps_stage_profile, "_synchronize") as synchronize: + with redirect_stdout(output), mps_stage_profile.request(processing_stub()): + with mps_stage_profile.stage("sampler"): + result = mps_stage_profile.unet_call(lambda value: value + 1, source) + + assert torch.equal(result, source + 1) + assert output.getvalue() == "" + synchronize.assert_not_called() + + +def test_enabled_profile_reports_stages_unet_shapes_and_memory(): + output = io.StringIO() + source = torch.zeros((2, 4, 80, 48)) + memory = { + "process_rss": 100, + "system_available": 200, + "mps_current_allocated_memory": 300, + } + + with mock.patch.dict(os.environ, {mps_stage_profile.ENVIRONMENT_VARIABLE: "1"}), \ + mock.patch.object(mps_stage_profile, "_mps_available", return_value=False), \ + mock.patch.object(mps_stage_profile, "_memory_snapshot", return_value=memory), \ + redirect_stdout(output): + with mps_stage_profile.request(processing_stub()): + with mps_stage_profile.stage("sampler"): + result = mps_stage_profile.unet_call(lambda value: value + 1, source) + + report = report_from_output(output.getvalue()) + assert torch.equal(result, source + 1) + assert report["metadata"]["model_hash"] == "8ecad70a19" + assert report["stages"]["sampler"]["calls"] == 1 + assert report["unet"] == { + "calls": 1, + "input_shapes": {"2x4x80x48": 1}, + "total_batch_elements": 2, + } + assert report["memory_bytes"]["sampled_peaks"] == memory + + +def test_enabled_mps_profile_synchronizes_only_at_coarse_boundaries(): + output = io.StringIO() + memory = {"process_rss": 100, "system_available": 200} + + with mock.patch.dict(os.environ, {mps_stage_profile.ENVIRONMENT_VARIABLE: "1"}), \ + mock.patch.object(mps_stage_profile, "_mps_available", return_value=True), \ + mock.patch.object(mps_stage_profile, "_memory_snapshot", return_value=memory), \ + mock.patch.object(mps_stage_profile, "_synchronize") as synchronize, \ + redirect_stdout(output): + with mps_stage_profile.request(processing_stub()): + with mps_stage_profile.stage("sampler"): + pass + + assert synchronize.call_count == 4 + assert report_from_output(output.getvalue())["stages"]["sampler"]["calls"] == 1 diff --git a/webui-macos-env.sh b/webui-macos-env.sh index 8e4a7711b80..3620ec1f295 100644 --- a/webui-macos-env.sh +++ b/webui-macos-env.sh @@ -5,12 +5,22 @@ #################################################################### export install_dir="$HOME" -export COMMANDLINE_ARGS="--skip-torch-cuda-test --no-half-vae --use-cpu interrogate" +export COMMANDLINE_ARGS="--skip-torch-cuda-test --use-cpu interrogate" export PYTORCH_ENABLE_MPS_FALLBACK=1 -if [[ "$(sysctl -n machdep.cpu.brand_string)" =~ ^.*"Intel".*$ ]]; then +mps_cpu_brand="$(sysctl -n machdep.cpu.brand_string)" + +if [[ "${mps_cpu_brand}" =~ ^.*"Intel".*$ ]]; then + export COMMANDLINE_ARGS="${COMMANDLINE_ARGS} --no-half-vae" export TORCH_COMMAND="pip install torch==2.1.2 torchvision==0.16.2" else + # The FP16 VAE path is measurably faster and has been validated on M1. + # Retain the conservative FP32 VAE default on other Apple Silicon until + # each family has equivalent output-quality and stability coverage. + if [[ "${mps_cpu_brand}" != Apple\ M1* ]]; then + export COMMANDLINE_ARGS="${COMMANDLINE_ARGS} --no-half-vae" + fi + export PIP_CONSTRAINT="${SCRIPT_DIR}/requirements_macos.txt" # Direct Metal matrix multiplication is faster than MPSGraph for the # projection sizes used by Stable Diffusion 1.x on M1. From 9038d1686a9e951ba4a87e3dbfd8e905fe60c05a Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 20:03:20 -0500 Subject: [PATCH 07/17] Fuse exact GEGLU on Apple Silicon --- README.md | 40 ++++-- modules/mps_flash_attention.py | 17 ++- modules/mps_fused_ops.py | 57 ++++++++ modules/sd_hijack_unet.py | 16 +++ modules/shared_options.py | 3 +- scripts/benchmark_mps_geglu_probe.py | 173 +++++++++++++++++++++++++ scripts/install_mps_flash_attention.py | 4 +- scripts/mps_fused_group_norm.mm | 115 ++++++++++++++++ state-of-things-next.md | 22 ++-- test/test_mps_fused_ops.py | 31 +++++ 10 files changed, 455 insertions(+), 23 deletions(-) create mode 100644 scripts/benchmark_mps_geglu_probe.py diff --git a/README.md b/README.md index 1d743c84baf..d35b0a93bb1 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ An Apple Silicon performance fork of [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), focused on faster and more memory-aware inference through PyTorch MPS and native Metal kernels. -The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, a fused GroupNorm + SiLU kernel, unified-memory-aware attention fallback, an M1-validated FP16 VAE path, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. +The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, fused GroupNorm + SiLU and exact-parity GEGLU kernels, unified-memory-aware attention fallback, an M1-validated FP16 VAE path, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. > [!IMPORTANT] > This is an experimental performance fork, not a new Stable Diffusion engine. It favors measured M1 inference performance and safe fallback behavior over broad hardware tuning. If a native Metal path is unavailable or fails its startup test, the WebUI falls back to the corresponding PyTorch implementation. @@ -39,7 +39,7 @@ Most of the added lines are isolated Metal code, benchmark utilities, and tests. | --- | --- | --- | | Metal runtime | `modules/mps_flash_attention.py`
`modules/mps_fused_ops.py`
`modules/mps_utils.py` | `modules/mac_specific.py`
`modules/sd_hijack_optimizations.py`
`modules/sd_hijack_unet.py`
`modules/sub_quadratic_attention.py` | | Startup and defaults | `requirements_macos.txt` | `modules/launch_utils.py`
`modules/shared_options.py`
`requirements_versions.txt`
`webui-macos-env.sh` | -| Native build and benchmarks | `scripts/install_mps_flash_attention.py`
`scripts/mps_fused_group_norm.mm`
`scripts/benchmark_mps_attention.py`
`scripts/benchmark_mps_unet_ops.py` | — | +| Native build and benchmarks | `scripts/install_mps_flash_attention.py`
`scripts/mps_fused_group_norm.mm`
`scripts/benchmark_mps_attention.py`
`scripts/benchmark_mps_unet_ops.py`
`scripts/benchmark_mps_geglu_probe.py` | — | | Tests | `test/test_mps_flash_attention.py`
`test/test_mps_fused_ops.py`
`test/test_mps_utils.py`
`test/test_sub_quadratic_attention.py` | — | @@ -83,6 +83,16 @@ A native inference-only Metal kernel combines GroupNorm and SiLU in one dispatch This is a focused fusion; convolutions and residual additions still use PyTorch MPS. A larger block-level MPSGraph prototype was tested and deliberately rejected because it was about 1% slower end to end and produced a larger numerical delta without a speed benefit. +### Exact-parity fused GEGLU + +The SD 1.x transformer feed-forward path normally stores a GELU result and then launches a separate multiply. On Apple Silicon, the fork combines the lookup and multiply into one Metal dispatch. + +- The model's linear projection still runs normally, so active LoRAs and other projection hooks remain compatible. +- A one-time 65,536-entry FP16 table is generated with the installed PyTorch MPS GELU implementation. The table is 128 KB and maps every possible half-precision gate value to PyTorch's exact result. +- The fused output was byte-identical to PyTorch at all SD 1.x transformer shapes for batch one and batch two. +- CPU, FP32, training/autograd, incompatible layouts, disabled settings, and runtime failures use the original PyTorch implementation. +- Enabled by default through **Settings → Optimizations → Fuse GEGLU on Apple Silicon**. + ### Unified-memory-aware attention The fork treats system RAM and GPU memory as the same constrained resource instead of relying on a fixed attention threshold. @@ -131,6 +141,7 @@ The following defaults intentionally differ from the upstream `dev` branch: | `--no-half-vae` on M1-family Macs | On | Off | Runs VAE encode/decode in FP16; Automatic1111 still retries in FP32 if VAE decode produces NaNs | | Cross-attention Automatic choice on MPS | Sub-quadratic | Metal Flash Attention | Uses the measured native route when available | | Fused GroupNorm + SiLU | Not present | On | Reduces compatible normalization/activation dispatches | +| Fused GEGLU | Not present | On | Preserves PyTorch FP16 output while reducing transformer activation dispatches | NGMS is the largest user-visible behavioral change. It is recorded in PNG generation metadata when active. Set NGMS to `0` and disable **NGMS all steps** if a workflow expects upstream guidance behavior. @@ -147,6 +158,12 @@ That observed run was approximately **32% lower latency**, or **1.47× as fast** The two recorded generations used different seeds. This makes the table a throughput comparison at matching tensor shapes, not an image-parity A/B. +### M1 fused GEGLU validation + +A fixed-process API A/B used `hyperGlance` (`8ecad70a19`), prompt `a dog and a cat`, seed `158926638`, 5-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 512×512 output. Three alternating warm pairs measured a positive saving in every pair: approximately 0.12–0.27 seconds, with a median paired saving of about 0.27 seconds. All fusion-on and fusion-off PNG files had the same SHA-256 hash. + +An additional active-LoRA check used `a dog `, seed `784504668`, five Euler a steps, and the same CFG, size, Clip skip, and NGMS settings. Fusion-on and fusion-off output hashes were identical. The timing from that single LoRA pair is not reported as a speed result because its first run included LoRA activation overhead. + ### M1 FP16 VAE validation A later controlled A/B isolated VAE precision on a 16 GB Apple M1 Mac mini. Both paths used checkpoint `8ecad70a19`, prompt `a dog`, seed `3163229250`, 5-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 384×640 output. Each result below is the median of five warm runs with coarse MPS stage profiling enabled. @@ -232,13 +249,13 @@ M1-family Macs use the validated FP16 VAE path by default. Intel and other Apple A healthy optimized startup prints messages similar to: ```text -Metal self-test passed; deferred MFA and fused GroupNorm+SiLU routing enabled. +Metal self-test passed; deferred MFA, fused GroupNorm+SiLU, and fused GEGLU routing enabled. Applying attention optimization: mps-flash... done. ``` -The first compatible generation also reports the first native attention and GroupNorm dispatch. These messages are informational and print only once per process. +The first compatible generation also reports the first native attention, GroupNorm, and GEGLU dispatch. These messages are informational and print only once per process. -If the extension cannot build or fails its isolated self-test, startup continues with native PyTorch MPS attention. If the fused GroupNorm kernel fails at runtime, that fusion is disabled for the process and PyTorch handles subsequent operations. +If the extension cannot build or fails its isolated self-test, startup continues with native PyTorch MPS operations. If either fused activation kernel fails at runtime, that fusion is disabled for the process and PyTorch handles subsequent operations. ## Compatibility and output parity @@ -270,12 +287,13 @@ For an upstream-style comparison: 1. Set **Negative Guidance minimum sigma** to `0`. 2. Disable **Negative Guidance minimum sigma all steps**. 3. Disable **Fuse GroupNorm and SiLU on Apple Silicon**. -4. Select `sub-quadratic` under **Cross attention optimization**. -5. Add `--upcast-sampling` to `COMMANDLINE_ARGS` in `webui-user.sh`. -6. On M1, also add `--no-half-vae`. -7. Restart the WebUI after changing launch arguments. +4. Disable **Fuse GEGLU on Apple Silicon**. +5. Select `sub-quadratic` under **Cross attention optimization**. +6. Add `--upcast-sampling` to `COMMANDLINE_ARGS` in `webui-user.sh`. +7. On M1, also add `--no-half-vae`. +8. Restart the WebUI after changing launch arguments. -For diagnostics only, `A1111_MPS_FORCE_LEGACY_OPS=1` restores version-gated MPS safety copies, and `A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU=1` disables the native normalization fusion before startup. +For diagnostics only, `A1111_MPS_FORCE_LEGACY_OPS=1` restores version-gated MPS safety copies. `A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU=1` and `A1111_MPS_DISABLE_FUSED_GEGLU=1` disable the corresponding native fusion before startup. ## Troubleshooting @@ -315,6 +333,7 @@ Two standalone benchmark scripts are included: ```bash ./venv/bin/python scripts/benchmark_mps_attention.py ./venv/bin/python scripts/benchmark_mps_unet_ops.py --batch 2 +./venv/bin/python scripts/benchmark_mps_geglu_probe.py ``` The first compares PyTorch MPS scaled dot product attention with sliced attention. The second measures representative SD 1.x convolution, GroupNorm + SiLU, linear projection, and attention shapes. @@ -332,6 +351,7 @@ Focused tests cover: - Metal Flash Attention routing and PyTorch fallback - M1-specific FP16 VAE launch defaults with conservative Intel and newer-chip behavior - Native fused GroupNorm + SiLU correctness +- Native fused GEGLU exact parity, fallback routing, and active-LoRA compatibility - Opt-in MPS stage profiling and its zero-synchronization disabled path - Unified-memory attention budgeting and dynamic query tiles - Streaming online-softmax forward results and gradients diff --git a/modules/mps_flash_attention.py b/modules/mps_flash_attention.py index edb8bd196e7..c4e5bcf7d5f 100644 --- a/modules/mps_flash_attention.py +++ b/modules/mps_flash_attention.py @@ -47,7 +47,7 @@ def _run_isolated_self_test(): code = """ import torch import torch.nn.functional as F -from metal_flash_sdpa import MetalFlashAttentionForward, fused_group_norm_silu_forward +from metal_flash_sdpa import MetalFlashAttentionForward, fused_geglu_forward, fused_group_norm_silu_forward torch.manual_seed(1) source = torch.randn((1, 256, 320), device='mps', dtype=torch.float16) @@ -75,6 +75,17 @@ def _run_isolated_self_test(): assert torch.isfinite(actual_norm).all().item() assert norm_difference.max().item() < 0.02 assert norm_difference.mean().item() < 0.001 + +import numpy as np +geglu_source = torch.randn((1, 256, 2560), device='mps', dtype=torch.float16) +half_values = np.arange(65536, dtype=np.uint16).view(np.float16).copy() +geglu_lut = F.gelu(torch.from_numpy(half_values).to('mps')).contiguous() +value, gate = geglu_source.chunk(2, dim=-1) +expected_geglu = value * F.gelu(gate) +actual_geglu = fused_geglu_forward(geglu_source, geglu_lut) + 0 +torch.mps.synchronize() +assert torch.isfinite(actual_geglu).all().item() +assert torch.equal(actual_geglu, expected_geglu) """ environment = os.environ.copy() environment["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" @@ -115,6 +126,8 @@ def is_available(): raise RuntimeError("native extension is missing the A1111 deferred MPS commit patch") if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU", False): raise RuntimeError("native extension is missing fused GroupNorm+SiLU") + if not getattr(_extension, "A1111_MPS_FUSED_GEGLU", False): + raise RuntimeError("native extension is missing fused GEGLU") _run_isolated_self_test() except (ImportError, OSError, RuntimeError, subprocess.SubprocessError) as exc: _availability_error = str(exc) @@ -123,7 +136,7 @@ def is_available(): return False _availability = True - print("Metal self-test passed; deferred MFA and fused GroupNorm+SiLU routing enabled.") + print("Metal self-test passed; deferred MFA, fused GroupNorm+SiLU, and fused GEGLU routing enabled.") return True diff --git a/modules/mps_fused_ops.py b/modules/mps_fused_ops.py index 68452aebde3..288bd482a09 100644 --- a/modules/mps_fused_ops.py +++ b/modules/mps_fused_ops.py @@ -4,6 +4,7 @@ import os +import numpy as np import torch import torch.nn.functional as F @@ -15,6 +16,12 @@ _runtime_failure_warned = False _first_dispatch_logged = False _runtime_disabled = False +_geglu_dispatch_count = 0 +_geglu_fallback_count = 0 +_geglu_runtime_failure_warned = False +_geglu_first_dispatch_logged = False +_geglu_runtime_disabled = False +_geglu_lut = None def _can_dispatch(input_tensor, norm): @@ -75,8 +82,58 @@ def group_norm_silu(input_tensor, norm): return F.silu(norm(input_tensor)) +def _can_dispatch_geglu(projected): + if _geglu_runtime_disabled: + return False + if os.environ.get("A1111_MPS_DISABLE_FUSED_GEGLU") == "1": + return False + if projected.device.type != "mps" or projected.dtype != torch.float16: + return False + if projected.ndim != 3 or projected.shape[-1] % 2 != 0 or not projected.is_contiguous(): + return False + if torch.is_grad_enabled() and projected.requires_grad: + return False + from modules import shared + + if not getattr(shared.opts, "mps_fused_geglu", True): + return False + return mps_flash_attention.is_available() + + +def geglu(input_tensor, projection): + """Run the model's projection normally, then fuse GEGLU's GELU and multiply.""" + global _geglu_dispatch_count, _geglu_fallback_count + global _geglu_runtime_failure_warned, _geglu_first_dispatch_logged + global _geglu_runtime_disabled, _geglu_lut + + projected = projection(input_tensor) + if _can_dispatch_geglu(projected): + try: + if _geglu_lut is None or _geglu_lut.device != projected.device: + half_values = np.arange(65536, dtype=np.uint16).view(np.float16).copy() + half_values = torch.from_numpy(half_values).to(projected.device) + _geglu_lut = F.gelu(half_values).contiguous() + result = mps_flash_attention._extension.fused_geglu_forward(projected, _geglu_lut) + _geglu_dispatch_count += 1 + if not _geglu_first_dispatch_logged: + print(f"Fused Metal GEGLU first dispatch: {tuple(projected.shape)}") + _geglu_first_dispatch_logged = True + return result + except RuntimeError as exc: + _geglu_runtime_disabled = True + if not _geglu_runtime_failure_warned: + print(f"Fused Metal GEGLU failed; using PyTorch: {exc}") + _geglu_runtime_failure_warned = True + + _geglu_fallback_count += 1 + value, gate = projected.chunk(2, dim=-1) + return value * F.gelu(gate) + + def diagnostics(): return { "dispatches": _dispatch_count, "fallbacks": _fallback_count, + "geglu_dispatches": _geglu_dispatch_count, + "geglu_fallbacks": _geglu_fallback_count, } diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index 148351025dd..8be49303b06 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -104,6 +104,20 @@ def fused_vae_resnet_forward(_, self, x, temb): return x + h +def fused_geglu_condition(_, self, x): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 3 + and not self.training + and hasattr(self, "proj") + ) + + +def fused_geglu_forward(_, self, x): + return mps_fused_ops.geglu(x, self.proj) + + # Below are monkey patches to enable upcasting a float16 UNet for float32 sampling def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): """Always make sure inputs to unet are in correct dtype.""" @@ -197,6 +211,8 @@ def hijack_ddpm_edit(): CondFunc('sgm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) CondFunc('ldm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) CondFunc('sgm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) +CondFunc('ldm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) +CondFunc('sgm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): diff --git a/modules/shared_options.py b/modules/shared_options.py index 23e6f295b6b..6ad2561ecb8 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -231,8 +231,9 @@ })) options_templates.update(options_section(('optimizations', "Optimizations", "sd"), { - "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), + "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), "mps_fused_group_norm_silu": OptionInfo(True, "Fuse GroupNorm and SiLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), + "mps_fused_geglu": OptionInfo(True, "Fuse GEGLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), "s_min_uncond": OptionInfo(1.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), "s_min_uncond_all": OptionInfo(True, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), "token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"), diff --git a/scripts/benchmark_mps_geglu_probe.py b/scripts/benchmark_mps_geglu_probe.py new file mode 100644 index 00000000000..16309460650 --- /dev/null +++ b/scripts/benchmark_mps_geglu_probe.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Measure a standalone fused Metal GEGLU at exact SD 1.x UNet shapes.""" + +from __future__ import annotations + +import argparse +import importlib +import statistics +import time + +import numpy as np +import torch +import torch.nn.functional as F + + +SHAPES = ((4096, 320, 5), (1024, 640, 5), (256, 1280, 6)) + + +def compile_extension(): + module = importlib.import_module("metal_flash_sdpa") + if not getattr(module, "A1111_MPS_FUSED_GEGLU", False): + raise RuntimeError("installed Metal extension does not include fused GEGLU") + return module + + +def reference(projected): + value, gate = projected.chunk(2, dim=-1) + return value * F.gelu(gate) + + +def timed(operation): + torch.mps.synchronize() + started = time.perf_counter() + result = operation() + torch.mps.synchronize() + return (time.perf_counter() - started) * 1000, result + + +def timed_series(operation, calls): + torch.mps.synchronize() + started = time.perf_counter() + result = None + for _ in range(calls): + result = operation() + torch.mps.synchronize() + return (time.perf_counter() - started) * 1000, result + + +def measure(extension, gelu_lut, batch, tokens, channels, calls, repeats): + inner = channels * 4 + generator = torch.Generator(device="cpu").manual_seed(600100635 + batch + channels) + # A real linear projection has roughly unit-scale activations after trained weights. + projected = torch.randn((batch, tokens, inner * 2), generator=generator) + projected = projected.to(device="mps", dtype=torch.float16).contiguous() + + first_reference, expected = timed(lambda: reference(projected)) + first_fused, actual = timed(lambda: extension.fused_geglu_forward(projected, gelu_lut)) + expected_float = expected.float().cpu() + actual_float = actual.float().cpu() + difference = (expected_float - actual_float).abs() + parity = { + "mean_abs": difference.mean().item(), + "max_abs": difference.max().item(), + "cosine": F.cosine_similarity(expected_float.flatten(), actual_float.flatten(), dim=0).item(), + "finite": bool(torch.isfinite(actual_float).all()), + } + parity["passed"] = ( + parity["finite"] + and parity["mean_abs"] <= 0.001 + and parity["max_abs"] <= 0.02 + and parity["cosine"] >= 0.9999 + ) + + for _ in range(5): + reference(projected) + extension.fused_geglu_forward(projected, gelu_lut) + torch.mps.synchronize() + + reference_times = [] + fused_times = [] + for index in range(repeats): + if index % 2: + fused_times.append(timed(lambda: extension.fused_geglu_forward(projected, gelu_lut))[0]) + reference_times.append(timed(lambda: reference(projected))[0]) + else: + reference_times.append(timed(lambda: reference(projected))[0]) + fused_times.append(timed(lambda: extension.fused_geglu_forward(projected, gelu_lut))[0]) + + reference_median = statistics.median(reference_times) + fused_median = statistics.median(fused_times) + + series_reference = [] + series_fused = [] + series_repeats = max(8, repeats // 3) + for index in range(series_repeats): + if index % 2: + series_fused.append(timed_series(lambda: extension.fused_geglu_forward(projected, gelu_lut), calls)[0]) + series_reference.append(timed_series(lambda: reference(projected), calls)[0]) + else: + series_reference.append(timed_series(lambda: reference(projected), calls)[0]) + series_fused.append(timed_series(lambda: extension.fused_geglu_forward(projected, gelu_lut), calls)[0]) + series_reference_median = statistics.median(series_reference) + series_fused_median = statistics.median(series_fused) + return { + "batch": batch, + "tokens": tokens, + "channels": channels, + "first_reference": first_reference, + "first_fused": first_fused, + "reference_median": reference_median, + "fused_median": fused_median, + "saving": reference_median - fused_median, + "speedup": (reference_median - fused_median) / reference_median * 100, + "calls": calls, + "series_reference": series_reference_median, + "series_fused": series_fused_median, + "parity": parity, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--repeats", type=int, default=30) + arguments = parser.parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is required") + + build_started = time.perf_counter() + extension = compile_extension() + half_values = np.arange(65536, dtype=np.uint16).view(np.float16).copy() + gelu_lut = F.gelu(torch.from_numpy(half_values).to("mps")).contiguous() + print(f"Extension build/import: {time.perf_counter() - build_started:.2f} s") + print("Reference: projected.chunk(2) followed by value * torch.nn.functional.gelu(gate)") + + results = [] + for batch in (2, 1): + evaluations = 5 if batch == 2 else 4 + for tokens, channels, count in SHAPES: + result = measure(extension, gelu_lut, batch, tokens, channels, count * evaluations, arguments.repeats) + result["blocks"] = count + results.append(result) + parity = result["parity"] + print( + f"batch={batch} shape={tokens}x{channels} blocks={count} " + f"PyTorch={result['reference_median']:.3f}ms fused={result['fused_median']:.3f}ms " + f"speedup={result['speedup']:.1f}% " + f"mean_abs={parity['mean_abs']:.6f} max_abs={parity['max_abs']:.6f} " + f"cosine={parity['cosine']:.8f} {'PASS' if parity['passed'] else 'FAIL'}" + ) + print( + f" coalesced {result['calls']}-call workload: " + f"PyTorch={result['series_reference']:.2f}ms fused={result['series_fused']:.2f}ms" + ) + + pytorch_total = sum(item["series_reference"] for item in results) + fused_total = sum(item["series_fused"] for item in results) + saving = pytorch_total - fused_total + speedup = saving / pytorch_total * 100 + parity_passed = all(item["parity"]["passed"] for item in results) + # Integration is only worthwhile if the kernel itself is materially faster + # and saves at least 50 ms in the exact nine-call workload estimate. + gate = parity_passed and speedup >= 20 and saving >= 50 + print("\nNine-evaluation DPM++ SDE estimate") + print(f" PyTorch GEGLU activation total: {pytorch_total:.1f} ms") + print(f" Fused GEGLU activation total: {fused_total:.1f} ms") + print(f" Estimated saving: {saving:.1f} ms ({speedup:.1f}%)") + print(f" Numerical parity: {'PASS' if parity_passed else 'FAIL'}") + print(f" RESULT: {'PASS — run an end-to-end integration A/B' if gate else 'FAIL — do not integrate'}") + raise SystemExit(0 if gate else 2) + + +if __name__ == "__main__": + main() diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py index a5927d6de86..b61cc99dd61 100644 --- a/scripts/install_mps_flash_attention.py +++ b/scripts/install_mps_flash_attention.py @@ -103,12 +103,14 @@ def patch_source(source): f'__version__ = "{VERSION}"\n' 'A1111_MPS_STREAM_FIX = True\n' 'A1111_MPS_DEFERRED_COMMIT = True\n' - 'A1111_MPS_FUSED_GROUP_NORM_SILU = True\n', + 'A1111_MPS_FUSED_GROUP_NORM_SILU = True\n' + 'A1111_MPS_FUSED_GEGLU = True\n', ) replace_exact( package_init, "from metal_flash_sdpa._C import mfa_attention_forward, mfa_attention_backward\n", "from metal_flash_sdpa._C import (\n" + " fused_geglu_forward,\n" " fused_group_norm_silu_forward,\n" " mfa_attention_backward,\n" " mfa_attention_forward,\n" diff --git a/scripts/mps_fused_group_norm.mm b/scripts/mps_fused_group_norm.mm index be76a5925a0..8a91fe9fb3d 100644 --- a/scripts/mps_fused_group_norm.mm +++ b/scripts/mps_fused_group_norm.mm @@ -7,6 +7,8 @@ #import #import +#include + namespace { struct FusedGroupNormParams { @@ -17,6 +19,11 @@ float epsilon; }; +struct FusedGEGLUParams { + uint32_t rows; + uint32_t width; +}; + static inline id getMTLBufferStorage(const at::Tensor& tensor) { return __builtin_bit_cast(id, tensor.storage().data()); } @@ -94,6 +101,7 @@ kernel void fused_group_norm_silu_half( output[base + index] = half(value); } } + )METAL"; NSError* error = nil; @@ -113,6 +121,57 @@ kernel void fused_group_norm_silu_half( return pipeline; } +static id getFusedGEGLUPipeline() { + static id pipeline = nil; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +using namespace metal; + +struct FusedGEGLUParams { + uint rows; + uint width; +}; + +kernel void fused_geglu_half( + device const half* input [[buffer(0)]], + device half* output [[buffer(1)]], + device const half* gelu_lut [[buffer(2)]], + constant FusedGEGLUParams& params [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + const uint count = params.rows * params.width; + if (index >= count) { + return; + } + const uint row = index / params.width; + const uint column = index - row * params.width; + const uint input_base = row * params.width * 2; + const float value = float(input[input_base + column]); + device const ushort* input_bits = reinterpret_cast(input); + const ushort gate_bits = input_bits[input_base + params.width + column]; + output[index] = half(value * float(gelu_lut[gate_bits])); +} +)METAL"; + + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK( + library != nil, + "Failed to compile fused GEGLU Metal library: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + id function = [library newFunctionWithName:@"fused_geglu_half"]; + TORCH_CHECK(function != nil, "Fused GEGLU Metal function was not found"); + pipeline = [device newComputePipelineStateWithFunction:function error:&error]; + TORCH_CHECK( + pipeline != nil, + "Failed to create fused GEGLU pipeline: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + }); + return pipeline; +} + torch::Tensor fused_group_norm_silu_forward( const torch::Tensor& input, const torch::Tensor& weight, @@ -178,6 +237,56 @@ kernel void fused_group_norm_silu_half( return output; } +torch::Tensor fused_geglu_forward( + const torch::Tensor& input, + const torch::Tensor& gelu_lut) { + TORCH_CHECK(input.device().is_mps(), "input must be an MPS tensor"); + TORCH_CHECK(input.scalar_type() == at::kHalf, "input must be float16"); + TORCH_CHECK(input.dim() == 3, "input must have shape [batch, tokens, 2 * width]"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(input.size(2) % 2 == 0, "the last input dimension must be even"); + TORCH_CHECK(gelu_lut.device().is_mps(), "GELU lookup table must be an MPS tensor"); + TORCH_CHECK(gelu_lut.scalar_type() == at::kHalf, "GELU lookup table must be float16"); + TORCH_CHECK(gelu_lut.is_contiguous(), "GELU lookup table must be contiguous"); + TORCH_CHECK(gelu_lut.numel() == 65536, "GELU lookup table must contain 65536 values"); + + const int64_t width = input.size(2) / 2; + auto output = torch::empty({input.size(0), input.size(1), width}, input.options()); + FusedGEGLUParams params = { + static_cast(input.size(0) * input.size(1)), + static_cast(width), + }; + const uint32_t count = params.rows * params.width; + auto pipeline = getFusedGEGLUPipeline(); + + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id encoder = [command_buffer computeCommandEncoder]; + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:getMTLBufferStorage(input) + offset:getMTLBufferOffset(input) + atIndex:0]; + [encoder setBuffer:getMTLBufferStorage(output) + offset:getMTLBufferOffset(output) + atIndex:1]; + [encoder setBuffer:getMTLBufferStorage(gelu_lut) + offset:getMTLBufferOffset(gelu_lut) + atIndex:2]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:3]; + const NSUInteger threads = + std::min(pipeline.maxTotalThreadsPerThreadgroup, 256); + [encoder dispatchThreads:MTLSizeMake(count, 1, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [encoder endEncoding]; + } + }); + } + return output; +} + } // namespace void register_fused_ops(pybind11::module_& module) { @@ -190,4 +299,10 @@ void register_fused_ops(pybind11::module_& module) { pybind11::arg("bias"), pybind11::arg("groups"), pybind11::arg("epsilon")); + module.def( + "fused_geglu_forward", + &fused_geglu_forward, + "Fused Metal GEGLU forward pass", + pybind11::arg("input"), + pybind11::arg("gelu_lut")); } diff --git a/state-of-things-next.md b/state-of-things-next.md index 849eef6e0d4..46a8318ca4a 100644 --- a/state-of-things-next.md +++ b/state-of-things-next.md @@ -9,10 +9,13 @@ Last committed head before this sprint: `771259243a5a6e9a938dcedab80999512b78f5f ## Current result -This fork remains Automatic1111 with targeted MPS and Metal acceleration rather than a separate inference engine. The working tree now adds two measured changes: +This fork remains Automatic1111 with targeted MPS and Metal acceleration rather than a separate inference engine. The current optimization stack includes three measured changes from the latest sprints: 1. An opt-in, coarse MPS stage profiler enabled with `A1111_MPS_PROFILE=1`. 2. FP16 VAE as the tracked launch default only on M1-family Macs. +3. An exact-parity fused Metal GEGLU path enabled by default on compatible Mac inference. + +The GEGLU kernel uses a 128 KB table generated once by PyTorch MPS to preserve every possible FP16 GELU result, then combines table lookup and multiplication in one Metal dispatch. Exact SD 1.x batch-one and batch-two tests were byte-identical to PyTorch. A fixed-process 512×512 DPM++ SDE A/B produced identical PNG hashes and saved approximately 0.12–0.27 seconds in every matched pair. An active LCM LoRA output was also byte-identical with the fusion on and off. The normal path adds no profiler synchronization. Intel and non-M1 Apple Silicon retain `--no-half-vae` until separately validated. Automatic1111's existing NaN recovery remains enabled and retries VAE decode in FP32 if necessary. @@ -74,18 +77,19 @@ Do not revisit these rejected directions without new evidence: - FP8 on M1: there is no matching M1 hardware acceleration path. - Per-operator MPS timing events on PyTorch 2.3: isolated event synchronization hung on the tested system. - The previous block-level MPSGraph prototype: it was about 1% slower end to end and had a larger numerical delta. +- A later real-weight MPSGraph ResBlock/down-stage executable: after comparing against the fork's fused GroupNorm baseline, it improved the measured stage by only about 1.4% and failed the integration gate. +- TorchScript fixed-shape UNet tracing: warm results were inconsistent and lost their benefit after cache loss while retaining enough state to increase unified-memory pressure. ## Recommended next sprint -The next useful experiment is an opt-in static UNet executor, not another broad rewrite. Keep the normal Automatic1111 model and sampler interfaces, and cache a compiled path by checkpoint, latent shape, conditional batch shape, and active network state. Start with the exact SD 1.x reference workload and refuse unsupported inputs rather than silently changing behavior. +Do not immediately revisit static UNet tracing or block-level MPSGraph; both have now failed measured gates on this M1. The next compatibility-preserving experiments should remain narrow transformer micro-fusions, with fused LayerNorm as the leading candidate. Packed self-attention QKV or cross-attention KV projection is a larger follow-up only if LoRA and model-mutation invalidation can be made exact. Suggested gates: -1. Support only txt2img, SD 1.x, batch 1, 384×640 and 512×512, no ControlNet, and no live model mutation in the prototype. -2. Preserve all nine DPM++ SDE evaluations and both batch-one and batch-two call shapes. -3. Compare five warm alternating runs against the current PyTorch MPS path. -4. Require at least a 5% end-to-end improvement before expanding compatibility. -5. Require deterministic output and report image deviation; fall back to PyTorch for every unsupported or failed graph. -6. Add LoRA-aware cache invalidation before considering a default-on route. +1. Preserve all nine DPM++ SDE evaluations and both batch-one and batch-two call shapes. +2. Compare alternating warm runs against the current PyTorch MPS path. +3. Require exact tensor parity for lookup-based or algebraically identical fusions; otherwise report image deviation explicitly. +4. Require a positive end-to-end result, not only an isolated kernel win. +5. Preserve LoRA, ControlNet, dynamic resolution, model switching, and training fallbacks. -This is significant engine work. If the static executor cannot clear the 5% end-to-end gate on this M1, the current approximately 7.8-second profiled result is the practical stopping point for compatibility-preserving changes on the pinned PyTorch 2.3 runtime. +A separate whole-UNet Metal, MLX, or Core ML backend remains the only plausible route to a large additional gain. That is significant engine work and should be treated as a new backend rather than another Automatic1111 micro-optimization. diff --git a/test/test_mps_fused_ops.py b/test/test_mps_fused_ops.py index f175a73709f..db06e80670f 100644 --- a/test/test_mps_fused_ops.py +++ b/test/test_mps_fused_ops.py @@ -31,3 +31,34 @@ def test_native_fusion_matches_pytorch(): assert torch.isfinite(actual).all().item() assert difference.max().item() < 0.02 assert difference.mean().item() < 0.001 + + +def test_geglu_cpu_fallback_matches_pytorch(): + torch.manual_seed(2) + projection = torch.nn.Linear(8, 32) + source = torch.randn(2, 6, 8) + + actual = mps_fused_ops.geglu(source, projection) + projected = projection(source) + value, gate = projected.chunk(2, dim=-1) + expected = value * F.gelu(gate) + + assert torch.equal(actual, expected) + + +def test_native_geglu_matches_pytorch(): + if not torch.backends.mps.is_available(): + return + torch.manual_seed(2) + projection = torch.nn.Linear(320, 2560).eval().half().to("mps") + source = torch.randn(2, 256, 320, device="mps", dtype=torch.float16) + + with torch.no_grad(): + projected = projection(source) + value, gate = projected.chunk(2, dim=-1) + expected = value * F.gelu(gate) + actual = mps_fused_ops.geglu(source, projection) + 0 + torch.mps.synchronize() + + assert torch.isfinite(actual).all().item() + assert torch.equal(actual, expected) From 6eefbb402d177ec5166dbb364ea8e313d1bdb206 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 20:26:21 -0500 Subject: [PATCH 08/17] Default CLIP skip to 2 --- modules/shared_options.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/shared_options.py b/modules/shared_options.py index 6ad2561ecb8..47193448d2f 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -179,7 +179,7 @@ "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"), "comma_padding_backtrack": OptionInfo(20, "Prompt word wrap length limit", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1}).info("in tokens - for texts shorter than specified, if they don't fit into 75 token limit, move them to the next 75 token chunk"), "sdxl_clip_l_skip": OptionInfo(False, "Clip skip SDXL", gr.Checkbox).info("Enable Clip skip for the secondary clip model in sdxl. Has no effect on SD 1.5 or SD 2.0/2.1."), - "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}, infotext="Clip skip").link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer"), + "CLIP_stop_at_last_layers": OptionInfo(2, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}, infotext="Clip skip").link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer"), "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"), "randn_source": OptionInfo("GPU", "Random number generator source.", gr.Radio, {"choices": ["GPU", "CPU", "NV"]}, infotext="RNG").info("changes seeds drastically; use CPU to produce the same picture across different videocard vendors; use NV to produce same picture as on NVidia videocards"), "tiling": OptionInfo(False, "Tiling", infotext='Tiling').info("produce a tileable picture"), From 140f562f4937f4819b98ffd9605f4d960673e44a Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Tue, 11 Aug 2026 20:54:58 -0500 Subject: [PATCH 09/17] Update README and state-of-things-next for current project status and optimization details Signed-off-by: Derek Anderson --- README.md | 201 +++++++++++++++++++++++++++++++------ state-of-things-next.md | 213 +++++++++++++++++++++++++++++++--------- 2 files changed, 338 insertions(+), 76 deletions(-) diff --git a/README.md b/README.md index d35b0a93bb1..fa76aabfd5b 100644 --- a/README.md +++ b/README.md @@ -7,49 +7,53 @@ The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA synta > [!IMPORTANT] > This is an experimental performance fork, not a new Stable Diffusion engine. It favors measured M1 inference performance and safe fallback behavior over broad hardware tuning. If a native Metal path is unavailable or fails its startup test, the WebUI falls back to the corresponding PyTorch implementation. -## How far is this from Automatic1111? +## Current project snapshot -The Metal implementation at commit [`78c3fc98`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/78c3fc988011add8f75dc66af259215d7fc56d2c) has an intentionally small, auditable delta from the official Automatic1111 `dev` branch. Documentation-only changes to this README are excluded from the implementation counts below. +The current tested head is [`6eefbb40`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/6eefbb402d177ec5166dbb364ea8e313d1bdb206) on `dev`. It remains recognizably Automatic1111: the performance work is concentrated in MPS routing, a small native Metal extension, macOS launch defaults, profiling, benchmarks, and tests. | Measure | Value | | --- | ---: | | Automatic1111 base | [`1937682a`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/1937682a20f7f0442311a1ede68f9f0cb480163b) | | Base version | `v1.10.1-96-g1937682a` | -| Metal implementation | [`78c3fc98`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/78c3fc988011add8f75dc66af259215d7fc56d2c) | -| Implementation version | `v1.10.1-100-g78c3fc98` | -| Code relationship | 4 implementation commits ahead, 0 upstream commits behind official `dev` | -| Changed implementation paths | 20 of 329 tracked repository paths (6.1%) | -| New implementation files | 12 | -| Modified upstream implementation files | 8 | -| Implementation delta | 1,280 insertions, 46 deletions | - -The four fork commits are: - -1. Apple Silicon dependency, attention, memory, and benchmark foundation. -2. Removal of obsolete MPS safety copies on modern PyTorch. +| Current Metal head | [`6eefbb40`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/6eefbb402d177ec5166dbb364ea8e313d1bdb206) | +| Current version | `v1.10.1-104-g6eefbb40` | +| Code relationship | 8 fork commits ahead of the selected Automatic1111 base | +| Changed tracked paths | 29, including this README and the roadmap | +| Changed implementation/test paths | 27 | +| Total delta | 2,585 insertions, 271 deletions | + +The eight fork commits are: + +1. Apple Silicon dependencies, attention routing, unified-memory budgeting, and benchmark foundation. +2. Removal of obsolete MPS safety operations on modern PyTorch. 3. Metal Flash Attention command-buffer coalescing. 4. Native fused GroupNorm + SiLU for compatible inference blocks. +5. Documentation and launch configuration cleanup. +6. Coarse MPS stage profiling and the M1 FP16 VAE default. +7. Exact-parity native fused GEGLU. +8. Clip skip 2 as the built-in default. -Most of the added lines are isolated Metal code, benchmark utilities, and tests. The fork does **not** change checkpoint formats, prompt syntax, the REST API contract, or the core Gradio workflow. +Most added code is isolated Metal code, profiling, benchmark utilities, and tests. The fork does **not** require a new checkpoint format, prompt syntax, REST API contract, or UI workflow.
-Complete 20-path change surface +Current implementation surface | Area | Added | Modified | | --- | --- | --- | | Metal runtime | `modules/mps_flash_attention.py`
`modules/mps_fused_ops.py`
`modules/mps_utils.py` | `modules/mac_specific.py`
`modules/sd_hijack_optimizations.py`
`modules/sd_hijack_unet.py`
`modules/sub_quadratic_attention.py` | -| Startup and defaults | `requirements_macos.txt` | `modules/launch_utils.py`
`modules/shared_options.py`
`requirements_versions.txt`
`webui-macos-env.sh` | +| Profiling | `modules/mps_stage_profile.py` | `modules/processing.py`
`modules/sd_samplers_cfg_denoiser.py` | +| Startup and defaults | `requirements_macos.txt` | `modules/launch_utils.py`
`modules/shared_options.py`
`requirements_versions.txt`
`webui-macos-env.sh`
`webui-user.sh` | | Native build and benchmarks | `scripts/install_mps_flash_attention.py`
`scripts/mps_fused_group_norm.mm`
`scripts/benchmark_mps_attention.py`
`scripts/benchmark_mps_unet_ops.py`
`scripts/benchmark_mps_geglu_probe.py` | — | -| Tests | `test/test_mps_flash_attention.py`
`test/test_mps_fused_ops.py`
`test/test_mps_utils.py`
`test/test_sub_quadratic_attention.py` | — | +| Tests | `test/test_macos_launch_defaults.py`
`test/test_mps_flash_attention.py`
`test/test_mps_fused_ops.py`
`test/test_mps_stage_profile.py`
`test/test_mps_utils.py`
`test/test_sub_quadratic_attention.py` | — |
You can reproduce the comparison locally: ```bash -git rev-list --left-right --count 1937682a...78c3fc98 -git diff --shortstat 1937682a..78c3fc98 -git diff --name-status 1937682a..78c3fc98 +git rev-list --left-right --count 1937682a...6eefbb40 +git diff --shortstat 1937682a..6eefbb40 +git diff --name-status 1937682a..6eefbb40 ``` ## What is different? @@ -115,6 +119,17 @@ Several workarounds needed by early PyTorch MPS releases are now gated by runtim - Prefers direct Metal matrix multiplication for the SD 1.x projection sizes measured on M1. - Removes Automatic1111's default `--upcast-sampling` flag on Apple Silicon; it can be restored locally when exact upstream behavior is more important than speed. +### Coarse MPS stage profiler + +An opt-in profiler measures the parts of a complete generation that are large enough to guide optimization decisions without adding synchronization to normal inference. + +- Enable it with `A1111_MPS_PROFILE=1 ./webui.sh`. +- Reports conditioning, sampler, VAE decode/transfer, image processing, and request wall time. +- Records every UNet call shape and MPS allocation snapshots in a machine-readable `MPS_PROFILE_JSON` line. +- Adds no MPS synchronization points when disabled. + +On the M1 reference workload, the profiler established that the sampler/UNet consumes roughly 87% of generation time after enabling the FP16 VAE. This is why current roadmap work targets whole-UNet execution rather than PNG conversion, conditioning, or more VAE micro-tuning. + ### Apple Silicon dependency profile The default Apple Silicon environment is pinned to the combination verified for this fork: @@ -139,6 +154,7 @@ The following defaults intentionally differ from the upstream `dev` branch: | NGMS all steps | Off | On | Applies the configured NGMS rule on every eligible step | | `--upcast-sampling` on macOS | On | Off | Keeps more sampling work in FP16 for speed | | `--no-half-vae` on M1-family Macs | On | Off | Runs VAE encode/decode in FP16; Automatic1111 still retries in FP32 if VAE decode produces NaNs | +| Clip skip | `1` | `2` | Uses the common SD 1.x checkpoint default without depending on local `config.json` | | Cross-attention Automatic choice on MPS | Sub-quadratic | Metal Flash Attention | Uses the measured native route when available | | Fused GroupNorm + SiLU | Not present | On | Reduces compatible normalization/activation dispatches | | Fused GEGLU | Not present | On | Preserves PyTorch FP16 output while reducing transformer activation dispatches | @@ -154,7 +170,7 @@ One recorded Apple M1 Mac mini comparison during development used the same check | Automatic1111 `v1.10.1-96-g1937682a` | 5 steps, DPM++ SDE, Karras, CFG 1.15, 384×640, SD 1.x checkpoint `8ecad70a19`, Clip skip 2, NGMS 1/all steps | 12.8 s | | This fork `v1.10.1-99-g38ac556a` | Same sampler, schedule, dimensions, checkpoint hash, Clip skip, and NGMS settings | 8.7 s | -That observed run was approximately **32% lower latency**, or **1.47× as fast**. The current head adds the fused GroupNorm + SiLU path after that recorded comparison. +That observed run was approximately **32% lower latency**, or **1.47× as fast**. Later heads added fused GroupNorm + SiLU, the profiled FP16 VAE default, and exact-parity GEGLU after that recorded comparison. The two recorded generations used different seeds. This makes the table a throughput comparison at matching tensor shapes, not an image-parity A/B. @@ -177,6 +193,28 @@ FP16 reduced the measured VAE stage by about **37%** and end-to-end latency by a Output quality was checked across three fixed-seed generations at 384×640 and 512×512. Compared with FP32 VAE output, every changed 8-bit RGB channel differed by at most 1 value, PSNR was 64.0–64.6 dB, and 97.4–97.7% of channels were byte-identical. All FP16 runs were deterministic and free of NaN, green, or corrupted output. The default is therefore enabled only on the tested M1 family; other Apple Silicon generations retain FP32 VAE until separately validated. +### Current warm-run range + +A user-facing run at `v1.10.1-102-g58e63e9f` used `fast-model` (`8ecad70a19`), prompt `a dog`, seed `4017012032`, five-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 512×512 output. It completed in **8.3 seconds** on the 16 GB M1 Mac mini. + +Later controlled development A/B runs of the same 512×512 shape typically clustered around 9.2–9.3 seconds after warm-up. Background load, extension startup activity, thermal state, and unified-memory pressure therefore matter at the sub-second scale. Report medians and the full generation settings rather than treating a single fastest run as a guarantee. + +## Experiments that did not pass the gate + +Failed experiments are documented to prevent attractive microbenchmarks from being repeated without new evidence. + +| Experiment | Isolated result | End-to-end result | Decision | +| --- | --- | --- | --- | +| DPM++ 2M substitution | Fewer or cheaper operations in some paths | Did not reproduce the desired LCM/DPM++ SDE result | Rejected; preserve the requested sampler | +| Block-level MPSGraph | Working block prototype | About 1% slower with a larger numerical delta | Removed | +| Real-weight MPSGraph ResBlock/down stage | About 1.4% stage improvement versus the fused GroupNorm baseline | Too small to survive integration overhead | Removed | +| Fixed-shape TorchScript UNet | Some warm runs improved | Benefit was inconsistent and disappeared after cache loss while retaining extra unified memory | Removed | +| Native fused LayerNorm | Exact SD1 workload projection suggested about 110 ms potential savings | Baseline median 11.512 s versus 11.503 s enabled; paired median regressed by 0.026 s. Only 56.1% of RGB channels were identical, with PSNR 47.53 dB | Removed because there was no speed gain and output drifted | +| Cross-attention K/V reuse | Reused 112 of 144 projections | Baseline median 9.258 s versus 9.262 s cached; paired median regressed by 0.016 s while retaining 11 MiB. PNG hashes were identical | Removed because the projections were already too cheap | +| FP8 on M1 | Reduced theoretical weight storage | No matching M1 FP8 execution path; conversion/unpacking would dominate | Not implemented | + +The LayerNorm and K/V probes were completely removed after testing. They are not hidden options and do not remain in the native extension. The repository returned to a clean state after each rejected sprint. + Treat these numbers as a development result, not a universal guarantee. Timing varies with: - Apple Silicon generation and GPU core count @@ -368,13 +406,120 @@ With `pytest` installed in the virtual environment: test/test_sub_quadratic_attention.py ``` -## Deliberately not included +## Future roadmap: native ggml/Metal UNet + +The remaining material opportunity is engine-level work. The best incremental direction is inspired by [stable-diffusion.cpp](https://github.com/leejet/stable-diffusion.cpp) and ggml: execute the complete UNet as one planned Metal graph with a reusable memory arena instead of crossing the Python/PyTorch boundary for individual kernels. + +This is not a plan to replace Automatic1111 wholesale. Prompt parsing, conditioning, the selected A1111/k-diffusion sampler, CFG and NGMS behavior, seed handling, extensions, VAE, image processing, metadata, API, and UI remain in the existing application. Only a compatible UNet evaluation may be delegated to the native backend. + +```text +Automatic1111 prompt, LoRA, and conditioning setup + | +Existing DPM++ SDE / Karras sampler and NGMS logic + | + Native ggml/Metal UNet evaluation + | +Existing CFG combination, VAE, image pipeline, API, and UI +``` + +stable-diffusion.cpp is relevant because its current implementation already provides a complete [SD 1.x UNet graph runner](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/src/model/diffusion/unet.hpp#L748), a [reusable graph allocator](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/src/core/ggml_extend.hpp#L2212), whole-graph [Metal command encoding](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/ggml/src/ggml-metal/ggml-metal-context.m#L438), safetensors/GGUF loading, LoRA support, Flash Attention, and fused quantized matrix kernels. Its advantage comes from owning the graph, buffers, weights, and submission lifecycle together—not from one operator that can be dropped into PyTorch. + +### Phase 0: reproducible captured-tensor corpus + +Capture the inputs and reference outputs of every UNet evaluation from the existing M1 workload: + +- Five calls with batch two and four calls with batch one under five-step DPM++ SDE plus NGMS. +- Both `512×512` and `384×640` latent shapes. +- Latent input, timestep, text conditioning, model hash, precision, and PyTorch output. +- A plain prompt, scheduled prompt, active LoRA, and a deliberately unsupported request for fallback testing. + +The capture path must be diagnostic-only and must not alter normal timing or output when disabled. + +### Phase 1: standalone native UNet shootout + +Build a small C/C++ harness around stable-diffusion.cpp's `UNetModelRunner`. Load the same SD 1.x safetensors checkpoint and replay the captured calls outside WebUI. + +Measure: + +- Per-call and complete nine-call latency after warm-up. +- Batch-one and batch-two behavior separately. +- Peak and retained unified memory. +- Mean, maximum, and percentile tensor deviation from PyTorch MPS. +- Determinism across repeated runs. + +Proceed only if the native nine-call workload is at least **20–25% faster** than the current PyTorch MPS UNet. A smaller isolated advantage is unlikely to survive framework-bridge synchronization and compatibility handling. + +### Phase 2: copied-buffer A1111 prototype + +Expose a minimal native interface that accepts latent, timestep, and conditioning buffers and returns the UNet prediction. Keep A1111's current sampler in control, initially accepting one synchronization and copy boundary per UNet evaluation. + +The first supported route should be intentionally narrow: + +- Apple M1 and SD 1.x only. +- FP16 inference. +- Txt2img, batch one, tested resolutions. +- No ControlNet, hypernetwork, training, or high-resolution pass. +- No active LoRA until mutation/invalidation is explicitly implemented. + +Every unsupported request must automatically use the existing PyTorch UNet. The native backend should be an optional `SdUnetOption`, never a global monkey patch with no escape path. + +### Phase 3: unified-memory zero-copy proof + +If the copied prototype remains faster, investigate sharing the underlying Metal storage rather than copying through CPU memory. PyTorch MPS tensors and ggml Metal tensors ultimately reside in `MTLBuffer` objects, but safe sharing requires explicit work on: + +- Buffer offsets, strides, dtype, and NCHW layout agreement. +- Ownership and lifetime across Python, PyTorch, and the native runner. +- Command-queue ordering and synchronization. +- Error recovery without leaving either backend in a poisoned state. + +This phase should begin with one captured UNet call. Do not attempt full sampling until the shared-buffer output matches the copied native path. + +### Phase 4: compatibility expansion + +Add features one at a time, with a PyTorch fallback and an output test for each: + +1. Dynamic SD 1.x resolutions and cached arenas per batch/shape regime. +2. Active LoRA application, model-mutation generation counters, and exact cache invalidation. +3. Img2img and inpainting conditioning. +4. High-resolution pass and model switching. +5. ControlNet where native semantics can match the installed A1111 extension. +6. Other model families only after SD 1.x is stable. + +Extension compatibility is a routing problem: requests using unsupported hooks should remain fully functional on PyTorch rather than partially executing through native code. + +### Phase 5: optional GGUF quantization + +Quantization follows a successful FP16 engine; it is not the first step. ggml gains from quantized weights because dequantization is fused into its Metal matrix kernels. Merely storing quantized tensors in PyTorch would not reproduce that behavior. + +Suggested order for the M1: + +1. FP16 native backend establishes the execution-engine benefit and parity baseline. +2. Q8_0 evaluates memory reduction with the smallest expected quality risk. +3. Q6_K or Q5_K may become optional balanced modes. +4. Q4 remains an explicit low-memory choice, not the default. + +Each format requires fixed-seed image comparisons, tensor statistics, LoRA checks, and end-to-end timing. A smaller model file alone is not a speed result. + +### Roadmap acceptance gates + +A native backend is eligible for default use only when it: + +1. Improves multiple alternating warm end-to-end pairs, not just an operator microbenchmark. +2. Preserves all nine DPM++ SDE evaluations and the current sampler's result. +3. Reports deterministic output and quantified deviation from the PyTorch path. +4. Does not retain enough extra unified memory to erase warm-run stability. +5. Falls back cleanly for LoRA, ControlNet, dynamic shapes, training, and extensions it cannot reproduce. +6. Can be disabled without changing checkpoint files or local configuration. + +The initial target is to determine whether native UNet execution can move the 16 GB M1 from the current roughly 8–9 second warm range toward 7–8 seconds. Phase 1 is deliberately a bounded proof: if the raw native UNet cannot clear its 20–25% gate, the integration project stops before modifying WebUI. + +### What not to borrow incrementally -- Block-level MPSGraph execution: implemented and benchmarked, but rejected after a small regression. -- FP8 acceleration on M1: there is no matching M1 hardware fast path, so conversion would primarily add unpacking overhead. -- Core ML/ANE conversion: this would introduce a separate static execution engine and materially reduce Automatic1111 compatibility. -- Whole-UNet static graphs: potentially higher upside, but a much larger project with difficult LoRA, ControlNet, model-switching, and dynamic-resolution tradeoffs. -- Model-format changes or required quantization: existing Automatic1111 checkpoints are used directly. +- Individual ggml convolutions or matrix kernels called from PyTorch. Repeated framework and command-queue boundaries would likely erase their benefit. +- Another standalone Flash Attention implementation. The fork already has a measured native MFA route. +- Required GGUF conversion. Existing Automatic1111 checkpoints remain the default input until an optional native backend proves itself. +- VAE tiling at ordinary 512-pixel resolutions. It reduces peak memory but normally increases latency. +- Sampler substitution. stable-diffusion.cpp supports related DPM++ samplers, but this project must retain the exact A1111 DPM++ SDE behavior already chosen for LCM output. ## Upstream features and documentation diff --git a/state-of-things-next.md b/state-of-things-next.md index 46a8318ca4a..dc5fc646922 100644 --- a/state-of-things-next.md +++ b/state-of-things-next.md @@ -1,25 +1,35 @@ -# State of Things and Next Work +# State of Things and Native UNet Roadmap Last updated: 2026-08-11 Target machine: 16 GB Apple M1 Mac mini Branch: `dev` -Last committed head before this sprint: `771259243a5a6e9a938dcedab80999512b78f5fb` -## Current result +Current committed head before this documentation update: `6eefbb402d177ec5166dbb364ea8e313d1bdb206` -This fork remains Automatic1111 with targeted MPS and Metal acceleration rather than a separate inference engine. The current optimization stack includes three measured changes from the latest sprints: +Automatic1111 base: `1937682a20f7f0442311a1ede68f9f0cb480163b` -1. An opt-in, coarse MPS stage profiler enabled with `A1111_MPS_PROFILE=1`. -2. FP16 VAE as the tracked launch default only on M1-family Macs. -3. An exact-parity fused Metal GEGLU path enabled by default on compatible Mac inference. +## Current state -The GEGLU kernel uses a 128 KB table generated once by PyTorch MPS to preserve every possible FP16 GELU result, then combines table lookup and multiplication in one Metal dispatch. Exact SD 1.x batch-one and batch-two tests were byte-identical to PyTorch. A fixed-process 512×512 DPM++ SDE A/B produced identical PNG hashes and saved approximately 0.12–0.27 seconds in every matched pair. An active LCM LoRA output was also byte-identical with the fusion on and off. +This remains an Automatic1111 fork with targeted MPS and native Metal acceleration. It does not currently contain a separate diffusion engine. -The normal path adds no profiler synchronization. Intel and non-M1 Apple Silicon retain `--no-half-vae` until separately validated. Automatic1111's existing NaN recovery remains enabled and retries VAE decode in FP32 if necessary. +The active optimization stack is: -## Reference workload +1. Selective Draw Things-style Metal Flash Attention for measured SD 1.x shapes, encoded on PyTorch's current MPS command buffer. +2. Unified-memory-aware routing to native or sub-quadratic attention. +3. Native fused GroupNorm + SiLU for compatible FP16 UNet and VAE blocks. +4. Exact-parity fused GEGLU using a 65,536-entry, 128 KB PyTorch-generated FP16 GELU table. +5. Modern-PyTorch removal of obsolete MPS clones and FP32 LayerNorm workarounds. +6. FP16 VAE as the tracked default on the tested M1 family, with Automatic1111's FP32 NaN retry retained. +7. NGMS 1.0/all steps and Clip skip 2 as built-in defaults. +8. An opt-in coarse profiler enabled by `A1111_MPS_PROFILE=1` with no synchronization in the disabled path. + +The native extension performs an isolated MPS startup test. Unsupported inputs and runtime failures retain PyTorch fallbacks. + +## Reference workloads + +Primary profiling workload: - Prompt: `a dog` - Negative prompt: empty @@ -34,62 +44,169 @@ The normal path adds no profiler synchronization. Intel and non-M1 Apple Silicon - NGMS: 1.0, all steps - Batch: 1 -Five warm profiled runs produced these medians: +The sampler makes nine UNet evaluations: five calls at batch two and four calls at batch one. At 384×640 the latent inputs are `2×4×80×48` and `1×4×80×48`. NGMS creates the batch-one regime. + +The recurring 512×512 validation workload uses prompt `a dog`, seed `4017012032`, and the same model, sampler, schedule, CFG, Clip skip, and NGMS settings. + +## Measured results + +### Fork versus Automatic1111 baseline -| VAE precision | Client wall time | Sampler | VAE decode + transfer | +At 384×640, the original Automatic1111 base recorded 12.8 seconds and an earlier fork head recorded 8.7 seconds at matching model hash and tensor shape: approximately 32% lower latency or 1.47× throughput. The paired runs used different seeds, so this is a throughput comparison rather than output parity. + +### FP16 VAE + +Five warm profiled 384×640 runs produced: + +| VAE precision | End-to-end median | Sampler | VAE decode + transfer | | --- | ---: | ---: | ---: | | FP32 | 8.450 s | 6.715 s | 1.536 s | | FP16 | 7.795 s | 6.666 s | 0.972 s | -FP16 VAE saved about 0.65 seconds end to end and reduced VAE time by about 37%. The nine UNet calls were unchanged: five calls with input `2×4×80×48` and four with `1×4×80×48`, for 14 total batch elements. NGMS is responsible for the four batch-one calls. +FP16 saved about 0.65 seconds end to end and reduced the VAE stage by about 37%. Across three fixed-seed cases, PSNR versus FP32 was 64.0–64.6 dB, every changed RGB channel moved by at most one 8-bit value, and 97.4–97.7% of channels were identical. + +### Exact GEGLU + +A 512×512 alternating A/B produced identical PNG hashes and positive paired savings of approximately 0.12–0.27 seconds. An active LCM LoRA output hash was also identical with fusion enabled and disabled. + +### Current range + +A normal user run at 512×512 recorded 8.3 seconds. Later controlled warm A/B sessions commonly clustered around 9.2–9.3 seconds. Background work, extensions, thermal state, and unified-memory pressure are material at this scale. + +## Rejected experiments + +Do not repeat these without a new mechanism or new evidence: + +| Experiment | Evidence | Result | +| --- | --- | --- | +| DPM++ 2M substitution | Changed the desired LCM/DPM++ SDE image behavior | Reject sampler substitution | +| FP8 on M1 | No M1 FP8 hardware execution path | Reject conversion/unpacking overhead | +| Per-operator MPS events on PyTorch 2.3 | Isolated synchronization hung | Use coarse profiling | +| Block-level MPSGraph | About 1% slower end to end with more numerical drift | Removed | +| Real-weight MPSGraph block | About 1.4% isolated stage gain versus fused GroupNorm | Failed integration gate | +| Fixed-shape TorchScript UNet | Inconsistent warm gain, lost after cache loss, increased retained memory | Removed | +| Fused LayerNorm | Projected 110 ms microbenchmark gain; paired end-to-end median regressed 0.026 s and only 56.1% of RGB channels matched | Removed | +| Cross-attention K/V reuse | Reused 112/144 projections and retained 11 MiB; paired end-to-end median regressed 0.016 s | Removed | + +The LayerNorm and K/V code and saved probe settings were removed completely. The current repository and native extension contain neither path. + +## What the profile says + +After the FP16 VAE improvement, sampling/UNet consumes roughly 87% of measured generation time. Conditioning, image conversion, metadata, PNG creation, and additional VAE micro-tuning cannot provide the next material gain. + +The failed LayerNorm and K/V experiments also show that transformer micro-operations are now below the useful granularity. The next work must reduce framework overhead across a large portion of the UNet or execute the complete UNet more efficiently. + +## Chosen direction: native ggml/Metal UNet sidecar + +Take architectural inspiration from stable-diffusion.cpp and ggml without replacing Automatic1111. + +Keep in Automatic1111: + +- Prompt parsing and conditioning. +- Existing DPM++ SDE/Karras sampler. +- CFG and NGMS decisions. +- Seed and RNG behavior. +- LoRA/extension activation and request routing. +- VAE, image pipeline, metadata, API, and UI. + +Delegate only a supported UNet evaluation to a native graph runner. Unsupported requests continue through the current PyTorch MPS UNet. + +stable-diffusion.cpp already demonstrates the relevant components: a complete SD 1.x UNet graph, graph-planned reusable buffers, whole-graph Metal encoding, safetensors/GGUF loading, LoRA support, Flash Attention, and fused quantized matrix kernels. The useful lesson is ownership of the complete graph and memory lifecycle, not copying individual kernels. + +## Phase 0: capture the existing UNet contract + +Create a diagnostic-only capture of all nine real calls for 384×640 and 512×512: + +- Latent input and reference output. +- Timestep. +- Text conditioning. +- Shape, dtype, model hash, and request settings. +- Batch-two and batch-one regimes. + +Include plain prompt, scheduled-prompt, active-LoRA, and unsupported/fallback fixtures. Disabled capture must add no synchronization or normal-path overhead. + +Deliverable: a reproducible tensor corpus and a PyTorch replay test. + +## Phase 1: standalone native shootout + +Build a small harness around stable-diffusion.cpp's `UNetModelRunner`, load the same SD 1.x checkpoint, and replay the captures outside WebUI. + +Measure complete nine-call warm latency, per-shape latency, retained/peak memory, determinism, and tensor deviation. + +Gate: the native nine-call workload must be at least 20–25% faster than current PyTorch MPS. Stop the project here if it does not clear the gate; smaller gains will likely disappear behind bridge synchronization and compatibility work. + +## Phase 2: copied-buffer WebUI prototype + +Expose a minimal native interface for latent, timestep, conditioning, and UNet output. Keep A1111's sampler in control. A first implementation may synchronize and copy once per UNet call to prove integration. + +Initial supported route: + +- Apple M1. +- SD 1.x FP16. +- Txt2img, batch one. +- Tested 384×640 and 512×512 shapes. +- No active LoRA, ControlNet, hypernetwork, training, or high-resolution pass. + +Implement it as an optional `SdUnetOption`. Every unsupported condition routes to PyTorch. + +Gate: multiple alternating end-to-end pairs must remain materially faster, deterministic, and within an explicitly approved output-deviation envelope. + +## Phase 3: zero-copy unified-memory proof + +If the copied bridge wins, share the underlying Metal storage between PyTorch MPS and ggml. + +Solve and test: -## Output validation +- `MTLBuffer` ownership and lifetime. +- Buffer offsets, strides, NCHW layout, and dtype agreement. +- PyTorch and ggml command-queue ordering. +- Error recovery and backend reset. -FP32 and FP16 VAE output was compared for three fixed-seed generations at 384×640 and 512×512: +Start with one captured UNet call. Do not attempt full sampling until shared-buffer output matches the copied native implementation. -| Case | Mean absolute RGB delta | PSNR | Largest 8-bit channel delta | Byte-identical channels | -| --- | ---: | ---: | ---: | ---: | -| `a dog` | 0.0257 | 64.02 dB | 1 | 97.43% | -| `a dog and a cat` | 0.0229 | 64.54 dB | 1 | 97.71% | -| `1man, batman, looking out across the city` | 0.0227 | 64.58 dB | 1 | 97.73% | +## Phase 4: compatibility expansion -All repeated FP16 runs were deterministic. No NaN, green, black, or corrupted images were observed. A normal `./webui.sh` launch reproduced the validated FP16 output hash. +Add independently gated support in this order: -## Working-tree changes +1. Dynamic SD 1.x resolutions and reusable arenas per shape/batch regime. +2. LoRA weight application plus explicit model-mutation generation counters. +3. Img2img and inpainting. +4. High-resolution pass and model switching. +5. ControlNet where native semantics can match A1111. +6. Additional model families. -- `modules/mps_stage_profile.py`: request/stage timing, memory snapshots, UNet call accounting, JSON report. -- `modules/processing.py`: coarse generation-stage boundaries. -- `modules/sd_samplers_cfg_denoiser.py`: profiler-only UNet call/shape accounting. -- `webui-macos-env.sh`: M1-family FP16 VAE default; conservative fallback elsewhere. -- `test/test_mps_stage_profile.py`: profiler behavior and disabled-path checks. -- `test/test_macos_launch_defaults.py`: M1, M1 Max, M3, and Intel launch behavior. -- `README.md`: launch, quality, performance, profiling, and troubleshooting documentation. +Never silently ignore an installed extension hook. Fall back to PyTorch for any request whose semantics the native backend cannot reproduce. -The focused suite currently passes 21 tests. Python compilation, shell syntax, `git diff --check`, native Metal self-tests, an API generation, and a normal WebUI launch also pass. +## Phase 5: optional quantization -## What the profile says next +Only after FP16 proves the native engine: -With FP16 VAE enabled, the sampler/UNet is now roughly 87% of measured generation time. Image conversion and orchestration are negligible. Another material gain cannot come from unified-memory copies, PNG conversion, conditioning, or more VAE tuning; it must reduce UNet work or execute the UNet more efficiently. +1. Q8_0 for the lowest-risk memory experiment. +2. Q6_K/Q5_K as optional balanced modes. +3. Q4 as an explicit low-memory mode, not a default. -Do not revisit these rejected directions without new evidence: +Quantization must use native fused dequantization/matrix kernels. Do not add quantized PyTorch storage with per-call unpacking. Require tensor, fixed-seed image, LoRA, memory, and timing validation for every format. -- DPM++ 2M substitution: it changes the desired LCM result. -- FP8 on M1: there is no matching M1 hardware acceleration path. -- Per-operator MPS timing events on PyTorch 2.3: isolated event synchronization hung on the tested system. -- The previous block-level MPSGraph prototype: it was about 1% slower end to end and had a larger numerical delta. -- A later real-weight MPSGraph ResBlock/down-stage executable: after comparing against the fork's fused GroupNorm baseline, it improved the measured stage by only about 1.4% and failed the integration gate. -- TorchScript fixed-shape UNet tracing: warm results were inconsistent and lost their benefit after cache loss while retaining enough state to increase unified-memory pressure. +## Global gates -## Recommended next sprint +1. Preserve A1111's exact DPM++ SDE evaluation sequence and both batch regimes. +2. Benchmark alternating warm pairs, never a single best run. +3. Report tensor and final-image deviation. +4. Measure retained unified memory and cache-loss behavior. +5. Preserve deterministic output within each path. +6. Keep automatic PyTorch fallback for unsupported features and runtime errors. +7. Never require checkpoint conversion for the normal PyTorch path. +8. Keep the native backend independently disableable. -Do not immediately revisit static UNet tracing or block-level MPSGraph; both have now failed measured gates on this M1. The next compatibility-preserving experiments should remain narrow transformer micro-fusions, with fused LayerNorm as the leading candidate. Packed self-attention QKV or cross-attention KV projection is a larger follow-up only if LoRA and model-mutation invalidation can be made exact. +## Explicit non-goals for the first sprint -Suggested gates: +- Replacing the A1111 UI, API, sampler, VAE, or extension ecosystem. +- Calling individual ggml convolutions or matrix kernels from PyTorch. +- Adding another Flash Attention implementation. +- Making GGUF mandatory. +- Supporting every model family or extension before the SD 1.x proof. +- Committing a backend before the standalone 20–25% gate passes. -1. Preserve all nine DPM++ SDE evaluations and both batch-one and batch-two call shapes. -2. Compare alternating warm runs against the current PyTorch MPS path. -3. Require exact tensor parity for lookup-based or algebraically identical fusions; otherwise report image deviation explicitly. -4. Require a positive end-to-end result, not only an isolated kernel win. -5. Preserve LoRA, ControlNet, dynamic resolution, model switching, and training fallbacks. +## Immediate next task -A separate whole-UNet Metal, MLX, or Core ML backend remains the only plausible route to a large additional gain. That is significant engine work and should be treated as a new backend rather than another Automatic1111 micro-optimization. +Implement Phase 0 only: capture and replay the exact nine-call PyTorch UNet contract. Then create the standalone native replay harness. Do not begin WebUI integration until the raw native shootout produces a clear result. From d3ca0fd695a3101fbc6d6380327ac4b6be731dab Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 07:34:06 -0500 Subject: [PATCH 10/17] Add MPS UNet capture functionality and validation scripts - Implement opt-in capture of UNet calls for native-backend experiments in mps_unet_capture.py. - Enhance sd_hijack_unet.py to utilize the new capture functionality for validating outputs. - Create inspect_mps_unet_capture.py for inspecting captured data and validating replay accuracy. - Add unit tests for capture functionality in test_mps_unet_capture.py to ensure correct behavior. - Update state-of-things-next.md with findings from initial capture and performance benchmarks. Signed-off-by: Derek Anderson --- modules/mps_unet_capture.py | 186 ++++++++++++++++++++++++++++ modules/sd_hijack_unet.py | 13 +- scripts/inspect_mps_unet_capture.py | 46 +++++++ state-of-things-next.md | 28 ++++- test/test_mps_unet_capture.py | 84 +++++++++++++ 5 files changed, 354 insertions(+), 3 deletions(-) create mode 100644 modules/mps_unet_capture.py create mode 100644 scripts/inspect_mps_unet_capture.py create mode 100644 test/test_mps_unet_capture.py diff --git a/modules/mps_unet_capture.py b/modules/mps_unet_capture.py new file mode 100644 index 00000000000..1b3d366a437 --- /dev/null +++ b/modules/mps_unet_capture.py @@ -0,0 +1,186 @@ +"""Opt-in capture of one real SD1 UNet call for native-backend experiments.""" + +import json +import os +from pathlib import Path +import statistics +import threading +import time + +import torch +from safetensors.torch import save_file + + +ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET" +BATCH_ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET_BATCH" +BENCHMARK_RUNS_ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET_BENCHMARK_RUNS" +FORMAT_VERSION = 1 + +_capture_lock = threading.Lock() +_capture_claimed = False + + +def enabled(): + return bool(os.environ.get(ENVIRONMENT_VARIABLE)) + + +def _claim(input_tensor): + global _capture_claimed + + if not enabled(): + return False + + try: + requested_batch = int(os.environ.get(BATCH_ENVIRONMENT_VARIABLE, "2")) + except ValueError: + requested_batch = 2 + + if input_tensor.ndim == 0 or input_tensor.shape[0] != requested_batch: + return False + + with _capture_lock: + if _capture_claimed: + return False + _capture_claimed = True + return True + + +def _tensor_key(path): + return "condition." + ".".join(str(part).replace("%", "%25").replace(".", "%2E") for part in path) + + +def _flatten(value, path, tensors): + if isinstance(value, torch.Tensor): + key = _tensor_key(path) + tensors[key] = value.detach().to("cpu").contiguous() + return {"type": "tensor", "key": key} + if isinstance(value, dict): + return { + "type": "dict", + "items": [[str(key), _flatten(item, (*path, key), tensors)] for key, item in value.items()], + } + if isinstance(value, (list, tuple)): + return { + "type": "tuple" if isinstance(value, tuple) else "list", + "items": [_flatten(item, (*path, index), tensors) for index, item in enumerate(value)], + } + if value is None or isinstance(value, (bool, int, float, str)): + return {"type": "literal", "value": value} + return {"type": "unsupported", "python_type": type(value).__qualname__, "repr": repr(value)} + + +def _model_metadata(): + try: + from modules import shared + + model = shared.sd_model + checkpoint = getattr(model, "sd_checkpoint_info", None) + return { + "checkpoint_filename": getattr(checkpoint, "filename", None), + "checkpoint_sha256": getattr(checkpoint, "sha256", None), + "checkpoint_shorthash": getattr(checkpoint, "shorthash", None), + "model_class": type(model).__qualname__ if model is not None else None, + } + except Exception as error: + return {"metadata_error": f"{type(error).__name__}: {error}"} + + +def _validation(reference, replay): + reference_float = reference.detach().float().to("cpu") + replay_float = replay.detach().float().to("cpu") + difference = (reference_float - replay_float).abs() + return { + "exact": bool(torch.equal(reference.detach().to("cpu"), replay.detach().to("cpu"))), + "mean_absolute_error": float(difference.mean().item()), + "maximum_absolute_error": float(difference.max().item()), + } + + +def _synchronize(device): + if device.type == "mps" and torch.backends.mps.is_available(): + torch.mps.synchronize() + + +def _benchmark(run_again, device): + try: + measured_runs = int(os.environ.get(BENCHMARK_RUNS_ENVIRONMENT_VARIABLE, "0")) + except ValueError: + measured_runs = 0 + if measured_runs < 1: + return None + + run_again() + _synchronize(device) + milliseconds = [] + for _ in range(measured_runs): + _synchronize(device) + started = time.perf_counter() + run_again() + _synchronize(device) + milliseconds.append((time.perf_counter() - started) * 1000) + + return { + "runs": measured_runs, + "median_ms": statistics.median(milliseconds), + "minimum_ms": min(milliseconds), + "maximum_ms": max(milliseconds), + "all_ms": milliseconds, + } + + +def capture_and_validate(run_again, input_tensor, timestep, condition, reference_output): + """Capture the first requested batch and immediately prove it replays in PyTorch.""" + if not _claim(input_tensor): + return + + destination = Path(os.environ[ENVIRONMENT_VARIABLE]).expanduser() + if destination.suffix != ".safetensors": + print(f"UNet capture skipped: {ENVIRONMENT_VARIABLE} must name a .safetensors file") + return + if destination.exists(): + print(f"UNet capture skipped: destination already exists: {destination}") + return + + try: + replay_output = run_again() + benchmark = _benchmark(run_again, input_tensor.device) + tensors = { + "input.latent": input_tensor.detach().to("cpu").contiguous(), + "input.timestep": timestep.detach().to("cpu").contiguous(), + "output.reference": reference_output.detach().to("cpu").contiguous(), + "output.replay": replay_output.detach().to("cpu").contiguous(), + } + condition_descriptor = _flatten(condition, (), tensors) + validation = _validation(reference_output, replay_output) + metadata = { + "format_version": str(FORMAT_VERSION), + "captured_unix_time": str(time.time()), + "condition_descriptor": json.dumps(condition_descriptor, separators=(",", ":")), + "model": json.dumps(_model_metadata(), separators=(",", ":")), + "validation": json.dumps(validation, separators=(",", ":")), + } + if benchmark is not None: + metadata["pytorch_benchmark"] = json.dumps(benchmark, separators=(",", ":")) + + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(destination.name + ".tmp") + save_file(tensors, str(temporary), metadata=metadata) + os.replace(temporary, destination) + print( + f"UNet capture saved: {destination} " + f"(exact_replay={validation['exact']}, max_abs={validation['maximum_absolute_error']:.8g})" + ) + if benchmark is not None: + print( + f"PyTorch UNet probe: median={benchmark['median_ms']:.3f}ms " + f"min={benchmark['minimum_ms']:.3f}ms max={benchmark['maximum_ms']:.3f}ms " + f"runs={benchmark['runs']}" + ) + except Exception as error: + print(f"UNet capture failed without affecting generation: {type(error).__name__}: {error}") + + +def reset_for_tests(): + global _capture_claimed + with _capture_lock: + _capture_claimed = False diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index 8be49303b06..ef6469b1697 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -3,7 +3,7 @@ from einops import repeat import math -from modules import devices, mps_fused_ops +from modules import devices, mps_fused_ops, mps_unet_capture from modules.sd_hijack_utils import CondFunc @@ -129,7 +129,16 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y] with devices.autocast(): - result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs) + unet_input = x_noisy.to(devices.dtype_unet) + unet_timestep = t.to(devices.dtype_unet) + result = orig_func(self, unet_input, unet_timestep, cond, **kwargs) + mps_unet_capture.capture_and_validate( + lambda: orig_func(self, unet_input, unet_timestep, cond, **kwargs), + unet_input, + unet_timestep, + cond, + result, + ) if devices.unet_needs_upcast: return result.float() else: diff --git a/scripts/inspect_mps_unet_capture.py b/scripts/inspect_mps_unet_capture.py new file mode 100644 index 00000000000..84d01797515 --- /dev/null +++ b/scripts/inspect_mps_unet_capture.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""Inspect and validate an A1111_MPS_CAPTURE_UNET fixture.""" + +import argparse +import json + +from safetensors import safe_open +from safetensors.torch import load_file +import torch + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("capture", help="Path to the captured .safetensors file") + args = parser.parse_args() + + tensors = load_file(args.capture) + with safe_open(args.capture, framework="pt", device="cpu") as source: + metadata = source.metadata() + + reference = tensors["output.reference"] + replay = tensors["output.replay"] + difference = (reference.float() - replay.float()).abs() + result = { + "model": json.loads(metadata["model"]), + "latent": {"shape": list(tensors["input.latent"].shape), "dtype": str(tensors["input.latent"].dtype)}, + "timestep": {"shape": list(tensors["input.timestep"].shape), "values": tensors["input.timestep"].tolist()}, + "condition_tensors": { + key: {"shape": list(value.shape), "dtype": str(value.dtype)} + for key, value in tensors.items() + if key.startswith("condition.") + }, + "output": {"shape": list(reference.shape), "dtype": str(reference.dtype)}, + "validation": { + "exact": bool(torch.equal(reference, replay)), + "mean_absolute_error": float(difference.mean().item()), + "maximum_absolute_error": float(difference.max().item()), + }, + } + if "pytorch_benchmark" in metadata: + result["pytorch_benchmark"] = json.loads(metadata["pytorch_benchmark"]) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/state-of-things-next.md b/state-of-things-next.md index dc5fc646922..97a467d0729 100644 --- a/state-of-things-next.md +++ b/state-of-things-next.md @@ -135,6 +135,32 @@ Measure complete nine-call warm latency, per-shape latency, retained/peak memory Gate: the native nine-call workload must be at least 20–25% faster than current PyTorch MPS. Stop the project here if it does not clear the gate; smaller gains will likely disappear behind bridge synchronization and compatibility work. +### M1 single-call probe result: stopped at the gate + +On 2026-08-11, the first bounded probe captured the real first batch-two call from the 512×512 reference request. The fixture contains an FP16 `2×4×64×64` latent, timestep `[999, 999]`, FP16 cross-attention context `2×77×768`, and FP16 reference output. Replaying the captured inputs immediately through the existing PyTorch UNet produced a bit-for-bit identical output with zero mean and maximum absolute error. + +Ten synchronized warm PyTorch MPS replays measured: + +| Runner | Median | Minimum | Maximum | +| --- | ---: | ---: | ---: | +| Current PyTorch MPS UNet | 874.597 ms | 868.626 ms | 880.409 ms | + +A fresh upstream stable-diffusion.cpp checkout at `bcc7e29` was built with its Metal backend and a temporary direct `UNetModelRunner` probe. With native Flash Attention enabled, three final measured calls produced: + +| Runner | Median | Minimum | Maximum | +| --- | ---: | ---: | ---: | +| stable-diffusion.cpp Metal UNet | 2,192.908 ms | 2,187.225 ms | 2,203.125 ms | + +The native call was approximately 2.51× slower than the current PyTorch MPS path before any Automatic1111 bridge or buffer-transfer overhead. It also returned 16,384 non-finite values out of 32,768 outputs—exactly one batch element—while the PyTorch reference contained none. Among finite values, mean absolute error was `0.0002314` and maximum absolute error was `0.0017264`. + +Additional findings: + +- Enabling mmap for the native Metal weights crashed in `ggml_metal_buffer_get_id`; disabling mmap allowed the probe to complete. +- Disabling native Flash Attention increased median latency to approximately 10.30 seconds per call and did not eliminate the non-finite output. +- The capture and immediate PyTorch replay were exact, so the input corpus itself passed its accuracy check. + +Decision: do not proceed to a copied-buffer WebUI integration with this native runner. It misses the required speed gate by a wide margin and currently fails numerical validity. Retain the opt-in capture tooling as a small reusable test for a materially different future engine, but treat the stable-diffusion.cpp sidecar described below as rejected on the tested M1 implementation. + ## Phase 2: copied-buffer WebUI prototype Expose a minimal native interface for latent, timestep, conditioning, and UNet output. Keep A1111's sampler in control. A first implementation may synchronize and copy once per UNet call to prove integration. @@ -209,4 +235,4 @@ Quantization must use native fused dequantization/matrix kernels. Do not add qua ## Immediate next task -Implement Phase 0 only: capture and replay the exact nine-call PyTorch UNet contract. Then create the standalone native replay harness. Do not begin WebUI integration until the raw native shootout produces a clear result. +Do not begin native WebUI integration. The single-call native shootout already failed the speed and validity gates. If a materially different engine becomes available, reuse the opt-in capture tooling and require it to beat the synchronized 874.597 ms batch-two PyTorch reference by at least 20–25% before expanding to batch one or the complete nine-call sequence. diff --git a/test/test_mps_unet_capture.py b/test/test_mps_unet_capture.py new file mode 100644 index 00000000000..8d6ae557da3 --- /dev/null +++ b/test/test_mps_unet_capture.py @@ -0,0 +1,84 @@ +import json +import os +from unittest import mock + +from safetensors import safe_open +from safetensors.torch import load_file +import torch + +from modules import mps_unet_capture + + +def test_disabled_capture_does_not_replay(tmp_path): + destination = tmp_path / "capture.safetensors" + environment = {key: value for key, value in os.environ.items() if key != mps_unet_capture.ENVIRONMENT_VARIABLE} + replay = mock.Mock(return_value=torch.ones((2, 4, 8, 8))) + + with mock.patch.dict(os.environ, environment, clear=True): + mps_unet_capture.reset_for_tests() + mps_unet_capture.capture_and_validate( + replay, + torch.zeros((2, 4, 8, 8)), + torch.ones(2), + {"c_crossattn": [torch.zeros((2, 77, 768))]}, + torch.ones((2, 4, 8, 8)), + ) + + replay.assert_not_called() + assert not destination.exists() + + +def test_capture_saves_complete_exact_replay_fixture(tmp_path): + destination = tmp_path / "capture.safetensors" + latent = torch.arange(2 * 4 * 8 * 8, dtype=torch.float16).reshape(2, 4, 8, 8) + timestep = torch.tensor([1.5, 1.5], dtype=torch.float16) + condition = {"c_crossattn": [torch.zeros((2, 77, 768), dtype=torch.float16)], "c_concat": []} + reference = latent + 1 + + with mock.patch.dict(os.environ, {mps_unet_capture.ENVIRONMENT_VARIABLE: str(destination)}, clear=False): + mps_unet_capture.reset_for_tests() + mps_unet_capture.capture_and_validate(lambda: latent + 1, latent, timestep, condition, reference) + + tensors = load_file(destination) + with safe_open(destination, framework="pt", device="cpu") as source: + metadata = source.metadata() + + assert torch.equal(tensors["input.latent"], latent) + assert torch.equal(tensors["input.timestep"], timestep) + assert tensors["condition.c_crossattn.0"].shape == (2, 77, 768) + assert torch.equal(tensors["output.reference"], tensors["output.replay"]) + assert json.loads(metadata["validation"]) == { + "exact": True, + "mean_absolute_error": 0.0, + "maximum_absolute_error": 0.0, + } + assert "pytorch_benchmark" not in metadata + + +def test_capture_waits_for_requested_batch(tmp_path): + destination = tmp_path / "capture.safetensors" + environment = { + mps_unet_capture.ENVIRONMENT_VARIABLE: str(destination), + mps_unet_capture.BATCH_ENVIRONMENT_VARIABLE: "2", + } + + with mock.patch.dict(os.environ, environment, clear=False): + mps_unet_capture.reset_for_tests() + mps_unet_capture.capture_and_validate( + lambda: torch.ones((1, 4, 8, 8)), + torch.zeros((1, 4, 8, 8)), + torch.ones(1), + {}, + torch.ones((1, 4, 8, 8)), + ) + assert not destination.exists() + + mps_unet_capture.capture_and_validate( + lambda: torch.ones((2, 4, 8, 8)), + torch.zeros((2, 4, 8, 8)), + torch.ones(2), + {}, + torch.ones((2, 4, 8, 8)), + ) + + assert destination.exists() From 89ee3349c43736dbbe22672a8b73680dc2db6c9e Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 10:33:05 -0500 Subject: [PATCH 11/17] Fix MPS benchmark lint --- scripts/benchmark_mps_unet_ops.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/scripts/benchmark_mps_unet_ops.py b/scripts/benchmark_mps_unet_ops.py index 4221a66c3e5..57d1bc5e790 100644 --- a/scripts/benchmark_mps_unet_ops.py +++ b/scripts/benchmark_mps_unet_ops.py @@ -40,7 +40,7 @@ def measure(operation, warmup, repeats): return statistics.median(timings) -def benchmark_shape(batch, tokens, channels, warmup, repeats): +def measure_shape(batch, tokens, channels, warmup, repeats): side = int(tokens**0.5) image = torch.randn((batch, channels, side, side), device="mps", dtype=torch.float16) convolution_weight = torch.randn((channels, channels, 3, 3), device="mps", dtype=torch.float16) @@ -51,7 +51,7 @@ def benchmark_shape(batch, tokens, channels, warmup, repeats): heads = 8 query = sequence.view(batch, tokens, heads, channels // heads).transpose(1, 2) - results = { + return { "conv3x3": measure( lambda: F.conv2d(image, convolution_weight, convolution_bias, padding=1), warmup, @@ -74,7 +74,9 @@ def benchmark_shape(batch, tokens, channels, warmup, repeats): ), } - del image, convolution_weight, convolution_bias, sequence, projection_weight, query + +def benchmark_shape(batch, tokens, channels, warmup, repeats): + results = measure_shape(batch, tokens, channels, warmup, repeats) torch.mps.empty_cache() return results From 7e80fdf2d4ce14199170a16018459139ba1d2e95 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 11:06:56 -0500 Subject: [PATCH 12/17] Remove native UNet experiment --- modules/mps_unet_capture.py | 186 ---------------------- modules/sd_hijack_unet.py | 13 +- scripts/inspect_mps_unet_capture.py | 46 ------ state-of-things-next.md | 238 ---------------------------- test/test_mps_unet_capture.py | 84 ---------- 5 files changed, 2 insertions(+), 565 deletions(-) delete mode 100644 modules/mps_unet_capture.py delete mode 100644 scripts/inspect_mps_unet_capture.py delete mode 100644 state-of-things-next.md delete mode 100644 test/test_mps_unet_capture.py diff --git a/modules/mps_unet_capture.py b/modules/mps_unet_capture.py deleted file mode 100644 index 1b3d366a437..00000000000 --- a/modules/mps_unet_capture.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Opt-in capture of one real SD1 UNet call for native-backend experiments.""" - -import json -import os -from pathlib import Path -import statistics -import threading -import time - -import torch -from safetensors.torch import save_file - - -ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET" -BATCH_ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET_BATCH" -BENCHMARK_RUNS_ENVIRONMENT_VARIABLE = "A1111_MPS_CAPTURE_UNET_BENCHMARK_RUNS" -FORMAT_VERSION = 1 - -_capture_lock = threading.Lock() -_capture_claimed = False - - -def enabled(): - return bool(os.environ.get(ENVIRONMENT_VARIABLE)) - - -def _claim(input_tensor): - global _capture_claimed - - if not enabled(): - return False - - try: - requested_batch = int(os.environ.get(BATCH_ENVIRONMENT_VARIABLE, "2")) - except ValueError: - requested_batch = 2 - - if input_tensor.ndim == 0 or input_tensor.shape[0] != requested_batch: - return False - - with _capture_lock: - if _capture_claimed: - return False - _capture_claimed = True - return True - - -def _tensor_key(path): - return "condition." + ".".join(str(part).replace("%", "%25").replace(".", "%2E") for part in path) - - -def _flatten(value, path, tensors): - if isinstance(value, torch.Tensor): - key = _tensor_key(path) - tensors[key] = value.detach().to("cpu").contiguous() - return {"type": "tensor", "key": key} - if isinstance(value, dict): - return { - "type": "dict", - "items": [[str(key), _flatten(item, (*path, key), tensors)] for key, item in value.items()], - } - if isinstance(value, (list, tuple)): - return { - "type": "tuple" if isinstance(value, tuple) else "list", - "items": [_flatten(item, (*path, index), tensors) for index, item in enumerate(value)], - } - if value is None or isinstance(value, (bool, int, float, str)): - return {"type": "literal", "value": value} - return {"type": "unsupported", "python_type": type(value).__qualname__, "repr": repr(value)} - - -def _model_metadata(): - try: - from modules import shared - - model = shared.sd_model - checkpoint = getattr(model, "sd_checkpoint_info", None) - return { - "checkpoint_filename": getattr(checkpoint, "filename", None), - "checkpoint_sha256": getattr(checkpoint, "sha256", None), - "checkpoint_shorthash": getattr(checkpoint, "shorthash", None), - "model_class": type(model).__qualname__ if model is not None else None, - } - except Exception as error: - return {"metadata_error": f"{type(error).__name__}: {error}"} - - -def _validation(reference, replay): - reference_float = reference.detach().float().to("cpu") - replay_float = replay.detach().float().to("cpu") - difference = (reference_float - replay_float).abs() - return { - "exact": bool(torch.equal(reference.detach().to("cpu"), replay.detach().to("cpu"))), - "mean_absolute_error": float(difference.mean().item()), - "maximum_absolute_error": float(difference.max().item()), - } - - -def _synchronize(device): - if device.type == "mps" and torch.backends.mps.is_available(): - torch.mps.synchronize() - - -def _benchmark(run_again, device): - try: - measured_runs = int(os.environ.get(BENCHMARK_RUNS_ENVIRONMENT_VARIABLE, "0")) - except ValueError: - measured_runs = 0 - if measured_runs < 1: - return None - - run_again() - _synchronize(device) - milliseconds = [] - for _ in range(measured_runs): - _synchronize(device) - started = time.perf_counter() - run_again() - _synchronize(device) - milliseconds.append((time.perf_counter() - started) * 1000) - - return { - "runs": measured_runs, - "median_ms": statistics.median(milliseconds), - "minimum_ms": min(milliseconds), - "maximum_ms": max(milliseconds), - "all_ms": milliseconds, - } - - -def capture_and_validate(run_again, input_tensor, timestep, condition, reference_output): - """Capture the first requested batch and immediately prove it replays in PyTorch.""" - if not _claim(input_tensor): - return - - destination = Path(os.environ[ENVIRONMENT_VARIABLE]).expanduser() - if destination.suffix != ".safetensors": - print(f"UNet capture skipped: {ENVIRONMENT_VARIABLE} must name a .safetensors file") - return - if destination.exists(): - print(f"UNet capture skipped: destination already exists: {destination}") - return - - try: - replay_output = run_again() - benchmark = _benchmark(run_again, input_tensor.device) - tensors = { - "input.latent": input_tensor.detach().to("cpu").contiguous(), - "input.timestep": timestep.detach().to("cpu").contiguous(), - "output.reference": reference_output.detach().to("cpu").contiguous(), - "output.replay": replay_output.detach().to("cpu").contiguous(), - } - condition_descriptor = _flatten(condition, (), tensors) - validation = _validation(reference_output, replay_output) - metadata = { - "format_version": str(FORMAT_VERSION), - "captured_unix_time": str(time.time()), - "condition_descriptor": json.dumps(condition_descriptor, separators=(",", ":")), - "model": json.dumps(_model_metadata(), separators=(",", ":")), - "validation": json.dumps(validation, separators=(",", ":")), - } - if benchmark is not None: - metadata["pytorch_benchmark"] = json.dumps(benchmark, separators=(",", ":")) - - destination.parent.mkdir(parents=True, exist_ok=True) - temporary = destination.with_name(destination.name + ".tmp") - save_file(tensors, str(temporary), metadata=metadata) - os.replace(temporary, destination) - print( - f"UNet capture saved: {destination} " - f"(exact_replay={validation['exact']}, max_abs={validation['maximum_absolute_error']:.8g})" - ) - if benchmark is not None: - print( - f"PyTorch UNet probe: median={benchmark['median_ms']:.3f}ms " - f"min={benchmark['minimum_ms']:.3f}ms max={benchmark['maximum_ms']:.3f}ms " - f"runs={benchmark['runs']}" - ) - except Exception as error: - print(f"UNet capture failed without affecting generation: {type(error).__name__}: {error}") - - -def reset_for_tests(): - global _capture_claimed - with _capture_lock: - _capture_claimed = False diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index ef6469b1697..cef1e4c361f 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -3,7 +3,7 @@ from einops import repeat import math -from modules import devices, mps_fused_ops, mps_unet_capture +from modules import devices, mps_fused_ops from modules.sd_hijack_utils import CondFunc @@ -129,16 +129,7 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y] with devices.autocast(): - unet_input = x_noisy.to(devices.dtype_unet) - unet_timestep = t.to(devices.dtype_unet) - result = orig_func(self, unet_input, unet_timestep, cond, **kwargs) - mps_unet_capture.capture_and_validate( - lambda: orig_func(self, unet_input, unet_timestep, cond, **kwargs), - unet_input, - unet_timestep, - cond, - result, - ) + result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs) if devices.unet_needs_upcast: return result.float() else: diff --git a/scripts/inspect_mps_unet_capture.py b/scripts/inspect_mps_unet_capture.py deleted file mode 100644 index 84d01797515..00000000000 --- a/scripts/inspect_mps_unet_capture.py +++ /dev/null @@ -1,46 +0,0 @@ -#!/usr/bin/env python3 -"""Inspect and validate an A1111_MPS_CAPTURE_UNET fixture.""" - -import argparse -import json - -from safetensors import safe_open -from safetensors.torch import load_file -import torch - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("capture", help="Path to the captured .safetensors file") - args = parser.parse_args() - - tensors = load_file(args.capture) - with safe_open(args.capture, framework="pt", device="cpu") as source: - metadata = source.metadata() - - reference = tensors["output.reference"] - replay = tensors["output.replay"] - difference = (reference.float() - replay.float()).abs() - result = { - "model": json.loads(metadata["model"]), - "latent": {"shape": list(tensors["input.latent"].shape), "dtype": str(tensors["input.latent"].dtype)}, - "timestep": {"shape": list(tensors["input.timestep"].shape), "values": tensors["input.timestep"].tolist()}, - "condition_tensors": { - key: {"shape": list(value.shape), "dtype": str(value.dtype)} - for key, value in tensors.items() - if key.startswith("condition.") - }, - "output": {"shape": list(reference.shape), "dtype": str(reference.dtype)}, - "validation": { - "exact": bool(torch.equal(reference, replay)), - "mean_absolute_error": float(difference.mean().item()), - "maximum_absolute_error": float(difference.max().item()), - }, - } - if "pytorch_benchmark" in metadata: - result["pytorch_benchmark"] = json.loads(metadata["pytorch_benchmark"]) - print(json.dumps(result, indent=2)) - - -if __name__ == "__main__": - main() diff --git a/state-of-things-next.md b/state-of-things-next.md deleted file mode 100644 index 97a467d0729..00000000000 --- a/state-of-things-next.md +++ /dev/null @@ -1,238 +0,0 @@ -# State of Things and Native UNet Roadmap - -Last updated: 2026-08-11 - -Target machine: 16 GB Apple M1 Mac mini - -Branch: `dev` - -Current committed head before this documentation update: `6eefbb402d177ec5166dbb364ea8e313d1bdb206` - -Automatic1111 base: `1937682a20f7f0442311a1ede68f9f0cb480163b` - -## Current state - -This remains an Automatic1111 fork with targeted MPS and native Metal acceleration. It does not currently contain a separate diffusion engine. - -The active optimization stack is: - -1. Selective Draw Things-style Metal Flash Attention for measured SD 1.x shapes, encoded on PyTorch's current MPS command buffer. -2. Unified-memory-aware routing to native or sub-quadratic attention. -3. Native fused GroupNorm + SiLU for compatible FP16 UNet and VAE blocks. -4. Exact-parity fused GEGLU using a 65,536-entry, 128 KB PyTorch-generated FP16 GELU table. -5. Modern-PyTorch removal of obsolete MPS clones and FP32 LayerNorm workarounds. -6. FP16 VAE as the tracked default on the tested M1 family, with Automatic1111's FP32 NaN retry retained. -7. NGMS 1.0/all steps and Clip skip 2 as built-in defaults. -8. An opt-in coarse profiler enabled by `A1111_MPS_PROFILE=1` with no synchronization in the disabled path. - -The native extension performs an isolated MPS startup test. Unsupported inputs and runtime failures retain PyTorch fallbacks. - -## Reference workloads - -Primary profiling workload: - -- Prompt: `a dog` -- Negative prompt: empty -- Checkpoint hash: `8ecad70a19` -- Steps: 5 -- Sampler: DPM++ SDE -- Schedule: Karras -- CFG: 1.15 -- Seed: `3163229250` -- Size: 384×640 -- Clip skip: 2 -- NGMS: 1.0, all steps -- Batch: 1 - -The sampler makes nine UNet evaluations: five calls at batch two and four calls at batch one. At 384×640 the latent inputs are `2×4×80×48` and `1×4×80×48`. NGMS creates the batch-one regime. - -The recurring 512×512 validation workload uses prompt `a dog`, seed `4017012032`, and the same model, sampler, schedule, CFG, Clip skip, and NGMS settings. - -## Measured results - -### Fork versus Automatic1111 baseline - -At 384×640, the original Automatic1111 base recorded 12.8 seconds and an earlier fork head recorded 8.7 seconds at matching model hash and tensor shape: approximately 32% lower latency or 1.47× throughput. The paired runs used different seeds, so this is a throughput comparison rather than output parity. - -### FP16 VAE - -Five warm profiled 384×640 runs produced: - -| VAE precision | End-to-end median | Sampler | VAE decode + transfer | -| --- | ---: | ---: | ---: | -| FP32 | 8.450 s | 6.715 s | 1.536 s | -| FP16 | 7.795 s | 6.666 s | 0.972 s | - -FP16 saved about 0.65 seconds end to end and reduced the VAE stage by about 37%. Across three fixed-seed cases, PSNR versus FP32 was 64.0–64.6 dB, every changed RGB channel moved by at most one 8-bit value, and 97.4–97.7% of channels were identical. - -### Exact GEGLU - -A 512×512 alternating A/B produced identical PNG hashes and positive paired savings of approximately 0.12–0.27 seconds. An active LCM LoRA output hash was also identical with fusion enabled and disabled. - -### Current range - -A normal user run at 512×512 recorded 8.3 seconds. Later controlled warm A/B sessions commonly clustered around 9.2–9.3 seconds. Background work, extensions, thermal state, and unified-memory pressure are material at this scale. - -## Rejected experiments - -Do not repeat these without a new mechanism or new evidence: - -| Experiment | Evidence | Result | -| --- | --- | --- | -| DPM++ 2M substitution | Changed the desired LCM/DPM++ SDE image behavior | Reject sampler substitution | -| FP8 on M1 | No M1 FP8 hardware execution path | Reject conversion/unpacking overhead | -| Per-operator MPS events on PyTorch 2.3 | Isolated synchronization hung | Use coarse profiling | -| Block-level MPSGraph | About 1% slower end to end with more numerical drift | Removed | -| Real-weight MPSGraph block | About 1.4% isolated stage gain versus fused GroupNorm | Failed integration gate | -| Fixed-shape TorchScript UNet | Inconsistent warm gain, lost after cache loss, increased retained memory | Removed | -| Fused LayerNorm | Projected 110 ms microbenchmark gain; paired end-to-end median regressed 0.026 s and only 56.1% of RGB channels matched | Removed | -| Cross-attention K/V reuse | Reused 112/144 projections and retained 11 MiB; paired end-to-end median regressed 0.016 s | Removed | - -The LayerNorm and K/V code and saved probe settings were removed completely. The current repository and native extension contain neither path. - -## What the profile says - -After the FP16 VAE improvement, sampling/UNet consumes roughly 87% of measured generation time. Conditioning, image conversion, metadata, PNG creation, and additional VAE micro-tuning cannot provide the next material gain. - -The failed LayerNorm and K/V experiments also show that transformer micro-operations are now below the useful granularity. The next work must reduce framework overhead across a large portion of the UNet or execute the complete UNet more efficiently. - -## Chosen direction: native ggml/Metal UNet sidecar - -Take architectural inspiration from stable-diffusion.cpp and ggml without replacing Automatic1111. - -Keep in Automatic1111: - -- Prompt parsing and conditioning. -- Existing DPM++ SDE/Karras sampler. -- CFG and NGMS decisions. -- Seed and RNG behavior. -- LoRA/extension activation and request routing. -- VAE, image pipeline, metadata, API, and UI. - -Delegate only a supported UNet evaluation to a native graph runner. Unsupported requests continue through the current PyTorch MPS UNet. - -stable-diffusion.cpp already demonstrates the relevant components: a complete SD 1.x UNet graph, graph-planned reusable buffers, whole-graph Metal encoding, safetensors/GGUF loading, LoRA support, Flash Attention, and fused quantized matrix kernels. The useful lesson is ownership of the complete graph and memory lifecycle, not copying individual kernels. - -## Phase 0: capture the existing UNet contract - -Create a diagnostic-only capture of all nine real calls for 384×640 and 512×512: - -- Latent input and reference output. -- Timestep. -- Text conditioning. -- Shape, dtype, model hash, and request settings. -- Batch-two and batch-one regimes. - -Include plain prompt, scheduled-prompt, active-LoRA, and unsupported/fallback fixtures. Disabled capture must add no synchronization or normal-path overhead. - -Deliverable: a reproducible tensor corpus and a PyTorch replay test. - -## Phase 1: standalone native shootout - -Build a small harness around stable-diffusion.cpp's `UNetModelRunner`, load the same SD 1.x checkpoint, and replay the captures outside WebUI. - -Measure complete nine-call warm latency, per-shape latency, retained/peak memory, determinism, and tensor deviation. - -Gate: the native nine-call workload must be at least 20–25% faster than current PyTorch MPS. Stop the project here if it does not clear the gate; smaller gains will likely disappear behind bridge synchronization and compatibility work. - -### M1 single-call probe result: stopped at the gate - -On 2026-08-11, the first bounded probe captured the real first batch-two call from the 512×512 reference request. The fixture contains an FP16 `2×4×64×64` latent, timestep `[999, 999]`, FP16 cross-attention context `2×77×768`, and FP16 reference output. Replaying the captured inputs immediately through the existing PyTorch UNet produced a bit-for-bit identical output with zero mean and maximum absolute error. - -Ten synchronized warm PyTorch MPS replays measured: - -| Runner | Median | Minimum | Maximum | -| --- | ---: | ---: | ---: | -| Current PyTorch MPS UNet | 874.597 ms | 868.626 ms | 880.409 ms | - -A fresh upstream stable-diffusion.cpp checkout at `bcc7e29` was built with its Metal backend and a temporary direct `UNetModelRunner` probe. With native Flash Attention enabled, three final measured calls produced: - -| Runner | Median | Minimum | Maximum | -| --- | ---: | ---: | ---: | -| stable-diffusion.cpp Metal UNet | 2,192.908 ms | 2,187.225 ms | 2,203.125 ms | - -The native call was approximately 2.51× slower than the current PyTorch MPS path before any Automatic1111 bridge or buffer-transfer overhead. It also returned 16,384 non-finite values out of 32,768 outputs—exactly one batch element—while the PyTorch reference contained none. Among finite values, mean absolute error was `0.0002314` and maximum absolute error was `0.0017264`. - -Additional findings: - -- Enabling mmap for the native Metal weights crashed in `ggml_metal_buffer_get_id`; disabling mmap allowed the probe to complete. -- Disabling native Flash Attention increased median latency to approximately 10.30 seconds per call and did not eliminate the non-finite output. -- The capture and immediate PyTorch replay were exact, so the input corpus itself passed its accuracy check. - -Decision: do not proceed to a copied-buffer WebUI integration with this native runner. It misses the required speed gate by a wide margin and currently fails numerical validity. Retain the opt-in capture tooling as a small reusable test for a materially different future engine, but treat the stable-diffusion.cpp sidecar described below as rejected on the tested M1 implementation. - -## Phase 2: copied-buffer WebUI prototype - -Expose a minimal native interface for latent, timestep, conditioning, and UNet output. Keep A1111's sampler in control. A first implementation may synchronize and copy once per UNet call to prove integration. - -Initial supported route: - -- Apple M1. -- SD 1.x FP16. -- Txt2img, batch one. -- Tested 384×640 and 512×512 shapes. -- No active LoRA, ControlNet, hypernetwork, training, or high-resolution pass. - -Implement it as an optional `SdUnetOption`. Every unsupported condition routes to PyTorch. - -Gate: multiple alternating end-to-end pairs must remain materially faster, deterministic, and within an explicitly approved output-deviation envelope. - -## Phase 3: zero-copy unified-memory proof - -If the copied bridge wins, share the underlying Metal storage between PyTorch MPS and ggml. - -Solve and test: - -- `MTLBuffer` ownership and lifetime. -- Buffer offsets, strides, NCHW layout, and dtype agreement. -- PyTorch and ggml command-queue ordering. -- Error recovery and backend reset. - -Start with one captured UNet call. Do not attempt full sampling until shared-buffer output matches the copied native implementation. - -## Phase 4: compatibility expansion - -Add independently gated support in this order: - -1. Dynamic SD 1.x resolutions and reusable arenas per shape/batch regime. -2. LoRA weight application plus explicit model-mutation generation counters. -3. Img2img and inpainting. -4. High-resolution pass and model switching. -5. ControlNet where native semantics can match A1111. -6. Additional model families. - -Never silently ignore an installed extension hook. Fall back to PyTorch for any request whose semantics the native backend cannot reproduce. - -## Phase 5: optional quantization - -Only after FP16 proves the native engine: - -1. Q8_0 for the lowest-risk memory experiment. -2. Q6_K/Q5_K as optional balanced modes. -3. Q4 as an explicit low-memory mode, not a default. - -Quantization must use native fused dequantization/matrix kernels. Do not add quantized PyTorch storage with per-call unpacking. Require tensor, fixed-seed image, LoRA, memory, and timing validation for every format. - -## Global gates - -1. Preserve A1111's exact DPM++ SDE evaluation sequence and both batch regimes. -2. Benchmark alternating warm pairs, never a single best run. -3. Report tensor and final-image deviation. -4. Measure retained unified memory and cache-loss behavior. -5. Preserve deterministic output within each path. -6. Keep automatic PyTorch fallback for unsupported features and runtime errors. -7. Never require checkpoint conversion for the normal PyTorch path. -8. Keep the native backend independently disableable. - -## Explicit non-goals for the first sprint - -- Replacing the A1111 UI, API, sampler, VAE, or extension ecosystem. -- Calling individual ggml convolutions or matrix kernels from PyTorch. -- Adding another Flash Attention implementation. -- Making GGUF mandatory. -- Supporting every model family or extension before the SD 1.x proof. -- Committing a backend before the standalone 20–25% gate passes. - -## Immediate next task - -Do not begin native WebUI integration. The single-call native shootout already failed the speed and validity gates. If a materially different engine becomes available, reuse the opt-in capture tooling and require it to beat the synchronized 874.597 ms batch-two PyTorch reference by at least 20–25% before expanding to batch one or the complete nine-call sequence. diff --git a/test/test_mps_unet_capture.py b/test/test_mps_unet_capture.py deleted file mode 100644 index 8d6ae557da3..00000000000 --- a/test/test_mps_unet_capture.py +++ /dev/null @@ -1,84 +0,0 @@ -import json -import os -from unittest import mock - -from safetensors import safe_open -from safetensors.torch import load_file -import torch - -from modules import mps_unet_capture - - -def test_disabled_capture_does_not_replay(tmp_path): - destination = tmp_path / "capture.safetensors" - environment = {key: value for key, value in os.environ.items() if key != mps_unet_capture.ENVIRONMENT_VARIABLE} - replay = mock.Mock(return_value=torch.ones((2, 4, 8, 8))) - - with mock.patch.dict(os.environ, environment, clear=True): - mps_unet_capture.reset_for_tests() - mps_unet_capture.capture_and_validate( - replay, - torch.zeros((2, 4, 8, 8)), - torch.ones(2), - {"c_crossattn": [torch.zeros((2, 77, 768))]}, - torch.ones((2, 4, 8, 8)), - ) - - replay.assert_not_called() - assert not destination.exists() - - -def test_capture_saves_complete_exact_replay_fixture(tmp_path): - destination = tmp_path / "capture.safetensors" - latent = torch.arange(2 * 4 * 8 * 8, dtype=torch.float16).reshape(2, 4, 8, 8) - timestep = torch.tensor([1.5, 1.5], dtype=torch.float16) - condition = {"c_crossattn": [torch.zeros((2, 77, 768), dtype=torch.float16)], "c_concat": []} - reference = latent + 1 - - with mock.patch.dict(os.environ, {mps_unet_capture.ENVIRONMENT_VARIABLE: str(destination)}, clear=False): - mps_unet_capture.reset_for_tests() - mps_unet_capture.capture_and_validate(lambda: latent + 1, latent, timestep, condition, reference) - - tensors = load_file(destination) - with safe_open(destination, framework="pt", device="cpu") as source: - metadata = source.metadata() - - assert torch.equal(tensors["input.latent"], latent) - assert torch.equal(tensors["input.timestep"], timestep) - assert tensors["condition.c_crossattn.0"].shape == (2, 77, 768) - assert torch.equal(tensors["output.reference"], tensors["output.replay"]) - assert json.loads(metadata["validation"]) == { - "exact": True, - "mean_absolute_error": 0.0, - "maximum_absolute_error": 0.0, - } - assert "pytorch_benchmark" not in metadata - - -def test_capture_waits_for_requested_batch(tmp_path): - destination = tmp_path / "capture.safetensors" - environment = { - mps_unet_capture.ENVIRONMENT_VARIABLE: str(destination), - mps_unet_capture.BATCH_ENVIRONMENT_VARIABLE: "2", - } - - with mock.patch.dict(os.environ, environment, clear=False): - mps_unet_capture.reset_for_tests() - mps_unet_capture.capture_and_validate( - lambda: torch.ones((1, 4, 8, 8)), - torch.zeros((1, 4, 8, 8)), - torch.ones(1), - {}, - torch.ones((1, 4, 8, 8)), - ) - assert not destination.exists() - - mps_unet_capture.capture_and_validate( - lambda: torch.ones((2, 4, 8, 8)), - torch.zeros((2, 4, 8, 8)), - torch.ones(2), - {}, - torch.ones((2, 4, 8, 8)), - ) - - assert destination.exists() From b5aa18127189b7fa7fa9e154c4f0f2593ffd8a87 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 11:09:20 -0500 Subject: [PATCH 13/17] Align README with Automatic1111 dev --- README.md | 776 ++++++++++++++++-------------------------------------- 1 file changed, 224 insertions(+), 552 deletions(-) diff --git a/README.md b/README.md index fa76aabfd5b..fcab35d041f 100644 --- a/README.md +++ b/README.md @@ -1,552 +1,224 @@ -# Stable Diffusion WebUI Metal - -An Apple Silicon performance fork of [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), focused on faster and more memory-aware inference through PyTorch MPS and native Metal kernels. - -The normal Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, and extension structure are preserved. The fork adds a selective Metal attention path, fused GroupNorm + SiLU and exact-parity GEGLU kernels, unified-memory-aware attention fallback, an M1-validated FP16 VAE path, and tested macOS dependency defaults. Stable Diffusion 1.x inference—particularly short DPM++ SDE runs—is the primary optimization target. - -> [!IMPORTANT] -> This is an experimental performance fork, not a new Stable Diffusion engine. It favors measured M1 inference performance and safe fallback behavior over broad hardware tuning. If a native Metal path is unavailable or fails its startup test, the WebUI falls back to the corresponding PyTorch implementation. - -## Current project snapshot - -The current tested head is [`6eefbb40`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/6eefbb402d177ec5166dbb364ea8e313d1bdb206) on `dev`. It remains recognizably Automatic1111: the performance work is concentrated in MPS routing, a small native Metal extension, macOS launch defaults, profiling, benchmarks, and tests. - -| Measure | Value | -| --- | ---: | -| Automatic1111 base | [`1937682a`](https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/1937682a20f7f0442311a1ede68f9f0cb480163b) | -| Base version | `v1.10.1-96-g1937682a` | -| Current Metal head | [`6eefbb40`](https://github.com/dmikey/stable-diffusion-webui-metal/commit/6eefbb402d177ec5166dbb364ea8e313d1bdb206) | -| Current version | `v1.10.1-104-g6eefbb40` | -| Code relationship | 8 fork commits ahead of the selected Automatic1111 base | -| Changed tracked paths | 29, including this README and the roadmap | -| Changed implementation/test paths | 27 | -| Total delta | 2,585 insertions, 271 deletions | - -The eight fork commits are: - -1. Apple Silicon dependencies, attention routing, unified-memory budgeting, and benchmark foundation. -2. Removal of obsolete MPS safety operations on modern PyTorch. -3. Metal Flash Attention command-buffer coalescing. -4. Native fused GroupNorm + SiLU for compatible inference blocks. -5. Documentation and launch configuration cleanup. -6. Coarse MPS stage profiling and the M1 FP16 VAE default. -7. Exact-parity native fused GEGLU. -8. Clip skip 2 as the built-in default. - -Most added code is isolated Metal code, profiling, benchmark utilities, and tests. The fork does **not** require a new checkpoint format, prompt syntax, REST API contract, or UI workflow. - -
-Current implementation surface - -| Area | Added | Modified | -| --- | --- | --- | -| Metal runtime | `modules/mps_flash_attention.py`
`modules/mps_fused_ops.py`
`modules/mps_utils.py` | `modules/mac_specific.py`
`modules/sd_hijack_optimizations.py`
`modules/sd_hijack_unet.py`
`modules/sub_quadratic_attention.py` | -| Profiling | `modules/mps_stage_profile.py` | `modules/processing.py`
`modules/sd_samplers_cfg_denoiser.py` | -| Startup and defaults | `requirements_macos.txt` | `modules/launch_utils.py`
`modules/shared_options.py`
`requirements_versions.txt`
`webui-macos-env.sh`
`webui-user.sh` | -| Native build and benchmarks | `scripts/install_mps_flash_attention.py`
`scripts/mps_fused_group_norm.mm`
`scripts/benchmark_mps_attention.py`
`scripts/benchmark_mps_unet_ops.py`
`scripts/benchmark_mps_geglu_probe.py` | — | -| Tests | `test/test_macos_launch_defaults.py`
`test/test_mps_flash_attention.py`
`test/test_mps_fused_ops.py`
`test/test_mps_stage_profile.py`
`test/test_mps_utils.py`
`test/test_sub_quadratic_attention.py` | — | - -
- -You can reproduce the comparison locally: - -```bash -git rev-list --left-right --count 1937682a...6eefbb40 -git diff --shortstat 1937682a..6eefbb40 -git diff --name-status 1937682a..6eefbb40 -``` - -## What is different? - -### Selective Metal Flash Attention - -On supported Apple Silicon inference shapes, Automatic mode prefers a native Metal Flash Attention implementation derived from the `mps-flash-sdpa` package. - -- Routes measured SD 1.x head dimensions (`40`, `80`, and `160`) with at least 192 query tokens to the native kernel. -- Supports both self-attention and cross-attention on the measured path. -- Encodes work on PyTorch's current MPS command buffer instead of forcing a submission after every attention call. -- Uses PyTorch scaled dot product attention for unsupported shapes, training, masks, grouped-query attention, non-FP16 tensors, or runtime failure. -- Builds from `mps-flash-sdpa==0.1.0` source on first launch, checks the downloaded artifact against PyPI's published SHA-256 metadata, and installs only into the local virtual environment. -- Runs an isolated GPU self-test before enabling the native route, so a native crash cannot take down the main WebUI process during capability detection. - -The attention choices appear under **Settings → Optimizations → Cross attention optimization**: - -- `Automatic`: use Metal Flash Attention when its startup test succeeds. -- `mps-flash`: explicitly select the native Metal route with PyTorch fallback. -- `mps-adaptive`: use native PyTorch attention while it fits a unified-memory budget, then fall back to sub-quadratic attention. -- `sub-quadratic`, `sdp`, and the other upstream implementations remain available for comparison and compatibility. - -### Fused GroupNorm + SiLU - -A native inference-only Metal kernel combines GroupNorm and SiLU in one dispatch for compatible contiguous FP16 tensors. - -- Used in compatible SD/SGM UNet residual blocks. -- Used in compatible VAE residual blocks when the VAE is running in FP16. -- Preserves the normal PyTorch path for CPU, non-FP16 tensors, training/autograd, incompatible layouts, missing affine parameters, or runtime errors. -- Enabled by default through **Settings → Optimizations → Fuse GroupNorm and SiLU on Apple Silicon**. - -This is a focused fusion; convolutions and residual additions still use PyTorch MPS. A larger block-level MPSGraph prototype was tested and deliberately rejected because it was about 1% slower end to end and produced a larger numerical delta without a speed benefit. - -### Exact-parity fused GEGLU - -The SD 1.x transformer feed-forward path normally stores a GELU result and then launches a separate multiply. On Apple Silicon, the fork combines the lookup and multiply into one Metal dispatch. - -- The model's linear projection still runs normally, so active LoRAs and other projection hooks remain compatible. -- A one-time 65,536-entry FP16 table is generated with the installed PyTorch MPS GELU implementation. The table is 128 KB and maps every possible half-precision gate value to PyTorch's exact result. -- The fused output was byte-identical to PyTorch at all SD 1.x transformer shapes for batch one and batch two. -- CPU, FP32, training/autograd, incompatible layouts, disabled settings, and runtime failures use the original PyTorch implementation. -- Enabled by default through **Settings → Optimizations → Fuse GEGLU on Apple Silicon**. - -### Unified-memory-aware attention - -The fork treats system RAM and GPU memory as the same constrained resource instead of relying on a fixed attention threshold. - -- Estimates scaled dot product attention's temporary memory from batch, heads, token counts, and element size. -- Limits native attention to a fraction of total and currently available unified memory. -- Dynamically reduces sub-quadratic query tiles for large self-attention workloads. -- Uses streaming online softmax when K/V attention is chunked, merging one tile at a time instead of stacking all partial outputs in memory. -- Keeps cross-attention on the fast path when its short key sequence remains inexpensive. - -The adaptive path is especially useful for high resolutions and lower-memory Macs. The default Metal Flash Attention path remains the measured choice for normal SD 1.x shapes. - -### Modern MPS runtime cleanup - -Several workarounds needed by early PyTorch MPS releases are now gated by runtime version: - -- Avoids cloning every `narrow()` result on PyTorch versions where the underlying MPS bug is fixed. -- Avoids unconditional FP32 LayerNorm conversion on modern runtimes. -- Keeps an environment switch for diagnosing regressions with legacy behavior. -- Prefers direct Metal matrix multiplication for the SD 1.x projection sizes measured on M1. -- Removes Automatic1111's default `--upcast-sampling` flag on Apple Silicon; it can be restored locally when exact upstream behavior is more important than speed. - -### Coarse MPS stage profiler - -An opt-in profiler measures the parts of a complete generation that are large enough to guide optimization decisions without adding synchronization to normal inference. - -- Enable it with `A1111_MPS_PROFILE=1 ./webui.sh`. -- Reports conditioning, sampler, VAE decode/transfer, image processing, and request wall time. -- Records every UNet call shape and MPS allocation snapshots in a machine-readable `MPS_PROFILE_JSON` line. -- Adds no MPS synchronization points when disabled. - -On the M1 reference workload, the profiler established that the sampler/UNet consumes roughly 87% of generation time after enabling the FP16 VAE. This is why current roadmap work targets whole-UNet execution rather than PNG conversion, conditioning, or more VAE micro-tuning. - -### Apple Silicon dependency profile - -The default Apple Silicon environment is pinned to the combination verified for this fork: - -| Dependency | Version | -| --- | --- | -| Python | 3.10 recommended; 3.10.20 used during development | -| PyTorch | 2.3.1 | -| torchvision | 0.18.1 | -| SciPy | 1.13.1 | -| Native extension | `mps-flash-sdpa` 0.1.0 with local stream-safety and fusion patches | - -SciPy is constrained on Apple Silicon because newer wheels encountered loader problems on the macOS beta used during development. Requirements parsing was also updated to understand platform markers and normal Python package specifiers correctly. - -### Changed defaults - -The following defaults intentionally differ from the upstream `dev` branch: - -| Setting | Upstream | This fork | Effect | -| --- | --- | --- | --- | -| Negative Guidance minimum sigma (NGMS) | `0.0` | `1.0` | May skip unconditional guidance late in sampling | -| NGMS all steps | Off | On | Applies the configured NGMS rule on every eligible step | -| `--upcast-sampling` on macOS | On | Off | Keeps more sampling work in FP16 for speed | -| `--no-half-vae` on M1-family Macs | On | Off | Runs VAE encode/decode in FP16; Automatic1111 still retries in FP32 if VAE decode produces NaNs | -| Clip skip | `1` | `2` | Uses the common SD 1.x checkpoint default without depending on local `config.json` | -| Cross-attention Automatic choice on MPS | Sub-quadratic | Metal Flash Attention | Uses the measured native route when available | -| Fused GroupNorm + SiLU | Not present | On | Reduces compatible normalization/activation dispatches | -| Fused GEGLU | Not present | On | Preserves PyTorch FP16 output while reducing transformer activation dispatches | - -NGMS is the largest user-visible behavioral change. It is recorded in PNG generation metadata when active. Set NGMS to `0` and disable **NGMS all steps** if a workflow expects upstream guidance behavior. - -## Measured performance - -One recorded Apple M1 Mac mini comparison during development used the same checkpoint hash and compute shape: - -| Build | Workload | Time | -| --- | --- | ---: | -| Automatic1111 `v1.10.1-96-g1937682a` | 5 steps, DPM++ SDE, Karras, CFG 1.15, 384×640, SD 1.x checkpoint `8ecad70a19`, Clip skip 2, NGMS 1/all steps | 12.8 s | -| This fork `v1.10.1-99-g38ac556a` | Same sampler, schedule, dimensions, checkpoint hash, Clip skip, and NGMS settings | 8.7 s | - -That observed run was approximately **32% lower latency**, or **1.47× as fast**. Later heads added fused GroupNorm + SiLU, the profiled FP16 VAE default, and exact-parity GEGLU after that recorded comparison. - -The two recorded generations used different seeds. This makes the table a throughput comparison at matching tensor shapes, not an image-parity A/B. - -### M1 fused GEGLU validation - -A fixed-process API A/B used `hyperGlance` (`8ecad70a19`), prompt `a dog and a cat`, seed `158926638`, 5-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 512×512 output. Three alternating warm pairs measured a positive saving in every pair: approximately 0.12–0.27 seconds, with a median paired saving of about 0.27 seconds. All fusion-on and fusion-off PNG files had the same SHA-256 hash. - -An additional active-LoRA check used `a dog `, seed `784504668`, five Euler a steps, and the same CFG, size, Clip skip, and NGMS settings. Fusion-on and fusion-off output hashes were identical. The timing from that single LoRA pair is not reported as a speed result because its first run included LoRA activation overhead. - -### M1 FP16 VAE validation - -A later controlled A/B isolated VAE precision on a 16 GB Apple M1 Mac mini. Both paths used checkpoint `8ecad70a19`, prompt `a dog`, seed `3163229250`, 5-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 384×640 output. Each result below is the median of five warm runs with coarse MPS stage profiling enabled. - -| VAE path | End-to-end client time | Sampler stage | VAE decode + transfer | -| --- | ---: | ---: | ---: | -| FP32 (`--no-half-vae`) | 8.450 s | 6.715 s | 1.536 s | -| FP16 | 7.795 s | 6.666 s | 0.972 s | - -FP16 reduced the measured VAE stage by about **37%** and end-to-end latency by about **7.8%**. The sampler time remained effectively unchanged, which is the expected result when only decode precision changes. - -Output quality was checked across three fixed-seed generations at 384×640 and 512×512. Compared with FP32 VAE output, every changed 8-bit RGB channel differed by at most 1 value, PSNR was 64.0–64.6 dB, and 97.4–97.7% of channels were byte-identical. All FP16 runs were deterministic and free of NaN, green, or corrupted output. The default is therefore enabled only on the tested M1 family; other Apple Silicon generations retain FP32 VAE until separately validated. - -### Current warm-run range - -A user-facing run at `v1.10.1-102-g58e63e9f` used `fast-model` (`8ecad70a19`), prompt `a dog`, seed `4017012032`, five-step DPM++ SDE with Karras, CFG 1.15, Clip skip 2, NGMS 1/all steps, and 512×512 output. It completed in **8.3 seconds** on the 16 GB M1 Mac mini. - -Later controlled development A/B runs of the same 512×512 shape typically clustered around 9.2–9.3 seconds after warm-up. Background load, extension startup activity, thermal state, and unified-memory pressure therefore matter at the sub-second scale. Report medians and the full generation settings rather than treating a single fastest run as a guarantee. - -## Experiments that did not pass the gate - -Failed experiments are documented to prevent attractive microbenchmarks from being repeated without new evidence. - -| Experiment | Isolated result | End-to-end result | Decision | -| --- | --- | --- | --- | -| DPM++ 2M substitution | Fewer or cheaper operations in some paths | Did not reproduce the desired LCM/DPM++ SDE result | Rejected; preserve the requested sampler | -| Block-level MPSGraph | Working block prototype | About 1% slower with a larger numerical delta | Removed | -| Real-weight MPSGraph ResBlock/down stage | About 1.4% stage improvement versus the fused GroupNorm baseline | Too small to survive integration overhead | Removed | -| Fixed-shape TorchScript UNet | Some warm runs improved | Benefit was inconsistent and disappeared after cache loss while retaining extra unified memory | Removed | -| Native fused LayerNorm | Exact SD1 workload projection suggested about 110 ms potential savings | Baseline median 11.512 s versus 11.503 s enabled; paired median regressed by 0.026 s. Only 56.1% of RGB channels were identical, with PSNR 47.53 dB | Removed because there was no speed gain and output drifted | -| Cross-attention K/V reuse | Reused 112 of 144 projections | Baseline median 9.258 s versus 9.262 s cached; paired median regressed by 0.016 s while retaining 11 MiB. PNG hashes were identical | Removed because the projections were already too cheap | -| FP8 on M1 | Reduced theoretical weight storage | No matching M1 FP8 execution path; conversion/unpacking would dominate | Not implemented | - -The LayerNorm and K/V probes were completely removed after testing. They are not hidden options and do not remain in the native extension. The repository returned to a clean state after each rejected sprint. - -Treat these numbers as a development result, not a universal guarantee. Timing varies with: - -- Apple Silicon generation and GPU core count -- Unified-memory capacity and pressure from other applications -- Model architecture and attention dimensions -- Resolution, batch size, sampler, and step count -- First-run shader compilation and warm-up -- VAE, LoRA, ControlNet, extensions, and live preview configuration - -For a meaningful comparison, use the same model hash, prompt, negative prompt, seed, sampler, schedule, steps, CFG, dimensions, Clip skip, VAE, and optimization settings. Run at least two warm-ups, then compare the median of several generations. - -## Installation - -### Requirements - -- An Apple Silicon Mac for the optimized native path -- macOS with Metal Performance Shaders support -- Python 3.10 -- Git -- Xcode Command Line Tools, required to compile the Objective-C++/Metal extension - -Install the command-line tools if needed: - -```bash -xcode-select --install -``` - -### Fresh installation - -```bash -git clone --branch dev https://github.com/dmikey/stable-diffusion-webui-metal.git -cd stable-diffusion-webui-metal -./webui.sh -``` - -The first launch creates the virtual environment, installs the pinned Apple Silicon dependencies, downloads and builds the native extension, runs its isolated self-test, and starts the normal Automatic1111 interface. Native extension compilation can make the first launch noticeably longer than later launches. - -Put checkpoints in: - -```text -models/Stable-diffusion/ -``` - -LoRAs, VAEs, embeddings, extensions, and outputs use the usual Automatic1111 directories. - -### Updating - -The repository's default branch is `dev`: - -```bash -git switch dev -git pull --ff-only origin dev -./webui.sh -``` - -Do not commit generated `config.json`, `ui-config.json`, `params.txt`, models, outputs, the virtual environment, or extension installations. They are local runtime state and are ignored by Git. - -### Local launch options - -`webui-macos-env.sh` contains the fork's tracked defaults. Put personal overrides in `webui-user.sh`, which is loaded afterward. For example, to restore sampling upcast while retaining the rest of the Metal work: - -```bash -export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate" -``` - -M1-family Macs use the validated FP16 VAE path by default. Intel and other Apple Silicon generations retain `--no-half-vae`. Add `--no-half-vae` to a local `COMMANDLINE_ARGS` override at any time to force the conservative FP32 VAE path. Automatic1111's enabled-by-default VAE precision recovery also converts the VAE to FP32 and retries if an FP16 decode produces NaNs. - -## Startup messages and fallback behavior - -A healthy optimized startup prints messages similar to: - -```text -Metal self-test passed; deferred MFA, fused GroupNorm+SiLU, and fused GEGLU routing enabled. -Applying attention optimization: mps-flash... done. -``` - -The first compatible generation also reports the first native attention, GroupNorm, and GEGLU dispatch. These messages are informational and print only once per process. - -If the extension cannot build or fails its isolated self-test, startup continues with native PyTorch MPS operations. If either fused activation kernel fails at runtime, that fusion is disabled for the process and PyTorch handles subsequent operations. - -## Compatibility and output parity - -### What should remain compatible - -- Automatic1111's txt2img, img2img, inpainting, high-resolution pass, API, and metadata workflow -- Existing `.safetensors` and `.ckpt` checkpoints -- Standard LoRA, embedding, VAE, and extension directory layouts -- Existing sampler names and generation parameter syntax -- CPU and non-MPS fallback implementations - -The optimized target is SD 1.x inference. Other architectures supported by this Automatic1111 base may run, but unsupported attention shapes fall back to PyTorch and may receive little or no speed benefit. Test model-specific extensions individually. - -### Why the same seed may differ from upstream - -Pixel-identical output is not guaranteed. Differences can come from: - -- NGMS being enabled by default -- Sampling no longer being upcast by default -- Native Flash Attention and fused GroupNorm changing FP16 reduction/rounding order -- A different selected attention implementation - -Small FP16 numerical differences can grow over multiple denoising evaluations even when both paths are deterministic. - -### Closest upstream behavior - -For an upstream-style comparison: - -1. Set **Negative Guidance minimum sigma** to `0`. -2. Disable **Negative Guidance minimum sigma all steps**. -3. Disable **Fuse GroupNorm and SiLU on Apple Silicon**. -4. Disable **Fuse GEGLU on Apple Silicon**. -5. Select `sub-quadratic` under **Cross attention optimization**. -6. Add `--upcast-sampling` to `COMMANDLINE_ARGS` in `webui-user.sh`. -7. On M1, also add `--no-half-vae`. -8. Restart the WebUI after changing launch arguments. - -For diagnostics only, `A1111_MPS_FORCE_LEGACY_OPS=1` restores version-gated MPS safety copies. `A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU=1` and `A1111_MPS_DISABLE_FUSED_GEGLU=1` disable the corresponding native fusion before startup. - -## Troubleshooting - -### Native extension does not build - -Confirm that Xcode Command Line Tools and the local virtual environment are available: - -```bash -xcode-select -p -./venv/bin/python scripts/install_mps_flash_attention.py -``` - -Then restart with `./webui.sh`. The installer intentionally rebuilds the package from source for the active Python and PyTorch environment. - -### Metal self-test fails - -The WebUI should continue on PyTorch MPS. Keep the final `Metal Flash Attention unavailable:` message when reporting the issue. Also include: - -- Mac model and memory capacity -- macOS version -- `./venv/bin/python -c "import torch; print(torch.__version__)"` -- The selected cross-attention optimization -- Model family, resolution, and batch size - -### Green, black, or corrupted output - -Add `--no-half-vae` to the local launch options and restart first. Also compare with NGMS disabled, sampling upcast restored, `sub-quadratic` attention selected, and the fused GroupNorm option disabled. That separates model/VAE precision issues from the native Metal paths. - -### High-resolution out-of-memory errors - -Select `mps-adaptive - native Metal attention with a memory-safe fallback` or `sub-quadratic` in the optimization settings. Reduce batch size before reducing attention chunk limits manually. - -## Benchmarks and tests - -Two standalone benchmark scripts are included: - -```bash -./venv/bin/python scripts/benchmark_mps_attention.py -./venv/bin/python scripts/benchmark_mps_unet_ops.py --batch 2 -./venv/bin/python scripts/benchmark_mps_geglu_probe.py -``` - -The first compares PyTorch MPS scaled dot product attention with sliced attention. The second measures representative SD 1.x convolution, GroupNorm + SiLU, linear projection, and attention shapes. - -For an end-to-end stage breakdown, launch with the opt-in profiler: - -```bash -A1111_MPS_PROFILE=1 ./webui.sh -``` - -Each generation reports synchronized wall time for conditioning, sampling, VAE decode/transfer, and image processing; it also records UNet call shapes and MPS allocation snapshots in a machine-readable `MPS_PROFILE_JSON` line. Profiling is intentionally coarse because PyTorch 2.3 MPS timing events are unreliable on the tested runtime. When the environment variable is absent, the profiler adds no MPS synchronization points. - -Focused tests cover: - -- Metal Flash Attention routing and PyTorch fallback -- M1-specific FP16 VAE launch defaults with conservative Intel and newer-chip behavior -- Native fused GroupNorm + SiLU correctness -- Native fused GEGLU exact parity, fallback routing, and active-LoRA compatibility -- Opt-in MPS stage profiling and its zero-synchronization disabled path -- Unified-memory attention budgeting and dynamic query tiles -- Streaming online-softmax forward results and gradients - -With `pytest` installed in the virtual environment: - -```bash -./venv/bin/python -m pytest -q \ - test/test_macos_launch_defaults.py \ - test/test_mps_flash_attention.py \ - test/test_mps_fused_ops.py \ - test/test_mps_stage_profile.py \ - test/test_mps_utils.py \ - test/test_sub_quadratic_attention.py -``` - -## Future roadmap: native ggml/Metal UNet - -The remaining material opportunity is engine-level work. The best incremental direction is inspired by [stable-diffusion.cpp](https://github.com/leejet/stable-diffusion.cpp) and ggml: execute the complete UNet as one planned Metal graph with a reusable memory arena instead of crossing the Python/PyTorch boundary for individual kernels. - -This is not a plan to replace Automatic1111 wholesale. Prompt parsing, conditioning, the selected A1111/k-diffusion sampler, CFG and NGMS behavior, seed handling, extensions, VAE, image processing, metadata, API, and UI remain in the existing application. Only a compatible UNet evaluation may be delegated to the native backend. - -```text -Automatic1111 prompt, LoRA, and conditioning setup - | -Existing DPM++ SDE / Karras sampler and NGMS logic - | - Native ggml/Metal UNet evaluation - | -Existing CFG combination, VAE, image pipeline, API, and UI -``` - -stable-diffusion.cpp is relevant because its current implementation already provides a complete [SD 1.x UNet graph runner](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/src/model/diffusion/unet.hpp#L748), a [reusable graph allocator](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/src/core/ggml_extend.hpp#L2212), whole-graph [Metal command encoding](https://github.com/leejet/stable-diffusion.cpp/blob/bcc7e29568b94a25f78e99d34a8fa048d77536b1/ggml/src/ggml-metal/ggml-metal-context.m#L438), safetensors/GGUF loading, LoRA support, Flash Attention, and fused quantized matrix kernels. Its advantage comes from owning the graph, buffers, weights, and submission lifecycle together—not from one operator that can be dropped into PyTorch. - -### Phase 0: reproducible captured-tensor corpus - -Capture the inputs and reference outputs of every UNet evaluation from the existing M1 workload: - -- Five calls with batch two and four calls with batch one under five-step DPM++ SDE plus NGMS. -- Both `512×512` and `384×640` latent shapes. -- Latent input, timestep, text conditioning, model hash, precision, and PyTorch output. -- A plain prompt, scheduled prompt, active LoRA, and a deliberately unsupported request for fallback testing. - -The capture path must be diagnostic-only and must not alter normal timing or output when disabled. - -### Phase 1: standalone native UNet shootout - -Build a small C/C++ harness around stable-diffusion.cpp's `UNetModelRunner`. Load the same SD 1.x safetensors checkpoint and replay the captured calls outside WebUI. - -Measure: - -- Per-call and complete nine-call latency after warm-up. -- Batch-one and batch-two behavior separately. -- Peak and retained unified memory. -- Mean, maximum, and percentile tensor deviation from PyTorch MPS. -- Determinism across repeated runs. - -Proceed only if the native nine-call workload is at least **20–25% faster** than the current PyTorch MPS UNet. A smaller isolated advantage is unlikely to survive framework-bridge synchronization and compatibility handling. - -### Phase 2: copied-buffer A1111 prototype - -Expose a minimal native interface that accepts latent, timestep, and conditioning buffers and returns the UNet prediction. Keep A1111's current sampler in control, initially accepting one synchronization and copy boundary per UNet evaluation. - -The first supported route should be intentionally narrow: - -- Apple M1 and SD 1.x only. -- FP16 inference. -- Txt2img, batch one, tested resolutions. -- No ControlNet, hypernetwork, training, or high-resolution pass. -- No active LoRA until mutation/invalidation is explicitly implemented. - -Every unsupported request must automatically use the existing PyTorch UNet. The native backend should be an optional `SdUnetOption`, never a global monkey patch with no escape path. - -### Phase 3: unified-memory zero-copy proof - -If the copied prototype remains faster, investigate sharing the underlying Metal storage rather than copying through CPU memory. PyTorch MPS tensors and ggml Metal tensors ultimately reside in `MTLBuffer` objects, but safe sharing requires explicit work on: - -- Buffer offsets, strides, dtype, and NCHW layout agreement. -- Ownership and lifetime across Python, PyTorch, and the native runner. -- Command-queue ordering and synchronization. -- Error recovery without leaving either backend in a poisoned state. - -This phase should begin with one captured UNet call. Do not attempt full sampling until the shared-buffer output matches the copied native path. - -### Phase 4: compatibility expansion - -Add features one at a time, with a PyTorch fallback and an output test for each: - -1. Dynamic SD 1.x resolutions and cached arenas per batch/shape regime. -2. Active LoRA application, model-mutation generation counters, and exact cache invalidation. -3. Img2img and inpainting conditioning. -4. High-resolution pass and model switching. -5. ControlNet where native semantics can match the installed A1111 extension. -6. Other model families only after SD 1.x is stable. - -Extension compatibility is a routing problem: requests using unsupported hooks should remain fully functional on PyTorch rather than partially executing through native code. - -### Phase 5: optional GGUF quantization - -Quantization follows a successful FP16 engine; it is not the first step. ggml gains from quantized weights because dequantization is fused into its Metal matrix kernels. Merely storing quantized tensors in PyTorch would not reproduce that behavior. - -Suggested order for the M1: - -1. FP16 native backend establishes the execution-engine benefit and parity baseline. -2. Q8_0 evaluates memory reduction with the smallest expected quality risk. -3. Q6_K or Q5_K may become optional balanced modes. -4. Q4 remains an explicit low-memory choice, not the default. - -Each format requires fixed-seed image comparisons, tensor statistics, LoRA checks, and end-to-end timing. A smaller model file alone is not a speed result. - -### Roadmap acceptance gates - -A native backend is eligible for default use only when it: - -1. Improves multiple alternating warm end-to-end pairs, not just an operator microbenchmark. -2. Preserves all nine DPM++ SDE evaluations and the current sampler's result. -3. Reports deterministic output and quantified deviation from the PyTorch path. -4. Does not retain enough extra unified memory to erase warm-run stability. -5. Falls back cleanly for LoRA, ControlNet, dynamic shapes, training, and extensions it cannot reproduce. -6. Can be disabled without changing checkpoint files or local configuration. - -The initial target is to determine whether native UNet execution can move the 16 GB M1 from the current roughly 8–9 second warm range toward 7–8 seconds. Phase 1 is deliberately a bounded proof: if the raw native UNet cannot clear its 20–25% gate, the integration project stops before modifying WebUI. - -### What not to borrow incrementally - -- Individual ggml convolutions or matrix kernels called from PyTorch. Repeated framework and command-queue boundaries would likely erase their benefit. -- Another standalone Flash Attention implementation. The fork already has a measured native MFA route. -- Required GGUF conversion. Existing Automatic1111 checkpoints remain the default input until an optional native backend proves itself. -- VAE tiling at ordinary 512-pixel resolutions. It reduces peak memory but normally increases latency. -- Sampler substitution. stable-diffusion.cpp supports related DPM++ samplers, but this project must retain the exact A1111 DPM++ SDE behavior already chosen for LCM output. - -## Upstream features and documentation - -This README focuses on the fork. For the complete WebUI feature set, usage documentation, and extension ecosystem, see: - -- [Automatic1111 feature overview](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features) -- [Automatic1111 wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki) -- [Automatic1111 API documentation](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API) -- [Automatic1111 troubleshooting](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Troubleshooting) - -## Contributing - -Keep changes narrow, measurable, and safe to fall back from. - -For performance work: - -1. Record the exact model hash and generation settings. -2. Warm up both paths. -3. Compare multiple alternating runs rather than a single best time. -4. Verify deterministic behavior within each path. -5. Measure output deviation as well as latency and memory. -6. Retain the upstream PyTorch path for unsupported inputs and runtime failure. - -Changes that improve an isolated operator but do not improve an end-to-end generation should not be enabled by default. - -## License and credits - -This fork retains Automatic1111's license and third-party notices. Licenses for bundled and borrowed components are available under **Settings → Licenses** and in `html/licenses.html`. - -Primary credit remains with the [Automatic1111 project](https://github.com/AUTOMATIC1111/stable-diffusion-webui) and its contributors. The native attention work builds on [`mps-flash-sdpa`](https://pypi.org/project/mps-flash-sdpa/) and ideas explored by Draw Things, adapted here for Automatic1111's PyTorch MPS execution path. +# Stable Diffusion web UI +A web interface for Stable Diffusion, implemented using Gradio library. + +![](screenshot.png) + +## Apple Silicon Metal enhancements + +This branch adds targeted Apple Silicon inference optimizations while retaining the Automatic1111 interface, API, checkpoint layout, samplers, LoRA syntax, extensions, and PyTorch fallbacks. Stable Diffusion 1.x inference—especially short DPM++ SDE runs—is the primary measured workload. + +- Shape-selective Metal Flash Attention for measured SD 1.x FP16 inference shapes, encoded on PyTorch's current MPS command buffer. +- Unified-memory-aware attention routing with dynamically sized sub-quadratic chunks and streaming online softmax. +- Native fused FP16 GroupNorm + SiLU and exact-parity GEGLU operations for compatible inference tensors. +- Modern-PyTorch removal of obsolete MPS clones and unconditional FP32 LayerNorm workarounds. +- An M1-validated FP16 VAE default with Automatic1111's FP32 NaN retry retained. +- An opt-in coarse profiler enabled with `A1111_MPS_PROFILE=1 ./webui.sh`. +- Local benchmarks and focused correctness tests for the Metal paths. + +Native paths perform isolated startup checks and fall back to PyTorch for unsupported shapes, dtypes, training, masks, runtime failures, and incompatible configurations. The regular safetensors format remains supported; no checkpoint conversion is required. + +Changed Apple Silicon defaults include NGMS 1.0/all steps, Clip skip 2, FP16 sampling without the upstream sampling-upcast default, and the validated FP16 VAE route on M1-family Macs. These settings can change same-seed output compared with upstream defaults. They can be changed through the existing settings or local launch overrides. + +A recorded 16 GB M1 comparison at the same checkpoint hash and compute shape improved a five-step 384×640 DPM++ SDE/Karras request from 12.8 seconds to 8.7 seconds. The runs used different seeds, so this is a throughput observation rather than an image-parity comparison. + +## Features +[Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features): +- Original txt2img and img2img modes +- One click install and run script (but you still must install python and git) +- Outpainting +- Inpainting +- Color Sketch +- Prompt Matrix +- Stable Diffusion Upscale +- Attention, specify parts of text that the model should pay more attention to + - a man in a `((tuxedo))` - will pay more attention to tuxedo + - a man in a `(tuxedo:1.21)` - alternative syntax + - select text and press `Ctrl+Up` or `Ctrl+Down` (or `Command+Up` or `Command+Down` if you're on a MacOS) to automatically adjust attention to selected text (code contributed by anonymous user) +- Loopback, run img2img processing multiple times +- X/Y/Z plot, a way to draw a 3 dimensional plot of images with different parameters +- Textual Inversion + - have as many embeddings as you want and use any names you like for them + - use multiple embeddings with different numbers of vectors per token + - works with half precision floating point numbers + - train embeddings on 8GB (also reports of 6GB working) +- Extras tab with: + - GFPGAN, neural network that fixes faces + - CodeFormer, face restoration tool as an alternative to GFPGAN + - RealESRGAN, neural network upscaler + - ESRGAN, neural network upscaler with a lot of third party models + - SwinIR and Swin2SR ([see here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/2092)), neural network upscalers + - LDSR, Latent diffusion super resolution upscaling +- Resizing aspect ratio options +- Sampling method selection + - Adjust sampler eta values (noise multiplier) + - More advanced noise setting options +- Interrupt processing at any time +- 4GB video card support (also reports of 2GB working) +- Correct seeds for batches +- Live prompt token length validation +- Generation parameters + - parameters you used to generate images are saved with that image + - in PNG chunks for PNG, in EXIF for JPEG + - can drag the image to PNG info tab to restore generation parameters and automatically copy them into UI + - can be disabled in settings + - drag and drop an image/text-parameters to promptbox +- Read Generation Parameters Button, loads parameters in promptbox to UI +- Settings page +- Running arbitrary python code from UI (must run with `--allow-code` to enable) +- Mouseover hints for most UI elements +- Possible to change defaults/mix/max/step values for UI elements via text config +- Tiling support, a checkbox to create images that can be tiled like textures +- Progress bar and live image generation preview + - Can use a separate neural network to produce previews with almost none VRAM or compute requirement +- Negative prompt, an extra text field that allows you to list what you don't want to see in generated image +- Styles, a way to save part of prompt and easily apply them via dropdown later +- Variations, a way to generate same image but with tiny differences +- Seed resizing, a way to generate same image but at slightly different resolution +- CLIP interrogator, a button that tries to guess prompt from an image +- Prompt Editing, a way to change prompt mid-generation, say to start making a watermelon and switch to anime girl midway +- Batch Processing, process a group of files using img2img +- Img2img Alternative, reverse Euler method of cross attention control +- Highres Fix, a convenience option to produce high resolution pictures in one click without usual distortions +- Reloading checkpoints on the fly +- Checkpoint Merger, a tab that allows you to merge up to 3 checkpoints into one +- [Custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) with many extensions from community +- [Composable-Diffusion](https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/), a way to use multiple prompts at once + - separate prompts using uppercase `AND` + - also supports weights for prompts: `a cat :1.2 AND a dog AND a penguin :2.2` +- No token limit for prompts (original stable diffusion lets you use up to 75 tokens) +- DeepDanbooru integration, creates danbooru style tags for anime prompts +- [xformers](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Xformers), major speed increase for select cards: (add `--xformers` to commandline args) +- via extension: [History tab](https://github.com/yfszzx/stable-diffusion-webui-images-browser): view, direct and delete images conveniently within the UI +- Generate forever option +- Training tab + - hypernetworks and embeddings options + - Preprocessing images: cropping, mirroring, autotagging using BLIP or deepdanbooru (for anime) +- Clip skip +- Hypernetworks +- Loras (same as Hypernetworks but more pretty) +- A separate UI where you can choose, with preview, which embeddings, hypernetworks or Loras to add to your prompt +- Can select to load a different VAE from settings screen +- Estimated completion time in progress bar +- API +- Support for dedicated [inpainting model](https://github.com/runwayml/stable-diffusion#inpainting-with-stable-diffusion) by RunwayML +- via extension: [Aesthetic Gradients](https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients), a way to generate images with a specific aesthetic by using clip images embeds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients)) +- [Stable Diffusion 2.0](https://github.com/Stability-AI/stablediffusion) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20) for instructions +- [Alt-Diffusion](https://arxiv.org/abs/2211.06679) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#alt-diffusion) for instructions +- Now without any bad letters! +- Load checkpoints in safetensors format +- Eased resolution restriction: generated image's dimensions must be a multiple of 8 rather than 64 +- Now with a license! +- Reorder elements in the UI from settings screen +- [Segmind Stable Diffusion](https://huggingface.co/segmind/SSD-1B) support + +## Installation and Running +Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for: +- [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) +- [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs. +- [Intel CPUs, Intel GPUs (both integrated and discrete)](https://github.com/openvinotoolkit/stable-diffusion-webui/wiki/Installation-on-Intel-Silicon) (external wiki page) +- [Ascend NPUs](https://github.com/wangshuai09/stable-diffusion-webui/wiki/Install-and-run-on-Ascend-NPUs) (external wiki page) + +Alternatively, use online services (like Google Colab): + +- [List of Online Services](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Online-Services) + +### Installation on Windows 10/11 with NVidia-GPUs using release package +1. Download `sd.webui.zip` from [v1.0.0-pre](https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/tag/v1.0.0-pre) and extract its contents. +2. Run `update.bat`. +3. Run `run.bat`. +> For more details see [Install-and-Run-on-NVidia-GPUs](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) + +### Automatic Installation on Windows +1. Install [Python 3.10.6](https://www.python.org/downloads/release/python-3106/) (Newer version of Python does not support torch), checking "Add Python to PATH". +2. Install [git](https://git-scm.com/download/win). +3. Download the stable-diffusion-webui repository, for example by running `git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git`. +4. Run `webui-user.bat` from Windows Explorer as normal, non-administrator, user. + +### Automatic Installation on Linux +1. Install the dependencies: +```bash +# Debian-based: +sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0 +# Red Hat-based: +sudo dnf install wget git python3 gperftools-libs libglvnd-glx +# openSUSE-based: +sudo zypper install wget git python3 libtcmalloc4 libglvnd +# Arch-based: +sudo pacman -S wget git python3 +``` +If your system is very new, you need to install python3.11 or python3.10: +```bash +# Ubuntu 24.04 +sudo add-apt-repository ppa:deadsnakes/ppa +sudo apt update +sudo apt install python3.11 python3.11-venv + +# Manjaro/Arch +sudo pacman -S yay +yay -S python311 # do not confuse with python3.11 package + +# Only for 3.11 +# Then set up env variable in launch script +export python_cmd="python3.11" +# or in webui-user.sh +python_cmd="python3.11" +``` +2. Navigate to the directory you would like the webui to be installed and execute the following command: +```bash +wget -q https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh +chmod +x webui.sh +``` +Or just clone the repo wherever you want: +```bash +git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui +``` + +3. Run `webui.sh`. +4. Check `webui-user.sh` for options. +### Installation on Apple Silicon + +Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon). + +## Contributing +Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) + +## Documentation + +The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki). + +For the purposes of getting Google and other search engines to crawl the wiki, here's a link to the (not for humans) [crawlable wiki](https://github-wiki-see.page/m/AUTOMATIC1111/stable-diffusion-webui/wiki). + +## Credits +Licenses for borrowed code can be found in `Settings -> Licenses` screen, and also in `html/licenses.html` file. + +- Stable Diffusion - https://github.com/Stability-AI/stablediffusion, https://github.com/CompVis/taming-transformers, https://github.com/mcmonkey4eva/sd3-ref +- k-diffusion - https://github.com/crowsonkb/k-diffusion.git +- Spandrel - https://github.com/chaiNNer-org/spandrel implementing + - GFPGAN - https://github.com/TencentARC/GFPGAN.git + - CodeFormer - https://github.com/sczhou/CodeFormer + - ESRGAN - https://github.com/xinntao/ESRGAN + - SwinIR - https://github.com/JingyunLiang/SwinIR + - Swin2SR - https://github.com/mv-lab/swin2sr +- LDSR - https://github.com/Hafiidz/latent-diffusion +- MiDaS - https://github.com/isl-org/MiDaS +- Ideas for optimizations - https://github.com/basujindal/stable-diffusion +- Cross Attention layer optimization - Doggettx - https://github.com/Doggettx/stable-diffusion, original idea for prompt editing. +- Cross Attention layer optimization - InvokeAI, lstein - https://github.com/invoke-ai/InvokeAI (originally http://github.com/lstein/stable-diffusion) +- Sub-quadratic Cross Attention layer optimization - Alex Birch (https://github.com/Birch-san/diffusers/pull/1), Amin Rezaei (https://github.com/AminRezaei0x443/memory-efficient-attention) +- Textual Inversion - Rinon Gal - https://github.com/rinongal/textual_inversion (we're not using his code, but we are using his ideas). +- Idea for SD upscale - https://github.com/jquesnelle/txt2imghd +- Noise generation for outpainting mk2 - https://github.com/parlance-zz/g-diffuser-bot +- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator +- Idea for Composable Diffusion - https://github.com/energy-based-model/Compositional-Visual-Generation-with-Composable-Diffusion-Models-PyTorch +- xformers - https://github.com/facebookresearch/xformers +- DeepDanbooru - interrogator for anime diffusers https://github.com/KichangKim/DeepDanbooru +- Sampling in float32 precision from a float16 UNet - marunine for the idea, Birch-san for the example Diffusers implementation (https://github.com/Birch-san/diffusers-play/tree/92feee6) +- Instruct pix2pix - Tim Brooks (star), Aleksander Holynski (star), Alexei A. Efros (no star) - https://github.com/timothybrooks/instruct-pix2pix +- Security advice - RyotaK +- UniPC sampler - Wenliang Zhao - https://github.com/wl-zhao/UniPC +- TAESD - Ollin Boer Bohan - https://github.com/madebyollin/taesd +- LyCORIS - KohakuBlueleaf +- Restart sampling - lambertae - https://github.com/Newbeeer/diffusion_restart_sampling +- Hypertile - tfernd - https://github.com/tfernd/HyperTile +- Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user. +- (You) From c3ab0514707f403534a029c827ec2bb7c3680fe0 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 14:19:48 -0500 Subject: [PATCH 14/17] Fuse timestep embedding with MPS GroupNorm --- modules/mps_flash_attention.py | 13 ++- modules/mps_fused_ops.py | 68 ++++++++++++ modules/sd_hijack_unet.py | 6 +- modules/shared_options.py | 1 + scripts/benchmark_mps_unet_ops.py | 22 ++++ scripts/install_mps_flash_attention.py | 4 +- scripts/mps_fused_group_norm.mm | 138 +++++++++++++++++++++++++ test/test_mps_fused_ops.py | 31 ++++++ 8 files changed, 276 insertions(+), 7 deletions(-) diff --git a/modules/mps_flash_attention.py b/modules/mps_flash_attention.py index c4e5bcf7d5f..9a9ea946c6e 100644 --- a/modules/mps_flash_attention.py +++ b/modules/mps_flash_attention.py @@ -47,7 +47,7 @@ def _run_isolated_self_test(): code = """ import torch import torch.nn.functional as F -from metal_flash_sdpa import MetalFlashAttentionForward, fused_geglu_forward, fused_group_norm_silu_forward +from metal_flash_sdpa import MetalFlashAttentionForward, fused_geglu_forward, fused_group_norm_silu_add_embedding_forward, fused_group_norm_silu_forward torch.manual_seed(1) source = torch.randn((1, 256, 320), device='mps', dtype=torch.float16) @@ -76,6 +76,15 @@ def _run_isolated_self_test(): assert norm_difference.max().item() < 0.02 assert norm_difference.mean().item() < 0.001 +embedding = torch.randn((1, 320), device='mps', dtype=torch.float16) +expected_embedding_norm = F.silu(F.group_norm(norm_source + embedding[:, :, None, None], 32, norm_weight, norm_bias, 1e-5)) +actual_embedding_norm = fused_group_norm_silu_add_embedding_forward(norm_source, embedding, norm_weight, norm_bias, 32, 1e-5) + 0 +torch.mps.synchronize() +embedding_difference = (actual_embedding_norm.float() - expected_embedding_norm.float()).abs() +assert torch.isfinite(actual_embedding_norm).all().item() +assert embedding_difference.max().item() < 0.02 +assert embedding_difference.mean().item() < 0.001 + import numpy as np geglu_source = torch.randn((1, 256, 2560), device='mps', dtype=torch.float16) half_values = np.arange(65536, dtype=np.uint16).view(np.float16).copy() @@ -126,6 +135,8 @@ def is_available(): raise RuntimeError("native extension is missing the A1111 deferred MPS commit patch") if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU", False): raise RuntimeError("native extension is missing fused GroupNorm+SiLU") + if not getattr(_extension, "A1111_MPS_FUSED_GROUP_NORM_SILU_EMBEDDING", False): + raise RuntimeError("native extension is missing fused GroupNorm+SiLU+embedding") if not getattr(_extension, "A1111_MPS_FUSED_GEGLU", False): raise RuntimeError("native extension is missing fused GEGLU") _run_isolated_self_test() diff --git a/modules/mps_fused_ops.py b/modules/mps_fused_ops.py index 288bd482a09..78c4dd95b8c 100644 --- a/modules/mps_fused_ops.py +++ b/modules/mps_fused_ops.py @@ -16,6 +16,11 @@ _runtime_failure_warned = False _first_dispatch_logged = False _runtime_disabled = False +_embedding_dispatch_count = 0 +_embedding_fallback_count = 0 +_embedding_runtime_failure_warned = False +_embedding_first_dispatch_logged = False +_embedding_runtime_disabled = False _geglu_dispatch_count = 0 _geglu_fallback_count = 0 _geglu_runtime_failure_warned = False @@ -82,6 +87,67 @@ def group_norm_silu(input_tensor, norm): return F.silu(norm(input_tensor)) +def group_norm_silu_add_embedding(input_tensor, embedding, norm): + """Fuse a broadcast timestep embedding, GroupNorm, and SiLU when supported.""" + from modules import shared + + global _embedding_dispatch_count, _embedding_fallback_count + global _embedding_runtime_failure_warned, _embedding_first_dispatch_logged + global _embedding_runtime_disabled + + can_dispatch = ( + not _embedding_runtime_disabled + and os.environ.get("A1111_MPS_DISABLE_FUSED_GROUP_NORM_SILU_EMBEDDING") != "1" + and input_tensor.device.type == "mps" + and input_tensor.dtype == torch.float16 + and embedding.device == input_tensor.device + and embedding.dtype == input_tensor.dtype + and input_tensor.ndim == 4 + and embedding.ndim == 2 + and input_tensor.is_contiguous() + and embedding.is_contiguous() + and embedding.shape == (input_tensor.shape[0], input_tensor.shape[1]) + and norm.weight is not None + and norm.bias is not None + and norm.weight.device == input_tensor.device + and norm.bias.device == input_tensor.device + and norm.weight.dtype == input_tensor.dtype + and norm.bias.dtype == input_tensor.dtype + and norm.weight.is_contiguous() + and norm.bias.is_contiguous() + and input_tensor.shape[1] % norm.num_groups == 0 + and getattr(shared.opts, "mps_fused_group_norm_silu_embedding", True) + and not (torch.is_grad_enabled() and (input_tensor.requires_grad or embedding.requires_grad or norm.weight.requires_grad or norm.bias.requires_grad)) + and mps_flash_attention.is_available() + ) + + if can_dispatch: + try: + result = mps_flash_attention._extension.fused_group_norm_silu_add_embedding_forward( + input_tensor, + embedding, + norm.weight, + norm.bias, + norm.num_groups, + norm.eps, + ) + _embedding_dispatch_count += 1 + if not _embedding_first_dispatch_logged: + print(f"Fused Metal GroupNorm+SiLU+embedding first dispatch: {tuple(input_tensor.shape)}") + _embedding_first_dispatch_logged = True + return result + except RuntimeError as exc: + _embedding_runtime_disabled = True + if not _embedding_runtime_failure_warned: + print(f"Fused Metal GroupNorm+SiLU+embedding failed; using PyTorch: {exc}") + _embedding_runtime_failure_warned = True + + _embedding_fallback_count += 1 + while embedding.ndim < input_tensor.ndim: + embedding = embedding[..., None] + return F.silu(norm(input_tensor + embedding)) + + def _can_dispatch_geglu(projected): if _geglu_runtime_disabled: return False @@ -134,6 +200,8 @@ def diagnostics(): return { "dispatches": _dispatch_count, "fallbacks": _fallback_count, + "embedding_dispatches": _embedding_dispatch_count, + "embedding_fallbacks": _embedding_fallback_count, "geglu_dispatches": _geglu_dispatch_count, "geglu_fallbacks": _geglu_fallback_count, } diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index cef1e4c361f..89f47bb6c6a 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -65,11 +65,7 @@ def fused_resblock_forward(_, self, x, emb): h = self.in_layers[2](mps_fused_ops.group_norm_silu(x, self.in_layers[0])) emb_out = self.emb_layers(emb).type(h.dtype) - while len(emb_out.shape) < len(h.shape): - emb_out = emb_out[..., None] - - h = h + emb_out - h = mps_fused_ops.group_norm_silu(h, self.out_layers[0]) + h = mps_fused_ops.group_norm_silu_add_embedding(h, emb_out, self.out_layers[0]) h = self.out_layers[2](h) h = self.out_layers[3](h) return self.skip_connection(x) + h diff --git a/modules/shared_options.py b/modules/shared_options.py index 47193448d2f..0587c5ab72e 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -233,6 +233,7 @@ options_templates.update(options_section(('optimizations', "Optimizations", "sd"), { "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), "mps_fused_group_norm_silu": OptionInfo(True, "Fuse GroupNorm and SiLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), + "mps_fused_group_norm_silu_embedding": OptionInfo(True, "Fuse timestep embedding with GroupNorm and SiLU on Apple Silicon").info("uses the native Metal glue kernel when supported; disable to isolate its end-to-end impact"), "mps_fused_geglu": OptionInfo(True, "Fuse GEGLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), "s_min_uncond": OptionInfo(1.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), "s_min_uncond_all": OptionInfo(True, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), diff --git a/scripts/benchmark_mps_unet_ops.py b/scripts/benchmark_mps_unet_ops.py index 57d1bc5e790..27b30e9d675 100644 --- a/scripts/benchmark_mps_unet_ops.py +++ b/scripts/benchmark_mps_unet_ops.py @@ -4,12 +4,18 @@ from __future__ import annotations import argparse +import pathlib import statistics +import sys import time import torch import torch.nn.functional as F +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +from modules import mps_fused_ops + SD1_SHAPES = ( (4096, 320), @@ -47,6 +53,12 @@ def measure_shape(batch, tokens, channels, warmup, repeats): convolution_bias = torch.randn((channels,), device="mps", dtype=torch.float16) sequence = image.flatten(2).transpose(1, 2) projection_weight = torch.randn((channels, channels), device="mps", dtype=torch.float16) + norm_weight = torch.randn((channels,), device="mps", dtype=torch.float16) + norm_bias = torch.randn((channels,), device="mps", dtype=torch.float16) + embedding = torch.randn((batch, channels), device="mps", dtype=torch.float16) + norm = torch.nn.GroupNorm(32, channels).to("mps").half() + norm.weight.data.copy_(norm_weight) + norm.bias.data.copy_(norm_bias) heads = 8 query = sequence.view(batch, tokens, heads, channels // heads).transpose(1, 2) @@ -62,6 +74,16 @@ def measure_shape(batch, tokens, channels, warmup, repeats): warmup, repeats, ), + "groupnorm+silu+embedding": measure( + lambda: F.silu(F.group_norm(image + embedding[:, :, None, None], 32, norm_weight, norm_bias)), + warmup, + repeats, + ), + "fused_groupnorm+silu+embedding": measure( + lambda: mps_fused_ops.group_norm_silu_add_embedding(image, embedding, norm), + warmup, + repeats, + ), "linear": measure( lambda: F.linear(sequence, projection_weight), warmup, diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py index b61cc99dd61..8ebaec58f16 100644 --- a/scripts/install_mps_flash_attention.py +++ b/scripts/install_mps_flash_attention.py @@ -104,13 +104,15 @@ def patch_source(source): 'A1111_MPS_STREAM_FIX = True\n' 'A1111_MPS_DEFERRED_COMMIT = True\n' 'A1111_MPS_FUSED_GROUP_NORM_SILU = True\n' - 'A1111_MPS_FUSED_GEGLU = True\n', + 'A1111_MPS_FUSED_GEGLU = True\n' + 'A1111_MPS_FUSED_GROUP_NORM_SILU_EMBEDDING = True\n', ) replace_exact( package_init, "from metal_flash_sdpa._C import mfa_attention_forward, mfa_attention_backward\n", "from metal_flash_sdpa._C import (\n" " fused_geglu_forward,\n" + " fused_group_norm_silu_add_embedding_forward,\n" " fused_group_norm_silu_forward,\n" " mfa_attention_backward,\n" " mfa_attention_forward,\n" diff --git a/scripts/mps_fused_group_norm.mm b/scripts/mps_fused_group_norm.mm index 8a91fe9fb3d..2f0c97586ca 100644 --- a/scripts/mps_fused_group_norm.mm +++ b/scripts/mps_fused_group_norm.mm @@ -153,6 +153,7 @@ kernel void fused_geglu_half( const ushort gate_bits = input_bits[input_base + params.width + column]; output[index] = half(value * float(gelu_lut[gate_bits])); } + )METAL"; NSError* error = nil; @@ -172,6 +173,85 @@ kernel void fused_geglu_half( return pipeline; } +static id getFusedGroupNormSiLUAddEmbeddingPipeline() { + static id pipeline = nil; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +using namespace metal; + +struct FusedGroupNormParams { + uint batch; + uint channels; + uint spatial; + uint groups; + float epsilon; +}; + +kernel void fused_group_norm_silu_add_embedding_half( + device const half* input [[buffer(0)]], + device const half* embedding [[buffer(1)]], + device const half* weight [[buffer(2)]], + device const half* bias [[buffer(3)]], + device half* output [[buffer(4)]], + constant FusedGroupNormParams& params [[buffer(5)]], + uint tid [[thread_index_in_threadgroup]], + uint group_index [[threadgroup_position_in_grid]], + uint threads [[threads_per_threadgroup]]) { + threadgroup float partial_sum[256]; + threadgroup float partial_square_sum[256]; + const uint channels_per_group = params.channels / params.groups; + const uint group_elements = channels_per_group * params.spatial; + const uint batch_index = group_index / params.groups; + const uint channel_group = group_index - batch_index * params.groups; + const uint base = (batch_index * params.channels + channel_group * channels_per_group) * params.spatial; + + float sum = 0.0f; + float square_sum = 0.0f; + for (uint index = tid; index < group_elements; index += threads) { + const uint local_channel = index / params.spatial; + const uint channel = channel_group * channels_per_group + local_channel; + const half value = half(input[base + index] + embedding[batch_index * params.channels + channel]); + sum += float(value); + square_sum += float(value) * float(value); + } + partial_sum[tid] = sum; + partial_square_sum[tid] = square_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = threads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + partial_sum[tid] += partial_sum[tid + stride]; + partial_square_sum[tid] += partial_square_sum[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float mean = partial_sum[0] / float(group_elements); + const float variance = max(partial_square_sum[0] / float(group_elements) - mean * mean, 0.0f); + const float inverse_stddev = rsqrt(variance + params.epsilon); + for (uint index = tid; index < group_elements; index += threads) { + const uint local_channel = index / params.spatial; + const uint channel = channel_group * channels_per_group + local_channel; + const half combined = half(input[base + index] + embedding[batch_index * params.channels + channel]); + float value = (float(combined) - mean) * inverse_stddev; + value = value * float(weight[channel]) + float(bias[channel]); + value = value / (1.0f + exp(-value)); + output[base + index] = half(value); + } +} +)METAL"; + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK(library != nil, "Failed to compile fused GroupNorm+SiLU+embedding Metal library: ", error ? [[error localizedDescription] UTF8String] : "unknown error"); + id function = [library newFunctionWithName:@"fused_group_norm_silu_add_embedding_half"]; + TORCH_CHECK(function != nil, "Fused GroupNorm+SiLU+embedding Metal function was not found"); + pipeline = [device newComputePipelineStateWithFunction:function error:&error]; + TORCH_CHECK(pipeline != nil, "Failed to create fused GroupNorm+SiLU+embedding pipeline: ", error ? [[error localizedDescription] UTF8String] : "unknown error"); + }); + return pipeline; +} + torch::Tensor fused_group_norm_silu_forward( const torch::Tensor& input, const torch::Tensor& weight, @@ -287,6 +367,54 @@ kernel void fused_geglu_half( return output; } +torch::Tensor fused_group_norm_silu_add_embedding_forward( + const torch::Tensor& input, + const torch::Tensor& embedding, + const torch::Tensor& weight, + const torch::Tensor& bias, + int64_t groups, + double epsilon) { + TORCH_CHECK(input.device().is_mps() && embedding.device().is_mps(), "input and embedding must be MPS tensors"); + TORCH_CHECK(weight.device().is_mps() && bias.device().is_mps(), "weight and bias must be MPS tensors"); + TORCH_CHECK(input.scalar_type() == at::kHalf && embedding.scalar_type() == at::kHalf, "input and embedding must be float16"); + TORCH_CHECK(weight.scalar_type() == at::kHalf && bias.scalar_type() == at::kHalf, "weight and bias must be float16"); + TORCH_CHECK(input.dim() == 4 && embedding.dim() == 2, "invalid input dimensions"); + TORCH_CHECK(input.is_contiguous() && embedding.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(embedding.size(0) == input.size(0) && embedding.size(1) == input.size(1), "embedding must have shape [batch, channels]"); + TORCH_CHECK(groups > 0 && input.size(1) % groups == 0, "channels must be divisible by groups"); + TORCH_CHECK(weight.numel() == input.size(1) && bias.numel() == input.size(1), "weight and bias must match the channel count"); + + auto output = torch::empty_like(input); + FusedGroupNormParams params = { + static_cast(input.size(0)), + static_cast(input.size(1)), + static_cast(input.size(2) * input.size(3)), + static_cast(groups), + static_cast(epsilon), + }; + auto pipeline = getFusedGroupNormSiLUAddEmbeddingPipeline(); + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id encoder = [command_buffer computeCommandEncoder]; + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:getMTLBufferStorage(input) offset:getMTLBufferOffset(input) atIndex:0]; + [encoder setBuffer:getMTLBufferStorage(embedding) offset:getMTLBufferOffset(embedding) atIndex:1]; + [encoder setBuffer:getMTLBufferStorage(weight) offset:getMTLBufferOffset(weight) atIndex:2]; + [encoder setBuffer:getMTLBufferStorage(bias) offset:getMTLBufferOffset(bias) atIndex:3]; + [encoder setBuffer:getMTLBufferStorage(output) offset:getMTLBufferOffset(output) atIndex:4]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:5]; + [encoder dispatchThreadgroups:MTLSizeMake(params.batch * params.groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [encoder endEncoding]; + } + }); + } + return output; +} + } // namespace void register_fused_ops(pybind11::module_& module) { @@ -299,6 +427,16 @@ void register_fused_ops(pybind11::module_& module) { pybind11::arg("bias"), pybind11::arg("groups"), pybind11::arg("epsilon")); + module.def( + "fused_group_norm_silu_add_embedding_forward", + &fused_group_norm_silu_add_embedding_forward, + "Fused Metal GroupNorm, SiLU, and timestep embedding addition", + pybind11::arg("input"), + pybind11::arg("embedding"), + pybind11::arg("weight"), + pybind11::arg("bias"), + pybind11::arg("groups"), + pybind11::arg("epsilon")); module.def( "fused_geglu_forward", &fused_geglu_forward, diff --git a/test/test_mps_fused_ops.py b/test/test_mps_fused_ops.py index db06e80670f..00f79628ebc 100644 --- a/test/test_mps_fused_ops.py +++ b/test/test_mps_fused_ops.py @@ -33,6 +33,37 @@ def test_native_fusion_matches_pytorch(): assert difference.mean().item() < 0.001 +def test_cpu_embedding_fallback_matches_pytorch(): + torch.manual_seed(3) + norm = torch.nn.GroupNorm(4, 8) + source = torch.randn(1, 8, 6, 10) + embedding = torch.randn(1, 8) + + actual = mps_fused_ops.group_norm_silu_add_embedding(source, embedding, norm) + expected = F.silu(norm(source + embedding[:, :, None, None])) + + assert torch.equal(actual, expected) + + +def test_native_embedding_fusion_matches_pytorch(): + if not torch.backends.mps.is_available(): + return + torch.manual_seed(3) + norm = torch.nn.GroupNorm(32, 320).eval().half().to("mps") + source = torch.randn(1, 320, 48, 80, device="mps", dtype=torch.float16) + embedding = torch.randn(1, 320, device="mps", dtype=torch.float16) + + with torch.no_grad(): + expected = F.silu(norm(source + embedding[:, :, None, None])) + actual = mps_fused_ops.group_norm_silu_add_embedding(source, embedding, norm) + 0 + torch.mps.synchronize() + + difference = (actual.float() - expected.float()).abs() + assert torch.isfinite(actual).all().item() + assert difference.max().item() < 0.02 + assert difference.mean().item() < 0.001 + + def test_geglu_cpu_fallback_matches_pytorch(): torch.manual_seed(2) projection = torch.nn.Linear(8, 32) From ad4c6738280d483536635ad500d82e3ec943f0d5 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 15:07:35 -0500 Subject: [PATCH 15/17] Add benchmarks for MPS quantized linear and FP16 attention projections Signed-off-by: Derek Anderson --- modules/sd_hijack_unet.py | 172 ++++++------- modules/shared_options.py | 12 +- scripts/benchmark_mps_projection_fusion.py | 126 ++++++++++ scripts/benchmark_mps_quantized_linear.py | 120 +++++++++ scripts/mps_quantized_linear.mm | 277 +++++++++++++++++++++ 5 files changed, 615 insertions(+), 92 deletions(-) create mode 100644 scripts/benchmark_mps_projection_fusion.py create mode 100644 scripts/benchmark_mps_quantized_linear.py create mode 100644 scripts/mps_quantized_linear.mm diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index 89f47bb6c6a..8b1815c2361 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -3,7 +3,7 @@ from einops import repeat import math -from modules import devices, mps_fused_ops +from modules import devices, mps_fused_ops from modules.sd_hijack_utils import CondFunc @@ -36,84 +36,84 @@ def cat(self, tensors, *args, **kwargs): th = TorchHijackForUnet() -def fused_resblock_condition(_, self, x, emb): - return ( - x.device.type == "mps" - and x.dtype == torch.float16 - and x.ndim == 4 - and emb is not None - and not self.training - and not self.use_scale_shift_norm - and not getattr(self, "skip_t_emb", False) - and not getattr(self, "exchange_temb_dims", False) - and len(self.in_layers) == 3 - and len(self.out_layers) == 4 - and isinstance(self.in_layers[0], torch.nn.GroupNorm) - and isinstance(self.in_layers[1], torch.nn.SiLU) - and isinstance(self.out_layers[0], torch.nn.GroupNorm) - and isinstance(self.out_layers[1], torch.nn.SiLU) - ) - - -def fused_resblock_forward(_, self, x, emb): - if self.updown: - h = mps_fused_ops.group_norm_silu(x, self.in_layers[0]) - h = self.h_upd(h) - x = self.x_upd(x) - h = self.in_layers[2](h) - else: - h = self.in_layers[2](mps_fused_ops.group_norm_silu(x, self.in_layers[0])) - - emb_out = self.emb_layers(emb).type(h.dtype) - h = mps_fused_ops.group_norm_silu_add_embedding(h, emb_out, self.out_layers[0]) - h = self.out_layers[2](h) - h = self.out_layers[3](h) - return self.skip_connection(x) + h - - -def fused_vae_resnet_condition(_, self, x, temb): - return ( - x.device.type == "mps" - and x.dtype == torch.float16 - and x.ndim == 4 - and not self.training - and isinstance(self.norm1, torch.nn.GroupNorm) - and isinstance(self.norm2, torch.nn.GroupNorm) - ) - - -def fused_vae_resnet_forward(_, self, x, temb): - h = self.conv1(mps_fused_ops.group_norm_silu(x, self.norm1)) - - if temb is not None: - h = h + self.temb_proj(torch.nn.functional.silu(temb))[:, :, None, None] - - h = self.dropout(mps_fused_ops.group_norm_silu(h, self.norm2)) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - x = self.conv_shortcut(x) - else: - x = self.nin_shortcut(x) - - return x + h - - -def fused_geglu_condition(_, self, x): - return ( - x.device.type == "mps" - and x.dtype == torch.float16 - and x.ndim == 3 - and not self.training - and hasattr(self, "proj") - ) - - -def fused_geglu_forward(_, self, x): - return mps_fused_ops.geglu(x, self.proj) - - +def fused_resblock_condition(_, self, x, emb): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 4 + and emb is not None + and not self.training + and not self.use_scale_shift_norm + and not getattr(self, "skip_t_emb", False) + and not getattr(self, "exchange_temb_dims", False) + and len(self.in_layers) == 3 + and len(self.out_layers) == 4 + and isinstance(self.in_layers[0], torch.nn.GroupNorm) + and isinstance(self.in_layers[1], torch.nn.SiLU) + and isinstance(self.out_layers[0], torch.nn.GroupNorm) + and isinstance(self.out_layers[1], torch.nn.SiLU) + ) + + +def fused_resblock_forward(_, self, x, emb): + if self.updown: + h = mps_fused_ops.group_norm_silu(x, self.in_layers[0]) + h = self.h_upd(h) + x = self.x_upd(x) + h = self.in_layers[2](h) + else: + h = self.in_layers[2](mps_fused_ops.group_norm_silu(x, self.in_layers[0])) + + emb_out = self.emb_layers(emb).type(h.dtype) + h = mps_fused_ops.group_norm_silu_add_embedding(h, emb_out, self.out_layers[0]) + h = self.out_layers[2](h) + h = self.out_layers[3](h) + return self.skip_connection(x) + h + + +def fused_vae_resnet_condition(_, self, x, temb): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 4 + and not self.training + and isinstance(self.norm1, torch.nn.GroupNorm) + and isinstance(self.norm2, torch.nn.GroupNorm) + ) + + +def fused_vae_resnet_forward(_, self, x, temb): + h = self.conv1(mps_fused_ops.group_norm_silu(x, self.norm1)) + + if temb is not None: + h = h + self.temb_proj(torch.nn.functional.silu(temb))[:, :, None, None] + + h = self.dropout(mps_fused_ops.group_norm_silu(h, self.norm2)) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x + h + + +def fused_geglu_condition(_, self, x): + return ( + x.device.type == "mps" + and x.dtype == torch.float16 + and x.ndim == 3 + and not self.training + and hasattr(self, "proj") + ) + + +def fused_geglu_forward(_, self, x): + return mps_fused_ops.geglu(x, self.proj) + + # Below are monkey patches to enable upcasting a float16 UNet for float32 sampling def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): """Always make sure inputs to unet are in correct dtype.""" @@ -125,7 +125,7 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y] with devices.autocast(): - result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs) + result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs) if devices.unet_needs_upcast: return result.float() else: @@ -203,12 +203,12 @@ def hijack_ddpm_edit(): CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding) CondFunc('ldm.modules.attention.SpatialTransformer.forward', spatial_transformer_forward) -CondFunc('ldm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) -CondFunc('sgm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) -CondFunc('ldm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) -CondFunc('sgm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) -CondFunc('ldm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) -CondFunc('sgm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) +CondFunc('ldm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) +CondFunc('sgm.modules.diffusionmodules.openaimodel.ResBlock._forward', fused_resblock_forward, fused_resblock_condition) +CondFunc('ldm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) +CondFunc('sgm.modules.diffusionmodules.model.ResnetBlock.forward', fused_vae_resnet_forward, fused_vae_resnet_condition) +CondFunc('ldm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) +CondFunc('sgm.modules.attention.GEGLU.forward', fused_geglu_forward, fused_geglu_condition) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): diff --git a/modules/shared_options.py b/modules/shared_options.py index 0587c5ab72e..ec8dfa856c9 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -179,7 +179,7 @@ "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"), "comma_padding_backtrack": OptionInfo(20, "Prompt word wrap length limit", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1}).info("in tokens - for texts shorter than specified, if they don't fit into 75 token limit, move them to the next 75 token chunk"), "sdxl_clip_l_skip": OptionInfo(False, "Clip skip SDXL", gr.Checkbox).info("Enable Clip skip for the secondary clip model in sdxl. Has no effect on SD 1.5 or SD 2.0/2.1."), - "CLIP_stop_at_last_layers": OptionInfo(2, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}, infotext="Clip skip").link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer"), + "CLIP_stop_at_last_layers": OptionInfo(2, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}, infotext="Clip skip").link("wiki", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#clip-skip").info("ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer"), "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"), "randn_source": OptionInfo("GPU", "Random number generator source.", gr.Radio, {"choices": ["GPU", "CPU", "NV"]}, infotext="RNG").info("changes seeds drastically; use CPU to produce the same picture across different videocard vendors; use NV to produce same picture as on NVidia videocards"), "tiling": OptionInfo(False, "Tiling", infotext='Tiling').info("produce a tileable picture"), @@ -231,12 +231,12 @@ })) options_templates.update(options_section(('optimizations', "Optimizations", "sd"), { - "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), - "mps_fused_group_norm_silu": OptionInfo(True, "Fuse GroupNorm and SiLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), + "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}), + "mps_fused_group_norm_silu": OptionInfo(True, "Fuse GroupNorm and SiLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), "mps_fused_group_norm_silu_embedding": OptionInfo(True, "Fuse timestep embedding with GroupNorm and SiLU on Apple Silicon").info("uses the native Metal glue kernel when supported; disable to isolate its end-to-end impact"), - "mps_fused_geglu": OptionInfo(True, "Fuse GEGLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), - "s_min_uncond": OptionInfo(1.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), - "s_min_uncond_all": OptionInfo(True, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), + "mps_fused_geglu": OptionInfo(True, "Fuse GEGLU on Apple Silicon").info("uses the native Metal inference kernel when supported; disable to compare with PyTorch"), + "s_min_uncond": OptionInfo(1.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"), + "s_min_uncond_all": OptionInfo(True, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"), "token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"), "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"), "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for high-res pass", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio hr').info("only applies if non-zero and overrides above"), diff --git a/scripts/benchmark_mps_projection_fusion.py b/scripts/benchmark_mps_projection_fusion.py new file mode 100644 index 00000000000..a8c960ce671 --- /dev/null +++ b/scripts/benchmark_mps_projection_fusion.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Benchmark merged FP16 attention projections without changing model execution.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import torch +import torch.nn.functional as F + + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +SD1_SHAPES = ( + (4096, 320), + (1024, 640), + (256, 1280), +) +CONTEXT_TOKENS = 77 +CONTEXT_DIM = 768 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--warmup", type=int, default=4) + parser.add_argument("--repeats", type=int, default=12) + return parser.parse_args() + + +def measure(operation, warmup, repeats): + for _ in range(warmup): + operation() + torch.mps.synchronize() + + timings = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + torch.mps.synchronize() + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings) + + +def max_difference(separate, merged): + return max( + (left.float() - right.float()).abs().max().item() + for left, right in zip(separate, merged) + ) + + +def benchmark_self_attention(batch, tokens, channels, warmup, repeats): + sequence = torch.randn((batch, tokens, channels), device="mps", dtype=torch.float16) + weights = [torch.randn((channels, channels), device="mps", dtype=torch.float16) for _ in range(3)] + merged_weight = torch.cat(weights, dim=0).contiguous() + + def separate(): + return tuple(F.linear(sequence, weight) for weight in weights) + + def merged(): + return tuple(F.linear(sequence, merged_weight).chunk(3, dim=-1)) + + separate_output = separate() + merged_output = merged() + torch.mps.synchronize() + return { + "separate_ms": measure(separate, warmup, repeats), + "merged_ms": measure(merged, warmup, repeats), + "max_error": max_difference(separate_output, merged_output), + "exact": all(torch.equal(left, right) for left, right in zip(separate_output, merged_output)), + } + + +def benchmark_cross_attention(batch, tokens, channels, warmup, repeats): + context = torch.randn((batch, CONTEXT_TOKENS, CONTEXT_DIM), device="mps", dtype=torch.float16) + weights = [torch.randn((channels, CONTEXT_DIM), device="mps", dtype=torch.float16) for _ in range(2)] + merged_weight = torch.cat(weights, dim=0).contiguous() + + def separate(): + return tuple(F.linear(context, weight) for weight in weights) + + def merged(): + return tuple(F.linear(context, merged_weight).chunk(2, dim=-1)) + + separate_output = separate() + merged_output = merged() + torch.mps.synchronize() + return { + "separate_ms": measure(separate, warmup, repeats), + "merged_ms": measure(merged, warmup, repeats), + "max_error": max_difference(separate_output, merged_output), + "exact": all(torch.equal(left, right) for left, right in zip(separate_output, merged_output)), + } + + +def report(kind, tokens, channels, result): + speedup = result["separate_ms"] / result["merged_ms"] + print( + f"{kind:5s} tokens={tokens:4d} channels={channels:4d} " + f"separate={result['separate_ms']:.3f}ms merged={result['merged_ms']:.3f}ms " + f"speedup={speedup:.3f}x max_error={result['max_error']:.6f} exact={result['exact']}" + ) + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + torch.manual_seed(1) + print(f"PyTorch {torch.__version__}; FP16 MPS projection fusion; batch={args.batch}") + print("Benchmark only: merged weights are prepared before timing and no model path is modified.") + for tokens, channels in SD1_SHAPES: + self_result = benchmark_self_attention(args.batch, tokens, channels, args.warmup, args.repeats) + cross_result = benchmark_cross_attention(args.batch, tokens, channels, args.warmup, args.repeats) + report("self", tokens, channels, self_result) + report("cross", tokens, channels, cross_result) + torch.mps.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_mps_quantized_linear.py b/scripts/benchmark_mps_quantized_linear.py new file mode 100644 index 00000000000..28cd5f9dcdc --- /dev/null +++ b/scripts/benchmark_mps_quantized_linear.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +"""Benchmark direct int8-weight Metal projections against MPS FP16 linear.""" + +from __future__ import annotations + +import argparse +import statistics +import time + +import torch +import torch.nn.functional as F + + +SD1_SHAPES = ( + (4096, 320), + (1024, 640), + (256, 1280), +) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup", type=int, default=4) + parser.add_argument("--repeats", type=int, default=12) + return parser.parse_args() + + +def quantize_per_output_channel(weight): + scale = weight.abs().amax(dim=1).clamp_min(1e-8) / 127.0 + quantized = torch.round(weight / scale[:, None]).clamp(-128, 127).to(torch.int8) + return quantized.contiguous(), scale.to(dtype=torch.float16).contiguous() + + +def measure(operation, warmup, repeats): + for _ in range(warmup): + operation() + torch.mps.synchronize() + + timings = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + torch.mps.synchronize() + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings) + + +def benchmark_shape(rows, features, warmup, repeats, extension): + torch.manual_seed(rows + features) + activation = torch.randn((rows, features), device="mps", dtype=torch.float16) + weight = torch.randn((features, features), device="mps", dtype=torch.float16) + quantized_weight, scale = quantize_per_output_channel(weight) + + expected = F.linear(activation, weight) + quantized_expected = F.linear(activation, quantized_weight.half() * scale[:, None]) + scalar_actual = extension.quantized_linear_forward(activation, quantized_weight, scale) + simd_actual = extension.quantized_linear_simd_forward(activation, quantized_weight, scale) + torch.mps.synchronize() + scalar_kernel_difference = (scalar_actual.float() - quantized_expected.float()).abs() + simd_kernel_difference = (simd_actual.float() - quantized_expected.float()).abs() + quantization_difference = (simd_actual.float() - expected.float()).abs() + + fp16_ms = measure(lambda: F.linear(activation, weight), warmup, repeats) + scalar_int8_ms = measure( + lambda: extension.quantized_linear_forward(activation, quantized_weight, scale), + warmup, + repeats, + ) + simd_int8_ms = measure( + lambda: extension.quantized_linear_simd_forward(activation, quantized_weight, scale), + warmup, + repeats, + ) + return { + "fp16_ms": fp16_ms, + "scalar_int8_ms": scalar_int8_ms, + "simd_int8_ms": simd_int8_ms, + "scalar_kernel_max_error": scalar_kernel_difference.max().item(), + "simd_kernel_max_error": simd_kernel_difference.max().item(), + "quantization_max_error": quantization_difference.max().item(), + "quantization_mean_error": quantization_difference.mean().item(), + "fp16_weight_bytes": weight.numel() * weight.element_size(), + "int8_weight_bytes": quantized_weight.numel() * quantized_weight.element_size() + scale.numel() * scale.element_size(), + } + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + try: + import metal_flash_sdpa + except ImportError as exc: + raise SystemExit("Install the optional Metal extension before running this benchmark.") from exc + + if not getattr(metal_flash_sdpa, "A1111_MPS_QUANTIZED_LINEAR", False): + raise SystemExit("The installed Metal extension does not contain the quantized linear proof kernel.") + + extension = metal_flash_sdpa + print(f"PyTorch {torch.__version__}; int8 weights with per-output-channel fp16 scales; MPS") + print("Acceptance gate: low error and simd_int8_ms approaches fp16_ms on representative shapes.") + for rows, features in SD1_SHAPES: + result = benchmark_shape(rows, features, args.warmup, args.repeats, extension) + compression = result["fp16_weight_bytes"] / result["int8_weight_bytes"] + print( + f"rows={rows:4d} features={features:4d} " + f"fp16={result['fp16_ms']:.3f}ms " + f"scalar_int8={result['scalar_int8_ms']:.3f}ms " + f"simd_int8={result['simd_int8_ms']:.3f}ms " + f"simd_slowdown={result['simd_int8_ms'] / result['fp16_ms']:.2f}x " + f"simd_kernel_max_error={result['simd_kernel_max_error']:.4f} " + f"quantization_max_error={result['quantization_max_error']:.4f} " + f"quantization_mean_error={result['quantization_mean_error']:.4f} " + f"weight_compression={compression:.2f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/mps_quantized_linear.mm b/scripts/mps_quantized_linear.mm new file mode 100644 index 00000000000..a1d8a07d98d --- /dev/null +++ b/scripts/mps_quantized_linear.mm @@ -0,0 +1,277 @@ +#include +#include +#include +#include + +#import +#import + +#include + +namespace { + +struct QuantizedLinearParams { + uint32_t rows; + uint32_t input_features; + uint32_t output_features; +}; + +static inline id getMTLBufferStorage(const at::Tensor& tensor) { + return __builtin_bit_cast(id, tensor.storage().data()); +} + +static inline size_t getMTLBufferOffset(const at::Tensor& tensor) { + return tensor.storage_offset() * tensor.element_size(); +} + +static id getQuantizedLinearPipeline() { + static id pipeline = nil; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +using namespace metal; + +struct QuantizedLinearParams { + uint rows; + uint input_features; + uint output_features; +}; + +kernel void quantized_linear_int8_half( + device const half* input [[buffer(0)]], + device const char* weight [[buffer(1)]], + device const half* scale [[buffer(2)]], + device half* output [[buffer(3)]], + constant QuantizedLinearParams& params [[buffer(4)]], + uint index [[thread_position_in_grid]]) { + const uint count = params.rows * params.output_features; + if (index >= count) { + return; + } + + const uint row = index / params.output_features; + const uint column = index - row * params.output_features; + const uint input_base = row * params.input_features; + const uint weight_base = column * params.input_features; + float sum = 0.0f; + for (uint feature = 0; feature < params.input_features; ++feature) { + sum += float(input[input_base + feature]) * float(weight[weight_base + feature]); + } + + output[index] = half(sum * float(scale[column])); +} +)METAL"; + + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK( + library != nil, + "Failed to compile quantized linear Metal library: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + id function = [library newFunctionWithName:@"quantized_linear_int8_half"]; + TORCH_CHECK(function != nil, "Quantized linear Metal function was not found"); + pipeline = [device newComputePipelineStateWithFunction:function error:&error]; + TORCH_CHECK( + pipeline != nil, + "Failed to create quantized linear Metal pipeline: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + }); + return pipeline; +} + +static id getQuantizedLinearSIMDPipeline() { + static id pipeline = nil; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +#include +using namespace metal; + +struct QuantizedLinearParams { + uint rows; + uint input_features; + uint output_features; +}; + +kernel void quantized_linear_int8_simd_half( + device const half* input [[buffer(0)]], + device const char* weight [[buffer(1)]], + device const half* scale [[buffer(2)]], + device half* output [[buffer(3)]], + constant QuantizedLinearParams& params [[buffer(4)]], + ushort lane [[thread_index_in_simdgroup]], + uint2 tile [[threadgroup_position_in_grid]]) { + threadgroup half input_tile[64]; + threadgroup half weight_tile[64]; + threadgroup float output_tile[64]; + + const uint output_column = tile.x * 8; + const uint output_row = tile.y * 8; + simdgroup_float8x8 accumulator = make_filled_simdgroup_matrix(0.0f); + + for (uint feature_base = 0; feature_base < params.input_features; feature_base += 8) { + for (uint element = lane; element < 64; element += 32) { + const uint local_row = element / 8; + const uint local_feature = element - local_row * 8; + input_tile[element] = input[ + (output_row + local_row) * params.input_features + feature_base + local_feature]; + + const uint local_output = element % 8; + const uint feature = element / 8; + const uint weight_output = output_column + local_output; + weight_tile[element] = half(weight[ + weight_output * params.input_features + feature_base + feature]) * scale[weight_output]; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_half8x8 input_matrix; + simdgroup_half8x8 weight_matrix; + simdgroup_load(input_matrix, input_tile, 8); + simdgroup_load(weight_matrix, weight_tile, 8); + simdgroup_multiply_accumulate(accumulator, input_matrix, weight_matrix, accumulator); + simdgroup_barrier(mem_flags::mem_threadgroup); + } + + simdgroup_store(accumulator, output_tile, 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint element = lane; element < 64; element += 32) { + const uint local_row = element / 8; + const uint local_output = element - local_row * 8; + output[(output_row + local_row) * params.output_features + output_column + local_output] = + half(output_tile[element]); + } +} +)METAL"; + + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK( + library != nil, + "Failed to compile SIMD quantized linear Metal library: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + id function = [library newFunctionWithName:@"quantized_linear_int8_simd_half"]; + TORCH_CHECK(function != nil, "SIMD quantized linear Metal function was not found"); + pipeline = [device newComputePipelineStateWithFunction:function error:&error]; + TORCH_CHECK( + pipeline != nil, + "Failed to create SIMD quantized linear Metal pipeline: ", + error ? [[error localizedDescription] UTF8String] : "unknown error"); + }); + return pipeline; +} + +} // namespace + +torch::Tensor quantized_linear_forward( + const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& scale) { + TORCH_CHECK(input.device().is_mps(), "input must be an MPS tensor"); + TORCH_CHECK(weight.device().is_mps() && scale.device().is_mps(), "weights must be MPS tensors"); + TORCH_CHECK(input.scalar_type() == at::kHalf, "input must be float16"); + TORCH_CHECK(weight.scalar_type() == at::kChar, "weight must be int8"); + TORCH_CHECK(scale.scalar_type() == at::kHalf, "scale must be float16"); + TORCH_CHECK(input.dim() == 2, "input must have shape [rows, input_features]"); + TORCH_CHECK(weight.dim() == 2, "weight must have shape [output_features, input_features]"); + TORCH_CHECK(scale.dim() == 1, "scale must have shape [output_features]"); + TORCH_CHECK(input.is_contiguous() && weight.is_contiguous() && scale.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.size(1) == weight.size(1), "input and weight feature dimensions must match"); + TORCH_CHECK(scale.size(0) == weight.size(0), "scale and weight output dimensions must match"); + + auto output = torch::empty({input.size(0), weight.size(0)}, input.options()); + QuantizedLinearParams params = { + static_cast(input.size(0)), + static_cast(input.size(1)), + static_cast(weight.size(0)), + }; + const uint32_t count = params.rows * params.output_features; + auto pipeline = getQuantizedLinearPipeline(); + + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id encoder = [command_buffer computeCommandEncoder]; + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:getMTLBufferStorage(input) offset:getMTLBufferOffset(input) atIndex:0]; + [encoder setBuffer:getMTLBufferStorage(weight) offset:getMTLBufferOffset(weight) atIndex:1]; + [encoder setBuffer:getMTLBufferStorage(scale) offset:getMTLBufferOffset(scale) atIndex:2]; + [encoder setBuffer:getMTLBufferStorage(output) offset:getMTLBufferOffset(output) atIndex:3]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:4]; + const NSUInteger threads = std::min(pipeline.maxTotalThreadsPerThreadgroup, 256); + [encoder dispatchThreads:MTLSizeMake(count, 1, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [encoder endEncoding]; + } + }); + } + return output; +} + +torch::Tensor quantized_linear_simd_forward( + const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& scale) { + TORCH_CHECK(input.device().is_mps(), "input must be an MPS tensor"); + TORCH_CHECK(weight.device().is_mps() && scale.device().is_mps(), "weights must be MPS tensors"); + TORCH_CHECK(input.scalar_type() == at::kHalf, "input must be float16"); + TORCH_CHECK(weight.scalar_type() == at::kChar, "weight must be int8"); + TORCH_CHECK(scale.scalar_type() == at::kHalf, "scale must be float16"); + TORCH_CHECK(input.dim() == 2 && weight.dim() == 2 && scale.dim() == 1, "invalid tensor dimensions"); + TORCH_CHECK(input.is_contiguous() && weight.is_contiguous() && scale.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.size(1) == weight.size(1), "input and weight feature dimensions must match"); + TORCH_CHECK(scale.size(0) == weight.size(0), "scale and weight output dimensions must match"); + TORCH_CHECK(input.size(0) % 8 == 0, "input rows must be divisible by 8"); + TORCH_CHECK(input.size(1) % 8 == 0, "input features must be divisible by 8"); + TORCH_CHECK(weight.size(0) % 8 == 0, "output features must be divisible by 8"); + + auto output = torch::empty({input.size(0), weight.size(0)}, input.options()); + QuantizedLinearParams params = { + static_cast(input.size(0)), + static_cast(input.size(1)), + static_cast(weight.size(0)), + }; + auto pipeline = getQuantizedLinearSIMDPipeline(); + + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id encoder = [command_buffer computeCommandEncoder]; + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:getMTLBufferStorage(input) offset:getMTLBufferOffset(input) atIndex:0]; + [encoder setBuffer:getMTLBufferStorage(weight) offset:getMTLBufferOffset(weight) atIndex:1]; + [encoder setBuffer:getMTLBufferStorage(scale) offset:getMTLBufferOffset(scale) atIndex:2]; + [encoder setBuffer:getMTLBufferStorage(output) offset:getMTLBufferOffset(output) atIndex:3]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:4]; + [encoder dispatchThreadgroups:MTLSizeMake(params.output_features / 8, params.rows / 8, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [encoder endEncoding]; + } + }); + } + return output; +} + +void register_quantized_linear_op(pybind11::module_& module) { + module.def( + "quantized_linear_forward", + &quantized_linear_forward, + "Direct Metal int8-weight linear forward pass", + pybind11::arg("input"), + pybind11::arg("weight"), + pybind11::arg("scale")); + module.def( + "quantized_linear_simd_forward", + &quantized_linear_simd_forward, + "Tiled SIMD-group Metal int8-weight linear forward pass", + pybind11::arg("input"), + pybind11::arg("weight"), + pybind11::arg("scale")); +} From c0318dc58aa6f79eea74952a9df411e568124f59 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 17:50:35 -0500 Subject: [PATCH 16/17] Profile VAE decode and document M1 FP16 path --- README.md | 4 + modules/processing.py | 42 +++--- scripts/benchmark_postprocessing.py | 205 ++++++++++++++++++++++++++++ scripts/benchmark_vae_attention.py | 140 +++++++++++++++++++ scripts/benchmark_vae_fp16.py | 115 ++++++++++++++++ 5 files changed, 486 insertions(+), 20 deletions(-) create mode 100644 scripts/benchmark_postprocessing.py create mode 100644 scripts/benchmark_vae_attention.py create mode 100644 scripts/benchmark_vae_fp16.py diff --git a/README.md b/README.md index fcab35d041f..e6c0cf41512 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,10 @@ This branch adds targeted Apple Silicon inference optimizations while retaining Native paths perform isolated startup checks and fall back to PyTorch for unsupported shapes, dtypes, training, masks, runtime failures, and incompatible configurations. The regular safetensors format remains supported; no checkpoint conversion is required. +On Apple M1, the macOS launcher uses the FP16 VAE path by default. A controlled 512x512, five-step DPM++ SDE/Karras comparison measured approximately 9.16 seconds with FP16 VAE versus 10.37 seconds with `--no-half-vae`, an approximately 1.21 second or 11.7% median improvement. The measured comparison had no NaN fallback; decoded RGB pixels had a mean absolute difference of 0.0227, a maximum difference of 2, and 6.548% changed pixels. Small FP16 rounding differences are expected. + +Automatic1111 retains its existing VAE NaN detection and retry behavior: if FP16 decoding produces non-finite output, the VAE is converted to FP32 and decoding is retried. Use `--no-half-vae` to force the conservative FP32 path. The FP16 VAE path was smoke-tested with txt2img, img2img, inpainting, and Hires Fix; external VAE coverage requires an external VAE asset. + Changed Apple Silicon defaults include NGMS 1.0/all steps, Clip skip 2, FP16 sampling without the upstream sampling-upcast default, and the validated FP16 VAE route on M1-family Macs. These settings can change same-seed output compared with upstream defaults. They can be changed through the existing settings or local launch overrides. A recorded 16 GB M1 comparison at the same checkpoint hash and compute shape improved a five-step 384×640 DPM++ SDE/Karras request from 12.8 seconds to 8.7 seconds. The runs used different seeds, so this is a throughput observation rather than an image-parity comparison. diff --git a/modules/processing.py b/modules/processing.py index 462eaedb8cc..dc6d8a07142 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -16,7 +16,7 @@ from typing import Any import modules.sd_hijack -from modules import devices, prompt_parser, masking, sd_samplers, lowvram, infotext_utils, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors, rng, profiling, mps_stage_profile +from modules import devices, prompt_parser, masking, sd_samplers, lowvram, infotext_utils, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors, rng, profiling, mps_stage_profile from modules.rng import slerp # noqa: F401 from modules.sd_hijack import model_hijack from modules.sd_samplers_common import images_tensor_to_samples, decode_first_stage, approximation_indexes @@ -629,7 +629,8 @@ def decode_latent_batch(model, batch, target_device=None, check_for_nans=False): devices.test_for_nans(batch, "unet") for i in range(batch.shape[0]): - sample = decode_first_stage(model, batch[i:i + 1])[0] + with mps_stage_profile.stage("vae_decode"): + sample = decode_first_stage(model, batch[i:i + 1])[0] if check_for_nans: @@ -662,10 +663,12 @@ def decode_latent_batch(model, batch, target_device=None, check_for_nans=False): model.first_stage_model.to(devices.dtype_vae) batch = batch.to(devices.dtype_vae) - sample = decode_first_stage(model, batch[i:i + 1])[0] + with mps_stage_profile.stage("vae_decode"): + sample = decode_first_stage(model, batch[i:i + 1])[0] if target_device is not None: - sample = sample.to(target_device) + with mps_stage_profile.stage("vae_decode_transfer"): + sample = sample.to(target_device) samples.append(sample) @@ -843,9 +846,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed: # backwards compatibility, fix sampler and scheduler if invalid sd_samplers.fix_p_invalid_sampler_and_scheduler(p) - with mps_stage_profile.request(p): - with profiling.Profiler(): - res = process_images_inner(p) + with mps_stage_profile.request(p): + with profiling.Profiler(): + res = process_images_inner(p) finally: sd_models.apply_token_merging(p.sd_model, 0) @@ -869,8 +872,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: else: assert p.prompt is not None - with mps_stage_profile.stage("initial_gc"): - devices.torch_gc() + with mps_stage_profile.stage("initial_gc"): + devices.torch_gc() seed = get_fixed_seed(p.seed) subseed = get_fixed_seed(p.subseed) @@ -919,7 +922,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: infotexts = [] output_images = [] with torch.no_grad(), p.sd_model.ema_scope(): - with devices.autocast(), mps_stage_profile.stage("initialization"): + with devices.autocast(), mps_stage_profile.stage("initialization"): p.init(p.all_prompts, p.all_seeds, p.all_subseeds) # for OSX, loading the model during sampling changes the generated picture, so it is loaded here @@ -965,8 +968,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if p.scripts is not None: p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds) - with mps_stage_profile.stage("conditioning"): - p.setup_conds() + with mps_stage_profile.stage("conditioning"): + p.setup_conds() p.extra_generation_params.update(model_hijack.extra_generation_params) @@ -987,7 +990,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: sd_models.apply_alpha_schedule_override(p.sd_model, p) - with mps_stage_profile.stage("sampler"), devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(): + with mps_stage_profile.stage("sampler"), devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(): samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts) if p.scripts is not None: @@ -1002,20 +1005,19 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if opts.sd_vae_decode_method != 'Full': p.extra_generation_params['VAE Decoder'] = opts.sd_vae_decode_method - with mps_stage_profile.stage("vae_decode_and_transfer"): - x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True) + x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True) - with mps_stage_profile.stage("image_tensor_processing"): - x_samples_ddim = torch.stack(x_samples_ddim).float() - x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) + with mps_stage_profile.stage("image_tensor_processing"): + x_samples_ddim = torch.stack(x_samples_ddim).float() + x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) del samples_ddim if lowvram.is_enabled(shared.sd_model): lowvram.send_everything_to_cpu() - with mps_stage_profile.stage("post_decode_gc"): - devices.torch_gc() + with mps_stage_profile.stage("post_decode_gc"): + devices.torch_gc() state.nextjob() diff --git a/scripts/benchmark_postprocessing.py b/scripts/benchmark_postprocessing.py new file mode 100644 index 00000000000..71131a32914 --- /dev/null +++ b/scripts/benchmark_postprocessing.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +"""Profile postprocessing stages without changing the WebUI runtime.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import numpy as np +import torch +from PIL import Image + +_benchmark_parser = argparse.ArgumentParser(add_help=False) +_benchmark_parser.add_argument("--width", type=int, default=512) +_benchmark_parser.add_argument("--height", type=int, default=512) +_benchmark_parser.add_argument("--tile-size", type=int, default=256) +_benchmark_parser.add_argument("--tile-overlap", type=int, default=32) +_benchmark_parser.add_argument("--warmup", type=int, default=2) +_benchmark_parser.add_argument("--repeats", type=int, default=5) +_benchmark_parser.add_argument("--realesrgan-model", default="models/RealESRGAN/RealESRGAN_x4plus_anime_6B.pth") +_benchmark_parser.add_argument("--skip-realesrgan", action="store_true") +_benchmark_parser.add_argument("--vae-channels-last", action="store_true") +_benchmark_parser.add_argument("--vae-repeats", type=int, default=6) +_benchmark_parser.add_argument("--vae-warmup", type=int, default=2) +_benchmark_args, _webui_args = _benchmark_parser.parse_known_args() +sys.argv = [sys.argv[0], *_webui_args] + +repository_root = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(repository_root)) +sys.path.insert(0, str(repository_root / "repositories" / "k-diffusion")) +sys.path.insert(0, str(repository_root / "repositories" / "stable-diffusion-stability-ai")) +sys.path.insert(0, str(repository_root / "repositories" / "generative-models")) + +import webui # noqa: E402, F401 + +from modules import images, modelloader, upscaler_utils # noqa: E402 + + +def parse_args(): + return _benchmark_args + + +def synchronize(device): + if device.type == "mps": + torch.mps.synchronize() + elif device.type == "cuda": + torch.cuda.synchronize(device) + + +def measure(operation, device, warmup, repeats): + with torch.inference_mode(): + for _ in range(warmup): + operation() + synchronize(device) + timings = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + synchronize(device) + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings) + + +def make_image(width, height): + values = np.arange(width * height * 3, dtype=np.uint32).reshape(height, width, 3) + return Image.fromarray((values % 256).astype(np.uint8), "RGB") + + +def report(name, milliseconds): + print(f"{name:28s} {milliseconds:9.3f} ms") + + +def distribution(values): + values = sorted(values) + return { + "median": statistics.median(values), + "p25": values[max(0, len(values) // 4)], + "p75": values[min(len(values) - 1, (len(values) * 3) // 4)], + } + + +def run_vae_channels_last(args): + from modules import devices, initialize, sd_models, sd_samplers_common, shared + + initialize.initialize() + if shared.sd_model is None: + sd_models.reload_model_weights() + model = shared.sd_model + if model is None or model.first_stage_model is None: + raise RuntimeError("WebUI did not load a VAE model") + + latent = torch.randn((1, 4, args.height // 8, args.width // 8), device=devices.device, dtype=devices.dtype_vae) + decoder = model.first_stage_model + original_format = next(decoder.parameters()).data.stride() + + def decode(): + return sd_samplers_common.decode_first_stage(model, latent)[0] + + def set_layout(channels_last): + decoder.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format) + latent_layout = latent.contiguous(memory_format=torch.channels_last) if channels_last else latent.contiguous() + return latent_layout + + results = {"contiguous": [], "channels_last": []} + outputs = {} + for index in range(args.vae_warmup + args.vae_repeats * 2): + layout = "channels_last" if index % 2 else "contiguous" + latent_layout = set_layout(layout == "channels_last") + if index < args.vae_warmup: + with torch.inference_mode(): + sd_samplers_common.decode_first_stage(model, latent_layout) + continue + started = time.perf_counter() + with torch.inference_mode(): + output = sd_samplers_common.decode_first_stage(model, latent_layout)[0] + synchronize(devices.device) + elapsed = (time.perf_counter() - started) * 1000 + results[layout].append(elapsed) + outputs.setdefault(layout, output.detach().float().cpu()) + + reference = outputs["contiguous"] + difference = (outputs["channels_last"] - reference).abs() + parameter_layouts = { + "4d_channels_last": 0, + "4d_contiguous": 0, + "4d_other": 0, + "non4d_contiguous": 0, + "non4d_total": 0, + } + for parameter in decoder.parameters(): + memory_format = parameter.data + if memory_format.ndim == 4 and memory_format.is_contiguous(memory_format=torch.channels_last): + parameter_layouts["4d_channels_last"] += 1 + elif memory_format.ndim == 4 and memory_format.is_contiguous(): + parameter_layouts["4d_contiguous"] += 1 + elif memory_format.ndim == 4: + parameter_layouts["4d_other"] += 1 + else: + parameter_layouts["non4d_contiguous"] += int(memory_format.is_contiguous()) + parameter_layouts["non4d_total"] += 1 + + print("VAE channels-last A/B") + print(f"original_first_parameter_stride={original_format}") + print(f"parameter_layouts={parameter_layouts}") + print(f"input_channels_last={latent.contiguous(memory_format=torch.channels_last).is_contiguous(memory_format=torch.channels_last)}") + for layout, values in results.items(): + print(f"{layout}={distribution(values)}") + print(f"max_pixel_difference={difference.max().item():.6f} mean_pixel_difference={difference.mean().item():.6f} exact={torch.equal(outputs['channels_last'], reference)}") + print(f"nan_contiguous={not torch.isfinite(outputs['contiguous']).all().item()} nan_channels_last={not torch.isfinite(outputs['channels_last']).all().item()}") + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + if args.vae_channels_last: + run_vae_channels_last(args) + return + + device = torch.device("mps") + image = make_image(args.width, args.height) + print(f"PyTorch {torch.__version__}; postprocessing profile; image={args.width}x{args.height}; device={device}") + print("Warm medians; model loading and first-time compilation are excluded from timings.") + + cpu_tensor = lambda: upscaler_utils.pil_image_to_torch_bgr(image).unsqueeze(0) + cpu_image_tensor = cpu_tensor() + conversion_to_mps = lambda: cpu_image_tensor.to(device=device, dtype=torch.float16) + tensor = conversion_to_mps() + conversion_to_pil = lambda: upscaler_utils.torch_bgr_to_pil_image(tensor) + report("PIL -> CPU tensor", measure(cpu_tensor, torch.device("cpu"), args.warmup, args.repeats)) + report("CPU tensor -> MPS", measure(conversion_to_mps, device, args.warmup, args.repeats)) + report("MPS tensor -> CPU/PIL", measure(conversion_to_pil, device, args.warmup, args.repeats)) + + identity_model = torch.nn.Identity().to(device) + gpu_tile = lambda: upscaler_utils.tiled_upscale_2(tensor, identity_model, tile_size=args.tile_size, tile_overlap=args.tile_overlap, scale=1, device=device, desc="profile") + report("GPU tile compose", measure(gpu_tile, device, args.warmup, args.repeats)) + + grid = images.split_grid(image, args.tile_size, args.tile_size, args.tile_overlap) + pil_tiles = [tile for _, _, row in grid.tiles for _, _, tile in row] + pil_tile_process = lambda: [tile.copy() for tile in pil_tiles] + report("PIL tile traversal/copy", measure(pil_tile_process, torch.device("cpu"), args.warmup, args.repeats)) + report("PIL tile composition", measure(lambda: images.combine_grid(grid), torch.device("cpu"), args.warmup, args.repeats)) + + if args.skip_realesrgan: + print("RealESRGAN inference skipped") + else: + model_path = pathlib.Path(args.realesrgan_model) + if not model_path.exists(): + print(f"RealESRGAN inference unavailable ({model_path})") + else: + model = modelloader.load_spandrel_model(str(model_path), device=device, prefer_half=True, expected_architecture="ESRGAN") + model_input = tensor + report("RealESRGAN full inference", measure(lambda: model(model_input), device, args.warmup, args.repeats)) + tiled_model = lambda: upscaler_utils.tiled_upscale_2(model_input, model, tile_size=args.tile_size, tile_overlap=args.tile_overlap, scale=4, device=device, desc="profile") + report("RealESRGAN GPU tiled inference", measure(tiled_model, device, args.warmup, args.repeats)) + + print("VAE decode unavailable: run this profile with a loaded WebUI VAE stage") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_vae_attention.py b/scripts/benchmark_vae_attention.py new file mode 100644 index 00000000000..208f3cbd7db --- /dev/null +++ b/scripts/benchmark_vae_attention.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Benchmark the SD1 VAE single-head attention path without changing routing.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import torch +import torch.nn.functional as F + +_benchmark_parser = argparse.ArgumentParser(add_help=False) +_benchmark_parser.add_argument("--warmup", type=int, default=2) +_benchmark_parser.add_argument("--repeats", type=int, default=5) +_benchmark_parser.add_argument("--decode-repeats", type=int, default=2) +_benchmark_args, _webui_args = _benchmark_parser.parse_known_args() +sys.argv = [sys.argv[0], *_webui_args] + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "repositories" / "k-diffusion")) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "repositories" / "stable-diffusion-stability-ai")) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "repositories" / "generative-models")) + +import webui # noqa: E402, F401 + +from modules import devices, initialize, mps_flash_attention, sd_models, sd_samplers_common, shared # noqa: E402 + + +def parse_args(): + return _benchmark_args + + +def synchronize(): + torch.mps.synchronize() + + +def measure(operation, warmup, repeats): + with torch.inference_mode(): + for _ in range(warmup): + operation() + synchronize() + values = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + synchronize() + values.append((time.perf_counter() - started) * 1000) + return { + "median": statistics.median(values), + "p25": sorted(values)[max(0, len(values) // 4)], + "p75": sorted(values)[min(len(values) - 1, (len(values) * 3) // 4)], + } + + +def attention_path(norm, q_conv, k_conv, v_conv, output_conv, source, mode): + normalized = norm(source) + q = q_conv(normalized) + k = k_conv(normalized) + v = v_conv(normalized) + batch, channels, height, width = q.shape + tokens = height * width + q = q.reshape(batch, channels, tokens).transpose(1, 2).unsqueeze(1) + k = k.reshape(batch, channels, tokens).transpose(1, 2).unsqueeze(1) + v = v.reshape(batch, channels, tokens).transpose(1, 2).unsqueeze(1) + if mode == "fp16": + q, k, v = q.half(), k.half(), v.half() + attended = F.scaled_dot_product_attention(q, k, v).float() + elif mode == "mfa": + attended = mps_flash_attention._extension.MetalFlashAttentionForward.apply(q.half(), k.half(), v.half(), channels**-0.5, False).float() + else: + attended = F.scaled_dot_product_attention(q, k, v) + attended = attended.squeeze(1).transpose(1, 2).reshape(batch, channels, height, width) + return source + output_conv(attended.to(dtype=source.dtype)) + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + initialize.initialize() + if shared.sd_model is None: + sd_models.reload_model_weights() + model = shared.sd_model + decoder = model.first_stage_model + attention = decoder.decoder.mid.attn_1 + source = torch.randn((1, attention.in_channels, 64, 64), device=devices.device, dtype=devices.dtype_vae) + latent = torch.randn((1, 4, 64, 64), device=devices.device, dtype=devices.dtype_vae) + print(f"PyTorch {torch.__version__}; VAE attention channels={attention.in_channels}; Q/K/V=[1,1,4096,{attention.in_channels}]") + + paths = { + "fp32": lambda: attention_path(attention.norm, attention.q, attention.k, attention.v, attention.proj_out, source, "fp32"), + "fp16": lambda: attention_path(attention.norm, attention.q, attention.k, attention.v, attention.proj_out, source, "fp16"), + "mfa": lambda: attention_path(attention.norm, attention.q, attention.k, attention.v, attention.proj_out, source, "mfa"), + } + outputs = {} + for name, operation in paths.items(): + try: + result = measure(operation, args.warmup, args.repeats) + print(f"{name} attention={result}") + with torch.inference_mode(): + outputs[name] = operation().float() + synchronize() + except Exception as exc: + print(f"{name} unavailable={type(exc).__name__}: {exc}") + + if "fp32" in outputs and "fp16" in outputs: + difference = (outputs["fp32"] - outputs["fp16"]).abs() + print(f"fp16_attention_error max={difference.max().item():.6f} mean={difference.mean().item():.6f} exact={torch.equal(outputs['fp32'], outputs['fp16'])}") + if "fp32" in outputs and "mfa" in outputs: + difference = (outputs["fp32"] - outputs["mfa"]).abs() + print(f"mfa_attention_error max={difference.max().item():.6f} mean={difference.mean().item():.6f} exact={torch.equal(outputs['fp32'], outputs['mfa'])}") + + original_forward = attention.forward + decode_outputs = {} + for name in ("fp32", "fp16", "mfa"): + attention.forward = lambda value, mode=name: attention_path( + attention.norm, attention.q, attention.k, attention.v, attention.proj_out, value, mode + ) + try: + with torch.inference_mode(): + decoded = sd_samplers_common.decode_first_stage(model, latent)[0].float().cpu() + decode_outputs[name] = decoded + print(f"full_decode_{name}=ok nan={not torch.isfinite(decoded).all().item()}") + except Exception as exc: + print(f"full_decode_{name}=unavailable {type(exc).__name__}: {exc}") + attention.forward = original_forward + if "fp32" in decode_outputs: + for name in ("fp16", "mfa"): + if name in decode_outputs: + difference = (decode_outputs[name] - decode_outputs["fp32"]).abs() + print(f"full_decode_{name}_error max={difference.max().item():.6f} mean={difference.mean().item():.6f} exact={torch.equal(decode_outputs[name], decode_outputs['fp32'])}") + print(f"decoder_parameters={sum(parameter.numel() for parameter in decoder.parameters())}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_vae_fp16.py b/scripts/benchmark_vae_fp16.py new file mode 100644 index 00000000000..62b6734ec3f --- /dev/null +++ b/scripts/benchmark_vae_fp16.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Benchmark full FP16 SD1 VAE decode against the loaded FP32 decoder.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import torch + +_parser = argparse.ArgumentParser(add_help=False) +_parser.add_argument("--warmup", type=int, default=2) +_parser.add_argument("--repeats", type=int, default=6) +_benchmark_args, _webui_args = _parser.parse_known_args() +sys.argv = [sys.argv[0], *_webui_args] + +repository_root = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(repository_root)) +sys.path.insert(0, str(repository_root / "repositories" / "k-diffusion")) +sys.path.insert(0, str(repository_root / "repositories" / "stable-diffusion-stability-ai")) +sys.path.insert(0, str(repository_root / "repositories" / "generative-models")) + +import webui # noqa: E402, F401 + +from modules import devices, initialize, mps_fused_ops, sd_models, shared # noqa: E402 + + +def synchronize(): + torch.mps.synchronize() + + +def distribution(values): + values = sorted(values) + return { + "median": round(statistics.median(values), 3), + "p25": round(values[max(0, len(values) // 4)], 3), + "p75": round(values[min(len(values) - 1, (len(values) * 3) // 4)], 3), + } + + +def as_uint8(image): + image = ((image.float() + 1.0) / 2.0).clamp(0.0, 1.0) + return (image * 255.0).round().to(torch.uint8) + + +def parameter_bytes(module): + return sum(parameter.numel() * parameter.element_size() for parameter in module.parameters()) + + +def decode_with_timing(model, latent, warmup, repeats): + def decode(): + return model.decode_first_stage(latent) + + with torch.inference_mode(): + for _ in range(warmup): + decode() + synchronize() + timings = [] + output = None + for _ in range(repeats): + started = time.perf_counter() + output = decode() + synchronize() + timings.append((time.perf_counter() - started) * 1000) + return distribution(timings), output.detach().float().cpu() + + +def main(): + args = _benchmark_args + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + initialize.initialize() + if shared.sd_model is None: + sd_models.reload_model_weights() + model = shared.sd_model + decoder = model.first_stage_model + fp32_state = {name: value.detach().cpu().clone() for name, value in decoder.state_dict().items()} + latent_fp32 = torch.randn((1, 4, 64, 64), device=devices.device, dtype=torch.float32) + baseline_dispatches = mps_fused_ops.diagnostics().copy() + + decoder.float() + decoder.load_state_dict(fp32_state) + latent_fp32 = latent_fp32.to(dtype=torch.float32) + fp32_timing, fp32_output = decode_with_timing(model, latent_fp32, args.warmup, args.repeats) + fp32_memory = parameter_bytes(decoder) + fp32_dispatches = mps_fused_ops.diagnostics().copy() + + decoder.half() + latent_fp16 = latent_fp32.half() + fp16_timing, fp16_output = decode_with_timing(model, latent_fp16, args.warmup, args.repeats) + fp16_memory = parameter_bytes(decoder) + fp16_dispatches = mps_fused_ops.diagnostics().copy() + + float_difference = (fp16_output - fp32_output).abs() + uint8_difference = (as_uint8(fp16_output).to(torch.int16) - as_uint8(fp32_output).to(torch.int16)).abs() + print(f"PyTorch {torch.__version__}; full FP16 VAE decode; latent={tuple(latent_fp32.shape)}") + print(f"fp32={fp32_timing} fp16={fp16_timing}") + print(f"median_gain_ms={fp32_timing['median'] - fp16_timing['median']:.3f}") + print(f"fp32_parameter_bytes={fp32_memory} fp16_parameter_bytes={fp16_memory} reduction={(1 - fp16_memory / fp32_memory) * 100:.1f}%") + print(f"float_error_max={float_difference.max().item():.6f} float_error_mean={float_difference.mean().item():.6f}") + print(f"uint8_mae={uint8_difference.float().mean().item():.6f} uint8_max={uint8_difference.max().item()} changed_pixels={(uint8_difference > 0).float().mean().item() * 100:.3f}%") + print(f"nan_fp32={not torch.isfinite(fp32_output).all().item()} nan_fp16={not torch.isfinite(fp16_output).all().item()}") + print(f"groupnorm_dispatch_delta_fp32={fp32_dispatches['dispatches'] - baseline_dispatches['dispatches']} embedding_dispatch_delta_fp32={fp32_dispatches['embedding_dispatches'] - baseline_dispatches['embedding_dispatches']}") + print(f"groupnorm_dispatch_delta_fp16={fp16_dispatches['dispatches'] - fp32_dispatches['dispatches']} embedding_dispatch_delta_fp16={fp16_dispatches['embedding_dispatches'] - fp32_dispatches['embedding_dispatches']}") + + decoder.float() + decoder.load_state_dict(fp32_state) + + +if __name__ == "__main__": + main() From 8ed6ac96874a53a87c76ba1822c6b1f837a33382 Mon Sep 17 00:00:00 2001 From: Derek Anderson Date: Wed, 12 Aug 2026 18:22:47 -0500 Subject: [PATCH 17/17] Add benchmarks for FP16 SD1 ResBlock runner and native second-half operation Signed-off-by: Derek Anderson --- scripts/benchmark_mps_resblock_runner.py | 139 +++++++++++ scripts/benchmark_mps_resblock_second_half.py | 91 +++++++ scripts/mps_fused_resblock.mm | 226 ++++++++++++++++++ 3 files changed, 456 insertions(+) create mode 100644 scripts/benchmark_mps_resblock_runner.py create mode 100644 scripts/benchmark_mps_resblock_second_half.py create mode 100644 scripts/mps_fused_resblock.mm diff --git a/scripts/benchmark_mps_resblock_runner.py b/scripts/benchmark_mps_resblock_runner.py new file mode 100644 index 00000000000..db2382104fc --- /dev/null +++ b/scripts/benchmark_mps_resblock_runner.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +"""Benchmark an isolated FP16 SD1 ResBlock runner against the existing fused path.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +from modules import mps_fused_ops + + +SD1_SHAPES = ( + (4096, 320), + (1024, 640), + (256, 1280), +) +EMBEDDING_CHANNELS = 1280 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--warmup", type=int, default=4) + parser.add_argument("--repeats", type=int, default=12) + return parser.parse_args() + + +def measure(operation, warmup, repeats): + with torch.inference_mode(): + for _ in range(warmup): + operation() + torch.mps.synchronize() + + timings = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + torch.mps.synchronize() + timings.append((time.perf_counter() - started) * 1000) + return statistics.median(timings) + + +class ReferenceResBlock(torch.nn.Module): + """Module-style equivalent of the existing fused SD1 inference path.""" + + def __init__(self, channels): + super().__init__() + self.in_norm = torch.nn.GroupNorm(32, channels) + self.in_conv = torch.nn.Conv2d(channels, channels, 3, padding=1) + self.emb_layers = torch.nn.Sequential( + torch.nn.SiLU(), + torch.nn.Linear(EMBEDDING_CHANNELS, channels), + ) + self.out_norm = torch.nn.GroupNorm(32, channels) + self.out_conv = torch.nn.Conv2d(channels, channels, 3, padding=1) + + def forward(self, source, embedding): + hidden = self.in_conv(mps_fused_ops.group_norm_silu(source, self.in_norm)) + embedding = self.emb_layers(embedding).to(dtype=hidden.dtype) + hidden = mps_fused_ops.group_norm_silu_add_embedding(hidden, embedding, self.out_norm) + return source + self.out_conv(hidden) + + +class ParameterBoundResBlockRunner: + """First runner boundary: bind tensors once, retain the same fused operators.""" + + def __init__(self, reference): + self.in_norm = reference.in_norm + self.in_conv_weight = reference.in_conv.weight + self.in_conv_bias = reference.in_conv.bias + self.emb_activation = reference.emb_layers[0] + self.emb_weight = reference.emb_layers[1].weight + self.emb_bias = reference.emb_layers[1].bias + self.out_norm = reference.out_norm + self.out_conv_weight = reference.out_conv.weight + self.out_conv_bias = reference.out_conv.bias + + def __call__(self, source, embedding): + hidden = mps_fused_ops.group_norm_silu(source, self.in_norm) + hidden = F.conv2d(hidden, self.in_conv_weight, self.in_conv_bias, padding=1) + embedding = F.silu(embedding) + embedding = F.linear(embedding, self.emb_weight, self.emb_bias).to(dtype=hidden.dtype) + hidden = mps_fused_ops.group_norm_silu_add_embedding(hidden, embedding, self.out_norm) + hidden = F.conv2d(hidden, self.out_conv_weight, self.out_conv_bias, padding=1) + return source + hidden + + +def benchmark_shape(batch, tokens, channels, warmup, repeats): + side = int(tokens**0.5) + reference = ReferenceResBlock(channels).eval().half().to("mps") + runner = ParameterBoundResBlockRunner(reference) + source = torch.randn((batch, channels, side, side), device="mps", dtype=torch.float16) + embedding = torch.randn((batch, EMBEDDING_CHANNELS), device="mps", dtype=torch.float16) + + with torch.inference_mode(): + expected = reference(source, embedding) + 0 + actual = runner(source, embedding) + 0 + torch.mps.synchronize() + + difference = (actual.float() - expected.float()).abs() + return { + "reference_ms": measure(lambda: reference(source, embedding), warmup, repeats), + "runner_ms": measure(lambda: runner(source, embedding), warmup, repeats), + "max_error": difference.max().item(), + "mean_error": difference.mean().item(), + "exact": torch.equal(actual, expected), + } + + +def main(): + args = parse_args() + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + + torch.manual_seed(1) + print(f"PyTorch {torch.__version__}; FP16 MPS ResBlock runner; batch={args.batch}") + print("Benchmark only: both paths use the existing fused GroupNorm/SiLU kernels.") + for tokens, channels in SD1_SHAPES: + result = benchmark_shape(args.batch, tokens, channels, args.warmup, args.repeats) + speedup = result["reference_ms"] / result["runner_ms"] + print( + f"tokens={tokens:4d} channels={channels:4d} " + f"reference={result['reference_ms']:.3f}ms runner={result['runner_ms']:.3f}ms " + f"speedup={speedup:.3f}x max_error={result['max_error']:.6f} " + f"mean_error={result['mean_error']:.6f} exact={result['exact']}" + ) + torch.mps.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_mps_resblock_second_half.py b/scripts/benchmark_mps_resblock_second_half.py new file mode 100644 index 00000000000..aef872e2820 --- /dev/null +++ b/scripts/benchmark_mps_resblock_second_half.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Benchmark the isolated native SD1 ResBlock second-half operation.""" + +from __future__ import annotations + +import argparse +import pathlib +import statistics +import sys +import time + +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +from modules import mps_fused_ops + + +SD1_SHAPES = ((4096, 320), (1024, 640), (256, 1280)) + + +def measure(operation, warmup, repeats): + with torch.inference_mode(): + for _ in range(warmup): + operation() + torch.mps.synchronize() + values = [] + for _ in range(repeats): + started = time.perf_counter() + operation() + torch.mps.synchronize() + values.append((time.perf_counter() - started) * 1000) + return statistics.median(values) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--warmup", type=int, default=4) + parser.add_argument("--repeats", type=int, default=12) + args = parser.parse_args() + + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available in this PyTorch installation.") + import metal_flash_sdpa + + if not getattr(metal_flash_sdpa, "A1111_MPS_FUSED_RESBLOCK", False): + raise SystemExit("The installed Metal extension does not contain the isolated ResBlock operation.") + + torch.manual_seed(1) + print(f"PyTorch {torch.__version__}; FP16 MPS ResBlock second half; batch={args.batch}") + print("Reference: existing fused GroupNorm+SiLU+embedding, PyTorch convolution, residual add.") + print("Native: same command buffer, Metal statistics, convolution, and residual output.") + + for tokens, channels in SD1_SHAPES: + side = int(tokens**0.5) + hidden = torch.randn((args.batch, channels, side, side), device="mps", dtype=torch.float16) + embedding = torch.randn((args.batch, channels), device="mps", dtype=torch.float16) + norm = torch.nn.GroupNorm(32, channels).eval().half().to("mps") + conv = torch.nn.Conv2d(channels, channels, 3, padding=1).eval().half().to("mps") + residual = torch.randn_like(hidden) + + def reference(hidden=hidden, embedding=embedding, norm=norm, residual=residual, conv=conv): + normalized = mps_fused_ops.group_norm_silu_add_embedding(hidden, embedding, norm) + return residual + F.conv2d(normalized, conv.weight, conv.bias, padding=1) + + def native(hidden=hidden, embedding=embedding, norm=norm, residual=residual, conv=conv): + return metal_flash_sdpa.fused_resblock_forward( + hidden, embedding, norm.weight, norm.bias, + conv.weight, conv.bias, residual, 32, norm.eps, + ) + + with torch.inference_mode(): + expected = reference() + 0 + actual = native() + 0 + torch.mps.synchronize() + difference = (actual.float() - expected.float()).abs() + reference_ms = measure(reference, args.warmup, args.repeats) + native_ms = measure(native, args.warmup, args.repeats) + print( + f"tokens={tokens:4d} channels={channels:4d} reference={reference_ms:.3f}ms " + f"native={native_ms:.3f}ms speedup={reference_ms / native_ms:.3f}x " + f"max_error={difference.max().item():.6f} mean_error={difference.mean().item():.6f} " + f"exact={torch.equal(actual, expected)}" + ) + torch.mps.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/scripts/mps_fused_resblock.mm b/scripts/mps_fused_resblock.mm new file mode 100644 index 00000000000..e36a38702d1 --- /dev/null +++ b/scripts/mps_fused_resblock.mm @@ -0,0 +1,226 @@ +#include +#include +#include +#include + +#import +#import + +#include + +namespace { + +struct ResBlockParams { + uint32_t batch; + uint32_t channels; + uint32_t height; + uint32_t width; + uint32_t groups; + float epsilon; +}; + +struct ResBlockPipelines { + id stats; + id block; +}; + +static inline id buffer_storage(const at::Tensor& tensor) { + return __builtin_bit_cast(id, tensor.storage().data()); +} + +static inline size_t buffer_offset(const at::Tensor& tensor) { + return tensor.storage_offset() * tensor.element_size(); +} + +static ResBlockPipelines get_pipelines() { + static ResBlockPipelines pipelines; + static dispatch_once_t once; + dispatch_once(&once, ^{ + id device = at::mps::MPSDevice::getInstance()->device(); + NSString* source = @R"METAL( +#include +using namespace metal; + +struct ResBlockParams { + uint batch; + uint channels; + uint height; + uint width; + uint groups; + float epsilon; +}; + +kernel void resblock_stats( + device const half* input [[buffer(0)]], + device const half* embedding [[buffer(1)]], + device float* stats [[buffer(2)]], + constant ResBlockParams& params [[buffer(3)]], + uint tid [[thread_index_in_threadgroup]], + uint group_index [[threadgroup_position_in_grid]], + uint threads [[threads_per_threadgroup]]) { + threadgroup float sums[256]; + threadgroup float squares[256]; + const uint spatial = params.height * params.width; + const uint channels_per_group = params.channels / params.groups; + const uint elements = channels_per_group * spatial; + const uint batch_index = group_index / params.groups; + const uint group = group_index % params.groups; + const uint base = (batch_index * params.channels + group * channels_per_group) * spatial; + float sum = 0.0f; + float square = 0.0f; + for (uint index = tid; index < elements; index += threads) { + const uint channel = group * channels_per_group + index / spatial; + const half value = half(input[base + index] + embedding[batch_index * params.channels + channel]); + sum += float(value); + square += float(value) * float(value); + } + sums[tid] = sum; + squares[tid] = square; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = threads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + sums[tid] += sums[tid + stride]; + squares[tid] += squares[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float mean = sums[0] / float(elements); + const float variance = max(squares[0] / float(elements) - mean * mean, 0.0f); + stats[group_index * 2] = mean; + stats[group_index * 2 + 1] = rsqrt(variance + params.epsilon); +} + +kernel void resblock_conv( + device const half* input [[buffer(0)]], + device const half* embedding [[buffer(1)]], + device const float* stats [[buffer(2)]], + device const half* norm_weight [[buffer(3)]], + device const half* norm_bias [[buffer(4)]], + device const half* conv_weight [[buffer(5)]], + device const half* conv_bias [[buffer(6)]], + device const half* residual [[buffer(7)]], + device half* output [[buffer(8)]], + constant ResBlockParams& params [[buffer(9)]], + uint index [[thread_position_in_grid]]) { + const uint spatial = params.height * params.width; + const uint count = params.batch * params.channels * spatial; + if (index >= count) return; + const uint pixel = index % spatial; + const uint x = pixel % params.width; + const uint y = pixel / params.width; + const uint output_channel = (index / spatial) % params.channels; + const uint batch_index = index / (params.channels * spatial); + const uint channels_per_group = params.channels / params.groups; + float result = float(conv_bias[output_channel]); + for (int ky = -1; ky <= 1; ++ky) { + for (int kx = -1; kx <= 1; ++kx) { + const int source_y = int(y) + ky; + const int source_x = int(x) + kx; + if (source_y < 0 || source_y >= int(params.height) || source_x < 0 || source_x >= int(params.width)) continue; + const uint kernel_index = uint((ky + 1) * 3 + kx + 1); + for (uint input_channel = 0; input_channel < params.channels; ++input_channel) { + const uint source_index = (batch_index * params.channels + input_channel) * spatial + uint(source_y) * params.width + uint(source_x); + const uint group = input_channel / channels_per_group; + const float mean = stats[(batch_index * params.groups + group) * 2]; + const float inverse = stats[(batch_index * params.groups + group) * 2 + 1]; + const half combined = half(input[source_index] + embedding[batch_index * params.channels + input_channel]); + float value = (float(combined) - mean) * inverse; + value = value * float(norm_weight[input_channel]) + float(norm_bias[input_channel]); + value = value / (1.0f + exp(-value)); + const uint weight_index = ((output_channel * params.channels + input_channel) * 9) + kernel_index; + result += value * float(conv_weight[weight_index]); + } + } + } + output[index] = half(result + float(residual[index])); +} +)METAL"; + NSError* error = nil; + id library = [device newLibraryWithSource:source options:nil error:&error]; + TORCH_CHECK(library != nil, "Failed to compile ResBlock Metal library: ", error ? [[error localizedDescription] UTF8String] : "unknown error"); + id stats_function = [library newFunctionWithName:@"resblock_stats"]; + id block_function = [library newFunctionWithName:@"resblock_conv"]; + TORCH_CHECK(stats_function != nil && block_function != nil, "ResBlock Metal functions were not found"); + pipelines.stats = [device newComputePipelineStateWithFunction:stats_function error:&error]; + pipelines.block = [device newComputePipelineStateWithFunction:block_function error:&error]; + TORCH_CHECK(pipelines.stats != nil && pipelines.block != nil, "Failed to create ResBlock Metal pipelines"); + }); + return pipelines; +} + +} // namespace + +torch::Tensor fused_resblock_forward( + const torch::Tensor& input, + const torch::Tensor& embedding, + const torch::Tensor& norm_weight, + const torch::Tensor& norm_bias, + const torch::Tensor& conv_weight, + const torch::Tensor& conv_bias, + const torch::Tensor& residual, + int64_t groups, + double epsilon) { + TORCH_CHECK(input.device().is_mps() && embedding.device().is_mps() && residual.device().is_mps(), "activations must be MPS tensors"); + TORCH_CHECK(norm_weight.device().is_mps() && norm_bias.device().is_mps() && conv_weight.device().is_mps() && conv_bias.device().is_mps(), "weights must be MPS tensors"); + TORCH_CHECK(input.scalar_type() == at::kHalf && embedding.scalar_type() == at::kHalf && residual.scalar_type() == at::kHalf, "activations must be float16"); + TORCH_CHECK(norm_weight.scalar_type() == at::kHalf && norm_bias.scalar_type() == at::kHalf && conv_weight.scalar_type() == at::kHalf && conv_bias.scalar_type() == at::kHalf, "weights must be float16"); + TORCH_CHECK(input.dim() == 4 && embedding.dim() == 2 && residual.sizes() == input.sizes(), "invalid activation shapes"); + TORCH_CHECK(input.is_contiguous() && embedding.is_contiguous() && residual.is_contiguous(), "activations must be contiguous"); + TORCH_CHECK(norm_weight.is_contiguous() && norm_bias.is_contiguous() && conv_weight.is_contiguous() && conv_bias.is_contiguous(), "weights must be contiguous"); + TORCH_CHECK(embedding.size(0) == input.size(0) && embedding.size(1) == input.size(1), "embedding must be [batch, channels]"); + TORCH_CHECK(norm_weight.numel() == input.size(1) && norm_bias.numel() == input.size(1), "normalization weights must match channels"); + TORCH_CHECK(conv_weight.dim() == 4 && conv_weight.size(0) == input.size(1) && conv_weight.size(1) == input.size(1) && conv_weight.size(2) == 3 && conv_weight.size(3) == 3 && conv_bias.numel() == input.size(1), "convolution weights must be [channels, channels, 3, 3]"); + TORCH_CHECK(groups > 0 && input.size(1) % groups == 0, "channels must be divisible by groups"); + + auto output = torch::empty_like(input); + auto stats = torch::empty({input.size(0), groups, 2}, input.options().dtype(torch::kFloat)); + ResBlockParams params = { + static_cast(input.size(0)), static_cast(input.size(1)), + static_cast(input.size(2)), static_cast(input.size(3)), + static_cast(groups), static_cast(epsilon), + }; + const auto pipelines = get_pipelines(); + @autoreleasepool { + dispatch_sync(torch::mps::get_dispatch_queue(), ^{ + @autoreleasepool { + at::mps::getCurrentMPSStream()->endKernelCoalescing(); + id command_buffer = torch::mps::get_command_buffer(); + id stats_encoder = [command_buffer computeCommandEncoder]; + [stats_encoder setComputePipelineState:pipelines.stats]; + [stats_encoder setBuffer:buffer_storage(input) offset:buffer_offset(input) atIndex:0]; + [stats_encoder setBuffer:buffer_storage(embedding) offset:buffer_offset(embedding) atIndex:1]; + [stats_encoder setBuffer:buffer_storage(stats) offset:buffer_offset(stats) atIndex:2]; + [stats_encoder setBytes:¶ms length:sizeof(params) atIndex:3]; + [stats_encoder dispatchThreadgroups:MTLSizeMake(params.batch * params.groups, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [stats_encoder endEncoding]; + id block_encoder = [command_buffer computeCommandEncoder]; + [block_encoder setComputePipelineState:pipelines.block]; + [block_encoder setBuffer:buffer_storage(input) offset:buffer_offset(input) atIndex:0]; + [block_encoder setBuffer:buffer_storage(embedding) offset:buffer_offset(embedding) atIndex:1]; + [block_encoder setBuffer:buffer_storage(stats) offset:buffer_offset(stats) atIndex:2]; + [block_encoder setBuffer:buffer_storage(norm_weight) offset:buffer_offset(norm_weight) atIndex:3]; + [block_encoder setBuffer:buffer_storage(norm_bias) offset:buffer_offset(norm_bias) atIndex:4]; + [block_encoder setBuffer:buffer_storage(conv_weight) offset:buffer_offset(conv_weight) atIndex:5]; + [block_encoder setBuffer:buffer_storage(conv_bias) offset:buffer_offset(conv_bias) atIndex:6]; + [block_encoder setBuffer:buffer_storage(residual) offset:buffer_offset(residual) atIndex:7]; + [block_encoder setBuffer:buffer_storage(output) offset:buffer_offset(output) atIndex:8]; + [block_encoder setBytes:¶ms length:sizeof(params) atIndex:9]; + const NSUInteger count = params.batch * params.channels * params.height * params.width; + const NSUInteger threads = std::min(pipelines.block.maxTotalThreadsPerThreadgroup, 256); + [block_encoder dispatchThreads:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [block_encoder endEncoding]; + } + }); + } + return output; +} + +void register_resblock_op(pybind11::module_& module) { + module.def( + "fused_resblock_forward", &fused_resblock_forward, + "Fused experimental ResBlock second half", + pybind11::arg("input"), pybind11::arg("embedding"), + pybind11::arg("norm_weight"), pybind11::arg("norm_bias"), + pybind11::arg("conv_weight"), pybind11::arg("conv_bias"), + pybind11::arg("residual"), pybind11::arg("groups"), pybind11::arg("epsilon")); +} \ No newline at end of file