| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ApplyRotaryEmb CustomOp dispatch behavior. |
| 4 | |
| 5 | This test ensures that RotaryEmbedding classes correctly call the appropriate |
| 6 | ApplyRotaryEmb methods based on the calling context: |
| 7 | |
| 8 | 1. RotaryEmbedding.forward_native() -> ApplyRotaryEmb.forward_native() |
| 9 | 2. RotaryEmbedding.forward_cuda() -> ApplyRotaryEmb.forward() (auto-dispatch) |
| 10 | 3. RotaryEmbedding.forward_hip() -> ApplyRotaryEmb.forward() (auto-dispatch) |
| 11 | """ |
| 12 | |
| 13 | from dataclasses import dataclass |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.config import ( |
| 19 | CompilationConfig, |
| 20 | VllmConfig, |
| 21 | get_cached_compilation_config, |
| 22 | set_current_vllm_config, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | |
| 26 | CUDA_DEVICES = ["cuda:0"] |
| 27 | |
| 28 | |
| 29 | @dataclass |
| 30 | class RotaryEmbeddingTestCase: |
| 31 | """Test case configuration for RotaryEmbedding dispatch tests.""" |
| 32 | |
| 33 | name: str |
| 34 | rope_class: type |
| 35 | rope_kwargs: dict |
| 36 | method_name: str # forward_native, forward_cuda, forward |
| 37 | positions_shape: tuple # (num_tokens,) or (3, num_tokens) or (4, num_tokens) |
| 38 | expect_forward_native: bool # Should call ApplyRotaryEmb.forward_native() |
| 39 | expect_forward: bool # Should call ApplyRotaryEmb.forward() |
| 40 | |
| 41 | |
| 42 | def get_test_cases() -> list[RotaryEmbeddingTestCase]: |
| 43 | """Generate test cases for all RotaryEmbedding classes.""" |
| 44 | from vllm.model_executor.layers.rotary_embedding.ernie45_vl_rope import ( |
| 45 | Ernie4_5_VLRotaryEmbedding, |
| 46 | ) |
| 47 | from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding |
| 48 | |
| 49 | common_kwargs = { |
| 50 | "head_size": 128, |
| 51 | "rotary_dim": 128, |
| 52 | "max_position_embeddings": 4096, |
| 53 | "base": 10000, |
| 54 | "is_neox_style": True, |
| 55 | "dtype": torch.bfloat16, |
| 56 | } |
| 57 | |
| 58 | return [ |
| 59 | # MRotaryEmbedding tests |
| 60 | RotaryEmbeddingTestCase( |
| 61 | name="MRotaryEmbedding.forward_native", |
| 62 | rope_class=MRotaryEmbedding, |
| 63 | rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]}, |
| 64 | method_name="forward_native", |
| 65 | positions_shape=(3, 32), # 2D for multimodal |
| 66 | expect_forward_native=True, |
| 67 | expect_forward=False, |
| 68 | ), |
| 69 | RotaryEmbeddingTestCase( |
| 70 | name="MRotaryEmbedding.forward_cuda_1d", |
| 71 | rope_class=MRotaryEmbedding, |
| 72 | rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]}, |
| 73 | method_name="forward_cuda", |
| 74 | positions_shape=(32,), # 1D triggers apply_rotary_emb path |
| 75 | expect_forward_native=False, |
| 76 | expect_forward=True, |
| 77 | ), |
| 78 | # Ernie4_5_VLRotaryEmbedding tests |
| 79 | RotaryEmbeddingTestCase( |
| 80 | name="Ernie4_5_VLRotaryEmbedding.forward_native", |
| 81 | rope_class=Ernie4_5_VLRotaryEmbedding, |
| 82 | rope_kwargs={**common_kwargs, "mrope_section": [22, 22, 20]}, |
| 83 | method_name="forward_native", |
| 84 | positions_shape=(3, 32), # 2D for multimodal |
| 85 | expect_forward_native=True, |
| 86 | expect_forward=False, |
| 87 | ), |
| 88 | ] |
| 89 | |
| 90 | |
| 91 | def run_dispatch_test( |
| 92 | test_case: RotaryEmbeddingTestCase, |
| 93 | device: str, |
| 94 | ): |
| 95 | """Run a dispatch test for a RotaryEmbedding class.""" |
| 96 | vllm_config = VllmConfig( |
| 97 | compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"]) |
| 98 | ) |
| 99 | get_cached_compilation_config.cache_clear() |
| 100 | |
| 101 | with set_current_vllm_config(vllm_config): |
| 102 | rope = test_case.rope_class(**test_case.rope_kwargs).to(device=device) |
| 103 | |
| 104 | apply_rotary_emb = rope.apply_rotary_emb |
| 105 | |
| 106 | # Verify custom op is enabled |
| 107 | if test_case.expect_forward_native: |
| 108 | assert ( |
| 109 | apply_rotary_emb._forward_method != apply_rotary_emb.forward_native |
| 110 | ), "Test setup error: ApplyRotaryEmb custom op should be enabled" |
| 111 | |
| 112 | # Setup call tracking |
| 113 | call_tracker = {"forward_native_called": False, "forward_called": False} |
| 114 | original_forward_native = apply_rotary_emb.forward_native |
| 115 | original_forward = apply_rotary_emb.forward |
| 116 | |
| 117 | def tracked_forward_native(*args, **kwargs): |
| 118 | call_tracker["forward_native_called"] = True |
| 119 | return original_forward_native(*args, **kwargs) |
| 120 | |
| 121 | def tracked_forward(*args, **kwargs): |
| 122 | call_tracker["forward_called"] = True |
| 123 | return original_forward(*args, **kwargs) |
| 124 | |
| 125 | apply_rotary_emb.forward_native = tracked_forward_native |
| 126 | apply_rotary_emb.forward = tracked_forward |
| 127 | |
| 128 | try: |
| 129 | num_tokens = test_case.positions_shape[-1] |
| 130 | num_q_heads = 8 |
| 131 | num_kv_heads = 2 |
| 132 | head_size = test_case.rope_kwargs["head_size"] |
| 133 | max_position = test_case.rope_kwargs["max_position_embeddings"] |
| 134 | |
| 135 | positions = torch.randint( |
| 136 | 0, max_position // 4, test_case.positions_shape, device=device |
| 137 | ) |
| 138 | query = torch.randn( |
| 139 | num_tokens, num_q_heads * head_size, dtype=torch.bfloat16, device=device |
| 140 | ) |
| 141 | key = torch.randn( |
| 142 | num_tokens, |
| 143 | num_kv_heads * head_size, |
| 144 | dtype=torch.bfloat16, |
| 145 | device=device, |
| 146 | ) |
| 147 | |
| 148 | # Call the method under test |
| 149 | method = getattr(rope, test_case.method_name) |
| 150 | method(positions, query.clone(), key.clone()) |
| 151 | |
| 152 | # Verify expectations |
| 153 | if test_case.expect_forward_native: |
| 154 | assert call_tracker["forward_native_called"], ( |
| 155 | f"{test_case.name} should call ApplyRotaryEmb.forward_native()" |
| 156 | ) |
| 157 | if not test_case.expect_forward: |
| 158 | assert not call_tracker["forward_called"], ( |
| 159 | f"{test_case.name} should NOT call ApplyRotaryEmb.forward(). " |
| 160 | "Bug: when +apply_rotary_emb is enabled, forward_native() " |
| 161 | "incorrectly dispatches to CUDA/HIP kernels." |
| 162 | ) |
| 163 | if test_case.expect_forward: |
| 164 | assert call_tracker["forward_called"], ( |
| 165 | f"{test_case.name} should call ApplyRotaryEmb.forward()" |
| 166 | ) |
| 167 | finally: |
| 168 | apply_rotary_emb.forward_native = original_forward_native |
| 169 | apply_rotary_emb.forward = original_forward |
| 170 | |
| 171 | |
| 172 | @pytest.mark.skipif( |
| 173 | not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests." |
| 174 | ) |
| 175 | @pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda tc: tc.name) |
| 176 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 177 | def test_rotary_embedding_dispatch( |
| 178 | test_case: RotaryEmbeddingTestCase, |
| 179 | device: str, |
| 180 | ): |
| 181 | """Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method. |
| 182 | |
| 183 | - forward_native methods should call ApplyRotaryEmb.forward_native() |
| 184 | - forward_cuda/forward methods should call ApplyRotaryEmb.forward() |
| 185 | """ |
| 186 | run_dispatch_test(test_case, device) |
| 187 | |
| 188 | |
| 189 | @pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU only test.") |
| 190 | @pytest.mark.parametrize( |
| 191 | "num_tokens,num_heads,head_size,rot_dim,is_neox_style", |
| 192 | [ |
| 193 | (128, 32, 128, 128, True), # typical text shape |
| 194 | (4096, 16, 80, 80, False), # GPT-J style (interleaved) |
| 195 | (4096, 32, 128, 64, True), # partial rotary (rot_dim < head_size) |
| 196 | ], |
| 197 | ids=["text", "gptj_interleaved", "partial_rotary"], |
| 198 | ) |
| 199 | def test_apply_rotary_emb_xpu_matches_native( |
| 200 | num_tokens: int, |
| 201 | num_heads: int, |
| 202 | head_size: int, |
| 203 | rot_dim: int, |
| 204 | is_neox_style: bool, |
| 205 | ): |
| 206 | """ApplyRotaryEmb.forward_xpu (vllm_xpu_kernels SYCL op) must match |
| 207 | forward_native (the CustomOp default fallback) exactly.""" |
| 208 | from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb |
| 209 | |
| 210 | vllm_config = VllmConfig( |
| 211 | compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"]) |
| 212 | ) |
| 213 | get_cached_compilation_config.cache_clear() |
| 214 | |
| 215 | with set_current_vllm_config(vllm_config): |
| 216 | op = ApplyRotaryEmb(enforce_enable=True, is_neox_style=is_neox_style) |
| 217 | |
| 218 | x = torch.randn( |
| 219 | num_tokens, num_heads, head_size, device="xpu", dtype=torch.bfloat16 |
| 220 | ) |
| 221 | x = x[..., :rot_dim].contiguous() |
| 222 | cos = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16) |
| 223 | sin = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16) |
| 224 | |
| 225 | expected = op.forward_native(x, cos, sin) |
| 226 | actual = op.forward_xpu(x, cos, sin) |
| 227 | torch.testing.assert_close(actual, expected, atol=0, rtol=0) |
| 228 | |
| 229 | |
| 230 | @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA/ROCm only test.") |
| 231 | @pytest.mark.parametrize("num_tokens", [257, 2304]) |
| 232 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 233 | def test_packed_qk_rope_correctness( |
| 234 | num_tokens: int, dtype: torch.dtype, default_vllm_config |
| 235 | ): |
| 236 | """packed_qk_rope_ must match the per-tensor ApplyRotaryEmb path with |
| 237 | enable_fp32_compute=True to within one ULP, and must leave the V slice |
| 238 | untouched. |
| 239 | |
| 240 | Both sides compute in fp32 on identical inputs. The only gap is which |
| 241 | product the backend keeps exact inside the o1 multiply-add, and that |
| 242 | follows the register layout, so it is decided by codegen rather than by |
| 243 | the source and holds neither across backends nor across Triton versions. |
| 244 | """ |
| 245 | from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb |
| 246 | from vllm.model_executor.layers.rotary_embedding.packed_qk_rope import ( |
| 247 | packed_qk_rope_, |
| 248 | ) |
| 249 | |
| 250 | num_heads, head_dim = 12, 128 |
| 251 | rng = torch.Generator(device="cuda").manual_seed(0) |
| 252 | |
| 253 | angles = ( |
| 254 | torch.rand(num_tokens, head_dim // 2, device="cuda", generator=rng) * torch.pi |
| 255 | ) |
| 256 | freqs_cis = torch.polar(torch.ones_like(angles), angles) |
| 257 | cos = freqs_cis.real.contiguous() |
| 258 | sin = freqs_cis.imag.contiguous() |
| 259 | |
| 260 | xqkv = torch.randn( |
| 261 | num_tokens, 3, num_heads, head_dim, dtype=dtype, device="cuda", generator=rng |
| 262 | ) |
| 263 | xqkv_ref = xqkv.clone() |
| 264 | xq, xk, xv = torch.unbind(xqkv_ref, dim=-3) |
| 265 | |
| 266 | op = ApplyRotaryEmb( |
| 267 | enforce_enable=True, is_neox_style=False, enable_fp32_compute=True |
| 268 | ) |
| 269 | # op() dispatches to forward_cuda / forward_hip per platform. |
| 270 | xq_ref = op(xq, cos, sin) |
| 271 | xk_ref = op(xk, cos, sin) |
| 272 | |
| 273 | packed_qk_rope_(xqkv, freqs_cis) |
| 274 | |
| 275 | # One ULP at the largest magnitude present, not relative: o1 can cancel, |
| 276 | # so the error scales with the inputs to the sum, not with the result. |
| 277 | eps = torch.finfo(dtype).eps |
| 278 | q_atol = eps * xq_ref.abs().max().item() |
| 279 | k_atol = eps * xk_ref.abs().max().item() |
| 280 | torch.testing.assert_close(xqkv[:, 0], xq_ref, atol=q_atol, rtol=0) |
| 281 | torch.testing.assert_close(xqkv[:, 1], xk_ref, atol=k_atol, rtol=0) |
| 282 | # V is never read or written by the kernel, so it stays exact. |
| 283 | torch.testing.assert_close(xqkv[:, 2], xv, atol=0, rtol=0) |