lmstudio-ai / lmstudio-ai/mlx-engine
DEBT: Prefill Trickery
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
PR https://github.com/lmstudio-ai/mlx-engine/pull/224 adds a "rachet" to the prompt progress callback. This is needed so that prompt progress does not appear to go backwards when we perform prefill in the `CacheWrapper` (the component that enables cache re-use for text-only prompts).
When prefill is done by the `CacheWrapper`, the progress callback is called as expected. The callback and a single input token are then passed to the `mlx_lm.stream_generate` method, which calls the progress callback again with "0 out of 1" tokens processed. This causes progress to appear to regress briefly. The ratchet patches this by ignoring any progress callback that is not monotonically increasing.
This is more of a band-aid than a proper fix. A proper solution would be to unify the prefill logic so that we don't need to "trick" `mlx_lm.stream_generate` into thinking there is only a single prompt token to process.
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 PR #224 and the CacheWrapper, then trace how its prefill progress callback reaches mlx_lm.stream_generate. Unify the prefill path so stream_generate does not receive a fabricated single-token prompt and progress no longer needs the monotonic ratchet workaround.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100