Repository navigation
[KV Offload] Load small CPU-tier hits synchronously within the step - #61961
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Documentation preview: https://vllm--61961.org.readthedocs.build/en/61961/ |
mergify
Bot
added
documentation
robertgshaw2-redhat
added
the
ready
|
/ci run |
|
✅ @kfirtoledo, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #95718 for commit |
Etelis
left a comment
Etelis
left a comment
There was a problem hiding this comment.
I agree this is a real issue in the specific scenario we’re seeing.
However, for a general-purpose inference engine, I’m concerned that sync_load_max_tokens adds more tuning complexity than value.
Users would need to choose an appropriate threshold for each deployment, and that value may become suboptimal as the workload changes.
For example, many individually small loads could all qualify for synchronous loading while collectively causing a substantial batch-wide delay.
I’d suggest either keeping the existing async behavior and accepting the latency cost, or exposing a simplesync_loads flag, defaulting to async. That would make KV loading synchronous before the forward pass and give users a clear, predictable tradeoff.
Another direction might worth exploring is reducing the async completion round trip from two scheduling steps to one where possible.
I think @orozery needs to decide.
|
Thanks @Etelis. The main idea is to keep long hits async and only load small ones sync. So some cut-off is needed, but it doesn't really need tuning: hits up to ~4096 tokens copy in about 1 ms, which is negligible against a 300+ ms step. Agreed that a per-step budget would bound the total better than a per-request limit, but it still needs tuning. |
d0144db to
473da74
Compare
| success = self.worker.submit_load(job_id, entry.src_spec, entry.dst_spec) | ||
| assert success | ||
| if sync_job_ids: | ||
| self.worker.wait(sync_job_ids) |
There was a problem hiding this comment.
We still submit the sync jobs to the stream containing async jobs.
We should give sync jobs their own handler / stream.
There was a problem hiding this comment.
Done. Sync loads now go through submit_sync_load, which CPUOffloadingWorker sends to a separate handler with its own stream. Other workers fall back to submit_load.
| # Keep at least one token to compute, or the request can't be scheduled. | ||
| req_status.sync_load = ( | ||
| num_hit_tokens is not None | ||
| and 0 < num_hit_tokens <= self.config.sync_load_max_tokens |
There was a problem hiding this comment.
We should have a per-step budget for sync loading instead of per-request.
There was a problem hiding this comment.
Done. One budget per step, reset at the end of the step. Hits that don't fit stay async.
| retention_interval=retention_interval, | ||
| dcp_world_size=vllm_config.parallel_config.decode_context_parallel_size, | ||
| # Hits up to this many tokens load within the step (0 disables). | ||
| sync_load_max_tokens=int(spec.extra_config.get("sync_load_max_tokens", 0)), |
There was a problem hiding this comment.
I think a limit in bytes is easier to set, since it does not depend on the model KV size per token.
There was a problem hiding this comment.
Done. It's now sync_load_budget_bytes, per GPU. A hit's size is estimated as hit tokens × KV bytes per token. That's exact for full attention and an overestimate for sliding-window and Mamba groups, so the budget is never exceeded.
| is_sync_load=req_status.sync_load, | ||
| ) | ||
|
|
||
| if self._chunks_being_loaded is not None: |
There was a problem hiding this comment.
sync load keys should not be added to _chunks_being_loaded
There was a problem hiding this comment.
Done. Sync-load keys are no longer added.
| # time.monotonic() of this request's first deferred offload lookup; | ||
| # None once consumed (observed) or while no lookup is pending. | ||
| deferred_lookup_start_time: float | None = None | ||
| sync_load: bool = False |
There was a problem hiding this comment.
Let's introduce _current_batch_sync_loads instead.
There was a problem hiding this comment.
Done. The per-request flag is replaced by _current_batch_sync_loads, a set reset each step.
|
This pull request has merge conflicts that must be resolved before it can be |
f034057 to
66ec8bd
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
| ) -> bool: | ||
| """Async offloaded medium -> GPU.""" | ||
|
|
||
| def submit_sync_load( |
There was a problem hiding this comment.
How about we introduce expedited: bool = False in submit_load instead? (I used expedited instead of sync to not confuse the implementor to not mistakenly block in the submit function itself).
There was a problem hiding this comment.
Done. submit_sync_load is replaced by submit_load(..., expedited=False); the CPU worker sends expedited loads to their own handler and stream. The connector passes expedited=True only for sync loads, so existing workers are unaffected while the option is off.
66ec8bd to
4ee6980
Compare
4ee6980 to
8200e0d
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #97165 for commit ℹ️ This PR is 3 commits behind upstream |
An async load parks the request until the load is reported finished, about two engine steps later. For a small hit the copy takes ~1 ms, so the wait dominates. With sync_load_max_tokens set in kv_connector_extra_config, hits up to that size load within the step and the request runs immediately. Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>
Address review: replace sync_load_max_tokens with a per-step sync_load_budget_bytes, track sync loads in _current_batch_sync_loads, keep their keys out of _chunks_being_loaded, and submit them through a separate CPU handler so they don't queue behind async loads. Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>
Address review: replace OffloadingWorker.submit_sync_load with an expedited flag on submit_load, so the API does not suggest blocking. Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>
fec16a8 to
e26409f
Compare
|
/ci retry |
|
✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/97165 |
|
/ci retry |
|
✅ Triggered Buildkite CI #97167 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #97227 for commit |
A sync load (vllm-project#61961) is waited on before its request's forward pass, and the request is already running, so failed_recving cannot reschedule it: the scheduler only acts on it for requests waiting for remote KV. Raise instead of letting the request return output computed on unloaded KV. Async loads and stores still survive their faults. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
…llm-project#61961) Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Christopher Mackenzie <christopher.mackenzie.software@gmail.com>
vllm-project#61961 built the handlers through make_handler and added an expedited load handler. Calibration still runs once after the store handler, and both load handlers now take its result. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UZAzEvT3BdxhAbW9FwkeTR Signed-off-by: hyunnnchoi <hyunmok.choi@gmail.com>
…llm-project#61961) Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Christopher Mackenzie <christopher.mackenzie.software@gmail.com>
A sync load (vllm-project#61961) is waited on before its request's forward pass, and the request is already running, so failed_recving cannot reschedule it: the scheduler only acts on it for requests waiting for remote KV. Raise instead of letting the request return output computed on unloaded KV. Async loads and stores still survive their faults. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
A sync load (vllm-project#61961) is waited on before its request's forward pass, and the request is already running, so failed_recving cannot reschedule it: the scheduler only acts on it for requests waiting for remote KV. Raise instead of letting the request return output computed on unloaded KV. Async loads and stores still survive their faults. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Oxana Korzh <okorzh@amd.com>
Overview
On wide-EP MoE prefill (such as GLM-5.3), a step takes 0.3 to 0.5 s. Almost every request has a small CPU-tier hit, and an async load parks it for about 2 steps although the copy takes ~1 ms.
This adds an opt-in
sync_load_budget_bytesoption (kv_connector_extra_config, default 0, off): small hits load in the same step, up to that many KV bytes per GPU per step. The change stays in the offloading connector, theOffloadingWorkerinterface, and the CPU offloading worker.Changes
scheduler.py: a hit that fits the step's remaining budget returnsload_async=False, so the request is scheduled in this step; the load targets its blocks by position and is marked sync._current_batch_sync_loadsand the byte total reset every step.worker.py: submits sync loads withsubmit_load(..., expedited=True), waits for them before the forward pass, and does not report them infinished_recving.kv_offload/base.py,cpu/gpu_worker.py: anexpeditedflag onOffloadingWorker.submit_load; the CPU worker runs expedited loads on their own stream.common.py: aTransferJob.syncflag carries that decision from scheduler to worker.Claims
sync_load_budget_bytesof KV per GPU per step; the rest stay async. The budget exists because the whole batch waits for the copy. A hit that would leave nothing to compute stays async.Validation
Qwen3-30B-A3B on H100, OffloadingConnector + TieringOffloadingSpec. 12k-token prefix in GPU cache, 2k-token tail only in the CPU tier. Median extra TTFT vs a GPU-hit control:
These numbers were measured with the earlier per-request token limit (
sync_load_max_tokens=4096) and a shared load stream; the per-step byte budget and separate stream are not yet run end to end.GLM-5.3, one prefill group (4x8 H100, wide EP), request latency with a CPU-tier hit: ~1.9 s without the fix or above the limit, ~0.85 s under it (hits of 2k to 15k tokens).
Correctness: greedy continuation with the tail served from the CPU tier is correct 7/7 (Qwen3 and GLM-5.3).
Details
Not a duplicate: no open PR adds sync offload loads (closest: #60574, #58198).
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.