[bug] Lfm2 issue in second prefill stage when asking followup question
Author: joeldushouyu
Pull request: no
Labels: name: bug

name: Code agent slop

### System Info


```bash

- `transformers` version: 4.57.1
- Platform: Linux-6.14.0-33-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.11.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.0+cu128 (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
```




### Who can help?

@ArthurZucker @Cyrilvallez

### Information

- [ ] The official example scripts
- [x] 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

Reproduction: Running the following python script

```python3
import torch
from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List, Optional, Tuple, Dict, Any

# Suppress a specific UserWarning from transformers about custom models
from warnings import filterwarnings
filterwarnings("ignore", message="You are using a model of type lfm2 to instantiate a model of type .*", category=UserWarning)


def get_eos_ids(model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> List[int]:
    """Collects all known EOS token IDs from the model config and tokenizer."""
    eos_ids = set()
    
    # From model config
    cfg_eos = getattr(model.config, "eos_token_id", None)
    if isinstance(cfg_eos, (list, tuple)):
        eos_ids.update(int(x) for x in cfg_eos)
    elif isinstance(cfg_eos, int):
        eos_ids.add(cfg_eos)
        
    # From tokenizer
    tok_eos = getattr(tokenizer, "eos_token_id", None)
    if isinstance(tok_eos, (list, tuple)):
        eos_ids.update(int(x) for x in tok_eos)
    elif isinstance(tok_eos, int):
        eos_ids.add(tok_eos)
        
    return list(eos_ids)


def generate_response(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    past_key_values: Optional[Tuple] = None,
    eos_ids: Optional[List[int]] = None,
    max_new_tokens: int = 512
) -> Tuple[torch.Tensor, torch.Tensor, Tuple, List[int]]:
    """
    Performs a manual greedy decoding loop, reusing and returning the KV cache.
    
    Args:
        model: The causal LM.
        input_ids: The *full* sequence of token IDs generated so far.
        attention_mask: The *full* attention mask for the entire sequence.
        past_key_values: The KV cache from the previous step, if any.
        eos_ids: A list of token IDs that signal the end of a sequence.
        max_new_tokens: The maximum number of new tokens to generate.

    Returns:
        A tuple containing:
        - final_sequence: The complete generated sequence (inputs + new tokens).
        - final_mask: The updated full attention mask.
        - final_cache: The updated KV cache.
        - new_token_ids: A list of *only* the new token IDs that were generated.
    """
    
    generated_sequence = input_ids
    current_mask = attention_mask
    current_cache = past_key_values
    
    new_token_ids = []  # Track only the newly generated tokens

    with torch.no_grad():
        for _ in range(max_new_tokens):
            
            # Determine what to pass to the model:
            # 1. If no cache (first pass), send the whole prompt.
            # 2. If cache exists, send only the *last* token.
            if current_cache is None:
                model_input_ids = generated_sequence
            else:
                model_input_ids = generated_sequence[:, -1:]

            # Prepare model inputs
            model_inputs = {
                "input_ids": model_input_ids,
                "attention_mask": current_mask,
                "use_cache": True,
                "past_key_values": current_cache,
            }

            outputs = model(**model_inputs)

            # Greedy: pick argmax of last-token logits
            next_token_logits = outputs.logits[:, -1, :]
            next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            # Append the new token to our trackers
            generated_sequence = torch.cat([generated_sequence, next_token_id], dim=-1)
            current_mask = torch.cat([current_mask, torch.ones_like(next_token_id)], dim=-1)
            current_cache = outputs.past_key_values

            new_token_int = int(next_token_id.item())
            new_token_ids.append(new_token_int)

            # Stop at EOS
            if eos_ids and new_token_int in eos_ids:
                print("[EOS reached]")
                break
                
    return generated_sequence, current_mask, current_cache, new_token_ids


def main():
    # Settings
    model_id = "LiquidAI/LFM2-1.2B"
    max_new_tokens = 512
    
    # --- 1. Load Model and Tokenizer ---
    print(f"Loading model: {model_id}...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    #    attn_implementation="flash_attention_2"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    
    eos_ids = get_eos_ids(model, tokenizer)
    print(f"Found EOS token IDs: {eos_ids}")

    # --- 2. First Question ---
    prompt = "How are you?"
    input_ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
    ).to(model.device)
    
    attention_mask = torch.ones_like(input_ids)
    
    # Call the reusable generation function
    # We start with no KV cache (past_key_values=None)
    generated_sequence, attention_mask, past_key_values, new_tokens = generate_response(
        model=model,
        input_ids=input_ids,
        attention_mask=attention_mask,
        past_key_values=None,
        eos_ids=eos_ids,
        max_new_tokens=max_new_tokens
    )

    print("\n=== First Response ===")
    print(tokenizer.decode(generated_sequence[0], skip_special_tokens=False))
    
    summary = {
        "new_tokens": new_tokens,
        "num_new_tokens": len(new_tokens),
        "total_length": int(generated_sequence.shape[-1]),
    }
    print(summary)

    # --- 3. Follow-up Question ---
    print("\n=== Follow-up Question ===")
    
    followup_prompt = "What is C. elegans?"
    
    # Manually construct the follow-up text *including* the template tokens
    # This text will be "fed" to the model to update the cache
    followup_text = f"
user\n{followup_prompt}
\n
assistant\n"
    followup_ids = tokenizer.encode(
        followup_text, 
        add_special_tokens=False, 
        return_tensors="pt"
    ).to(model.device)
    
    # --- 4. "Warm up" the cache ---
    # Run a forward pass on the follow-up prompt to update the cache
    # We do *not* use the logits, only the new past_key_values
    
    # Create the attention mask for *just* the follow-up prompt
    followup_attention_mask = torch.ones_like(followup_ids)
    
    # The *full* mask for this pass includes the old history AND the new prompt
    updated_attention_mask = torch.cat([attention_mask, followup_attention_mask], dim=-1)

    with torch.no_grad():
        model_inputs = {
            "input_ids": followup_ids,
            "attention_mask": updated_attention_mask, # Pass the *full* mask
            "use_cache": True,
            "past_key_values": past_key_values, # Pass the *old* cache
        }
        outputs = model(**model_inputs)
        
        # Save the *new* cache
        past_key_values = outputs.past_key_values
    
    # Update our trackers to include the follow-up prompt
    generated_sequence = torch.cat([generated_sequence, followup_ids], dim=-1)
    attention_mask = updated_attention_mask


    # --- 5. Generate the follow-up answer ---
    # Call the *same* generation function, but this time pass the *updated*
    # sequence, mask, and cache.
    full_conversation_seq, _, _, followup_tokens = generate_response(
        model=model,
        input_ids=generated_sequence,  # The full history + new prompt
        attention_mask=attention_mask, # The full, updated mask
        past_key_values=past_key_values, # The warmed-up cache
        eos_ids=eos_ids,
        max_new_tokens=max_new_tokens
    )
    
    print("\n=== Full Conversation ===")
    print(tokenizer.decode(full_conversation_seq[0], skip_special_tokens=False))

    followup_summary = {
        "followup_new_tokens": followup_tokens,
        "num_followup_tokens": len(followup_tokens),
        "total_conversation_length": int(full_conversation_seq.shape[-1]),
    }
    print(followup_summary)


if __name__ == "__main__":
    main()

```

