| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Monitor unexpected kernel JIT compilation during inference. |
| 4 | |
| 5 | After server warmup completes, any kernel JIT compilation or autotuning event |
| 6 | indicates a cache miss or unexpected input shape that causes a latency spike. |
| 7 | This module registers hooks in supported runtimes to detect such events so |
| 8 | they can be investigated. |
| 9 | |
| 10 | Set ``--jit-monitor-mode=error`` to fail fast on unexpected runtime |
| 11 | compilation. Set ``--jit-monitor-verbose`` to log every JIT compile with |
| 12 | additional runtime details. Verbose logging is intentionally opt-in because it |
| 13 | can emit many logs and add overhead. |
| 14 | |
| 15 | Currently monitors: |
| 16 | - CuTeDSL cute.compile calls |
| 17 | - Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``) |
| 18 | - Triton ``@triton.jit`` first-time compilations |
| 19 | (via ``knobs.runtime.jit_post_compile_hook``) |
| 20 | - TileLang ``@tilelang.jit`` first-time compilations |
| 21 | """ |
| 22 | |
| 23 | import contextlib |
| 24 | import functools |
| 25 | import importlib |
| 26 | import os |
| 27 | from collections.abc import Iterator, Mapping |
| 28 | from contextlib import suppress |
| 29 | from typing import Any, Literal, cast |
| 30 | |
| 31 | from vllm.logger import init_logger |
| 32 | from vllm.platforms import current_platform |
| 33 | from vllm.triton_utils.importing import HAS_TRITON |
| 34 | |
| 35 | logger = init_logger(__name__) |
| 36 | JitMonitorMode = Literal["warn", "error"] |
| 37 | |
| 38 | _active: bool = False |
| 39 | _mode: JitMonitorMode = "warn" |
| 40 | _verbose: bool = False |
| 41 | _cutedsl_hook_installed: bool = False |
| 42 | _tilelang_hook_installed: bool = False |
| 43 | _tilelang_jitimpl_compile_depth: int = 0 |
| 44 | |
| 45 | |
| 46 | def is_active() -> bool: |
| 47 | """Return whether the JIT compilation monitor is currently active.""" |
| 48 | return _active |
| 49 | |
| 50 | |
| 51 | def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None: |
| 52 | """Enable JIT compilation monitoring after warmup. |
| 53 | |
| 54 | Call once per worker process at the end of |
| 55 | :func:`compile_or_warm_up_model`. After activation every monitored kernel |
| 56 | compilation or autotuning benchmark that happens during inference will be |
| 57 | logged as a warning or raised as an error, depending on ``mode``. |
| 58 | |
| 59 | Safe to call multiple times; subsequent calls are no-ops. |
| 60 | |
| 61 | If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in |
| 62 | their environment, autotuning printing is left disabled; the JIT |
| 63 | compilation hook is still registered regardless. |
| 64 | """ |
| 65 | global _active, _mode, _verbose |
| 66 | if _active: |
| 67 | return |
| 68 | if mode not in ("warn", "error"): |
| 69 | raise ValueError(f"Unsupported JIT monitor mode: {mode!r}") |
| 70 | _active = True |
| 71 | _mode = mode |
| 72 | _verbose = verbose |
| 73 | |
| 74 | _setup_triton_autotuning_print() |
| 75 | _setup_triton_jit_hook() |
| 76 | _setup_cutedsl_jit_hook() |
| 77 | |
| 78 | # Refer to #51159. tilelang ships broken symbols |
| 79 | # on rocm. |
| 80 | # TODO: Remove the guard once tilelang upstream is fixed. |
| 81 | if not current_platform.is_rocm(): |
| 82 | _setup_tilelang_jit_hook() |
| 83 | |
| 84 | logger.info( |
| 85 | "Kernel JIT monitor activated; monitored JIT compilations during " |
| 86 | "inference will use mode=%s.", |
| 87 | mode, |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | # ------------------------------------------------------------------ |
| 92 | # Triton autotuning print |
| 93 | # ------------------------------------------------------------------ |
| 94 | |
| 95 | |
| 96 | def _setup_triton_autotuning_print() -> None: |
| 97 | """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out.""" |
| 98 | if not HAS_TRITON: |
| 99 | return |
| 100 | from triton import knobs # type: ignore[import-untyped] |
| 101 | |
| 102 | user_val = os.environ.get("TRITON_PRINT_AUTOTUNING") |
| 103 | if user_val == "0": |
| 104 | logger.debug( |
| 105 | "TRITON_PRINT_AUTOTUNING=0 set by user; " |
| 106 | "autotuning messages will stay suppressed." |
| 107 | ) |
| 108 | return |
| 109 | |
| 110 | knobs.autotuning.print = True |
| 111 | |
| 112 | |
| 113 | # ------------------------------------------------------------------ |
| 114 | # Triton JIT compilation hook |
| 115 | # ------------------------------------------------------------------ |
| 116 | |
| 117 | |
| 118 | def _handle_jit_event( |
| 119 | *, |
| 120 | backend: str, |
| 121 | event: str, |
| 122 | fn_name: str, |
| 123 | detail: str | None = None, |
| 124 | ) -> None: |
| 125 | message = ( |
| 126 | "%s %s during inference: %s%s. " |
| 127 | "This causes a latency spike; consider extending warmup " |
| 128 | "to cover this shape/config." |
| 129 | ) |
| 130 | detail_suffix = f" ({detail})" if detail else "" |
| 131 | args = (backend, event, fn_name, detail_suffix) |
| 132 | |
| 133 | if _mode == "error": |
| 134 | raise RuntimeError(message % args) |
| 135 | |
| 136 | if _verbose: |
| 137 | logger.warning(message, *args) |
| 138 | return |
| 139 | |
| 140 | logger.warning_once(message, *args) |
| 141 | |
| 142 | |
| 143 | def _safe_repr(value: object, *, max_len: int = 120) -> str: |
| 144 | try: |
| 145 | text = repr(value) |
| 146 | except Exception: |
| 147 | text = f"<{type(value).__name__}>" |
| 148 | if len(text) > max_len: |
| 149 | return text[: max_len - 3] + "..." |
| 150 | return text |
| 151 | |
| 152 | |
| 153 | def _get_compile_info(kwargs: Mapping[str, object]) -> dict: |
| 154 | compile_info = kwargs.get("compile") |
| 155 | if isinstance(compile_info, dict): |
| 156 | return compile_info |
| 157 | return {} |
| 158 | |
| 159 | |
| 160 | def _constant_name(fn: object, path: object) -> str: |
| 161 | jit_function = getattr(fn, "jit_function", None) |
| 162 | params = getattr(jit_function, "params", ()) |
| 163 | if isinstance(path, tuple) and path and isinstance(path[0], int): |
| 164 | idx = path[0] |
| 165 | if idx < len(params): |
| 166 | param_name = getattr(params[idx], "name", None) |
| 167 | if param_name is not None: |
| 168 | if len(path) == 1: |
| 169 | return param_name |
| 170 | suffix = "".join(f"[{part!r}]" for part in path[1:]) |
| 171 | return f"{param_name}{suffix}" |
| 172 | return str(path) |
| 173 | |
| 174 | |
| 175 | def _format_constants(fn: object, compile_info: Mapping[str, object]) -> str: |
| 176 | constants = compile_info.get("constants") |
| 177 | if not isinstance(constants, Mapping) or not constants: |
| 178 | return "{}" |
| 179 | |
| 180 | items = sorted( |
| 181 | ( |
| 182 | (_constant_name(fn, path), _safe_repr(value)) |
| 183 | for path, value in constants.items() |
| 184 | ), |
| 185 | key=lambda item: item[0], |
| 186 | ) |
| 187 | return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}" |
| 188 | |
| 189 | |
| 190 | def _format_signature(compile_info: Mapping[str, object]) -> str: |
| 191 | signature = compile_info.get("signature") |
| 192 | if not isinstance(signature, Mapping) or not signature: |
| 193 | return "{}" |
| 194 | items = sorted((str(k), _safe_repr(v)) for k, v in signature.items()) |
| 195 | return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}" |
| 196 | |
| 197 | |
| 198 | def _format_extra_compile_info(compile_info: Mapping[str, object]) -> str: |
| 199 | skip_keys = frozenset( |
| 200 | { |
| 201 | "constants", |
| 202 | "signature", |
| 203 | "key", |
| 204 | "fn", |
| 205 | "name", |
| 206 | } |
| 207 | ) |
| 208 | items = [ |
| 209 | f"{name}={_safe_repr(value)}" |
| 210 | for name, value in sorted(compile_info.items()) |
| 211 | if name not in skip_keys |
| 212 | ] |
| 213 | return "{" + ", ".join(items) + "}" |
| 214 | |
| 215 | |
| 216 | def _format_verbose_triton_compile_details(kwargs: Mapping[str, object]) -> str: |
| 217 | compile_info = _get_compile_info(kwargs) |
| 218 | fn = kwargs.get("fn") |
| 219 | key = compile_info.get("key") or kwargs.get("key") |
| 220 | return ( |
| 221 | f"constexprs={_format_constants(fn, compile_info)}; " |
| 222 | f"signature={_format_signature(compile_info)}; " |
| 223 | f"extra_compile_info={_format_extra_compile_info(compile_info)}; " |
| 224 | f"key={_safe_repr(key)}" |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | def _log_triton_jit_compile(fn_name: str, kwargs) -> None: |
| 229 | detail = _format_verbose_triton_compile_details(kwargs) if _verbose else None |
| 230 | _handle_jit_event( |
| 231 | backend="Triton", |
| 232 | event="kernel JIT compilation", |
| 233 | fn_name=fn_name, |
| 234 | detail=detail, |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | def _setup_triton_jit_hook() -> None: |
| 239 | """Register a jit_post_compile_hook that warns on compilation.""" |
| 240 | if not HAS_TRITON: |
| 241 | return |
| 242 | from triton import knobs |
| 243 | from triton.runtime import jit as _triton_jit |
| 244 | |
| 245 | # kernels pass non-JSON-serializable constexprs |
| 246 | # make that serialization non-fatal |
| 247 | _orig = _triton_jit.serialize_specialization_data |
| 248 | if not getattr(_orig, "_vllm_guarded", False): |
| 249 | |
| 250 | @functools.wraps(_orig) |
| 251 | def _guarded(*args, **kwargs): |
| 252 | try: |
| 253 | return _orig(*args, **kwargs) |
| 254 | except (TypeError, ValueError): |
| 255 | return None # best-effort metadata; monitor ignores it |
| 256 | |
| 257 | cast(Any, _guarded)._vllm_guarded = True |
| 258 | _triton_jit.serialize_specialization_data = _guarded |
| 259 | |
| 260 | existing_hook = knobs.runtime.jit_post_compile_hook |
| 261 | |
| 262 | def _on_jit_compile(**kwargs): |
| 263 | # `jit_post_compile_hook` is Triton internal API and its |
| 264 | # signature has changed across releases (kwargs added/renamed). |
| 265 | # Accept **kwargs so an upstream change cannot crash this hook |
| 266 | # with TypeError, and forward the full kwarg set to any |
| 267 | # pre-existing hook unchanged. |
| 268 | fn = kwargs.get("fn") |
| 269 | fn_name = getattr(fn, "name", "<unknown>") |
| 270 | _log_triton_jit_compile(fn_name, kwargs) |
| 271 | if existing_hook is not None: |
| 272 | return existing_hook(**kwargs) |
| 273 | return None |
| 274 | |
| 275 | knobs.runtime.jit_post_compile_hook = _on_jit_compile |
| 276 | |
| 277 | |
| 278 | # ------------------------------------------------------------------ |
| 279 | # CuTeDSL JIT compilation hook |
| 280 | # ------------------------------------------------------------------ |
| 281 | |
| 282 | |
| 283 | def _log_cutedsl_jit_compile(fn_name: str) -> None: |
| 284 | _handle_jit_event( |
| 285 | backend="CuTeDSL", |
| 286 | event="JIT compilation", |
| 287 | fn_name=fn_name, |
| 288 | ) |
| 289 | |
| 290 | |
| 291 | class _MonitoredCuteCompile: |
| 292 | """Logs JIT compilations; a plain function would break ``cute.compile[opts]``.""" |
| 293 | |
| 294 | def __init__(self, inner): |
| 295 | self._inner = inner |
| 296 | |
| 297 | def __getitem__(self, options) -> "_MonitoredCuteCompile": |
| 298 | return _MonitoredCuteCompile(self._inner[options]) |
| 299 | |
| 300 | def __call__(self, *args, **kwargs): |
| 301 | kernel = args[0] if args else kwargs.get("function") |
| 302 | kernel_name = getattr(kernel, "__name__", None) |
| 303 | if kernel_name is None: |
| 304 | kernel_name = ( |
| 305 | kernel.__class__.__name__ if kernel is not None else "<unknown>" |
| 306 | ) |
| 307 | _log_cutedsl_jit_compile(kernel_name) |
| 308 | return self._inner(*args, **kwargs) |
| 309 | |
| 310 | |
| 311 | def _setup_cutedsl_jit_hook() -> None: |
| 312 | """Wrap ``cutlass.cute.compile`` to warn on compilation.""" |
| 313 | global _cutedsl_hook_installed |
| 314 | if _cutedsl_hook_installed: |
| 315 | return |
| 316 | |
| 317 | try: |
| 318 | import cutlass.cute as cute |
| 319 | except Exception: |
| 320 | logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.") |
| 321 | return |
| 322 | |
| 323 | cute.compile = _MonitoredCuteCompile(cute.compile) |
| 324 | _cutedsl_hook_installed = True |
| 325 | |
| 326 | |
| 327 | # ------------------------------------------------------------------ |
| 328 | # TileLang JIT compilation hook |
| 329 | # ------------------------------------------------------------------ |
| 330 | |
| 331 | |
| 332 | def _tilelang_arg( |
| 333 | args: tuple[object, ...], |
| 334 | kwargs: Mapping[str, object], |
| 335 | index: int, |
| 336 | name: str, |
| 337 | default: object = None, |
| 338 | ) -> object: |
| 339 | if len(args) > index: |
| 340 | return args[index] |
| 341 | return kwargs.get(name, default) |
| 342 | |
| 343 | |
| 344 | def _tilelang_kernel_name(func: object) -> str: |
| 345 | attrs = getattr(func, "attrs", None) |
| 346 | global_symbol = None |
| 347 | get_attr = getattr(attrs, "get", None) |
| 348 | if callable(get_attr): |
| 349 | try: |
| 350 | global_symbol = get_attr("global_symbol") |
| 351 | except Exception: |
| 352 | global_symbol = None |
| 353 | elif isinstance(attrs, Mapping): |
| 354 | global_symbol = attrs.get("global_symbol") |
| 355 | |
| 356 | if global_symbol is not None: |
| 357 | return str(global_symbol) |
| 358 | |
| 359 | name = getattr(func, "__name__", None) |
| 360 | if name is not None: |
| 361 | return str(name) |
| 362 | return func.__class__.__name__ if func is not None else "<unknown>" |
| 363 | |
| 364 | |
| 365 | def _tilelang_call_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: |
| 366 | call_kwargs = dict(kwargs) |
| 367 | tune_params = call_kwargs.pop("__tune_params", {}) |
| 368 | if isinstance(tune_params, Mapping): |
| 369 | call_kwargs.update(tune_params) |
| 370 | return call_kwargs |
| 371 | |
| 372 | |
| 373 | def _tilelang_cache_miss_key( |
| 374 | jit_impl: object, |
| 375 | args: tuple[object, ...], |
| 376 | kwargs: Mapping[str, object], |
| 377 | ) -> object | None: |
| 378 | if kwargs.get("__return_compile_arguments", False): |
| 379 | return None |
| 380 | |
| 381 | call_kwargs = _tilelang_call_kwargs(kwargs) |
| 382 | func = getattr(jit_impl, "func", None) |
| 383 | parse_args = getattr(func, "parse_args", None) |
| 384 | if not callable(parse_args): |
| 385 | return None |
| 386 | |
| 387 | try: |
| 388 | if getattr(jit_impl, "mode", None) == "auto": |
| 389 | impl = cast(Any, jit_impl) |
| 390 | mode = impl._infer_jit_mode(*args, **call_kwargs) |
| 391 | impl.mode = mode |
| 392 | if func is not None: |
| 393 | func.set_mode(mode) |
| 394 | key, _ = parse_args(*args, **call_kwargs) |
| 395 | except Exception: |
| 396 | return None |
| 397 | |
| 398 | cache = getattr(jit_impl, "_kernel_cache", {}) |
| 399 | if isinstance(cache, Mapping): |
| 400 | return key if key not in cache else None |
| 401 | |
| 402 | get = getattr(cache, "get", None) |
| 403 | if not callable(get): |
| 404 | return key |
| 405 | try: |
| 406 | return key if get(key) is None else None |
| 407 | except Exception: |
| 408 | return key |
| 409 | |
| 410 | |
| 411 | def _format_tilelang_runtime_shapes( |
| 412 | jit_impl: object, |
| 413 | args: tuple[object, ...], |
| 414 | kwargs: Mapping[str, object], |
| 415 | ) -> str: |
| 416 | signature = getattr(jit_impl, "signature", None) |
| 417 | bind_partial = getattr(signature, "bind_partial", None) |
| 418 | if not callable(bind_partial): |
| 419 | return "{}" |
| 420 | |
| 421 | try: |
| 422 | bound = bind_partial(*args, **_tilelang_call_kwargs(kwargs)) |
| 423 | bound.apply_defaults() |
| 424 | except Exception: |
| 425 | return "{}" |
| 426 | |
| 427 | items = [] |
| 428 | for name, value in bound.arguments.items(): |
| 429 | if str(name).startswith("__"): |
| 430 | continue |
| 431 | if not (hasattr(value, "shape") and hasattr(value, "dtype")): |
| 432 | continue |
| 433 | shape = getattr(value, "shape", None) |
| 434 | if shape is not None: |
| 435 | with suppress(TypeError): |
| 436 | shape = tuple(shape) |
| 437 | items.append(f"{name}={_safe_repr(shape, max_len=80)}") |
| 438 | return "{" + ", ".join(items) + "}" |
| 439 | |
| 440 | |
| 441 | def _format_verbose_tilelang_compile_details( |
| 442 | jit_impl: object, |
| 443 | args: tuple[object, ...], |
| 444 | kwargs: Mapping[str, object], |
| 445 | cache_key: object, |
| 446 | ) -> str: |
| 447 | return ( |
| 448 | f"cache_key={_safe_repr(cache_key, max_len=400)}; " |
| 449 | f"runtime_shapes={_format_tilelang_runtime_shapes(jit_impl, args, kwargs)}" |
| 450 | ) |
| 451 | |
| 452 | |
| 453 | def _log_tilelang_jit_compile( |
| 454 | fn_name: str, |
| 455 | detail: str | None = None, |
| 456 | ) -> None: |
| 457 | _handle_jit_event( |
| 458 | backend="TileLang", |
| 459 | event="JIT compilation", |
| 460 | fn_name=fn_name, |
| 461 | detail=detail, |
| 462 | ) |
| 463 | |
| 464 | |
| 465 | def _setup_tilelang_jit_hook() -> None: |
| 466 | """Wrap TileLang JIT entry points to warn on compilation.""" |
| 467 | global _tilelang_hook_installed |
| 468 | if _tilelang_hook_installed: |
| 469 | return |
| 470 | |
| 471 | try: |
| 472 | tilelang_kernel = importlib.import_module("tilelang.jit.kernel") |
| 473 | except Exception: |
| 474 | logger.debug("TileLang is not available; skipping TileLang JIT monitor.") |
| 475 | return |
| 476 | |
| 477 | jit_kernel_cls = getattr(tilelang_kernel, "JITKernel", None) |
| 478 | if jit_kernel_cls is None: |
| 479 | logger.debug( |
| 480 | "TileLang JITKernel is unavailable; skipping TileLang JIT monitor." |
| 481 | ) |
| 482 | return |
| 483 | |
| 484 | try: |
| 485 | tilelang_jit = importlib.import_module("tilelang.jit") |
| 486 | except Exception: |
| 487 | tilelang_jit = None |
| 488 | jit_impl_cls = getattr(tilelang_jit, "JITImpl", None) |
| 489 | original_init = jit_kernel_cls.__init__ |
| 490 | |
| 491 | @functools.wraps(original_init) |
| 492 | def _init_with_monitor(self, *args, **kwargs): |
| 493 | from_database = bool(_tilelang_arg(args, kwargs, 7, "from_database", False)) |
| 494 | if not from_database and _tilelang_jitimpl_compile_depth == 0: |
| 495 | func = _tilelang_arg(args, kwargs, 0, "func") |
| 496 | _log_tilelang_jit_compile(_tilelang_kernel_name(func)) |
| 497 | return original_init(self, *args, **kwargs) |
| 498 | |
| 499 | jit_kernel_cls.__init__ = _init_with_monitor |
| 500 | |
| 501 | if jit_impl_cls is not None: |
| 502 | original_call = jit_impl_cls.__call__ |
| 503 | |
| 504 | @functools.wraps(original_call) |
| 505 | def _call_with_monitor(self, *args, **kwargs): |
| 506 | global _tilelang_jitimpl_compile_depth |
| 507 | cache_key = _tilelang_cache_miss_key(self, args, kwargs) |
| 508 | if cache_key is None: |
| 509 | return original_call(self, *args, **kwargs) |
| 510 | |
| 511 | _tilelang_jitimpl_compile_depth += 1 |
| 512 | try: |
| 513 | detail = None |
| 514 | if _verbose: |
| 515 | detail = _format_verbose_tilelang_compile_details( |
| 516 | self, args, kwargs, cache_key |
| 517 | ) |
| 518 | func = getattr(self, "func", None) |
| 519 | orig_func = getattr(func, "orig_func", None) |
| 520 | _log_tilelang_jit_compile( |
| 521 | _tilelang_kernel_name(orig_func or func), detail |
| 522 | ) |
| 523 | return original_call(self, *args, **kwargs) |
| 524 | finally: |
| 525 | _tilelang_jitimpl_compile_depth -= 1 |
| 526 | |
| 527 | jit_impl_cls.__call__ = _call_with_monitor |
| 528 | |
| 529 | _tilelang_hook_installed = True |
| 530 | |
| 531 | |
| 532 | @contextlib.contextmanager |
| 533 | def numba_workqueue_threading_layer() -> Iterator[None]: |
| 534 | """Force numba's fork-safe `workqueue` threading layer for this block. |
| 535 | |
| 536 | GNU OpenMP (numba's default `omp` threading layer) aborts the process |
| 537 | if a forked child re-enters an OpenMP-active runtime. vLLM forks the |
| 538 | EngineCore subprocess from a process that may already have launched |
| 539 | numba's parallel accelerator, so the first call to any |
| 540 | `@njit(parallel=True)` function must happen under `workqueue` instead. |
| 541 | The threading layer choice is sticky for the life of the process once |
| 542 | launched, so restoring the config on exit does not undo the effect. |
| 543 | """ |
| 544 | import numba |
| 545 | |
| 546 | key = "NUMBA_THREADING_LAYER" |
| 547 | previous_env = os.environ.get(key) |
| 548 | previous_config = numba.config.THREADING_LAYER |
| 549 | os.environ[key] = "workqueue" |
| 550 | numba.config.THREADING_LAYER = "workqueue" |
| 551 | try: |
| 552 | yield |
| 553 | finally: |
| 554 | if previous_env is None: |
| 555 | os.environ.pop(key, None) |
| 556 | else: |
| 557 | os.environ[key] = previous_env |
| 558 | numba.config.THREADING_LAYER = previous_config |