[RFC] Unified KVCache and Model Weight Management in Mooncake Store

# [RFC] Unified KVCache and Model Weight Management in Mooncake Store

## Summary

This RFC proposes extending Mooncake Store from a KVCache-centric distributed store into a unified inference data resource layer that can manage both KVCache and model weights.

The first milestone is intentionally small: import model weights into Mooncake Store as `ObjectDataType::WEIGHT` objects, hard pin them, and let an inference engine load a model from Mooncake Store during startup.

Longer term, Mooncake can provide a type-aware resource view across KVCache and Weight objects, including quota, eviction priority, pin semantics, tiering, metrics, and scheduler-facing query APIs.

## Motivation

Large model serving consumes two major categories of reusable data:

1. KVCache, which is generated at runtime and reused across requests or instances.
2. Model weights, which are stored as model artifacts before serving and repeatedly loaded during instance startup.

Both data types consume the same physical resources: DRAM, SSD, network bandwidth, and potentially VRAM. If KVCache and model weights are managed by separate systems, each system only sees a partial resource view. This can cause resource fragmentation and local optimization. For example, one system may preserve cold model weights while another system evicts valuable KVCache because it cannot see the full resource picture.

Mooncake Store already has many primitives needed for this direction: object/tensor APIs, RDMA/GDR transfer, replicas, leases, soft/hard pin, SSD offload, copy/move tasks, and `ObjectDataType` values such as `KVCACHE` and `WEIGHT`.

The missing part is not the raw storage path. The missing part is a model-version-level Weight Management Layer:

```text
model_id / revision -> manifest -> shards / tensors -> readiness
-> serving lease -> import / query / delete / unpin
-> quota / eviction / scheduling policy
```

This RFC defines an incremental path to add that layer without building a separate resource management system.

## Goals

1. Treat model weights as first-class inference data in Mooncake Store.
2. Reuse existing Store primitives for object storage, tensor storage, transfer, replica, lease, pin, and offload.
3. Preserve the semantic differences between KVCache and model weights.
4. Start with a minimal end-to-end path: import weights, hard pin them, and load them from Store.
5. Add weight lifecycle management before adding dynamic KVCache/Weight resource competition policies.
6. Expose a unified resource view for schedulers and serving platforms in later phases.

## Non-Goals

The first version does not aim to implement:

- fallback to remote storage on Store miss
- full Weight CRUD management APIs
- automatic cold/hot weight migration
- dynamic KVCache/Weight eviction arbitration
- VRAM hot weight cache
- CUDA MPS or fine-grained GPU memory isolation
- direct GPU-buffer weight loading
- predictive preloading
- cross-region weight distribution

## Current Foundation

Mooncake Store can already store model weight bytes or tensors as regular objects. A weight object can be written with:

```text
data_type = ObjectDataType::WEIGHT
with_hard_pin = true
replica_num = N
```

This is enough to protect a weight object from normal eviction and make it readable through Store `Get`.

However, Store currently manages individual objects, not model revisions. It does not yet know whether all shards of a model revision are present, whether their checksums match, whether a revision is ready for serving, or whether it is safe to delete or unpin that revision.

## Architecture

The proposed architecture keeps Mooncake Store as the common data and resource substrate. KVCache and Weight share object storage, transfer, metadata, replicas, leases, pins, tiering, metrics, and resource accounting. Their lifecycle policies remain type-specific.

```mermaid
flowchart TB
Platform[Serving Platform / Scheduler] --> WML[Weight Management Layer]
Engine[Inference Engine] --> Loader[Weight Loader]
Engine --> KVRuntime[KVCache Runtime]

WML --> Master[Mooncake Master]
Loader --> Store[Mooncake Store]
KVRuntime --> Store
Store --> Master
Store --> TE[Transfer Engine]

Store --> ObjKV[KVCACHE Objects]
Store --> ObjWeight[WEIGHT Objects]
Store --> ObjMeta[METADATA Objects]

TE --> DRAM[DRAM Tier]
TE --> SSD[SSD Tier]
TE --> VRAM[Future VRAM Tier]

Source[Model Artifacts / Remote Storage] --> WML
```

The main components are:

1. **Weight Management Layer**: defines model-version-level metadata, manifests, readiness, import, query, delete, unpin, and serving protection.
2. **Mooncake Store**: stores `KVCACHE`, `WEIGHT`, `METADATA`, and other object types, and manages replicas, leases, pins, eviction, offload, and copy/move tasks.
3. **Inference Engine Loader**: loads model weights from Mooncake Store during startup. The first version may materialize Store objects back to local temporary files and reuse the existing model loader.
4. **KVCache Runtime**: continues using Mooncake Store as a KVCache backend through existing runtime cache paths.
5. **Scheduler or Serving Platform**: queries where model weights are ready, where KVCache may be useful, and how much resource remains on each node or tier.

## KVCache and Weight Semantics

KVCache and Weight should be unified at the resource layer, not forced into the same cache semantics.

| Dimension | KVCache | Weight |
| --- | --- | --- |
| Source | Generated by inference runtime | Imported from model artifacts |
| Lifecycle | Request/session related | Model-version related |
| Eviction impact | Lower hit rate and extra prefill cost | Startup or serving failure if required shards disappear |
| Integrity unit | Page/block/key can often miss independently | Readiness is checked at model-version level |
| Write pattern | Frequent dynamic writes | Low-frequency import, high-concurrency reads |
| Recovery | Recompute on miss | Re-import missing shards according to manifest |

The target model is:

```text
shared storage substrate + type-specific lifecycle policies
```

KVCache may be evicted to trade hit rate for capacity. Weight must not be silently and partially evicted while a model revision is `READY` or `SERVING`.

## Proposed Design

### Phase 1: Store-Backed Weight Loading

Phase 1 proves the minimal data path:

1. Import model weights into Mooncake Store.
2. Store each weight file or shard as a `WEIGHT` object.
3. Hard pin weight objects.
4. Generate a minimal manifest.
5. Let an inference engine load the model from Mooncake Store.

In the first version, model weights should be stored at safetensors file or file-shard granularity:

```text
weight:{model_id}:{revision}:file:{file_name}
weight:{model_id}:{revision}:manifest
```

Each weight object is written with:

```text
ObjectDataType::WEIGHT
with_hard_pin = true
replica_num = 1
```

The manifest should record:

- `model_id`
- `revision`
- source URI or source description
- file list
- Store object keys
- file sizes
- checksums or equivalent validation fields
- readiness status

The first import tool can be simple:

```text
mooncake_weight_import --model-id --revision --source
mooncake_weight_check --model-id --revision
```

The first inference-engine integration can use a URI or equivalent configuration:

```text
--model mooncake-weight:///
```

Startup flow:

```mermaid
flowchart LR
Source[Model Artifacts] --> Import[Weight Import Tool]
Import --> Manifest[Generate Manifest]
Import --> Put[Put WEIGHT Objects
with hard pin]
Manifest --> Store[Mooncake Store]
Put --> Store

Startup[Inference Engine Startup] --> Resolve[Resolve mooncake-weight URI]
Resolve --> ReadManifest[Read Manifest]
ReadManifest --> ReadWeights[Read Weight Objects]
Store --> ReadManifest
Store --> ReadWeights
ReadWeights --> LocalFiles[Materialize Local Temp Files]
LocalFiles --> ExistingLoader[Existing Model Loader]
ExistingLoader --> Ready[Model Ready]
```

The loader resolves the manifest, checks minimal readiness, reads required weight objects from Mooncake Store, writes them to a local temporary model directory, and then reuses the existing safetensors or model-file loader.

This is not the final performance path, but it minimizes integration risk and proves that Mooncake Store can act as the model startup data source.

Phase 1 success criteria:

- A model can be imported into Mooncake Store.
- Store metadata records weight objects as `ObjectDataType::WEIGHT`.
- Weight objects are hard pinned and not selected by normal eviction.
- An inference engine can start from Mooncake Store without reading the original model files directly.

### Phase 2: Weight Lifecycle Management

After the minimal data path works, Mooncake should promote Weight from raw Store objects into a model artifact resource.

Phase 2 should focus on correctness and operability:

- standardized Weight import
- list/query/delete/unpin by `model_id` and `revision`
- immutable revision semantics
- model-version readiness
- model-version state machine
- serving lease or reference protection
- missing shard re-import based on manifest
- weight-specific metrics
- scheduler-facing weight residency query

