perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache#714
Merged
Merged
Conversation
The MLX serving path implements continuous batching, batched prefill, and a prompt-prefix KV cache, but the shipped defaults disabled the first two: `--parallel` (and thus `max_batch_size`) defaulted to 1 and `--max-batch-prefill` to 1, so out of the box the server decoded and prefilled one sequence at a time. The prompt cache was audited and found already default-on (2 GiB budget, APC on), so only the two batching defaults were genuinely off; the third premise was stale. Change the defaults for `mlxcel-server` and `mlxcel serve` to `--parallel 4` and `--max-batch-prefill 4`, add `--no-prompt-cache` as a clean opt-out for the default-on prompt cache, and clamp the effective decode batch to 1 for families that cannot batch (SSM / hybrid / mixed-cache, `supports_batching() == false`) so the default is safe for every architecture. `--parallel 1`, `--no-batch`, `--max-batch-prefill 1`, and `--no-prompt-cache` restore the previous single-client behavior. Batched prefill only engages for families that opt into `supports_batched_prefill()`; others fall back to sequential prefill. Measured on Apple M1 Ultra (Metal, `meta-llama-3.1-8b-instruct-4bit`, `scripts/bench_serving_concurrency.py`): at 4 concurrent clients the new default delivers 1.90x the single-client aggregate throughput (107.9 vs 56.8 tok/s) and cuts mean time-to-first-token under load ~17x (3257 ms to 189 ms), with single-client throughput unchanged and no OOM at B=8. Batched prefill halves cold-prefill TTFT (4215 to 2142 ms). The GB10 `>= 2.5x`-at-4-clients acceptance target and `cargo test --features cuda` are backend-specific and remain pending a CUDA session; full data and the audit are in docs/benchmark_results/serving-throughput-defaults-m1u-2026-07-09.md. Docs (docs/CONTINUOUS_BATCHING.md tuning table, docs/environment-variables.md context-sizing note) and CHANGELOG updated; config-default unit tests adjusted and a default-parse plus no-prompt-cache test added. Closes #628
inureyes
added
type:performance
3 tasks
…mory guard Pairs the #628 batched-decode default (`--parallel 4`) with a default `--kv-cache-budget auto` so the concurrent batch cannot run four full-context sequences into an uncatchable OOM abort. The #122 paged block-budget admission now bounds KV and returns clean backpressure by default; on the dense decode backend the budget is inert, preserving pre-existing behaviour. The default is set both in the CLI arg (`default_value = "auto"` on `mlxcel-server` and `mlxcel serve`) and in `ServerConfig::default()` / `ServerStartupConfig::default()`. Adds an explicit escape hatch: `--kv-cache-budget none` (also `off` / `disabled` / `unbounded` / `0`) leaves the pool unbounded via a new `PagedBudgetDirective::Disabled` variant, resolved before the block-count path so it skips the "geometry unavailable" warning. `auto` and a raw byte count keep their meaning. Also folds in the review-batch fixes: refresh the stale `max_batch_size` doc comment (was "typically 1", now 4 with the non-batching clamp noted); relabel the docs/CONTINUOUS_BATCHING.md scaling table so the 8-client row is attributed to its `--parallel 8` measurement rather than the `--parallel 4` default; add the `--ctx-size` migration note to CHANGELOG (an explicit small ctx-size is now divided across 4 slots and can trip the 512-token/slot floor); and pin the new defaults with tests: `max_batch_prefill == 4` and `kv_cache_budget == Some(Auto)` in `ServerConfig::default()`, the `auto` CLI default plus the `none` / `0` / bytes parse paths, the `Disabled` FromStr keywords, and `resolve_paged_block_budget(Disabled) == None`. Verified on Apple M1 Ultra: lib config/startup/memory_estimate tests (85 passed) and `mlxcel-server` tests (9 passed) green; a live `qwen2.5-0.5b-bf16` run logs `Paged KV block budget: ... blocks` by default and `--kv-cache-budget disabled` with `none`, and `/v1/cache/stats` responds. Refs #628
The paged block-pool budget section still said the default was unbounded, which went stale when cb3ec59 changed the shipped default to `auto` so the batched-decode default cannot run concurrent full-context sequences into an OOM abort. Update the flag/env description to match config.rs and the already-updated CONTINUOUS_BATCHING.md / environment-variables.md, and note `none`/`0` as the escape hatch back to an unbounded pool. Validation: - cargo fmt --all -- --check - cargo clippy --lib --bins --tests --features metal,accelerate -- -D warnings Refs #628
5 tasks
inureyes
added a commit
that referenced
this pull request
Jul 9, 2026
…720) ## Summary Resolves the PR #709 review finding that `paged_decode_attention_pooled` has no `mlxcel-server` caller. The decision recorded here is to retire the pooled decode entry point (and its `select_pooled_paged_dispatch` selector) to a library-only API rather than wire it into the scheduler decode. ## Why retire (occupancy evidence, Apple M1 Ultra) The fused kernel's winning island is single-slab only. A slab is `POOL_SLAB_BLOCKS = 32` block rows (`src/lib/mlxcel-core/src/cache/paged.rs`) and a block holds `DEFAULT_PAGED_BLOCK_SIZE = 32` tokens (`src/server/batch/scheduler.rs`), so one slab is 32 x 32 = 1024 token rows per layer across every sequence resident in the shared pool (`PagedBlockPool::slab_count` returns the layer's whole slab-list length). The selector requires `slab_count <= 1` and `batch_size >= 4`, so the pool stays native only while `B * ceil(len / 32) <= 32`. At the shipped `--parallel 4` default (#714, `n_parallel = 4`) that is `len <= 256` total tokens per sequence, counting prompt, chat template, and generated tokens. That island is negligible in production. The #714 serving-throughput bench drives 512-token prompts, which at `B = 4` need `4 x 16 = 64` block rows = 2 slabs, so every layer is multi-slab from the first decode step (the #331 bench table in ADR 0001 shows every `B >= 4` row at `ctx >= 512` reading `declined`). The pool only appends slabs (#235: existing slabs are never freed), so a request that starts inside the island leaves it permanently within its first few dozen generated tokens. Wiring the pooled path in would thread a new batched-decode arm through every transformer family (high blast radius, jitter-class parity risk) for that transient sliver. Since the pooled path has no server caller, its measured production occupancy is definitionally zero today, so the closed-form geometry is stronger evidence than a thermally-noisy instrumented run. ## What changed - `src/lib/mlxcel-core/src/layers.rs`: library-only doc comments on `paged_decode_attention_pooled`, `select_pooled_paged_dispatch`, the `MLXCEL_PAGED_ATTENTION_NATIVE` override (`NativePagedOverride`), and the per-shape dispatch memo (`PagedDispatchCache`), stating they are not on the `mlxcel-server` decode path and are kept for external mlxcel-core consumers and `examples/paged_attention_kernel_bench.rs`. - `docs/adr/0001-paged-attention-gather-vs-fused-kernel.md`: new decision record with the occupancy derivation, the geometry math, the chosen exit, and what would reopen the question. - `docs/environment-variables.md`, `docs/turbo-kv-cache.md`: `MLXCEL_PAGED_ATTENTION_NATIVE` is now described as a library-consumer control and A/B pin for the bench, not a server knob. ## What was kept vs removed - Kept: the fused kernel (`PagedBlockPool::paged_decode_fused`), the selector, the memo, the env override, and `examples/paged_attention_kernel_bench.rs` (tested, benchmarked library surface). - Kept: `use_native_paged_kernel` scheduler request. It also gates the live `paged_decode_attention_dense_compat` block-table decode path (`src/models/llama3.rs` and the other families), so it is shared plumbing, not dead code. No plumbing was removed. - Removed: nothing. The pooled function, selector, kernel, and bench all remain. ## Test plan - [x] `cargo fmt --all -- --check` - [x] `cargo check --lib --tests --features metal,accelerate` - [x] `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` - [x] `cargo test --release --features metal,accelerate -p mlxcel-core layers::tests::` (53 passed: selector island, boundaries, override pinning) - [x] `cargo test --release --features metal,accelerate -p mlxcel-core ffi_tests::test_fused_paged_decode` (2 passed: gather parity over 200 steps, GQA + batched) Hardware: Apple M1 Ultra. Closes #710
This was referenced Jul 9, 2026
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
…ll (#723) ## Summary Raises the CUDA quantized-matmul (`qmm_sm80`) CTA tile M cap from 64 to 128 on Blackwell consumer GPUs (sm_120/121, `cc_major >= 12`), recovering +31-38% prefill throughput on GB10. Part of epic #623. Closes #637 ## Root cause `qmm_sm80`'s `make_cta_tiler` caps `tile_m` at 64 (Ampere-tuned). On GB10 (sm_121) the Ampere tile leaves the SMs idle: ncu shows **35-47% SM / 52-68% memory throughput**, and 4bit prefill runs at **0.55x** the bf16 ceiling. (The sm_90 Hopper kernel isn't compiled for arch `121`, so Blackwell falls to the Ampere `qmm_sm80`.) ## Fix Arch-gated `make_cta_tiler` (overlay `qmm_sm80.cu`): `tile_m` cap = 128 when `device.compute_capability_major() >= 12`, else 64 (stock). A tile sweep confirmed `tile_m` is the **only** safely-tunable axis (wider `tile_n` / deeper `tile_k` break the fixed-MMA shared-memory layout -> JIT failure). Adds `MLXCEL_QMM_TILE_M/N/K` env overrides as a tuning hatch. ## Benchmark (GB10, full detail committed under `docs/benchmark_results/`) | model | forced tile_m=64 | arch-gated default (128) | delta | |---|---|---|---| | llama-3.1-8b-4bit prefill @8192 | 2213 | 3054 | **+38%** | | qwen2.5-7b-4bit prefill @8192 | 2410 | 3168 | **+31%** | - Greedy parity: default(128) vs forced(64) generated tokens **byte-identical**. - Decode (m=1): **unaffected** (single-seq decode takes the qmv path, never reaches `qmm_sm80`). - 0.55x -> 0.72x of the bf16 ceiling. ## Scope / not addressed - **Addressed**: large-M prefill + batched prefill on sm_120/121 (+31-38%). - **NOT addressed**: small-M batched decode (`M*B` in [2,8)) still uses per-row qmv (no weight amortization); the `tile_m` cap doesn't raise small-M tiles. This is the gap behind #714's `--parallel 4` CUDA regression and needs a dedicated small-M kernel (upstream MLX). #714's `--parallel` default should be made backend-aware until then. ## Validation - C++ overlay change only (no Rust touched). Validated by greedy parity + prefill benchmark (the right gates for a GEMM tile change). - `make release-cuda` builds clean. The arch gate uses runtime `compute_capability_major()`, so it works for both `121` and `121a` builds (both cc 12.1 at runtime). ## Acceptance criteria (#637) - [x] Findings doc in `docs/benchmark_results/` (achieved-vs-ceiling, sm_121 dispatch map, ncu roofline, go/no-go). - [x] >= 25% prefill @8192 on llama-3.1-8b-4bit on GB10 (measured **+38%**), no decode regression, parity green. - [x] Both `121`/`121a` build (arch-gate is runtime cc, not compile-time arch string).
4 tasks
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
With --max-batch-prefill 4 now the default (#628), the server's padded batched-prefill path (run_padded_batched_prefill) engages out of the box for supports_batched_prefill() families and, for mixed-length prompts, runs a single unchunked [B, padded_len] forward that materializes a stacked [B, L, L] FP32 attention mask, an O(B*L^2) transient that ignores prefill_chunk_size. Four concurrent 8k prompts build a [4, 8192, 8192] FP32 mask (~1 GiB) far above the sequential chunked prefill working set, and on the fail-fast serving path an allocation failure is an uncatchable MLX C++ throw that aborts the whole server. This is an availability edge the --kv-cache-budget guard does not model: that budget bounds steady-state KV, not this prefill transient. Bound the drained batched-prefill window by total padded tokens (rows * max_len), not just row count. A cohort of B >= 2 rows padded to L keeps B*L within max_batch_prefill_tokens, so the [B, L, L] mask stays within budget^2 / 2 elements (~2*budget^2 bytes at FP32). Rows past the budget stay queued and prefill on a later tick: short ones re-batch, long ones take the chunked single-sequence path. A head prompt too long to join a two-row batch (2*head_len > budget) skips the batched path entirely via a dispatch-time guard, keeping the attention mask chunked to [chunk, L] instead of an unchunked [L, L]. The drain always takes the head so it makes forward progress (no livelock). The budget is configurable via --max-batch-prefill-tokens (CLI) and MLXCEL_MAX_BATCH_PREFILL_TOKENS (env), with the flag taking precedence, then the env, then the derived default. The default is max_batch_prefill * prefill_chunk_size (the shipped 4 * 512 = 2048), which bounds the FP32 mask to about 8 MiB while keeping a full batch of chunk-sized prompts eligible, so the #714 short-prompt concurrency case (4 x 512-token = 2048 tokens) still batches unchanged. 0 disables the cap for the pre-#715 unbounded behavior. The cap logic lives in pure, unit-tested functions in prefill_cohort.rs (batched_window_admits / batched_prefill_window_len / default_batched_prefill_token_budget) covering fits/spills/boundary/uncapped/long-head/default-budget. The scheduler drains via peek_prompt_len + batched_window_admits and gates dispatch with batched_prefill_admits_head; the value threads from ServerConfig through WorkerSchedulerConfig to a with_max_batch_prefill_tokens builder. Validated: cargo check --lib --tests, cargo test --lib server::batch::prefill_cohort/queue (42 passed, incl. 8 new cap tests), cargo clippy --lib --tests -D warnings, cargo fmt --all --check all clean. The bound is analytic and documented (code comment + docs/CONTINUOUS_BATCHING.md formula); the empirical M1 Ultra 8k-prompt peak-memory and short-prompt TTFT A/B are pending a measurement session (repro commands recorded in the docs). Closes #715
inureyes
added a commit
that referenced
this pull request
Jul 10, 2026
#722) ## Summary `--max-batch-prefill 4` (the #628 default) engages the server's padded batched-prefill path out of the box, and for mixed-length prompts it runs one unchunked `[B, padded_len]` forward that materializes a stacked `[B, L, L]` FP32 attention mask, an `O(B*L^2)` transient that ignores `prefill_chunk_size`. Four concurrent 8k prompts build a `[4, 8192, 8192]` FP32 mask (~1 GiB) far above the sequential chunked prefill working set, and on the fail-fast serving path an allocation failure is an uncatchable MLX C++ throw that aborts the whole server. This bounds that transient by a padded-token budget while keeping the #714 short-prompt concurrency case unchanged. ## Design and bound Cap the drained batched-prefill window by total padded tokens (`rows * max_len`), not just row count. A cohort of `B >= 2` rows padded to `L` keeps `B*L` within `max_batch_prefill_tokens`, and since `L <= (B*L)/2` the `[B, L, L]` mask stays within `budget^2 / 2` elements, i.e. `~2*budget^2` bytes at FP32. - Rows past the budget stay queued and prefill on a later tick: short ones re-batch, long ones take the chunked single-sequence path. - A head prompt too long to join a two-row batch (`2*head_len > budget`) skips the batched path entirely via a dispatch-time guard, keeping the mask chunked to `[chunk, L]` instead of an unchunked `[L, L]`. The drain always takes the head, so it makes forward progress (no livelock). - Default budget = `2 * max_batch_prefill * prefill_chunk_size` (the shipped `2 * 4 * 512 = 4096`), bounding the FP32 mask to ~34 MiB. The #714 case (4 x 512-token = 2048) still batches, with 2x headroom to spare for prompts that tokenize slightly over `prefill_chunk_size`. `0` disables the cap (pre-#715 unbounded behavior). - Precedence: `--max-batch-prefill-tokens` flag > `MLXCEL_MAX_BATCH_PREFILL_TOKENS` env > derived default. ## What changed - `src/server/batch/prefill_cohort.rs`: pure cap logic `batched_window_admits` / `batched_prefill_window_len` (test-only) / `default_batched_prefill_token_budget`, plus unit tests (fits/spills/boundary/uncapped/long-head/default-budget). - `src/server/batch/scheduler.rs`: token-bounded drain in `execute_batched_prefill` (peek + `batched_window_admits`), the `batched_prefill_admits_head` dispatch guard, `resolve_max_batch_prefill_tokens` (env/default), and the `with_max_batch_prefill_tokens` builder. Direct unit tests cover the `configured = Some(_)` precedence branch of `resolve_max_batch_prefill_tokens` and the `batched_prefill_admits_head` boundary (`2 * head_len == budget` admits, `+1` rejects). - `src/server/batch/queue.rs`: `peek_prompt_len` for the drain/guard. - Config plumbing: `--max-batch-prefill-tokens` on both binaries (`bin/mlx_server.rs`, `main.rs`, `commands/serve.rs`) threaded `ServerConfig` -> `WorkerSchedulerConfig` (`cli_input.rs`, `startup.rs`, `config.rs`, `model_provider.rs`, `model_worker.rs`) plus the two touched test literals. - 959bfe1: the empirical short-prompt A/B (see Test plan) caught a boundary regression at the exactly-tight `max_batch_prefill * prefill_chunk_size` default: real ~512-token prompts tokenize slightly over `prefill_chunk_size`, spilling the last row of a 4-client batch and doubling p95 TTFT. The derived default now carries 2x padding headroom (`2 * max_batch_prefill * prefill_chunk_size`, `4096` shipped); the FP32 mask bound stays negligible (~34 MiB) and the 8k-prompt spill path is unaffected. - Docs: `docs/CONTINUOUS_BATCHING.md` (bound formula + flag), `docs/environment-variables.md` (env var), `CHANGELOG.md`. ## Test plan Measured here (Apple M1 Ultra, macOS): - [x] `cargo test --lib -- server::batch::prefill_cohort server::batch::queue`: passing (incl. cap tests, the `resolve_max_batch_prefill_tokens` precedence test, and the `batched_prefill_admits_head` boundary tests). - [x] `cargo check --lib --tests --features metal,accelerate`: clean. - [x] `cargo clippy --lib --tests --features metal,accelerate -- -D warnings`: exit 0. - [x] `cargo fmt --all -- --check`: exit 0. - [x] Empirical peak-memory A/B: four concurrent ~9.8k-token prompts (`meta-llama-3.1-8b-instruct-4bit`), capped default vs `--max-batch-prefill-tokens 0`. Capped peak RSS 4115 MiB vs uncapped 4375 MiB, first response 3.5x earlier (17.6s vs 62s). Full numbers in the [measurement comment](#722 (comment)). - [x] Short-prompt TTFT A/B: `scripts/bench_serving_concurrency.py --concurrency 4 --prompt-tokens 512 --max-tokens 128` against main, the pre-headroom-fix budget, and 959bfe1. The pre-fix budget (`2048`) is the boundary regression this A/B caught (p95 TTFT 2.2x); 959bfe1 (`4096`) restores main parity (mean +0.5%, p95 uniform, aggregate -0.7%). Full numbers in the [measurement comment](#722 (comment)). Closes #715
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ship serving-throughput defaults for multi-client serving:
--parallel 4(batched decode) and--max-batch-prefill 4(batched prefill), clamped to single-slot for families that cannot batch, plus a--no-prompt-cacheopt-out. The batched-decode default is paired with a default--kv-cache-budget automemory guard. Defaults chosen from measurements on Apple M1 Ultra. This changes defaults and documentation, not machinery.Audit first (issue premises vs reality on
main)Per the audit-first mandate (we were burned by a stale premise on #330), each premise was checked against
9808b675before implementing:--parallel/max_batch_sizedefault_value_t = 1,unwrap_or(n_parallel) = 1)--max-batch-prefilldefault_value_t = 1)--prompt-cache-enableddefaulttrue,PromptCacheConfig::default().enabled = true, 2 GiB budget, APC on; satisfied by #228)--no-prompt-cacheopt-outThe prompt cache was already default-on. Only the two batching defaults were genuinely off, so only those changed.
What changed
--paralleldefault 1 -> 4 and--max-batch-prefilldefault 1 -> 4 on bothmlxcel-server(src/bin/mlx_server.rs) andmlxcel serve(src/main.rs), plus theServerConfig::default()/ServerStartupConfig::default()struct defaults.--no-prompt-cacheflag (both binaries) as the highest-precedence opt-out for the default-on prompt cache.supports_batching()gate: the worker (src/server/model_worker.rs) clamps the effective decode batch to 1 for SSM / hybrid / mixed-cache families, so the default is safe for every architecture and the scheduler log reports the effective value.--kv-cache-budgetnow defaults toauto, so the Phase 5: Block-budget admission, eviction, and preemption #122 paged block-budget admission bounds KV for the concurrent batch and returns clean backpressure instead of an OOM abort. A newPagedBudgetDirective::Disabledvariant plus--kv-cache-budget none/0is the escape hatch that leaves the pool unbounded; the guard is inert on the dense decode backend (pre-existing behavior).docs/CONTINUOUS_BATCHING.md(defaults + measured tuning table + budget guard),docs/environment-variables.md(context-sizing note,MLXCEL_KV_CACHE_BUDGETdefault),CHANGELOG.md(behavior change +--ctx-sizemigration note), and a dateddocs/benchmark_results/serving-throughput-defaults-m1u-2026-07-09.md.max_batch_prefill == 4,kv_cache_budget == Some(Auto)); default-parse,--no-prompt-cache,--kv-cache-budgetauto/none/0/bytes, andPagedBudgetDirective::Disabledtests added;ServeArgsfixture updated.Measured (Apple M1 Ultra, Metal;
meta-llama-3.1-8b-instruct-4bit, 512-token prompt, 128 tokens)-p1)-p1)-p8)-p8)At 4 clients: 1.90x aggregate throughput and ~17x lower mean TTFT under load, single-client unchanged. Batched prefill halves cold-prefill TTFT (4215 -> 2142 ms). B=8 on the 8B model: 0 failures, peak RSS 4.4 GB (no OOM on 128 GB). The 8-client batched row was measured at
--parallel 8; the shipped default is--parallel 4(batch caps at 4).Test plan
Measured / verified here (Apple M1 Ultra, Metal):
cargo fmt --all -- --checkclean;cargo clippy --lib --bins --tests -- -D warningsclean.server::config::tests+server::startup::tests+execution::memory_estimate::tests(85 passed).mlxcel-servertests: default-parse,--no-prompt-cache,--kv-cache-budgetauto/none/0/bytes (9 passed);mlxcelcommands::serve::tests(11 passed).max_batch_size=4, max_batch_prefill=4andPaged KV block budget: ... blocks(auto default);--kv-cache-budget nonelogs the clean disable; mamba2 (non-batching) clamps4 -> 1;--no-prompt-cachedisables the store;/v1/cache/statsresponds; concurrent requests served with 0 failures.mlxcel_prompt_cache_hits_totalincrements (28 in the scaling run).Pending a CUDA session (GB10 not available this session):
cargo test --features cuda(config default tests).Closes #628