ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

Stel je voor dat het internet een gigantische bibliotheek is waar de boeken elke dag slimmer worden. Deze "slimme boeken" worden Large Language Models (LLM's) genoemd en ze kunnen verhalen schrijven, wiskundige problemen oplossen en chatten als mensen. Maar om echt slim te worden, moeten deze modellen enorm zijn en miljarden feiten in hun geheugen opslaan. Recentelijk ontdekten wetenschappers een slimme truc genaamd "Mixture of Experts" (MoE). In plaats van één gigantisch brein dat alles probeert te onthouden, bouwt MoE een team van kleinere specialisten. Wanneer je een vraag stelt, wordt niet het hele team wakker; er worden slechts een paar experts opgeroepen die het antwoord weten. Het is alsof je een bibliothecaris vraagt om een boek te halen: in plaats van het hele bibliotheekpers wakker te maken, bel je alleen voor de geschiedenisafdeling.

Er is echter een addertje onder het gras. In de echte wereld is het oproepen van deze specialisten een rommelige aangelegenheid. De experts zijn verspreid overal en de bibliothecaris moet heen en weer rennen om de juiste boeken te halen, en dan weer terugrennen om de antwoorden samen te brengen. Dit "heen en weer rennen" creëert een verkeersopstopping. De computer besteedt meer tijd aan het wachten tot data aankomt dan aan het daadwerkelijk nadenken erover. Dit is de "memory wall", een bottleneck die zelfs de snelste supercomputers vertraagt. De vraag is: hoe bouwen we een bibliotheek waar de boeken direct naast de lezers staan en waar de renners een perfect snelwegsysteem hebben om verkeersopstoppingen te vermijden?

---

Het Probleem: Een Verkeersopstopping in het Brein

Het artikel introduceert een nieuwe architectuur genaamd ThAME (3D Memory-Enabled Heterogeneous Accelerator) om precies dit probleem op te lossen. De auteurs leggen uit dat huidige computers als een drukke stad zijn waar het "denkende" deel (de CPU of GPU) ver weg is van het "geheugen" deel (waar de data leeft). Wanneer een Large Language Model de Mixture of Experts-truc gebruikt, creëert dit een chaotische situatie.

Stel je een schoolkantine voor waar leerlingen (tokens) eten moeten halen bij verschillende lunchbalies (experts). In een normale kantine gaat iedereen naar dezelfde rij. Maar in een MoE-systeem wordt elke leerling op basis van wat hij wil eten naar een andere, willekeurige balie gestuurd. Sommige balies krijgen een enorme menigte, terwijl andere niemand krijgen. De leerlingen moeten door de kantine rennen om hun eten te halen, en dan weer terugrennen naar een centrale tafel om hun maaltijden te mengen. Dit creëert een "scatter-gather" verkeerspatroon: een chaotische mix van leerlingen die alle kanten op rennen, wat botsingen en lange wachttijden veroorzaakt. De auteurs stellen dat standaard computerchips niet gebouwd zijn voor dit soort onvoorspelbaar, rommelig verkeer. Ze raken vastgelopen in een gridlock, en het hele systeem vertraagt.

De Oplossing: Een 3D-Stad met een Speciaal Snelweg

Om dit op te lossen, stellen de auteurs ThAME voor, wat lijkt op het bouwen van een gloednieuwe, high-tech stad voor deze AI-modellen. Ze gebruiken drie hoofdonderdelen om het te laten werken:

1. De Juiste Tools voor de Juiste Taak (Heterogene Geheugen):
   Het artikel wijst erop dat niet alle geheugen hetzelfde is. Sommige delen van de AI moeten heel vaak beschreven worden (zoals een whiteboard), terwijl andere delen alleen gelezen hoeven te worden van een enorme boekenkast.

   • Voor de "whiteboard"-delen (het attention-mechanisme) gebruiken ze DRAM, wat snel is en gemakkelijk te herschrijven is, maar veel ruimte inneemt.
   • Voor de "gigantische boekenkast"-delen (de expert weights) gebruiken ze FeFET-NAND, een type 3D-geheugen dat ongelooflijk dicht is en niet vaak herschreven hoeft te worden.
   • Denk aan een bibliotheek waar de naslagwerken gestapeld zijn in een enorme, hoog-dichtheid toren (FeFET), terwijl de actieve werkstations een ander, sneller type bureau gebruiken (DRAM). Door deze lagen verticaal (3D) op elkaar te stapelen, kunnen ze een enorme hoeveelheid data direct naast de processoren plaatsen, zodat de "renners" niet ver hoeven te reizen.

2. Het Speciale Snelweg (De NoC):
   Dit is de grootste innovatie van het artikel. Zelfs met het geheugen dichtbij, moeten de "renners" (data) nog steeds tussen de verschillende expert-balies bewegen. De auteurs realiseerden zich dat standaard wegennetwerken (zoals een simpel raster of een ring) falen wanneer het verkeer onvoorspelbaar is.

   • Ze ontwierpen een hiërarchische Network-on-Chip (NoC). Stel je een stad voor waar lokale wijken kleine, privéwegen (crossbars) hebben om het lokale verkeer snel af te handelen. Vervolgens zijn deze wijken verbonden door een slim, boomstructuur-achtig snelwegsysteem dat de grote verkeersstromen tussen verschillende gebieden kan afhandelen.
   • Dit systeem werd geoptimaliseerd met een complexe wiskundige methode om ervoor te zorgen dat verkeersopstoppingen worden geminimaliseerd, ongeacht hoe de leerlingen verdeeld zijn (welke experts druk zijn). Het is als een verkeersregelsysteem dat automatisch auto's omleidt om ongelukken te voorkomen voordat ze gebeuren.

3. De Slimme Dispatcher:
   Het systeem bevat een scheduler die fungeert als een slimme verkeersregelaar. Hij kijkt naar de menigte en beslist precies hoeveel "renners" (computercores) hij aan elk expert-balie toewijst, zodat iedereen ongeveer tegelijk klaar is. Dit voorkomt dat één trage expert de hele groep ophoudt.

Wat Ze Hebben Gevonden

De auteurs hebben dit niet alleen op papier gebouwd; ze hebben gedetailleerde simulaties uitgevoerd om te zien hoe het zou presteren. Ze vergeleken ThAME met de beste bestaande systemen (zoals Stratum en H3D-T) en standaard GPU's.

• Snelheid: In hun simulaties was ThAME ongelooflijk snel. Het was tot wel 15,7 keer sneller dan de beste bestaande hardware bij het genereren van tekst. Dit betekent dat als een standaard computer 10 seconden nodig heeft om een paragraaf te schrijven, ThAME dit in minder dan een seconde kan doen.
• Energie: Het was ook veel efficiënter en verbruikte tot wel 9,8 keer minder energie voor dezelfde taak. Dit is enorm belangrijk, omdat het draaien van deze enorme modellen meestal veel elektriciteit opslokt.
• Robuustheid: Het artikel laat zien dat zelfs als de geheugentechnologie niet perfect is (bijvoorbeeld als de leessnelheid trager is dan verwacht), ThAME nog steeds goed presteert omdat het zoveel extra bandbreedte ingebouwd heeft.

Wat Ze Niet Hebben Gedaan

Het is belangrijk om op te merken wat dit artikel niet claimt. De auteurs hebben geen fysieke chip in een fabriek gebouwd en deze getest in een echte serverfarm. Al hun resultaten komen uit cycle-accurate simulaties — zeer gedetailleerde computermodellen die nabootsen hoe de hardware zou reageren. Ze claimen ook niet dat ze elk probleem in AI hebben opgelost; ze hebben zich specif specifiek gericht op de "decode"-fase (het genereren van tekst woord voor woord), wat het moeilijkste deel is voor huidige computers. Ze erkenden ook dat de "prefill"-fase (het lezen van de initiële prompt) wordt afgehandeld door een standaard, krachtige processor (een TPU), en hun innovatie is specifelijk bedoeld voor het rommelige, geheugen-intensieve deel dat daarop volgt.

De Kern van het Verhaal

Het artikel suggereert dat door verschillende soorten geheugen te mengen en een op verkeer gericht snelwegsysteem binnen de chip te bouwen, we de Mixture of Experts-modellen eindelijk snel en efficiënt kunnen laten draaien. Hoewel dit momenteel een simulatie is, zijn de resultaten veelbelovend genoeg om te suggereren dat deze "3D-stad met een slim snelweg"-aanpak de sleutel kan zijn tot het ontsluiten van de volgende generatie super-slimme AI zonder ons elektriciteitsnet te overbelasten.