### Expected behavior

When asking an followup question utilizing the kv-cache result from first prefill&decode stage, the model starts to give output non-related to the followup question, and eachtime the ouput different even when I manually doing a greedy inference.
```bash
Loading model: LiquidAI/LFM2-1.2B...
`torch_dtype` is deprecated! Use `dtype` instead!
Found EOS token IDs: [7]
[EOS reached]

=== First Response ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I'm functioning well and ready to help! How can I assist you today?

{'new_tokens': [550, 6217, 1804, 768, 5354, 2209, 521, 1101, 859, 2008, 1901, 1052, 11965, 797, 779, 1959, 6934, 1112, 521, 1203, 859, 6217, 18199, 1714, 810, 8177, 811, 1801, 510, 2213, 1011, 859, 6727, 1010, 4008, 540, 7], 'num_new_tokens': 37, 'total_length': 50}

=== Follow-up Question ===
[EOS reached]

=== Full Conversation ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I'm functioning well and ready to help! How can I assist you today?

user
What is C. elegans?


assistant
I'm doing well, thank you for asking! How can I assist you today?

{'followup_new_tokens': [550, 6217, 5296, 1714, 521, 14954, 1010, 875, 14065, 510, 2213, 1011, 859, 6727, 1010, 4008, 540, 7], 'num_followup_tokens': 18, 'total_conversation_length': 83}
```


