Matematyka cache KV: ile kosztuje długi kontekst i co daje 4-bitowy cache

Wagi modelu to koszt stały: ładujesz je raz i się nie zmieniają. Cache KV rośnie z każdym tokenem kontekstu i z każdą równoległą rozmową. Przy długim kontekście to on, a nie wagi, decyduje, co mieści się na karcie. Poniżej rachunek, policzony na modelu, na którym pracuję codziennie.

Czym jest cache KV

Podczas dekodowania każdy nowy token „patrzy” na wszystkie tokeny przed nim. Liczenie kluczy i wartości dla całego prefiksu na każdym kroku zrobiłoby z generowania proces kwadratowy. Silnik liczy więc wektory klucza i wartości każdego tokenu raz, w każdej warstwie atencji, i je przechowuje. Ten magazyn to cache KV.

Cache jest per sekwencja. Dwie rozmowy prowadzone równolegle potrzebują dwóch cache, nawet na tym samym modelu.

Bajty na token

kv_bajty_na_token = 2 × n_warstw_atencji × n_głowic_kv × wymiar_głowicy × bajty_na_wartość
                    └ K i V

Zmniejszają go trzy rzeczy:

• Grouped-query attention. Kilka głowic zapytań dzieli jedną głowicę klucza i wartości, więc n_głowic_kv to ułamek liczby głowic zapytań.
• Warstwy hybrydowe. Warstwy z atencją liniową, takie jak Gated DeltaNet, trzymają stan rekurencyjny o stałym rozmiarze zamiast cache na token. Do n_warstw_atencji liczą się tylko warstwy z pełną atencją.
• Kwantyzacja cache. To jedyna z tych trzech rzeczy, którą wybierasz przy serwowaniu, a nie przy treningu.

Przykład: model 27B przy 262K

Do tego wyniku nie potrzebuję konfiguracji modelu. README forka podaje 4224 MiB cache TBQ4 przy 262 144 tokenach, a TBQ4 zapisuje 128 wartości w 66 bajtach. To 2 bajty skali plus 64 bajty 4-bitowych indeksów, czyli 4,125 bita na wartość.

4224 MiB × 1 048 576 / 262 144 tokenów  = 16 896 bajtów na token
16 896 / 66 × 128                       = 32 768 wartości w cache na token

Te 32 768 wartości to 2 × n_warstw_atencji × n_głowic_kv × wymiar_głowicy dla tego modelu. Dalej każdy format to już arytmetyka:

Format KV	Bity na wartość	Na token	8K	32K	128K	262K
FP16	16	64 KiB	512 MiB	2 GiB	8 GiB	16 GiB
Q8_0 (34 B / 32)	8,5	34 KiB	272 MiB	1,06 GiB	4,25 GiB	8,5 GiB
Q4_0 (18 B / 32)	4,5	18 KiB	144 MiB	576 MiB	2,25 GiB	4,5 GiB
TBQ4 (66 B / 128)	4,125	16,5 KiB	132 MiB	528 MiB	2,06 GiB	4,125 GiB
3 bity (50 B / 128)	3,125	12,5 KiB	100 MiB	400 MiB	1,56 GiB	3,125 GiB

Wagi w Q4_K_M zajmują około 17 GB z 24 GB karty, co zostawia mniej więcej 7 GiB na cache i bufory obliczeń. W FP16 to około 110K tokenów przed buforami; w Q8_0 około 210K. Przy 4 bitach 262K mieści się z zapasem.

Równoległość mnoży wszystkie te liczby. Czterech agentów po 64K potrzebuje tyle samo cache, co jeden przy 256K.

Dlaczego cache trudniej skwantyzować niż wagi

Wagi kwantyzuje się raz, offline. Można je skalibrować na przykładowych danych, sprawdzić perplexity i wypuścić plik. Cache jest zapisywany w trakcie działania, token po tokenie, przez jądro, które musi być tanie. Niczego nie da się skalibrować pod konkretne wejście.

