Zdezagregowane wypełnianie i dekodowanie

Aby wygenerować wyjściowe tokeny z przetwarzania wejściowego, wnioskowanie LLM jest podzielone na dwa etapy: uzupełnianie i dekodowanie. Uzupełnianie działa na tokenach wejściowych, wypełniając pamięci podręczne KV, zanim przejdzie do etapu dekodowania, który generuje tokeny jeden po drugim.

Podczas gdy pojedynczy krok dekodowania zazwyczaj trwa dziesiątki milisekund, uzupełnianie zajmuje znacznie więcej czasu. Jeśli działają na tych samych urządzeniach, mieszanie uzupełniania z dekodowaniem pogarsza wydajność dekodowania. W tym artykule badamy ustalone rozwiązanie w postaci rozdzielonego uzupełniania i dekodowania, uruchamiając je na oddzielnych urządzeniach, aby zmaksymalizować zarówno przepustowość uzupełniania, jak i opóźnienia dekodowania.

Wydajność uzupełniania vs dekodowania

W typowym silniku obsługi LLM, harmonogram partii wybiera żądania do przetworzenia w każdym kroku wykonawczym modelu. Gdy działa na jednym urządzeniu lub węźle, żądania zarówno do uzupełniania, jak i dekodowania są grupowane razem w partie. Koszt uwagi, który sumuje się wzdłuż długości sekwencji, rośnie zarówno dla uzupełniania, jak i dekodowania, proporcjonalnie do długości wpisów w pamięci podręcznej KV (kv_len). Żądania dekodowania zwykle przekazują pojedynczy token (qo_len=1), przy minimalnym koszcie przez inne warstwy, które działają niezależnie na tokenach sekwencji. Żądania uzupełniania przekazują tysiące lub dziesiątki tysięcy tokenów przy znacznym koszcie przez warstwy gęste (duże qo_len).

Opóźnienie przepustowości jest silniej wpływane przez liczbę niezależnych tokenów przelatujących przez warstwy gęste (qo_len) niż przez liczbę tokenów pobieranych z pamięci podręcznej KV podczas uwagi (kv_len). Uwagi mogą być równoległe w szerokości żądań i kv_len proporcjonalnie do długości sekwencji, osiągając dobrą wydajność. Uzupełnianie jest ograniczone obliczeniowo: qo_len jest wysoka, więc jądra GEMM mogą alokować wystarczającą liczbę bloków wzdłuż wymiaru M do pełnego wykorzystania możliwości obliczeniowych nowoczesnych GPU. Dekodowanie jest zależne od pamięci: z powodu zazwyczaj małych rozmiarów paczek, liczba wejść wzdłuż M jest zazwyczaj mała, wystarczająca tylko na jeden blok. Podczas gdy jądra GEMM typu Split-K mogą poprawić wykorzystanie SM dla niewielkich rozmiarów paczek tokenów, pamięci podręczne i jednostki mnożenia macierzy pozostają zazwyczaj niewykorzystane.

Gdy są mieszane, partie zawierające żądania do uzupełniania powodują wyższe opóźnienia przez przebieg w przód, negatywnie wpływając na przepustowość dekodowania całego instancji. Chociaż mieszanie żądań uzupełniania z żądaniami dekodowania lub stosowanie uzupełniania w kawałkach może nieznacznie poprawić wydajność dekodowania, trudno jest utrzymać wystarczającą przepustowość uzupełniania do przetworzenia wystarczającej liczby żądań na instancji, aby zmaksymalizować przepustowość dekodowania. W przypadku dużych modeli, przy typowych długościach wyjścia, aby utrzymać duży rozmiar partii do dekodowania, uzupełnianie musi być wykonywane wystarczająco często, aby znacząco pogorszyć średnie opóźnienie i powodować zacinanie w wyjściu.

