fix(observability): restore speculative-decoding spans on the Spec V2 paths

[MicheleCampi]

Fixes #37128.

What broke

The EAGLE V2 workers emit no per-request speculative spans: neither eagle_worker_v2.py nor multi_layer_eagle_worker_v2.py calls the spec_draft or spec_verify setters. The spec_stage_span scopes in eagle_worker_v2.py resolve to profile_range, which records a torch profiler or NVTX range and is a no-op otherwise, so they never reach the request trace. NGRAM opens and closes spec_draft but never spec_verify.

Since this PR was opened, #40795 removed SchedulerReqTimeStats.set_spec_verify_end_time, which had no callers left in the tree. That left RequestStage.SPEC_VERIFY with no code path that emits it.

What this does

• Restores set_spec_verify_end_time, without the accepted_tokens alias that [misc] Remove deprecated endpoints, env vars and aliases past two releases #40795 retired on purpose. The verify span carries num_correct_drafts only.
• Adds set_spec_verify_end_time_batch next to set_time_batch. It closes the span for a batch, derives the drafts-only count as accept_lens - 1, and reads accept_lens to the host only when tracing is enabled.
• Opens and closes each span around the operation it measures, as NGRAM already does for spec_draft:
  • spec_draft: around draft() in both EAGLE workers. In EAGLEWorkerV2 it sits inside the branch that drafts: the PP pre-built-verify and zero-step branches do not draft, so they do not close a draft span they never opened.
  • spec_verify: inside EAGLEWorkerV2.verify() and MultiLayerEagleWorkerV2.verify(), around run_eagle_verify, and around the target-verify step in NGRAM.

run_eagle_verify itself is unchanged. uno_worker_v2 calls it directly and packs the result into its own tree output, so a close inside the shared function would have attached a verify span to UNO as a side effect; keeping it in the verify() methods leaves UNO untouched.

Paths that go through EAGLEWorkerV2.verify() gain the verify span too: the DP speculative prefill coordination path, and FrozenKVMTPWorkerV2, which inherits verify().

Left alone deliberately

• spec_draft_extend: there is no SPEC_DRAFT_EXTEND member in RequestStage on current main, so restoring it would add a stage rather than fix a regression.
• The accepted_tokens alias. [Bug] Spec V2 paths no longer emit speculative-decoding OpenTelemetry spans #37128 lists it among what NGRAM cannot emit; [misc] Remove deprecated endpoints, env vars and aliases past two releases #40795 has since retired it, and this PR follows that decision. Happy to add it back if you prefer.
• The PP tail draft (draft(..., with_topology=True)) is not instrumented.

Testing

test/registered/unit/observability/test_spec_decode_tracing.py, CPU-only, 5 tests. They run against the real SchedulerReqTimeStats with only trace_ctx replaced — the pattern test_req_time_stats.py uses — and assert on the emitted slices:

• the verify slice carries exactly {"num_correct_drafts": ...};
• with tracing off, accept_lens is never read to the host and no slice is emitted;
• each request in a batch gets its own count, from one host read and one timestamp;
• the setters the workers pass to set_time_batch by name exist and act.

Against req_time_stats.py from current main, 3 of the 5 fail with AttributeError (the restored method and the batch helper); the two setter-name tests pass there, since those setters exist.

Run locally on CPU: Python 3.10, the CPU CI job's uv pip install -e "python[dev]" line with torch from the CPU index (2.13.0+cpu), and flashinfer masked as unavailable because importing it queries the GPU. test_req_time_stats.py passes unchanged (13/13), and pre-commit passes on the changed files.

Runtime verification against a live server with tracing enabled is not included.

---

CI States

Latest PR Test (Base): ❌ Run #36349940952
Latest PR Test (Extra): ❌ Run #36349940591
Latest PR Test (AMD ROCm 10): ❌ Run #36349940737

[MicheleCampi]

@hnyls2002 @Qiaolin-Yu — this needs the run-ci label to get past pr-gate; no test job has run on it since it was opened.

