[CUDA] Fix alpha not respected when using bias epilogue (#2578)

mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -248,11 +248,19 @@ void CublasGemm::run(
     const array& b,
     const Shape& batch_shape,
     const Strides& a_batch_strides,
-    const Strides& b_batch_strides) {
+    const Strides& b_batch_strides,
+    float alpha) {
   int batch_count = out.size() / (M_ * N_);
   if (batch_count / batch_shape.back() > 1) {
     run_batched(
-        encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides);
+        encoder,
+        out,
+        a,
+        b,
+        batch_shape,
+        a_batch_strides,
+        b_batch_strides,
+        alpha);
     return;
   }
 
@@ -260,7 +268,13 @@ void CublasGemm::run(
   encoder.set_input_array(b);
   encoder.set_output_array(out);
 
-  execute(encoder, out.data<void>(), a.data<void>(), b.data<void>(), nullptr);
+  execute(
+      encoder,
+      out.data<void>(),
+      a.data<void>(),
+      b.data<void>(),
+      nullptr,
+      alpha);
 }
 
 void CublasGemm::run(

mlx/backend/cuda/gemms/cublas_gemm.h

@@ -64,7 +64,8 @@ class CublasGemm {
       const array& b,
       const Shape& batch_shape,
       const Strides& a_batch_strides,
-      const Strides& b_batch_strides);
+      const Strides& b_batch_strides,
+      float alpha = 1.0f);
 
   void run(
       cu::CommandEncoder& encoder,
@@ -87,7 +88,8 @@ class CublasGemm {
       const array& b,
       const Shape& batch_shape,
       const Strides& a_batch_strides,
-      const Strides& b_batch_strides);
+      const Strides& b_batch_strides,
+      float alpha);
 
   void run_batched(
       cu::CommandEncoder& encoder,

mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp

@@ -13,7 +13,8 @@ void CublasGemm::run_batched(
     const array& b,
     const Shape& batch_shape,
     const Strides& a_batch_strides,
-    const Strides& b_batch_strides) {
+    const Strides& b_batch_strides,
+    float alpha) {
   encoder.set_input_array(a);
   encoder.set_input_array(b);
   encoder.set_output_array(out);
@@ -27,7 +28,8 @@ void CublasGemm::run_batched(
         out.data<int8_t>() + out.itemsize() * i * batch_shape.back() * M_ * N_,
         a.data<int8_t>() + a.itemsize() * a_it.loc,
         b.data<int8_t>() + b.itemsize() * b_it.loc,
-        nullptr);
+        nullptr,
+        alpha);
     a_it.step();
     b_it.step();
   }

mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu

@@ -154,7 +154,8 @@ void CublasGemm::run_batched(
     const array& b,
     const Shape& batch_shape,
     const Strides& a_batch_strides,
-    const Strides& b_batch_strides) {
+    const Strides& b_batch_strides,
+    float alpha) {
   int batch_count = out.size() / (M_ * N_);
   set_pointer_mode(a_desc_, batch_count);
   set_pointer_mode(b_desc_, batch_count);
@@ -226,7 +227,8 @@ void CublasGemm::run_batched(
       reinterpret_cast<void*>(out_pointers),
       reinterpret_cast<void*>(a_pointers),
       reinterpret_cast<void*>(b_pointers),
-      nullptr);
+      nullptr,
+      alpha);
 }
 
 void CublasGemm::run_batched(

mlx/backend/cuda/matmul.cpp

@@ -41,7 +41,8 @@ void gemm_and_bias(
     array& out,
     const array& a,
     const array& b,
-    void* bias = nullptr) {
+    void* bias = nullptr,
+    float alpha = 1.0f) {
   // Check and collapse batch dimensions
   auto [batch_shape, a_batch_strides, b_batch_strides] = collapse_batches(a, b);
 
@@ -94,7 +95,8 @@ void gemm_and_bias(
   if (bias) {
     gemm.set_bias(bias);
   }
-  gemm.run(encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides);
+  gemm.run(
+      encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha);
 }
 
 } // namespace
@@ -169,7 +171,8 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
         out,
         a,
         b,
-        c.data<void>());
+        c.data<void>(),
+        alpha_);
     return;
   }
 

python/tests/test_blas.py

@@ -594,8 +594,7 @@ class TestBlas(mlx_tests.MLXTestCase):
         np.random.seed(0)
         # Batched matmul
         alpha = 0.5
-        beta = 2.0
+        for beta in (1.0, 2.0):
-
             # c must broadcast to the output shape
             with self.assertRaises(ValueError):
                 mx.addmm(mx.zeros((2, 2, 2)), mx.zeros((2, 2)), mx.zeros((2, 2)))
@@ -701,16 +700,16 @@ class TestBlas(mlx_tests.MLXTestCase):
             # Transposed c
             a = mx.ones((10, 5)).T
             b = mx.ones((5, 5))
-        out = mx.addmm(a, b, a, beta=1.5, alpha=0.5)
+            out = mx.addmm(a, b, a, beta=beta, alpha=alpha)
-        expected = 1.5 * a + 0.5 * (b @ a)
+            expected = beta * a + alpha * (b @ a)
             self.assertTrue(mx.allclose(expected, out))
 
             # Broadcast c
             a = mx.ones((5, 5))
             b = mx.ones((5, 5))
             c = mx.ones((1, 5))
-        out = mx.addmm(c, a, b, beta=1.5, alpha=0.5)
+            out = mx.addmm(c, a, b, beta=beta, alpha=alpha)
-        expected = 1.5 * c + 0.5 * (a @ b)
+            expected = beta * c + alpha * (a @ b)
             self.assertTrue(mx.allclose(expected, out))
 
     def test_addmm_grad(self):
@@ -724,8 +723,7 @@ class TestBlas(mlx_tests.MLXTestCase):
         shapes = ((1, 64, 32, 128), (4, 28, 24, 47), (1, 1, 24, 47))
 
         alpha = 2.0
-        beta = 0.5
+        for beta in (1.0, 0.5):
-
             f_test = make_addmm(alpha, beta)
             f_ref = make_ref_addmm(alpha, beta)