Investigation d'ingénierie : ce portage MLX peut-il devenir 10× plus rapide de nouveau ?

Date : 2026-09-19. Machine : Apple M3 Max, 40 cœurs GPU, 128 GiB de mémoire unifiée, macOS 27.2, MLX / MLX Metal 0.32.2, inférence FP16. Ce rapport contient de vraies expériences locales, y compris un kernel Metal écrit à la main. Il ne change pas le runtime de production et ne publie pas de poids quantizés.

Les changements d’ingénierie testés n’apportent pas 10×. Des mesures entrelacées soutiennent des améliorations modestes et dépendantes de la forme grâce à la compilation et à l’élagage des sorties inutilisées de la dernière couche de la tête de décision. Des cas sélectionnés se sont améliorés d’environ 3–8% en utilisant des médianes appariées par tour. Certains intervalles à plus grand lot n’incluent aucune amélioration. Un kernel personnalisé GELU/porte à erf exact a été numériquement réussi mais n’a pas apporté de bénéfice de bout en bout supplémentaire cohérent par rapport à la compilation MLX. Une quantization naïve 8 bits et 4 bits du backbone a réduit le stockage, n’a pas accéléré les charges pilotes plus grandes, et a changé les prédictions ou les probabilités calibrées.

Les limites mathématiques et les compromis d’approximation sont examinés séparément dans MATH_10X_RESEARCH.md. La revue d’implémentation d’origine est dans PERFORMANCE_RESEARCH.md ; le benchmark de checkpoint publié reste BENCHMARKS.md.

Contrôles et limites expérimentaux

Tout le travail GPU de recherche s’est exécuté en série. Le reste du travail des agents n’utilisait que le CPU/le système de fichiers/le réseau. La machine était sur secteur, sans avertissement thermique/de performance pmset enregistré et sans utilisation de swap rapportée pendant l’expérience. L’activité normale du bureau s’est poursuivie. Ce n’est pas une chambre thermique contrôlée ni une machine de benchmark dédiée autrement inactive.

Les premières exécutions de présélection exécutaient chaque candidat dans un processus neuf, avec 4–5 warmups et 12–16 échantillons. Elles ont révélé une dérive substantielle d’une exécution à l’autre. Par exemple, le pilote anglais à une question suggérait une amélioration de compilation de 1.24×, alors que l’expérience entrelacée suivante n’a trouvé qu’environ 1.03×. Les latences du pilote séquentiel sont donc des preuves de présélection, pas la prétention de gain de vitesse causale principale.

Le script de confirmation paired.py :

• Fait tourner l’ordre des candidats au sein de chaque tour et utilise les mêmes entrées pour chaque candidat de ce tour.
• Change le texte d’état réel entre les tours. Il génère jusqu’à 16 variantes d’état et retient les variantes de même forme de tenseur ; les cas courts multilingues ont 10 telles variantes, tandis que les autres cas rapportés en ont 16.
• Utilise des questions distinctes en langage naturel, dont 50 instructions différentes pour la plus grande charge courte. Il vérifie les hachages d’entrée et ne met pas en cache les réponses, ne déduplique pas les questions et ne réutilise pas les états d’encodeur contextuels.
• Évalue les résultats et synchronise le GPU avant d’arrêter chaque minuteur. Il mesure à la fois les appels de passe avant préparés et le chemin de prédiction public, incluant la tokenisation et le formatage de sortie. Le chargement du modèle est exclu.
• Exécute 32 tours mesurés pour l’expérience anglaise tête/compilation et 16 pour les expériences multilingue et Metal personnalisé, après warmup. Chaque candidat voit le même nombre de tours et la même séquence d’entrées.

Ces entrées diffèrent des jeux de référence publiés. Les comparaisons ci-dessous sont au sein de l’expérience de recherche, pas des comparaisons avant/après obtenues en divisant des tableaux sans rapport. La limite de lot de la recherche est 64, tandis que l’API publiée est par défaut à 16. Les variantes d’état répétées sont des mesures répétées intentionnelles ; il n’y a pas de cache de résultats.

analyze.py calcule par tour les rapports eager_time / candidate_time et des intervalles bootstrap de centiles exploratoires pour leur médiane, à l’aide de 2,000 rééchantillonnages des indices de tour. Ces intervalles ne tiennent pas compte de chaque source de bruit du système d’exploitation ni de la corrélation sérielle et ne remplacent pas une réplication multi-session. Un rapport de valeurs p50 calculées indépendamment peut différer du rapport apparié médian.

Le JSON brut inclut tous les temps, les hachages d’entrée, les métadonnées d’environnement, les métriques de parité et l’empreinte source enregistrée au moment de la mesure. Les scripts d’expérience ont ensuite été formatés et étendus avec des candidats optionnels disjoints ; les empreintes antérieures décrivent ces versions antérieures des scripts.

