[MLX] Honor --max-running-requests in the model runner stub - #30547
Conversation
MlxModelRunnerStub.initialize() hardcoded
max_running_requests = min(max_total_num_tokens // 2, 4096)
and never consulted server_args.max_running_requests. On MLX the scheduler
reads its concurrency cap from model_runner.max_running_requests (the base
CUDA path honors the flag in model_runner_kv_cache_mixin._resolve_max_num_reqs),
so the stub's value is the sole cap -- and --max-running-requests was silently
ignored. Users capping concurrency on a memory-constrained Mac, or setting it
to 1 for deterministic/debug runs, got a different server than requested
(--max-total-tokens is already honored via pool_size, making the gap
inconsistent).
Extract MlxModelRunnerStub._resolve_max_running_requests(), mirroring the base
clamp: honor the flag (split per dp worker, capped by the KV pool capacity, with
the same reduction warning) and fall back to the previous default when the flag
is unset.
Add a regression test that exercises the resolver directly (flag honored, split
per dp worker, clamped to capacity, default when unset).
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
hi, @yeahdongcn pls take a look when you are free, Thanks! |
test/registered/unit/README.md requires unit tests to inherit from sglang.test.test_utils.CustomTestCase (adds CI retry behavior) instead of raw unittest.TestCase.
No problem (just invited you to the slack channel, please check your email). |
The mlx-unit-test job enumerates its model-free test files explicitly; without this entry the new test never executes in CI (the cpu suite skips it because mlx is not installed there). The test is model-free (mocked stub, pure integer arithmetic), matching the job's contract.
Resolve .github/workflows/pr-test-mlx.yml by taking main's version: since #30121 the MLX lane dispatches test/run_suite.py (register_mlx_ci markers) instead of enumerating test files inline, superseding the workflow-list registration added earlier on this branch. The test is registered via the marker in the follow-up commit.
Since #30121 the MLX CI lane dispatches test/run_suite.py --hw mlx --suite stage-a-unit-test-mlx, which discovers tests via the register_mlx_ci marker. Add it (alongside the existing cpu marker, mirroring the sibling files) so the test actually executes on the Apple Silicon lane.
| return min(capacity_cap, 4096) | ||
|
|
||
| requested_per_worker = requested // self.dp_size | ||
| resolved = min(requested_per_worker, capacity_cap) |
There was a problem hiding this comment.
This clamp does not account for the separate auxiliary-state pool used by hybrid/linear-attention models. I reproduced this on the PR head with max_total_num_tokens=100, max_running_requests=4, and max_mamba_cache_size=2: initialization resolves four request slots backed by only two auxiliary slots, and allocating the third request raises AssertionError: Not enough MLX auxiliary state slots from MlxAuxiliaryStateReqToTokenPool.alloc().
The base resolver also bounds hybrid concurrency by max_mamba_cache_size // self._calculate_mamba_ratio() and rejects a zero result. Please add the MLX-appropriate auxiliary-state bound/error here and cover the hybrid initialize() plus request-allocation path. The current _stub() tests only exercise token capacity and DP splitting, so they cannot catch this mismatch between the two pools.
There was a problem hiding this comment.
Great catch, thanks for the repro. Fixed: _resolve_max_running_requests now mirrors the base resolver's mamba bound.
There was a problem hiding this comment.
Thanks, the update resolves the original hybrid allocation failure: concurrency is now bounded before pool construction, an infeasible auxiliary-state pool fails at startup, and the end-to-end test covers real request plus auxiliary-slot allocation.
I think one configuration path still needs correction. The new four-slot ratio represents headroom for radix-held snapshots, but it is also applied under --disable-radix-cache. With radix disabled there are no radix snapshots to reserve; MlxAuxiliaryStateReqToTokenPool.alloc() consumes exactly one auxiliary slot per live request, and the base resolver likewise uses a ratio of one in this mode.
There was a problem hiding this comment.
Fixed in the latest commit. The ratio is now selected on disable_radix_cache
Address review: on hybrid / linear-attention models each running request
allocates one auxiliary-state slot out of max_mamba_cache_size, so a
concurrency cap the pool cannot back crashed mid-serving
("AssertionError: Not enough MLX auxiliary state slots" on the third
allocation with max_running_requests=4, max_mamba_cache_size=2).
Mirror the base resolver's mamba bound in _resolve_max_running_requests:
cap the resolved concurrency (both the requested and the default path) at
max_mamba_cache_size // 4 and reject a zero result at startup with an
actionable error. Extract the ratio into
MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO, shared with the default
auxiliary-pool sizing in initialize() so the sizing and the bound cannot
drift apart.
Add hybrid coverage: resolver bound (requested + default paths, zero
reject, unset-flag and non-hybrid negative branches) and the end-to-end
initialize() + request-allocation path, including the reviewed repro.
|
|
||
| aux_state_size = self.server_args.max_mamba_cache_size | ||
| if self.mambaish_config is not None and aux_state_size is not None: | ||
| ratio = MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO |
There was a problem hiding this comment.
This fixed ratio over-reserves radix snapshot headroom even when server_args.disable_radix_cache is true. I reproduced the updated behavior with max_running_requests=8 and max_mamba_cache_size=8: the resolver returned 2, while a real MlxAuxiliaryStateReqToTokenPool(size=8, auxiliary_state_size=8) successfully allocated all eight requests because each live request uses one slot.
This also differs from ModelRunnerKVCacheMixin._calculate_mamba_ratio(), which returns 1 when radix caching is disabled. Please select the ratio based on disable_radix_cache and cover both resolver and initialize() behavior for the no-radix hybrid path.
There was a problem hiding this comment.
fixed. The ratio now comes from _aux_state_slots_per_request(): 1 when disable_radix_cache is set, MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO: 4 otherwise, mirroring _calculate_mamba_ratio.
Address review: the fixed 4x ratio reserves headroom for radix-held snapshots and chunk track buffers, but with --disable-radix-cache the MLX prefill path returns before any tracked-state store, so each live request holds exactly one auxiliary slot and the headroom can never be used (max_running_requests=8 with max_mamba_cache_size=8 was cut to 2 despite the pool backing all 8). Select the ratio via _aux_state_slots_per_request, mirroring ModelRunnerKVCacheMixin._calculate_mamba_ratio: 1 with radix disabled, 4 otherwise. MLX's auxiliary pool rejects the mamba extra-buffer modes at construction, so no additional ratio term applies. The helper feeds both the concurrency bound and the default pool sizing, keeping them in sync in both modes. Cover the no-radix hybrid path: resolver ratio selection, the end-to-end initialize() + allocation of every request slot one-to-one, and the default-sizing drift guard.
yeahdongcn
left a comment
yeahdongcn
left a comment
There was a problem hiding this comment.
There is another comment should be left to python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py:269:
Could this override call free_mamba_cache(req) before super().free(req) (or otherwise make the release contract explicit)? The current implementation leaves req.mamba_pool_idx set and leaves the auxiliary slot unavailable. That is harmless only when MlxAuxiliaryStateComponent.cleanup_after_caching_req() has already transferred or freed the state; it is not safe for the ChunkCache no-radix release path.
Please cover both cases: a request whose auxiliary state is still owned by the live request, and a request whose state has already been handed to the radix component. The latter must not double-free a slot.
| the mamba extra-buffer modes (``enable_mamba_extra_buffer`` is | ||
| rejected at pool construction), so no additional ratio term applies. | ||
| """ | ||
| if self.server_args.disable_radix_cache: |
There was a problem hiding this comment.
The ratio selection is correct for the number of concurrently live requests, but this path assumes the one auxiliary slot is returned when a request is released. That is not true for every no-radix cache implementation.
MlxAuxiliaryStateReqToTokenPool.free() currently only calls super().free(req). In the disable_radix_cache + chunked-prefill configuration, default_radix_cache_factory() selects ChunkCache; release_kv_cache() calls the pool's free(req) after freeing token KV, and the pool is not a HybridReqToTokenPool, so free_mamba_cache() is never called. The request row is reusable while req.mamba_pool_idx and the auxiliary slot remain occupied.
A real updated-head reproduction:
aux_before_free: 0
aux_after_free: 0
second_request_error: Not enough MLX auxiliary state slots
Please make ownership/release of the auxiliary slot explicit in this no-radix path and add a sequential allocate/free/reallocate assertion. The existing tests only allocate up to the cap and therefore cannot detect the leak.
There was a problem hiding this comment.
Great catch, I have fixed it by making ownership explicit.
I think it is a pre-existing bug as well, separate from this PR.
… them Address review: with --disable-radix-cache the tree cache is a plain ChunkCache, which frees token KV only, and release_kv_cache's free_mamba_cache fallback is gated on HybridReqToTokenPool -- which the MLX pool is not. pool.free(req) then freed only the request row, so every finished (or retracted) request permanently leaked its auxiliary slot and the (aux_size + 1)-th sequential request crashed with "Not enough MLX auxiliary state slots" even at concurrency 1. Make auxiliary-slot release ownership explicit: the stub constructs the pool with owns_auxiliary_state_release=disable_radix_cache, and free(req) releases the slot together with the row only when the pool owns it. With the radix cache enabled the ownership contract is unchanged: the unified radix component frees or adopts the slot (nulling req.mamba_pool_idx) before the row is freed, and free(req) keeps its hands off auxiliary slots entirely -- keyed on req.mamba_pool_idx, never on the req_index_to_auxiliary_state_index_mapping, which can reference tree-owned snapshots. Add the requested sequential allocate/free/reallocate regression (cycles of 3x the pool size at concurrency 1) plus a retention guard asserting free(req) never releases auxiliary slots in the radix-enabled configuration.
The mode is not rejected at pool construction: the pool silently sets enable_mamba_extra_buffer=False; ServerArgs.enable_mamba_extra_buffer() already requires the radix cache to be enabled, and on that path the MLX radix component (MlxAuxiliaryStateComponent) raises NotImplementedError. The ratio conclusion is unchanged.
|
/tag-and-rerun-ci |
|
@noob-se7en Please check the MLX UT failures. Thanks! |
Main retired the ModelRunner.mambaish_config property in favor of the configs.hybrid_arch.mambaish_config(model_config) function (which resolves through model_config.hf_config). The branch merge combined cleanly in text but not in semantics: the merged initialize() already used the function while _resolve_max_running_requests still read the removed property, which would raise AttributeError on any hybrid MLX startup, and the unit tests' property-shadowing subclasses no longer intercepted anything (the MLX CI failure on the merge commit). Use the function in the resolver, matching the rest of the file and the repo-wide convention. In the tests, patch the symbol the stub module imports instead of shadowing the retired property, and provide the two attributes the merged initialize() newly reads (device, model_config.use_ngram_embedding).
|
@Kangyan-Zhou The MLX PR test goes green. Could you please merge it? Thanks! |
…ect#30547) Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
3 participants
Motivation
--max-running-requestsis silently ignored on the MLX backend.MlxModelRunnerStub.initialize()hardcoded:and never consulted
server_args.max_running_requests. On MLX the scheduler reads its concurrency cap frommodel_runner.max_running_requests(the base CUDA path honors the flag inmodel_runner_kv_cache_mixin._resolve_max_num_reqs), so the stub's value is the only cap — and the flag had no effect. Users capping concurrency on a memory-constrained Mac, or setting it to1for deterministic / debug runs, silently got a different server than they asked for. (--max-total-tokensis already honored on MLX viapool_size, which makes the gap inconsistent.)Reproduced on
mlx-community/Qwen2.5-0.5B-Instruct-4bitwith--max-running-requests 1:Modifications
Extract
MlxModelRunnerStub._resolve_max_running_requests(), mirroring the base runner's clamp:--max-running-requestsis set, split it per dp worker and cap it by the KV pool capacity (max_total_num_tokens // 2), emitting the same reduction warning as the base runner when capacity forces a smaller value;min(max_total_num_tokens // 2, 4096)).initialize()now calls the helper instead of hardcoding the value. The change is confined to the MLX stub; no behavior change when the flag is unset.Accuracy Tests
Added
test/registered/unit/hardware_backend/mlx/test_max_running_requests.py, exercising the resolver directly (pure integer arithmetic, MLX-gated because importing the stub pulls inmlx.core):Verified the flag-honoring cases fail if the flag is ignored, and confirmed end-to-end (above) that a launched MLX server now reports the requested cap.
Speed Tests and Profiling
No performance impact: one-time integer computation during stub initialization.
Checklist
CI States
Latest PR Test (Base): ❌ Run #29557999719
Latest PR Test (Extra): ❌ Run #29557999581