Te problemy można rozwiązać, używając oddzielnego zestawu węzłów do wykonywania uzupełniania i dekodowania. Poprzez związanie węzła uzupełniającego z wieloma węzłami dekodującymi, można odpowiednio zaplanować liczbę żądań do uzupełniania, aby zmaksymalizować przepustowość i utrzymać wystarczającą liczbę jednoczesnych żądań na węzłach dekodujących, aby również zmaksymalizować przepustowość dekodowania. Węzły uzupełniające wypełniają pamięci podręczne KV, które są następnie przekazywane do węzłów dekodujących. Ponieważ dekodery nie muszą już przerywać na uzupełnianie, opóźnienia stają się o wiele bardziej deterministyczne, ponieważ ogólny wpływ rosnącego kv_len aktywnych żądań jest znacznie mniej wyraźny. Koszt ten jest płacony we wzroście Time to First Token (TTFT), ponieważ transfer pamięci podręcznych KV przez sieć może zająć dziesiątki do setek milisekund.

KV Messenger

W Perplexity, nasza implementacja rozdzielonego uzupełniania i dekodowania jest zbudowana wokół komunikatora KV, który współdziała z silnikiem LLM, aby orkiestrować transfery pamięci podręcznych KV z węzłów uzupełniających do węzłów dekodujących za pośrednictwem sieci. Po stronie uzupełniającej komunikator przyjmuje żądania z węzłów dekodujących, przekazując je harmonogramowi partii i śledząc wykonanie przebiegów w przód w celu przesyłania pamięci podręcznych KV z jak najmniejszym opóźnieniem. Po stronie dekodującej, po przydzieleniu niewymazywalnych stron, komunikator blokuje żądanie przed planowaniem dekodowania, dopóki nie zostanie powiadomiony o zakończeniu transferów pamięci podręcznej KV i kontekstu dekodera.

Rozdzielanie uzupełniania wymaga połączeń o wysokiej przepustowości i niskim opóźnieniu, dlatego nasza implementacja jest dostosowana do RDMA, wspierając zarówno kontrolery interfejsu sieciowego EFA, jak i ConnectX (NIC). Komunikator KV zbudowany jest na libfabric, używając naszych nakładek fabric-lib do zapewnienia wyższych abstrakcji niskiego opóźnienia w odniesieniu do prymitywów zdalnego bezpośredniego dostępu do pamięci (RDMA), implementując wydajne transfery stron i metadanych, wraz z sygnalizacją o niskim opóźnieniu. W tle fabric-lib koordynuje GPU i jego bezpośrednio połączone NIC do kopiowania danych z węzła uzupełniającego do węzła dekodującego.

Po otrzymaniu, węzeł uzupełniający przydziela odpowiedni zestaw źródłowych stron KV i planuje żądanie do uzupełniania za pomocą swojego lokalnego silnika. Aby zminimalizować opóźnienie, transfery nie czekają na przebieg w przód: zamiast tego, kopie stron KV są inicjowane, jak tylko model zakończy dodawanie wpisów pamięci podręcznych KV do pamięci podręcznej KV dla poszczególnych warstw. Ponieważ żądania uzupełniające mogą być dzielone na kawałki, harmonogram partii powiadamia komunikator KV o obecnie zaplanowanych kawałkach przed wykonaniem. Aby wspierać grafy CUDA przy jednoczesnej możliwości śledzenia warstw, komunikator utrzymuje dedykowany wątek, śledzący licznik zwiększany po projekcji wyjściowej uwagi. Licznik jest utrzymywany tylko na głównym węźle w środowisku dzielonym: mimo że wpisy pamięci podręcznej KV są ważne po dodaniu i przed uwagą, a projekcja wyjściowa jest eliminowana w poprzek pozycji, zsynchronizowanie ich jest domyślne. Po zaobserwowaniu zmiany w liczniku, komunikator jest powiadamiany i wywołuje fabric-lib w celu rozpoczęcia transferu warstwy.

Po zakończeniu transferu ostatniego kawałka, wszystkie dodatkowe metadane również są kopiowane: spekulacyjne dekodowanie lub MTP wymagają przeniesienia logitów i stanów ukrytych do dekodera. Te kopie są również wykonywane za pośrednictwem RDMA, do i z wcześniej przydzielonych buforów.

