minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-14 00:04:59 -07:00
committed by GitHub
parent 50993554d8
commit bb48405c31
9 changed files with 95 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/operations.py
+11 -21
View File
@@ -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 = {}
python/sglang/srt/disaggregation/decode.py
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+6 -1
View File
@@ -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
+5 -5
View File
@@ -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
+2 -2
View File
@@ -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
+2 -5
View File
@@ -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(
python/sglang/srt/speculative/spec_utils.py
+3 -4
View File
@@ -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(
python/sglang/srt/utils/nvtx_utils.py
+62 -41
View File
@@ -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)