Skip to content

feat(mp): add chunk-store server plugin - #1

Draft
yhl-amd wants to merge 1 commit into
codex/server-module-pluginsfrom
codex/chunk-store-on-plugins
Draft

yhl-amd wants to merge 1 commit into
codex/server-module-pluginsfrom
codex/chunk-store-on-plugins

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Sep 29, 2026

Copy link
Copy Markdown
Owner

What this PR does / why we need it:

Add an opt-in chunk-store server module that returns a completion event for each chunk, allowing explicit RPC clients to reuse source KV blocks as their transfers finish. The module loads through --server-module, reuses the existing transfer module, and exposes a typed RPC over ZMQ and gRPC. Includes unit tests, cross-process GPU tests, and usage/design documentation.

Special notes for your reviewers:

This PR is entirely within yhl-amd/LMCache. Its base, codex/server-module-plugins, is an exact copy of maobaolong's plugin branch at 352ea6bc56ff814c61df5be6332399465cbb0a29. The head adds one signed-off commit. Files changed contains only the chunk-store delta: 9 files, +723/-0 (190 lines of implementation/protocol, 453 test lines, 80 documentation lines). The feature patch is identical to the previously reviewed implementation.

Each chunk uses the ordinary store path independently: earlier chunks may remain committed if a later chunk fails. Event handles inherit the existing backend retention contract. Submission overhead has not been benchmarked.

Validation on head 74f63569e8fbbd94a9985fd33d6bb687d159a836, directly on the plugin baseline:

  • 150 tests passed, including 4 real cross-process GPU cases (ZMQ/gRPC, token offsets 0/16) checking exact KV retrieval after per-chunk source overwrite.
  • Ran on mi355-gpu-51 through Slurm job 7803, with one GPU; allocation released after completion. Python 3.10.12, PyTorch 2.10.0+rocm7.0, HIP 7.0.51831. Reused an existing native build; native extension sources are unchanged.
  • Ruff, formatting, isort, mypy for the changed production files, codespell, SPDX, protobuf generation, and diff checks passed. The chunk-store documentation page builds with strict Sphinx checks.
====================== 150 passed, 150 warnings in 18.37s ======================

The warnings are the existing OpenTelemetry LoggingHandler deprecation.

Reproduce the test selection
python -m lmcache.v1.multiprocess.transport.grpc_impl._proto_gen._generate
python -m pytest -q \
  tests/v1/multiprocess/test_chunk_store.py \
  tests/v1/multiprocess/test_chunk_store_gpu.py \
  tests/v1/multiprocess/test_server_module.py \
  tests/v1/multiprocess/test_grpc_transport.py \
  tests/v1/multiprocess/test_event_ipc_handle_path.py \
  tests/v1/multiprocess/test_config.py \
  tests/v1/multiprocess/test_transport_factory.py \
  tests/v1/multiprocess/test_protocols.py

If applicable:

  • this PR contains user facing changes - docs added
  • this PR contains unit tests

Signed-off-by: Honglie Yi <hyi@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant