Beam search cache reorder silently skipped for Mamba, XLNet, RWKV, and Reformer models (wrong generation output)
Author: wi5nuu
Pull request: no
Labels: name: bug

### System Info

transformers version: main/latest
Platform: all
Python version: 3.10+

### Who can help?

@Cyrilvallez 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

<html>
<body>
<!--StartFragment--><html><head></head><body><h2>Steps to Reproduce</h2><p>Run the following code using any model that stores generation cache under a key other than <code inline="">past_key_values</code>:</p><pre><code class="language-python">from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")

tokenizer.pad_token = tokenizer.eos_token

inputs = tokenizer(
    "Hello, my dog is cute",
    return_tensors="pt",
)

# Greedy decoding works correctly
greedy_out = model.generate(
    **inputs,
    num_beams=1,
)

print(
    "Greedy:",
    tokenizer.decode(
        greedy_out[0],
        skip_special_tokens=True,
    ),
)

# Beam search
beam_out = model.generate(
    **inputs,
    num_beams=3,
    num_return_sequences=3,
)

print(
    "Beam:",
    tokenizer.batch_decode(
        beam_out,
        skip_special_tokens=True,
    ),
)
</code></pre><h3>Expected Behavior</h3><p>Beam search should produce results that are consistent with the model's cache implementation.</p><p>During beam reordering, the active cache should be reordered regardless of which cache key name the model uses internally.</p><h3>Actual Behavior</h3><p>Beam search only checks for the presence of:</p><pre><code class="language-python">past_key_values
</code></pre><p>before triggering cache reordering.</p><p>Models that store generation state under alternative cache keys do not have their cache reordered during beam search.</p><p>As a result, beam hypotheses may continue using stale or incorrectly aligned cached states, producing incorrect or degraded generation results.</p><h3>Root Cause</h3><p>In:</p><pre><code class="language-text">src/transformers/generation/utils.py
</code></pre><p>(around line 3395)</p><p>the generation loop contains logic equivalent to:</p><pre><code class="language-python">if model_kwargs.get("past_key_values") is not None:
</code></pre><p>However, the library already defines multiple supported cache names through <code inline="">ALL_CACHE_NAMES</code> (around line 126):</p><pre><code class="language-python">ALL_CACHE_NAMES = (
    "past_key_values",
    "cache_params",
    "state",
    "mems",
    "past_buckets_states",
)
</code></pre><p>These correspond to different model families:</p>
Cache Key | Model Family
-- | --
past_key_values | Default transformer models
cache_params | Mamba-based models
state | RWKV
mems | XLNet
past_buckets_states | Reformer

<p>The beam-search reordering logic only considers <code inline="">past_key_values</code> and ignores the remaining cache types.</p><h3>Impact</h3><p>Models that rely on:</p><ul><li><p><code inline="">cache_params</code></p></li><li><p><code inline="">state</code></p></li><li><p><code inline="">mems</code></p></li><li><p><code inline="">past_buckets_states</code></p></li></ul><p>may not correctly reorder cached generation state during beam search.</p><p>This can lead to:</p><ul><li><p>incorrect beam search behavior,</p></li><li><p>degraded generation quality,</p></li><li><p>inconsistent outputs between greedy decoding and beam search,</p></li><li><p>model-specific failures that are difficult to diagnose.</p></li></ul><h3>Suggested Fix</h3><p>Instead of checking only:</p><pre><code class="language-python">model_kwargs.get("past_key_values")
</code></pre><p>the beam-search code should handle all cache keys listed in <code inline="">ALL_CACHE_NAMES</code>.</p><p>For example:</p><pre><code class="language-python">for cache_name in ALL_CACHE_NAMES:
    if model_kwargs.get(cache_name) is not None:
        ...
</code></pre><p>or reuse the existing cache abstraction already defined by the generation framework so that cache reordering remains model-agnostic.</p></body></html><!--EndFragment-->
</body>
</html>

### Expected behavior

The cache reorder at `src/transformers/generation/utils.py:3395` should use `ALL_CACHE_NAMES` to find the correct cache key, identical to the pattern already used at line 2864:

```python
if any(cache_key in model_kwargs for cache_key in ALL_CACHE_NAMES):
    cache_key = next(cache_key for cache_key in ALL_CACHE_NAMES if cache_key in model_kwargs)
    beam_idx = self._flatten_beam_dim(running_beam_indices[..., cur_len - decoder_prompt_len])
    if hasattr(self, "_reorder_cache"):
        model_kwargs[cache_key] = self._reorder_cache(model_kwargs[cache_key], beam_idx)
    else:
        model_kwargs[cache_key].reorder_cache(beam_idx)

Comments:
Took a look — the reorder block in `_beam_search` only checks `past_key_values`, so a model whose cache lives under another key (`cache_params` for mamba, `state` for rwkv, `mems` for xlnet, `past_buckets_states` for reformer) skips the reorder entirely. Resolving it through `ALL_CACHE_NAMES` instead fixes the lookup — I confirmed mamba beam search now matches the `use_cache=False` reference (it was degraded before), and gpt2 is unaffected.

RWKV looks like a separate, deeper issue: its state is a plain list with no reorder method, and the output stays degenerate even after adding one, so I'd leave that out of this fix. Happy to send a PR for the cache-key lookup if that sounds right, @Cyrilvallez.