Repository navigation
Conversation
Hash the token after each block so EAGLE-family drafters can reuse the last matched prefix-cache block instead of dropping it. KV connectors keep the legacy hashes; disable_lookahead_block_hashes opts out. Split from vllm-project#50897. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> # Conflicts: # vllm/v1/core/kv_cache_utils.py # vllm/v1/engine/core.py
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
test_mamba_prefix_cache_mrv2_async_flashinfer compared the checkpoint a 2B + 10 prefill exports with the final state of a B-token prefill. That holds only with the EAGLE drop: lookahead hashes place the checkpoint at 2B, so the comparison failed (rel_err 1.40) though the code is correct. Compute the checkpoint from the scheduler's placement instead. Signed-off-by: Dakai An <dakaian108@gmail.com>
A lookahead hash covers the token after its block, so the block can only be published once that token is sampled. Under the sync scheduler the next allocation published it, and the stop path covered finished requests, but an abort in between dropped it (a probe hit 2 tokens after an abort vs. 4 after a stop). Publish on every step that appends tokens instead; the stop path is the same case. The async scheduler already caches right after appending. Signed-off-by: Dakai An <dakaian108@gmail.com>
Multi-module MTP draft KV reads one more token per module: module j at
position p reads x_{p+j}, so a block's last positions depend on the K
tokens after it (K = num_speculative_tokens, the prefill lookahead).
Lookahead hashes covered only the next token, so these models were
excluded and kept the drop.
Cover num_prefill_lookahead_tokens tokens instead: 1 for EAGLE and
single-module MTP, K for multi-module MTP (Step3.5, Inkling, MiMo-V2).
The hash then waits for all K tokens. Later modules' KV for the last
K - 1 computed tokens is re-prefilled on the next step; a block ending
there cannot be hashed yet (its K tokens are not all known), so cached
blocks only hold final draft KV. Partial entries whose lookahead hash is
not computed yet are skipped.
get_num_hash_lookahead_tokens() derives this count once for the request
hasher, the block pool and the scheduler. It is passed as
num_hash_lookahead_tokens, apart from num_lookahead_tokens, the KV slots
reserved for draft tokens.
Signed-off-by: Dakai An <dakaian108@gmail.com>
…ater tokens Lookahead hashes cover the tokens after a block because EAGLE and MTP draft KV at a position also reads them. DFlash and DSpark only read the next token as the bonus query input, whose KV lands past the computed range; their cached context KV is projected from target hidden states. Gemma4 MTP caches no KV of its own. Their cached KV thus depends on the prefix alone, yet lookahead hashes made each hit also match the next token and delayed every block by one token. Return 0 from get_num_hash_lookahead_tokens() for them, so they use plain hashes, still without the drop. Only the lookahead-hash path changes; connectors and HiSparse keep the drop as before. Signed-off-by: Dakai An <dakaian108@gmail.com>
BlockStored.extra_keys carries each block's untagged extra keys (LoRA name, mm pairs, cache_salt, prompt embedding digests). KV-aware routers parse them: Dynamo maps mm identifiers and cache salts into its routing keys, and llm-d adds mm hashes to its block keys. Lookahead hashes replaced each entry with a (block keys, lookahead tokens, lookahead keys) triple, even (None, (x_e,), None) for plain text. Dynamo's extra-key item type has no nil, so it fails to decode such an event and drops the whole batch; llm-d no longer finds the mm identifier and drops mm hashes. Keep extra_keys as documented and publish what the hash additionally covers in two new optional fields, lookahead_token_ids and lookahead_extra_keys, whose docstrings give the hash input. Events are encoded as maps and these fields are omitted when unset, so consumers that do not know them skip them. Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com> # Conflicts: # vllm/config/speculative.py # vllm/v1/core/sched/scheduler.py
Hash tokens and ordinary extra keys over the expanded window, retaining multimodal features across overlapping blocks. Remove lookahead KV-event plumbing and keep events on legacy hashes. Defer synchronous post-output cache registration and simplify hash-boundary counting and configuration. Validation: 46 source-isolated tests passed; pre-commit and staged preflight passed. Full integration and GPU tests were not rerun. Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Stop producer target computation before a known prompt suffix while keeping the full prompt available to the drafter. Align cache lookup, Mamba allocation and transfer boundaries, and reject incompatible peer protocols. Refresh final NIXL host buffers from computed blocks even without new allocations. Cover scheduler, store and connector boundaries with focused regressions. Signed-off-by: Dakai An <dakaian108@gmail.com>
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
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.
Summary
Extend expanded-window lookahead hashing to Offloading, Mooncake Store, NIXL pull/push, and Mooncake P/D. Adapted from ZJY0516/vllm#21 to the hash format in #11.
N-L. Validate matching lookahead settings between P/D peers.Validation
CPU tests ran on macOS with the repository's
skip_global_cleanupmarker to avoid an environment-specific teardown crash; test assertions were unchanged.This PR includes the earlier core lookahead changes and targets
main. AI assistance was used for implementation and testing.