קוונטיזציה של KV Cache עבור LLMs במכשירים (On-Device)
הרצת Llama 3.2 3B על מכשיר Android עם 2 GB של RAM היא משימה קשה. רוב המפתחים מתמקדים במשקלי המודל (model weights). זו טעות. "הרוצח" האמיתי של הזיכרון הוא ה-KV cache.
ה-KV cache גדל ככל שמתנהלים בשיחה. אם משתמשים בדיוק (precision) סטנדרטי של FP16, ה-cache צורך מאות מגה-בייט. זה גורם לאפליקציה שלכם לקרוס לאחר מספר סיבובים בודדים בלבד.
ניתן לפתור זאת באמצעות שלושה צעדים ספציפיים.
- שימוש בקוונטיזציה בערבוב דיוק (Mixed-Precision Quantization) למפתחות (Keys) ולערכים (Values) אין צורך באותו דיוק. Key caches מתמודדים היטב עם דיוק נמוך. Value caches לא.
- השתמשו ב-INT4 עבור keys.
- השתמשו ב-INT8 עבור values.
גישה זו מפחיתה את גודל ה-cache ב-62%. עבור הקשר (context) של 2048 טוקנים, אתם יורדים מ-224 MB ל-84 MB בלבד. זה קורה מבלי לשנות את משקלי המודל.
יישום Sliding Window Eviction אי אפשר לשמור כל טוקן בזיכרון פעיל. השתמשו ב-sliding window כדי לשמור רק את 1536 הטוקנים האחרונים ביותר. שמרו את 64 הטוקנים הראשונים כעוגנים (anchors) כדי לשמר את ה-system prompt.
שימוש ב-Flash Spilling כאשר טוקנים יוצאים מה-sliding window, העבירו אותם לאחסון flash. השתמשו בקבצים ממופי זיכרון (memory-mapped files) ב-Android. אחסון UFS 4.0 מודרני הוא מהיר מספיק כדי להחזיר את הנתונים הללו לזיכרון (paging) ללא השהיה.
התוצאות משמעותיות. ב-Snapdragon 8 Gen 3:
- צריכת הזיכרון בשיא (Peak memory) יורדת מתחת למגבלת ה-2 GB.
- מספר סיבובים מקסימלי בשיחה עולה מ-4 ליותר מ-12.
- מהירות הטוקנים עולה מכיוון ש-caches קטנים יותר מנצלים טוב יותר את רוחב פס הזיכרון (memory bandwidth).
- איכות המודל נשארת כמעט ללא שינוי.
הימנעו מהטעויות הבאות:
- אל תבצעו קוונטיזציה ל-keys ול-values באותה רמה. אתם תאבדו איכות.
- אל תתעלמו מ-thermal throttling. פעולות inference ממושכות גורמות להתחממות. בדקו את ה-Android Thermal HAL כדי לנהל את הביצועים.
- אל תשכחו את מחזור החיים (lifecycle) של ה-cache. תמיד קשרו buffers ממופים (mapped buffers) לטווח (scope) מתאים כדי למנוע דליפות זיכרון.
בניתם את תקציב הזיכרון שלכם לפני שאתם בונים את הפיצ'רים שלכם.
מקור: https://dev.to/software_mvp-factory/kv-cache-quantization-for-on-device-llms-kf