Skip to content

[MRV2] Keep block table metadata across KV wake-up - #9

Merged
aoshen02 merged 1 commit into
Ronald1995:codex/sleep-mrv2-kv-scopefrom
aoshen02:codex/pr53508-remove-wake-rebuild
Aug 25, 2026
Merged

aoshen02 merged 1 commit into
Ronald1995:codex/sleep-mrv2-kv-scopefrom
aoshen02:codex/pr53508-remove-wake-rebuild

Conversation

@aoshen02

Copy link
Copy Markdown
Collaborator

Motivation

PR vllm-project#53508 correctly moves MRV2 block-table metadata outside the sleepable KV allocation scope. Recreating that metadata after wake-up is therefore unnecessary and changes tensor objects captured by fused multi-step CUDA graphs.

Changes

  • Do not run the legacy KV wake callback for Model Runner V2.
  • Remove MRV2's post-wake block-table reconstruction.
  • Keep the callback for Model Runner V1, where it still restores FP8 KV state.

Validation

Tested with the same reproducer before and after this change:

  • 4x H200, single node
  • Qwen/Qwen3.5-35B-A3B, TP=4
  • Model Runner V2
  • MTP with 3 speculative tokens
  • full CUDA Graph
  • level-2 sleep, wake weights, reload weights, wake KV cache, decode

PR vllm-project#53508 alone reproduced CUDA error: an illegal memory access was encountered on the first post-wake decode. With this follow-up:

  • layout_ptrs_stable=true
  • tokens_equal=true
  • no CUDA error
  • sleep: 0.507 s
  • wake weights: 0.455 s
  • wake KV: 1.132 s
  • first post-wake decode: 0.121 s

Pre-commit passes, including mypy and DCO.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

@aoshen02
aoshen02 merged commit b0aa9f4 into Ronald1995:codex/sleep-mrv2-kv-scope Aug 25, 2026
2 of 3 checks passed
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