lmstudio-ai / lmstudio-ai/mlx-engine
Feature: fused batched decode for concurrent requests (close the gap to mlx_lm.batch_generate)
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
Thanks for the continuous-batching work in v1.8.5 — clear step up. This requests the next lever: **fused batched decode** for concurrent requests, which is *orthogonal* to the prefill/caching improvements tracked in #245.
### Evidence (M5 Max, 128 GB)
Same 4-bit MLX weights, single-stream `mlx_lm.generate` vs offline `mlx_lm.batch_generate`:
| model | single-stream decode | `batch_generate` aggregate | batch speedup |
|---|---|---|---|
| gemma-4-e4b (4B dense) | ~130 tok/s | ~2,800 tok/s (B≈128) | **~21×** |
| Qwen3.6-35B-A3B (MoE) | ~137 tok/s | ~326 tok/s | ~2.4× |
| Qwen3.6-27B (dense) | ~33 tok/s | ~84 tok/s | ~2.5× |
On a bandwidth-bound chip, decoding one token ≈ reading the whole model once. A fused batch reads the weights **once** and applies them to N sequences, so small dense models gain ~20×. (Big/MoE models gain less because a batch activates a wider weight footprint — but they still leave a lot on the table.) The aggregate from a single fused step is far above what the server currently delivers under the same concurrency.
### Request
Route eligible in-flight **decode** steps (same model, compatible sampling params) into a single shared-weight batched step — i.e. vLLM/SGLang-style fused decode — instead of near-isolated per-request streams. This is distinct from:
- **#245** (unbounded prefill + multi-slot cache → targets TTFT and branching), and
- **#176** (Qwen3 think-block cache wastage).
It targets **aggregate decode throughput under concurrency**, which is the dominant cost for batch/agentic fan-out workloads.
### Upstream context / correctness
- `ml-explore/mlx-lm` #499 (support batching in `mlx_lm.server`) and #178 (vLLM/SGLang-style parallel) track the primitive.
- Please ensure any batched-KV decode path avoids `ml-explore/mlx-lm` #965 (cross-request KV contamination at ≥16 concurrency) and #754 (merge crash on mixed cached/empty batches).
Happy to share the benchmark scripts.
Contributor guide
Research direction
The issue names no repository files or tests; begin by locating the server's continuous-batching and decode entry points, then compare them with the cited mlx_lm.server batching work. Completion means compatible in-flight requests use one shared decode step, with correctness checks covering cross-request KV contamination and mixed cached/empty batches.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100