[benenzhu]

Overview

MoRIIO READ pairs prefill and decode sliding-window blocks by clipping each block list from its end. A lookahead tail that only one leg allocated shifts the decoder's window by one block, so the decoder reads the wrong KV for the end of the prompt. This PR trims both lists to the prompt's blocks before clipping.

Claims

• Fixes wrong outputs in MoRIIO READ P/D for sliding-window hybrid models when the two legs have different lookahead tails. On DeepSeek-V4-Pro, ~14% of prompt lengths were affected.
• Also fixes the full-attention pairing when only the prefiller has a lookahead tail.
• No change when both legs hold the same number of blocks.

Validation

Check	main	this PR
New unit test, 4 tail combinations	3/4 fail	4/4 pass
test_moriio_hma_scheduler.py + test_moriio_routing.py (macOS, CPU)	–	111 passed
+ test_moriio_connector.py (ROCm, MI355X)	139 passed	143 passed
DSv4-Pro FP4 1P1D MoRIIO READ*, 2x MI350X (P TP8, D TP8): key-value lookup at 102 prompt lengths, PD vs direct-to-prefill	fails at 251, 377, 384, 440, 447, 636, 832, 895, ... (~14%)	102/102

VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py \
  tests/v1/kv_connector/unit/test_moriio_routing.py

• The unit tests ran on macOS arm64 with CPU torch. They need no GPU and no mori runtime.
• On ROCm (MI355X, mori installed), test_moriio_connector.py + test_moriio_hma_scheduler.py + test_moriio_routing.py: main 139 passed; this PR 143 passed (the 4 new cases). 0 failed, 0 skipped on both.
• *The E2E row was measured on an internal build that carried this trim (as a separate helper with the same logic) together with other MoRIIO READ fixes. It was not re-run on this exact branch.

Details

Flow. Example: DSv4-Pro SWA cache, block_size=64, sliding_window=128, so blocks_per_sw = cdiv(128, 64) + 1 = 3. Prompt N = 440 tokens fills 7 blocks. The decoder reserves k lookahead slots and the prefiller none; the observed failure range (N - 1) % 64 ∈ [55, 63] implies k = 9.

1. P (request_finished): 440 slots, 7 blocks P0..P6. Clip keeps [P4, P5, P6] and sends them as remote_block_ids.
2. D admission: load_kv_async=False, so the scheduler allocates 440 + 9 slots, 8 blocks D0..D7. The async-load path sets lookahead to 0 to keep block counts equal, but synchronous READ does not take it.
3. D (update_state_after_alloc): clip keeps [D5, D6, D7]. _align_read_blocks sees equal lengths and pairs P4→D5, P5→D6, P6→D7.
4. The window of the last prompt token needs D4..D6. It gets wrong or missing KV, so the model answers as if the end of the prompt were missing.

The failure condition with no prefill-side lookahead is (N - 1) % block_size >= block_size - k. _align_read_blocks cannot detect the shift because it runs after the clip.

Fix. get_exchange_clipped_blocks takes an optional num_prompt_tokens. When it is set, every transfer group is first trimmed to cdiv(num_prompt_tokens, block_size) blocks, using the existing _group_block_sizes. Both exchange points pass it: request_finished on the producer and update_state_after_alloc on the consumer, both on the non-mamba path. The list lengths then depend only on the prompt, not on either leg's lookahead.

Safety.

• update_state_after_alloc receives kv_cache_manager.get_blocks(request_id). That is the full position-indexed list, including prefix-hit blocks, so the trim never drops a prompt block.
• The hybrid mamba path is unchanged. It does not clip sliding-window groups, and its prefiller drops the last prompt token.
• _clamp_to_prompt_blocks (WRITE save path) is not reused. Its non-mamba branch skips sliding-window groups, which are the groups this fix needs.

Limitation. _group_block_sizes and blocks_per_sw use the spec block size without the DCP factor. Under DCP > 1 the trim keeps too many blocks and does not fix the shift. It does not regress either.

Duplicate check. I searched gh pr list --repo vllm-project/vllm --state open --search "moriio", "moriio sliding window" and "moriio hybrid READ". No open PR trims before the SWA clip. Two PRs are related:

• [ROCm/MoRIIO] Route READ-mode remote KV load via async lifecycle (fix scheduler req_id_to_index KeyError) #54301 (async READ lifecycle) would remove the decoder-side admission tail, but not a prefill-side tail.
• [KV Connector][ROCm] MoRIIO: hybrid READ for several attention groups and shared-group speculative methods #60867 (multi-group hybrid READ) adds another get_exchange_clipped_blocks caller in split_block_groups. That caller can pass num_prompt_tokens the same way.

---

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.

AI assistance (Claude Code) was used to prepare this change.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[benenzhu]

previous:

now:

[github-actions]

✅ @benenzhu, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[benenzhu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95684 for commit e6f02f65f822.

ℹ️ This PR is 31 commits behind upstream main. Merged upstream main into this branch as e6f02f65f822 before starting CI. Pull the branch before pushing more changes.

[benenzhu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95806 for commit bf88a92435f2.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as bf88a92435f2 before starting CI. Pull the branch before pushing more changes.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benenzhu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[benenzhu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98900 for commit 4b5b859b0705.

ℹ️ This PR is 42 commits behind upstream main. Merged upstream main into this branch as 4b5b859b0705 before starting CI. Pull the branch before pushing more changes.