Profiling Enhancements [2/3]: detailed execution step annotations (#24911)

docs/docs/developer_guide/benchmark_and_profiling.mdx

@@ -198,6 +198,7 @@ curl -X POST http://127.0.0.1:30000/start_profile \
 - `start_step` (optional): Step number at which to start profiling (inclusive). Useful for skipping warmup iterations
 - `activities` (optional): List of activities to profile, e.g., `["CPU", "GPU"]`. Default is `["CPU", "GPU"]`
 - `merge_profiles` (optional): Whether to merge distributed traces. Default is `false`
+- `detailed_annotations` (optional): Whether to fold per-iteration request and KV-length aggregates into the trace's `step[...]` markers, for detailed analysis. Default is `false`. See [Detailed annotations](#detailed-annotations) below.
 
 **Note on step ranges:** Profiling starts at `start_step` (inclusive) and continues for `num_steps` iterations. For example, with `start_step=3` and `num_steps=10`, profiling captures steps 3, 4, 5, 6, 7, 8, 9, 10, 11, and 12 (10 steps total, starting from step 3).
 
@@ -222,6 +223,45 @@ curl -X POST http://127.0.0.1:30000/start_profile \
 curl -X POST http://127.0.0.1:30000/start_profile
 ```
 
+#### Detailed annotations
+
+Set `detailed_annotations` to `true` to fold per-iteration aggregates into SGLang's existing per-forward `step[...]` span. For every execution step that runs while profiling is active, SGLang augments that step's marker on the GPU stream with the request and KV-length distribution of the step, so you can reconstruct compute and memory bounds directly from the trace without per-request details.
+
+All four per-request aggregates are appended, prefixed by phase — `c_` for context (prefill) and `g_` for generation (decode). The per-phase `sq` is always emitted so each `step[...]` label is self-contained for roofline analysis, even where it duplicates the base label's `bs` (decode) or `toks` (prefill):
+
+- `sq`: total query tokens (`Σ N_Q`)
+- `sqsq`: sum of squared query tokens per request (`Σ N_Q²`)
+- `sqsk`: sum of query·KV tokens per request (`Σ N_Q·N_KV`)
+- `sk`: total KV tokens (`Σ N_KV`)
+
+A pure prefill (`EXTEND`) or decode (`DECODE`) forward emits a single group; a mixed forward emits both, with `c=`/`g=` request counts. Example labels:
+
+```text
+step[EXTEND bs=1 toks=1025 c_sq=1025 c_sqsq=1050625 c_sqsk=1050625 c_sk=1025]
+step[DECODE bs=64 g_sq=64 g_sqsq=64 g_sqsk=100032 g_sk=100032]
+step[MIXED bs=66 c=2 g=64 c_sq=2048 c_sk=2048 c_sqsq=2097152 c_sqsk=2097152 g_sq=64 g_sk=65600 g_sqsq=64 g_sqsk=65600]
+```
+
+With speculative decoding (EAGLE/MTP) each request contributes multiple query tokens per step, so `sq` no longer equals `bs`. Both draft-decode and target-verify (`TARGET_VERIFY`) steps are emitted in the **generation** group. For example, a verify step with 3 draft tokens across 2 requests (`seq_lens=[10, 20]`):
+
+```text
+step[TARGET_VERIFY bs=2 g_sq=6 g_sqsq=18 g_sqsk=90 g_sk=30]
+```
+
+```bash Command
+# Profile 10 steps with detailed annotations enabled
+curl -X POST http://127.0.0.1:30000/start_profile \
+  -H "Content-Type: application/json" \
+  -d '{
+    "output_dir": "/tmp/profiles",
+    "num_steps": 10,
+    "activities": ["CPU", "GPU"],
+    "detailed_annotations": true
+  }'
+```
+
+The annotations only appear when profiling is active with `detailed_annotations` enabled, so they add no overhead on the normal serving path. The behavior is identical in eager and CUDA graph modes. When viewing the trace (see [View traces](#view-traces)), the augmented `step[...]` markers appear on the GPU stream alongside the kernels for each step.
+
 #### Using `/stop_profile` endpoint
 
 The `/stop_profile` endpoint stops an ongoing profiling session and saves the trace file.

python/sglang/srt/managers/io_struct.py

@@ -2094,6 +2094,8 @@ class ProfileReq(BaseReq, kw_only=True):
     profile_prefix: Optional[str] = None
     # Only profile these stages and ignore others
     profile_stages: Optional[List[str]] = None
+    # Add iteration-level annotations (KV / request aggregates) for roofline-style analysis
+    detailed_annotations: bool = False
 
 
 class ProfileReqOutput(BaseReq, kw_only=True):

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -19,6 +19,7 @@ import torch
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.step_span_utils import set_detailed_annotations_enabled
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_device
 from sglang.srt.utils import is_mps, is_npu
@@ -78,6 +79,7 @@ class SchedulerProfilerManager:
         self.profile_by_stage: bool = False
         self.profile_in_progress: bool = False
         self.merge_profiles = False
+        self.detailed_annotations: bool = False
 
         # For ROCM
         self.rpd_profiler = None
@@ -94,9 +96,11 @@ class SchedulerProfilerManager:
         profile_id: str,
         merge_profiles: bool = False,
         profile_prefix: str = "",
+        detailed_annotations: bool = False,
         profile_stages: Optional[List[str]] = None,
     ) -> ProfileReqOutput:
         if envs.SGLANG_PROFILE_V2.get():
+            self.detailed_annotations = detailed_annotations
             return self._profile_manager.configure(
                 output_dir=output_dir,
                 start_step=start_step,
@@ -109,6 +113,7 @@ class SchedulerProfilerManager:
                 merge_profiles=merge_profiles,
                 profile_prefix=profile_prefix,
                 profile_stages=profile_stages,
+                detailed_annotations=detailed_annotations,
             )
 
         if self.profile_in_progress:
@@ -131,6 +136,7 @@ class SchedulerProfilerManager:
         self.profiler_activities = activities
         self.profile_id = profile_id
         self.profile_prefix = profile_prefix
+        self.detailed_annotations = detailed_annotations
 
         if start_step:
             self.profiler_start_forward_ct = max(start_step, self.get_forward_ct() + 1)
@@ -153,10 +159,17 @@ class SchedulerProfilerManager:
 
         return ProfileReqOutput(success=True, message="Succeeded")
 
+    def _apply_detailed_annotations(self, enabled: bool) -> None:
+        # Toggle the process-wide flag read by build_step_span_name; folds the
+        # per-phase sq/sqsq/sqsk/sk aggregates (context c_ / generation g_)
+        # into the step span while a detailed-annotation profile is active.
+        set_detailed_annotations_enabled(enabled)
+
     def _start_profile(
         self, stage: Optional[ForwardMode] = None
     ) -> ProfileReqOutput | None:
         if envs.SGLANG_PROFILE_V2.get():
+            self._apply_detailed_annotations(self.detailed_annotations)
             return self._profile_manager.manual_start()
 
         stage_str = f" for {stage.name}" if stage else ""
@@ -262,6 +275,7 @@ class SchedulerProfilerManager:
                 torch.cuda.cudart().cudaProfilerStart()
             self.profile_in_progress = True
 
+        self._apply_detailed_annotations(self.detailed_annotations)
         return ProfileReqOutput(success=True, message="Succeeded")
 
     def _merge_profile_traces(self) -> str:
@@ -300,6 +314,7 @@ class SchedulerProfilerManager:
         self, stage: Optional[ForwardMode] = None
     ) -> ProfileReqOutput | None:
         if envs.SGLANG_PROFILE_V2.get():
+            self._apply_detailed_annotations(False)
             return self._profile_manager.manual_stop()
 
         if not self.profile_in_progress:
@@ -387,6 +402,7 @@ class SchedulerProfilerManager:
         self.profile_in_progress = False
         self.profiler_start_forward_ct = None
 
+        self._apply_detailed_annotations(False)
         return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}")
 
     def _profile_batch_predicate(self, batch: ScheduleBatch):
@@ -443,6 +459,7 @@ class SchedulerProfilerManager:
                     recv_req.profile_id,
                     recv_req.merge_profiles,
                     recv_req.profile_prefix,
+                    recv_req.detailed_annotations,
                     recv_req.profile_stages,
                 )
             else:
@@ -457,6 +474,7 @@ class SchedulerProfilerManager:
                     recv_req.profile_id,
                     recv_req.merge_profiles,
                     recv_req.profile_prefix,
+                    recv_req.detailed_annotations,
                 )
                 return self._start_profile()
         else:

python/sglang/srt/model_executor/step_span_utils.py

@@ -0,0 +1,139 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Profile-trace step-span naming (kept dependency-light for CPU unit tests).
+
+The step span wraps each ``ModelRunner.forward`` in the torch/Perfetto trace.
+Its name carries the forward mode and batch shape; when detailed annotations
+are enabled it also folds in the per-iteration aggregates (for roofline-style
+analysis) so a single label describes both timing and the analytical work of
+that forward.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Tuple
+
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+# Process-wide toggle for detailed step-span annotations. Set by the scheduler's
+# profiler manager on profile start/stop and read by build_step_span_name, so no
+# per-runner flag or scheduler bridge is needed. Covers all model runners in the
+# process (e.g. both the EAGLE draft and target runners).
+_DETAILED_ANNOTATIONS_ENABLED = False
+
+
+def set_detailed_annotations_enabled(enabled: bool) -> None:
+    global _DETAILED_ANNOTATIONS_ENABLED
+    _DETAILED_ANNOTATIONS_ENABLED = bool(enabled)
+
+
+def detailed_annotations_enabled() -> bool:
+    return _DETAILED_ANNOTATIONS_ENABLED
+
+
+def _agg(nqs: List[int], nkvs: List[int]) -> Tuple[int, int, int, int]:
+    """Return (Σ N_Q, Σ N_KV, Σ N_Q², Σ N_Q·N_KV) for one request group."""
+    sq = sum(nqs)
+    sk = sum(nkvs)
+    sqsq = sum(nq * nq for nq in nqs)
+    sqsk = sum(nq * nkv for nq, nkv in zip(nqs, nkvs))
+    return sq, sk, sqsq, sqsk
+
+
+def _decode_query_width(forward_batch: ForwardBatch) -> int:
+    """Per-request query-token count (N_Q) for a decode-family forward.
+
+    Vanilla decode emits one token per request, but speculative decoding does
+    not: EAGLE/MTP draft-decode and target-verify process a uniform
+    ``num_tokens_per_req`` tokens per request (draft top-k for draft-decode,
+    the draft-token count for verify), so N_Q per request is that width, and
+    the step's total Σ N_Q is ``bs * num_tokens_per_req`` rather than ``bs``.
+
+    The width is read from ``forward_batch.spec_info.num_tokens_per_req`` and
+    falls back to 1 when there is no spec input or the width is unset (-1).
+    """
+    spec = getattr(forward_batch, "spec_info", None)
+    width = getattr(spec, "num_tokens_per_req", -1) if spec is not None else -1
+    return width if isinstance(width, int) and width > 0 else 1
+
+
+def build_detailed_annotation_suffix(forward_batch: ForwardBatch) -> str:
+    """Compute the detailed-annotation aggregates from the batch's CPU-side length mirrors.
+
+    All aggregates are emitted, prefixed by the roofline compute-shape bucket:
+    ``c_`` for context/extend-shaped work (EXTEND and DRAFT_EXTEND_V2) and ``g_`` for
+    single-query generation (DECODE, TARGET_VERIFY), with MIXED emitting both
+    groups.
+    """
+    mode = forward_batch.forward_mode
+    seq_lens_cpu = forward_batch.seq_lens_cpu
+
+    # DECODE (vanilla or spec draft-decode) and TARGET_VERIFY both key off
+    # ``seq_lens_cpu`` for N_KV and a uniform per-request query width N_Q, and
+    # are both classified as generation (``g_``) by request phase
+    #   * DECODE        -> N_Q is 1 (vanilla) or the spec draft-decode width.
+    #   * TARGET_VERIFY -> N_Q is ``num_tokens_per_req`` (the draft-token count);
+    #                      the request is past its prompt (generation phase)
+    if mode == ForwardMode.DECODE or mode == ForwardMode.TARGET_VERIFY:
+        if seq_lens_cpu is None:
+            return ""
+        nq = _decode_query_width(forward_batch)
+        nkvs = [int(x) for x in seq_lens_cpu.tolist()]
+        nqs = [nq] * len(nkvs)
+        sq, sk, sqsq, sqsk = _agg(nqs, nkvs)
+        # ``sq`` is always emitted (self-contained suffix): it equals ``bs``
+        # (vanilla decode) or ``bs * num_tokens_per_req`` (spec draft-decode /
+        # target-verify).
+        return f"g_sq={sq} g_sqsq={sqsq} g_sqsk={sqsk} g_sk={sk}"
+
+    ext_seq = forward_batch.extend_seq_lens_cpu
+    ext_prefix = forward_batch.extend_prefix_lens_cpu
+    if ext_seq is None or ext_prefix is None:
+        return ""
+
+    if mode == ForwardMode.EXTEND or mode == ForwardMode.DRAFT_EXTEND_V2:
+        # Both are extend-shaped, multi-query context
+        nqs = [int(q) for q in ext_seq]
+        nkvs = [int(p) + int(q) for p, q in zip(ext_prefix, ext_seq)]
+        sq, sk, sqsq, sqsk = _agg(nqs, nkvs)
+        return f"c_sq={sq} c_sqsq={sqsq} c_sqsk={sqsk} c_sk={sk}"
+
+    if mode == ForwardMode.MIXED:
+        # A running-decode request appears as a length-1 extend; everything
+        # else is a context (prefill) chunk.
+        c_nqs: List[int] = []
+        c_nkvs: List[int] = []
+        g_nqs: List[int] = []
+        g_nkvs: List[int] = []
+        for p, q in zip(ext_prefix, ext_seq):
+            nq, nkv = int(q), int(p) + int(q)
+            if nq == 1:
+                g_nqs.append(nq)
+                g_nkvs.append(nkv)
+            else:
+                c_nqs.append(nq)
+                c_nkvs.append(nkv)
+        c_sq, c_sk, c_sqsq, c_sqsk = _agg(c_nqs, c_nkvs)
+        g_sq, g_sk, g_sqsq, g_sqsk = _agg(g_nqs, g_nkvs)
+        return (
+            f"c={len(c_nqs)} g={len(g_nqs)} "
+            f"c_sq={c_sq} c_sk={c_sk} c_sqsq={c_sqsq} c_sqsk={c_sqsk} "
+            f"g_sq={g_sq} g_sk={g_sk} g_sqsq={g_sqsq} g_sqsk={g_sqsk}"
+        )
+
+    return ""

python/sglang/srt/utils/profile_utils.py

@@ -12,6 +12,11 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import ProfileReqOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.step_span_utils import (
+    build_detailed_annotation_suffix,
+    detailed_annotations_enabled,
+    set_detailed_annotations_enabled,
+)
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_device
 from sglang.srt.utils import is_npu
@@ -76,6 +81,7 @@ class ProfileManager:
         self.first_rank_in_node = ps.gpu_id == get_device().base_gpu_id
         self.profiler_kwargs = None
         self.profiler = None
+        self.detailed_annotations = False
 
     def step(self, forward_mode: ForwardMode):
         stage = _get_stage_from_forward_mode(forward_mode)
@@ -98,7 +104,9 @@ class ProfileManager:
         merge_profiles: bool,
         profile_prefix: str,
         profile_stages: Optional[List[str]] = None,
+        detailed_annotations: bool = False,
     ):
+        self.detailed_annotations = detailed_annotations
         # not supported yet
         assert start_step is None
         assert (
@@ -141,6 +149,9 @@ class ProfileManager:
         )
 
         assert self.profiler is None
+        # Fold the per-phase c_/g_ aggregates into the step span while this
+        # stage's profile is active (v2 auto-start path; reset in _do_stop).
+        set_detailed_annotations_enabled(self.detailed_annotations)
         self.profiler = _ProfilerBase.create(
             **self.profiler_kwargs,
             ps=self.ps,
@@ -157,6 +168,10 @@ class ProfileManager:
             f"Profiling done. Traces are saved to: {self.profiler_kwargs['output_dir']}"
         )
         self.profiler = None
+        # Clear the detailed step-span toggle here too: the v2 trigger auto-stop
+        # goes through _do_stop (not SchedulerProfilerManager._stop_profile), so
+        # this guarantees the flag resets on every stop path.
+        set_detailed_annotations_enabled(False)
 
 
 def _get_stage_from_forward_mode(forward_mode: ForwardMode):
@@ -445,11 +460,29 @@ class _ProfilerRPD(_ProfilerConcreteBase):
             rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
 
 
-def build_step_span_name(forward_batch: ForwardBatch) -> str:
+def build_step_span_name(
-    """Build a profile-trace span name for one forward step."""
+    forward_batch: ForwardBatch, detailed_annotations: bool | None = None
+) -> str:
+    """Build the profile-trace span name for one forward step.
+
+    Detailed annotations are folded into the label (via
+    build_detailed_annotation_suffix) when enabled. detailed_annotations
+    defaults to the process-wide toggle (detailed_annotations_enabled, set
+    by the profiler manager); pass an explicit bool to override (e.g. in tests).
+    """
+    if detailed_annotations is None:
+        detailed_annotations = detailed_annotations_enabled()
+
     mode = forward_batch.forward_mode
     bs = forward_batch.batch_size
     if mode == ForwardMode.EXTEND:
         ext_toks = forward_batch.extend_num_tokens or 0
-        return f"step[EXTEND bs={bs} toks={ext_toks}]"
+        base = f"step[EXTEND bs={bs} toks={ext_toks}"
-    return f"step[{mode.name} bs={bs}]"
+    else:
+        base = f"step[{mode.name} bs={bs}"
+
+    if detailed_annotations:
+        suffix = build_detailed_annotation_suffix(forward_batch)
+        if suffix:
+            base = f"{base} {suffix}"
+    return f"{base}]"

test/registered/unit/managers/test_detailed_annotations.py

@@ -0,0 +1,256 @@
+"""Unit tests for detailed profiling annotations (#24911) — no server, no model loading.
+
+The detailed-annotation aggregates are folded into SGLang's existing per-forward ``step[...]``
+span (see ``sglang.srt.utils.profile_utils.build_step_span_name``): the
+per-phase ``sq``/``sqsq``/``sqsk``/``sk`` terms (with the context/generation split
+for MIXED) are appended and are self-contained, so ``sq`` is emitted even where it
+duplicates the base label's ``bs``/``toks``. This also covers the
+``detailed_annotations`` plumbing on ``ProfileReq``.
+"""
+
+import json
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.managers.io_struct import ProfileReq
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.step_span_utils import (
+    detailed_annotations_enabled,
+    set_detailed_annotations_enabled,
+)
+from sglang.srt.utils.profile_utils import build_step_span_name
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+
+class _CpuMirror:
+    """Minimal stand-in for the ``seq_lens_cpu`` tensor (only ``tolist`` used)."""
+
+    def __init__(self, data):
+        self._data = list(data)
+
+    def tolist(self):
+        return list(self._data)
+
+
+def _fb(
+    forward_mode,
+    *,
+    batch_size,
+    extend_num_tokens=None,
+    seq_lens_cpu=None,
+    extend_seq_lens_cpu=None,
+    extend_prefix_lens_cpu=None,
+    num_tokens_per_req=None,
+):
+    # A spec input (EAGLE/MTP) only needs to expose ``num_tokens_per_req`` for
+    # the detailed-annotation suffix; None -> no spec_info (vanilla decode, N_Q == 1).
+    spec_info = (
+        None
+        if num_tokens_per_req is None
+        else SimpleNamespace(num_tokens_per_req=num_tokens_per_req)
+    )
+    return SimpleNamespace(
+        forward_mode=forward_mode,
+        batch_size=batch_size,
+        extend_num_tokens=extend_num_tokens,
+        seq_lens_cpu=None if seq_lens_cpu is None else _CpuMirror(seq_lens_cpu),
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        extend_prefix_lens_cpu=extend_prefix_lens_cpu,
+        spec_info=spec_info,
+    )
+
+
+class TestStepSpanDetailedAnnotations(CustomTestCase):
+    def _name(self, fb):
+        return build_step_span_name(fb, detailed_annotations=True)
+
+    def test_pure_decode_batch(self):
+        # Two decode reqs: each nq=1, nkv=seqlen.
+        # sk=30, sqsq=1+1=2, sqsk=1*10+1*20=30.
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=[10, 20])
+        self.assertEqual(
+            self._name(fb), "step[DECODE bs=2 g_sq=2 g_sqsq=2 g_sqsk=30 g_sk=30]"
+        )
+
+    def test_pure_prefill_batch(self):
+        # req a: nq=8, nkv=10 -> sqsq=64, sqsk=80
+        # req b: nq=4, nkv=10 -> sqsq=16, sqsk=40
+        # sk=20, sqsq=80, sqsk=120; toks left to the base label.
+        fb = _fb(
+            ForwardMode.EXTEND,
+            batch_size=2,
+            extend_num_tokens=12,
+            extend_seq_lens_cpu=[8, 4],
+            extend_prefix_lens_cpu=[2, 6],
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[EXTEND bs=2 toks=12 c_sq=12 c_sqsq=80 c_sqsk=120 c_sk=20]",
+        )
+
+    def test_mixed_batch_splits_context_and_generation(self):
+        # ctx: nq=5, nkv=8 -> sqsq=25, sqsk=40; gen (len-1 extend): nq=1, nkv=12.
+        fb = _fb(
+            ForwardMode.MIXED,
+            batch_size=2,
+            extend_seq_lens_cpu=[5, 1],
+            extend_prefix_lens_cpu=[3, 11],
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[MIXED bs=2 c=1 g=1 "
+            "c_sq=5 c_sk=8 c_sqsq=25 c_sqsk=40 "
+            "g_sq=1 g_sk=12 g_sqsq=1 g_sqsk=12]",
+        )
+
+    def test_mixed_batch_all_context(self):
+        fb = _fb(
+            ForwardMode.MIXED,
+            batch_size=1,
+            extend_seq_lens_cpu=[3],
+            extend_prefix_lens_cpu=[0],
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[MIXED bs=1 c=1 g=0 "
+            "c_sq=3 c_sk=3 c_sqsq=9 c_sqsk=9 "
+            "g_sq=0 g_sk=0 g_sqsq=0 g_sqsk=0]",
+        )
+
+    def test_spec_draft_decode_uses_num_tokens_per_req(self):
+        # EAGLE draft-decode: N_Q per req = topk (num_tokens_per_req), not 1.
+        # topk=4, seqs=[10,20]: sq=4*2=8, sk=30, sqsq=16+16=32,
+        # sqsk=4*10+4*20=120. g_sq is emitted because it != bs.
+        fb = _fb(
+            ForwardMode.DECODE,
+            batch_size=2,
+            seq_lens_cpu=[10, 20],
+            num_tokens_per_req=4,
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[DECODE bs=2 g_sq=8 g_sqsq=32 g_sqsk=120 g_sk=30]",
+        )
+
+    def test_target_verify_uses_draft_token_width(self):
+        # MTP/EAGLE target-verify: N_Q per req = num_draft_tokens. It is
+        # classified as generation (``g_``) by request phase;
+        # its quadratic self-attention is still captured in g_sqsq.
+        # ndt=3, seqs=[10,20]: sq=3*2=6, sk=30, sqsq=9+9=18,
+        # sqsk=3*10+3*20=90.
+        fb = _fb(
+            ForwardMode.TARGET_VERIFY,
+            batch_size=2,
+            seq_lens_cpu=[10, 20],
+            num_tokens_per_req=3,
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[TARGET_VERIFY bs=2 g_sq=6 g_sqsq=18 g_sqsk=90 g_sk=30]",
+        )
+
+    def test_target_verify_without_cpu_mirror_falls_back_to_base(self):
+        fb = _fb(
+            ForwardMode.TARGET_VERIFY,
+            batch_size=2,
+            seq_lens_cpu=None,
+            num_tokens_per_req=3,
+        )
+        self.assertEqual(self._name(fb), "step[TARGET_VERIFY bs=2]")
+
+    def test_draft_extend_v2_uses_extend_mirrors_with_context_prefix(self):
+        # EAGLE/MTP draft-extend is extend-shaped
+        # req a: nq=2, nkv=12 -> sqsq=4, sqsk=24
+        # req b: nq=3, nkv=23 -> sqsq=9, sqsk=69
+        # sq=5, sk=35, sqsq=13, sqsk=93.
+        fb = _fb(
+            ForwardMode.DRAFT_EXTEND_V2,
+            batch_size=2,
+            extend_seq_lens_cpu=[2, 3],
+            extend_prefix_lens_cpu=[10, 20],
+        )
+        self.assertEqual(
+            self._name(fb),
+            "step[DRAFT_EXTEND_V2 bs=2 c_sq=5 c_sqsq=13 c_sqsk=93 c_sk=35]",
+        )
+
+    def test_draft_extend_v2_without_extend_mirrors_falls_back_to_base(self):
+        fb = _fb(ForwardMode.DRAFT_EXTEND_V2, batch_size=2)
+        self.assertEqual(self._name(fb), "step[DRAFT_EXTEND_V2 bs=2]")
+
+    def test_missing_cpu_mirror_falls_back_to_base(self):
+        # No seq_lens_cpu (some overlap paths) -> emit the base label unchanged.
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=None)
+        self.assertEqual(self._name(fb), "step[DECODE bs=2]")
+
+
+class TestStepSpanGating(CustomTestCase):
+    def test_disabled_flag_emits_base_label(self):
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=[10, 20])
+        self.assertEqual(
+            build_step_span_name(fb, detailed_annotations=False), "step[DECODE bs=2]"
+        )
+
+    def test_disabled_flag_is_default(self):
+        fb = _fb(
+            ForwardMode.EXTEND,
+            batch_size=1,
+            extend_num_tokens=4,
+            extend_seq_lens_cpu=[4],
+            extend_prefix_lens_cpu=[0],
+        )
+        self.assertEqual(build_step_span_name(fb), "step[EXTEND bs=1 toks=4]")
+
+
+class TestDetailedAnnotationPlumbing(CustomTestCase):
+    def test_default_is_false(self):
+        self.assertFalse(ProfileReq().detailed_annotations)
+
+    def test_json_round_trip(self):
+        req = ProfileReq(output_dir="/tmp/x", detailed_annotations=True)
+        payload = {"detailed_annotations": req.detailed_annotations}
+        parsed = json.loads(json.dumps(payload))
+        self.assertTrue(parsed["detailed_annotations"])
+        self.assertTrue(ProfileReq(**parsed).detailed_annotations)
+
+
+class TestDetailedAnnotationsToggle(CustomTestCase):
+    """The process-wide toggle (set by the profiler manager) is the default source
+    for build_step_span_name when no explicit flag is passed."""
+
+    def tearDown(self):
+        set_detailed_annotations_enabled(False)
+
+    def test_default_off_no_suffix(self):
+        set_detailed_annotations_enabled(False)
+        self.assertFalse(detailed_annotations_enabled())
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=[10, 20])
+        self.assertEqual(build_step_span_name(fb), "step[DECODE bs=2]")
+
+    def test_toggle_on_folds_suffix(self):
+        set_detailed_annotations_enabled(True)
+        self.assertTrue(detailed_annotations_enabled())
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=[10, 20])
+        # sq=2, sk=30, sqsq=2, sqsk=30
+        self.assertEqual(
+            build_step_span_name(fb),
+            "step[DECODE bs=2 g_sq=2 g_sqsq=2 g_sqsk=30 g_sk=30]",
+        )
+
+    def test_explicit_arg_overrides_toggle(self):
+        set_detailed_annotations_enabled(True)
+        fb = _fb(ForwardMode.DECODE, batch_size=2, seq_lens_cpu=[10, 20])
+        # explicit False wins over the enabled toggle
+        self.assertEqual(
+            build_step_span_name(fb, detailed_annotations=False), "step[DECODE bs=2]"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()