Implement prewarm for MLXLanguageModel

[noorbhatia]

• Implements prewarm() for MLXLanguageModel that improves first response time.
• Prewarms the model with instructions, tools and prefixPrompt

[Copilot]

Pull request overview

Implements prewarm(for:promptPrefix:) for MLXLanguageModel to reduce first-response latency by loading the model context and priming the MLX processor with session instructions, tools, and an optional prompt prefix.

Changes:

• Add MLXLanguageModel.prewarm(for:promptPrefix:) implementation.
• Prewarm loads/caches ModelContext and calls context.processor.prepare(input:) with a minimal chat + tool specs.
• Include session instructions and optional prompt prefix in the prewarm input.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mattt]

Unless "." has special significance in MLX, this makes me think that promptPrefix should be non-optional (and maybe non-empty?)

What do you think?

[mattt]

This seems like the kind of thing that we'd want to parameterize in the method, rather than hard-code.

[mattt]

Thanks for opening this PR, @noorbhatia!

I think this kind of functionality gets into the realm of KV cache management, which so far this implementation hasn't attempted to support. At a high-level, I'd expect an API that has some concept of prewarming a common prefix of tokens, caching that, and then reusing for various suffixes. Most likely, cache selection and management would be automatic; I'm not sure yet what controls we'd want to expose.

Can you say more about how you understand the problem?

[noorbhatia]

Thanks for opening this PR, @noorbhatia!

I think this kind of functionality gets into the realm of KV cache management, which so far this implementation hasn't attempted to support. At a high-level, I'd expect an API that has some concept of prewarming a common prefix of tokens, caching that, and then reusing for various suffixes. Most likely, cache selection and management would be automatic; I'm not sure yet what controls we'd want to expose.

Can you say more about how you understand the problem?

My problem: The very first respond() call has significant latency because the model must be loaded from disk, transferred to GPU memory, and the processor initialized. This cold start happens every time the app launches or when a model is first used.

My understanding of prewarm is simply to have the model ready before the user sends their first query, so respond() can start generation immediately.

I'd like to understand your vision and would be happy to implement a KV cache based solution if you could point me to the right direction.

[mattt]

@noorbhatia Thanks for elaborating — that's really helpful.

Let's break the cold start into two parts:

• First, load the model
• Second, create + cache the context with a given prompt prefix.

I suspect that the first step—loading the model—is the bulk of the time spent waiting. So let's try solving that first.

Then, once we have that, we can implement all of the KV cache infrastructure needed to make the promptPrefix parameter do something.

Does that track with your mental model of the problem?

[noorbhatia]

Understood. And solving the first step, loading the model should be simply calling loadContext in prewarm?

Perhaps we can expose loadContext but that would require a change in LanguageModel's API contract. What do you suggest?

[mattt]

Understood. And solving the first step, loading the model should be simply calling loadContext in prewarm?

Yes, exactly.

Perhaps we can expose loadContext but that would require a change in LanguageModel's API contract. What do you suggest?

I'd be interested to see how far we can get within the constraints of the existing Foundation Models API abstraction before we expand the surface area. Let's revisit this when we dig into KV caching.

[noorbhatia]

Great, I'll update the PR. Thanks @mattt !

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.