Przyspieszanie Sonar poprzez spekulację

Dekodowanie spekulatywne przyspiesza generowanie przez duże modele językowe (LLM) przy użyciu szybkiego i niewielkiego modelu roboczego, który tworzy kandydatów do uzupełnienia, weryfikowanych następnie przez większy model docelowy. W tym schemacie zamiast kosztownego uruchamiania modelu docelowego w celu wygenerowania pojedynczego tokenu, w jednym kroku emitowanych jest ich kilka. Niniejszym przedstawiamy szczegóły implementacji różnych rodzajów dekodowania spekulatywnego, zastosowanych w Perplexity w celu zmniejszenia opóźnień międzytokenowych w modelach Sonar.

Dekodowanie spekulatywne

Dekodowanie spekulatywne wykorzystuje strukturę języków naturalnych oraz autoregresyjny charakter transformerów w celu przyspieszenia generowania tokenów. Choć większe modele, takie jak Llama-70B, dysponują szerszą wiedzą niż mniejsze, np. Llama-1B, w niektórych prostszych zadaniach działają podobnie. Ta zbieżność sugeruje, że określone sekwencje są lepiej generowane przez tańsze modele, co pozwala pozostawić złożone problemy modelom większym. Wyzwanie polega na określeniu, które uzupełnienia są lepsze oraz czy jakość generowania mniejszego modelu jest równa jakości większego.

Na szczęście modele LLM to autoregresyjne transformery: po otrzymaniu sekwencji tokenów wyprowadzają rozkład prawdopodobieństwa następnego tokenu. Ponadto logity wywiedzione z pośrednich cech powiązanych z tokenami w sekwencji wejściowej również wskazują, jak prawdopodobne jest wyemitowanie przez model dokładnie tych tokenów. Ta właściwość umożliwia spekulację: jeśli sekwencja tokenów jest generowana przez mniejszy model począwszy od prefiksu wejściowego, można ją przepuścić przez większy model, aby określić, jak dobrze pokrywa się ona z modelem docelowym. Każdy prefiks kandydatów jest oceniany prawdopodobieństwem, a wybierany jest najdłuższy z nich, który przekracza próg akceptacji. Jako dodatkowa korzyść, model docelowy dostarcza również kolejny token bezpłatnie: jeśli model roboczy generuje n tokenów, w jednym kroku można wyemitować do n + 1.

Powiększ

W czasie inferencji proces próbkowania spekulatywnego można podzielić z grubsza na 4 etapy:

Prefill: zarówno model docelowy, jak i roboczy muszą zostać uruchomione na sekwencji wejściowej, aby wypełnić wpisy w pamięci podręcznej KV. Choć niektóre schematy, takie jak Medusa, wykorzystują prostsze warstwy gęste do przewidywania, w tym wpisie koncentrujemy się na modelach roboczych opartych na transformerach, które wymagają własnych pamięci podręcznych KV.

Generowanie robocze: model roboczy iteruje w celu wytworzenia określonej liczby stałych tokenów. Sekwencja robocza może być liniowa lub model może przeszukiwać strukturę drzewiastą do określonej głębokości (EAGLE, Medusa). Tutaj skupiamy się na sekwencjach liniowych.

Akceptacja: model docelowy uruchamia się na sekwencji roboczej, tworząc logity odpowiadające każdemu tokenowi roboczemu. Określana jest długość najdłuższej akceptowalnej sekwencji.

Generowanie docelowe: ponieważ model docelowy wygenerował logity, w pozycji niedopasowania lub na końcu sekwencji logity odpowiadają kolejnemu tokenowi. Logity te mogą być próbkowane w celu zapewnienia solidnego tokenu z modelu docelowego, kończącego sekwencję.

Istnieją różne metody implementacji dekodowania spekulatywnego. W tym wpisie skupimy się na schematach, których użyliśmy do przyspieszenia modeli Sonar przy użyciu wewnętrznego modelu 1B, a także mechanizmach przewidywania, które rozwijamy w celu przyspieszenia modeli w skali DeepSeek.

Docelowo-roboczy

