Skip to content

[Weight] Add runtime weight manifest contracts - #3187

Open
Bo-Vincent wants to merge 9 commits into
kvcache-ai:mainfrom
Bo-Vincent:vin/heterogeneous-conversion-manifest
Open

[Weight] Add runtime weight manifest contracts#3187
Bo-Vincent wants to merge 9 commits into
kvcache-ai:mainfrom
Bo-Vincent:vin/heterogeneous-conversion-manifest

Conversation

@Bo-Vincent

@Bo-Vincent Bo-Vincent commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Description

Related RFC: #3111

This PR introduces the framework-neutral runtime weight manifest contract for
Mooncake model-weight reuse and heterogeneous resharding.

It is the first implementation phase of the RFC. This PR defines how frameworks
describe the logical placement and live physical location of model weights,
without coupling Mooncake Core to SGLang, vLLM, PyTorch, model-specific naming
rules, or a particular transfer backend.

What This PR Adds

1. Top-level Model Weight module

Add mooncake-model-weight as a top-level Mooncake module.

The public Python API is exposed through:

mooncake.model_weight

The tracked source of truth remains in:

mooncake-model-weight/python/mooncake/model_weight

During release-wheel construction, this package is temporarily staged into the
combined Mooncake wheel and removed after the build.

2. Logical placement contract

Add PlacementManifest to describe reusable and serializable model-weight
placement:

  • stable tensor identity;
  • global tensor shape, dtype, and item size;
  • layout fingerprint;
  • optional layer and expert identity;
  • N-D logical boxes defined by global_offset and local_shape;
  • TP, PP, EP, and DP ownership coordinates;
  • normalized partition_dim and shard_dims semantics;
  • canonical placement identity and digest.

Logical placement does not contain GPU addresses, endpoints, leases, or runtime
allocation owners.

3. Runtime binding contract

Add RuntimeBindingManifest to describe one live runtime generation:

  • worker and endpoint;
  • device and transferable GPU address;
  • contiguous byte range;
  • lease and generation fence;
  • placement digest binding;
  • optional process-local allocation owner.

A runtime restart can reuse the same logical placement while publishing a new
runtime binding.

4. Validated runtime snapshot

Add RuntimeManifest and framework-neutral inventory adapters for exporting
runtime tensor facts.

The contract validates:

  • integer and ordered-sequence inputs;
  • tensor dimensions and fragment geometry;
  • logical fragment bounds;
  • overlapping logical boxes for the same tensor and parallel rank;
  • contiguous runtime views;
  • storage-offset and address semantics;
  • unsigned 64-bit address-range overflow;
  • duplicate fragment and alias declarations;
  • placement identity and digest consistency;
  • lease-generation consistency;
  • placement-to-binding compatibility.

Mooncake does not infer layer, expert, partition, or layout semantics from
parameter names. These facts must be supplied by the framework adapter.

5. Projection and binding APIs

Add APIs to:

  • project a runtime snapshot into reusable logical placement;
  • project live runtime facts into an ephemeral runtime binding;
  • validate and bind placement and runtime state before planning or transfer.

Why Split Placement and Binding?

Logical placement is stable and reusable, while GPU addresses and allocations
belong to one process lifetime.

Separating them allows future reshard planners and transfer executors to reuse
the same logical placement while still requiring a fresh, generation-fenced
runtime binding before accessing memory.

Wheel Integration

The main Mooncake wheel now stages mooncake.model_weight from the top-level
module during scripts/build_wheel.sh.

The installation test runs the manifest contract tests against the installed
wheel. The staging directory is cleaned after both successful and failed
builds.

Non-Goals

This PR intentionally does not:

  • generate heterogeneous reshard transfer regions;
  • execute Transfer Engine operations;
  • persist weights in Mooncake Store;
  • infer model semantics from tensor names;
  • define discovery, activation, rollback, or control-plane policy;
  • integrate M2N, scatter batching, or another optimized executor.

N-D reshard planning, TE execution, Store integration, and framework adapters
will be submitted as follow-up changes built on this contract.

Module

  • Transfer Engine (mooncake-transfer-engine)
  • Mooncake Store (mooncake-store)
  • Mooncake EP (mooncake-ep)
  • Mooncake PG (mooncake-pg)
  • Integration (mooncake-integration)
  • P2P Store (mooncake-p2p-store)
  • Python Wheel (mooncake-wheel)
  • Common (mooncake-common)
  • Mooncake RL (mooncake-rl)
  • Model Weight (mooncake-model-weight)
  • CI/CD
  • Docs
  • Other

