Add file-level model weight cache control plane by xianzhiT · Pull Request #2829 · kvcache-ai/Mooncake

xianzhiT · 2026-07-10T03:11:11Z

Description

Add a file-level weight caching control plane on top of MooncakeStore, so model checkpoints can be imported into the store and served to inference engines (e.g. SGLang) as plain files — without a tensor-level, GPU-coupled save path.

The idea is to treat MooncakeStore as a generic model cache that is decoupled from the inference engine: weights are stored file-level (HuggingFace safetensors chunked with a per-model manifest), and engines load files exactly as they would from disk, only the source changes (disk → MooncakeStore). This control plane provides the management side — import / list / inspect / verify / delete — independently of any serving process.

The matching SGLang connector (data plane) is in sgl-project/sglang#30728. Design discussion / RFC: #2282 (Unified KVCache and Model Weight Management); this PR implements Phase 1 (store-backed weight loading) plus the import/manage control plane.

Module

Python Wheel (mooncake-wheel)
Docs

Type of Change

New feature

How Has This Been Tested?

Test commands:

cd mooncake-wheel
python -m pytest tests/test_weight_model_cache.py tests/test_weight_model_cli.py -v

End-to-end (cross-machine, one storage node + one 8×H20 node):

# import a checkpoint into MooncakeStore
python -m mooncake.weight_store.cli \
    --master-server-addr <master_ip:50051> \
    --metadata-server <metadata_ip:8080> \
    --protocol rdma --rdma-devices <mlx5_x> \
    --local-hostname <host_ip:port> \
    model import \
    --checkpoint-id Qwen3-235B-A22B \
    --model-id Qwen/Qwen3-235B-A22B \
    --revision main \
    --source /path/to/local/weights

# confirm it landed
python -m mooncake.weight_store.cli ... model list
python -m mooncake.weight_store.cli ... model inspect Qwen3-235B-A22B
python -m mooncake.weight_store.cli ... model verify  Qwen3-235B-A22B

Then loaded Qwen3-235B-A22B from the store via the SGLang connector in ~30s
with correct inference output.

Test results:

Unit tests pass
Integration tests pass (if applicable)
Manual testing done (cross-machine import + SGLang load + inference, above)

Checklist

