kvcache-ai / kvcache-ai/Mooncake

[RFC]: Manifest-driven heterogeneous model weight resharding and storage

Open
#3,111 1 comment 1 reaction 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Changes proposed

### Summary

This RFC proposes manifest-driven heterogeneous model-weight resharding and storage for Mooncake. A source and target may use different `TP`, `PP`, `EP`, and `DP` layouts while sharing one model-semantic manifest, N-D logical-box planner, and backend-neutral transfer-region contract.

The proposed contract supports three paths:

* live runtime GPU-to-GPU (G2G) transfer;

* runtime-to-Store upload; and

* Store-to-runtime load.

Live G2G and Store-to-runtime use a `LogicalTransferPlan`. Runtime-to-Store uses a separate `WeightUploadPlan`, while sharing tensor identity, DP replica selection, coverage, and fragment semantics with the load path.

The framework runtime remains the source of truth for model semantics and live GPU addresses. Mooncake consumes explicit placement and runtime-binding data, computes N-D overlaps in global logical-tensor coordinates, and lowers the result into Transfer Engine (TE), Store, or future executor operations.

The first version only accepts byte-compatible layouts. A mismatch in dtype, shape, item size, or `layout_fingerprint` fails closed.

### Motivation

RL rollouts, elastic scale-out, and inference replacement can require a conversion such as:

```text
TP4 / PP2 / EP8 / DP2 -> TP8 / PP4 / EP2 / DP4

```

The existing one-dimensional TP overlap model cannot express all of the following together:

* source and target sharding different tensor dimensions;

* pipeline-parallel layer or tensor ownership changes;

* expert-parallel and tensor-parallel changes in the same tensor; and

* a common load contract for live GPU and persistent Store sources.

Mooncake already provides TE, Store, Group, leases, memory registration, and ranged transfer. What is missing is a stable contract between framework-owned model semantics and Mooncake-owned physical transfer and storage operations.

### Goals and non-goals

Goals:

* Support joint TP, PP, EP, and DP changes, including different source and target shard dimensions.

* Persist source-native TP/PP/EP fragments as an immutable weight revision that can serve future target layouts.

* Avoid model-wide, tensor-wide, or all-expert all-gathers.

* Avoid a model-sized temporary GPU or CPU buffer.

* Bind one logical load plan to either live-runtime or Store fragments; keep a dedicated upload plan for the write path.

* Preserve legacy manifests, one-dimensional TP behavior, and normal Store/TE/KVCache paths.

* Preserve lease, generation, coverage, layout, and address-bounds checks.

Non-goals:

* Inferring tensor, layer, or expert semantics from parameter names in Mooncake core.

* Owning source discovery, Kubernetes scheduling, or target traffic activation.

* Requiring NCCL M2N, `ncclMemAlloc`, or any particular framework.

### Proposed design

```mermaid
flowchart TB
P["Source Placement + Target Placement"]
L["LogicalTransferPlan
address-independent"]
B["Source Binding + Target Binding"]
T["TransferPlan
executable"]
E["TE / Store executor"]

P -->|"N-D planning"| L
L --> T
B -->|"bind and validate"| T
T --> E
```

Placement is address-independent and is the only input to N-D planning. Runtime
Binding adds the current addresses, endpoints, generation, and lease immediately
before execution. Address changes therefore require rebinding, not replanning.
For a Store source, the stored fragment/object-range mapping acts as the source
binding; only the target requires a live Runtime Binding.

#### Manifest contract

| Manifest | Purpose | Key fields |
| --- | --- | --- |
| Placement | Address-independent planning | tensor identity, global/local shape, logical offset, dtype/layout, layer/expert owner, TP/PP/EP/DP rank |
| Runtime Binding | Binds one live allocation | placement ID, instance, generation, lease, address, size, device, endpoint |
| Store Weight Manifest | Publishes one persistent revision | model/revision, group, stored fragments, object key/offset, coverage, optional checksum |

```text
Placement + Runtime Binding = executable runtime manifest

```

Frameworks export Placement and Runtime Binding. Mooncake neither imports SGLang, vLLM, or slime nor maintains model-specific naming rules. `placement_id` is derived only from canonical logical content, never from an address, endpoint, or lease. A legacy `partition_dim=d` remains compatible as `shard_dims=[d]`.

The Store Weight Manifest never stores a live address. Group provides co-routing and resource coordination; the Weight Manifest defines a complete revision and indexes its payloads. A production multi-writer implementation requires revision compare-and-swap plus `OPEN -> COMMITTED/ABORTED` arbitration. The current prototype uses a single-writer commit protocol.

#### N-D reshard planner

Each fragment is an N-D box in a global logical tensor. The planner directly computes source/target box intersections:

* **TP:** source and target may shard different tensor dimensions.

* **PP:** routing follows the real layer/tensor owner recorded by Placement.

* **EP:** an expert is a leading logical coordinate; independent per-expert allocations remain valid fragments.

* **DP:** G2G selects one complete source replica; Store persists one complete DP replica.

