[Fix] Isolate Mooncake KV-store keys by cache identity

[wildoranges]

Local PR draft - F2 Mooncake KV-store key identity

Related issue: #4930

Motivation

Mooncake external KV-cache keys currently identify entries by model name, an
optional cache prefix, shard rank, and token-prefix hash. The key does not bind
the stored bytes to the resolved KV-cache format or the model weights lineage.
The default empty cache prefix also leaves deployments sharing a Mooncake store
without a required tenant partition.

This can advertise stale same-dtype KV after a weight revision or online weight
update. Mixed KV formats can produce false external-cache hits and unnecessary
transfer work, and deployments with colliding default configuration can share a
keyspace unintentionally. See #4930 for the production-code probe and impact
analysis.

Modification

• Require a non-empty cache_prefix for Mooncake deployments.
• Accept weights_version, using the engine model revision when available, and
  reject configurations without a usable weights identity.
• Compute a canonical resolved KV-format identity from device type, model dtype,
  quant_policy, and MLA KV-cache dtype.
• Hash the model name, tenant prefix, KV format, weights version, and a
  versioned weights generation into the Mooncake namespace.
• Rotate the namespace after each completed online weight update, including both
  serialized and distributed update paths, and propagate the new metadata to
  active Mooncake transfer workers.
• Add focused tests for namespace partitioning, configuration validation,
  generation rotation, and sender/receiver metadata propagation.
• Update CLI help with the required Mooncake identity configuration and example.

The change is intentionally limited to F2 external-store key identity. KV-head
shard schema negotiation (F8), local prefix-cache invalidation (F1), and
adapter-weight fingerprinting are not part of this PR.

The documented update-weights workflow drains/sleeps the engine before loading
new weights. This PR rotates the active connector identity after the completed
update; ordering for transfer tasks that are intentionally left queued across an
update remains part of the separate cache-lifecycle work.

BC-breaking (Optional)

Yes, for the unreleased Mooncake connector configuration. Enabling
MooncakeStoreConnector now requires an explicit non-empty cache_prefix and a
non-empty weights_version, unless the engine provides a model revision. Old
unscoped remote entries are no longer read by the new namespace, which is
intentional because their format and weights lineage cannot be verified.

Non-Mooncake connectors, local KV caching, and the public inference request API
are unchanged.

Use cases (Optional)

• PD disaggregation deployments sharing one Mooncake cluster with separate
  tenant prefixes.
• Restarted or hot-updated deployments that must not consume KV produced by an
  older weights generation.
• Deployments using different KV-cache dtypes or quantization policies against
  the same external store.

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
   • git diff --check passes.
   • pre-commit run --files <changed-files> passes.
2. The modification is covered by complete unit tests.
   • Focused tests were added under tests/pytorch/kv_connector/ and
     tests/pytorch/engine/.
   • Focused pytest passes in the installed CPU environment (125 passed).
3. If the modification has a dependency on downstream projects of a newer
   version, this PR should be tested with all supported versions of downstream
   projects.
   • No downstream API dependency is introduced.
4. The documentation has been modified accordingly, like docstring or example
   tutorials.
   • Mooncake CLI help and relevant implementation docstrings describe the
     new identity requirements.

Verification command

pytest -q tests/pytorch/kv_connector \
  tests/pytorch/engine/test_kv_connector_wiring.py \
  tests/pytorch/engine/test_model_agent_kv_connector.py

Current result: 125 passed in 6.31s in the local CPU environment and 125 passed in 11.32s on rtx4080 with CUDA_VISIBLE_DEVICES=0, torch 2.12.1+cu130, and 8x RTX 4080 available. Related engine regression tests also
pass on the GPU host: 14 passed in 8.68s. compileall, ruff check,
git diff --check, and all configured pre-commit hooks pass. A live Mooncake
Store integration was not run because no Mooncake Store service/configuration
was provided.

[Copilot]

🔵 Needs a closer look

It changes production-critical external KV namespace semantics and online weight-update behavior without an end-to-end Mooncake Store integration run to validate real backend interoperability.

Pull request overview

This PR hardens the Mooncake external KV-store keyspace by binding store keys to an explicit tenant namespace plus a resolved KV-cache format and weights lineage, and by rotating the namespace across online weight updates to avoid stale/external-cache reuse.

Changes:

• Require explicit cache_prefix and weights_version (or engine-provided model revision) for Mooncake connector configuration.
• Introduce a versioned, hashed Mooncake store namespace derived from tenant/model/kv-format/weights identity plus a monotonic weights_generation.
• Rotate and propagate weights_generation after online weight updates, and add focused unit tests covering partitioning, validation, and propagation.

File summaries

File	Description
tests/pytorch/kv_connector/test_mooncake_store_worker.py	Updates test configs for new identity requirements; adds tests for partitioning and generation rotation propagation to worker threads.
tests/pytorch/kv_connector/test_mooncake_store_connector.py	Extends connector delegation tests to cover set_weights_generation; updates required extra config fields.
tests/pytorch/engine/test_model_agent_kv_connector.py	Adds tests ensuring weight updates trigger Mooncake namespace rotation and metadata propagation.
tests/pytorch/engine/test_kv_connector_wiring.py	Validates config preparation enforces identity requirements and binds kv-format/generation to connector config.
lmdeploy/pytorch/kv_connector/mooncake/store/worker.py	Threads kv_cache_format, weights_version, and weights_generation into key metadata; adds runtime generation rotation API.
lmdeploy/pytorch/kv_connector/mooncake/store/data.py	Implements namespace_hash and a v2 namespace schema; enforces non-empty identity fields; updates key builder format.
lmdeploy/pytorch/kv_connector/mooncake/store/connector.py	Exposes set_weights_generation on the connector API and forwards to the worker.
lmdeploy/pytorch/kv_connector/factory.py	Enforces explicit tenant+weights identity; adds model-identity binding (kv_cache_format, weights_generation).
lmdeploy/pytorch/kv_connector/base.py	Adds a no-op default set_weights_generation hook to the base connector API.
lmdeploy/pytorch/kv_connector/init.py	Exports prepare_kv_connector_model_identity.
lmdeploy/pytorch/engine/model_agent/agent.py	Binds kv-format/generation before building the worker connector; rotates generation after weight updates and propagates to the connector.
lmdeploy/pytorch/engine/engine.py	Passes engine_config.revision through to connector config preparation for weights identity.
lmdeploy/cli/utils.py	Updates CLI help text/example to document required Mooncake identity configuration.

Review details

• Files reviewed: 13/13 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.