Riepilogo Tecnico: Inferenza Efficiente di Modelli di Diffusione LLM su Dispositivi tramite NPU Mobile

1. Definizione del Problema

Sebbene i modelli di diffusione Large Language Model (dLLM) rappresentino un'alternativa promettente ai modelli autoregressivi grazie alla generazione parallela dei token, il loro dispiegamento su dispositivi mobili affronta ostacoli significativi. A differenza della decodifica autoregressiva, che genera i token in sequenza, i dLLM denoisano iterativamente sequenze mascherate. Questo processo introduce tre sfide specifiche quando si punta alle NPU (Neural Processing Unit) mobili:

1. Impegno dei Token e Riduzione del Carico di Lavoro: Nella decodifica dLLM a blocchi, man mano che il denoising progredisce, il numero di token rimasti mascherati in un blocco diminuisce. Ciò porta a "step di coda" (tail steps) in cui il carico di lavoro effettivo è troppo piccolo per sfruttare appieno i motori tensoriali ad alto throughput delle NPU, causando una scarsa ammortizzazione dell'hardware.
2. Revisione dei Token e Riutilizzo della KV Cache: I dLLM richiedono un raffinamento iterativo degli stati dei token. I token impegnati precocemente in un blocco potrebbero ancora necessitare di revisione man mano che il contesto diventa più disponibile. Questo crea aggiornamenti sparsi e irregolari della cache Key-Value (KV), che interrompono i pattern di esecuzione densi e regolari preferiti dalle NPU.
3. Accesso alla Memoria e Vincoli dello Spazio di Indirizzamento: Le NPU mobili hanno uno spazio di indirizzamento virtuale (VA) limitato. L'inferenza dLLM richiede lo swapping frequente di pesi del modello, KV cache e attivazioni tra la memoria di sistema e la NPU. Ciò comporta costi di remapping e di trasferimento dati che possono annullare i benefici computazionali dell'accelerazione NPU.

Le soluzioni esistenti spesso si affidano a un offloading ingenuo verso la NPU, che non riesce ad affrontare la natura dinamica dell'inferenza dLLM rispetto al modello di esecuzione statico e basato su tile delle NPU mobili.

2. Metodologia: llada.cpp

Gli autori propongono llada.cpp, il primo framework di inferenza consapevole delle NPU (NPU-aware) progettato per allineare la natura di decodifica parallela dei dLLM con le caratteristiche esecutive delle Noche NPU mobili. Il framework integra tre tecniche core:

A. Decodifica Speculativa Multi-Blocco

Per affrontare il problema della riduzione del carico di lavoro nelle fasi finali della decodifica, llada.cpp espande dinamicamente la finestra di esecuzione.

• Meccanismo: Inveve di limitare il forward pass della NPU al blocco corrente, il sistema incorpora specularmente i token dai blocchi futuri nella finestra di decodifica attuale.
• Strategia: Utilizza lo Sliding Continuo (spostamento della finestra verso destra man mano che i token vengono impegnati) e l'Adaptive Lookahead (estensione proattiva della finestra per includere blocchi futuri quando il conteggio dei token mascherati nel blocco corrente scende sotto una determinata soglia).
• Beneficio: Questo mantiene un'alta densità di token attivi per la NPU, garantendo che gli oneri fissi (lancio dei kernel, trasferimenti di dati) siano ammortizzati su un carico di lavoro più ampio e produttivo senza violare la semantica originale della decodifica.

B. Revisione Progressiva a Doppio Percorso

Per gestire l'irregolarità delle revisioni dei token senza bloccare la NPU, il framework scinde la computazione densa dagli aggiornamenti sparsi.

• Stabilizzazione a Stadi dei Token: I token transitano attraverso tre stati: Invisibile (mascherato/bassa confidenza), Visibile (predizione concreta ma rivedibile) e Stabile (affidabile per il riutilizzo a lungo termine).
• Esecuzione a Doppio Percorso:
  • Percorso NPU: Dedicato al denoising denso del blocco corrente.
  • Percorso CPU: Gestisce gli aggiornamenti sparsi per i token "Visibili" che richiedono revisione. La CPU ricalcola in modo asincrono la confidenza, aggiorna gli stati KV e aggiorna le identità dei token utilizzando kernel specializzati che interpretano i layout a tile specifici della NPU (HMX).
• Salto Selettivo dei Logit: Per ridurre ulteriormente il carico della CPU, il sistema salta la proiezione dell'output e la stima della confidenza per i token che hanno già raggiunto lo stato "Stabile".
• Merge Ritardato: Le voci KV aggiornate dalla CPU vengono unite alla prefix cache solo ai confini dei passi (step boundaries), evitando la sincronizzazione fine.