Dekodowanie spekulatywne można osiągnąć poprzez połączenie istniejącego małego modelu LLM jako modelu roboczego z modelem docelowym w celu generowania sekwencji kandydatów. W produkcji przyspieszyliśmy Sonar przy użyciu modelu Llama-1B dostrojonego na tym samym zestawie danych co model docelowy. Choć to podejście nie wymagało trenowania modelu roboczego od zera, mały model nadal wykorzystuje znaczną pojemność pamięci podręcznej KV i wprowadza niewielki narzut fazy prefill, zwiększając TTFT.

W tym schemacie dekoder spekuluje tylko w partiach decode-only, generując tokeny poprzez standardowe próbkowanie podczas fazy prefill lub w mieszanych partiach prefill-decode. W fazie prefill logity modelu docelowego są natychmiast próbkowane, aby również wypełnić nowo wygenerowany token w pamięci podręcznej KV modelu roboczego. Model roboczy nie jest jeszcze próbkowany, ale produkowane przez niego logity są przenoszone do etapu dekodowania.

Powiększ

W dekodowaniu model roboczy jest rozwijany, próbkując najlepszy token na każdym etapie. Po osiągnięciu pożądanej długości roboczej tokeny są przepuszczane przez model docelowy w celu uzyskania logitów, na podstawie których sampler identyfikuje zaakceptowaną długość sekwencji. Akceptacja jest określana poprzez porównanie pełnych rozkładów prawdopodobieństwa z modelu roboczego i docelowego. Ponieważ model docelowy zawsze wyprowadza jeden zestaw logitów po zaakceptowanej sekwencji roboczej, jest on próbkowany w celu wyprodukowania dodatkowego wyjścia. Ponieważ model roboczy nie widział jeszcze tego zaakceptowanego tokenu, jest on uruchamiany ponownie w celu wypełnienia odpowiednich wpisów w pamięci podręcznej KV w przygotowaniu do następnego kroku dekodowania, ponownie przenosząc logity.

EAGLE

EAGLE to schemat dekodowania spekulatywnego, który bada wiele sekwencji roboczych, generowanych poprzez drzewiaste przeszukiwanie prawdopodobnych tokenów roboczych. Stałe (EAGLE) lub dynamicznie kształtowane (EAGLE-2) drzewo jest przeszukiwane przy użyciu kolejnych wykonań tokenów roboczych, biorąc pod uwagę kandydatów Top-K w każdym węźle zamiast podążania za tokenem o najwyższym wyniku w sekwencji liniowej. Następnie sekwencje są oceniane i wybierana jest najdłuższa odpowiednia, aby kontynuować, dołączając również dodatkowy token z modelu docelowego.

Powiększ

Aby osiągnąć dokładniejsze przewidywanie, model roboczy EAGLE przewiduje nie tylko na podstawie tokenów, ale także przy użyciu cech docelowych (stanów ukrytych ostatniej warstwy) modelu docelowego. Wadą EAGLE jest konieczność trenowania niestandardowych, małych modeli roboczych, które są wystarczająco dokładne, aby generować odpowiednich kandydatów w ramach niskiego budżetu opóźnień. Zazwyczaj model roboczy to pojedyncza warstwa transformera identyczna z warstwą dekodera modelu oryginalnego, która jest ściśle powiązana z modelem docelowym poprzez połączenie z jego embeddingami i projekcjami lm_head. Ponieważ wymaga to mniejszej pojemności pamięci podręcznej KV, EAGLE ma mniejszy ślad pamięciowy.

Aby zweryfikować sekwencje drzewiaste w modelu docelowym, należy użyć niestandardowych masek uwagi. Niestety, użycie niestandardowej maski uwagi dla całej sekwencji znacząco spowalnia uwagę dla realistycznych długości wejściowych (nawet o 50%), niwelując część przyspieszenia możliwego do osiągnięcia dzięki spekulacji. Z tego powodu nie wdrożyliśmy jeszcze pełnego przeszukiwania drzewa w produkcji, skupiając się zamiast tego na szczególnym przypadku przewidywania pojedynczego tokenu za pomocą schematów typu MTP przedstawionych w raporcie technicznym DeepSeek-V3.

MTP

Schemat ten jest podobny do dekodowania draft-target, z wyjątkiem użycia stanów ukrytych obok tokenów do przewidywania. Należy wykonać nieco więcej pracy zarówno w fazie prefill, jak i dekodowania w porównaniu ze zwykłą spekulacją draft-target. Model roboczy używa zarówno tokenów, jak i stanów ukrytych: token t_{i+1} jest próbkowany z logitów L_i odpowiadających tokenowi t_i, które z kolei pochodzą ze stanów ukrytych H_i. W rezultacie bufory tokenów wejściowych muszą zostać przesunięte o jeden krok w lewo względem wektorów stanów ukrytych wyprowadzonych przez model docelowy. Poniższa rycina oznacza korespondencje użyte do trenowania, jak również przesunięcie podczas inferencji.

Powiększ

Przepływ dekodowania jest dość podobny do dekodowania draft-target, z wyjątkiem przenoszenia zarówno stanów ukrytych, jak i logitów. Nasza implementacja współdzieli wszystkie powiązane logity próbkowania i przetwarzania logitów, specjalizując jedynie wywołania modelu w przód. Kiedy przewidywanych jest wiele tokenów, model roboczy używa roboczych stanów ukrytych do przewidywania, wypełniając również wpisy w pamięci podręcznej KV na podstawie własnych cech. W dłuższej perspektywie może to pogorszyć dokładność. Następnie, podczas uruchamiania modelu roboczego w celu wypełnienia wpisu w pamięci podręcznej KV dla przewidywania docelowego, uruchamiamy go na całej sekwencji, biorąc za dane wejściowe dokładniejsze docelowe stany ukryte. Ponieważ te modele robocze są małe, dodatkowy koszt przetwarzania dodatkowych tokenów jest pomijalny.

Trenowanie głowic MTP

Aby skorzystać z MTP, zbudowaliśmy infrastrukturę wymaganą do trenowania głowic MTP dołączonych do naszych dostrojonych modeli na zbiorach danych Perplexity, działających na jednym węźle z 8 urządzeniami H100. W około jeden dzień możemy zbudować głowice dla modeli od Llama-1B do Llama-70B i DeepSeek V2-Lite. W przypadku większych modeli polegamy na głowicach MTP zbudowanych podczas procesu dostrajania.

Celem trenowania MTP jest dopasowanie roboczych stanów ukrytych i logitów ekstrapolowanych ze stanów ukrytych modelu docelowego do logitów następnego tokenu i stanów ukrytych modelu docelowego. Ponieważ inferencja dla stanów ukrytych jest kosztowna, wstępnie obliczamy je przy użyciu naszej zoptymalizowanej pod kątem inferencji implementacji modelu docelowego, aby użyć ich podczas treningu. Jednak w celu walidacji implementacji MTP dla inferencji i zapewnienia, że różnice liczbowe wynikające z kwantyzacji lub optymalizacji nie utrudniają wyników, dla oceny utraty walidacji i dokładności w pełni ponownie wykorzystujemy implementację inferencji zarówno modelu docelowego, jak i roboczego.

Podczas skalowania ze zbioru danych ShareGPT użytego w oryginalnym artykule do większych próbek, zauważyliśmy, że architektura głowicy MTP nakreślona i zaimplementowana w artykule EAGLE nie powiodła się podczas trenowania dla modeli o rozmiarze 70B. W przeciwieństwie do ShareGPT, który zawierał większą liczbę krótszych sekwencji, trenujemy na nieco mniejszej liczbie znacznie dłuższych monitów. Ponieważ oryginalne głowice EAGLE nieznacznie odbiegały strukturą od typowego transformera, ponownie wprowadziliśmy niektóre warstwy normalizacji RMS, które zostały usunięte. Odkryliśmy, że pozwoliło to nie tylko na zbieżność treningu, ale także zwiększyło dokładność głowic o kilka punktów procentowych.

Powiększ

Normy warstw nie tylko ułatwiają trenowanie, ponowne wprowadzenie norm jest również matematycznie intuicyjne. Głowice MTP ponownie wykorzystują embeddingi i projekcje logitów modelu docelowego, ponieważ mogą one mieć znaczny rozmiar (około 2 GB dla Llama 70B). Podczas trenowania są one zamrożone, a oczekuje się, że warstwa MTP nauczy się osadzać przewidywania w tej samej przestrzeni wektorowej, której nauczyła się warstwa projekcyjna oryginalnego modelu podczas treningu. Porzucając normy, oczekuje się, że pojedynczy MLP nauczy się tej samej funkcji co MLP, po którym następuje norma, co utrudnia dopasowanie między stanami ukrytymi modelu roboczego i docelowego.

