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

python/sglang/srt/observability/req_time_stats.py

@@ -690,6 +690,18 @@ def set_spec_verify_start_time(self, ts=None):
         ts = ts or time.perf_counter()
         self.spec_verify_start_time = ts
 
+    def set_spec_verify_end_time(self, ts=None, num_correct_drafts: int = 0):
+        ts = ts or time.perf_counter()
+
+        if self.trace_ctx.tracing_enable:
+            stage = RequestStage.SPEC_VERIFY
+            self.trace_slice(
+                stage,
+                self.spec_verify_start_time,
+                ts,
+                {"num_correct_drafts": num_correct_drafts},
+            )
+
     def set_run_batch_cpu_start_time(self, ts=None, attrs=None):
         ts = ts or time.perf_counter()
         self.run_batch_cpu_start_time = ts
@@ -1260,6 +1272,27 @@ def set_time_batch(
             method(ts, attrs)
 
 
+def set_spec_verify_end_time_batch(reqs: List[Any], accept_lens: Any):
+    """Close the spec_verify span for every request in a verify batch.
+
+    accept_lens counts accepted tokens per request including the bonus token;
+    the drafts-only count excludes it. accept_lens may live on the device, so
+    converting it to a list forces a host synchronisation: it is read only when
+    tracing is enabled.
+    """
+    if reqs is None or len(reqs) == 0:
+        return
+    if not get_global_tracing_enabled():
+        return
+
+    num_correct_drafts = (accept_lens - 1).tolist()
+    ts = time.perf_counter()
+    for idx, req in enumerate(reqs):
+        req.time_stats.set_spec_verify_end_time(
+            ts, num_correct_drafts=num_correct_drafts[idx]
+        )
+
+
 def flush_trace_batch(reqs: List[Any]):
     """Proactively flush buffered trace ops for a batch of requests.
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -57,6 +57,10 @@
     DecodeCudaGraphRunner,
     get_batch_sizes_to_capture,
 )
+from sglang.srt.observability.req_time_stats import (
+    set_spec_verify_end_time_batch,
+    set_time_batch,
+)
 from sglang.srt.runtime_context import (
     get_context,
     get_device,
@@ -1457,6 +1461,7 @@ def forward_batch_generation(
                 # runs, keeping draft KV warm for when the batch shrinks.
                 verify_input = self._build_trivial_verify_input(batch)
             else:
+                set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
                 with (
                     self.draft_worker.draft_tp_context(
                         self.draft_worker.draft_runner.tp_group,
@@ -1467,6 +1472,7 @@ def forward_batch_generation(
                     spec_stage_span("draft"),
                 ):
                     verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
+                set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
             assert verify_input.is_verify_input()
             batch.spec_info = verify_input
             batch_output = self.verify(
@@ -1965,7 +1971,8 @@ def _override_worker_state(
             dw._rebuild_topk1_chain_buffers()
 
     def verify(self, batch: ScheduleBatch, pp_proxy_tensors=None, grammar_barrier=None):
-        return run_eagle_verify(
+        set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
+        result = run_eagle_verify(
             batch,
             pp_proxy_tensors=pp_proxy_tensors,
             target_worker=self.target_worker,
@@ -1980,3 +1987,5 @@ def verify(self, batch: ScheduleBatch, pp_proxy_tensors=None, grammar_barrier=No
             finalize_tree_path=True,
             grammar_barrier=grammar_barrier,
         )
+        set_spec_verify_end_time_batch(batch.reqs, result.accept_lens)
+        return result

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -43,6 +43,10 @@
     CaptureHiddenMode,
     ForwardBatch,
 )
+from sglang.srt.observability.req_time_stats import (
+    set_spec_verify_end_time_batch,
+    set_time_batch,
+)
 from sglang.srt.runtime_context import (
     get_device,
     get_parallel,
@@ -1117,7 +1121,9 @@ def forward_batch_generation(
                     topk=self.topk * self.speculative_num_steps,
                     capture_hidden_mode=capture_mode,
                 )
+            set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
             verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
+            set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
             assert verify_input.is_verify_input()
             batch.spec_info = verify_input
             staged = self.draft_worker._draft_extend_plan_for_decode(batch)
@@ -1131,7 +1137,8 @@ def forward_batch_generation(
             return batch_output
 
     def verify(self, batch: ScheduleBatch, grammar_barrier=None):
-        return run_eagle_verify(
+        set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
+        result = run_eagle_verify(
             batch,
             target_worker=self.target_worker,
             req_to_token_pool=self.req_to_token_pool,
@@ -1145,3 +1152,5 @@ def verify(self, batch: ScheduleBatch, grammar_barrier=None):
             finalize_tree_path=False,
             grammar_barrier=grammar_barrier,
         )
+        set_spec_verify_end_time_batch(batch.reqs, result.accept_lens)
+        return result

python/sglang/srt/speculative/ngram_worker.py

@@ -12,7 +12,10 @@
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.observability.req_time_stats import set_time_batch
+from sglang.srt.observability.req_time_stats import (
+    set_spec_verify_end_time_batch,
+    set_time_batch,
+)
 from sglang.srt.runtime_context import (
     get_device,
     get_schedule,
@@ -438,6 +441,7 @@ def forward_batch_generation(
         accept_lens = torch.ones(bs, dtype=torch.int32, device=self.device)
 
         if batch.forward_mode.is_target_verify():
+            set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
             batch_result = self.target_worker.forward_batch_generation(
                 batch, pp_proxy_tensors=pp_proxy_tensors, is_verify=True
             )
@@ -495,6 +499,7 @@ def forward_batch_generation(
             # The KV mover expects drafts-only counts. NGRAM's
             # accept_lens includes the bonus token, matching scheduler output.
             num_correct_drafts_per_req = accept_lens - 1
+            set_spec_verify_end_time_batch(batch.reqs, accept_lens)
             move_accept_tokens_to_target_kvcache(
                 batch,
                 accept_index,

test/registered/unit/observability/test_spec_decode_tracing.py

@@ -0,0 +1,145 @@
+"""Unit tests for the speculative-decoding verify span.
+
+set_spec_verify_end_time_batch closes the spec_verify span for a verify batch.
+These tests run it against the real SchedulerReqTimeStats and observe what
+reaches the tracing context: one SPEC_VERIFY slice per request carrying the
+drafts-only count (accept_lens minus the bonus token), and no host read of
+accept_lens when tracing is off. They also pin the setters that the speculative
+workers pass to set_time_batch by name, since a renamed or removed setter would
+otherwise fail only at runtime, and only with tracing enabled.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import sglang.srt.observability.req_time_stats as rts
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+
+class _FakeDeviceTensor:
+    """Stands in for accept_lens: records whether it was read to the host."""
+
+    def __init__(self, values):
+        self._values = list(values)
+        self._counter = [0]
+
+    def __sub__(self, other):
+        # Share the counter with the derived tensor: the host read happens on
+        # `accept_lens - 1`, not on accept_lens itself, so counting only the
+        # original would let a host read go unnoticed.
+        derived = _FakeDeviceTensor([v - other for v in self._values])
+        derived._counter = self._counter
+        return derived
+
+    def tolist(self):
+        self._counter[0] += 1
+        return list(self._values)
+
+    @property
+    def host_reads(self):
+        return self._counter[0]
+
+
+def _traced_stats():
+    stats = rts.SchedulerReqTimeStats()
+    stats.trace_ctx = mock.MagicMock()
+    stats.trace_ctx.tracing_enable = True
+    return stats
+
+
+def _emitted_slices(stats):
+    return [call.args[0] for call in stats.trace_ctx.trace_slice.call_args_list]
+
+
+class TestSpecVerifyEndTime(CustomTestCase):
+    def test_verify_slice_carries_drafts_only_count(self):
+        stats = _traced_stats()
+        stats.set_spec_verify_start_time(10.0)
+        stats.set_spec_verify_end_time(12.5, num_correct_drafts=3)
+
+        (verify_slice,) = _emitted_slices(stats)
+        self.assertEqual(
+            verify_slice.slice_name, rts.RequestStage.SPEC_VERIFY.stage_name
+        )
+        self.assertEqual(
+            verify_slice.start_time_ns, rts.convert_time_to_realtime_ns(10.0)
+        )
+        self.assertEqual(
+            verify_slice.end_time_ns, rts.convert_time_to_realtime_ns(12.5)
+        )
+        self.assertEqual(verify_slice.attrs, {"num_correct_drafts": 3})
+
+
+class TestSetSpecVerifyEndTimeBatch(CustomTestCase):
+    def _reqs(self, n, start):
+        reqs = []
+        for _ in range(n):
+            stats = _traced_stats()
+            stats.set_spec_verify_start_time(start)
+            reqs.append(SimpleNamespace(time_stats=stats))
+        return reqs
+
+    def test_tracing_off_leaves_accept_lens_on_device(self):
+        reqs = self._reqs(3, start=1.0)
+        accept_lens = _FakeDeviceTensor([3, 2, 4])
+
+        with mock.patch.object(rts, "get_global_tracing_enabled", return_value=False):
+            rts.set_spec_verify_end_time_batch(reqs, accept_lens)
+
+        self.assertEqual(accept_lens.host_reads, 0)
+        for req in reqs:
+            self.assertEqual(_emitted_slices(req.time_stats), [])
+
+    def test_each_request_gets_its_own_drafts_only_count(self):
+        reqs = self._reqs(3, start=1.0)
+        accept_lens = _FakeDeviceTensor([3, 2, 4])
+
+        with mock.patch.object(rts, "get_global_tracing_enabled", return_value=True):
+            rts.set_spec_verify_end_time_batch(reqs, accept_lens)
+
+        slices = [_emitted_slices(req.time_stats) for req in reqs]
+        self.assertEqual([len(s) for s in slices], [1, 1, 1])
+        self.assertEqual([s[0].attrs["num_correct_drafts"] for s in slices], [2, 1, 3])
+        # One timestamp for the batch, one host read for the batch.
+        self.assertEqual(len({s[0].end_time_ns for s in slices}), 1)
+        self.assertEqual(accept_lens.host_reads, 1)
+
+
+class TestSpecSettersDispatchedByName(CustomTestCase):
+    """The workers name these setters as strings; set_time_batch resolves them."""
+
+    def test_start_setters_stamp_their_field(self):
+        for setter, field in (
+            ("set_spec_draft_start_time", "spec_draft_start_time"),
+            ("set_spec_verify_start_time", "spec_verify_start_time"),
+        ):
+            with self.subTest(setter=setter):
+                stats = _traced_stats()
+                with mock.patch.object(
+                    rts, "get_global_tracing_enabled", return_value=True
+                ):
+                    rts.set_time_batch(
+                        [SimpleNamespace(time_stats=stats)], setter, trace_only=True
+                    )
+                self.assertGreater(getattr(stats, field), 0.0)
+
+    def test_draft_end_setter_emits_spec_draft_slice(self):
+        stats = _traced_stats()
+        stats.set_spec_draft_start_time(1.0)
+        with mock.patch.object(rts, "get_global_tracing_enabled", return_value=True):
+            rts.set_time_batch(
+                [SimpleNamespace(time_stats=stats)],
+                "set_spec_draft_end_time",
+                trace_only=True,
+            )
+
+        (draft_slice,) = _emitted_slices(stats)
+        self.assertEqual(draft_slice.slice_name, rts.RequestStage.SPEC_DRAFT.stage_name)
+
+
+if __name__ == "__main__":
+    unittest.main()