Compilation et élagage exact de la tête finale

Quatre chemins ont été comparés :

1. Eager : le DecisionModel FP16 publié.
2. Compilé : mx.compile autour du modèle chargé, évalué et figé, avec une spécialisation de forme normale.
3. Q sélectionné + compilé : dans la dernière couche de tête, préserver la projection QKV et les K/V pleine longueur, mais n’émettre que les requêtes d’attention CLS/marqueur d’option. Exécuter la projection de sortie et le FFN seulement sur ces sorties sélectionnées.
4. Attention complète + sorties sélectionnées + compilé : préserver l’appel QKV et SDPA pleine longueur d’origine, puis rassembler les sorties CLS/option avant la projection de sortie et le FFN. Cela conserve la forme du kernel d’attention d’origine tout en retirant la majeure partie du travail dense inutilisé de la tête finale.

Les deux prototypes d’élagage préservent les dépendances mathématiques du modèle. Ils calculent encore toutes les projections QKV ; ils ne réalisent pas l’économie supplémentaire de projection Q seule du calcul de borne supérieure mathématique. Changer les formes GEMM et SDPA peut changer l’arrondi à virgule flottante. Aucun des deux prototypes n’est un cache de décodeur, une sortie anticipée ou une approximation qui écarte des couches de transformer antérieures.

Latence p50 de bout en bout, en millisecondes :

Sources : données appariées anglaises et données appariées multilingues.

Pour le chemin attention complète/sorties sélectionnées, le gain de vitesse médian apparié et les intervalles exploratoires à 95% incluent :

Les intervalles anglais à 50 questions et multilingue à grand lot long incluent 1. Ils n’établissent pas une amélioration reproductible. Le chemin à Q sélectionné est un peu meilleur pour les cas multilingue court-16 et long-1, mais aucun chemin d’élagage ne domine toutes les formes. Tous les intervalles des candidats, les mesures de forward et les rapports bruts par tour sont dans paired_analysis.json.

La compilation correspondait exactement aux logits eager, aux logits d’action et aux probabilités calibrées sur les 1,530 comparaisons de questions à entrées modifiées à travers les deux familles de modèles dans cette expérience tête/compilation. Les deux chemins d’élagage étaient d’accord sur les 1,530 décisions d’argmax, avec une différence maximale de probabilité calibrée de 0.0001883. Le chemin d’élagage à attention complète a aussi passé la suite distincte de 63 questions pour chaque modèle : 126/126 d’accord, avec des différences maximales de probabilité de 4.31e-5 pour l’anglais et 6.48e-6 pour le multilingue. Ce sont des contrôles de régression, pas une prétention d’exactitude de tâche sur 1,530 exemples étiquetés indépendamment.

La compilation du modèle complet et par blocs a été présélectionnée. L’expérience par blocs a aussi préservé les 63 sorties du jeu anglais, mais n’a pas établi d’avantage matériel sur la compilation du modèle complet. La spécialisation de forme doit être bornée dans un service. Le modèle utilise des remises en forme et des masques Python dépendant de la forme, donc appliquer shapeless=True sans discernement est dangereux. Le guide officiel de compilation documente la spécialisation de forme et la capture d’état.

Le premier appel candidat anglais sur le modèle complet a pris 2,166.7 ms, suivi d’environ 12.75 ms de p50 de forward à chaud dans ce pilote ; un premier appel à nouvelle forme B16 a pris 272.4 ms. Le champ JSON est nommé cold_forward, mais il désigne le premier appel candidat après l’inférence de référence eager, pas une application entièrement froide ni un pilote Metal fraîchement initialisé. Les candidats suivants réutilisaient des kernels Metal déjà compilés, donc leurs temps de premier appel ne sont pas un classement contrôlé du coût de démarrage à froid. La mémoire MLX active/pointe de l’anglais court-1 compilé était d’environ 803.6/918.6 MiB dans le pilote ; le multilingue était d’environ 614.1/676.9 MiB. Ces mesures d’allocateur n’incluent pas toutes les allocations de compilateur côté hôte et n’établissent pas de limites de mémoire sous changement de forme non borné. Voir pilote de compilation anglais et pilote de compilation multilingue.

Quantization sélective : économies de stockage utiles, inadaptée comme prétention de vitesse

Le prototype appelle nn.quantize après avoir chargé le modèle dense FP16. Il ne sélectionne que les modules linéaires encoder.layers.*, avec une taille de groupe affine de 64, puis compile le modèle résultant. Les embeddings, les normes, la tête de décision, le scorer et la tête d’action restent en FP16. Cela évite de caster des poids entiers empaquetés à travers le loader dense actuel et évite la largeur d’entrée non divisible 1028/772 de la tête d’action. Aucun format de checkpoint quantizé ni contrat de chargement n’est livré. L’implémentation officielle de couche quantizée de MLX fournit ce mécanisme de sélection.

