[GH-ISSUE #336] 是否支持像 Transformers 那样可以自由保存和载入 kv cache #215

New issue

Closed

opened 2026-10-29 14:35:58 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 14:35:58 +00:00

Owner

Originally created by @hucorz on GitHub (Feb 15, 2025).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/336

类似这里面做的事情：https://huggingface.co/docs/transformers/main/kv_cache

就是可以自由存取每次推理的 kvCache 并且推理时也可以载入

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache
from accelerate.test_utils.testing import get_backend

DEVICE, _, _ = get_backend() # automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map=DEVICE)
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_batch_size=1, max_cache_len=1024, device=DEVICE, dtype=torch.bfloat16)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(DEVICE)
# This is the common prompt cached, we need to run forward without grad to be abel 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(DEVICE)
    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)

在2台机器上各自因为架构是Turning以及系统的libstdcxx版本太低导致没有运行成功，所以先来问一问

Originally created by @hucorz on GitHub (Feb 15, 2025). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/336 类似这里面做的事情：https://huggingface.co/docs/transformers/main/kv_cache 就是可以自由存取每次推理的 kvCache 并且推理时也可以载入 ```python import copy import torch from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache from accelerate.test_utils.testing import get_backend DEVICE, _, _ = get_backend() # automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.) model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map=DEVICE) 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_batch_size=1, max_cache_len=1024, device=DEVICE, dtype=torch.bfloat16) INITIAL_PROMPT = "You are a helpful assistant. " inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(DEVICE) # This is the common prompt cached, we need to run forward without grad to be abel 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(DEVICE) 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) ``` > 在2台机器上各自因为架构是Turning以及系统的libstdcxx版本太低导致没有运行成功，所以先来问一问

bsakaguchi closed this issue

2026-10-29 14:35:58 +00:00

bsakaguchi commented

2026-10-29 14:35:59 +00:00

Author

Owner

@Azure-Tang commented on GitHub (Feb 16, 2025):

目前不支持的～

@Azure-Tang commented on GitHub (Feb 16, 2025): 目前不支持的～