[Bugfix][UMBP] Store the attention page of Mamba hybrid boundaries

tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py

@@ -173,3 +173,53 @@ def test_restore_under_pool_eviction(umbp_llm):
     before = _external_tokens(umbp_llm)
     umbp_llm.generate(prompts[0], PARAMS, use_tqdm=False)
     assert _external_tokens(umbp_llm) == before
+
+
+HYBRID_MODEL = "ibm-granite/granite-4.0-h-350m"
+HYBRID_REVISION = "3b17b717b8f2f5d305b0a92c1491e239aeda19c8"
+
+
+@pytest.fixture
+def umbp_hybrid_llm(vllm_runner, socket_dir):
+    with vllm_runner(
+        HYBRID_MODEL,
+        revision=HYBRID_REVISION,
+        enforce_eager=True,
+        max_model_len=2048,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=True,
+        prefix_match_unit=16,
+        disable_log_stats=False,
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="UMBPStoreConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "mode": "embedded",
+                "capacity_bytes": 1 << 30,
+                "lookup_dir": socket_dir,
+                "enable_partial_hash_hits": True,
+            },
+        ),
+    ) as runner:
+        yield runner.llm
+
+
+def test_hybrid_restores_a_boundary_inside_an_attention_block(umbp_hybrid_llm):
+    """A Mamba checkpoint inside an attention block restores like a GPU hit.
+
+    Mamba hybrids have no batch-invariant kernels, so the reference is a local
+    prefix hit rather than a cold run.
+    """
+    llm = umbp_hybrid_llm
+    block_size = llm.llm_engine.vllm_config.cache_config.block_size
+    prompt = _prompt(0, 2 * block_size + 40)
+    llm.generate(prompt, PARAMS, use_tqdm=False)
+    local_hit = llm.generate(prompt, PARAMS, use_tqdm=False)
+    assert local_hit[0].num_cached_tokens == 2 * block_size + 32
+    _reset_gpu_prefix_cache(llm)
+
+    before = _external_tokens(llm)
+    restored = llm.generate(prompt, PARAMS, use_tqdm=False)
+
+    assert _external_tokens(llm) - before == 2 * block_size + 32
+    _assert_same_outputs(local_hit, restored)

tests/v1/kv_connector/unit/test_umbp_hybrid.py

@@ -772,6 +772,94 @@ def test_scheduler_stores_exact_hybrid_boundary_state():
     assert plan.key == scheduler.codec.key(b"a", 1)
 
 
