Clarify whether mlx_closure_compile(shapeless=true) should fuse C API closure bodies
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:
- Is
mlx_closure_compile(shapeless=true)intended to expose the same graph optimization/fusion behavior as Pythonmx.compile(shapeless=True)? - Are there known op families where C API closure compilation should not be expected to fuse, such as
quantized_matmulormlx_fast_rope? - Is there a supported way to inspect whether a compiled C API closure lowered to fewer Metal dispatches/kernels?
- Should downstream projects avoid documenting C API closure compile as equivalent to Python
mx.compileunless they have per-closure benchmark evidence?
Thanks!
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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