diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 983da2fbd..7f530651b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -721,6 +721,12 @@ class Envs: # (parity with flash-attn's ragged-aware launch). The feature checks _is_hip # explicitly in code; this env var allows override (0=force off, 1=force on). SGLANG_TRITON_COMPACT_EXTEND_ATTENTION = EnvBool(True) + # Raise if Triton loads a kernel after the engine starts serving. This + # verifies that startup warmup covers every kernel specialization used at + # serving time. + SGLANG_CRASH_ON_TRITON_LOAD_AFTER_READY = EnvBool(False) + SGLANG_TRITON_SLOW_COMPILE_THRESHOLD_SECS = EnvFloat(1.0) + SGLANG_TRITON_LOAD_WARNING_THRESHOLD_GB = EnvFloat(1.0) # Torch Compile SGLANG_ENABLE_TORCH_COMPILE = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 07aaf4f3e..d235a4440 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -301,6 +301,7 @@ from sglang.srt.utils import ( set_gpu_proc_affinity, set_random_seed, suppress_other_loggers, + triton_load_watch, ) from sglang.srt.utils.common import is_npu from sglang.srt.utils.hf_transformers_utils import ( @@ -1629,6 +1630,11 @@ class Scheduler( Sets up the schedule stream and dispatches to the appropriate event loop. The event loop blocks until shutdown. """ + # Engine init (graph capture, warmups) is done; from here on any + # Triton kernel device-load is a lazy first-use at serving time. + triton_load_watch.install() + triton_load_watch.mark_serving_started() + if use_mlx(): # MLX overlap uses mx.async_eval for CPU/GPU overlap, # not PyTorch MPS streams. diff --git a/python/sglang/srt/utils/triton_load_watch.py b/python/sglang/srt/utils/triton_load_watch.py new file mode 100644 index 000000000..27b74efea --- /dev/null +++ b/python/sglang/srt/utils/triton_load_watch.py @@ -0,0 +1,121 @@ +"""Detect Triton kernel device-loads after the engine starts serving. + +Triton loads each kernel specialization's cubin onto the GPU at its first +launch (``CompiledKernel._init_handles`` -> ``cuModuleLoadData``). That load +needs free device memory *outside* the torch caching allocator. Engines size +their pools to leave little post-init headroom, and the allocator's high-water +mark consumes the rest during early serving — so a specialization first used +mid-serving (e.g. a new adaptive speculative draft length, or a rare batch-size +bucket) can die in ``cuModuleLoadData`` with CUDA OOM, minutes or hours in. + +Once ``mark_serving_started()`` has been called, this module warns when an +uncached Triton compilation takes at least one second or a device-load starts +with less than 1 GiB of free device memory. Set +``SGLANG_CRASH_ON_TRITON_LOAD_AFTER_READY=1`` to raise on every late load +instead — for CI recipes that assert full startup warmup coverage. The hooks +only run for compilation and first-use loads, so steady-state cost is zero. + +Note: request-driven warmup (``--warmups``, the server warmup request) runs +*after* ``mark_serving_started()`` and is subject to the same diagnostics; +crash mode is only meant for deployments whose kernels are fully pre-loaded at +engine init. +""" + +from __future__ import annotations + +import logging + +import torch + +from sglang.srt.environ import envs +from sglang.srt.utils.common import get_available_gpu_memory + +logger = logging.getLogger(__name__) + +_serving_started = False +_prev_compile_listener = None +_installed = False + + +def install() -> None: + """Install the diagnostics (idempotent; chains pre-existing hooks).""" + global _installed, _prev_compile_listener + if _installed: + return + try: + import triton.knobs as knobs + except ImportError: + return + + # Triton 3.4 builds used by ROCm expose triton.knobs without these hooks. + runtime = getattr(knobs, "runtime", None) + compilation = getattr(knobs, "compilation", None) + hook = getattr(runtime, "kernel_load_start_hook", None) + if hook is None or compilation is None or not hasattr(compilation, "listener"): + return + + _prev_compile_listener = compilation.listener + hook.add(_on_kernel_load) + compilation.listener = _on_compilation + _installed = True + + +def mark_serving_started() -> None: + """Arm diagnostics for subsequent Triton compilations and device-loads.""" + global _serving_started + _serving_started = True + + +def _on_compilation(*, src, metadata, metadata_group, times, cache_hit) -> None: + if _prev_compile_listener is not None: + _prev_compile_listener( + src=src, + metadata=metadata, + metadata_group=metadata_group, + times=times, + cache_hit=cache_hit, + ) + if not _serving_started or cache_hit: + return + + compile_time_secs = times.total / 1e6 + if compile_time_secs < envs.SGLANG_TRITON_SLOW_COMPILE_THRESHOLD_SECS.get(): + return + + logger.warning( + "Triton kernel '%s' took %.2f s to compile after serving started. " + "Serving-time compilation can stall the engine; pre-compile it during " + "engine init.", + src.name, + compile_time_secs, + ) + + +def _on_kernel_load(module, function, name, metadata_group, hash) -> None: + if not _serving_started: + return + + free_gb = None + if torch.cuda.is_available(): + try: + free_gb = get_available_gpu_memory( + "cuda", torch.cuda.current_device(), empty_cache=False + ) + except RuntimeError: + logger.debug("Unable to query free device memory", exc_info=True) + + should_crash = envs.SGLANG_CRASH_ON_TRITON_LOAD_AFTER_READY.get() + if not should_crash and ( + free_gb is None or free_gb >= envs.SGLANG_TRITON_LOAD_WARNING_THRESHOLD_GB.get() + ): + return + + free_memory = f"{free_gb:.2f} GiB" if free_gb is not None else "unknown" + msg = ( + f"Triton kernel '{name}' device-loaded after serving started " + f"(free device mem: {free_memory}). Pre-load it during engine init " + f"to avoid CUDA OOM." + ) + if should_crash: + raise RuntimeError(msg) + logger.warning(msg)