| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for GPUModelRunner._gather_mm_embeddings (model runner V1). |
| 4 | |
| 5 | Mirrors tests/v1/worker/test_encoder_runner.py (the V2 runner): the EAGLE/MTP |
| 6 | drafter reads one position ahead of the target (shift_computed_tokens=1). The |
| 7 | +1 look-ahead feature past the processed boundary is used when its encoder |
| 8 | output is present and tolerated (token-embedding fallback) when it is not, |
| 9 | while a miss within the processed range still fails loudly. |
| 10 | |
| 11 | `_gather_mm_embeddings` only uses CPU-side state, so it is exercised against a |
| 12 | lightweight stub for `self` instead of a full (CUDA-only) runner. |
| 13 | """ |
| 14 | |
| 15 | from types import SimpleNamespace |
| 16 | |
| 17 | import pytest |
| 18 | import torch |
| 19 | |
| 20 | from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange |
| 21 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner |
| 22 | |
| 23 | pytestmark = pytest.mark.cpu_test |
| 24 | |
| 25 | HIDDEN = 4 |
| 26 | |
| 27 | |
| 28 | def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec: |
| 29 | return MultiModalFeatureSpec( |
| 30 | data=None, |
| 31 | modality="image", |
| 32 | identifier=identifier, |
| 33 | mm_position=PlaceholderRange(offset=offset, length=length), |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | def _gather(features, cached, *, num_scheduled, shift, num_computed=0): |
| 38 | encoder_cache = { |
| 39 | f.identifier: torch.arange( |
| 40 | f.mm_position.length * HIDDEN, dtype=torch.float32 |
| 41 | ).reshape(f.mm_position.length, HIDDEN) |
| 42 | for f in cached |
| 43 | } |
| 44 | req_state = SimpleNamespace(num_computed_tokens=num_computed, mm_features=features) |
| 45 | runner = SimpleNamespace( |
| 46 | input_batch=SimpleNamespace(req_ids=["req0"]), |
| 47 | requests={"req0": req_state}, |
| 48 | encoder_cache=encoder_cache, |
| 49 | _get_encoder_output_from_cache=lambda mm_hash: encoder_cache.get(mm_hash), |
| 50 | is_multimodal_pruning_enabled=False, |
| 51 | uses_mrope=False, |
| 52 | ) |
| 53 | scheduler_output = SimpleNamespace( |
| 54 | total_num_scheduled_tokens=num_scheduled, |
| 55 | num_scheduled_tokens={"req0": num_scheduled}, |
| 56 | ) |
| 57 | return GPUModelRunner._gather_mm_embeddings( |
| 58 | runner, scheduler_output, shift_computed_tokens=shift |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def test_draft_shift_uses_boundary_feature_when_cached(): |
| 63 | """The drafter's +1 look-ahead reaches the feature at offset == |
| 64 | processed_end; when it is already cached it is used for the look-ahead |
| 65 | position rather than ignored.""" |
| 66 | f0 = _feature("h0", offset=0, length=8) |
| 67 | f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end |
| 68 | mm_embeds, is_mm_embed = _gather([f0, f1], [f0, f1], num_scheduled=8, shift=1) |
| 69 | |
| 70 | # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7. |
| 71 | assert len(mm_embeds) == 2 |
| 72 | assert bool(is_mm_embed[7]) |
| 73 | assert int(is_mm_embed.sum()) == 8 |
| 74 | |
| 75 | |
| 76 | def test_draft_shift_tolerates_missing_boundary_feature(): |
| 77 | """When the +1 look-ahead feature past the processed boundary is not yet |
| 78 | encoded, fall back to the token embedding instead of raising.""" |
| 79 | f0 = _feature("h0", offset=0, length=8) |
| 80 | f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached |
| 81 | mm_embeds, is_mm_embed = _gather([f0, f1], [f0], num_scheduled=8, shift=1) |
| 82 | |
| 83 | assert len(mm_embeds) == 1 # only f0; f1's boundary position falls back |
| 84 | assert not bool(is_mm_embed[7]) |
| 85 | assert int(is_mm_embed.sum()) == 7 |
| 86 | |
| 87 | |
| 88 | def test_draft_shift_raises_on_interior_miss(): |
| 89 | """A miss for a feature within the processed range (not the look-ahead |
| 90 | boundary) is a real invariant violation, even on the drafter path.""" |
| 91 | f0 = _feature("h0", offset=0, length=8) # interior, within processed range |
| 92 | with pytest.raises(RuntimeError, match="Encoder cache miss"): |
| 93 | _gather([f0], [], num_scheduled=8, shift=1) |
| 94 | |
| 95 | |
| 96 | def test_target_path_raises_on_encoder_cache_miss(): |
| 97 | """On the target path (no shift) a miss is a real invariant violation.""" |
| 98 | f0 = _feature("h0", offset=0, length=8) |
| 99 | with pytest.raises(RuntimeError, match="Encoder cache miss"): |
| 100 | _gather([f0], [], num_scheduled=8, shift=0) |