A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
[Feature] Add Triton kernel JIT compilation monitor for inference (#40137)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Artem Perevedentsev committed
May 5, 2026 at 13:08 UTC
8b9ea2f881ea7319c13994ce34c7f1941e07f883
4 files changed
+363
-1
.buildkite/test_areas/engine.yaml
+2
-1
@@ -13,8 +13,9 @@ steps:
13
- tests/test_config
14
- tests/test_logger
15
- tests/test_vllm_port
16
+ - tests/test_jit_monitor.py
17
commands:
17
- - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
18
+ - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
19
20
- label: Engine (1 GPU)
21
key: engine-1-gpu
tests/test_jit_monitor.py
new
+240
@@ -0,0 +1,240 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+import sys
5
+from types import SimpleNamespace
6
+from unittest import mock
7
+
8
+import pytest
9
+
10
+from vllm.triton_utils import jit_monitor
11
+
12
+
13
+@pytest.fixture(autouse=True)
14
+def _reset_monitor():
15
+ """Reset global monitor state between tests."""
16
+ jit_monitor._active = False
17
+ yield
18
+ jit_monitor._active = False
19
+
20
+
21
+# ------------------------------------------------------------------
22
+# Helpers — lightweight stand-ins for triton.knobs
23
+# ------------------------------------------------------------------
24
+
25
+
26
+def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
27
+ """Build a minimal fake ``triton.knobs`` namespace."""
28
+ autotuning = SimpleNamespace(print=autotuning_print)
29
+ runtime = SimpleNamespace(jit_post_compile_hook=jit_hook)
30
+ return SimpleNamespace(autotuning=autotuning, runtime=runtime)
31
+
32
+
33
+def _patch_triton_knobs(fake_knobs):
34
+ """Context manager that makes ``from triton import knobs`` return *fake_knobs*."""
35
+ fake_triton = SimpleNamespace(knobs=fake_knobs)
36
+ return mock.patch.dict(sys.modules, {"triton": fake_triton})
37
+
38
+
39
+# ------------------------------------------------------------------
40
+# Unit tests (no GPU required, triton is mocked)
41
+# ------------------------------------------------------------------
42
+
43
+
44
+class TestActivateBasic:
45
+ def test_sets_active(self):
46
+ assert not jit_monitor.is_active()
47
+ with _patch_triton_knobs(_make_fake_knobs()):
48
+ jit_monitor.activate()
49
+ assert jit_monitor.is_active()
50
+
51
+ def test_idempotent(self):
52
+ fake = _make_fake_knobs()
53
+ with _patch_triton_knobs(fake):
54
+ jit_monitor.activate()
55
+ first_hook = fake.runtime.jit_post_compile_hook
56
+ jit_monitor.activate()
57
+ assert fake.runtime.jit_post_compile_hook is first_hook
58
+
59
+ def test_logs_info_on_activation(self):
60
+ with (
61
+ mock.patch.object(jit_monitor.logger, "info") as m,
62
+ _patch_triton_knobs(_make_fake_knobs()),
63
+ ):
64
+ jit_monitor.activate()
65
+ m.assert_called_once()
66
+ assert "Kernel JIT monitor activated" in m.call_args[0][0]
67
+
68
+
69
+class TestAutotuningPrint:
70
+ def test_enables_autotuning_print(self):
71
+ fake = _make_fake_knobs(autotuning_print=False)
72
+ with _patch_triton_knobs(fake):
73
+ jit_monitor.activate()
74
+ assert fake.autotuning.print is True
75
+
76
+ def test_respects_user_opt_out(self):
77
+ fake = _make_fake_knobs(autotuning_print=False)
78
+ with (
79
+ mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}),
80
+ _patch_triton_knobs(fake),
81
+ ):
82
+ jit_monitor.activate()
83
+ assert fake.autotuning.print is False
84
+
85
+ def test_noop_when_user_already_enabled(self):
86
+ fake = _make_fake_knobs(autotuning_print=True)
87
+ with (
88
+ mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}),
89
+ _patch_triton_knobs(fake),
90
+ ):
91
+ jit_monitor.activate()
92
+ assert fake.autotuning.print is True
93
+
94
+
95
+class TestJitHook:
96
+ def test_hook_registered(self):
97
+ fake = _make_fake_knobs()
98
+ assert fake.runtime.jit_post_compile_hook is None
99
+ with _patch_triton_knobs(fake):
100
+ jit_monitor.activate()
101
+ assert fake.runtime.jit_post_compile_hook is not None
102
+
103
+ def test_hook_logs_warning(self):
104
+ fake = _make_fake_knobs()
105
+ with _patch_triton_knobs(fake):
106
+ jit_monitor.activate()
107
+
108
+ hook = fake.runtime.jit_post_compile_hook
109
+ mock_fn = SimpleNamespace(name="test_kernel")
110
+
111
+ with mock.patch.object(jit_monitor.logger, "warning") as m:
112
+ hook(
113
+ key="some_key",
114
+ repr="some_repr",
115
+ fn=mock_fn,
116
+ compile=lambda: None,
117
+ is_manual_warmup=False,
118
+ already_compiled=False,
119
+ )
120
+
121
+ m.assert_called_once()
122
+ msg = m.call_args[0][0] % m.call_args[0][1:]
123
+ assert "Triton kernel JIT compilation during inference" in msg
124
+ assert "test_kernel" in msg
125
+
126
+ def test_hook_chains_existing_hook(self):
127
+ existing = mock.MagicMock(return_value="existing_result")
128
+ fake = _make_fake_knobs(jit_hook=existing)
129
+ with _patch_triton_knobs(fake):
130
+ jit_monitor.activate()
131
+
132
+ hook = fake.runtime.jit_post_compile_hook
133
+ mock_fn = SimpleNamespace(name="chained_kernel")
134
+ kwargs = dict(
135
+ key="k",
136
+ repr="r",
137
+ fn=mock_fn,
138
+ compile=lambda: None,
139
+ is_manual_warmup=False,
140
+ already_compiled=False,
141
+ )
142
+ result = hook(**kwargs)
143
+
144
+ existing.assert_called_once()
145
+ assert result == "existing_result"
146
+
147
+ def test_hook_works_without_existing_hook(self):
148
+ fake = _make_fake_knobs(jit_hook=None)
149
+ with _patch_triton_knobs(fake):
150
+ jit_monitor.activate()
151
+
152
+ hook = fake.runtime.jit_post_compile_hook
153
+ mock_fn = SimpleNamespace(name="solo_kernel")
154
+ result = hook(
155
+ key="k",
156
+ repr="r",
157
+ fn=mock_fn,
158
+ compile=lambda: None,
159
+ is_manual_warmup=False,
160
+ already_compiled=False,
161
+ )
162
+ assert result is None
163
+
164
+
165
+class TestNoTritonFallback:
166
+ def test_activate_without_triton(self):
167
+ with mock.patch.object(jit_monitor, "HAS_TRITON", False):
168
+ jit_monitor.activate()
169
+ assert jit_monitor.is_active()
170
+
171
+
172
+# ------------------------------------------------------------------
173
+# Integration tests (real Triton + GPU)
174
+# ------------------------------------------------------------------
175
+
176
+try:
177
+ import torch
178
+
179
+ _HAS_CUDA = torch.cuda.is_available()
180
+except ImportError:
181
+ _HAS_CUDA = False
182
+
183
+try:
184
+ import triton
185
+ import triton.language as tl
186
+
187
+ _HAS_TRITON = True
188
+except ImportError:
189
+ _HAS_TRITON = False
190
+
191
+_skip_no_gpu = pytest.mark.skipif(
192
+ not (_HAS_CUDA and _HAS_TRITON),
193
+ reason="Requires CUDA GPU and Triton",
194
+)
195
+
196
+
197
+if _HAS_TRITON:
198
+
199
+ @triton.jit
200
+ def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
201
+ pid = tl.program_id(0)
202
+ offs = pid * BLOCK + tl.arange(0, BLOCK)
203
+ mask = offs < n
204
+ x = tl.load(x_ptr + offs, mask=mask)
205
+ y = tl.load(y_ptr + offs, mask=mask)
206
+ tl.store(out_ptr + offs, x + y, mask=mask)
207
+
208
+
209
+def _run_add_kernel(n: int, block: int = 256) -> None:
210
+ """Launch ``_add_kernel`` with vectors of length *n*."""
211
+ x = torch.randn(n, device="cuda")
212
+ y = torch.randn(n, device="cuda")
213
+ out = torch.empty(n, device="cuda")
214
+ grid = ((n + block - 1) // block,)
215
+ _add_kernel[grid](x, y, out, n, BLOCK=block)
216
+ torch.accelerator.synchronize()
217
+
218
+
219
+@_skip_no_gpu
220
+class TestTritonJitHookIntegration:
221
+ """End-to-end: real Triton kernel, real GPU, real hook."""
222
+
223
+ def test_no_warning_on_cached_shape(self):
224
+ _run_add_kernel(1024)
225
+
226
+ jit_monitor.activate()
227
+ with mock.patch.object(jit_monitor.logger, "warning") as w:
228
+ _run_add_kernel(1024)
229
+ w.assert_not_called()
230
+
231
+ def test_warning_on_new_constexpr(self):
232
+ _run_add_kernel(1024, block=256)
233
+
234
+ jit_monitor.activate()
235
+ with mock.patch.object(jit_monitor.logger, "warning") as w:
236
+ # Different BLOCK (a tl.constexpr) forces recompilation.
237
+ _run_add_kernel(1024, block=512)
238
+ w.assert_called()
239
+ msg = w.call_args[0][0] % w.call_args[0][1:]
240
+ assert "_add_kernel" in msg
vllm/triton_utils/jit_monitor.py
new
+113
@@ -0,0 +1,113 @@
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
+Currently monitors:
12
+- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``)
13
+- Triton ``@triton.jit`` first-time compilations
14
+ (via ``knobs.runtime.jit_post_compile_hook``)
15
+"""
16
+
17
+import os
18
+
19
+from vllm.logger import init_logger
20
+from vllm.triton_utils.importing import HAS_TRITON
21
+
22
+logger = init_logger(__name__)
23
+
24
+_active: bool = False
25
+
26
+
27
+def is_active() -> bool:
28
+ """Return whether the JIT compilation monitor is currently active."""
29
+ return _active
30
+
31
+
32
+def activate() -> None:
33
+ """Enable JIT compilation monitoring after warmup.
34
+
35
+ Call once per worker process at the end of
36
+ :func:`compile_or_warm_up_model`. After activation every Triton
37
+ kernel compilation or autotuning benchmark that happens during
38
+ inference will be logged as a warning.
39
+
40
+ Safe to call multiple times — subsequent calls are no-ops.
41
+
42
+ If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in
43
+ their environment, autotuning printing is left disabled; the JIT
44
+ compilation hook is still registered regardless.
45
+ """
46
+ global _active
47
+ if _active:
48
+ return
49
+ _active = True
50
+
51
+ _setup_triton_autotuning_print()
52
+ _setup_triton_jit_hook()
53
+
54
+ logger.info(
55
+ "Kernel JIT monitor activated — Triton JIT compilations "
56
+ "during inference will be logged as warnings."
57
+ )
58
+
59
+
60
+# ------------------------------------------------------------------
61
+# Triton autotuning print
62
+# ------------------------------------------------------------------
63
+
64
+
65
+def _setup_triton_autotuning_print() -> None:
66
+ """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out."""
67
+ if not HAS_TRITON:
68
+ return
69
+ from triton import knobs # type: ignore[import-untyped]
70
+
71
+ user_val = os.environ.get("TRITON_PRINT_AUTOTUNING")
72
+ if user_val == "0":
73
+ logger.debug(
74
+ "TRITON_PRINT_AUTOTUNING=0 set by user — "
75
+ "autotuning messages will stay suppressed."
76
+ )
77
+ return
78
+
79
+ knobs.autotuning.print = True
80
+
81
+
82
+# ------------------------------------------------------------------
83
+# Triton JIT compilation hook
84
+# ------------------------------------------------------------------
85
+
86
+
87
+def _setup_triton_jit_hook() -> None:
88
+ """Register a ``jit_post_compile_hook`` that warns on compilation."""
89
+ if not HAS_TRITON:
90
+ return
91
+ from triton import knobs # type: ignore[import-untyped]
92
+
93
+ existing_hook = knobs.runtime.jit_post_compile_hook
94
+
95
+ def _on_jit_compile(**kwargs):
96
+ # `jit_post_compile_hook` is Triton internal API and its
97
+ # signature has changed across releases (kwargs added/renamed).
98
+ # Accept **kwargs so an upstream change cannot crash this hook
99
+ # with TypeError, and forward the full kwarg set to any
100
+ # pre-existing hook unchanged.
101
+ fn = kwargs.get("fn")
102
+ 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
+ )
109
+ if existing_hook is not None:
110
+ return existing_hook(**kwargs)
111
+ return None
112
+
113
+ knobs.runtime.jit_post_compile_hook = _on_jit_compile
vllm/v1/worker/gpu_worker.py
+8
@@ -711,6 +711,14 @@ class Worker(WorkerBase):
711
# the model initialization and profiling.
712
set_random_seed(self.model_config.seed)
713
714
+ # All warmup is done — start monitoring for unexpected JIT
715
+ # compilations that would cause latency spikes during inference.
716
+ from vllm.triton_utils.jit_monitor import (
717
+ activate as activate_triton_jit_monitor,
718
+ )
719
+
720
+ activate_triton_jit_monitor()
721
+
722
return CompilationTimes(
723
language_model=self.compilation_config.compilation_time,
724
encoder=self.compilation_config.encoder_compilation_time,