This PR contains the minimal release-wheel staging required to publish the new
module, but feature ownership remains under mooncake-model-weight.

Type of Change

  • Bug fix
  • New feature
  • Refactor
  • Breaking change
  • Documentation update
  • Performance improvement
  • Other

How Has This Been Tested?

Test commands

PYTHONPATH=mooncake-model-weight/python \
python -m pytest -q mooncake-model-weight/tests
python -m ruff check mooncake-model-weight
python -m ruff format --check mooncake-model-weight
git diff --check

Test results

140 passed
Ruff check passed
Ruff format check passed
git diff --check passed
  • Unit tests pass
  • Integration tests pass (not applicable to this contract-only PR)
  • Manual testing done

The tests cover:

  • placement and runtime manifest construction;
  • mapping- and attribute-based runtime inventories;
  • canonical identity and digest generation;
  • placement/runtime projection and binding;
  • N-D fragment bounds and overlap rejection;
  • aliases and tied-weight declarations;
  • address, size, contiguity, and overflow validation;
  • lease and generation consistency;
  • ordered-sequence input normalization;
  • public API exports.

Real framework integration, planner, Transfer Engine, Store, and SGLang serving
E2E coverage will be added in follow-up PRs.

Commit Structure

The change is organized into reviewable commits:

  1. Add the manifest contracts and validation.
  2. Expose the public mooncake.model_weight API.
  3. Add manifest contract tests.
  4. Stage the module in the release wheel.
  5. Add the module overview.
  6. Document framework integration and contract boundaries.
  7. Register module ownership and PR labeling.
  8. Split the manifest implementation and tests by responsibility.

