Riepilogo Tecnico: LYNX – Abilitazione di Inferenze MoE Efficienti Tramite Selezione Dinamica degli Esperti Consapevole del Batch

1. Enunciato del Problema

I modelli fondazionali moderni utilizzano sempre più architetture Mixture-of-Experts (MoE) per disaccoppiare il numero di parametri dal costo computazionale. Sebbene gli MoE attivino solo un piccolo sottoinsieme di parametri per token (ad esempio, 14 su 57 miliardi in Qwen2-57B), questa promessa di efficienza viene meno nel servizio di produzione a causa della tensione tra latenza e throughput.

Per ottenere un throughput elevato, i sistemi di inferenza raggruppano le richieste in batch. Tuttavia, poiché ogni token in un batch seleziona indipendentemente i propri top-$k$ esperti, l'insieme degli esperti che devono essere caricati dalla memoria ad alta larghezza di banda (HBM) cresce con la dimensione e la composizione del batch. A dimensioni di batch moderate richieste per soddisfare gli Obiettivi di Livello di Servizio (SLO) sulla latenza, questo pattern di accesso alla memoria dinamico e dipendente dai dati rende la decodifica MoE limitata dalla larghezza di banda della memoria. Di conseguenza, la latenza scala con il numero totale di esperti distinti attivati in tutto il batch, spesso saturando l'intero pool di esperti (ad esempio, un batch di sole 8 richieste diverse può attivare tutti i 64 esperti in Qwen2-57B). Questo annulla i benefici di sparsità degli MoE, facendoli comportare in modo simile a modelli densi di capacità equivalente ma con il sovraccarico aggiuntivo del movimento dinamico degli esperti.

Le soluzioni esistenti, come la potatura statica, la quantizzazione o il clustering, si basano su calibrazioni offline e alterano permanentemente il modello, rendendole poco flessibili ai cambiamenti del carico di lavoro. Gli approcci dinamici per token non affrontano il collo di bottiglia perché non riducono il numero aggregato di esperti recuperati a livello di batch.

2. Metodologia: LYNX

LYNX è un sistema agnostico rispetto al carico di lavoro progettato per alleviare il collo di bottiglia della larghezza di banda della memoria nell'inferenza MoE senza modificare l'architettura del modello o richiedere calibrazioni offline. Sfrutta un'idea chiave derivata dalle dinamiche di addestramento degli MoE: le perdite di bilanciamento del carico costringono i router a distribuire i token ampiamente tra gli esperti, anche quando la fiducia nelle selezioni secondarie è bassa. Questo crea una ridondanza a livello di batch in cui il router seleziona esperti a causa di vincoli di addestramento piuttosto che di una genuina affinità token-esperto.

LYNX sfrutta questa ridondanza attraverso una nuova tecnica di AffinityBinning e una pipeline a tre componenti che opera interamente a runtime:

2.1 Principi Fondamentali

Il design si basa su tre osservazioni empiriche:

1. Sensibilità alla Fase: Gli MoE sono limitati dal calcolo durante la fase di prefill, ma limitati dalla larghezza di banda della memoria durante la fase di decode. Il rimappaggio degli esperti è benefico solo durante il decode. Inoltre, la fase di prefill è altamente sensibile alla fedeltà degli esperti, mentre la fase di decode è resiliente alla riassegnazione.
2. Fiducia del Token: I punteggi di fiducia del router (derivati dal divario tra i logit degli esperti meglio classificati) identificano in modo affidabile quali mappature token-esperto sono critiche. I token a bassa fiducia possono essere rimappati con una degradazione minima della precisione.
3. Gerarchia del Rangho degli Esperti: L'esperto meglio classificato (rank-0) domina la qualità dell'output. Gli esperti di rango inferiore tra i top-$k$ mostrano un'alta ridondanza. Anche la sensibilità alla riassegnazione varia tra i livelli.

2.2 Componenti del Sistema

LYNX si integra con il motore di servizio (ad esempio, vLLM) tramite un ottimizzatore consapevole della fase e tre componenti leggeri di kernel CUDA:

