Repository: huggingface/transformers
Pull request: Fix Bad Outputs in Fast Path for GraniteMoeHybrid

This PR fixes garbage output being produced when running granite moe hybrid models on the fast path; this is currently caused by inconsistent handling in `has_previous_state` on the fast path because the cache param was updated directly in the forward call of the state space layer in the slow path of bamba + updated through modular instead of handling the cache params at the end of forward like other state space models. 

Quick repro case:
```python
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
import torch

model_path="ibm-granite/granite-4.0-tiny-preview"
device="cuda"
model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map=device,
        torch_dtype=torch.bfloat16,
    )
tokenizer = AutoTokenizer.from_pretrained(
        model_path
)

conv = [{"role": "user", "content":"What is a rattlesnake? "}]

input_ids = tokenizer.apply_chat_template(conv, return_tensors="pt", thinking=True, return_dict=True, add_generation_prompt=True).to(device)

set_seed(42)
output = model.generate(
    **input_ids,
    max_new_tokens=32,
)

prediction = tokenizer.decode(output[0, input_ids["input_ids"].shape[1]:], skip_special_tokens=True)
print(prediction)
```

Current output with `mamba_ssm` and `causal_conv1d` installed:
```
< the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the
```

vs without it
```
<think>A rattlesnake is a venomous snake species known for its distinctive rattle at the end of its tail, which it
```
## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

@ArthurZucker can you please take a look?

Changed files:
filename: src/transformers/models/bamba/modeling_bamba.py
patch: @@ -867,7 +867,6 @@ def torch_forward(
             # Init cache
             if ssm_state is not None and cache_params is not None:
                 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
-                cache_params.has_previous_state = True
 
         scan_output = self.norm(y, gate)
 

filename: src/transformers/models/bamba/modular_bamba.py
patch: @@ -666,7 +666,6 @@ def torch_forward(
             # Init cache
             if ssm_state is not None and cache_params is not None:
                 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
-                cache_params.has_previous_state = True
 
         scan_output = self.norm(y, gate)
 

filename: src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
patch: @@ -794,7 +794,6 @@ def torch_forward(
             # Init cache
             if ssm_state is not None and cache_params is not None:
                 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
-                cache_params.has_previous_state = True
 
         scan_output = self.norm(y, gate)
 
@@ -1376,6 +1375,9 @@ def forward(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
+        if past_key_values and not past_key_values.has_previous_state:
+            past_key_values.has_previous_state = True
+
         next_cache = next_decoder_cache if use_cache else None
 
         return MoeModelOutputWithPast(

filename: src/transformers/models/granitemoehybrid/modular_granitemoehybrid.py
patch: @@ -301,6 +301,9 @@ def forward(
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
+        if past_key_values and not past_key_values.has_previous_state:
+            past_key_values.has_previous_state = True
+
         next_cache = next_decoder_cache if use_cache else None
 
         return MoeModelOutputWithPast(