lmstudio-ai / lmstudio-ai/mlx-engine

Feature: fused batched decode for concurrent requests (close the gap to mlx_lm.batch_generate)

Open
#342 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.