Add periodic MLX.eval() in denoise loops to reduce memory pressure - #2
Add periodic MLX.eval() in denoise loops to reduce memory pressure#2ronaldmannak wants to merge 1 commit into
Conversation
Add evalInterval parameter (default 5) to denoiseLoop() and the Klein pipeline's denoise(). Every evalInterval steps, MLX.eval(current) is called to force materialization and free intermediate compute graph nodes, preventing unbounded memory growth during long denoising runs. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review the evaluation in denoise loop base on https://github.com/awni/mlx-skills/blob/main/mlx_skills/skills/fast-mlx/references/fast-mlx-guide.md, I was under the impression the eval in denoise loop would cause some performance issue |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4ad506be9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| guidance: MLXArray? = nil, | ||
| modelTimestepScale: Float = 0.001 | ||
| modelTimestepScale: Float = 0.001, | ||
| evalInterval: Int = 5 |
There was a problem hiding this comment.
Setting evalInterval to 5 by default means denoiseLoop now triggers MLX.eval periodically in its hot path for every caller that relies on defaults. Because MLX.eval forces materialization/synchronization, this introduces avoidable per-step sync overhead and can noticeably reduce throughput on memory-safe GPU runs; the periodic evaluation behavior should be opt-in rather than the default execution mode.
Useful? React with 👍 / 👎.
| guidance: guidance, | ||
| modelTimestepScale: modelTimestepScale | ||
| modelTimestepScale: modelTimestepScale, | ||
| evalInterval: 5 |
There was a problem hiding this comment.
This hard-coded evalInterval: 5 makes periodic MLX.eval unavoidable for the dev pipeline path, but Flux2DevPipeline.generate does not expose any parameter to disable or tune it. As a result, users cannot recover prior throughput behavior even when memory pressure is not a concern, so this should be configurable (or left to the shared pipeline default) instead of fixed here.
Useful? React with 👍 / 👎.
2 participants
Summary
MLX.eval()calls every 5 steps in the denoise loops of all three pipeline implementationsFiles changed
Sources/Flux2/Flux2Pipeline.swiftSources/Flux2/Flux2DevPipeline.swiftSources/Flux2/Flux2KleinPipeline.swiftTest plan
swift build🤖 Generated with Claude Code