Checklist

  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh (not applicable; no C/C++ code changed)
  • I have run pre-commit run --all-files
  • I have run the relevant Ruff checks
  • I have updated the documentation
  • I have added tests to prove my changes are effective
  • For changes >500 LOC: I have filed an RFC issue ([RFC]: Manifest-driven heterogeneous model weight resharding and storage #3111)

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used

AI assistance was used to inspect the existing Mooncake packaging and manifest
boundaries, implement the contract and tests, organize the commit history, and
run code review and validation. The submitter has reviewed the changes and is
responsible for understanding and maintaining the code.

@zxpdemonio zxpdemonio left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review — [Weight] Add runtime weight manifest contracts

I reviewed the contract module, docs, wheel staging, and tests against RFC #3111. Locally: 139 passed with PYTHONPATH=mooncake-model-weight/python.

Verdict

Solid first-phase contract: placement / binding split is clear, validation is fail-closed in the right places, and the PR stays within the stated non-goals. CI is currently red (spell check), and there are a few contract/docs nits worth fixing before merge. No deep semantic blockers beyond that if the deferred planner checks are intentional.

Blocker / merge gate

  1. Spell Check fails on nd_* identifiers
    Typos treats nd as and in test_model_weight_manifest.py (nd_descriptor, nd_placement, nd_runtime). CI Gate is red because of this.
    Fix by renaming (n_dim_* / multi_dim_*) or adding an allowlist entry in .typos.toml.

Contract review (what looks good)

  • Clear split: PlacementManifest (serializable, address-free) vs RuntimeBindingManifest (ephemeral, generation/lease-fenced).
  • Canonicalization of partition_dim / shard_dims, stable placement_id + digest, and bind-time digest checks.
  • Geometry checks: bounds, non-shard axis integrity, nbytes = prod(local_shape) * itemsize.
  • Same-rank N-D box overlap rejection (sweep + AABB) and runtime address-range overlap with explicit alias opt-in.
  • Contiguity / canonical stride / address_semantics="view" rules are documented and tested.
  • u64 address-range overflow checks; strict JSON schema for placement round-trip.
  • Wheel staging + EXIT cleanup; pkgutil.extend_path enables mooncake.model_weight as a split package.
  • Scope matches RFC phase-1 (no planner / TE / Store).

Issues / questions

  1. README local test recipe is wrong (mooncake-model-weight/README.md)
    It suggests:

    PYTHONPATH=mooncake-wheel python -m pytest ...

    That fails with No module named 'mooncake.model_weight' unless the package has been staged into the wheel. The working recipes are:

    PYTHONPATH=mooncake-model-weight/python python -m pytest -q mooncake-model-weight/tests
    # or, with extend_path:
    PYTHONPATH=mooncake-wheel:mooncake-model-weight/python ...
  2. Cross-rank logical overlap is intentionally unchecked
    Same tensor_id with overlapping boxes on different ParallelRank values is accepted (verified). That is fine if coverage/conflict belongs to the planner, but please call it out explicitly in docs/source/design/model-weight-manifest.md so adapters do not assume the contract enforces a global partition.

  3. Incomplete shard coverage is also accepted
    A TP shard that only covers part of global_shape validates today. Same ask: document as planner-owned, so this is not mistaken for a full placement invariant.

  4. Strict type(x) is int will reject NumPy integers
    Framework inventories often carry numpy.int64 extents. Torch Size is fine (Python int), but raw NumPy metadata will raise. Consider operator.index() / numbers.Integral (still rejecting bool), or document that adapters must coerce.

  5. address rejects 0 (minimum=1)
    Reasonable as a null sentinel for GPU views; worth one sentence in the design doc so CPU/testing harnesses do not treat it as accidental.

  6. Empty placement is allowed
    PlacementManifest(model_id=..., revision=..., tensors=(), fragments=()) succeeds and gets a digest. Harmless, but if empty manifests should be invalid at the control-plane boundary, reject them here or document “empty is a valid sentinel”.

  7. Generated placement_fragment_id is content-hashed, but explicit IDs become part of placement_id
    Adapters that sometimes omit and sometimes supply IDs for the same geometry will see identity churn. Document “IDs must be stable across restarts” for framework exporters.

Non-blocking nits

  • RuntimeManifest can exist without lease_id; only projection to binding requires it — good, already documented.
  • build_wheel.sh installs a second EXIT trap that replaces the first; safe today because the later trap also cleans staging, but a single composed cleanup would be clearer.
  • No public schema version on serialized placement beyond the internal "schema": "weight-placement" used for placement_id. Fine for v0; consider surfacing it in JSON before Store persistence lands.
  • PR checklist marks docs updated (good); spell-check + README fix should land with this PR.

Test coverage

Strong for construction, projection, bind, aliases, overflow, and type strictness. Missing (optional follow-ups):

  • README/PYTHONPATH smoke as documented
  • explicit “cross-rank overlap allowed” / “incomplete coverage allowed” tests if those are permanent contract guarantees
  • installed-wheel import mooncake.model_weight assertion in test_installation.sh (tests help, but a one-liner import check is cheaper signal)

Summary

Approve directionally after fixing the typos CI failure and correcting the README import path. The remaining items are mostly “make deferred invariants explicit” so the next planner PR does not rediscover them.

@Bo-Vincent
Bo-Vincent force-pushed the vin/heterogeneous-conversion-manifest branch from b849312 to c71dc62 Compare July 30, 2026 03:00
@Bo-Vincent
Bo-Vincent force-pushed the vin/heterogeneous-conversion-manifest branch from c71dc62 to 6a19b22 Compare July 30, 2026 03:36
@codecov-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' 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

Aionw commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

manifest.py is already approximately 1,300 lines and covers several distinct responsibilities:

  • public data types and fragments;
  • placement serialization and identity generation;
  • runtime inventory adapters;
  • geometry, overlap, address, and alias validation;
  • placement/runtime projection and binding.

The corresponding test file is also approximately 1,700 lines. Since future planner, Store, and framework-integration work will likely extend this module, it may be worth splitting it along these existing responsibility boundaries to keep the implementation easier to review and maintain.

@Bo-Vincent

Copy link
Copy Markdown
Collaborator Author

manifest.py is already approximately 1,300 lines and covers several distinct responsibilities:

  • public data types and fragments;
  • placement serialization and identity generation;
  • runtime inventory adapters;
  • geometry, overlap, address, and alias validation;
  • placement/runtime projection and binding.

The corresponding test file is also approximately 1,700 lines. Since future planner, Store, and framework-integration work will likely extend this module, it may be worth splitting it along these existing responsibility boundaries to keep the implementation easier to review and maintain.

Thanks, this was a very helpful suggestion. I have split both the implementation and tests along the proposed responsibility boundaries, while keeping the public import surface unchanged.

@Hubert-Zhu Hubert-Zhu mentioned this pull request Aug 1, 2026
22 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI/Build documentation Improvements or additions to documentation Installation run-ci Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants