Step-dLLM: attenzione sparsa adattiva e sensibile ai passaggi per un’inferenza efficiente degli LLM a diffusione

Abstract degli autori

I modelli linguistici di grandi dimensioni a diffusione (dLLM) generano token in parallelo tramite denoising iterativo, ma la loro attenzione bidirezionale impedisce di usare la cache KV adottata dai modelli autoregressivi e impone di ricalcolare l’intera mappa di attenzione a ogni passaggio, con un costo quadratico per passaggio. I metodi esistenti di attenzione sparsa per i dLLM costruiscono una maschera sparsa dopo una fase iniziale di warm-up con attenzione completa e la riutilizzano nei passaggi successivi, distribuendo uniformemente il budget di sparsità tra le teste. Sosteniamo che entrambe le scelte ignorino proprietà strutturali chiave dell’attenzione nei dLLM: empiricamente, gli schemi di attenzione sono coerenti solo localmente tra i passaggi e cambiano bruscamente in corrispondenza di alcune transizioni; inoltre, le diverse teste concentrano la propria massa di attenzione in misura molto diversa. A questo scopo, progettiamo Step-dLLM, che aggiorna periodicamente una maschera sparsa a blocchi nei passaggi di ancoraggio lungo la traiettoria di denoising e la riutilizza all’interno di ciascun intervallo localmente coerente, impiegando al contempo un budget globale adattivo per testa che comprime in modo più aggressivo quelle con attenzione concentrata. Grazie a un kernel Triton fuso e consapevole dell’I/O per l’attenzione sparsa a blocchi, Step-dLLM si applica a LLaDA-1.5 e Dream-7B-Instruct preaddestrati, senza riaddestramento, ottenendo un’accelerazione end-to-end di 2,78× rispetto a FlashAttention densa su LLaDA-1.5 con un contesto di 8K, 256 passaggi di denoising e una sparsità del 90%, e migliorando il compromesso tra accuratezza e latenza rispetto ai precedenti metodi di attenzione sparsa.