KV-välimuisti LLM-malleissa: ymmärtäminen ja toteutus alusta alkaen

KV-välimuisti tallentaa aiemmista attention-laskuista saadut avain- (K) ja arvovektorit (V), jotta niitä voidaan käyttää uudelleen seuraavien tokenien tuottamisessa, mikä nopeuttaa päättelyä välttämällä uudelleenlaskennan. Ilman KV-välimuistia jokainen generointivaihe laskee avaimet ja arvot uudelleen kaikille aiemmille tokeneille; välimuistin kanssa lasketaan ja liitetään välimuistiin vain uuden tokenin vektorit. Artikkeli tarjoaa kooditoteutuksen, joka perustuu kirjoittajan kirjan GPT-tyyppiseen malliin. Keskeiset muutokset ovat: välimuistipuskurien (cache_k ja cache_v) lisääminen MultiHeadAttention-luokkaan, forward-metodin muokkaaminen käyttämään välimuistia ehdollisesti, reset-metodin lisääminen ja use_cache-lippujen välittäminen mallin läpi. Generoinnissa, kun use_cache on tosi, malli käsittelee vain uuden tokenin alkuperäisen kehotteen jälkeen, kun taas ilman välimuistia se käsittelee koko sekvenssin jokaisessa vaiheessa. Yksinkertainen suorituskykyvertailu osoittaa, että KV-välimuisti noin kaksinkertaistaa generointinopeuden testatussa esimerkissä.