Conversation
…s in one wheel Adds a pinned, source-built Mooncake ROCm SDK and wheel, and makes the existing LMCache wheel carry BOTH native extensions: the HIP `lmcache.c_ops` and the host-only `lmcache.lmcache_mooncake`. Packaging only: no runtime service, compose option, default or documented behavior changes. Why a compiler change was needed LMCache collects the common C++ extensions, the ROCm c_ops extension and the optional storage-backend extensions into a single `build_ext`, so the whole wheel was built under one global `CXX=hipcc`. That also handed hipcc the host-only Mooncake sources (ASIO plus coroutine headers), which do not build that way. Splitting the wheel, repacking it or forcing host mode with extra compile flags were all rejected: they either ship two artifacts or hide the mismatch instead of fixing it. patches/lmcache/15-lmcache-per-extension-cxx.patch makes the compiler a per-extension property. The base compiler is g++; only `lmcache.c_ops` is tagged with a private attribute naming its compiler (hipcc, overridable via `LMCACHE_ROCM_CXX`). While that one extension builds, the `CXX` environment variable (read by torch when it writes the ninja file) and the command head of every compiler/linker executable array present on the distutils compiler object are switched together and restored in a `finally` block, including on failure. Both are process global, so a parallel `build_ext` is rejected while a tagged extension exists; `MAX_JOBS` and ninja parallelism inside one extension are unaffected. The patch carries unit tests for the switch, the restoration path and the parallel rejection. Mooncake An independent builder stage builds Mooncake at the exact pinned commit with the ROCm recipe from that revision's own CI (USE_HIP, no CUDA, no EP, HTTP and etcd metadata, no unit tests, hipify-perl on PATH), plus BUILD_SHARED_LIBS and a private install prefix so a third-party consumer can link the store library. The SDK is staged explicitly, because `cmake --install` does not carry the store headers, the vendored allocator tree, the generated headers or yalantinglibs. The official HIP wheel is built by upstream's own script with a committed pip constraints file, so the build frontend is pinned rather than whatever the index serves that day. Only the SDK and the wheel cross into the main stage, so the builder's system-wide side effects stay behind. NIXL is unchanged, and that is enforced rather than assumed NIXL enables its Mooncake backend by auto-detection: its plugin list is opt-out and the plugin turns on when the compiler resolves -ltransfer_engine. The meson arguments are untouched, the SDK is staged off the default link path, and nothing exports LIBRARY_PATH, CPATH or PKG_CONFIG_PATH for it. A probe before the NIXL build reproduces meson's own link test and fails the build if it ever starts resolving; a check after the build fails if a Mooncake plugin appears anyway. Gates (fail closed, in the image, re-runnable against a built image via /usr/local/bin/aic-mooncake-gates.sh) - wheel archive membership, checked before anything is installed - fresh forced install of that wheel - imports of lmcache.c_ops, lmcache.lmcache_mooncake, mooncake.engine, mooncake.store - the MP adapters resolve to lmcache modules, not to a bundled fallback - the packaged master binary reports its version - ldd on both extensions, both mooncake modules and both store libraries, with no missing libraries - pip check, failing on findings that name these packages Locally verified: the patch queue applies in lexical order against the pinned LMCache release, the new unit tests pass, byte compilation is clean, and the shell and Dockerfile linters report no new findings. The image build itself needs a ROCm host and has not been run here.
…lusion Follow-up to the previous commit, tightening the parts that could pass while the property they check is false. Mooncake builder stage - Install the exact Go toolchain the pinned revision expects from go.dev and verify its published checksum before the dependency helper runs. The helper otherwise pulls Go from several mirrors with no verification. - Pin the remaining build switches instead of relying on upstream defaults: examples ON (upstream's own wheel script packages the transfer engine benchmark binary from that tree), benchmarks OFF, Rust store bindings OFF, compiler cache OFF. NIXL - Disable the preview Mooncake plugin explicitly in the meson arguments (`disable_plugins` plus `disable_mooncake_backend`) rather than relying on the SDK staying off the default link path. This image owns Mooncake through LMCache, so the plugin must never build, whatever else lands in the image. - The post-build check stays and is narrowed to plugin objects. Gates - Check BOTH wheel archives before anything is installed, and assert the Mooncake wheel actually carries the engine, the store and the master binary. - Assert the installed LMCache version alongside the extension imports. - Resolve the MP adapters through the real connector module and require the exact expected adapter module, so a bundled fallback cannot satisfy it. - Extend the shared-object check to the common, asio and etcd-wrapper libraries and to the packaged master binary. - Make `pip check` strict: any inconsistency fails the build. - Fail the build unless exactly one wheel of each kind is present, so a stale wheel in the directory cannot silently change what is gated or installed. Linters stay at their pre-existing findings; the new wheel-count assertion carries an explicit ignore for the array warning the rest of this file already triggers.
A filtered clone does not guarantee that an arbitrary full commit SHA is available locally, so checkout interpreted the pin as a path and failed. Initialize the repository, fetch the exact pinned commit at depth one, detach at FETCH_HEAD, and retain the existing full-SHA verification.
The PR merge SHA remains visible through GitHub's API but is no longer advertised by the Git repository. Pin the final reviewed PR head, fetch it through the durable pull-request ref with blob filtering, and continue to verify the exact full commit SHA before building.
Install the constrained Python build frontend with --ignore-installed before invoking Mooncake's upstream wheel script. This avoids an attempted uninstall of Ubuntu's record-less Debian pip package while keeping the upstream ROCm wheel flow unchanged.
Rewrite blank context lines in the embedded LMCache patch as equivalent remove/add pairs so the ROCm-AIC repository passes git diff --check. The applied LMCache source is unchanged, and the patch still applies cleanly to v0.5.3.
Contributor
Author
|
@sbates130272 I couldn’t complete the runtime image build today because the shared host ran low on root space; could you run /run-ci-fast or the normal dist-build and share the first failing step? |
1 task
2imi9
force-pushed
the
feature/mooncake-rocm-packaging
branch
from
August 15, 2026 05:50
9578c94 to
a088f73
Compare
2imi9
marked this pull request as ready for review
August 15, 2026 07:04
2imi9
requested review from
amd-ivaganev,
gaoikawa,
sbates130272 and
vcave
as code owners
August 15, 2026 07:04
Contributor
Author
|
@john00003 I just finished checking this PR. It’s ready for review. |
Collaborator
|
/run-ci-fast |
There was a problem hiding this comment.
Pull request overview
Adds source-built Mooncake ROCm packaging as an optional LMCache backend while keeping the NIXL integration disabled.
Changes:
- Builds and validates Mooncake SDK and ROCm wheels.
- Adds per-extension LMCache compiler selection and wheel-relative linking.
- Publishes Mooncake through nightly and release workflows.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
README.md |
Updates LMCache patch count. |
patches/lmcache/15-lmcache-per-extension-cxx.patch |
Adds per-extension compiler selection. |
patches/lmcache/16-lmcache-mooncake-wheel-rpath.patch |
Adds configurable runtime library paths. |
docs/PIP_WHEELS.md |
Documents Mooncake wheel installation. |
docker/scripts/mooncake-gates.sh |
Adds packaging and runtime validation gates. |
docker/scripts/build-nixl.sh |
Disables NIXL’s Mooncake integration. |
docker/Dockerfile |
Builds, packages, installs, and validates Mooncake. |
docker/docker-compose.yml |
Provides a tracked fallback wheel context. |
docker/constraints/mooncake-wheel-build.txt |
Pins Mooncake build tooling. |
.github/workflows/aic-release.yml |
Verifies Mooncake release artifacts. |
.github/workflows/aic-nightly-wheels.yml |
Publishes nightly Mooncake wheels. |
.github/workflows/aic-export.yml |
Extends release-note validation. |
.github/scripts/render-release-notes.sh |
Adds Mooncake installation metadata. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| NIXL_PLUGIN_DIR=/opt/nixl/lib/x86_64-linux-gnu/plugins \ | ||
| PYTHONPATH=/opt/rocm/share/amd_smi:/opt/nixl/lib/python3/dist-packages \ | ||
| LD_LIBRARY_PATH=/opt/nixl/lib/x86_64-linux-gnu:/opt/nixl/lib:/opt/rocnixl-ucx/lib:/opt/rocm/lib \ | ||
| LD_LIBRARY_PATH=/opt/nixl/lib/x86_64-linux-gnu:/opt/nixl/lib:/opt/rocnixl-ucx/lib:/opt/rocm/lib:/usr/local/lib/python3.12/dist-packages/mooncake:/usr/local/lib/python3.12/dist-packages/mooncake_transfer_engine_rocm.libs \ |
| # upstream ROCm/HIP wheel + CI recipe). Bump only with evidence: the build | ||
| # below follows that revision's own ROCm recipe. | ||
| ARG MOONCAKE_GIT_URL=https://github.com/kvcache-ai/Mooncake.git | ||
| ARG MOONCAKE_REF=e164c96b3c184b34f7979c47c176c2e51a926f66 |
|
✅ Fast hardware CI passed for ff3da34 (dist-build-fast + smoke-test-fast + tiny-test). This is a single-arch dev check — run |
CI pins shellcheck 0.9.0 and hadolint 2.12.0, both older than the local
toolchain, so three checks failed on green local runs.
shellcheck 0.9.0 reports SC2317 on the body of the trap-only cleanup
function rather than the SC2329 that 0.10+ reports on its definition, so
disable both codes. Replace the Mooncake SHA check's && test ... || {}
with an if block: it silences hadolint SC2015 and stops an earlier git
failure from printing the checkout-mismatch error. Add Mooncake, CUDA and
PyPI's to the aspell wordlist for the new docs/PIP_WHEELS.md text.
Verified against the CI-pinned versions, not the local ones: shellcheck
0.9.0 and 0.11.0, hadolint 2.12.0, and pyspelling 2.12.1 all clean.
The image-wide LD_LIBRARY_PATH named the two private directories of the Mooncake wheel. That path precedes the LMCache extension's wheel-relative RUNPATH, so an LMCache installed into a virtualenv or another target could bind to the image copy instead of its own companion wheel and mix two glog builds. The clean-install gate already proves the exported pair loads with only Torch and ROCm exposed, so the entries were redundant. Drop them and keep both libraries private to their RUNPATH.
Contributor
Author
|
Lint failures were version skew. CI pins shellcheck 0.9.0 and hadolint 2.12.0, mine are newer, so I checked against the pinned ones this time. Merged main too, which took care of the conflict. Runs are waiting on workflow approval before they'll go, then |
Collaborator
|
@2imi9 can we wait till the PR gets merged and there is an official release tag of mooncake with ROCm support please? |
Contributor
Author
Yes sure |
The pin rode refs/pull/3184/head, the head of the pull request that added upstream's ROCm/HIP wheel and CI recipe. That ref is rewritable, deletable and not reachable from any release, so the image identity depended on a revision upstream never promised to keep. kvcache-ai/Mooncake#3184 is merged and released, so pin the tag that carries it. v0.3.13.post1 rather than v0.3.13: the .post1 release fixes published non-CUDA wheels linking libcudart because Store treated a merely present CUDA toolkit as permission to enable CUDA staging, and the same change gates HIP staging on the project option instead. The ROCm build takes that path. The revision's own ROCm recipe (.github/workflows/ci_rocm.yml) is unchanged, so the build flags here still follow it. One upstream change does reach the SDK staging: dependencies.sh no longer builds and installs yalantinglibs into /usr/local, because CMake now fetches it into the build tree with EXCLUDE_FROM_ALL (mooncake-common/FindYLT.cmake), and the store client compiles against those headers. Stage them from the build tree as well, and leave the existing include/ylt assertion to prove one of the two paths landed.
Codex review caught both: the stage header and the dependencies.sh comment still described yalantinglibs as built and installed into /usr/local, which stopped being true at v0.3.13.post1.
main moved LMCache to v0.5.4, which renamed the GPU extension from lmcache.c_ops to lmcache.cuda_ops, moved its sources under csrc/cuda and changed the ROCm profile's hipify helper signature. Re-derive 15-lmcache-per-extension-cxx.patch against that layout (patch 16 still applies unchanged; the 24 compiler-switch tests pass) and follow the rename in the wheel and runtime gates. The gates used to hard-code the LMCache version. The build now records LMCACHE_REF in /etc/aic/lmcache-ref and the gate reads it (an explicit LMCACHE_REF still wins), so the standalone re-run inside a built image checks the pin the image was built with instead of a literal in the script. The export workflow's release-notes fixture follows main's vLLM v0.27.1 pin.
…d tree At the v0.3.13.post1 pin, Mooncake fetches yalantinglibs with CMake FetchContent instead of installing it into /usr/local. The old install flattened include/ylt/thirdparty and include/ylt/standalone into the include root, and Mooncake's store headers rely on that layout: real_client.h includes <async_simple/coro/Lazy.h> directly. Mirror the flattening when staging the SDK from the fetched tree, and assert the async_simple header so an incomplete SDK fails in this stage rather than in the LMCache compile.
Upstream's scripts/build_wheel.sh now installs patchelf from PyPI next to auditwheel (auditwheel >= 6.8 needs patchelf >= 0.14.5). The build frontend is pinned through PIP_CONSTRAINT, so add patchelf to the constraints file instead of letting it float.
Contributor
Author
|
@sbates130272 Mooncake now has a release with ROCm support, |
Contributor
Author
|
Two follow-up changes came with the rebase. Patch 15 was redone for LMCache 0.5.4, which renamed the GPU extension to |
Contributor
Author
|
All GitHub checks pass on the new head. The image build and GPU gates still need |
Contributor
Author
|
Side note: nightly wheels has failed on main since Aug 12 because the |
Collaborator
|
/run-ci |
The SDK's shared libraries need each other by soname (libmooncake_store -> libtransfer_engine -> libmooncake_common -> libasio, plus libetcd_wrapper), but CMake installs them with no RUNPATH of their own, and the build-tree copies that overwrite them point at /src/Mooncake/build, which no later stage has. The runtime ldd gate only ever passed because the wheel's site-packages directories used to sit on LD_LIBRARY_PATH; a6f38a8 removed them, correctly, and the first /run-ci on this head (run 34011579638) failed with every sibling reported "not found". Set RUNPATH=$ORIGIN:/opt/rocm/lib on every staged library, the treatment upstream's build_wheel.sh gives the wheel copies plus the ROCm directory CMake already chose for the transfer engine, install patchelf explicitly in the SDK stage, and assert at staging time that each library resolves with nothing else on the loader path.
2imi9
force-pushed
the
feature/mooncake-rocm-packaging
branch
from
September 6, 2026 14:47
96ce5ba to
867fad2
Compare
Contributor
Author
|
@sbates130272 thanks for running it. The build itself was fine; it failed in the runtime |
Contributor
Author
|
The checks on this push need workflow approval. Could you approve them and run |
Collaborator
|
/run-ci |
|
✅ Hardware CI passed for 867fad2 (dist-build + smoke-test + tiny-test). |
…cm-packaging Conflict in docker/Dockerfile: main now carries the empty default `pip-wheels` stage from ROCm#151, which this branch had introduced in the same place. Keep main's copy verbatim and drop this branch's; the Mooncake build arguments stay ahead of the first FROM so they remain global. No image instruction changes.
Contributor
Author
|
Merged main to clear the conflict from #151; nothing in the image changed, but the checks need approval again. |
Collaborator
|
/run-ci |
Contributor
Author
|
@sbates130272 Looks like a node hiccup rather than the PR: the smoke job died about ten seconds in, well before the image could even finish loading, and the merge did not touch anything the smoke test runs (same image passed all three stages on 867fad2 earlier today). Mind kicking |
|
✅ Hardware CI passed for 32c652c (dist-build + smoke-test + tiny-test). |
…cm-packaging Conflict in docker/scripts/aic-image-tag.sh: ROCm#153 bumped hsa-snoop to v1.1.0 and updated the example tag in the header comment, on the line where this branch inserts the mooncake segment. Keep this branch's example and take main's hsasnoop1.1.0 suffix. Comment only: the script body is unchanged, and against the merged Dockerfile it emits exactly the example, 0.1.0-rocm7.14.0-vllm0.27.1-lmcache0.5.4-mooncake7197358-nixl1.3.2-hsasnoop1.1.0. The Dockerfile, README and monitoring hunks merged automatically as main has them.
Collaborator
|
@2imi9 i am going to hold off on adding mooncake for now. We need to decide what frameworks are priority. I'll work on that this week and we can go from there. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Package Mooncake upstream ROCm support as a prerequisite for an optional KV transfer backend, without changing AIC runtime defaults.
Technical Details
Pin Mooncake to the
v0.3.13.post1release tag (commit7197358), the first release with the ROCm wheel recipe from kvcache-ai/Mooncake#3184. Source-build its SDK and wheel, then build one LMCache wheel containing both the HIPcuda_opsextension and the host-only Mooncake extension. Keep the NIXL preview Mooncake plugin disabled so this change does not add a second integration path.Test Plan
Complete the canonical runtime image build and its archive, import, adapter, loader, dependency, and NIXL-isolation gates before marking this PR ready.
Test Result
2026-09-06: full hardware CI (
/run-ci) passed on867fad2; the local results below are from an earlier head,ff3da34.The canonical image build from exact PR head
ff3da34completed on MI308X/gfx942. Build-time SDK, wheel, loader, dependency, adapter, and NIXL-isolation gates passed, as did all 24 compiler-switch tests, Bash syntax, ShellCheck, and diff checks.The exact OCI archive (
sha256:669cc6bdd09a6adb46f855ccd654553dcabd5b419178cc3d0eb5fb9e10ac6fac) exported and imported into an empty isolated containerd store on/data; its 12 pinned image labels, post-import runtime gate, clean LMCache/Mooncake wheel-pair gate, and all four wheel CRC/SHA-256 checks passed.With
HIP_VISIBLE_DEVICES=0, PyTorch2.13.0+rocm7.2saw one AMD Instinct MI308X (gfx942); a real 64x64 GPU matmul, native vLLM/LMCache/Mooncake imports, LMCache connector provenance, and atorch.compile-generated HIP Triton kernel all passed.Scope: packaging/runtime/GPU smoke only; no model-serving or distributed-KV end-to-end job was run.
Submission Checklist