استراتيجيات التخزين المؤقت

تُستخدم متجهات المفتاح-القيمة (KV) لحساب درجات الانتباه. بالنسبة للنماذج التوليدية الذاتية، تُحسب درجات KV في كل مرة لأن النموذج يتنبأ برمز واحد في كل مرة. يعتمد كل تنبؤ على الرموز السابقة، مما يعني أن النموذج يؤدي نفس العمليات الحسابية في كل مرة.

تُخزّن ذاكرة cache هذه الحسابات حتى يمكن إعادة استخدامها دون الحاجة إلى إعادة حسابها. يُعدّ التخزين المؤقت الفعّال أمراً بالغ الأهمية لتحسين أداء النموذج، لأنه يقلّل من زمن الحساب ويزيد من معدلات الاستجابة. للاطلاع على شرح أكثر تفصيلاً حول آلية عمل ذاكرة التخزين المؤقت، راجع وثيقة Caching.

توفر مكتبة Transformers عدة فئات Cache تنفّذ آليات تخزين مؤقت مختلفة. بعض فئات Cache هذه مُحسّنة لتوفير الذاكرة، بينما صُمّمت فئات أخرى لتعظيم سرعة التوليد. راجع الجدول أدناه لمقارنة أنواع التخزين المؤقت، واستخدمه للمساعدة في اختيار أفضل تخزين مؤقت يناسب حالتك الاستخدامية.

يُعرّفك هذا الدليل على فئات Cache المختلفة ويوضح لك كيفية استخدامها للتوليد.

الذاكرة المؤقتة الافتراضية

يُعدّ DynamicCache فئة التخزين المؤقت الافتراضية لجميع النماذج. فهو يتيح للتخزين المؤقت أن ينمو ديناميكيًا من أجل تخزين عدد متزايد من المفاتيح والقيم مع تقدم عملية التوليد.

لاحظ أنه في النماذج التي تستخدم انتباه النافذة المنزلقة (Mistral, Gemma2,…) أو الانتباه المقطعي (Llama4)، سيتوقف نمو الكاش عندما تصل الطبقات التي تستخدم هذه الأنواع من الانتباه إلى حجمها الأقصى (حجم النافذة المنزلقة أو حجم المقطع).

عطّل ذاكرة التخزين المؤقت من خلال ضبط use_cache=False في generate().

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

يمكن أيضًا تهيئة فئات الكاش (Cache) أولاً قبل استدعاء النموذج ومرورها إلى معامل past_key_values. وقد يكون ذلك مفيدًا للحصول على تحكم أدق، أو لاستخدامات متقدمة مثل كاش السياق.

في معظم الحالات، يكون من الأسهل تحديد استراتيجية التخزين المؤقت في معامل cache_implementation.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache(config=model.config)
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

ذاكرة مؤقتة ذات حجم ثابت

يمنعك DynamicCache الافتراضي من الاستفادة من معظم تحسينات الترجمة الفورية (JIT) لأن حجم ذاكرة التخزين المؤقت غير ثابت. تتيح لك تحسينات JIT تقليل زمن الاستجابة على حساب استهلاك الذاكرة. جميع أنواع ذاكرة التخزين المؤقت التالية متوافقة مع تحسينات JIT مثل torch.compile لتسريع عملية التوليد.

يخصص ذاكرة مؤقتة بحجم ثابت (StaticCache) حداً أقصى محدداً مسبقاً لحجم ذاكرة المؤقتة لزوجات kv. يمكنك توليد ما يصل إلى الحد الأقصى لحجم ذاكرة المؤقتة دون الحاجة إلى تعديله. ومع ذلك، فإن وجود حجم ثابت (عادةً كبير) لحالات المفتاح/القيمة يعني أنه أثناء التوليد، سيتم إخفاء عدد كبير من الرموز فعلياً لأنها لا ينبغي أن تشارك في الانتباه. لذا، تتيح هذه الحيلة بسهولة compile مرحلة فك الترميز، لكنها تسبب هدرًا للرموز في حساب الانتباه. وككل الأمور، فإن الأمر يتعلق بمفاضلة، وهي ستكون جيدة جداً إذا كنت تولد عدة تسلسلات بطول متقارب إلى حد ما، لكنها قد تكون غير مثالية إذا كان لديك على سبيل المثال تسلسل واحد كبير جداً، ثم تسلسلات قصيرة فقط (لأن حجم ذاكرة المؤقتة الثابتة سيكون كبيراً، فسيُهدر الكثير للتسلسلات القصيرة). تأكد من فهم الأثر إذا كنت تستخدمها!

