| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | import vllm.envs as envs |
| 8 | from vllm.compilation.decorators import support_torch_compile |
| 9 | from vllm.config import ( |
| 10 | CompilationConfig, |
| 11 | ModelConfig, |
| 12 | VllmConfig, |
| 13 | set_current_vllm_config, |
| 14 | ) |
| 15 | from vllm.config.compilation import CompilationMode, CUDAGraphMode |
| 16 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | DEVICE_TYPE = current_platform.device_type |
| 20 | |
| 21 | |
| 22 | @support_torch_compile |
| 23 | class RotaryEmbeddingCompileModule(torch.nn.Module): |
| 24 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 25 | super().__init__() |
| 26 | self.rotary_emb = get_rope( |
| 27 | head_size=32, |
| 28 | max_position=128, |
| 29 | dtype=torch.float32, |
| 30 | rope_parameters={"rope_type": "default", "rope_theta": 10000}, |
| 31 | is_neox_style=True, |
| 32 | ) |
| 33 | |
| 34 | def forward( |
| 35 | self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor |
| 36 | ) -> torch.Tensor: |
| 37 | q_rot, k_rot = self.rotary_emb(positions, query, key) |
| 38 | return q_rot + k_rot |
| 39 | |
| 40 | |
| 41 | @pytest.mark.skipif(current_platform.is_cpu(), reason="Requires GPU for torch.compile") |
| 42 | def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch): |
| 43 | # Ensure env toggles take effect for this test only. |
| 44 | # The bytecode hook is required to detect buffer mutation in compiled code, |
| 45 | # and AOT compile bypasses that hook entirely. |
| 46 | envs.disable_envs_cache() |
| 47 | monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1") |
| 48 | monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0") |
| 49 | |
| 50 | device = DEVICE_TYPE |
| 51 | positions = torch.arange(16, device=device) |
| 52 | query = torch.randn(16, 32, device=device, dtype=torch.bfloat16) |
| 53 | key = torch.randn(16, 32, device=device, dtype=torch.bfloat16) |
| 54 | |
| 55 | vllm_config = VllmConfig( |
| 56 | model_config=ModelConfig(dtype=torch.bfloat16), |
| 57 | compilation_config=CompilationConfig( |
| 58 | mode=CompilationMode.VLLM_COMPILE, |
| 59 | backend="inductor", |
| 60 | custom_ops=["+rotary_embedding"], |
| 61 | cudagraph_mode=CUDAGraphMode.NONE, |
| 62 | cudagraph_num_of_warmups=0, |
| 63 | ), |
| 64 | ) |
| 65 | |
| 66 | with set_current_vllm_config(vllm_config): |
| 67 | model = RotaryEmbeddingCompileModule(vllm_config=vllm_config) |
| 68 | model(positions, query, key) |
| 69 | assert model._compiled_bytecode is not None |
| 70 | assert "update" not in model._compiled_bytecode.co_names |