adding-support-for-mamba2

ACKNOWLEDGMENTS.md

@@ -14,4 +14,4 @@ MLX Examples was developed with contributions from the following individuals:
 - Markus Enzweiler: Added the `cvae` examples.
 - Prince Canuma: Helped add support for `Starcoder2` models.
 - Shiyu Li: Added the `Segment Anything Model`.
-- Gökdeniz Gülmez: Added support for `MiniCPM`, `Mamba` and support for `full-fine-tuning`.
+- Gökdeniz Gülmez: Added support for `MiniCPM`, `Mamba version 1`, `Mamba version 2` and support for `full-fine-tuning`.

llms/mlx_lm/models/cache.py

@@ -324,6 +324,7 @@ def trim(self, n):
 class MambaCache(_BaseCache):
     def __init__(self):
         self.cache = [None, None]
+        self.offset = 0
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -338,3 +339,15 @@ def state(self):
     @state.setter
     def state(self, v):
         self.cache = v
+
+
+class Mamba2Cache:
+    def __init__(self, batch_size, conv_dim, kernel_size, num_heads, head_dim, state_size):
+        self.conv_states = mx.zeros((batch_size, conv_dim, kernel_size - 1))
+        self.ssm_states = mx.zeros((batch_size, num_heads, head_dim, state_size))
+        self.seqlen_offset = 0
+
+    def update(self, new_conv_state, new_ssm_state):
+        self.conv_states = new_conv_state
+        self.ssm_states = new_ssm_state
+        self.seqlen_offset += 1