أما بالنسبة إلى DynamicCache، فانتبه إلى أنه في النماذج التي تستخدم انتباه النافذة المنزلقة (Mistral، Gemma2،…) أو الانتباه المقطعي (Llama4)، لن يتجاوز حجم الكاش حجم النافذة المنزلقة/القِطع في الطبقات التي تستخدم هذه الأنواع من الانتباه، حتى لو كان الطول الأقصى المحدد أكبر.

يمكنك تفعيل StaticCache من خلال ضبط cache_implementation="static" في generate(). كما سيؤدي ذلك إلى تفعيل compilation التلقائي لمراحل فك الترميز لاستراتيجيات فك الترميز الجشعة والعشوائية.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

إبقاء تينسورات التوليد على وحدة المعالجة المركزية

تتتبع وحدات المعالجة الخلفية للمترجمات مثل Neuron وTPU نموذجك لتحويله إلى رسم حسابي ثابت. يحافظ حلقة التوليد على متجهات (output sequence, attention_mask, position_ids) تنمو بمقدار رمز واحد في كل خطوة. يعيد المترجم تتبع الرسم كل مرة تتغير فيها أبعاد هذه المتجهات على المعالج التسارعي، مما يُبطئ عملية التوليد.

تنقل generate() فقط الموترات التي يستهلكها forward إلى جهاز النموذج، مباشرةً قبل كل استدعاء لـforward. ثم يُعاد نقل المخرجات إلى جهاز المدخلات لمطابقته. مرّر مدخلاتك على وحدة المعالجة المركزية (CPU) للحفاظ على سجل الموترات المتنامي في الحلقة بعيدًا عن معالج التسريع. تبقى المخطط المُجمَّع مستقرًا، وبما أن المخرجات تتبع جهاز المدخلات، فإن المخرجات المولّدة تبقى على وحدة المعالجة المركزية (CPU).

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="auto")

# leave the inputs on CPU instead of calling .to(model.device).
inputs = tokenizer("The French Bread Law states", return_tensors="pt")

# generate runs forward on the model device but returns the output on the input_ids device.
output = model.generate(**inputs, do_sample=False, max_new_tokens=20)
print(output.device)
cpu

إخلاء ذاكرة التخزين المؤقت

يمكن أن يشغل ذاكرة KV حصةً كبيرة من الذاكرة، ويصبح عنقاً زجاجة في توليد السياقات الطويلة. تركز الذاكرات الفعّالة من حيث استهلاك الذاكرة على الموازنة بين السرعة وتقليل استهلاك الذاكرة. وهذا أمرٌ مهمّ بشكل خاص بالنسبة لنماذج اللغة الكبيرة (LLMs) وفي حال كانت موارد الذاكرة في جهازك محدودة.

يوفر نقل ذاكرة التخزين المؤقت (Offloading) مساحة ذاكرة وحدة معالجة الرسوميات (GPU) من خلال نقل ذاكرة KV cache الخاصة بطبقات النموذج، باستثناء طبقة واحدة، إلى وحدة المعالجة المركزية (CPU). أثناء تكرار forward للنموذج عبر الطبقات، تُحافظ على ذاكرة التخزين المؤقت للطبقة الحالية على وحدة معالجة الرسوميات (GPU) فقط. كما ستقوم بجلب ذاكرة التخزين المؤقت للطبقة التالية مسبقًا بشكل غير متزامن، وإرجاع ذاكرة التخزين المؤقت للطبقة الحالية إلى وحدة المعالجة المركزية (CPU) بعد حساب الانتباه (attention).

قد ترغب في النظر في خيار النقل (Offloading) إذا كان لديك بطاقة GPU صغيرة وتواجه أخطاء نفاد الذاكرة (OOM).

قد تلاحظ انخفاضًا طفيفًا في معدل إنتاج النص مقارنةً بذاكرة التخزين المؤقت الكاملة على الجهاز، وذلك اعتمادًا على النموذج وخيارات التوليد (حجم السياق، عدد الرموز المولّدة، عدد الحزم، إلخ). والسبب في ذلك أن نقل حالات المفتاح/القيمة ذهابًا وإيابًا يتطلب بعض المعالجة.

يتوفر خيار التفريغ (Offloading) لكلٍّ من DynamicCache وStaticCache. يمكنك تفعيله من خلال ضبط cache_implementation="offloaded" للإصدار الديناميكي، أو cache_implementation="offloaded_static" للإصدار الثابت، وذلك في GenerationConfig أو generate(). علاوة على ذلك، يمكنك أيضًا إنشاء كائن خاص بك من DynamicCache أو StaticCache باستخدام الخيار offloading=True، ثم تمرير هذا الكاش في generate أو في دالة forward الخاصة بنموذجك (على سبيل المثال، past_key_values=DynamicCache(config=model.config, offloading=True) لكاش ديناميكي).

