From fc0b95e7badd94198925722445b93f063905a71c Mon Sep 17 00:00:00 2001 From: mohbasit Date: Tue, 18 Aug 2026 04:09:07 -0400 Subject: [PATCH] Profiling Enhancements [2/3]: detailed execution step annotations (#24911) --- .../benchmark_and_profiling.mdx | 40 +++ python/sglang/srt/managers/io_struct.py | 2 + .../scheduler_components/profiler_manager.py | 18 ++ .../srt/model_executor/step_span_utils.py | 139 ++++++++++ python/sglang/srt/utils/profile_utils.py | 41 ++- .../managers/test_detailed_annotations.py | 256 ++++++++++++++++++ 6 files changed, 492 insertions(+), 4 deletions(-) create mode 100644 python/sglang/srt/model_executor/step_span_utils.py create mode 100644 test/registered/unit/managers/test_detailed_annotations.py diff --git a/docs/docs/developer_guide/benchmark_and_profiling.mdx b/docs/docs/developer_guide/benchmark_and_profiling.mdx index bd4508b6b..fa9963d22 100644 --- a/docs/docs/developer_guide/benchmark_and_profiling.mdx +++ b/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. diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 736642c7f..8f84a18fe 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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): diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 67a74bdbf..63aa5aac3 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/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: diff --git a/python/sglang/srt/model_executor/step_span_utils.py b/python/sglang/srt/model_executor/step_span_utils.py new file mode 100644 index 000000000..0e57fe20e --- /dev/null +++ b/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 "" diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 705f1010b..4d1a80c72 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/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: - """Build a profile-trace span name for one forward step.""" +def build_step_span_name( + 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}]" - return f"step[{mode.name} bs={bs}]" + base = f"step[EXTEND bs={bs} toks={ext_toks}" + 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}]" diff --git a/test/registered/unit/managers/test_detailed_annotations.py b/test/registered/unit/managers/test_detailed_annotations.py new file mode 100644 index 000000000..c5ad6a98c --- /dev/null +++ b/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()