Transformers documentation

Chiến lược bộ nhớ đệm

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Chiến lược bộ nhớ đệm

Các vector key-value (KV) được sử dụng để tính toán điểm chú ý (attention scores). Đối với các mô hình tự hồi quy (autoregressive), điểm KV được tính toán mỗi lần vì mô hình dự đoán từng token một. Mỗi lần dự đoán phụ thuộc vào các token trước đó, điều này có nghĩa là mô hình thực hiện cùng các phép tính mỗi lần.

Cache KV lưu trữ các phép tính này để có thể tái sử dụng mà không cần tính lại. Việc cache hiệu quả là rất quan trọng để tối ưu hóa hiệu suất mô hình vì nó giảm thời gian tính toán và cải thiện tốc độ phản hồi. Tham khảo tài liệu Caching để có giải thích chi tiết hơn về cách hoạt động của cache.

Transformers cung cấp một số lớp Cache triển khai các cơ chế lưu trữ bộ nhớ đệm khác nhau. Một số lớp Cache này được tối ưu hóa để tiết kiệm bộ nhớ trong khi những lớp khác được thiết kế để tối đa hóa tốc độ sinh. Hãy tham khảo bảng dưới đây để so sánh các loại bộ nhớ đệm và sử dụng nó để giúp bạn chọn bộ nhớ đệm phù hợp nhất cho trường hợp sử dụng của mình.

Loại bộ nhớ đệmHỗ trợ các lớp trượtHỗ trợ offloadingHỗ trợ torch.compile()Lượng bộ nhớ dự kiến
Bộ đệm độngCóCóKhôngTrung bình
Bộ đệm tĩnhCóCóCóCao
Bộ đệm lượng tử hóaKhôngKhôngKhôngThấp

Hướng dẫn này giới thiệu cho bạn các lớp Cache khác nhau và cho bạn thấy cách sử dụng chúng để tạo sinh.

Bộ nhớ đệm mặc định

DynamicCache là lớp cache mặc định cho tất cả các mô hình. Nó cho phép kích thước cache tăng lên một cách động để lưu trữ ngày càng nhiều key và value khi quá trình sinh tiến triển.

Lưu ý rằng đối với các mô hình sử dụng chú ý cửa sổ trượt (Mistral, Gemma2,…) hoặc chú ý phân mảnh (Llama4), bộ nhớ đệm sẽ ngừng tăng lên khi các lớp sử dụng các loại chú ý này đạt đến kích thước tối đa (kích thước cửa sổ trượt hoặc kích thước phân mảnh).

Tắt bộ nhớ đệm bằng cách cấu hình use_cache=False trong 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)

Các lớp Cache cũng có thể được khởi tạo trước khi gọi và truyền vào tham số past_key_values của các model. Điều này có thể hữu ích để kiểm soát chi tiết hơn, hoặc cho các trường hợp sử dụng nâng cao hơn như lưu trữ ngữ cảnh (context caching).

Trong hầu hết các trường hợp, sẽ dễ dàng hơn khi định nghĩa chiến lược cache trong tham số 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)

Bộ nhớ đệm có kích thước cố định

DynamicCache mặc định ngăn bạn tận dụng hầu hết các tối ưu hóa just-in-time (JIT) vì kích thước cache không cố định. Các tối ưu hóa JIT cho phép bạn giảm thiểu độ trễ với cái giá là mức sử dụng bộ nhớ. Tất cả các loại cache sau đây đều tương thích với các tối ưu hóa JIT như torch.compile để tăng tốc quá trình sinh.

