[Bugfix][KV Connector][ROCm] Fix MoRI-IO KV block-offset for speculative decoding (MTP) (#55053)

Signed-off-by: Mohit Deopujari <Mohit.Deopujari@amd.com>
Signed-off-by: MohitAMD <mohit.deopujari@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -3,6 +3,8 @@
 import importlib.util
 import socket
 import uuid
+from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import msgspec
@@ -20,11 +22,13 @@ from vllm.config import (
     set_current_vllm_config,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
+    ROLE,
     MoRIIOAgentMetadata,
     MoRIIOConnectorMetadata,
     MoRIIOConstants,
     MoRIIOMode,
     resolve_host_ip,
+    set_role,
     zmq_ctx,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
@@ -147,6 +151,50 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2):
     return scheduler
 
 
+def _write_producer_scheduler(block_size: int = 1) -> Any:
+    """Bare WRITE-mode PRODUCER scheduler exercising the save path without a
+    real engine/RDMA stack. Uses a single full-attention KV group (index 0)."""
+    scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler)
+    scheduler.mode = MoRIIOMode.WRITE
+    scheduler.is_producer = True
+    scheduler.block_size = block_size
+    scheduler.blocks_per_sw = [0]
+    scheduler._group_block_sizes = [block_size]
+    scheduler.transfer_id_to_request_id = {}
+    scheduler._reqs_need_recv = {}
+    scheduler._reqs_need_save = {}
+    scheduler._reqs_need_pending_save = {}
+    scheduler._req_kv_params = {}
+    scheduler._reqs_need_send = {}
+    return scheduler
+
+
+def _spec_kv_params(transfer_id: str = "xfer-spec") -> dict[str, Any]:
+    # request_id embeds no zmq address, so add_new_req resolves the peer here.
+    return {
+        "transfer_id": transfer_id,
+        "remote_engine_id": "remote-engine",
+        "remote_block_ids": [],
+        "remote_host": "127.0.0.1",
+        "remote_handshake_port": 5000,
+        "remote_notify_port": 5001,
+    }
+
+
+def _build_meta(
+    scheduler: Any, req_ids=None, new_block_ids=None, num_scheduled_tokens=None
+) -> Any:
+    # Final-chunk detection reads num_scheduled_tokens, so expose it.
+    scheduler_output = SimpleNamespace(
+        scheduled_cached_reqs=SimpleNamespace(
+            req_ids=req_ids or [],
+            new_block_ids=new_block_ids or [],
+        ),
+        num_scheduled_tokens=num_scheduled_tokens or {},
+    )
+    return scheduler.build_connector_meta(scheduler_output)
+
+
 class FakeMoRIIOWrapper:
     # A fake MoRIIOWrapper for testing purposes
     def __init__(self, *args, **kwargs):
@@ -700,6 +748,203 @@ def test_resolve_host_ip_prefers_extra_config():
     assert resolve_host_ip({"host_ip": ""}) == fallback
 
 
