server: park unified requests with proactive KV watermarks - #188
Draft
danielhanchen wants to merge 2 commits into
Draft
server: park unified requests with proactive KV watermarks#188danielhanchen wants to merge 2 commits into
danielhanchen wants to merge 2 commits into
Conversation
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.
Summary
A full unified KV pool currently drives decode retries that can fail active chats, including MTP sub-batch errors. Park independent server tasks in host memory before allocation fails and resume their existing streams when capacity returns.
Policy
Default HIGH/LOW thresholds are 94/80 percent, with an 8192 MiB snapshot budget. Account for next-token and speculative cells and the shared prefill batch. Prefer newest victims, protect the largest request, and pass over repeat victims when alternatives exist. Restore most-parked and longest-waiting requests that fit below LOW. Keep task, sampler, speculative state and stream offsets alive. Add metrics, slot fields and a forced-parking test interval.
Results
On GPU 1 with Qwen3.5-4B MTP, the default policy completed 24/24 pressure requests across 8192/16384 contexts, versus 12/24 on master. Across all bands the branch completed 72/72 without errors. Four simultaneous requests each reached 8191 logical tokens at context 8192 without truncation, using existing draft clipping near the limit.
Exactness
Two individual 1000-token prompts match bytes and token IDs with four forced parks each. Controlled concurrent prompts diverge at token indices 257 and 418. CUDA batching changes numerical reductions; a separate matmul control and byte-preserving sequence-state tests provide evidence. Universal concurrent identity is not claimed.
Cost when it does not fire
Mean solo throughput was 277.51 tok/s on master and 275.84 tok/s with the change, a 0.60% decrease across three pairs. There are no snapshot copies or GPU synchronization below the watermark, but exact zero overhead is not established.
Tests
Nine new server integration cases pass and fail on master. Expanded fragmented-state assertions pass on stories260K CPU/CUDA and Qwen3.5-4B CUDA. CUDA Release builds and clean application of the full patch series pass.
Limitations
RAM denial can still lead to existing context errors. Parent/child completions and unsupported server/memory modes are excluded. Restore staging uses up to 64 MiB beyond the snapshot budget; larger fragmented tensors can restore slowly. Concurrent exactness, strict zero throughput regression, other MTP architectures and complete external tool loops remain limitations.
Relation to #184, #185 and #186
One of several server-side designs written independently for the same problem so they can be compared on the same hardware. #184 parks at the moment a decode does not fit; this branch parks earlier, at a high watermark, and restores only below a low one, which trades unused cells for fewer pauses (measured in the report above). Its batched fragmented-restore change to the sequence-state API is independent of the policy and applies to #184 as well.