Inform the gather mm regarding the sorted indices (#100)

* Inform the gather mm regarding the sorted indices

* Change the condition for sorting and passing sorted_indices

* Change sorting again

* Remove num_experts member variable

* version bump

---------

Co-authored-by: Awni Hannun <[email protected]>

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.22.4"
+__version__ = "0.23.0"

mlx_lm/models/switch_layers.py

@@ -71,7 +71,7 @@ class QuantizedSwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
+    def __call__(self, x, indices, sorted_indices=False):
         x = mx.gather_qmm(
             x,
             self["weight"],
@@ -81,6 +81,7 @@ class QuantizedSwitchLinear(nn.Module):
             transpose=True,
             group_size=self.group_size,
             bits=self.bits,
+            sorted_indices=sorted_indices,
         )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
@@ -114,8 +115,13 @@ class SwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
-        x = mx.gather_mm(x, self["weight"].swapaxes(-1, -2), rhs_indices=indices)
+    def __call__(self, x, indices, sorted_indices=False):
+        x = mx.gather_mm(
+            x,
+            self["weight"].swapaxes(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=sorted_indices,
+        )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
         return x
@@ -150,17 +156,23 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        should_sort = (x.size // x.shape[-1]) >= 128
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
         idx = indices
         inv_order = None
-        if should_sort:
+        if do_sort:
             x, idx, inv_order = _gather_sort(x, indices)
 
-        x_up = self.up_proj(x, idx)
-        x_gate = self.gate_proj(x, idx)
-        x = self.down_proj(self.activation(x_gate) * x_up, idx)
+        x_up = self.up_proj(x, idx, sorted_indices=do_sort)
+        x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
+        x = self.down_proj(
+            self.activation(x_gate) * x_up,
+            idx,
+            sorted_indices=do_sort,
+        )
 
-        if should_sort:
+        if do_sort:
             x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)
@@ -184,17 +196,19 @@ class SwitchMLP(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        should_sort = (x.size // x.shape[-1]) >= 128
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
         idx = indices
         inv_order = None
-        if should_sort:
+        if do_sort:
             x, idx, inv_order = _gather_sort(x, indices)
 
-        x = self.fc1(x, idx)
+        x = self.fc1(x, idx, sorted_indices=do_sort)
         x = self.activation(x)
-        x = self.fc2(x, idx)
+        x = self.fc2(x, idx, sorted_indices=do_sort)
 
-        if should_sort:
+        if do_sort:
             x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.24.2
+mlx>=0.25.0
 numpy
 transformers[sentencepiece]>=4.39.3
 protobuf