diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 28d1beae5..40a7ce104 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 41f28a2e3..bd5f3d810 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/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) diff --git a/test/registered/unit/managers/test_detailed_annotations.py b/test/registered/unit/managers/test_detailed_annotations.py index 5d8a12d19..604af818b 100644 --- a/test/registered/unit/managers/test_detailed_annotations.py +++ b/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