[sohom-cs]

Overview

With store_threshold >= 2, CPUOffloadingManager admits a chunk to the CPU tier only after it has been offered that many times. When prepare_store() can't allocate, it returns None, and the offloading scheduler offers the same keys again on a later step. Each retry counted as another offer, so a chunk that only one request ever had could be admitted by that request's own retry. A preempted request's re-offers counted the same way: the scheduler restarts its stores from chunk 0 through the same request context. With this PR, each key counts once per request.

Claims

• A request counts each key toward store_threshold once. When it offers a key again, as a retry of a failed prepare_store() call or after preemption, the key doesn't count again; it still moves to the most recent end of the tracker, as before. A cache reset starts the request's counting over.
• vllm:kv_offload_stores_skipped counts a key's skip once per request too.
• Offers from other requests count as before. Nothing changes when store_threshold < 2.

Validation

VLLM_TARGET_DEVICE=cpu python -m pytest tests/v1/kv_offload/cpu/test_manager.py -q
VLLM_TARGET_DEVICE=cpu python -m pytest tests/v1/kv_offload -q
VLLM_TARGET_DEVICE=cpu python -m pytest tests/v1/kv_connector/unit/offloading_connector -q
pre-commit run --files vllm/v1/kv_offload/cpu/manager.py tests/v1/kv_offload/cpu/test_manager.py

test_store_threshold_does_not_count_retries_of_a_failed_offer runs with store_threshold 2 and 3:

1. It fills the CPU tier with in-flight stores.
2. One request offers an eligible key and a cold key twice. Both calls return None.
3. The stores complete, and the request offers the two keys again.
4. That offer must store only the eligible key. The cold key's count must still be 1, and STORES_SKIPPED must be 1.

test_store_threshold_counts_a_key_once_per_request (also 2 and 3) is the preemption case: the request stores the eligible key, then offers both keys again through the same context, as a resumed request does. The cold key must not be stored, its count must stay 1, STORES_SKIPPED must be 1, and another request's offer must still count it.

With main's manager.py all four fail, at the store; with this PR's first commit (bb7b8b6) the two preemption cases fail the same way:

E       AssertionError: assert [b'1\x00\x00\...\x00\x00\x00'] == [b'1\x00\x00\x00\x00']
E         Left contains one more item: b'2\x00\x00\x00\x00'

b'2...' is the cold key: main stores it, although only this request ever offered it.

Through the scheduler. I ran the offloading unit tests' RequestRunner, which drives the real Scheduler and OffloadingConnectorScheduler, with a real CPUOffloadingManager(num_chunks=2, store_threshold=2) behind it:

1. Request A (one block) is offered once.
2. Both CPU chunks are filled with stores that stay in flight.
3. Request B (A's block plus one of its own) finishes in one step. Its offer fails.
4. The stores complete, and B's finish-time retry offers both blocks again.

	B's retry stores	B's own block in the CPU tier
main 4440f0e263	both blocks	yes
this PR	A's block only	no

The same runner with one request preempted once (free GPU blocks set to 0, as test_request_preemption does) and resumed, with sync and async scheduling: the request offers its two blocks, and after the resume offers them again through the same ReqContext. On main dc45a5f552 and on bb7b8b6 that re-offer stores both of its own blocks; with this PR it stores neither.

The script (copy it into tests/v1/kv_connector/unit/offloading_connector/ and run it with pytest -s)

# It drives the real Scheduler and OffloadingConnectorScheduler (the offloading
# unit tests' RequestRunner) with a real CPUOffloadingManager(store_threshold=2).
# A thin proxy hands the runner's mock worker its usual MockLoadStoreSpec.
from tests.v1.kv_connector.unit.offloading_connector.utils import (
    EOS_TOKEN_ID,
    MockLoadStoreSpec,
)
from vllm.v1.kv_offload.base import PrepareStoreOutput, ReqContext, make_offload_key
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager


class _Proxy:
    def __init__(self, real: CPUOffloadingManager):
        self._real = real
        self.offers: list[tuple[str, list, list | None]] = []

    def __getattr__(self, name):
        return getattr(self._real, name)

    def prepare_store(self, keys, req_context):
        keys = list(keys)
        out = self._real.prepare_store(keys, req_context)
        self.offers.append(
            (req_context.req_id, keys, None if out is None else list(out.keys_to_store))
        )
        if out is None:
            return None
        return PrepareStoreOutput(
            keys_to_store=out.keys_to_store,
            store_spec=MockLoadStoreSpec(out.keys_to_store),
            evicted_keys=out.evicted_keys,
        )

    def prepare_load(self, keys, req_context):
        self._real.prepare_load(keys, req_context)
        return MockLoadStoreSpec(keys)


def test_single_request_block_is_not_stored_by_its_own_retry(request_runner):
    runner = request_runner(block_size=4, num_gpu_blocks=20, async_scheduling=False)
    real = CPUOffloadingManager(num_chunks=2, cache_policy="lru", store_threshold=2)
    proxy = _Proxy(real)
    runner.connector_scheduler.manager = proxy

    # Request A: one block. Its offer only counts the block (count 1).
    runner.new_request(token_ids=[1] * 4)
    runner._run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=True)

    # Fill both CPU chunks with stores that stay in flight.
    pinned = [make_offload_key(bytes([200 + i]) * 8, 0) for i in range(2)]
    real.prepare_store(pinned, ReqContext(req_id="pin-warmup", kv_transfer_params=None))
    pin_ctx = ReqContext(req_id="pin", kv_transfer_params=None)
    out = real.prepare_store(pinned, pin_ctx)
    assert out is not None and out.keys_to_store == pinned

    freed = []

    def free_pins_once():
        if not freed:
            real.complete_store(pinned, pin_ctx)
            freed.append(True)

    # Request B: A's block plus a block only B has. Its first offer fails
    # (no free or evictable chunk); the stores then complete, and B's
    # finish-time retry offers both blocks again.
    runner.new_request(token_ids=[1] * 4 + [2] * 4)
    b_id = str(runner.req_id)
    runner._run(
        decoded_tokens=[EOS_TOKEN_ID], complete_transfers=True, post_step_fn=free_pins_once
    )

    b_offers = [o for o in proxy.offers if o[0] == b_id]
    print("\nB offers (keys offered -> keys admitted, None = allocation failed):")
    for _, keys, admitted in b_offers:
        print("  ", len(keys), "keys ->", None if admitted is None else len(admitted))
    shared, own = b_offers[0][1]
    print("B-only block resident in CPU tier:", real._policy.get(own) is not None)
    print("shared block resident in CPU tier:", real._policy.get(shared) is not None)
    assert real._policy.get(shared) is not None
    assert real._policy.get(own) is None, "B's own block was stored by B's retry"

On this branch (Linux CPU):

• test_manager.py: 72 passed.
• tests/v1/kv_offload: 893 passed, 34 failed. Main fails the same 34: 32 in cpu/test_gpu_worker.py, which needs a GPU, and 2 in tiering/p2p/test_manager.py::TestIdleWhileServingPeers.
• The offloading connector unit tests: 679 passed, the same as bb7b8b6.

CI: both tests are in tests/v1/kv_offload, which the "V1 KV Offload" step runs. I have no GPU, so I haven't run vLLM end to end. This changes which chunks are admitted, not how they are written or read, and the scheduler run above covers the path from the scheduler's retry to the manager. If you'd like an end-to-end run before merging, I'll do it when I have GPU time. Model outputs are unaffected, so no evals.

Details

Root cause:

1. prepare_store() counts every offered key toward store_threshold (_record_accesses) before it allocates.
2. If no chunk is free or evictable, it returns None. _build_store_jobs counts vllm:kv_offload_allocation_failure and moves on without advancing the request's store position. So the same keys are offered again on a later step, and again when the request finishes (test_scheduler_reports_allocation_failure asserts both attempts).
3. Each retry counted again. With store_threshold=2, one failed allocation and its retry admit every cold chunk in the offer. That happens under CPU-tier pressure, where the filter matters most.

Why a retry isn't another offer: #52227 moved the counter from lookups to store offers. It relied on prepare_store() receiving "the request's whole offerable key set", so that "a block crosses store_threshold after that many repetitions of the content". A retry resends that same offer, and so does a preempted request when it resumes: no new request, no new repetition. The metric is likewise documented as skips under the "reuse threshold". The same function already counts a request's keys once across calls in two places: the policy's on_store_miss (through store_miss_keys) and the request's access record (through seen_keys). Before this PR, the threshold counter and its skip metric were the per-key counts that a retry or a resumed request incremented again.

Scope: the keys a request has counted are kept in counted_keys on its _RequestCacheAccess, next to store_miss_keys and seen_keys, only when store_threshold >= 2. The scheduler keeps a request's ReqContext across preemption and drops it once the request is finished, so the set lasts as long as the request; a cache reset starts it over, as it does the rest of that state.

The first commit skipped only retries of failed offers. After review, every key counts once per request, as store_miss_keys does, so a request's re-offers after preemption don't count either. Three existing tests offered twice through one context to stand for two requests: test_filter_reused_manager, test_filter_reused_manager_oversized_offer_makes_progress and test_store_threshold_does_not_hide_ready_reuse. They now use one context per request.

Not a duplicate: no open PR changes _record_accesses or the threshold filter. I searched open PRs for store_threshold, _record_accesses and prepare_store. #63686 rewords the store_threshold docs. #57813, #63900 and #64216 change other parts of manager.py.

AI assistance

AI-assisted; reviewed and tested by me.

---

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.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-29T18:53:24.034958Z We couldn’t complete this review. Comment @codex review to try again.	bb7b8b6	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[orozery]

Thanks @sohom-cs !
Can we adapt so that we also don't recount blocks due to preemption?
Each key should be counted once per request.

[sohom-cs]

@orozery Done in 6304ad3.
The failed-offer set is now counted_keys, the keys a request has counted toward store_threshold: each key counts once per request, after a retry or after preemption, and so does its skip. It sits on the request's _RequestCacheAccess in its ReqContext, which the scheduler keeps across preemption and drops once the request is finished; a cache reset starts it over. test_store_threshold_counts_a_key_once_per_request covers the preemption case. Three tests that used one context for two requests now use two.

[orozery]

/ci run

[github-actions]

✅ @sohom-cs, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #99139 for commit 63e8fe54419d.

[orozery]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99303 for commit 26bb9487a004.

[orozery]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99476 for commit 6f15805b8d70.

ℹ️ This PR is 33 commits behind upstream main. Merged upstream main into this branch as 6f15805b8d70 before starting CI. Pull the branch before pushing more changes.