Een gefuseerde forward-kernel voor lineaire cross-entropy op Apple-GPU's

Denis Ineshin · 2026-07-19

Geheugenbesparing, grenzen aan tiling en lessen uit mislukte performancemodellen

Het trainen van een taalmodel alloceert vaak een enorme logits-matrix alleen om de cross-entropy te berekenen. Voor N tokens en een vocabulaire van grootte V heeft die matrix de vorm (N, V). Toch heeft de loss per rij alleen een log-sum-exp en één target-logit nodig. We bouwden een gefuseerde MLX/Metal-forward-kernel die deze waarden in tiles berekent en de volledige matrix nooit opslaat. Bij N=8192, V=151936 en hidden dimension D=4096 duurde de gefuseerde forward 2.105 s tegenover 1.286 s voor gematerialiseerd MLX. Een publieke geheugenrun noteerde ook een veel lagere piek voor het gefuseerde pad. Ongelijke warm-up-baselines maken echter een zuivere geheugenverhouding per aanroep onmogelijk; paragraaf 3 geeft de volledige metingen en legt de beperking uit.

Het optimalisatieproces leverde twee minder voor de hand liggende resultaten op. Ten eerste verhoogde het forceren van mx.eval na elke pure-MLX-chunk het geheugengebruik, omdat MLX de getracete tussenresultaten vasthield. Ten tweede voorspelde een load-reuse-model één snellere register-tile en faalde meteen bij de volgende. Grotere tiles werden trager. Vervolgcontroles sloten verschillende verklaringen uit, maar konden geen enkele oorzaak isoleren. Het resterende bewijs wijst op een afweging tussen hergebruik van loads op broncodeniveau en resourcegebruik van de gecompileerde kernel — geen van beide hebben we rechtstreeks op hardwareniveau gemeten.

Dit artikel documenteert de implementatie en die metingen. Het introduceert geen nieuw cross-entropy-algoritme. Logit-vrije lineaire cross-entropy en streaming-softmax-reducties zijn bestaand werk, waaronder Cut Cross-Entropy. De optimalisatieladder is bovendien een retrospectief verslag, geen reproduceerbare benchmark-suite: de historische broncoderevisies en volledige launch-instellingen zijn niet vastgepind. Het forward-resultaat impliceert niet dezelfde reductie voor een volledige trainingsstap.

1. De loss streamen zonder volledige logits

Laat H een (N, D)-matrix van hidden states zijn en W een (V, D)-outputprojectie. Gewone cross-entropy berekent eerst alle logits

Z = H Wᵀ

en berekent dan, voor target y_i,

loss_i = logsumexp(Z_i) - Z_i,y_i

De loss vereist geen willekeurige toegang tot de volledige rij logits. Ze vereist slechts twee statistieken per rij: de log-sum-exp en de target-logit. Dat maakt een streaming-formulering mogelijk.

De kernel voegt elke vocabulary-tile samen in een lopende online log-sum-exp. Laat (m, s) het huidige maximum en de door dat maximum geschaalde som van exponenten bevatten. Voor een nieuwe tile met maximum m_t is de update

