Skip to content

Commit

Permalink
[Mamba2] Fix caching, slow path, and multi-gpu (#35154)
Browse files Browse the repository at this point in the history
* fixup mamba2 - caching and several other small fixes

* fixup cached forward

* correct fix this time

* fixup cache - we do not need to extend the attn mask it's handled by generate (gives total ids + mask at each step)

* remove unnecessary (un)squeeze

* fixup cache position

* simplify a few things

* [run-slow] mamba2

* multi gpu attempt two

* [run-slow] mamba2

* [run-slow] mamba2

* [run-slow] mamba2

* [run-slow] mamba2

* add newer slow path fix

* [run-slow] mamba2
  • Loading branch information
vasqu authored Dec 20, 2024
1 parent ff9141b commit 5a2aedc
Show file tree
Hide file tree
Showing 2 changed files with 265 additions and 179 deletions.
Loading

0 comments on commit 5a2aedc

Please sign in to comment.