LLM inference with static kv-cache example gives different generations depending on the batch examples
Author: jpiabrantes
Pull request: no

### System Info

transformers.__version__ = 4.41.0.dev0
Python 3.10.12
unix

### Who can help?

@gante

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

When using the code documented [here](https://huggingface.co/docs/transformers/main/en/llm_optims?static-kv=Static+Cache#static-kv-cache-and-torchcompile), we get different generations for the same prompt when generating it in a different batch.

For example if you run the code bellow 1st generation for the prompt hey is different than the 2nd generation.

```python
import os
from transformers import LlamaTokenizer, LlamaForCausalLM
from transformers import StaticCache
import torch


def decode_one_tokens(model, cur_token, input_pos, cache_position, past_key_values):
    logits = model(
        cur_token,
        position_ids=input_pos,
        cache_position=cache_position,
        past_key_values=past_key_values,
        return_dict=False,
        use_cache=True
    )[0]
    new_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
    return new_token


def generate(prompts: list[str], model: LlamaForCausalLM, tokenizer: LlamaTokenizer, num_tokens_to_generate: int = 40) ->  list[str]:
    global decode_one_tokens
    inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
    batch_size, seq_length = inputs["input_ids"].shape
    with torch.no_grad():
        past_key_values = StaticCache(
            config=model.config, max_batch_size=batch_size, max_cache_len=4096, device=torch_device, dtype=model.dtype
        )
        cache_position = torch.arange(seq_length, device=torch_device)
        generated_ids = torch.zeros(
            batch_size, seq_length + num_tokens_to_generate + 1, dtype=torch.int, device=torch_device
        )
        generated_ids[:, cache_position] = inputs["input_ids"].to(torch_device).to(torch.int)
        with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
            logits = model(
                **inputs, cache_position=cache_position, past_key_values=past_key_values,return_dict=False, use_cache=True
            )[0]
        next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
        generated_ids[:, seq_length] = next_token[:, 0]
        
        # Not using torch.compile to simplify debugging
        # decode_one_tokens = torch.compile(decode_one_tokens, mode="reduce-overhead", fullgraph=True)
        cache_position = torch.tensor([seq_length + 1], device=torch_device)
        for _ in range(1, num_tokens_to_generate):
            with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
                next_token = decode_one_tokens(model, next_token.clone(), None, cache_position, past_key_values)
                generated_ids[:, cache_position] = next_token.int()
            cache_position += 1

    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)


if __name__ == "__main__":
    torch_device = "cuda"
    tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", pad_token="
", padding_side="right")
    model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16).to(torch_device)
    model.eval()
    print(generate(["hey", "yo"], model, tokenizer)[0])
    # prints: hey, i'm alex.\ni'm a 20-something year old photographer and filmmaker based in los angeles. i'm a lover of all things cre
    print('-'*50 + '\n'*2)
    print(generate(["hey", "yo code math"], model, tokenizer)[0])
    #prints: hey The 2018-19 season is the 10th season of the National Women's Soccer League (NWSL), the top division of women's soccer
    print('-'*50 + '\n'*2)
```

### Expected behavior

The generation for each prompt should not depend on the other examples in the batch.

Comments:
I think one problem is related with this line:
```python
next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
```

If the prompt is not the longest in the batch then logits[:, -1] is the logit corresponding to a padding token instead of the last token. This can be fixed with left padding or with the following code:

```python
last_token = inputs['attention_mask'].sum(-1) - 1
idx = torch.arange(batch_size).to(model.device)
next_token = torch.argmax(logits[idx, last_token], dim=-1)[:, None]
```

This indeed makes the first sampled token to be the same, however, there are still problems after the first token. I think it has something to do with `cache_position`?

@jpiabrantes hey!

The model in the second cases is generating gibberish because it did not receive attention mask as the input. In other words, you have to pass in the attention mask in the `decode_one_token` function so that the model knows to ignore padded inputs.

Something like below would work:

```python
def decode_one_tokens(model, cur_token, input_pos, attention_mask, cache_position, past_key_values):
    logits = model(
        cur_token,
        attention_mask=attention_mask,
        position_ids=input_pos,
        cache_position=cache_position,
        past_key_values=past_key_values,
        return_dict=False,
        use_cache=True
    )[0]
    new_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
    return new_token


def generate(prompts: List[str], model: LlamaForCausalLM, tokenizer: LlamaTokenizer, num_tokens_to_generate: int = 40) ->  List[str]:
    global decode_one_tokens
    inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
    batch_size, seq_length = inputs["input_ids"].shape

    attention_mask = inputs["attention_mask"]
    with torch.no_grad():
        past_key_values = StaticCache(
            config=model.config, max_batch_size=batch_size, max_cache_len=4096, device=torch_device, dtype=model.dtype
        )
        cache_position = torch.arange(seq_length, device=torch_device)
        generated_ids = torch.zeros(
            batch_size, seq_length + num_tokens_to_generate + 1, dtype=torch.int, device=torch_device
        )
        generated_ids[:, cache_position] = inputs["input_ids"].to(torch_device).to(torch.int)
        with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
            logits = model(
                **inputs, cache_position=cache_position, past_key_values=past_key_values, return_dict=False, use_cache=True
            )[0]
        next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
        generated_ids[:, seq_length] = next_token[:, 0]

        # Not using torch.compile to simplify debugging
        # decode_one_tokens = torch.compile(decode_one_tokens, mode="reduce-overhead", fullgraph=True)
        cache_position = torch.tensor([seq_length + 1], device=torch_device)
        attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) # update and pass every step
        for _ in range(1, num_tokens_to_generate):
            with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
                next_token = decode_one_tokens(model, next_token.clone(), None, attention_mask, cache_position, past_key_values)
                generated_ids[:, cache_position] = next_token.int()
            cache_position += 1
            position_ids = position_ids[:, -1:] + 1
            attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)

    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

```

@zucchini-nlp were you able to make your code work? Note that `decode_one_tokens` receives only one token at a time and should never receive a padded token.

@jpiabrantes yes, the suggested code works as expected as the results match the dynamic cache generation :)