Recommended model-version states:

```text
IMPORTING -> READY -> SERVING -> IDLE
READY / IDLE -> EVICTING -> EVICTED
EVICTED / partial missing -> REHYDRATING -> READY
```

Phase 2 ensures Weight itself is not partially corrupted, can be recovered, and can be safely deleted. It does not need to solve dynamic resource competition between KVCache and Weight.

If a weight object is merely offloaded from DRAM to SSD and Store `Get` can still read it, the model revision should remain `READY`. `EVICTED` should mean Mooncake no longer holds a complete readable copy of the required weight data.

### Phase 3: Unified Resource Management

Once Weight lifecycle management is available, Mooncake Master can manage KVCache and Weight in a unified resource view.

Phase 3 should add:

- resource accounting by `ObjectDataType`
- type-aware quota
- type-aware eviction priority
- clear pin behavior under resource pressure
- deterministic pressure handling
- unified resource view APIs
- scheduler integration
- type-aware metrics and alerts

The important policy boundary is:

```text
KVCache eviction can reduce hit rate.
Weight eviction must happen at model-revision granularity.
READY or SERVING Weight must not be silently partially evicted.
```

When capacity is insufficient, Mooncake should follow a predictable order, for example:

```text
expired / unleased KVCache
-> over-quota low-priority objects
-> cold IDLE Weight revisions through model-version eviction
-> explicit failure
```

### Future Work

Future phases can explore:

- direct host/GPU buffer loading
- tensor-, layer-, expert-, and TP-shard-level weight storage
- dynamic weight swapping
- automatic replica adjustment
- predictive model preloading
- VRAM hot weight cache
- fine-grained GPU sharing

These are optimization and serving-density features. They should build on the correctness guarantees from Phases 1 to 3.

## Compatibility

This design should remain backward compatible:

1. Existing KVCache clients continue using Store APIs unchanged.
2. Objects without explicit type can keep using `UNKNOWN` or existing defaults.
3. Weight management APIs are additive.
4. Phase 1 can be implemented on top of existing Store `Put` and `Get` paths.
5. Type-aware quota and eviction should be opt-in at first.

## Alternatives Considered

### Separate Weight Store

Keeping weights in a separate store avoids changes to Mooncake Store, but it keeps resource management fragmented. KVCache and Weight would still compete for the same DRAM/SSD/VRAM resources without a unified quota, eviction, and scheduling view.

### P2P-Only Distribution

P2P distribution is useful for fast transfer from existing replicas, but it does not provide long-lived object residency, global capacity accounting, eviction policy, SSD offload management, or a unified KVCache/Weight resource view.

### Tensor-Level Storage First

Tensor-level storage is attractive for future direct loading, expert-level preloading, and dynamic swapping. However, it requires deeper inference-engine integration. File/shard-level storage is a smaller first step because it can reuse existing model loaders.

## Open Questions

1. What should be the default key format for model weight objects and manifests?
2. Should the manifest be stored as a normal `METADATA` object, a structured object, or a dedicated model-version metadata entry?
3. Should Phase 1 require hard pin for all weight objects, or allow configurable pin policies?
4. What is the first inference engine integration target and loader interface?
5. How should type-aware quota interact with hard-pinned weights in later phases?

## Conclusion

Mooncake already has the core Store primitives needed to hold model weights: object storage, tensor APIs, high-speed transfer, replicas, lease, hard pin, SSD offload, and object data types. The next step is to add a model-version-level Weight Management Layer and then integrate Weight into the same resource control plane as KVCache.

The recommended path is incremental:

1. First, prove Store-backed weight import and startup.
2. Then, make Weight lifecycle safe and manageable.
3. Finally, add type-aware unified resource management for KVCache and Weight.

This keeps the first implementation small while aligning Mooncake with multi-model serving and fast instance startup requirements.

Start by reading the existing Store Put/Get paths and ObjectDataType handling, then trace how the inference engine currently loads model files. A Phase 1 implementation is done when weights can be imported as hard-pinned WEIGHT objects with a manifest and an inference engine can start from the Mooncake Store without reading the original files directly.

Written by the indexing model from the issue text.