DynamicCache的cache_position越界错误原因及无额外字符修复方案

场景描述 #

使用Hugging Face Transformers库的DynamicCache缓存初始提示词"I have a dream"，后续计划每次生成输出后追加到当前提示词，再执行生成任务。运行代码时触发IndexError: index -1 is out of bounds for dimension 0 with size 0，错误源于访问cache_position[-1]；但在初始提示词后添加空格，代码可正常运行。需明确问题根源及无需额外字符的解决方法。

复现代码 #

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

# 假设已初始化model和tokenizer
max_length = 6
prompt_cache = DynamicCache()
INITIAL_PROMPT = "I have a dream"

inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt", padding=True).to("cuda")
# 缓存初始提示词的KV
with torch.no_grad():
    prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values.to("cuda")

input_text = INITIAL_PROMPT
# input_text = INITIAL_PROMPT + " "  # 添加空格可正常运行

responses = []
for  _ in range(max_length):
    new_inputs = tokenizer(input_text, return_tensors="pt", padding=True).to("cuda")
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20,pad_token_id=tokenizer.eos_token_id, do_sample=True, temperature=0.7)
    output_text = tokenizer.decode(outputs[0])
    print(output_text[len(input_text):])
    input_text = output_text
    print("#"*24)

错误栈 #

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[61], line 20
     18 new_inputs = tokenizer(input_text, return_tensors="pt", padding=True).to("cuda")
     19 past_key_values = copy.deepcopy(prompt_cache)
---> 20 outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20,pad_token_id=tokenizer.eos_token_id, do_sample=True, temperature=0.7)
     21 output_text = tokenizer.decode(outputs[0])
     22 print(output_text[len(input_text):])

File /media/data1/haque/.conda/envs/llm/lib/python3.12/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File /media/data1/haque/.conda/envs/llm/lib/python3.12/site-packages/transformers/generation/utils.py:2252, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2244     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2245         input_ids=input_ids,
   2246         expand_size=generation_config.num_return_sequences,
   2247         is_encoder_decoder=self.config.is_encoder_decoder,
   2248         **model_kwargs,
   2249     )
   2251     # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
--> 2252     result = self._sample(
   2253         input_ids,
   2254         logits_processor=prepared_logits_processor,
   2255         stopping_criteria=prepared_stopping_criteria,
   2256         generation_config=generation_config,
   2257         synced_gpus=synced_gpus,
   2258         streamer=streamer,
   2259         **model_kwargs,
   2260     )
   2262 elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   2263     # 11. prepare beam search scorer
   2264     beam_scorer = BeamSearchScorer(
   2265         batch_size=batch_size,
   2266         num_beams=generation_config.num_beams,
   (...)
   2271         max_length=generation_config.max_length,
   2272     )

File /media/data1/haque/.conda/envs/llm/lib/python3.12/site-packages/transformers/generation/utils.py:3244, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   3239 is_prefill = True
   3240 while self._has_unfinished_sequences(
   3241     this_peer_finished, synced_gpus, device=input_ids.device, cur_len=cur_len, max_length=max_length
   3242 ):
   3243     # prepare model inputs
--> 3244     model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
   3246     # prepare variable output controls (note: some models won't accept all output controls)
   3247     model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})

File /media/data1/haque/.conda/envs/llm/lib/python3.12/site-packages/transformers/generation/utils.py:388, in GenerationMixin.prepare_inputs_for_generation(self, input_ids, past_key_values, attention_mask, inputs_embeds, cache_position, **kwargs)
    384 if past_key_values is not None:
    385     model_inputs["past_key_values"] = past_key_values
    386     if (
    387         inputs_embeds is not None  # Exception 1
--> 388         or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])  # Exception 3
    389     ):
    390         input_ids = input_ids[:, -cache_position.shape[0] :]
    391     elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)

IndexError: index -1 is out of bounds for dimension 0 with size 0

问题根源 #

• 核心矛盾：初始提示词处理后，DynamicCache的cache_position为空张量（size 0），但generate阶段尝试访问cache_position[-1]导致越界。
• 当直接使用初始提示词时，模型第一次forward处理完所有token后，cache_position没有被初始化对应长度的位置索引；而添加空格后，提示词的token数量变化，模型会正确生成匹配长度的cache_position，避免了空张量访问。
• 错误栈中代码cache_position[-1] >= input_ids.shape[1]是触发点：空张量无法通过索引-1获取元素，直接抛出越界错误。

解决方法 #

方法1：手动传入匹配长度的cache_position #

在调用generate时，根据初始提示词的token长度手动生成cache_position：

# 提前获取初始提示词的token长度
prompt_token_len = inputs_initial_prompt.input_ids.shape[1]

# 修改循环内的generate调用
outputs = model.generate(
    **new_inputs,
    past_key_values=past_key_values,
    cache_position=torch.arange(prompt_token_len, device="cuda"),  # 传入对应长度的位置索引
    max_new_tokens=20,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=True,
    temperature=0.7
)

方法2：保存并复用模型返回的cache_position #

第一次处理初始提示词时，同时保存模型返回的cache_position，后续generate时直接传入：

# 初始缓存阶段，同时保存cache_position
with torch.no_grad():
    outputs = model(**inputs_initial_prompt, past_key_values=prompt_cache)
    prompt_cache = outputs.past_key_values.to("cuda")
    prompt_cache_pos = outputs.cache_position.to("cuda")  # 保存cache_position

# 循环内传入保存的cache_position
outputs = model.generate(
    **new_inputs,
    past_key_values=past_key_values,
    cache_position=prompt_cache_pos,
    max_new_tokens=20,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=True,
    temperature=0.7
)

方法3：优化生成逻辑，避免重复编码完整提示词 #

更高效的方式是每次只传入新增的token，而非重新编码整个提示词，同时自动复用更新后的缓存：

# 初始处理
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model(**inputs_initial_prompt, past_key_values=prompt_cache)
    prompt_cache = outputs.past_key_values.to("cuda")

# 初始化当前输入为初始提示词的token
current_input_ids = inputs_initial_prompt.input_ids
current_past_cache = copy.deepcopy(prompt_cache)

responses = []
for _ in range(max_length):
    outputs = model.generate(
        input_ids=current_input_ids,
        past_key_values=current_past_cache,
        max_new_tokens=20,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.7
    )
    # 提取本次新增的token
    new_tokens = outputs[0][len(current_input_ids[0]):]
    output_text = tokenizer.decode(new_tokens)
    print(output_text)
    # 更新输入为新增token，缓存自动更新为最新状态
    current_input_ids = new_tokens.unsqueeze(0)
    current_past_cache = outputs.past_key_values
    print("#"*24)

内容的提问来源于stack exchange，提问作者melatonin15