[AdvayMonga]

Overview

get_external_cache_hit_sources works out which tier each offload hit token came from. It runs on the scheduler thread for every request admitted with a hit, and it sorts and sweeps every loaded chunk, so it gets slower as prompts get longer. For a single attention group the chunks are already in order and don't overlap, so this PR adds them up in one pass. Follow-up to #56318.

Claims

• Same output.
• About 3.5x faster: 5.4 ms to 1.5 ms per 128k-token hit, on the scheduler thread.
• I don't expect req/s or TTFT to change, since async scheduling mostly hides this time. I didn't run an end-to-end benchmark.

Validation

Microbench, one attention group, 16 tokens per chunk, median of 200 calls on an M4 Pro. Three runs each side were within 0.05 ms.

hit length	before	after
8k tokens	0.31 ms	0.09 ms
32k tokens	1.33 ms	0.35 ms
128k tokens	5.4 ms	1.5 ms

Microbench script

import time
from types import SimpleNamespace

from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
    OffloadingConnectorScheduler,
)
from vllm.v1.metrics.cache_hit_source import CacheHitSource

TOKENS_PER_CHUNK = 16


def median_ms(fn, iters=200):
    ts = []
    for _ in range(iters):
        t0 = time.perf_counter_ns()
        fn()
        ts.append(time.perf_counter_ns() - t0)
    ts.sort()
    return ts[len(ts) // 2] / 1e6


for num_chunks in (512, 2048, 8192):
    # One attention group: ascending, disjoint ranges, all from the CPU tier.
    ranges = [
        (i * TOKENS_PER_CHUNK, (i + 1) * TOKENS_PER_CHUNK, b"key")
        for i in range(num_chunks)
    ]
    state = SimpleNamespace(
        num_locally_computed_tokens=0, load_key_ranges=ranges, req_context=None
    )
    scheduler = SimpleNamespace(
        _req_status={"r": state},
        manager=SimpleNamespace(get_load_source=lambda key, ctx: CacheHitSource.HOST),
    )
    request = SimpleNamespace(request_id="r")
    tokens = num_chunks * TOKENS_PER_CHUNK
    fn = lambda: OffloadingConnectorScheduler.get_external_cache_hit_sources(
        scheduler, request, tokens
    )
    assert fn() == {CacheHitSource.HOST: tokens}
    print(f"{tokens // 1024}k tokens ({num_chunks} chunks): {median_ms(fn):.2f} ms")

pytest tests/v1/kv_connector/unit/offloading_connector/ tests/v1/kv_connector/unit/test_multi_connector.py tests/v1/core/test_scheduler.py -k "hit_source or cache_hit"
101 passed

Added one test case with mixed tiers and a gap, which goes through the new path. pre-commit is clean.

Details

The chunks are in order because the scheduler adds one range per chunk, front to back, when it plans the load. The new loop adds each range's length to its tier and counts any gaps as external_unspecified. If a range starts before the previous one ends, it falls back to the existing sweep. That covers recurrent state and models with more than one KV cache group, so those behave exactly as before.

I also removed the CacheHitSource(...) wrap around get_load_source. It already returns a CacheHitSource, and the wrap was 0.7 ms of the 128k case.

I couldn't find an open PR touching this function.

AI assistance (Claude) was used. I reviewed every line and ran the tests and benchmark above.

---

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.

[orozery]

IIUC this optimization will likely only benefit dense models, which are rarely used today.
I think this does not justify the added code complexity.