Bộ đệm có kích thước cố định (StaticCache) phân bổ trước một kích thước bộ đệm tối đa cụ thể cho các cặp kv. Bạn có thể tạo ra số lượng token lên đến kích thước tối đa của bộ đệm mà không cần phải sửa đổi nó. Tuy nhiên, việc có kích thước cố định (thường là lớn) cho các trạng thái key/value có nghĩa là trong quá trình tạo, rất nhiều token thực tế sẽ bị che (masked) vì chúng không nên tham gia vào cơ chế chú ý. Vì vậy, thủ thuật này cho phép dễ dàng compile giai đoạn giải mã, nhưng nó gây ra sự lãng phí token trong tính toán chú ý. Như mọi thứ, đây là một sự đánh đổi, sẽ rất tốt nếu bạn tạo ra nhiều chuỗi có độ dài xấp xỉ nhau, nhưng có thể không tối ưu nếu bạn có, ví dụ, 1 chuỗi rất lớn và sau đó chỉ có các chuỗi ngắn (vì kích thước bộ đệm cố định sẽ lớn, một lượng lớn sẽ bị lãng phí cho các chuỗi ngắn). Hãy đảm bảo bạn hiểu rõ tác động nếu bạn sử dụng nó!

Đối với DynamicCache, lưu ý rằng đối với các mô hình sử dụng chú ý cửa sổ trượt (Mistral, Gemma2,…) hoặc chú ý phân mảnh (Llama4), bộ nhớ đệm sẽ không bao giờ lớn hơn kích thước cửa sổ trượt/phân mảnh trên các lớp sử dụng các loại chú ý này, ngay cả khi độ dài tối đa được chỉ định là lớn hơn.

Bạn có thể bật StaticCache bằng cách cấu hình cache_implementation="static" trong generate(). Thao tác này cũng sẽ tự động bật compilation cho giai đoạn giải mã đối với các chiến lược giải mã tham lam và lấy mẫu.

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"

Giữ các tensor sinh trên CPU

Các backend trình biên dịch như Neuron và TPU sẽ trace mô hình của bạn thành một đồ thị tính toán cố định. Vòng lặp sinh (generation loop) duy trì các tensor (chuỗi đầu ra, attention_mask, position_ids) tăng thêm một token ở mỗi bước. Trình biên dịch sẽ trace lại đồ thị mỗi khi các tensor này thay đổi hình dạng trên bộ tăng tốc, khiến quá trình sinh chậm đi.

generate() chỉ di chuyển các tensor mà forward sử dụng sang thiết bị của mô hình, ngay trước mỗi lần gọi forward. Kết quả đầu ra được di chuyển trở lại để khớp với thiết bị của đầu vào. Hãy truyền các đầu vào của bạn trên CPU để giữ việc theo dõi các tensor ngày càng tăng của vòng lặp nằm ngoài bộ tăng tốc. Đồ thị đã biên dịch vẫn ổn định, và vì đầu ra tuân theo thiết bị của đầu vào nên kết quả tạo ra vẫn nằm trên 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

Chuyển cache

Bộ nhớ đệm KV có thể chiếm một phần đáng kể bộ nhớ và trở thành điểm nghẽn cho việc tạo ngữ cảnh dài. Các bộ nhớ đệm hiệu quả về mặt bộ nhớ tập trung vào việc đánh đổi tốc độ để giảm mức sử dụng bộ nhớ. Điều này đặc biệt quan trọng đối với các mô hình ngôn ngữ lớn (LLM) và khi phần cứng của bạn bị giới hạn về bộ nhớ.

Việc offload bộ nhớ cache giúp tiết kiệm bộ nhớ GPU bằng cách chuyển KV cache của các lớp mô hình, ngoại trừ một lớp, sang CPU. Chỉ bộ nhớ cache của lớp hiện tại được duy trì trên GPU trong quá trình forward của mô hình khi lặp qua các lớp. Hệ thống sẽ prefetch bất đồng bộ bộ nhớ cache của lớp tiếp theo, và gửi bộ nhớ cache của lớp hiện tại trở lại CPU sau khi tính toán attention.

Bạn có thể cân nhắc offloading nếu bạn có một GPU nhỏ và đang gặp lỗi hết bộ nhớ (OOM).

Bạn có thể nhận thấy hiệu suất tạo (generation throughput) giảm nhẹ so với bộ nhớ đệm (cache) hoàn toàn trên thiết bị, tùy thuộc vào mô hình và các lựa chọn tạo của bạn (kích thước ngữ cảnh, số lượng token được tạo, số lượng beam, v.v.). Điều này là do việc di chuyển trạng thái key/value qua lại đòi hỏi một số công việc xử lý.

