TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs

Riepilogo Tecnico: TileFuse

Problematica

L'implementazione di Large Language Models (LLM) su dispositivi edge dipende sempre più dalle Neural Processing Unit (NPU) per soddisfare rigorosi budget di potenza e termici. Tuttavia, l'inferenza pratica di LLM su NPU client, specificamente l'architettura AMD XDNA2, affronta barriere significative. L'ostacolo principale è la frammentazione degli stack software e la mancanza di programmabilità a basso livello. La maggior parte dei framework NPU (ad es., Apple Core ML, Intel OpenVINO, Qualcomm QNN) espone solo interfacce di alto livello o percorsi di deployment proprietari che non supportano nativamente formati di quantizzazione standard e ampiamente utilizzati come AWQ (Activation-Weight Quantization).

Di conseguenza, gli sviluppatori sono spesso costretti a rimodellare i modelli per adattarli a schemi di quantizzazione specifici per NPU o a fare affidamento su un'esecuzione a precisione piena meno efficiente. Sebbene lavori precedenti abbiano esplorato la quantizzazione hardware-aware per piattaforme specifiche (ad es., ScalingLLM di Qualcomm), essi spesso non riescono a supportare formati standard come $W4A16$ o non possono gestire le grandi dimensioni delle matrici (ad es., >8K) comuni nei moderni LLM. Inoltre, le ottimizzazioni NPU esistenti spesso mancano della capacità di fondere l'unpacking e la dequantizzazione direttamente nel kernel di calcolo, portando a traffico di memoria non necessario e a un movimento inefficiente dei dati.

Metodologia

TileFuse è una libreria di kernel a precisione mista "close-to-metal" progettata per le NPU AMD XDNA2, mirata ai layer lineari (GEMM e GEMV) che dominano l'inferenza degli LLM. Il sistema sfrutta il flusso di programmazione aperto MLIR-AIE (IRON) per co-progettare il layout dei pesi, il posizionamento dei metadati, i microkernel e i dataflow a livello di array.

1. Microkernel a Precisione Mista Fusi

TileFuse fonde tre operazioni distinte in un singolo flusso di kernel:

• Unpacking: Estrazione dei pesi INT4/INT8 impacchettati.
• Dequantizzazione: Conversione degli interi a basso bit in BF16 utilizzando scale e zero-point per gruppo o per canale.
• Esecuzione: Esecuzione della moltiplicazione di matrici (GEMM) o della moltiplicazione matrice-vettore (GEMV).

Mantenendo i pesi in formato compresso fino a quando non raggiungono il core di calcolo, TileFuse evita il round-trip nella memoria globale richiesto per materializzare i pesi a precisione piena (BF16). I microkernel sono progettati per scompattare e dequantizzare ogni tile di peso una sola volta, memorizzare i valori BF16 risultanti nella memoria locale del core di calcolo e riutilizzarli per le successive operazioni di matrice, preservando così i pattern di riutilizzo essenziali per le prestazioni del GEMM a tile.

2. Pre-Tiling Interleaved

Per affrontare i limiti dei vincoli di streaming DMA a runtime (che tipicamente limitano le dimensioni supportate sotto gli 8K), TileFuse impiega una strategia di pre-tiling interleaved offline.

• Metadata Packing: I metadati di quantizzazione (scale e zero-point) sono impacchettati direttamente nei tile di peso in memoria, garantendo che tutti i dati necessari per la dequantizzazione siano disponibili in modo contiguo.
• Riorganizzazione del Layout: La matrice dei pesi viene riorganizzata in un layout interleaved column-major. Ciò assicura che i tile di peso assegnati alla stessa colonna AIE siano contigui in memoria, eliminando i grandi stride di memoria che altrimenti costringerebbero un'unica operazione GEMM a essere suddivisa in più invocazioni di kernel. Questo approccio supporta dimensioni GEMM fino a 32K.

3. Co-Design del Dataflow GEMV

La generazione di token negli LLM edge è tipicamente un'operazione GEMV con batch-size-1, che spesso è memory-bound in precisione piena ma diventa compute-bound quando vengono introdotti la quantizzazione e la dequantizzazione.

• Limitazione di Base: Le mappature GEMV standard utilizzano spesso solo una singola riga dell'array di calcolo AIE 4×8, lasciando la maggior parte dell'hardware inattivo.
• Ottimizzazione: TileFuse ridisegna il dataflow per distribuire i pesi attraverso l'intero array 4×8. Utilizza i core di memoria come uno strato di distribuzione intermedio, trasmettendo bundle di tile di peso dai core shim ai core di memoria, che poi li distribuiscono a tutte e quattro le righe di calcolo. Ciò massimizza l'utilizzo spaziale e il throughput aritmetico per la generazione di token.

