Skip to content

Optimize BF16 quantized matmul with FP16 tiles - #4216

Open
justinchuby wants to merge 3 commits into
ml-explore:mainfrom
justinchuby:optimize/bf16-qmm-fp16-tiles
Open

Optimize BF16 quantized matmul with FP16 tiles#4216
justinchuby wants to merge 3 commits into
ml-explore:mainfrom
justinchuby:optimize/bf16-qmm-fp16-tiles

Conversation

@justinchuby

Copy link
Copy Markdown

Summary

  • add an opt-in affine INT4/group-32 QMM path with BF16 device I/O, FP16 threadgroup/MMA tiles, and FP32 accumulation
  • vectorize BF16-to-FP16 activation loads and use 64x64 output tiling
  • keep AOT and Metal JIT tile instantiations consistent and test both wide/narrow variants

Enable with MLX_BF16_QMM_FP16=1; MLX_BF16_QMM_FP16_TILE=32 selects the narrow comparison path.

Performance

M1 Max, M=512, K=6656, N=19968, affine INT4 group size 32:

Path Time
Stock BF16 30.6 ms
Stock FP16 17.2 ms
BF16 I/O + FP16 tiles 17.32 ms

The mixed path is ~1.77x faster than stock BF16 and effectively matches pure FP16 while preserving BF16 I/O.

Validation

  • python/tests/test_quantized.py: 33 tests / 2896 subtests pass
  • AOT and MLX_METAL_JIT=ON builds
  • wide/narrow tiles, small M/N, non-64 K, and batched inputs
  • JIT regression test reproduces the pre-fix tile mismatch and passes with the fix

Use vectorized BF16-to-FP16 loads and wider tiles for affine INT4 matrix-matrix workloads while preserving BF16 I/O. Keep the path opt-in and cover AOT/JIT tile dispatch.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 5f4a437b-39c5-44ac-9bd7-e726eed87db0
@justinchuby
justinchuby force-pushed the optimize/bf16-qmm-fp16-tiles branch from 39835ce to 5c42446 Compare August 12, 2026 22:04
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 12, 2026
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 68eb37ce-1aa5-4ef3-a9e9-ef0176f326e2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants