[ROCm][AMD][Installation] add LMCache kv-connector installation and runtime packages to docker image

[hongxiayang]

Purpose

Ship the LMCache KV connector in the ROCm image so it works out of the box.

Note:
vLLM's CUDA docker/Dockerfile installs LMCache when INSTALL_KV_CONNECTORS=true,
and the release pipeline sets that for every published vllm/vllm-openai tag.

docker/Dockerfile.rocm never referenced that arg, so vllm/vllm-openai-rocm ships without LMCache and users must build and install it themselves. This adds it so the ROCm release pipeline can enable it the same way.

What this does

• Adds a build_lmcache stage that obtains the wheel one of two ways, mirroring
  the CUDA path (docker/Dockerfile tries a prebuilt wheel first and falls back
  to a source build):

  1. Published ROCm wheel. LMCache tags ROCm builds on GitHub Releases
     (v0.5.3-rocm), which PyPI does not carry. When the effective arch set is
     covered by that wheel it is downloaded — no clone, no compile. This is what
     the release build does.
     --no-index is deliberate: with a plain --extra-index-url, pip silently
     falls back to the CUDA wheel on PyPI.
  2. Source build. A wider arch set falls back to compiling for HIP
     (BUILD_WITH_HIP=1, CXX=hipcc, TORCH_DONT_CHECK_COMPILER_ABI=1), so
     coverage is never silently narrowed.

  PyPI remains CUDA-only and ships no sdist for >=0.4, so --no-binary cannot
  force a HIP build — hence a source path is still required.

• Architecture selection. The image's PYTORCH_ROCM_ARCH is inherited, minus
  the consumer archs vLLM itself strips for its own wheel
  (Dockerfile.rocm_base), then intersected with the archs LMCache publishes
  ROCm wheels for. Anything dropped is named in the build log, and a set with no
  supported arch fails the build rather than shipping an unusable connector.

  On the release base this resolves to gfx942;gfx950 — MI300X/MI325X and
  MI350X/MI355X.

  gfx90a (MI200) is not covered. vLLM builds for it, LMCache publishes no
  ROCm wheel for it. Rather than ship non-runnable device code, it is skipped
  and logged. To build it from source instead:
  --build-arg LMCACHE_ROCM_ARCH="gfx90a;gfx942;gfx950".

• The wheel is bind-mounted for install while its sha256 digest is COPYd, so
  a changed LMCACHE_REF or arch set invalidates the install layer instead of
  reusing a stale wheel — without committing the 7.8 MB wheel blob to a layer.

• The install is placed before the vLLM wheel install, so vLLM source changes
  do not invalidate this otherwise stable connector layer.

• The build asserts the wheel actually carries code objects for every requested
  arch, rather than only that some .so exists.

• The source build opts into the sccache HIP wrappers when USE_SCCACHE=1 (the
  contract documented in Dockerfile.rocm_base) and shares the ccache mount.

• Installs the wheel with --no-deps plus four runtime packages.

  • sortedcontainers — hard requirement; without it
    import lmcache.integration.vllm.lmcache_mp_connector fails outright.
  • opentelemetry-exporter-prometheus — lmcache server calls
    init_observability() unconditionally, which imports PrometheusMetricReader
    (mp_observability/otel_init.py:96). Omitting it builds fine but the server
    dies at startup.
  • cupy-rocm-7-0 — GPU cache-registration path.
  • aiofile — local-disk storage tier.

  --no-deps is required, not a convenience. The published ROCm wheel declares
  numpy<=2.2.6, opentelemetry-api<=1.40.0 and
  opentelemetry-exporter-prometheus<=0.61b0. The numpy ceiling is
  unsatisfiable — vLLM requires 2.3.5 — so resolving it would downgrade vLLM's
  own dependency. The two OpenTelemetry ceilings are upper bounds rather than
  known breakages; the image runs 1.44.0/0.65b0 and lmcache server starts and
  serves through init_observability() on those versions (see Test Result 2).
  Worth raising upstream so the declared ranges match what is tested.

• Gates everything behind INSTALL_LMCACHE, defaulting to off.

Build-arg selection

INSTALL_LMCACHE is an independent boolean rather than being derived from an
umbrella arg, so it is not an input to the cache key of unrelated stages. The
release pipeline sets it explicitly.

Build args	Result
(none)	nothing installed (default)
INSTALL_LMCACHE=true	LMCache installed

A boolean rather than a list because the build stage is selected via
FROM <stage>_${ARG}, which needs the value in the stage name. With a list the
stage would run on every build and be discarded.

Arch lists (LMCACHE_ROCM_ARCH, LMCACHE_SUPPORTED_ARCHS,
LMCACHE_PREBUILT_ARCHS) use PYTORCH_ROCM_ARCH-style ; separators; space and
, forms are normalised so a mismatched separator cannot silently drop archs.

Test Plan

# (1) release default — inherits the base arch set, resolves to gfx942;gfx950
docker buildx build -f docker/Dockerfile.rocm --target vllm-openai \
  --build-arg INSTALL_LMCACHE=true \
  -t vllm-rocm-lmcache:release --load .

# (2) source-build path — explicit arch set wider than the published wheel
docker buildx build -f docker/Dockerfile.rocm --target vllm-openai \
  --build-arg INSTALL_LMCACHE=true \
  --build-arg LMCACHE_ROCM_ARCH="gfx90a;gfx942;gfx950" \
  -t vllm-rocm-lmcache:src --load .

# (3) start the LMCache server from the image, installing nothing
docker run --rm --network host --ipc host \
  --device /dev/kfd --device /dev/dri --group-add video \
  -e LMCACHE_L1_SIZE_GB=906 --entrypoint bash vllm-rocm-lmcache:release -c '
lmcache server --host 127.0.0.1 --port 5555 \
  --http-host 127.0.0.1 --http-port 8080 \
  --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 20 \
  --l1-read-ttl-seconds 7200 --chunk-size 1536 --max-workers 8 \
  --eviction-trigger-watermark 0.85 --eviction-ratio 0.10 \
  --eviction-policy LRU'

Test Result

Built and verified on 8x MI355X (gfx950), ROCm 7.2.3, torch 2.11.

(1) Arch selection and both wheel paths.

The release default inherits the base's nine architectures and resolves them
down, taking the published wheel:

[LMCACHE] not supported upstream, skipping: gfx90a
[LMCACHE] effective PYTORCH_ROCM_ARCH: gfx942;gfx950
[LMCACHE] using published ROCm wheel

An explicit wider set falls back to a source build:

[LMCACHE] effective PYTORCH_ROCM_ARCH: gfx90a;gfx942;gfx950
[LMCACHE] building from source for gfx90a;gfx942;gfx950

	published wheel	source build
wheel	lmcache-0.5.3-cp312-cp312-manylinux_2_35_x86_64.whl	...-linux_x86_64.whl
stage time	~1.3 s (download)	full hipcc compile
GPU archs in c_ops	gfx942, gfx950	gfx90a, gfx942, gfx950
DT_NEEDED	libamdhip64 — no CUDA libs	same
digest check	...whl: OK	...whl: OK

The arch assertion was negative-tested: forcing the published wheel to be
selected for a gfx90a build fails the build with
wheel has no code objects for: gfx90a (exit 1, nothing exported).
All three separator forms (;, ,, space) resolve identically.

(2) LMCache server starts out of the box, nothing installed at runtime:

zmq  5555 -> accepting
http 8080 -> HTTP 200

Startup log shows shm capacity checked against the L1 request, L1 + L2 eviction
controllers started, SessionManager cleanup thread, DeviceHostFuncDispatcher,
mp-worker-reaper, and an 8-slot affinity pool matching --max-workers 8.
This exercises init_observability() on opentelemetry 1.44.0 / exporter 0.65b0.

(3) No dependency downgrades — the reason for --no-deps:

Package	Before	After
numpy	2.3.5	2.3.5
transformers	5.14.1	5.14.1
opentelemetry-api / sdk	1.44.0	1.44.0
opentelemetry-semantic-conventions	0.65b0	0.65b0

Added: lmcache 0.5.3, sortedcontainers 2.4.0, aiofile 3.12.3,
cupy-rocm-7-0 14.1.1, opentelemetry-exporter-prometheus 0.65b0.
Unchanged after moving the install ahead of the vLLM wheel.

(4) Connectors register: ['LMCacheConnectorV1', 'LMCacheMPConnector'].

(5) Layer and cache behaviour.

• The wheel no longer lands in a layer: docker history shows the LMCache
  install layer at 118 B (the digest) rather than 7.82 MB.
• Cache correctness verified against the failure it prevents: build a wheel,
  rebuild unchanged (layers cached), then change the wheel content — the install
  layer re-runs and picks up the new wheel.
• With INSTALL_LMCACHE=false the stage is pruned entirely
  (--target export_lmcache produces empty output — no clone, no download,
  no compile).

(6) Image size: 42.00 → 42.92 GB, +921 MB (+2.19%) — ~366 MB of that is CuPy,
required by the GPU cache-registration path; the compiled extension is <8 MB.

AI-assisted.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hongxiayang]

cc @vllmellm Can you help me to verify on your end?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hongxiayang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jamesETsmith]

Would it also be worth adding a comment about the PYTHONHASHSEED=0 here? I'm not an lmcache expert, but for people running with TP>1 they might have poor performance out of the box here. I'm going of this blog about that env variable. I'm trying to test it locally now

[jamesETsmith]

Should we set this to true here or update the buildkite release pipelines here to set it to true?

[hongxiayang]

yes, added to the release pipeline now

[mergify]

Documentation preview: https://vllm--51208.org.readthedocs.build/en/51208/

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83990 for commit f411fe7e6f58.

[hongxiayang]

A note, I used to do -no-deps so it keeps our vllm's numpy 2.3.5 (non-pined). The way the build work now after @AndreasKaratzas 's latest two commits is to installed LMCache numpy pined version, which is 2.2.6. So basically, the numpy consumers may see the difference. Thoughts @AndreasKaratzas ?

[AndreasKaratzas]

Yeah, I can restore that part. Indeed numpy is too central. Give me a sec and will be following up

[mergify]

Hi @hongxiayang, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84254 for commit a50eeb658d79.