Le résultat multilingue 4 bits illustre pourquoi le petit jeu de test seul est insuffisant : ses 63 argmax de jeu sont restés les mêmes, mais 7 des 26 décisions de charge distincte ont changé. Ce sont des mesures d’accord contre FP16, pas des mesures d’exactitude de vérité terrain. Un changement absolu de probabilité de 0.8008 représente 80.08 points de pourcentage.

Sur les entrées pilotes anglaises court-16, le p50 de bout en bout eager/compilé FP16 était de 91.26/87.94 ms ; le 8 bits/4 bits compilé était de 96.66/93.20 ms. La quantization court-1 semblait un peu plus rapide dans cette exécution de présélection, tandis que les formes plus grandes non. La présélection multilingue à grande forme n’a pas non plus montré de gain de vitesse, mais ses exécutions séquentielles avaient une dérive substantielle. Ces observations justifient de rejeter une prétention de gain de vitesse ou de publication non qualifiée, pas d’attribuer des facteurs de ralentissement précis sans réplication quantizée entrelacée. Les travaux de quantization ultérieurs nécessitent une calibration sensible aux activations ou un fine-tuning, et une suite de qualité étiquetée représentative.

Sources brutes : anglais 8 bits, anglais 4 bits, multilingue 8 bits, multilingue 4 bits.

Metal écrit à la main : une fusion exacte GELU/porte implémentée et testée

kernels.py implémente un vrai kernel Metal personnalisé qui lit les deux branches MLP concaténées, calcule le même GELU fondé sur erf, multiplie par la porte et écrit une seule sortie. Il ne substitue ni tanh-GELU ni une approximation sigmoïde. Le kernel utilise les propres helpers erf et expm1 de MLX v0.32.2, en préservant leurs licences et mentions dans vendor/README.md. Il prend explicitement en charge uniquement FP16 et utilise le mode mathématique Metal sûr. Le guide officiel des kernels personnalisés décrit cette API et ses contrôles de mode mathématique.

Sur huit formes d’activation représentatives, 27,958,016 éléments de sortie FP16 générés aléatoirement avaient exactement les mêmes valeurs que l’opération d’origine. Le microbenchmark compare l’égalité numérique, pas le bit de signe de zéro. Les tests sur modèle complet à entrées modifiées correspondaient aussi exactement : 474/474 comparaisons de questions à travers les deux familles de modèles, plus les deux suites de 63 questions, avec zéro différence de logit, de logit d’action ou de probabilité calibrée.

Ce résultat de correction ne s’est pas traduit par un avantage de vitesse cohérent sur l’expression compilée fusionnée de MLX. Par exemple, à 1,312 tokens et une largeur intermédiaire de 2,624, le temps d’activation synchronisé par appel était de 0.378 ms pour le GELU-puis-porte eager, 0.268 ms pour mx.compile et 0.280 ms pour le kernel personnalisé. À 8,192 tokens et largeur 1,152, les valeurs correspondantes étaient 0.846/0.764/0.714 ms. Ces microbenchmarks incluent le surcoût de dispatch et de synchronisation et sont des sondes de présélection ; ce ne sont pas des mesures du temps d’exécution isolé de l’appareil. Les entrées complètes, les temps bruts et les contrôles d’égalité sont dans microbench.json.

Le kernel personnalisé a ensuite été installé dans chaque MLP d’encodeur et mesuré dans le modèle complet avec un ordre de candidats rotatif et des entrées changeantes :

Les exécutions appariées du kernel personnalisé complet utilisent une seconde instance de modèle avec des poids identiques pour que les implémentations non modifiée et personnalisée coexistent sans mutation ni capture compilée périmée. Leurs temps absolus ne doivent pas être comparés avec l’exécution antérieure d’élagage de tête. Les résultats mitigés modestes ne soutiennent pas la publication du kernel personnalisé comme amélioration générale de performance. Sources : données appariées Metal anglaises et données appariées Metal multilingues.

Où l’ingénierie personnalisée mériterait une investigation plus poussée

Le modèle appelle déjà mx.fast.scaled_dot_product_attention, mx.fast.rope et une normalisation de couche optimisée. Son chemin SDPA D64 à masque booléen est fusionné ; il n’y a pas d’interrupteur Flash Attention manquant qui explique un écart de 10×. L’attention locale parcourt encore des tuiles clé/valeur denses. Un vrai kernel à fenêtre bidirectionnelle pourrait sauter ces tuiles tout en préservant une distance inclusive <=64 et la sémantique de padding, mais son opportunité arithmétique sur le modèle complet est faible sur les entrées courtes et bornée sur les formes longues publiées. La revue de code existante et le rapport mathématique quantifient cette distinction.

