Skip to content

Fix matmul precision with broadcast RHS - #3983

Open
Sarves1911 wants to merge 1 commit into
ml-explore:mainfrom
Sarves1911:fix/matmul-broadcast-k1-precision
Open

Fix matmul precision with broadcast RHS#3983
Sarves1911 wants to merge 1 commit into
ml-explore:mainfrom
Sarves1911:fix/matmul-broadcast-k1-precision

Conversation

@Sarves1911

Copy link
Copy Markdown

Proposed changes

Fixes #3953.

The Metal matmul backend was collapsing a broadcast batch dimension into M even when the original operation had M == 1.

For a float32 batched matrix-vector product with a stride-zero broadcasted right-hand operand, this changed the operation from the GEMV path to the TF32-enabled GEMM path. As a result, multiplying by exactly 1.0 could reduce the precision of the right-hand operand.

This change prevents the batch-into-M collapse when the original M dimension is 1, allowing the operation to remain on the existing GEMV path.

A regression test is included that runs the operation with TF32 enabled in a separate process.

Testing

  • Focused regression test passed:
    PYTHONPATH=python/tests python -m unittest test_blas.TestBlas.test_matmul_broadcast_rhs_k1_precision -v
  • Complete BLAS test suite passed:
    PYTHONPATH=python/tests python -m unittest test_blas -v
  • 29 BLAS tests passed.
  • Pre-commit hooks passed for the modified files.
  • The original issue reproduction now reports a maximum error of 0.0 with MLX_ENABLE_TF32=1.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

float32 matmul loses precision to float16 when the right operand has a broadcast batch dimension (also corrupts batch-1 SDPA fallback)

1 participant