• Ottimizzatore Consapevole della Fase: Determina se l'iterazione corrente è limitata dalla memoria (decode) o dal calcolo (prefill). Ignora la logica di LYNX durante le fasi di prefill o prefill a blocchi per evitare sovraccarichi non necessari.
• Analizzatore di Fiducia: Intercetta i logit del router a ogni livello. Discretizza la fiducia del router di ogni token in bin utilizzando AffinityBinning. La larghezza e il numero dei bin sono determinati esclusivamente dal rapporto di sparsità del modello ($k/N$), rendendo il sistema auto-calibrante e indipendente dal carico di lavoro. Identifica i token con bassa affinità rispetto agli esperti assegnati.
• Valutatore Adattivo degli Esperti: Calcola un punteggio a livello di batch per ogni esperto. Utilizza uno schema di ponderazione esponenziale basato sui valori dei bin assegnati dall'Analizzatore di Fiducia. Questo garantisce che gli esperti preferiti dai token ad alta convinzione ricevano punteggi sproporzionatamente più alti. Determina l'insieme minimo di esperti "critici" richiesti per il batch.
• Rimappatore degli Esperti: Reindirizza le assegnazioni token-esperto a bassa affinità all'insieme critico di esperti sopravvissuti. Crucialmente, LYNX preserva la semantica di attivazione top-$k$; ogni token attiva ancora esattamente $k$ esperti, ma il batch collettivamente attiva un insieme più piccolo e consolidato di esperti unici.

2.3 Implementazione

LYNX è implementato come modulo plug-and-play su vLLM. La sua logica principale è realizzata utilizzando quattro kernel fusi Triton/CUDA che intercettano l'output del router. Questo design elimina il movimento intermedio dei dati tensoriali e riduce il sovraccarico di avvio dei kernel, garantendo che il sistema aggiunga una latenza trascurabile (<4% di sovraccarico) operando nel percorso critico.

3. Contributi Chiave

1. Primo Sistema Agnostico rispetto al Carico di Lavoro: LYNX è il primo sistema ad alleviare il collo di bottiglia della larghezza di banda della memoria nell'inferenza MoE senza richiedere dati di calibrazione o modifiche permanenti al modello.
2. Tecnica AffinityBinning: Un metodo innovativo per discretizzare la fiducia del router basato sui rapporti di sparsità architetturali, consentendo l'adattamento automatico a qualsiasi architettura MoE.
3. Ottimizzazione a Livello di Batch: A differenza dei precedenti approcci dinamici per token, LYNX opera a livello di batch per ridurre il numero totale di esperti unici recuperati dall'HBM, affrontando direttamente il collo di bottiglia della larghezza di banda della memoria.
4. Complementarità: LYNX è progettato per funzionare in conjunction con ottimizzazioni esistenti, inclusi l'offloading e la quantizzazione.

4. Risultati Sperimentali

LYNX è stato valutato su modelli all'avanguardia di quattro famiglie (Qwen, DeepSeek, Mixtral, Llama) attraverso otto benchmark (generazione di codice, matematica, ragionamento e visione).

• Riduzione della Latenza: LYNX raggiunge una latenza fino a 1,30× inferiore (Tempo per Token di Output) rispetto alle linee di base. In scenari di servizio disaggregato (dove prefill e decode sono separati), le riduzioni raggiungono fino a 1,35×.
• Impatto sulla Precisione: Il sistema comporta una perdita di precisione inferiore all'1% nel caso peggiore. Notabilmente, LYNX migliora frequentemente la precisione in media (ad esempio, +1,9% su HumanEval per DeepSeek-Coder-V2), suggerendo che il bilanciamento del carico durante l'addestramento forza selezioni di esperti subottimali che LYNX corregge.
• Throughput: Con budget SLO fissi (ad esempio, 25ms o 50ms), LYNX migliora il throughput del sistema fino a 2,1×.
• Compatibilità:
  • Offloading: Quando combinato con tecniche di offloading (ad esempio, Fiddler), LYNX aumenta le prestazioni del 31%.
  • Quantizzazione: Quando applicato sopra la quantizzazione INT4, LYNX fornisce un'ulteriore accelerazione del 10%.
  • Parallelismo: LYNX fornisce guadagni di throughput additivi sopra il Parallelismo degli Esperti (EP) e il Parallelismo Tensoriale (TP).
• Robustezza: I guadagni di prestazioni sono coerenti tra dimensioni di batch variabili (saturando a dimensioni moderate) e tracce reali con tassi di arrivo irregolari.

5. Significato

Il documento sostiene che, poiché la potenza di calcolo delle GPU continua a superare la crescita della larghezza di banda della memoria, il divario di intensità aritmetica per la decodifica MoE si allargherà, rendendo i colli di bottiglia limitati dalla memoria sempre più critici. LYNX affronta questo problema cambiando fondamentalmente il modo in cui gli esperti vengono selezionati a livello di batch. Sfruttando la ridondanza introdotta dal bilanciamento del carico durante l'addestramento, LYNX permette alle architetture MoE di realizzare la loro promessa di efficienza sull'hardware futuro senza richiedere il riaddestramento del modello o il profiling offline. Il sistema dimostra che la gestione dinamica e consapevole del runtime degli esperti può disaccoppiare significativamente il throughput dai vincoli di larghezza di banda della memoria intrinseci agli attuali paradigmi di servizio MoE.