diff --git a/README.md b/README.md index a93079fd19b..e6c0cf41512 100644 --- a/README.md +++ b/README.md @@ -3,6 +3,28 @@ 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. + +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. + ## Features [Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features): - Original txt2img and img2img modes diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 804b802057a..af3ee590495 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,39 @@ 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_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}"', + "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/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/modules/mps_flash_attention.py b/modules/mps_flash_attention.py new file mode 100644 index 00000000000..9a9ea946c6e --- /dev/null +++ b/modules/mps_flash_attention.py @@ -0,0 +1,221 @@ +"""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 < 192: + 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(): + code = """ +import torch +import torch.nn.functional as F +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) +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.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 + +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() +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" + 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_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") + 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() + 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 self-test passed; deferred MFA, fused GroupNorm+SiLU, and fused GEGLU 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_fused_ops.py b/modules/mps_fused_ops.py new file mode 100644 index 00000000000..78c4dd95b8c --- /dev/null +++ b/modules/mps_fused_ops.py @@ -0,0 +1,207 @@ +"""Measured native Metal fusions for Apple Silicon inference.""" + +from __future__ import annotations + +import os + +import numpy as np +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 +_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 +_geglu_first_dispatch_logged = False +_geglu_runtime_disabled = False +_geglu_lut = None + + +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 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 + 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, + "embedding_dispatches": _embedding_dispatch_count, + "embedding_fallbacks": _embedding_fallback_count, + "geglu_dispatches": _geglu_dispatch_count, + "geglu_fallbacks": _geglu_fallback_count, + } 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/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/processing.py b/modules/processing.py index 92c3582cc66..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 +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,8 +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 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 +872,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 +922,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 +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) - p.setup_conds() + with mps_stage_profile.stage("conditioning"): + p.setup_conds() p.extra_generation_params.update(model_hijack.extra_generation_params) @@ -984,7 +990,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: @@ -1001,15 +1007,17 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: 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) - 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_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/sd_hijack_unet.py b/modules/sd_hijack_unet.py index b4f03b138a4..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 +from modules import devices, mps_fused_ops from modules.sd_hijack_utils import CondFunc @@ -36,6 +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) + + # 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 +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.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/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/modules/shared_options.py b/modules/shared_options.py index 03632ecc050..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(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"), @@ -232,8 +232,11 @@ 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"), + "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"), "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/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/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/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/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/benchmark_mps_unet_ops.py b/scripts/benchmark_mps_unet_ops.py new file mode 100644 index 00000000000..27b30e9d675 --- /dev/null +++ b/scripts/benchmark_mps_unet_ops.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +"""Benchmark the main SD 1.x UNet operation shapes on Apple MPS.""" + +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 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 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) + 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) + + return { + "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, + ), + "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, + repeats, + ), + "sdpa": measure( + lambda: F.scaled_dot_product_attention(query, query, query, dropout_p=0.0), + warmup, + repeats, + ), + } + + +def benchmark_shape(batch, tokens, channels, warmup, repeats): + results = measure_shape(batch, tokens, channels, warmup, repeats) + 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() 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() diff --git a/scripts/install_mps_flash_attention.py b/scripts/install_mps_flash_attention.py new file mode 100644 index 00000000000..8ebaec58f16 --- /dev/null +++ b/scripts/install_mps_flash_attention.py @@ -0,0 +1,145 @@ +#!/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 shutil +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, + ) + # 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( + setup, + "'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( + package_init, + f'__version__ = "{VERSION}"\n', + 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_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" + ")\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/scripts/mps_fused_group_norm.mm b/scripts/mps_fused_group_norm.mm new file mode 100644 index 00000000000..2f0c97586ca --- /dev/null +++ b/scripts/mps_fused_group_norm.mm @@ -0,0 +1,446 @@ +// Fused inference-only GroupNorm + SiLU for contiguous float16 MPS tensors. +#include +#include +#include +#include + +#import +#import + +#include + +namespace { + +struct FusedGroupNormParams { + uint32_t batch; + uint32_t channels; + uint32_t spatial; + uint32_t groups; + 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()); +} + +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; +} + +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; +} + +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, + 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; +} + +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; +} + +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) { + 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")); + 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, + "Fused Metal GEGLU forward pass", + pybind11::arg("input"), + pybind11::arg("gelu_lut")); +} 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 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")); +} 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_flash_attention.py b/test/test_mps_flash_attention.py new file mode 100644 index 00000000000..7e70e225ccd --- /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_measured_sd1_dimensions_route_to_mfa(): + assert not should_use_mfa_shape(4096, 4096, 64) + 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(): + assert not should_use_mfa_shape(128, 128, 40) diff --git a/test/test_mps_fused_ops.py b/test/test_mps_fused_ops.py new file mode 100644 index 00000000000..00f79628ebc --- /dev/null +++ b/test/test_mps_fused_ops.py @@ -0,0 +1,95 @@ +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 + + +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) + 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) 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/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..3620ec1f295 100644 --- a/webui-macos-env.sh +++ b/webui-macos-env.sh @@ -5,13 +5,30 @@ #################################################################### 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 --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. + 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 #################################################################### 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" + ###########################################