[PL] TEMPO — balansowanie obciążenia w expert-parallel MoE, które widzi oba reżimy pracy GPU

Serwowanie modeli Mixture-of-Experts (MoE) w trybie expert-parallel (EP) ma uporczywy problem: każda warstwa synchronizuje się na najwolniejszym GPU. Jeden przeciążony ekspert i cały batch czeka na barierze. Rozwiązaniem mają być dispatchery, które rozdzielają pracę między GPU – tylko że wszystkie popularne heurystyki balansują jedną wielkość: liczbę tokenów (EPLB, LPLB, UltraEP) albo liczbę aktywowanych ekspertów (METRO), przy cichym założeniu, że czas eksperta zależy od niej liniowo.

Zespół KlingAI Research (Kuaishou) zmierzył to na dwóch generacjach GPU do centrów danych i w pracy TEMPO: Makespan-Aware Expert-Parallel Load Balancing Across Memory- and Compute-Bound Regimes (arXiv:2608.13057) pokazuje, że oba założenia są fałszywe. Czas eksperta w dekodowaniu nie jest liniowy ani w tokenach, ani w ekspertach – i to zmienia to, jak powinien wyglądać dispatch. Na nagranych batchach wybór heurystyki zmienia modelowany czas bloku o 1,4-1,6x (p95 do 1,7x), a to, która heurystyka wygrywa, zależy od reżimu pracy GPU.

Dwa reżimy, jeden próg: 156-168 tokenów

Autorzy mierzyli czas ekspertów w funkcji liczby tokenów i liczby aktywowanych replik (repliki to kopie tego samego eksperta rozłożone na wiele GPU). Wynik: są dwa reżimy z wyraźnym progiem n* około 156-168 tokenów:

• Poniżej progu – reżim pamięciowy (memory-bound). Dominuje streaming wag z HBM. Koszt wykonania zależy od liczby aktywowanych replik: każda replika czyta pełne wagi eksperta z pamięci, niezależnie od tego, ile tokenów do niej trafi. Balansowanie tokenów nie ma sensu, bo koszt jest stały per replika.
• Powyżej progu – reżim obliczeniowy (compute-bound). Grouped GEMM na tensor core’ach zaokrągla wymiar M (tokeny) do kafelków po 128. Rozbicie eksperta na więcej GPU dodaje więc paddingu: każdy kafelek 128 trzeba wypełnić, więc podział eksperta na dwie części może kosztować więcej niż trzymanie go w jednym miejscu.

Oba reżimy da się opisać jednym profilem max-affine: t = max(a + bG, c + betaN), gdzie G to liczba aktywowanych replik, a N liczba tokenów. Prosty model, ale robi robotę: dla danego batcha wiadomo, po której stronie progu się jest.

Dlaczego dotychczasowe heurystyki zawodzą

W prawdziwym dekodowaniu batch trzyma jednocześnie gorące eksperty (dużo tokenów, reżim liniowy) i zimne (mało tokenów, reżim płaski). Heurystyka balansująca tokeny świetnie działa na gorących i szkodzi na zimnych, a heurystyka balansująca repliki – odwrotnie. Dlatego wyniki proxy tak się rozjeżdżają: na nagranych batchach decyzje różnych heurystyk różnią się 1,4-1,6x w modelowanym czasie bloku.

Dispatch jako problem makespan

TEMPO formalizuje dispatch per batch jako problem makespan z kosztem stałym (fixed-charge): każdy ekspert ma koszt aktywacji (stały) plus koszt zależny od tokenów. To problem NP-trudny już na dwóch w pełni zreplikowanych GPU, ale w granicznych przypadkach wielomianowy, a przy pełnej replikacji obowiązuje addytywna gwarancja aproksymacji. W praktyce dispatcher rozwiązuje instancję w milisekundach, poza ścieżką krytyczną – nie dodaje latencji.

Integracja z SGLang działa out-of-process: dispatch i zbieranie liczników są scalone w jeden kernel w grafie, więc narzut jest minimalny.

Wyniki: fazowy diagram zamiast uniwersalnego zwycięstwa

Na mikrobenchmarku (8 GPU, Testbed A) TEMPO wszędzie trzyma się w granicach 1% najlepszego stałego baseline’u, a tam, gdzie reżimy się mieszają, wygrywa do 15,5%. End-to-end na Testbed B:

• Qwen3-235B (model w regionie wygranej): +4-6% przepustowości, p99 latencji w dół o około 15,6%;
• DeepSeek-V3 (poza regionem, zdominowany komunikacją): tylko koszt mechanizmu, bez zysków.

To uczciwe podejście: autorzy nie sprzedają TEMPO jako panaceum. Dostarczają diagram fazowy, który przed wdrożeniem przewiduje, czy twój model i profil batchy będą w regionie wygranej. Jeśli serwujesz duży model MoE z gorącymi i zimnymi ekspertami (reżimy mieszają się wtedy na co dzień), jest to warte przetestowania.

Co to oznacza dla inżyniera

• Założenie “czas eksperta jest liniowy w tokenach” jest po prostu nieprawdziwe na nowoczesnych GPU. Każdy własny dispatcher albo balanser powinien uwzględniać próg około 160 tokenów.
• Replikacja ekspertów ma ukryty koszt w reżimie pamięciowym: każda replika czyta pełne wagi. To argument przeciwko przesadnej replikacji zimnych ekspertów.
• Grouped GEMM z kafelkami 128 oznacza, że dzielenie eksperta między GPU ma sens tylko powyżej progu – poniżej dodajesz tylko padding.
• SGLang to dojrzały punkt integracji: dispatcher działający out-of-process z fuzją kernela to wzorzec, który można powtórzyć w innych stackach (vLLM, TensorRT-LLM).

Treść została wygenerowana przy użyciu sztucznej inteligencji i zweryfikowana przez redaktora.

Źródła

• TEMPO: Makespan-Aware Expert-Parallel Load Balancing Across Memory- and Compute-Bound Regimes – arXiv:2608.13057 (KlingAI Research, opublikowana 13.08.2026)
• China AI Weekly Brief, okres 9-16.08.2026 (job 3d03e2d4fae9)