From df57219f02ed59f4fdc03c19626f0170bd471cc6 Mon Sep 17 00:00:00 2001
From: Wei Zhao <weizha@nvidia.com>
Date: Fri, 23 Oct 2026 10:12:36 -0700
Subject: [PATCH 1/6] [KV Connector] Mooncake store: opt-in selective leases
 for prefix lookups

A prefix lookup checks every candidate key with batch_is_exist, which
grants each key it finds a read lease, and a lease both blocks eviction
and renews eviction order. Every lookup therefore renews the whole
prompt history. For hybrid models that includes the Mamba state at
every earlier block boundary, though a hit loads only the one at its
tail.

With VLLM_MOONCAKE_STORE_SELECTIVE_LEASE=1 the lookup probes the
candidates with the lease-free batch_probe_key, computes the hit as
before, and then leases only the keys the load will read. If a key was
evicted in between, the shorter hit is leased in turn; a hit that keeps
shrinking, or a failed lease, is reported as a miss. The lease RPC is
recorded as the lookup_lease operation.

The flag needs Mooncake >= 0.3.14 on the master and every client. A
worker checks at startup that the master serves the probe and fails
otherwise. With the flag off, lookups are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Wei Zhao <weizha@nvidia.com>
---
 .../mooncake_store_connector_usage.md         |   1 +
 .../unit/test_mooncake_store_hma_e2e.py       | 288 ++++++++++++++++++
 .../unit/test_mooncake_store_worker.py        |   1 +
 .../kv_connector/v1/mooncake/store/metrics.py |   6 +-
 .../v1/mooncake/store/selective_lease.py      | 217 +++++++++++++
 .../kv_connector/v1/mooncake/store/worker.py  |  39 ++-
 vllm/envs.py                                  |   8 +
 7 files changed, 556 insertions(+), 4 deletions(-)
 create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py

diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md
index 79e200bf9307..6996003ba5a1 100644
--- a/docs/features/mooncake_store_connector_usage.md
+++ b/docs/features/mooncake_store_connector_usage.md
@@ -261,6 +261,7 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
 | `MOONCAKE_PREFERRED_SEGMENT` | Pin this rank's replicas to a specific owner segment (`host:port`); used in `standalone-store` mode | — |
 | `MOONCAKE_REQUESTER_LOCAL_HOSTNAME` | Override the hostname the vLLM rank registers with Mooncake as a requester. Defaults to the rank's resolved IP. | — |
 | `VLLM_MOONCAKE_STORE_TIER_LOG` | When `1`, logs a per-batch tier summary (memory vs disk hits) for observability | disabled |
+| `VLLM_MOONCAKE_STORE_SELECTIVE_LEASE` | When `1`, prefix lookups probe keys without a read lease and lease only the keys the hit loads, so Mamba states at earlier block boundaries stay evictable. Only helps hybrid (Mamba/SWA) models. Needs Mooncake >= 0.3.14 on the master and every client. | disabled |
 | `VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO` | Fraction of the owner's DirectIO staging buffer that the requester will fill in a single `batch_get_into_multi_buffers` call. Lower → more conservative pre-split, more round trips. | 0.9 |
 
 ## KV Transfer Config
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
index 0682f9f465f1..854b93944ae0 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
@@ -8,6 +8,7 @@
 import types
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 
 from tests.v1.attention.utils import dense_kv_cache_views
@@ -22,6 +23,7 @@
     ChunkedTokenDatabase,
     KeyMetadata,
     LoadSpec,
