Mac (MPS) support: CUDA templates run unchanged on Apple Silicon

CLAUDE.md

@@ -23,7 +23,7 @@ python -m dw.run workflows/templates/text-to-image.json prompt="a cat" num_image
 # Validate a workflow against schema
 python -m dw.validate workflows/models/z-image.json
 
-# Basic system test (torch, diffusers import check)
+# System test - downloads SD 1.5 (a few GB) and generates one image
 python -m dw.test
 
 # Interactive REPL
@@ -235,7 +235,7 @@ SDNQ pre-quantized models use a different pattern: `pre_load_modules` imports sd
 
 `dw/__init__.py` handles device detection (CUDA > MPS > CPU) and platform-specific optimizations:
 - **CUDA**: TF32 matmul, cuDNN benchmark, deterministic mode (configurable via settings)
-- **MPS**: `PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0` (use all unified memory), autocast warnings suppressed, attention slicing enabled by default
+- **MPS**: `PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0` (use all unified memory), autocast warnings suppressed; attention slicing automatic (faster for SD 1.5's head dims, 2.4x slower for SDXL's - `disable_attention_slicing`)
 - **CPU**: Warning displayed
 
 Detection is overridden by the `DW_DEVICE` environment variable (single run) or the `device` setting (standing), either of which can name a specific accelerator such as `cuda:1`. Device placement is explicit throughout — no default torch device is set, since that would build models directly in VRAM and defeat offloading. Compare backends with `get_device_type()` rather than `== "cuda"`, which a device like `cuda:1` would fail.
@@ -244,6 +244,8 @@ A step can override the device it runs on: `device` in a pipeline `configuration
 
 Every device a workflow names passes through `resolve_device()`, which translates a backend this machine does not have into the one it does and warns — a `cuda` workflow runs on a Mac and an `mps` one runs on a CUDA box. Only the backend is translated: an index survives when the backend matches (`cuda:1` on a single-GPU CUDA box stays a genuine error) and is dropped when it does not. `cpu` is never rewritten, since pinning a step to the CPU is how a GPU-specific problem gets ruled out. Translation happens before anything reads the backend, so the MPS accommodations (the sequential-offload downgrade, attention slicing, the compile skip) fire for a translated device too.
 
+The same translation reaches the settings that carry a device or a CUDA-only feature, so a template written on the CUDA box runs unchanged on a Mac and a caller (an MCP agent included) never has to know the backend: SDNQ `quantization_device`/`return_device` go through `resolve_device()`, and `use_quantized_matmul(_conv)` is turned off on MPS, where it falls back to `torch._int_mm`, ~500x slower (`portable_quantization_arguments`, `config_objects.py`). Group-offload `use_stream`/`record_stream` are dropped when the onload device is not CUDA/XPU. `vram_estimate` checks against the serving device's own `cost` entries, else its `device_capacity_gb()` (Metal's recommended working set on a Mac), else every entry. `device_memory_stats()` reports real unified-memory figures on MPS. Each adaptation logs a warning, and torch's MPS CPU-fallback warning is let through the blanket `UserWarning` filter. `HF_ENABLE_PARALLEL_LOADING` defaults to off on macOS (`_parallel_loading_default`): diffusers' per-shard loader threads copying onto MPS at once segfaulted the LTX-2.5 SDNQ transformer load; it is chosen by platform because diffusers reads it at import, before dw can ask torch for a device.
+
 A `components` entry can additionally set `residency: "on_demand"`, which rests the component on the CPU and wraps its `forward`/`encode`/`decode` to move it to the device around each call (`apply_on_demand_placement` in `pipeline.py`). The wrappers use `functools.wraps` because callers introspect the signature — MiniMax H3's denoiser picks its arguments from `signature(transformer.forward)`. It is mutually exclusive with `group_offload` on the same component, and like `group_offload` it suppresses the wholesale `pipeline.to(device)` at load.
 
 Settings in `~/.diffusers_helper/settings.json`: `device`, `enable_tf32`, `cudnn_benchmark`, `cudnn_deterministic`, `log_level`, `log_filename`.
@@ -321,7 +323,7 @@ same reason - default setup cannot load a pack.
 - **Cartesian product explosion** — multiple `previous_result` references multiply: 4 images × 3 masks = 12 iterations
 - **Component sharing requires exact key matching** between `shared_components` and `reused_components`
 - **Built-in workflows** need explicit argument mapping: `"prompt": "variable:prompt"`
-- **MPS differences from CUDA**: no autocast, no bitsandbytes, no flash_attn, no triton, no torch.compile. Model offloading has less benefit on unified memory, and `"offload": "sequential"` is downgraded to `"model"` with a warning there (`place_component`) — per-submodule streaming hands back no residency when the CPU and the accelerator share one pool. `exclude_from_cpu_offload` is sequential-only and does not survive the downgrade.
+- **MPS differences from CUDA**: no bitsandbytes, no flash_attn, no triton, no torch.compile (`torch.autocast("mps")` works on torch 2.14, but dw does not use it). Model offloading has less benefit on unified memory, and `"offload": "sequential"` is downgraded to `"model"` with a warning there (`place_component`) — per-submodule streaming hands back no residency when the CPU and the accelerator share one pool. `exclude_from_cpu_offload` is sequential-only and does not survive the downgrade.
 - **`{}`-escaped strings** in JSON arguments: `"{nf4}"` stays as string `"nf4"`, without braces it would try to load as a type
 - **A stored prompt's `text` may not begin with a reference prefix** (`variable:`, `previous_result:`, `constant:`, `asset:`, `output:`, `prompt:`) — the engine rejects it to prevent double resolution or iteration expansion
 - **Audio+video muxing**: pipelines that generate audio alongside video (LTX-2) have the two muxed into one `video/mp4` file with PyAV in `result.py`
@@ -573,7 +575,7 @@ same reason - default setup cannot load a pack.
   stops the next instance. The half that fixes this one: every template
   whose deliverable ends in a `pair_audio` mux (`music-video`,
   `assemble-and-score`, `dissolve-between-shots`) normalizes to **-3 dBFS**;
-  `music`, an mp3, keeps -1
+  `music`, an mp3, normalizes to -3 as well (#362)
 - **A variable's bound is declared by the author, checked three times** — a
   model's own rule about a value (H3's `num_frames` is `17 * n + 5` from 124
   to 345) is a property of the model, so it lives in the workflow rather than

README.md

@@ -32,8 +32,8 @@ source ./activate
 .\venv\scripts\activate
 ```
 
