Automatic1111 Metal laten 'spreken' op een M1

Ik maak veel gebruik van Draw Things op Apple-hardware, en er is één ding dat me altijd heeft gestoord aan Automatic1111: het voelt trager dan het zou moeten zijn. Niet onbruikbaar traag, maar net traag genoeg om het op te merken.

Op mijn M3 Pro duurde een korte generatie van vijf stappen met DPM++ SDE in Automatic1111 doorgaans zo'n 8 tot 10 seconden. Draw Things had me al laten zien dat Stable Diffusion op Apple Silicon veel directer kon aanvoelen. Daarom wilde ik onderzoeken hoeveel van dat gat daadwerkelijk noodzakelijk was.

Voor dit project heb ik gebruikgemaakt van de volgende repository: dmikey/stable-diffusion-webui-metal, een gefinetunede versie van Automatic1111 voor Apple Silicon.

De randvoorwaarden

Er was één belangrijke beperking: ik wilde Automatic1111 niet vervangen. Ik wilde dezelfde WebUI, checkpoints, LoRAs, samplers, extensies, API, prompt-syntaxis en algemene workflow behouden. Ik was niet geïnteresseerd in het converteren van alles naar Core ML om daar een nieuwe inference engine omheen te bouwen. Het doel was veel specifieker:

Hoe snel kan Automatic1111 worden als we de cruciale onderdelen laten werken als native Apple-software?

Het antwoord is, althans voor mijn workloads, aanzienlijk sneller. De generaties die op mijn M3 Pro ongeveer 8–10 seconden duurden, landen nu over het algemeen tussen de 3 en 7 seconden. Op mijn M1 Mac Mini ging dit van 13–20 seconden naar 8–10 seconden.

Dit zijn geobserveerde ranges in mijn huidige workloads, geen gecontroleerde benchmark die een universele tweevoudige verbetering claimt. Er is ook een belangrijk onderscheid tussen de runtime-verbeteringen en NGMS (Negative Guidance minimum sigma), wat daadwerkelijk de hoeveelheid guidance-werk vermindert. Toch is het verschil in het dagelijks gebruik substantieel.

Wat interessanter was dan het uiteindelijke getal, was wat er nodig was om daar te komen. Het was niet één enkele optimalisatie.

Begin bij de workload, niet bij de benchmark

De workload waar ik me op richtte was zeer specifiek:

  • Stable Diffusion 1.x
  • DPM++ SDE
  • Karras
  • 5 stappen
  • CFG rond de 1.15
  • Resoluties van 384×640 en 512×512
  • FP16 UNet op MPS (Metal Performance Shaders)
  • FP32 VAE standaard

Die specificiteit is belangrijk. In het begin leek DPM++ 2M een gemakkelijke manier om tijd te besparen. Het was sneller, maar produceerde niet het resultaat dat ik wilde bij een kort schema. Dat is geen optimalisatie; dat is een andere workload. Dit werd de regel voor alles wat volgde: als een optimalisatie in isolatie geweldig lijkt, maar de werkelijke generatie niet versnelt terwijl het gewenste resultaat behouden blijft, telt het niet.

Selectieve Metal Flash Attention

Attention was het meest logische startpunt. De MPS-backend van PyTorch is aanzienlijk verbeterd, maar er zijn nog steeds Stable Diffusion attention shapes waarbij direct overschakelen naar Metal zinvol is. Een fout zou zijn geweest om een custom Metal-implementatie als universeel sneller te beschouwen; dat is het namelijk niet.

In plaats daarvan heb ik een Metal Flash Attention-pad toegevoegd, specifiek voor de SD 1.x shapes waar het in tests daadwerkelijk won. De router werkt conceptueel als volgt:

if inference and fp16_mps and query_tokens >= 192 and head_dim in (40, 80, 160):
    return metal_flash_attention(q, k, v)
return pytorch_sdpa(q, k, v)

Er zijn aanvullende controles voor masks, training, dropout, tensor-layout, grouped-query attention en ondersteunde typen, maar dat is het basisidee. Metal is niet de standaard omdat "Metal sneller klinkt", maar krijgt de operatie toegewezen wanneer we hebben gemeten dat die specifieke shape daar baat bij heeft.

Alles wat niet aan de criteria voldoet, gaat terug via PyTorch. Deze fallback is essentieel, omdat Automatic1111 veel meer configuraties ondersteunt dan mijn vijfstaps SD 1.x workflow.

De kernel was niet het enige probleem

Het verplaatsen van attention naar Metal hielp, maar legde iets anders bloot: de native extensie committe de MPS command buffer na elke attention-aanroep. Stable Diffusion roept attention herhaaldelijk aan binnen elke UNet-evaluatie. Bij een korte generatie van vijf stappen wordt het herhaaldelijk indienen van kleine stukjes werk een betekenisend deel van de totale runtime.