Po zakończeniu wszystkich oczekujących transferów ostatniego kawałka, węzeł uzupełniający dezallocuje strony KV i kończy żądanie. Węzeł dekodujący nie jest wyraźnie powiadamiany: zamiast tego, używa bezpośrednich liczników do śledzenia liczby zakończonych operacji. Liczba operacji RDMA po stronie uzupełniającej jest proporcjonalna do liczby przesyłanych stron. Po zakończeniu znanej liczby kopii strony i kontekstu, fabric-lib wywołuje komunikator KV, aby wskazać, że żądanie jest gotowe do dekodowania. Komunikator dezallocuje wszelki kontekst i przekazuje żądanie do silnika LLM.

Przenoszenie dzielonych pamięci podręcznych KV

Jeśli uzupełniający i dekodujący polegają na Paralelizmie Tensora (TP) i identycznie dzielą lub replikują pamięci podręczne KV, pojedynczy silnik transferowy koordynuje wiele urządzeń w celu wysyłania i odbierania stron wszystkich replik. Aby móc korzystać z jednego komunikatora i silnika transferowego mimo faktu, że wykonawca modelu jest replikowany na wielu urządzeniach i procesach, używane są cuMem i cuMemImportFromShareableHandle, aby przydzielić pamięć urządzenia wspierającą pamięci podręczne KV i mapować ją do głównego procesu. Silnik transferowy inspekcjonuje topologię węzła, aby znaleźć NIC i CPU w najbliższym węźle NUMA do wykorzystania podczas przenoszenia każdej z części pamięci podręcznych KV.

Jeśli źródło i miejsce przeznaczenia dzielą jegó identycznie, transfery są proste, ponieważ istnieje odwzorowanie jeden do jednego z urządzeń i stron źródła i miejsca przeznaczenia. W tej sytuacji sharding automatycznie ułatwia transfery: używając więcej GPU, można użyć więcej powiązanych NIC, osiągając bliżej maksymalne wykorzystanie przepustowości. Jednak gdy istnieje rozbieżność, silnik transferowy musi dzielić lub rekonstruować strony w zależności od stosunku między źródłem a miejscem przeznaczenia.

Jeśli uzupełniający dzieli pamięć podręczną KV na więcej urządzeń, pełne strony są rekonstruowane na dekodującym poprzez wysłanie odpowiednich połówek z urządzeń uzupełniających. Jeśli dekodujący ma więcej cząstek, odbiera strony z wielu źródeł. Dekodujący musi znać schemat podziału uzupełniającego, aby móc obliczyć liczbę zapisów RDMA, które spodziewa się otrzymać. Jeśli wchodzi w grę replikacja, uzupełniający grupuje urządzenia w zestawy replik, które replikują pełną pamięć podręczną KV wewnątrz siebie. Zestawy replik docelowych są losowo przydzielane jednemu z zestawów źródłowych, aby wykorzystać wszystkie dostępne urządzenia do inicjowania zapisów RDMA.

Dzielone transfery wymagają niewielkiej zmiany w pamięciach podręcznych KV. Domyślnie FlashInfer polega na układzie NHD, który porządkuje tokeny w ramach strony w obrębie nagłówków. Ponieważ pamięci podręczne są najczęściej dzielone wzdłuż liczby nagłówków uwagi, tworzy to nieciągłość w ramach nagłówka. Tranfers RDMA nie wspierają bezpośrednio zapisów ze stridem, wymagając jednej operacji na nagłówek, aby wykonać transfer. Zamiast tego, aby zredukować liczbę interakcji z libfabric, organizujemy pamięci podręczne KV używając układu HND, który umieszcza wymiar nagłówka przed liczbą tokenów. Zapewnia to ciągłość, pozwalając na skopiowanie strony jednym zapisem.

Spekulacyjne dekodowanie

Spekulacyjne dekodowanie wymaga niewielkich zmian w rozdzielonym uzupełnianiu-dekodowaniu. W naszej implementacji, węzły uzupełniające nie mogą próbować tokenów. Ponieważ modele Sonar firmy Perplexity wspierają strukturalne wyjście, nie chcemy popadać w złożoność synchronizowania implementacji procesorów schematów między uzupełniaczami a dekodującymi. W mechanizmach MTP i spekulacyjnego dekodowania, wypełnianie modelu roboczego do ostatniego tokena obejmuje próbki tokenów z modelu docelowego.

Aby obejść te problemy, uzupełnianie nie obejmuje ostatniego tokena sekwencji wejściowej. Zamiast tego, stany ukryte lub logity z uzupełniania poprzedzającego ostatni token są przesyłane i są traktowane jako token dekodowany w następnym kroku na dekodującym. Chociaż nieco zwiększa to opóźnienia, ponieważ pełny krok dekodowania musi zostać wykonany po uzupełnianiu, aby wygenerować pierwszy token, złożoność implementacji jest znacznie zredukowana.

Rozdzielone wdrożenia

Wdrożyliśmy lub eksperymentowaliśmy z wieloma rozdzielonymi konfiguracjami z różnymi modelami, aby wspierać ruch produkcyjny lub prace ewaluacyjne wewnętrzne. Na podstawie rozmiaru i mechanizmu uwagi modeli, wybieraliśmy odpowiednie schematy podziału węzłów uzupełniających i dekodujących, aby najlepiej wykorzystać GPU.

DeepSeek-R1

Z DeepSeek rozważaliśmy zarówno wdrożenia równoległe tensora (TP), jak i równoległe danych (DP). Jak omawiano w poprzednich postach na blogu, wdrożenia TP zapewniają lepsze opóźnienie przy większych kosztach niższej przepustowości, wymagając więcej GPU do obsługi ciężkiego ruchu. Wdrożenia DP lepiej skalują się z obciążeniem, jednak ich maksymalna przepustowość jest niższa z powodu kosztów komunikacji między urządzeniami lub węzłami.

DeepSeek polega na wielogłowowej latencji uwagi, kompresując pamięci podręczne KV. Ponieważ wszystkie głowice KV są kompresowane w jeden wektor latentny, TP nie może dzielić pamięci podręcznych KV, ponieważ musi zamiast tego replikować wektory latentne na wszystkich szczeblach. Podział następuje po dekompresji, ponieważ każda ranga może wyodrębnić inne głowice z tego samego reprezentacji latentnej. W związku z tym wszystkie cząstki pamięci podręcznych KV są identyczne w przypadku uzuepłaniającego i dekodującego.

Przy wewnętrznym ustawieniu TP zarówno uzupełniacze, jak i dekodery są dzielone identycznie. Przesyłania są uruchamiane ze wszystkich szczebli w celu pełnego wykorzystania wszystkich dostępnych NIC. Jednak w przypadku wdrożenia DP, gdzie rozmiar rangi TP jest mniejszy lub każda ranga DP przypisana jest do jednego GPU, każde urządzenie uzupełniające, które ma replikowaną kopię pamięci podręcznej KV, może ją przesłać. Aby zrównoważyć żądania wśród wszystkich dostępnych NIC, losowo wybieramy GPU i NIC do przesłania pamięci podręcznej KV z uzupełniacza do dekodera.

W przypadku mieszanego uzupełniania-dekodowania, nasze wdrożenie R1 borykało się z problemem konsekwentnego przekraczania 50 TPS z powodu częstych przerw uzupełniania w rzędzie setek milisekund. W porównaniu z tym, przez rozdzielenie uzupełniania ponieśliśmy karę około 100 ms do TTFT dla każdej prośby, ale pojedynczy węzeł uzupełniający mógł utrzymać stałe rozmiary paczek na 3 węzły dekodujące, dostarczając przepustowość powyżej 90 TPS przy jednoczesnym obciążeniu około 1 QPS na węzeł dekodera. W przypadku wdrożeń równoległych danych TPS była nieco niższa na poziomie około 50, jednak instancje te mogły obsługiwać obciążenie 1 QPS na rangę, z 8 rangami przypisanymi do jednego węzła.

Qwen3-Coder

Ten model 480B używa grupowanej uwagi zapytań (GQA), więc uwagi można łatwo dzielić i mogą korzystać z równoległości tensora bez poświęcania pamięci na pamięć podręczną KV. W związku z tym mogliśmy podzielić model na 8 GPU zarówno do uzupełniania, jak i dekodowania, parując około 3 węzły dekodujące z jednym węzłem uzupełniającym. Ponieważ uwaga jest dzielona, polegamy na układzie pamięci podręcznych HND do podziału pamięci uzupełniających i dekodujących, parując rangi uzupełniające z rangą dekodującą i w pełni wykorzystując wszystkie NIC, aby przesyłać części równolegle.