[profiler] Label draft-runner steps DRAFT and target verify VERIFY in step spans (#38630)

Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

python/sglang/srt/model_executor/model_runner.py

@@ -1648,7 +1648,9 @@ class ModelRunner:
             self.msprobe_debugger.start(model=self.model, rank_id=rank_id)
 
         # Step span
-        step_span_ctx = profile_range(build_step_span_name(forward_batch))
+        step_span_ctx = profile_range(
+            build_step_span_name(forward_batch, is_draft_worker=self.is_draft_worker)
+        )
 
         canary_ctx = (
             context_tuple(

python/sglang/srt/utils/profile_utils.py

@@ -461,7 +461,10 @@ class _ProfilerRPD(_ProfilerConcreteBase):
 
 
 def build_step_span_name(
-    forward_batch: ForwardBatch, detailed_annotations: bool | None = None
+    forward_batch: ForwardBatch,
+    detailed_annotations: bool | None = None,
+    *,
+    is_draft_worker: bool = False,
 ) -> str:
     """Build the profile-trace span name for one forward step.
 
@@ -469,17 +472,27 @@ def build_step_span_name(
     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).
+
+    The target-verify step is labeled ``VERIFY``; every step a draft model
+    runner emits is labeled ``DRAFT`` (some draft paths borrow the TARGET_VERIFY
+    mode, so the mode name alone cannot tell the two models apart).
     """
     if detailed_annotations is None:
         detailed_annotations = detailed_annotations_enabled()
 
     mode = forward_batch.forward_mode
     bs = forward_batch.batch_size
+    if is_draft_worker:
+        stage = "DRAFT"
+    elif mode == ForwardMode.TARGET_VERIFY:
+        stage = "VERIFY"
+    else:
+        stage = mode.name
     if mode == ForwardMode.EXTEND:
         ext_toks = forward_batch.extend_num_tokens or 0
-        base = f"step[EXTEND bs={bs} toks={ext_toks}"
+        base = f"step[{stage} bs={bs} toks={ext_toks}"
     else:
-        base = f"step[{mode.name} bs={bs}"
+        base = f"step[{stage} bs={bs}"
 
     if detailed_annotations:
         suffix = build_detailed_annotation_suffix(forward_batch)

test/registered/unit/managers/test_detailed_annotations.py

@@ -152,7 +152,7 @@ class TestStepSpanDetailedAnnotations(CustomTestCase):
         )
         self.assertEqual(
             self._name(fb),
-            "step[TARGET_VERIFY bs=2 g_sq=6 g_sqsq=18 g_sqsk=90 g_sk=30]",
+            "step[VERIFY bs=2 g_sq=6 g_sqsq=18 g_sqsk=90 g_sk=30]",
         )
 
     def test_target_verify_without_cpu_mirror_falls_back_to_base(self):
@@ -162,7 +162,16 @@ class TestStepSpanDetailedAnnotations(CustomTestCase):
             seq_lens_cpu=None,
             num_tokens_per_req=3,
         )
-        self.assertEqual(self._name(fb), "step[TARGET_VERIFY bs=2]")
+        self.assertEqual(self._name(fb), "step[VERIFY bs=2]")
+
+    def test_draft_worker_prefixes_stage(self):
+        # A draft runner can run under TARGET_VERIFY; its span must read as
+        # the draft's, not the target's.
+        fb = _fb(ForwardMode.TARGET_VERIFY, batch_size=2)
+        self.assertEqual(
+            build_step_span_name(fb, detailed_annotations=False, is_draft_worker=True),
+            "step[DRAFT bs=2]",
+        )
 
     def test_draft_extend_v2_uses_extend_mirrors_with_context_prefix(self):
         # EAGLE/MTP draft-extend is extend-shaped