DropKV: disaccoppiare la perturbazione dell’output residuo per un’espulsione quasi ottimale dalla cache KV

Abstract degli autori

La gestione della cache delle chiavi e dei valori (KV) è diventata un ostacolo cruciale all’uso dei modelli linguistici di grandi dimensioni (LLM) in scenari con contesti lunghi, a causa della crescita lineare della memoria necessaria. Sebbene i metodi di espulsione più recenti siano andati oltre le semplici euristiche basate sui punteggi di attenzione, includendo le rappresentazioni dei valori, devono in genere affrontare un compromesso fondamentale: l’ottimizzazione necessaria per ridurre al minimo la perturbazione dell’output è un problema combinatorio NP-difficile, che porta a euristiche matematicamente subottimali o a implementazioni difficili da realizzare tramite kernel. In questo lavoro proponiamo DropKV, un semplice metodo di espulsione fondato sul disaccoppiamento della perturbazione dell’output residuo. Scomponendo la decisione congiunta di espulsione in valutazioni indipendenti per ciascun token, DropKV aggira la complessità combinatoria e offre una garanzia dimostrabile di approssimazione entro un fattore costante, a condizione che sia soddisfatta una condizione di cono che verifichiamo empiricamente su tre LLM per contesti lunghi. Per garantirne l’applicabilità pratica, forniamo un’implementazione ad alte prestazioni basata su kernel Triton fusi che evitano di materializzare l’intera matrice di attenzione. Otteniamo così un’accelerazione di 19,8 volte del kernel di valutazione, che si traduce in un’accelerazione end-to-end del prefill fino al 9,9% a parità di batch. Valutazioni approfondite sui benchmark RULER e LongBench mostrano che DropKV supera sistematicamente i metodi di riferimento più avanzati a parità di spazio disponibile nella cache, ottenendo la latenza di inferenza più bassa tra i metodi di espulsione e rimanendo al tempo stesso più veloce del metodo di riferimento con cache dinamica.