+def test_write_mode_excludes_spec_lookahead_blocks():
+    # The save path must record only the prompt blocks, clamping off the
+    # trailing lookahead blocks decode never allocates.
+    set_role(ROLE.PRODUCER)
+    scheduler = _write_producer_scheduler(block_size=1)
+
+    req_id = "req-spec"
+    prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17]
+    lookahead_blocks = [18, 19, 20]
+    local_block_ids = [prompt_blocks + lookahead_blocks]
+
+    # block_size=1 => num_prompt_tokens == prompt block count, all scheduled in
+    # this single chunk, so final-chunk detection fires.
+    num_prompt_tokens = len(prompt_blocks)
+    req = SimpleNamespace(
+        request_id=req_id,
+        num_prompt_tokens=num_prompt_tokens,
+        num_computed_tokens=0,
+        kv_transfer_params=_spec_kv_params(),
+    )
+    scheduler._reqs_need_save[req_id] = (req, local_block_ids)
+    scheduler._req_kv_params[req_id] = _spec_kv_params()
+
+    meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens})
+
+    assert isinstance(meta, MoRIIOConnectorMetadata)
+    assert req_id in meta.reqs_to_save
+    saved = meta.reqs_to_save[req_id]
+    assert saved.local_block_ids == [prompt_blocks]
+
+
+def test_write_mode_chunked_prefill_clamps_spec_lookahead_blocks():
+    # Non-final chunks are buffered untouched; lookahead blocks are only clamped
+    # on the final-chunk save.
+    set_role(ROLE.PRODUCER)
+    scheduler = _write_producer_scheduler(block_size=1)
+
+    req_id = "req-spec-chunked"
+    num_prompt_tokens = 8
+    first_chunk = [10, 11, 12, 13]
+    final_chunk = [14, 15, 16, 17, 18, 19, 20]  # 4 prompt + 3 lookahead
+    prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17]
+
+    req = SimpleNamespace(
+        request_id=req_id,
+        num_prompt_tokens=num_prompt_tokens,
+        num_computed_tokens=0,
+        kv_transfer_params=_spec_kv_params(),
+    )
+
+    # Step 1: first chunk (4 of 8 tokens) is buffered, not saved.
+    scheduler._reqs_need_save[req_id] = (req, [first_chunk])
+    scheduler._req_kv_params[req_id] = _spec_kv_params()
+    meta_step1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4})
+    assert req_id not in meta_step1.reqs_to_save
+    assert req_id in scheduler._reqs_need_pending_save
+    assert scheduler._reqs_need_pending_save[req_id][1] == [first_chunk]
+
+    # Step 2: final chunk (computed 4 + scheduled 4 >= 8) assembles the full
+    # local set, clamps the lookahead, and records only the prompt blocks.
+    req.num_computed_tokens = 4
+    meta_step2 = _build_meta(
+        scheduler,
+        req_ids=[req_id],
+        new_block_ids=[[final_chunk]],
+        num_scheduled_tokens={req_id: 4},
+    )
+    assert req_id in meta_step2.reqs_to_save
+    assert meta_step2.reqs_to_save[req_id].local_block_ids == [prompt_blocks]
+    assert req_id not in scheduler._reqs_need_pending_save
+
+
+@pytest.mark.parametrize(
+    ("block_size", "num_prompt_tokens", "num_prompt_blocks"),
+    [
+        # block_size=4: exact multiple (2 full blocks) and non-multiple (3 blocks,
+        # last one partial).
+        pytest.param(4, 8, 2, id="bs4-exact-multiple"),
+        pytest.param(4, 9, 3, id="bs4-non-multiple"),
+        # block_size=16: exact multiple (2 full blocks) and non-multiple (3
+        # blocks, last one partial).
+        pytest.param(16, 32, 2, id="bs16-exact-multiple"),
+        pytest.param(16, 40, 3, id="bs16-non-multiple"),
+    ],
+)
+def test_write_mode_clamps_prompt_blocks_with_block_size_gt_1(
+    block_size, num_prompt_tokens, num_prompt_blocks
+):
+    # With block_size > 1 the clamp must keep exactly
+    # ceil(num_prompt_tokens / block_size) prompt blocks and drop the tail.
+    set_role(ROLE.PRODUCER)
+    scheduler = _write_producer_scheduler(block_size=block_size)
+
+    req_id = f"req-bs{block_size}-{num_prompt_tokens}"
+    prompt_blocks = list(range(10, 10 + num_prompt_blocks))
+    lookahead_blocks = [10 + num_prompt_blocks]
+    local_block_ids = [prompt_blocks + lookahead_blocks]
+
+    req = SimpleNamespace(
+        request_id=req_id,
+        num_prompt_tokens=num_prompt_tokens,
+        num_computed_tokens=0,
+        kv_transfer_params=_spec_kv_params(),
+    )
+    scheduler._reqs_need_save[req_id] = (req, local_block_ids)
+    scheduler._req_kv_params[req_id] = _spec_kv_params()
+
+    meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens})
+
+    assert isinstance(meta, MoRIIOConnectorMetadata)
+    assert req_id in meta.reqs_to_save
+    assert len(prompt_blocks) == num_prompt_blocks
+    assert meta.reqs_to_save[req_id].local_block_ids == [prompt_blocks]
+
+
+def test_write_mode_defers_save_until_prompt_complete_with_lookahead():
+    # Under speculative decoding the producer reserves lookahead blocks beyond
+    # the last-scheduled prompt token, inflating the accumulated block count on
+    # a non-final chunk. Token-progress detection must keep deferring until the
+    # prompt is complete, where the lookahead is clamped off.
+    #
+    # block_size=1, num_prompt_tokens=10, chunk=4, num_lookahead=3: after 8
+    # prompt tokens the producer holds 8 + 3 = 11 blocks, but token progress
+    # 8 < 10 must keep waiting until the final chunk completes tokens 9-10.
+    set_role(ROLE.PRODUCER)
+    scheduler = _write_producer_scheduler(block_size=1)
+
+    req_id = "req-lookahead-final-chunk"
+    num_prompt_tokens = 10
+
+    chunk1_blocks = [100, 101, 102, 103]  # prompt tokens 1-4
+    chunk2_blocks = [104, 105, 106, 107, 108, 109, 110]  # 4 prompt + 3 lookahead
+    prompt_blocks = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109]
+
+    req = SimpleNamespace(
+        request_id=req_id,
+        num_prompt_tokens=num_prompt_tokens,
+        num_computed_tokens=0,
+        kv_transfer_params=_spec_kv_params(),
+    )
+    scheduler._req_kv_params[req_id] = _spec_kv_params()
+
+    # Chunk 1 (tokens 1-4): computed 0 + scheduled 4 < 10 -> buffered.
+    scheduler._reqs_need_save[req_id] = (req, [chunk1_blocks])
+    meta1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4})
+    assert req_id not in meta1.reqs_to_save, "chunk 1 must not be saved"
+    assert req_id in scheduler._reqs_need_pending_save
+    assert scheduler._reqs_need_pending_save[req_id][1] == [chunk1_blocks]
+
+    # Chunk 2 (tokens 5-8) inflates the buffered count to 11 via lookahead, but
+    # token progress 4 + 4 = 8 < 10 keeps the request deferred.
+    req.num_computed_tokens = 4
+    meta2 = _build_meta(
+        scheduler,
+        req_ids=[req_id],
+        new_block_ids=[[chunk2_blocks]],
+        num_scheduled_tokens={req_id: 4},
+    )
+    assert req_id not in meta2.reqs_to_save, (
+        "second-to-last chunk saved prematurely despite incomplete prompt"
+    )
+    assert req_id in scheduler._reqs_need_pending_save
+    assert scheduler._reqs_need_pending_save[req_id][1] == [
+        chunk1_blocks + chunk2_blocks
+    ]
+
+    # Final chunk (tokens 9-10, no new blocks): computed 8 + scheduled 2 >= 10
+    # emits the save with the trailing lookahead block clamped off.
+    req.num_computed_tokens = 8
+    meta3 = _build_meta(
+        scheduler,
+        req_ids=[req_id],
+        new_block_ids=[None],
+        num_scheduled_tokens={req_id: 2},
+    )
+    assert req_id in meta3.reqs_to_save, "final chunk must emit the WRITE save"
+    assert meta3.reqs_to_save[req_id].local_block_ids == [prompt_blocks]
+    assert req_id not in scheduler._reqs_need_pending_save
+
+
+def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size():
+    # Two full-attention groups with different block sizes are each clamped to
+    # their own prompt-block count; a sliding-window group is left untouched.
+    scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler)
+    scheduler.blocks_per_sw = [0, 0, 3]
+    scheduler._group_block_sizes = [16, 32, 16]
+
+    req = SimpleNamespace(num_prompt_tokens=40)
+    block_ids = [
+        [0, 1, 2, 3, 4],  # ceil(40/16)=3 -> clamp to 3
+        [10, 11, 12],  # ceil(40/32)=2 -> clamp to 2
+        [20, 21, 22, 23, 24],  # sliding-window -> untouched
+    ]
+    clamped = scheduler._clamp_to_prompt_blocks(req, block_ids)
+    assert clamped == [[0, 1, 2], [10, 11], [20, 21, 22, 23, 24]]
+
+
 def _make_hybrid_kv_cache_config() -> KVCacheConfig:
     """One full-attention group + one sliding-window group (Gemma-like)."""
     full_spec = FullAttentionSpec(
@@ -731,6 +976,35 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig:
     )
 
 
