Optimalizace nezměnila výpočet, změnila předpoklad o datech
Prompt lookup decoding je spekulativní dekódování s tím nejhloupějším možným draft modelem: n-gramovým počítadlem. Žádná druhá neuronová síť. Jen tabulka, která říká, jaký token po téhle sekvenci nejčastěji následoval, a hlavní model pak hromadu takových odhadů ověří naráz.
V llama.cpp šlo tohle hledání z 165 mikrosekund na token na 1,18. Čtyřicetdvakrát, se Lemierovým příspěvkem sto čtyřicetkrát. Ani jeden krok nezměnil, co se počítá. Všechny změnily, kde ta data leží a v jakém pořadí se na ně sahá.
Nejdřív se přestalo kopírovat
Největší jednotlivý skok, 4,5× až 25,6× podle velikosti korpusu, přišel z toho, že se přestaly kopírovat mapy. Ne z lepšího hledání. Z toho, že se hodnota, která už v paměti byla, nepřenášela znovu.
Stojí za to si to zapamatovat jako pořadí, v jakém se má profilovat. Tři cache tu drží různě dlouhověká data: kontextová n-gramy z právě běžící konverzace, dynamická počty z předchozích běhů, statická předpočítané dvojice z referenčního korpusu, které se po inicializaci nikdy nezmění. Rozdíl mezi nimi není v algoritmu, je v tom, kdo je smí přepsat. A právě u té statické, kterou nikdo přepsat nesmí, se kopie dělaly úplně zbytečně.
Šedesát čtyři procent n-gramů má jediného následníka
Tohle je to load-bearing zjištění celé práce. Autor se podíval na rozložení a zjistil, že asi 64 % dvojic má v korpusu právě jednoho pokračovatele. Distribuce má sice těžký chvost, kde po nějaké běžné fázi následují tisíce různých tokenů, ale ten typický případ je jednoprvkový.
Hashmapa je pro jednoprvkovou kolekci absurdně drahá. Nese bucket array, režii na kolize a rozptýlené odkazy, aby v konstantním čase našla to jedno. Setřídený std::vector má logaritmické hledání, které na jednom prvku nic neznamená, a leží celý v jednom kusu paměti. Výsledek: paměť až 1,97× menší, latence 2,09× lepší bez cache.
Následoval krok, který vypadá jako detail a není. Binární hledání s pevnou délkou, kde počet iterací nezávisí na výsledku porovnání. Klasické binární hledání musí po každém srovnání vědět, kam skočit, takže procesor čeká na cache miss, než smí pokračovat. Když je délka smyčky daná dopředu, může načítání běžet napřed. Pevný počet kroků je horší algoritmus a rychlejší kód. Statická cache přešla na constmap a doba načtení spadla z 3,76 s na 0,23 s.
Levný test patří před drahý výpočet
Poslední krok přinesl Daniel Lemire a je celý o pořadí. Původní kód spočítal skóre všem kandidátům a pak se podíval, jestli některý překonal prahy. Nová verze se nejdřív zeptá, jestli vůbec součet výskytů daného n-gramu dosáhne minima. Pokud ne, končí. Pak se zeptá, jestli i ten nejčastější následník splní požadovaný podíl. Pokud nesplní on, nesplní ho nikdo, a skóre se nepočítá vůbec.
Žádné skóre se nezmění. Změní se jen množina těch, které se vyplatilo spočítat. 4,2× rychlejší s cache.
Tahle transformace má obecnou podobu: když má filtr monotónní horní mez, dá se ta mez otestovat dřív než jednotlivé prvky. Vyplatí se to přesně tehdy, když většina vstupů filtrem neprojde. A to tady platí, protože většina n-gramů v běžném textu je vzácná.
Co si z toho odnést
Tři kroky ze čtyř nezměnily výpočet, změnily rozvržení paměti. Čtvrtý nezměnil ani to, jen pořadí testů. To je typické: v horké smyčce už algoritmus obvykle správný je a rozdíl dělá, jak se data dotýkají procesoru.
Rozložení dat si nejdřív změř, teprve pak vyber strukturu. Hashmapa je rozumný default pro neznámé rozložení, ale jednoprvkové kolekce jí platí režii, kterou nikdy nevyužijí. A když se v profilu nudíš, zkontroluj, jestli se něco nekopíruje a jestli nejdražší test neběží první.
Predikovatelnost přitom někdy vyhraje nad asymptotikou. Binární hledání s pevnou délkou udělá víc práce a je rychlejší, protože nechá procesor pracovat napřed.