minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-29 19:22:58 -07:00
committed by GitHub
parent 2aa86e9130
commit 313a518bee
1 changed files with 51 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+51 -46
View File
@@ -82,6 +82,8 @@ from sglang.srt.utils.async_probe import (
maybe_detect_oob, maybe_detect_oob,
) )
from sglang.srt.utils.common import empty_context, fast_topk 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_npu = is_npu()
_is_cpu = is_cpu() _is_cpu = is_cpu()
@@ -738,53 +740,56 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
next_token_ids_backup = batch_result.next_token_ids.clone() next_token_ids_backup = batch_result.next_token_ids.clone()
if can_cuda_graph: if can_cuda_graph:
cgr = self.cuda_graph_runner_for_draft_extend # Graph replay bypasses ModelRunner.forward, which emits the
# Populate the single shared buffer set once; each step replays # step[...] trace span for every other phase; emit it here.
# against it and the chain is advanced in place between steps. with profile_range(build_step_span_name(forward_batch)):
cgr.prepare(forward_batch) cgr = self.cuda_graph_runner_for_draft_extend
rotates_in_graph = cgr.rotates_in_graph # Populate the single shared buffer set once; each step replays
for step in range(self.speculative_num_steps): # against it and the chain is advanced in place between steps.
_out, ret_topk_p, ret_topk_index = cgr.replay(step) cgr.prepare(forward_batch)
# Rejection sampling with the per-step runner re-picks X ~ q rotates_in_graph = cgr.rotates_in_graph
# worker-side so the worker rotation carries it to step N+1; the for step in range(self.speculative_num_steps):
# single-CG runner samples in-graph (q cloned after the loop). _out, ret_topk_p, ret_topk_index = cgr.replay(step)
if ( # Rejection sampling with the per-step runner re-picks X ~ q
self.use_rejection_sampling # worker-side so the worker rotation carries it to step N+1; the
and self.topk == 1 # single-CG runner samples in-graph (q cloned after the loop).
and not rotates_in_graph if (
): self.use_rejection_sampling
if cgr.prune_draft_extend_logits: and self.topk == 1
step_logits = _out.next_token_logits and not rotates_in_graph
):
if cgr.prune_draft_extend_logits:
step_logits = _out.next_token_logits
else:
sel = cgr.buffers.select_index[: cgr.raw_bs]
step_logits = _out.next_token_logits[sel]
probs, ret_topk_p, ret_topk_index = sample_draft_proposal(
step_logits,
forward_batch.sampling_info.temperatures,
)
ret_draft_probs_list.append(probs)
if rotates_in_graph:
# Single-CG step outputs coexist until the trailing cat.
ret_topk_p_list.append(ret_topk_p)
ret_topk_index_list.append(ret_topk_index)
else: else:
sel = cgr.buffers.select_index[: cgr.raw_bs] # Per-step graphs share the global graph pool; snapshot
step_logits = _out.next_token_logits[sel] # before the next step's replay can reuse the buffer.
probs, ret_topk_p, ret_topk_index = sample_draft_proposal( ret_topk_p_list.append(ret_topk_p.clone())
step_logits, ret_topk_index_list.append(ret_topk_index.clone())
forward_batch.sampling_info.temperatures, # Advance the draft chain by rotating the shared input_ids window
) # in place; step N+1's graph then reads the rotated values. The
ret_draft_probs_list.append(probs) # single-CG runner rotates in-graph, so skip the worker-side rotate.
if rotates_in_graph: if step < self.speculative_num_steps - 1 and not rotates_in_graph:
# Single-CG step outputs coexist until the trailing cat. rotate_input_ids(
ret_topk_p_list.append(ret_topk_p) cgr.buffers.input_ids[: cgr.raw_num_tokens],
ret_topk_index_list.append(ret_topk_index) cgr.buffers.extend_start_loc[: cgr.raw_bs],
else: cgr.buffers.extend_seq_lens[: cgr.raw_bs],
# Per-step graphs share the global graph pool; snapshot ret_topk_index,
# before the next step's replay can reuse the buffer. cgr.buffers.select_index[: cgr.raw_bs],
ret_topk_p_list.append(ret_topk_p.clone()) )
ret_topk_index_list.append(ret_topk_index.clone()) if self.use_rejection_sampling and self.topk == 1 and rotates_in_graph:
# Advance the draft chain by rotating the shared input_ids window ret_draft_probs = cgr.clone_draft_probs()
# in place; step N+1's graph then reads the rotated values. The
# single-CG runner rotates in-graph, so skip the worker-side rotate.
if step < self.speculative_num_steps - 1 and not rotates_in_graph:
rotate_input_ids(
cgr.buffers.input_ids[: cgr.raw_num_tokens],
cgr.buffers.extend_start_loc[: cgr.raw_bs],
cgr.buffers.extend_seq_lens[: cgr.raw_bs],
ret_topk_index,
cgr.buffers.select_index[: cgr.raw_bs],
)
if self.use_rejection_sampling and self.topk == 1 and rotates_in_graph:
ret_draft_probs = cgr.clone_draft_probs()
else: else:
logger.warning_once( logger.warning_once(
"can't use cuda graph for draft extend! may have correctness issue!" "can't use cuda graph for draft extend! may have correctness issue!"