perf(server): cap the batched-prefill transient memory (O(B*L^2) padded mask, unchunked forward) #715

Problem / Background

• With PR perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714 (issue perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #628) shipping --max-batch-prefill 4 as the default, execute_batched_prefill newly engages out of the box for families with supports_batched_prefill(). Its padded path run_padded_batched_prefill (src/server/batch/scheduler.rs around lines 3339-3380) runs a single unchunked [B, padded_len] forward and, for mixed-length prompts, materializes a stacked [B, L, L] attention mask. Four long prompts arriving together therefore produce an O(B*L^2) transient allocation that ignores prefill_chunk_size (default 512), a peak far above the old chunked single-sequence prefill working set.
• Flagged as a MEDIUM-adjacent resource concern in the PR perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714 security review: the serving path is fail-fast (run_core_thread_or_abort; MLX C++ allocation failure is an uncatchable throw), so a memory spike aborts the whole server rather than 503ing one request. The new default --kv-cache-budget auto guard bounds steady-state KV, but the prefill mask spike is a transient that admission does not model.

Proposed direction

• Cap the drained batched-prefill window by TOTAL tokens (not just row count), e.g. spill rows beyond a max_batch_prefill_tokens budget to the next tick, so B*L stays bounded.
• Or chunk the batched prefill itself (respect prefill_chunk_size inside the padded cohort), which also removes the unchunked-forward memory cliff.
• Consider a mask representation that avoids materializing [B, L, L] for the mixed-length case (per-row lengths / additive mask broadcast), if MLX supports it cheaply.

Acceptance criteria

• A bound exists on batched-prefill transient memory as a function of configuration (documented formula).
• Four concurrent long prompts (e.g. 8k each) on an 8B model no longer spike above the equivalent sequential chunked prefill peak by more than a documented factor.
• No TTFT regression for the short-prompt concurrent case that motivated the default (PR perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714 benchmark note).

References

• PR perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #714 / issue perf(server): serving-throughput defaults: parallel decode, batched prefill, prompt cache #628 (defaults), issue Phase 5: Block-budget admission, eviction, and preemption #122 (paged budget admission), src/server/batch/scheduler.rs (execute_batched_prefill, run_padded_batched_prefill, plan_prefill_cohorts).