Refactor `MambaCache` to `modeling_mamba.py`
Author: manueldeprada
Pull request: yes

This PR moves the specialized `MambaCache` class from `cache_utils.py` to `src/transformers/models/mamba/modeling_mamba.py`. This is preliminary work for #38077.

**Changes:**
- Moved `MambaCache` to its own file, aligning with Zamba, Bamba, etc.
- Removed unnecessary Mamba-specific code from `generate`.
- Added Mamba cache init to `preparing_inits_from_generation` from `forward()`. See [this comment](https://github.com/huggingface/transformers/pull/38086#discussion_r2107581297). 
  - Why? Bamba, Jamba, GraniteMoeHybrid, Zamba, and Zamba2 had settled into initializing custom caches in `preparing_inits_from_generation` and only Mamba, Mamba2, and FalconMamba were doing it in `forward`, which is bad for torch.compile. 
- We dont break BC with any import (thanks joao for the idea!)
- Cleaned up some Mamba and FalconMamba slow tests, which had been failing on main for a long time. 
- Removed DDP Mamba tests. I had a DDP implementation for Mamba in 66b7162c7c5c5ce8a73ccf48cffc8a96343ebb33 so tests passed but removed it since DDP is not needed in Mamba, as per Joaos's instructions.

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

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

@manueldeprada have you run `slow` mamba (and falcon mamba) tests to ensure there are no regressions?

(after we confirm slow tests are okay, let's tag arthur)

> @manueldeprada have you run `slow` mamba (and falcon mamba) tests to ensure there are no regressions?
> 
> (after we confirm slow tests are okay, let's tag arthur)

Ran the tests locally, there are some failures that **already happened before the PR**:
```
FAILED tests/models/mamba/test_modeling_mamba.py::MambaModelTest::test_multi_gpu_data_parallel_forward - TypeError: 'MambaCache' object is not iterable
FAILED tests/models/mamba/test_modeling_mamba.py::MambaIntegrationTests::test_compile_mamba_cache - AssertionError: Attempt to trace forbidden callable <function mark_static_address at 0x7fd585001440>
---
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaModelTest::test_multi_gpu_data_parallel_forward - TypeError: 'MambaCache' object is not iterable
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_batched_generation - AssertionError: Lists differ: ['Hello today I am going to be talking abo[148 chars]The"] != ["Hello today I'm going to show you how to[149 chars].\n']
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_4bit - AssertionError: "Hello today Iava,\n\nI'm sorry to hear that you're having trouble with the " != 'Hello today I\'m going to talk about the "C" in the "C-I-'
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_bf16 - AssertionError: "Hello today Iava,\n\nI'm sorry to hear t[31 chars]the " != 'Hello today I am going to show you how t[47 chars]Step'
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_torch_compile - AssertionError: "Hello today Iava,\n\nI'm sorry to hear t[31 chars]the " != 'Hello today I am going to show you how t[47 chars]Step'
```
will run in CI once I have permissions.

Same failures on CI than locally @gante . Let me know if those slow tests should be fixed, ignored or deleted.

all set, tests clear. @ArthurZucker could you have a look?

(@manueldeprada there were some failures, are they the same as on `main`? )

> Ran the slow tests locally, there are some failures that **already happened before the PR**:
> 
> ```
> FAILED tests/models/mamba/test_modeling_mamba.py::MambaModelTest::test_multi_gpu_data_parallel_forward - TypeError: 'MambaCache' object is not iterable
> FAILED tests/models/mamba/test_modeling_mamba.py::MambaIntegrationTests::test_compile_mamba_cache - AssertionError: Attempt to trace forbidden callable <function mark_static_address at 0x7fd585001440>
> ---
> FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaModelTest::test_multi_gpu_data_parallel_forward - TypeError: 'MambaCache' object is not iterable
> FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_batched_generation - AssertionError: Lists differ: ['Hello today I am going to be talking abo[148 chars]The"] != ["Hello today I'm going to show you how to[149 chars].\n']
> FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_4bit - AssertionError: "Hello today Iava,\n\nI'm sorry to hear that you're having trouble with the " != 'Hello today I\'m going to talk about the "C" in the "C-I-'
> FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_bf16 - AssertionError: "Hello today Iava,\n\nI'm sorry to hear t[31 chars]the " != 'Hello today I am going to show you how t[47 chars]Step'
> FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_torch_compile - AssertionError: "Hello today Iava,\n\nI'm sorry to hear t[31 chars]the " != 'Hello today I am going to show you how t[47 chars]Step'
> ```
> 
> will run in CI once I have permissions.



> Same failures on CI than locally @gante . Let me know if those slow tests should be fixed, ignored or deleted.

yes @gante !! those test were broken. Should I delete/fix? Specially the falcon assertion failures...

if low effort -> fix, otherwise leave them be :)

@gante done. The falcon_mamba slow assertion tests were failing since their first commit. I wrote new assertions.

~~Also I coded quickly torch.DDP support in MambaCache mimicking  DynamicCache in 66b7162c7c5c5ce8a73ccf48cffc8a96343ebb33.~~ EDIT: removed since DDP support on inference is not needed.

run-slow: mamba, falcon_mamba

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

models: ['models/falcon_mamba', 'models/mamba']
quantizations: [] ...

> This is a little bit breaking, as we cannot import the MambaCache from the same place as before! I would maybe deprecate it?

EDIT: BC break solved, see below :)

@ArthurZucker (and @manueldeprada ) replying to both Arthur's comment here at the top level since they are oddly related 👀 

If we align mamba with other mamba-like models, where their caches are defined in their modeling files, `utils.py` (where `generate` lives) can't import `MambaCache`. This is a circular import: `modeling_mamba.py` imports `GenerationMixin` from `utils.py`, so `utils.py` can't import `MambaCache` from `modeling_mamba.py`. 

Breaking `from transformers.cache_utils import MambaCache` was also my [question above](https://github.com/huggingface/transformers/pull/38086#discussion_r2086389849), and it is also because of a circular import.

________________

I think we can solve both circular import issues by moving `MambaCache` to a separate file, e.g. `cache_mamba.py`. `MambaCache` has minimal requirements, all files should be able to import it. @ArthurZucker WDYT? Will it play well with modular?

As for the initialization of the cache: I think it makes more sense for it to stay in the model itself (as opposed to `generate`), to prevent circular dependencies. Otherwise, the mixin needs to know about the classes that inherit it, which is a bad programming pattern.

> I think we can solve both circular import issues by moving `MambaCache` to a separate file, e.g. `cache_mamba.py`. `MambaCache` has minimal requirements, all files should be able to import it. @ArthurZucker WDYT? Will it play well with modular?

That's a great idea!! I have coded it on the fly and restored BC :)

EDIT: cleaned everything up and updated original post with TLDR info for easy review @ArthurZucker 

run-slow: mamba, falcon_mamba, mamba2

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

models: ['models/falcon_mamba', 'models/mamba', 'models/mamba2']
quantizations: [] ...

@ArthurZucker gentle ping :)

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

run-slow: falcon_mamba, mamba, mamba2

Thanks, @ArthurZucker! I ended up refactoring the entire `prepare_inputs_for_generation` method to make it clearer, rather than just making minimal changes to get the tests passing. I’ll do this from the start going forward!

lmk if theres anything left!