[Bugfix][KV Offload] Preserve new loads during same-step resume

[Vegetog]

Overview

Fix an OffloadingConnector assertion when a prefix-cache reset preempts a request and its next scheduling step both reports that preemption and creates a new KV load. Exclude only loads created in the current step from the preemption flush, while preserving the store-only check for older jobs.

Related to the serving path exercised by #60552. This deliberately overlaps the same-step resume fix in #57810; the scope and implementation differences are described below.

Claims

• Preserve newly created CPU-to-GPU loads through the preemption notification, allowing the request to resume and finish.
• Continue flushing earlier stores and checking their transfer type before block reuse.

Validation

Standalone patch on 0f112d180033f3757cf283723f5593a40b02ed82:

.venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py \
  -k 'preempt or reset' -v

10 passed. Pre-commit passed for both changed files.

The two new cases reproduce the assertion before the fix in both synchronous and asynchronous scheduling. They start with a running request and no pending transfers, then verify that its new load is not flushed, completes normally, and the resumed request finishes. Existing cases cover preemption with pending stores and explicit connector reset.

Combined serving validation on the same base, with this patch and the separate #60552 return-value fix applied: one RTX 3070, Qwen3-0.6B, eager execution with TRITON_ATTN, 320 GPU blocks and 4 GB CPU offload. The Python checkout reused the official v0.31.0+cu129 CUDA extensions.

• Strict replay, using ignore_eos=True and streamed usage checks: 3 flagged and 3 control resets returned 200 {"success": false}; all 12 streams completed 64 tokens, all 6 follow-up health/generation checks passed, and every trial recorded positive offload-load bytes.
• The original [Bug]: POST /reset_prefix_cache?reset_running_requests=true returns HTTP 500 while an OffloadingConnector load is in flight #60552 client then completed 5 flagged and 5 control trials on the same server, all HTTP 200 and healthy, without the connector assertion.

These are combined-patch serving results. This PR does not change the HTTP reset failure contract; the standalone regression coverage is reported separately above.

No model-quality or performance claim is made. The serving checks validate HTTP responses, stream completion and continued generation; upstream CI has not been run for this proposed PR.

Details

reset_prefix_cache() queues a preemption notification for the next schedule(). That step may already create a load to restore the same request from the offloaded cache. The preemption flush incorrectly assumes every pending transfer predates the preemption and must be a store, so it asserts on the new load.

Subtract self._current_batch_load_jobs.keys() only when collecting jobs for the preemption flush. These loads have not yet been submitted to the worker; they remain in the load metadata and retain their normal completion and cleanup path. Older jobs still pass through the existing assertion and flush logic.

Overlap with #57810: that open PR already fixes the same root cause as part of its pause/sleep cache-retention feature. This is proposed as a small standalone alternative because the crash is reproducible on current main without that feature. Unlike filtering all pending jobs by is_store, this version excludes only current-step loads and retains the assertion for older jobs. It introduces no cache-retention option or pause/sleep API changes. Author/maintainer coordination is pending; the overlapping fix should be consolidated into one implementation rather than merged twice.

Developed and validated with OpenAI Codex assistance.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.