Static Cache was moved to be a standalone object in #30476. You have to init StaticCache outside the model and pass in every forward call, similar to following:

past_key_values = StaticCache(model.config, bs, max_cache_length, model.device, model.dtype)
for i in range(max_new_tokens):
    out = model(next_token, past_key_values=past_key_values, return_dict=True, **model_kwargs)
    past_key_values = out.past_key_values
    next_token = sample_next_token(out)
    model_kwargs = update_model_kwargs(model_kwargs)

Also you can simply pass-in cache_implementation="static" to generate() (see docs):

model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
model.generation_config.max_new_tokens = 20
out = model.generate(input_ids, cache_implementation="static)

Static Cache was moved to be a standalone object in #30476. You have to init StaticCache outside the model and pass in every forward call, similar to following:

Thanks, will try that!

model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
model.generation_config.max_new_tokens = 20
out = model.generate(input_ids, cache_implementation="static)

That's not the same thing. This will compile the whole forward pass, which will force re-compilation every time the prompt length changes. I only want to compile the decoding part and leave the prefill phase uncompiled.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.