The four axes are planned jointly: ownership and replica selection are followed by one N-D intersection calculation. The planner never creates an intermediate TP, PP, or EP model layout.

Each logical region contains the overlap offset and shape, source and target base offsets, `inner_bytes`, outer-loop counts, and source/target strides. Physical addresses, object keys, leases, and registrations are attached only during binding.

Lowering must be bounded. The prototype caps per-region segments and lazily batches operations with `max_batch_operations`; the production contract must also cap total-plan operations and per-batch bytes. It must not expand a region into one operation per row or element.

#### Weight storage

Store retains source-native TP/PP/EP fragments and their logical coordinates. It does not all-gather or rewrite them into a fixed parallel layout before persisting. One generation-consistent, complete DP replica is selected. Consequently, stored payload is approximately one logical model: TP, PP, and EP alter fragment count and shape, while DP does not multiply storage capacity.

```text
weights////
payload// source-native bytes
manifest immutable revision visibility point
sessions//decision commit or abort

```

The write path is:

1. `prepare_upload()` validates model/revision, tensor descriptors, and full coverage; selects one complete DP replica; and returns a `WeightUploadPlan`.

2. Workers call `batch_put_from()` to upload registered GPU fragments as independent `ObjectDataType.WEIGHT` payloads. Independent experts remain independent objects.

3. The coordinator collects complete receipts, verifies payload existence, and commits the session decision.

4. It publishes an immutable `WeightManifest`. Payloads without a published manifest are not readable as a weight revision.

The read path combines `WeightManifest + Target Placement` into an N-D load plan and uses `get_into_ranges()` to fill target final buffers directly:

```text
TP4/PP2/EP8/DP2 runtime
-> Store: source-native TP4/PP2/EP8 fragments, one DP replica
-> TP8/PP4/EP2/DP4 runtime

```

Plans and manifests contain metadata, not model bytes. Upload reads registered source-GPU ranges; load writes target final buffers. `max_range_bytes` and `max_ranges_per_request` bound staging and submission, so no model-scale temporary buffer is required.

`group_id` groups a revision's payloads, manifest, and session metadata. Group infrastructure provides physical co-routing, lease refresh, hard pinning, and eviction coordination. The manifest provides the logical tensor directory and revision visibility. A Group is neither a cross-object transaction nor a model serving state machine.

The prototype implements upload, commit, abort, load, session decisions, conflict detection, coverage validation, and an optional checksum field. It does not yet enforce checksum production/verification, create-only revision CAS, cross-node `READY/SERVING` state, serving leases, or a tier-residency catalog. Those lifecycle states should extend the Model Weight Management (WML) control plane proposed by RFC #2282:

```text
IMPORTING -> READY -> SERVING -> IDLE -> EVICTING/EVICTED

```

RAM/L3 paths are covered by the prototype. SSD/OSS still require ranged-read, bounded-staging, and recovery validation.

#### Execution paths

**Live G2G.** The source is snapshotted, target final buffers are bound, and TE executes the plan. The target activates only after completion validation and a world barrier. The source may continue read-only inference while snapshotted, but cannot mutate, free, or unregister the weights.

**Runtime-to-Store.** One complete source DP replica uploads source-native fragments to immutable payloads. A `WeightManifest` becomes visible only after all payload checks pass. This path uses `WeightUploadPlan` and does not assume that upload and load share the same physical operation form.

**Store-to-runtime.** A Store Weight Manifest and a target Placement produce a logical load plan. Ranged reads write target final buffers directly. RAM/L3 can serve ranges directly; SSD/OSS may use fixed chunks or bounded staging.

A checkpoint or object file can later be another persistent source, but its model semantics and post-load transforms remain framework-loader work.

### Compatibility and safety

Only an explicit Weight Transfer API enters this planner. Existing KVCache and ordinary TE operations retain their current paths.

Before any data operation, Mooncake validates model/revision, schema, placement digest, tensor layout, generation, lease, registered address ranges, selected source replica, and target coverage. A stale generation, revoked lease, layout mismatch, coverage gap, or out-of-bounds address fails closed.

When DMA completion is uncertain, source ownership, registration, and the snapshot remain retained. Legacy `runtime_v1`, one-dimensional TP, and existing Store/TE paths remain available through compatibility adapters. Unknown capabilities must not silently fall back to an unchecked copy.

### Prototype status and validation