We have to use the full attn mask because in every attention layer we concatenate previous keys/values with the current one-token key/value, thus obtaining the full attn matrix for the whole sequence length. So we have to apply the correct mask to get correct scores 

Thank you!

`attention_mask` together with left padding and everything is working :)

not sure if `torch.compile` will work well with an input that keeps changing its size but I will leave that for now.

@jpiabrantes 

> not sure if torch.compile will work well with an input that keeps changing its size but I will leave that for now.

Correct, prefill needs special treatment to be efficient with `torch.compile` (which we are not doing at the moment in `generate`):
1. Ideally, all inputs at prefill time have the same length 👉 no problem, constant shapes
2. We left-pad to a limited set of values (`pad_to_multiple_of` flag of the tokenizer), to ensure there are few possible shapes
3. We don't compile at all

@gante thank you!

Have you been able to generate with `torch.compile` when batch_size > 1?

@jpiabrantes yeah, with no issues

```py
# `torch.compile`-enabled Llama 3
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch, time, os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", padding_side="left")
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct", device_map="auto", torch_dtype=torch.float16
)
model.forward = torch.compile(model.forward, fullgraph=True, mode="reduce-overhead")

# The first iteration is slow (compilation overhead). Subsequent iterations are
# faster than the uncompiled call.
inputs = tokenizer(["The quick brown", "Foo bar"], return_tensors="pt", padding=True).to(model.device)
for i in range(10):
    start = time.time()
    gen_out = model.generate(
        **inputs, do_sample=False, cache_implementation="static", max_new_tokens=100
    )
    print(f"Time taken: {time.time() - start:.2f}s")
print(tokenizer.batch_decode(gen_out, skip_special_tokens=True))
```

