L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Implement gather_mm_rhs (#2902)

Browse Source
This commit is contained in:
Cheng
2025-12-24 09:42:56 +09:00
committed by GitHub
parent 1eef1d155c
commit 1d21d0e696
7 changed files with 458 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/cuda_skip.py
+1 -1
View File
@@ -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
+29 -5
View File
@@ -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],