Reference implementation: [vin/heterogeneous-weight-transfer-v3](https://github.com/Bo-Vincent/Mooncake/tree/vin/heterogeneous-weight-transfer-v3)at `a2467783ed01556227bee25956ca47a590ec6384`.

Companion SGLang integration: [vin/runtime-weight-manifest-v3](https://github.com/Bo-Vincent/vin-sglang/tree/vin/runtime-weight-manifest-v3)at `288705ffdd0ea16d38b98676dac3412984cd3583`.

The prototype includes Placement/Binding, v1 compatibility, a TP/PP/EP/DP N-D planner, bounded TE lazy lowering, runtime G2G, Store upload/load, and lease, generation, bounds, and completion protection. Mooncake core contains no model naming rules.

"Supported" below means that a framework has an adapter, live runtime binding, and a validation loop. It does not mean that the generic planner alone supports every model.

| Framework | Prototype status | Boundary |
| --- | --- | --- |
| SGLang | Personal branch; live G2G serving E2E completed; Store contract/library tests exist | Not upstream; production revision barrier, discovery, and traffic activation remain open |
| vLLM | Not integrated | Needs Placement/Binding export, snapshot lifecycle, and target loader |
| slime / verl | Not directly integrated | Can orchestrate an SGLang rollout but still need discovery, barrier, rollback, and activation integration |
| Other frameworks | Protocol-level only | Require Placement/Binding export and conformance validation |

| Model/layout in SGLang | Prototype status | Evidence |
| --- | --- | --- |
| Qwen3 dense / MoE | Adapter implemented | Unit/contract tests; Qwen3 MoE 235B FP8 TP4/PP1/EP4/DP1 -> TP2/PP2/EP2/DP1 serving E2E |
| Qwen3.5 dense / MoE text | Adapter implemented | Qwen3.5-0.8B TP1 -> TP2 serving E2E; MoE cross-repository contract |
| Qwen3.5 multimodal / multimodal MoE | Adapter implemented | Unit/manifest contract only; no VL serving E2E yet |
| Qwen3-Next | Adapter implemented; canonical Triton MoE runner required | Unit/contract and TP4/PP1/EP4/DP1 -> TP2/PP1/EP2/DP1 serving E2E |
| DeepSeek, Llama, Mistral/Mixtral, and others | Not adapted | |

The Mooncake test matrix covers bidirectional TP/PP/EP changes, EP/TP cross-dimension cases, a four-axis combination, Store/G2G sources, and legacy manifest regression. Five dual-host H20 RDMA 4/8 GiB cases passed content verification.  A dim0-to-dim2 case currently produces 2,097,152 `CopyRange` operations. Scatter batching from PR #3000 can reduce submission overhead but cannot remove this region count; compressed strided lowering needs a separate design.

### Alternatives considered

* **Framework-only conversion.** Frameworks may retain their own planners, but Store could not independently reshard-load a revision for a target Placement.

* **Model adapters in Mooncake core.** This duplicates and drifts model rules; semantics must come from a framework manifest.

* **All-gather into one storage layout before upload.** This increases peak memory and network traffic; Store instead preserves source-native fragments.

* **Hard dependency on NCCL M2N.** It may be an optimized executor when its topology and allocation constraints hold, but it is not a correctness dependency.

### Phasing

| Phase | Scope |
| --- | --- |
| 1. Core contract | Stabilize Placement, Binding, N-D regions, legacy-manifest compatibility, and TE fallback |
| 2. Weight storage | Integrate `WeightStore` with Group/WML; add revision CAS, mandatory integrity, READY/serving lifecycle, and SSD/OSS recovery |
| 3. Executor | Integrate scatter batching from PR #3000, design compressed strided lowering, and evaluate an optional M2N executor |
| 4. L3 weight write | Add the Weight Write API, GPU-range direct write, bounded batching/concurrency, receipt recovery, and revision publish/abort semantics |
| 5. Frameworks | Upstream SGLang integration, then vLLM; add discovery, barriers, rollback, and activation integration for slime/verl |
| 6. Models | Keep Qwen3/Qwen3.5/Qwen3-Next as the baseline; add DeepSeek-V2/V3/R1, Llama/Mistral, Mixtral, and multimodal conformance/E2E |

Model-semantic adapters belong in SGLang, vLLM, and other framework repositories. Mooncake owns the schema, planner, Store/TE executors, and cross-framework conformance fixtures; it must not duplicate model rules in core.

### Open questions

1. Which component should persist and publish Store Weight Manifest and the RFC #2282 WML revision metadata, and what should the complete weight lifecycle look like?

2. Should M2N be added as an optional fast path? Which topology and allocation conditions should enable an NCCL M2N executor?

3. Can this architecture later support additional forms of conversion, for example FP16-to-FP8 conversion?

### Related work

* [Unified KVCache and Model Weight Management RFC #2282](https://github.com/kvcache-ai/Mooncake/issues/2282)

* [Scatter transfer batching PR #3000](https://github.com/kvcache-ai/Mooncake/pull/3000)

* [NVIDIA NCCL M2N](https://github.com/NVIDIA/nccl/tree/master/contrib/nccl_m2n)

* SGLang Runtime Weight Manifest companion RFC (issue link to be added)

### Before submitting a new issue...

- [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

Start with the referenced prototype branch at vin/heterogeneous-weight-transfer-v3 and compare it with Mooncake's current implementation; also review RFC #2282 for the proposed lifecycle boundary. The issue is complete only when the manifest, N-D planning, runtime and Store paths, validation, and production lifecycle requirements have an agreed implementation scope.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend-api-design, distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.