Inferencja z dekodowaniem spekulatywnym

W silniku inferencji, aby generować tokeny dla sekwencji wejściowych, muszą one zostać najpierw pogrupowane w partie o rozsądnym rozmiarze, a następnie strony muszą zostać przydzielone w pamięci podręcznej KV dla następnych tokenów. Tokeny wejściowe i informacje o stronie KV są następnie pakowane do bufora rozsyłanego do wszystkich równoległych rang uruchamiających model. Wreszcie, metadane są kopiowane do pamięci GPU i model jest wykonywany w celu wytworzenia logitów, z których próbkowany jest następny token.

W przeciwieństwie do niektórych implementacji, które luźno łączą serwer inferencji roboczej i docelowej za pomocą opakowania orkiestrującego żądania między nimi, nasze pary docelowo-robocze są ściśle sprzężone i przechodzą przez generowanie w sposób zsynchronizowany. Harmonogramowanie partii i przydział stron KV są współdzielone między modelami dla wszystkich form dekodowania spekulatywnego: unifikuje to logikę łączącą model z nadrzędnym serwerem inferencji, ponieważ wszystkie udostępniają ten sam interfejs.

Środowisko uruchomieniowe inferencji w Perplexity jest oparte na FlashInfer, który określa metadane, które muszą zostać zbudowane w celu skonfigurowania i zaplanowania jądra uwagi. Biorąc pod uwagę niektóre sekwencje wejściowe tworzące partię, dla fazy prefill, dekodowania lub weryfikacji, praca po stronie CPU musi zostać wykonana w celu przydzielenia buforów pośrednich i wypełnienia niektórych stałych buforów używanych w uwadze. Praca ta jest dodatkiem do kosztu harmonogramowania partii i przydziału stron KV, które również generują opóźnienia, które muszą zostać ukryte w celu maksymalizacji wykorzystania GPU.

Podczas gdy w pełni zrównolegliliśmy pracę po stronie CPU i GPU dla inferencji bez spekulacji, odkryliśmy, że równowaga CPU-GPU dla dekodowania spekulatywnego jest bardziej złożona. Główne wyzwanie wynika z faktu, że liczba zaakceptowanych tokenów określa długość sekwencji dla kolejnego uruchomienia, wprowadzając trudny do uniknięcia punkt synchronizacji GPU-CPU. Eksperymentowaliśmy z różnymi schematami harmonogramowania, aby najlepiej ukryć opóźnienie pracy CPU.

Harmonogram docelowo-roboczy

Mimo że jest mniejszy niż model docelowy, gdy cały LLM jest używany jako roboczy, nadal wprowadza znaczne opóźnienie na GPU, zapewniając pewną przestrzeń do ukrycia kosztownych operacji CPU. Ponieważ mniejsze modele nie korzystają z paralelizmu tensorowego, istnieje niedopasowanie między liczbą rang, na które podzielony jest model docelowy i roboczy. W naszej implementacji model roboczy działa tylko na wiodącej randze grupy TP.

Powiększ

Jak wskazano wcześniej, krok dekodowania przenosi logity do następnego uruchomienia. Pozwala nam to na nałożenie jednego wykonania modelu roboczego na pracę harmonogramowania partii po stronie CPU. Po zestawieniu partii, powtarzane wywołania samplera i modelu roboczego produkują tokeny robocze. Równolegle partia do weryfikacji jest zestawiana dla modelu docelowego i synchronizowana z pracownikami równoległymi. Logity docelowe są weryfikowane i próbkowane w celu określenia zaakceptowanych długości sekwencji. W tym momencie synchronizacja GPU-CPU jest niezbędna w celu określenia kolejnych długości sekwencji. Ponieważ model roboczy jest uruchamiany tylko na węźle wiodącym, jego partia jest ustawiana sekwencyjnie, a jego wykonanie jest uruchamiane w celu wypełnienia wpisów w pamięci podręcznej KV dodatkowym tokenem wyprodukowanym przez model docelowy. Logity wyprodukowane przez ten model roboczy w bieżącym uruchomieniu zostaną użyte do spróbkowania pierwszego tokenu roboczego w kolejnym uruchomieniu. Co najważniejsze, podczas gdy model roboczy działa, następna partia może być planowana.

