Revert "Warn on risky serving-time Triton work" (#33826)

2026-08-05 23:36:10 -07:00
parent 32e5d788bd
commit 5d1a0c7129
4 changed files with 0 additions and 187 deletions
@@ -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)
@@ -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.
@@ -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)
@@ -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()