+def _make_sliding_window_only_kv_cache_config() -> KVCacheConfig:
+    """A uniform sliding-window model (e.g. Mistral-7B): the single group is a
+    SlidingWindowSpec, so HMA is not required once hybrid management is
+    disabled, yet the spec guard still applies."""
+    sw_spec = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=32,
+    )
+    num_blocks = 2
+    page = sw_spec.page_size_bytes
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_blocks * page,
+                layers=["sw0"],
+                layer_stride=num_blocks * page,
+                block_stride=page,
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec),
+        ],
+    )
+
+
 def _read_scheduler(
     kv_cache_config: KVCacheConfig, disable_hma: bool = False
 ) -> MoRIIOConnectorScheduler:
@@ -820,15 +1094,6 @@ def test_non_sliding_window_hybrid_is_rejected():
         _read_scheduler(config)
 
 
-def test_token_count_basis_uses_full_attention_group():
-    """Chunked-prefill token counting must use an unclipped full-attention
-    group (blocks_per_sw == 0), not a clipped sliding-window group."""
-    scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
-    # Group 0 is the full-attention group
-    assert scheduler._full_attn_group_idx == 0
-    assert scheduler._full_attn_block_size == 16
-
-
 def test_get_exchange_clipped_blocks_clips_only_sw_group():
     """get_exchange_clipped_blocks keeps the full attn group intact and clips the
     sliding-window group to its window tail."""
