lmstudio-ai / lmstudio-ai/mlx-engine

[Feature Request] "Agent Mode": Unbounded Prefill & Multi-Slot Caching for High-End Hardware (M3 ultras)

Open
#245 2 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

Hey team! First off, MLX-engine is fantastic for UI-focused tasks and ensuring stability on standard hardware.

However, on high-bandwidth hardware (e.g., M3 Ultra, Mac Studio), the current architecture acts as a governor on performance specifically for Agentic workflows (RAG, Tree-of-Thoughts, ReAct).

I was wondering if I maybe an optional configuration (e.g., --agent-mode or separate flags) could be leveraged to unlock the full throughput of these chips.

**The Problem**

1. Bandwidth Throttling via Chunking The current PROMPT_PROCESSING_CHUNK_SIZE (often 512 tokens) is excellent for UI progress bars. However, on an M3 Ultra (800GB/s bandwidth), this creates significant latency.

* Issue: The GPU spends excessive time syncing with Python and stopping/starting the graph execution every 512 tokens.
* Bottleneck: We are leaving massive amounts of bandwidth on the table by not saturating the hardware with the full prompt in one go.

2. Cache Thrashing (Single-Slot) The current CacheWrapper effectively assumes a linear conversation.

* Issue: Complex agents often branch (e.g., generating "Thought A", backtracking, then generating "Thought B").
* Bottleneck: The engine detects a divergence, trims/clears the cache, and re-computes the prompt. On a 128GB+ RAM machine, we have the memory to store these branches, but the engine discards them.

I was wondering if maybe one way this could be tackled could be having a separate "mode" or run variant ? Maybe for larger M3 ultras we could handle Prefill and slotting differently ? Example:

1. Unbounded Prefill
* Bypass the progress-bar chunking loop.
* Feed the entire prompt_tokens array (up to context limit) to the model in a single forward pass.
* Potential imapact: reduction in Time-To-First-Token (TTFT) for large RAG contexts on Ultra chips.

2. Multi-Slot LRU Caching
* Switch CacheWrapper to use a Dict[PromptHash, KVCache].
* Implement an LRU eviction policy (configurable N slots).
* Potential impact: Instant backtracking for branching agents, utilizing idle RAM to save compute.

The above would improve a lot the life when using LM Studio as a sever with agent CLI's like opencode when having access to large models locally

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating PROMPT_PROCESSING_CHUNK_SIZE and CacheWrapper in the Python engine, then trace the prefill loop and cache divergence handling. Compare the current behavior with the proposed unbounded prefill and multi-slot LRU design, including memory and branching implications. Done should include an agreed configuration, implementation, and performance or regression tests for large prompts and branched conversations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.