However, the problem is solved once I modify [this if statement](https://github.com/huggingface/transformers/blob/6efc1799c1d0665c41717c555f3551d974e4d61a/src/transformers/models/lfm2/modular_lfm2.py#L344) to
```python3
 if past_key_values is not None and seqlen == 1:  
```

which then the model start to give some reasonable response and the output is the same every time I run the script.
```bash

(venv) shouyud@ser11-NucBox-EVO-X1:~/liquid-mega-kernel-npu/python_code/liquid-ai$ python3 test.py 
Loading model: LiquidAI/LFM2-1.2B...
`torch_dtype` is deprecated! Use `dtype` instead!
Found EOS token IDs: [7]
[EOS reached]

=== First Response ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I'm functioning well and ready to help! How can I assist you today?

{'new_tokens': [550, 6217, 1804, 768, 5354, 2209, 521, 1101, 859, 2008, 1901, 1052, 11965, 797, 779, 1959, 6934, 1112, 521, 1203, 859, 6217, 18199, 1714, 810, 8177, 811, 1801, 510, 2213, 1011, 859, 6727, 1010, 4008, 540, 7], 'num_new_tokens': 37, 'total_length': 50}

=== Follow-up Question ===
[EOS reached]

=== Full Conversation ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I'm functioning well and ready to help! How can I assist you today?

user
What is C. elegans?


assistant
C. elegans, also known as Caenorhabditis elegans, is a small, free-living nematode worm (roundworm) that has become a widely used model organism in scientific research, particularly in the fields of genetics, development, and neuroscience. It's a fascinating creature, measuring only about 0.5 to 1 millimeter in length, and has a simple yet well-understood genome, making it an ideal organism for studying basic biological processes.

Researchers have been studying C. elegans for decades because of its unique characteristics: it has a short lifespan, produces only a few offspring, and its development is relatively simple and well-characterized. This makes it an excellent model for understanding fundamental biological processes, such as cell differentiation, neural development, and aging.

Many discoveries made in C. elegans have been applied to understanding human diseases, including cancer, neurodegenerative disorders, and developmental disorders. Its genome was even fully sequenced in 2004, which has greatly facilitated genetic studies. As a result, C. elegans continues to be an invaluable tool for scientists seeking to unravel the complexities of life.

{'followup_new_tokens': [544, 523, 14230, 998, 521, 1236, 2538, 906, 11379, 771, 777, 7707, 577, 6160, 14230, 998, 521, 856, 768, 2211, 521, 3262, 2811, 2242, 58831, 1584, 36548, 843, 1562, 33637, 518, 896, 1178, 3117, 768, 9087, 1593, 2801, 19817, 797, 6799, 2098, 521, 5777, 797, 779, 7387, 803, 28327, 521, 2917, 521, 810, 58637, 523, 1311, 1090, 768, 23776, 28231, 521, 16554, 1550, 1337, 730, 525, 523, 530, 811, 730, 526, 52912, 797, 5397, 521, 810, 1178, 768, 4048, 4763, 1714, 522, 7755, 10959, 16121, 521, 3525, 936, 902, 8977, 19817, 875, 13144, 5975, 10874, 6895, 523, 509, 26101, 1052, 1381, 13144, 841, 523, 14230, 998, 875, 9229, 1929, 803, 1352, 5746, 3338, 4642, 535, 936, 1178, 768, 3290, 43212, 521, 13240, 1550, 768, 2674, 29854, 521, 810, 1352, 2917, 856, 7932, 4048, 810, 1714, 522, 37778, 1884, 523, 1470, 4201, 936, 902, 9991, 2801, 875, 5139, 10262, 10874, 6895, 521, 1559, 906, 2009, 1847, 4149, 521, 16898, 2917, 521, 810, 19179, 523, 509, 10428, 29529, 2047, 797, 841, 523, 14230, 998, 1052, 1381, 6944, 811, 5139, 2495, 6615, 521, 2452, 3588, 521, 57322, 1539, 9546, 521, 810, 21925, 9546, 523, 8874, 16121, 953, 1919, 7408, 5169, 8470, 797, 730, 1086, 529, 521, 1144, 1178, 13450, 47966, 8047, 3609, 523, 1783, 768, 1725, 521, 841, 523, 14230, 998, 10199, 811, 874, 902, 49656, 6690, 875, 6599, 14688, 811, 946, 46347, 779, 4121, 1535, 803, 2105, 523, 7], 'num_followup_tokens': 237, 'total_conversation_length': 302}
```

Thus, it seem to me this is a implement bug?


Comments:
It looks like the model is not answering the follow-up question because of how the conversation history is being passed to the `generate_response` function.
I've updated the `generate_response` function to correctly handle the initial input and the subsequent token generation when a KV cache is provided. I also adjusted the `main` function to properly prepare the input for the follow-up question, ensuring only the new turn's tokens are passed to `generate_response`. This should help the model generate the correct response for the second question while utilizing the cache.

``` python

import torch
from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List, Optional, Tuple, Dict, Any

# Suppress a specific UserWarning from transformers about custom models
from warnings import filterwarnings
filterwarnings("ignore", message="You are using a model of type lfm2 to instantiate a model of type .*", category=UserWarning)

def get_eos_ids(model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> List[int]:
    """Collects all known EOS token IDs from the model config and tokenizer."""
    eos_ids = set()
    
    # From model config
    cfg_eos = getattr(model.config, "eos_token_id", None)
    if isinstance(cfg_eos, (list, tuple)):
        eos_ids.update(int(x) for x in cfg_eos)
    elif isinstance(cfg_eos, int):
        eos_ids.add(cfg_eos)
        
    # From tokenizer
    tok_eos = getattr(tokenizer, "eos_token_id", None)
    if isinstance(tok_eos, (list, tuple)):
        eos_ids.update(int(x) for x in tok_eos)
    elif isinstance(tok_eos, int):
        eos_ids.add(tok_eos)
        
    return list(eos_ids)


def generate_response(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    past_key_values: Optional[Tuple] = None,
    eos_ids: Optional[List[int]] = None,
    max_new_tokens: int = 512
) -> Tuple[torch.Tensor, torch.Tensor, Tuple, List[int]]:
    """
    Performs a manual greedy decoding loop, reusing and returning the KV cache.

    Args:
        model: The causal LM.
        input_ids: The *initial* token IDs for this step (either the full prompt for the first step, or just the start of the new turn for subsequent steps).
        attention_mask: The *full* attention mask for the entire sequence up to and including the *current* step.
        past_key_values: The KV cache from the previous step, if any.
        eos_ids: A list of token IDs that signal the end of a sequence.
        max_new_tokens: The maximum number of new tokens to generate.

    Returns:
        A tuple containing:
        - final_sequence: The complete generated sequence (inputs + new tokens).
        - final_mask: The updated full attention mask.
        - final_cache: The updated KV cache.
        - new_token_ids: A list of *only* the new token IDs that were generated.
    """

    generated_sequence = input_ids
    current_mask = attention_mask
    current_cache = past_key_values

    new_token_ids = []  # Track only the newly generated tokens

    with torch.no_grad():
        # Process the initial input_ids to update the cache and get the first output
        outputs = model(
            input_ids=generated_sequence,
            attention_mask=current_mask,
            use_cache=True,
            past_key_values=current_cache,
        )
        current_cache = outputs.past_key_values

        # Start generating from the next token
        next_token_logits = outputs.logits[:, -1, :]
        next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)

        # Append the first generated token
        generated_sequence = torch.cat([generated_sequence, next_token_id], dim=-1)
        current_mask = torch.cat([current_mask, torch.ones_like(next_token_id)], dim=-1)
        new_token_ids.append(int(next_token_id.item()))

        # Continue generating for the remaining tokens
        for _ in range(max_new_tokens - 1): # -1 because we already generated one token

            # Prepare model inputs (only the last generated token for subsequent steps)
            model_input_ids = generated_sequence[:, -1:]

            model_inputs = {
                "input_ids": model_input_ids,
                "attention_mask": current_mask,
                "use_cache": True,
                "past_key_values": current_cache,
            }

            outputs = model(**model_inputs)

            # Greedy: pick argmax of last-token logits
            next_token_logits = outputs.logits[:, -1, :]
            next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            # Append the new token to our trackers
            generated_sequence = torch.cat([generated_sequence, next_token_id], dim=-1)
            current_mask = torch.cat([current_mask, torch.ones_like(next_token_id)], dim=-1)
            current_cache = outputs.past_key_values

            new_token_int = int(next_token_id.item())
            new_token_ids.append(new_token_int)

            # Stop at EOS
            if eos_ids and new_token_int in eos_ids:
                print("[EOS reached]")
                break

        return generated_sequence, current_mask, current_cache, new_token_ids


def main():
    # Settings
    model_id = "LiquidAI/LFM2-1.2B"
    max_new_tokens = 512

    # --- 1. Load Model and Tokenizer ---
    print(f"Loading model: {model_id}...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    #    attn_implementation="flash_attention_2"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id)

    eos_ids = get_eos_ids(model, tokenizer)
    print(f"Found EOS token IDs: {eos_ids}")

    # --- 2. First Question ---
    prompt = "How are you?"
    input_ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
    ).to(model.device)

    attention_mask = torch.ones_like(input_ids)

    # Call the reusable generation function
    # We start with no KV cache (past_key_values=None)
    # This first call processes the full initial prompt.
    generated_sequence, attention_mask, past_key_values, new_tokens = generate_response(
        model=model,
        input_ids=input_ids,
        attention_mask=attention_mask,
        past_key_values=None,
        eos_ids=eos_ids,
        max_new_tokens=max_new_tokens
    )

    print("\n=== First Response ===")
    # Decode the full generated sequence, including the prompt and special tokens
    print(tokenizer.decode(generated_sequence[0], skip_special_tokens=False))

    summary = {
        "new_tokens": new_tokens,
        "num_new_tokens": len(new_tokens),
        "total_length": int(generated_sequence.shape[-1]),
    }
    print(summary)

    # --- 3. Follow-up Question ---
    print("\n=== Follow-up Question ===")

    followup_prompt = "What is C. elegans?"

    # Tokenize the follow-up prompt, including the assistant turn start
    # This is needed to determine the input_ids for the current turn based on the template.
    full_history_input_ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}, {"role": "assistant", "content": tokenizer.decode(new_tokens, skip_special_tokens=True)}, {"role": "user", "content": followup_prompt}],
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
    ).to(model.device)

    # We need to identify the tokens that are *new* in this turn.
    # This is everything after the previous conversation's tokens.
    # The length of the previous conversation's tokens is the length of `generated_sequence`.
    current_turn_input_ids = full_history_input_ids[:, generated_sequence.shape[-1]:]

    # Create the attention mask for *just* the current turn's input_ids
    current_turn_attention_mask = torch.ones_like(current_turn_input_ids)

    # The *full* mask for this pass includes the old history AND the new prompt
    updated_attention_mask = torch.cat([attention_mask, current_turn_attention_mask], dim=-1)

    # --- 4. Generate the follow-up answer ---
    # Call the *same* generation function, but this time pass the *updated*
    # mask and the warmed-up cache. Crucially, we pass only the input_ids
    # for the *current turn*.
    generated_followup_response_tokens, final_attention_mask, final_cache, followup_tokens = generate_response(
        model=model,
        input_ids=current_turn_input_ids, # Only the input for the current turn
        attention_mask=updated_attention_mask, # The full, updated mask
        past_key_values=past_key_values, # The warmed-up cache from the previous turns
        eos_ids=eos_ids,
        max_new_tokens=max_new_tokens
    )

    print("\n=== Full Conversation ===")
    # To print the full conversation, we concatenate the original generated_sequence
    # with the newly generated response tokens for the follow-up.
    # The `generated_followup_response_tokens` returned by generate_response contains the current turn's
    # input plus the newly generated tokens. We need to prepend the history.
    full_conversation_history_and_response = torch.cat([generated_sequence[:, :-current_turn_input_ids.shape[-1]], generated_followup_response_tokens], dim=-1)

    print(tokenizer.decode(full_conversation_history_and_response[0], skip_special_tokens=False))

    followup_summary = {
        "followup_new_tokens": followup_tokens,
        "num_followup_tokens": len(followup_tokens),
        "total_conversation_length": int(full_conversation_history_and_response.shape[-1]),
    }
    print(followup_summary)


if __name__ == "__main__":
    main()

```


Output: 

```console

Loading model: LiquidAI/LFM2-1.2B...
Found EOS token IDs: [7]
[EOS reached]

=== First Response ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I'm functioning well and ready to help! How can I assist you today?

{'new_tokens': [550, 6217, 1804, 768, 5354, 2209, 521, 1101, 859, 2008, 1901, 1052, 11965, 797, 779, 1959, 6934, 1112, 521, 1203, 859, 6217, 18199, 1714, 810, 8177, 811, 1801, 510, 2213, 1011, 859, 6727, 1010, 4008, 540, 7], 'num_new_tokens': 37, 'total_length': 50}

=== Follow-up Question ===
[EOS reached]

=== Full Conversation ===


user
How are you?


assistant
I'm just a computer program, so I don't have feelings in the way humans do, but I

user
What is C. elegans?


assistant
I'm here to provide information, answer questions, and help with tasks. How can I assist you further?

{'followup_new_tokens': [550, 6217, 2712, 811, 3176, 2203, 521, 5642, 4639, 521, 810, 1801, 916, 10003, 523, 2213, 1011, 859, 6727, 1010, 3759, 540, 7], 'num_followup_tokens': 23, 'total_conversation_length': 73}

```

I tried the updated script, but I'm confused about the result. Even with the cache updated, the model failed to answer my follow-up question about 'C. elegans'. My understanding was that the update would allow it to answer this. Is my expectation incorrect, or is something else wrong?

You are correct to expect the model to answer the follow-up question using the updated cache. The fact that it's not suggests there's still an issue with how the follow-up prompt is being processed or how the `generate_response` function is handling the input with the updated cache. 

This is entirely caused by bad code written by your LLM, rather than an issue in Transformers - please don't ping our core maintainers to debug your LLM code for you!