@@ -897,6 +1162,71 @@ def test_single_group_path_unchanged():
     assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
 
 
+class _FakeSpecConfig(SimpleNamespace):
+    """Minimal speculative_config stand-in. Any drafter-capability probe
+    (use_dflash/use_eagle/use_dspark/uses_draft_model/...) defaults to False so
+    VllmConfig.num_lookahead_tokens resolves cleanly as upstream adds methods."""
+
+    def __getattr__(self, name: str):
+        if name.startswith(("use_", "uses_")):
+            return lambda *a, **k: False
+        raise AttributeError(name)
+
+
+def _fake_speculative_config(num_speculative_tokens: int = 1) -> SimpleNamespace:
+    return _FakeSpecConfig(num_speculative_tokens=num_speculative_tokens)
+
+
+def _spec_vllm_config() -> VllmConfig:
+    vllm_config = create_vllm_config(role="kv_producer", read_mode=False)
+    vllm_config.speculative_config = _fake_speculative_config()
+    # HMA off => the WRITE+hybrid guard is skipped, so construction reaches the
+    # spec-decode guard (the only path that hits it).
+    vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True
+    return vllm_config
+
+
+def test_spec_decode_rejected_with_sliding_window_groups():
+    """MoRIIO + speculative decoding fails closed on sliding-window models,
+    since the prompt-block clamp only covers full-attention groups. The only
+    reachable path is a uniform SWA model with HMA disabled."""
+    vllm_config = _spec_vllm_config()
+    with (
+        set_current_vllm_config(vllm_config),
+        pytest.raises(NotImplementedError, match="speculative decoding"),
+    ):
+        MoRIIOConnector(
+            vllm_config,
+            KVConnectorRole.SCHEDULER,
+            _make_sliding_window_only_kv_cache_config(),
+        )
+
+
+def test_spec_decode_allowed_with_full_attention_only():
+    """Full-attention-only models (e.g. GLM) do not trip the spec guard."""
+    vllm_config = _spec_vllm_config()
+    with set_current_vllm_config(vllm_config):
+        connector = MoRIIOConnector(
+            vllm_config,
+            KVConnectorRole.SCHEDULER,
+            _make_test_kv_cache_config(),
+        )
+    assert connector.connector_scheduler is not None
+
+
+def test_spec_decode_allowed_with_sliding_window_in_read_mode():
+    """The spec guard is WRITE-mode only; READ mode is unaffected."""
+    vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
+    vllm_config.speculative_config = _fake_speculative_config()
+    with set_current_vllm_config(vllm_config):
+        connector = MoRIIOConnector(
+            vllm_config,
+            KVConnectorRole.SCHEDULER,
+            _make_hybrid_kv_cache_config(),
+        )
+    assert connector.connector_scheduler is not None
+
+
 def test_hybrid_write_mode_rejected():
     """Hybrid KV cache groups are unsupported in WRITE mode and fail closed."""
     vllm_config = create_vllm_config(role="kv_producer", read_mode=False)

tests/v1/kv_connector/unit/test_moriio_kv_layout.py

@@ -661,16 +661,34 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match):
         wrapper._handle_message(payload)
 
 
-def test_local_block_ids_longer_than_remote_raises_value_error():
-    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+def test_local_block_ids_longer_than_remote_raises():
+    # A longer local list is a genuine bug here and must fail loudly; the WRITE
+    # producer save path trims lookahead blocks before they reach this layer.
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)  # interleaved
     worker = _worker({"layer": cache}, {"layer": _full_spec()})
 
     with pytest.raises(ValueError, match="longer than remote_block_ids"):
         moriio_layout.compute_block_transfer_offsets(
-            "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks
+            "layer",
+            cache,
+            worker.layer_to_spec,
+            [1, 3, 6, 7],
+            [4, 5],
+            _remote_meta().num_blocks,
         )
 
 
+def test_shorter_local_block_ids_still_transfers_only_what_decode_allocated():
+    # READ-mode partial transfer keeps its shorter-local behavior.
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    short = moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1], [4, 5], _remote_meta().num_blocks
+    )
+    assert len(short[0]) == 1
+
+
 @pytest.mark.parametrize(
     ("local_block_ids", "remote_block_ids", "match"),
     [

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -581,12 +581,31 @@ class MoRIIOConnectorScheduler:
             else (0, self.block_size)
             for g in kv_cache_config.transfer_groups
         ]
-        # add 1 to conservatively account for boundary overlap eg window isn't fully
-        # aligned with blocks.
+        # +1 conservatively covers window/block boundary overlap.
         self.blocks_per_sw = [
             cdiv(n_tokens, block_size) + 1 if n_tokens else 0
             for n_tokens, block_size in sw_sizes_tokens
         ]
+        # Per-group block size over transfer_groups (matching blocks_per_sw and
+        # the connector's block payload), used to clamp each full-attention group
+        # to its prompt length on the WRITE save path.
+        self._group_block_sizes = [
+            getattr(g.kv_cache_spec, "block_size", self.block_size)
+            for g in kv_cache_config.transfer_groups
+        ]
+        # The WRITE-mode prompt-block clamp only covers full-attention groups, so
+        # MoRIIO + speculative decoding is not yet safe on sliding-window /
+        # hybrid models. TODO: lift once SWA + spec is validated end-to-end.
+        if (
+            self.mode == MoRIIOMode.WRITE
+            and vllm_config.speculative_config is not None
+            and any(bps != 0 for bps in self.blocks_per_sw)
+        ):
+            raise NotImplementedError(
+                "MoRIIO with speculative decoding is not yet supported for "
+                "sliding-window / hybrid models: the prompt-block clamp is only "
+                "validated for full-attention KV groups."
+            )
         # In WRITE mode and chunked prefill, we perform the write after the last chunk.
         # Hence we need to track once we are in the last chunk. We do this by computing
         # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups
@@ -1178,6 +1197,45 @@ class MoRIIOConnectorScheduler:
 
             params["do_remote_prefill"] = False
 
+    def _clamp_to_prompt_blocks(self, req: "Request", block_ids: BlockIds) -> BlockIds:
+        """Drop trailing speculative-lookahead blocks the consumer never
+        allocated, keeping each full-attention group's prompt blocks.
+
+        Args:
+            req: Request being saved.
+            block_ids: Per-KV-cache-group local block ids.
+
+        Returns:
+            ``block_ids`` with every full-attention group truncated to
+            ``ceil(num_prompt_tokens / block_size)`` leading blocks;
+            sliding-window groups are left unchanged.
+
+        """
+        clamped = list(block_ids)
+        changed = False
+        for g, group_blocks in enumerate(block_ids):
+            if self.blocks_per_sw[g] != 0:
+                continue
+            num_prompt_blocks = math.ceil(
+                req.num_prompt_tokens / self._group_block_sizes[g]
+            )
+            if len(group_blocks) > num_prompt_blocks:
+                clamped[g] = group_blocks[:num_prompt_blocks]
+                changed = True
+        return clamped if changed else block_ids
+
+    def _is_final_prefill_chunk(
+        self, req: "Request", scheduler_output: SchedulerOutput
+    ) -> bool:
+        """Whether this step completes the prompt.
+
+        ``build_connector_meta`` runs before ``_update_after_schedule``, so
+        ``num_computed_tokens`` excludes the current chunk; add this step's
+        scheduled tokens. Token progress is lookahead-safe unlike block counts.
+        """
+        num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0)
+        return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens
+
     def build_connector_meta(
         self,
         scheduler_output: SchedulerOutput,
@@ -1191,38 +1249,36 @@ class MoRIIOConnectorScheduler:
             # It places the request metadata into the saving queue.
 
             for i, req_id in enumerate(scheduler_output.scheduled_cached_reqs.req_ids):
+                # Non-disagg requests aren't registered here; skip them
+                # instead of KeyError-ing.
+                if req_id not in self._reqs_need_pending_save:
+                    continue
+                req, existing_blocks = self._reqs_need_pending_save[req_id]
+                # The final chunk may allocate no new blocks (the previous chunk
+                # reserved speculative-lookahead slack), so new_block_ids is None;
+                # still fall through to the final-chunk check below.
                 new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]
