الـ KV Cache: ليه أول توكن بطيء وباقي التوكنات بتطير

هذا المقال يتطلب مستوى: محترف

لو خدمتك بتولّد نص من نموذج لغة، وأول توكن بياخد وقت محسوس وباقي التوكنات بتطلع بسرعة، ده مش عشوائي. ده الـ KV Cache شغّال. المقال ده هيوريك بالظبط ليه بيحصل كده، إزاي توفّر تكلفة حسابية بترتيب O(n) بدل O(n²)، وإيه الثمن اللي بتدفعه في الذاكرة.

الـ KV Cache: ليه أول توكن بطيء وباقي التوكنات بتطير

المشكلة باختصار

نموذج اللغة بيولّد توكن واحد في المرة. عشان يطلّع التوكن الجديد، الانتباه الذاتي (Self-Attention) محتاج يبص على كل التوكنات اللي قبله. لو كل خطوة أعادت حساب المفاتيح (Keys) والقيم (Values) لكل التوكنات السابقة من الأول، تكون بتكرّر نفس الشغل ملايين المرات في نص طويل. النتيجة: التكلفة بتكبر تربيعيًا مع طول النص، والـ GPU بيتخنق من غير سبب حقيقي.

الفكرة ببساطة: المذيع اللي بيقرأ نشرة

تخيّل مذيع بيقرأ نشرة أخبار جملة ورا جملة. عشان ينطق الكلمة الجاية صح، هو محتاج يفتكر اللي قاله قبل كده. فيه طريقتين. الأولى الغبية: قبل كل كلمة جديدة، يرجع يقرأ النشرة كلها من أول سطر لغاية مكانه، وبعدين ينطق الكلمة الواحدة الجديدة. الثانية العاقلة: هو فاكر خلاص كل اللي فات، فبيضيف الكلمة الجديدة على طول من غير ما يعيد قراءة أي حاجة.

الطريقة الأولى هي التوليد بدون KV Cache. الثانية هي التوليد بالـ KV Cache. الفرق مش في الجودة، الفرق في إنك بتوفّر إعادة قراءة كل اللي فات في كل خطوة.

المفهوم علميًا

في طبقة الانتباه، كل توكن بيتحوّل لثلاث متجهات: Query و Key و Value. التوكن الجديد بيقارن الـ Query بتاعه بكل الـ Keys السابقة عشان يحسب أوزان الانتباه، وبعدين يجمّع الـ Values حسب الأوزان دي. الملاحظة المفتاحية: الـ Keys والـ Values بتاعة التوكنات القديمة ما بتتغيّرش لما نضيف توكن جديد. فبدل ما نحسبها تاني كل مرة، نحسبها مرة واحدة ونخزّنها في ذاكرة الـ GPU. ده هو الـ KV Cache.

عشان كده الاستدلال بينقسم لمرحلتين. الأولى Prefill: بنحسب K و V لكل توكنات الـ prompt دفعة واحدة، وده بياخد الوقت الأطول (أول توكن). الثانية Decode: كل توكن جديد بيحسب K و V بتاعه هو بس، بيضيفهم للـ cache، ويقرأ الباقي جاهز. ده اللي بيحصل فعلاً لما تحس إن أول توكن تقيل وباقي التوكنات خفيفة.

القياس بالكود

الكود ده بيقيس الفرق فعليًا على أي نموذج من Hugging Face. جرّبه بنفسك:

import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "meta-llama/Llama-2-7b-hf"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.float16).cuda()

prompt = tok("اشرح لي فكرة الـ KV Cache باختصار:", return_tensors="pt").to("cuda")

for use_cache in (True, False):
    torch.cuda.synchronize(); t0 = time.time()
    model.generate(**prompt, max_new_tokens=200, use_cache=use_cache)
    torch.cuda.synchronize()
    print(f"use_cache={use_cache}: {time.time()-t0:.2f}s")

على GPU من فئة A100 وسياق قصير، النتيجة التقريبية بتبان كده (الأرقام توضيحية للاتجاه، وبتختلف حسب العتاد وطول النص):

لاحظ الاتجاه: من غير cache التكلفة بتنمو تربيعيًا (كل خطوة بتعيد حساب كل اللي فات)، فالفجوة بتتّسع كل ما النص يطول. مع cache بتنمو خطيًا. ده الفرق بين O(n²) و O(n) على مستوى إجمالي التوليد.

الـ trade-off هنا: الذاكرة

مفيش حاجة ببلاش. الـ KV Cache بيوفّر حساب مقابل إنه بياكل VRAM، والاستهلاك بيكبر مع طول السياق. الحجم بيتحسب تقريبًا بالمعادلة دي:

حجم الـ cache = 2 (K و V) × عدد الطبقات × البُعد المخفي × طول السياق × عدد البتات لكل قيمة

مثال واقعي على Llama‑2‑7B بدقة fp16: عدد الطبقات 32، البُعد المخفي 4096. ده بيطلّع حوالي 0.5 ميجابايت لكل توكن. يعني سياق 4096 توكن بياخد قرابة 2 جيجابايت من الـ VRAM لكل طلب، فوق حجم النموذج نفسه. لو بتخدم عدة طلبات بالتوازي، اضرب في عدد الطلبات. الافتراض إنك على نموذج 7B وسياق ≤ 4K؛ في النماذج الأكبر أو السياقات الطويلة الرقم بيقفز بسرعة، وساعتها بتلجأ لتقنيات زي Grouped-Query Attention أو PagedAttention (اللي في vLLM) لتقليل بصمة الذاكرة.

متى لا تشغّل بالك

لو بتعمل تلخيص أو تصنيف بمخرج قصير جدًا (توكن أو اتنين)، مكسب الـ cache هيبقى ضئيل لأن مفيش خطوات decode كتير أصلًا. ولو بتشتغل embeddings بس (تشفير من غير توليد)، مفيش KV cache من الأساس لأن مفيش توليد تتابعي. كمان لو الذاكرة عندك مخنوقة والسياق طويل جدًا، ساعات تعطيل الـ cache أو تقليصه بيكون مقايضة مقصودة: بتخسر سرعة عشان تكسب ذاكرة تكفّي الطلب يشتغل أصلًا.

الخطوة التالية

افتح أي سكربت توليد عندك ودوّر على وسيط use_cache. لو مش موجود صريحًا، اعرف إنه بيبقى True افتراضيًا في generate. بعدين قِس زمن التوليد لـ 200 توكن مرة بـ True ومرة بـ False، وقارن. لو الفرق مكانش واضح، غالبًا نصك قصير أوي علشان تحس بالفرق؛ طوّل الـ max_new_tokens وأعد القياس.

المصادر

• Vaswani et al., "Attention Is All You Need" (2017) — أصل آلية الانتباه الذاتي و Q/K/V.
• Hugging Face Transformers — توثيق generate ووسيط use_cache وpast_key_values.
• Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models" (2023) — تقليل حجم الـ KV Cache.
• Kwon et al., "Efficient Memory Management for LLM Serving with PagedAttention" (2023) — ورقة vLLM.