NVIDIA wprowadza Skip Softmax – rewolucja w szybkości LLM-ów

Czy naprawdę trzeba liczyć uwagę do każdego słowa w 128 tysiącach tokenów? NVIDIA twierdzi, że nie – i pokazuje, jak zrobić to szybciej, nie rozwalając modelu.

Skip Softmax w TensorRT-LLM to sprytna „wycinka” w jądrze FlashAttention, która dynamicznie omija nieistotne bloki uwagi. Efekt: do 1,4x krótszy czas do pierwszego tokenu (TTFT) i do 1,4x szybszy czas na token wyjściowy (TPOT), bez trenowania od zera. Działa na Hopperze i Blackwellu, współgra z istniejącymi modelami i szczególnie lubi długie konteksty.

Jeśli budujesz RAG, agentów lub generujesz długie formy, znasz ból: koszt atencji rośnie wraz z długością kontekstu. Dekodowanie dławi pamięć (KV cache lata tam i z powrotem), prefill mieli FLOPs. Wiele „cudownych” terapii wymaga modyfikacji architektury. Skip Softmax gra inaczej: to sprzętowo przyjazna, drop-in metoda rzadkiej atencji, która odchyla się minimalnie od standardu – i za to płaci szybkością.

O co chodzi

Zwykły przebieg FlashAttention wygląda tak: liczysz logity QK, normalizujesz softmaksem, ważysz wartości V. W praktyce duża część tych obliczeń nie wnosi nic: sporo bloków (przetwarzanych kafelków Q i K) ma tak niskie logity względem dominujących tokenów, że ich wkład jest statystycznie pomijalny.

Skip Softmax przenosi tę intuicję prosto do jądra obliczeń. Dla każdego bloku wylicza lokalne maksimum logitów i porównuje je z „bieżącym” globalnym maksimum. Jeśli różnica przekracza skalibrowany próg, blok zostaje pominięty: nie wykonuje się softmax ani drugie mnożenie macierzy (BMM2), a co kluczowe – blok nie jest nawet ładowany z HBM. To jak bramkarz w klubie, który na oko rozpoznaje gości bez zaproszeń i nie marnuje czasu, by sprawdzać ich listę. [1]

Dlaczego to przyspiesza

• Dekodowanie (generacja) jest zwykle ograniczone przepustowością pamięci. GPU częściej taszczy dane KV z HBM, niż liczy. Skip Softmax pozwala niczego nie taszczyć, jeśli blok i tak nie „zapunktuje” – realnie zmniejsza transfery.
• Prefill (przetwarzanie promptu) bywa ograniczony obliczeniami. Pomijając softmax i BMM2 dla wielu bloków, ucinasz FLOPs wprost.

NVIDIA podaje liczby dla Llama 3.3 70B na platformie GB200 NVL72: w dekodowaniu zysk rzędu 1,36x end-to-end; w prefillu przy 128K kontekstu – ok. 1,4x. Co ważne, zyski rosną wraz z długością sekwencji. Próg „skipowania” jest matematycznie powiązany z długością kontekstu, więc im dłużej, tym więcej okazji, by bezpiecznie odpuścić nieistotne bloki.

Jak to działa pod maską

• Lokalny maksimum logitów: dla bieżącego kafelka QK.
• Porównanie z globalnym maksimum: różnica większa od progu? Odcinamy.
• Konsekwencja: nie ładujemy bloku z HBM i nie liczymy softmax/BMM2 dla tego fragmentu.

Ten drobny zwrot w pętli FlashAttention ma dwie duże zalety. Po pierwsze, nie wymaga żadnych zmian w architekturze modelu: działa z MHA, GQA, MLA – czyli z tym, co już masz wytrenowane. Po drugie, jest zoptymalizowany pod Tensor Cores i hierarchię pamięci w Hopperze i Blackwellu, więc nie trzeba „dokręcać” runtime’u ręcznie.

Szerszy kontekst: nie jeszcze jedna egzotyczna atencja

Na rynku przewinęły się już linearyzacje atencji i różne sztuczki z kompresją, ale zwykle okupione zmianami w modelu lub jakości. Skip Softmax jest bardziej pragmatyczny: model zostaje w spokoju, a my przyspieszamy na poziomie runtime’u, korzystając z naturalnej rzadkości rozkładu uwagi. To także metoda modularna: w prefillu można używać XAttention, a w dekodowaniu odpalić Skip Softmax – NVIDIA pokazuje, że takie zestawienie daje duże przyspieszenia bez strat jakości. [1]

Fakty i dane, bez waty

• Do 1,4x szybciej: zarówno TTFT (czas do pierwszego tokenu), jak i TPOT (czas na każdy token wyjściowy).
• Llama 3.3 70B na GB200 NVL72: ok. 1,36x w dekodowaniu end-to-end; ok. 1,4x w prefillu przy 128K.
• Skala działa na korzyść: im dłuższy kontekst, tym lepszy efekt z „wycinki”.
• Kompatybilność: żadnego retrainingu; działa z klasyczną atencją i popularnymi wariantami głowic.
• Hardware: zoptymalizowane pod Hopper/Blackwell i pamięć HBM, gdzie koszt I/O jest realnym ograniczeniem.

A co z jakością?

Jest próg, który decyduje, jak agresywnie tniemy. To parametr do kalibracji: im wyższy, tym bezpieczniej dla jakości, ale mniejszy zysk; im niższy – tym szybciej, z potencjalnym ryzykiem błędów. NVIDIA pokazuje, że w praktyce można dobrać ustawienia, które nie degradują jakości, a w duecie z XAttention w prefillu nadal zapewniają solidny wzrost szybkości. Elegancja polega na tym, że dobierasz kompromis na poziomie runtime’u, bez ruszania wag.

Gdzie to się wpisuje w trend

Długie konteksty – 64K, 128K i dalej – to dziś nie tyle „nice to have”, co wymóg: RAG, narzędzia-agenci, długie dokumenty, multi-turn z pamięcią. Problem w tym, że tradycyjna atencja skaluje się kwadratowo, a GPU, choć szybkie, mają realne ograniczenia przepustowości. Optymalizacje w stylu Skip Softmax – oszczędzające zarówno FLOPs, jak i I/O – to trzeźwa odpowiedź rynku: zamiast wymyślać nowe modele, najpierw wyciskamy to, co da się z hardware’u i kerneli.

Krótki komentarz z dystansem

To nie jest magiczny skrót do atencji o złożoności liniowej. Ale to rzecz, która naprawdę daje procenty – dokładnie tam, gdzie boli najbardziej. Jeśli ładujesz 128K kontekstu, a dekoder jest ograniczony przepustowością HBM, to lepiej, żeby nie ładował bloków, które i tak wylądują pod progiem softmaksu. Czasem inżynieria wydajności to po prostu odwaga, by powiedzieć „nie” nieistotnym obliczeniom.

Na wynos

Skip Softmax to sensowny, praktyczny upgrade dla długo-kontekstowych LLM-ów w TensorRT-LLM: bezboleśnie skraca TTFT i TPOT, skaluje się z długością sekwencji, nie wymaga zmian w modelu i korzysta z realiów Hopper/Blackwell (pamięć, kerneli). Jeśli Twoja aplikacja opiera się na RAG-ach i agentach, to jedno z tych usprawnień, które po prostu warto włączyć i skalibrować. A potem spokojnie patrzeć, jak pierwsze tokeny pojawiają się szybciej. [1]

Źródła

• [1] https://developer.nvidia.com/blog/accelerating-long-context-inference-with-skip-softmax-in-nvidia-tensorrt-llm/