RSH-SpMM: A Row-Structured Hybrid Kernel for Sparse Matrix-Matrix Multiplication on GPUs

Het artikel introduceert RSH-SpMM, een hybride kernel voor GPU's die door middel van adaptieve rij-partitionering en een RS-Tile-representatie de prestaties van Sparse Matrix-Matrix Multiplication (SpMM) bij onregelmatige sparsiteit aanzienlijk verbetert met een versnelling van 1,27x tot 6,13x ten opzichte van bestaande methoden.

De kern

Stel je voor dat je een enorme bibliotheek hebt met miljoenen boeken, maar de meeste boeken zijn eigenlijk blanco pagina's. Je wilt een lijst maken van alle boeken die een bepaald woord bevatten. Dit is wat computers doen bij SpMM (Sparse Matrix-Matrix Multiplication): ze zoeken naar de "belangrijke" stukjes informatie in een zee van "lege" ruimte.

Het probleem is dat deze "belangrijke" stukjes heel willekeurig verspreid liggen. Soms zitten ze in een rijtje, soms springen ze over de hele bibliotheek.

Het Probleem: De Twee Types Werknemers

De moderne computerschijven (GPUs) die dit moeten doen, hebben twee soorten werknemers:

1. De "Tensor Cores" (De Super-Snelheidswagens): Deze zijn ontworpen om enorme blokken gegevens tegelijk te verwerken. Ze zijn razendsnel, maar ze zijn extreem kieskeurig. Ze willen alleen werken als de boeken perfect in een vierkant blokje passen. Als er een blanco pagina in het blokje zit, stopt de wagen en moet hij wachten. Bij willekeurige data (zoals in echte bibliotheken) zitten er vaak veel blanco pagina's, waardoor deze super-snelheidswagens vaak stil staan.
2. De "CUDA Cores" (De Slimme Handwerkslieden): Deze zijn iets langzamer, maar ze zijn flexibel. Ze kunnen prima werken met losse, willekeurige boeken. Ze hoeven geen perfecte blokken.

De huidige situatie:
De meeste software probeert ofwel alles te doen met de Super-Snelheidswagens (wat vaak vastloopt omdat de blokken niet passen) of alles met de Handwerkslieden (wat te langzaam is). Soms proberen ze een mix, maar dan verdelen ze het werk op een grof niveau: "Deze hele sectie gaat naar de wagens, die andere naar de handwerkslieden." Het probleem is dat binnen één sectie soms wel en soms geen perfecte blokken zijn. Het resultaat is dat de wagens vaak stilstaan en de handwerkslieden overbelast raken.

De Oplossing: RSH-SpMM (De Slimme Bibliothecaris)

De onderzoekers van deze paper hebben een nieuwe methode bedacht, genaamd RSH-SpMM. Stel je voor dat dit een super-slimme bibliothecaris is die het werk op een heel slimme manier verdeelt.

Hier is hoe het werkt, stap voor stap:

1. De "Reorganiserende" Stap (Locality-aware Reordering)

Voordat het werk begint, kijkt de bibliothecaris naar de hele bibliotheek. Hij merkt op dat sommige boeken vaak samen voorkomen (bijvoorbeeld: boeken over "katten" en "honden" staan vaak dicht bij elkaar).

• De analogie: In plaats van de boeken in willekeurige volgorde te laten staan, schuift hij ze een beetje op. Hij zet boeken die op elkaar lijken, naast elkaar.
• Het effect: Hierdoor ontstaan er meer "dichte" blokken waar de Super-Snelheidswagens (Tensor Cores) perfect kunnen werken, zonder dat ze hoeven te wachten op blanco pagina's.

2. De "Adaptieve" Verdeling (Adaptive Partitioning)

Nu de boeken netjes zijn gerangschikt, kijkt de bibliothecaris naar elke rij boeken.

• De slimme keuze: Als een rij boeken perfect past in een blokje voor de Super-Snelheidswagen, stuurt hij die rij daar naartoe.
• De uitzondering: Maar! Als een rij boeken heel kort is, of heel raar verspreid zit (bijvoorbeeld: één boek hier, één boek daar), dan zegt hij: "Nee, jij past niet in het blokje. Jij gaat naar de Handwerkslieden (CUDA Cores)."
• Het geheim: Dit gebeurt niet op grote schaal (bijvoorbeeld "hele afdelingen"), maar per rij. Hij maakt een heel fijnmazige verdeling. Dit zorgt ervoor dat de Super-Snelheidswagens altijd volle blokken krijgen en nooit hoeven te wachten.

3. De "Pipelijn" (Pipelined Execution)

Terwijl de Super-Snelheidswagens aan het werk zijn, zorgt de bibliothecaris ervoor dat de volgende lading boeken al klaarstaat.

