[Feat] Add runtime monitor for post-warmup CuTeDSL compilation (#46167)
Roberto L. Castro committed
Jun 24, 2026 at 01:33 UTC
d86c66c981006b2e7f6eef3b32145a4ddc67425a
7 files changed
+366
-172
tests/engine/test_arg_utils.py
+11
@@ -214,6 +214,17 @@ def test_jit_monitor_verbose_arg():
214
assert EngineArgs(model="test", jit_monitor_verbose=True).jit_monitor_verbose
215
216
217
+@pytest.mark.parametrize("mode", ["warn", "error"])
218
+def test_jit_monitor_mode_arg(mode):
219
+ parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
220
+ args = parser.parse_args(["--jit-monitor-mode", mode])
221
+
222
+ assert args.jit_monitor_mode == mode
223
+ engine_args = EngineArgs(model="test", jit_monitor_mode=mode)
224
+ assert engine_args.jit_monitor_mode == mode
225
+ assert engine_args.create_observability_config().jit_monitor_mode == mode
226
+
227
+
228
def test_hf_token_get_kwargs():
229
kwargs = get_kwargs(ModelConfig)["hf_token"]
230
tests/test_jit_monitor.py
+113
-18
@@ -3,26 +3,31 @@
3
import os
4
import sys
5
from contextlib import contextmanager
6
-from types import SimpleNamespace
6
+from types import ModuleType, SimpleNamespace
7
+from typing import Any, cast
8
from unittest import mock
9
10
import pytest
11
11
-from vllm.triton_utils import jit_monitor
12
+from vllm.utils import jit_monitor
13
14
15
@pytest.fixture(autouse=True)
16
def _reset_monitor():
17
"""Reset global monitor state between tests."""
18
jit_monitor._active = False
19
+ jit_monitor._mode = "warn"
20
jit_monitor._verbose = False
21
+ jit_monitor._cutedsl_hook_installed = False
22
yield
23
jit_monitor._active = False
24
+ jit_monitor._mode = "warn"
25
jit_monitor._verbose = False
26
+ jit_monitor._cutedsl_hook_installed = False
27
28
29
# ------------------------------------------------------------------
25
-# Helpers — lightweight stand-ins for triton.knobs
30
+# Helpers — lightweight stand-ins for the modules ``activate()`` patches
31
# ------------------------------------------------------------------
32
33
@@ -33,12 +38,36 @@ def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
38
return SimpleNamespace(autotuning=autotuning, runtime=runtime)
39
40
41
+def _fake_cute_import_modules(compile_fn):
42
+ """Fake Python's parent package + submodule for ``import cutlass.cute``."""
43
+ fake_cute = cast(Any, ModuleType("cutlass.cute"))
44
+ fake_cute.compile = compile_fn
45
+ fake_parent_package = cast(Any, ModuleType("cutlass"))
46
+ fake_parent_package.__path__ = []
47
+ fake_parent_package.cute = fake_cute
48
+ return {
49
+ "cutlass": fake_parent_package,
50
+ "cutlass.cute": fake_cute,
51
+ }
52
+
53
+
54
+def _fake_cute_compile(*args, **kwargs):
55
+ return "compiled"
56
+
57
+
58
@contextmanager
37
-def _patch_triton_knobs(fake_knobs):
38
- """Context manager that makes ``from triton import knobs`` return *fake_knobs*."""
39
- fake_triton = SimpleNamespace(knobs=fake_knobs)
59
+def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
60
+ """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``."""
61
+ fake_triton = cast(Any, ModuleType("triton"))
62
+ fake_triton.knobs = fake_knobs
63
with (
41
- mock.patch.dict(sys.modules, {"triton": fake_triton}),
64
+ mock.patch.dict(
65
+ sys.modules,
66
+ {
67
+ "triton": fake_triton,
68
+ **_fake_cute_import_modules(cute_compile),
69
+ },
70
+ ),
71
mock.patch.object(jit_monitor, "HAS_TRITON", True),
72
):
73
yield
@@ -52,13 +81,13 @@ def _patch_triton_knobs(fake_knobs):
81
class TestActivateBasic:
82
def test_sets_active(self):
83
assert not jit_monitor.is_active()
55
- with _patch_triton_knobs(_make_fake_knobs()):
84
+ with _patch_jit_modules(_make_fake_knobs()):
85
jit_monitor.activate()
86
assert jit_monitor.is_active()
87
88
def test_idempotent(self):
89
fake = _make_fake_knobs()
61
- with _patch_triton_knobs(fake):
90
+ with _patch_jit_modules(fake):
91
jit_monitor.activate()
92
first_hook = fake.runtime.jit_post_compile_hook
93
jit_monitor.activate()
@@ -67,17 +96,21 @@ class TestActivateBasic:
96
def test_logs_info_on_activation(self):
97
with (
98
mock.patch.object(jit_monitor.logger, "info") as m,
70
- _patch_triton_knobs(_make_fake_knobs()),
99
+ _patch_jit_modules(_make_fake_knobs()),
100
):
101
jit_monitor.activate()
102
m.assert_called_once()
103
assert "Kernel JIT monitor activated" in m.call_args[0][0]
104
105
+ def test_rejects_unknown_mode(self):
106
+ with pytest.raises(ValueError, match="Unsupported JIT monitor mode"):
107
+ jit_monitor.activate(mode="panic") # type: ignore[arg-type]
108
+
109
110
class TestAutotuningPrint:
111
def test_enables_autotuning_print(self):
112
fake = _make_fake_knobs(autotuning_print=False)
80
- with _patch_triton_knobs(fake):
113
+ with _patch_jit_modules(fake):
114
jit_monitor.activate()
115
assert fake.autotuning.print is True
116
@@ -85,7 +118,7 @@ class TestAutotuningPrint:
118
fake = _make_fake_knobs(autotuning_print=False)
119
with (
120
mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}),
88
- _patch_triton_knobs(fake),
121
+ _patch_jit_modules(fake),
122
):
123
jit_monitor.activate()
124
assert fake.autotuning.print is False
@@ -94,23 +127,23 @@ class TestAutotuningPrint:
127
fake = _make_fake_knobs(autotuning_print=True)
128
with (
129
mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}),
97
- _patch_triton_knobs(fake),
130
+ _patch_jit_modules(fake),
131
):
132
jit_monitor.activate()
133
assert fake.autotuning.print is True
134
135
103
-class TestJitHook:
136
+class TestTritonJitHook:
137
def test_hook_registered(self):
138
fake = _make_fake_knobs()
139
assert fake.runtime.jit_post_compile_hook is None
107
- with _patch_triton_knobs(fake):
140
+ with _patch_jit_modules(fake):
141
jit_monitor.activate()
142
assert fake.runtime.jit_post_compile_hook is not None
143
144
def test_hook_logs_warning(self):
145
fake = _make_fake_knobs()
113
- with _patch_triton_knobs(fake):
146
+ with _patch_jit_modules(fake):
147
jit_monitor.activate()
148
149
hook = fake.runtime.jit_post_compile_hook
@@ -138,7 +171,7 @@ class TestJitHook:
171
def test_hook_chains_existing_hook(self):
172
existing = mock.MagicMock(return_value="existing_result")
173
fake = _make_fake_knobs(jit_hook=existing)
141
- with _patch_triton_knobs(fake):
174
+ with _patch_jit_modules(fake):
175
jit_monitor.activate()
176
177
hook = fake.runtime.jit_post_compile_hook
@@ -158,7 +191,7 @@ class TestJitHook:
191
192
def test_hook_works_without_existing_hook(self):
193
fake = _make_fake_knobs(jit_hook=None)
161
- with _patch_triton_knobs(fake):
194
+ with _patch_jit_modules(fake):
195
jit_monitor.activate()
196
197
hook = fake.runtime.jit_post_compile_hook
@@ -173,6 +206,23 @@ class TestJitHook:
206
)
207
assert result is None
208
209
+ def test_error_mode_raises(self):
210
+ fake = _make_fake_knobs()
211
+ with _patch_jit_modules(fake):
212
+ jit_monitor.activate(mode="error")
213
+
214
+ hook = fake.runtime.jit_post_compile_hook
215
+ mock_fn = SimpleNamespace(name="error_kernel")
216
+ with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"):
217
+ hook(
218
+ key="k",
219
+ repr="r",
220
+ fn=mock_fn,
221
+ compile=lambda: None,
222
+ is_manual_warmup=False,
223
+ already_compiled=False,
224
+ )
225
+
226
227
class TestNoTritonFallback:
228
def test_activate_without_triton(self):
@@ -181,6 +231,51 @@ class TestNoTritonFallback:
231
assert jit_monitor.is_active()
232
233
234
+class TestCuTeDSLHook:
235
+ def test_compile_logs_warning(self):
236
+ def compile_fn(*args, **kwargs):
237
+ return "compiled"
238
+
239
+ with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
240
+ import cutlass.cute as cute
241
+
242
+ jit_monitor.activate()
243
+ with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
244
+ result = cute.compile(lambda: None, "arg", option=True)
245
+
246
+ assert result == "compiled"
247
+ warning_once.assert_called_once()
248
+ msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
249
+ assert "CuTeDSL JIT compilation during inference" in msg
250
+
251
+ def test_compile_logs_verbose_warning(self):
252
+ def compile_fn(*args, **kwargs):
253
+ return "compiled"
254
+
255
+ with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
256
+ import cutlass.cute as cute
257
+
258
+ jit_monitor.activate(verbose=True)
259
+ with mock.patch.object(jit_monitor.logger, "warning") as warning:
260
+ result = cute.compile(lambda: None, "arg", option=True)
261
+
262
+ assert result == "compiled"
263
+ warning.assert_called_once()
264
+ msg = warning.call_args[0][0] % warning.call_args[0][1:]
265
+ assert "CuTeDSL JIT compilation during inference" in msg
266
+
267
+ def test_error_mode_raises(self):
268
+ def compile_fn(*args, **kwargs):
269
+ return "compiled"
270
+
271
+ with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
272
+ import cutlass.cute as cute
273
+
274
+ jit_monitor.activate(mode="error")
275
+ with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"):
276
+ cute.compile(lambda: None, "arg", option=True)
277
+
278
+
279
# ------------------------------------------------------------------
280
# Integration tests (real Triton + GPU)
281
# ------------------------------------------------------------------
vllm/config/observability.py
+4
-1
@@ -76,8 +76,11 @@ class ObservabilityConfig:
76
This includes number of context/generation requests and tokens
77
and the elapsed cpu time for the iteration."""
78
79
+ jit_monitor_mode: Literal["warn", "error"] = "warn"
80
+ """How to handle post-warmup JIT compilation events."""
81
+
82
jit_monitor_verbose: bool = False
80
- """Log every Triton JIT compile with its dispatch key. This can emit many
83
+ """Log every monitored JIT compile with runtime details. This can emit many
84
logs and add overhead, so it is intended for debugging."""
85
86
@cached_property
vllm/engine/arg_utils.py
+22
-13
@@ -643,6 +643,7 @@ class EngineArgs:
643
enable_logging_iteration_details: bool = (
644
ObservabilityConfig.enable_logging_iteration_details
645
)
646
+ jit_monitor_mode: Literal["warn", "error"] = ObservabilityConfig.jit_monitor_mode
647
jit_monitor_verbose: bool = ObservabilityConfig.jit_monitor_verbose
648
enable_mm_processor_stats: bool = ObservabilityConfig.enable_mm_processor_stats
649
scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
@@ -1378,6 +1379,10 @@ class EngineArgs:
1379
"--enable-logging-iteration-details",
1380
**observability_kwargs["enable_logging_iteration_details"],
1381
)
1382
+ observability_group.add_argument(
1383
+ "--jit-monitor-mode",
1384
+ **observability_kwargs["jit_monitor_mode"],
1385
+ )
1386
observability_group.add_argument(
1387
"--jit-monitor-verbose",
1388
**observability_kwargs["jit_monitor_verbose"],
@@ -1781,6 +1786,22 @@ class EngineArgs:
1786
cfg = json.loads(cfg)
1787
return DiffusionConfig(**cfg)
1788
1789
+ def create_observability_config(self) -> ObservabilityConfig:
1790
+ return ObservabilityConfig(
1791
+ show_hidden_metrics_for_version=self.show_hidden_metrics_for_version,
1792
+ otlp_traces_endpoint=self.otlp_traces_endpoint,
1793
+ collect_detailed_traces=self.collect_detailed_traces,
1794
+ kv_cache_metrics=self.kv_cache_metrics,
1795
+ kv_cache_metrics_sample=self.kv_cache_metrics_sample,
1796
+ cudagraph_metrics=self.cudagraph_metrics,
1797
+ enable_layerwise_nvtx_tracing=self.enable_layerwise_nvtx_tracing,
1798
+ enable_mfu_metrics=self.enable_mfu_metrics,
1799
+ enable_mm_processor_stats=self.enable_mm_processor_stats,
1800
+ enable_logging_iteration_details=self.enable_logging_iteration_details,
1801
+ jit_monitor_mode=self.jit_monitor_mode,
1802
+ jit_monitor_verbose=self.jit_monitor_verbose,
1803
+ )
1804
+
1805
def create_engine_config(
1806
self,
1807
usage_context: UsageContext | None = None,
@@ -2263,19 +2284,7 @@ class EngineArgs:
2284
self.reasoning_parser_plugin
2285
)
2286
2266
- observability_config = ObservabilityConfig(
2267
- show_hidden_metrics_for_version=self.show_hidden_metrics_for_version,
2268
- otlp_traces_endpoint=self.otlp_traces_endpoint,
2269
- collect_detailed_traces=self.collect_detailed_traces,
2270
- kv_cache_metrics=self.kv_cache_metrics,
2271
- kv_cache_metrics_sample=self.kv_cache_metrics_sample,
2272
- cudagraph_metrics=self.cudagraph_metrics,
2273
- enable_layerwise_nvtx_tracing=self.enable_layerwise_nvtx_tracing,
2274
- enable_mfu_metrics=self.enable_mfu_metrics,
2275
- enable_mm_processor_stats=self.enable_mm_processor_stats,
2276
- enable_logging_iteration_details=self.enable_logging_iteration_details,
2277
- jit_monitor_verbose=self.jit_monitor_verbose,
2278
- )
2287
+ observability_config = self.create_observability_config()
2288
2289
# Compilation config overrides
2290
compilation_config = copy.deepcopy(self.compilation_config)
vllm/triton_utils/jit_monitor.py
deleted
-135
@@ -1,135 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Monitor unexpected Triton kernel JIT compilation during inference.
5
-
6
-After server warmup completes, any Triton JIT compilation or autotuning
7
-event indicates a cache miss or unexpected input shape that causes a
8
-latency spike. This module registers hooks in the Triton runtime to
9
-detect and log such events so they can be investigated.
10
-
11
-Set ``--jit-monitor-verbose`` to log every Triton JIT compile with its
12
-dispatch key. This is intentionally opt-in because it can emit many logs and
13
-add overhead.
14
-
15
-Currently monitors:
16
-- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``)
17
-- Triton ``@triton.jit`` first-time compilations
18
- (via ``knobs.runtime.jit_post_compile_hook``)
19
-"""
20
-
21
-import os
22
-
23
-from vllm.logger import init_logger
24
-from vllm.triton_utils.importing import HAS_TRITON
25
-
26
-logger = init_logger(__name__)
27
-
28
-_active: bool = False
29
-_verbose: bool = False
30
-
31
-
32
-def is_active() -> bool:
33
- """Return whether the JIT compilation monitor is currently active."""
34
- return _active
35
-
36
-
37
-def activate(*, verbose: bool = False) -> None:
38
- """Enable JIT compilation monitoring after warmup.
39
-
40
- Call once per worker process at the end of
41
- :func:`compile_or_warm_up_model`. After activation every Triton
42
- kernel compilation or autotuning benchmark that happens during
43
- inference will be logged as a warning.
44
-
45
- Safe to call multiple times — subsequent calls are no-ops.
46
-
47
- If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in
48
- their environment, autotuning printing is left disabled; the JIT
49
- compilation hook is still registered regardless.
50
- """
51
- global _active, _verbose
52
- if _active:
53
- return
54
- _active = True
55
- _verbose = verbose
56
-
57
- _setup_triton_autotuning_print()
58
- _setup_triton_jit_hook()
59
-
60
- logger.info(
61
- "Kernel JIT monitor activated — Triton JIT compilations "
62
- "during inference will be logged as warnings."
63
- )
64
-
65
-
66
-# ------------------------------------------------------------------
67
-# Triton autotuning print
68
-# ------------------------------------------------------------------
69
-
70
-
71
-def _setup_triton_autotuning_print() -> None:
72
- """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out."""
73
- if not HAS_TRITON:
74
- return
75
- from triton import knobs # type: ignore[import-untyped]
76
-
77
- user_val = os.environ.get("TRITON_PRINT_AUTOTUNING")
78
- if user_val == "0":
79
- logger.debug(
80
- "TRITON_PRINT_AUTOTUNING=0 set by user — "
81
- "autotuning messages will stay suppressed."
82
- )
83
- return
84
-
85
- knobs.autotuning.print = True
86
-
87
-
88
-# ------------------------------------------------------------------
89
-# Triton JIT compilation hook
90
-# ------------------------------------------------------------------
91
-
92
-
93
-def _log_jit_compile(fn_name: str, kwargs) -> None:
94
- if _verbose:
95
- compile_info = kwargs.get("compile")
96
- if not isinstance(compile_info, dict):
97
- compile_info = {}
98
- logger.warning(
99
- "Triton %sJIT compilation during inference: %s (key=%s).",
100
- "autotune/warmup candidate " if kwargs.get("warmup") else "kernel ",
101
- fn_name,
102
- compile_info.get("key") or kwargs.get("key"),
103
- )
104
- return
105
-
106
- logger.warning_once(
107
- "Triton kernel JIT compilation during inference: %s. "
108
- "This causes a latency spike; consider extending warmup "
109
- "to cover this shape/config.",
110
- fn_name,
111
- )
112
-
113
-
114
-def _setup_triton_jit_hook() -> None:
115
- """Register a ``jit_post_compile_hook`` that warns on compilation."""
116
- if not HAS_TRITON:
117
- return
118
- from triton import knobs # type: ignore[import-untyped]
119
-
120
- existing_hook = knobs.runtime.jit_post_compile_hook
121
-
122
- def _on_jit_compile(**kwargs):
123
- # `jit_post_compile_hook` is Triton internal API and its
124
- # signature has changed across releases (kwargs added/renamed).
125
- # Accept **kwargs so an upstream change cannot crash this hook
126
- # with TypeError, and forward the full kwarg set to any
127
- # pre-existing hook unchanged.
128
- fn = kwargs.get("fn")
129
- fn_name = getattr(fn, "name", "<unknown>")
130
- _log_jit_compile(fn_name, kwargs)
131
- if existing_hook is not None:
132
- return existing_hook(**kwargs)
133
- return None
134
-
135
- knobs.runtime.jit_post_compile_hook = _on_jit_compile
vllm/utils/jit_monitor.py
new
+212
@@ -0,0 +1,212 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Monitor unexpected kernel JIT compilation during inference.
5
+
6
+After server warmup completes, any kernel JIT compilation or autotuning event
7
+indicates a cache miss or unexpected input shape that causes a latency spike.
8
+This module registers hooks in supported runtimes to detect such events so
9
+they can be investigated.
10
+
11
+Set ``--jit-monitor-mode=error`` to fail fast on unexpected runtime
12
+compilation. Set ``--jit-monitor-verbose`` to log every JIT compile with
13
+additional runtime details. Verbose logging is intentionally opt-in because it
14
+can emit many logs and add overhead.
15
+
16
+Currently monitors:
17
+- CuTeDSL cute.compile calls
18
+- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``)
19
+- Triton ``@triton.jit`` first-time compilations
20
+ (via ``knobs.runtime.jit_post_compile_hook``)
21
+"""
22
+
23
+import functools
24
+import os
25
+from typing import Literal
26
+
27
+from vllm.logger import init_logger
28
+from vllm.triton_utils.importing import HAS_TRITON
29
+
30
+logger = init_logger(__name__)
31
+JitMonitorMode = Literal["warn", "error"]
32
+
33
+_active: bool = False
34
+_mode: JitMonitorMode = "warn"
35
+_verbose: bool = False
36
+_cutedsl_hook_installed: bool = False
37
+
38
+
39
+def is_active() -> bool:
40
+ """Return whether the JIT compilation monitor is currently active."""
41
+ return _active
42
+
43
+
44
+def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None:
45
+ """Enable JIT compilation monitoring after warmup.
46
+
47
+ Call once per worker process at the end of
48
+ :func:`compile_or_warm_up_model`. After activation every monitored kernel
49
+ compilation or autotuning benchmark that happens during inference will be
50
+ logged as a warning or raised as an error, depending on ``mode``.
51
+
52
+ Safe to call multiple times; subsequent calls are no-ops.
53
+
54
+ If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in
55
+ their environment, autotuning printing is left disabled; the JIT
56
+ compilation hook is still registered regardless.
57
+ """
58
+ global _active, _mode, _verbose
59
+ if _active:
60
+ return
61
+ if mode not in ("warn", "error"):
62
+ raise ValueError(f"Unsupported JIT monitor mode: {mode!r}")
63
+ _active = True
64
+ _mode = mode
65
+ _verbose = verbose
66
+
67
+ _setup_triton_autotuning_print()
68
+ _setup_triton_jit_hook()
69
+ _setup_cutedsl_jit_hook()
70
+
71
+ logger.info(
72
+ "Kernel JIT monitor activated; monitored JIT compilations during "
73
+ "inference will use mode=%s.",
74
+ mode,
75
+ )
76
+
77
+
78
+# ------------------------------------------------------------------
79
+# Triton autotuning print
80
+# ------------------------------------------------------------------
81
+
82
+
83
+def _setup_triton_autotuning_print() -> None:
84
+ """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out."""
85
+ if not HAS_TRITON:
86
+ return
87
+ from triton import knobs # type: ignore[import-untyped]
88
+
89
+ user_val = os.environ.get("TRITON_PRINT_AUTOTUNING")
90
+ if user_val == "0":
91
+ logger.debug(
92
+ "TRITON_PRINT_AUTOTUNING=0 set by user; "
93
+ "autotuning messages will stay suppressed."
94
+ )
95
+ return
96
+
97
+ knobs.autotuning.print = True
98
+
99
+
100
+# ------------------------------------------------------------------
101
+# Triton JIT compilation hook
102
+# ------------------------------------------------------------------
103
+
104
+
105
+def _handle_jit_event(
106
+ *,
107
+ backend: str,
108
+ event: str,
109
+ fn_name: str,
110
+ detail: str | None = None,
111
+) -> None:
112
+ message = (
113
+ "%s %s during inference: %s%s. "
114
+ "This causes a latency spike; consider extending warmup "
115
+ "to cover this shape/config."
116
+ )
117
+ detail_suffix = f" ({detail})" if detail else ""
118
+ args = (backend, event, fn_name, detail_suffix)
119
+
120
+ if _mode == "error":
121
+ raise RuntimeError(message % args)
122
+
123
+ if _verbose:
124
+ logger.warning(message, *args)
125
+ return
126
+
127
+ logger.warning_once(message, *args)
128
+
129
+
130
+def _log_triton_jit_compile(fn_name: str, kwargs) -> None:
131
+ compile_info = kwargs.get("compile")
132
+ if not isinstance(compile_info, dict):
133
+ compile_info = {}
134
+ key = compile_info.get("key") or kwargs.get("key")
135
+ detail = f"key={key}" if _verbose and key is not None else None
136
+ event = (
137
+ "autotune/warmup candidate JIT compilation"
138
+ if kwargs.get("warmup")
139
+ else "kernel JIT compilation"
140
+ )
141
+ _handle_jit_event(
142
+ backend="Triton",
143
+ event=event,
144
+ fn_name=fn_name,
145
+ detail=detail,
146
+ )
147
+
148
+
149
+def _setup_triton_jit_hook() -> None:
150
+ """Register a ``jit_post_compile_hook`` that warns on compilation."""
151
+ if not HAS_TRITON:
152
+ return
153
+ from triton import knobs # type: ignore[import-untyped]
154
+
155
+ existing_hook = knobs.runtime.jit_post_compile_hook
156
+
157
+ def _on_jit_compile(**kwargs):
158
+ # `jit_post_compile_hook` is Triton internal API and its
159
+ # signature has changed across releases (kwargs added/renamed).
160
+ # Accept **kwargs so an upstream change cannot crash this hook
161
+ # with TypeError, and forward the full kwarg set to any
162
+ # pre-existing hook unchanged.
163
+ fn = kwargs.get("fn")
164
+ fn_name = getattr(fn, "name", "<unknown>")
165
+ _log_triton_jit_compile(fn_name, kwargs)
166
+ if existing_hook is not None:
167
+ return existing_hook(**kwargs)
168
+ return None
169
+
170
+ knobs.runtime.jit_post_compile_hook = _on_jit_compile
171
+
172
+
173
+# ------------------------------------------------------------------
174
+# CuTeDSL JIT compilation hook
175
+# ------------------------------------------------------------------
176
+
177
+
178
+def _log_cutedsl_jit_compile(fn_name: str) -> None:
179
+ _handle_jit_event(
180
+ backend="CuTeDSL",
181
+ event="JIT compilation",
182
+ fn_name=fn_name,
183
+ )
184
+
185
+
186
+def _setup_cutedsl_jit_hook() -> None:
187
+ """Wrap ``cutlass.cute.compile`` to warn on compilation."""
188
+ global _cutedsl_hook_installed
189
+ if _cutedsl_hook_installed:
190
+ return
191
+
192
+ try:
193
+ import cutlass.cute as cute
194
+ except Exception:
195
+ logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.")
196
+ return
197
+
198
+ original_compile = cute.compile
199
+
200
+ @functools.wraps(original_compile)
201
+ def _compile_with_monitor(*args, **kwargs):
202
+ kernel = args[0] if args else kwargs.get("function")
203
+ kernel_name = getattr(kernel, "__name__", None)
204
+ if kernel_name is None:
205
+ kernel_name = (
206
+ kernel.__class__.__name__ if kernel is not None else "<unknown>"
207
+ )
208
+ _log_cutedsl_jit_compile(kernel_name)
209
+ return original_compile(*args, **kwargs)
210
+
211
+ cute.compile = _compile_with_monitor
212
+ _cutedsl_hook_installed = True
vllm/v1/worker/gpu_worker.py
+4
-5
@@ -761,12 +761,11 @@ class Worker(WorkerBase):
761
762
# All warmup is done — start monitoring for unexpected JIT
763
# compilations that would cause latency spikes during inference.
764
- from vllm.triton_utils.jit_monitor import (
765
- activate as activate_triton_jit_monitor,
766
- )
764
+ from vllm.utils.jit_monitor import activate as activate_jit_monitor
765
768
- activate_triton_jit_monitor(
769
- verbose=self.observability_config.jit_monitor_verbose
766
+ activate_jit_monitor(
767
+ mode=self.observability_config.jit_monitor_mode,
768
+ verbose=self.observability_config.jit_monitor_verbose,
769
)
770
771
# Freeze the worker heap so the GC won't scan static objects