wehub-resource-sync
531 lines
16 KiB
Python
531 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Monitor unexpected kernel JIT compilation during inference.
|
|
|
|
After server warmup completes, any kernel JIT compilation or autotuning event
|
|
indicates a cache miss or unexpected input shape that causes a latency spike.
|
|
This module registers hooks in supported runtimes to detect such events so
|
|
they can be investigated.
|
|
|
|
Set ``--jit-monitor-mode=error`` to fail fast on unexpected runtime
|
|
compilation. Set ``--jit-monitor-verbose`` to log every JIT compile with
|
|
additional runtime details. Verbose logging is intentionally opt-in because it
|
|
can emit many logs and add overhead.
|
|
|
|
Currently monitors:
|
|
- CuTeDSL cute.compile calls
|
|
- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``)
|
|
- Triton ``@triton.jit`` first-time compilations
|
|
(via ``knobs.runtime.jit_post_compile_hook``)
|
|
- TileLang ``@tilelang.jit`` first-time compilations
|
|
"""
|
|
|
|
import contextlib
|
|
import functools
|
|
import importlib
|
|
import os
|
|
from collections.abc import Iterator, Mapping
|
|
from contextlib import suppress
|
|
from typing import Any, Literal, cast
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.triton_utils.importing import HAS_TRITON
|
|
|
|
logger = init_logger(__name__)
|
|
JitMonitorMode = Literal["warn", "error"]
|
|
|
|
_active: bool = False
|
|
_mode: JitMonitorMode = "warn"
|
|
_verbose: bool = False
|
|
_cutedsl_hook_installed: bool = False
|
|
_tilelang_hook_installed: bool = False
|
|
_tilelang_jitimpl_compile_depth: int = 0
|
|
|
|
|
|
def is_active() -> bool:
|
|
"""Return whether the JIT compilation monitor is currently active."""
|
|
return _active
|
|
|
|
|
|
def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None:
|
|
"""Enable JIT compilation monitoring after warmup.
|
|
|
|
Call once per worker process at the end of
|
|
:func:`compile_or_warm_up_model`. After activation every monitored kernel
|
|
compilation or autotuning benchmark that happens during inference will be
|
|
logged as a warning or raised as an error, depending on ``mode``.
|
|
|
|
Safe to call multiple times; subsequent calls are no-ops.
|
|
|
|
If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in
|
|
their environment, autotuning printing is left disabled; the JIT
|
|
compilation hook is still registered regardless.
|
|
"""
|
|
global _active, _mode, _verbose
|
|
if _active:
|
|
return
|
|
if mode not in ("warn", "error"):
|
|
raise ValueError(f"Unsupported JIT monitor mode: {mode!r}")
|
|
_active = True
|
|
_mode = mode
|
|
_verbose = verbose
|
|
|
|
_setup_triton_autotuning_print()
|
|
_setup_triton_jit_hook()
|
|
_setup_cutedsl_jit_hook()
|
|
_setup_tilelang_jit_hook()
|
|
|
|
logger.info(
|
|
"Kernel JIT monitor activated; monitored JIT compilations during "
|
|
"inference will use mode=%s.",
|
|
mode,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Triton autotuning print
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def _setup_triton_autotuning_print() -> None:
|
|
"""Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out."""
|
|
if not HAS_TRITON:
|
|
return
|
|
from triton import knobs # type: ignore[import-untyped]
|
|
|
|
user_val = os.environ.get("TRITON_PRINT_AUTOTUNING")
|
|
if user_val == "0":
|
|
logger.debug(
|
|
"TRITON_PRINT_AUTOTUNING=0 set by user; "
|
|
"autotuning messages will stay suppressed."
|
|
)
|
|
return
|
|
|
|
knobs.autotuning.print = True
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Triton JIT compilation hook
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def _handle_jit_event(
|
|
*,
|
|
backend: str,
|
|
event: str,
|
|
fn_name: str,
|
|
detail: str | None = None,
|
|
) -> None:
|
|
message = (
|
|
"%s %s during inference: %s%s. "
|
|
"This causes a latency spike; consider extending warmup "
|
|
"to cover this shape/config."
|
|
)
|
|
detail_suffix = f" ({detail})" if detail else ""
|
|
args = (backend, event, fn_name, detail_suffix)
|
|
|
|
if _mode == "error":
|
|
raise RuntimeError(message % args)
|
|
|
|
if _verbose:
|
|
logger.warning(message, *args)
|
|
return
|
|
|
|
logger.warning_once(message, *args)
|
|
|
|
|
|
def _safe_repr(value: object, *, max_len: int = 120) -> str:
|
|
try:
|
|
text = repr(value)
|
|
except Exception:
|
|
text = f"<{type(value).__name__}>"
|
|
if len(text) > max_len:
|
|
return text[: max_len - 3] + "..."
|
|
return text
|
|
|
|
|
|
def _get_compile_info(kwargs: Mapping[str, object]) -> dict:
|
|
compile_info = kwargs.get("compile")
|
|
if isinstance(compile_info, dict):
|
|
return compile_info
|
|
return {}
|
|
|
|
|
|
def _constant_name(fn: object, path: object) -> str:
|
|
jit_function = getattr(fn, "jit_function", None)
|
|
params = getattr(jit_function, "params", ())
|
|
if isinstance(path, tuple) and path and isinstance(path[0], int):
|
|
idx = path[0]
|
|
if idx < len(params):
|
|
param_name = getattr(params[idx], "name", None)
|
|
if param_name is not None:
|
|
if len(path) == 1:
|
|
return param_name
|
|
suffix = "".join(f"[{part!r}]" for part in path[1:])
|
|
return f"{param_name}{suffix}"
|
|
return str(path)
|
|
|
|
|
|
def _format_constants(fn: object, compile_info: Mapping[str, object]) -> str:
|
|
constants = compile_info.get("constants")
|
|
if not isinstance(constants, Mapping) or not constants:
|
|
return "{}"
|
|
|
|
items = sorted(
|
|
(
|
|
(_constant_name(fn, path), _safe_repr(value))
|
|
for path, value in constants.items()
|
|
),
|
|
key=lambda item: item[0],
|
|
)
|
|
return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}"
|
|
|
|
|
|
def _format_signature(compile_info: Mapping[str, object]) -> str:
|
|
signature = compile_info.get("signature")
|
|
if not isinstance(signature, Mapping) or not signature:
|
|
return "{}"
|
|
items = sorted((str(k), _safe_repr(v)) for k, v in signature.items())
|
|
return "{" + ", ".join(f"{name}={value}" for name, value in items) + "}"
|
|
|
|
|
|
def _format_extra_compile_info(compile_info: Mapping[str, object]) -> str:
|
|
skip_keys = frozenset(
|
|
{
|
|
"constants",
|
|
"signature",
|
|
"key",
|
|
"fn",
|
|
"name",
|
|
}
|
|
)
|
|
items = [
|
|
f"{name}={_safe_repr(value)}"
|
|
for name, value in sorted(compile_info.items())
|
|
if name not in skip_keys
|
|
]
|
|
return "{" + ", ".join(items) + "}"
|
|
|
|
|
|
def _format_verbose_triton_compile_details(kwargs: Mapping[str, object]) -> str:
|
|
compile_info = _get_compile_info(kwargs)
|
|
fn = kwargs.get("fn")
|
|
key = compile_info.get("key") or kwargs.get("key")
|
|
return (
|
|
f"constexprs={_format_constants(fn, compile_info)}; "
|
|
f"signature={_format_signature(compile_info)}; "
|
|
f"extra_compile_info={_format_extra_compile_info(compile_info)}; "
|
|
f"key={_safe_repr(key)}"
|
|
)
|
|
|
|
|
|
def _log_triton_jit_compile(fn_name: str, kwargs) -> None:
|
|
detail = _format_verbose_triton_compile_details(kwargs) if _verbose else None
|
|
_handle_jit_event(
|
|
backend="Triton",
|
|
event="kernel JIT compilation",
|
|
fn_name=fn_name,
|
|
detail=detail,
|
|
)
|
|
|
|
|
|
def _setup_triton_jit_hook() -> None:
|
|
"""Register a ``jit_post_compile_hook`` that warns on compilation."""
|
|
if not HAS_TRITON:
|
|
return
|
|
from triton import knobs # type: ignore[import-untyped]
|
|
|
|
existing_hook = knobs.runtime.jit_post_compile_hook
|
|
|
|
def _on_jit_compile(**kwargs):
|
|
# `jit_post_compile_hook` is Triton internal API and its
|
|
# signature has changed across releases (kwargs added/renamed).
|
|
# Accept **kwargs so an upstream change cannot crash this hook
|
|
# with TypeError, and forward the full kwarg set to any
|
|
# pre-existing hook unchanged.
|
|
fn = kwargs.get("fn")
|
|
fn_name = getattr(fn, "name", "<unknown>")
|
|
_log_triton_jit_compile(fn_name, kwargs)
|
|
if existing_hook is not None:
|
|
return existing_hook(**kwargs)
|
|
return None
|
|
|
|
knobs.runtime.jit_post_compile_hook = _on_jit_compile
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# CuTeDSL JIT compilation hook
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def _log_cutedsl_jit_compile(fn_name: str) -> None:
|
|
_handle_jit_event(
|
|
backend="CuTeDSL",
|
|
event="JIT compilation",
|
|
fn_name=fn_name,
|
|
)
|
|
|
|
|
|
def _setup_cutedsl_jit_hook() -> None:
|
|
"""Wrap ``cutlass.cute.compile`` to warn on compilation."""
|
|
global _cutedsl_hook_installed
|
|
if _cutedsl_hook_installed:
|
|
return
|
|
|
|
try:
|
|
import cutlass.cute as cute
|
|
except Exception:
|
|
logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.")
|
|
return
|
|
|
|
original_compile = cute.compile
|
|
|
|
@functools.wraps(original_compile)
|
|
def _compile_with_monitor(*args, **kwargs):
|
|
kernel = args[0] if args else kwargs.get("function")
|
|
kernel_name = getattr(kernel, "__name__", None)
|
|
if kernel_name is None:
|
|
kernel_name = (
|
|
kernel.__class__.__name__ if kernel is not None else "<unknown>"
|
|
)
|
|
_log_cutedsl_jit_compile(kernel_name)
|
|
return original_compile(*args, **kwargs)
|
|
|
|
cute.compile = _compile_with_monitor
|
|
_cutedsl_hook_installed = True
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# TileLang JIT compilation hook
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def _tilelang_arg(
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
index: int,
|
|
name: str,
|
|
default: object = None,
|
|
) -> object:
|
|
if len(args) > index:
|
|
return args[index]
|
|
return kwargs.get(name, default)
|
|
|
|
|
|
def _tilelang_kernel_name(func: object) -> str:
|
|
attrs = getattr(func, "attrs", None)
|
|
global_symbol = None
|
|
get_attr = getattr(attrs, "get", None)
|
|
if callable(get_attr):
|
|
try:
|
|
global_symbol = get_attr("global_symbol")
|
|
except Exception:
|
|
global_symbol = None
|
|
elif isinstance(attrs, Mapping):
|
|
global_symbol = attrs.get("global_symbol")
|
|
|
|
if global_symbol is not None:
|
|
return str(global_symbol)
|
|
|
|
name = getattr(func, "__name__", None)
|
|
if name is not None:
|
|
return str(name)
|
|
return func.__class__.__name__ if func is not None else "<unknown>"
|
|
|
|
|
|
def _tilelang_call_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
|
|
call_kwargs = dict(kwargs)
|
|
tune_params = call_kwargs.pop("__tune_params", {})
|
|
if isinstance(tune_params, Mapping):
|
|
call_kwargs.update(tune_params)
|
|
return call_kwargs
|
|
|
|
|
|
def _tilelang_cache_miss_key(
|
|
jit_impl: object,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
) -> object | None:
|
|
if kwargs.get("__return_compile_arguments", False):
|
|
return None
|
|
|
|
call_kwargs = _tilelang_call_kwargs(kwargs)
|
|
func = getattr(jit_impl, "func", None)
|
|
parse_args = getattr(func, "parse_args", None)
|
|
if not callable(parse_args):
|
|
return None
|
|
|
|
try:
|
|
if getattr(jit_impl, "mode", None) == "auto":
|
|
impl = cast(Any, jit_impl)
|
|
mode = impl._infer_jit_mode(*args, **call_kwargs)
|
|
impl.mode = mode
|
|
if func is not None:
|
|
func.set_mode(mode)
|
|
key, _ = parse_args(*args, **call_kwargs)
|
|
except Exception:
|
|
return None
|
|
|
|
cache = getattr(jit_impl, "_kernel_cache", {})
|
|
if isinstance(cache, Mapping):
|
|
return key if key not in cache else None
|
|
|
|
get = getattr(cache, "get", None)
|
|
if not callable(get):
|
|
return key
|
|
try:
|
|
return key if get(key) is None else None
|
|
except Exception:
|
|
return key
|
|
|
|
|
|
def _format_tilelang_runtime_shapes(
|
|
jit_impl: object,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
) -> str:
|
|
signature = getattr(jit_impl, "signature", None)
|
|
bind_partial = getattr(signature, "bind_partial", None)
|
|
if not callable(bind_partial):
|
|
return "{}"
|
|
|
|
try:
|
|
bound = bind_partial(*args, **_tilelang_call_kwargs(kwargs))
|
|
bound.apply_defaults()
|
|
except Exception:
|
|
return "{}"
|
|
|
|
items = []
|
|
for name, value in bound.arguments.items():
|
|
if str(name).startswith("__"):
|
|
continue
|
|
if not (hasattr(value, "shape") and hasattr(value, "dtype")):
|
|
continue
|
|
shape = getattr(value, "shape", None)
|
|
if shape is not None:
|
|
with suppress(TypeError):
|
|
shape = tuple(shape)
|
|
items.append(f"{name}={_safe_repr(shape, max_len=80)}")
|
|
return "{" + ", ".join(items) + "}"
|
|
|
|
|
|
def _format_verbose_tilelang_compile_details(
|
|
jit_impl: object,
|
|
args: tuple[object, ...],
|
|
kwargs: Mapping[str, object],
|
|
cache_key: object,
|
|
) -> str:
|
|
return (
|
|
f"cache_key={_safe_repr(cache_key, max_len=400)}; "
|
|
f"runtime_shapes={_format_tilelang_runtime_shapes(jit_impl, args, kwargs)}"
|
|
)
|
|
|
|
|
|
def _log_tilelang_jit_compile(
|
|
fn_name: str,
|
|
detail: str | None = None,
|
|
) -> None:
|
|
_handle_jit_event(
|
|
backend="TileLang",
|
|
event="JIT compilation",
|
|
fn_name=fn_name,
|
|
detail=detail,
|
|
)
|
|
|
|
|
|
def _setup_tilelang_jit_hook() -> None:
|
|
"""Wrap TileLang JIT entry points to warn on compilation."""
|
|
global _tilelang_hook_installed
|
|
if _tilelang_hook_installed:
|
|
return
|
|
|
|
try:
|
|
tilelang_kernel = importlib.import_module("tilelang.jit.kernel")
|
|
except Exception:
|
|
logger.debug("TileLang is not available; skipping TileLang JIT monitor.")
|
|
return
|
|
|
|
jit_kernel_cls = getattr(tilelang_kernel, "JITKernel", None)
|
|
if jit_kernel_cls is None:
|
|
logger.debug(
|
|
"TileLang JITKernel is unavailable; skipping TileLang JIT monitor."
|
|
)
|
|
return
|
|
|
|
try:
|
|
tilelang_jit = importlib.import_module("tilelang.jit")
|
|
except Exception:
|
|
tilelang_jit = None
|
|
jit_impl_cls = getattr(tilelang_jit, "JITImpl", None)
|
|
original_init = jit_kernel_cls.__init__
|
|
|
|
@functools.wraps(original_init)
|
|
def _init_with_monitor(self, *args, **kwargs):
|
|
from_database = bool(_tilelang_arg(args, kwargs, 7, "from_database", False))
|
|
if not from_database and _tilelang_jitimpl_compile_depth == 0:
|
|
func = _tilelang_arg(args, kwargs, 0, "func")
|
|
_log_tilelang_jit_compile(_tilelang_kernel_name(func))
|
|
return original_init(self, *args, **kwargs)
|
|
|
|
jit_kernel_cls.__init__ = _init_with_monitor
|
|
|
|
if jit_impl_cls is not None:
|
|
original_call = jit_impl_cls.__call__
|
|
|
|
@functools.wraps(original_call)
|
|
def _call_with_monitor(self, *args, **kwargs):
|
|
global _tilelang_jitimpl_compile_depth
|
|
cache_key = _tilelang_cache_miss_key(self, args, kwargs)
|
|
if cache_key is None:
|
|
return original_call(self, *args, **kwargs)
|
|
|
|
_tilelang_jitimpl_compile_depth += 1
|
|
try:
|
|
detail = None
|
|
if _verbose:
|
|
detail = _format_verbose_tilelang_compile_details(
|
|
self, args, kwargs, cache_key
|
|
)
|
|
func = getattr(self, "func", None)
|
|
orig_func = getattr(func, "orig_func", None)
|
|
_log_tilelang_jit_compile(
|
|
_tilelang_kernel_name(orig_func or func), detail
|
|
)
|
|
return original_call(self, *args, **kwargs)
|
|
finally:
|
|
_tilelang_jitimpl_compile_depth -= 1
|
|
|
|
jit_impl_cls.__call__ = _call_with_monitor
|
|
|
|
_tilelang_hook_installed = True
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def numba_workqueue_threading_layer() -> Iterator[None]:
|
|
"""Force numba's fork-safe `workqueue` threading layer for this block.
|
|
|
|
GNU OpenMP (numba's default `omp` threading layer) aborts the process
|
|
if a forked child re-enters an OpenMP-active runtime. vLLM forks the
|
|
EngineCore subprocess from a process that may already have launched
|
|
numba's parallel accelerator, so the first call to any
|
|
`@njit(parallel=True)` function must happen under `workqueue` instead.
|
|
The threading layer choice is sticky for the life of the process once
|
|
launched, so restoring the config on exit does not undo the effect.
|
|
"""
|
|
import numba
|
|
|
|
key = "NUMBA_THREADING_LAYER"
|
|
previous_env = os.environ.get(key)
|
|
previous_config = numba.config.THREADING_LAYER
|
|
os.environ[key] = "workqueue"
|
|
numba.config.THREADING_LAYER = "workqueue"
|
|
try:
|
|
yield
|
|
finally:
|
|
if previous_env is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = previous_env
|
|
numba.config.THREADING_LAYER = previous_config
|