KV-cache-quantisatie voor on-device LLM's
Het draaien van Llama 3.2 3B op een Android-toestel met 2 GB RAM is lastig. De meeste ontwikkelaars richten zich op modelgewichten. Dit is een fout. De echte geheugenvreter is de KV-cache.
De KV-cache groeit naarmate je chat. Als je standaard FP16-precisie gebruikt, verbruikt de cache honderden megabytes. Dit zorgt ervoor dat je app na slechts enkele beurten crasht.
Je kunt dit oplossen met drie specifieke stappen.
- Gebruik mixed-precision quantisatie Keys en values hebben niet dezelfde precisie nodig. Key-caches gaan goed om met lage precisie. Value-caches niet.
- Gebruik INT4 voor keys.
- Gebruik INT8 voor values.
Deze aanpak verkleint je cachegrootte met 62%. Voor een context van 2048 tokens daal je van 224 MB naar 84 MB. Dit gebeurt zonder de modelgewichten te wijzigen.
Implementeer sliding window eviction Je kunt niet elk token in het actieve geheugen houden. Gebruik een sliding window om alleen de meest recente 1536 tokens te bewaren. Houd de eerste 64 tokens aan als ankers om de system prompt te behouden.
Gebruik flash spilling Wanneer tokens het sliding window verlaten, verplaats ze dan naar flash-opslag. Gebruik memory-mapped files op Android. Moderne UFS 4.0-opslag is snel genoeg om deze gegevens zonder vertraging terug in het geheugen te laden.
De resultaten zijn aanzienlijk. Op een Snapdragon 8 Gen 3:
- Het piekgeheugen daalt onder de grens van 2 GB.
- Het maximale aantal gespreksbeurten neemt toe van 4 naar meer dan 12.
- De tokensnelheid neemt toe omdat kleinere caches de geheugenbandbreedte beter benutten.
- De modelkwaliteit blijft bijna gelijk.
Vermijd deze fouten:
- Quantiseer keys en values niet naar hetzelfde niveau. Je zult kwaliteit verliezen.
- Negeer thermal throttling niet. Voortdurende inferentie zorgt voor hitte. Controleer de Android Thermal HAL om de prestaties te beheren.
- Vergeet de levenscyclus van de cache niet. Koppel gemapte buffers altijd aan een juiste scope om geheugenlekken te voorkomen.
Stel je geheugenbudget vast voordat je je functies bouwt.
Bron: https://dev.to/software_mvp-factory/kv-cache-quantization-for-on-device-llms-kf