لاحظ أن صفحتي Cache المذكورتين أعلاه تحتويان على خيار إضافي عند إنشاء كائناتهما مباشرة، وهو offload_only_non_sliding. يحدد هذا الحجة الإضافية ما إذا كانت الطبقات التي تستخدم الانتباه النافذة المنزلقة/المجزأة (إن وُجدت) ستُحمَّل أيضاً. وبما أن هذه الطبقات قصيرة عادةً، فقد يكون من الأفضل تجنب تحميلها، لأن التحميل قد يسبب عقوبة في السرعة. افتراضياً، يكون هذا الخيار False لـ DynamicCache، وTrue لـ StaticCache.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

يُظهر المثال أدناه كيفية الرجوع إلى ذاكرة تخزين مؤقت مُحمّلة خارجياً إذا نفدت الذاكرة:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

def resilient_generate(model, *args, **kwargs):
    oom = False
    device = Accelerator().device
    torch_device_module = getattr(torch, device, torch.cuda)
    try:
        return model.generate(*args, **kwargs)
    except torch.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch_device_module.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_return_sequences": 20, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

ذاكرة مؤقتة مُكمّاة

يقلل QuantizedCache من متطلبات الذاكرة من خلال تكميم قيم KV إلى دقة أدنى. يدعم QuantizedCache حاليًا ثلاثة محركات خلفية للتكميم:

• يدعم hqq أنواع البيانات int1 وint2 وint3 وint4 وint8.
• يدعم quanto أنواع البيانات int2 و int4. وهو واجهة التكميم الافتراضية.
• fp8 يخزّن قيم KV بنقطة عائمة 8-بت (float8_e4m3fn). لا يتطلب أي اعتمادية إضافية، ويقوم بتكميم على مستوى المصفوفة (per-tensor) بدلاً من تكميم على مستوى المجموعة (per-group)، لذا يرفض nbits وaxis_key وaxis_value وq_group_size وresidual_length.

قد يُسبّب تكميم ذاكرة التخزين المؤقت تدهور زمن الاستجابة إذا كانت طول السياق قصيرًا وكان هناك ذاكرة GPU كافية للتوليد دون تفعيل تكميم ذاكرة التخزين المؤقت. حاول إيجاد توازن بين كفاءة استخدام الذاكرة وزمن الاستجابة.

فعّل QuantizedCache من خلال ضبط cache_implementation="quantized" في GenerationConfig. يتم تحديد خلفية التكميم، إلى جانب أي معلمات تكميم إضافية، في قاموس cache_config المُمرَّر إلى generate(). ينبغي استخدام القيم الافتراضية لهذه المعلمات الإضافية ما لم تكن تواجه مشكلة نفاد الذاكرة. في تلك الحالة، فكّر في تقليل طول المتبقي.

لواجهة hqq، ننصح بإعداد المعاملين axis_key وaxis_value على القيمة 1.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "hqq",
        "axis_key": 1,
        "axis_value": 1
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

نوصي، في حالة استخدام واجهة quanto، بإعداد المعاملين axis_key وaxis_value على القيمة 0.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "quanto",
        "axis_key": 0,
        "axis_value": 0
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

لا يقبل الواجهة الخلفية fp8 أي معامل إضافي، وتُطلق استثناءً إذا تم تمرير أحد المعاملات الخاصة بالمجموعات.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy

ذاكرة مؤقت للمشفّر-المفكّك

صُمم EncoderDecoderCache لنماذج المُشفِّر-المُفكِّك. وهو يدير كاشات الانتباه الذاتي والانتباه المتقاطع معًا لضمان تخزين أزواج kv السابقة واسترجاعها. ومن الممكن تعيين نوع كاش مختلف بشكل مستقل لكل من المُشفِّر والمُفكِّك.

لا تتطلب هذه الفئة من الكاش أي إعداد مسبق. فهي غلاف بسيط حول كاشين من نوع Cache كما هو موضح أعلاه، وسيستخدمهما النموذج بشكل مستقل ومباشر.

الذاكرة المؤقتة الخاصة بالنموذج

تتميز بعض النماذج بطريقة فريدة لتخزين أزواج kv السابقة أو الحالات، وهي طريقة غير متوافقة مع أي من فئات التخزين المؤقت الأخرى.

تتطلب نماذج Mamba، مثل Mamba، ذاكرة مؤقتة محددة لأن النموذج لا يحتوي على آلية انتباه أو حالات kv. وبالتالي، فهي غير متوافقة مع فئات Cache المذكورة أعلاه.

التوليد التكراري