I'm re-opening this issue because I can't get static cache to work with `torch.compile` as described in the [documentation](https://huggingface.co/docs/transformers/main/en/llm_optims?static-kv=Static+Cache#static-kv-cache-and-torchcompile) when batch_size > 1.

@gante when I run the code of your last comment I get the following error:
<details>

```sh
Traceback (most recent call last):
  File "/home/joaoabrantis/rosetta/tmp3.py", line 20, in <module>
    gen_out = model.generate(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/transformers/generation/utils.py", line 1785, in generate
    result = self._sample(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/transformers/generation/utils.py", line 2424, in _sample
    outputs = self(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
    return fn(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 921, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
    return _compile(
  File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
    transformations(instructions, code_options)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
    return fn(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
    tracer.run()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
    super().run()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
    and self.step()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
    getattr(self, inst.opname)(inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper
    return inner_fn(self, inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1260, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 674, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 94, in realize_and_forward
    return getattr(self.realize(), name)(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 759, in call_function
    return self.func.call_function(tx, merged_args, merged_kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 289, in call_function
    return super().call_function(tx, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 680, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2285, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2399, in inline_call_
    tracer.run()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
    and self.step()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
    getattr(self, inst.opname)(inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper
    return inner_fn(self, inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1260, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 674, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 335, in call_function
    return super().call_function(tx, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 289, in call_function
    return super().call_function(tx, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 680, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2285, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2399, in inline_call_
    tracer.run()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
    and self.step()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
    getattr(self, inst.opname)(inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper
    return inner_fn(self, inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1272, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 674, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 289, in call_function
    return super().call_function(tx, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 680, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2285, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2399, in inline_call_
    tracer.run()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
    and self.step()
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
    getattr(self, inst.opname)(inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper
    return inner_fn(self, inst)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1219, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 674, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 687, in call_function
    result = handler(tx, *args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 1160, in call_hasattr
    return obj.call_hasattr(tx, name)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 344, in call_hasattr
    unimplemented(f"hasattr {self.__class__.__name__} {name}")
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: hasattr ConstDictVariable to

from user code:
   File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
    return fn(*args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/accelerate/hooks.py", line 161, in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/accelerate/hooks.py", line 356, in pre_forward
    return send_to_device(args, self.execution_device), send_to_device(
  File "/home/joaoabrantis/rosetta/.rosetta/lib/python3.10/site-packages/accelerate/utils/operations.py", line 148, in send_to_device
    if is_torch_tensor(tensor) or hasattr(tensor, "to"):

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

</details>

The version of transformers I have installed: 4.42.0.dev0

@jpiabrantes 👀 I can indeed reproduce your issue on colab ([link](https://colab.research.google.com/drive/1DAKrXuU3NrHL_DS3SoRDspXHL4Ow6R8k?usp=sharing)), but running locally it works fine.

In both Colab and local machine:
- accelerate 0.30.1
- transformers@5c88253556b7f15cf7d7e9793d7b2a39b4aa588a (`4.42.0.dev0`)
- script as pasted above

Differences:
- local has `torch==2.3.0`, colab has a specific 2.3.0 wheel (both are built for the same cuda version)
- local has an RTX3090, colab has a T4 -- different GPU architectures and VRAM

👉 my local run doesn't hit the accelerate lines in the stack trace. Looking in `accelerate/hooks.py` line 356, [we can see that is triggered when offloading is active](https://github.com/huggingface/accelerate/blob/b24a0ef5dbb4e0408074f92f84683608eab474af/src/accelerate/hooks.py#L318) [i.e. when model size > GPU size], which is enabled by the `device_map="auto"` flag in the `.from_pretrained()` call. In other words, `torch.compile` seems to be limited to the case where the model doesn't need to be offloaded to the CPU :) For Llama3-8B in FP16, that is a GPU with >16GB VRAM.

Were you running on a GPU with >16GB VRAM? If yes, we should dig deeper. If not, the mystery is solved!


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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.