DFlash Speculative Decoding su Apple Silicio: un balzo nelle prestazioni con MLX

Un recente sviluppo nel campo dell'inference di Large Language Models (LLM) su hardware locale ha mostrato progressi significativi. Un utente ha implementato il DFlash speculative decoding, una tecnica avanzata per accelerare la generazione di testo, specificamente per l'architettura Apple Silicio. Questa implementazione nativa, realizzata con il framework MLX di Apple, promette di sbloccare nuove capacità per i professionisti che operano con deployment on-premise e self-hosted.

I test iniziali, condotti su un sistema M5 Max equipaggiato con 64GB di memoria unificata, hanno rivelato un notevole incremento delle prestazioni. Per modelli come Qwen3.5-9B in formato bf16, la velocità di generazione ha raggiunto gli 85 token al secondo, rappresentando un miglioramento di 3.3 volte rispetto al baseline senza DFlash. Questi risultati sottolineano il potenziale delle ottimizzazioni software mirate a sfruttare appieno le peculiarità dell'hardware.

Dettagli tecnici e benchmark di performance

L'implementazione di DFlash su Apple Silicio si basa su un modello "draft" di dimensioni ridotte che genera in parallelo un blocco di 16 token, i quali vengono poi verificati dal modello "target" in un singolo passaggio forward. Questo processo garantisce un output identico al metodo greedy standard, ma con un'efficienza superiore. Le misurazioni escludono il tempo di prefill, concentrandosi unicamente sulla fase di generazione, con un'accettazione dei token tra l'80% e l'87% per tutti i modelli.

I benchmark dettagliati mostrano prestazioni variabili a seconda del modello e della lunghezza del contesto:
* Qwen3.5-9B (bf16): Fino a 85 token/s (3.3x di speedup) per 1024 token e 80 token/s (3.1x) per 2048 token.
* Qwen3.5-4B (bf16): Ha raggiunto 109 token/s (2.7x) per 1024 token e 133 token/s (3.2x) per 2048 token, mostrando una peculiarità interessante: la velocità aumenta con contesti più lunghi, poiché il modello è sufficientemente piccolo da mantenere un equilibrio ottimale tra draft e verifica.
* Qwen3.5-27B (quantized): Con quantization a 8bit, si sono registrati 35 token/s (2.5x) per 1024 token. Con quantization a 4bit, le prestazioni sono salite a 44 token/s (1.9x) per 1024 token. È interessante notare che la quantization a 8bit ha fornito un rapporto di speedup migliore rispetto alla 4bit, poiché quest'ultima rende la fase di verifica così rapida da spostare il collo di bottiglia sul modello draft in bf16.

Le ottimizzazioni chiave che hanno permesso questi risultati includono una patch per supportare head_dim=256 nell'attenzione di MLX, l'eliminazione di una sincronizzazione GPU-CPU per ciclo e l'impacchettamento delle proiezioni QKV per ridurre il numero di kernel dispatch.

Lezioni dall'architettura Apple Silicio

L'esperienza di sviluppo ha rivelato aspetti cruciali dell'architettura a memoria unificata di Apple Silicio. Su questi sistemi, le operazioni sono prevalentemente legate alla banda passante della memoria, un fattore che altera le dinamiche del speculative decoding. Contrariamente alle aspettative, l'implementazione di kernel Metal personalizzati per operazioni come batched-GEMV o fused gated SiLU si è dimostrata più lenta (0.5-0.8x) rispetto alle implementazioni GEMM standard di MLX, portando alla decisione di ripristinare le soluzioni di base.

Un'altra osservazione importante riguarda il costo della fase di verifica, che rimane quasi costante (57ms vs 59ms) passando da 4 a 16 token. Questo suggerisce che il caricamento dei pesi del modello domina il costo, piuttosto che il numero di token da verificare. Di conseguenza, strategie come "verificare meno token quando la confidenza è bassa" potrebbero non portare i benefici attesi in questo contesto. Per i modelli quantizzati, il panorama delle ottimizzazioni si ribalta: il modello draft (bf16) diventa più lento del modello target verificato (int4/int8), una limitazione strutturale del speculative decoding su hardware vincolato dalla banda passante con target quantizzati.

Prospettive future e implicazioni per il deployment on-premise

Il progetto è ancora in fase di sviluppo, con lavori in corso su diverse aree. Tra queste, la compressione o distillazione del modello draft per il Qwen3.5-27B, al fine di mitigare il collo di bottiglia del draft bf16 sui target quantizzati. Si sta anche lavorando sulla stabilità delle prestazioni con contesti lunghi, dove lo speedup tende a degradare oltre i 2048 token a causa della crescita della KV cache. L'esplorazione dei modelli MoE (Mixture of Experts) è un'altra priorità, con l'obiettivo di combinare il costo di verifica di un modello piccolo con la qualità di uno grande.

Questi sviluppi sono particolarmente rilevanti per le organizzazioni che considerano il deployment di LLM on-premise o in ambienti air-gapped. La capacità di ottenere prestazioni elevate su hardware locale, come Apple Silicio, offre alternative concrete alle soluzioni basate su cloud, con implicazioni positive per la sovranità dei dati e il TCO. Per chi valuta framework analitici per confrontare i trade-off tra deployment on-premise e cloud, AI-RADAR offre risorse e approfondimenti su /llm-onpremise. L'autore ha dichiarato l'intenzione di rendere il codice open source una volta che sarà pronto, il che potrebbe accelerare ulteriormente l'adozione di queste tecniche.