[hxu296]

Proposed changes

buffer_sizes is compared against a megabyte threshold (can be set by user with MLX_MAX_MB_PER_BUFFER), but was accumulating data_size() which is in elements, not bytes. This made the commit heuristic dtype-blind, a 4MB float32 buffer counted the same as a 2MB float16 buffer with the same number of elements.

Checklist

Put an x in the boxes that apply.

• 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 (please lmk if new tests are needed)
• I have updated the necessary documentation (if needed)

[angeloskath]

Yep not necessarily a bad change. We 'll have to do some perf testing before merging this though (and adjust the heuristics). Given that this is a heuristic for the time an op will take it isn't necessarily bad that it is in elements, so if it does work better we might have to change the env variable name.

[hxu296]

Makes sense. The change will make command buffer commits more frequent for larger dtypes, so small workloads may not fully hide command buffer encode/schedule overhead. Definitely worth benchmarking and potentially tuning the thresholds / just change env var name for now.

[zcbenz]

The MLX_MAX_MB_PER_BUFFER env is already commonly used by users and apps so we are not going to merge this, I think it would make sense adding a new env using bytes but then the critical part would be re-tuning the heuristics. So I'm closing this PR.