awni/mlx-skills/fast-mlx

Fast MLX

Workflow

• Looks for opportunities to compile functions of mostly elementwise operations.
• For models with fixed shape inputs or where the shapes don't change much, compile the entire graph
• Replace slow implementations with MLX fast ops
• Identify evaluation boundaries and unintended sync points (mx.eval, item(), NumPy conversions).
• Check dtype promotion and scalar usage; keep precision consistent with intent.
• Review compilation strategy; avoid unnecessary recompiles and closure captures.
• Reduce peak memory via lazy loading order and releasing temporaries before mx.eval.
• Suggest profiling steps if the bottleneck is unclear.

References

• Read references/fast-mlx-guide.md for detailed tips and examples. Use it as the source of truth.

Output expectations

• Provide concrete code changes with brief rationale
• Call out changes that need user confirmation (e.g., enabling async eval or shapeless compile).