[Spec] Emit step trace span for multi-layer draft-extend graph replays (#32850)
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user