@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 5.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 from types import SimpleNamespace
6
7 import pytest
8 import torch
9
10 from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
11 from vllm.model_executor.layers.rotary_embedding import get_rope
12 from vllm.model_executor.layers.rotary_embedding.bailing_mrope import (
13 BailingMRotaryEmbedding,
14 )
15 from vllm.model_executor.models.bailing_moe_v3 import _build_mla_rotary_embedding
16 from vllm.platforms import current_platform
17
18
19 def test_bailing_mrope_selects_alternating_spatial_then_temporal():
20 with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))):
21 rotary_emb = BailingMRotaryEmbedding(
22 head_size=12,
23 rotary_dim=12,
24 max_position_embeddings=8,
25 base=10_000,
26 is_neox_style=False,
27 dtype=torch.float32,
28 mrope_section=[2, 2, 2],
29 )
30
31 positions = torch.tensor([[1, 2], [3, 4], [5, 6]])
32 frequencies = torch.arange(6).expand(8, -1)
33 position_values = torch.arange(8).unsqueeze(1) * 10
34 cos = position_values + frequencies
35 sin = cos + 100
36 cos_sin_cache = torch.cat((cos, sin), dim=-1)
37
38 actual_cos, actual_sin = rotary_emb.select_cos_sin(
39 positions,
40 cos_sin_cache,
41 )
42
43 expected_cos = torch.tensor(
44 [
45 [30, 51, 32, 53, 14, 15],
46 [40, 61, 42, 63, 24, 25],
47 ]
48 )
49 torch.testing.assert_close(actual_cos, expected_cos)
50 torch.testing.assert_close(actual_sin, expected_cos + 100)
51
52
53 def test_bailing_factory_preserves_mla_width_and_config():
54 config = SimpleNamespace(
55 max_position_embeddings=128,
56 rope_parameters={
57 "rope_type": "default",
58 "rope_theta": 6_000_000,
59 "partial_rotary_factor": 0.5,
60 "mrope_section": [8, 12, 12],
61 },
62 rope_scaling={"partial_rotary_factor": 0.5},
63 )
64 original = copy.deepcopy(vars(config))
65 with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))):
66 rope = _build_mla_rotary_embedding(config, head_size=64)
67 assert rope is _build_mla_rotary_embedding(config, head_size=64)
68 generic = get_rope(
69 head_size=64,
70 max_position=128,
71 is_neox_style=False,
72 rope_parameters={
73 "rope_type": "default",
74 "rope_theta": 6_000_000,
75 "mrope_section": [8, 12, 12],
76 },
77 )
78 assert isinstance(rope, BailingMRotaryEmbedding)
79 assert rope is not generic
80 assert rope.rotary_dim == 64
81 assert vars(config) == original
82
83
84 def test_bailing_factory_rejects_unsupported_scaling():
85 config = SimpleNamespace(
86 rope_parameters={
87 "rope_type": "yarn",
88 "factor": 2.0,
89 "mrope_section": [8, 12, 12],
90 }
91 )
92 with pytest.raises(ValueError, match="only supports rope_type='default'"):
93 _build_mla_rotary_embedding(config, head_size=64)
94
95
96 @pytest.mark.skipif(
97 not current_platform.is_cuda_alike(),
98 reason="Skipping CUDA/ROCm-only test.",
99 )
100 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
101 @pytest.mark.parametrize("position_kind", ["text", "multimodal"])
102 @torch.inference_mode()
103 def test_bailing_mrope_cuda_matches_native_for_strided_mla_views(
104 dtype: torch.dtype,
105 position_kind: str,
106 ):
107 device = torch.device("cuda")
108 num_tokens = 17
109 num_query_heads = 8
110 with set_current_vllm_config(VllmConfig()):
111 rotary_emb = BailingMRotaryEmbedding(
112 head_size=64,
113 rotary_dim=64,
114 max_position_embeddings=64,
115 base=6_000_000,
116 is_neox_style=False,
117 dtype=dtype,
118 mrope_section=[8, 12, 12],
119 ).to(device)
120
121 torch.manual_seed(0)
122 multimodal_positions = torch.randint(
123 0,
124 64,
125 (3, 2 * num_tokens),
126 dtype=torch.int32,
127 device=device,
128 )[:, ::2]
129 if position_kind == "multimodal":
130 positions = multimodal_positions
131 reference_positions = positions
132 else:
133 positions = torch.randint(
134 0,
135 64,
136 (2 * num_tokens,),
137 dtype=torch.int64,
138 device=device,
139 )[::2]
140 reference_positions = positions.expand(3, -1)
141 query_buffer = torch.randn(
142 num_tokens,
143 num_query_heads,
144 192,
145 dtype=dtype,
146 device=device,
147 )
148 key_buffer = torch.randn(
149 num_tokens,
150 576,
151 dtype=dtype,
152 device=device,
153 )
154 query = query_buffer[..., 128:]
155 key = key_buffer[:, 512:].unsqueeze(1)
156 assert not query.is_contiguous()
157 assert not key.is_contiguous()
158
159 query_prefix = query_buffer[..., :128].clone()
160 key_prefix = key_buffer[:, :512].clone()
161 expected_query, expected_key = rotary_emb.forward_native(
162 reference_positions,
163 query.clone(),
164 key.clone(),
165 )
166 actual_query, actual_key = rotary_emb.forward_cuda(positions, query, key)
167 assert expected_key is not None
168 assert actual_key is not None
169
170 torch.testing.assert_close(actual_query, expected_query, atol=1e-2, rtol=1.6e-2)
171 torch.testing.assert_close(actual_key, expected_key, atol=1e-2, rtol=1.6e-2)
172 assert torch.equal(query_buffer[..., :128], query_prefix)
173 assert torch.equal(key_buffer[:, :512], key_prefix)
174 if position_kind == "multimodal":
175 assert actual_query.data_ptr() == query.data_ptr()
176 assert actual_key.data_ptr() == key.data_ptr()