Mise à jour Shmoosh :
C’est le résultat de vitesse le plus encourageant jusqu’ici.
Pour le contexte, Shmoosh est mon projet de recherche actuel pour accélérer les modèles de diffusion en compressant des parties de l’attention de manière à préserver le comportement du modèle. Le but est de trouver des techniques qui comptent vraiment pour les outils créatifs locaux comme Underpaint et Comfy, où les gens font tourner des flux de génération et d’édition d’images sur des GPU grand public.
Et aujourd’hui, le kernel a enfin bougé.
J’ai poussé la tranche suivante :
1f0ee4e Add TF32 packed attention fast mode.
En bref, j’ai ajouté le support optionnel de `dot_precision` pour l’attention compactée fusionnée.
Le défaut reste `ieee`, parce que la fidélité compte et que je ne veux pas modifier silencieusement la politique de haute qualité. Mais il y a maintenant un mode rapide `tf32` explicite pour les cas où on veut de la vitesse et peut tolérer un très léger compromis sur la qualité.
Les changements mesurés
Sur des tests de packed_t synthétiques 1024px K7/no-QJL sur RTX 4070, l’attention compactée est passée de :
0.6232ms à 0.2831ms
La phase du kernel d’attention compactée est donc plus de 2x plus rapide.
Le chemin total encode+attention est passé de :
0.8177ms à 0.4761ms
Et le RMSE relatif a quasiment pas bougé :
0.023998 à 0.024101
C’est le genre de résultat qui fait qu’on arrête de plisser les yeux devant le profiler pour commencer à faire attention.
La suite d’images 1024px a aussi bien tenu le coup :
Accélération moyenne : 1.068x
PSNR moyen : 54.17 dB
PSNR min : 51.70 dB
La phase d’attention compactée dans la suite d’images est passée de :
1.2806ms à 0.7293ms
Les résultats synthétiques sont utiles, mais ils peuvent mentir. Voir la phase d’attention compactée bouger significativement dans la génération d’images réelle est plus intéressant.
Une politique de mode rapide explicite
Je n’ai pas promu le TF32 dans la politique par défaut de haute fidélité, parce que l’un des cas de test, reading-nook, a payé une petite taxe sur la qualité :
52.07 dB à 51.70 dB
C’est encore très élevé, mais c’est une différence réelle, et ça mérite d’être traité comme un choix de politique.
J’ai donc ajouté une config de mode rapide séparée :
configs/underpaint-juggernaut-sdxl-up0-self-attn1-firstblocks-gated70pct-k7-noqjl-tf32-policy.json
La validation est propre :
pytest: 87 passed
compileall: passed
jq: passed
git diff --check: passed
repo clean and synced
La prochaine forme du système
Je suis en train de cerner lentement la bonne forme.
La compression K-only s’est révélée plus prometteuse que de toucher à V.
La représentation packée a battu la voie directe rotated-K, parce que la densité, ça compte. La self-attention late-gated semble toujours être la voie la plus sensée.
K7/no-QJL tient assez bien la route pour être intéressant.
Et maintenant, TF32 donne un vrai coup de boost au kernel de packed attention.
Avant ça, Shmoosh était dans la zone « scientifiquement prometteur, modestement plus rapide ». Mais un outil créatif local a besoin de victoires qui survivent au contact des workflows d’image réels, des vrais GPU et des exigences de qualité.
Je bosse encore dessus.
C’est le premier résultat où la voie de la packed attention est devenue significativement plus rapide.
Je suis excitée.