Les prochains projets utiles, avec leurs exigences de preuve, sont :

• Attention à fenêtre sur entrée longue : spécialiser les bornes de tuiles pour D64, la vraie fenêtre bidirectionnelle et les lots paddés. Comparer au SDPA dense fusionné à 512/1024 tokens, puis dans le modèle complet. Ce kernel n’a pas été construit ni mesuré dans ce rapport.
• Épilogues et ordonnancement des kernels denses : étudier la fusion de l’épilogue MLP à porte dans le GEMM ou l’amélioration de l’ordonnancement matriciel à petit M. MLX utilise déjà des implémentations GEMM Metal spécialisées, donc les remplacer exige un vrai profil de dispatch/kernel et des gains mesurés pour les formes M/N/K exactes. Le résultat d’activation autonome montre pourquoi un autre kernel élément par élément seul est insuffisant.
• Mise en lots sensible à la longueur et préparation CPU partagée : préserver des IDs d’entrée exacts tout en tokenisant une fois le texte d’état partagé avant de construire chaque séquence de question, et éviter de padd des petits éléments vers des éléments longs sans rapport. Le pilote multilingue long-8 a passé environ 13.1 ms à préparer les entrées, contre des centaines de millisecondes de bout en bout. Même éliminer entièrement cette préparation ne produirait pas 10× sur cette charge. La latence de file d’attente et le nombre d’inférences uniques doivent faire partie de toute prétention sur la mise en lots.
• Un élève plus petit qui répond conjointement : si 10× est une exigence produit, distille ou repense le modèle pour retirer la majeure partie du travail dense ou répondre à beaucoup de questions fixes avec un seul encodage contextuel. Cela change le modèle appris et nécessite un entraînement/évaluation étiqueté représentatif ; ce n’est pas une optimisation exacte de portage. Réutiliser un état contextuel/KV arbitraire entre les questions dans l’encodeur bidirectionnel actuel est invalide.

Huit sondes GEMM autonomes de projection d’entrée d’encodeur FP16 ont atteint 0.66–11.55 TFLOP/s, synchronisation par appel incluse. La grande sonde anglaise M=4096, N=5248, K=1024 a atteint 11.55 TFLOP/s ; la sonde multilingue M=8192, N=2304, K=768 a atteint 7.96 TFLOP/s. Ce sont des valeurs de débit observées, pas des spécifications de pic matériel ni des bornes supérieures du débit du graphe complet. Les mesures à petit M sont particulièrement dominées par les coûts de soumission et de synchronisation ; un graphe streamé les amortit différemment. Elles montrent quelles formes méritent un profilage, pas une preuve qu’aucun meilleur kernel ne peut exister. Les budgets de débit 10× à travail égal du rapport mathématique restent des exigences théoriques plutôt que des capacités mesurées de l’appareil.

Reproduction et décision de publication

Les scripts utilisent le .venv existant et les checkpoints épinglés locaux. Exécute les commandes GPU séquentiellement, jamais en parallèle du benchmark formel :

# Screening: repeat for eager, compiled, blocks, q8, q4, selected-compiled.
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant compiled --iterations 12 --warmup 4 --quality \
  --output experiments/engineering/reproduced-compiled.json

# Primary confirmation, including 50 genuinely different questions.
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 32 \
  --output experiments/engineering/reproduced-laya-paired.json
.venv/bin/python -m experiments.engineering.paired \
  --model laya-multilingual --iterations 16 --cases short1,short16,long1,long8 \
  --output experiments/engineering/reproduced-multilingual-paired.json

# Hand-written kernel microbench and complete-model comparison.
.venv/bin/python -m experiments.engineering.microbench
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 16 --cases short1,short16,long1 --metal \
  --output experiments/engineering/reproduced-metal-paired.json
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant metal-compiled --iterations 5 --warmup 3 \
  --cases short1 --quality --output experiments/engineering/reproduced-metal-quality.json

# CPU-only paired analysis.
.venv/bin/python -m experiments.engineering.analyze

Tous les fichiers Python expérimentaux passent le formatage et le lint Ruff. Le runtime stable, les résultats de benchmark d’origine et les checkpoints FP16 publiés restent les artefacts de release. La compilation et l’élagage exact de la tête finale sont des optimisations futures optionnelles crédibles après une politique de forme froide/cache et une validation de qualité plus large ; les gains mesurés ne justifient pas d’ajouter silencieusement une latence de compilation ou un kernel personnalisé au chemin par défaut. Aucun gain de vitesse de 10×, aucun checkpoint quantizé prêt pour la production et aucun gain mesuré de kernel à fenêtre locale n’est revendiqué.