4. Ottimizzazione Consapevole della Riconfigurabilità

Il sistema sfrutta esplicitamente la riconfigurabilità dell'architettura XDNA2 per personalizzare i microkernel dei core di calcolo e i pattern di accesso DMA per carichi di lavoro specifici. Tuttavia, gli autori riconoscono che questa flessibilità introduce overhead di dispatch a runtime e di riconfigurazione del fabric. TileFuse è ottimizzato per ammortizzare questo overhead durante le grandi GEMM di prefilling, ma ne riconosce i limiti per le operazioni GEMV su piccola scala e sensibili alla latenza.

Contributi Chiave

1. Kernel Fusi per la Quantizzazione Pratica: Implementazione di kernel AIE fusi che consumano direttamente pesi in stile AWQ $W4A16$ e $W8A16$, eliminando la necessità di rimodellamento dei modelli o schemi di quantizzazione proprietari.
2. Pre-Tiling Interleaved: Una strategia di layout che supporta dimensioni GEMM fino a 3K e consente uno streaming efficiente a runtime allineando l'accesso alla memoria con i vincoli dei descrittori dei buffer DMA.
3. Redesign del Dataflow GEMV: Un nuovo dataflow che distribuisce i pesi attraverso l'intero array AIE 4×8, migliorando significativamente il throughput della generazione di token.
4. Efficienza Edge End-to-End: Integrazione di questi kernel in un runtime ibrido (NPU per i layer lineari, iGPU per l'attenzione/layer non lineari) che dimostra guadagni sostanziali in latenza ed efficienza energetica.

Risultati

Le valutazioni sono state condotte su laptop Ryzen AI (Ryzen AI 7 350 e Ryzen AI 9 HX 370) utilizzando modelli come Llama3-8B, Gemma-2B e Qwen2.5-3B.

• Prestazioni dei Kernel:
  • GEMM: TileFuse ottiene un miglioramento delle prestazioni fino al 121,6% rispetto alle baseline a precisione piena.
  • GEMV: TileFuse ottiene un miglioramento delle prestazioni fino al 281% rispetto alle baseline a precisione piena.
  • Confronto: Su GEMM, TileFuse offre guadagni di prestazioni ed efficienza energetica superiori a 2× rispetto alle forti baseline iGPU (kernel fusi di llama.cpp).
• Inferenza LLM End-to-End:
  • Prefilling: Su Ryzen AI 7 350, TileFuse raggiunge fino a 2,0× di latenza di prefilling inferiore (ad es., 2,00× di speedup per Llama3-8B a 2048 token) rispetto alla baseline iGPU.
  • Energia: Il percorso NPU riduce il consumo energetico di oltre il 64,6% su Ryzen AI 7 350 e del 52,8% su Ryzen AI 9 HX 370.
  • Generazione di Token: La NPU è meno performante della iGPU per la generazione di token a causa dell'overhead di dispatch del kernel e della riconfigurazione su scala millisecondo rispetto all'esecuzione su scala microsecondo degli operatori GEMV.
• Confronto con il Lavoro Precedente: TileFuse supera gli altri framework guidati da NPU come FastFlowLM e Scaling NPU nel throughput di prefilling (ad es., rispettivamente 1,45× e 1,61× di speedup su Llama3-8B e Qwen2.5-3B).

Significato e Claim

Il paper sostiene che TileFuse dimostri che le NPU AMD XDNA2 sono un target pratico per un'efficiente inferenza LLM quantizzata "off-the-shelf". Portando formati standard come AWQ direttamente sull'hardware attraverso l'ottimizzazione close-to-metal, questo lavoro colma il divario tra il potenziale grezzo delle NPU e il deployment reale sui client.

Gli autori sottolineano che il successo di TileFuse si basa su una strategia di deployment ibrida: utilizzare l'array spaziale riconfigurabile della NPU per il prefilling a lungo contesto (dove l'overhead è ammortizzato), affidandosi alla iGPU per la generazione di token sensibile alla latenza. Questo approccio valida il fatto che le NPU possono migliorare sostanzialmente l'usabilità degli LLM edge quando l'hardware viene adattato ai formati di quantizzazione utilizzati dagli utenti, piuttosto che forzare i modelli a conformarsi a vincoli proprietari. Il lavoro stabilisce che, con un appropriato co-design di kernel e dataflow, le NPU possono offrire significativi benefici di prestazione ed energia per l'AI edge, anche in presenza di frammentazione dello stack software.