يمكن أن يعمل الكاش أيضًا في إعدادات التوليد التكراري التي تتضمن تفاعلًا متبادلًا مع النموذج (روبوتات المحادثة). وبالمثل للتوليد العادي، يتيح التوليد التكراري باستخدام الكاش للنموذج معالجة المحادثات المستمرة بكفاءة دون إعادة حساب السياق الكامل في كل خطوة.

للتوليد التكراري مع استخدام ذاكرة مؤقتة، ابدأ بتهيئة فئة ذاكرة مؤقتة فارغة، ثم يمكنك إدخال الموجهات الجديدة. حافظ على سجل المحادثة باستخدام قالب المحادثة.

يُوضّح المثال التالي Llama-2-7b-chat-hf. إذا كنت تستخدم نموذجاً آخر بأسلوب المحادثة، فقد يعالج apply_chat_template() الرسائل بشكل مختلف. وقد يحذف رموزاً مهمة اعتماداً على كيفية كتابة قالب Jinja. بالنسبة لنماذج المحادثة متعددة الوسائط، راجع دليل iterative chatting with cache لمعرفة كيفية معالجة الصور أو الصوت.

على سبيل المثال، تستخدم بعض النماذج رموزًا خاصة <think> ... </think> أثناء الاستدلال. وقد تُفقد هذه الرموز أثناء إعادة الترميز، مما يسبب مشكلات في الفهرسة. قد تحتاج إلى إزالة أو تعديل الرموز الإضافية من الإكمالات يدويًا للحفاظ على استقرار الأمور.

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache(config=model.config)

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

قصّ ذاكرة التخزين المؤقت

استخدم crop() للتراجع عن الرموز من ذاكرة التخزين المؤقت، وذلك عند رفض استمرارية مرشحة أثناء التوليد التكراري. مرّر عدد الرموز المراد إزالتها كعدد صحيح سالب.

# Remove the last three tokens from the cache.
past_key_values.crop(-3)

تتخلّى طبقات الانتباه الخطي وطبقات نافذة الانزلاق عن الحالات السابقة أثناء المعالجة، لذا لا يوجد ما يمكن التراجع إليه افتراضيًا. استدعِ activate_past_recording() قبل تمرير الأمامي لجعلها تحتفظ بتلك الحالات. بدونه، يرمي crop استثناء RuntimeError بمجرد امتلاء نافذة طبقة نافذة الانزلاق، وبشكل فوري لطبقات الانتباه الخطي.

# Enable this before the forward pass if you need to roll back tokens.
past_key_values.activate_past_recording()

crop(0) لا يمسح ذاكرة التخزين المؤقت. فهو لا يزيل أي رموز من التسلسل المخزّن، وهو إجراء لا يؤثر على طبقات الانتباه الكامل. أما في طبقات النافذة المنزلقة والانتباه الخطي، فقد يتخلّى عن معلومات مخزّنة لن يستخدمها استدعاء النموذج التالي. وهذا يقلّل من استهلاك الذاكرة دون التراجع عن الرموز.

                    cached positions
                    ┌───┬───┬───┬───┬───┬───┐
before crop(-3)     │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │
                    └───┴───┴───┴───┴───┴───┘
                                ╰─────┬─────╯
                                 rolled back

after crop(-3)
                    ┌───┬───┬───┐
full attention      │ 0 │ 1 │ 2 │    keeps the whole prefix
                    └───┴───┴───┘
                            ┌───┐
sliding_window=2            │ 2 │    keeps sliding_window - 1
                            └───┘

تُبلّغ كلا الطبقتين الآن عن طول تسلسل يساوي 3، وتُستأنف عملية التوليد من الموضع 3.

إن تمرير 0 لا يُزيل أي رموز من التسلسل. فهو إجراء لا يؤثر على طبقات الانتباه الكامل، لكن بالنسبة لطبقات النافذة المنزلقة والانتباه الخطي، فإنه لا يزال يتخلص من الحالات المسجلة التي لن تقرأها عملية التمرير الأمامي التالية، مما يُحرر الذاكرة دون التراجع عن أي شيء.

تعبئة ذاكرة مؤقتة مسبقاً (تخزين بادئات مؤقت)

في بعض الحالات، قد ترغب في ملء Cache بأزواج kv لبادئة معينة وإعادة استخدامه لتوليد تسلسلات مختلفة.

يُهيّئ المثال أدناه StaticCache، ثم يخزّن مبدئيًا موقفاً أولياً. والآن يمكنك توليد عدة تسلسلات من الموقف المُعبّأ مسبقًا.

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map={"": 0})
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_cache_len=1024)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(model.device.type)
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(model.device.type)
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)