ml-explore / ml-explore/mlx-c

Clarify whether mlx_closure_compile(shapeless=true) should fuse C API closure bodies

Open
#116 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
243
Forks
89
Avg merge
14d 43m
Merged PRs (30d)
2

Description

Question

Could you clarify the expected optimization contract for mlx_closure_compile(..., shapeless=true) in the C API?

Specifically: should C API users expect it to perform graph optimization / kernel fusion comparable to Python mx.compile(..., shapeless=True) for a closure body, or should it be treated primarily as a compiled callable/cache boundary where fusion behavior must be validated per op family?

Why I am asking

In AX Engine we tested a few opt-in prefill experiments with MLX 0.31.x / mlx-c 0.6.x. The closure bodies were correct and engaged, but they did not produce measurable throughput improvement:

Probe scope Result
one quantized_matmul inside mlx_closure_compile(shapeless=true) about -0.6% prefill, within noise
dense FFN body: quantized_matmul + activation + multiply + quantized_matmul about -0.5% prefill, within noise
Q+K mlx_fast_rope pair about -0.5% to +0.5% prefill, within noise

This surprised us because the Python MLX compile documentation says compile() can build and optimize the compute graph and fuse certain operations. We may have incorrectly assumed the C API closure compile path has the same practical fusion behavior for these closure bodies.

Related note

I found #104, but that issue appears to be about a shapeless=True correctness/stale-result case involving reductions. This issue is only asking about expected optimization/fusion semantics and what C API users should document or rely on.

What clarification would help

A short answer to any of these would be enough:

  1. Is mlx_closure_compile(shapeless=true) intended to expose the same graph optimization/fusion behavior as Python mx.compile(shapeless=True)?
  2. Are there known op families where C API closure compilation should not be expected to fuse, such as quantized_matmul or mlx_fast_rope?
  3. Is there a supported way to inspect whether a compiled C API closure lowered to fewer Metal dispatches/kernels?
  4. Should downstream projects avoid documenting C API closure compile as equivalent to Python mx.compile unless they have per-closure benchmark evidence?

Thanks!

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the mlx_closure_compile(shapeless=true) C API entry point and compare its documented contract with Python mx.compile(shapeless=True). Review the reported quantized_matmul, activation, and mlx_fast_rope probes; done means documenting whether fusion is guaranteed, operation-dependent, and inspectable for C API users.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
api, documentation, performance
Issue type
Documentation
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.