perf(server): cap the batched-prefill transient memory by token budget

[inureyes]

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 perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #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-perf(server): cap the batched-prefill transient memory (O(B*L^2) padded mask, unchunked forward) #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):

• 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).
• cargo check --lib --tests --features metal,accelerate: clean.
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings: exit 0.
• cargo fmt --all -- --check: exit 0.
• 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.
• 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.

Closes #715

[inureyes]

Orchestrator empirical validation (Apple M1 Ultra, meta-llama-3.1-8b-instruct-4bit, real mlxcel-server), completing the two pending acceptance measurements.

1. Long-prompt transient A/B (4 concurrent ~9.8k-token prompts, max_tokens 16), capped default vs --max-batch-prefill-tokens 0 on the same binary:

scenario	peak RSS	wall	completion pattern
capped (default)	4115 MiB	65s	staggered: 17.6 / 33.7 / 49.8 / 65.3s (chunked sequential per long head)
uncapped (0)	4375 MiB	~62s	all four at ~62s (one padded batch)

Peak RSS with the cap is lower and total wall time is within ~5%; the cap also delivers the first response 3.5x earlier (17.6s vs 62s). Caveat: 0.5s ps RSS sampling understates short-lived Metal transients, so the RSS delta (~260 MiB) is a floor; the analytic mask bound (uncapped [4, 9772, 9772] FP32 ~= 1.5 GiB vs capped chunk masks ~16 MiB one at a time) remains the primary sizing evidence. Both scenarios completed 4/4 HTTP 200 with prompt_tokens=9772.

2. Short-prompt TTFT no-regression (bench_serving_concurrency.py --concurrency 4 --prompt-tokens 512 --max-tokens 128, identical command on main-built vs branch-built binaries, cold server each):

binary	ttft mean	ttft p95	aggregate tok/s
main (8b2329cf2)	2941.5 ms	2942.4 ms	68.2
branch @ 38bf918 (budget 2048)	3299.4 ms	6557.6 ms	63.4
branch @ 959bfe1 (budget 4096)	2954.8 ms	2955.1 ms	67.7

The middle row is the boundary regression this A/B caught: real "~512-token" prompts tokenize slightly over prefill_chunk_size, so the exactly-tight 4*512=2048 default spilled the last row and staggered prefill (p95 2.2x). Commit 959bfe1 gives the derived default 2x padding headroom (2 * max_batch_prefill * prefill_chunk_size = 4096, mask bound still ~34 MiB); the re-run restores main parity (mean +0.5%, p95 uniform, aggregate -0.7%, all within run noise). The 8k spill behavior is unchanged by the bump (2*9772 > 4096 still routes long heads to the chunked path, and 4*9772 padded tokens still exceed any single window).