| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for Helion ConfigManager and ConfigSet. |
| 4 | |
| 5 | Tests the simplified configuration management system for Helion custom kernels. |
| 6 | """ |
| 7 | |
| 8 | import json |
| 9 | import tempfile |
| 10 | from pathlib import Path |
| 11 | |
| 12 | import pytest |
| 13 | |
| 14 | from vllm.utils.import_utils import has_helion |
| 15 | |
| 16 | # Skip entire module if helion is not available |
| 17 | if not has_helion(): |
| 18 | pytest.skip( |
| 19 | "Helion is not installed. Install with: pip install vllm[helion]", |
| 20 | allow_module_level=True, |
| 21 | ) |
| 22 | |
| 23 | import helion |
| 24 | |
| 25 | from vllm.kernels.helion.case_key import CaseKey |
| 26 | from vllm.kernels.helion.config_manager import ( |
| 27 | ConfigManager, |
| 28 | ConfigSet, |
| 29 | ) |
| 30 | |
| 31 | |
| 32 | @pytest.fixture(autouse=True) |
| 33 | def reset_config_manager_singleton(): |
| 34 | """Reset ConfigManager singleton before each test.""" |
| 35 | ConfigManager.reset_instance() |
| 36 | yield |
| 37 | ConfigManager.reset_instance() |
| 38 | |
| 39 | |
| 40 | class TestConfigSet: |
| 41 | """Test suite for ConfigSet class.""" |
| 42 | |
| 43 | def test_config_set_creation(self): |
| 44 | """Test creating an empty ConfigSet.""" |
| 45 | config_set = ConfigSet("test_kernel") |
| 46 | |
| 47 | assert config_set.kernel_name == "test_kernel" |
| 48 | assert config_set.get_platforms() == [] |
| 49 | |
| 50 | def test_config_set_from_dict(self): |
| 51 | """Test creating ConfigSet from dictionary data.""" |
| 52 | config_data = { |
| 53 | "block_sizes": [32, 16], |
| 54 | "num_warps": 4, |
| 55 | "num_stages": 3, |
| 56 | "pid_type": "persistent_interleaved", |
| 57 | } |
| 58 | data = { |
| 59 | "h100": [ |
| 60 | {"key": {"batch": 32, "hidden": 4096}, "config": config_data}, |
| 61 | ] |
| 62 | } |
| 63 | |
| 64 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 65 | |
| 66 | assert config_set.kernel_name == "test_kernel" |
| 67 | assert config_set.get_platforms() == ["h100"] |
| 68 | |
| 69 | internal_key = CaseKey({"batch": 32, "hidden": 4096}) |
| 70 | config = config_set.get_config("h100", internal_key) |
| 71 | assert isinstance(config, helion.Config) |
| 72 | assert config.block_sizes == [32, 16] |
| 73 | assert config.num_warps == 4 |
| 74 | assert config.num_stages == 3 |
| 75 | assert config.pid_type == "persistent_interleaved" |
| 76 | |
| 77 | def test_config_set_get_config_keyerror(self): |
| 78 | """Test that accessing non-existent configs raises informative KeyErrors.""" |
| 79 | config_set = ConfigSet("test_kernel") |
| 80 | |
| 81 | with pytest.raises(KeyError, match="platform 'h100' not found"): |
| 82 | config_set.get_config("h100", CaseKey({"batch": 32})) |
| 83 | |
| 84 | config_data = {"num_warps": 8, "num_stages": 4} |
| 85 | data = { |
| 86 | "h100": [ |
| 87 | {"key": {"batch": 64, "hidden": 2048}, "config": config_data}, |
| 88 | ] |
| 89 | } |
| 90 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 91 | |
| 92 | nonexistent_key = CaseKey({"batch": 32, "hidden": 4096}) |
| 93 | with pytest.raises(KeyError, match="config_key .* not found"): |
| 94 | config_set.get_config("h100", nonexistent_key) |
| 95 | |
| 96 | def test_config_set_get_platforms(self): |
| 97 | """Test get_platforms method.""" |
| 98 | # Use realistic config data |
| 99 | config1 = {"num_warps": 4, "num_stages": 3} |
| 100 | config2 = {"num_warps": 8, "num_stages": 5} |
| 101 | |
| 102 | data = { |
| 103 | "h100": [ |
| 104 | {"key": {"batch": 32, "hidden": 4096}, "config": config1}, |
| 105 | ], |
| 106 | "a100": [ |
| 107 | {"key": {"batch": 16, "hidden": 2048}, "config": config2}, |
| 108 | ], |
| 109 | } |
| 110 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 111 | |
| 112 | platforms = config_set.get_platforms() |
| 113 | assert platforms == ["a100", "h100"] # Should be sorted |
| 114 | |
| 115 | def test_config_set_get_config_keys(self): |
| 116 | """Test get_config_keys method.""" |
| 117 | config1 = {"num_warps": 4, "num_stages": 3} |
| 118 | config2 = {"num_warps": 8, "num_stages": 5} |
| 119 | |
| 120 | data = { |
| 121 | "h100": [ |
| 122 | {"key": {"batch": 32, "hidden": 4096}, "config": config1}, |
| 123 | {"key": {"batch": 64, "hidden": 2048}, "config": config2}, |
| 124 | ] |
| 125 | } |
| 126 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 127 | |
| 128 | config_keys = config_set.get_config_keys("h100") |
| 129 | expected_keys = sorted( |
| 130 | [ |
| 131 | CaseKey({"batch": 32, "hidden": 4096}), |
| 132 | CaseKey({"batch": 64, "hidden": 2048}), |
| 133 | ], |
| 134 | key=lambda k: str(k) if k is not None else "", |
| 135 | ) |
| 136 | assert config_keys == expected_keys |
| 137 | |
| 138 | assert config_set.get_config_keys("v100") == [] |
| 139 | |
| 140 | def test_config_set_to_dict(self): |
| 141 | """Test converting ConfigSet to dictionary.""" |
| 142 | original_config = { |
| 143 | "block_sizes": [64, 32], |
| 144 | "num_warps": 16, |
| 145 | "num_stages": 4, |
| 146 | "pid_type": "persistent_blocked", |
| 147 | } |
| 148 | original_data = { |
| 149 | "h100": [ |
| 150 | {"key": {"batch": 32, "hidden": 4096}, "config": original_config}, |
| 151 | ] |
| 152 | } |
| 153 | |
| 154 | config_set = ConfigSet.from_dict("test_kernel", original_data) |
| 155 | result_data = config_set.to_dict() |
| 156 | |
| 157 | internal_key = CaseKey({"batch": 32, "hidden": 4096}) |
| 158 | assert internal_key in result_data["h100"] |
| 159 | assert result_data["h100"][internal_key] == original_config |
| 160 | |
| 161 | |
| 162 | class TestConfigManager: |
| 163 | """Test suite for ConfigManager class.""" |
| 164 | |
| 165 | def test_config_manager_creation_default_base_dir(self): |
| 166 | """Test creating ConfigManager with default base directory.""" |
| 167 | manager = ConfigManager() |
| 168 | assert manager._base_dir.name == "configs" |
| 169 | |
| 170 | def test_config_manager_creation_custom_base_dir(self): |
| 171 | """Test creating ConfigManager with custom base directory.""" |
| 172 | custom_dir = "/tmp/custom_configs" |
| 173 | manager = ConfigManager(base_dir=custom_dir) |
| 174 | |
| 175 | # Paths are resolved, so compare with resolved path |
| 176 | assert manager._base_dir == Path(custom_dir).resolve() |
| 177 | |
| 178 | def test_get_config_file_path(self): |
| 179 | """Test getting config file path for a kernel.""" |
| 180 | manager = ConfigManager(base_dir="/tmp") |
| 181 | |
| 182 | dir_path = manager.get_config_file_path("silu_mul_fp8") |
| 183 | assert dir_path == Path("/tmp/silu_mul_fp8") |
| 184 | |
| 185 | file_path = manager.get_config_file_path("silu_mul_fp8", "nvidia_h100") |
| 186 | assert file_path == Path("/tmp/silu_mul_fp8/nvidia_h100.json") |
| 187 | |
| 188 | def test_ensure_base_dir_exists(self): |
| 189 | """Test ensuring base directory exists.""" |
| 190 | with tempfile.TemporaryDirectory() as temp_dir: |
| 191 | base_dir = Path(temp_dir) / "non_existent" / "configs" |
| 192 | manager = ConfigManager(base_dir=base_dir) |
| 193 | assert not base_dir.exists() |
| 194 | |
| 195 | returned_path = manager.ensure_base_dir_exists() |
| 196 | |
| 197 | assert base_dir.exists() |
| 198 | assert base_dir.is_dir() |
| 199 | assert returned_path == base_dir |
| 200 | |
| 201 | def test_load_config_set_file_not_exists(self): |
| 202 | """Test loading config set when file doesn't exist.""" |
| 203 | with tempfile.TemporaryDirectory() as temp_dir: |
| 204 | manager = ConfigManager(base_dir=temp_dir) |
| 205 | config_set = manager.load_config_set("non_existent_kernel") |
| 206 | |
| 207 | assert isinstance(config_set, ConfigSet) |
| 208 | assert config_set.kernel_name == "non_existent_kernel" |
| 209 | assert config_set.get_platforms() == [] |
| 210 | |
| 211 | def test_load_config_set_valid_file(self): |
| 212 | """Test loading config set from per-platform files.""" |
| 213 | with tempfile.TemporaryDirectory() as temp_dir: |
| 214 | kernel_config = { |
| 215 | "block_sizes": [128, 64], |
| 216 | "num_warps": 8, |
| 217 | "num_stages": 6, |
| 218 | "pid_type": "persistent_interleaved", |
| 219 | } |
| 220 | kernel_dir = Path(temp_dir) / "test_kernel" |
| 221 | kernel_dir.mkdir() |
| 222 | platform_file = kernel_dir / "h100.json" |
| 223 | with open(platform_file, "w") as f: |
| 224 | json.dump( |
| 225 | [{"key": {"batch": 32, "hidden": 4096}, "config": kernel_config}], |
| 226 | f, |
| 227 | ) |
| 228 | |
| 229 | manager = ConfigManager(base_dir=temp_dir) |
| 230 | config_set = manager.load_config_set("test_kernel") |
| 231 | |
| 232 | assert isinstance(config_set, ConfigSet) |
| 233 | assert config_set.kernel_name == "test_kernel" |
| 234 | assert config_set.get_platforms() == ["h100"] |
| 235 | |
| 236 | internal_key = CaseKey({"batch": 32, "hidden": 4096}) |
| 237 | config = config_set.get_config("h100", internal_key) |
| 238 | assert isinstance(config, helion.Config) |
| 239 | assert config.block_sizes == [128, 64] |
| 240 | assert config.num_warps == 8 |
| 241 | |
| 242 | def test_load_config_set_invalid_json(self): |
| 243 | """Test loading config set from file with invalid JSON.""" |
| 244 | with tempfile.TemporaryDirectory() as temp_dir: |
| 245 | kernel_dir = Path(temp_dir) / "test_kernel" |
| 246 | kernel_dir.mkdir() |
| 247 | config_file = kernel_dir / "h100.json" |
| 248 | with open(config_file, "w") as f: |
| 249 | f.write("invalid json content {") |
| 250 | |
| 251 | manager = ConfigManager(base_dir=temp_dir) |
| 252 | config_set = manager.load_config_set("test_kernel") |
| 253 | |
| 254 | assert isinstance(config_set, ConfigSet) |
| 255 | assert config_set.kernel_name == "test_kernel" |
| 256 | assert config_set.get_platforms() == [] |
| 257 | |
| 258 | def test_save_config_set(self): |
| 259 | """Test saving ConfigSet to per-platform files.""" |
| 260 | with tempfile.TemporaryDirectory() as temp_dir: |
| 261 | kernel_config = { |
| 262 | "block_sizes": [256, 128], |
| 263 | "num_warps": 16, |
| 264 | "num_stages": 8, |
| 265 | "pid_type": "persistent_blocked", |
| 266 | } |
| 267 | data = { |
| 268 | "h100": [ |
| 269 | {"key": {"batch": 32, "hidden": 4096}, "config": kernel_config}, |
| 270 | ] |
| 271 | } |
| 272 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 273 | |
| 274 | manager = ConfigManager(base_dir=temp_dir) |
| 275 | saved_path = manager.save_config_set(config_set) |
| 276 | |
| 277 | expected_dir = Path(temp_dir) / "test_kernel" |
| 278 | assert saved_path == expected_dir |
| 279 | assert saved_path.is_dir() |
| 280 | |
| 281 | platform_file = expected_dir / "h100.json" |
| 282 | assert platform_file.exists() |
| 283 | with open(platform_file) as f: |
| 284 | loaded_data = json.load(f) |
| 285 | assert isinstance(loaded_data, list) |
| 286 | assert len(loaded_data) == 1 |
| 287 | entry = loaded_data[0] |
| 288 | assert entry["key"] == {"batch": 32, "hidden": 4096} |
| 289 | assert entry["config"] == kernel_config |
| 290 | |
| 291 | def test_save_config_set_creates_directory(self): |
| 292 | """Test that save_config_set creates parent directories if needed.""" |
| 293 | with tempfile.TemporaryDirectory() as temp_dir: |
| 294 | nested_dir = Path(temp_dir) / "nested" / "configs" |
| 295 | data = { |
| 296 | "h100": [ |
| 297 | {"key": {}, "config": {"num_warps": 4}}, |
| 298 | ] |
| 299 | } |
| 300 | config_set = ConfigSet.from_dict("test_kernel", data) |
| 301 | |
| 302 | manager = ConfigManager(base_dir=nested_dir) |
| 303 | saved_path = manager.save_config_set(config_set) |
| 304 | |
| 305 | assert nested_dir.exists() |
| 306 | assert nested_dir.is_dir() |
| 307 | assert saved_path.is_dir() |
| 308 | assert (saved_path / "h100.json").exists() |
| 309 | |
| 310 | def test_get_platform_configs(self): |
| 311 | """Test getting all configs for a specific platform.""" |
| 312 | with tempfile.TemporaryDirectory() as temp_dir: |
| 313 | config_1 = {"num_warps": 4, "num_stages": 3, "block_sizes": [64, 32]} |
| 314 | config_2 = {"num_warps": 8, "num_stages": 5, "block_sizes": [128, 64]} |
| 315 | default_config = { |
| 316 | "num_warps": 16, |
| 317 | "num_stages": 7, |
| 318 | "block_sizes": [256, 128], |
| 319 | } |
| 320 | config_3 = {"num_warps": 2, "num_stages": 2, "block_sizes": [32, 16]} |
| 321 | |
| 322 | kernel_dir = Path(temp_dir) / "test_kernel" |
| 323 | kernel_dir.mkdir() |
| 324 | with open(kernel_dir / "h100.json", "w") as f: |
| 325 | json.dump( |
| 326 | [ |
| 327 | {"key": {"batch": 32, "hidden": 4096}, "config": config_1}, |
| 328 | {"key": {"batch": 64, "hidden": 2048}, "config": config_2}, |
| 329 | {"key": {}, "config": default_config}, |
| 330 | ], |
| 331 | f, |
| 332 | ) |
| 333 | with open(kernel_dir / "a100.json", "w") as f: |
| 334 | json.dump( |
| 335 | [{"key": {"batch": 16, "hidden": 1024}, "config": config_3}], |
| 336 | f, |
| 337 | ) |
| 338 | |
| 339 | manager = ConfigManager(base_dir=temp_dir) |
| 340 | |
| 341 | key_b32_h4096 = CaseKey({"batch": 32, "hidden": 4096}) |
| 342 | key_b64_h2048 = CaseKey({"batch": 64, "hidden": 2048}) |
| 343 | key_b16_h1024 = CaseKey({"batch": 16, "hidden": 1024}) |
| 344 | |
| 345 | h100_configs = manager.get_platform_configs("test_kernel", "h100") |
| 346 | assert len(h100_configs) == 3 |
| 347 | assert key_b32_h4096 in h100_configs |
| 348 | assert key_b64_h2048 in h100_configs |
| 349 | assert CaseKey.default() in h100_configs |
| 350 | for config in h100_configs.values(): |
| 351 | assert isinstance(config, helion.Config) |
| 352 | |
| 353 | assert h100_configs[key_b32_h4096].num_warps == 4 |
| 354 | assert h100_configs[CaseKey.default()].num_stages == 7 |
| 355 | |
| 356 | a100_configs = manager.get_platform_configs("test_kernel", "a100") |
| 357 | assert len(a100_configs) == 1 |
| 358 | assert key_b16_h1024 in a100_configs |
| 359 | assert isinstance(a100_configs[key_b16_h1024], helion.Config) |
| 360 | assert a100_configs[key_b16_h1024].num_warps == 2 |
| 361 | |
| 362 | nonexistent_configs = manager.get_platform_configs("test_kernel", "v100") |
| 363 | assert len(nonexistent_configs) == 0 |
| 364 | |
| 365 | def test_singleton_returns_same_instance(self): |
| 366 | """Test that ConfigManager returns the same instance on repeated calls.""" |
| 367 | manager1 = ConfigManager(base_dir="/tmp/test_singleton") |
| 368 | manager2 = ConfigManager(base_dir="/tmp/test_singleton") |
| 369 | |
| 370 | assert manager1 is manager2 |
| 371 | |
| 372 | def test_singleton_with_default_base_dir(self): |
| 373 | """Test singleton behavior with default base directory.""" |
| 374 | manager1 = ConfigManager() |
| 375 | manager2 = ConfigManager() |
| 376 | |
| 377 | assert manager1 is manager2 |
| 378 | assert manager1._base_dir == manager2._base_dir |
| 379 | |
| 380 | def test_singleton_error_on_different_base_dir(self): |
| 381 | """Test that ConfigManager raises error when created with different base_dir.""" |
| 382 | ConfigManager(base_dir="/tmp/first_dir") |
| 383 | |
| 384 | with pytest.raises(ValueError, match="singleton already exists"): |
| 385 | ConfigManager(base_dir="/tmp/different_dir") |
| 386 | |
| 387 | def test_reset_instance_allows_new_base_dir(self): |
| 388 | """Test that reset_instance allows creating with a new base_dir.""" |
| 389 | manager1 = ConfigManager(base_dir="/tmp/first_dir") |
| 390 | assert manager1._base_dir == Path("/tmp/first_dir").resolve() |
| 391 | |
| 392 | ConfigManager.reset_instance() |
| 393 | |
| 394 | manager2 = ConfigManager(base_dir="/tmp/second_dir") |
| 395 | assert manager2._base_dir == Path("/tmp/second_dir").resolve() |
| 396 | assert manager1 is not manager2 |
| 397 | |
| 398 | def test_get_instance_returns_existing(self): |
| 399 | """Test that get_instance returns the existing singleton.""" |
| 400 | manager1 = ConfigManager(base_dir="/tmp/test_get_instance") |
| 401 | manager2 = ConfigManager.get_instance() |
| 402 | |
| 403 | assert manager1 is manager2 |
| 404 | |
| 405 | def test_get_instance_raises_if_not_initialized(self): |
| 406 | """Test that get_instance raises RuntimeError if no instance exists.""" |
| 407 | with pytest.raises(RuntimeError, match="has not been created"): |
| 408 | ConfigManager.get_instance() |