m' = max(m, m_t)
s' = s · exp(m - m') + Σ_j exp(z_j - m')

Na de laatste tile, met (m, s) als de uiteindelijke lopende toestand, geldt logsumexp = m + log(s). De kernel legt de target-logit vast wanneer de vocabulaire-index ervan binnen de huidige tile valt. Alleen de lopende toestand en de uiteindelijke per-token-waarden blijven over. De (N, V)-matrix bestaat nooit.

In de hier bestudeerde implementatie construeert Python één lazy Metal-dispatch per vocabulary-tile. De kernel wordt één keer JIT-gecompileerd en de gecachete pipeline wordt over de hele keten hergebruikt. Bij V=151936 en een tile van 8192 kolommen vormt de productievorm 19 afhankelijke dispatches. Elke aanroep verbruikt de N-element-accumulators van de vorige en produceert de volgende. Eén enkele evaluatie voert de voltooide keten uit en synchroniseert die.

De scope is alleen de forward. Het uitgeleverde trainingspad bewaart de forward-log-sum-exp en gebruikt een aparte, in chunks opgedeelde pure-MLX-backward die de benodigde logits tile voor tile opnieuw genereert. Het gebruikt nog geen gefuseerde Metal-backward, dus de geïsoleerde metingen in dit artikel mogen niet als forward-plus-backward-resultaten worden gelezen.

2. Waarom chunking in puur MLX toch tekortschoot

Het eerste prototype implementeerde dezelfde vocabulaire-decompositie met gewone MLX-operaties. Het was numeriek correct, ondersteunde dense en gekwantiseerde output heads en toonde aan dat een grote loss verwerkt kon worden zonder één monolithische logits-tensor te vormen.

Het legde ook een contra-intuïtieve evaluatieregel bloot. mx.eval aanroepen na elke chunk terwijl MLX de gradiënt tracete, hield de getracete tussenresultaten vast in plaats van elke chunk vrij te geven. In een vroeg forward-plus-backward-experiment bij N=8192 bereikte geforceerde evaluatie per chunk 55.98 GiB, tegenover 9.13 GiB toen de chunks lazy bleven en de graaf één keer aan het eind werd geëvalueerd. De eager-versie gebruikte ongeveer zes keer zoveel geheugen.

Door de chunks lazy te laten, daalde het piekgeheugen van forward-plus-backward bij de gemeten vormen tot ruwweg de helft. Bij N=4096 kostte het echter 1.31 keer zoveel als het dense pad bij de dense head en 1.71 keer zoveel bij de gekwantiseerde head. Deze verkennende resultaten motiveerden kernelfusie. Het zijn geen definitieve library-benchmarks en ze zijn niet rechtstreeks vergelijkbaar met de forward-only-getallen hieronder.

Chunking drukte de juiste wiskundige decompositie uit, maar kon de levensduur en fusie van elk tussenresultaat niet sturen. Een eigen Metal-kernel kon tile-lokale logits en reducties in registers houden, waarmee de beoogde geheugengrens een eigenschap van de implementatie werd in plaats van een hoop op goede scheduling.

3. Experimentele opzet

Alle gerapporteerde metingen kwamen van één Apple M1 Max met 32 GB unified memory (gedeeld geheugen). De verkennende kernelladder gebruikte MLX 0.31.2 en macOS 26.5.1. De vastgelegde reproductie van de loss-laag gebruikte MLX 0.32.0 en macOS 26.5.2. Beide gebruikten bfloat16-invoer met V=151936 en D=4096.

Elke getimede conditie draaide in een vers proces. Vóór het timen initialiseerde en evalueerde het proces zijn invoer en warmde daarna de JIT-compiler en de allocator op. De tabel rapporteert de mediaan van drie gesynchroniseerde runs.

Vóór het benchmarken moest elke kernelvariant overeenkomen met een fp32-forward-referentie. De verkennende suites dekten 10 tot 19 gevallen per variant, inclusief staart- en uitlijningsgevallen. De grootste opgeslagen absolute verschillen in per-token-loss waren 4.8e-6 voor fp32-invoer en 1.9e-6 voor bfloat16-invoer. Deze controles dekken alleen de forward, niet de correctheid van de gradiënt.

De verkennende artefacten van de referentievergelijking zijn niet openbaar. De uiteindelijke library publiceert wél een forward-benchmark met de MLX- en OS-versies, de pakket-/bronidentiteit, de vorm, het aantal herhalingen en de ruwe wall-times. Dat artefact dateert van vóór het script_sha-veld van de driver, dus het kan de exacte scriptrevisie die de rijen voortbracht niet identificeren.

De vastgelegde run met MLX 0.32.0 leverde het volgende geïsoleerde forward-resultaat op:

De timing binnen dezelfde sessie geeft een forward-vertraging van 1.64×. De geheugenkolommen vergen meer zorg. De benchmark warmt elke implementatie op, houdt de warm-up-loss in leven, wist de allocator-cache en reset de piekteller. Het gematerialiseerde pad begint zijn meetvenster daardoor met 4.6367 GiB meer actief geheugen dan het gefuseerde pad.

Dat verschil in baseline domineert een verhouding tussen de twee marginale kolommen, terwijl de waarde 0.0006 GiB bovendien op vier decimalen is afgerond. De tabel toont verschillend geheugengedrag, maar levert geen zuivere reductieverhouding per aanroep. Zo'n verhouding zou ruwe byte-tellingen vereisen en een verse run vanaf een gemeenschappelijke baseline met alleen invoer en zonder levende warm-up-graaf.

4. Het load-reuse-model werkt één keer

Het eerste Metal-ontwerp wees aan meerdere rijen samen één SIMD-group (groep van 32 threads, Apples tegenhanger van een warp) toe. Elke lane liep door de hidden dimension, accumuleerde partiële dotproducten in fp32 en gebruikte simd_sum om de lanes te combineren. De belangrijkste experimentele variabele was de register-tile: hoeveel rijen (R) en vocabulaire-kolommen (C) één SIMD-group samen verwerkte.

Een R × C-tile met bfloat16-vector-loads van vier elementen voert 4RC multiply-accumulates (MAC) uit en vraagt daarbij (R+C) vectoren op uit de binnenlus op broncodeniveau. We definieerden de nominale load-reuse-verhouding

load reuse = 4RC / (8(R+C)) MAC per requested byte.

Deze verhouding suggereerde dat de overstap van een 4 × 1-tile naar 4 × 4 de throughput met ongeveer 2.5× zou kunnen verhogen. Bij N=8192 steeg de gemeten throughput van 294.8 naar 814.9 G MAC/s, een winst van 2.76×. De richting en de schaal pasten bij dit eerste paar. De verhouding is geen arithmetische intensiteit op hardwareniveau: ze negeert stores, het werk van de online-reductie, verkeer op cacheniveau, accumulator-verkeer en de instructiemix. Toch was het verleidelijk om de eerste overeenkomst als een voorspellend model te behandelen.

De volgende tile was trager. De overstap van 4 × 4 naar 4 × 8 verhoogde de nominale load-reuse met 33%, maar verlaagde de throughput bij beide gemeten sequentielengtes.

De tabel gebruikt de historische labels voor elke experimentele trede. Ze biedt geen openbare, onveranderlijke koppeling van elk label aan zijn broncoderevisie en volledige launch-instellingen. In het bijzonder zijn de vocabulary-tile- en dispatch-geometrie niet per rij bewaard. De tabel ondersteunt een retrospectieve vergelijking, geen onafhankelijke herhaling van de ladder.

De 4 × 8-tile had 33% meer nominale load-reuse dan 4 × 4, en was toch 8% trager bij N=8192 en ongeveer 20% trager bij N=2048. De 8 × 8-versie stortte nog ernstiger in: in een vergelijking met dezelfde vorm en dezelfde vocabulary-tile van 4096 kolommen bij N=2048 haalde ze 182.1 G MAC/s tegen 635.0 voor 4 × 8.

5. Vier controles perken de verklaringen in

Verschillende verklaringen leken plausibel. De row-stride (rijafstand) was precies 8192 bytes. Dat is gelijk aan de L1-datacachegrootte van de M1-GPU zoals gerapporteerd in de door reverse-engineering verkregen architectuurnotities van metal-benchmarks, dus extra gelijktijdige streams hadden cache-set-conflicten kunnen veroorzaken. Sommige varianten gebruikten oorspronkelijk een inline kopie van de benchmark-harness. Het 8 × 8-experiment wijzigde ook zijn tilevorm en broncodestijl. De eerste vergelijking bij kleine vorm gaf beide varianten niet dezelfde kans om de GPU te bezetten.

Vier controles perkten het veld in:

1. D wijzigen van 4096 naar 4160 veranderde de row-stride van 8192 naar 8320 bytes. De throughput-verhouding 4 × 8/4 × 4 bleef 0.80 bij N=2048. Dit weerlegde het voorgestelde macht-van-twee-cache-set-mechanisme voor de 4 × 8-regressie.
2. Beide varianten achter elkaar door het gedeelde script draaien reproduceerde de eerdere snelheden tot op 1%, wat de gekopieerde harness als oorzaak uitsloot.
3. De 4 × 4-kernel herschrijven in de array-en-unroll-stijl van de grotere tiles reproduceerde de implementatie met expliciete scalairen. Het broncode-idioom was niet verantwoordelijk.
4. 8 × 8 en 4 × 8 vergelijken bij dezelfde N=2048 en vocabulary-tile-grootte liet een resterende vertraging van 3.5× over. Dit nam het verschil in aantal dispatches weg en verkleinde het verzadigingsprobleem bij kleine vorm, maar de varianten hadden nog steeds verschillende rijblokgeometrie.

We compileerden de gegenereerde Metal-bron ook via het Metal-framework en inspecteerden MTLComputePipelineState.maxTotalThreadsPerThreadgroup. Het maximum van het apparaat was 1024 threads. De 4 × 4-kernel compileerde met een plafond van 512, terwijl 4 × 8 en 8 × 8 beide op 384 compileerden. Deze pipeline-legaliteitsgrens correleert met zwaarder resourcegebruik van de gecompileerde kernel. Ze meet geen behaalde occupancy (bezettingsgraad van de GPU), residente SIMD-groups of registers. Het experiment toonde ook niet aan dat het lagere plafond de residency voor de daadwerkelijke dispatch veranderde.

De controles stellen echte regressies vast en correleren ze met zwaarder resourcegebruik van de gecompileerde kernel. Ze wijzen geen oorzaak aan. Verminderde occupancy is één kandidaat voor 4 × 8. Voor 8 × 8 overschreed de geschatte levende toestand het door reverse-engineering verkregen model met 128 GPR's van 32 bits in Dougall Johnsons Apple-GPU-notities. Na de op elkaar afgestemde controles bleef een grote vertraging over. Register spill (registers naar geheugen uitgeweken) is een kandidaat, maar er is geen compiler-ISA, spill-statistiek, geheugenverkeersteller of meting van behaalde occupancy vastgelegd. Geen van beide mechanismen is hier causaal aangetoond.

6. Matrix-tiles verminderen scalaire toestand

Het volgende ontwerp gebruikte simdgroup_matrix-operaties om hergebruik te verhogen zonder de gewone scalaire toestand op dezelfde manier te laten groeien. De vroege treden waren trager dan de beste register-array-kernel. Grotere matrix-tiles dichtten dat gat uiteindelijk.

De prestaties verbeterden tot aan de 32 × 32-tile, die ruwweg 32 fp32-accumulator-elementen per lane gebruikte. Eén tiledimensie verdubbelen bracht de throughput vervolgens terug van 2423.7 naar 1403.2 G MAC/s. In beide ontwerpfamilies hielp meer datahergebruik alleen zolang de extra toestand per lane de gecompileerde kernel niet wezenlijk zwaarder maakte.

Een compact model voor het waargenomen gedrag is

throughput kan beperkt worden door zowel data-aanvoer als resourcegebruik van de gecompileerde kernel.

Dit is een kwalitatieve ontwerpheuristiek, geen gefit model of universele Apple-GPU-wet. Het experiment mat geen geheugenverkeer, bandbreedteplafonds, registeraantallen, spills of behaalde occupancy, dus het kan bandbreedte, throughput van load-instructies, cachegedrag en latency-hiding niet onderscheiden. De conclusie is beperkter: betere nominale load-reuse kan alsnog verliezen wanneer ze de kernel op een andere, ongemeten manier zwaarder maakt.

7. Forward-besparingen voorspellen geen besparingen over de hele stap

De gefuseerde kernel haalt de logits-allocatie per constructie uit de forward-loss-laag. De publieke geheugenrun is consistent met een grote reductie, maar de ongelijke opgewarmde baselines ondersteunen geen zuivere verhouding. In een volledige trainingsstap blijven de transformer, de optimizer, de modelgewichten en de attention-backward over. Eén allocatie verwijderen garandeert geen evenredige toename van de maximale context.

De releasemetingen van versie 0.1.0 van het project laten zien waarom scope ertoe doet. Bij gelijkblijvende stock attention in de loss-vergelijking rapporteert het releaseverslag dat gefuseerde cross-entropy op zichzelf de maximale trainbare context niet wezenlijk verlengde. Ze schrijft de piek bij lange context toe aan de backward van stock attention. Omdat de ruwe historische sweep niet openbaar is, blijft dit een door de bron gerapporteerd resultaat in plaats van een onafhankelijk reconstrueerbare vergelijking.

De loss-kernel verwijdert dus verspilling in de loss-laag zonder het geheugenprobleem van de training als geheel op te lossen. Hij verhoogt de maximale context niet automatisch.

Ook de snelheid verandert met de scope. De geïsoleerde gefuseerde forward was 1.64× trager dan het sterk geoptimaliseerde gematerialiseerde MLX-pad, maar de loss-laag is slechts een deel van een trainingsstap. Latere end-to-end-metingen vonden een veel kleinere vertraging op stapniveau. Die vertraging hangt af van het model, de sequentielengte, de attention-implementatie en het backward-pad. Ze kan niet uit de geïsoleerde tabel alleen worden afgeleid.

De publieke bench_train_step.py documenteert de latere methode op stapniveau. De conclusie over context voor alleen de loss-laag staat onder de bekende beperkingen van 0.1.0 in de changelog. De huidige northstar_context_sweep.py wijzigt zowel attention als loss. Het documenteert daarom de latere vergelijking van het hele product, niet dit resultaat voor alleen de loss.

8. Lessen uit de mislukte modellen

Voor toekomstig MLX-kernelwerk:

• Beheer evaluatiegrenzen expliciet. In een lazy, differentieerbaar systeem kan geforceerde evaluatie binnen een lus levensduren juist verlengen in plaats van verkorten.
• Vertrouw een performancemodel niet na één geslaagde voorspelling. De nominale load-reuse-verhouding paste bij het 4 × 4-paar en faalde meteen daarna.
• Vergelijk overeenkomende vormen, tilegroottes en kansen om de GPU te bezetten. De eerste 8 × 8-vergelijking mengde een mogelijk resource-druk-effect met inactieve cores en verschillende aantallen dispatches.
• Behandel pipeline-grenzen als proxy's. Een verlaagde maxTotalThreadsPerThreadgroup correleert met een zwaardere gecompileerde kernel, maar onthult geen registeraantal, residency-niveau of spill-gebeurtenis.
• Leg mislukte voorspellingen vast. De hypotheses over stride, harness en broncode-idioom waren redelijk, en dat ze verworpen werden, verklaart waarom het uiteindelijke ontwerp veranderde.
• Inspecteer de baseline achter een geheugenverschil. Marginale pieken zijn alleen vergelijkbaar wanneer ook de persistente toestand op de resetgrens vergelijkbaar is.

9. Beperkingen

Deze studie gebruikte één M1 Max, één primaire productievorm, bfloat16-invoer en vastgepinde MLX-versies. Absolute snelheden kunnen veranderen met de machine, het OS, de MLX-release, de thermische toestand en de JIT-compiler. Medianen over drie runs leggen grote ontwerpverschillen vast, niet kleine variantie.

De experimenten bevestigden de verklaringen op basis van occupancy en register spill niet met ISA- of tellerbewijs. Het publieke geheugenprotocol gebruikt bovendien ongelijke warm-up-baselines, dus het kan geen zuivere geheugenverhouding per aanroep ondersteunen. Dit artikel rapporteert die beperking in plaats van het experiment opnieuw uit te voeren. Het dekt ook alleen de forward-kernel. Het uitgeleverde backward-pad en de varianten met gekwantiseerde head vallen buiten de scope.

10. Reproduceerbaarheid en bronnen

De uiteindelijke loss-laag-benchmark komt uit het script bench_loss_layer.py. Het vastgelegde JSON-artefact legt de conditie-identiteiten en alle drie de wall-time-metingen vast. De verkennende ladder heeft geen gelijkwaardige openbare bundel, dus de resultaten ervan blijven door de bron gerapporteerd en kunnen niet onafhankelijk worden herhaald. De openbare JSON identificeert de gemeten pakketbron, maar dateert van vóór het script_sha-veld van de driver. De exacte historische driverrevisie is daarom onbekend.

Cut Cross-Entropy ontwikkelt de bredere logit-vrije lineaire cross-entropy-aanpak. Liger-Kernel biedt gefuseerde trainingskernels voor andere accelerator-stacks.

11. Conclusie

De gefuseerde forward verwijdert de (N, V)-logits-matrix, maar bespaart geheugen ten koste van tijd. Bij de op de M1 Max gemeten vorm draaide hij 1.64× trager dan gematerialiseerd MLX. Het publieke geheugenresultaat wijst op een grote reductie, maar de ongelijke warm-up-baselines verhinderen een betrouwbare verhouding per aanroep.

De optimalisatieladder is om een andere reden van belang. Een load-reuse-model voorspelde de verbetering bij 4 × 4 correct en faalde bij 4 × 8; de 32 × 32-matrix-tile herstelde later de best waargenomen throughput voordat een grotere tile opnieuw regresseerde. Zonder register- of occupancy-tellers blijft de oorzaak open. Het onderbouwde resultaat is bescheidener en nuttiger: datahergebruik helpt alleen zolang de resourcekosten van de gecompileerde kernel onder controle blijven, en hergebruik op broncodeniveau alleen kan die grens niet voorspellen.

Opgesteld 2026-07-14. Laatst bijgewerkt 2026-07-19. Denis Ineshin.