+    PoolKey,
     ReqMeta,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (  # noqa: E501
@@ -932,3 +934,289 @@ def _fake_thread_init(*args, **kwargs):
 
     assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
     assert not any("@group:1" in key for key in store._data)
+
+
+class _LeaseTrackingStore(_DictStore):
+    """Dict store that tells lease-free probes from leasing existence checks.
+
+    ``batch_probe_key`` grants no lease; ``batch_is_exist`` leases every key it
+    finds. Each leasing call is recorded separately, and ``evictions`` holds
+    keys to drop just before each successive leasing call, as evictions between
+    probe and lease would.
+    """
+
+    def __init__(self, probe_result=None):
+        super().__init__()
+        self.probe_result = probe_result
+        self.probed: list[str] = []
+        self.lease_calls: list[list[str]] = []
+        self.evictions: list[set[str]] = []
+
+    def batch_probe_key(self, keys):
+        self.probed.extend(keys)
+        if self.probe_result is not None:
+            return [self.probe_result] * len(keys)
+        return [1 if k in self._data else 0 for k in keys]
+
+    def batch_is_exist(self, keys):
+        if self.evictions:
+            for key in self.evictions.pop(0):
+                self._data.pop(key, None)
+        self.lease_calls.append(list(keys))
+        return super().batch_is_exist(keys)
+
+
+def _full_plus_mamba_worker(store, *, partial_hash_hits=False):
+    """Full-attention + align-mode Mamba worker on ``store``."""
+    full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
+    mamba = MambaSpec(
+        block_size=16,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    cfg = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=8 * full.page_size_bytes,
+                layers=["L0"],
+                layer_stride=8 * full.page_size_bytes,
+                block_stride=full.page_size_bytes,
+            ),
+            KVCacheTensor(
+                size=8 * mamba.page_size_bytes,
+                layers=["L1"],
+                layer_stride=8 * mamba.page_size_bytes,
+                block_stride=mamba.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["L0"], full),
+            KVCacheGroupSpec(["L1"], mamba),
+        ],
+        # The hit alignment the engine core resolves for 4-token hash units.
+        **({"cache_hit_alignment_tokens": 4} if partial_hash_hits else {}),
+    )
+    vllm_config = _minimal_vllm_config(cache_block_size=16)
+    if partial_hash_hits:
+        # Hash unit 4 < Mamba block 16 enables partial hash hits.
+        vllm_config.cache_config.prefix_match_unit = 4
+    worker = _build_worker_with_dict_store(vllm_config, cfg, store)
+    for g_idx, db in enumerate(worker.token_dbs):
+        db.set_kv_caches_base_addr([g_idx * 10_000])
+        db.set_block_len([512])
+    return worker
+
+
+def _store_keys(worker, g_idx, block_hash):
+    return [
+        PoolKey.build_key_string(prefix, block_hash.hex())
+        for prefix in worker._lookup_key_prefixes[g_idx]
+    ]
+
+
+def _store_full_history(worker, store):
+    """Store every block and every block-boundary Mamba state of a 64-token
+    prefix, as earlier turns of a conversation would."""
+    hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
+    for g_idx in range(2):
+        for h in hs:
+            for key in _store_keys(worker, g_idx, h):
+                store._data[key] = b"x"
+    return hs
+
+
+def test_lookup_leases_every_probed_key_by_default():
+    """Without the flag, the lookup leases every candidate it finds, including
+    the Mamba states at earlier block boundaries."""
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
+
+    assert not store.probed
+    assert set(_store_keys(worker, 1, hs[0])) <= set(store.lease_calls[0])
+
+
+def test_selective_lease_leases_only_loaded_keys(monkeypatch):
+    """The lookup probes the whole prefix lease-free and leases only what the
+    hit loads: every full-attention block and the Mamba state at the tail."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
+
+    full_keys = [k for h in hs for k in _store_keys(worker, 0, h)]
+    mamba_keys = [k for h in hs for k in _store_keys(worker, 1, h)]
+    assert set(full_keys + mamba_keys) <= set(store.probed)
+    assert len(store.lease_calls) == 1
+    assert sorted(store.lease_calls[0]) == sorted(
+        full_keys + _store_keys(worker, 1, hs[-1])
+    )
+    stats = worker.get_kv_connector_stats().reduce()
+    assert stats["lookup_lease_count"] == 1
+    assert stats["lookup_lease_error_count"] == 0
+
+
+def test_selective_lease_leases_shorter_hit_after_eviction(monkeypatch):
+    """A tail state evicted between probe and lease shortens the hit, and the
+    shorter hit is leased in a second call."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+    store.evictions = [set(_store_keys(worker, 1, hs[3]))]
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 48
+
+    assert len(store.lease_calls) == 2
+    assert sorted(store.lease_calls[1]) == sorted(
+        [k for h in hs[:3] for k in _store_keys(worker, 0, h)]
+        + _store_keys(worker, 1, hs[2])
+    )
+
+
+def test_selective_lease_returns_no_hit_while_evictions_continue(monkeypatch):
+    """When every lease round finds an eviction, the lookup reports no hit
+    rather than a hit whose tail was never leased."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+    store.evictions = [
+        set(_store_keys(worker, 1, hs[3])),
+        set(_store_keys(worker, 1, hs[2])),
+    ]
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 0
+
+
+def test_selective_lease_treats_error_results_as_failed_lease(monkeypatch):
+    """Per-key error codes from the lease RPC report no hit and an error,
+    instead of shortening the hit as evictions would."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+    store.batch_is_exist = lambda keys: [-900] * len(keys)
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 0
+    assert worker.get_kv_connector_stats().reduce()["lookup_lease_error_count"] == 1
+
+
+def test_selective_lease_reports_no_hit_when_lease_raises(monkeypatch):
+    """A lease RPC that raises reports no hit, as a failed lookup does."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+
+    def _fail(_keys):
+        raise RuntimeError("master unreachable")
+
+    store.batch_is_exist = _fail
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 0
+
+
+def test_selective_lease_covers_every_key_the_load_reads(monkeypatch):
+    """Every key the receive thread reads for the hit was leased by the
+    lookup, so the two key walks cannot drift apart."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store)
+    hs = _store_full_history(worker, store)
+    hit = worker.lookup(num_tokens=65, block_hashes=hs)
+
+    read_keys: list[str] = []
+
+    class _CapturingStore:
+        def batch_get_into_multi_buffers(self, keys, addrs, sizes):
+            read_keys.extend(keys)
+            return [0] * len(keys)
+
+    recv = KVCacheStoreRecvingThread(
+        store=_CapturingStore(),
+        token_databases=worker.token_dbs,
+        block_size=worker.block_size,
+        tp_rank=0,
+        ready_event=threading.Event(),
+        coord=worker.coord,
+    )
+    recv.request_queue.put(
+        ReqMeta(
+            req_id="r0",
+            token_len_chunk=hit.hit_length,
+            block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
+            block_hashes=hs,
+            load_spec=LoadSpec(
+                vllm_cached_tokens=0,
+                kvpool_cached_tokens=hit.hit_length,
+                can_load=True,
+                token_len=hit.hit_length,
+                tail_key_boundaries=hit.tail_key_boundaries,
+            ),
+        )
+    )
+    recv._handle_request(recv.request_queue.get())
+
+    assert read_keys
+    assert set(read_keys) <= set(store.lease_calls[0])
+
+
+def test_selective_lease_with_partial_hash_hits(monkeypatch):
+    """With sub-block hash units, the lease keys the partial tail chunk by its
+    stored boundary, exactly as it was saved."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    store = _LeaseTrackingStore()
+    worker = _full_plus_mamba_worker(store, partial_hash_hits=True)
+    assert worker.coord.enable_partial_hash_hits
+
+    send_thread = KVCacheStoreSendingThread(
+        store=store,
+        token_databases=worker.token_dbs,
+        block_size=worker.block_size,
+        coord=worker.coord,
+        tp_rank=0,
+        group_put_steps=[1, 1],
+        kv_role="kv_both",
+        ready_event=threading.Event(),
+        replicate_config=MagicMock(),
+    )
+    # A 12-token prompt leaves a resend-safe sub-block checkpoint at 8.
+    hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)]
+    req = ReqMeta(
+        req_id="r0",
+        token_len_chunk=0,
+        block_ids=([1], [2]),
+        block_hashes=hs,
+        can_save=True,
+        num_prompt_tokens=12,
+        completed_token_len=12,
+        boundary_puts=[(1, 7, 8)],
+    )
+    send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req))
+    store.lease_calls.clear()
+
+    assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
+    assert sorted(store.lease_calls[0]) == sorted(store._data)
+
+
+def test_selective_lease_requires_probe_support(monkeypatch):
+    """Enabling the flag on a Mooncake without ``batch_probe_key`` fails at
+    startup instead of silently leasing every key."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    with pytest.raises(RuntimeError, match="batch_probe_key"):
+        _full_plus_mamba_worker(_DictStore())
+
+
+def test_selective_lease_requires_master_probe_support(monkeypatch):
+    """A master that rejects the probe RPC answers with per-key errors; the
+    startup check turns that into an error instead of silent lookup misses."""
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
+    with pytest.raises(RuntimeError, match="master"):
+        _full_plus_mamba_worker(_LeaseTrackingStore(probe_result=-900))
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
index 09d27eb947a0..5b91aed0d989 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
@@ -3711,6 +3711,7 @@ def _make_bare_worker(
     worker.kv_role = kv_role
     worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache
     worker._capacity_only = False
+    worker._selective_lease = None
     worker.block_size = block_size
     worker._is_hma_required = False
     worker.tp_rank = 0
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
index 8dc949cf326b..7df2e902fd33 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
@@ -3,9 +3,9 @@
 """Per-operation telemetry for MooncakeStoreConnector.
 
 Records one row per Mooncake RPC (``save_exists``, ``save_put``, ``load_get``,
-``lookup_exists``) with duration, key/byte counts, status, and failed-key
-count. Exposed to the logger via ``KVConnectorLogging`` and to Prometheus
-via ``MooncakeStorePromMetrics``.
+``lookup_exists``, ``lookup_lease``) with duration, key/byte counts, status,
+and failed-key count. Exposed to the logger via ``KVConnectorLogging`` and to
+Prometheus via ``MooncakeStorePromMetrics``.
 """
 
 from dataclasses import dataclass
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
new file mode 100644
index 000000000000..b07b6237b375
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Selective read leases for Mooncake store prefix lookups.
+
+``batch_is_exist`` leases every key it finds, so a lookup renews the whole
+prompt history, including the Mamba state at every earlier block boundary of a
+hybrid model. With ``VLLM_MOONCAKE_STORE_SELECTIVE_LEASE`` the lookup probes
+lease-free (``batch_probe_key``) and leases only the keys the hit will load.
+"""
+
+import time
+from collections.abc import Callable, Sequence
+from importlib import metadata
+from typing import Any, Protocol
+
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import (  # noqa: E501
+    MooncakeStoreCoordinator,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    ChunkedTokenDatabase,
+    MooncakeLookupResult,
+    PoolKey,
+)
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_interface import KVCacheGroupSpec
+
+logger = init_logger(__name__)
+
+# Lease rounds before giving up on a hit whose keys keep being evicted between
+# probe and lease. Only a hit whose lease found every key is returned.
+_LEASE_ROUNDS = 2
+# Key probed at startup to check that the master serves batch_probe_key.
+_SENTINEL_KEY = "vllm-selective-lease-probe"
+_MOONCAKE_DISTRIBUTIONS = (
+    "mooncake-transfer-engine",
+    "mooncake-transfer-engine-cuda13",
+    "mooncake-transfer-engine-non-cuda",
+)
+
+
+class _RecordOperation(Protocol):
+    def __call__(
+        self,
+        operation: str,
+        duration_seconds: float,
+        num_keys: int,
+        *,
+        num_bytes: int = 0,
+        status: str = "ok",
+        num_failed_keys: int = 0,
+    ) -> None: ...
+
+
+def _installed_mooncake() -> str:
+    for dist in _MOONCAKE_DISTRIBUTIONS:
+        try:
+            return f"{dist} {metadata.version(dist)}"
+        except metadata.PackageNotFoundError:
+            continue
+    return "unknown version"
+
+
+class SelectiveLookupLease:
+    """Leases only the keys a prefix-lookup hit will load.
+
+    The lookup probes its candidates with :meth:`probe`, computes the hit as
+    usual, and passes it through :meth:`lease_hit`.
+    """
+
+    def __init__(
+        self,
+        store: Any,
+        coord: MooncakeStoreCoordinator,
+        token_dbs: Sequence[ChunkedTokenDatabase],
+        kv_cache_groups: Sequence[KVCacheGroupSpec],
+        lookup_key_prefixes: Sequence[Sequence[str]],
+        record_operation: _RecordOperation,
+    ):
+        self._check_supported(store)
+        self.store = store
+        self.coord = coord
+        self.token_dbs = token_dbs
+        self.kv_cache_groups = kv_cache_groups
+        self.lookup_key_prefixes = lookup_key_prefixes
+        self.record_operation = record_operation
+
+    @staticmethod
+    def _check_supported(store: Any) -> None:
+        """Raise unless the client and the master both serve lease-free probes.
+
+        A master without the probe RPC answers each key with an error code
+        instead of raising, which would turn every lookup into a silent miss.
+        """
+        if not callable(getattr(store, "batch_probe_key", None)):
+            raise RuntimeError(
+                "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE needs Mooncake 0.3.14 or "
+                f"newer with batch_probe_key; installed: {_installed_mooncake()}"
+            )
+        res = store.batch_probe_key([_SENTINEL_KEY])
+        if len(res) != 1 or res[0] < 0:
+            raise RuntimeError(
+                "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE: the Mooncake master "
+                f"rejected batch_probe_key (result {list(res)}); the master and "
+                "every client must run the same Mooncake version "
+                f"(installed client: {_installed_mooncake()})"
+            )
+
+    def probe(self, keys: list[str]) -> list[int]:
+        """Existence of ``keys``, without granting a lease."""
+        return self.store.batch_probe_key(keys)
+
+    def lease_hit(
+        self,
+        block_hashes: Sequence[BlockHash],
+        hit: MooncakeLookupResult,
+        exists_set: set[tuple[int, bytes]],
+        longest_hit: Callable[[set[tuple[int, bytes]]], MooncakeLookupResult],
+    ) -> MooncakeLookupResult:
+        """Lease the keys a load of ``hit`` reads and return the hit to load.
+
+        ``exists_set`` holds the probed ``(group, hash)`` pairs and
+        ``longest_hit`` recomputes the hit from them. A key evicted since the
+        probe shortens the hit, whose tail keys then differ, so the shorter hit
+        is leased in turn. Returns an empty hit if the lease fails or the hit
+        keeps shrinking.
+        """
+        try:
+            for _ in range(_LEASE_ROUNDS):
+                if hit.hit_length <= 0:
+                    return hit
+                evicted = self._lease(block_hashes, hit)
+                if evicted is None:
+                    break
+                if not evicted:
+                    return hit
+                exists_set = exists_set - evicted
+                hit = longest_hit(exists_set)
+        except Exception:
+            # The lookup server has no handler of its own; report a miss.
+            logger.exception("Failed to lease a Mooncake lookup hit")
+        return MooncakeLookupResult(0)
+
+    def loaded_keys(
+        self,
+        block_hashes: Sequence[BlockHash],
+        hit: MooncakeLookupResult,
+    ) -> tuple[list[str], list[tuple[int, bytes]]]:
+        """Keys a load of ``hit`` reads, in every rank namespace, and their
+        ``(group, hash)`` pairs.
+
+        Mirrors ``KVCacheStoreRecvingThread._handle_request``: each group's
+        load mask, with the tail chunk keyed by the hit's tail boundary.
+        """
+        hit_length = hit.hit_length
+        load_masks = self.coord.load_mask(block_hashes, hit_length)
+        tail_tokens = {b.group_id: b.num_tokens for b in hit.tail_key_boundaries}
+        keys: list[str] = []
+        key_meta: list[tuple[int, bytes]] = []
+        for g_idx, db in enumerate(self.token_dbs):
+            if not self.kv_cache_groups[g_idx].kv_cache_spec.prefix_cacheable:
+                continue
+            mask = load_masks[g_idx]
+            key_prefixes = self.lookup_key_prefixes[g_idx]
+            for start, end, block_hash in db.process_tokens(hit_length, block_hashes):
+                chunk_idx = start // db.block_size
+                if chunk_idx >= len(mask) or not mask[chunk_idx]:
+                    continue
+                if end == hit_length and g_idx in tail_tokens:
+                    block_hash = block_hashes[
+                        tail_tokens[g_idx] // db.hash_block_size - 1
+                    ]
+                keys.extend(
+                    PoolKey.build_key_string(key_prefix, block_hash.hex())
+                    for key_prefix in key_prefixes
+                )
+                key_meta.extend([(g_idx, bytes(block_hash))] * len(key_prefixes))
+        return keys, key_meta
+
+    def _lease(
+        self,
+        block_hashes: Sequence[BlockHash],
+        hit: MooncakeLookupResult,
+    ) -> set[tuple[int, bytes]] | None:
+        """Lease the keys a load of ``hit`` reads.
+
+        Returns the ``(group, hash)`` pairs no longer in the store, or None
+        when the lease RPC fails, including per-key error results.
+        """
+        keys, key_meta = self.loaded_keys(block_hashes, hit)
+        if not keys:
+            return set()
+        start = time.perf_counter()
+        try:
+            res = self.store.batch_is_exist(keys)
+        except Exception as e:
+            logger.error("Remote connection failed leasing a lookup hit: %s", e)
+            res = [-1] * len(keys)
+        num_errors = sum(1 for r in res if r < 0)
+        if num_errors:
+            self.record_operation(
+                "lookup_lease",
+                time.perf_counter() - start,
+                len(keys),
+                status="error",
+                num_failed_keys=num_errors,
+            )
+            return None
+        evicted = [meta for meta, r in zip(key_meta, res, strict=True) if r != 1]
+        self.record_operation(
+            "lookup_lease",
+            time.perf_counter() - start,
+            len(keys),
+            status="partial_failure" if evicted else "ok",
+            num_failed_keys=len(evicted),
+        )
+        return set(evicted)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
index 004b69286ea3..c8c310fa659a 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
@@ -11,6 +11,7 @@
 """
 
 import dataclasses
+import functools
 import math
 import os
 import queue
@@ -65,6 +66,9 @@
     decode_lookup_response,
     encode_lookup_response,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.selective_lease import (  # noqa: E501
+    SelectiveLookupLease,
+)
 from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
@@ -1498,6 +1502,7 @@ def __init__(
 
         self._kv_cache_config = kv_cache_config
         self.token_dbs: list[ChunkedTokenDatabase] = []
+        self._selective_lease: SelectiveLookupLease | None = None
 
         # a capacity-only instance does not need below utils
         if self._capacity_only:
@@ -1533,6 +1538,15 @@ def __init__(
         )
         self.token_dbs = self._build_token_databases(metadata, store_layout_cls)
         self._init_lookup_key_prefixes()
+        if envs.VLLM_MOONCAKE_STORE_SELECTIVE_LEASE:
+            self._selective_lease = SelectiveLookupLease(
+                self.store,
+                self.coord,
+                self.token_dbs,
+                self._kv_cache_groups,
+                self._lookup_key_prefixes,
+                self._record_kv_connector_operation,
+            )
 
     def _supports_tp_sharded_store_layout(
         self,
@@ -2151,9 +2165,14 @@ def lookup(
         if not candidate_keys:
             return MooncakeLookupResult(0)
 
+        exists_fn = (
+            self.store.batch_is_exist
+            if self._selective_lease is None
+            else self._selective_lease.probe
+        )
         lookup_start = time.perf_counter()
         try:
-            res = self.store.batch_is_exist(candidate_keys)
+            res = exists_fn(candidate_keys)
             self._record_kv_connector_operation(
                 "lookup_exists",
                 time.perf_counter() - lookup_start,
@@ -2181,6 +2200,24 @@ def lookup(
                 exists_set.add((g_idx, hash_bytes))
             pos += count
 
+        hit = self._longest_hit(num_tokens, token_len, block_hashes, exists_set)
+        if self._selective_lease is None:
+            return hit
+        return self._selective_lease.lease_hit(
+            block_hashes,
+            hit,
+            exists_set,
+            functools.partial(self._longest_hit, num_tokens, token_len, block_hashes),
+        )
+
+    def _longest_hit(
+        self,
+        num_tokens: int,
+        token_len: int,
+        block_hashes: Sequence[BlockHash],
+        exists_set: set[tuple[int, bytes]],
+    ) -> MooncakeLookupResult:
+        """Longest prefix hit over the ``(group, hash)`` pairs in the store."""
         cached_block_pool = ExternalCachedBlockPool(
             self.hash_block_size,
             exists_set,
diff --git a/vllm/envs.py b/vllm/envs.py
index 1d6a4aa598c1..e71b6ad61202 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -239,6 +239,7 @@
     VLLM_EC_SIDE_CHANNEL_PORT: int = 5601
     VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
     VLLM_MOONCAKE_STORE_TIER_LOG: bool = False
+    VLLM_MOONCAKE_STORE_SELECTIVE_LEASE: bool = False
     VLLM_MOONCAKE_LOAD_RECV_THREADS: int = 1
     VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9
     MOONCAKE_PREFERRED_SEGMENT: str | None = None
@@ -1771,6 +1772,13 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: (
         os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1")
     ),
+    # Probe prefix-lookup keys without a read lease and lease only the keys the
+    # hit loads, so unused Mamba states at past block boundaries stay
+    # evictable. Needs Mooncake >= 0.3.14 on the master and every client.
+    "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE": lambda: (
+        os.getenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "False").lower()
+        in ("true", "1")
+    ),
     # Number of parallel KV-load receive threads per worker rank. Lets the
     # per-request control overhead (Python prep + master key lookup) of one
     # request overlap with the RDMA transfer of another, keeping the transfer

From 9b1feea6ff3942e97d153631d98e6a8c7c05a236 Mon Sep 17 00:00:00 2001
From: Wei Zhao <weizha@nvidia.com>
Date: Fri, 23 Oct 2026 10:55:12 -0700
Subject: [PATCH 2/6] [KV Connector] Mooncake selective lease: simplify the
 support check and lease failures

Check only that the installed Mooncake provides batch_probe_key, and
raise when it does not.

A failed lease RPC now raises an internal _LeaseFailedError instead of
returning None, so _lease always returns the evicted (group, hash)
pairs, and lease_hit reports a miss when the lease fails or the hit
keeps shrinking.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Wei Zhao <weizha@nvidia.com>
---
 .../unit/test_mooncake_store_hma_e2e.py       | 13 +---
 .../v1/mooncake/store/selective_lease.py      | 74 +++++++------------
 2 files changed, 29 insertions(+), 58 deletions(-)

diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
index 854b93944ae0..65ab704f946c 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
@@ -945,17 +945,14 @@ class _LeaseTrackingStore(_DictStore):
     probe and lease would.
     """
 
-    def __init__(self, probe_result=None):
+    def __init__(self):
         super().__init__()
-        self.probe_result = probe_result
         self.probed: list[str] = []
         self.lease_calls: list[list[str]] = []
         self.evictions: list[set[str]] = []
 
     def batch_probe_key(self, keys):
         self.probed.extend(keys)
-        if self.probe_result is not None:
-            return [self.probe_result] * len(keys)
         return [1 if k in self._data else 0 for k in keys]
 
     def batch_is_exist(self, keys):
@@ -1212,11 +1209,3 @@ def test_selective_lease_requires_probe_support(monkeypatch):
     monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
     with pytest.raises(RuntimeError, match="batch_probe_key"):
         _full_plus_mamba_worker(_DictStore())
-
-
-def test_selective_lease_requires_master_probe_support(monkeypatch):
-    """A master that rejects the probe RPC answers with per-key errors; the
-    startup check turns that into an error instead of silent lookup misses."""
-    monkeypatch.setenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "1")
-    with pytest.raises(RuntimeError, match="master"):
-        _full_plus_mamba_worker(_LeaseTrackingStore(probe_result=-900))
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
index b07b6237b375..4775500045ef 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
@@ -10,7 +10,6 @@
 
 import time
 from collections.abc import Callable, Sequence
-from importlib import metadata
 from typing import Any, Protocol
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import (  # noqa: E501
@@ -30,13 +29,10 @@
 # Lease rounds before giving up on a hit whose keys keep being evicted between
 # probe and lease. Only a hit whose lease found every key is returned.
 _LEASE_ROUNDS = 2
-# Key probed at startup to check that the master serves batch_probe_key.
-_SENTINEL_KEY = "vllm-selective-lease-probe"
-_MOONCAKE_DISTRIBUTIONS = (
-    "mooncake-transfer-engine",
-    "mooncake-transfer-engine-cuda13",
-    "mooncake-transfer-engine-non-cuda",
-)
+
+
+class _LeaseFailedError(Exception):
+    """The lease RPC raised or returned per-key error results."""
 
 
 class _RecordOperation(Protocol):
@@ -52,15 +48,6 @@ def __call__(
     ) -> None: ...
 
 
-def _installed_mooncake() -> str:
-    for dist in _MOONCAKE_DISTRIBUTIONS:
-        try:
-            return f"{dist} {metadata.version(dist)}"
-        except metadata.PackageNotFoundError:
-            continue
-    return "unknown version"
-
-
 class SelectiveLookupLease:
     """Leases only the keys a prefix-lookup hit will load.
 
@@ -87,23 +74,11 @@ def __init__(
 
     @staticmethod
     def _check_supported(store: Any) -> None:
-        """Raise unless the client and the master both serve lease-free probes.
-
-        A master without the probe RPC answers each key with an error code
-        instead of raising, which would turn every lookup into a silent miss.
-        """
         if not callable(getattr(store, "batch_probe_key", None)):
             raise RuntimeError(
-                "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE needs Mooncake 0.3.14 or "
-                f"newer with batch_probe_key; installed: {_installed_mooncake()}"
-            )
-        res = store.batch_probe_key([_SENTINEL_KEY])
-        if len(res) != 1 or res[0] < 0:
-            raise RuntimeError(
-                "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE: the Mooncake master "
-                f"rejected batch_probe_key (result {list(res)}); the master and "
-                "every client must run the same Mooncake version "
-                f"(installed client: {_installed_mooncake()})"
+                "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE requires "
+                "MooncakeDistributedStore.batch_probe_key, which the installed "
+                "Mooncake version does not provide."
             )
 
     def probe(self, keys: list[str]) -> list[int]:
@@ -130,15 +105,16 @@ def lease_hit(
                 if hit.hit_length <= 0:
                     return hit
                 evicted = self._lease(block_hashes, hit)
-                if evicted is None:
-                    break
                 if not evicted:
                     return hit
                 exists_set = exists_set - evicted
                 hit = longest_hit(exists_set)
+        except _LeaseFailedError as e:
+            logger.error("Failed to lease a Mooncake lookup hit: %s", e)
         except Exception:
             # The lookup server has no handler of its own; report a miss.
             logger.exception("Failed to lease a Mooncake lookup hit")
+        # The lease failed, or keys kept being evicted until rounds ran out.
         return MooncakeLookupResult(0)
 
     def loaded_keys(
@@ -181,11 +157,12 @@ def _lease(
         self,
         block_hashes: Sequence[BlockHash],
         hit: MooncakeLookupResult,
-    ) -> set[tuple[int, bytes]] | None:
+    ) -> set[tuple[int, bytes]]:
         """Lease the keys a load of ``hit`` reads.
 
-        Returns the ``(group, hash)`` pairs no longer in the store, or None
-        when the lease RPC fails, including per-key error results.
+        Returns the ``(group, hash)`` pairs no longer in the store. Raises
+        ``_LeaseFailedError`` when the lease RPC fails, including per-key error
+        results.
         """
         keys, key_meta = self.loaded_keys(block_hashes, hit)
         if not keys:
@@ -194,18 +171,14 @@ def _lease(
         try:
             res = self.store.batch_is_exist(keys)
         except Exception as e:
-            logger.error("Remote connection failed leasing a lookup hit: %s", e)
-            res = [-1] * len(keys)
+            self._record_lease_error(start, len(keys), len(keys))
+            raise _LeaseFailedError(f"lease RPC raised: {e}") from e
         num_errors = sum(1 for r in res if r < 0)
         if num_errors:
-            self.record_operation(
-                "lookup_lease",
-                time.perf_counter() - start,
-                len(keys),
-                status="error",
-                num_failed_keys=num_errors,
+            self._record_lease_error(start, len(keys), num_errors)
+            raise _LeaseFailedError(
+                f"lease RPC returned errors for {num_errors} of {len(keys)} keys"
             )
-            return None
         evicted = [meta for meta, r in zip(key_meta, res, strict=True) if r != 1]
         self.record_operation(
             "lookup_lease",
@@ -215,3 +188,12 @@ def _lease(
             num_failed_keys=len(evicted),
         )
         return set(evicted)
+
+    def _record_lease_error(self, start: float, num_keys: int, num_failed: int) -> None:
+        self.record_operation(
+            "lookup_lease",
+            time.perf_counter() - start,
+            num_keys,
+            status="error",
+            num_failed_keys=num_failed,
+        )

From 7b5f88264f506281e87572df53641ab22cbad8d5 Mon Sep 17 00:00:00 2001
From: Wei Zhao <weizha@nvidia.com>
Date: Fri, 23 Oct 2026 15:21:57 -0700
Subject: [PATCH 3/6] [KV Connector] Mooncake selective lease: generic wording,
 drop docs and metrics edits

Describe the feature without tying it to Mamba: a lookup lease covers the
keys the hit loads, and a hit that shrinks after an eviction can end at a
key the first lease did not cover, hence the second round. Revert the
usage-doc row and the metrics docstring edit.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Wei Zhao <weizha@nvidia.com>
---
 .../mooncake_store_connector_usage.md         |  1 -
 .../kv_connector/v1/mooncake/store/metrics.py |  6 +++---
 .../v1/mooncake/store/selective_lease.py      | 19 +++++++++++--------
 vllm/envs.py                                  |  3 +--
 4 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md
index 6996003ba5a1..79e200bf9307 100644
--- a/docs/features/mooncake_store_connector_usage.md
+++ b/docs/features/mooncake_store_connector_usage.md
@@ -261,7 +261,6 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
 | `MOONCAKE_PREFERRED_SEGMENT` | Pin this rank's replicas to a specific owner segment (`host:port`); used in `standalone-store` mode | — |
 | `MOONCAKE_REQUESTER_LOCAL_HOSTNAME` | Override the hostname the vLLM rank registers with Mooncake as a requester. Defaults to the rank's resolved IP. | — |
 | `VLLM_MOONCAKE_STORE_TIER_LOG` | When `1`, logs a per-batch tier summary (memory vs disk hits) for observability | disabled |
-| `VLLM_MOONCAKE_STORE_SELECTIVE_LEASE` | When `1`, prefix lookups probe keys without a read lease and lease only the keys the hit loads, so Mamba states at earlier block boundaries stay evictable. Only helps hybrid (Mamba/SWA) models. Needs Mooncake >= 0.3.14 on the master and every client. | disabled |
 | `VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO` | Fraction of the owner's DirectIO staging buffer that the requester will fill in a single `batch_get_into_multi_buffers` call. Lower → more conservative pre-split, more round trips. | 0.9 |
 
 ## KV Transfer Config
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
index 7df2e902fd33..8dc949cf326b 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
@@ -3,9 +3,9 @@
 """Per-operation telemetry for MooncakeStoreConnector.
 
 Records one row per Mooncake RPC (``save_exists``, ``save_put``, ``load_get``,
-``lookup_exists``, ``lookup_lease``) with duration, key/byte counts, status,
-and failed-key count. Exposed to the logger via ``KVConnectorLogging`` and to
-Prometheus via ``MooncakeStorePromMetrics``.
+``lookup_exists``) with duration, key/byte counts, status, and failed-key
+count. Exposed to the logger via ``KVConnectorLogging`` and to Prometheus
+via ``MooncakeStorePromMetrics``.
 """
 
 from dataclasses import dataclass
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
index 4775500045ef..6cbbbda1ddb2 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
@@ -2,10 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Selective read leases for Mooncake store prefix lookups.
 
-``batch_is_exist`` leases every key it finds, so a lookup renews the whole
-prompt history, including the Mamba state at every earlier block boundary of a
-hybrid model. With ``VLLM_MOONCAKE_STORE_SELECTIVE_LEASE`` the lookup probes
-lease-free (``batch_probe_key``) and leases only the keys the hit will load.
+``batch_is_exist`` leases every key it finds, so a lookup renews every stored
+key of the prompt, including keys a load of the hit skips, such as blocks
+outside a sliding window or states at earlier block boundaries. With
+``VLLM_MOONCAKE_STORE_SELECTIVE_LEASE`` the lookup probes lease-free
+(``batch_probe_key``) and leases only the keys the hit will load.
 """
 
 import time
@@ -95,10 +96,12 @@ def lease_hit(
         """Lease the keys a load of ``hit`` reads and return the hit to load.
 
         ``exists_set`` holds the probed ``(group, hash)`` pairs and
-        ``longest_hit`` recomputes the hit from them. A key evicted since the
-        probe shortens the hit, whose tail keys then differ, so the shorter hit
-        is leased in turn. Returns an empty hit if the lease fails or the hit
-        keeps shrinking.
+        ``longest_hit`` recomputes the hit from them. The lease covers only the
+        keys this hit loads, e.g. the state at its tail boundary but not states
+        at earlier boundaries. If a key was evicted since the probe, the shorter
+        hit can end at one of those unleased keys, so it is leased in another
+        round. Returns an empty hit if the lease fails or the hit keeps
+        shrinking.
         """
         try:
             for _ in range(_LEASE_ROUNDS):
diff --git a/vllm/envs.py b/vllm/envs.py
index e71b6ad61202..e639f22b6713 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -1773,8 +1773,7 @@ def _resolve_rust_cli_path() -> str | None:
         os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1")
     ),
     # Probe prefix-lookup keys without a read lease and lease only the keys the
-    # hit loads, so unused Mamba states at past block boundaries stay
-    # evictable. Needs Mooncake >= 0.3.14 on the master and every client.
+    # hit loads. Requires Mooncake >= 0.3.14.
     "VLLM_MOONCAKE_STORE_SELECTIVE_LEASE": lambda: (
         os.getenv("VLLM_MOONCAKE_STORE_SELECTIVE_LEASE", "False").lower()
         in ("true", "1")

From d4fa8ebf4ea4a006da2a844e419a0d550bf3a244 Mon Sep 17 00:00:00 2001
From: Wei Zhao <weizha@nvidia.com>
Date: Sun, 25 Oct 2026 23:54:42 -0400
Subject: [PATCH 4/6] [KV Connector] Mooncake selective lease: fix mypy

- ChunkedTokenDatabase.process_tokens takes any Sequence of block hashes, as
  the lookup path passes one.
- Patch the store's batch_is_exist with monkeypatch in the lease tests.

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py     | 4 ++--
 .../kv_transfer/kv_connector/v1/mooncake/store/data.py        | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
index 65ab704f946c..80848f3150d4 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
@@ -1099,7 +1099,7 @@ def test_selective_lease_treats_error_results_as_failed_lease(monkeypatch):
     store = _LeaseTrackingStore()
     worker = _full_plus_mamba_worker(store)
     hs = _store_full_history(worker, store)
-    store.batch_is_exist = lambda keys: [-900] * len(keys)
+    monkeypatch.setattr(store, "batch_is_exist", lambda keys: [-900] * len(keys))
 
     assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 0
     assert worker.get_kv_connector_stats().reduce()["lookup_lease_error_count"] == 1
@@ -1115,7 +1115,7 @@ def test_selective_lease_reports_no_hit_when_lease_raises(monkeypatch):
     def _fail(_keys):
         raise RuntimeError("master unreachable")
 
-    store.batch_is_exist = _fail
+    monkeypatch.setattr(store, "batch_is_exist", _fail)
 
     assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 0
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
index 4cd0402b7b4f..ece941a52956 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
@@ -605,7 +605,7 @@ def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:
     def process_tokens(
         self,
         token_len: int,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
         mask_num: int = 0,
         *,
         chunk_mask: list[bool] | None = None,

From d988b93233d38278fa514ca4b4c9db943e65ded3 Mon Sep 17 00:00:00 2001
From: Yifan Qiao <yifanqiao@inferact.ai>
Date: Thu, 29 Oct 2026 01:41:56 +0000
Subject: [PATCH 5/6] [Mooncake] Reuse lookup masks and share lease/load key
 selection

Retain the final hit masks instead of repeating the coordinator search
when leasing. Use the token database iterator for mask filtering and
stored tail keys in both lease and receive paths, preserving suffix-mask
semantics for saves without copying receive masks.

Keep bounded eviction retries and refresh their masks with the hit.
Targeted connector tests and changed-file pre-commit checks pass.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
---
 .../unit/test_mooncake_store_coordinator.py   | 11 +++--
 .../unit/test_mooncake_store_hma_e2e.py       | 13 ++++++
 .../unit/test_mooncake_store_worker.py        | 11 ++++-
 .../kv_connector/v1/mooncake/store/data.py    | 18 ++++++--
 .../v1/mooncake/store/selective_lease.py      | 41 ++++++++---------
 .../kv_connector/v1/mooncake/store/worker.py  | 46 +++++++++----------
 6 files changed, 85 insertions(+), 55 deletions(-)

diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
index 5a1a00158691..2dc6272cb932 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
@@ -589,12 +589,13 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention():
     coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
     hs = _hashes(4)
     cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
-    _masks, hit = coord.find_longest_cache_hit(
+    lookup_masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
     assert hit == 48  # eagle popped 1 block
 
     masks = coord.load_mask(hs, token_len=hit)
+    assert lookup_masks == masks  # Lookup leases and receives need the same blocks.
     # Every chunk that process_tokens(token_len=48, ...) would yield must
     # have a corresponding mask slot. process_tokens emits chunk_id 0..2
     # (start=0, 16, 32), so the mask must be length 3, all True.
@@ -613,13 +614,14 @@ def test_load_mask_with_eagle_hybrid_full_plus_swa():
     hs = _hashes(4)
     exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
     cmap = ExternalCachedBlockPool(16, exists)
-    _masks, hit = coord.find_longest_cache_hit(
+    lookup_masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
     # FullAttn dictates the convergence; eagle pops one block off it.
     assert hit == 48
 
     masks = coord.load_mask(hs, token_len=hit)
+    assert lookup_masks == masks  # Lookup leases and receives need the same blocks.
     # FullAttn: all chunks populated locally.
     assert masks[0] == [True, True, True]
     # SWA: tail-window only (ceil((32-1)/16) = 2 trailing blocks).
@@ -668,12 +670,11 @@ def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun():
     hs = _hashes(4)
     exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
     cmap = ExternalCachedBlockPool(16, exists)
-    _masks, hit = coord.find_longest_cache_hit(
-        hs, max_length=64, cached_block_pool=cmap
-    )
+    masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
     # FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The
     # Mamba group must serve its state@48 snapshot, not peek to state@64.
     assert hit == 48
+    assert masks == coord.load_mask(hs, token_len=hit)
 
 
 def test_eagle_flag_propagates_to_all_merged_swa_groups():
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
index 80848f3150d4..09cbcbd7936b 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
@@ -383,6 +383,19 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
     assert out[0][2].hex() == fine_hashes[3].hex()
     assert out[1][2].hex() == fine_hashes[6].hex()
 
+    # Loading a suffix can use a stored boundary beyond the reusable tail.
+    out = list(
+        db.process_tokens(
+            token_len=28,
+            block_hashes=fine_hashes,
+            mask_num=16,
+            chunk_mask=[False, True],
+            chunk_mask_start=0,
+            tail_key_boundary=32,
+        )
+    )
+    assert out == [(16, 28, fine_hashes[7])]
+
 
 def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta:
     """Add the tail's non-Mamba puts as the scheduler does before the worker."""
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
index 5b91aed0d989..4785afb5ee91 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
@@ -780,7 +780,13 @@ def test_process_tokens_uses_mask_num_as_start_chunk():
     assert results == [(64, 96, bytes([11]))]
 
 
-def test_process_tokens_applies_chunk_mask_before_hash_access():
+@pytest.mark.parametrize(
+    ("chunk_mask", "chunk_mask_start"),
+    [([False, True], None), ([True, True, False, True], 0)],
+)
+def test_process_tokens_applies_chunk_mask_before_hash_access(
+    chunk_mask, chunk_mask_start
+):
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0),
         block_size=32,
@@ -793,7 +799,8 @@ def test_process_tokens_applies_chunk_mask_before_hash_access():
             token_len=128,
             block_hashes=block_hashes,
             mask_num=64,
-            chunk_mask=[False, True],
+            chunk_mask=chunk_mask,
+            chunk_mask_start=chunk_mask_start,
         )
     )
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
index ece941a52956..6c77399883cf 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
@@ -609,6 +609,8 @@ def process_tokens(
         mask_num: int = 0,
         *,
         chunk_mask: list[bool] | None = None,
+        chunk_mask_start: int | None = None,
+        tail_key_boundary: int | None = None,
         put_step: int = 1,
         put_step_rank: int = 0,
     ) -> Iterable[tuple[int, int, BlockHash]]:
@@ -628,6 +630,9 @@ def process_tokens(
             mask_num: Number of tokens to skip from the beginning.
             chunk_mask: Optional mask relative to the first chunk after
                 ``mask_num``. False entries are skipped before hash access.
+            chunk_mask_start: Absolute chunk index of the first mask entry.
+                Defaults to the first chunk after ``mask_num``.
+            tail_key_boundary: Token boundary identifying the stored tail key.
             put_step: Stride for distributing chunks across ranks.
             put_step_rank: ``chunk_id % put_step`` value this rank stores.
 
@@ -639,16 +644,23 @@ def process_tokens(
         assert token_len // self.hash_block_size <= len(block_hashes)
         start_chunk = max(0, cdiv(mask_num, self.block_size))
         max_chunks = cdiv(token_len, self.block_size)
+        if chunk_mask_start is None:
+            chunk_mask_start = start_chunk
         if chunk_mask is not None:
-            max_chunks = min(max_chunks, start_chunk + len(chunk_mask))
+            max_chunks = min(max_chunks, chunk_mask_start + len(chunk_mask))
         for chunk_id in range(start_chunk, max_chunks):
-            if chunk_mask is not None and not chunk_mask[chunk_id - start_chunk]:
+            if chunk_mask is not None and not chunk_mask[chunk_id - chunk_mask_start]:
                 continue
             if chunk_id % put_step != put_step_rank:
                 continue
             start_idx = chunk_id * self.block_size
             end_idx = min(start_idx + self.block_size, token_len)
-            h = block_hashes[end_idx // self.hash_block_size - 1]
+            key_end = (
+                tail_key_boundary
+                if end_idx == token_len and tail_key_boundary is not None
+                else end_idx
+            )
+            h = block_hashes[key_end // self.hash_block_size - 1]
             yield start_idx, end_idx, h
 
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
index 6cbbbda1ddb2..15aa5efef764 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
@@ -13,9 +13,6 @@
 from collections.abc import Callable, Sequence
 from typing import Any, Protocol
 
-from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import (  # noqa: E501
-    MooncakeStoreCoordinator,
-)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     ChunkedTokenDatabase,
     MooncakeLookupResult,
@@ -59,7 +56,6 @@ class SelectiveLookupLease:
     def __init__(
         self,
         store: Any,
-        coord: MooncakeStoreCoordinator,
         token_dbs: Sequence[ChunkedTokenDatabase],
         kv_cache_groups: Sequence[KVCacheGroupSpec],
         lookup_key_prefixes: Sequence[Sequence[str]],
@@ -67,7 +63,6 @@ def __init__(
     ):
         self._check_supported(store)
         self.store = store
-        self.coord = coord
         self.token_dbs = token_dbs
         self.kv_cache_groups = kv_cache_groups
         self.lookup_key_prefixes = lookup_key_prefixes
@@ -90,8 +85,12 @@ def lease_hit(
         self,
         block_hashes: Sequence[BlockHash],
         hit: MooncakeLookupResult,
+        load_masks: tuple[list[bool], ...],
         exists_set: set[tuple[int, bytes]],
-        longest_hit: Callable[[set[tuple[int, bytes]]], MooncakeLookupResult],
+        longest_hit: Callable[
+            [set[tuple[int, bytes]]],
+            tuple[MooncakeLookupResult, tuple[list[bool], ...]],
+        ],
     ) -> MooncakeLookupResult:
         """Lease the keys a load of ``hit`` reads and return the hit to load.
 
@@ -107,11 +106,11 @@ def lease_hit(
             for _ in range(_LEASE_ROUNDS):
                 if hit.hit_length <= 0:
                     return hit
-                evicted = self._lease(block_hashes, hit)
+                evicted = self._lease(block_hashes, hit, load_masks)
                 if not evicted:
                     return hit
                 exists_set = exists_set - evicted
-                hit = longest_hit(exists_set)
+                hit, load_masks = longest_hit(exists_set)
         except _LeaseFailedError as e:
             logger.error("Failed to lease a Mooncake lookup hit: %s", e)
         except Exception:
@@ -124,33 +123,30 @@ def loaded_keys(
         self,
         block_hashes: Sequence[BlockHash],
         hit: MooncakeLookupResult,
+        load_masks: tuple[list[bool], ...],
     ) -> tuple[list[str], list[tuple[int, bytes]]]:
         """Keys a load of ``hit`` reads, in every rank namespace, and their
         ``(group, hash)`` pairs.
 
-        Mirrors ``KVCacheStoreRecvingThread._handle_request``: each group's
-        load mask, with the tail chunk keyed by the hit's tail boundary.
+        Uses the same chunk iterator and tail keys as the receive path.
         """
         hit_length = hit.hit_length
-        load_masks = self.coord.load_mask(block_hashes, hit_length)
         tail_tokens = {b.group_id: b.num_tokens for b in hit.tail_key_boundaries}
         keys: list[str] = []
         key_meta: list[tuple[int, bytes]] = []
         for g_idx, db in enumerate(self.token_dbs):
             if not self.kv_cache_groups[g_idx].kv_cache_spec.prefix_cacheable:
                 continue
-            mask = load_masks[g_idx]
             key_prefixes = self.lookup_key_prefixes[g_idx]
-            for start, end, block_hash in db.process_tokens(hit_length, block_hashes):
-                chunk_idx = start // db.block_size
-                if chunk_idx >= len(mask) or not mask[chunk_idx]:
-                    continue
-                if end == hit_length and g_idx in tail_tokens:
-                    block_hash = block_hashes[
-                        tail_tokens[g_idx] // db.hash_block_size - 1
-                    ]
+            for _, _, block_hash in db.process_tokens(
+                hit_length,
+                block_hashes,
+                chunk_mask=load_masks[g_idx],
+                tail_key_boundary=tail_tokens.get(g_idx),
+            ):
+                hash_hex = block_hash.hex()
                 keys.extend(
-                    PoolKey.build_key_string(key_prefix, block_hash.hex())
+                    PoolKey.build_key_string(key_prefix, hash_hex)
                     for key_prefix in key_prefixes
                 )
                 key_meta.extend([(g_idx, bytes(block_hash))] * len(key_prefixes))
@@ -160,6 +156,7 @@ def _lease(
         self,
         block_hashes: Sequence[BlockHash],
         hit: MooncakeLookupResult,
+        load_masks: tuple[list[bool], ...],
     ) -> set[tuple[int, bytes]]:
         """Lease the keys a load of ``hit`` reads.
 
@@ -167,7 +164,7 @@ def _lease(
         ``_LeaseFailedError`` when the lease RPC fails, including per-key error
         results.
         """
-        keys, key_meta = self.loaded_keys(block_hashes, hit)
+        keys, key_meta = self.loaded_keys(block_hashes, hit, load_masks)
         if not keys:
             return set()
         start = time.perf_counter()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
index c8c310fa659a..87b9cc519c1e 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
@@ -1154,18 +1154,13 @@ def _handle_request(self, req_meta: ReqMeta):
             chunks: list[tuple[int, int]] = []
             store_shard_ids: list[StoreShardId] = []
             for start, end, block_hash in db.process_tokens(
-                token_len, req_meta.block_hashes, mask_num
+                token_len,
+                req_meta.block_hashes,
+                mask_num,
+                chunk_mask=mask,
+                chunk_mask_start=0,
+                tail_key_boundary=tail_key_boundaries.get(g_idx),
             ):
-                chunk_idx = start // db.block_size
-                if chunk_idx >= len(mask) or not mask[chunk_idx]:
-                    continue
-                boundary_tokens = (
-                    tail_key_boundaries.get(g_idx) if end == token_len else None
-                )
-                if boundary_tokens is not None:
-                    block_hash = req_meta.block_hashes[
-                        boundary_tokens // db.hash_block_size - 1
-                    ]
                 for store_shard_id in db.store_layout.local_shard_ids:
                     key_list.append(db.store_layout.key_for(store_shard_id, block_hash))
                     chunks.append((start, end))
@@ -1541,7 +1536,6 @@ def __init__(
         if envs.VLLM_MOONCAKE_STORE_SELECTIVE_LEASE:
             self._selective_lease = SelectiveLookupLease(
                 self.store,
-                self.coord,
                 self.token_dbs,
                 self._kv_cache_groups,
                 self._lookup_key_prefixes,
@@ -2200,12 +2194,15 @@ def lookup(
                 exists_set.add((g_idx, hash_bytes))
             pos += count
 
-        hit = self._longest_hit(num_tokens, token_len, block_hashes, exists_set)
+        hit, load_masks = self._longest_hit(
+            num_tokens, token_len, block_hashes, exists_set
+        )
         if self._selective_lease is None:
             return hit
         return self._selective_lease.lease_hit(
             block_hashes,
             hit,
+            load_masks,
             exists_set,
             functools.partial(self._longest_hit, num_tokens, token_len, block_hashes),
         )
@@ -2216,13 +2213,13 @@ def _longest_hit(
         token_len: int,
         block_hashes: Sequence[BlockHash],
         exists_set: set[tuple[int, bytes]],
-    ) -> MooncakeLookupResult:
-        """Longest prefix hit over the ``(group, hash)`` pairs in the store."""
+    ) -> tuple[MooncakeLookupResult, tuple[list[bool], ...]]:
+        """Longest prefix hit and its load masks over the stored group/hash pairs."""
         cached_block_pool = ExternalCachedBlockPool(
             self.hash_block_size,
             exists_set,
         )
-        _, hit_length = self.coord.find_longest_cache_hit(
+        load_masks, hit_length = self.coord.find_longest_cache_hit(
             block_hashes,
             token_len,
             cached_block_pool,
@@ -2230,19 +2227,22 @@ def _longest_hit(
         if hit_length >= num_tokens:
             usable_length = self.coord.align_lookup_length(num_tokens - 1)
             if usable_length <= 0:
-                return MooncakeLookupResult(0)
-            _, hit_length = self.coord.find_longest_cache_hit(
+                return MooncakeLookupResult(0), ()
+            load_masks, hit_length = self.coord.find_longest_cache_hit(
                 block_hashes,
                 usable_length,
                 cached_block_pool,
             )
-        return MooncakeLookupResult(
-            hit_length,
-            self._tail_key_boundaries(
-                block_hashes,
+        return (
+            MooncakeLookupResult(
                 hit_length,
-                cached_block_pool,
+                self._tail_key_boundaries(
+                    block_hashes,
+                    hit_length,
+                    cached_block_pool,
+                ),
             ),
+            load_masks,
         )
 
     def _tail_key_boundaries(

From 26b5dcce4dfaa768c8dfe833d43068fc36f92d8d Mon Sep 17 00:00:00 2001
From: wzhao18 <wzhao18.sz@gmail.com>
Date: Mon, 2 Nov 2026 15:28:11 -0500
Subject: [PATCH 6/6] [Mooncake] Fix mypy: process_tokens takes list[BlockHash]
 on main

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
---
 .../kv_connector/v1/mooncake/store/selective_lease.py         | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
index 6318cdaa2f72..a78ad5af363e 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py
@@ -11,7 +11,7 @@
 
 import time
 from collections.abc import Callable, Sequence
-from typing import Any, Protocol
+from typing import Any, Protocol, cast
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     ChunkedTokenDatabase,
@@ -146,7 +146,7 @@ def loaded_keys(
             key_prefixes = self.lookup_key_prefixes[g_idx]
             for start, _, block_hash in db.process_tokens(
                 hit_length,
-                block_hashes,
+                cast(list[BlockHash], block_hashes),
                 tail_key_boundary=tail_tokens.get(g_idx),
             ):
                 block_index = start // store_block_size