Fix use_cache with seq_len > 1 ( #46032)
Author: Ramshankar07
Pull request: yes



# What does this PR do?

Add a third dispatch path to Mamba2Mixer for the case where a prior cached state exists and seq_len > 1 (chunked prefill / speculative decode verification).

Previously:
   - cuda_kernels_forward: .squeeze(1) crashed or corrupted tensors for seq_len > 1 when has_previous_state was True.
   - torch_forward: dt[:, 0, :] silently dropped all tokens after index 0, and previous_states was always zero-initialised ignoring the cache.

Both paths now:
   - Gate the single-step kernel on seq_len == 1.
   - For seq_len > 1, prepend the cached conv buffer (last K-1 entries) to the chunk input, run a full causal conv, drop the prepended prefix from the output, and pass the cached recurrent_state as initial_states  to mamba_chunk_scan_combined / as previous_states in the naive SSD.

This PR follows the same pattern applied to GDN-based models in #45513,  as discussed in the issues I added test_mamba2_chunked_prefill (CPU/torch path, always runs) and test_mamba2_chunked_prefill_cuda_path (skipped without mamba-ssm).

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #46032 

## Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result, 
**we ask that new users do not submit pure code agent PRs** at this time. 
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously. 

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result, 
this policy is likely to be updated regularly in the near future. For more information, please read [`CONTRIBUTING.md`](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md).

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] 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. #46032 
- [ ] 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).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@zucchini-nlp @Cyrilvallez 🤗@vasqu @Clara2911
<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
I didn't run the cuda tests, can anybody test it?

I was working to make it look closer like bamba script. I'll rework on it today

Sorry I was to eager then :grimacing: feel free to ping me when ready, ignoring the thread for now then

Sure thing, I'll let you know

Should I change `initial_states=cache_params.layers[self.layer_idx].conv_states[:, :, 1:],` to  `initial_states=conv_state[:, :, 1:],` to match the qwen's style?

I changed for now, @vasqu lmk if that's right

Hey @vasqu , Got bit busy lately. I have modified based on your suggestion. Can you take look and lmk if I'm missing any. I couldn't run `` test_mamba2_chunked_prefill_cuda_path(), test_mamba2_slow_vs_fast_forward()
, test_mamba2_slow_vs_fast_forward_grouped()`` cause of GPU. if every thing looks good. I'll try to rent a gpu to run. 

Made changes based on your suggestion, let me know if I miss any

I think it looks clean now! imo

@vasqu have look when you're free

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: mamba2

run-slow: mamba2

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28033791168)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/mamba2"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28033791168)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [6155764d](https://github.com/huggingface/transformers/commit/6155764d99c9ddd15dad0e1dd728e89fec4eabe5) | workflow commit (merge commit) |
| PR | [bb795032](https://github.com/Ramshankar07/transformers/commit/bb795032ea5810311fbf88f6e7bdf2ab2f080338) | branch commit (from PR) |
| main | [5bbed98a](https://github.com/huggingface/transformers/commit/5bbed98a97f1d3b3ef5e15936a126e719c98573a) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46084). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Also checked locally and seems fine, merging now thanks a lot!

Thanks for your suggestions! Got many lessons on mistakes that I make.