Offloading có sẵn cho cả DynamicCache và StaticCache. Bạn có thể bật tính năng này bằng cách cấu hình cache_implementation="offloaded" cho phiên bản động, hoặc cache_implementation="offloaded_static" cho phiên bản tĩnh, trong GenerationConfig hoặc generate(). Ngoài ra, bạn cũng có thể khởi tạo DynamicCache hoặc StaticCache của riêng mình với tùy chọn offloading=True, và truyền cache này vào generate hoặc forward của model (ví dụ, past_key_values=DynamicCache(config=model.config, offloading=True) cho cache động).

Lưu ý rằng hai lớp Cache được đề cập ở trên có một tùy chọn bổ sung khi khởi tạo trực tiếp, đó là offload_only_non_sliding. Tham số bổ sung này quyết định xem các lớp sử dụng chú ý cửa sổ trượt/chunk (nếu có) có bị offload hay không. Vì các lớp này thường ngắn ngay từ đầu, nên có thể tốt hơn là tránh offload chúng, vì offload có thể gây ra tổn thất về tốc độ. Mặc định, tùy chọn này là False đối với DynamicCache, và True đối với 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.

Ví dụ dưới đây cho thấy cách bạn có thể chuyển sang bộ nhớ đệm đã chuyển (offloaded cache) nếu hết bộ nhớ:

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)

Bộ nhớ đệm lượng tử hóa

QuantizedCache giảm yêu cầu về bộ nhớ bằng cách lượng tử hóa các giá trị KV xuống độ chính xác thấp hơn. QuantizedCache hiện hỗ trợ ba backend lượng tử hóa:

  • hqq hỗ trợ các kiểu dữ liệu int1, int2, int3, int4 và int8.
  • quanto hỗ trợ các kiểu dữ liệu int2 và int4. Đây là backend lượng tử hóa mặc định.
  • fp8 lưu các giá trị KV ở dạng số thực 8-bit (float8_e4m3fn). Chế độ này không yêu cầu phụ thuộc bổ sung nào và thực hiện lượng tử hóa theo tensor thay vì theo nhóm, do đó từ chối các tham số nbits, axis_key, axis_value, q_group_size và residual_length.

Việc lượng tử hóa cache có thể gây hại cho độ trễ nếu độ dài ngữ cảnh ngắn và có đủ bộ nhớ GPU để tạo ra nội dung mà không cần bật lượng tử hóa cache. Hãy cố gắng tìm sự cân bằng giữa hiệu quả bộ nhớ và độ trễ.

Bật QuantizedCache bằng cách cấu hình cache_implementation="quantized" trong GenerationConfig. Backend lượng tử hóa, cùng với các tham số lượng tử hóa bổ sung, được chỉ định trong dict cache_config được truyền vào generate(). Bạn nên sử dụng các giá trị mặc định cho các tham số bổ sung này trừ khi bạn đang gặp lỗi hết bộ nhớ. Trong trường hợp đó, hãy cân nhắc giảm độ dài phần dư.

Đối với backend hqq, chúng tôi khuyến nghị đặt các tham số axis_key và axis_value thành 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

Đối với backend quanto, chúng tôi khuyến nghị đặt các tham số axis_key và axis_value thành 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

Backend fp8 không nhận thêm tham số nào và sẽ gây ra lỗi nếu được truyền một trong các tham số theo nhóm.

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

Bộ nhớ đệm encoder-decoder

EncoderDecoderCache được thiết kế cho các mô hình encoder-decoder. Nó quản lý cả bộ nhớ đệm self-attention và cross-attention để đảm bảo việc lưu trữ và truy xuất các cặp kv trước đó. Bạn có thể thiết lập loại bộ nhớ đệm khác nhau cho riêng phần encoder và phần decoder.

Loại cache này không yêu cầu bất kỳ cấu hình nào. Đây là một lớp bọc đơn giản quanh 2 Cache như đã mô tả ở trên, sẽ được mô hình sử dụng độc lập trực tiếp.

Bộ nhớ đệm cụ thể cho từng mô hình

Một số mô hình có cách lưu trữ các cặp kv quá khứ hoặc trạng thái riêng biệt, không tương thích với bất kỳ lớp cache nào khác.

