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

2026-08-18 01:09:07 -07:00
parent 667389c50f
commit fc0b95e7ba
6 changed files with 492 additions and 4 deletions
@@ -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.