@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 225 lines 6.16 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5 from types import SimpleNamespace
6
7 import pytest
8 import torch
9
10 from vllm.config import CacheConfig, KVTransferConfig, VllmConfig
11 from vllm.platforms.cpu import CpuPlatform
12
13
14 class _FakeModel:
15 @classmethod
16 def get_mamba_state_dtype_from_config(
17 cls, vllm_config: VllmConfig
18 ) -> tuple[torch.dtype, torch.dtype]:
19 state_dtype = {
20 "float16": torch.float16,
21 "bfloat16": torch.bfloat16,
22 }.get(vllm_config.cache_config.mamba_ssm_cache_dtype, torch.float32)
23 return torch.float32, state_dtype
24
25
26 class _FakeRegistry:
27 @staticmethod
28 def resolve_model_cls(
29 architecture: str,
30 model_config: SimpleNamespace,
31 ) -> tuple[type[_FakeModel], str]:
32 return _FakeModel, architecture
33
34
35 def _cpu_config(
36 cache_config: CacheConfig,
37 *,
38 model_type: str,
39 resolved_dtype: str,
40 architecture: str | None = None,
41 layer_types: tuple[str, ...] = ("linear_attention",),
42 ) -> SimpleNamespace:
43 model_config = SimpleNamespace(
44 disable_cascade_attn=False,
45 architecture=architecture or model_type,
46 has_inner_state=True,
47 use_mla=False,
48 hf_config=SimpleNamespace(model_type=model_type),
49 hf_text_config=SimpleNamespace(
50 model_type=model_type,
51 layer_types=layer_types,
52 ),
53 registry=_FakeRegistry,
54 )
55 config = VllmConfig(cache_config=cache_config)
56 config.model_config = model_config
57 cache_config.mamba_ssm_cache_dtype = resolved_dtype
58 return config
59
60
61 @pytest.mark.parametrize(
62 (
63 "model_type",
64 "requested_dtype",
65 "resolved_dtype",
66 "layer_types",
67 "expected_dtype",
68 ),
69 [
70 pytest.param(
71 "qwen3_5",
72 "bfloat16",
73 "bfloat16",
74 ("linear_attention",),
75 "bfloat16",
76 id="gdn-explicit-bf16",
77 ),
78 pytest.param(
79 "qwen3_5",
80 "float16",
81 "float16",
82 ("linear_attention",),
83 "float16",
84 id="gdn-explicit-fp16",
85 ),
86 pytest.param(
87 "qwen3_5",
88 "auto",
89 "bfloat16",
90 ("linear_attention",),
91 "bfloat16",
92 id="gdn-model-bf16",
93 ),
94 pytest.param(
95 "qwen3_5",
96 "auto",
97 "float16",
98 ("linear_attention",),
99 "float16",
100 id="gdn-model-fp16",
101 ),
102 pytest.param(
103 "nemotron_h",
104 "auto",
105 "bfloat16",
106 ("mamba",),
107 "float32",
108 id="unsupported-model-bf16",
109 ),
110 pytest.param(
111 "nemotron_h",
112 "auto",
113 "float16",
114 ("mamba",),
115 "float32",
116 id="unsupported-model-fp16",
117 ),
118 pytest.param(
119 "nemotron_h",
120 "bfloat16",
121 "bfloat16",
122 ("mamba",),
123 "float32",
124 id="unsupported-explicit-bf16",
125 ),
126 pytest.param(
127 "nemotron_h",
128 "float16",
129 "float16",
130 ("mamba",),
131 "float32",
132 id="unsupported-explicit-fp16",
133 ),
134 ],
135 )
136 def test_cpu_accelerated_gdn_dtype_policy(
137 monkeypatch: pytest.MonkeyPatch,
138 model_type: str,
139 requested_dtype: str,
140 resolved_dtype: str,
141 layer_types: tuple[str, ...],
142 expected_dtype: str,
143 ) -> None:
144 monkeypatch.setattr("torch.cpu._is_avx512_bf16_supported", lambda: True)
145 cache_config = CacheConfig(mamba_ssm_cache_dtype=requested_dtype)
146 config = _cpu_config(
147 cache_config,
148 model_type=model_type,
149 resolved_dtype=resolved_dtype,
150 layer_types=layer_types,
151 )
152
153 CpuPlatform.check_and_update_config(config)
154 assert cache_config.mamba_ssm_cache_dtype == expected_dtype
155
156
157 @pytest.mark.parametrize(
158 (
159 "connector",
160 "extra_config",
161 "explicit_layout",
162 "avx512_bf16_supported",
163 "expected_layout",
164 ),
165 [
166 pytest.param("NixlConnector", None, None, True, "DS", id="nixl"),
167 pytest.param("NixlPullConnector", None, None, True, "DS", id="nixl-pull"),
168 pytest.param("NixlPushConnector", None, None, True, "DS", id="nixl-push"),
169 pytest.param(
170 "MultiConnector",
171 {
172 "connectors": [
173 {
174 "kv_connector": "ExampleConnector",
175 "kv_connector_extra_config": {},
176 },
177 {
178 "kv_connector": "NixlConnector",
179 "kv_connector_extra_config": {},
180 },
181 ]
182 },
183 None,
184 True,
185 "DS",
186 id="multi-nixl",
187 ),
188 pytest.param("NixlConnector", None, "SD", True, "SD", id="explicit-override"),
189 pytest.param("OffloadingConnector", None, None, True, "SD", id="non-nixl"),
190 pytest.param(
191 "OffloadingConnector", None, None, False, None, id="non-nixl-no-avx"
192 ),
193 ],
194 )
195 def test_cpu_conv_state_layout_selection(
196 monkeypatch: pytest.MonkeyPatch,
197 connector: str,
198 extra_config: dict | None,
199 explicit_layout: str | None,
200 avx512_bf16_supported: bool,
201 expected_layout: str | None,
202 ) -> None:
203 layout_env = "VLLM_SSM_CONV_STATE_LAYOUT"
204 monkeypatch.setattr(
205 "torch.cpu._is_avx512_bf16_supported",
206 lambda: avx512_bf16_supported,
207 )
208
209 kv_transfer_config = KVTransferConfig(
210 kv_connector=connector,
211 kv_connector_extra_config=extra_config or {},
212 kv_role="kv_both",
213 )
214 config = _cpu_config(
215 CacheConfig(mamba_ssm_cache_dtype="float32"),
216 model_type="qwen3_5",
217 resolved_dtype="float32",
218 )
219 config.kv_transfer_config = kv_transfer_config
220 monkeypatch.delenv(layout_env, raising=False)
221 if explicit_layout is not None:
222 monkeypatch.setenv(layout_env, explicit_layout)
223
224 CpuPlatform.check_and_update_config(config)
225 assert os.environ.get(layout_env) == expected_layout