[whx-sjtu]

Overview

MoRIIO hybrid (KDA/Mamba) READ re-reads attention pages the decoder already caches; this reads only external-token pages.

Not duplicates: #48471 (non-hybrid length mismatch), #62165 (SWA tail clipping). Rebased on #60867.

Claims

• With a decoder prefix-cache hit, hybrid READ moves only the full-attention pages for tokens [N-1-E, N-1) (N prompt length, E final num_external_tokens). KDA/Mamba and state attention groups are unchanged.
• Full-attention tables that do not cover that range raise MoRIIOError instead of being tail-aligned.
• No throughput claim: the tested workload is prefill-queue bound.

Validation

• test_moriio_hma_scheduler.py: main @ 5689c27 110 passed, PR 119 passed, new tests on main connector 13 failed. All MoRIIO unit tests: 438 passed, 33 skipped. ruff 0.14.0 clean.
• Kimi-K3 1P1D TP8/DCP8, 2 x 8 MI355X: on all 8 decode ranks, 32 sampled prefix-hit transfers READ exactly the missing pages; attention READ 45.77 GB to 12.11 GB (pre-rebase code, same for one MLA group).
• 3600 s agentic replay, concurrency 48: 4273 succeeded, 9 empty-content errors (unpatched: 9).
• No accuracy eval: skipped pages are the decoder's own cached pages.

Commands, revert-red and GPU setup

PYTHONDONTWRITEBYTECODE=1 VLLM_TARGET_DEVICE=cpu PYTHONPATH=$PWD \
  python -m pytest -q -p no:cacheprovider \
  tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
uvx ruff@0.14.0 check <changed files>
uvx ruff@0.14.0 format --check <changed files>

Revert-red, new tests with the main connector: 6 x test_scheduler_accepts_dspark_with_wrapped_full_attention and test_scheduler_accepts_glm5_next_groups_with_mtp (main has no _attn_read_block_sizes), test_glm5_next_mtp_read_pairs_every_group[476-1], 2 x test_update_state_rejects_incomplete_attention_tables, test_update_state_selects_missing_attention_pages[128] and [192], test_update_state_uses_final_external_token_count[128].

GPU build: vllm/vllm-openai-rocm:nightly-81198e97 plus this fix before the #60867 rebase, which only made the selection per attention group. MoRIIO READ, prefix caching on both sides, prefill-side MultiConnector (MoRIIO + SimpleCPUOffloadConnector), DSpark speculative decoding, FP8 KV. One attention page is 12288 tokens after DCP sharding. The decoder's local prefix hit was 89.7% of prompt tokens.

Sampled transfers came after warmup, with E > 0 and at least one complete cached decoder page. A per-rank observer recorded READ submissions and completions; bytes are logical payload.

	attention pages	attention GB	KDA GB	total GB
old selection, same transfers (offline)	223	45.77	15.03	60.80
this PR (observed)	59	12.11	15.03	27.14

Throughput 10557 tok/s/chip vs 10882 for the unpatched run, which used a different cluster, so not a controlled A/B. Decoder request_prefill_time avg 0.12 s unpatched vs 0.11 s patched; prefill-side queue avg 13.8 s vs 16.2 s.

Details

Root cause on the decoder:

1. The local prefix hit covers [0, L), so E = N - 1 - L.
2. update_state_after_alloc gets complete attention tables, cached pages included.
3. _align_read_blocks paired P and D tables by length and tail, so READ also covered [0, L).
4. Fix: per full-attention group, slice both by page [(N-1-E) // B, ceil((N-1) / B)), B = its block size x DCP when sharded; window-clipped and state groups keep _align_read_blocks.

 for local, remote, read_block_size in zip(..., self._attn_read_block_sizes):
-    self._align_read_blocks(local, remote, self._max_decode_tail_blocks)
+    if read_block_size is None:  # window-clipped or state group
+        self._align_read_blocks(local, remote, self._max_decode_tail_blocks)
+    else:
+        self._select_read_attention(request, E, local, remote, read_block_size)

• Contract change: a shorter decoder table (the removed test's receive suffix) now raises; a real KVCacheManager case and the GPU run pass complete tables.

---

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

[mergify]

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

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

[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.

[kliuae]

Nice work, thanks for the optimization!

[kliuae]

Nit: the prefix blocks skipped are computed from the request wide num_external_tokens. For a EAGLE drafter that can have a different cache group from the main model, its prefix cache lookup drops the last matched page, and its cache hit is one page shorter than the main group. Since the scheduler counts the main group's hit, the dropped page is neither recomputed nor copied where previously the full read would fetch it from the prefill side, but it gets zeroed and may lead to lowered draft acceptance. Should we consider gating this case off this path? Not sure if this configuration is encountered often though.