In plaats van de Metal-kernel als een eigen kleine applicatie te behandelen, heb ik deze geïntegreerd in de huidige MPS-stream van PyTorch. De extensie beëindigt de huidige kernel-coalescing van PyTorch, codeert de Metal Flash Attention-operatie in de huidige command buffer en laat het overige PyTorch MPS-werk vanaf daar doorgaan.

De expliciete commit na elke attention-aanroep is hiermee verwijderd. Dit bleek een van de belangrijkste lessen: zelfs de snelste kernel verliest als je de command buffer na elke aanroep moet indienen. Bij deze generatietijden telt overhead zwaar mee; je optimaliseert niet langer alleen hoe snel de GPU matrices kan vermenigvuldigen, maar hoe vaak Python, PyTorch, MPSGraph en Metal met elkaar moeten coördineren.

Tevens een belangrijke herinnering om de timer niet blindelings te vertrouwen: een van de vroege versies produceerde een volledig groen beeld. Het was razendsnel, maar het resultaat was onbruikbaar. Daarom voert het Metal-pad nu een geïsoleerde attention-plus-projection correctietest uit voordat de WebUI het activeert.

Unified Memory verandert de regels

Het volgende probleem was het geheugen. Apple Silicon heeft geen discrete VRAM naast het systeemgeheugen; de GPU en de rest van de machine concurreren om hetzelfde fysieke geheugen. Dit maakt traditionele GPU-aannames onjuist.

Een attention-matrix kan technisch gezien in het geheugen passen, maar toch een slecht idee zijn als macOS onder druk staat, de allocator begint te 'thrashen' of de machine gaat swappen. In plaats van een vaste VRAM-drempel te gebruiken, schat deze fork de kosten van native attention af tegenover zowel het totale als het momenteel beschikbare geheugen.

Conceptueel:

  • attentionbytes = batch × heads × querytokens × keytokens × elementsize
  • estimatedpeak = attentionbytes × 2.5
  • budget = min(10% van totaal geheugen, 20% van beschikbaar geheugen, 1.5 GiB)

Als de geschatte piek binnen dit budget past, kan native SDPA draaien. Zo niet, dan gaat het verzoek via het geheugenbeperkte sub-kwadratische pad. De chunk-grootte voor die fallback is ook dynamisch; een Mac met 8 GB moet immers andere beslissingen nemen dan een Mac met 32 GB, rekening houdend met andere actieve programma's zoals Chrome of Xcode.

Stop met het bewaren van elke attention chunk

Ik heb ook veranderd hoe de sub-kwadratische fallback omgaat met K/V chunks. De bestaande aanpak berekent partiële attention-resultaten en bewaart de teller, het normalisatiegewicht en het maximum voor elke chunk, om alles aan het einde samen te voegen. Dat is onnodig.

In plaats daarvan houdt de fork een lopend maximum, een normalisatiesom en een gewogen output bij. Elke nieuwe K/V chunk wordt samengevoegd in deze status en kan daarna worden weggegooid. De recurrentie ziet er basismatig zo uit:

new_max = max(running_max, chunk_max)
running_scale = exp(running_max - new_max)
chunk_scale = exp(chunk_max - new_max)
running_values = running_values × running_scale + chunk_values × chunk_scale
running_weights = running_weights × running_scale + chunk_weights × chunk_scale

Dit is hetzelfde algemene online-softmax idee dat Flash Attention geheugenefficiënt maakt. Het geheugen schaalt nu rond de huidige chunk in plaats van alle partiële resultaten te accumuleren tot het einde.

Verouderde MPS-workarounds verwijderen

Er was ook een minder glamoureuze categorie optimalisaties: het verwijderen van oude workarounds. De PyTorch-backend van Apple is sterk veranderd, maar Automatic1111 had defensief gedrag opgebouwd voor oudere MPS-implementaties (zoals het clonen van torch.narrow() resultaten en het pushen van LayerNorm via FP32).

Op nieuwere versies van PyTorch zijn deze fixes simpelweg overbodige kopieën, allocaties en conversies. Deze behaviors zijn nu gekoppeld aan de runtime-versie in plaats van onvoorwaardelijk te worden toegepast. Er is nog steeds een A1111MPSFORCELEGACYOPS=1 escape hatch voor wie het oude gedrag nodig heeft.

Daarnaast heb ik PYTORCHMPSPREFER_METAL=1 ingeschakeld omdat directe Metal-matrixvermenigvuldiging beter testte voor de SD 1.x projectiegrootten, en de standaard sampling upcast verwijderd zodat meer van het korte samplingpad in FP16 blijft. Dat laatste is een trade-off; het kan de output bij dezelfde seed beïnvloeden, maar voor deze workflow is dat acceptabel.

Het fuseren van GroupNorm en SiLU

Nadat het onnodige werk was verminderd, zocht ik naar operaties die noodzakelijk waren maar constant werden herhaald. GroupNorm gevolgd door SiLU komt overal voor in de SD 1.x UNet. Normaal zijn dit aparte PyTorch-operaties, wat betekent: aparte dispatches en een tussenliggende activatie die wordt weggeschreven om direct daarna weer gelezen te worden.

