WiSP: A Working-Set View of Mixture-of-Experts Serving on Extremely Low-Resource Hardware

Sintesi Tecnica: WiSP – Una visione del Working-Set per il serving di Mixture-of-Experts su hardware con risorse estremamente limitate

Enunciato del Problema
I moderni modelli Mixture-of-Experts (MoE), come Qwen3-30B-A3B e Mixtral-8×7B, collocano la maggior parte dei loro parametri nei layer degli esperti, pur avendo solo una piccola frazione di questi esperti attiva per ogni dato token. Su GPU consumer con VRAM limitata, questi modelli spesso non possono entrare interamente in memoria. La soluzione standard, l'offloading statico a livello di layer su CPU (ad esempio, --cpu-offload-gb in vLLM), mantiene basso l'uso della memoria ma trasferisce tutti gli esperti di un layer attraverso il bus PCIe ad ogni passaggio di forward pass. Questo approccio non riesce a sfruttare la sparsità dei modelli MoE, annullando di fatto i loro benefici di efficienza.

Gli autori inquadrano il serving di MoE a basse risorse come un problema di gestione del working-set. In questo regime, la memoria GPU è lo store veloce limitato, mentre la memoria host funge da backing store. Due flussi distinti di domanda di memoria competono per lo stesso budget di VRAM:

1. Pesi degli Esperti: Uno stream dipendente dal routing dove solo un sottoinsieme di esperti è necessario per ogni token.
2. KV Cache: Uno stream dipendente dal token richiesto per la generazione autoregressiva.

I sistemi esistenti spesso trattano questi come problemi separati o si rivolgono a setting di datacenter ad alta concorrenza dove la computazione può sovrapporsi al trasferimento dati. Tuttavia, nel serving interattivo a bassa concorrenza (ad esempio, assistenti di coding locali), la computazione è minima per ogni step, rendendo la larghezza di banda PCIe il collo di bottiglia primario piuttosto che la computazione o l'accuratezza della predizione.

Metodologia: WiSP (Working-Set Paging)
Il paper introduce WiSP, un sistema che tratta l'insieme di esperti residenti come una cache sopra il comportamento di routing del modello, gestita congiuntamente alla KV cache. WiSP è composto da due componenti principali:

1. Routing-Aware Expert Pager:

   • Meccanismo: WiSP si inserisce in un motore di serving non modificato (ad esempio, vLLM) sfruttando un esistente hook di indirezione dell'expert-map (π). Invece di caricare interi layer, mantiene un piccolo spazio scratch sulla GPU contenente solo il "working set" di esperti probabilmente riutilizzati.
   • Operazione: Ad ogni step, il pager recupera solo gli esperti routati mancanti dalla memoria host (H2D) ed effettua l'eviction degli esperti meno di recente utilizzati (LRU) per fare spazio.
   • Garanzie: Poiché i kernel e i pesi sottostanti rimangono invariati (viene modificata solo la mappatura), WiSP produce output byte-identici rispetto al motore di base. Non richiede nuovi kernel, cambiamenti di quantizzazione o riaddestramento del modello.

2. MV-WSA (Marginal-Value Working-Set Allocation):

   • Il Problema di Allocazione: Poiché i pesi degli esperti e la KV cache condividono lo stesso budget di VRAM, aumentare l'insieme di esperti residenti riduce la capacità della KV cache, e viceversa.
   • Strategia: MV-WSA determina la suddivisione ottimale ($f$) tra il budget degli esperti e il budget della KV cache. Mira a equalizzare il beneficio marginale di latenza per byte tra entrambi i pool.
   • Admission Floor: Fondamentalmente, l'allocatore impone una dimensione minima della KV cache ($\kappa_{min}$) necessaria per ammettere il livello di concorrenza attuale. Se la suddivisione ottimale non vincolata svuoterebbe la KV cache, il sistema la blocca per garantire che lo scheduler possa inizializzarsi ed eseguire.
   • Modalità di Implementazione:
     • Configuratore Offline: Analizza una traccia rappresentativa per impostare una suddivisione fissa all'avvio.
     • Controller Online: Ridimensiona dinamicamente entrambi i pool durante il serving in corrispondenza di "drained barriers" (punti in cui nessuna richiesta attiva detiene blocchi di memoria), spostando fisicamente i byte tra lo scratch degli esperti e il pool della KV senza superare il budget totale.

Risultati Chiave sul Prefetching
Una scoperta empirica centrale del paper è che il prefetching basato sul routing fornisce poco o nessun beneficio nel decode a singolo stream e bassa concorrenza.

• Collo di Bottiglia: Le prestazioni sono limitate dalla larghezza di banda PCIe, non dall'accuratezza della predizione. Nel decode batch-1, la computazione per step è troppo piccola per nascondere i trasferimenti PCIe dietro di essa.
• Risultato: Il prefetching speculativo spesso degrada le prestazioni (riducendo il throughput di circa il 50% nei test) perché compete per la larghezza di banda ed effettua l'eviction di esperti necessari per lo step immediatamente successivo.
• Riformulazione: Il segnale di routing è prezioso non per nascondere la latenza, ma per l'allocazione della memoria. Aiuta a identificare la dimensione minima "sicura" del set residente, permettendo al sistema di dimensionare correttamente la cache degli esperti per utente e allocare la VRAM liberata alla KV cache.

Risultati
Gli autori valutano WiSP su varie architetture MoE (LLM, VLM, Ibridi SSM-MoE e modelli di Diffusione) utilizzando una NVIDIA H100 (con alcuni esperimenti che emulano budget di memoria più piccoli).

• Guadagni di Throughput: Sotto budget di VRAM corrispondenti in cui il modello non entra in memoria, WiSP raggiunge fino a 1.95× il throughput di decode rispetto all'offloading statico su CPU. Il guadagno è massimo a budget di memoria intermedi, dove WiSP può fare il caching di una parte significativa degli esperti riutilizzati mentre il baseline trasmette tutto.
• Capacità KV: Riducendo l'impronta degli esperti, WiSP aumenta direttamente la capacità della KV cache. Gli esperimenti mostrano che restringere il limite degli esperti può aumentare la concorrenza massima di oltre il 40% allo stesso budget di memoria totale.
• Performance di Allocazione:
  • Offline: Il configuratore MV-WSA funziona bene sia nel prefill che nel decode, rimanendo entro pochi punti percentuali di un oracle per workflow in simulazione, mentre le suddivisioni fisse (ad esempio, 50/50) sono circa il 20% peggiori.
  • Online: Il controller di ridimensionamento duale dinamico migliora il throughput di 1.20× rispetto al miglior split statico offline, adattandosi ai cambiamenti del carico di lavoro (ad esempio, passando da turni a lungo contesto a burst di tool-call).
• Generalità: WiSP serve con successo modelli che non possono entrare in una singola GPU (ad esempio, Jamba-v0.1, MiniMax-M2) e funziona attraverso varie modalità, inclusi i modelli Vision-Language (Kimi-VL) e i modelli di Diffusione (DiT-MoE, LLaDA-MoE), mantenendo output byte-identici.

Significato e Rivendicazioni
Il paper sostiene che WiSP fornisca un framework unificato per il serving di MoE a basse risorse, trattando la residenza degli esperti e il caching della KV come un unico problema di allocazione del working-set.

• Cambio di Paradigma del Design: Argomentano che, per il serving a bassa concorrenza, l'attenzione dovrebbe spostarsi da "come prefetchare gli esperti più velocemente" a "come allocare la VRAM in modo ottimale tra esperti e KV cache".
• Impatto Pratico: WiSP consente il serving di modelli MoE di scala frontier su hardware consumer senza sacrificare la correttezza dell'output o richiedere hardware specializzato (come Intel AMX).
• Limitazioni: Gli autori notano esplicitamente che WiSP è progettato per il regime in cui il modello non entra in VRAM; se il modello entra, la residenza statica rimane più veloce poiché il paging aggiunge un'indirezione non necessaria. Inoltre, l'ottimizzazione del prefetching è mostrata come inefficace nel decode a singolo stream, sebbene possa essere rilevante per sessioni ad alto batch o multi-turn agent dove è possibile il sovrapporsi della computazione.

Il lavoro conclude che, nel decode a singolo stream limitato dalla memoria, il segnale di routing è una leva di memoria (determina la dimensione del working set) piuttosto che una leva di velocità (nasconde la latenza), rendendo l'allocazione intelligente il driver principale delle prestazioni.