server: park unified requests with proactive KV watermarks

[danielhanchen]

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.

[danielhanchen]

Closing as answered. This branch has two commits. The first, b925a5f, batches fragmented sequence restores through bounded host staging, and that is a byte-identical patch to 75944db, which is already on the #184 branch, so there is nothing to carry across. The second is the watermark policy itself, which was opened as one of several server-side designs written independently for the same problem so they could be compared. That comparison has happened: #192 answers it with preempt_n_margin(), eight decode steps of lookahead, and that is what the rest of the stack is built on. Proactive watermarks are not a second mechanism worth maintaining beside it.

[danielhanchen]

Reopened. I closed this earlier today as part of a sweep over the open PRs on the fork and I should not have. It is not mine to retire and the call belongs to whoever owns the preemption work.

The two observations from the close comment are still worth having, and I would rather leave them here than take them away with the PR:

• b925a5f, which batches fragmented sequence restores through bounded host staging, is a byte-identical patch to 75944db already on the server: preempt a slot instead of ending every conversation when the KV pool fills #184 branch. I diffed the two patches. So that commit has nothing left to carry across, whatever happens to the rest.
• The watermark policy itself is the part still deciding something. This PR was opened as one of several server-side designs written independently for the same problem so they could be compared, and server: take the park and restore copies off the decode loop #192 took a different route with preempt_n_margin(). Whether proactive watermarks are still worth having beside that is a real question, and it is yours to answer rather than mine.