Speculative decoding na Trainium2 tnie koszt tokenów

AWS zademonstrował, jak speculative decoding na układach Trainium2 w połączeniu z vLLM redukuje koszt inferencji w zadaniach z dominującą fazą dekodowania.

Jak to działa technicznie

Speculative decoding polega na tym, że mały, tani model szkicowy (draft model) generuje kilka tokenów naraz, a duży model weryfikuje je w jednym przebiegu. Jeśli propozycje się zgadzają — akceptuje wszystkie. Jeśli nie — odrzuca od pierwszego błędnego i przejmuje stery. Efekt: ten sam output, mniej wywołań dużego modelu.

Trainium2 to własny chip AWS zaprojektowany do trenowania i inferencji LLM. AWS twierdzi, że jego architektura pamięci i przepustowość HBM robią z niego sensowne środowisko dla speculative decoding — techniki, która jest wrażliwa na latencję przesyłania danych między modelem szkicowym a modelem docelowym.

Dlaczego akurat decode-heavy?

Nie każde zadanie skorzysta tak samo. Speculative decoding najbardziej opłaca się tam, gdzie model generuje długie odpowiedzi — chatboty, asystenci kodu, systemy RAG zwracające rozbudowane teksty. W zadaniach z krótkim outputem zysk jest marginalny, a narzut weryfikacji może nawet podnieść latencję.

AWS celuje właśnie w ten segment: długie odpowiedzi, wysokie obciążenie, wrażliwość na koszt per token. To naturalne środowisko dla enterprise’owych wdrożeń, gdzie firma płaci za każdy wygenerowany token w produkcji.

vLLM jako warstwa orkiestracji

vLLM obsługuje tu kolejkowanie żądań, zarządzanie pamięcią KV cache i integrację z mechanizmem spekulatywnego dekodowania. AWS wrzucił obsługę Trainium2 do stosu vLLM, co oznacza, że istniejące pipeline’y oparte na tym frameworku można odpalić na nowym sprzęcie bez przepisywania kodu.

To istotne z praktycznego punktu widzenia: vLLM ma już ugruntowaną pozycję w produkcyjnych deploymentach LLM, więc adopcja Trainium2 nie wymaga od zespołów inżynierskich nauki nowego ekosystemu od zera.

Czy Trainium2 realnie konkuruje z H100?

AWS nie podaje bezpośredniego porównania Trainium2 vs NVIDIA H100 w tym poście — i to samo w sobie mówi coś o strategii komunikacyjnej. Zamiast benchmarku head-to-head, firma skupia się na metryce koszt/token, gdzie własny sprzęt w chmurze AWS może wypaść korzystniej niż wynajmowane instancje GPU od innych providerów.

Trainium2 debiutował w instancjach trn2.48xlarge z 16 układami i 512 GB pamięci HBM. AWS chwali się przepustowością 20,8 petaflopsów przy BF16 na jedną instancję. Dla porównania: jedna karta H100 SXM5 wykręca 1,979 petaflopsa przy BF16 — instancja trn2 ma więc zbliżony teoretyczny pułap do klastra 10-12 kart H100, przy potencjalnie niższym koszcie na AWS.

Co to oznacza dla zespołów MLOps?

Praktyczny wniosek jest prosty: jeśli już używasz vLLM i rozważasz skalowanie inferencji na AWS, Trainium2 ze speculative decoding staje się opcją wartą przetestowania — szczególnie przy modelach generujących długie teksty.

Konfiguracja wymaga dobrania właściwego modelu szkicowego. Zbyt duży draft model zje oszczędności na weryfikacji; zbyt słaby będzie często odrzucany i nie przyspieszy nic. AWS rekomenduje modele z tej samej rodziny co model docelowy, tylko mniejsze — np. Llama 3.1 8B jako szkicowy dla Llama 3.1 70B.

Jeden nieoczywisty haczyk: speculative decoding zmienia profil wykorzystania pamięci. Trzeba utrzymać dwa modele jednocześnie w VRAM, więc instancja musi mieć odpowiedni zapas pamięci HBM — coś, o czym łatwo zapomnieć przy planowaniu wdrożenia.

AWS nie podał konkretnych liczb procentowych przyspieszenia w publicznym poście — szczegóły benchmarków są dostępne w pełnej dokumentacji technicznej na blogu AWS Machine Learning.