RBS-Attention: Przyspieszenie przetwarzania długich kontekstów w dużych modelach językowych

Podczas przetwarzania długiego tekstu przez modele LLM, faza prefill wymaga zastosowania gęstej self-attention do całej sekwencji, co stanowi znaczące wąskie gardło obliczeniowe. RBS-Attention rozwiązuje ten problem poprzez selektywne przetwarzanie tylko najbardziej istotnych fragmentów tekstu bez konieczności retrainingu modelu.

Metoda operuje na zasadzie dwóch komplementarnych gałęzi selekcji. Pierwsza gałąź (centroid base) identyfikuje bloki o średniej istotności, podczas gdy druga (rescue branch) wykorzystuje maksymalny promień klucza oraz statystykę zależną od prompt'a, warstwy i głowicy uwagi, aby wyłapać bloki zagrożone niedowartościowaniem - problem zwany mean dilution. Niezależne progi dla obu gałęzi pozwalają kontrolować wkład rescue bloków, zachowując jednocześnie wydajność sparse FlashAttention.

Wyniki są imponujące: na GPU H100 RBS-Attention osiąga 20.65-krotne przyspieszenie samo prefill-attention, 11.92-krotne przyspieszenie w kontekście vLLM i 5.97-krotne przyspieszenie end-to-end czasu do pierwszego tokenu dla modelu Qwen3-30B-A3B-Instruct na kontekście 128K. Na gęstym modelu Qwen3-32B dokładność RULER wyniosła 88.65% w porównaniu z 89.52% dla gęstej uwagi, a dodatkowe testy na LongBench-v2, InfiniteBench i Video-MME potwierdzają zachowanie jakości generacji przy radykalnym zmniejszeniu kosztów obliczeniowych.