C. Runtime di Memoria Ottimizzato per lo Swap

Per mitigare lo spazio di indirizzamento limitato della NPU e gli oneri di trasferimento, il framework ottimizza la gestione della memoria.

• Mappatura dei Buffer Guidata dal Grafo: Il runtime analizza il grafo computazionale per classificare i tensori come "a lunga durata" (pesi, KV cache) o "a breve durata" (attivazioni). Utilizza metriche profilate per determinare quali oggetti rimangono residenti nello spazio VA visibile alla NPU e quali vengono mappati temporaneamente, minimizzando la frammentazione.
• Trasferimento Dati Basato su Pipeline: Il sistema impiega il double-buffering. Mentre la NPU esegue gli operatori, l'host prepara i descrittori dei buffer e trasferisce i dati per l'operatore successivo in uno slot di staging separato. Questo sovrappone i costi di movimento dei dati e di mappatura con la computazione, nascondendo la latenza.

3. Contributi Chiave

Il documento rivendica i seguenti contributi:

• Identificazione del Disallineamento: Caratterizza il disallineamento fondamentale tra i modelli di esecuzione statici delle NPU mobili e la decodifica dinamica a blocchi dei dLLM, evidenziando specificamente le sfide nell'impegno dei token, nella revisione e nell'accesso alla memoria.
• Framework llada.cpp: Introduce il primo framework di inferenza consapevole delle NPU per i dLLM sugli smartphone, combinando innovazioni algoritmiche (decodifica speculativa, revisione progressiva) con ottimizzazioni a livello di sistema (runtime di memoria).
• Valutazione Completa: Valida il framework su diverse piattaforme hardware (Qualcomm SM8650, SM8750, SM8850) e carichi di lavoro (LLaDA-8B, Dream-7B), dimostrando una conversione efficace del parallelismo dei dLLM in una riduzione della latenza.

4. Risultati

La valutazione è stata condotta su smartphone tra cui il OnePlus Ace5 Pro (SM8750). Le scoperte chiave includono:

• Riduzione della Latenza: llada.cpp riduce la latenza di generazione end-to-end di LLaDA-8B da 17× a 42× rispetto a un baseline CPU con riutilizzo della cache KV del prefisso.
• Confronto con i Modelli Autoregressivi: Sui nuovi SoC (SM8750, SM8850), llada.cpp ottiene un'accelerazione fino a 3.9× rispetto a un modello autoregressivo di pari dimensioni (Llama-3-8B) eseguito sullo stesso hardware, evidenziando il vantaggio del parallelismo di diffusione sulle NPU mobili.
• Preservazione della Qualità: Il framework mantiene una qualità di generazione comparabile ai baseline CPU. Gli studi di ablazione mostrano che, sebbene la decodifica speculativa multi-blocco da sola possa ridurre leggermente l'accuratezza a causa dell'uso prematuro dei token futuri, il meccanismo di revisione progressiva a doppio percorso recupera la qualità mantenendo i primi token rivedibili.
• Efficienza Energetica: Fornendo un carico di lavoro parallelo sufficiente alla NPU, llada.cpp riduce sia la potenza media che l'energia per richiesta. Al contrario, l'esecuzione di modelli autoregressivi sulle NPU ha mostrato un consumo energetico più elevato a causa dell'overhead di pianificazione di carichi di lavoro sequenziali piccoli.
• Scalabilità: Il sistema gestisce con successo esecuzioni a lungo output (fino a 512 token) e diverse architetture dLLM (Dream-7B), provando la robustezza della gestione della memoria e delle strategie di decodifica.

5. Significato e Rivendicazioni

Il documento posiziona llada.cpp come un passo critico per rendere i dLLM praticabili per le applicazioni mobili sensibili alla latenza. Sostiene che la "computazione parallela ripetuta" dei dLLM, spesso vista come un limite, può essere convertita in un vantaggio sfruttando le capacità di calcolo tensoriale delle moderne NPU mobili.

Gli autori sottolineano che il semplice offloading dei dLLM alle NPU è insufficiente; il successo richiede un approccio di co-design che rispetti i vincoli dell'hardware (tile statici, spazio di indirizzamento limitato) pur sfruttando i punti di forza dell'algoritmo (denoising parallelo). Il lavoro dimostra che, con le giuste ottimizzazioni di sistema, i dLLM possono superare i tradizionali modelli autoregressivi sui dispositivi mobili, offrendo una via verso un'IA a bassa latenza, rispettosa della privacy e accessibile sugli smartphone.