| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test that M-RoPE position initialization handles prompt_embeds-only inputs. |
| 4 | |
| 5 | Regression test for GHSA-33cg-gxv8-3p8g: sending /v1/completions with |
| 6 | prompt_embeds and no prompt_token_ids on M-RoPE models crashed the |
| 7 | EngineCore via an assertion failure. |
| 8 | """ |
| 9 | |
| 10 | from unittest.mock import Mock |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm.model_executor.models.interfaces import SupportsMRoPE |
| 16 | from vllm.v1.worker.gpu_input_batch import CachedRequestState |
| 17 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner |
| 18 | |
| 19 | |
| 20 | class FakeMRoPEModel(SupportsMRoPE): |
| 21 | """Minimal model that passes supports_mrope() check.""" |
| 22 | |
| 23 | def get_mrope_input_positions(self, input_tokens, mm_features): |
| 24 | seq_len = len(input_tokens) |
| 25 | positions = torch.arange(seq_len).unsqueeze(0).expand(3, -1) |
| 26 | return positions.clone(), 0 |
| 27 | |
| 28 | |
| 29 | def _make_runner_and_req(prompt_token_ids, prompt_embeds): |
| 30 | """Create a minimal GPUModelRunner instance and request state.""" |
| 31 | model = FakeMRoPEModel() |
| 32 | instance = object.__new__(GPUModelRunner) |
| 33 | instance.get_model = lambda: model |
| 34 | |
| 35 | req_state = Mock(spec=CachedRequestState) |
| 36 | req_state.prompt_token_ids = prompt_token_ids |
| 37 | req_state.prompt_embeds = prompt_embeds |
| 38 | req_state.mm_features = [] |
| 39 | req_state.mrope_positions = None |
| 40 | req_state.mrope_position_delta = None |
| 41 | return instance, req_state |
| 42 | |
| 43 | |
| 44 | class TestMRopePromptEmbeds: |
| 45 | """Verify _init_mrope_positions handles prompt_embeds-only inputs.""" |
| 46 | |
| 47 | def test_prompt_embeds_only_does_not_crash(self): |
| 48 | """Prompt-embeds-only request must not raise AssertionError.""" |
| 49 | instance, req_state = _make_runner_and_req( |
| 50 | prompt_token_ids=None, |
| 51 | prompt_embeds=torch.randn(15, 896), |
| 52 | ) |
| 53 | |
| 54 | instance._init_mrope_positions(req_state) |
| 55 | |
| 56 | assert req_state.mrope_positions is not None |
| 57 | assert req_state.mrope_positions.shape == (3, 15) |
| 58 | |
| 59 | def test_prompt_token_ids_still_works(self): |
| 60 | """Normal path with prompt_token_ids continues working.""" |
| 61 | instance, req_state = _make_runner_and_req( |
| 62 | prompt_token_ids=[1, 2, 3, 4, 5], |
| 63 | prompt_embeds=None, |
| 64 | ) |
| 65 | |
| 66 | instance._init_mrope_positions(req_state) |
| 67 | |
| 68 | assert req_state.mrope_positions is not None |
| 69 | assert req_state.mrope_positions.shape == (3, 5) |
| 70 | |
| 71 | def test_neither_token_ids_nor_embeds_raises(self): |
| 72 | """When both are None, a ValueError should be raised.""" |
| 73 | instance, req_state = _make_runner_and_req( |
| 74 | prompt_token_ids=None, |
| 75 | prompt_embeds=None, |
| 76 | ) |
| 77 | |
| 78 | with pytest.raises(ValueError, match="prompt_token_ids or prompt_embeds"): |
| 79 | instance._init_mrope_positions(req_state) |