• De analogie: Het is alsof een assemblagelijn werkt. Terwijl de ene wagen het werk doet, wordt de volgende lading alvast ingeladen. Hierdoor is er nooit een moment waarop de wagen stilstaat omdat hij wacht op materiaal.

Waarom is dit zo goed?

In de tests hebben ze gekeken naar echte, chaotische datasets (zoals sociale netwerken of wetenschappelijke simulaties).

• Resultaat: De nieuwe methode was 1,2 tot 6 keer sneller dan de beste bestaande methoden.
• Stabiliteit: Het maakt niet uit hoe chaotisch de data is; de methode past zich altijd aan. Soms gaat het heel snel, soms iets minder, maar het blijft stabiel en snel.

Samenvatting in één zin

RSH-SpMM is als een slimme bibliothecaris die eerst de boeken netjes schuift, en dan voor elke rij precies beslist of die naar de snelle, kieskeurige robots gaat of naar de flexibele handwerkslieden, zodat niemand ooit hoeft te wachten en alles razendsnel gaat.

Dit betekent voor de toekomst dat AI-systemen (zoals Chatbots of aanbevelingssystemen) veel sneller kunnen leren en werken, zelfs als de onderliggende data heel onregelmatig is.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "RSH-SpMM: A Row-Structured Hybrid Kernel for Sparse Matrix-Matrix Multiplication on GPUs" in het Nederlands.

Probleemstelling

