From a3a1ebc7b7ee7ab185213b1a58ad6e482b7e15b1 Mon Sep 17 00:00:00 2001 From: cctry Date: Wed, 5 Aug 2026 14:55:03 -0700 Subject: [PATCH] Warn on risky serving-time Triton work (#33120) --- python/sglang/srt/environ.py | 6 + python/sglang/srt/managers/scheduler.py | 6 + python/sglang/srt/utils/triton_load_watch.py | 113 ++++++++++++++++++ .../unit/utils/test_triton_load_watch.py | 62 ++++++++++ 4 files changed, 187 insertions(+) create mode 100644 python/sglang/srt/utils/triton_load_watch.py create mode 100644 test/registered/unit/utils/test_triton_load_watch.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3ba3e809a..72081aa1b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1043,6 +1043,12 @@ class Envs: SGLANG_PYSPY_DUMP_BEFORE_CRASH = EnvBool(True) SGLANG_CUDA_COREDUMP_BEFORE_CRASH = EnvBool(True) SGLANG_CUDA_COREDUMP_BEFORE_CRASH_WAIT_SECS = EnvFloat(60.0) + # 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) # Encoder gRPC SGLANG_ENCODER_GRPC_TIMEOUT_SECS = EnvInt(60) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 42c4de72b..8365066b5 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 ( @@ -1640,6 +1641,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..ffff4744b --- /dev/null +++ b/python/sglang/srt/utils/triton_load_watch.py @@ -0,0 +1,113 @@ +"""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 + _prev_compile_listener = knobs.compilation.listener + knobs.runtime.kernel_load_start_hook.add(_on_kernel_load) + knobs.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) diff --git a/test/registered/unit/utils/test_triton_load_watch.py b/test/registered/unit/utils/test_triton_load_watch.py new file mode 100644 index 000000000..92e674c4e --- /dev/null +++ b/test/registered/unit/utils/test_triton_load_watch.py @@ -0,0 +1,62 @@ +"""Unit tests for triton_load_watch — no server, no model loading.""" + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") + +import unittest +from unittest.mock import patch + +import torch +import triton +import triton.language as tl + +from sglang.srt.environ import envs +from sglang.srt.utils import triton_load_watch +from sglang.test.test_utils import CustomTestCase + + +@triton.jit +def _probe_kernel(x_ptr, C: tl.constexpr): + # Each constexpr C is a distinct specialization -> a fresh device load. + tl.store(x_ptr + tl.program_id(0), C) + + +class TestTritonLoadWatch(CustomTestCase): + def tearDown(self): + # The watch is process-global; disarm so later tests in the same + # pytest process don't warn on their own first-use kernel loads. + triton_load_watch._serving_started = False + + def test_load_after_ready_warns_and_crashes(self): + triton_load_watch.install() + x = torch.zeros(4, device="cuda", dtype=torch.int32) + + # Loads during init (before serving starts) are silent. + with self.assertNoLogs(triton_load_watch.logger, level="WARNING"): + _probe_kernel[(1,)](x, C=1) + + triton_load_watch.mark_serving_started() + + # First use of a new specialization after ready warns with the name. + with ( + patch.object( + torch.cuda, "mem_get_info", return_value=(128 << 20, 80 << 30) + ), + self.assertLogs(triton_load_watch.logger, level="WARNING") as logs, + ): + _probe_kernel[(1,)](x, C=2) + self.assertTrue(any("free device mem" in line for line in logs.output)) + + # Already-loaded specializations stay silent. + with self.assertNoLogs(triton_load_watch.logger, level="WARNING"): + _probe_kernel[(1,)](x, C=2) + + # Crash mode turns the next late load into a hard error. + with envs.SGLANG_CRASH_ON_TRITON_LOAD_AFTER_READY.override(True): + with self.assertRaises(RuntimeError): + _probe_kernel[(1,)](x, C=3) + + +if __name__ == "__main__": + unittest.main()