feat(pytorch): support mooncake store by caikun-pjlab · Pull Request #4880 · InternLM/lmdeploy

caikun-pjlab · 2026-08-18T10:14:27Z

Motivation

In production LLM serving, a large fraction of requests share common prompt prefixes — system prompts, few-shot examples, RAG context, multi-turn conversation history, etc. LMDeploy's local prefix caching reuses KV cache within a single engine instance, but it cannot help when:

Cross-instance reuse: multiple LMDeploy instances serving similar traffic recompute the same KV blocks independently.
Post-eviction recomputation: after KV cache eviction under memory pressure, the same blocks must be recomputed from scratch.
Cold-start warm-up: a newly launched instance has an empty KV cache and must compute every block, even those computed elsewhere.

This PR addresses these gaps by integrating Mooncake's distributed store as a shared KV-cache pool. Computed KV blocks are stored with content-addressable keys (block hashes); before computing a prefill, the engine queries the store and, on a hit, loads the blocks directly via zero-copy RDMA, skipping redundant prefill computation. The capability is orthogonal to PD disaggregation and benefits any deployment topology.

The overall design is referenced from vLLM RFC #38474 and the vllm implementation.

Deploy

This section shows a complete deployment of GLM-5.2-FP8 with TP=8, Ray as the distributed executor, prefix caching enabled, and Mooncake Store acting as a shared KV-cache pool (kv_role=kv_both).

1. Start the Mooncake master

The master coordinates the distributed store metadata. Start it on a host reachable from every LMDeploy worker:

mooncake_master \
  -rpc_port=50051 \
  -rpc_thread_num=4 \
  -default_kv_lease_ttl=30000 \
  -eviction_high_watermark_ratio=0.95 \
  -eviction_ratio=0.1 \
  -logtostderr

2. Prepare the Mooncake config

Each LMDeploy worker reads this JSON config. It declares the embedded store mode, the master address, the RDMA devices to use, and the global / local buffer sizes:

{
  "mode": "embedded",
  "metadata_server": "P2PHANDSHAKE",
  "master_server_address": "127.0.0.1:50051",
  "global_segment_size": "160GB",
  "local_buffer_size": "4GB",
  "protocol": "rdma",
  "device_name": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7",
  "enable_offload": false
}

3. Launch LMDeploy with `--kv-transfer-config`

export MOONCAKE_CONFIG_PATH=mooncake_config.json

lmdeploy serve api_server /path/to/GLM-5.2-FP8 \
    --tp 8 \
    --backend pytorch \
    --server-name 0.0.0.0 \
    --model-name GLM-5.2-FP8 \
    --server-port 9000 \
    --enable-prefix-caching \
    --log-level INFO \
    --quant-policy fp8 \
    --max-batch-size 64 \
    --cache-max-entry-count 0.9 \
    --reasoning-parser default \
    --tool-call-parser glm47 \
    --trust-remote-code \
    --max-prefill-token-num 8192 \
    --session-len 262144 \
    --distributed-executor-backend ray \
    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"mooncake_config_path":"mooncake_config.json"}}'

Test

Verified on GLM-5.2-FP8, PyTorch backend, TP=8, Ray executor, prefix caching enabled, Mooncake kv_both.

1. Precision — GSM8K

Dataset	mean_acc
GSM8K	0.9750

2. Performance

2.1 GPU KV cache sufficient (no eviction pressure)

Test method: multi-turn dialogue, SWE_Smith dataset, 8 requests, first turn length 32K, generating 512 tokens per turn, increasing by 4K per turn, for a total of 5 turns.

Metric	Mooncake OFF	Mooncake ON	ON vs OFF
Success / failed	40 / 0	40 / 0	—
Duration	99.35 s	99.38 s	+0.03%
Total throughput	16,716.60 tok/s	16,712.73 tok/s	-0.02%
Avg latency	19.70 s	19.70 s	0%
TTFT	4,314.63 ms	4,315.32 ms	0%
TPOT	30.10 ms	30.10 ms	0%

Conclusion: with sufficient GPU KV cache the connector only saves blocks asynchronously and never loads — overhead is ~0.02%, within measurement noise.

2.2 GPU KV cache insufficient (eviction pressure)

Test method: multi-turn dialogue, SWE_Smith dataset, 32 requests, first turn length 32K, generating 512 tokens per turn, increasing by 4K per turn, for a total of 5 turns.

Metric	Mooncake OFF	Mooncake ON	ON vs OFF
Success / failed	160 / 0	160 / 0	—
Duration	581.92 s	286.24 s	-50.81%
Request throughput	0.28 req/s	0.56 req/s	+100%
Total throughput	11,417.69 tok/s	23,212.21 tok/s	+103.30%
Avg latency	104.97 s	53.61 s	-48.93%
TTFT	40,499.25 ms	23,391.21 ms	-42.24%
TPOT	126.17 ms	59.13 ms	-53.13%
KV cache hit rate	27.97%	76.38%	+48.41%
Subsequent-turn TTFT	37,357.02 ms	16,128.65 ms	-56.83%

