Axon: ett programspråk som låter AI-modeller köras överallt – och snabbare

arXiv cs.AI

Forskare har tagit fram Axon, ett domänspecifikt programspråk (ett språk designat för ett specifikt användningsområde) som låter utvecklare skriva en LLM-arkitektur en enda gång och sedan automatiskt kompilera den till populära ramverk som PyTorch, JAX och MLX. I 467 prestandatest på modeller mellan 135 miljoner och 32 miljarder parametrar visade Axon medianförbättringar på upp till 107% jämfört med standardimplementationer – störst förbättring på MLX. Det intressanta är egentligen bakgrundsargumentet: hela open source-ekosystemet för språkmodeller lutar tungt mot en enda plattform, och Axon är ett försök att minska det beroendet.

Djupdykning

Hela AI-världens öppna modellekosystem hänger på Hugging Face – en enda plattform vars bibliotek "Transformers" är den defacto-standarden för att definiera och köra språkmodeller. Axon är ett försök att bryta det beroendet genom att introducera ett eget programmeringsspråk (en "domain-specific language", alltså ett språk byggt för ett specifikt syfte) där du skriver din modellarkitektur én gång och sedan låter kompilatorn automatiskt generera optimerad kod för PyTorch, JAX, MLX eller vLLM – ungefär som hur Java-kod kan köras på vilken dator som helst utan att skrivas om. Haskellinspirationen är inte slumpmässig: Haskell är känt för sitt starka typsystem som fångar fel vid kompilering snarare än vid körning, vilket här används för att garantera att tensorernas dimensioner alltid stämmer – något som annars är en klassisk källa till frustration och buggar i ML-kod. Det de flesta missar är att den verkliga risken inte är att Hugging Face stänger ned imorgon, utan att hela forskningsvärlden omedvetet har delegerat kontroll över optimeringsinfrastrukturen till ett privat företags designbeslut. När Axon visar 107% snabbhetsökning på MLX (Apples ramverk för Mac-chips) beror det till stor del på att Transformers-implementationerna aldrig var optimerade för den plattformen från början – de porterades i efterhand. Det är symptomatiskt: när ett bibliotek blir monopol prioriteras populära hårdvarumål, och alla andra får suboptimal prestanda. Om Axon eller något liknande faktiskt får genomslag förskjuts maktbalansen från "vilket ramverk stödjer din modell" till "vilket ramverk kompilerar Axon bäst" – en subtil men fundamental skillnad som öppnar för hårdvarukonkurrens på allvar.