KV Cache للمتوسط: ليه Context طويل في LLM بياكل GPU وازاي تقلله 4x

المستوى المطلوب: متوسط — تحتاج فهمًا أساسيًا لـ Transformers و Self-Attention، مع خبرة عملية في تشغيل LLMs محليًا أو سحابيًا (vLLM، Ollama، أو HuggingFace Transformers).

لو حاولت تشغّل Llama 3 70B على A100 80GB بـ 32K context وجاتلك CUDA out of memory، المشكلة مش في وزن الموديل. وزن الموديل في FP16 بيبقى 140GB، ولو quantized لـ INT4 بينزل لـ 40GB. الباقي من ذاكرة الـ 80GB بيتأكل في حاجة اسمها KV Cache، وبتكبر خطيًا مع طول الـ context. كل توكن إضافي بياخد ذاكرة محسوبة بدقة، ولما توصل لـ 32K توكن الـ Cache لوحده ممكن يوصل 10GB على Llama 3 70B بـ GQA، أو 80GB لو الموديل بـ Multi-Head Attention الكاملة.

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

الـ Self-Attention في Transformer لازم لكل توكن جديد يحسب علاقاته بكل التوكنز اللي قبله. عشان متعيدش الحسابات في كل مرة، الموديل بيخزّن مصفوفتي Key و Value لكل التوكنز السابقة في الذاكرة. ده اللي اسمه KV Cache. النتيجة: كل توكن إضافي بيضيف للذاكرة بشكل ثابت، ومع contexts طويلة الـ Cache بيستهلك ذاكرة أكتر من الموديل نفسه.

مثال للمبتدئ: السكرتير اللي بيحفظ الاجتماع

تخيّل سكرتير في اجتماع طويل بيكتب محضر. كل ما حد يقول جملة، السكرتير محتاج يفهم علاقتها بكل جملة قيلت قبلها. لو في كل مرة هيرجع يقرأ المحضر كله من الأول، الاجتماع هيقعد ساعات بدون ما يخلص.

السكرتير الشاطر بيعمل حاجة تانية: لكل جملة، بيكتب جنبها "أهم نقطتين فيها" (ده الـ Key) و"ملخصها التنفيذي" (ده الـ Value). لما تيجي جملة جديدة، بيفتح ملخصاته السابقة بدل ما يعيد قراءة كل المحضر. ده الـ KV Cache بالظبط.

المشكلة: المحضر بيكبر، ودفتر الملخصات بيكبر معاه. لو الاجتماع طول 10 ساعات، دفتر الملخصات وحده بقى أتقل من المحضر الأصلي. ده اللي بيحصل لما توصل لـ 32K توكن: الـ KV Cache يتعدّى وزن الموديل نفسه.

التعريف العلمي الدقيق

في كل layer من الـ Transformer، الـ Self-Attention بيحسب ثلاث مصفوفات: Query (Q)، Key (K)، Value (V). للتوكن الجديد بيتم حساب Q جديد، لكن K و V لكل التوكنز السابقة لازم يكونوا موجودين في الذاكرة عشان عملية الـ attention dot product تشتغل بدون إعادة حساب.

الصيغة الرياضية لحجم الـ KV Cache:

KV_size = 2 × num_layers × seq_length × num_kv_heads × head_dim × precision_bytes

الرقم 2 لأن عندك Key و Value. على Llama 3 70B (FP16):

• num_layers = 80
• num_kv_heads = 8 (بفضل Grouped Query Attention)
• head_dim = 128
• precision = 2 bytes (FP16)

الحساب لكل توكن: 2 × 80 × 8 × 128 × 2 = 327,680 بايت = 320 KB.

لـ 32,768 توكن: 320 KB × 32,768 = 10.24 GB.

على Llama 2 70B الأصلي بـ Multi-Head Attention الكامل (64 KV head)، نفس الحساب بيطلع 81.92 GB لـ 32K توكن. ده بيتجاوز سعة A100 80GB لوحده، يعني الموديل أصلًا مش هيتحمّل.

قياس الاستهلاك بكود Python شغّال

الكود ده بيشتغل على Llama 3 8B وبيقيس الـ KV Cache فعليًا لطول contexts مختلفة:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.float16, device_map="cuda"
)

