Profiling Enhancements [2/3]: detailed execution step annotations (#24911)

python/sglang/srt/managers/io_struct.py

@@ -2094,6 +2094,8 @@ class ProfileReq(BaseReq, kw_only=True):
     profile_prefix: Optional[str] = None
     # Only profile these stages and ignore others
     profile_stages: Optional[List[str]] = None
+    # Add iteration-level annotations (KV / request aggregates) for roofline-style analysis
+    detailed_annotations: bool = False
 
 
 class ProfileReqOutput(BaseReq, kw_only=True):

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -19,6 +19,7 @@ import torch
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.step_span_utils import set_detailed_annotations_enabled
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_device
 from sglang.srt.utils import is_mps, is_npu
@@ -78,6 +79,7 @@ class SchedulerProfilerManager:
         self.profile_by_stage: bool = False
         self.profile_in_progress: bool = False
         self.merge_profiles = False
+        self.detailed_annotations: bool = False
 
         # For ROCM
         self.rpd_profiler = None
@@ -94,9 +96,11 @@ class SchedulerProfilerManager:
         profile_id: str,
         merge_profiles: bool = False,
         profile_prefix: str = "",
+        detailed_annotations: bool = False,
         profile_stages: Optional[List[str]] = None,
     ) -> ProfileReqOutput:
         if envs.SGLANG_PROFILE_V2.get():
+            self.detailed_annotations = detailed_annotations
             return self._profile_manager.configure(
                 output_dir=output_dir,
                 start_step=start_step,
@@ -109,6 +113,7 @@ class SchedulerProfilerManager:
                 merge_profiles=merge_profiles,
                 profile_prefix=profile_prefix,
                 profile_stages=profile_stages,
+                detailed_annotations=detailed_annotations,
             )
 
         if self.profile_in_progress:
@@ -131,6 +136,7 @@ class SchedulerProfilerManager:
         self.profiler_activities = activities
         self.profile_id = profile_id
         self.profile_prefix = profile_prefix
+        self.detailed_annotations = detailed_annotations
 
         if start_step:
             self.profiler_start_forward_ct = max(start_step, self.get_forward_ct() + 1)
@@ -153,10 +159,17 @@ class SchedulerProfilerManager:
 
         return ProfileReqOutput(success=True, message="Succeeded")
 
+    def _apply_detailed_annotations(self, enabled: bool) -> None:
+        # Toggle the process-wide flag read by build_step_span_name; folds the
+        # per-phase sq/sqsq/sqsk/sk aggregates (context c_ / generation g_)
+        # into the step span while a detailed-annotation profile is active.
+        set_detailed_annotations_enabled(enabled)
+
     def _start_profile(
         self, stage: Optional[ForwardMode] = None
     ) -> ProfileReqOutput | None:
         if envs.SGLANG_PROFILE_V2.get():
+            self._apply_detailed_annotations(self.detailed_annotations)
             return self._profile_manager.manual_start()
 
         stage_str = f" for {stage.name}" if stage else ""
@@ -262,6 +275,7 @@ class SchedulerProfilerManager:
                 torch.cuda.cudart().cudaProfilerStart()
             self.profile_in_progress = True
 
+        self._apply_detailed_annotations(self.detailed_annotations)
         return ProfileReqOutput(success=True, message="Succeeded")
 
     def _merge_profile_traces(self) -> str:
@@ -300,6 +314,7 @@ class SchedulerProfilerManager:
         self, stage: Optional[ForwardMode] = None
     ) -> ProfileReqOutput | None:
         if envs.SGLANG_PROFILE_V2.get():
+            self._apply_detailed_annotations(False)
             return self._profile_manager.manual_stop()
 
         if not self.profile_in_progress:
@@ -387,6 +402,7 @@ class SchedulerProfilerManager:
         self.profile_in_progress = False
         self.profiler_start_forward_ct = None
 
+        self._apply_detailed_annotations(False)
         return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}")
 
     def _profile_batch_predicate(self, batch: ScheduleBatch):
@@ -443,6 +459,7 @@ class SchedulerProfilerManager:
                     recv_req.profile_id,
                     recv_req.merge_profiles,
                     recv_req.profile_prefix,
+                    recv_req.detailed_annotations,
                     recv_req.profile_stages,
                 )
             else:
@@ -457,6 +474,7 @@ class SchedulerProfilerManager:
                     recv_req.profile_id,
                     recv_req.merge_profiles,
                     recv_req.profile_prefix,
+                    recv_req.detailed_annotations,
                 )
                 return self._start_profile()
         else:

python/sglang/srt/model_executor/step_span_utils.py