I have performed a self-review of my own code
I have formatted my code using ./scripts/code_format.sh
I have run pre-commit run --all-files and all hooks pass
I have updated the documentation (if applicable)
I have added tests to prove my changes are effective
For changes >500 LOC: I have filed an RFC issue ([RFC] Unified KVCache and Model Weight Management in Mooncake Store #2282)

AI Assistance Disclosure

AI tools were used (specify below)

AI assistance (Claude Code) was used for parts of the implementation, tests, and documentation. All changes have been reviewed and are understood by the human submitter.

Add a file-level weight caching control plane on top of MooncakeStore, so model checkpoints can be imported into the store and served to inference engines (e.g. SGLang) without a tensor-level, GPU-coupled save path. - mooncake/weight_store/: import/list/verify/delete/inspect a HuggingFace safetensors checkpoint as file-level chunks with a per-model manifest (cli.py, model.py, model_keyspace.py, __init__.py). - Store weights as fixed-size chunks with a manifest describing files, sizes, sha256 digests and chunk keys; supports a config file so server addresses need not be passed on every call. - pyproject.toml: expose the `mooncake_weight_store` console entry point. - tests: weight store model cache and CLI tests. - docs: file-level weight cache control-plane design and the SGLang file connector contract.

gemini-code-assist

Code Review

This pull request introduces an experimental file-level model weight cache control plane for Mooncake Store, adding a command-line interface, a model cache client, keyspace helpers, and comprehensive unit tests. The implementation is solid, but several improvements should be made to enhance robustness, correctness, and performance. Specifically, the store setup result check should safely handle None returns, and the size parser should support TB and PB units. Additionally, file materialization should write to a temporary file and atomically rename it to prevent corruption, index updates should be made atomic to avoid race conditions during concurrent imports, redundant store existence checks should be removed to optimize round-trips, and chunk cleanup during import failures should be wrapped in try-except blocks to ensure all orphaned chunks are removed.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-10T03:13:32Z

+    except TypeError:
+        result = store.setup(config)
+
+    if result != 0:


The setup method of MooncakeDistributedStore might return None on success in some environments or wrappers. Checking result != 0 could lead to a false-positive RuntimeError if None is returned. It is safer and more consistent with the rest of the codebase to check result not in (0, None).

Suggested change

if result != 0:

if result not in (0, None):

gemini-code-assist · 2026-07-10T03:13:32Z

+    units = {
+        "KB": 1024,
+        "MB": 1024**2,
+        "GB": 1024**3,
+    }


Adding support for Terabytes (TB) and Petabytes (PB) to the size parser makes it more robust and future-proof, especially when dealing with extremely large models or storage segments.

units = { "KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4, "PB": 1024**5, }

gemini-code-assist · 2026-07-10T03:13:32Z

+WEIGHT_SUFFIXES = {
+    ".bin",
+    ".gguf",
+    ".pt",
+    ".pth",
+    ".safetensors",
+}


Consider adding .ckpt to the list of recognized weight suffixes. It is a very common format used by many models (e.g., Stable Diffusion checkpoints).

Suggested change

WEIGHT_SUFFIXES = {

".bin",

".gguf",

".pt",

".pth",

".safetensors",

}

WEIGHT_SUFFIXES = {

".bin",

".ckpt",

".gguf",

".pt",

".pth",

".safetensors",

}

gemini-code-assist · 2026-07-10T03:13:32Z

+        with output.open("wb") as output_file:
+            for chunk_key in self._record_chunks(record):
+                chunk = self._get(chunk_key)
+                output_file.write(chunk)
+                digest.update(chunk)
+                size += len(chunk)
+        if size != record.size:
+            raise ValueError(
+                f"file size mismatch for {record.path}: "
+                f"expected {record.size}, got {size}"
+            )
+        if digest.hexdigest() != record.sha256:
+            raise ValueError(f"sha256 mismatch for {record.path}")


Writing directly to the target output_path can leave a partially written or corrupted file if the operation is interrupted or fails halfway through. A more robust pattern is to write to a temporary file in the same directory and atomically rename it to the target path once the size and checksum verification succeeds.

Suggested change

with output.open("wb") as output_file:

for chunk_key in self._record_chunks(record):

chunk = self._get(chunk_key)

output_file.write(chunk)

digest.update(chunk)

size += len(chunk)

if size != record.size:

raise ValueError(

f"file size mismatch for {record.path}: "

f"expected {record.size}, got {size}"

)

if digest.hexdigest() != record.sha256:

raise ValueError(f"sha256 mismatch for {record.path}")

temp_output = output.with_suffix(f".{time.time_ns()}.tmp")

try:

with temp_output.open("wb") as output_file:

for chunk_key in self._record_chunks(record):

chunk = self._get(chunk_key)

output_file.write(chunk)

digest.update(chunk)

size += len(chunk)

if size != record.size:

raise ValueError(

f"file size mismatch for {record.path}: "

f"expected {record.size}, got {size}"

)

if digest.hexdigest() != record.sha256:

raise ValueError(f"sha256 mismatch for {record.path}")

temp_output.rename(output)

except BaseException:

if temp_output.exists():

temp_output.unlink()

raise

gemini-code-assist · 2026-07-10T03:13:32Z

+    def _add_to_indexes(self, manifest: ModelFileManifest) -> None:
+        all_models = self._read_index(model_index_key())
+        all_models.add(manifest.checkpoint_id)
+        self._write_index(model_index_key(), all_models)
+
+        per_model = self._read_index(model_id_index_key(manifest.model_id))
+        per_model.add(manifest.checkpoint_id)
+        self._write_index(model_id_index_key(manifest.model_id), per_model)


The read-modify-write cycle for updating index files (weight/index/models and weight/index/model/{sha256}) is not atomic. If multiple clients import checkpoints concurrently, they can overwrite each other's updates, leading to lost index entries. Consider using a proper metadata database/catalog backend or prefix-based key listing if supported by the store to avoid a single global index object.

gemini-code-assist · 2026-07-10T03:13:32Z

+    def _read_index(self, key: str) -> set[str]:
+        if not self._exists(key):
+            return set()
+        value = self.store.get(key)
+        if value is None:
+            return set()
+        if isinstance(value, str):
+            payload = value.encode("utf-8")
+        else:
+            payload = bytes(value)
+        if not payload:
+            return set()
+        return set(json.loads(payload.decode("utf-8")))


Calling _exists before get causes an unnecessary extra round-trip to the store. Since self.store.get(key) already returns None if the key does not exist, we can directly call get and handle the None case, which is much more efficient.

def _read_index(self, key: str) -> set[str]: value = self.store.get(key) if value is None: return set() if isinstance(value, str): payload = value.encode("utf-8") else: payload = bytes(value) if not payload: return set() return set(json.loads(payload.decode("utf-8")))

gemini-code-assist · 2026-07-10T03:13:32Z

+            except BaseException:
+                for chunk_key in chunks:
+                    self._remove(chunk_key, force=True)
+                raise


If self._remove raises an exception during cleanup (e.g., due to a temporary network issue or store error), the loop will be interrupted, leaving the remaining chunks orphaned in the store. Wrapping the individual chunk deletion in a try...except block ensures that we attempt to clean up as many chunks as possible.

except BaseException: for chunk_key in chunks: try: self._remove(chunk_key, force=True) except Exception: pass raise

codecov-commenter · 2026-07-10T04:07:36Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Aionw · 2026-07-13T08:04:36Z

+    return hashlib.sha256(value.encode("utf-8")).hexdigest()
+
+
+def _validate_id(name: str, value: str) -> None:


Is a//b a valid value here?

xianzhiT requested review from ShangmingCai, stmatengss and ykwd as code owners July 10, 2026 03:11

xianzhiT mentioned this pull request Jul 10, 2026

Add Mooncake Store file-level weight connector for SGLang sgl-project/sglang#30728

Open

5 tasks

gemini-code-assist Bot reviewed Jul 10, 2026

View reviewed changes

github-actions Bot added documentation Improvements or additions to documentation run-ci Installation Tests labels Jul 10, 2026

ykwd requested a review from Aionw July 10, 2026 03:32

Aionw reviewed Jul 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add file-level model weight cache control plane#2829

Add file-level model weight cache control plane#2829
xianzhiT wants to merge 1 commit into
kvcache-ai:mainfrom
xianzhiT:feature/file-level-weight-cache-v2

xianzhiT commented Jul 10, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

codecov-commenter commented Jul 10, 2026

Uh oh!

Aionw Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return hashlib.sha256(value.encode("utf-8")).hexdigest()


		def _validate_id(name: str, value: str) -> None:

Uh oh!

Conversation

xianzhiT commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Checklist

AI Assistance Disclosure

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jul 10, 2026

Codecov Report

Uh oh!

Aionw Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

xianzhiT commented Jul 10, 2026 •

edited

Loading