Fix Falcon-Mamba long prompt resource exhaustion

mlx_lm/models/mamba.py

@@ -63,10 +63,6 @@ def __init__(self, args: ModelArgs):
         self.time_step_rank = int(args.time_step_rank)
         self.use_conv_bias = args.use_conv_bias
         self.use_bcdt_rms = args.use_bcdt_rms
-        if self.use_bcdt_rms:
-            self.mixer_norm = lambda x: mx.fast.rms_norm(
-                x, mx.ones(x.shape[-1], x.dtype), eps=args.mixer_rms_eps
-            )
 
         self.in_proj = nn.Linear(
             self.hidden_size, self.intermediate_size * 2, bias=args.use_bias
@@ -103,16 +99,16 @@ def __init__(self, args: ModelArgs):
     def ssm_step(self, x, A, state=None):
         D = self.D
         deltaBC = self.x_proj(x)
-        delta, B, C = map(
-            self.mixer_norm if self.use_bcdt_rms else lambda x: x,
-            mx.split(
-                deltaBC,
-                [self.time_step_rank, self.time_step_rank + self.ssm_state_size],
-                axis=-1,
-            ),
+        delta, B, C = mx.split(
+            deltaBC,
+            [self.time_step_rank, self.time_step_rank + self.ssm_state_size],
+            axis=-1,
         )
         if self.use_bcdt_rms:
-            delta, B, C = map(self.mixer_norm, (delta, B, C))
+            eps = self.args.mixer_rms_eps
+            delta = mx.fast.rms_norm(delta, weight=None, eps=eps)
+            B = mx.fast.rms_norm(B, weight=None, eps=eps)
+            C = mx.fast.rms_norm(C, weight=None, eps=eps)
         delta = nn.softplus(self.dt_proj(delta))
         new_state = mx.expand_dims(delta * x, -1) * mx.expand_dims(B, 1)
         if state is not None:

tests/test_models.py

@@ -2,6 +2,7 @@
 import copy
 import importlib
 import unittest
+from unittest import mock
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -1216,6 +1217,40 @@ def test_mamba(self):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_falcon_mamba_bcdt_normalization(self):
+        from mlx_lm.models import mamba
+
+        args = mamba.ModelArgs(
+            model_type="falcon_mamba",
+            vocab_size=32,
+            use_bias=False,
+            use_conv_bias=True,
+            conv_kernel=4,
+            hidden_size=4,
+            num_hidden_layers=1,
+            state_size=2,
+            intermediate_size=8,
+            time_step_rank=2,
+        )
+        block = mamba.MambaBlock(args)
+
+        x = mx.ones((1, args.intermediate_size))
+        A = -mx.ones((args.intermediate_size, args.state_size))
+
+        with (
+            mock.patch.object(mx.fast, "rms_norm", wraps=mx.fast.rms_norm) as rms_norm,
+            mock.patch.object(mx, "ones", wraps=mx.ones) as ones,
+        ):
+            y, state = block.ssm_step(x, A)
+
+        mx.eval(y, state)
+        self.assertEqual(rms_norm.call_count, 3)
+        for call in rms_norm.call_args_list:
+            self.assertIsNone(call.kwargs["weight"])
+        self.assertEqual(ones.call_count, 0)
+        self.assertEqual(y.shape, (1, args.intermediate_size))
+        self.assertEqual(state.shape, (1, args.intermediate_size, args.state_size))
+
     def test_falcon_h1(self):
         from mlx_lm.models import falcon_h1