@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a76df87db8

[MooncakeStore] Re-derive full external hits on stored boundaries (#49481)

Signed-off-by: Dao Le <Dao007forever@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com>

Dao007forever committed Jul 23, 2026 at 03:12 UTC a76df87db89bb87213f60cff817d5789796ee158

8 files changed +110 -38

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

+3 -3

@@ -410,7 +410,7 @@ def test_lookup_key_client_lookup_prepends_typed_tag():

410

411 # Blocking lookup (non_block defaults to False) runs on the executor and

412 # returns the resolved hit length.

413 - assert client.lookup("req0", token_len=128, block_hashes=[]) == 5

413 + assert client.lookup("req0", num_tokens=128, block_hashes=[]) == 5

414

415 sent_frames = fake_socket.send_multipart.call_args[0][0]

416 assert sent_frames[0] == protocol.LOOKUP_MSG

@@ -439,11 +439,11 @@ def test_lookup_key_client_reset_uses_typed_protocol():

439 assert client.reset() is False

440

441

442 -def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0):

442 +def _poll_lookup(client, req_id, num_tokens=128, block_hashes=(), timeout=5.0):

443 """Drive non-blocking lookup until the executor completes it."""

444 deadline = time.monotonic() + timeout

445 while time.monotonic() < deadline:

446 - result = client.lookup(req_id, token_len, list(block_hashes), non_block=True)

446 + result = client.lookup(req_id, num_tokens, list(block_hashes), non_block=True)

447 if result is not None:

448 return result

449 time.sleep(0.005)

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

+10 -2

@@ -240,7 +240,10 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

240 worker.store = store

241

242 # Both groups stored all 4 blocks -> full hit.

243 - assert worker.lookup(token_len=64, block_hashes=hs) == 64

243 + assert worker.lookup(num_tokens=65, block_hashes=hs) == 64

244 + # Exact-multiple prompt: the full hit is re-derived one block lower,

245 + # where both groups' stored blocks still cover the SWA window.

246 + assert worker.lookup(num_tokens=64, block_hashes=hs) == 48

247

248 # Evict SWA's first two blocks (outside its window of 32 tokens = 2 blocks).

