diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index d0e3d9f497f..89e25f23bea 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -33,12 +33,12 @@ steps: - tests/test_config - tests/test_logger - tests/test_vllm_port - - tests/test_jit_monitor.py + - tests/jit_monitor/test_hooks.py commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py' + pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py' - label: Engine (1 GPU) timeout_in_minutes: 30 diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 3e631fe5d28..e1cdacf02c3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -971,7 +971,7 @@ steps: - tests/test_logger - tests/test_vllm_port commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py #-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------# diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index ce4fc590eec..6099f09ab37 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -23,9 +23,10 @@ steps: - tests/test_config - tests/test_logger - tests/test_vllm_port - - tests/test_jit_monitor.py + - tests/jit_monitor/test_hooks.py + - tests/jit_monitor/test_hooks_gpu.py commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py mirror: amd: dind: false diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml new file mode 100644 index 00000000000..7119c489c01 --- /dev/null +++ b/.buildkite/test_areas/jit_monitor.yaml @@ -0,0 +1,22 @@ +group: JIT Monitor +depends_on: + - image-build +steps: +- label: No Runtime JITs e2e tests + key: jit-monitor-no-runtime-jit + device: h200_35gb + timeout_in_minutes: 45 + source_file_dependencies: + - vllm/utils/jit_monitor.py + - vllm/v1/worker/gpu_worker.py + - vllm/model_executor/warmup/ + - vllm/config/observability.py + - tests/jit_monitor/test_no_runtime_jit.py + - tests/models/registry.py + commands: + # Boot a curated JIT-heavy model set with the JIT monitor in "error" mode + # and run generation; any post-warmup JIT compilation fails the test. + # Per-test watchdog so a wedged engine/CUDA init fails with a traceback + # instead of running until the build timeout. + - export PYTHONFAULTHANDLER=1 + - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread diff --git a/tests/jit_monitor/__init__.py b/tests/jit_monitor/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/jit_monitor/conftest.py b/tests/jit_monitor/conftest.py new file mode 100644 index 00000000000..3f869995a99 --- /dev/null +++ b/tests/jit_monitor/conftest.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.utils import jit_monitor + + +@pytest.fixture(autouse=True) +def _reset_monitor(): + """Reset global monitor state between tests. + + ``activate()`` installs process-global hooks and flips module globals, so + without this every test would inherit the previous test's monitor state. + """ + + def reset(): + jit_monitor._active = False + jit_monitor._mode = "warn" + jit_monitor._verbose = False + jit_monitor._cutedsl_hook_installed = False + jit_monitor._tilelang_hook_installed = False + jit_monitor._tilelang_jitimpl_compile_depth = 0 + + reset() + yield + reset() diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py new file mode 100644 index 00000000000..1285b4c89e8 --- /dev/null +++ b/tests/jit_monitor/test_hooks.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the JIT monitor hooks. Backends are mocked, so no GPU.""" + +import inspect +import os +import sys +from contextlib import contextmanager +from types import ModuleType, SimpleNamespace +from typing import Any, cast +from unittest import mock + +import pytest + +from vllm.utils import jit_monitor + +pytestmark = pytest.mark.cpu_test + + +# ------------------------------------------------------------------ +# Helpers — lightweight stand-ins for the modules ``activate()`` patches +# ------------------------------------------------------------------ + + +def _make_fake_knobs(*, autotuning_print=False, jit_hook=None): + """Build a minimal fake ``triton.knobs`` namespace.""" + autotuning = SimpleNamespace(print=autotuning_print) + runtime = SimpleNamespace(jit_post_compile_hook=jit_hook) + return SimpleNamespace(autotuning=autotuning, runtime=runtime) + + +def _fake_cute_import_modules(compile_fn): + """Fake Python's parent package + submodule for ``import cutlass.cute``.""" + fake_cute = cast(Any, ModuleType("cutlass.cute")) + fake_cute.compile = compile_fn + fake_parent_package = cast(Any, ModuleType("cutlass")) + fake_parent_package.__path__ = [] + fake_parent_package.cute = fake_cute + return { + "cutlass": fake_parent_package, + "cutlass.cute": fake_cute, + } + + +def _fake_cute_compile(*args, **kwargs): + return "compiled" + + +def _fake_tilelang_import_modules(): + """Fake Python's TileLang modules touched by ``jit_monitor.activate``.""" + + class FakeJITKernel: + def __init__(self, *args, **kwargs): + pass + + class FakeJITImpl: + def __init__(self, func, signature): + self.func = func + self.signature = signature + self.mode = "lazy" + self._kernel_cache = {} + + def __call__(self, *args, **kwargs): + key, _ = self.func.parse_args(*args, **kwargs) + kernel = self._kernel_cache.get(key) + if kernel is None: + kernel = "compiled" + self._kernel_cache[key] = kernel + return kernel + + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) + fake_kernel.JITKernel = FakeJITKernel + + fake_jit = cast(Any, ModuleType("tilelang.jit")) + fake_jit.JITImpl = FakeJITImpl + fake_jit.kernel = fake_kernel + + fake_tilelang = cast(Any, ModuleType("tilelang")) + fake_tilelang.jit = fake_jit + + return { + "tilelang": fake_tilelang, + "tilelang.jit": fake_jit, + "tilelang.jit.kernel": fake_kernel, + } + + +@contextmanager +def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile): + """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``.""" + fake_triton = cast(Any, ModuleType("triton")) + fake_triton.knobs = fake_knobs + with ( + mock.patch.dict( + sys.modules, + { + "triton": fake_triton, + **_fake_cute_import_modules(cute_compile), + **_fake_tilelang_import_modules(), + }, + ), + mock.patch.object(jit_monitor, "HAS_TRITON", True), + ): + yield + + +def _triton_hook_kwargs(name: str): + return dict( + key="k", + repr="r", + fn=SimpleNamespace(name=name), + compile=lambda: None, + is_manual_warmup=False, + already_compiled=False, + ) + + +# ------------------------------------------------------------------ +# activate() +# ------------------------------------------------------------------ + + +def test_activate_sets_active(): + assert not jit_monitor.is_active() + with _patch_jit_modules(_make_fake_knobs()): + jit_monitor.activate() + assert jit_monitor.is_active() + + +def test_activate_is_idempotent(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate() + first_hook = fake.runtime.jit_post_compile_hook + jit_monitor.activate() + assert fake.runtime.jit_post_compile_hook is first_hook + + +def test_activate_logs_info(): + with ( + mock.patch.object(jit_monitor.logger, "info") as m, + _patch_jit_modules(_make_fake_knobs()), + ): + jit_monitor.activate() + m.assert_called_once() + assert "Kernel JIT monitor activated" in m.call_args[0][0] + + +def test_activate_rejects_unknown_mode(): + with pytest.raises(ValueError, match="Unsupported JIT monitor mode"): + jit_monitor.activate(mode="panic") # type: ignore[arg-type] + + +def test_activate_without_triton(): + with mock.patch.object(jit_monitor, "HAS_TRITON", False): + jit_monitor.activate() + assert jit_monitor.is_active() + + +# ------------------------------------------------------------------ +# Triton autotuning print +# ------------------------------------------------------------------ + + +def test_autotuning_print_is_enabled(): + fake = _make_fake_knobs(autotuning_print=False) + with _patch_jit_modules(fake): + jit_monitor.activate() + assert fake.autotuning.print is True + + +def test_autotuning_print_respects_user_opt_out(): + fake = _make_fake_knobs(autotuning_print=False) + with ( + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}), + _patch_jit_modules(fake), + ): + jit_monitor.activate() + assert fake.autotuning.print is False + + +def test_autotuning_print_noop_when_user_already_enabled(): + fake = _make_fake_knobs(autotuning_print=True) + with ( + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}), + _patch_jit_modules(fake), + ): + jit_monitor.activate() + assert fake.autotuning.print is True + + +# ------------------------------------------------------------------ +# Triton JIT hook +# ------------------------------------------------------------------ + + +def test_triton_hook_is_registered(): + fake = _make_fake_knobs() + assert fake.runtime.jit_post_compile_hook is None + with _patch_jit_modules(fake): + jit_monitor.activate() + assert fake.runtime.jit_post_compile_hook is not None + + +def test_triton_hook_logs_warning(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + + with ( + mock.patch.object(jit_monitor.logger, "warning_once") as m, + mock.patch.object(jit_monitor.logger, "warning") as warning, + ): + hook(**_triton_hook_kwargs("test_kernel")) + + m.assert_called_once() + warning.assert_not_called() + msg = m.call_args[0][0] % m.call_args[0][1:] + assert "Triton kernel JIT compilation during inference" in msg + assert "test_kernel" in msg + + +def test_triton_hook_chains_existing_hook(): + existing = mock.MagicMock(return_value="existing_result") + fake = _make_fake_knobs(jit_hook=existing) + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + result = hook(**_triton_hook_kwargs("chained_kernel")) + + existing.assert_called_once() + assert result == "existing_result" + + +def test_triton_hook_works_without_existing_hook(): + fake = _make_fake_knobs(jit_hook=None) + with _patch_jit_modules(fake): + jit_monitor.activate() + + hook = fake.runtime.jit_post_compile_hook + assert hook(**_triton_hook_kwargs("solo_kernel")) is None + + +def test_triton_hook_error_mode_raises(): + fake = _make_fake_knobs() + with _patch_jit_modules(fake): + jit_monitor.activate(mode="error") + + hook = fake.runtime.jit_post_compile_hook + with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"): + hook(**_triton_hook_kwargs("error_kernel")) + + +# ------------------------------------------------------------------ +# CuTeDSL hook +# ------------------------------------------------------------------ + + +def test_cutedsl_compile_logs_warning(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile(lambda: None, "arg", option=True) + + assert result == "compiled" + warning_once.assert_called_once() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "CuTeDSL JIT compilation during inference" in msg + + +def test_cutedsl_compile_logs_verbose_warning(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate(verbose=True) + with mock.patch.object(jit_monitor.logger, "warning") as warning: + result = cute.compile(lambda: None, "arg", option=True) + + assert result == "compiled" + warning.assert_called_once() + msg = warning.call_args[0][0] % warning.call_args[0][1:] + assert "CuTeDSL JIT compilation during inference" in msg + + +def test_cutedsl_error_mode_raises(): + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): + import cutlass.cute as cute + + jit_monitor.activate(mode="error") + with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): + cute.compile(lambda: None, "arg", option=True) + + +def test_cutedsl_subscripted_compile_is_monitored(): + """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" + + class FakeCompileCallable: + def __getitem__(self, options): + return self + + def __call__(self, *args, **kwargs): + return "compiled" + + with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile[("opt_level", 3)](lambda: None, "arg") + + assert result == "compiled" + warning_once.assert_called_once() + + +# ------------------------------------------------------------------ +# TileLang hook +# ------------------------------------------------------------------ + + +def test_tilelang_jit_kernel_logs_warning(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"}) + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi") + + warning_once.assert_called_once() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "TileLang JIT compilation during inference" in msg + assert "tl_kernel" in msg + + +def test_tilelang_jit_impl_logs_warning(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit import JITImpl + + def tilelang_fn( + gemm_out_mul, + hidden_size: int, + n_splits: int = 1, + hc_mult: int = 4, + ): + return None + + class FakeFunc: + orig_func = tilelang_fn + + def parse_args(self, *args, **kwargs): + return ( + ( + "tilelang_key", + kwargs["hidden_size"], + kwargs.get("n_splits", 1), + ), + {}, + ) + + def set_mode(self, mode): + self.mode = mode + + tensor = SimpleNamespace( + shape=(2, 16, 24), + dtype="float32", + device="cuda:0", + ) + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) + + jit_monitor.activate() + with ( + mock.patch.object(jit_monitor.logger, "warning_once") as warning_once, + mock.patch.object(jit_monitor.logger, "warning") as warning, + ): + impl(tensor, hidden_size=7168, n_splits=2) + + warning_once.assert_called_once() + warning.assert_not_called() + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] + assert "TileLang JIT compilation during inference" in msg + assert "tilelang_fn" in msg + + +def test_tilelang_jit_impl_does_not_log_on_cache_hit(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit import JITImpl + + def tilelang_fn(gemm_out_mul, n_splits: int = 1): + return None + + class FakeFunc: + orig_func = tilelang_fn + + def parse_args(self, *args, **kwargs): + return (("tilelang_key", kwargs.get("n_splits", 1)), {}) + + def set_mode(self, mode): + self.mode = mode + + tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32") + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + impl(tensor, n_splits=2) + impl(tensor, n_splits=2) + + warning_once.assert_called_once() + + +def test_tilelang_from_database_does_not_log(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"}) + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + JITKernel(func=func, from_database=True) + + warning_once.assert_not_called() + + +def test_tilelang_error_mode_raises(): + with _patch_jit_modules(_make_fake_knobs()): + from tilelang.jit.kernel import JITKernel + + func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"}) + jit_monitor.activate(mode="error") + with pytest.raises(RuntimeError, match="TileLang JIT compilation"): + JITKernel(func=func) diff --git a/tests/jit_monitor/test_hooks_gpu.py b/tests/jit_monitor/test_hooks_gpu.py new file mode 100644 index 00000000000..7c4b93abcd3 --- /dev/null +++ b/tests/jit_monitor/test_hooks_gpu.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end JIT monitor tests: real Triton kernel, real GPU, real hook.""" + +from unittest import mock + +import pytest + +from vllm.utils import jit_monitor + +try: + import torch + + _HAS_CUDA = torch.cuda.is_available() +except ImportError: + _HAS_CUDA = False + +try: + import triton + import triton.language as tl + + _HAS_TRITON = True +except ImportError: + _HAS_TRITON = False + +pytestmark = pytest.mark.skipif( + not (_HAS_CUDA and _HAS_TRITON), + reason="Requires CUDA GPU and Triton", +) + + +if _HAS_TRITON: + + @triton.jit + def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + y = tl.load(y_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x + y, mask=mask) + + +def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None: + """Launch ``_add_kernel`` with vectors of length *n*.""" + x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment + y = torch.randn(n, device="cuda") + out = torch.empty(n, device="cuda") + grid = ((n + block - 1) // block,) + _add_kernel[grid](x, y, out, n, BLOCK=block) + torch.accelerator.synchronize() + + +def test_no_warning_on_cached_shape(): + _run_add_kernel(1024) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as w: + _run_add_kernel(1024) + w.assert_not_called() + + +def test_warning_on_new_constexpr(): + _run_add_kernel(1024, block=256) + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as w: + # Different BLOCK (a tl.constexpr) forces recompilation. + _run_add_kernel(1024, block=512) + w.assert_called() + msg = w.call_args[0][0] % w.call_args[0][1:] + assert "_add_kernel" in msg + + +def test_verbose_warning_on_each_new_pointer_alignment(): + _run_add_kernel(1024) + + jit_monitor.activate(verbose=True) + with ( + mock.patch.object(jit_monitor.logger, "warning") as w, + mock.patch.object(jit_monitor.logger, "warning_once") as w_once, + ): + _run_add_kernel(1024, offset=1) + assert w.called + w_once.assert_not_called() diff --git a/tests/jit_monitor/test_no_runtime_jit.py b/tests/jit_monitor/test_no_runtime_jit.py new file mode 100644 index 00000000000..86afd626b34 --- /dev/null +++ b/tests/jit_monitor/test_no_runtime_jit.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Catch runtime (post-warmup) JIT compilations for JIT-heavy backends running +e2e tests on popular models, for which we only load a few blocks for performance. + +NOTE(NickLucche) With cuda graphs on, kernels fully covered by graphs captured during +warmup do not re-trigger the Python JIT hooks. The targeted paths (prefill +MoE/MLA/SSM and the sampler) run mixed, so they are unaffected. +""" + +from dataclasses import dataclass + +import pytest + +from vllm import LLM, SamplingParams +from vllm.inputs import TokensPrompt + +from ..models.utils import dummy_hf_overrides +from ..utils import create_new_process_for_each_test + +# Warmup coverage is still incomplete for these backends, so the monitor fires +# during inference. Tracked in https://github.com/vllm-project/vllm/issues/49349; +# drop this once the warmup contract migrations land. +pytestmark = pytest.mark.skip(reason="Kernel warmup coverage is still incomplete") + + +@dataclass(frozen=True) +class JitModel: + model: str + draft: str | None = None + trust_remote_code: bool = False + + +JIT_MONITOR_MODELS = [ + JitModel("Qwen/Qwen3-0.6B"), + JitModel("deepseek-ai/DeepSeek-V2-Lite-Chat", trust_remote_code=True), + JitModel("deepseek-ai/DeepSeek-V3", trust_remote_code=True), + JitModel("ibm-granite/granite-4.0-tiny-preview"), + JitModel( + "luccafong/deepseek_mtp_main_random", + draft="luccafong/deepseek_mtp_draft_random", + trust_remote_code=True, + ), + JitModel( + "eagle618/deepseek-v3-random", + draft="eagle618/eagle-deepseek-v3-random", + trust_remote_code=True, + ), +] + + +def _run_shape_battery(llm: LLM) -> None: + """Exercise diverse compile keys so missing warmup keys surface. + + Token-id prompts keep shapes exact and avoid depending on a tokenizer. + Outputs are meaningless under dummy weights; we assert only that no JIT + fired. + """ + short = TokensPrompt(prompt_token_ids=[1, 2, 3, 4]) + medium = TokensPrompt(prompt_token_ids=list(range(1, 33))) + long = TokensPrompt(prompt_token_ids=list(range(1, 129))) + + # Greedy single-sequence multi-step decode: prefill + autoregressive decode + # + greedy sampler. + llm.generate(medium, SamplingParams(temperature=0.0, max_tokens=16)) + + # Batched prefill with mixed lengths: varlen prefill + padded decode. + llm.generate([short, medium, long], SamplingParams(temperature=0.0, max_tokens=8)) + + # Triton sampler kernels: top_k / top_p / min_p each specialize. + for sampling_params in ( + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0), + SamplingParams( + temperature=0.8, top_k=20, top_p=0.9, min_p=0.1, max_tokens=8, seed=0 + ), + ): + llm.generate(medium, sampling_params) + + # Heterogeneous SamplingParams in one step, where missing sampler warmup + # keys most often hide. + llm.generate( + [medium] * 4, + [ + SamplingParams(temperature=0.0, max_tokens=8), + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0), + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0), + ], + ) + + +@create_new_process_for_each_test("spawn") +def can_run_without_jit(spec: JitModel): + """Boot ``spec`` with the monitor armed and run the shape battery. + + A subprocess per model is required: the monitor's hooks are process-global + and, once armed in ``error`` mode, stay armed. It must be spawned rather + than forked, since forking a pytest process that already initialized CUDA + poisons the child. + """ + llm = LLM( + spec.model, + trust_remote_code=spec.trust_remote_code, + max_model_len=2048, + max_num_seqs=8, + gpu_memory_utilization=0.80, + load_format="dummy", + hf_overrides=dummy_hf_overrides, + # cuda graphs cover captured decode shapes, run eager. + enforce_eager=False, + jit_monitor_mode="error", + speculative_config={ + "model": spec.draft, + "num_speculative_tokens": 2, + } + if spec.draft + else None, + ) + + try: + _run_shape_battery(llm) + except Exception as e: + # The monitor's message contains "during inference"; distinguish a real + # JIT miss from an unrelated crash. + if "during inference" in str(e): + pytest.fail( + f"{spec.model}: post-warmup JIT compilation detected - a warmup " + f"key is missing for a shape in the battery.\n{e}" + ) + raise + + +@pytest.mark.parametrize("spec", JIT_MONITOR_MODELS, ids=lambda s: s.model) +def test_no_runtime_jit(spec: JitModel, monkeypatch: pytest.MonkeyPatch): + """Assert JIT-heavy backends do not JIT-compile during inference.""" + # Set here rather than in the child so the spawned process inherits it: + # the engine core must not be forked once the test process has CUDA up. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + can_run_without_jit(spec) diff --git a/tests/test_jit_monitor.py b/tests/test_jit_monitor.py deleted file mode 100644 index 50261a479d9..00000000000 --- a/tests/test_jit_monitor.py +++ /dev/null @@ -1,534 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import inspect -import os -import sys -from contextlib import contextmanager -from types import ModuleType, SimpleNamespace -from typing import Any, cast -from unittest import mock - -import pytest - -from vllm.utils import jit_monitor - - -@pytest.fixture(autouse=True) -def _reset_monitor(): - """Reset global monitor state between tests.""" - jit_monitor._active = False - jit_monitor._mode = "warn" - jit_monitor._verbose = False - jit_monitor._cutedsl_hook_installed = False - jit_monitor._tilelang_hook_installed = False - jit_monitor._tilelang_jitimpl_compile_depth = 0 - yield - jit_monitor._active = False - jit_monitor._mode = "warn" - jit_monitor._verbose = False - jit_monitor._cutedsl_hook_installed = False - jit_monitor._tilelang_hook_installed = False - jit_monitor._tilelang_jitimpl_compile_depth = 0 - - -# ------------------------------------------------------------------ -# Helpers — lightweight stand-ins for the modules ``activate()`` patches -# ------------------------------------------------------------------ - - -def _make_fake_knobs(*, autotuning_print=False, jit_hook=None): - """Build a minimal fake ``triton.knobs`` namespace.""" - autotuning = SimpleNamespace(print=autotuning_print) - runtime = SimpleNamespace(jit_post_compile_hook=jit_hook) - return SimpleNamespace(autotuning=autotuning, runtime=runtime) - - -def _fake_cute_import_modules(compile_fn): - """Fake Python's parent package + submodule for ``import cutlass.cute``.""" - fake_cute = cast(Any, ModuleType("cutlass.cute")) - fake_cute.compile = compile_fn - fake_parent_package = cast(Any, ModuleType("cutlass")) - fake_parent_package.__path__ = [] - fake_parent_package.cute = fake_cute - return { - "cutlass": fake_parent_package, - "cutlass.cute": fake_cute, - } - - -def _fake_cute_compile(*args, **kwargs): - return "compiled" - - -def _fake_tilelang_import_modules(): - """Fake Python's TileLang modules touched by ``jit_monitor.activate``.""" - - class FakeJITKernel: - def __init__(self, *args, **kwargs): - pass - - class FakeJITImpl: - def __init__(self, func, signature): - self.func = func - self.signature = signature - self.mode = "lazy" - self._kernel_cache = {} - - def __call__(self, *args, **kwargs): - key, _ = self.func.parse_args(*args, **kwargs) - kernel = self._kernel_cache.get(key) - if kernel is None: - kernel = "compiled" - self._kernel_cache[key] = kernel - return kernel - - fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) - fake_kernel.JITKernel = FakeJITKernel - - fake_jit = cast(Any, ModuleType("tilelang.jit")) - fake_jit.JITImpl = FakeJITImpl - fake_jit.kernel = fake_kernel - - fake_tilelang = cast(Any, ModuleType("tilelang")) - fake_tilelang.jit = fake_jit - - return { - "tilelang": fake_tilelang, - "tilelang.jit": fake_jit, - "tilelang.jit.kernel": fake_kernel, - } - - -@contextmanager -def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile): - """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``.""" - fake_triton = cast(Any, ModuleType("triton")) - fake_triton.knobs = fake_knobs - with ( - mock.patch.dict( - sys.modules, - { - "triton": fake_triton, - **_fake_cute_import_modules(cute_compile), - **_fake_tilelang_import_modules(), - }, - ), - mock.patch.object(jit_monitor, "HAS_TRITON", True), - ): - yield - - -# ------------------------------------------------------------------ -# Unit tests (no GPU required, triton is mocked) -# ------------------------------------------------------------------ - - -class TestActivateBasic: - def test_sets_active(self): - assert not jit_monitor.is_active() - with _patch_jit_modules(_make_fake_knobs()): - jit_monitor.activate() - assert jit_monitor.is_active() - - def test_idempotent(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate() - first_hook = fake.runtime.jit_post_compile_hook - jit_monitor.activate() - assert fake.runtime.jit_post_compile_hook is first_hook - - def test_logs_info_on_activation(self): - with ( - mock.patch.object(jit_monitor.logger, "info") as m, - _patch_jit_modules(_make_fake_knobs()), - ): - jit_monitor.activate() - m.assert_called_once() - assert "Kernel JIT monitor activated" in m.call_args[0][0] - - def test_rejects_unknown_mode(self): - with pytest.raises(ValueError, match="Unsupported JIT monitor mode"): - jit_monitor.activate(mode="panic") # type: ignore[arg-type] - - -class TestAutotuningPrint: - def test_enables_autotuning_print(self): - fake = _make_fake_knobs(autotuning_print=False) - with _patch_jit_modules(fake): - jit_monitor.activate() - assert fake.autotuning.print is True - - def test_respects_user_opt_out(self): - fake = _make_fake_knobs(autotuning_print=False) - with ( - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}), - _patch_jit_modules(fake), - ): - jit_monitor.activate() - assert fake.autotuning.print is False - - def test_noop_when_user_already_enabled(self): - fake = _make_fake_knobs(autotuning_print=True) - with ( - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}), - _patch_jit_modules(fake), - ): - jit_monitor.activate() - assert fake.autotuning.print is True - - -class TestTritonJitHook: - def test_hook_registered(self): - fake = _make_fake_knobs() - assert fake.runtime.jit_post_compile_hook is None - with _patch_jit_modules(fake): - jit_monitor.activate() - assert fake.runtime.jit_post_compile_hook is not None - - def test_hook_logs_warning(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="test_kernel") - - with ( - mock.patch.object(jit_monitor.logger, "warning_once") as m, - mock.patch.object(jit_monitor.logger, "warning") as warning, - ): - hook( - key="some_key", - repr="some_repr", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - - m.assert_called_once() - warning.assert_not_called() - msg = m.call_args[0][0] % m.call_args[0][1:] - assert "Triton kernel JIT compilation during inference" in msg - assert "test_kernel" in msg - - def test_hook_chains_existing_hook(self): - existing = mock.MagicMock(return_value="existing_result") - fake = _make_fake_knobs(jit_hook=existing) - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="chained_kernel") - kwargs = dict( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - result = hook(**kwargs) - - existing.assert_called_once() - assert result == "existing_result" - - def test_hook_works_without_existing_hook(self): - fake = _make_fake_knobs(jit_hook=None) - with _patch_jit_modules(fake): - jit_monitor.activate() - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="solo_kernel") - result = hook( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - assert result is None - - def test_error_mode_raises(self): - fake = _make_fake_knobs() - with _patch_jit_modules(fake): - jit_monitor.activate(mode="error") - - hook = fake.runtime.jit_post_compile_hook - mock_fn = SimpleNamespace(name="error_kernel") - with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"): - hook( - key="k", - repr="r", - fn=mock_fn, - compile=lambda: None, - is_manual_warmup=False, - already_compiled=False, - ) - - -class TestNoTritonFallback: - def test_activate_without_triton(self): - with mock.patch.object(jit_monitor, "HAS_TRITON", False): - jit_monitor.activate() - assert jit_monitor.is_active() - - -class TestCuTeDSLHook: - def test_compile_logs_warning(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - result = cute.compile(lambda: None, "arg", option=True) - - assert result == "compiled" - warning_once.assert_called_once() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "CuTeDSL JIT compilation during inference" in msg - - def test_compile_logs_verbose_warning(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate(verbose=True) - with mock.patch.object(jit_monitor.logger, "warning") as warning: - result = cute.compile(lambda: None, "arg", option=True) - - assert result == "compiled" - warning.assert_called_once() - msg = warning.call_args[0][0] % warning.call_args[0][1:] - assert "CuTeDSL JIT compilation during inference" in msg - - def test_error_mode_raises(self): - def compile_fn(*args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn): - import cutlass.cute as cute - - jit_monitor.activate(mode="error") - with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): - cute.compile(lambda: None, "arg", option=True) - - def test_subscripted_compile_is_monitored(self): - """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" - - class FakeCompileCallable: - def __getitem__(self, options): - return self - - def __call__(self, *args, **kwargs): - return "compiled" - - with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): - import cutlass.cute as cute - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - result = cute.compile[("opt_level", 3)](lambda: None, "arg") - - assert result == "compiled" - warning_once.assert_called_once() - - -class TestTileLangHook: - def test_jit_kernel_logs_warning(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"}) - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi") - - warning_once.assert_called_once() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "TileLang JIT compilation during inference" in msg - assert "tl_kernel" in msg - - def test_jit_impl_logs_warning(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit import JITImpl - - def tilelang_fn( - gemm_out_mul, - hidden_size: int, - n_splits: int = 1, - hc_mult: int = 4, - ): - return None - - class FakeFunc: - orig_func = tilelang_fn - - def parse_args(self, *args, **kwargs): - return ( - ( - "tilelang_key", - kwargs["hidden_size"], - kwargs.get("n_splits", 1), - ), - {}, - ) - - def set_mode(self, mode): - self.mode = mode - - tensor = SimpleNamespace( - shape=(2, 16, 24), - dtype="float32", - device="cuda:0", - ) - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) - - jit_monitor.activate() - with ( - mock.patch.object(jit_monitor.logger, "warning_once") as warning_once, - mock.patch.object(jit_monitor.logger, "warning") as warning, - ): - impl(tensor, hidden_size=7168, n_splits=2) - - warning_once.assert_called_once() - warning.assert_not_called() - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] - assert "TileLang JIT compilation during inference" in msg - assert "tilelang_fn" in msg - - def test_jit_impl_does_not_log_on_cache_hit(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit import JITImpl - - def tilelang_fn(gemm_out_mul, n_splits: int = 1): - return None - - class FakeFunc: - orig_func = tilelang_fn - - def parse_args(self, *args, **kwargs): - return (("tilelang_key", kwargs.get("n_splits", 1)), {}) - - def set_mode(self, mode): - self.mode = mode - - tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32") - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - impl(tensor, n_splits=2) - impl(tensor, n_splits=2) - - warning_once.assert_called_once() - - def test_from_database_does_not_log(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"}) - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: - JITKernel(func=func, from_database=True) - - warning_once.assert_not_called() - - def test_error_mode_raises(self): - with _patch_jit_modules(_make_fake_knobs()): - from tilelang.jit.kernel import JITKernel - - func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"}) - jit_monitor.activate(mode="error") - with pytest.raises(RuntimeError, match="TileLang JIT compilation"): - JITKernel(func=func) - - -# ------------------------------------------------------------------ -# Integration tests (real Triton + GPU) -# ------------------------------------------------------------------ - -try: - import torch - - _HAS_CUDA = torch.cuda.is_available() -except ImportError: - _HAS_CUDA = False - -try: - import triton - import triton.language as tl - - _HAS_TRITON = True -except ImportError: - _HAS_TRITON = False - -_skip_no_gpu = pytest.mark.skipif( - not (_HAS_CUDA and _HAS_TRITON), - reason="Requires CUDA GPU and Triton", -) - - -if _HAS_TRITON: - - @triton.jit - def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n - x = tl.load(x_ptr + offs, mask=mask) - y = tl.load(y_ptr + offs, mask=mask) - tl.store(out_ptr + offs, x + y, mask=mask) - - -def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None: - """Launch ``_add_kernel`` with vectors of length *n*.""" - x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment - y = torch.randn(n, device="cuda") - out = torch.empty(n, device="cuda") - grid = ((n + block - 1) // block,) - _add_kernel[grid](x, y, out, n, BLOCK=block) - torch.accelerator.synchronize() - - -@_skip_no_gpu -class TestTritonJitHookIntegration: - """End-to-end: real Triton kernel, real GPU, real hook.""" - - def test_no_warning_on_cached_shape(self): - _run_add_kernel(1024) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as w: - _run_add_kernel(1024) - w.assert_not_called() - - def test_warning_on_new_constexpr(self): - _run_add_kernel(1024, block=256) - - jit_monitor.activate() - with mock.patch.object(jit_monitor.logger, "warning_once") as w: - # Different BLOCK (a tl.constexpr) forces recompilation. - _run_add_kernel(1024, block=512) - w.assert_called() - msg = w.call_args[0][0] % w.call_args[0][1:] - assert "_add_kernel" in msg - - def test_verbose_warning_on_each_new_pointer_alignment(self): - _run_add_kernel(1024) - - jit_monitor.activate(verbose=True) - with ( - mock.patch.object(jit_monitor.logger, "warning") as w, - mock.patch.object(jit_monitor.logger, "warning_once") as w_once, - ): - _run_add_kernel(1024, offset=1) - assert w.called - w_once.assert_not_called()