kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #2868] [RFC]: Unified Artifact I/O for Model Weights, KV Cache, and RL Checkpoints #747

New issue

Closed

opened 2026-10-29 21:19:01 +00:00 by bsakaguchi · 3 comments

bsakaguchi commented

2026-10-29 21:19:01 +00:00

Owner

Originally created by @cursor[bot] on GitHub (Jul 12, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2868

Originally assigned to: @stmatengss on GitHub.

Motivation

Mooncake Store already supports in-cluster tensor storage and high-bandwidth weight distribution (Transfer Engine, P2P Store, checkpoint-engine). However, loading artifacts from external sources is fragmented today:

Source	Status on `main`
Local safetensors	Partial — `load_tensor_from_safetensor` / `save_tensor_to_safetensor` (local paths only, C++ binding)
S3	Not for weights — S3 exists only for master HA snapshot storage (`S3SnapshotObjectStore`)
OSS (Alibaba Cloud)	Not supported
Hugging Face Hub	Not supported — HF appears only in downstream integration docs (vLLM/SGLang model download)

Use cases that need a unified path:

LLM serving cold start — load model weights from S3/OSS/HF into Mooncake Store, then RDMA-distribute to inference nodes.
KV cache persistence — save/load prefill KV tensors to remote object storage for HiCache L3 or cross-session reuse.
RL checkpoint handoff — persist training weights to object storage and load into rollout engines (complements existing P2P hot-path sync).

A prototype exists on branch feature/kv-cache-file-persistence (store_file_io.py + design docs). This RFC proposes merging and extending that work into a stable, documented API.

Goals

Unified read/write API for KV cache, RL checkpoints, and model weights with consistent parameters.
Multi-backend paths via a single save_path / load_path parameter:
- Local: /data/ckpt/, file:///data/ckpt/
- S3: s3://bucket/prefix/
- OSS: oss://bucket/prefix/ (via ossfs / fsspec)
- Hugging Face Hub: hf://org/repo/path/ (via huggingface_hub)
Format support: safetensors (default/recommended), torch (.pt/.pth, weights_only=True by default).
Tensor-parallel sharding: tp_rank / tp_size for multi-file layouts (manifest.json + rank_00000.safetensors, ...).
Backward compatibility with existing load_tensor_from_safetensor / save_tensor_to_safetensor (local C++ path).
Optional P2P hot path after cold load — integrate with P2P Store / checkpoint-engine for cluster-wide distribution.

Non-Goals (initial phase)

Replacing Hugging Face from_pretrained() inside vLLM/SGLang — Mooncake provides storage/transfer, not model class instantiation.
Full WeightManager with automatic model.load_state_dict() — callers remain responsible for applying loaded tensors.
OSS/HF authentication policy beyond storage_options passthrough.

Current State

`main` — local safetensors only

store.load_tensor_from_safetensor("weights/key", "/local/path/model.safetensors")
# → safetensors.torch.load_file → put_tensor → available for RDMA get

Implemented in mooncake-integration/store/store_py.cpp; tested in mooncake-wheel/tests/test_safetensor_functions.py.

Prototype (`feature/kv-cache-file-persistence`)

mooncake-wheel/mooncake/store_file_io.py monkey-patches MooncakeDistributedStore at import:

save_tensor_to_file / load_tensor_from_file — fsspec-backed, torch + safetensors
save_kv_cache_to_file / load_kv_cache_from_file — aliases
Upgraded save_tensor_from_safetensor / load_tensor_from_safetensor — delegate to fsspec layer

Tested: local paths, file:// URIs, .pt and .safetensors. Not tested: s3://, oss://, hf://.

Design docs (prototype branch, commit ca9100d8):

docs/source/design/usage_unified_artifact_io.md
docs/source/design/usage_model_weights.md
docs/source/design/usage_kvcache.md

Proposed Design

1. Unified API (Python)

from typing import Any, Dict, Literal, Optional

ArtifactKind = Literal["kv_cache", "rl_checkpoint", "model_weights"]

class MooncakeArtifactIO:
    """Facade over MooncakeDistributedStore; optional separate class or mixin."""

    def save_artifact(
        self,
        kind: ArtifactKind,
        name: str,
        obj: Any,                    # dict[str, Tensor] | Tensor | bytes manifest
        save_path: str,
        *,
        format: str = "auto",        # auto | safetensors | torch
        storage_options: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        tp_rank: Optional[int] = None,
        tp_size: int = 1,
        use_p2p: str = "auto",       # auto | true | false
    ) -> Dict[str, Any]:
        ...

    def load_artifact(
        self,
        kind: ArtifactKind,
        name: str,
        load_path: str,
        *,
        format: str = "auto",
        storage_options: Optional[Dict[str, Any]] = None,
        map_location: str = "cpu",
        weights_only: bool = True,     # torch format only
        tp_rank: Optional[int] = None,
        tp_size: int = 1,
        use_p2p: str = "auto",
    ) -> Any:
        ...

2. Path resolver

Scheme	Backend	Notes
(none) / `file://`	`fsspec` local	Default
`s3://`	`fsspec` + `s3fs`	Reuse patterns from `S3Helper`; `storage_options` for credentials
`oss://`	`fsspec` + `ossfs`	Alibaba OSS; document required extra dep
`hf://`	`huggingface_hub`	Download/upload; `storage_options`: `token`, `revision`, `repo_type`

format="auto": if path ends with .safetensors → safetensors; else default safetensors for directories, torch for .pt/.pth.

3. TP shard layout (directory paths)

save_path/
  manifest.json          # { "tp_size": 4, "format": "safetensors", "keys": [...] }
  rank_00000.safetensors
  rank_00001.safetensors
  ...

load_artifact with tp_rank reads the matching shard and put_tensor into Store (optionally with TensorParallelism / existing unified tensor IO).

4. End-to-end workflow

flowchart TD
    A[External source<br/>S3 / OSS / HF / local] -->|load_artifact| B[Path resolver]
    B --> C{fsspec / huggingface_hub}
    C --> D[Decode safetensors or torch]
    D --> E[put_tensor / batch_put into Mooncake Store]
    E --> F{use_p2p?}
    F -->|yes| G[P2P Store Register / GetReplica]
    F -->|no| H[get_tensor / RDMA to peers]
    G --> H
    H --> I[Inference / training consumer]

5. Compatibility layer

Keep existing methods; delegate to unified API:

Legacy	Unified
`load_tensor_from_safetensor`	`load_artifact` or `load_tensor_from_file(format="safetensors")`
`save_tensor_to_safetensor`	`save_artifact` or `save_tensor_to_file`
`load_kv_cache_from_file`	`load_artifact(kind="kv_cache")`
(future) `WeightManager.load_checkpoint`	`load_artifact(kind="model_weights")`

C++ load_tensor_from_safetensor remains for local-only, zero-extra-dep use; Python fsspec layer preferred for remote paths.

6. Dependencies

# pyproject.toml
dependencies = ["fsspec", ...]

[project.optional-dependencies]
file-io = ["safetensors", "torch"]
s3 = ["s3fs"]
oss = ["ossfs"]
hf = ["huggingface_hub"]
artifact-io = ["safetensors", "torch", "s3fs", "huggingface_hub"]

Import-time patch (patch_store_file_io_support) stays best-effort: missing optional deps → clear ValueError with install hint.

Implementation Plan

Phase	Scope	Est. size
P0	Merge `store_file_io.py`; document local + `file://` + `s3://`; CI unit tests with moto/minio	~400 LOC
P1	`hf://` adapter; `save_artifact` / `load_artifact` facade; TP manifest layout	~300 LOC
P2	`oss://` adapter; integration tests; user docs under `docs/source/python-api-reference/`	~200 LOC
P3	P2P auto-register after load; benchmark vs checkpoint-engine cold path	~200 LOC

Testing Strategy

Unit: format auto-detection, path resolver, tensor entry picking, weights_only, map_location
Integration: local round-trip (existing); MinIO for s3://; mock HF hub for hf://
TP: save tp_size=4, load per-rank, byte-compare against golden
Regression: existing test_safetensor_functions.py unchanged on main behavior

Open Questions

Should hf:// use hf_hub_download (read-only) only in P0, deferring upload to P1?
Reuse S3Helper (C++/AWS SDK) for Store master paths vs pure Python s3fs for artifact I/O — acceptable duplication?
Should loaded weights default to ObjectDataType::WEIGHT + with_hard_pin when kind="model_weights"?
Is oss:// required for Alibaba production, or is S3-compatible endpoint sufficient initially?

References

Prototype branch: feature/kv-cache-file-persistence on stmatengss/mooncake
Existing local API: mooncake-integration/store/store_py.cpp (load_tensor_from_safetensor)
P2P weight sync: checkpoint-engine, SGLang P2P weight transfer blog
CONTRIBUTING.md: RFC expected for architectural changes >500 LOC

Request for Feedback

Is the unified save_artifact / load_artifact API the right surface, or should we only extend load_tensor_from_file?
Priority order: S3 vs HF vs OSS for your deployments?
Any hard requirements on zero-copy / mmap for large safetensors from object storage?

Originally created by @cursor[bot] on GitHub (Jul 12, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2868 Originally assigned to: @stmatengss on GitHub. ## Motivation Mooncake Store already supports in-cluster tensor storage and high-bandwidth weight distribution (Transfer Engine, P2P Store, checkpoint-engine). However, **loading artifacts from external sources** is fragmented today: | Source | Status on `main` | |--------|------------------| | **Local safetensors** | Partial — `load_tensor_from_safetensor` / `save_tensor_to_safetensor` (local paths only, C++ binding) | | **S3** | Not for weights — S3 exists only for master HA snapshot storage (`S3SnapshotObjectStore`) | | **OSS (Alibaba Cloud)** | Not supported | | **Hugging Face Hub** | Not supported — HF appears only in downstream integration docs (vLLM/SGLang model download) | Use cases that need a unified path: 1. **LLM serving cold start** — load model weights from S3/OSS/HF into Mooncake Store, then RDMA-distribute to inference nodes. 2. **KV cache persistence** — save/load prefill KV tensors to remote object storage for HiCache L3 or cross-session reuse. 3. **RL checkpoint handoff** — persist training weights to object storage and load into rollout engines (complements existing P2P hot-path sync). A prototype exists on branch `feature/kv-cache-file-persistence` (`store_file_io.py` + design docs). This RFC proposes merging and extending that work into a stable, documented API. --- ## Goals 1. **Unified read/write API** for KV cache, RL checkpoints, and model weights with consistent parameters. 2. **Multi-backend paths** via a single `save_path` / `load_path` parameter: - Local: `/data/ckpt/`, `file:///data/ckpt/` - S3: `s3://bucket/prefix/` - OSS: `oss://bucket/prefix/` (via `ossfs` / fsspec) - Hugging Face Hub: `hf://org/repo/path/` (via `huggingface_hub`) 3. **Format support**: `safetensors` (default/recommended), `torch` (`.pt`/`.pth`, `weights_only=True` by default). 4. **Tensor-parallel sharding**: `tp_rank` / `tp_size` for multi-file layouts (`manifest.json` + `rank_00000.safetensors`, ...). 5. **Backward compatibility** with existing `load_tensor_from_safetensor` / `save_tensor_to_safetensor` (local C++ path). 6. **Optional P2P hot path** after cold load — integrate with P2P Store / checkpoint-engine for cluster-wide distribution. ## Non-Goals (initial phase) - Replacing Hugging Face `from_pretrained()` inside vLLM/SGLang — Mooncake provides storage/transfer, not model class instantiation. - Full `WeightManager` with automatic `model.load_state_dict()` — callers remain responsible for applying loaded tensors. - OSS/HF authentication policy beyond `storage_options` passthrough. --- ## Current State ### `main` — local safetensors only ```python store.load_tensor_from_safetensor("weights/key", "/local/path/model.safetensors") # → safetensors.torch.load_file → put_tensor → available for RDMA get ``` Implemented in `mooncake-integration/store/store_py.cpp`; tested in `mooncake-wheel/tests/test_safetensor_functions.py`. ### Prototype (`feature/kv-cache-file-persistence`) `mooncake-wheel/mooncake/store_file_io.py` monkey-patches `MooncakeDistributedStore` at import: - `save_tensor_to_file` / `load_tensor_from_file` — fsspec-backed, `torch` + `safetensors` - `save_kv_cache_to_file` / `load_kv_cache_from_file` — aliases - Upgraded `save_tensor_from_safetensor` / `load_tensor_from_safetensor` — delegate to fsspec layer Tested: local paths, `file://` URIs, `.pt` and `.safetensors`. **Not tested**: `s3://`, `oss://`, `hf://`. Design docs (prototype branch, commit `ca9100d8`): - `docs/source/design/usage_unified_artifact_io.md` - `docs/source/design/usage_model_weights.md` - `docs/source/design/usage_kvcache.md` --- ## Proposed Design ### 1. Unified API (Python) ```python from typing import Any, Dict, Literal, Optional ArtifactKind = Literal["kv_cache", "rl_checkpoint", "model_weights"] class MooncakeArtifactIO: """Facade over MooncakeDistributedStore; optional separate class or mixin.""" def save_artifact( self, kind: ArtifactKind, name: str, obj: Any, # dict[str, Tensor] | Tensor | bytes manifest save_path: str, *, format: str = "auto", # auto | safetensors | torch storage_options: Optional[Dict[str, Any]] = None, metadata: Optional[Dict[str, Any]] = None, tp_rank: Optional[int] = None, tp_size: int = 1, use_p2p: str = "auto", # auto | true | false ) -> Dict[str, Any]: ... def load_artifact( self, kind: ArtifactKind, name: str, load_path: str, *, format: str = "auto", storage_options: Optional[Dict[str, Any]] = None, map_location: str = "cpu", weights_only: bool = True, # torch format only tp_rank: Optional[int] = None, tp_size: int = 1, use_p2p: str = "auto", ) -> Any: ... ``` ### 2. Path resolver | Scheme | Backend | Notes | |--------|---------|-------| | (none) / `file://` | `fsspec` local | Default | | `s3://` | `fsspec` + `s3fs` | Reuse patterns from `S3Helper`; `storage_options` for credentials | | `oss://` | `fsspec` + `ossfs` | Alibaba OSS; document required extra dep | | `hf://` | `huggingface_hub` | Download/upload; `storage_options`: `token`, `revision`, `repo_type` | `format="auto"`: if path ends with `.safetensors` → safetensors; else default safetensors for directories, torch for `.pt`/`.pth`. ### 3. TP shard layout (directory paths) ``` save_path/ manifest.json # { "tp_size": 4, "format": "safetensors", "keys": [...] } rank_00000.safetensors rank_00001.safetensors ... ``` `load_artifact` with `tp_rank` reads the matching shard and `put_tensor` into Store (optionally with `TensorParallelism` / existing unified tensor IO). ### 4. End-to-end workflow ```mermaid flowchart TD A[External source<br/>S3 / OSS / HF / local] -->|load_artifact| B[Path resolver] B --> C{fsspec / huggingface_hub} C --> D[Decode safetensors or torch] D --> E[put_tensor / batch_put into Mooncake Store] E --> F{use_p2p?} F -->|yes| G[P2P Store Register / GetReplica] F -->|no| H[get_tensor / RDMA to peers] G --> H H --> I[Inference / training consumer] ``` ### 5. Compatibility layer Keep existing methods; delegate to unified API: | Legacy | Unified | |--------|---------| | `load_tensor_from_safetensor` | `load_artifact` or `load_tensor_from_file(format="safetensors")` | | `save_tensor_to_safetensor` | `save_artifact` or `save_tensor_to_file` | | `load_kv_cache_from_file` | `load_artifact(kind="kv_cache")` | | (future) `WeightManager.load_checkpoint` | `load_artifact(kind="model_weights")` | C++ `load_tensor_from_safetensor` remains for local-only, zero-extra-dep use; Python fsspec layer preferred for remote paths. ### 6. Dependencies ```toml # pyproject.toml dependencies = ["fsspec", ...] [project.optional-dependencies] file-io = ["safetensors", "torch"] s3 = ["s3fs"] oss = ["ossfs"] hf = ["huggingface_hub"] artifact-io = ["safetensors", "torch", "s3fs", "huggingface_hub"] ``` Import-time patch (`patch_store_file_io_support`) stays best-effort: missing optional deps → clear `ValueError` with install hint. --- ## Implementation Plan | Phase | Scope | Est. size | |-------|-------|-----------| | **P0** | Merge `store_file_io.py`; document local + `file://` + `s3://`; CI unit tests with moto/minio | ~400 LOC | | **P1** | `hf://` adapter; `save_artifact` / `load_artifact` facade; TP manifest layout | ~300 LOC | | **P2** | `oss://` adapter; integration tests; user docs under `docs/source/python-api-reference/` | ~200 LOC | | **P3** | P2P auto-register after load; benchmark vs checkpoint-engine cold path | ~200 LOC | --- ## Testing Strategy - **Unit**: format auto-detection, path resolver, tensor entry picking, `weights_only`, `map_location` - **Integration**: local round-trip (existing); MinIO for `s3://`; mock HF hub for `hf://` - **TP**: save `tp_size=4`, load per-rank, byte-compare against golden - **Regression**: existing `test_safetensor_functions.py` unchanged on `main` behavior --- ## Open Questions 1. Should `hf://` use `hf_hub_download` (read-only) only in P0, deferring `upload` to P1? 2. Reuse `S3Helper` (C++/AWS SDK) for Store master paths vs pure Python `s3fs` for artifact I/O — acceptable duplication? 3. Should loaded weights default to `ObjectDataType::WEIGHT` + `with_hard_pin` when `kind="model_weights"`? 4. Is `oss://` required for Alibaba production, or is S3-compatible endpoint sufficient initially? --- ## References - Prototype branch: `feature/kv-cache-file-persistence` on `stmatengss/mooncake` - Existing local API: `mooncake-integration/store/store_py.cpp` (`load_tensor_from_safetensor`) - P2P weight sync: [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/), [SGLang P2P weight transfer blog](https://lmsys.org/blog/2026-04-29-p2p-update/) - CONTRIBUTING.md: RFC expected for architectural changes >500 LOC --- ## Request for Feedback - Is the unified `save_artifact` / `load_artifact` API the right surface, or should we only extend `load_tensor_from_file`? - Priority order: S3 vs HF vs OSS for your deployments? - Any hard requirements on zero-copy / mmap for large safetensors from object storage?

bsakaguchi

2026-10-29 21:19:01 +00:00

closed this issue
added the
stale

auto-closed
labels

bsakaguchi commented

2026-10-29 21:19:02 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Jul 12, 2026):

Thanks for opening this issue, @cursor[bot]!

Field	Value
Issue	#2868
GitHub user ID	`206951365`
Reporter	@cursor[bot]

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Jul 12, 2026): Thanks for opening this issue, @cursor[bot]! | Field | Value | |-------|-------| | **Issue** | #2868 | | **GitHub user ID** | `206951365` | | **Reporter** | @cursor[bot] | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:19:02 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 11, 2026):

This issue has had no activity for 90 days and will be closed in 7 days if there is no further activity. Please comment or react if it should stay open.

@github-actions[bot] commented on GitHub (Oct 11, 2026): This issue has had no activity for 90 days and will be closed in 7 days if there is no further activity. Please comment or react if it should stay open.

bsakaguchi commented

2026-10-29 21:19:02 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 19, 2026):

Closing due to 3 months of inactivity. If this is still relevant, please comment and we can reopen.

@github-actions[bot] commented on GitHub (Oct 19, 2026): Closing due to 3 months of inactivity. If this is still relevant, please comment and we can reopen.

bsakaguchi referenced this issue

2026-10-29 21:30:23 +00:00

[PR #747] [MERGED] [Docs] Minor Update: Explain Return Value of batch_put_from #1792