Conclusion: total throughput reaches 2.03× the OFF baseline (+103.3%). The gain comes from external cache hits on subsequent turns — once the local L1 is evicted, blocks are reloaded from Mooncake instead of being recomputed.

BC-breaking

No backward-compatibility break for existing users:

kv_transfer_config defaults to None and the connector is never imported or built when disabled — existing engine configurations and CLI invocations are unchanged.
num_replicate_key_value_heads defaults to 1 and is only consumed when a connector is enabled.
The new SchedulerOutput fields, Scheduler / ExecutorBase / BaseModelAgent methods, and CacheEngine.connector_kv_caches property are additive.
The only behavioral change visible without a connector is a longer worker-shutdown timeout in RayExecutor and ZMQMPEngine, and only when kv_transfer_config is set.

Downstream projects that vendor lmdeploy.pytorch.paging.scheduler.SchedulerOutput or subclass KVConnectorBase will see new fields / methods, all of which have defaults or no-op base implementations.

Use cases (Optional)

Cross-instance prefix reuse: a fleet of LMDeploy instances serving the same model + system prompt share a Mooncake store; each newly computed block is published once and reused by every instance.
Post-eviction recovery: when local KV cache is evicted under memory pressure, the next request with the same prefix loads blocks from the store instead of recomputing them.
Cold-start warm-up: a new instance joins the fleet and immediately benefits from blocks already in the store, eliminating cold-start prefill latency for common prefixes.

grimoire · 2026-08-18T12:38:27Z


 from lmdeploy.pytorch.backends.selector import init_backend
 from lmdeploy.pytorch.config import BackendConfig, CacheConfig, DistConfig, MiscConfig, ModelConfig, SpecDecodeConfig
+from lmdeploy.pytorch.kv_connector.base import KVConnectorMetadata, KVConnectorOutput


MPExecutor is deprecated, we don't need to update it.

grimoire · 2026-08-18T12:54:13Z

+
+    def _kv_connector_poll_acknowledgements(self) -> tuple[set[int], set[int]]:
+        """Return save/load completions to acknowledge in the next poll."""
+        sending = set(getattr(self, '_kv_connector_acknowledged_sending', set()))


_kv_connector_acknowledged_sending is initialized in the init, use getattr is too defensive here. we can just just self._kv_connector_acknowledged_sending

grimoire · 2026-08-18T13:26:05Z

+        connector = getattr(self, 'kv_connector', None)
+        if connector is None:
+            return
+        has_saves = getattr(connector, 'has_pending_step_saves', None)


These getattr are too defensive.

we can use something like

if connector.has_pending_step_saves(): event = torch.cuda.Event() event.record(self.stream) connector.submit_saves(save_ready_event=event)

right?

Yes, you are right. This version of the code contains some difficult-to-read and overly protected content, which will be fixed later.

grimoire · 2026-08-18T13:33:42Z

+                        ),
+                    )
+                    if self.kv_connector is not None:
+                        self.kv_connector.register_kv_caches(self.cache_engine.connector_kv_caches)


Do we support register_kv_caches for spec decoding? Spec decoding use different cache manager for now.

We can explicit reject connector for specdecoding for this PR and add it in the future.

OK，spec decoding is not yet supported; this version will explicitly refuse to enable spec decoding.

grimoire · 2026-08-19T03:00:23Z

+        checker = getattr(getattr(self, 'executor', None), 'has_pending_kv_connector_ack', None)
+        return checker is not None and checker()
+
+    def _has_pending_kv_connector_work(self) -> bool:


This function has not been used

grimoire · 2026-08-19T03:12:38Z

+
+    def _has_pending_kv_connector_ack(self) -> bool:
+        """Return whether the executor owes workers a sticky-completion ACK."""
+        checker = getattr(getattr(self, 'executor', None), 'has_pending_kv_connector_ack', None)


Too many getattr, I think we need a cleanup to remove unnecessary getattr.

grimoire · 2026-08-19T03:19:46Z

+        token_lens = None
+        if inputs is not None and not inputs.is_decoding:
+            ready_token_lens = inputs.history_lengths + inputs.seq_length
+            token_lens = tuple(int(token_len) for token_len in ready_token_lens.tolist())


Can we use tuple(ready_token_lens.tolist())?

grimoire · 2026-08-19T08:02:23Z

+                    0, target_blocks - int(seq.num_blocks))
+                extra_blocks = max(
+                    0, target_required_blocks - required_blocks)
+                evict_alloc_size += extra_blocks * seq.block_size


Do we have to reserve all blocks for request without remote kv?

grimoire · 2026-08-19T08:19:12Z

    def _admit_resources(self):
        if self.scheduler.block_trie.enabled:
            return self._admit_prefix_cache_resources()
+        lookup_result = self._query_external_prefix(None)


Is this lookup a "prefix matching"? What if we disable prefix caching?

grimoire · 2026-08-19T08:26:05Z

+                    int(block_id)
+                    for block_id in self.block_manager.get_block_table(seq)
+                )
+                connector.request_finished(seq, block_ids)


The docstring of request_finished:

The boolean is true when the connector temporarily takes ownership of the blocks for an asynchronous save. In that case the scheduler must defer freeing them until the request is reported by ``get_finished``. The optional dictionary is reserved for connector-specific response metadata.

but the return of request_finished has never been read here.

grimoire · 2026-08-19T08:38:53Z

+        self,
+        request: 'SchedulerSequence',
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:


Do we need load_async in this interface? None already represents an asynchronous lookup, and the scheduler always enters WAITING_FOR_REMOTE_KVS for every positive hit. A positive result with load_async=False is not actually supported. Could we remove the boolean and define every positive hit as requiring the existing asynchronous load lifecycle?

grimoire · 2026-08-19T08:58:10Z

+            # Keep accepting an explicit value while generating vLLM-style
+            # keys from the model basename by default.
+            model_name = extra_config.get('model_namespace', model_name)
+            extra_config.setdefault('model_name', model_name)


Is it safe If we launch 2 engine with different config(tp/block_size,....)

grimoire · 2026-08-19T09:03:25Z

+        return bool(self.completed_save_ids or self.completed_load_ids or self.failed_load_ids)
+
+    @property
+    def finished_sending(self) -> set[RequestId]:


is these compatible property necessary?

grimoire · 2026-08-19T09:10:04Z

+    def __init__(self, role: KVConnectorRole) -> None:
+        if not isinstance(role, KVConnectorRole):
+            raise TypeError(f'role must be a KVConnectorRole, got {type(role).__name__}')
+        self._role = role


How many of the apis are shared between different role? Is it a good idea to split them info different class by role? will we have other role in the future?

grimoire · 2026-08-19T10:07:03Z

+    contain hashes already computed for the same request identity and token
+    prefix; only newly completed blocks are encoded and hashed.
+    """
+    if isinstance(block_size, bool) or not isinstance(block_size, int) or block_size <= 0:


is this check necessary?

grimoire · 2026-08-19T10:18:23Z

+
+
+def _get_local_hostname() -> str:
+    """Resolve a routable local address, including on offline hosts."""


We can add a environment variable in pytorch/envs.py so users can choose to manually set the address. It might be helpful for debug.

grimoire · 2026-08-19T10:25:38Z

+
+    save_requests: tuple[MooncakeStoreSaveRequest, ...] = field(default_factory=tuple)
+    load_requests: tuple[MooncakeStoreLoadRequest, ...] = field(default_factory=tuple)
+    preempted_save_ids: tuple[int, ...] = field(default_factory=tuple)


has this been used?

grimoire · 2026-08-19T10:51:51Z

+    return timeout_ms
+
+
+def _make_zmq_socket(


Maybe we can reuse our executor rpc so we don't have to add another ipc here.

I can accept the zmq in this PR, and I think we can optimize it in the future.

caikun-pjlab added 9 commits August 16, 2026 11:08

feat: add KV connector base interface

76758e7

feat: add kv transfer config

409b92f

feat: add Mooncake Store connector skeleton

640798d

feat: add mooncake store setup and register_kv_caches

07fd844

feat: add mooncake LookupKey client and server

50a0be5

feat: save kvcache into mooncake store

abd22d3

feat: load kvcache from mooncake store

e4a3d93

fix: format comments

632dfad

refactor: optimize scheduler and mooncake store code

df584ec

caikun-pjlab marked this pull request as draft August 18, 2026 10:14

grimoire reviewed Aug 18, 2026

View reviewed changes

grimoire reviewed Aug 19, 2026

View reviewed changes



		def _get_local_hostname() -> str:
		"""Resolve a routable local address, including on offline hosts."""

Uh oh!

Conversation

caikun-pjlab commented Aug 18, 2026

Motivation

Deploy

1. Start the Mooncake master

2. Prepare the Mooncake config

3. Launch LMDeploy with --kv-transfer-config

Test

1. Precision — GSM8K

2. Performance

2.1 GPU KV cache sufficient (no eviction pressure)

2.2 GPU KV cache insufficient (eviction pressure)

BC-breaking

Use cases (Optional)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

3. Launch LMDeploy with `--kv-transfer-config`