[KVConnector] Prevent async loads from exploding the KV Cache

[orozery]

Previously, async loaded requests were unbounded,
both in terms of number of requests and number of tokens.
This PR accounts async requests for both max_num_running_reqs and token_budget.
This will prevent cases where async loaded requests will unnecessarily fill up the entire GPU KV cache.

[gemini-code-assist]

Code Review

This pull request introduces important resource management controls for asynchronous KV cache loading in the scheduler. By accounting for async requests in both max_num_running_reqs and token_budget, it effectively prevents them from overwhelming the GPU KV cache. The implementation is mostly correct and adds a new counter for async loading requests, which is correctly updated during the scheduling lifecycle.

I have one suggestion to improve the token budget logic to ensure full utilization. Overall, this is a valuable improvement for stability and performance.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[robertgshaw2-redhat]

Can you explain what you were observing that trigger this issue?

[orozery]

Can you explain what you were observing that trigger this issue?

I'm testing with the OffloadingConnector, though every connector which uses async load can hit this.
Running a workload with many concurrent requests yielding external tokens, the scheduler tries to allocate GPU blocks to accommodate all of these external tokens.
There is no limit the scheduler imposes on that.
Say you have 10000 requests in self.waiting, all of them yielding external tokens.
In a single scheduler.schedule() step, you will allocate GPU blocks for all of these requests, until the KVCacheManager fails to allocate.
So potentially you allocated 1000's of requests.
But you cannot decode them all at the same time since you don't have enough GPU KV cache.
So you will get preemptions, and even worse, you will preempt requests that finished async loading (and turned into running state), while still asynchronously loading new requests.

[robertgshaw2-redhat]

@njhill - can you review this one?

[orozery]

I feel like the current implementation of this PR is not sufficient, this is why I turned it into a draft.
Still thinking it over.
Would be happy to brainstorm over it offline, but in any case will update here.

[orozery]

With the recent support for resuming preempted requests from CPU, it seems the case describe here is not so bad.
There are possible heuristic optimizations to balance off KV cache usage between async loads and running requests.
But I think we can defer this until we have a real justified use-case.