Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Technische Samenvatting: Versnelling van GPU-inferentie van Large Language Models met Matig Ongestructureerde Sparse Gewichtsmatrices

Probleemstelling
De implementatie van Large Language Models (LLMs) staat voor een kritieke bottleneck in de inferentiekosten. Hoewel technieken voor modelpruning (bijv. Wanda, RIA) de computationele last kunnen verminderen door sparsity (ijlheid) in gewichtsmatrices te introduceren, worden ze beperkt door eisen aan de modelkwaliteit. In tegenstelling tot kleinschalige modellen zijn LLMs zeer gevoelig voor hoge niveaus van sparsity; het behouden van prestaties beperkt pruning doorgaans tot matig ongestructureerde sparsity (rond de 50%).

Op dit specifieke sparsityniveau presteren bestaande GPU-kernels voor Sparse Matrix Multiplication (SpMM) niet beter dan dense matrixvermenigvuldigingsbibliotheken zoals cuBLAS. Dit prestatiegat komt voort uit drie fundamentele uitdagingen:

1. Tensor Core Incompatibiliteit: Moderne GPU's bieden hoogefficiënte sparse tensor cores, maar deze ondersteunen strikt 2:4 gestructureerde sparsity. Ze kunnen irreguliere, ongestructureerde pruning-patronen niet direct versnellen.
2. Metadata Overhead: Traditionele opslagformaten (bijv. CSR) vereisen het opslaan van volledige positionele metadata voor elk niet-nul element. Bij 50% sparsity kan deze metadata evenveel ruimte innemen als de waarden zelf, wat de voordelen van geheugenbandbreedte tenietdoet.
3. On-Chip Decoding Bottlenecks: Compressietechnieken die metadata verminderen (bijv. bitmap-encoding) verschuiven de werklast naar on-chip decoding. Op moderne GPU's met High-Bandwidth Memory (HBM) kan de toegang tot het globale geheugen de decoding-taken die door tragere CUDA-cores worden uitgevoerd, niet volledig overlappen, wat leidt tot een prestatie-stalling.

Methodologie
De auteurs stellen een holistische oplossing voor die bestaat uit een nieuw drielaags matrixopslagformaat en een co-geoptimaliseerde SpMM-kernel, ontworpen voor matig ongestructureerde sparsity.

1. Drielaags Opslagformaat
Het voorgestelde formaat deelt de sparse matrix op in drie afzonderlijke lagen om een balans te vinden tussen hardwarecompatibiliteit en opslagefficiëntie:

• Sparse-TC Laag: Deze laag extraheert 2:4 gestructureerde sparsity-patronen uit de irreguliere distributie. Het behoudt twee niet-nul waarden van elke vier opeenvolgende gewichten. Omdat het patroon vaststaat, is er slechts 2 bits per waarde aan metadata nodig, wat directe compatibiliteit met sparse tensor cores garandeert zonder decoding-overhead.
• Slot-Filling Laag: Overschot aan niet-nul elementen die niet in het 2:4 patroon passen, worden herverdeeld in de lege slots die door de Sparse-TC laag zijn achtergelaten. Om de metadata voor deze ~17,7% van de elementen te minimaliseren, introduceren de auteurs Parallel Differential Distance (PDD) encoding. In plaats van absolute indices op te slaan, codeert PDD de relatieve afstanden tussen de herverdeelde elementen. De meeste afstanden zijn klein genoeg om in 4 bits gecodeerd te worden. Voor grotere gaten worden dummy-punten ingevoegd om de span te breken, zodat geen enkele afstand de encoding-limiet overschrijdt.
• Residual Laag: De resterende niet-nul elementen (meestal <1% van het totaal) die niet efficiënt door de eerste twee lagen kunnen worden afgehandeld, worden opgeslagen in een traditioneel Compressed Sparse Row (CSR) formaat. Vanwege het minimale volume aan data is de metadata-overhead hier verwaarloosbaar.

2. Co-Geoptimaliseerde SpMM Kernel
Op basis van dit formaat ontwerpen de auteurs een GPU-kernel die gezamenlijk sparse tensor cores en CUDA-cores gebruikt binnen een gepipelined uitvoeringsmodel:

• Parallel Decoding: CUDA-cores zijn belast met het decoderen van de PDD-metadata van de Slot-Filling laag in parallel. Omdat deze laag slechts een klein deel van de totale data vertegenwoordigt, wordt de decoding-overhead effectief beheerd.
• Gestructureerde Versnelling: Sparse tensor cores verwerken de Sparse-TC laag met behulp van wgmma-instructies (Hopper-architectuur), waarbij gebruik wordt gemaakt van de 2:4 structuur voor een hoge throughput.
• Residual Afhandeling: CUDA-cores verwerken ook de sparse Residual laag en schrijven de resultaten naar een aparte on-chip buffer om ze aan het einde van de pipeline te reduceren met de outputs van de tensor cores.
• Pipeline Overlap: De kernel implementeert een fijnmazige pipeline die drie kritieke operaties overlapt:
  • Het laden van dense inputmatrices (gecached in L2) en sparse weight tiles (uit het Global Memory).
  • Het decoderen van PDD-metadata op de CUDA-cores.
  • Het uitvoeren van matrixvermenigvuldiging op de sparse tensor cores.
    Deze "load-as-sparse, compute-as-sparse" strategie zorgt ervoor dat on-chip berekeningen overlappen met de toegang tot het globale geheugen, waardoor de HBM-benutting wordt gemaximaliseerd.

Belangrijkste Bijdragen

• Nieuw Opslagformaat: Een drielaags formaat dat is afgestemd op matig ongestructureerde sparsity, wat compatibiliteit met sparse tensor cores mogelijk maakt, de metadata-overhead drastisch vermindert via PDD-encoding en de on-chip decoding-kosten minimaliseert.
• Co-Geoptimaliseerde Kernel: Een SpMM-kernel die efficiënt sparse tensor cores en CUDA-cores coördineert, waarbij pipelining wordt gebruikt om geheugentoegang te overlappen met berekening.
• Eerste die 'Dense' verslaat: De implementatie is de eerste die aantoont dat matig ongestructureerde sparse SpMM de dense matrixvermenigvuldiging op moderne GPU's met HBM kan overtreffen.

Resultaten
Evaluaties werden uitgevoerd op NVIDIA H100 SXM5 GPU's met OPT-30B en OPT-66B modellen met sparsity-niveaus van 50%, 60% en 70%.

• Kernel-Level Prestaties: Bij 50% sparsity bereikt de voorgestelde methode een versnelling van tot wel 1,64× ten opzichte van SpInfer (EuroSys'25 Best Paper) en presteert aanzienlijk beter dan andere baselines zoals SparTA en FlashLLM. Het overtreft ook de dense cuBLAS-prestaties.
• End-to-End Inferentie: De methode levert een end-to-end versnelling van tot wel 1,41× ten opzichte van FlashLLM (VLDB'24) en een tot 41,4% hogere throughput vergeleken met FlashLLM op OPT-30B met een batchgrootte van 16.
• Geheugenefficiëntie: De aanpak vermindert het piekgebruik van het globale geheugen met ongeveer 21,4% vergeleken met dense baselines bij 50% sparsity, hoewel het iets meer geheugen gebruikt dan SpInfer vanwege het ontbreken van bitmap-encoding voor metadata.
• Schaalbaarheid: De methode behoudt haar voordelen over de sparsity-niveaus van 50% tot 70%.

Betekenis en Claims
Het artikel beweert een kritiek gat in de optimalisatie van LLM-inferentie aan te pakken: het onvermogen van bestaande sparse kernels om matig ongestructureerde sparsity effectief te benutten. Door een opslagformaat te introduceren dat de kloof overbrugt tussen irreguliere pruning-patronen en gestructureerde hardwareversnelling, laten de auteurs zien dat pruning daadwerkelijke inferentie-versnellingen kan opleveren zonder de modelkwaliteit op te offeren.

Het werk benadrukt dat terwijl hoge sparsity (>70%) de modelkwaliteit verslechtert, matige sparsity (~50%) het "sweet spot" is voor LLMs. De voorgestelde oplossing ontsluit het prestatiepotentieel van dit sweet spot, dat voorheen ontoegankelijk was door de inefficiënties van bestaande sparse kernels. De auteurs merken op dat hun optimalisatie specifief gericht is op de decode-fase van LLM-inferentie (waar batchgroottes klein zijn en geheugen-gebonden zijn), wat het zeer geschikt maakt voor schaalbare implementatie in productieomgevingen waar prefill- en decode-fasen vaak ontkoppeld zijn. De broncode wordt beschikbaar gesteld om verdere research en adoptie te faciliteren.