@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 211 lines 6.48 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for the JSON-based config loader added to selective_state_update.
4
5 Tests cover:
6 - Flat MoE-style filename generation
7 - VLLM_TUNED_CONFIG_FOLDER env-var override
8 - Fallback to heuristic when no config file exists
9 - Nearest effective_batch interpolation
10 - Edge cases: non-dict JSON, empty config
11 """
12
13 import json
14
15 from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
16 _get_default_ssm_launch_config,
17 _try_get_optimal_ssm_config_cached,
18 get_ssm_config_file_name,
19 get_ssm_configs,
20 get_ssm_device_name,
21 try_get_optimal_ssm_config,
22 )
23
24 # Common kwargs for try_get_optimal_ssm_config. Tests pick (batch, nheads) so
25 # their product (effective_batch) matches the value being probed.
26 _HEADDIM = 64
27 _CACHE_DTYPE = "float32"
28
29
30 def _clear_caches() -> None:
31 get_ssm_configs.cache_clear()
32 _try_get_optimal_ssm_config_cached.cache_clear()
33
34
35 def _write_config(tmp_path, dstate: int, payload: dict) -> None:
36 """Write payload as the bundled config for (headdim, dstate, cache_dtype)."""
37 device_name = get_ssm_device_name()
38 config_path = tmp_path / get_ssm_config_file_name(
39 _HEADDIM, dstate, _CACHE_DTYPE, device_name
40 )
41 with open(config_path, "w") as f:
42 json.dump(payload, f)
43
44
45 # ---------------------------------------------------------------------------
46 # Config filename generation
47 # ---------------------------------------------------------------------------
48
49
50 def test_config_file_name_format():
51 name = get_ssm_config_file_name(
52 headdim=64, dstate=128, cache_dtype="float32", device_name="NVIDIA_B200"
53 )
54 assert name == (
55 "headdim=64,dstate=128,device_name=NVIDIA_B200,cache_dtype=float32.json"
56 )
57
58
59 # ---------------------------------------------------------------------------
60 # VLLM_TUNED_CONFIG_FOLDER override
61 # ---------------------------------------------------------------------------
62
63
64 def test_env_override_loads_custom_config(monkeypatch, tmp_path):
65 """VLLM_TUNED_CONFIG_FOLDER should take precedence over the bundled dir."""
66 _write_config(
67 tmp_path,
68 dstate=16,
69 payload={
70 "1": {"BLOCK_SIZE_M": 4, "num_warps": 1},
71 },
72 )
73
74 monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
75 _clear_caches()
76
77 cfg = get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE)
78 assert cfg is not None
79 assert cfg[1] == {"BLOCK_SIZE_M": 4, "num_warps": 1}
80
81 _clear_caches()
82
83
84 # ---------------------------------------------------------------------------
85 # Fallback to heuristic when no config file exists
86 # ---------------------------------------------------------------------------
87
88
89 def test_fallback_when_no_config(monkeypatch, tmp_path):
90 """try_get_optimal_ssm_config must fall back to _get_default_ssm_launch_config
91 when no JSON file is found for the current
92 (device, headdim, dstate, cache_dtype) combination.
93 """
94 monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
95 monkeypatch.setattr(
96 "vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
97 str(tmp_path),
98 )
99
100 for dstate in (8, 16, 32, 64, 128, 256):
101 for is_blackwell in (False, True):
102 _clear_caches()
103 block_m, warps = try_get_optimal_ssm_config(
104 headdim=_HEADDIM,
105 dstate=dstate,
106 batch=1,
107 nheads=1,
108 cache_dtype=_CACHE_DTYPE,
109 is_blackwell=is_blackwell,
110 )
111 assert (block_m, warps) == _get_default_ssm_launch_config(
112 dstate, is_blackwell=is_blackwell
113 )
114
115 _clear_caches()
116
117
118 # ---------------------------------------------------------------------------
119 # Nearest effective_batch interpolation
120 # ---------------------------------------------------------------------------
121
122
123 def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path):
124 """When effective_batch = batch*nheads is not an exact key, the closest
125 key should be selected."""
126 _write_config(
127 tmp_path,
128 dstate=32,
129 payload={
130 "64": {"BLOCK_SIZE_M": 8, "num_warps": 1},
131 "4096": {"BLOCK_SIZE_M": 32, "num_warps": 4},
132 },
133 )
134
135 monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
136 _clear_caches()
137
138 # effective_batch = 1*128 = 128 -> closer to 64 than to 4096
139 block_m, warps = try_get_optimal_ssm_config(
140 headdim=_HEADDIM,
141 dstate=32,
142 batch=1,
143 nheads=128,
144 cache_dtype=_CACHE_DTYPE,
145 is_blackwell=False,
146 )
147 assert block_m == 8 and warps == 1
148
149 # effective_batch = 4*1024 = 4096 -> exact match on 4096
150 block_m, warps = try_get_optimal_ssm_config(
151 headdim=_HEADDIM,
152 dstate=32,
153 batch=4,
154 nheads=1024,
155 cache_dtype=_CACHE_DTYPE,
156 is_blackwell=False,
157 )
158 assert block_m == 32 and warps == 4
159
160 _clear_caches()
161
162
163 # ---------------------------------------------------------------------------
164 # Edge cases: malformed / empty config files
165 # ---------------------------------------------------------------------------
166
167
168 def test_non_dict_json_returns_none(monkeypatch, tmp_path):
169 """A valid JSON file that is not a dict (e.g. a list) must be ignored
170 and return None rather than raising AttributeError."""
171 device_name = get_ssm_device_name()
172 config_path = tmp_path / get_ssm_config_file_name(
173 _HEADDIM, 16, _CACHE_DTYPE, device_name
174 )
175 with open(config_path, "w") as f:
176 json.dump([1, 2, 3], f)
177
178 monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
179 monkeypatch.setattr(
180 "vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
181 str(tmp_path),
182 )
183 _clear_caches()
184
185 assert get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE) is None
186
187 _clear_caches()
188
189
190 def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path):
191 """An empty JSON object {} must not crash min() — should fall back
192 to the hard-coded heuristic."""
193 _write_config(tmp_path, dstate=64, payload={})
194
195 monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
196 _clear_caches()
197
198 dstate = 64
199 block_m, warps = try_get_optimal_ssm_config(
200 headdim=_HEADDIM,
201 dstate=dstate,
202 batch=1,
203 nheads=64,
204 cache_dtype=_CACHE_DTYPE,
205 is_blackwell=False,
206 )
207 assert (block_m, warps) == _get_default_ssm_launch_config(
208 dstate=dstate, is_blackwell=False
209 )
210
211 _clear_caches()