-
                 if new_block_ids is not None:
-                    # A non-disagg request (no kv_transfer_params, e.g. smoke
-                    # test) is never registered in _reqs_need_pending_save;
-                    # indexing it unconditionally would KeyError and crash the
-                    # EngineCore. Skip it silently.
-                    if req_id not in self._reqs_need_pending_save:
-                        continue
-                    req, existing_blocks = self._reqs_need_pending_save[req_id]
-                    updated_blocks = [
+                    existing_blocks = [
                         existing_blocks[g] + new_block_ids[g]
                         for g in range(len(new_block_ids))
                     ]
-                    self._reqs_need_pending_save[req_id] = (req, updated_blocks)
-                    saved_tokens = (
-                        len(updated_blocks[self._full_attn_group_idx])
-                        * self._full_attn_block_size
-                    )
-                    if saved_tokens >= req.num_prompt_tokens:
-                        # Final chunk: live kv_transfer_params may be cleared,
-                        # so prefer the snapshot from update_state_after_alloc.
-                        kv_params = self._req_kv_params.pop(
-                            req_id, req.kv_transfer_params or {}
-                        )
-                        meta.add_new_req(
-                            request_id=req_id,
-                            local_block_ids=self._reqs_need_pending_save[req_id][1],
-                            kv_transfer_params=kv_params,
-                            write_mode=True,
-                        )
-                        del self._reqs_need_pending_save[req_id]
+                    self._reqs_need_pending_save[req_id] = (req, existing_blocks)
+                if not self._is_final_prefill_chunk(req, scheduler_output):
+                    continue
+                # Live kv_transfer_params may be cleared; prefer the
+                # update_state_after_alloc snapshot.
+                kv_params = self._req_kv_params.pop(
+                    req_id, req.kv_transfer_params or {}
+                )
+                save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks)
+                meta.add_new_req(
+                    request_id=req_id,
+                    local_block_ids=save_block_ids,
+                    kv_transfer_params=kv_params,
+                    write_mode=True,
+                )
+                del self._reqs_need_pending_save[req_id]
 
         # Loop through scheduled reqs and convert to ReqMeta.
         for req_id, (req, block_ids) in self._reqs_need_recv.items():
@@ -1253,16 +1309,13 @@ class MoRIIOConnectorScheduler:
 
         for req_id, (req, block_ids) in self._reqs_need_save.items():
             kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
-            tokens_covered = (
-                len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size
-            )
-            if req.num_prompt_tokens > tokens_covered:
-                # not last chunk prefill
+            if not self._is_final_prefill_chunk(req, scheduler_output):
                 self._reqs_need_pending_save[req_id] = (req, block_ids)
                 continue
+            save_block_ids = self._clamp_to_prompt_blocks(req, block_ids)
             meta.add_new_req(
                 request_id=req_id,
-                local_block_ids=block_ids,
+                local_block_ids=save_block_ids,
                 kv_transfer_params=kv_params,
                 write_mode=True,
             )

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py

@@ -377,10 +377,8 @@ def iter_layer_registration_regions(
     )
     region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
     if geometry.regions_per_block == 1:
-        # With padded or interleaved pages the block stride exceeds the
-        # meaningful block_len; register the strided span. The span ends with
-        # the last block's meaningful bytes, not a whole stride, so it never
-        # runs past the backing allocation of a strided layer view.
+        # Padded/interleaved pages stride past block_len: register the strided
+        # span, ending at the last block's bytes so it stays in-allocation.
         block_stride_bytes = geometry.block_stride * kv_cache.element_size()
         region_len = max(
             region_len,
@@ -431,12 +429,8 @@ def compute_block_transfer_offsets(
         [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]
     ] = merge_contiguous_offsets,
 ) -> tuple[list[int], list[int], list[int]]:
-    # A shorter (or empty) local list is the READ-mode "drop the transfer, just
-    # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled):
-    # decode pulls fewer blocks than the prefill holds. The zip loop below pairs
-    # local[i]<->remote[i] and sizes by len(local), so a short local transfers
-    # only what decode allocated and an empty local is a no-op. A longer local
-    # list is a genuine bug and still fails loudly.
+    # A shorter/empty local list transfers only what decode allocated (READ-mode
+    # partial or full-prefix-hit); a longer local list is a genuine bug.
     if len(local_block_ids) > len(remote_block_ids):
         raise ValueError(
             "local_block_ids longer than remote_block_ids: "