@@ -0,0 +1,139 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Profile-trace step-span naming (kept dependency-light for CPU unit tests).
+
+The step span wraps each ``ModelRunner.forward`` in the torch/Perfetto trace.
+Its name carries the forward mode and batch shape; when detailed annotations
+are enabled it also folds in the per-iteration aggregates (for roofline-style
+analysis) so a single label describes both timing and the analytical work of
+that forward.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Tuple
+
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+# Process-wide toggle for detailed step-span annotations. Set by the scheduler's
+# profiler manager on profile start/stop and read by build_step_span_name, so no
+# per-runner flag or scheduler bridge is needed. Covers all model runners in the
+# process (e.g. both the EAGLE draft and target runners).
+_DETAILED_ANNOTATIONS_ENABLED = False
+
+
+def set_detailed_annotations_enabled(enabled: bool) -> None:
+    global _DETAILED_ANNOTATIONS_ENABLED
+    _DETAILED_ANNOTATIONS_ENABLED = bool(enabled)
+
+
+def detailed_annotations_enabled() -> bool:
+    return _DETAILED_ANNOTATIONS_ENABLED
+
+
+def _agg(nqs: List[int], nkvs: List[int]) -> Tuple[int, int, int, int]:
+    """Return (Σ N_Q, Σ N_KV, Σ N_Q², Σ N_Q·N_KV) for one request group."""
+    sq = sum(nqs)
+    sk = sum(nkvs)
+    sqsq = sum(nq * nq for nq in nqs)
+    sqsk = sum(nq * nkv for nq, nkv in zip(nqs, nkvs))
+    return sq, sk, sqsq, sqsk
+
+
+def _decode_query_width(forward_batch: ForwardBatch) -> int:
+    """Per-request query-token count (N_Q) for a decode-family forward.
+
+    Vanilla decode emits one token per request, but speculative decoding does
+    not: EAGLE/MTP draft-decode and target-verify process a uniform
+    ``num_tokens_per_req`` tokens per request (draft top-k for draft-decode,
+    the draft-token count for verify), so N_Q per request is that width, and
+    the step's total Σ N_Q is ``bs * num_tokens_per_req`` rather than ``bs``.
+
+    The width is read from ``forward_batch.spec_info.num_tokens_per_req`` and
+    falls back to 1 when there is no spec input or the width is unset (-1).
+    """
+    spec = getattr(forward_batch, "spec_info", None)
+    width = getattr(spec, "num_tokens_per_req", -1) if spec is not None else -1
+    return width if isinstance(width, int) and width > 0 else 1
+
+
+def build_detailed_annotation_suffix(forward_batch: ForwardBatch) -> str:
+    """Compute the detailed-annotation aggregates from the batch's CPU-side length mirrors.
+
+    All aggregates are emitted, prefixed by the roofline compute-shape bucket:
+    ``c_`` for context/extend-shaped work (EXTEND and DRAFT_EXTEND_V2) and ``g_`` for
+    single-query generation (DECODE, TARGET_VERIFY), with MIXED emitting both
+    groups.
+    """
+    mode = forward_batch.forward_mode
+    seq_lens_cpu = forward_batch.seq_lens_cpu
+
+    # DECODE (vanilla or spec draft-decode) and TARGET_VERIFY both key off
+    # ``seq_lens_cpu`` for N_KV and a uniform per-request query width N_Q, and
+    # are both classified as generation (``g_``) by request phase
+    #   * DECODE        -> N_Q is 1 (vanilla) or the spec draft-decode width.
+    #   * TARGET_VERIFY -> N_Q is ``num_tokens_per_req`` (the draft-token count);
+    #                      the request is past its prompt (generation phase)
+    if mode == ForwardMode.DECODE or mode == ForwardMode.TARGET_VERIFY:
+        if seq_lens_cpu is None:
+            return ""
+        nq = _decode_query_width(forward_batch)
+        nkvs = [int(x) for x in seq_lens_cpu.tolist()]
+        nqs = [nq] * len(nkvs)
+        sq, sk, sqsq, sqsk = _agg(nqs, nkvs)
+        # ``sq`` is always emitted (self-contained suffix): it equals ``bs``
+        # (vanilla decode) or ``bs * num_tokens_per_req`` (spec draft-decode /
+        # target-verify).
+        return f"g_sq={sq} g_sqsq={sqsq} g_sqsk={sqsk} g_sk={sk}"
+
+    ext_seq = forward_batch.extend_seq_lens_cpu
+    ext_prefix = forward_batch.extend_prefix_lens_cpu
+    if ext_seq is None or ext_prefix is None:
+        return ""
+
+    if mode == ForwardMode.EXTEND or mode == ForwardMode.DRAFT_EXTEND_V2:
+        # Both are extend-shaped, multi-query context
+        nqs = [int(q) for q in ext_seq]
+        nkvs = [int(p) + int(q) for p, q in zip(ext_prefix, ext_seq)]
+        sq, sk, sqsq, sqsk = _agg(nqs, nkvs)
+        return f"c_sq={sq} c_sqsq={sqsq} c_sqsk={sqsk} c_sk={sk}"
+
+    if mode == ForwardMode.MIXED:
+        # A running-decode request appears as a length-1 extend; everything
+        # else is a context (prefill) chunk.
+        c_nqs: List[int] = []
+        c_nkvs: List[int] = []
+        g_nqs: List[int] = []
+        g_nkvs: List[int] = []
+        for p, q in zip(ext_prefix, ext_seq):
+            nq, nkv = int(q), int(p) + int(q)
+            if nq == 1:
+                g_nqs.append(nq)
+                g_nkvs.append(nkv)
+            else:
+                c_nqs.append(nq)
+                c_nkvs.append(nkv)
+        c_sq, c_sk, c_sqsq, c_sqsk = _agg(c_nqs, c_nkvs)
+        g_sq, g_sk, g_sqsq, g_sqsk = _agg(g_nqs, g_nkvs)
+        return (
+            f"c={len(c_nqs)} g={len(g_nqs)} "
+            f"c_sq={c_sq} c_sk={c_sk} c_sqsq={c_sqsq} c_sqsk={c_sqsk} "
+            f"g_sq={g_sq} g_sk={g_sk} g_sqsq={g_sqsq} g_sqsk={g_sqsk}"
+        )
+
+    return ""

