[wzhao18]

Purpose

MooncakeStoreConnector prefix lookups call batch_is_exist, which grants a read lease on every key it finds and moves it to the hot end of the eviction order. A lookup therefore renews every stored key of the prompt, including keys that loading the hit never reads: blocks outside a sliding window, or the per-block-boundary states of linear-attention / Mamba layers before the hit's tail. Once the store is full, those never-loaded keys look hot and survive eviction, while keys that later requests would actually load are evicted.

This PR adds an opt-in mode, VLLM_MOONCAKE_STORE_SELECTIVE_LEASE=1, in which a lookup:

1. probes the candidate keys with batch_probe_key, which checks existence without leasing (kvcache-ai/Mooncake#3801);
2. computes the hit as before;
3. leases only the keys that loading that hit reads (batch_is_exist on the hit's load set, per group load mask and tail key).

If a key was evicted between the probe and the lease, the hit gets shorter, and the shorter hit can end at a key the first lease did not cover (for example, an earlier boundary state). So the shorter hit is leased in a second round. A hit is returned only when the lease on all of its keys succeeded; otherwise the lookup reports a miss. A failing lease RPC also reports a miss and never breaks the lookup server.

The default path is unchanged.

Changes

• New vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/selective_lease.py: SelectiveLookupLease (probe, compute the load set, lease, record the lookup_lease operation in the connector metrics).
• store/worker.py: when the env var is set, the lookup probes through SelectiveLookupLease and passes its hit to lease_hit. The longest-hit computation moves into _longest_hit so it can be reused after an eviction.
• envs.py: VLLM_MOONCAKE_STORE_SELECTIVE_LEASE (default off).
• Startup fails with a clear error if the installed Mooncake lacks MooncakeDistributedStore.batch_probe_key.

Requirements

Requires Mooncake >= 0.3.14 on the master and every client. batch_probe_key is not in a tagged release yet; this was tested with the nightly 0.3.14.dev20261022.

Test Plan

Unit tests in tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py, on a full-attention + Mamba (align mode) worker over an in-memory store:

• default mode still leases every probed key;
• selective mode leases only the loaded keys, and they cover every key the load reads (including partial hash hits);
• a key evicted between probe and lease yields a shorter, fully leased hit; continued evictions yield no hit;
• error results and a raising lease RPC report a miss;
• startup fails without batch_probe_key.

End-to-end A/B on Kimi-K3 (hybrid KDA + MLA), disaggregated 4 prefill workers (TP8) + 1 decode worker (DEP16) on GB300, agentic coding traces at concurrency 480, 60 min. Mooncake store with 160 GB per rank (7.5 TB total), eviction at 95% high-watermark. Both arms run the same build and the same Mooncake nightly; only the env var differs.

Test Result

Disaggregated, 4P1D (DEP16 decode), concurrency 480

metric	lease all (default)	selective lease	change
total throughput / GPU (tok/s)	27,491	28,305	+3.0%
P90 ITL	35.8 ms	36.4 ms	+1.7%
TTFT P50	2.35 s	1.73 s	−26%
TTFT P90	12.8 s	3.83 s	−70%
TTFT P95	20.9 s	6.39 s	−69%
E2E latency P90	79.0 s	75.8 s	−4%
prefix cache hit rate (benchmark)	93.4%	95.0%	+1.6 pt

The store's raw capacity is the same in both arms: about 523k keys at the watermark, 10% evicted per pass, no failed evictions. What changes is which keys survive. A per-request audit against the trace (which prefixes had a saved state available when the request arrived) shows misses on states saved 2–10 min earlier falling from 1.13 to 0.21 points of prompt tokens, and on states older than 10 min from 1.55 to 0.91. Misses on states younger than 2 min are unchanged (0.46 vs 0.45), since nothing that young is evicted.

Cost: the extra lease RPC averages 1.07 ms (P90 1.9 ms) per lookup, with no errors and no keys evicted between probe and lease during the run. The lease-free probe was itself slightly faster than the leasing one (36.8 ms vs 43.1 ms average).

Aggregated, concurrency sweep

One aggregated worker (TP8, DCP8, FP8 KV cache) on 8 GB300 GPUs, the same agentic coding traces, 60 min per run, Mooncake store with 160 GB per rank (1.25 TB total). Both arms run the same build; only the env var differs. Prefix cache hit rate is the benchmark's: cached prompt tokens / prompt tokens over the measured requests.

concurrency	lease	prefix cache hit rate	throughput / GPU (tok/s)	TTFT P50 / P90	P90 ITL	E2E latency P90
48	all (default)	96.50%	19,921	1.13 / 2.58 s	36.1 ms	61.3 s
48	selective	96.22%	19,531 (−2.0%)	1.10 / 2.81 s	36.4 ms	63.3 s
64	all (default)	95.05%	19,574	1.42 / 4.87 s	55.9 ms	94.6 s
64	selective	95.18%	20,341 (+3.9%)	1.36 / 3.96 s	52.6 ms	84.8 s
80	all (default)	93.33%	18,104	6.09 / 21.2 s	80.6 ms	127.0 s
80	selective	93.93%	19,254 (+6.4%)	5.30 / 19.4 s	72.2 ms	117.9 s
96	all (default)	92.48%	16,599	27.1 / 53.1 s	84.8 ms	156.9 s
96	selective	93.97%	18,080 (+8.9%)	24.6 / 43.2 s	73.5 ms	145.9 s

The gain grows with concurrency, i.e. with pressure on the store: at concurrency 96, misses on states that were saved and later evicted fall from 2.7 to 1.0 points of prompt tokens. Since the GPU KV cache is full from concurrency 64 on, the extra hits also save prefill work in the shared worker, which shows up as lower ITL.

At concurrency 48 the store is under little pressure, and selective lease is slightly worse. The difference comes from a few long requests (150k–690k tokens) that resumed from a state saved at an earlier boundary and not loaded since: the default mode kept it alive because later lookups over the same prefix renewed it, while selective lease let it age out. This is the trade-off of the mode: states that are only passed over by lookups are no longer protected.

---

🤖 Generated with Claude Code

[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.