Repository navigation
Serve concurrent requests with a continuous-batching scheduler - #367
Merged
stikves merged 13 commits intoOct 23, 2026
Merged
Conversation
BatchScheduler serves several requests on one batched CoreAISequentialEngine session. Each request takes a free row, is prefilled into it with prefillRow, and then advances one token per batched decode step at its own cursor. A row retires on EOS, a stop sequence or its token limit, and the next waiting request takes its slot. - SchedulerSession presents one scheduled request as an InferenceEngine, so decoding strategies and stop handling work unchanged. - IsolatedSessionEngine runs a request on its own session of the same engine, for requests that need per-step logits. - BatchServingRouter picks between the two. - submit throws queueFull once the scheduler holds maxConcurrent + maxQueued requests. Free rows absorb a burst before the queue limit applies. Empty prompts are rejected. - acquireExclusiveAccess returns a lease that pauses the scheduler at a step boundary, so a request on its own session never runs on the engine at the same time as the batch. - CoreAISequentialEngine.supportsBatching reports whether a graph can serve several requests at once.
--concurrency N serves chat requests through the batch scheduler when the model supports batching, meaning it has an attn_mask input and a batch dimension above 1. Otherwise the server warns and serves one request at a time. Requests that need per-step logits, chat with a response schema and completions scoring, run on an isolated session while holding an exclusive lease, so the batch pauses for them instead of sharing the engine. A full scheduler returns HTTP 429. Streaming requests start decoding before the SSE response is sent, so they get the 429 too. --replay-concurrent sends all replay requests at once, so replay runs can exercise the scheduler. Without --concurrency, requests take the same path as before.
With a batch scheduler running, /v1/completions scoring no longer pauses the batch with an exclusive lease. BatchScheduler.score(_:) queues the prompts, and between decode steps the pump scores them as one padded batch on a separate session. - Prompts are queued shortest first, and a batch stops before a third of its positions would be padding. - The batch advances 128 positions per forward. Each chunk's logits are reduced to log-probabilities, rows in parallel, before the next chunk runs. - CoreAISequentialEngine.allPositionLogits returns every position's logits from one batched forward. decodeBatch now slices the same forward. All prompts in one completions request score together. Scoring requests still take the server permit, so one runs at a time. Without a scheduler, scoring runs on the shared engine as before.
Doc comments on the batch scheduler, llm-server, engine selection and llm-runner --batch-prototype now state what each type and function does, its inputs, outputs and throw conditions, and how to call it. Two small code changes come with it: - BatchScheduler(engine:maxConcurrent:maxQueued:) requires engine.supportsBatching when maxConcurrent is above 1. - ServerState holds the batch scheduler and its model config as one optional, so one cannot be set without the other.
- Fallback requests (chat with a JSON schema) run on the shared engine under the exclusive lease. IsolatedSessionEngine is removed. - runBatchedForward returns the logits NDArray. decodeBatch copies only each row's last position, and allPositionLogits returns one flat [N, q, vocab] buffer. - LogProbabilities computes from a contiguous [positions, vocab] buffer and reuses its Float32 scratch buffers across positions. Scoring passes buffer slices instead of per-position arrays. - BatchScheduler.swift lists BatchScheduler first, then the types it uses. The class doc has a diagram of how requests fill the batch rows. - llm-benchmark help and errors no longer ask for --inference-engine-variant coreai-sequential, which the default now picks for attn_mask graphs. - BatchSchedulerTests is trimmed from 19 tests to 12 with the same coverage.
BatchScheduler+Scorer.swift holds Scorer, which runs padded scoring batches and picks how many queued requests go into each batch, and ScoreRequest. BatchScheduler keeps the scoring queue and calls into its scorer.
stikves
marked this pull request as ready for review
October 22, 2026 23:16
stikves
requested review from
Lewis300,
alejandro-isaza,
carinapeng,
kevchengcodes,
pkmandke and
srjoglekar246
October 22, 2026 23:17
tjia1818
reviewed
Oct 23, 2026
carinapeng
reviewed
Oct 23, 2026
A dropped stream, or setStopReason on the scheduler's output sequence (as the decoding strategy does on a stop-sequence match), now cancels the request. The scheduler retires a cancelled row before the next decode step and skips cancelled requests in the queue.
carinapeng
reviewed
Oct 23, 2026
carinapeng
reviewed
Oct 23, 2026
Scoring ran each pass at [requests, positions], so on a dynamic-shape graph every new prompt length or batch size was a new input shape, and a new shape can cost a compile or an allocation. The last chunk now pads to 16, 32, 64 or 128 positions, never past the context length, and a dynamic graph's rows round up to a power of two, at most maxConcurrent. Padded positions repeat the row's last token and padded rows are never read, so scores do not change.
Batched decoding sums in a different order than one request at a time, so the same seed can give slightly different logits. The fingerprint descriptor now carries batch=<concurrent requests>, so a batched server and a single-request server report different fingerprints.
carinapeng
approved these changes
Oct 23, 2026
3 of 8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a continuous-batching scheduler so
llm-servercan serve several requests on one batched model.BatchScheduler: requests take free rows, decode together, and leave independently. The scheduler returns 429 when full.SchedulerSession. JSON-schema requests pause the batch and run alone./v1/completionsscoring: prompts are scored in padded batches between decode steps.--concurrent-requests, nothing changes.Validation
Qwen3-0.6B on an M2 Max:
--concurrent-requests 4is identical to serial and about 1.5× faster (8 requests).Live HTTP was not tested. Replay covers the same handlers.