[Bug][MP] num_stored_tokens double counted when the scheduler retries get_num_new_matched_tokens, later stores unreachable
- Dominant language
- Python
- Stars
- 11.9k
- Forks
- 1.9k
- Avg merge
- 4d 4h
- Merged PRs (30d)
- 141
Description
## Environment
- LMCache dev @ 00f763c1
- vLLM 0.23.0
## Bug
`LMCacheMPConnector.get_num_new_matched_tokens` advances the tracker's stored-token watermark on every call (`lmcache/integration/vllm/lmcache_mp_connector.py:733`):
```python
ret = self.scheduler_adapter.check_lookup_result(request.request_id)
...
# Update num stored tokens for the tracker
tracker.increase_num_stored_tokens(ret)
```
The scheduler can call this more than once for the same waiting request. When `allocate_slots` returns `None` (vllm `v1/core/sched/scheduler.py:774`, "The request cannot be scheduled"), the request stays in the waiting queue with `num_computed_tokens == 0` and the next scheduling step queries the connector again. The cached lookup result is only cleared by `cleanup_lookup_result` inside `update_state_after_alloc`, which has not run at that point, so `check_lookup_result` returns the same cached hit and the watermark compounds: after n attempts, `num_stored_tokens == n * ret`.
## Consequence
Store metadata is produced from the inflated watermark (`GetStoreMetadata` starts at `num_stored_tokens`), so the token range `[ret, n * ret)` is never stored. Retrieval stops at the first missing chunk, so every chunk the request stores past the hole is unreachable. The trigger is a lookup hit combined with an allocation failure, which is exactly the memory-pressure regime.
## Repro
Two tests against dev @ 00f763c1, run with vLLM 0.23.0 installed. Both assert the correct behavior and fail:
```
E AssertionError: belief corrupted: num_stored_tokens 1024 != server truth 512
E AssertionError: store starts at 1024, leaving [512, 1024) never stored: all later chunks are unreachable
```
test_issue_num_stored_tokens_double_count.py
```python
# SPDX-License-Identifier: Apache-2.0
"""Repro for the num_stored_tokens double count.
Both tests assert the correct behavior; a failure demonstrates the bug.
"""
# Standard
from types import SimpleNamespace
# Third Party
import pytest
pytest.importorskip("vllm", reason="MP connector imports vLLM at module top")
# Third Party
from vllm.v1.request import RequestStatus # noqa: E402
# First Party
from lmcache.integration.vllm.lmcache_mp_connector import ( # noqa: E402
LMCacheMPConnector,
)
from lmcache.integration.vllm.lmcache_mp_metadata import ( # noqa: E402
LMCacheMPRequestMetadata,
)
CHUNK = 256
TOKENS_PER_BLOCK = 16
class _FakeLookupAdapter:
"""Mimics LMCacheMPSchedulerAdapter lookup caching semantics.
check_lookup_result returns the cached aggregate on every call after the
first: production caches results in _finished_lookup_results and only
cleanup_lookup_result (called from update_state_after_alloc) clears them.
"""
lmcache_tokens_per_chunk = CHUNK
def __init__(self, hit_tokens: int) -> None:
self._hit = hit_tokens
def maybe_submit_lookup_request(
self, request_id: str, token_ids: list[int], cache_salt: str = ""
) -> None:
pass
def check_lookup_result(self, request_id: str) -> int | None:
return self._hit
def _make_connector(hit_tokens: int) -> LMCacheMPConnector:
connector = LMCacheMPConnector.__new__(LMCacheMPConnector)
connector.request_trackers = {}
connector.scheduler_adapter = _FakeLookupAdapter(hit_tokens) # type: ignore
connector._hit_alignment_tokens = TOKENS_PER_BLOCK
return connector
def _make_request(num_tokens: int, request_id: str = "req-0") -> SimpleNamespace:
return SimpleNamespace(
request_id=request_id,
cache_salt=None,
all_token_ids=list(range(num_tokens)),
status=RequestStatus.WAITING,
)
def test_repeated_matched_tokens_calls_keep_stored_belief_consistent() -> None:
hit = 2 * CHUNK # 512-token LMCache hit
connector = _make_connector(hit)
request = _make_request(num_tokens=600)
# First scheduling attempt: lookup hit, then allocate_slots fails.
need, is_async = connector.get_num_new_matched_tokens(request, 0)
assert need == hit and is_async
tracker = connector.request_trackers["req-0"]
assert tracker.num_stored_tokens == hit
# Retry on the next step (same request, still num_computed_tokens == 0).
need2, _ = connector.get_num_new_matched_tokens(request, 0)
assert need2 == hit
assert tracker.num_stored_tokens == hit, (
"belief corrupted: num_stored_tokens "
f"{tracker.num_stored_tokens} != server truth {hit}"
)
def test_store_after_retry_starts_at_true_watermark() -> None:
"""Consequence: the first store op produced after the request finally
schedules must start at the server-stored watermark (512). Starting
beyond it leaves [512, 1024) never stored, and the prefix lookup stops
there on retrieval, so every chunk stored after the hole is unreachable.
"""
hit = 2 * CHUNK
connector = _make_connector(hit)
request = _make_request(num_tokens=1600)
connector.get_num_new_matched_tokens(request, 0)
connector.get_num_new_matched_tokens(request, 0) # allocate_slots retry
tracker = connector.request_trackers["req-0"]
# Allocation finally succeeds; the whole prompt is scheduled.
tracker.append_block_ids(([*range(1, 101)],)) # 100 blocks = 1600 tokens
tracker.increase_num_scheduled_tokens(1600 - hit)
tracker.num_vllm_hit_tokens = 0
tracker.num_lmcache_hit_tokens = hit
meta = LMCacheMPRequestMetadata.GetStoreMetadata(
tracker, CHUNK, [TOKENS_PER_BLOCK]
)
assert meta is not None, "no store metadata produced at all (silent hole)"
assert meta.op.start == hit, (
f"store starts at {meta.op.start}, leaving [{hit}, {meta.op.start}) "
"never stored: all later chunks are unreachable"
)
```
## Suggested fix
Apply the lookup accounting exactly once per tracker lifetime: either move `increase_num_stored_tokens` into `update_state_after_alloc` (the commit point that already calls `cleanup_lookup_result`), or keep it in `get_num_new_matched_tokens` behind a once-per-tracker flag.
Contributor guide
Assessment
This issue has not been assessed yet.