AKV z Tsinghua: agenci LLM sami wycinają fragmenty KV cache zamiast streszczać historię

Długie sesje agentowe mają prosty problem: historia rośnie, a każda edycja jej początku unieważnia obliczone wcześniej stany KV cache. Typowe rozwiązanie to streszczenie starszych wiadomości, po którym trzeba ponownie przeliczyć (prefill) wszystko, co zostało. Zespół z Tsinghua University i CUHK proponuje inne podejście i publikuje kod AKV na licencji Apache-2.0. Opisuje je MarkTechPost, a szczegóły są w pracy na arXiv.

Cache engineering zamiast context engineering

Autorzy nazywają to „cache engineering”. Tekstowa trajektoria agenta zostaje nietknięta. Zmienia się tylko to, które wpisy KV są widoczne dla mechanizmu uwagi. Zachowane wpisy są używane ponownie bez przeliczania i nadal mogą nieść informację z wiadomości, które wyrzucono. W klasycznym podejściu harness streszcza historię późno i agresywnie, przez co mogą zginąć szczegóły ze środkowych rund. Motywacja jest też finansowa: według analizy kosztów z pracy odczyty KV cache odpowiadały za 82,4% całkowitego kosztu API.

AKV definiuje trzy operacje:

• Append – standardowe dopisanie nowych tokenów;
• Drop – usunięcie wybranych wpisów z przyszłej uwagi bez zmiany pozycji;
• Repos – przepisanie zachowanych kluczy na pozycje od 0 do r-1 za pomocą rotacji RoPE. Wartości zostają bez zmian, więc nie trzeba liczyć stanów ukrytych od nowa.

Trudniejsza jest strona serwerowa. Te same tokeny mogą dawać różne stany cache zależnie od historii operacji, więc zwykłe dopasowanie prefiksu mogłoby użyć niewłaściwych stanów. AKV rozszerza drzewo radix w SGLang o węzły drop i reposition i zwalnia usunięte KV dopiero wtedy, gdy nie potrzebuje ich żadne aktywne zapytanie. Z kolei „segment-batch attention” przywraca wypchnięte stany w jednym przebiegu prefill.

API jest bezstanowe i zgodne z OpenAI. Klient wysyła pełną historię oraz dwa opcjonalne pola: drop_message (które wiadomości usunąć po danym wyzwalaczu) i reposition (po których indeksach zagęścić pozycje). Pola budują polityki CachePolicy – oparte na turach, tokenach lub modelu. Narzędzie CacheTool pozwala też samemu modelowi wywołać drop_context.

Wyniki

Testy objęły trzy modele: GPT-OSS-120B (5,1 mld aktywnych parametrów), MiniMax-M2.7 (230 mld łącznie, 10 mld aktywnych) i MiroThinker-1.7-Mini (30 mld łącznie, 3 mld aktywnych). Sprawdzano je na czterech benchmarkach: BrowseComp-Plus, GAIA-Text, SWE-bench Verified i xbench-DeepSearch-2510. Główna polityka to kroczące usuwanie odpowiedzi narzędzi, z zachowaniem 12 ostatnich. Punktem odniesienia było streszczanie po przekroczeniu 96 tys. tokenów (192 tys. dla MiroThinkera).

• Średnio koszt spada o 27,7%, a wynik rośnie o 1,7 punktu względem streszczania progowego.
• W porównaniu z wariantem edycji tekstowej (Text Engineering) koszt jest niższy o 21,9%.
• Najlepszy przypadek: GPT-OSS-120B na SWE-bench Verified – koszt zadania spadł z 20,0 do 9,7 centa, a wynik wzrósł z 40,0 do 43,2.
• Na BrowseComp-Plus przy kontekście od 96 tys. tokenów GPT-OSS-120B osiąga z AKV 37%, przy 24% dla wersji standardowej i 32% dla Text Engineering.
• W ablacji dodanie Repos podniosło wskaźnik zaliczeń GPT-OSS-120B na BrowseComp-Plus z 55% do 62%.
• Przy współbieżności 16 opóźnienie P95 spadło z 229,90 do 30,63 s.

Nie wszędzie AKV wygrywa. Najgorszy wynik to MiniMax-M2.7 na BrowseComp-Plus: 58,9, czyli 3,4 punktu poniżej 62,3 dla Text Engineering.

Wymagania i zastrzeżenia

AKV to system serwowania i biblioteka polityk, bez wag modeli. Działa na GPU NVIDIA przez SGLang. W pracy użyto ośmiu kart A800-80GB w BF16, a szybki start z README serwuje GPT-OSS-120B na dwóch GPU (tensor parallel 2). Na razie obsługiwane są tylko tekstowe modele z rodzin Qwen, GPT-OSS i MiniMax-M2.

Koszty to szacunki „API-equivalent” oparte na opublikowanych cennikach. Tokeny po repozycjonowaniu rozliczano po stawce cache-hit, co autorzy uznają za celowo ostrożne założenie. To nie są więc rachunki z produkcyjnego wdrożenia. Pomysł traktowania KV cache jako warstwy runtime dla agentów nie jest zresztą odosobniony – podobny kierunek opisywał Yandex Research.