python/sglang/srt/utils/profile_utils.py

@@ -12,6 +12,11 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import ProfileReqOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.step_span_utils import (
+    build_detailed_annotation_suffix,
+    detailed_annotations_enabled,
+    set_detailed_annotations_enabled,
+)
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_device
 from sglang.srt.utils import is_npu
@@ -76,6 +81,7 @@ class ProfileManager:
         self.first_rank_in_node = ps.gpu_id == get_device().base_gpu_id
         self.profiler_kwargs = None
         self.profiler = None
+        self.detailed_annotations = False
 
     def step(self, forward_mode: ForwardMode):
         stage = _get_stage_from_forward_mode(forward_mode)
@@ -98,7 +104,9 @@ class ProfileManager:
         merge_profiles: bool,
         profile_prefix: str,
         profile_stages: Optional[List[str]] = None,
+        detailed_annotations: bool = False,
     ):
+        self.detailed_annotations = detailed_annotations
         # not supported yet
         assert start_step is None
         assert (
@@ -141,6 +149,9 @@ class ProfileManager:
         )
 
         assert self.profiler is None
+        # Fold the per-phase c_/g_ aggregates into the step span while this
+        # stage's profile is active (v2 auto-start path; reset in _do_stop).
+        set_detailed_annotations_enabled(self.detailed_annotations)
         self.profiler = _ProfilerBase.create(
             **self.profiler_kwargs,
             ps=self.ps,
@@ -157,6 +168,10 @@ class ProfileManager:
             f"Profiling done. Traces are saved to: {self.profiler_kwargs['output_dir']}"
         )
         self.profiler = None
+        # Clear the detailed step-span toggle here too: the v2 trigger auto-stop
+        # goes through _do_stop (not SchedulerProfilerManager._stop_profile), so
+        # this guarantees the flag resets on every stop path.
+        set_detailed_annotations_enabled(False)
 
 
 def _get_stage_from_forward_mode(forward_mode: ForwardMode):
@@ -445,11 +460,29 @@ class _ProfilerRPD(_ProfilerConcreteBase):
             rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
 
 
-def build_step_span_name(forward_batch: ForwardBatch) -> str:
-    """Build a profile-trace span name for one forward step."""
+def build_step_span_name(
+    forward_batch: ForwardBatch, detailed_annotations: bool | None = None
+) -> str:
+    """Build the profile-trace span name for one forward step.
+
+    Detailed annotations are folded into the label (via
+    build_detailed_annotation_suffix) when enabled. detailed_annotations
+    defaults to the process-wide toggle (detailed_annotations_enabled, set
+    by the profiler manager); pass an explicit bool to override (e.g. in tests).
+    """
+    if detailed_annotations is None:
+        detailed_annotations = detailed_annotations_enabled()
+
     mode = forward_batch.forward_mode
     bs = forward_batch.batch_size
     if mode == ForwardMode.EXTEND:
         ext_toks = forward_batch.extend_num_tokens or 0
-        return f"step[EXTEND bs={bs} toks={ext_toks}]"
-    return f"step[{mode.name} bs={bs}]"
+        base = f"step[EXTEND bs={bs} toks={ext_toks}"
+    else:
+        base = f"step[{mode.name} bs={bs}"
+
+    if detailed_annotations:
+        suffix = build_detailed_annotation_suffix(forward_batch)
+        if suffix:
+            base = f"{base} {suffix}"
+    return f"{base}]"