server: priority-aware preemption with an explicit slot park and unpark API - #189
Draft
danielhanchen wants to merge 1 commit into
Draft
server: priority-aware preemption with an explicit slot park and unpark API#189danielhanchen wants to merge 1 commit into
danielhanchen wants to merge 1 commit into
Conversation
Keep live request and speculative state while swapping sequence memory to bounded host RAM. Add client priorities, explicit park and unpark actions, monitoring, and preemption regression tests.
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
Prevent independent text requests from failing when a unified KV pool fills. Park live sequences in bounded host RAM and restore their existing tasks, samplers and streams when capacity returns. Add request priorities, slot park/unpark controls and metrics.
Policy
Choose victims from the lowest effective-priority tier, protect the longest sequence unless it is the only candidate in that tier, and evict newest tasks on remaining ties. Exclude parent/child slots. Promote every third park and restore by priority, park count, wait age and fit. Keep one pressure survivor running to completion.
Results
Across three four-chat runs per context, completion counts improved from 6/12 to 12/12 at 8K and 0/12 to 12/12 at 16K. The priority-10 request was never parked and finished first. Four separate requests each reached 8,191 resident tokens in the 8,192-cell boundary test.
Exactness
Two 800-token comparisons with forced parks every 200 tokens matched byte-for-byte and token-for-token. An explicit concurrent CUDA pause preserved the 204-token emitted prefix, then first diverged at token 232. CUDA scheduling invariance is not guaranteed.
Cost when it does not fire
One O(P) scan; no snapshot copies or allocations. Paired solo throughput averaged 31.20 tok/s on master and 31.28 tok/s with the change. No consistent regression was observed; shared-GPU drift limits precision.
Tests
12 new tests pass and fail on base. Both builds pass 57 selected existing tests, with one existing skip. CUDA builds and clean patch application pass.
Limitations
Insufficient RAM or protected groups can stall progress. Multimodal/non-unified contexts and automatic sleep are outside the supported configuration. Other speculative architectures and external tool-approval integrations remain unvalidated. Concurrent CUDA output can diverge after restoration.
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. The mechanism is the swap of #184; what differs is who decides: the client sets a priority per request and may park or unpark a slot explicitly, which is what a chat waiting on a tool approval needs.