-`python -m dw.test` confirms torch and diffusers import and shows which
-accelerator was found.
+`python -m dw.test` runs a small built-in workflow end to end - it downloads
+SD 1.5 (a few GB) and generates one image on whichever accelerator was found.
 
 **2. Start the engine.** Leave it running; everything else talks to it.
 

docs/ACCELERATION.md

@@ -215,7 +215,7 @@ Prefer the pinned form for a compiled component - the per-call context manager s
 }
 ```
 
-Processes attention in slices to reduce memory at some cost to speed. Enabled automatically on MPS (unified memory benefits from slicing) unless `disable_attention_slicing` is set. Modular pipelines have no `enable_attention_slicing()` method - the setting is silently skipped rather than failing when the pipeline doesn't support it.
+Processes attention in slices to reduce memory at some cost to speed. Automatic on MPS unless `disable_attention_slicing` is set, opt-in elsewhere. Only UNet/ControlNet models take it. On MPS, whether it helps depends on the model's attention head dim: MPS SDPA is slow at 40, 48 and 160 (slicing made SD 1.5 ~20% faster end to end) and fast at 64-128 (slicing made SDXL-shaped attention 2.4x slower - set `disable_attention_slicing` for SDXL-family models). Modular pipelines have no `enable_attention_slicing()` method - the setting is silently skipped rather than failing when the pipeline doesn't support it.
 
 ## torch.compile
 
@@ -344,9 +344,13 @@ For faster model downloads, optionally `pip install hf_transfer` and set `HF_HUB
 
 Apple Silicon has narrower acceleration support than CUDA:
 
