Unify NVTX annotation helpers and split the enable gate per subsystem (#28165)

python/sglang/srt/batch_overlap/operations.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-import os
+from contextlib import nullcontext
-from contextlib import contextmanager, 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 = {}

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]:

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]:

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)

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,

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]]:

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 = (
+        step_span_ctx = profile_range(_build_step_span_name(forward_batch))
-            torch.profiler.record_function(_build_step_span_name(forward_batch))
-            if torch.autograd._profiler_enabled()
-            else contextlib.nullcontext()
-        )
 
         canary_ctx = (
             context_tuple(

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 profile_range(name)
-        return torch.profiler.record_function(name)
-    return nullcontext()
 
 
 def move_accept_tokens_to_target_kvcache(

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).
+A span has two independent emitters:
-When disabled, the decorator/context manager add zero runtime overhead so they
+
-are safe to leave on hot scheduler paths.
+* ``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 contextlib import ExitStack, contextmanager, nullcontext
-from functools import wraps
+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. "
+            "An SGLANG_ENABLE_NVTX_* flag is set, but the `nvtx` package is "
-            "NVTX annotations are disabled."
+            "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
+# Default nvtx colors for statically-named spans (only used on the nvtx path).
-# distinguish in Nsight Systems.
 _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):
+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)
-    # record_function carries a non-trivial (~microseconds) cost per call even
+            stack.enter_context(_nvtx_module.annotate(debug_name, color=color))
-    # 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
 
 
-if NVTX_ENABLED:
+def profile_range(
-    nvtx_range = _nvtx_range_enabled
+    debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False
-else:
+):
-    # When NVTX is disabled, hand back a shared no-op context manager so hot
+    """Context manager emitting a profiler span for ``debug_name``.
-    # paths using `with nvtx_range(...)` pay no per-call generator overhead.
-    _NULL_CONTEXT = nullcontext()
 
-    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):
+def profile_method(
-    # Decide at decoration time. When NVTX is disabled this returns the
+    debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False
-    # original function untouched, so decorated methods on hot paths have zero
+):
-    # runtime cost.
+    """Decorator form of ``profile_range``."""
-    if not NVTX_ENABLED:
-        return lambda func: func
 
     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)