From c3805ea3739f0dcca1e1f8f685d55666e547a433 Mon Sep 17 00:00:00 2001 From: yutongwu Date: Wed, 4 Nov 2026 07:40:08 +0000 Subject: [PATCH] [Bugfix][UMBP] Store the attention page of Mamba hybrid boundaries Core hands the connector only the Mamba state page of a prompt-tail boundary. When the boundary falls inside an attention block, resuming the prefix there also needs that block's page, which UMBP never stored, so a later turn could only resume from an older boundary inside fully stored blocks. Store the attention page with each handed-off boundary. With EAGLE's block drop the lookup coordinator matches attention one lookup unit past a candidate before dropping that unit, so the page is keyed one unit after the boundary, where core registers its own partial tail. The request-finished path now keys its attention page the same way. Signed-off-by: yutongwu --- .../umbp_integration/test_mori_token_e2e.py | 50 +++++++ .../v1/kv_connector/unit/test_umbp_hybrid.py | 135 ++++++++++++++++++ .../kv_connector/v1/umbp/scheduler.py | 125 +++++++++++++--- 3 files changed, 294 insertions(+), 16 deletions(-) diff --git a/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py b/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py index ed25edf076fb..13dd0ccc373a 100644 --- a/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py +++ b/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) diff --git a/tests/v1/kv_connector/unit/test_umbp_hybrid.py b/tests/v1/kv_connector/unit/test_umbp_hybrid.py index 6d211f2a68a8..a3b2c28a85bb 100644 --- a/tests/v1/kv_connector/unit/test_umbp_hybrid.py +++ b/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", [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py index a7e88e3a4308..73b4b1249750 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py +++ b/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.