Server: add --parallel N interleaved sessions (default 1)#516
Open
iCreil wants to merge 3 commits into
Open
Conversation
Split ds4_gpu_graph into per-session state (KV caches, compressor and indexer frontiers, MTP verification state) and a new engine-owned ds4_gpu_scratch holding the transient work buffers: one-token decode tensors, per-layer stage tensors, MTP draft tensors and the batched prefill buffers. Everything in the scratch is written and consumed inside a single sync/eval call, so sessions sharing an engine can share one scratch as long as they are driven from one thread at a time -- which is exactly the serialization the server's single graph worker already provides. This is what makes a second live session affordable: the prefill batch buffers dominate a session's non-KV footprint (~1.8 MiB per prefill-chunk token), and they are now allocated once per engine instead of once per session. The scratch grows on demand to the union of the attached sessions' capacities and is freed with the engine. No behavior change with a single session: tensors, sizes and the execution order are unchanged, only their ownership moved. Add tests/two_sessions_smoke: drives two sessions of one engine interleaved (one greedy token each, alternating on one thread) and checks both streams match isolated single-session runs token by token. Verified on CUDA (RTX PRO 6000, DeepSeek V4 Flash q2): MATCH for both sessions, and isolated outputs identical to the pre-refactor build.
Move the request-scoped locals of generate_job() into a gen_state struct and split the body into job_begin / job_prefill / job_start_decode / job_decode_round_init / job_decode_step / job_finish. The decode_again tool-recovery label becomes a GEN_STEP_REDECODE result from job_finish. generate_job() is now a small driver that performs exactly the same calls in exactly the same order, so behavior is unchanged; this is groundwork for a scheduler that drives several jobs stepwise (one prefill chunk or a few decode tokens at a time) from the single graph worker instead of owning it until a job completes. Verified against the previous build on CUDA (DeepSeek V4 Flash q2): greedy chat completion is byte-identical in both stream and non-stream mode, and the SSE frame flow is intact.
Replace the single live session with an array of slots, each owning one ds4_session plus the protocol live state (Responses/Anthropic/thinking bindings) tied to that timeline. The graph worker becomes a scheduler: it hands queued jobs to idle slots (longest-common-prefix affinity, then LRU) and steps one slot at a time through the gen_state phases -- a whole begin/finish transition, one prefill slice, or a burst of decode tokens per turn. Prefill preemption reuses the cooperative cancel callback: syncs are interrupted at chunk boundaries once the slice deadline passes and another job is waiting, and resume from the session checkpoint on the slot's next turn. Decode interleaves round-robin between slots, so every connected client keeps streaming while another prompt prefills. The scheduler swaps the kvstore continued-checkpoint frontier in and out per slot; parser-side live call-id lookups scan every slot under the existing tool mutex. --parallel stays opt-in: the default is one slot, the yield callback never fires (no contention), and the phase sequence is exactly the old generate_job order -- verified byte-identical against the previous build on a greedy chat completion (stream and non-stream). N>1 requires resident experts and is refused with --ssd-streaming. Timeslices are tunable via DS4_SERVER_SCHED_PREFILL_MS (default 1500) and DS4_SERVER_SCHED_DECODE_TOKENS (default 6). Measured on CUDA (RTX PRO 6000 96GB, DeepSeek V4 Flash q2, --parallel 2 --ctx 32768 --prefill-chunk 1024, 94.9 GiB VRAM): with a 350-token generation in flight, a second client got its first token in 0.7 s and the first stream's worst inter-token gap during the newcomer's prefill was 0.15 s; a third concurrent request queued and completed once a slot freed.
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
1 participant
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.
What
Add
--parallel Nto ds4-server: N live sessions interleaved on the single graph worker, so several clients can be served concurrently — everyone keeps streaming while another prompt prefills. Default is--parallel 1, which keeps today's behavior exactly (verified byte-identical, see Testing).Stacked on #515 (engine-shared work buffers), which is what makes the second live session affordable in VRAM; the first commit here is that PR.
Why
The README describes the current model honestly: "concurrent requests wait their turn on the single live graph/session", and sketches a future "stateful session-based protocol". This PR is a bounded, opt-in step in that direction that fits the existing architecture: still one graph worker, still cooperative scheduling, no threads added, no batching — just interleaving at the token/chunk granularity the engine already exposes.
The concrete pain it removes: with one live session, a user who arrives while a long generation is running waits for all of it before seeing a first token. With
--parallel 2, on our box (RTX PRO 6000 96 GB, V4 Flash q2) a second client got its first token in 0.7 s while a 350-token generation was in flight, and that first stream's worst inter-token gap during the newcomer's prefill was 0.15 s. A third concurrent request queues FIFO as before and completes when a slot frees.How
struct server_slotowns oneds4_sessionplus the protocol live state bound to that timeline (Responses/Anthropic/thinking bindings) and a per-slot copy of the kvstore continued-checkpoint frontier, swapped in/out around each turn.server.activepoints at the slot being stepped; per-request code reaches its session through it, which keeps the single-worker ownership model explicit. Parser-side live call-id lookups scan every slot under the existing tool mutex.generate_job()intojob_begin/job_prefill/job_start_decode/job_decode_round_init/job_decode_step/job_finishwith the old stack frame moved into agen_statestruct. The olddecode_again:tool-recovery label becomes aGEN_STEP_REDECODEresult. With one slot the driver performs exactly the same calls in the same order.DS4_SERVER_SCHED_DECODE_TOKENS).DS4_SERVER_SCHED_PREFILL_MS) passes and another job is actually waiting; it resumes from the session checkpoint on the slot's next turn — the same resume path prefill cancellation already uses. With--parallel 1the callback never fires.--parallel N>1is refused together with--ssd-streaming(interleaving would thrash the streaming expert cache); the flag requires resident experts.Testing
On CUDA (RTX PRO 6000 Blackwell 96 GB, DeepSeek V4 Flash q2-imatrix):
--parallel 2 --ctx 32768 --prefill-chunk 1024, 94.9 GiB VRAM): client A streaming 350 tokens; client B arrives at +3 s → B's TTFT 0.7 s, A's worst inter-token gap during B's prefill 0.15 s, both complete correctly; a third concurrent request queues and completes.tests/two_sessions_smoke(from the first commit) passes: interleaved greedy streams match isolated runs token by token.Not covered yet (called out for review): multi-slot interaction with MTP speculative decode is exercised only lightly (state is fully per-graph, so isolation holds by construction); Metal untested beyond compile paths — no Apple hardware here.