diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 92421fae1..cb68793d1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1041,12 +1041,6 @@ 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 3aed2f55e..33643749b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -301,7 +301,6 @@ 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,11 +1628,6 @@ 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 deleted file mode 100644 index ffff4744b..000000000 --- a/python/sglang/srt/utils/triton_load_watch.py +++ /dev/null @@ -1,113 +0,0 @@ -"""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 deleted file mode 100644 index 92e674c4e..000000000 --- a/test/registered/unit/utils/test_triton_load_watch.py +++ /dev/null @@ -1,62 +0,0 @@ -"""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()