[Spec] Emit step trace span for multi-layer draft-extend graph replays (#32850)

This commit is contained in:
Liangsheng Yin
2026-07-29 19:22:58 -07:00
committed by GitHub
parent 2aa86e9130
commit 313a518bee
@@ -82,6 +82,8 @@ from sglang.srt.utils.async_probe import (
maybe_detect_oob,
)
from sglang.srt.utils.common import empty_context, fast_topk
from sglang.srt.utils.nvtx_utils import profile_range
from sglang.srt.utils.profile_utils import build_step_span_name
_is_npu = is_npu()
_is_cpu = is_cpu()
@@ -738,6 +740,9 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
next_token_ids_backup = batch_result.next_token_ids.clone()
if can_cuda_graph:
# Graph replay bypasses ModelRunner.forward, which emits the
# step[...] trace span for every other phase; emit it here.
with profile_range(build_step_span_name(forward_batch)):
cgr = self.cuda_graph_runner_for_draft_extend
# Populate the single shared buffer set once; each step replays
# against it and the chain is advanced in place between steps.