diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index f95e12c601..cb33bfbcf0 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/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) diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 9860a7b7e1..5ec6fb5755 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/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"), [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 5a8a917664..2d17d7a7ac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/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, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index b00692dc51..8c55837dcd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/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: "