Skip to content

perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #628

Description

@inureyes

Part of #623. Depends on #624 (TTFT/decode-rate telemetry and the concurrency load script are needed to justify the defaults).

Context

The MLX serving path has full continuous batching, chunked prefill, batched prefill, prompt-prefix KV caching, and paged KV storage implemented, but the shipped defaults disable nearly all of it:

  • Effective decode batch size is max_batch_size.unwrap_or(n_parallel).max(1); --parallel defaults to 1 and --max-batch-size to None (src/server/startup.rs:507,792, src/bin/mlx_server.rs:234,272, src/server/config.rs:538). Out of the box the server decodes one sequence at a time.
  • --max-batch-prefill defaults to 1 (src/bin/mlx_server.rs:391), so batched prefill never engages.
  • The prompt cache (radix-trie prefix reuse, src/server/prompt_cache/) only activates when configured (state.prompt_cache.is_some(), gate at src/server/routes/chat.rs:364); it is off by default, so repeated system prompts are re-prefilled every request.

On a bandwidth-bound decode target like GB10, batched decode is the single largest aggregate-throughput lever available (weights are read once per step regardless of B).

Scope

Choose and ship sane defaults for multi-client serving, with measurements, docs, and clean opt-outs. This issue changes defaults and documentation, not machinery.

Implementation plan

  1. Measure on GB10 with perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624's scripts/bench_serving_concurrency.py on llama-3.1-8b-4bit and qwen3-30b-a3b-4bit:
    • --parallel in {1, 2, 4, 8}: aggregate tok/s, per-request decode tok/s, TTFT p50/p95.
    • --max-batch-prefill 1 vs 4 with concurrent arrivals (TTFT effect).
    • Prompt cache on vs off with a shared 1-2k-token system prompt (TTFT effect on cache hits).
    • Watch memory: KV for B sequences must fit; record peak memory per configuration.
  2. Set defaults from the data. Suggested starting point, to be confirmed by measurement:
    • default --parallel (or max_batch_size) = 4 when the loaded model supports batching (supports_batching()), 1 otherwise;
    • default --max-batch-prefill = 4 where supports_batched_prefill();
    • prompt cache enabled by default with a bounded budget, disabled by --no-prompt-cache.
      Implement the defaults in ServerConfig::default() / CLI arg defaults so --parallel 1 --no-batch remains a full escape hatch. Preserve the documented semantics of --no-batch (legacy sequential worker, src/server/model_provider.rs:223).
  3. Guard rails: when --max-kv-size or the paged budget (--kv-cache-budget) is set, admission already handles pressure; verify OOM behavior at B=8 on a 122 GB GB10 with an 8B model, and document sizing guidance.
  4. Update docs/ server pages (flags, defaults, tuning guide section with measured tables) and CHANGELOG.md (behavior change).

Acceptance criteria

  • Defaults changed with measured justification committed to docs/benchmark_results/ (Apple M1 Ultra concurrency scaling table: aggregate tok/s and TTFT for B in {1,2,4,8}).
  • Aggregate decode tok/s at 4 concurrent clients is at least 2.5x the B=1 configuration on llama-3.1-8b-4bit on GB10. Pending CUDA session. Measured 1.90x on Apple M1 Ultra (107.9 vs 56.8 tok/s); GB10's higher bandwidth should exceed 2.5x.
  • Prompt-cache hit path verified end to end: hit counter increments (mlxcel_prompt_cache_hits_total 28, mlxcel_prompt_cache_prefix_tokens_reused_total 14336 in the scaling run). Cache was already default-on (premise stale, satisfied by feat(server): make the default serve config actually reuse prefix KV and validate footprint end-to-end #228).
  • Opt-outs work: --parallel 1 (unit test) and --no-prompt-cache (added; verified no cache-store log) restore current behavior; --no-batch legacy worker preserved.
  • Docs (CONTINUOUS_BATCHING.md, environment-variables.md) and CHANGELOG updated; config default tests updated and passing under --features metal,accelerate. cargo test --features cuda pending CUDA session.

Validation

cargo build --release --features cuda
./target/release/mlxcel-server --model ./models/llama-3.1-8b-4bit &        # new defaults
python3 scripts/bench_serving_concurrency.py --concurrency 4
curl -s localhost:8080/metrics | grep -E "prompt_cache|ttft"

References

  • Defaults plumbing: src/bin/mlx_server.rs:234,272,391, src/server/startup.rs:507,792, src/server/config.rs:538-547.
  • Scheduler admission/interleaving: src/server/batch/scheduler.rs:2585-2622.
  • Prompt cache: src/server/prompt_cache/, adoption at src/server/batch/scheduler.rs:2394-2397, gate at src/server/routes/chat.rs:364.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions