[Feat] Improve Triton JIT diagnostics (#46621)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Roberto L. Castro committed
Jun 30, 2026 at 20:50 UTC
20434c472e07d6e4116b5dd02058697331f3d1af
1 file changed
+86
-9
vllm/utils/jit_monitor.py
+86
-9
@@ -22,6 +22,7 @@ Currently monitors:
22
23
import functools
24
import os
25
+from collections.abc import Mapping
26
from typing import Literal
27
28
from vllm.logger import init_logger
@@ -127,20 +128,96 @@ def _handle_jit_event(
128
logger.warning_once(message, *args)
129
130
130
-def _log_triton_jit_compile(fn_name: str, kwargs) -> None:
131
+def _safe_repr(value: object, *, max_len: int = 120) -> str:
132
+ try:
133
+ text = repr(value)
134
+ except Exception:
135
+ text = f"<{type(value).__name__}>"
136
+ if len(text) > max_len:
137
+ return text[: max_len - 3] + "..."
138
+ return text
139
+
140
+
141
+def _get_compile_info(kwargs: Mapping[str, object]) -> dict:
142
compile_info = kwargs.get("compile")
132
- if not isinstance(compile_info, dict):
133
- compile_info = {}
143
+ if isinstance(compile_info, dict):
144
+ return compile_info
145
+ return {}
146
+
147
+
148
+def _constant_name(fn: object, path: object) -> str:
149
+ jit_function = getattr(fn, "jit_function", None)
150
+ params = getattr(jit_function, "params", ())
151
+ if isinstance(path, tuple) and path and isinstance(path[0], int):
152
+ idx = path[0]
153
+ if idx < len(params):
154
+ param_name = getattr(params[idx], "name", None)
155
+ if param_name is not None:
156
+ if len(path) == 1:
157
+ return param_name
158
+ suffix = "".join(f"[{part!r}]" for part in path[1:])
159
+ return f"{param_name}{suffix}"
160
+ return str(path)
161
+
162
+
163
+def _format_constants(fn: object, compile_info: Mapping[str, object]) -> str:
164
+ constants = compile_info.get("constants")
165
+ if not isinstance(constants, Mapping) or not constants:
166
+ return "{}"
167
+
168
+ items = sorted(
169
+ (
170
+ (_constant_name(fn, path), _safe_repr(value))
171
+ for path, value in constants.items()
172
+ ),
173
+ key=lambda item: item[0],
174
+ )
175
+ return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}"
176
+
177
+
178
+def _format_signature(compile_info: Mapping[str, object]) -> str:
179
+ signature = compile_info.get("signature")
180
+ if not isinstance(signature, Mapping) or not signature:
181
+ return "{}"
182
+ items = sorted((str(k), _safe_repr(v)) for k, v in signature.items())
183
+ return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}"
184
+
185
+
186
+def _format_extra_compile_info(compile_info: Mapping[str, object]) -> str:
187
+ skip_keys = frozenset(
188
+ {
189
+ "constants",
190
+ "signature",
191
+ "key",
192
+ "fn",
193
+ "name",
194
+ }
195
+ )
196
+ items = [
197
+ f"{name}={_safe_repr(value)}"
198
+ for name, value in sorted(compile_info.items())
199
+ if name not in skip_keys
200
+ ]
201
+ return "{" + ", ".join(items) + "}"
202
+
203
+
204
+def _format_verbose_triton_compile_details(kwargs: Mapping[str, object]) -> str:
205
+ compile_info = _get_compile_info(kwargs)
206
+ fn = kwargs.get("fn")
207
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"
208
+ return (
209
+ f"constexprs={_format_constants(fn, compile_info)}; "
210
+ f"signature={_format_signature(compile_info)}; "
211
+ f"extra_compile_info={_format_extra_compile_info(compile_info)}; "
212
+ f"key={_safe_repr(key)}"
213
)
214
+
215
+
216
+def _log_triton_jit_compile(fn_name: str, kwargs) -> None:
217
+ detail = _format_verbose_triton_compile_details(kwargs) if _verbose else None
218
_handle_jit_event(
219
backend="Triton",
143
- event=event,
220
+ event="kernel JIT compilation",
221
fn_name=fn_name,
222
detail=detail,
223
)