Skip to content

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

Merged
hnyls2002 merged 1 commit into
mainfrom
lsyin/draft-extend-step-span
Jul 30, 2026
Merged

hnyls2002 merged 1 commit into
mainfrom
lsyin/draft-extend-step-span

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Jul 30, 2026 •

Copy link
Copy Markdown
Collaborator

Every forward phase gets a step[<MODE> bs=N] trace span from ModelRunner.forward, but the multi-layer EAGLE decode draft-extend replays its cuda graphs directly through the graph runner (cgr.prepare + cgr.replay(step)), bypassing that entry. In profiler traces the draft phase's kernels therefore show up under scheduler.run_batch with no step marker, while TARGET_VERIFY and EXTEND have theirs.

Wrap the graph-replay block in profile_range(build_step_span_name(forward_batch)) so traces now carry step[DRAFT_EXTEND_V2 bs=N] alongside the existing spans. Pure observability: the wrapper resolves to a shared no-op context when no profiler is active, and the prefill-side chain already goes through ModelRunner.forward so it needs nothing.

Verified on a tp=8 EAGLE run: captured trace shows step[DRAFT_EXTEND_V2 bs=1] spans matching the draft replay windows; generated output unchanged.

image

CI States

Latest PR Test (Base): 🚫 Run #30507587448
Latest PR Test (Extra): ❌ Run #30507587349

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Jul 30, 2026
@hnyls2002
hnyls2002 merged commit 313a518 into main Jul 30, 2026
160 of 196 checks passed
@hnyls2002
hnyls2002 deleted the lsyin/draft-extend-step-span branch July 30, 2026 02:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant