| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Targeted unit tests for VLLM_TRITON_FORCE_FIRST_CONFIG. |
| 4 | |
| 5 | These tests exercise only the patched `Autotuner.run` logic installed by |
| 6 | `vllm.triton_utils.force_first_config.install`. The wrapped kernel is a |
| 7 | plain callable so the tests run on CPU-only hosts (no GPU, no actual |
| 8 | kernel launch) as long as the `triton` package is importable. |
| 9 | """ |
| 10 | |
| 11 | from types import SimpleNamespace |
| 12 | |
| 13 | import pytest |
| 14 | |
| 15 | from vllm.triton_utils import HAS_TRITON, triton |
| 16 | |
| 17 | if not HAS_TRITON: |
| 18 | pytest.skip("triton not available", allow_module_level=True) |
| 19 | |
| 20 | from vllm.triton_utils import force_first_config # noqa: E402 |
| 21 | |
| 22 | OutOfResources = triton.runtime.errors.OutOfResources |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def patched_autotuner(monkeypatch: pytest.MonkeyPatch): |
| 27 | """Install the first-valid-config patch and restore after. |
| 28 | |
| 29 | The env-var gate lives in vllm.env_override; install() itself does not |
| 30 | read the environment, so the test calls it directly. |
| 31 | """ |
| 32 | Autotuner = triton.runtime.autotuner.Autotuner |
| 33 | original_run = Autotuner.run |
| 34 | # Reset the once-only guard so install() re-runs for each test. |
| 35 | monkeypatch.setattr(force_first_config, "_installed", False) |
| 36 | force_first_config.install() |
| 37 | yield Autotuner |
| 38 | Autotuner.run = original_run |
| 39 | |
| 40 | |
| 41 | def _make_fake_self(configs, fn): |
| 42 | """Minimal stand-in for an Autotuner instance.""" |
| 43 | return SimpleNamespace( |
| 44 | configs=configs, |
| 45 | keys=[], |
| 46 | arg_names=[], |
| 47 | base_fn=fn, |
| 48 | fn=fn, |
| 49 | best_config=None, |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | def test_skips_invalid_first_config_and_caches_second(patched_autotuner): |
| 54 | bad = triton.Config({"BLOCK": 1024}) |
| 55 | good = triton.Config({"BLOCK": 64}) |
| 56 | calls = [] |
| 57 | |
| 58 | def fake_fn(*args, **kwargs): |
| 59 | calls.append(kwargs["BLOCK"]) |
| 60 | if kwargs["BLOCK"] == 1024: |
| 61 | raise OutOfResources(required=99999, limit=1, name="shared memory") |
| 62 | return "ok" |
| 63 | |
| 64 | fake_self = _make_fake_self([bad, good], fake_fn) |
| 65 | |
| 66 | # First call: walks past the invalid config, picks the second. |
| 67 | assert patched_autotuner.run(fake_self) == "ok" |
| 68 | assert calls == [1024, 64] |
| 69 | assert fake_self.best_config is good |
| 70 | |
| 71 | # Second call: cached index is reused, invalid config is NOT retried. |
| 72 | calls.clear() |
| 73 | assert patched_autotuner.run(fake_self) == "ok" |
| 74 | assert calls == [64] |
| 75 | |
| 76 | |
| 77 | def test_empty_configs_falls_back_to_direct_call(patched_autotuner): |
| 78 | def fake_fn(*args, **kwargs): |
| 79 | return "direct" |
| 80 | |
| 81 | fake_self = _make_fake_self([], fake_fn) |
| 82 | assert patched_autotuner.run(fake_self) == "direct" |
| 83 | |
| 84 | |
| 85 | def test_all_configs_invalid_raises_runtime_error(patched_autotuner): |
| 86 | cfgs = [triton.Config({"BLOCK": 1024}), triton.Config({"BLOCK": 2048})] |
| 87 | |
| 88 | def always_oor(*args, **kwargs): |
| 89 | raise OutOfResources(required=99999, limit=1, name="shared memory") |
| 90 | |
| 91 | fake_self = _make_fake_self(cfgs, always_oor) |
| 92 | with pytest.raises(RuntimeError, match="[Nn]o valid config"): |
| 93 | patched_autotuner.run(fake_self) |