Repository navigation
Conversation
Let Offloading, Mooncake Store, Nixl, Mooncake P2P and MultiConnector opt into lookahead block hashes via supports_lookahead_block_hashes. Split from vllm-project#50897. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35d824bb30
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| req_meta = ReqMeta.from_request_tracker( | ||
| tracker, | ||
| self._block_size, | ||
| block_hashes=request.block_hashes, | ||
| max_save_tokens=self._max_save_tokens(request), | ||
| num_prompt_tokens=request.num_prompt_tokens, | ||
| ) |
There was a problem hiding this comment.
Preserve Mamba state in the finish-time lookahead save
When a hybrid Mamba request finishes on the first token after a hash boundary (for example, an aligned prompt with max_tokens=1), that sampled token publishes the boundary's lookahead hash only at finish. This new ReqMeta contains no exact Mamba boundary_puts, while the normal positional store path always excludes Mamba groups, so only the attention KV is saved under the newly published hash. A later Mooncake Store lookup must therefore shorten the hybrid hit at that boundary. Include the corresponding pinned Mamba boundary state in this finish-time job, as is done for other Mamba boundary saves.
Useful? React with 👍 / 👎.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
# Conflicts: # tests/v1/core/test_kv_cache_utils.py # vllm/distributed/kv_transfer/kv_connector/v1/prefix_cache.py Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> # Conflicts: # vllm/distributed/kv_transfer/kv_connector/v1/prefix_cache.py # vllm/model_executor/layers/mamba/checkpoint.py
Overview
Part 2 of 2 split from vllm-project#50897, stacked on #20. Extends lookahead block hashes to KV connectors: Offloading, Mooncake Store, Nixl, Mooncake P2P, and MultiConnector when all of its children support it. Other connectors keep the EAGLE last-block drop.
Claims
request.block_hashes, so an external hit carries the same lookahead-token proof as a local one.Validation
Unit tests on this branch (CPU, local):
python -m pytest tests/v1/kv_connector/unit/test_multi_connector.py \ tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py \ tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py \ tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \ tests/v1/kv_connector/unit/test_mooncake_store_connector.py \ tests/v1/kv_connector/unit/test_mooncake_store_worker.py \ tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py \ tests/v1/kv_connector/unit/test_mooncake_connector.py \ tests/v1/kv_connector/unit/test_nixl_connector.py \ tests/v1/kv_connector/unit/test_nixl_connector_hma.py \ tests/v1/core/test_kv_cache_utils.py tests/v1/core/test_scheduler.py -k "not ray"Result: 1,568 passed, 1 deselected.
From the combined PR (vllm-project#50897, earlier revision; not yet re-run on this split): with
OffloadingConnectoron Qwen/Qwen3.6-27B + MTP k=4 (single run), hit rate 72.7% → 85.0% and throughput 8.93 → 11.26 req/s, with the same GSM8K accuracy and acceptance length. Llama-3.2-1B + EAGLE3 overNixlConnectorP/D, GSM8K 1,319 questions: 0.3215 / 0.3207 across two passes, acceptance length 2.077.GPU tests added here, to be re-run before review:
tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py::test_qwen3_5_mtp_prefix_cache_reuses_last_safe_block(OffloadingConnector) andtests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py::test_mtp_pd_reuses_last_safe_prefix_block.Details
Protocol.
KVConnectorBase_V1.supports_lookahead_block_hashes(default False) declares that a connector stays correct with lookahead hashes.is_lookahead_block_hashing_enabledtakes the connector and uses this flag when a KV transfer config is present.KVConnectorFactorycallsset_lookahead_block_hasheson every connector it builds, on both the scheduler and the worker side, and MultiConnector forwards it to its children. The scheduler,EngineCoreand the Mamba checkpoint builder (viaget_kv_transfer_group()) evaluate the same function on the same connector class, so they agree.Content-addressed connectors. Lookahead hashes lag the computed tokens by up to one block.
len(block_hashes) * tokens_per_hash, and the EAGLE tail exclusion is skipped for lookup, store and the reachable mask.max_save_tokens, and lookup and store masks are built withapply_eagle=False. At request finish,register_finished_lookahead_savepins a store job for the block hashed by the last sampled token and emits it on the next step.Direct-transfer connectors. Nixl and Mooncake P2P transfer a request's own blocks rather than cache hashes, so they opt in unconditionally. The capability does not depend on local APC, so P and D pick the same protocol when another connector shares hashes. For Mamba models they keep truncating the prefill side at
h(N-1).Limitations. Engines sharing a content-addressed store must use the same mode (and the same
disable_lookahead_block_hashes); a mismatch shows up only as a lower hit rate.AI assistance (Claude Code) was used to split, implement, test and describe this change. The submitting human reviews every changed line and is responsible for it end to end.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code