[Perf][KVConnector][Mooncake] Vectorize prepare_value on the KV load path - #48531
Merged
youkaichao merged 3 commits intoJul 23, 2026
Merged
Conversation
…path For a 289K-token request the per-chunk prepare_value loop runs ~2.3K chunks x ~100 cache regions of pure-Python integer arithmetic — ~35 ms of GIL-held time per request. On the load path this serializes the receive-thread pool (vllm-project#45971): the key-building phases contend on the GIL, capping effective concurrency at ~3x and slowing co-resident threads. Add ChunkedTokenDatabase.prepare_values(), a numpy batch variant that computes the whole request's address/size matrix in one broadcast (~4.4x cheaper; 34.2 ms -> 7.7 ms measured for a 2,258-block request), and switch KVCacheStoreRecvingThread to one call per group. prepare_value stays as a single-chunk wrapper so there is only one implementation. Equivalence is enforced by tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py, which pins the old loop as the oracle (including block_len cycling and native-int output types). Signed-off-by: girasoley <girasoley@inferact.ai> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
GirasoleY
requested review from
ApostaC,
NickLucche,
ivanium,
orozery and
xuechendi
as code owners
July 13, 2026 21:09
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ivanium
reviewed
Jul 21, 2026
ivanium
left a comment
ivanium
left a comment
Collaborator
There was a problem hiding this comment.
LGTM overall. Sorry for the delayed review. nit: can we also simplify the comments?
Batch missing store chunks by KV-cache group before computing their addresses, and cover mixed-group deduplication and ordering in the worker unit suite. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>
Contributor
Author
|
ac: simplified comment + use same prepare_values in send loop. |
ivanium
approved these changes
Jul 22, 2026
ivanium
enabled auto-merge (squash)
July 22, 2026 23:31
github-actions
Bot
added
the
ready
youkaichao
disabled auto-merge
July 23, 2026 10:11
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
3 participants
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.
Purpose
For a long-context request, 289K tokens is about 2,258 blocks.
ChunkedTokenDatabase.prepare_valuespends about 35 ms of GIL-held time per request computing RDMA target addresses on the KV load path.This cost is amplified by the receive-thread pool in #45971: four concurrent key builds take 208 ms wall time, compared with 174 ms when run serially.
This PR vectorizes the computation with NumPy and keeps
prepare_valueas a single-chunk wrapper around the batch implementation. Measured time for the 2,258-block shape improves from 34.2 ms to 7.7 ms (about 4.4×). Following review feedback, the store/send path also batches address preparation once per KV-cache group.Non-duplication
Open-PR searches for
Mooncake prepare_valueandChunkedTokenDatabase prepare_valuesfound no overlapping work other than this PR.Test Plan
.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py tests/v1/kv_connector/unit/test_mooncake_store_worker.py::test_store_sending_thread_prepares_missing_chunks_once_per_group -q— 7 passed.VLLM_RPC_BASE_PATH=/tmp/vllm-pr48531-rpc-20260722 .venv/bin/python -m pytest tests/v1/kv_connector/unit/test_mooncake_store_worker.py -v— 75 passed..venv/bin/pre-commit run --files vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py tests/v1/kv_connector/unit/test_mooncake_store_worker.py— all applicable hooks passed, including ruff and mypy.Model evaluation
Not applicable. This changes KV-transfer address preparation without changing model outputs, accuracy, or serving behavior.
AI assistance
AI assistance was used to address review feedback, add focused unit coverage, and run validation.