Skip to content

[KV Connector][Spec Decode] Lookahead block hashes for KV connectors - #21

Closed
ZJY0516 wants to merge 6 commits into
lookahead-hash-corefrom
lookahead-hash-connector
Closed

ZJY0516 wants to merge 6 commits into
lookahead-hash-corefrom
lookahead-hash-connector

Conversation

@ZJY0516

@ZJY0516 ZJY0516 commented Oct 12, 2026

Copy link
Copy Markdown
Owner

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

  • With a supporting connector, EAGLE-family prefix-cache hits (local and external) reach the same length as without speculative decoding, instead of one cache unit less.
  • Content-addressed connectors store and look up only blocks covered by request.block_hashes, so an external hit carries the same lookahead-token proof as a local one.
  • Connectors that do not opt in are unchanged.

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 OffloadingConnector on 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 over NixlConnector P/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) and tests/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_enabled takes the connector and uses this flag when a KV transfer config is present. KVConnectorFactory calls set_lookahead_block_hashes on every connector it builds, on both the scheduler and the worker side, and MultiConnector forwards it to its children. The scheduler, EngineCore and the Mamba checkpoint builder (via get_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.

  • Offloading: the offloadable range is capped at len(block_hashes) * tokens_per_hash, and the EAGLE tail exclusion is skipped for lookup, store and the reachable mask.
  • Mooncake Store: saves are capped by max_save_tokens, and lookup and store masks are built with apply_eagle=False. At request finish, register_finished_lookahead_save pins 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-checklist skill. (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

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>
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-12T08:56:31.241946Z 35d824b PR opened
ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment on lines +638 to +644
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,
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

ZJY0516 and others added 5 commits October 13, 2026 09:43
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
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