def measure_kv_cache(seq_length: int) -> float:
    text = "هذا نص اختباري للقياس " * (seq_length // 6)
    inputs = tokenizer(
        text, return_tensors="pt",
        truncation=True, max_length=seq_length
    ).to("cuda")

    with torch.no_grad():
        outputs = model(**inputs, use_cache=True)

    kv_size_gb = sum(
        k.numel() * k.element_size() + v.numel() * v.element_size()
        for k, v in outputs.past_key_values
    ) / 1e9

    del outputs
    torch.cuda.empty_cache()
    return kv_size_gb

for length in [512, 4096, 16384, 32768]:
    kv = measure_kv_cache(length)
    print(f"seq_len={length:5d}: KV cache = {kv:.2f} GB")

المخرج المتوقع على Llama 3 8B (32 layer، 8 KV heads):

seq_len=  512: KV cache = 0.13 GB
seq_len= 4096: KV cache = 1.07 GB
seq_len=16384: KV cache = 4.29 GB
seq_len=32768: KV cache = 8.59 GB

ركز في الرقم الأخير: على موديل 8B، الـ Cache لوحده تجاوز حجم الموديل (16 GB في FP16). على Llama 3 70B هيكون 4x أكتر من ده.

الحلول الأربعة لتقليل الـ KV Cache

1. Multi-Query Attention (MQA)

بدل ما كل attention head يكون له K و V خاصين، كل الـ heads بيشتركوا في زوج K/V واحد بس. التوفير بيوصل 8x لكن الجودة بتفقد محسوسة في موديلات صغيرة، خصوصًا في tasks الـ reasoning الطويل.

الاستخدام: PaLM، Falcon، StarCoder.

2. Grouped Query Attention (GQA)

توازن وسط بين MHA الكاملة و MQA. مجموعة من الـ heads بتشترك في K/V واحد. Llama 3 70B عنده 64 query head لكن 8 KV heads فقط — توفير 8x. الجودة بتفقد 0.3% إلى 0.7% بس على MMLU benchmark، وده مقبول جدًا في 99% من الحالات.

الاستخدام: Llama 3، Mixtral، Qwen 2.5، DeepSeek-V3.

3. PagedAttention (vLLM)

المشكلة الأساسية في الـ Cache التقليدي إنه بيحجز ذاكرة متصلة (contiguous) بطول max_context لكل request، حتى لو الـ request الفعلي 200 توكن بس. PagedAttention بياخد فكرة Virtual Memory من الـ OS: بيقسم الـ Cache لـ blocks صغيرة (16 توكن لكل block)، وبيحجز Block على الطلب. التوفير بيوصل 60% إلى 80% في scenarios متعدد المستخدمين.

vllm serve meta-llama/Meta-Llama-3-70B-Instruct \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.95 \
  --enable-prefix-caching \
  --tensor-parallel-size 2

على إنتاج فعلي بـ 100 user متوازي على نفس الـ GPU، الفرق بيكون: throughput من 14 req/s لـ 47 req/s — تحسن 3.4x.

4. KV Cache Quantization

تخزين الـ K و V بـ INT8 أو INT4 بدل FP16. بيوفر 2x إلى 4x في الذاكرة بفقد جودة أقل من 1% لو استخدمت per-channel quantization مع HQQ backend.

from transformers import AutoModelForCausalLM, QuantizedCacheConfig

cache_config = QuantizedCacheConfig(
    backend="hqq",
    nbits=4,
    axis_key=0,
    axis_value=0,
    compute_dtype=torch.float16,
)

outputs = model.generate(
    **inputs,
    cache_implementation="quantized",
    cache_config=cache_config,
    max_new_tokens=512,
)

الـ Trade-offs الحقيقية

• MQA: توفير 8x في الذاكرة، خسارة 1% إلى 2% في accuracy على tasks reasoning. الافتراض إن الموديل أكبر من 13B parameter — على موديلات أصغر الخسارة بتكبر.
• GQA: توفير 4x إلى 8x، خسارة 0.3% إلى 0.7% بس. الأكثر استخدامًا في 2026 وافتراضي في معظم موديلات SOTA.
• PagedAttention: توفير 60% إلى 80% في الـ multi-user serving، لكن في single-user inference بيضيف overhead بدون فايدة. بيحتاج vLLM (مش بيشتغل في raw transformers).
• KV Quantization: توفير 2x إلى 4x، بس latency إضافي 5% إلى 15% بسبب dequantization. ممكن يضرّ بـ long-context retrieval لو نزلت لـ INT4 بدون testing.

متى لا تستخدم هذه التقنيات

لو الـ context بتاعك أقل من 4K توكن وعندك GPU بـ 24GB، الـ KV Cache مش هيتعدى 1GB حتى على Llama 3 70B. المشكلة دي أصلاً غير موجودة عندك. ركّز في تحسينات تانية أهم زي batching، speculative decoding، أو quantization على الموديل نفسه.

كذلك لو بتشتغل على single-user inference (مش API serving)، الـ PagedAttention بيضيف complexity مش لازم. استخدم vanilla generation أو llama.cpp.

لو الموديل بتاعك صغير (أقل من 7B parameter)، KV quantization لـ INT4 ممكن يضرّ الجودة بشكل ملحوظ — اختبر دايمًا قبل ما تنزل لـ INT4.

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

لو بتشغّل LLM في إنتاج بـ contexts طويلة:

1. قيس الاستهلاك الحالي بالكود في الأعلى — اعرف هل الـ KV Cache فعلاً المشكلة ولا حاجة تانية.
2. لو الـ Cache بيتعدى 30% من ذاكرة GPU، انقل لـ vLLM مع --enable-prefix-caching. ده الحل الأسرع تطبيقًا.
3. لو لسه مش كافي، فعّل KV quantization بـ INT8 أولًا (متجربش INT4 قبل ما تختبر الجودة على dataset حقيقي عندك).
4. لو الموديل عندك بـ MHA كاملة، فكر في الـ migrate لموديل بـ GQA (Llama 3 70B، Qwen 2.5 72B، Mixtral). الانتقال ده غالبًا بيحل المشكلة بدون تعقيد إضافي.

المصادر

• Ainslie et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" — EMNLP 2023.
• Shazeer "Fast Transformer Decoding: One Write-Head is All You Need" — ورقة MQA الأصلية، 2019.
• Kwon et al. "Efficient Memory Management for Large Language Model Serving with PagedAttention" — SOSP 2023.
• Meta Llama 3 Model Card — التوثيق الرسمي لـ 8 KV heads على 70B.
• HuggingFace Transformers documentation: KV Cache Quantization guide.
• vLLM official documentation — docs.vllm.ai.
• Badri & Shaji "Half-Quadratic Quantization (HQQ)" — Mobius Labs technical report، 2024.