Harmonogram MTP dla pojedynczego tokenu

Chociaż środowisko uruchomieniowe nie zapewnia jeszcze przeszukiwania drzewa roboczego w stylu Eagle, zaimplementowaliśmy specjalny przypadek tego schematu, biorąc pod uwagę liniową sekwencję tokenów roboczych wyprodukowanych przez model wielkości pojedynczej warstwy dekodera transformera. Schemat ten może być używany do przewidywania roboczego przy użyciu wag open-source DeepSeek R1. Podprzypadek przewidywania pojedynczego tokenu jest interesujący, ponieważ duże warstwy MTP osiągają wystarczająco wysokie wskaźniki akceptacji, aby uzasadnić swój narzut.

Harmonogramowanie MTP jest nieco bardziej złożone, ponieważ model roboczy jest znacznie szybszy, ukrywając mniej opóźnień po stronie CPU. Dodatkowo, model roboczy jest podzielony na fragmenty wraz z modelem docelowym, co wymaga transferów współdzielonej pamięci dla informacji o partii. Uruchomienie rozpoczyna się od przesłania informacji o partii i spróbkowania pierwszego tokenu z przeniesionych logitów, podobnie jak w poprzednim schemacie. Następnie model docelowy jest uruchamiany w celu weryfikacji tokenów, przetwarzając 2 * D tokenów, gdzie D jest rozmiarem partii dekodowania. Jest to idealne rozwiązanie dla mikrobaczowania w modelach typu Mixture-of-Experts (MoE) przez wolniejsze połączenia, takie jak InfiniBand, ponieważ partia dzieli się równo na dwie połowy. Stany ukryte modelu docelowego przechodzą do następnego uruchomienia modelu roboczego, podczas gdy logity są przekazywane do samplera w celu weryfikacji.

Powiększ

Wykonując ograniczoną ilość dodatkowej pracy na GPU, unikamy synchronizacji CPU-GPU po akceptacji sekwencji roboczej. Po przesunięciu tokenów wejściowych modeli docelowych, jądro wstawia następne tokeny docelowe w ich odpowiednie lokalizacje. Model roboczy jest następnie ponownie uruchamiany z tymi samymi informacjami o partii co model docelowy, wypełniając wpisy pamięci podręcznej KV oraz budując logity i stany ukryte dla następnego uruchomienia, wykonując pewną redundantną pracę na tokenach, które nie zostały zaakceptowane. W takich sytuacjach opóźnienie nieużywanej pracy jest ledwo mierzalne ze względu na mały rozmiar modelu roboczego. Równolegle z uruchomieniem modelu roboczego określane są długości sekwencji na CPU i rozpoczyna się harmonogramowanie następnej partii, bez konieczności czekania na zakończenie pracy GPU.

Narzut dodatkowej pracy w warstwie roboczej nie jest zauważalny w uwadze, jednak warstwy MLP są bardziej problematyczne. Ponieważ instrukcje mnożenia macierzy są wyrównane do granicy 64 wzdłuż wymiaru liczby tokenów, jeśli podwojenie nie wymaga znacząco więcej bloków, narzut jest ukryty. W przypadku dłuższych sekwencji roboczych narzut jest bardziej kosztowny, a schemat używany dla zwykłych modeli docelowo-roboczych działa lepiej.

Bibliografia

Szybka inferencja z transformerów poprzez dekodowanie spekulatywne

EAGLE: Próbkowanie spekulatywne wymaga przemyślenia niepewności cech

EAGLE-2: Szybsza inferencja modeli językowych z dynamicznymi drzewami roboczymi

EAGLE-3: Skalowanie przyspieszenia inferencji dużych modeli językowych poprzez testy w czasie trenowania

Medusa: Prosta struktura przyspieszania inferencji LLM z wieloma głowicami dekodującymi

Normalizacja warstw

FlashInfer: Wydajny i konfigurowalny silnik uwagi do obsługi inferencji LLM

Raport techniczny DeepSeek-V3

Cytat

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}