Skip to content

server: priority-aware preemption with an explicit slot park and unpark API - #189

Draft
danielhanchen wants to merge 1 commit into
masterfrom
astra-5-priority
Draft

server: priority-aware preemption with an explicit slot park and unpark API#189
danielhanchen wants to merge 1 commit into
masterfrom
astra-5-priority

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

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.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant