Cache prefiksu: ponowne użycie cache KV między turami i zapytaniami

Przy każdej turze agent wysyła serwerowi całą rozmowę od nowa: prompt systemowy, definicje narzędzi, historię i nową wiadomość na końcu. Prefill tego wszystkiego od zera to największe pojedyncze marnotrawstwo przy serwowaniu agentów. README mojego forka podaje, że ponowny prefill tury z długim kontekstem trwa około 46 sekund, a po ponownym użyciu cache około 460 ms. Poniżej to, jak działa ponowne użycie i co je po cichu psuje.

Dlaczego ponowne użycie jest możliwe

W modelu przyczynowym klucze i wartości na pozycji N zależą tylko od tokenów od 0 do N. Jeśli więc nowe zapytanie zaczyna się od dokładnie tych samych tokenów co zapytanie, które serwer już przetworzył, wpisy KV dla tego wspólnego prefiksu są identyczne. Serwer może je zachować, znaleźć najdłuższy wspólny prefiks i zrobić prefill tylko dla tokenów po nim.

Dla agenta to prawie wszystko. Tura 40 dzieli z turą 39 pierwsze 99% tokenów.

Co psuje ponowne użycie

Ponowne użycie wymaga, żeby prefiks był identyczny token w token. Jeden zmieniony token unieważnia wszystko po nim, więc najwięcej kosztuje to, co zmienia się wcześnie:

• Znacznik czasu w prompcie systemowym. Zmienia się przy każdym zapytaniu, więc nic po nim nigdy nie zostaje użyte ponownie.
• Definicje narzędzi w innej kolejności albo wygenerowane od nowa z innymi białymi znakami. Tokeny się różnią, nawet jeśli znaczenie nie.
• Historia przepisywana w miejscu, na przykład streszczanie starych tur na początku. Każda tura po miejscu edycji jest liczona od nowa.
• Przycinanie od początku, żeby zmieścić się w kontekście. Cały prefiks się przesuwa, a ponowne użycie spada do zera.

Stąd zasada projektowa: kontekst tylko dopisuj, stabilną treść umieszczaj na początku, a zmienną na końcu.

W obrębie serwera i między zapytaniami

• Serwer llama.cpp ponownie używa cache slotu, gdy nowy prompt ma wspólny prefiks z tym, co ten slot ostatnio przetworzył (cache_prompt). Potrafi też zapisać stan slotu na dysk i go przywrócić (--slot-save-path).
• Automatyczny cache prefiksu w vLLM haszuje bloki KV według zawartości. Każde zapytanie, którego prefiks pasuje do bloków w cache, współdzieli je między użytkownikami, na bazie tabel bloków PagedAttention (te tabele omawia notatka o batchingu).
• Dostawcy API stosują tę samą zasadę i rozliczają tokeny prefiksu z cache ze zniżką. Zasady projektowania promptów są identyczne.

Modele hybrydowe potrzebują checkpointów

W czystym transformerze cofnięcie cache do wspólnego prefiksu oznacza po prostu odrzucenie wpisów po nim. Modele hybrydowe, takie jak te z Gated DeltaNet, trzymają też stan rekurencyjny, a tego stanu nie da się cofnąć: stanu na pozycji N nie da się odzyskać ze stanu na późniejszej pozycji. Ponowne użycie prefiksu wymaga migawki zrobionej na tej pozycji (dlaczego, wyjaśnia notatka o modelach hybrydowych).

Mój fork trzyma do 32 checkpointów promptu po około 149 MiB, czyli około 4,7 GB w RAM hosta. Ich przechwytywanie wcześniej blokowało prefill, bo każde przechwycenie synchronizowało cały kontekst. Przerobiłem przechwytywanie na asynchroniczne, przez przypięty, pierścieniowy bufor pośredni z dwoma slotami; konstrukcję opisuje notatka o forku.

Zasady, których się trzymam

1. Projektuj prompty pod cache. Najpierw stabilny prompt systemowy i narzędzia, żadnych znaczników czasu ani identyfikatorów zapytań blisko początku, historia dopisywana, a nie przepisywana.
2. Kompaktuj na końcu, nie na początku. Jeśli historię trzeba streścić, zrób to na wyraźnych granicach i zaakceptuj jedno chybienie cache, zamiast przepisywać początek w każdej turze.
3. Mierz współczynnik trafień. Liczba tokenów promptu przetworzonych w turze powinna być małym ułamkiem liczby wysłanych tokenów promptu. Jeśli nie jest, coś na początku promptu się zmienia.
4. Zaplanuj budżet pamięci hosta na checkpointy w modelach hybrydowych i ogranicz, ile ich trzymasz.

Źródła: liczby 46 s i 460 ms oraz rozmiary checkpointów pochodzą z README mojego forka. Flagi serwera pochodzą z dokumentacji serwera llama.cpp. Kwon i in., PagedAttention (SOSP 2023), współdzielenie na poziomie bloków.