[Kernel][CI] --jit-monitor-mode error e2e tests for kernel warmup infra (#50109)

Signed-off-by: NickLucche <[email protected]>

.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

.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  --------------------------------------------------------#
 

.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

.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

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()

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)

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()

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)

tests/test_jit_monitor.py

@@ -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()