[Aionw]

Description

Implements one module-sized Phase 2 increment of #3534 for the shared-segment Python helper.

• move mooncake.shared_segment from mooncake-integration/shared_segment.py to python/mooncake/shared_segment.py as its sole tracked source;
• preserve the public mooncake.shared_segment import and native mooncake.engine boundary;
• keep the native capability probe usable without the optional Torch extra, while reporting a focused error when tensor creation is requested without Torch;
• update direct CMake installation, the transitional legacy wheel-builder input, ownership, and labeling paths;
• add focused coverage for layout creation, mmap/VMM and host-registration capability errors, cleanup after failed completion, and optional Torch behavior;
• add source-tree and isolated installed-wheel checks for the canonical packaged location and removal of the legacy tracked source.

This PR does not rename the native engine extension or add the Phase 3 mooncake.engine facade.

Overlap

No open PR migrates the shared-segment helper, but #4152 (open) adds a hugetlb option to mooncake-integration/shared_segment.py, the file this PR moves; whichever lands second must carry that change to python/mooncake/shared_segment.py and its tests. Since this PR was opened, #3586, #3624, #3626, #3630, #3632, and #3635 have landed on main; this branch includes them through its main syncs, most recently the 2026-09-28 merge of main (d9152a4). #3585 (Reshard) is still open and remains out of scope. This PR does not duplicate any of those migrations.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Reshard (mooncake-reshard)
• 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)
• CI/CD
• Docs
• Other

Type of Change

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

How Has This Been Tested?

Original submission (before merging main on 2026-09-28; superseded by "After merging main" below):

Test commands:

MOONCAKE_TEST_WHEEL=/tmp/mooncake-shared-segment-wheel.VDwWPN/mooncake_transfer_engine-0.3.12.post1-cp310-cp310-linux_x86_64.whl \
  python -m pytest -q python/tests/engine/test_shared_segment.py python/tests/packaging
CMAKE_BUILD_PARALLEL_LEVEL=128 python -m build --wheel \
  --outdir /tmp/mooncake-shared-segment-wheel.VDwWPN
python -m twine check \
  /tmp/mooncake-shared-segment-wheel.VDwWPN/mooncake_transfer_engine-0.3.12.post1-cp310-cp310-linux_x86_64.whl
pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD)

Test results:

• Unit tests pass (13 focused and packaging tests passed)
• Integration tests pass (root PEP 517 wheel built successfully; twine check passed; the wheel was installed with --no-deps into a fresh virtual environment and imported under python -I without repository PYTHONPATH)
• Manual testing done (the isolated installed-wheel smoke test invoked the real native shared-segment capability probe)

The available environment had no CUDA/NPU runtime, so live VMM/HostRegister allocation was not exercised. Native capability branches and cleanup are covered with focused binding fakes, and the native shared-segment binding compiled in the root non-CUDA wheel build.

After merging main (d9152a4, 2026-09-28):

PYTHONPATH=python:mooncake-wheel python -m pytest -q python/tests/engine/test_shared_segment.py python/tests/packaging
pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD)

• 23 passed, 2 skipped; the skips are the installed-wheel smoke tests, which need MOONCAKE_TEST_WHEEL / MOONCAKE_TEST_LEGACY_WHEEL.
• The merge's conflicts in .github/labeler.yml, python/tests/packaging/test_installed_wheel.py, python/tests/packaging/test_project.py were resolved by keeping both sides.
• CI's changed-lines Python check (ruff check plus scripts/ci/check_changed_python_format.py, ruff 0.6.9) passes.
• PR-scoped pre-commit passes every hook except cmake-format, which only rewrites the tent_metrics block in mooncake-integration/CMakeLists.txt inherited from main; main at d9152a4 fails the same way, and this PR's own CMake lines pass.
• The root PEP 517 wheel build, twine check, and installed-wheel smoke runs from the original submission were not repeated after the merge. CI does not run them either (it builds only the legacy scripts/build_wheel.sh wheel and runs scripts/test_installation.sh), so they still need a rerun.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh (no C/C++ source changes; all relevant pre-commit formatting hooks pass)
• I have run pre-commit on the files changed in this PR and all hooks pass (after merging main, except the pre-existing cmake-format issue noted under "After merging main")
• I have updated the documentation (not applicable: the public import and command behavior are unchanged)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed RFC issue [RFC]: Unify Mooncake's Python package, build, and release architecture #3534

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specified below)

OpenAI Codex assisted with the scoped module migration, reference audit, test implementation, packaging validation, and PR preparation. The human submitter remains responsible for reviewing and defending every changed line.

Claude Code (Anthropic) merged main into this branch on 2026-09-28, resolved the conflicts, re-ran the focused tests listed under "After merging main", and updated this description.

[catyans]

LGTM. The wrapper now has one canonical source under python/mooncake, and the PEP 517, legacy-wheel, and CMake paths all install it. Optional Torch handling preserves base-wheel importability; the focused lifecycle and packaging tests pass, and CI is green.

[catyans]

LGTM. Re-reviewed current head 2b7ba4f: the merge from main only resolves the three reported conflicts by retaining both the shared-segment ownership/packaging checks and the new EP module checks. The previously reviewed migration behavior is unchanged. CI is currently queued.