[Bugfix][Core] Keep every multimodal feature in the partial-block KV event (#58288)
Signed-off-by: haosenwang1018 <haosenwang1018@users.noreply.github.com> Co-authored-by: haosenwang1018 <haosenwang1018@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Sense_wang committed
Sep 25, 2026 at 07:25 UTC
f6aa2919bc4afb93e286be55dbe2a4b1d778fbc3
2 files changed
+78
-1
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py
+69
@@ -7,6 +7,11 @@ import pytest
7
8
import vllm.v1.core.kv_cache_utils as kv_cache_utils
9
from vllm.distributed.kv_events import BlockRemoved, BlockStored
10
+from vllm.multimodal.inputs import (
11
+ MultiModalFeatureSpec,
12
+ MultiModalKwargsItem,
13
+ PlaceholderRange,
14
+)
15
from vllm.sampling_params import SamplingParams
16
from vllm.utils.hashing import sha256
17
from vllm.v1.core.block_pool import BlockPool
@@ -34,12 +39,28 @@ def make_request(
39
hash_block_size: int,
40
hash_fn: Callable,
41
session_id: str | None = None,
42
+ mm_positions: list[PlaceholderRange] | None = None,
43
+ mm_hashes: list[str] | None = None,
44
) -> Request:
45
+ mm_features = []
46
+ if mm_positions is not None:
47
+ for j, position in enumerate(mm_positions):
48
+ identifier = mm_hashes[j] if mm_hashes else f"hash_{j}"
49
+ mm_features.append(
50
+ MultiModalFeatureSpec(
51
+ data=MultiModalKwargsItem.dummy(),
52
+ mm_position=position,
53
+ identifier=identifier,
54
+ modality="image",
55
+ )
56
+ )
57
+
58
sampling_params = SamplingParams(max_tokens=17)
59
sampling_params.update_from_generation_config({}, eos_token_id=100)
60
return Request(
61
request_id=request_id,
62
prompt_token_ids=prompt_token_ids,
63
+ mm_features=mm_features if mm_features else None,
64
sampling_params=sampling_params,
65
pooling_params=None,
66
block_hasher=get_request_block_hasher(hash_block_size, hash_fn),
@@ -173,6 +194,54 @@ def test_cache_partial_block_kv_cache_events():
194
assert removed_event.group_idx == kv_cache_group_id
195
196
197
+def test_cache_partial_block_event_keeps_every_mm_feature():
198
+ """The event window can reach into more than one multimodal feature.
199
+
200
+ ``cache_partial_block`` builds the ``BlockStored`` extra keys over
201
+ ``[(num_hash_blocks - 1) * hash_block_size, num_tokens)``. Starting that
202
+ scan at the last feature drops any earlier one that also reaches into the
203
+ window, so two requests differing only in that feature would be published
204
+ under the same key.
205
+ """
206
+ hash_block_size = 4
207
+ block_size = 8
208
+ num_tokens = 12
209
+ # The event window is tokens [8, 12): "A" ends inside it and "B" sits
210
+ # entirely within it.
211
+ mm_positions = [
212
+ PlaceholderRange(offset=6, length=3),
213
+ PlaceholderRange(offset=10, length=2),
214
+ ]
215
+
216
+ pool = BlockPool(
217
+ num_gpu_blocks=2,
218
+ enable_caching=True,
219
+ hash_block_size=hash_block_size,
220
+ enable_kv_cache_events=True,
221
+ )
222
+ req = make_request(
223
+ "req_partial_mm",
224
+ prompt_token_ids=list(range(num_tokens)),
225
+ hash_block_size=hash_block_size,
226
+ hash_fn=sha256,
227
+ mm_positions=mm_positions,
228
+ mm_hashes=["A", "B"],
229
+ )
230
+
231
+ assert pool.cache_partial_block(
232
+ request=req,
233
+ block=pool.get_new_blocks(1)[0],
234
+ num_tokens=num_tokens,
235
+ kv_cache_group_id=0,
236
+ block_size=block_size,
237
+ )
238
+
239
+ events = pool.take_events()
240
+ assert len(events) == 1
241
+ assert isinstance(events[0], BlockStored)
242
+ assert events[0].extra_keys == [(("A", -2), ("B", 2))]
243
+
244
+
245
def test_partial_block_replacement_emits_remove_then_store_events():
246
hash_block_size = 2
247
block_size = 6
vllm/v1/core/block_pool.py
+9
-1
@@ -11,6 +11,7 @@ from vllm.distributed.kv_events import (
11
KVCacheEvent,
12
)
13
from vllm.logger import init_logger
14
+from vllm.multimodal.utils import get_mm_features_in_window
15
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
16
from vllm.v1.core.kv_cache_utils import (
17
BlockHash,
@@ -532,7 +533,14 @@ class BlockPool:
533
else None
534
)
535
block_end = num_tokens
535
- curr_mm_idx = -1 if block_start > 0 else 0
536
+ curr_mm_idx = 0
537
+ mm_features = request.mm_features
538
+ if block_start > 0 and mm_features:
539
+ last_mm_pos = mm_features[-1].mm_position
540
+ if last_mm_pos.offset + last_mm_pos.length > block_start:
541
+ curr_mm_idx, _ = get_mm_features_in_window(
542
+ mm_features, block_start, block_end
543
+ )
544
extra_keys, _ = generate_block_hash_extra_keys(
545
request, block_start, block_end, curr_mm_idx
546
)