Jak Cerebras uruchamia bilionowe modele AI na wafer-scale chipach

Cerebras opublikował szczegółowy techniczny wpis wyjaśniający, jak firma uruchamia modele liczące biliony parametrów na swoich procesorach WSE-3T wielkości całej płytki krzemowej (wafer-scale). Problem jest fundamentalny: pojedynczy wafer WSE-3T dysponuje 44 GB pamięci on-chip, podczas gdy same wagi modeli frontierowych o skali bilionów parametrów zajmują terabajty. Cerebras pokazuje krok po kroku, jak dzieli taki model na dziesiątki wafferów, by uzyskać przepustowość rzędu 1000–2000+ tokenów na sekundę.

Jako przykład posłużył otwartoźródłowy model Kimi K3, liczący około 3 bilionów parametrów, zbudowany z 69 warstw linear-attention i 24 warstw gęstych (dense).

Kluczowe wnioski

• Kimi K3, model open source z ~3 bilionami parametrów, osiąga 88,3% w Terminal-Bench 2.1 (vs 88,8% dla GPT-5.6 Sol) i 91,2% w BrowseComp (vs 90,4% dla GPT-5.6 Sol) – wyniki na poziomie frontierowym w kodowaniu i zadaniach agentowych.
• Przy 4-bitowej kwantyzacji wag każda z 93 warstw modelu zajmuje około 16 GB, co przy 44 GB dostępnej pamięci na wafferze pozwala umieścić całe warstwy bezpośrednio na chipie, bez ich dzielenia.
• Do obsługi pełnego modelu Kimi K3 Cerebras potrzebuje łącznie 59 wafferów: 24 dla warstw gęstych (po jednej na wafer, z miejscem na pamięć podręczną KV) oraz 35 dla warstw linear-attention (po dwie warstwy na wafer).
• Pomiędzy wafferami przesyłane są wyłącznie aktywacje (nie wagi modelu), za pomocą RDMA przez Ethernet, co zajmuje około 2 mikrosekund na jedno „przeskoczenie” między chipami.
• System obsługuje wiele zapytań jednocześnie dzięki pipeline'owaniu – każdy wafer pracuje nad innym zapytaniem w tym samym momencie, a maksymalną liczbę równoległych zapytań wyznacza dostępna pamięć na cache KV.

Jak policzyć, ile wafferów potrzeba

Punktem wyjścia jest obliczenie wielkości pojedynczej warstwy modelu. Kimi K3 ma 3 biliony parametrów. Jeśli każdy parametr przechowywany jest na 4 bitach (tak jak w implementacji Hugging Face), łączny „ładunek” wag po podzieleniu przez 93 warstwy daje około 16 GB na warstwę.

Skoro każda warstwa zajmuje ok. 16 GB, a pojedynczy wafer ma 44 GB, kluczowe staje się pytanie, ile warstw zmieści się na jednym chipie. Odpowiedź zależy od typu warstwy. Warstwy gęste (dense, pełna uwaga) wymagają dodatkowego miejsca na pamięć podręczną KV, ponieważ przechowują klucze i wartości dla każdego tokenu w każdym aktywnym zapytaniu – dlatego Cerebras umieszcza jedną warstwę dense na wafer, zostawiając resztę SRAM-u na cache KV. Warstwy linear-attention działają inaczej: zamiast przechowywać klucze i wartości dla każdego tokenu, utrzymują niewielki, stały rozmiarowo stan na zapytanie, niezależnie od długości kontekstu. Dzięki temu na jednym wafferze mieszczą się dwie takie warstwy (2 × 16 GB = 32 GB).

Przy 69 warstwach linear-attention i 24 warstwach dense w Kimi K3 daje to 35 wafferów na warstwy liniowe i 24 na warstwy gęste – łącznie 59 wafferów.

Pipeline zamiast przesuwania wag

Cerebras wykonuje modele w architekturze pipeline. Po tym, jak kompilator Cerebras rozmieści warstwy modelu na obszarach wafferów, runtime wyznacza ścieżkę przepływu tokenów między chipami. Dla ultraszybkiej latencji między wafferami przekazywane są wyłącznie aktywacje – odbywa się to przez RDMA over Ethernet, co zajmuje około 2 mikrosekundy na jeden przeskok.

Istotne jest to, że ponieważ każdy wafer Cerebras dysponuje 44 GB pamięci SRAM, a każda warstwa ma tylko 16 GB, firma zazwyczaj może zmieścić całe warstwy modelu bezpośrednio na chipie. Oznacza to, że w większości przypadków wagi modelu nigdy się nie przemieszczają, a warstwy nie muszą być dzielone między waffery – przez chipy przechodzą wyłącznie aktywacje.

Każdy wafer działa jak jedno stanowisko w linii montażowej: wagi modelu dla danego etapu pozostają w miejscu, a przez kolejne etapy przepływają wyniki pośrednie zapytania. Przykładowo wafer 1 wykonuje warstwy 1–2 i wysyła aktywacje do wafferu 2, ten wykonuje warstwy 3–4 i przekazuje wynik do wafferu 3, i tak dalej, aż ostatni wafer wygeneruje kolejny token.

Obsługa wielu zapytań jednocześnie

Pojawia się tu częste nieporozumienie: skoro token musi przejść przez każdy wafer po kolei, czy każdy wafer nie czeka bezczynnie na swoją kolej, co de facto oznaczałoby pracę przy rozmiarze batcha równym jeden? Cerebras wyjaśnia, że tak by było, gdyby system wysyłał jedno zapytanie na raz i czekał na jego zakończenie przed rozpoczęciem kolejnego – byłoby to nieefektywne wykorzystanie pipeline'u.

Zamiast tego cały pipeline jest stale wypełniony różnymi zapytaniami. Ponieważ model jest już podzielony na etapy pipeline'u, każdy etap może jednocześnie pracować nad innym zapytaniem – gdy etap 3 przetwarza token dla zapytania A, etap 2 robi to dla zapytania B, a etap 1 dla zapytania C. W przykładzie z artykułu każdy etap to jedna warstwa, ale w produkcji Cerebras stosuje znacznie bardziej drobnoziarniste etapy i jest ich dużo więcej.

Każdy etap może też przetwarzać mikro-batch, czyli niewielką grupę zapytań, zamiast pojedynczego zapytania – dzięki temu liczba zapytań „w locie” to w przybliżeniu liczba etapów pomnożona przez rozmiar mikro-batcha. Mniejsze mikro-batche przyspieszają pojedyncze zapytanie, większe zwiększają całkowitą przepustowość systemu.

O maksymalnej liczbie równoległych zapytań decyduje ilość pamięci on-chip przeznaczonej na cache KV. Każde zapytanie w locie potrzebuje własnego cache KV na każdym wafferze z warstwą gęstą, więc to rozmiar pamięci na chipie wyznacza górną granicę liczby jednoczesnych zapytań – a ponieważ warstwy linear-attention nie wymagają dodatkowego cache KV, obliczenia te dotyczą wyłącznie warstw dense. W przeciwieństwie do klasycznego batcha na GPU, zapytania w architekturze pipeline Cerebras nie czekają na siebie nawzajem, co pozwala utrzymać wysoką przepustowość przy jednocześnie niskiej latencji dla pojedynczego zapytania.

Opisana przez Cerebras metoda pokazuje, jak firma radzi sobie z wyzwaniem skalowania modeli frontierowych rzędu bilionów parametrów bez konieczności dzielenia pojedynczych warstw między chipy – podejście konkurencyjne wobec typowych rozwiązań opartych na klastrach GPU, gdzie takie problemy z pamięcią i przepustowością rozwiązuje się inaczej.