I rebased it onto 5c2de3f35 just now (was CONFLICTING, now MERGEABLE). One conflict, in multi_layer_eagle_worker_v2.py: main had added staged = self.draft_worker._draft_extend_plan_for_decode(batch) on the same line where this PR inserts set_time_batch(..., "set_spec_verify_start_time", ...). Both are needed — staged is consumed four lines down — so I kept both, with staged first so the verify span does not cover draft-extend planning.

I could not run the test locally (no torch on that machine), so the four auto-merged files are unverified beyond the merge itself. CI is the check that matters here.

@RichardoMrMu — this restores the stages you added in #19545, on the paths you reported in #37128.

[RichardoMrMu]

Thanks for picking this up — I filed #37128 and went through this PR against all three gaps in that report; it looks complete and correct to me:

1. EAGLE v2 — both eagle_worker_v2 and multi_layer_eagle_worker_v2 now emit spec_draft (start/end) and spec_verify (start) via set_time_batch(..., trace_only=True), and the shared run_eagle_verify closes spec_verify with num_correct_drafts. That restores the per-request spans EAGLE lost.
2. NGRAM — the target-verify path now calls set_spec_verify_start_time / set_spec_verify_end_time(num_correct_drafts=...). The accepted_tokens compat attribute is emitted automatically by set_spec_verify_end_time in req_time_stats, so it is covered without extra work here.
3. spec_draft_extend — correctly not revived; as noted in the issue it was intentionally removed as dead code ([Spec] Clear dead DRAFT_EXTEND objects left after EAGLE v1 removal #28133), so leaving it out is the right call.

One detail worth calling out: the accept_lens.tolist() host reads are correctly gated behind get_global_tracing_enabled(), so tracing-disabled runs do not pay a device to host sync on the hot path — and the test pins exactly that property. This matches the "tracing-disabled execution should not add per-request tracing-only work" expectation from the issue.

LGTM from a reporter's standpoint. Would be great to get a maintainer look (cc @sufeng-buaa) so this can land.

[MicheleCampi]

@hnyls2002 — rebased onto current main, which took more than a conflict fix: #40795 removed SchedulerReqTimeStats.set_spec_verify_end_time while this PR was open, since nothing in the tree called it yet. This version restores it without the accepted_tokens alias you retired there, so the verify span carries num_correct_drafts only.

Two things that are your call:

1. Not restoring accepted_tokens. [Bug] Spec V2 paths no longer emit speculative-decoding OpenTelemetry spans #37128 lists it among what NGRAM cannot emit; I followed [misc] Remove deprecated endpoints, env vars and aliases past two releases #40795 instead. Easy to add back if you'd rather keep it.
2. The EAGLE V2 mapping [Bug] Spec V2 paths no longer emit speculative-decoding OpenTelemetry spans #37128 asked maintainers to confirm: spec_draft is draft(), spec_verify is the worker's verify(), no spec_draft_extend. Each span now opens and closes around the operation it measures, as NGRAM already does for spec_draft, so run_eagle_verify is untouched and uno_worker_v2, which calls it directly, is unaffected.

Could you add run-ci? pr-gate fails and every test stage is skipped, base-a-test-cpu included, so nothing has been tested in CI on this version. Locally the new tests and test_req_time_stats.py pass on CPU and pre-commit is clean; the updated description has the details. The earlier review in this thread was against the previous version.

[MicheleCampi]

@hnyls2002 — this still merges cleanly onto today's main (21ed9fbf4), and the new tests plus test_req_time_stats.py pass on the merged tree. Could you comment /tag-and-rerun-ci so CI can run on it?

[MicheleCampi]

@Qiaolin-Yu — you're the other speculative-decoding Merge Oncall, and #41175 touched eagle_worker_v2.py, which this PR changes, and eagle_worker_common.py, where run_eagle_verify lives. This still merges cleanly onto today's main (1e490772e), which includes #41175, and the new tests plus test_req_time_stats.py pass on the merged tree. Could you comment /tag-and-rerun-ci so CI can run on it?