WaveFront Decoding: decodifica auto-speculativa parallelizzata per modelli linguistici con cicli

Abstract degli autori

I modelli linguistici con cicli applicano ripetutamente un blocco con pesi condivisi per aumentare la profondità effettiva senza aumentare il numero di parametri. Le conseguenti T chiamate sequenziali al blocco ricorrente per ogni token generato, però, aumentano notevolmente la latenza della decodifica. Per affrontare il problema, presentiamo Wavefront Decoding (WFD), un framework di decodifica auto-speculativa che non richiede addestramento, progettato per i modelli linguistici con cicli. WFD sfrutta due proprietà di queste architetture: gli output intermedi delle ricorrenze forniscono previsioni preliminari efficaci e la condivisione dei pesi consente di elaborare gli stati dei token relativi a posizioni e profondità di ricorrenza diverse in un’unica chiamata in batch al blocco ricorrente. WFD dispone questi stati a profondità diverse in un fronte d’onda diagonale: genera continuamente previsioni preliminari per nuove posizioni a bassa profondità, mentre porta le posizioni precedenti verso la verifica alla profondità massima. A differenza dello schema in cui la generazione delle previsioni preliminari e la verifica avvengono in fasi separate, WFD le raggruppa quindi simultaneamente in batch nelle stesse chiamate al blocco ricorrente; le previsioni preliminari respinte vengono corrette usando quelle ottenute alla profondità massima. Nelle sei categorie di attività di Spec-Bench, WFD raggiunge un’accelerazione di 2,42x su Ouro-2.6B e di 3,54x su Huginn-3.5B rispetto alla decodifica autoregressiva, superando sistematicamente lo schema di generazione preliminare seguita dalla verifica. La condivisione KV tra ricorrenze riduce ulteriormente il traffico KV del fronte d’onda e porta l’accelerazione di WFD a 4,81x su Huginn-3.5B. Il codice è disponibile all’indirizzo https://github.com/summerbro-hhj/wavefront-decoding.