Continuous batching i PagedAttention: serwowanie dla wielu użytkowników z jednego GPU

Jeden użytkownik na mojej maszynie dostaje najwyżej około 59 przebiegów na sekundę, bo każdy przebieg czyta 17 GB wag (matematyka jest w notatce o dekodowaniu). Obsłużenie dziesięciu użytkowników nie wymaga dziesięciokrotnej przepustowości pamięci. Na tym opiera się cały batching (partie) i stąd też biorą się jego granice.

Dlaczego batching jest prawie darmowy

Krok dekodowania czyta wagi raz, niezależnie od tego, czy liczy jedną sekwencję, czy szesnaście. Włóż B sekwencji do tego samego przebiegu, a wszystkie podzielą się tym odczytem. Łączna przepustowość rośnie z B, dopóki obliczenia nie dogonią transferu z pamięci.

Przy zgrubnych założeniach (54 GFLOP na token, do wykorzystania połowa przepustowości rdzeni tensorowych karty, tylko wagi):

Partia	Czas przebiegu	Tokeny/s łącznie
1	16,9 ms	59
8	16,9 ms	474
16	16,9 ms	949
26	17,0 ms	1528
32	21,0 ms	1528
64	41,9 ms	1528

Do około 26 sekwencji przebieg nie zwalnia: każdy użytkownik zachowuje szybkość pojedynczego strumienia, a serwer wykonuje 26× więcej pracy. Powyżej tego punktu przebieg staje się ograniczony obliczeniami, łączna przepustowość się wypłaszcza, a szybkość każdego użytkownika zaczyna spadać.

Tabela pomija jedną istotną rzecz: atencja każdej sekwencji czyta jej własny cache KV. Długie konteksty dokładają bajtów na sekwencję, więc prawdziwy punkt przejścia wypada wcześniej.

Partię ogranicza pamięć, nie obliczenia

Każda sekwencja potrzebuje własnego cache KV. Po 17 GB wag na karcie 24 GB zostaje mniej więcej 7 GiB na cache. To twardy limit tego, ile tokenów może być naraz w obsłudze (rozmiary na token są w notatce o cache KV):

Format cache	Ile tokenów się mieści	Po 8K na sekwencję	Po 32K na sekwencję
FP16	~115K	14 sekwencji	3
Q8_0	~216K	26 sekwencji	6
TBQ4, 4 bity	~445K	54 sekwencje	13

Skwantyzowany cache nie służy więc tylko temu, żeby jeden użytkownik miał dłuższy kontekst. Często to różnica między 14 a 54 równoległymi użytkownikami.

Statyczny batching marnuje GPU

Naiwny serwer zbiera partię, przetwarza ją do końca i zaczyna następną. Sekwencje kończą się na różnych długościach, więc krótkie stoją bezczynnie jako dopełnienie, aż skończy się najdłuższa, a nowe zapytania czekają na całą partię.

Continuous batching, czyli ciągłe partie (Orca, Yu i in., 2022), szereguje za to na poziomie pojedynczego kroku dekodowania. Ukończone sekwencje wychodzą po dowolnym kroku, a czekające zapytania dołączają w następnym. Partia pozostaje pełna i nikt nie czeka na długą odpowiedź kogoś obcego. Dziś robi tak każdy poważny serwer. Serwer llama.cpp uruchamia równoległe sloty (-np) i ma domyślnie włączony continuous batching.

Fragmentacja: PagedAttention

Drugie marnotrawstwo dotyczy pamięci. Serwer, który dla każdej sekwencji rezerwuje ciągły obszar cache na jej maksymalną długość, marnuje wszystko, czego nie użyje. Zarezerwuj 32K dla zapytania, które kończy się na 2K, a 94% tego obszaru przeleży odłogiem.

PagedAttention (Kwon i in., 2023, rdzeń vLLM) zapożycza z systemów operacyjnych pamięć wirtualną:

• cache jest przechowywany w blokach o stałym rozmiarze, po 16 tokenów;
• każda sekwencja ma tablicę bloków, która mapuje jej pozycje logiczne na bloki fizyczne;
• sekwencja marnuje tylko niezapełnioną część swojego ostatniego bloku.

Artykuł podaje, że istniejące systemy wykorzystują na faktyczne tokeny tylko 20–38% swojej pamięci KV, wobec marnotrawstwa poniżej 4% przy stronicowaniu. Bloki sprawiają też, że współdzielenie jest tanie. Równoległe próbki dla tego samego promptu albo zapytania ze wspólnym prefiksem mogą wskazywać na te same bloki fizyczne, dopóki się nie rozejdą, co jest podstawą prefix cachingu.

Prefill wchodzi w drogę

Prompt nowego zapytania musi przejść prefill, czyli duży przebieg ograniczony obliczeniami. Puść prefill 100K tokenów za jednym razem, a każdy użytkownik w trakcie dekodowania stanie, dopóki prefill się nie skończy. Chunked prefill dzieli prompt na kawałki i przeplata je z krokami dekodowania. Nowy użytkownik dostaje pierwszy token trochę później, a u wszystkich pozostałych strumień płynie dalej.

Zasady, których się trzymam

1. Zdecyduj, co optymalizujesz. Za punktem przejścia przepustowość i opóźnienie dla pojedynczego użytkownika ciągną w przeciwne strony, więc rozmiar partii dobieraj pod docelowe opóźnienie, a nie odwrotnie.
2. Skwantyzuj cache, zanim dołożysz GPU. Na jednej karcie to najtańszy sposób na zwiększenie równoległości.
3. Włącz chunked prefill, jeśli serwujesz długie prompty więcej niż jednemu użytkownikowi.
4. Mierz przy realistycznych długościach kontekstu. Test obciążeniowy z promptami po 500 tokenów mówi niewiele o sesji agenta na 50K tokenów.

Źródła: Yu i in., Orca (OSDI 2022). Kwon i in., Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023). Tabela przepustowości to model, a nie pomiar: 54 GFLOP na token, 50% z ~165 TFLOPS, 1008 GB/s, transfer samych wag.