Ik heb daarom een fused Metal kernel geschreven. Voor compatibele FP16 inference tensors handelt één 256-thread Metal threadgroup elk batch/group paar af. De kernel accumuleert de som en kwadraatsom in FP32, reduceert deze naar gemiddelde en variantie, past normalisatie en de affiene parameters toe, voert SiLU uit en schrijft het FP16 resultaat weg.

Eén dispatch. Geen tussenliggende activatie. Als de tensor niet compatibel is of als er gradients zijn ingeschakeld, valt het systeem direct terug op: F.silu(norm(input_tensor)).

NGMS is iets anders

Een deel van de uiteindelijke snelheidswinst moet worden gescheiden van het engine-werk: NGMS (Negative Guidance minimum sigma). NGMS kan onvoorwaardelijke guidance overslaan tijdens geschikte delen van het samplingproces. Bij een lage CFG zoals 1.15 op een vijfstaps schema, kan het overslaan van dit werk een aanzienlijke hoeveelheid berekeningen besparen.

Dit is inherent snel omdat de GPU simpelweg minder werk doet. De fork stelt NGMS standaard in op 1.0 met all-steps behavior voor deze tuned workflow. Maar dit is niet dezelfde categorie als het sneller maken van attention of GroupNorm; NGMS verandert de denoising-berekening en kan de compositie en details beïnvloeden.

Er zijn dus twee verhalen:

  1. De bestaande engine goedkoper maken (Metal attention, minder command-buffer submissions, beter geheugenbeheer, fused operations).
  2. De engine vragen om minder werk te doen via NGMS.

De meest nuttige optimalisaties waren degene die ik verwijderde

Veel van dit project bestond uit het proberen van zaken die theoretisch zouden moeten werken en deze vervolgens weer te verwijderen.

  • Packed QKV projections: Ik implementeerde ze, het beeld was identiek, maar de performance ging van 8,988 seconden naar 9,011 seconden (0,26% trager). Verwijderd.
  • MPSGraph voor residual blocks: Het idee was om GroupNorm, SiLU, convoluties en meer in één graph te plaatsen. De microbenchmarks waren bemoedigend (sommige blokken 1–9% sneller). Maar bij een volledige image-generatie was de MPSGraph-versie (9,6533s) trager dan het bestaande pad (9,5556s). Verwijderd.

Deze laatste experimenten leerden me dat correcte code niet noodzakelijkerwijs nuttige code is. Microbenchmarks nomineren wijzigingen; volledige generaties kiezen ze.

Wat er uiteindelijk is overgebleven

De uiteindelijke implementatie raakt 20 van de 329 tracked paths aan. De onderdelen die zijn blijven staan:

  • Shape-selectieve Metal Flash Attention
  • Uitgestelde Metal command-buffer submission
  • Unified-memory-bewuste attention routing
  • Dynamische query chunk sizing
  • Streaming online softmax
  • Verwijdering van verouderde MPS kopieën en FP32 omwegen
  • Voorkeur voor directe Metal matmul
  • Minder sampling upcasting
  • Fused GroupNorm plus SiLU
  • NGMS voor de tuned low-CFG workflow

Niet-ondersteunde inputs vallen nog steeds terug op PyTorch. De native extensies testen zichzelf bij het opstarten in een subprocess; ze worden alleen ingeschakeld als de tests slagen. Dit is cruciaal, omdat fouten in native GPU-code soms de hele Python-interpreter kunnen meesleuren.

Resultaten

Op de M3 Pro verschoof de workload van ongeveer 8–10 seconden naar 3–7 seconden. Op een M1 Mac mini zijn de resultaten meetbaarder (met dezelfde model hash, shape, sampler en configuratie):

BuildTijd
Automatic1111 v1.10.1-96-g1937682a12.8 s
Metal fork v1.10.1-99-g38ac556a88.7 s

Dit is een latency-vermindering van ongeveer 32%, of een doorvoersnelheid van circa 1,47x.

Waarom dit nog steeds geen "Draw Things" is

Draw Things heeft een fundamenteel voordeel: het beheert de volledige execution environment en kan alles (modelrepresentatie, geheugen, scheduling) ontwerpen rondom Apple-hardware.

Automatic1111 kan dat niet zonder op te offeren wat het nuttig maakt. Het is dynamische Python-software waar mensen monkey patches voor maken, extensies in haken en modellen/LoRAs wisselen terwijl het draait. Dat ecosysteem wilde ik behouden.

Dit project richt zich in plaats daarvan op de "naden" (the seams) waar native Apple-executie kan binnenkomen en vertrekken zonder dat Automatic1111 een andere applicatie hoeft te worden. We hebben daarmee een aanzienlijke prestatiewinst behaald. Voor nu is dit nog steeds Automatic1111: dezelfde checkpoints, LoRAs, UI en extensies. Het brengt alleen veel minder tijd door met wachten op de grenzen tussen Python, PyTorch, MPSGraph en Metal.