[Spec] Add nvtx to spec regions (#27615)

Co-authored-by: Cursor <cursoragent@cursor.com>

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -71,6 +71,7 @@ from sglang.srt.speculative.spec_utils import (
     record_stream_each,
     record_stream_for_v2_verify,
     select_top_k_tokens,
+    spec_stage_span,
 )
 from sglang.srt.utils.async_probe import (
     maybe_detect_inf,
@@ -944,6 +945,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 ),
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft_extend"),
             ):
                 batch_output.next_draft_input = (
                     self.draft_worker._draft_extend_for_prefill(
@@ -976,6 +978,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 ),
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft"),
             ):
                 verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
             assert verify_input.is_verify_input()
@@ -990,6 +993,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 ),
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft_extend"),
             ):
                 self.draft_worker._draft_extend_for_decode(batch, batch_output)
 

python/sglang/srt/speculative/frozen_kv_mtp_worker.py

@@ -71,6 +71,7 @@ from sglang.srt.speculative.spec_utils import (
     fast_topk,
     generate_token_bitmask,
     select_top_k_tokens,
+    spec_stage_span,
 )
 from sglang.srt.utils import empty_context
 from sglang.srt.utils.async_probe import (
@@ -414,6 +415,7 @@ class FrozenKVMTPWorker(TpModelWorker):
                 self.draft_tp_context(self.draft_model_runner.tp_group),
                 speculative_moe_backend_context(),
                 speculative_moe_a2a_backend_context(),
+                spec_stage_span("draft_extend"),
             ):
                 self.forward_draft_extend(
                     batch,
@@ -434,6 +436,7 @@ class FrozenKVMTPWorker(TpModelWorker):
             self.draft_tp_context(self.draft_model_runner.tp_group),
             speculative_moe_backend_context(),
             speculative_moe_a2a_backend_context(),
+            spec_stage_span("draft"),
         ):
             verify_input = self.draft(batch)
         set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
@@ -465,7 +468,8 @@ class FrozenKVMTPWorker(TpModelWorker):
                 # step; `_run_assistant_seed_step` replaces it with a fresh
                 # `FrozenKVMTPDraftInput` for next iter.
                 batch.spec_info = draft_extend_input
-                self.forward_draft_extend_after_decode(batch)
+                with spec_stage_span("draft_extend"):
+                    self.forward_draft_extend_after_decode(batch)
             else:
                 # All reqs finished and dp_attention isn't forcing extend.
                 # Install an idle FrozenKVMTPDraftInput so next iter's scheduler

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
+from contextlib import contextmanager, nullcontext
 from typing import TYPE_CHECKING, List, Optional
 
 import torch
@@ -464,3 +464,12 @@ def draft_tp_context(tp_group: GroupCoordinator):
     # We disable mscclpp now because it doesn't support 2 comm groups.
     with patch_tensor_parallel_group(tp_group):
         yield
+
+
+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()