From bb48405c3121f23fd89a40ede45725dfe378cd79 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 14 Jun 2026 00:04:59 -0700 Subject: [PATCH] Unify NVTX annotation helpers and split the enable gate per subsystem (#28165) --- python/sglang/srt/batch_overlap/operations.py | 32 ++---- python/sglang/srt/disaggregation/decode.py | 4 +- python/sglang/srt/disaggregation/prefill.py | 4 +- python/sglang/srt/environ.py | 7 +- python/sglang/srt/managers/scheduler.py | 10 +- .../scheduler_components/request_receiver.py | 4 +- .../sglang/srt/model_executor/model_runner.py | 7 +- python/sglang/srt/speculative/spec_utils.py | 7 +- python/sglang/srt/utils/nvtx_utils.py | 103 +++++++++++------- 9 files changed, 95 insertions(+), 83 deletions(-) diff --git a/python/sglang/srt/batch_overlap/operations.py b/python/sglang/srt/batch_overlap/operations.py index ea9090e9c..2b853e4ec 100644 --- a/python/sglang/srt/batch_overlap/operations.py +++ b/python/sglang/srt/batch_overlap/operations.py @@ -1,7 +1,6 @@ from __future__ import annotations -import os -from contextlib import contextmanager, nullcontext +from contextlib import nullcontext from dataclasses import dataclass, replace from typing import ( TYPE_CHECKING, @@ -15,23 +14,17 @@ from typing import ( Union, ) -import torch - from sglang.srt.layers.dp_attention import set_dp_buffer_len from sglang.srt.model_executor.forward_context import ( forward_context, get_forward_context, ) +from sglang.srt.utils.nvtx_utils import operations_nvtx_range if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ForwardContext -_ENABLE_PROFILE = bool(int(os.environ.get("SGLANG_OPERATIONS_ENABLE_PROFILE", "0"))) - -if _ENABLE_PROFILE: - import nvtx - def execute_operations(inputs, operations): stages = _convert_operations_to_stages(operations) @@ -157,9 +150,16 @@ class _StageExecutor: if self._child_ctx is not None else nullcontext() ) - with ctx_mgr, _annotate_region(debug_name=f"{self._debug_name}{self._index}"): + stage_range = operations_nvtx_range( + debug_name=f"{self._debug_name}{self._index}", + color="orange", + ) + with ctx_mgr, stage_range: for op in stage: - with _annotate_region(debug_name=op.debug_name): + with operations_nvtx_range( + debug_name=op.debug_name, + color="yellow", + ): self._stage_output = op.fn( state=self._stage_state, **( @@ -183,16 +183,6 @@ class _StageExecutor: return len(self._stages) -@contextmanager -def _annotate_region(debug_name): - if _ENABLE_PROFILE: - with torch.autograd.profiler.record_function(debug_name): - with nvtx.annotate(debug_name): - yield - else: - yield - - class _StateDict: def __init__(self): self._data = {} diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 208d0976e..7a751077c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -85,7 +85,7 @@ from sglang.srt.observability.req_time_stats import ( ) from sglang.srt.utils import get_num_new_pages from sglang.srt.utils.network import NetworkAddress -from sglang.srt.utils.nvtx_utils import nvtx_annotated_method +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) @@ -1825,7 +1825,7 @@ class SchedulerDisaggregationDecodeMixin: return GenerationBatchResult() - @nvtx_annotated_method("scheduler.get_next_batch_to_run") + @scheduler_nvtx_method("scheduler.get_next_batch_to_run") def get_next_disagg_decode_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index eee6fd133..a57bedb36 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -61,7 +61,7 @@ from sglang.srt.mem_cache.common import ( ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.observability.req_time_stats import set_schedule_time_batch -from sglang.srt.utils.nvtx_utils import nvtx_annotated_method +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method if TYPE_CHECKING: from torch.distributed import ProcessGroup @@ -406,7 +406,7 @@ class SchedulerDisaggregationPrefillMixin: if room is not None and room in kv_mgr.transfer_infos: prefetch(room) - @nvtx_annotated_method("scheduler.get_next_batch_to_run") + @scheduler_nvtx_method("scheduler.get_next_batch_to_run") def get_next_disagg_prefill_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f7bf43b19..79d2a922a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -243,7 +243,12 @@ class Envs: SGLANG_PROFILE_WITH_STACK = EnvBool(True) SGLANG_PROFILE_RECORD_SHAPES = EnvBool(True) SGLANG_PROFILE_V2 = EnvBool(False) - SGLANG_ENABLE_NVTX = EnvBool(False) + SGLANG_ENABLE_NVTX_SCHEDULER = EnvBoolWithAlias( + False, deprecated_name="SGLANG_ENABLE_NVTX" + ) + SGLANG_ENABLE_NVTX_OPERATIONS = EnvBoolWithAlias( + False, deprecated_name="SGLANG_OPERATIONS_ENABLE_PROFILE" + ) SGLANG_RECORD_STEP_TIME = EnvBool(False) SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 98ba64f2d..953c4fc45 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -263,7 +263,7 @@ from sglang.srt.utils.hf_transformers_utils import ( get_tokenizer_from_processor, ) from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node -from sglang.srt.utils.nvtx_utils import nvtx_annotated_method +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method from sglang.srt.utils.tensor_bridge import use_mlx from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.utils import TypeBasedDispatcher, get_exception_traceback @@ -1586,7 +1586,7 @@ class Scheduler( return disable_overlap_for_batch or need_grammar_sync - @nvtx_annotated_method("scheduler.process_input_requests") + @scheduler_nvtx_method("scheduler.process_input_requests") def process_input_requests(self, recv_reqs: List): now = time.monotonic() self.session_controller.maybe_reap(now) @@ -2463,7 +2463,7 @@ class Scheduler( # todo hisparse, maybe other info to contain for the new batch return batch - @nvtx_annotated_method("scheduler.get_next_batch_to_run") + @scheduler_nvtx_method("scheduler.get_next_batch_to_run") def get_next_batch_to_run(self) -> Optional[ScheduleBatch]: if self.enable_fpm: self._fpm_batch_t0 = time.monotonic() @@ -3051,7 +3051,7 @@ class Scheduler( else: batch.sampling_info = sched_sampling_info - @nvtx_annotated_method("scheduler.run_batch") + @scheduler_nvtx_method("scheduler.run_batch") def run_batch( self, batch: ScheduleBatch, @@ -3273,7 +3273,7 @@ class Scheduler( if batch_result.logits_output is not None: batch_result.logits_output.next_token_logits = None - @nvtx_annotated_method("scheduler.process_batch_result") + @scheduler_nvtx_method("scheduler.process_batch_result") def process_batch_result( self, batch: ScheduleBatch, diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 1d06d0750..4c0f6fa91 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -29,7 +29,7 @@ from sglang.srt.utils import ( broadcast_pyobj, point_to_point_pyobj, ) -from sglang.srt.utils.nvtx_utils import nvtx_annotated_method +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig @@ -68,7 +68,7 @@ class SchedulerRequestReceiver: return False return num_recv_reqs >= self.max_recv_per_poll - @nvtx_annotated_method("scheduler.recv_requests") + @scheduler_nvtx_method("scheduler.recv_requests") def recv_requests( self, ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e418ffd47..7fdbc82e2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -232,6 +232,7 @@ from sglang.srt.utils import ( from sglang.srt.utils.common import ceil_align, next_power_of_2, require_mlp_sync from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks +from sglang.srt.utils.nvtx_utils import profile_range from sglang.srt.utils.offloader import ( create_offloader_from_server_args, get_offloader, @@ -3454,11 +3455,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.msprobe_debugger.start(model=self.model, rank_id=rank_id) # Step span - step_span_ctx = ( - torch.profiler.record_function(_build_step_span_name(forward_batch)) - if torch.autograd._profiler_enabled() - else contextlib.nullcontext() - ) + step_span_ctx = profile_range(_build_step_span_name(forward_batch)) canary_ctx = ( context_tuple( diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 604db2faf..eb7c410f1 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging import os import time -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from typing import TYPE_CHECKING, List, Optional import torch @@ -45,6 +45,7 @@ from sglang.srt.speculative.triton_ops.eagle import ( ) from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, next_power_of_2 from sglang.srt.utils.async_probe import maybe_detect_oob +from sglang.srt.utils.nvtx_utils import profile_range _is_cuda = is_cuda() _is_hip = is_hip() @@ -477,9 +478,7 @@ def spec_stage_span(name: str): """Profiler span for a coarse speculative-decoding stage (``draft`` / ``draft_extend`` / ``verify``). """ - if torch.autograd._profiler_enabled(): - return torch.profiler.record_function(name) - return nullcontext() + return profile_range(name) def move_accept_tokens_to_target_kvcache( diff --git a/python/sglang/srt/utils/nvtx_utils.py b/python/sglang/srt/utils/nvtx_utils.py index 834180387..86f1edba0 100644 --- a/python/sglang/srt/utils/nvtx_utils.py +++ b/python/sglang/srt/utils/nvtx_utils.py @@ -11,16 +11,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Lightweight NVTX annotations for the scheduler main loop. +"""Profiler span helpers for hot SGLang code paths. -Enabled via the ``SGLANG_ENABLE_NVTX`` environment variable (off by default). -When disabled, the decorator/context manager add zero runtime overhead so they -are safe to leave on hot scheduler paths. +A span has two independent emitters: + +* ``record_function`` -- emitted whenever a torch profiler is active, so spans + show up in torch/Perfetto traces for free (no env, no extra package). +* ``nvtx`` range -- emitted only when the caller opts in via ``nvtx_enabled`` + (wired to a per-subsystem ``SGLANG_ENABLE_NVTX_*`` gate) and the ``nvtx`` + package is importable, for Nsight Systems timelines. + +Decoupling the two lets every annotation site -- scheduler stages, batch-overlap +ops, and the speculative-decoding / forward spans -- share one primitive. """ import logging -from contextlib import contextmanager, nullcontext -from functools import wraps +from contextlib import ExitStack, contextmanager, nullcontext +from functools import partial, wraps +from typing import Optional import torch @@ -28,23 +36,27 @@ from sglang.srt.environ import envs logger = logging.getLogger(__name__) -_NVTX_ENV_ENABLED = envs.SGLANG_ENABLE_NVTX.get() +_SCHEDULER_NVTX = envs.SGLANG_ENABLE_NVTX_SCHEDULER.get() +_OPERATIONS_NVTX = envs.SGLANG_ENABLE_NVTX_OPERATIONS.get() + _nvtx_module = None -if _NVTX_ENV_ENABLED: +if _SCHEDULER_NVTX or _OPERATIONS_NVTX: try: import nvtx as _nvtx_module # type: ignore except ImportError: logger.warning( - "SGLANG_ENABLE_NVTX is set, but the `nvtx` package is missing. " - "NVTX annotations are disabled." + "An SGLANG_ENABLE_NVTX_* flag is set, but the `nvtx` package is " + "missing. NVTX markers are disabled; torch profiler spans still emit." ) -NVTX_ENABLED = _nvtx_module is not None +NVTX_AVAILABLE = _nvtx_module is not None +# Per-subsystem nvtx gates: emit nvtx ranges only when the flag is set AND the +# package is importable. The record_function path is independent of both. +NVTX_SCHEDULER_ENABLED = _SCHEDULER_NVTX and NVTX_AVAILABLE +NVTX_OPERATIONS_ENABLED = _OPERATIONS_NVTX and NVTX_AVAILABLE -# Colors are assigned per scheduler main-loop stage so the markers are easy to -# distinguish in Nsight Systems. +# Default nvtx colors for statically-named spans (only used on the nvtx path). _NVTX_COLOR_MAP = { - # === Scheduler main loop (pipeline order) === "scheduler.recv_requests": "blue", "scheduler.process_input_requests": "purple", "scheduler.get_next_batch_to_run": "green", @@ -52,47 +64,56 @@ _NVTX_COLOR_MAP = { "scheduler.process_batch_result": "cyan", } +_NULL_CONTEXT = nullcontext() + @contextmanager -def _nvtx_range_enabled(debug_name: str): - color = _NVTX_COLOR_MAP.get(debug_name) - # record_function carries a non-trivial (~microseconds) cost per call even - # when no PyTorch profiler is collecting, so only pay it when one is active - # (e.g. Chrome-trace export). For Nsight-only runs the nvtx.annotate marker - # alone is enough. - if torch.autograd._profiler_enabled(): - with torch.autograd.profiler.record_function(debug_name): - with _nvtx_module.annotate(debug_name, color=color): - yield - else: - with _nvtx_module.annotate(debug_name, color=color): - yield +def _profile_range_impl( + debug_name: str, color: Optional[str], record: bool, nvtx_enabled: bool +): + with ExitStack() as stack: + if record: + stack.enter_context(torch.profiler.record_function(debug_name)) + if nvtx_enabled: + if color is None: + color = _NVTX_COLOR_MAP.get(debug_name) + stack.enter_context(_nvtx_module.annotate(debug_name, color=color)) + yield -if NVTX_ENABLED: - nvtx_range = _nvtx_range_enabled -else: - # When NVTX is disabled, hand back a shared no-op context manager so hot - # paths using `with nvtx_range(...)` pay no per-call generator overhead. - _NULL_CONTEXT = nullcontext() +def profile_range( + debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False +): + """Context manager emitting a profiler span for ``debug_name``. - def nvtx_range(debug_name: str): + A torch ``record_function`` is emitted whenever a torch profiler is active; + an nvtx range is emitted additionally when ``nvtx_enabled`` is true. Returns a + shared no-op when neither applies, so off-profile hot paths pay only one + ``_profiler_enabled()`` check. + """ + record = torch.autograd._profiler_enabled() + if not record and not nvtx_enabled: return _NULL_CONTEXT + return _profile_range_impl(debug_name, color, record, nvtx_enabled) -def nvtx_annotated_method(debug_name: str): - # Decide at decoration time. When NVTX is disabled this returns the - # original function untouched, so decorated methods on hot paths have zero - # runtime cost. - if not NVTX_ENABLED: - return lambda func: func +def profile_method( + debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False +): + """Decorator form of ``profile_range``.""" def decorator(func): @wraps(func) def wrapper(*args, **kwargs): - with nvtx_range(debug_name): + with profile_range(debug_name, color=color, nvtx_enabled=nvtx_enabled): return func(*args, **kwargs) return wrapper return decorator + + +# Pre-bound per-subsystem helpers: torch spans always (under a profiler), nvtx +# ranges only when that subsystem's gate is on. +scheduler_nvtx_method = partial(profile_method, nvtx_enabled=NVTX_SCHEDULER_ENABLED) +operations_nvtx_range = partial(profile_range, nvtx_enabled=NVTX_OPERATIONS_ENABLED)