EarlySelectKV: selezione anticipatoria di KV tramite query proxy tra livelli

Abstract degli autori

La selezione sparsa di KV condizionata alla query riduce il costo dell’inferenza su contesti lunghi, concentrando l’attenzione su un sottoinsieme delle voci KV storiche a ogni passaggio di decodifica. I selettori esistenti devono attendere la query del livello bersaglio prima di assegnare punteggi ai candidati e generare gli indici, collocando il routing nel percorso immediatamente precedente all’attenzione. Presentiamo EarlySelectKV, una trasformazione che non richiede addestramento e usa query proxy ricavate da stati nascosti precedenti per precalcolare gli indici KV selezionati, mantenendo la query effettiva del livello bersaglio per l’attenzione finale. I test con RocketKV, Quest e Loki non mostrano un andamento coerente di peggioramento della qualità nelle impostazioni valutate. Nella valutazione principale di RocketKV, ES-RocketKV si discosta di +0,018 punti su LongBench, circa -0,01 punti su RULER e +0,25 punti percentuali nelle impostazioni rappresentative NIAH/Passkey. In una distribuzione di TensorRT-LLM con routing delegato alla CPU, ES-RocketKV aumenta il throughput end-to-end dell’output fino al 12,4% e riduce la memoria residente sulla GPU di un valore pari all’11,1% del pool combinato KV/K-summary. Questo rende i tempi del routing un obiettivo esplicito di ottimizzazione, al pari della selezione del sottoinsieme KV nella decodifica sparsa su contesti lunghi. Il codice è disponibile all’indirizzo https://github.com/sea-with-sakura/EarlySelectKV.