[MLX] Fix startup after logical token capacity change by Manfredss · Pull Request #41314 · sgl-project/sglang · GitHub
Skip to content

[MLX] Fix startup after logical token capacity change - #41314

Merged
Kangyan-Zhou merged 2 commits into
sgl-project:mainfrom
Manfredss:fix/mlx-logical-token-capacity
Sep 30, 2026
Merged

Kangyan-Zhou merged 2 commits into
sgl-project:mainfrom
Manfredss:fix/mlx-logical-token-capacity

Conversation

@Manfredss

@Manfredss Manfredss commented Sep 26, 2026 •

Copy link
Copy Markdown

Motivation

MLX servers fail during startup after #39731 with AttributeError: 'MlxModelRunnerStub' object has no attribute 'kv_cache_configurator'. The new capacity query reaches the Torch configurator, which the MLX stub does not create.

Modifications

Report logical token capacity from the MLX pool, keeping the request pool's capacity adjustment. Add a regression covering worker startup and the request limits reported to the scheduler.

Accuracy Tests

  • The new regression fails before the fix; 30 related tests and 3 subtests pass after it on current main.
  • On Apple M4, mlx-community/Qwen3-0.6B-4bit now starts and generates text. With default overlap scheduling, both radix-cache modes pass short and chunked-prefill requests, followed by a healthy /health response.
  • Changed-file pre-commit checks pass. These are local results; full remote test CI still needs maintainer activation.

Speed Tests and Profiling

No model forward or sampling changes; no throughput claim.

Checklist

  • Run pre-commit on the changed files.
  • Add a regression that fails before the fix.
  • Document the MLX capacity behavior in the property docstring.
  • Provide relevant real-model startup and generation checks.
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ⏳ Run #36650163801
Latest PR Test (Extra): ❌ Run #36650163569
Latest PR Test (AMD ROCm 10): ⏳ Run #36650163553

@Manfredss

Copy link
Copy Markdown
Author

@SelamatAndagi

Copy link
Copy Markdown

Thanks for fixing this. I independently verified the patch on Apple Silicon
against the latest main (c311bc9611).

Environment:

  • macOS 26.6.2, arm64
  • Python 3.12.13
  • PyTorch 2.13.0
  • MLX 0.32.2
  • mlx-lm 0.31.3
  • Model: mlx-community/Qwen3-0.6B-4bit

I confirmed that the server starts successfully and an end-to-end /generate
request completes after applying this fix.

I found one small rebase issue caused by df1fc6ab64
([kv-shard 3/4] Enable Control Plane B). The regression test's fake runner
now also needs token_to_kv_pool_allocator, because
TpModelWorker.get_worker_info() calls page_interleave_shard_size().

Results on the latest main:

  • Before updating the test fixture: 33 passed, 1 failed
  • After adding the allocator stub: 34 passed

I left an inline suggestion for the test adjustment. Happy to send a follow-up
or superseding PR if that would be more convenient.

device="cpu",
enable_memory_saver=False,
)
runner.token_to_kv_pool = SimpleNamespace(size=64)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The latest main reads token_to_kv_pool_allocator when reporting worker
capacity, so the fake runner needs a stock allocator stub.

Suggested change
runner.token_to_kv_pool = SimpleNamespace(size=64)
runner.token_to_kv_pool = SimpleNamespace(size=64)
runner.token_to_kv_pool_allocator = SimpleNamespace()

@Manfredss
Manfredss force-pushed the fix/mlx-logical-token-capacity branch from 7e21f3b to acd56c9 Compare September 29, 2026 07:19
@yeahdongcn

Copy link
Copy Markdown
Collaborator

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 30, 2026
@Kangyan-Zhou
Kangyan-Zhou merged commit 7fb85a9 into sgl-project:main Sep 30, 2026
205 of 255 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

apple-silicon run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants