KV Cache

Model dopisuje odpowiedź kawałek po kawałku. Wiele informacji o wcześniejszym tekście już obliczył, więc liczenie ich od początku przy każdym kroku byłoby powtarzaniem pracy. Chcesz zapamiętać wyniki, które pozostają ważne.

KV Cache przechowuje key i value z warstw attention. Są to listy liczb służące odpowiednio do dopasowania fragmentów i przekazywania ich informacji. Nie przechowuje gotowej odpowiedzi ani streszczenia rozmowy.

Po przetworzeniu A, B i C nowy fragment D korzysta z zapisanych opisów tych pozycji i dodaje własne. W modelu z blokadą dostępu do przyszłości nowe fragmenty nie zmieniają wcześniejszych opisów, więc można je wykorzystać ponownie przy zgodnych ustawieniach.

Oszczędzamy powtarzanie obliczeń, ale płacimy pamięcią na przechowane liczby. Przy długim kontekście ten koszt rośnie. Prefill przygotowuje cache znanego wejścia, a Decode uzupełnia go podczas odpowiedzi. Nadal trzeba odczytywać i mieszać potrzebne informacje.

Mechanizm i szczegóły

W modelu z Causal Masking późniejsze tokeny nie zmieniają reprezentacji wcześniejszych pozycji. To umożliwia ponowne użycie cache dla tego samego prefiksu i ustawień modelu. Podstawę przepływu informacji opisuje Attention Is All You Need, §3.1 i §3.2.3, a aktualizację cache w każdej warstwie — dokumentacja Hugging Face: How caching works.

Co zostaje z poprzedniego kroku

Własna ilustracja: model przetworzył tokeny A, B i C. Gdy przetwarza D, korzysta z zachowanych K i V dla trzech wcześniejszych pozycji i dopisuje parę dla D. Query bieżącej pozycji porównuje z dostępnymi key; wynik wykorzystuje odpowiednie value. Cache nie jest gotową odpowiedzią ani tekstowym streszczeniem rozmowy.

Przechowywanie wyników ma koszt. Własny mały przykład dla jednej warstwy: cztery pozycje, dwie KV heads, po trzy liczby w każdym wektorze K i V, cztery bajty na liczbę:

$$2\cdot4\cdot2\cdot3\cdot4=192\ \text{bajty}$$

Pierwszy czynnik 2 oznacza osobne key i value. Wynik obejmuje same elementy tensorów, bez narzutu alokacji. Pełny dynamiczny cache rośnie z liczbą przechowywanych pozycji; warianty z ograniczonym oknem mogą usuwać starsze pozycje.

Starszy artykuł o skalowaniu KV cache rozwija temat kosztów pamięci. Artykuł o przenoszeniu cache między modelami dotyczy odrębnego problemu: zgodne wymiary tensorów nie zapewniają zgodności ich znaczenia. Zmiana modelu lub prefiksu wymaga ponownego ustalenia, która część cache pozostaje poprawna. Zobacz także: PagedAttention opisuje blokową organizację pamięci, a Prefix Caching — ponowne użycie KV wspólnego prefiksu.