+@pytest.mark.parametrize(
+    ("use_eagle", "prompt_tokens", "page_hash_index"),
+    [(False, 22, 4), (True, 26, 5)],
+)
+def test_boundary_inside_attention_block_stores_its_attention_page(
+    monkeypatch, use_eagle, prompt_tokens, page_hash_index
+):
+    """Core hands off only the state page of a boundary; resuming inside an
+    attention block also needs that block's page. EAGLE matches attention one
+    hash unit past the boundary, so its page is keyed there."""
+    monkeypatch.setattr(
+        umbp_scheduler, "resolve_kv_cache_block_sizes", lambda *args: (16, 4)
+    )
+    hashes = [bytes([index + 1]) * 4 for index in range(8)]
+    codec = BlockIdentityCodec(UMBPNamespace("attention-page"))
+
+    def make_scheduler(hits):
+        vllm_config = _vllm_config(
+            {"mode": "embedded", "enable_partial_hash_hits": True},
+            prefix_match_unit=4,
+        )
+        if use_eagle:
+            vllm_config.speculative_config = SimpleNamespace(
+                use_eagle_block_drop=lambda: True
+            )
+        return UMBPStoreConnectorScheduler(
+            vllm_config, _hybrid_kv_cache_config(), _SchedulerHandle(hits), codec
+        )
+
+    first = make_scheduler({})
+    request = make_request(
+        "turn-1",
+        num_tokens=prompt_tokens,
+        block_hashes=hashes[: prompt_tokens // 4],
+        block_ids=([3, 4], [6, 7]),
+    )
+    try:
+        first.update_state_after_alloc(
+            request,
+            SimpleNamespace(get_block_ids=lambda group_ids: ([3, 4], [6, 7])),
+            0,
+        )
+        metadata = first.build_connector_meta(
+            make_scheduler_output(
+                new_reqs=[
+                    SimpleNamespace(
+                        req_id="turn-1",
+                        num_computed_tokens=0,
+                        block_ids=([3, 4], [6, 7]),
+                    )
+                ],
+                scheduled_tokens={"turn-1": prompt_tokens},
+                block_state=SimpleNamespace(
+                    boundary_state_offloads={"turn-1": [(1, 9, 20)]}
+                ),
+            )
+        )
+    finally:
+        first.close()
+    stored = {
+        plan.key: (plan.group_id, plan.block_id) for plan in metadata.iter_store_plans()
+    }
+    # The full block [0, 16), the attention page holding [16, 20), the state.
+    assert stored == {
+        codec.key(hashes[3], 0): (0, 3),
+        codec.key(hashes[page_hash_index], 0): (0, 4),
+        codec.key(hashes[4], 1): (1, 9),
+    }
+
+    second = make_scheduler(dict.fromkeys(stored, True))
+    try:
+        next_turn = make_request("turn-2", num_tokens=32, block_hashes=hashes)
+        assert second.get_num_new_matched_tokens(next_turn, 0) == (20, True)
+        second.update_state_after_alloc(
+            next_turn,
+            SimpleNamespace(get_block_ids=lambda group_ids: ([11, 12], [13, 14])),
+            20,
+        )
+        loads = second._pending_loads["turn-2"]
+        assert dict(zip(loads.keys, loads.block_ids)) == {
+            codec.key(hashes[3], 0): 11,
+            codec.key(hashes[page_hash_index], 0): 12,
+            codec.key(hashes[4], 1): 14,
+        }
+    finally:
+        second.close()
+
+
 @pytest.mark.parametrize("local_tokens", [0, 16])
 def test_scheduler_partial_prefix_loads_full_page_at_absolute_index(
     local_tokens, monkeypatch
@@ -897,6 +985,53 @@ def test_finished_hybrid_tail_pins_full_pages_until_all_workers_finish():
     assert not scheduler.has_pending_push_work()
 
 
+@pytest.mark.parametrize("use_eagle", [False, True])
+def test_finished_hybrid_tail_keys_its_attention_page_like_core(use_eagle):
+    """With EAGLE the attention page of a boundary is matched one hash unit
+    past it, where core registers its own partial tail."""
+    vllm_config = _vllm_config(
+        {"mode": "embedded", "load_async": False, "enable_partial_hash_hits": True},
+        prefix_match_unit=4,
+    )
+    if use_eagle:
+        vllm_config.speculative_config = SimpleNamespace(
+            use_eagle_block_drop=lambda: True
+        )
+    scheduler = UMBPStoreConnectorScheduler(
+        vllm_config,
+        _hybrid_kv_cache_config(),
+        _SchedulerHandle([]),
+        BlockIdentityCodec(UMBPNamespace("finished-tail-page")),
+    )
+    request = make_request(
+        "finished",
+        num_tokens=16,
+        block_hashes=[b"a", b"b", b"c", b"d"],
+        block_ids=([7], [9]),
+        num_computed_tokens=8,
+    )
+    try:
+        scheduler.update_state_after_alloc(
+            request, SimpleNamespace(get_block_ids=lambda group_ids: ([7], [9])), 0
+        )
+        scheduler.bind_gpu_block_pool(BlockPool(16, True, 4))
+        assert not scheduler.register_finished_partial_tail(
+            request, ([7], [9]), [(1, 9, 8)]
+        )
+        metadata = scheduler.build_connector_meta(
+            make_scheduler_output(finished={request.request_id})
+        )
+        attention_hash = b"c" if use_eagle else b"b"
+        assert {
+            (plan.group_id, plan.block_id, plan.key) for plan in metadata.store_plans
+        } == {
+            (0, 7, scheduler.codec.key(attention_hash, 0)),
+            (1, 9, scheduler.codec.key(b"b", 1)),
+        }
+    finally:
+        scheduler.close()
+
+
 @pytest.mark.parametrize(
     "device",
     [

vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py

@@ -205,6 +205,16 @@ def __init__(
             if isinstance(group.kv_cache_spec, MambaSpec)
             and getattr(group.kv_cache_spec, "mamba_cache_mode", None) == "align"
         )
+        self._attention_groups = frozenset(
+            group_id
+            for group_id in kv_cache_config.prefix_cacheable_group_ids
+            if isinstance(
+                _unwrap_store_spec(
+                    kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
+                ),
+                FullAttentionSpec,
+            )
+        )
         self.runtime = runtime
         self.enable_kv_cache_events = bool(
             vllm_config.kv_events_config
@@ -327,6 +337,18 @@ def __init__(
             )
             for group_id in kv_cache_config.prefix_cacheable_group_ids
         }
+        # The coordinator matches these groups one lookup unit past a hit and
+        # then drops that unit (EAGLE).
+        self._eagle_lookup_groups = frozenset(
+            group_id
+            for group_id, manager in zip(
+                kv_cache_config.prefix_cacheable_group_ids,
+                self._lookup_coordinator.single_type_managers
+                if self._lookup_coordinator is not None
+                else (),
+            )
+            if manager.use_eagle
+        )
         self._pending_lookups: dict[str, _PendingLookup] = {}
         self._lookup_executor = ThreadPoolExecutor(
             max_workers=int(extra.get("lookup_workers", 2)),
@@ -841,6 +863,7 @@ def _add_boundary_state_plans(
                 )
                 continue
             plans: list[BlockTransferPlan] = []
+            boundaries: set[int] = set()
             for group_id, block_id, boundary_tokens in entries:
                 if group_id not in prefix_groups or block_id == NULL_BLOCK_ID:
                     continue
@@ -865,6 +888,11 @@ def _add_boundary_state_plans(
                         block_size=self.group_block_sizes[group_id],
                     )
                 )
+                boundaries.add(boundary_tokens)
+            for boundary in sorted(boundaries):
+                plans.extend(
+                    self._attention_page_plans(request, tracker.block_ids, boundary)
+                )
             logger.debug(
                 "UMBP boundary offloads request=%s offered=%d stored=%d boundaries=%s",
                 request_id,
@@ -875,6 +903,65 @@ def _add_boundary_state_plans(
             if plans:
                 meta.store_plans.extend(plans)
 
+    def _attention_page_plans(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+        boundary: int,
+    ) -> list[BlockTransferPlan]:
+        """Attention pages a hybrid prefix needs to resume at ``boundary``.
+
+        Core hands off only the state pages of a boundary. When the boundary
+        falls inside an attention block, that block's page holds the tokens
+        before it, and without it the prefix cannot resume there.
+        ``block_ids`` holds one list per prefix-cacheable group.
+        """
+        assert self.codec is not None
+        plans: list[BlockTransferPlan] = []
+        for index, group_id in enumerate(
+            self.kv_cache_config.prefix_cacheable_group_ids
+        ):
+            if group_id not in self._attention_groups or index >= len(block_ids):
+                continue
+            block_size = self.group_block_sizes[group_id]
+            end = self._attention_page_end(group_id, boundary)
+            # A block-aligned end is a full block, stored with its prefix.
+            if end % block_size == 0:
+                continue
+            hash_index = end // self.hash_block_size - 1
+            if not 0 <= hash_index < len(request.block_hashes):
+                continue
+            group_block_ids = block_ids[index]
+            block_index = end // block_size
+            if (
+                block_index >= len(group_block_ids)
+                or group_block_ids[block_index] == NULL_BLOCK_ID
+            ):
+                continue
+            block_hash = request.block_hashes[hash_index]
+            plans.append(
+                BlockTransferPlan(
+                    key=self.codec.key(block_hash, group_id),
+                    block_id=group_block_ids[block_index],
+                    request_id=request.request_id,
+                    group_id=group_id,
+                    block_hash=block_hash,
+                    parent_block_hash=(
+                        request.block_hashes[hash_index - 1]
+                        if self.enable_kv_cache_events and hash_index > 0
+                        else None
+                    ),
+                    block_size=block_size,
+                )
+            )
+        return plans
+
+    def _attention_page_end(self, group_id: int, boundary: int) -> int:
+        """Token position the attention page restoring ``boundary`` is keyed at."""
+        if group_id in self._eagle_lookup_groups:
+            return boundary + self._lookup_units[group_id]
+        return boundary
+
     def _tracker_for_request(self, request_id: str) -> RequestTracker:
         return self._request_trackers.setdefault(request_id, RequestTracker())
 
@@ -1198,8 +1285,7 @@ def register_finished_partial_tail(
             or bool(getattr(request, "num_in_flight_tokens", 0))
         ):
             return False
-        hash_index = boundary // self.hash_block_size - 1
-        if hash_index >= len(request.block_hashes):
+        if boundary // self.hash_block_size > len(request.block_hashes):
             return False
         sources = {
             group_id: block_id for group_id, block_id, _ in partial_tail_offloads
@@ -1208,31 +1294,38 @@ def register_finished_partial_tail(
             return False
         # Core hands off Mamba state pages. Include the matching FA tail page
         # so every group can resolve the same fine-grained prefix boundary.
-        for group_id in prefix_groups:
-            spec = self.kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
-            if (
-                isinstance(spec, FullAttentionSpec)
-                and boundary % self.group_block_sizes[group_id]
-            ):
-                index = boundary // self.group_block_sizes[group_id]
-                if index >= len(block_ids[group_id]):
-                    return False
-                sources.setdefault(group_id, block_ids[group_id][index])
+        key_ends = dict.fromkeys(sources, boundary)
+        for group_id in self.kv_cache_config.prefix_cacheable_group_ids:
+            if group_id not in self._attention_groups:
+                continue
+            end = self._attention_page_end(group_id, boundary)
+            group_block_size = self.group_block_sizes[group_id]
+            if end % group_block_size == 0:
+                continue
+            index = end // group_block_size
+            if index >= len(block_ids[group_id]):
+                return False
+            if end // self.hash_block_size > len(request.block_hashes):
+                return False
+            sources.setdefault(group_id, block_ids[group_id][index])
+            key_ends.setdefault(group_id, end)
         plans: list[BlockTransferPlan] = []
         for group_id, block_id in sources.items():
             if block_id == NULL_BLOCK_ID or not 0 <= block_id < len(pool.blocks):
                 return False
             if pool.blocks[block_id].is_null:
                 return False
-            group_block_size = self.group_block_sizes[group_id]
+            block_hash = request.block_hashes[
+                key_ends[group_id] // self.hash_block_size - 1
+            ]
             plans.append(
                 BlockTransferPlan(
                     request_id=request.request_id,
                     block_id=block_id,
                     group_id=group_id,
-                    key=self.codec.key(request.block_hashes[hash_index], group_id),
-                    block_hash=request.block_hashes[hash_index],
-                    block_size=group_block_size,
+                    key=self.codec.key(block_hash, group_id),
+                    block_hash=block_hash,
+                    block_size=self.group_block_sizes[group_id],
                 )
             )
         # A Mamba page is a state snapshot, not a token-addressable byte array.