249 swa_keys_outside_window = [

@@ -253,7 +256,12 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

256

257 # SWA window=32 -> only last 2 blocks must be present in SWA group.

258 # Full has all 4. Coordinator should still return 64.

256 - assert worker.lookup(token_len=64, block_hashes=hs) == 64

259 + assert worker.lookup(num_tokens=65, block_hashes=hs) == 64

260 + # Exact-multiple prompt after eviction: the boundary one block lower

261 + # needs SWA block 1, which is gone — no usable stored boundary remains

262 + # (the pre-fix arithmetic clamp would have returned 48 and livelocked

263 + # on load failure -> recompute -> same lookup).

264 + assert worker.lookup(num_tokens=64, block_hashes=hs) == 0

265

266

267 def test_recv_skips_swa_blocks_before_window():

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+8 -9

@@ -473,27 +473,25 @@ def test_from_request_tracker_no_load_saves_normally():

473 class _StubLookupClient:

474 def __init__(self, hit_tokens: int) -> None:

475 self._hit_tokens = hit_tokens

476 + self.num_tokens: list[int] = []

477

478 def lookup(

479 self,

480 req_id: str,

480 - token_len: int,

481 + num_tokens: int,

482 block_hashes: list[bytes],

483 non_block: bool = False,

484 ) -> int:

485 + self.num_tokens.append(num_tokens)

486 return self._hit_tokens

487

488

489 def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned():

488 - # When the external store hits the entire prompt, scheduler must leave at

489 - # least one token uncomputed for sampling but stay on a block boundary.

490 - # Otherwise the recv-side load mask floors token_len to

491 - # (num_tokens-1)//block_size, the tail partial chunk is dropped, and -- if

492 - # the local cache covers the aligned prefix -- key_list ends up empty

493 - # (ZeroDivisionError in the recv thread's `tp_rank % len(key_list)`).

490 + # The worker re-derives a full external hit below the request end on an

491 + # existing boundary, so the scheduler receives the usable aligned hit.

492 scheduler = _make_bare_scheduler()

493 scheduler.load_async = True

496 - scheduler.client = _StubLookupClient(hit_tokens=48) # full hit on 48-token prompt

494 + scheduler.client = _StubLookupClient(hit_tokens=32)

495

496 request = SimpleNamespace(

497 request_id="req-0",

@@ -510,6 +508,7 @@ def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned():

508 assert need_to_allocate == 16

509 assert load_async is True

510 load_spec = scheduler.load_specs["req-0"]

511 + assert scheduler.client.num_tokens == [48]

512 assert load_spec.vllm_cached_tokens == 16

513 assert load_spec.kvpool_cached_tokens == 32

514 assert load_spec.kvpool_cached_tokens % 16 == 0

@@ -522,7 +521,7 @@ def test_full_external_hit_with_full_local_hit_skips_load():

521 # into any block-aligned key.

522 scheduler = _make_bare_scheduler()

523 scheduler.load_async = True

525 - scheduler.client = _StubLookupClient(hit_tokens=48)

524 + scheduler.client = _StubLookupClient(hit_tokens=32)

525

526 request = SimpleNamespace(

527 request_id="req-0",

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+57 -2

@@ -1737,6 +1737,61 @@ def test_lookup_partial_prefix_returns_first_hit_length():

1737 assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32

1738

1739

1740 +def test_lookup_full_hit_reuses_existing_boundary():

1741 + """A full hit is re-derived below the request end without another RPC."""

1742 + worker = _make_bare_worker(block_size=16)

1743 + worker.store.batch_is_exist.return_value = [1, 1]

1744 +

1745 + assert worker.lookup(32, [b"h0", b"h1"]) == 16

1746 + assert worker.store.batch_is_exist.call_count == 1

1747 +

1748 +

1749 +def test_lookup_full_hit_with_eagle_pops_once_not_twice():

1750 + """Eagle already leaves the last block for the drafter, so a

1751 + full-prompt re-derivation must never fire for eagle-governed hits:

1752 + firing would anchor the search one block lower and pop a second

1753 + block, regressing the hit by an extra producer boundary."""

1754 + worker = _make_bare_worker(block_size=16)

1755 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

1756 + worker._kv_cache_groups,

1757 + scheduler_block_size=16,

1758 + hash_block_size=16,

1759 + use_eagle=True,

1760 + )

1761 + worker.store.batch_is_exist.return_value = [1, 1, 1, 1]

1762 +

1763 + # 64-token exact-multiple prompt, all 4 blocks stored: one eagle pop

1764 + # gives 48; a spurious re-derivation (anchored at 48) would pop again

1765 + # and return 32.

1766 + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 48

1767 + assert worker.store.batch_is_exist.call_count == 1

1768 +

1769 +

1770 +def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable():

1771 + """The motivating livelock: the producer of a 64-token prompt stored

1772 + only its SWA tail window (blocks 2-3). The old arithmetic clamp turned

1773 + the full hit into 48, whose SWA window needs the never-written block 1,

1774 + so every load failed and the recompute re-entered the same lookup. The

1775 + re-derivation must report that no stored boundary below the request end

1776 + is usable."""

1777 + from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec

1778 +

1779 + worker = _make_bare_worker(block_size=16)

1780 + swa = SlidingWindowSpec(

1781 + block_size=16, num_kv_heads=8, head_size=64, dtype=None, sliding_window=32

1782 + )

1783 + worker._kv_cache_groups = [KVCacheGroupSpec(["layer0"], swa)]

1784 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

1785 + worker._kv_cache_groups,

1786 + scheduler_block_size=worker.hash_block_size,

1787 + hash_block_size=worker.hash_block_size,

1788 + )

1789 + worker.store.batch_is_exist.return_value = [0, 0, 1, 1]

1790 +

1791 + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 0

1792 + assert worker.store.batch_is_exist.call_count == 1

1793 +

1794 +

1795 def test_lookup_swa_single_group_returns_full_when_tail_window_present():

1796 """Single-SWA, sliding_window=32 (= 2 blocks): producer stored only the

1797 tail. Coordinator-driven lookup returns full prefix even though the

@@ -1754,7 +1809,7 @@ def test_lookup_swa_single_group_returns_full_when_tail_window_present():

1809 hash_block_size=worker.hash_block_size,

1810 )

1811 worker.store.batch_is_exist.return_value = [0, 0, 1, 1]

1757 - assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 64

1812 + assert worker.lookup(65, [b"h0", b"h1", b"h2", b"h3"]) == 64

1813

1814

1815 def test_lookup_checks_all_potential_swa_hit_boundaries():

@@ -2219,7 +2274,7 @@ def test_lookup_records_mooncake_metrics():

2274 worker = _make_bare_worker()

2275 worker.store.batch_is_exist.return_value = [1, 1]

2276

2222 - result = worker.lookup(32, [b"a0", b"a1"])

2277 + result = worker.lookup(33, [b"a0", b"a1"])

2278 stats = worker.get_kv_connector_stats()

2279

2280 assert result == 32

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -93,6 +93,9 @@ class MooncakeStoreCoordinator:

93 self.eagle_group_ids = set(range(len(kv_cache_groups)))

94 self._verify_and_split_kv_cache_groups()

96 + def align_lookup_length(self, length: int) -> int:

97 + return length // self.lcm_block_size * self.lcm_block_size

98 +

99 def _verify_and_split_kv_cache_groups(self) -> None:

100 """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but

101 dispatches via spec_manager_map (we don't allocate managers).

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py

+2 -1

@@ -10,7 +10,8 @@ Wire format (REQ/REP over IPC):

10 Request: [msg_type: bytes] [payload_frames...]

12 msg_type == LOOKUP_MSG:

13 - frame 1: token_len (u32 big-endian, 4 bytes)

13 + frame 1: num_tokens (u32 big-endian, 4 bytes); the worker derives

14 + the aligned lookup length

15 frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each

16 fixed-size block hash (0 when there are no hashes)

17 frame 3: raw block hashes concatenated back-to-back (each hash_len

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

+2 -12

@@ -80,14 +80,12 @@ class MooncakeStoreScheduler:

80 Returns ``(None, False)`` when an async lookup is still in flight,

81 signaling the scheduler to retry this request on a later step.

82 """

83 - # Look up against the full prefill range, not just the prompt.

84 - token_len = request.num_tokens // self._block_size * self._block_size

85 - if token_len < self._block_size:

83 + if request.num_tokens < self._block_size:

84 return 0, False

86 num_external_hit_tokens = self.client.lookup(

87 request.request_id,

90 - token_len,

88 + request.num_tokens,

89 request.block_hashes,

90 non_block=self.lookup_async,

91 )

@@ -95,14 +93,6 @@ class MooncakeStoreScheduler:

93 # Lookup not ready yet; scheduler will retry on a later step.

94 return None, False

98 - if num_external_hit_tokens == request.num_tokens:

99 - # Leave a sub-block tail uncomputed for sampling, on a block

100 - # boundary so the recv-side load mask covers every yielded chunk.

101 - num_external_hit_tokens = max(

102 - 0,

103 - (request.num_tokens - 1) // self._block_size * self._block_size,

104 - )

105 -

96 if num_external_hit_tokens < num_computed_tokens:

97 need_to_allocate = 0

98 else:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+25 -9

@@ -1467,11 +1467,14 @@ class MooncakeStoreWorker:

1467

1468 return finished_sending

1469

1470 - def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:

1470 + def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int:

1471 """Check how many prefix tokens exist in the store.

1472

1473 - Checks across all rank-specific key namespaces that may be loaded.

1473 + Checks across all rank-specific key namespaces that may be loaded. A

1474 + hit covering all ``num_tokens`` is re-derived below the request end so

1475 + the last token is recomputed for sampling.

1476 """

1477 + token_len = self.coord.align_lookup_length(num_tokens)

1478 if not block_hashes or token_len <= 0:

1479 return 0

1480

@@ -1535,11 +1538,24 @@ class MooncakeStoreWorker:

1538 )

1539 }

1540

1541 + cached_block_pool = ExternalCachedBlockPool(

1542 + self.hash_block_size,

1543 + exists_set,

1544 + )

1545 _masks, hit_length = self.coord.find_longest_cache_hit(

1546 block_hashes,

1547 token_len,

1541 - ExternalCachedBlockPool(self.hash_block_size, exists_set),

1548 + cached_block_pool,

1549 )

1550 + if hit_length >= num_tokens:

1551 + usable_length = self.coord.align_lookup_length(num_tokens - 1)

1552 + if usable_length <= 0:

1553 + return 0

1554 + _masks, hit_length = self.coord.find_longest_cache_hit(

1555 + block_hashes,

1556 + usable_length,

1557 + cached_block_pool,

1558 + )

1559 return hit_length

1560

1561 def get_kv_events(self) -> list[BlockStored]:

@@ -1605,11 +1621,11 @@ class LookupKeyServer:

1621 msg_type = bytes(all_frames[0])

1622

1623 if msg_type == LOOKUP_MSG:

1608 - token_len = int.from_bytes(all_frames[1], byteorder="big")

1624 + num_tokens = int.from_bytes(all_frames[1], byteorder="big")

1625 hash_len = int.from_bytes(all_frames[2], byteorder="big")

1626 blob = all_frames[3].buffer

1627 block_hashes = BlobBlockHashes(blob, hash_len)

1612 - result = self.store_worker.lookup(token_len, block_hashes)

1628 + result = self.store_worker.lookup(num_tokens, block_hashes)

1629 self.socket.send(result.to_bytes(4, "big"))

1630

1631 elif msg_type == RESET_MSG:

@@ -1672,11 +1688,11 @@ class LookupKeyClient:

1688 )

1689 self.futures: dict[str, Future[int]] = {}

1690

1675 - def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:

1691 + def _lookup(self, num_tokens: int, block_hashes: list[BlockHash]) -> int:

1692 hash_len = len(block_hashes[0]) if block_hashes else 0

1693 all_frames = (

1694 LOOKUP_MSG,

1679 - token_len.to_bytes(4, byteorder="big"),

1695 + num_tokens.to_bytes(4, byteorder="big"),

1696 hash_len.to_bytes(2, byteorder="big"),

1697 b"".join(block_hashes),

1698 )

@@ -1687,7 +1703,7 @@ class LookupKeyClient:

1703 def lookup(

1704 self,

1705 req_id: str,

1690 - token_len: int,

1706 + num_tokens: int,

1707 block_hashes: list[BlockHash],

1708 non_block: bool = False,

1709 ) -> int | None:

@@ -1695,7 +1711,7 @@ class LookupKeyClient:

1711 so the caller retries on a later step."""

1712 future = self.futures.get(req_id)

1713 if future is None:

1698 - future = self.executor.submit(self._lookup, token_len, list(block_hashes))

1714 + future = self.executor.submit(self._lookup, num_tokens, list(block_hashes))

1715 self.futures[req_id] = future

1716 if non_block and not future.done():

1717 return None