diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 6bf5d6182..e7b9e4688 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index adb556452..ee2fc439f 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 48d07334a..a6f3c7467 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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()