Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

Il paper presenta Helios, un acceleratore hardware-software basato su DRAM 3D e bonding ibrido che, attraverso una gestione dinamica della cache KV e un'esecuzione distribuita dell'attenzione, supera le limitazioni delle architetture NMP esistenti per offrire un servizio LLM più veloce ed energeticamente efficiente.

L'essenziale

Immagina di dover gestire un gigantesco ristorante di intelligenza artificiale, dove i clienti (le persone che chiedono risposte a un chatbot) arrivano in modo imprevedibile: a volte sono pochi e ordinano piatti semplici, altre volte sono una folla che richiede ricette lunghissime e complesse.

Questo è il problema che affronta la ricerca "Helios", presentata da un team di ricercatori (tra cui l'Università di Pechino e HKUST). Ecco come funziona, spiegato in modo semplice.

1. Il Problema: Il Collo di Bottiglia della Memoria

Oggi, i modelli linguistici (come quelli che usi per scrivere email o creare storie) hanno due fasi:

• La fase di "pensiero" (Prefill): Analizzano la tua domanda. È veloce e richiede molta potenza di calcolo.
• La fase di "risposta" (Decoding): Scrivono la risposta parola per parola. Qui il problema è la memoria. Per scrivere ogni nuova parola, il modello deve "ricordare" tutto ciò che ha scritto prima. Questo ricordo si chiama KV Cache.

I computer attuali (come le GPU) sono bravissimi a calcolare, ma lenti a recuperare queste memorie. È come avere un cuoco velocissimo che però deve correre in un magazzino enorme e disordinato ogni volta che gli serve un ingrediente. Inoltre, i clienti arrivano in modo caotico: alcuni ordinano un panino (richiesta breve), altri un banchetto (richiesta lunga). I sistemi attuali sono rigidi: assegnano un "ripiano" fisso del magazzino a ogni cliente. Se il cliente finisce il panino, il ripiano rimane vuoto ma bloccato. Se il cliente vuole un banchetto, non c'è spazio e deve aspettare. Questo spreca spazio e tempo.

2. La Soluzione: Helios, il "Cucina-Magazzino" Intelligente

Gli autori propongono Helios, un nuovo acceleratore basato su una tecnologia chiamata Hybrid Bonding (unione ibrida).

Immagina Helios non come un cuoco e un magazzino separati, ma come una cucina dove i fornelli (i processori) sono costruiti direttamente sopra gli scaffali dei ripiani (la memoria). Non devi più correre avanti e indietro; l'ingrediente è già sotto il tuo dito.

Ma la vera magia non è solo la vicinanza, è come gestiscono gli ingredienti:

• Il vecchio metodo (Rigido): Assegnavano un intero scaffale a un cliente, anche se il cliente ne usava solo un angolo. Se arrivava un altro cliente, lo scaffale era "pieno" anche se c'era spazio, perché era tutto bloccato.
• Il metodo Helios (Flessibile): Usano un sistema a blocchi. Immagina che il magazzino sia fatto di mattoncini LEGO. Se un cliente ha una richiesta breve, Helios gli dà solo 2 mattoncini. Se un cliente ha una richiesta lunga, gli dà 100 mattoncini, sparsi ovunque nel magazzino, non tutti in un unico posto.
  • Vantaggio: Non c'è più spazio sprecato. Il magazzino è sempre pieno al 100% della sua capacità utile.
  • Vantaggio: Se un cliente finisce e lascia i mattoncini, questi tornano subito disponibili per il prossimo.

3. Come Collaborano i "Cuochi" (I Processori)

Helios ha molti piccoli cuochi (chiamati Processing Engines o PE) che lavorano insieme.
Invece di far lavorare ogni cuoco su un compito intero e isolato, Helios divide il lavoro in modo intelligente:

• Se devi scrivere una frase lunga, la dividete in pezzi. Ogni cuoco scrive una parte, poi si passano i pezzi di carta (i dati) velocemente l'uno all'altro.
• Usano una mappa intelligente per decidere quale cuoco prende quale pezzo di memoria. Non scelgono a caso, ma guardano chi è più vicino e chi ha meno lavoro, per evitare che un cuoco resti fermo mentre un altro è sommerso di lavoro.

4. I Risultati: Più Veloce ed Efficiente

Grazie a questo approccio "co-progettato" (dove hardware e software sono pensati insieme fin dal primo giorno), Helios ottiene risultati straordinari:

• Velocità: È circa 3,25 volte più veloce dei sistemi attuali (come le GPU Nvidia) quando gestisce molte richieste contemporaneamente.
• Risparmio Energetico: Consuma 3,36 volte meno energia per fare lo stesso lavoro.
• Tempi di Attesa: Riduce drasticamente il tempo tra una parola e l'altra (anche del 72-76%), rendendo la conversazione molto più fluida per l'utente.

In Sintesi

Pensa a Helios come alla differenza tra un vecchio magazzino con scaffali fissi e un magazzino moderno automatizzato con robot che spostano le scatole dove servono.
Mentre i computer attuali si bloccano perché cercano di adattarsi a un sistema rigido, Helios è fluido: si adatta alla forma della richiesta, riempiendo ogni singolo centimetro di spazio disponibile e facendo lavorare tutti i processori in perfetta armonia. È un passo avanti fondamentale per rendere le intelligenze artificiali più veloci, economiche e capaci di gestire milioni di utenti contemporaneamente.

Sommario tecnico

Titolo: Co-design Hardware-Software per un Acceleratore di Servizio LLM basato su 3D-DRAM

Autore: HeLioS (Hybrid-bonding-based LLM Serving accelerator)

1. Il Problema

I Large Language Models (LLM) sono ampiamente utilizzati per servizi generativi online, caratterizzati da carichi di lavoro altamente dinamici:

• Fluttuazione delle richieste: Tassi di arrivo variabili (da picchi a bassi carichi) e lunghezze di sequenza non uniformi (da poche centinaia a 10.000+ token).
• Natura dei carichi di lavoro: Le fasi di inferenza (prefill e decoding) richiedono sia operazioni computazionalmente intensive (prefill) che operazioni limitate dalla memoria (decoding, specialmente per la gestione della KV cache).

Limitazioni delle soluzioni esistenti (NMP - Near-Memory Processing):
Le architetture NMP esistenti, che integrano motori di elaborazione (PE) nei moduli di memoria (es. HBM), soffrono di due limitazioni fondamentali che ne impediscono l'efficienza in scenari di servizio dinamici:

1. Gestione della KV Cache a grana grossa: Le soluzioni attuali assegnano staticamente l'intera KV cache di una testa di attenzione a un modulo NMP fisso. Questo causa:
   • Squilibrio del carico computazionale: Se il numero di teste di attenzione attive è inferiore al numero di moduli, alcune risorse rimangono inattive.
   • Frammentazione della memoria: Non è possibile allocare richieste che superano la capacità di un singolo modulo, anche se la capacità totale del sistema è sufficiente. Inoltre, si spreca memoria allocando lo spazio massimo possibile per ogni richiesta invece di quello effettivo.
2. Flusso di esecuzione dell'attenzione inflessibile: L'attenzione viene eseguita come due GEMM (General Matrix Multiply) indipendenti su matrici intere. Questo impedisce la partizione fine della KV cache necessaria per gestire lunghezze di sequenza variabili e richiede che l'intera cache risieda nello stesso modulo, limitando la flessibilità.

2. Metodologia: HeLioS

HeLioS è un acceleratore basato su Hybrid Bonding (HB) che colma il divario tra la natura dinamica del servizio LLM e l'astrazione della memoria non uniforme delle architetture NMP.

A. Architettura Hardware

• Tecnologia Hybrid Bonding: Utilizza il bonding diretto Cu-Cu per impilare die DRAM su un die logico. Questo offre un'alta densità di I/O (110.000/mm²), bassa latenza e un'efficienza energetica superiore rispetto a HBM/GDDR.
• Array di PE Distribuiti: Il die logico ospita una griglia di Processing Engines (PE) (es. 4x4 mesh). Ogni PE ha accesso diretto alla memoria locale (DRAM) e comunica con gli altri tramite una rete on-chip (NoC).
• Unità Specializzate: Ogni PE contiene unità per GEMM (Matrice), Softmax Online, vettori e riduzione (Reduction), ottimizzate per l'esecuzione distribuita.

B. Esecuzione degli Operatori (Software/Hardware Co-design)

• Attenzione Tiled Distribuita: Invece di calcolare l'attenzione su sequenze intere, HeLioS divide la KV cache in blocchi (es. 64 token). L'attenzione viene eseguita in modo iterativo su questi blocchi (Tiled Attention), permettendo a più PE di lavorare su parti diverse della stessa sequenza.
• Flusso di Esecuzione Ottimizzato:
  • Intra-PE: Sovrapposizione di operazioni vettoriali (softmax) e matriciali (GEMM) per nascondere la latenza.
  • Inter-PE: Primitive di comunicazione collettive ottimizzate (All-Reduce, All-Gather, Reduce-Scatter) su topologia Mesh per gestire la distribuzione delle chiavi, valori e somme parziali.
• Supporto per Varianti LLM: L'architettura supporta nativamente MHA, GQA, MQA e MLA (Multi-Head Latent Attention) tramite tecniche di assorbimento delle matrici per evitare la ricostruzione esplicita di tensori KV completi.

C. Gestione Dinamica della KV Cache (Sistema)

• Allocazione Consapevole dello Spazio (Spatially-Aware): Viene proposto un algoritmo di allocazione che distribuisce i blocchi KV su tutti i PE disponibili, non su un modulo fisso.
  • Bilanciamento del carico: Assegna i blocchi ai PE con meno token totali per bilanciare il carico computazionale.
  • Minimizzazione della comunicazione: Posiziona gli ultimi blocchi (necessari per il prossimo step di decoding) su PE centrali o simmetrici per ridurre il numero di hop nella rete NoC durante le operazioni di scatter/gather.
• Gestione a Blocchi: Simile a vLLM, ma adattata all'architettura distribuita NMP, permettendo l'allocazione dinamica di blocchi liberi man mano che le richieste arrivano e terminano.

3. Contributi Chiave

1. Analisi delle Limitazioni: Identificazione del divario fondamentale tra l'astrazione della memoria non uniforme NMP e gli algoritmi di gestione dinamica della KV cache basati su GPU.
2. Architettura HeLioS: Progettazione di un acceleratore HB-native che abilita l'attenzione tiled distribuita end-to-end.
3. Gestione Dinamica della KV Cache: Implementazione di un meccanismo di allocazione spazialmente consapevole che bilancia il carico di calcolo e minimizza il trasferimento dati tra PE.
4. Co-design Completo: Integrazione di strategie di esecuzione intra-PE, primitive di comunicazione inter-PE e schedulazione di sistema.

4. Risultati Sperimentali

Le valutazioni sono state condotte su modelli come OPT 66B, LLaMA3 70B, Mixtral, Qwen3 e DeepSeek 236B, confrontando HeLioS con GPU (A100), soluzioni NMP in-die (AttAcc), soluzioni Duplex e varianti non-tiled.

• Velocità (Speedup): HeLioS raggiunge un speedup geometrico medio di 3.25x rispetto alle migliori soluzioni GPU/NMP esistenti.
• Efficienza Energetica: Miglioramento di 3.36x nell'efficienza energetica (geometric mean).
• Prestazioni di Servizio (Serving):
  • Riduzione del tempo tra i token (Time-Between-Tokens, TBT) fino al 72% (P50) e 76% (P99) rispetto alle baseline.
  • Migliore capacità di servizio (throughput) grazie alla riduzione della frammentazione della memoria.
• Analisi della Frammentazione: Le soluzioni NMP esistenti mostrano un tasso di frammentazione della memoria 4.7x - 23x superiore a HeLioS, che mantiene la frammentazione sotto il 15% anche sotto carico elevato.
• Overhead di Comunicazione: La comunicazione inter-PE rappresenta solo il 7.08% della latenza totale di decoding, dimostrando l'efficienza del design della rete e delle primitive di comunicazione.

5. Significato e Impatto

Il lavoro di HeLioS è significativo perché:

• Risoluzione del Collo di Bottiglia della Memoria: Dimostra che l'uso di tecnologie di bonding ibrido (3D-DRAM) combinato con una gestione software intelligente della memoria può superare i limiti di banda delle GPU tradizionali per l'inferenza LLM.
• Adattabilità Dinamica: Offre una soluzione scalabile che gestisce efficacemente la variabilità delle richieste (lunghezza e arrivo) tipica dei servizi cloud reali, un punto debole delle architetture NMP statiche precedenti.
• Nuovo Paradigma di Co-design: Stabilisce un nuovo standard per l'accelerazione LLM, dove l'architettura hardware (mesh di PE, bonding ibrido) e le strategie software (allocazione spaziale, attenzione tiled) sono progettate congiuntamente per massimizzare l'efficienza in scenari di servizio dinamici.

In sintesi, HeLioS rappresenta un passo avanti cruciale verso acceleratori LLM dedicati che possono fornire servizi generativi ad alta velocità, bassa latenza ed elevata efficienza energetica, superando le limitazioni delle architetture basate su GPU e delle prime generazioni di acceleratori Near-Memory.