Skip to content

server: park unified requests with proactive KV watermarks - #188

Draft
danielhanchen wants to merge 2 commits into
masterfrom
astra-1-watermarks
Draft

server: park unified requests with proactive KV watermarks#188
danielhanchen wants to merge 2 commits into
masterfrom
astra-1-watermarks

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

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.

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