TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

1. Probleemstelling

Distributed inference van Large Language Models (LLMs) met Tensor Parallelism (TP) ondervindt aanzienlijke prestatieknelpunten door communicatie-overhead, zelfs op high-speed interconnects zoals NVLink.

• Communicatie-Overhead: In TP-implementaties staat communicatie (specifiek AllReduce) voor 9–23% van de end-to-end inferentielatentie.
• De "Small Batch"-Uitdaging: Bestaande technieken om communicatie te overlappen met berekening (bijvoorbeeld door werk op te splitsen in kleinere tegels of tokens) zijn niet standaard ingeschakeld in grote serving-systemen (vLLM, SGLang, TensorRT-LLM). Dit komt omdat het opsplitpen van kleine werklasten (typisch bij low-latency inferentie, bijvoorbeeld <2K tokens) vaak berekeningsinefficiënties introduceert (wave quantization effects) die de winst van het overlappen van communicatie tenietdoen.
• Resource Contention: Communicatie-operaties verbruiken vaak een groot aantal Streaming Multiprocessors (SMs) (16–20+), waardoor compute-kernels worden uitgehongerd van resources.
• RMSNorm Redundantie: De RMSNorm-operatie, vaak over het hoofd gezien, voegt 4–9% overhead toe. In standaard TP wordt deze redundant uitgevoerd op elke GPU na AllReduce, ondanks dat alle GPU's op dat moment identieke data hebben.

2. Methodologie: TokenWeave

TokenWeave is een systeem ontworpen om efficiënt compute-communicatie-overlap mogelijk te maken voor tensor-parallelle inferentie, zelfs voor kleine batchgroottes (tot 1024 tokens). Het maakt gebruik van drie kerninnovaties:

A. Wave-Aware Smart Token Splitting

In plaats van een batch naïef in twee gelijke helften te splitsen, gebruikt TokenWeave een wave-aware splitting-strategie.

• Mechanisme: Het partitioneert de input-batch in twee subsets (prefix en suffix) zodat het totale aantal uitvoerings-"waves" (groepen CTAs) dat nodig is voor beide subsets, niet meer bedraagt dan het aantal waves dat nodig is voor de originele, niet-gesplitste batch.
• Doel: Dit minimaliseert wave quantization-overheads, zodat het splitsen van het werk de berekeningsefficiëntie niet verslechtert.
• Uitvoering: Het systeem gebruikt twee CUDA-streams (Compute en Communication). Terwijl de AllReduce (gefuseerd met RMSNorm) draait op de eerste split, draait de attention-berekening op de tweede split, en vice versa.

B. RMSNorm Reordering en Fusion

TokenWeave identificeert RMSNorm als een kritiek optimalisatiedoel en herschikt de uitvoeringsflow.

• Reordering: In plaats van AllReduce gevolgd door RMSNorm uit te voeren, splitst TokenWeave AllReduce op in ReduceScatter en AllGather.
  1. ReduceScatter: Elke GPU ontvangt zijn shard van de tensor.
  2. Lokale RMSNorm: Elke GPU voert RMSNorm uit op zijn lokale shard voordat de definitieve aggregatie plaatsvindt. Omdat de data al gepartitioneerd is, elimineert dit redundante berekening over GPU's heen.
  3. AllGather: De genormaliseerde shards worden verzameld.
• Gefuseerde Kernel: Een nieuwe Fused AllReduce–RMSNorm kernel is geïmplementeerd. Deze fuseert ReduceScatter, RMSNorm en AllGather tot één enkele operatie.
  • Memory-optimalisatie: Door variantie direct te berekenen op het ReduceScatter-resultaat in SM-registers, wordt één HBM-lezen geëlimineerd. Door genormaliseerde waarden direct te schrijven naar Multimem-adressen voor AllGather, wordt één HBM-schrijven geëlimineerd.
  • SM-efficiëntie: Deze kernel maakt gebruik van NVSHARP/Multimem (hardware-functies op Hopper/Blackwell GPU's) om communicatie en normalisatie uit te voeren met slechts 2–8 SMs, waardoor het merendeel van de SMs vrij blijft voor berekening.

C. Adaptieve Overlap-strategie

TokenWeave beslist dynamisch of volledige splitting/overlap moet worden gebruikt of alleen de gefuseerde kernel, afhankelijk van de batchgrootte:

• Grote Batches: Schakelt volledige tweeweg-splitting en overlap in.
• Kleine Batches: Schakelt splitting uit (om overhead te vermijden) maar past toch de Fused AllReduce–RMSNorm kernel toe om memory-traffic te verminderen en redundante berekening te elimineren.

3. Belangrijkste Bijdragen

1. RMSNorm-optimalisatie: Eerste identificatie en optimalisatie van RMSNorm in TP-inferentie via reordering en een aangepaste gefuseerde kernel, wat HBM-traffic en redundante berekening reduceert.
2. Wave-Aware Splitting: Een nieuw token-splitting-algoritme dat verlies van berekeningsefficiëntie door wave quantization voorkomt, waardoor overlap mogelijk wordt zelfs voor kleine token-batches (1K–2K).
3. Hardware-bewuste Efficiëntie: Benutting van NVSHARP/Multimem om communicatie-SM-gebruik te reduceren van ~20 naar 2–8, waardoor resource-uitputting wordt voorkomen.
4. Systeemintegratie: Het eerste systeem dat succesvol compute-communicatie-overlap mogelijk maakt voor tensor-parallelle serving in low-latency omgevingen, geïntegreerd in het vLLM-framework.

4. Experimentele Resultaten

Evaluaties zijn uitgevoerd op 8×H100 DGX-systemen met modellen zoals Llama-3.3-70B, Qwen2.5-72B en Mixtral-8x22B.

• Latentieverbetering:
  • Bereikt tot 1,28× versnelling in latentie vergeleken met een geoptimaliseerde baseline (vLLM met Multimem).
  • Aanzienlijke winst zelfs bij korte sequentielengtes (bijvoorbeeld 1,2× bij 1K tokens).
  • Presteert beter dan state-of-the-art overlap-oplossingen zoals TileLink, die alleen winst tonen bij grotere batchgroottes (≥4K) en overhead veroorzaken bij kleinere maten.
• Throughputverbetering:
  • Levert tot 1,19× hogere throughput op real-world traces (ShareGPT, arXiv).
  • In sommige settings presteert TokenWeave beter dan een theoretische baseline met zero communication (vLLM-nocomm), omdat zijn gefuseerde kernel de memory-bound RMSNorm-operatie beter optimaliseert dan de baseline pure berekening aankan.
• Kernel-efficiëntie:
  • De gefuseerde AllReduce–RMSNorm-kernel biedt 1,34–1,39× versnelling ten opzichte van sequentiële uitvoering over alle token-groottes heen.
  • Het bereikt bijna optimale prestaties met slechts 4–8 SMs.

5. Betekenis

• Maakt Low-Latency Distributed Inference mogelijk: Door het "small batch"-probleem op te lossen, maakt TokenWeave distributed TP-inferentie haalbaar voor interactieve toepassingen waar latentie kritiek is, zonder enorme batchgroottes te vereisen.
• Hardware-gebruik: Het maximaliseert het nut van moderne GPU-functies (NVSHARP, Multimem) en verschuift de bottleneck van communicatie-overhead naar pure berekening.
• Praktische Impact: De broncode is open-source gemaakt en geïntegreerd in vLLM, wat directe prestatiewinst biedt voor productie-LLM-serving-systemen zonder wijzigingen in de modelarchitectuur te vereisen.
• Paradigmaverschuiving: Het daagt de aanname uit dat communicatie-overlap alleen gunstig is voor grote trainingsbatches of expert-parallelism (MoE)-opstellingen, en bewijst dat het essentieel is voor standaard tensor-parallelle inferentie.