Optimize Mamba2 memory usage by replacing broadcast with einsum
Author: SohamWalam11
Pull request: yes

# What does this PR do?

<!--
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 # (issue)


## Before submitting
- [x] 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.
- [x] 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.

<!-- 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: @zach-huggingface @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker @S1ro1
- CIs: @ydshieh

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber
- kernels: @MekkCyber @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:
This pull request resolves a critical bug where the Mamba2 model's output was inconsistent between its optimized 'fast' inference path and its standard path. The discrepancy caused multiple test cases to fail. By debugging the low-level implementation, I identified and fixed a numerical inconsistency, ensuring the model's output is now reliable and all validation tests pass

I will always support a beautiful `einsum` instead of a wasteful and ugly intermediate array :angel: 

But @SohamWalam11 I'm a bit confused, I ran those tests on `main` and the output difference between `train` and `eval` was 0. Can you clarify if this is a PR to optimize memory usage or to fix numerical discrepancies?

@Rocketknight1 Thank you for reviewing.
This PR focuses on memory optimization rather than correcting numerical discrepancies. The goal is to eliminate unnecessary intermediate tensors and reduce the overall memory footprint while maintaining functional equivalence.

Summary of Changes

Memory usage reduced from gigabytes to megabytes in large-scale runs.

Replaced inefficient broadcasting with a compact `torch.eins

cc @arthurzucker - I'll have to do some testing, but this seems like something we should accept if you're okay with it. The old code does an element-wise broadcast-multiplication and then a sum on two tensors, and that seems like code that should clearly be replaced by a matmul or einsum because it creates a **huge** intermediate array!

We could also use a similar trick to eliminate the computation of `M_intermediate` below. Those intermediate arrays are massive and I don't think the compiler compiles them out.

Thanks, @ArthurZucker. I agree — the previous element-wise broadcast-multiplication was unnecessarily creating large intermediate arrays, which led to excessive memory usage without numerical benefits.

The updated implementation replaces that pattern with a direct einsum, which performs the same operation in a single, memory-efficient step. I’ll also look into applying a similar optimization to the M_intermediate computation as suggested — that should further reduce memory overhead and improve runtime efficiency.

I’ll run additional benchmarks and share the profiling results here once done.


These operations were deliberately kept to torch without einsum. If you really want to have an efficient implementation, then there is no way around the kernels. I have a complete einsum version at https://github.com/vasqu/mamba2-torch/blob/main/tests/ssd_minimal.py for example

cc @molbap 

If einsum is totally not allowed then I think at least one of these could become a transpose-matmul-transpose, which at least would not create the intermediate array (transpose just alters the stride metadata, so no extra memory is allocated)

Totally open for any improvements. Maybe I sounded to harsh what I wanted to tell is that:
- we need convincing numbers on einsum vs torch (or other changes)
- everyone onboard with this especially @ArthurZucker :D 

Yes, the mamba2 algorithms are based on nd tensor contractions and as said above, array-based impl is deliberate, welcome to optimize it of course, we already merged some (non-einsum) optimizations for the array ops in a couple mamba-like models so if an equivalent formulation without einsum is better, I'm all for it!

:fire: :fist_raised: the people demand einsums :fist_raised: :fire: 

https://github.com/vasqu/mamba2-torch/blob/main/tests/ssd_minimal.py has everything with einsums, just needs benchmarking now :P 

@arthurzucker I tried quickly but I can't even get a performance benchmark at a normal seq_len because of this line in the original, lol:

```python
G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]  # shape: (b, c, l, s, h, n)
```

This is an outer multiplication of two attention tensors, creating a `(b, c, l, s, h, n)` intermediate when the final output we want is only `(b, c, l, s, h)`. `G_intermediate` is massive! It's `n` times larger than the actual tensor we want so my GPU goes OOM on this line! `einsum` computes the output tensor directly without creating `G_intermediate` and works fine.

I can keep attempting, maybe I can get a benchmark on a different GPU or with shorter seq lengths, but the memory issue is enough to justify the change imo.

Yeah that makes sense, IDK if you were training or not, it would probably have been worse. 
Fine with a supeer well commentedd einsum that explains what we go from and to !

yeah I remember this line @Rocketknight1, took me some time to figure out how to reproduce this computation without tensor contraction and couldn't without an intermediate :sweat_smile: (maybe there's a way?) however I was able to run it/measure memory usage on a standard machine, surprised it OOMs immediately

@molbap I think you'd have to transpose the inputs to `(b, c, h, l, n)` and `(b, c, h, n, s)`, then do a matmul to get `(b, c, h, l, s)` and finally transpose that to get the `(b, c, l, s, h)` you want, lol. It would be extremely confusing and probably not very performant! 

Also looking at this more, this whole block (not just `G`) feels like one or two enormous tensor contractions - I'm going to see if I can replace it all with a single big einsum to completely wipe out all of those intermediates, and then leave the original path in as comments. That should make it possible to follow for people who don't like einsums, but should give us some nice memory + performance savings.

@Rocketknight1 see https://github.com/vasqu/mamba2-torch/blob/main/tests/ssd_minimal.py, the torch only version was mostly derived from this. It's complete einsum with least intermediate tensors

@ArthurZucker permission to revert to the no-intermediates big einsum approach to save memory + time and add the non-einsum path as comments instead? :angel: 

@ArthurZucker The tests has been passed perfectly should I share the screenshot of that peculiarly

Since I got approval, I committed the full einsum approach to speed things up and hugely cut memory usage!

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

run-slow: mamba2

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

That might add linebreaks and I'm not sure it'd end up clearer!

I stayed vague because just "some space" is good haha. many letters indeed. Just wanted some oxygen in the middle of the einsum :eyes: 
though +1 for benches

@SohamWalam11 are you still willing to run some benchmarks to check on the latest version? It's just four big einsums now, so hopefully it should be a lot more performant

Yes @Rocketknight1