Sparse Matrix-Matrix Multiplicatie (SpMM) is een fundamentele berekening in toepassingen zoals Graph Neural Networks (GNN's), wetenschappelijke simulaties en sparse deep learning. Hoewel moderne GPU's krachtige Tensor Cores (TC) bieden voor hoge doorvoer, kampen bestaande methoden met ernstige prestatieproblemen bij real-world sparse matrices.

De kern van het probleem ligt in de extreme onregelmatigheid van deze matrices:

1. Zware staartverdeling: De meeste rijen bevatten slechts een paar niet-nul-elementen, terwijl een klein aantal rijen extreem veel elementen bevat.
2. Lokale dichtheidsvariatie: De dichtheid van niet-nul-elementen verandert snel binnen en tussen matrices.
3. Misalignement met GPU-architectuur:
   • CUDA-cores zijn flexibel maar hebben een lagere reken-doorvoer dan Tensor Cores.
   • Tensor Cores vereisen dichte, op tegels (tiles) uitgelijnde data voor maximale efficiëntie. Bestaande TC-methoden gebruiken vaste venstergroottes, wat leidt tot veel "padding" (opvulling) en lage gebruikte rekenkracht (vaak <20% effectieve MMA-cycli) wanneer de lokale sparsiteit varieert.
   • Huidige hybride methoden maken vaak grove, matrix-brede beslissingen over welke rijen naar welke core gaan, waardoor ze de fijne-granulariteit van structurele coherentie binnen rijgroepen missen.

Methodologie: RSH-SpMM

RSH-SpMM is een fijn-granulair, rij-gestructureerd hybride framework dat de onregelmatige sparsiteit beter afstemt op de heterogene capaciteiten van moderne GPU's. De architectuur bestaat uit vier hoofdbestanddelen:

1. RS-Tile: Gecomprimeerd Formaat

Het paper introduceert een nieuw opslagformaat genaamd RS-Tile (Row-Structured Tile). In plaats van de hele matrix als één eenheid te behandelen, wordt deze dynamisch opgesplitst in twee disjuncte delen:

• TC Block Part: Rijgroepen met voldoende structuur en dichtheid worden samengevoegd tot vaste 8x8 blokken (geschikt voor Tensor Cores). Deze worden gecomprimeerd met een bitmap en kolom-ID's, wat de metadata- overhead verlaagt en de tegeldichtheid maximaliseert.
• CUDA Residual Part: Rijen met zeer weinig niet-nul-elementen of structureel incompatibele patronen worden uitgesloten van de TC-pad. Deze worden opgeslagen in een compact formaat voor CUDA-cores (zonder zware tile-opbouw), wat de verwerkingskosten minimaliseert.

2. Fijn-granulair Hybride Kernel

RSH-SpMM gebruikt een adaptieve strategie om rijen toe te wijzen:

• Adaptieve Rij-partitie: Voordat vensters worden gevormd, scant het systeem rijen en beoordeelt deze op basis van het aantal niet-nul-elementen en de structurele overlap met buren. Rijen die de dichtheid van een TC-venster zouden verlagen, worden direct naar de CUDA-pad gestuurd.
• Gepipelined TC Kernel: Voor de TC-deel wordt een dubbele-buffer pipeline gebruikt die globale geheugenoverdracht, gedeeld geheugen-decodering (via bitmap) en MMA-berekeningen overlapt. Dit verbergt latentie en houdt de Tensor Cores constant gevuld.
• Lichtgewicht CUDA Kernel: Voor de resterende rijen wordt een eenvoudige, niet-gepipelined multiply-add kern gebruikt die geen tile-conversie vereist, waardoor registerdruk en synchronisatie-overhead minimaal blijven.

3. Locality-Aware Reordering

Om de kans op het vormen van kwalitatief goede TC-vensters te vergroten, wordt een lokaliteitsbewuste herschikking toegepast vóór de format-conversie:

• Er wordt een gewogen Jaccard-similariteit gebruikt om rijen met vergelijkbare kolompatronen te groeperen.
• Een Minimum Spanning Tree (MST) wordt gebruikt om een globale volgorde te bepalen die de structurele continuïteit maximaliseert.
• Een lokale verfijning (2-opt swaps) en isolatie-herkenning zorgen ervoor dat "geïsoleerde" rijen niet de coherentie van vensters verstoren.

4. Adaptieve Load Balancing

In tegenstelling tot methoden die vaste venstergroottes gebruiken, past RSH-SpMM de belasting dynamisch aan. "Super-lange" rijen worden optioneel opgesplitst, terwijl korte rijen worden geïsoleerd. Dit voorkomt dat één lange rij de latency van een heel venster bepaalt en zorgt voor een stabiele uitlast van de Streaming Multiprocessors (SM's).

Belangrijkste Bijdragen

1. RS-Tile Representatie: Een compact formaat dat Tensor-Core-vriendelijke fragmenten blootlegt terwijl het onregelmatige rijen naar een efficiënte CUDA-pad routeert, wat metadata-overhead vermindert.
2. Fijn-granulair Hybride Strategie: Een adaptieve partitionering die de nadruk legt op structurele coherentie binnen rijgroepen, zonder de noodzaak van dubbele sparse indexstructuren.
3. Locality-Aware Reordering: Een techniek die rijen groepeert op basis van structurele gelijkenis, wat fragmentatie vermindert en de MMA-uitlast verbetert.
4. Uitgebreide Validatie: Demonstratie van robuuste prestaties over een breed scala aan werklasten, inclusief GNN's en diverse SuiteSparse-matrices.

Resultaten

De auteurs hebben RSH-SpMM geëvalueerd op NVIDIA RTX 4090 en RTX 3090 GPU's en vergeleken met state-of-the-art baselines (zoals cuSPARSE, TC-GNN, DTC-SpMM, Acc-SpMM, en hybride methoden).

• Snelheidswinst: RSH-SpMM levert consistent een versnelling van 1,27x tot 6,13x ten opzichte van de beste bestaande methoden.
• Vergelijking met cuSPARSE: Gemiddeld 2,35x sneller op RTX 4090 en 2,86x sneller op RTX 3090.
• Vergelijking met Tensor Core-methoden: Overtreft TC-GNN met tot 6,13x (door het vermijden van runtime-fouten en inefficiënte padding) en DTC-SpMM/Acc-SpMM met respectievelijk 1,91x en 1,61x.
• Tensor Core Utilisatie: De Tensor Core-uitlast stijgt van gemiddeld 5,6% (bij Acc-SpMM) naar 8,8% bij RSH-SpMM, wat aantoont dat de pipeline minder vaak stilstaat door ongeschikte data.
• End-to-End Training: In een GCN-trainingsscenario (6 lagen) op de RTX 4090 was RSH-SpMM 1,49x sneller dan cuSPARSE en 1,09x sneller dan DTC-SpMM, wat aantoont dat de kernel-prestaties direct vertalen naar applicatie-snelheid.
• Robuustheid: Het framework presteert stabiel over matrices met zeer verschillende sparsiteitspatronen, waarbij meer dan 80% van de SuiteSparse-matrices een versnelling van 1,24x tot 8,2x behaalde.

Betekenis

RSH-SpMM is een doorbraak in het optimaliseren van SpMM op GPU's omdat het de fundamentele beperking van bestaande hybride methoden oplost: het gebrek aan fijn-granulair inzicht in de lokale structuur van sparse matrices. Door rijen dynamisch te partitioneren op basis van hun bijdrage aan de tegeldichtheid en door een slimme herschikking toe te passen, combineert het het beste van twee werelden: de hoge doorvoer van Tensor Cores voor dichte gebieden en de flexibiliteit van CUDA-cores voor onregelmatige gebieden.

Dit werk is van groot belang voor de toekomst van Graph Neural Networks en Large Language Models (LLM's), waar sparse matrixvermenigvuldiging een kritieke bottleneck vormt. Het biedt een schaalbare oplossing die de prestaties van deze AI-werklasten aanzienlijk kan verbeteren, zelfs bij complexe en onvoorspelbare datastructuren.