[BUG] matmul with size greater than 2^31 gives undermined results - appears to be memory leak #3591

Activity

1. vansickledavid-sketch commented on May 26, 2026

   

   vansickledavid-sketch

   on May 26, 2026

   Author

   More actions

   matmul at [12_347, 174_000] produces intermittently incorrect results. Some runs are ok, but others are wildly inaccurate. This appears to be a downcast memory address from UInt64 to Int32 somewhere in the kernel. MacBook Pro M5 Max 128 Gb.

    for n in [12000, 24000, 48000, 96000, 144000, 168000, 174000, 192000, 216000, 240000, 264000, 288000] {
           let A = MLXRandom.normal([12347, n], dtype: .float32)
           let v = MLXRandom.normal([n, 1], dtype: .float32)
           let direct = matmul(A, v)
           direct.eval()
           
           var chunked = MLXArray.zeros([12347, 1], dtype: .float32)
           let chunks = 4
           let chunkSize = n / chunks
           for i in 0..<chunks {
               let lo = i * chunkSize
               let hi = (i == chunks-1) ? n : (i+1) * chunkSize
               chunked = chunked + matmul(A[0..., lo..<hi], v[lo..<hi])
           }
           chunked.eval()
           let rel = abs(direct - chunked).max().item(Float.self) / abs(direct).max().item(Float.self)
           print("n=\(n), rel_diff=\(rel)")
       }

   Typical output:

   n=12000,  rel_diff=3.6e-07
   n=24000,  rel_diff=4.7e-07
   n=48000,  rel_diff=7.1e-07
   n=96000,  rel_diff=1.0e-06
   n=144000, rel_diff=1.5e-06
   n=168000, rel_diff=1.5e-06
   n=174000, rel_diff=0.065    ← intermittent; varies from ~1e-6 to ~0.25 across runs
   n=192000, rel_diff=1.6e-06
   n=216000, rel_diff=1.1e-06
   n=240000, rel_diff=9.5e-07
   n=264000, rel_diff=8.4e-07
   n=288000, rel_diff=1.1e-06
   
2. vansickledavid-sketch commented on May 27, 2026

   

   vansickledavid-sketch

   on May 27, 2026

   Author

   More actions

   Built a workaround using manual block multiplications. Key finding beyond the original report: sliced inputs to matmul also produce incorrect results unless explicitly materialized via slice + 0 (or equivalent). matmul(left[0..., lo..<hi], right[lo..<hi, 0...]) produces wrong results; matmul((left[0..., lo..<hi] + 0), (right[lo..<hi, 0...] + 0)) is correct.
   Both conditions (large operand size and strided-view-into-large-parent) appear to trigger the same offset-arithmetic issue in the matmul kernel. The workaround below handles both by chunking and materializing each chunk.

   infix operator ×: MultiplicationPrecedence
   func ×(left: MLXArray, right: MLXArray) -> MLXArray {
       let safeLimit = 2_000_000_000
       if left.size <= safeLimit && right.size <= safeLimit {
           return matmul(left, right)
       }
       
       assert(left.ndim == 2 && right.ndim == 2, "Only for matrix multiplication")
       
       let m = left.shape[0]
       let k = left.shape[1]
       let n = right.shape[1]
       
       assert(k == right.shape[0], "Dimension Mismatch")
       assert(n <= safeLimit && m <= safeLimit, "Single matrix dimension too large")
    
       let kChunk = safeLimit/max(m, n)
       let numChunks = (k + kChunk - 1)/kChunk
       
       var result = MLXArray.zeros([m, n], dtype: left.dtype)
       for i in 0..<numChunks {
           let lo = i * kChunk
           let hi = min(k, lo + kChunk)
           result = result + matmul(left[0..., lo..<hi] + 0, right[lo..<hi, 0...] + 0)
           result.eval()
       }
       return result
   }
3. vansickledavid-sketch commented on May 27, 2026

   

   vansickledavid-sketch

   on May 27, 2026

   Author

   More actions

   Related
   #3248: same overflow pattern in conv_general
   #3327: root cause identified as ShapeElem = int32_t in mlx/array.h, affecting multiple operations
4. zcbenz commented on May 27, 2026

   

   zcbenz

   on May 27, 2026

   Member

   More actions

   What is the version of MLX you are using and can you also try some old versions of MLX?
5. vansickledavid-sketch commented on May 28, 2026

   

   vansickledavid-sketch

   on May 28, 2026

   Author

   More actions

   What is the version of MLX you are using and can you also try some old versions of MLX?

   Currently using 0.31.3. Switched to 0.25.1 and produced similar output...

   Typical Output:

   n=12000, rel_diff=4.385549e-07
   n=24000, rel_diff=5.5819635e-07
   n=48000, rel_diff=7.3384103e-07
   n=96000, rel_diff=1.0258304e-06
   n=144000, rel_diff=1.6888195e-06
   n=168000, rel_diff=1.6149158e-06
   n=174000, rel_diff=0.1652007
   n=192000, rel_diff=1.8757138e-06
   n=216000, rel_diff=1.7218825e-06
   n=240000, rel_diff=1.9794297e-06
   n=264000, rel_diff=1.8857528e-06
   n=288000, rel_diff=2.141796e-06
   
6. zcbenz

   added

   bug

   on May 29, 2026
7. zcbenz

   closed this as completedin #3609on Jun 4, 2026