[CUDA] Implement gather_mm_rhs (#2902)

python/tests/cuda_skip.py

@@ -6,7 +6,7 @@ cuda_skip = {
     # Gather matmul NYI
     "TestBlas.test_gather_matmul",
     "TestBlas.test_gather_matmul_grad",
-    "TestBlas.test_gather_mm_sorted",
+    "TestBlas.test_gather_mm_sorted_vjp",
     # Segmented matmul NYI
     "TestBlas.test_segmented_mm",
     # Hadamard NYI

python/tests/test_blas.py

@@ -1235,15 +1235,39 @@ class TestBlas(mlx_tests.MLXTestCase):
         def gather_mm_test(a, b, rhs):
             return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
 
+        dtypes = [(mx.float32, 1e-4)]
+        if mx.cuda.is_available():
+            dtypes += [
+                (mx.float16, 1e-3),
+                (mx.bfloat16, 1e-2),
+            ]
+
+        for b_transposed in (True, False):
+            for dtype, tol in dtypes:
+                with self.subTest(b_transposed=b_transposed, dtype=dtype):
+                    a = mx.random.normal((100, 1, 100), dtype=dtype)
+                    b = mx.random.normal((8, 100, 100), dtype=dtype)
+                    if b_transposed:
+                        b = b.swapaxes(-1, -2)
+                    rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
+
+                    c1 = gather_mm_ref(a, b, rhs)
+                    c2 = gather_mm_test(a, b, rhs)
+                    self.assertTrue(mx.allclose(c1, c2, rtol=tol, atol=tol))
+
+    def test_gather_mm_sorted_vjp(self):
+        def gather_mm_ref(a, b, rhs):
+            b = b[rhs]
+            return a @ b
+
+        def gather_mm_test(a, b, rhs):
+            return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
+
         a = mx.random.normal((100, 1, 100))
         b = mx.random.normal((8, 100, 100))
         rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
 
-        c1 = gather_mm_ref(a, b, rhs)
-        c2 = gather_mm_test(a, b, rhs)
-        self.assertTrue(mx.allclose(c1, c2, atol=1e-4))
-
-        cotan = mx.random.normal(c1.shape)
+        cotan = mx.random.normal((100, 1, 100))
         c1, dc1 = mx.vjp(
             lambda a, b: gather_mm_ref(a, b, rhs),
             [a, b],