Grouped Value Attention: caching KV efficiente tramite ricostruzione delle chiavi su richiesta

Abstract degli autori

La cache KV è uno dei principali colli di bottiglia nella decodifica dei Transformer: la memoria che occupa e il traffico di lettura dalla cache crescono con la lunghezza della sequenza. L’attenzione a query raggruppate (GQA) riduce questo costo condividendo le teste di chiavi e valori, ma continua a memorizzare una chiave e un valore a ogni passaggio. Presentiamo Grouped Value Attention (GVA), che memorizza valori raggruppati e ricostruisce le chiavi di contenuto mediante una trasformazione lineare appresa. Durante l’inferenza, la trasformazione può essere incorporata nella query, eliminando la necessità di materializzare le chiavi di contenuto nel percorso di decodifica previsto. Un piccolo canale RoPE disaccoppiato e condiviso conserva le informazioni posizionali tramite una chiave posizionale memorizzata separatamente nella cache. Per le configurazioni studiate, questa rappresentazione riduce di circa il 45-47% il numero di scalari conservati stabilmente nella cache rispetto a configurazioni GQA equivalenti. Su una scala di 350 milioni di parametri, con 30 miliardi di token FineWeb-Edu, la variante posizionale a 16 dimensioni raggiunge una precisione media di 44,18 su cinque compiti, contro 44,36 per GQA e 43,88 per MLA. Questi risultati mostrano una precisione nei benchmark quasi pari a quella di GQA con una rappresentazione della cache più compatta. Per tradurre questa rappresentazione compatta in un’inferenza autoregressiva più rapida, abbiamo sviluppato kernel di decodifica personalizzati e stiamo valutando le loro prestazioni di inferenza end-to-end, con un rilascio open source previsto a breve.