Các mô hình Mamba, chẳng hạn như Mamba, yêu cầu một bộ đệm cụ thể vì mô hình không có cơ chế chú ý hay trạng thái kv. Do đó, chúng không tương thích với các lớp Cache ở trên.

Tạo lặp

Cache cũng có thể hoạt động trong các cài đặt tạo lặp, nơi có sự tương tác qua lại với mô hình (chatbot). Giống như tạo thông thường, tạo lặp với cache cho phép mô hình xử lý hiệu quả các cuộc trò chuyện liên tục mà không cần tính toán lại toàn bộ ngữ cảnh ở mỗi bước.

Để tạo nội dung lặp lại với bộ nhớ đệm, hãy bắt đầu bằng cách khởi tạo một lớp bộ nhớ đệm trống, sau đó bạn có thể đưa vào các prompt mới. Theo dõi lịch sử hội thoại bằng chat template.

Ví dụ sau minh họa Llama-2-7b-chat-hf. Nếu bạn đang sử dụng một mô hình chat khác, apply_chat_template() có thể xử lý các tin nhắn theo cách khác nhau. Nó có thể cắt bỏ các token quan trọng tùy thuộc vào cách viết template Jinja. Đối với các mô hình chat đa phương thức, hãy xem hướng dẫn iterative chatting with cache để biết cách xử lý hình ảnh hoặc âm thanh.

Ví dụ, một số mô hình sử dụng các token đặc biệt <think> ... </think> trong quá trình suy luận. Những token này có thể bị mất trong quá trình mã hóa lại, gây ra các vấn đề về chỉ mục. Bạn có thể cần phải tự động xóa hoặc điều chỉnh các token thừa từ các phần hoàn thành để giữ cho hệ thống ổn định.

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})

Cắt một bộ nhớ đệm

Sử dụng crop() để hoàn tác các token từ bộ nhớ đệm, chẳng hạn khi một phần tiếp theo ứng viên bị từ chối trong quá trình tạo lặp. Truyền số lượng token cần xóa dưới dạng số nguyên âm.

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

Các lớp sliding-window và linear-attention sẽ loại bỏ các trạng thái quá khứ trong khi xử lý, nên mặc định không có gì để hoàn tác. Hãy gọi activate_past_recording() trước bước forward để các lớp này giữ lại những trạng thái đó. Nếu không, crop sẽ gây ra RuntimeError ngay khi một lớp sliding-window lấp đầy cửa sổ của nó, và ngay lập tức đối với các lớp linear-attention.

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

crop(0) không xóa bộ nhớ đệm. Nó không loại bỏ token nào khỏi chuỗi được lưu trong bộ nhớ đệm và là một thao tác không có tác dụng đối với các lớp full-attention. Đối với các lớp sliding-window và linear-attention, nó có thể loại bỏ thông tin đã được lưu trong bộ nhớ đệm mà lần gọi mô hình tiếp theo sẽ không sử dụng. Điều này giúp giảm mức sử dụng bộ nhớ mà không cần hoàn tác các token.

                    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
                            └───┘

Cả hai lớp giờ đều báo độ dài chuỗi là 3, và quá trình sinh tiếp tục từ vị trí 3.

Truyền 0 sẽ không loại bỏ token nào khỏi chuỗi. Đây là thao tác không làm gì (no-op) đối với các lớp full-attention, nhưng đối với các lớp sliding-window và linear-attention, nó vẫn loại bỏ các trạng thái đã ghi mà lần forward pass tiếp theo sẽ không đọc, qua đó giải phóng bộ nhớ mà không cần hoàn tác bất kỳ điều gì.

Điền trước bộ nhớ đệm (tiền đệm tiền tố)

Trong một số trường hợp, bạn có thể muốn điền một Cache với các cặp kv cho một prompt tiền tố nhất định và tái sử dụng nó để tạo ra các chuỗi khác nhau.

Ví dụ bên dưới khởi tạo một StaticCache, sau đó lưu vào bộ nhớ đệm một prompt ban đầu. Bây giờ bạn có thể tạo ra một vài chuỗi từ prompt đã được nạp trước.

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)
Update on GitHub