-- No flash-attn, no Triton, no bitsandbytes - `attention_backend` is effectively CUDA-only; use `"native"`-family backends or leave it unset on MPS. `compile` is skipped with a warning (inductor support on MPS is immature).
-- No `torch.autocast` support - autocast-related warnings from other libraries are suppressed automatically rather than surfaced.
-- `enable_attention_slicing` is on by default (set `disable_attention_slicing` to turn it off).
-- `float16` produces NaN values on Apple Silicon - use `float32` or `bfloat16` for `torch_dtype` instead; dw only warns, it doesn't override the dtype for you.
+- No flash-attn, no Triton, no bitsandbytes - `attention_backend` is effectively CUDA-only; use `"native"`-family backends or leave it unset on MPS. `compile` is skipped with a warning (inductor support on MPS is immature). A workflow that loads a bitsandbytes checkpoint (`flux2-dev`, `multi-image-reference`) or a TorchAO int4 config (`flux-torchao`) has no Mac path yet.
+- Workflows written for CUDA are adapted rather than refused: a `cuda` device becomes `mps`, SDNQ's `quantization_device` follows it and its quantized matmul is switched off (it runs through `torch._int_mm`, ~500x slower on MPS), group-offload CUDA streams are dropped, and `"offload": "sequential"` becomes `"model"`. Each change is logged as a warning.
+- Attention slicing is automatic; set `disable_attention_slicing` for SDXL-family models (head dim 64), where it is 2.4x slower in attention. It helps SD 1.5.
+- `float16` can produce NaN values on Apple Silicon - use `bfloat16`; dw only warns, it doesn't override the dtype for you.
 - `PYTORCH_MPS_HIGH_WATERMARK_RATIO` defaults to `0.0` (use all unified memory) unless already set in the environment.
+- Ops the MPS backend lacks: with `PYTORCH_ENABLE_MPS_FALLBACK=1` (set by the `fp4-fp8-for-torch-mps` package install.sh adds) they run on the CPU, and dw logs torch's warning when one does. Without it they raise.
+- Memory figures are real: allocated, driver-reserved, and Metal's recommended working set as the total. `vram_estimate` checks a Mac against that total, since the catalog's `cost` entries are CUDA cards; the per-voxel figures were calibrated on CUDA.
+- Cost estimates quoted on a Mac are the CUDA figures (`basis: "other_device"`) until the Mac's own runs build observed history - expect them to be optimistic.
+- Parallel checkpoint loading (`HF_ENABLE_PARALLEL_LOADING`) defaults to off on macOS - concurrent loader threads copying onto MPS segfaulted a quantized load. Set it explicitly to override.
 - Offloading has less benefit than on CUDA, since unified memory is already shared between CPU and GPU.

docs/DEPENDENCIES.md

@@ -55,7 +55,7 @@ pyproject; install.ps1 installs it explicitly on Windows.
 
 **Windows (CUDA):** bitsandbytes, kernels
 
-**macOS (MPS):** fp4-fp8-for-torch-mps (FP8/FP4 dtype support for Metal), fluidtop
+**macOS (MPS):** fp4-fp8-for-torch-mps (FP8/FP4 dtypes for Metal; autoloads on every torch import, overrides mm/linear/copy on the MPS dispatch key and sets PYTORCH_ENABLE_MPS_FALLBACK=1), fluidtop (a GPU/power TUI; needs sudo)
 
 ## Optional
 

docs/TESTING.md

@@ -63,7 +63,7 @@ An autouse `_clear_task_model_cache` fixture clears `dw.tasks.model_cache` befor
 ## Quick Validation
 
 ```bash
-# Verify torch and diffusers are working
+# End-to-end check: downloads SD 1.5 and generates one image
 python -m dw.test
 
 # Validate a workflow against schema

docs/WORKFLOW_GUIDE.md

@@ -1289,8 +1289,9 @@ stays alive for the steps that reuse it.
 }
 ```
 
-- `enable_attention_slicing` — process attention in slices to reduce memory. Enabled
-  automatically on MPS unless `disable_attention_slicing` is set.
+- `enable_attention_slicing` — process attention in slices to reduce memory at some
+  cost to speed. Automatic on MPS unless `disable_attention_slicing` is set
+  (UNet/ControlNet models only; worth disabling for SDXL on a Mac).
 - `attention_backend` — selects a diffusers attention backend (e.g. `"flash_hub"`) for
   the duration of each pipeline call.
 - `prompt_weighting` — enables A1111-style prompt weighting (`(word:1.5)`, `[word]`,