Same wartości też są mniej przyjazne. Klucze mają kanały odstające: kilka wymiarów ma dużo większe wartości niż reszta. Dlatego KIVI kwantyzuje klucze per kanał, a wartości per token. Przy jednej skali na blok skala rozciąga się pod wartość odstającą i każda inna wartość w bloku traci rozdzielczość.

Obie połowy psują się też inaczej. Błąd w kluczu przesuwa wyniki atencji, a softmax go wzmacnia. Błąd w wartości przesuwa wynik liniowo.

Moje własne testy na Qwen3.5-4B, wszystkie na CPU z tym samym promptem:

• 4 bity: 4-bitowe klucze i wartości dały czysty tekst.
• 3-bitowe klucze i wartości: śmieci, a winę wskazały wartości.
• Przyczyna: częściowo był to błąd pakowania w kodeku 3-bitowym. Dla bitów od pozycji 8 w górę zapisywał ujemne przesunięcie, czyli niezdefiniowane zachowanie, i psuł wartości 3–7 w każdej 8-elementowej grupie.
• Po poprawce: 3-bitowe wartości dają czysty tekst, ale są z natury bardziej stratne. Inna implementacja TurboQuant podaje około +5,8% perplexity przy 3,5 bita.

Dlatego w moim forku domyślne są 4 bity. 3-bitowe wartości są odrzucane, chyba że ustawisz LLAMA_ALLOW_TBQ3_KV=1.

Jak 4-bitowemu cache to uchodzi: najpierw obrót

TBQ4 przepuszcza każdy 128-elementowy blok przez randomizowaną transformację Walsha–Hadamarda przed kwantyzacją: zmiana znaków, transformacja, znowu zmiana znaków. Transformacja jest ortonormalna, więc nie traci informacji. Rozkłada energię wartości odstającej na wszystkie 128 współrzędnych, a wynik wygląda prawie jak rozkład normalny. Rozkład normalny łatwo dobrze skwantyzować: 16 centroidów Lloyda–Maxa dopasowanych do N(0, 1/√128) daje 4-bitowy indeks na wartość plus jedną normę FP16 na blok.

Ortonormalność opłaca się drugi raz w atencji. Iloczyny skalarne przeżywają obrót: q · k = (Hq) · (Hk). Dlatego połączone jądro z upstreamu nigdy nie obraca cache z powrotem:

• obraca zapytanie raz;
• czyta klucze i wartości wprost jako centroid × norma z surowych 4-bitowych bloków, dwa odczyty z tablicy na bajt;
• na końcu obraca wynik atencji z powrotem, też raz.

Nie ma osobnego przebiegu dekwantyzacji ani kopii cache w FP16. Na Metal mój port wciąż dekwantyzuje do pierwotnej dziedziny przed atencją. Działa z prędkością około 90% Q4_0, a różnica to właśnie ta odwrotna transformacja.

Zasady, których się trzymam

1. Policz cache, zanim dobierzesz GPU. Bajty na token × kontekst × liczba równoległych sekwencji. Wagi to łatwa część budżetu.
2. Q8_0 to bezpieczny domyślny wybór, jeśli się mieści. Na obrócone 4 bity przechodź wtedy, gdy ograniczeniem jest kontekst. 3 bity traktuj jak eksperyment i przy nich ponownie uruchom swoje ewaluacje.
3. Sprawdzaj jakość na własnych zadaniach, nie na perplexity z Wikipedii. Cache, który wystarcza do czatu, potrafi zepsuć JSON wywołań narzędzi.
4. Pamiętaj, że cache kosztuje przepustowość, nie tylko pamięć. Każdy krok dekodowania czyta cały cache w warstwach atencji. Przy 262K 4-bitowy cache to też optymalizacja szybkości. Więcej w notatce o dekodowaniu.

Źródła: formaty bloków i wyniki jakości są w docs/kv-types.md w forku. Połączone jądro w obróconej dziedzinie to praca upstreamu z Indras-Mirror/llama.cpp-turboq-mtp. Kwantyzacja kluczy per kanał: Liu i in., KIVI (2024). Obrót przed kwantyzacją: Ashkboos i in., QuaRot (2024).