Fix silently incorrect CPU matmul for a stride-zero broadcast batch

The baseline CPU backend folded a batched matmul with a broadcast lhs into a
single gemm over (1, m, b * n, k). That fold is not valid: rhs is batch-major,
so its batches are not adjacent columns, and dst_rs stays n while each merged
row is b * n wide. The gemm wrote over itself, leaving only (m - 1) * n + b * n
elements written and the rest at their zero-initialised value.

Drop the fold. The per-batch loop below already handles a_skip == 0 by feeding
the same lhs to every step, which is the correct broadcast semantics.

Fixes #3744

Author: akashjainn

candle-core/src/cpu_backend/mod.rs

@@ -1485,12 +1485,17 @@ impl Map2 for MatMul {
             Parallelism::None
         };
         let (b, m, n, k) = if b_skip == 0 && a_skip == m * k {
+            // A batch-invariant rhs lets the batches of lhs stack into the rows of a single
+            // (b * m, k) matrix, which the destination already matches row for row.
             // a_skip and c_skip should be updated but step is always 0 so
             // it wouldn't matter.
             (1, b * m, n, k)
-        } else if a_skip == 0 && b_skip == n * k {
-            (1, m, b * n, k)
         } else {
+            // There is deliberately no mirrored fold for a batch-invariant lhs. Stacking the
+            // batches of rhs into the columns of a (k, b * n) matrix would need those columns
+            // to be adjacent, but rhs is batch-major, and the destination rows would still be
+            // n apart while each merged row is b * n wide. Both make the merged gemm write
+            // over itself, so a stride-zero batch on lhs goes through the loop below.
             (b, m, n, k)
         };
         for step in 0..b {

candle-core/tests/matmul_tests.rs

@@ -82,6 +82,30 @@ fn broadcast_matmul(device: &Device) -> Result<()> {
     Ok(())
 }
 
+// A stride-zero batch dim on the lhs, as produced by `broadcast_as`, must give the same
+// result as materializing that batch. The batches of rhs cannot be folded into the columns
+// of one matmul the way a batch-invariant rhs folds into the rows, so this exercises the
+// per-batch loop.
+fn broadcast_matmul_stride_zero_lhs(device: &Device) -> Result<()> {
+    for (b, m, n, k) in [(32, 32, 32, 32), (3, 1, 2, 5), (4, 2, 3, 1), (2, 5, 1, 3)] {
+        let lhs = Tensor::randn(0f32, 1f32, (1, m, k), device)?;
+        let rhs = Tensor::randn(0f32, 1f32, (b, k, n), device)?;
+
+        let out = lhs.broadcast_as((b, m, k))?.matmul(&rhs)?;
+        assert_eq!(out.dims(), &[b, m, n]);
+
+        // Every batch is the same lhs against that batch of rhs.
+        let lhs = lhs.i(0)?;
+        for idx in 0..b {
+            let diff = (out.i(idx)? - lhs.matmul(&rhs.i(idx)?)?)?
+                .sqr()?
+                .sum_all()?;
+            assert!(diff.to_vec0::<f32>()? < 1e-6, "batch {idx} differs");
+        }
+    }
+    Ok(())
+}
+
 fn zero_matmul(device: &Device) -> Result<()> {
     let lhs = Tensor::zeros((2, 0), DType::F32, device)?;
     let rhs = Tensor::zeros((0, 3), DType::F32, device)?;
@@ -238,6 +262,12 @@ test_device!(
     broadcast_matmul_gpu,
     broadcast_matmul_metal
 );
+test_device!(
+    broadcast_matmul_stride_zero_lhs,
+    broadcast_matmul_stride_zero_lhs_cpu,
+    broadcast_matmul_stride_zero_lhs_gpu,
+    broadcast_matmul_stride_zero_lhs_metal
+);
 test_device!(
     zero_matmul,
     zero_matmul_cpu,