A high-throughput and memory-efficient inference and serving engine for LLMs
Add Triton recompile detection (#45631)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Jun 16, 2026 at 18:25 UTC
b2cfae777dbad80096e5969212da58ff01cc432e
6 files changed
+81
-15
tests/engine/test_arg_utils.py
+8
@@ -206,6 +206,14 @@ def test_get_kwargs():
206
assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2) # type: ignore[call-arg]
207
208
209
+def test_jit_monitor_verbose_arg():
210
+ parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
211
+ args = parser.parse_args(["--jit-monitor-verbose"])
212
+
213
+ assert args.jit_monitor_verbose
214
+ assert EngineArgs(model="test", jit_monitor_verbose=True).jit_monitor_verbose
215
+
216
+
217
def test_hf_token_get_kwargs():
218
kwargs = get_kwargs(ModelConfig)["hf_token"]
219
tests/test_jit_monitor.py
+30
-6
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import os
4
import sys
5
+from contextlib import contextmanager
6
from types import SimpleNamespace
7
from unittest import mock
8
@@ -14,8 +15,10 @@ from vllm.triton_utils import jit_monitor
15
def _reset_monitor():
16
"""Reset global monitor state between tests."""
17
jit_monitor._active = False
18
+ jit_monitor._verbose = False
19
yield
20
jit_monitor._active = False
21
+ jit_monitor._verbose = False
22
23
24
# ------------------------------------------------------------------
@@ -30,10 +33,15 @@ def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
33
return SimpleNamespace(autotuning=autotuning, runtime=runtime)
34
35
36
+@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)
36
- return mock.patch.dict(sys.modules, {"triton": fake_triton})
40
+ with (
41
+ mock.patch.dict(sys.modules, {"triton": fake_triton}),
42
+ mock.patch.object(jit_monitor, "HAS_TRITON", True),
43
+ ):
44
+ yield
45
46
47
# ------------------------------------------------------------------
@@ -108,7 +116,10 @@ class TestJitHook:
116
hook = fake.runtime.jit_post_compile_hook
117
mock_fn = SimpleNamespace(name="test_kernel")
118
111
- with mock.patch.object(jit_monitor.logger, "warning") as m:
119
+ with (
120
+ mock.patch.object(jit_monitor.logger, "warning_once") as m,
121
+ mock.patch.object(jit_monitor.logger, "warning") as warning,
122
+ ):
123
hook(
124
key="some_key",
125
repr="some_repr",
@@ -119,6 +130,7 @@ class TestJitHook:
130
)
131
132
m.assert_called_once()
133
+ warning.assert_not_called()
134
msg = m.call_args[0][0] % m.call_args[0][1:]
135
assert "Triton kernel JIT compilation during inference" in msg
136
assert "test_kernel" in msg
@@ -206,9 +218,9 @@ if _HAS_TRITON:
218
tl.store(out_ptr + offs, x + y, mask=mask)
219
220
209
-def _run_add_kernel(n: int, block: int = 256) -> None:
221
+def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None:
222
"""Launch ``_add_kernel`` with vectors of length *n*."""
211
- x = torch.randn(n, device="cuda")
223
+ x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment
224
y = torch.randn(n, device="cuda")
225
out = torch.empty(n, device="cuda")
226
grid = ((n + block - 1) // block,)
@@ -224,7 +236,7 @@ class TestTritonJitHookIntegration:
236
_run_add_kernel(1024)
237
238
jit_monitor.activate()
227
- with mock.patch.object(jit_monitor.logger, "warning") as w:
239
+ with mock.patch.object(jit_monitor.logger, "warning_once") as w:
240
_run_add_kernel(1024)
241
w.assert_not_called()
242
@@ -232,9 +244,21 @@ class TestTritonJitHookIntegration:
244
_run_add_kernel(1024, block=256)
245
246
jit_monitor.activate()
235
- with mock.patch.object(jit_monitor.logger, "warning") as w:
247
+ with mock.patch.object(jit_monitor.logger, "warning_once") as w:
248
# Different BLOCK (a tl.constexpr) forces recompilation.
249
_run_add_kernel(1024, block=512)
250
w.assert_called()
251
msg = w.call_args[0][0] % w.call_args[0][1:]
252
assert "_add_kernel" in msg
253
+
254
+ def test_verbose_warning_on_each_new_pointer_alignment(self):
255
+ _run_add_kernel(1024)
256
+
257
+ jit_monitor.activate(verbose=True)
258
+ with (
259
+ mock.patch.object(jit_monitor.logger, "warning") as w,
260
+ mock.patch.object(jit_monitor.logger, "warning_once") as w_once,
261
+ ):
262
+ _run_add_kernel(1024, offset=1)
263
+ assert w.called
264
+ w_once.assert_not_called()
vllm/config/observability.py
+4
@@ -76,6 +76,10 @@ 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_verbose: bool = False
80
+ """Log every Triton JIT compile with its dispatch key. This can emit many
81
+ logs and add overhead, so it is intended for debugging."""
82
+
83
@cached_property
84
def collect_model_forward_time(self) -> bool:
85
"""Whether to collect model forward time for the request."""
vllm/engine/arg_utils.py
+6
@@ -637,6 +637,7 @@ class EngineArgs:
637
enable_logging_iteration_details: bool = (
638
ObservabilityConfig.enable_logging_iteration_details
639
)
640
+ jit_monitor_verbose: bool = ObservabilityConfig.jit_monitor_verbose
641
enable_mm_processor_stats: bool = ObservabilityConfig.enable_mm_processor_stats
642
scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
643
scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls
@@ -1357,6 +1358,10 @@ class EngineArgs:
1358
"--enable-logging-iteration-details",
1359
**observability_kwargs["enable_logging_iteration_details"],
1360
)
1361
+ observability_group.add_argument(
1362
+ "--jit-monitor-verbose",
1363
+ **observability_kwargs["jit_monitor_verbose"],
1364
+ )
1365
1366
# Scheduler arguments
1367
scheduler_kwargs = get_kwargs(SchedulerConfig)
@@ -2202,6 +2207,7 @@ class EngineArgs:
2207
enable_mfu_metrics=self.enable_mfu_metrics,
2208
enable_mm_processor_stats=self.enable_mm_processor_stats,
2209
enable_logging_iteration_details=self.enable_logging_iteration_details,
2210
+ jit_monitor_verbose=self.jit_monitor_verbose,
2211
)
2212
2213
# Compilation config overrides
vllm/triton_utils/jit_monitor.py
+30
-8
@@ -8,6 +8,10 @@ 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
@@ -22,6 +26,7 @@ from vllm.triton_utils.importing import HAS_TRITON
26
logger = init_logger(__name__)
27
28
_active: bool = False
29
+_verbose: bool = False
30
31
32
def is_active() -> bool:
@@ -29,7 +34,7 @@ def is_active() -> bool:
34
return _active
35
36
32
-def activate() -> None:
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
@@ -43,10 +48,11 @@ def activate() -> None:
48
their environment, autotuning printing is left disabled; the JIT
49
compilation hook is still registered regardless.
50
"""
46
- global _active
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()
@@ -84,6 +90,27 @@ def _setup_triton_autotuning_print() -> None:
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:
@@ -100,12 +127,7 @@ def _setup_triton_jit_hook() -> None:
127
# pre-existing hook unchanged.
128
fn = kwargs.get("fn")
129
fn_name = getattr(fn, "name", "<unknown>")
103
- logger.warning_once(
104
- "Triton kernel JIT compilation during inference: %s. "
105
- "This causes a latency spike; consider extending warmup "
106
- "to cover this shape/config.",
107
- fn_name,
108
- )
130
+ _log_jit_compile(fn_name, kwargs)
131
if existing_hook is not None:
132
return existing_hook(**kwargs)
133
return None
vllm/v1/worker/gpu_worker.py
+3
-1
@@ -737,7 +737,9 @@ class Worker(WorkerBase):
737
activate as activate_triton_jit_monitor,
738
)
739
740
- activate_triton_jit_monitor()
740
+ activate_triton_jit_monitor(
741
+ verbose=self.observability_config.jit_monitor_verbose
742
+ )
743
744
# Freeze the worker heap so the GC won't scan static objects
745
# (model weights, KV caches, CUDA graphs) during inference.