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.

[Copilot]

Task { ... } inherits the caller’s actor context. If prewarm() is called from the main actor/UI, the model load + tokenization work inside this task can end up running on the main actor and cause UI hitches. Prefer running this as a detached/background task (e.g., Task.detached or explicitly hopping off the main actor) and consider setting an appropriate priority for prewarming work.

[Copilot]

The Task body can throw (loadContext / prepare), but the error is never handled or surfaced. Since prewarm is fire-and-forget, it should catch and intentionally ignore (or otherwise report) failures so prewarm errors don’t get silently lost in a failed task.

Suggested change

	let context = try await loadContext(modelId: modelId, hub: hub, directory: directory)
	// Build chat history similar to respond() to prime the cache effectively
	var chat: [MLXLMCommon.Chat.Message] = []
	// Add system instructions if present
	if let instructions, !instructions.isEmpty {
	chat.append(.init(role: .system, content: instructions))
	}
	// Add prompt prefix or minimal user message
	let promptText = promptPrefix?.description ?? "."
	chat.append(.init(role: .user, content: promptText))
	// Convert tools to MLX format
	let toolSpecs: [ToolSpec]? =
	tools.isEmpty
	? nil
	: tools.map { convertToolToMLXSpec($0) }
	let userInput = MLXLMCommon.UserInput(
	chat: chat,
	processing: .init(resize: .init(width: 512, height: 512)),
	tools: toolSpecs
	)
	// Prepare input - triggers tokenization and processor initialization
	_ = try await context.processor.prepare(input: userInput)
	do {
	let context = try await loadContext(modelId: modelId, hub: hub, directory: directory)
	// Build chat history similar to respond() to prime the cache effectively
	var chat: [MLXLMCommon.Chat.Message] = []
	// Add system instructions if present
	if let instructions, !instructions.isEmpty {
	chat.append(.init(role: .system, content: instructions))
	}
	// Add prompt prefix or minimal user message
	let promptText = promptPrefix?.description ?? "."
	chat.append(.init(role: .user, content: promptText))
	// Convert tools to MLX format
	let toolSpecs: [ToolSpec]? =
	tools.isEmpty
	? nil
	: tools.map { convertToolToMLXSpec($0) }
	let userInput = MLXLMCommon.UserInput(
	chat: chat,
	processing: .init(resize: .init(width: 512, height: 512)),
	tools: toolSpecs
	)
	// Prepare input - triggers tokenization and processor initialization
	_ = try await context.processor.prepare(input: userInput)
	} catch {
	// Intentionally ignore prewarm failures (model will be loaded on demand)
	// You may replace this with a more sophisticated logging mechanism if desired.
	print("MLXLanguageModel prewarm failed for modelId \(modelId): \(error)")
	}

[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?