wehub-resource-sync
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Targeted unit tests for VLLM_TRITON_FORCE_FIRST_CONFIG.
|
|
|
|
These tests exercise only the patched `Autotuner.run` logic installed by
|
|
`vllm.triton_utils.force_first_config.install`. The wrapped kernel is a
|
|
plain callable so the tests run on CPU-only hosts (no GPU, no actual
|
|
kernel launch) as long as the `triton` package is importable.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm.triton_utils import HAS_TRITON, triton
|
|
|
|
if not HAS_TRITON:
|
|
pytest.skip("triton not available", allow_module_level=True)
|
|
|
|
from vllm.triton_utils import force_first_config # noqa: E402
|
|
|
|
OutOfResources = triton.runtime.errors.OutOfResources
|
|
|
|
|
|
@pytest.fixture
|
|
def patched_autotuner(monkeypatch: pytest.MonkeyPatch):
|
|
"""Install the first-valid-config patch and restore after.
|
|
|
|
The env-var gate lives in vllm.env_override; install() itself does not
|
|
read the environment, so the test calls it directly.
|
|
"""
|
|
Autotuner = triton.runtime.autotuner.Autotuner
|
|
original_run = Autotuner.run
|
|
# Reset the once-only guard so install() re-runs for each test.
|
|
monkeypatch.setattr(force_first_config, "_installed", False)
|
|
force_first_config.install()
|
|
yield Autotuner
|
|
Autotuner.run = original_run
|
|
|
|
|
|
def _make_fake_self(configs, fn):
|
|
"""Minimal stand-in for an Autotuner instance."""
|
|
return SimpleNamespace(
|
|
configs=configs,
|
|
keys=[],
|
|
arg_names=[],
|
|
base_fn=fn,
|
|
fn=fn,
|
|
best_config=None,
|
|
)
|
|
|
|
|
|
def test_skips_invalid_first_config_and_caches_second(patched_autotuner):
|
|
bad = triton.Config({"BLOCK": 1024})
|
|
good = triton.Config({"BLOCK": 64})
|
|
calls = []
|
|
|
|
def fake_fn(*args, **kwargs):
|
|
calls.append(kwargs["BLOCK"])
|
|
if kwargs["BLOCK"] == 1024:
|
|
raise OutOfResources(required=99999, limit=1, name="shared memory")
|
|
return "ok"
|
|
|
|
fake_self = _make_fake_self([bad, good], fake_fn)
|
|
|
|
# First call: walks past the invalid config, picks the second.
|
|
assert patched_autotuner.run(fake_self) == "ok"
|
|
assert calls == [1024, 64]
|
|
assert fake_self.best_config is good
|
|
|
|
# Second call: cached index is reused, invalid config is NOT retried.
|
|
calls.clear()
|
|
assert patched_autotuner.run(fake_self) == "ok"
|
|
assert calls == [64]
|
|
|
|
|
|
def test_empty_configs_falls_back_to_direct_call(patched_autotuner):
|
|
def fake_fn(*args, **kwargs):
|
|
return "direct"
|
|
|
|
fake_self = _make_fake_self([], fake_fn)
|
|
assert patched_autotuner.run(fake_self) == "direct"
|
|
|
|
|
|
def test_all_configs_invalid_raises_runtime_error(patched_autotuner):
|
|
cfgs = [triton.Config({"BLOCK": 1024}), triton.Config({"BLOCK": 2048})]
|
|
|
|
def always_oor(*args, **kwargs):
|
|
raise OutOfResources(required=99999, limit=1, name="shared memory")
|
|
|
|
fake_self = _make_fake_self(cfgs, always_oor)
|
|
with pytest.raises(RuntimeError, match="[Nn]o valid config"):
|
|
patched_autotuner.run(fake_self)
|