[Apple Silicon][MLX] Fix scheduler OOM on long prompts and concurrent requests - #29439
LarrySimingDeng wants to merge 9 commits into
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
Follow-up evidence on the high-memory behavior of the dynamic chunk resize (radix off, ctx 32768, Qwen3-30B-A3B-4bit shape). This is a mock simulation of the real
Without the resize, the static chunk is sized for the worst-case concurrency, so when the running cap auto-grows to fill the budget it stays around 460 to 576 tokens no matter how large the machine is. The dynamic chunk tracks live headroom, so an idle or lightly loaded box uses the large chunk that fits, the same a user would otherwise get only by setting Caveats: mock plus arithmetic validated against the 24 GB live evidence in the description, not a large-memory Mac run; 24 GB still fails fast at this context (the dynamic path does not weaken the OOM fix); and 64 GB shows no gain because its concurrency is already 1. |
|
Retested on an Apple M4 Pro (24 GB), macOS 26.3, MLX 0.31.2, radix off — independent repro. 🙌 Unit tests (GPU-free): 81/81 pass. Serving (Qwen3-30B-A3B-4bit): my box is a bit tighter than your M5 Pro, so the full 40960-ctx config fail-fasts at startup with your clean
No |
changminbark
left a comment
changminbark
left a comment
There was a problem hiding this comment.
Overall LGTM, just a few comments
| # Memory-safe chunked prefill on Apple Silicon. A long prefill whose activation peak | ||
| # exceeds the Metal working set aborts the scheduler with an uncatchable command-buffer | ||
| # OOM, so the chunk is sized to keep that peak below the limit. | ||
| _MLX_PREFILL_ACT_BYTES_PER_QHEAD = 32 # conservative per-(token x ctx) cost per q-head |
There was a problem hiding this comment.
Would be nice if you could comment how you derived these numbers.
There was a problem hiding this comment.
Good call. Added a derivation block above the constants and grouped them: ACT_BYTES_PER_QHEAD and SAFETY are conservative startup seeds the in-server probe refines with a real measurement, POOL_SLOTS_SLACK is a ratio, and MIN_PREFILL_CHUNK / MIN_POOL_SLOTS / the PROBE_* values are hard floors and probe-replay bounds. Each line now says which kind it is and where the number comes from.
| pool_size = kv_budget // bytes_per_slot if bytes_per_slot else 0 | ||
| if need_slots is not None: | ||
| pool_size = min(pool_size, need_slots) | ||
| if context_length is None: |
There was a problem hiding this comment.
Where does the 256 come from? and doesn't this if statement execute mutually exclusively with the if statement above (so can't we make this an else statement)
There was a problem hiding this comment.
256 is a KV-pool floor in tokens for the branch where the context length is unknown, so the pool can't collapse to near-zero. It happens to equal _MLX_MIN_PREFILL_CHUNK, but that one is a chunk floor and is easy to conflate, so I split it out as a named _MLX_MIN_POOL_SLOTS. And you're right on the if statement, the two are mutually exclusive. Folded into if/else.
|
@LarrySimingDeng Could you check whether this server crash is related to the issue I opened earlier? It seems @karanb192 hasn't had much time to investigate it. |
|
@yeahdongcn Yes, and I ran into it myself: the machine reboot you describe in #21443 hit me while I was developing this PR. I dug into it on my machine, and it turns out to be what this PR fixes. The reboot is a kernel panic in Apple's GPU driver, not a system-RAM or paging death. The three lines that matter from that session's panic: IOGPUFamily tried to release a GPU memory object it couldn't find in its own table, and the compressor and swap were both healthy, so this is a GPU memory-bookkeeping fault, not the machine running out of RAM. That panic and the process crash this PR fixes ( Worth flagging that this was never the "treat system RAM as GPU memory" problem from #21443: the MLX backend already caps to The one piece I can't close out myself: I verified the sizing and the process-crash path on 24 GB, but at that size the tight configs now fail-fast, so I can't run the long-prompt trajectory that would show the panic path is actually gone. Thanks @LijuanTang94, independently reproduced the cap / gate / queue behavior on an M4 Pro 24 GB, but neither of us can reach the OOM edge safely. Nailing that last piece really wants a larger-memory Mac (64 GB+) running this branch under load with |
Document each constant's derivation, extract the 256 pool floor as _MLX_MIN_POOL_SLOTS, and fold the mutually-exclusive pool branch into if/else. No behavior change.
|
@yeahdongcn The reporting gap #21443 describes was still live on main; fixed now in #29949 (all three callsites capped to One correction to my earlier "yes": your March reboots predate the MLX backend (#20342 merged the same day you filed the issue), so they were almost certainly #21443's reporting overcommit on the torch-MPS side, not the MLX-path bug this PR fixes. Same failure mode (crossing the Metal working set), two different causes, each fixed in its own PR. |
# Conflicts: # python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
# Conflicts: # python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
pr-test-mlx.yml selects stage A tests via run_suite.py suite registration since sgl-project#30121.
…cap composition initialize() reads _mlx_max_running_requests unconditionally, but the resolver tests in test_max_running_requests.py construct the stub via __new__, bypassing __init__. Default the field to None as a class attribute (same pattern as canary_manager / prefill_aware_swa) so those harnesses keep working. Add two initialize()-level cases pinning the scheduler-cap composition: the runner's memory-safe value caps _resolve_max_running_requests when set, and the resolver default stands when the runner did not auto-size the pool (also exercising the class-attribute default).
|
@yeahdongcn This one is still waiting for a first review, could you take a look when you have time? The branch is up to date with the latest main: |
4 participants
Motivation
Part of the Apple device support roadmap (#19137). This fix builds on @yeahdongcn's MLX backend work, in particular the native MLX execution backend (#20342) and radix cache support (#21509); it targets the radix-off path of that backend, which is the default MLX serving config.
Serving a larger model with the MLX backend on Apple Silicon can crash the whole server, because GPU memory crests the Metal recommended working set and the scheduler subprocess is killed by an uncatchable Metal command-buffer out-of-memory (
kIOGPUCommandBufferCallbackErrorOutOfMemory, thenstd::terminate, thenSIGABRT, scheduler exits with -6). The error is raised on a GPU completion-callback thread, so it cannot be caught in Python and surfaces asynchronously on a latereval; it has to be prevented up front, not handled. On a 24 GB M5 Pro running Qwen3-30B-A3B-4bit (radix cache disabled, the default MLX serving config), there are two independent triggers.Trigger 1, a single long prompt. The backend auto-sizes its attention KV pool against the Metal working-set limit (about 17.76 GiB on this box), which on Apple Silicon is only a no-pressure recommendation, and it sizes that pool from a memory snapshot taken before any forward has run. That snapshot undercounts the steady-state resident set: with radix off, each request holds a private
ContiguousAttentionKVCacheallocated to its full span on first write, plus the overlap pipeline's caches and MLX runtime buffers. Measuringmx.get_active_memory()inside the running server puts the real baseline at about 17.12 GiB, versus the about 16.5 GiB the pre-forward snapshot implies (a roughly 0.6 GiB gap). With that undercount the prefill chunk stays effectively unbounded, and a chunk near 400 tokens peaks about 0.18 GiB under the limit, close enough that allocator jitter tips it over; a 256-token chunk leaves about 0.33 GiB and is stable. The safe chunk therefore has to be sized from the real resident baseline, which only a forward reveals.Figure 1. Measured prefill peak memory vs chunk size (Qwen3-30B-A3B-4bit, radix off, mem-fraction 0.95, 24 GB M5 Pro). The pre-fix unbounded prefill peaks just 0.18 GiB under the Metal working-set limit at chunk about 398 and crashes on allocator jitter; a 256-token chunk leaves 0.33 GiB and is stable; the probe caps the chunk at 196, well clear of the limit.
Trigger 2, request concurrency. Each radix-off request pre-allocates a fixed full-span KV cache (about 0.4 GiB for this model) regardless of prompt length, so the box holds only about four of them above the weights. But the scheduler's running cap came from the MLX stub, which computed
max_running_requests = min(max_total_num_tokens // 2, 4096)(1280 here) and ignored--max-running-requests. Launching with--max-running-requests 1still ran 3 concurrent (#running-req: 3) with the token pool nearly empty, while GPU memory was already pushed to about 17.15 GiB by the three caches. A handful of concurrent requests, even short ones, then crest the working set and trigger the same uncatchable OOM.Modifications
Long-prompt prefill:
model_runner.py_compute_pool_size: clamp the KV budget to real free memory, cap the radix-off pool at the running set's context need so the remainder stays as prefill headroom, model the real resident KV (private per-request caches, not the token pool), and derive a conservative safe prefill chunk. Raise a clear, actionable startup error if even a minimum-size prefill cannot fit.model_runner.py_calibrate_prefill_chunk/_run_prefill_probe: after load and before serving, replay a real chunked prefill to measure the steady-state resident baseline and the per-token transient, then refine the chunk down from the measured baseline (never up). This self-calibrates per model and per machine, so there are no box-specific magic constants; any measurement failure keeps the conservative startup estimate.scheduler_mixin.py_mlx_reject_unfittable_prefills: in the overlap loop, re-check each queued prefill against live device plus system memory and reject one that no longer fits via the existing prompt-length path (set_finish_with_abort, client gets HTTP 400) instead of letting it abort the scheduler.tp_worker.py: capchunked_prefill_sizeto the runner's safe chunk.model_runner.pylive_safe_prefill_chunkandscheduler_mixin.py_mlx_resize_chunked_prefill_budget: the static chunk above is sized for the worst case (allmax_running_requestsprivate caches resident), so when the concurrency cap auto-grows to fill the budget it stays low even on a large machine. In the overlap loop, before the admission gate, recompute the safe chunk from live headroom and raisechunked_prefill_sizeto it, so a long prompt prefills in a bigger chunk and fewer forwards when load is light. It reuses the same formula, measured cost, and conservative full-context span as the static cap and the gate, so it only raises the ceiling to what live memory allows; the static cap and the gate stay the backstops. Radix off and overlap loop only.Request concurrency:
model_runner.py_memory_safe_max_running: resolve the scheduler concurrency cap from the same memory budget. Honor an explicit--max-running-requests; otherwise default to the most requests whose private caches (plus the overlap pipeline's spare caches) still leave a minimum prefill's activation headroom. This is the inverse of the chunk-sizing math, so the cap and the chunk stay consistent, and it floors at 1. Exposed via themax_running_requestsproperty.tp_worker.py/model_runner_stub.py: the worker passes the runner's memory-safe value to the stub, which uses it as the scheduler's running cap instead ofmin(max_total_num_tokens // 2, 4096), falling back to the old heuristic only when the runner did not auto-size it (radix on or explicit pool).Scope: this models the radix-off path, the default MLX serving config and where the crashes were reported. Radix on uses a single greedy, pre-allocated shared pool whose memory profile differs from the radix-off private caches the model assumes, so it keeps its prior pool sizing and concurrency behavior unchanged. Added
test/registered/unit/hardware_backend/mlx/test_mlx_pool_sizing.py, a GPU-free suite (mocked device and system memory, mocked probe forwards) covering startup sizing, the concurrency cap, the live admission gate, the scheduler gate, probe calibration, the dynamic chunk resize, and the radix-on pass-through.Accuracy Tests
Unit tests (no GPU; mocked memory and probe).
pool_sizingis the new suite (37 cases); the others guard the surrounding MLX paths:Live serving with
SGLANG_USE_MLX=1on M5 Pro 24 GB (Metal working set about 17.76 GiB):#running-reqcapped at 1,#queue-reqpeaks at 3 (the rest wait), 4/4 return 200, zero kIOGPU. Before the fix the same load ran#running-req: 3and piled caches toward the limit.max_running_requests1280 to 1, chunk capped to 196, prompt served in 11 chunks, 200, zero kIOGPU.RuntimeError(KV pool 0 tokens, hint to raise mem-fraction to about 0.93), zero kIOGPU, process exits cleanly rather than crashing on the first long prompt.Startup log on the 30B run (mf 0.95) showing the probe-measured baseline and the two caps:
Speed Tests and Profiling
Mostly a stability fix: the memory-safe caps add no slowdown in the constrained case (about 1000 to 1130 tokens/s input at steady state on the 24 GB 30B run; concurrency above the safe cap queues instead of overflowing). The one intentional speed change is the dynamic chunk resize, which recomputes the prefill chunk from live headroom each step: because the static cap is sized for the worst case (all max_running_requests caches resident), a machine with spare memory or under light load would otherwise prefill long prompts in needlessly small chunks, and the resize lets those cases use a larger chunk and fewer forwards. It reuses the same memory-safety formula, so it never raises the chunk beyond what live memory allows or trades away the OOM protection. The startup probe adds one short replay (a few hundred ms) once, before serving.
Known limitations
--disable-overlap-schedulethe startup caps (chunk size and concurrency) are the sole guard; the concurrency cap is enforced where the prefill batch is built, which both loops share, so it still applies. That non-default path also has a separate pre-existing MLX issue (a torch device mismatch in the generic loop'sfuture_map) unrelated to this change.Checklist
CI States
Latest PR Test (Base): ❌ Run #30196851174
Latest PR Test (Extra): ❌ Run #30196851034