Profiling Enhancements [1/3]: cuda graph profile traces (#24370)

Co-authored-by: Basit <mohbasit@ctr2-alola-ctrl-01.amd.com>
Co-authored-by: HAI <hixiao@gmail.com>

docs/docs/developer_guide/benchmark_and_profiling.mdx

@@ -296,6 +296,37 @@ The profile merger generates:
 - Individual rank trace files: `{profile_id}-TP-{tp}-DP-{dp}-PP-{pp}-EP-{ep}.trace.json.gz`
 - Merged trace file: `merged-{profile_id}.trace.json.gz`
 
+### Profile the CUDA graph capture phase
+
+The tools above profile the steady-state runtime (prefill / decode). To instead profile the **CUDA graph capture phase** that runs once at server startup, launch the server with `--enable-profile-cuda-graph`. This runs a PyTorch Profiler pass over the decode CUDA-graph capture, which is useful for diagnosing slow or memory-heavy graph capture.
+
+`--enable-profile-cuda-graph` (server arg) builds the capture profiler and always emits the per-kernel CPU/CUDA time summary tables and a CUDA memory snapshot. Persisting Chrome traces to disk is opt-in via one of two env vars (both no-ops unless `--enable-profile-cuda-graph` is also set):
+
+- `SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE=1` — writes **one combined trace per tensor-parallel rank** for the whole capture pass, named `cuda_graph_capture-{runner}-TP-{tp_rank}.json.gz`.
+- `SGLANG_GRAPH_BATCH_CAPTURE=1` — writes **one trace per captured batch size per rank**, named `{runner}_bs_{bs}_rank{tp_rank}.json.gz`. The profiler runs on a `wait=2, warmup=0, active=1` schedule (the two dummy runs before each capture are skipped) with `record_shapes`, `with_stack`, `with_flops`, and `profile_memory` enabled, giving per-shape kernel identities, input shapes, FLOPs, and memory.
+
+If both env vars are set, `SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE` (the single combined trace) takes precedence.
+
+```bash Command
+# set trace path
+export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log
+
+# opt in to per-batch-size capture traces (or set
+# SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE=1 for a single combined trace per rank)
+export SGLANG_GRAPH_BATCH_CAPTURE=1
+
+# launch the server with CUDA graph capture profiling enabled
+python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-profile-cuda-graph
+```
+
+Behavior and output:
+
+- All traces are written to `${SGLANG_TORCH_PROFILER_DIR}/graph_capture_profile/` (defaults to `/tmp/graph_capture_profile/` if the variable is unset). Files are namespaced by runner class and TP rank so concurrent capture passes (e.g. EAGLE target/draft/draft-extend) and ranks don't collide.
+- A CUDA memory snapshot (`cuda_graph_runner_memory_usage.pickle`) and per-kernel CPU/CUDA time summary tables are always emitted for the capture phase (independent of the env vars above).
+- Only the decode CUDA-graph runner is profiled.
+
+The capture traces are viewed the same way as other PyTorch Profiler traces (see [View traces](#view-traces)).
+
 ### Possible PyTorch bugs
 If in any cases you encounter the following error (for example, using qwen 2.5 VL):
 ```bash Command

python/sglang/srt/environ.py

@@ -328,6 +328,10 @@ class Envs:
     )
     SGLANG_RECORD_STEP_TIME = EnvBool(False)
     SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE = EnvBool(False)
+    # Opt-in: emit one CUDA-graph capture trace per captured batch size (per-bs).
+    # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE (single combined trace) takes
+    # precedence when both are set.
+    SGLANG_GRAPH_BATCH_CAPTURE = EnvBool(False)
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER = EnvBool(False)

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -28,6 +28,7 @@ from __future__ import annotations
 import contextlib
 import inspect
 import logging
+import os
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
@@ -102,7 +103,10 @@ from sglang.srt.utils import (
     require_mlp_tp_gather,
 )
 from sglang.srt.utils.device_timer import device_timer_ctx
-from sglang.srt.utils.profile_utils import export_cuda_graph_capture_trace
+from sglang.srt.utils.profile_utils import (
+    export_cuda_graph_capture_trace,
+    graph_capture_profile_dir,
+)
 
 try:
     from kt_kernel import KTMoEWrapper
@@ -682,11 +686,63 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             and capture_hidden_mode_matches
         )
 
-    def _init_profile_context_and_memory_record(self):
-        profile_context = profile(
-            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-            record_shapes=True,
+    def _graph_batch_capture_active(self) -> bool:
+        """Whether the per-batch-size capture-trace feature is active.
+
+        Gated by SGLANG_GRAPH_BATCH_CAPTURE. The original single-trace export
+        (SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE) takes precedence: when both are
+        set we fall back to the original behavior.
+        """
+        return (
+            envs.SGLANG_GRAPH_BATCH_CAPTURE.get()
+            and not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get()
         )
+
+    def _init_profile_context_and_memory_record(self):
+        if self._graph_batch_capture_active():
+            # Per-batch-size capture traces (SGLANG_GRAPH_BATCH_CAPTURE): a
+            # scheduled profiler is stepped once per batch size (see
+            # FullCudaGraphBackend.capture_one) and on_trace_ready writes one
+            # chrome trace per bs.
+            rank = get_parallel().tp_rank
+            runner_name = type(self).__name__
+            trace_dir = graph_capture_profile_dir()
+            os.makedirs(trace_dir, exist_ok=True)
+
+            # Track which BS is currently being captured for trace file naming
+            self._profile_bs_list = list(reversed(self.capture_bs))
+            self._profile_bs_idx = 0
+
+            def on_trace_ready(prof):
+                bs = self._profile_bs_list[self._profile_bs_idx]
+                trace_file = os.path.join(
+                    trace_dir, f"{runner_name}_bs_{bs}_rank{rank}.json.gz"
+                )
+                prof.export_chrome_trace(trace_file)
+                logger.info(f"Saved trace for bs={bs} to {trace_file}")
+                self._profile_bs_idx += 1
+
+            profile_context = profile(
+                activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+                # Schedule: wait=2 (skip 2 dummy runs), warmup=0, active=1
+                # (capture run); repeat=0 repeats the cycle so each batch size
+                # gets its own trace.
+                schedule=torch.profiler.schedule(wait=2, warmup=0, active=1, repeat=0),
+                record_shapes=True,
+                with_stack=True,
+                with_flops=True,
+                profile_memory=True,
+                on_trace_ready=on_trace_ready,
+            )
+        else:
+            # a single unscheduled pass over the whole
+            # capture. The combined trace (if any) is exported in
+            # _post_process_after_profile via export_cuda_graph_capture_trace,
+            # gated by SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.
+            profile_context = profile(
+                activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+                record_shapes=True,
+            )
         torch.cuda.memory._record_memory_history()
         return profile_context
 
@@ -706,10 +762,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         )
         logger.info(log_message)
 
-        # Optionally persist the shaped capture trace (record_shapes=True) for
-        # offline per-kernel analysis -- opt-in via
-        # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE; the in-log tables above are
-        # unchanged.
+        # single-trace export for the whole capture pass; no-op unless
+        # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE is set. In per-bs mode
+        # (SGLANG_GRAPH_BATCH_CAPTURE) that env is unset, so this stays a no-op
+        # and the per-bs on_trace_ready handles export instead.
         export_cuda_graph_capture_trace(
             prof_context,
             runner_name=type(self).__name__,
@@ -886,8 +942,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 self.model_runner, self.captured_req_width
             )
         profile_context = empty_context()
+        # Holds the active torch profiler during capture so the backend can
+        # advance its schedule (profiler.step()) per batch size. Only the
+        # scheduled per-bs profiler (SGLANG_GRAPH_BATCH_CAPTURE) needs stepping;
+        # the original unscheduled pass leaves this None.
+        self._profiler = None
         if self.enable_profile_cuda_graph:
             profile_context = self._init_profile_context_and_memory_record()
+            if self._graph_batch_capture_active():
+                self._profiler = profile_context
 
         # share_buffers() coalesces seq_lens / seq_lens_cpu through the process-
         # wide pool, so they may alias a buffer seeded by an earlier runner (the
@@ -924,6 +987,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
 
         if self.enable_profile_cuda_graph:
             self._post_process_after_profile(prof)
+        self._profiler = None
 
         # No pool-side pin to clear: the captured full-physical write loc rides the
         # backend's `ForwardMetadata.out_cache_loc_full_physical` (-> KVWriteLoc.full_loc).

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -58,6 +58,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {}
         self._outputs: Dict[Any, Any] = {}
         self._pool = None
+        self._cuda_graph_runner = cuda_graph_runner
         self._device_module = cuda_graph_runner.device_module
         self._tp_group = cuda_graph_runner.model_runner.tp_group
         self._capture_stream: Optional[torch.cuda.Stream] = None
@@ -84,12 +85,29 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
+        # When per-bs capture traces are enabled (--enable-profile-cuda-graph +
+        # SGLANG_GRAPH_BATCH_CAPTURE), the runner created a scheduled
+        # torch profiler (wait=2, active=1) and exposed it as _profiler. We step()
+        # past the two warmup runs so only the capture run is recorded, and each
+        # batch size produces its own trace via the profiler's on_trace_ready.
+        # With --enable-profile-cuda-graph alone the runner leaves _profiler None
+        # (its unscheduled profiler records the whole capture in one pass), so no
+        # stepping happens here.
+        runner = self._cuda_graph_runner
+        profiler = (
+            getattr(runner, "_profiler", None)
+            if getattr(runner, "enable_profile_cuda_graph", False)
+            else None
+        )
+
         # Two warmups so kernels are loaded and one-time setup is paid before capture.
         # post_warmup_hook lets the attention backend reset state that warmup mutated.
         for _ in range(2):
             self._device_module.synchronize()
             self._tp_group.barrier()
             forward_fn()
+            if profiler is not None:
+                profiler.step()
             if post_warmup_hook is not None:
                 post_warmup_hook()
 
@@ -110,6 +128,9 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream):
             out = forward_fn()
 
+        if profiler is not None:
+            profiler.step()
+
         self._graphs[shape_key] = graph
         self._outputs[shape_key] = out
 

python/sglang/srt/utils/profile_utils.py

@@ -30,6 +30,19 @@ if _is_npu:
 
 logger = logging.getLogger(__name__)
 
+# Single source of truth for the CUDA-graph capture trace output directory,
+# shared by both the original single-trace export and the per-batch-size
+# (SGLANG_GRAPH_BATCH_CAPTURE) traces so they land in the same place.
+GRAPH_CAPTURE_PROFILE_DIRNAME = "graph_capture_profile"
+
+
+def graph_capture_profile_dir() -> str:
+    """``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile`` — the one directory
+    both capture-trace modes write to. Change the location here only."""
+    return os.path.join(
+        envs.SGLANG_TORCH_PROFILER_DIR.get(), GRAPH_CAPTURE_PROFILE_DIRNAME
+    )
+
 
 def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: int):
     """Persist a CUDA-graph capture profiler trace (chrome trace) to disk.
@@ -37,15 +50,13 @@ def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank:
     Opt-in via ``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE`` (no-op otherwise). The
     capture profiler must have run with ``record_shapes=True`` so the trace can
     be inspected offline as a per-kernel shape/identity record. The file lands in
-    ``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile/`` and is namespaced by
-    runner class and TP rank so concurrent capture passes (e.g. EAGLE3
-    target/draft/draft-extend) and ranks don't overwrite each other.
+    ``graph_capture_profile_dir()`` and is namespaced by runner class and TP rank
+    so concurrent capture passes (e.g. EAGLE3 target/draft/draft-extend) and
+    ranks don't overwrite each other.
     """
     if not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get():
         return
-    output_dir = os.path.join(
-        envs.SGLANG_TORCH_PROFILER_DIR.get(), "graph_capture_profile"
-    )
+    output_dir = graph_capture_profile_dir()
     os.makedirs(output_dir, exist_ok=True)
     path = os.path.join(
         output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz"

test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py

@@ -0,0 +1,294 @@
+"""Unit tests for ``DecodeCudaGraphRunner`` capture-phase profiling — CPU-only.
+
+Two capture-trace modes plus their precedence:
+
+  * **Original single-trace** (``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE``):
+    ``_init_profile_context_and_memory_record`` builds an *unscheduled* profiler
+    (``record_shapes`` only, no schedule / no ``on_trace_ready``); the combined
+    trace is exported in ``_post_process_after_profile`` via
+    ``export_cuda_graph_capture_trace``.
+  * **Per-batch-size traces** (``SGLANG_GRAPH_BATCH_CAPTURE``): a *scheduled*
+    profiler (``wait=2, warmup=0, active=1, repeat=0``) with the trace-export
+    knobs (record_shapes / with_stack / with_flops / profile_memory) and an
+    ``on_trace_ready`` hook that writes one trace per batch size to
+    ``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile/`` named
+    ``{runner_name}_bs_{bs}_rank{rank}.json.gz``.
+  * **Precedence**: when both env vars are set, the original single-trace path
+    wins (no per-bs schedule / dir / bookkeeping).
+
+The profiler / CUDA-memory APIs are mocked; the directory + naming + schedule
+logic is pure-Python and runs on CPU. The method is invoked unbound against a
+lightweight stand-in (with the real precedence helper bound) so no model or
+server is constructed.
+"""
+
+import os
+import tempfile
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from sglang.srt.model_executor.runner import decode_cuda_graph_runner as mod
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    DecodeCudaGraphRunner,
+)
+from sglang.srt.utils import profile_utils as putils
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+_CAPTURE_TRACE = "SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE"
+_BATCH_CAPTURE = "SGLANG_GRAPH_BATCH_CAPTURE"
+
+
+def _make_fake_self(capture_bs):
+    """Stand-in ``self`` with the real precedence helper bound so the env-var
+    gating in ``_init_profile_context_and_memory_record`` applies."""
+    fake_self = SimpleNamespace(capture_bs=list(capture_bs))
+    fake_self._graph_batch_capture_active = (
+        DecodeCudaGraphRunner._graph_batch_capture_active.__get__(fake_self)
+    )
+    return fake_self
+
+
+class TestInitProfileBatchMode(CustomTestCase):
+    """SGLANG_GRAPH_BATCH_CAPTURE -> scheduled per-bs profiler."""
+
+    def _invoke(self, *, capture_bs, rank=0, profiler_dir=None):
+        fake_self = _make_fake_self(capture_bs)
+        env = {_BATCH_CAPTURE: "1"}
+        if profiler_dir is not None:
+            env["SGLANG_TORCH_PROFILER_DIR"] = profiler_dir
+        with mock.patch.dict(os.environ, env, clear=False), mock.patch.object(
+            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
+        ), mock.patch.object(mod, "profile") as mock_profile, mock.patch(
+            "torch.profiler.schedule"
+        ) as mock_schedule, mock.patch(
+            "torch.cuda.memory._record_memory_history"
+        ) as mock_record_history:
+            os.environ.pop(_CAPTURE_TRACE, None)  # original flag off
+            if profiler_dir is None:
+                os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None)
+            ctx = DecodeCudaGraphRunner._init_profile_context_and_memory_record(
+                fake_self
+            )
+        self.assertIs(ctx, mock_profile.return_value)
+        return fake_self, mock_profile, mock_schedule, mock_record_history
+
+    def test_creates_graph_capture_profile_dir(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            self._invoke(capture_bs=[1, 2, 4], profiler_dir=tmp)
+            self.assertTrue(os.path.isdir(os.path.join(tmp, "graph_capture_profile")))
+
+    def test_primes_reversed_bs_list_and_zero_index(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            fake_self, *_ = self._invoke(capture_bs=[1, 2, 4, 8], profiler_dir=tmp)
+            # Capture iterates large -> small, so the bs list is reversed.
+            self.assertEqual(fake_self._profile_bs_list, [8, 4, 2, 1])
+            self.assertEqual(fake_self._profile_bs_idx, 0)
+
+    def test_profiler_built_with_trace_export_knobs(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            _, mock_profile, mock_schedule, mock_record_history = self._invoke(
+                capture_bs=[1, 2], profiler_dir=tmp
+            )
+            self.assertEqual(mock_profile.call_count, 1)
+            kwargs = mock_profile.call_args.kwargs
+            self.assertTrue(kwargs["record_shapes"])
+            self.assertTrue(kwargs["with_stack"])
+            self.assertTrue(kwargs["with_flops"])
+            self.assertTrue(kwargs["profile_memory"])
+            self.assertTrue(callable(kwargs["on_trace_ready"]))
+            # Schedule skips the two dummy/warmup runs and records the capture.
+            mock_schedule.assert_called_once_with(wait=2, warmup=0, active=1, repeat=0)
+            self.assertIs(kwargs["schedule"], mock_schedule.return_value)
+            # Memory history recording is armed alongside the profiler.
+            mock_record_history.assert_called_once()
+
+    def test_default_dir_used_when_profiler_dir_env_unset(self):
+        # No SGLANG_TORCH_PROFILER_DIR -> falls back to the envs default base dir.
+        # Patch makedirs so the test never writes to the cwd.
+        fake_self = _make_fake_self([1])
+        with mock.patch.dict(
+            os.environ, {_BATCH_CAPTURE: "1"}, clear=False
+        ), mock.patch.object(
+            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
+        ), mock.patch.object(
+            mod, "profile"
+        ), mock.patch(
+            "torch.profiler.schedule"
+        ), mock.patch(
+            "torch.cuda.memory._record_memory_history"
+        ), mock.patch.object(
+            mod.os, "makedirs"
+        ) as mock_makedirs:
+            os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None)
+            os.environ.pop(_CAPTURE_TRACE, None)
+            DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self)
+
+        mock_makedirs.assert_called_once()
+        self.assertEqual(
+            mock_makedirs.call_args.args[0],
+            os.path.join("/tmp", "graph_capture_profile"),
+        )
+
+
+class TestInitProfileOriginalMode(CustomTestCase):
+    """No flag, original flag only, or both (precedence) -> unscheduled pass with
+    no per-bs schedule / directory / bookkeeping."""
+
+    def _invoke_original(self, *, env):
+        fake_self = _make_fake_self([1, 2])
+        with tempfile.TemporaryDirectory() as tmp:
+            environ = dict(env)
+            environ["SGLANG_TORCH_PROFILER_DIR"] = tmp
+            with mock.patch.dict(os.environ, environ, clear=False), mock.patch.object(
+                mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
+            ), mock.patch.object(mod, "profile") as mock_profile, mock.patch(
+                "torch.profiler.schedule"
+            ) as mock_schedule, mock.patch(
+                "torch.cuda.memory._record_memory_history"
+            ):
+                for k in (_CAPTURE_TRACE, _BATCH_CAPTURE):
+                    if k not in environ:
+                        os.environ.pop(k, None)
+                DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self)
+            kwargs = mock_profile.call_args.kwargs
+            # Unscheduled pass: record_shapes only, no schedule / on_trace_ready.
+            self.assertTrue(kwargs["record_shapes"])
+            self.assertIsNone(kwargs.get("schedule"))
+            self.assertIsNone(kwargs.get("on_trace_ready"))
+            mock_schedule.assert_not_called()
+            self.assertFalse(os.path.isdir(os.path.join(tmp, "graph_capture_profile")))
+            self.assertFalse(hasattr(fake_self, "_profile_bs_list"))
+
+    def test_no_flags(self):
+        self._invoke_original(env={})
+
+    def test_original_flag_only(self):
+        self._invoke_original(env={_CAPTURE_TRACE: "1"})
+
+    def test_both_flags_original_takes_precedence(self):
+        self._invoke_original(env={_CAPTURE_TRACE: "1", _BATCH_CAPTURE: "1"})
+
+
+class TestOnTraceReadyNaming(CustomTestCase):
+    def _build_on_trace_ready(self, *, capture_bs, rank, tmp):
+        fake_self = _make_fake_self(capture_bs)
+        with mock.patch.dict(
+            os.environ,
+            {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"},
+            clear=False,
+        ), mock.patch.object(
+            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
+        ), mock.patch.object(
+            mod, "profile"
+        ) as mock_profile, mock.patch(
+            "torch.profiler.schedule"
+        ), mock.patch(
+            "torch.cuda.memory._record_memory_history"
+        ):
+            os.environ.pop(_CAPTURE_TRACE, None)
+            DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self)
+        on_trace_ready = mock_profile.call_args.kwargs["on_trace_ready"]
+        return fake_self, on_trace_ready
+
+    def test_exports_one_named_trace_per_bs_and_advances_index(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            capture_bs = [1, 2, 4]  # reversed -> [4, 2, 1]
+            fake_self, on_trace_ready = self._build_on_trace_ready(
+                capture_bs=capture_bs, rank=0, tmp=tmp
+            )
+            trace_dir = os.path.join(tmp, "graph_capture_profile")
+            runner = type(fake_self).__name__
+
+            exported = []
+            for expected_bs in [4, 2, 1]:
+                prof = mock.Mock()
+                prof.export_chrome_trace.side_effect = lambda p: exported.append(p)
+                on_trace_ready(prof)
+                prof.export_chrome_trace.assert_called_once_with(
+                    os.path.join(trace_dir, f"{runner}_bs_{expected_bs}_rank0.json.gz")
+                )
+
+            self.assertEqual(
+                exported,
+                [
+                    os.path.join(trace_dir, f"{runner}_bs_4_rank0.json.gz"),
+                    os.path.join(trace_dir, f"{runner}_bs_2_rank0.json.gz"),
+                    os.path.join(trace_dir, f"{runner}_bs_1_rank0.json.gz"),
+                ],
+            )
+            # Index advanced once per flush.
+            self.assertEqual(fake_self._profile_bs_idx, 3)
+
+    def test_rank_in_trace_filename(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            fake_self, on_trace_ready = self._build_on_trace_ready(
+                capture_bs=[8], rank=3, tmp=tmp
+            )
+            runner = type(fake_self).__name__
+            prof = mock.Mock()
+            on_trace_ready(prof)
+            prof.export_chrome_trace.assert_called_once_with(
+                os.path.join(
+                    tmp, "graph_capture_profile", f"{runner}_bs_8_rank3.json.gz"
+                )
+            )
+
+
+class TestOriginalTraceExport(CustomTestCase):
+    """export_cuda_graph_capture_trace (original single combined trace per rank),
+    gated by SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE, and the shared dir helper.
+    Both trace modes land under graph_capture_profile/."""
+
+    def test_writes_named_trace_when_flag_set(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            with mock.patch.dict(
+                os.environ,
+                {"SGLANG_TORCH_PROFILER_DIR": tmp, _CAPTURE_TRACE: "1"},
+                clear=False,
+            ):
+                prof = mock.Mock()
+                putils.export_cuda_graph_capture_trace(
+                    prof, runner_name="DecodeCudaGraphRunner", tp_rank=2
+                )
+                expected = os.path.join(
+                    tmp,
+                    "graph_capture_profile",
+                    "cuda_graph_capture-DecodeCudaGraphRunner-TP-2.json.gz",
+                )
+                prof.export_chrome_trace.assert_called_once_with(expected)
+                self.assertTrue(
+                    os.path.isdir(os.path.join(tmp, "graph_capture_profile"))
+                )
+
+    def test_noop_when_flag_unset(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            with mock.patch.dict(
+                os.environ, {"SGLANG_TORCH_PROFILER_DIR": tmp}, clear=False
+            ):
+                os.environ.pop(_CAPTURE_TRACE, None)
+                prof = mock.Mock()
+                putils.export_cuda_graph_capture_trace(
+                    prof, runner_name="DecodeCudaGraphRunner", tp_rank=0
+                )
+                prof.export_chrome_trace.assert_not_called()
+                self.assertFalse(
+                    os.path.isdir(os.path.join(tmp, "graph_capture_profile"))
+                )
+
+    def test_dir_helper_uses_profiler_dir(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            with mock.patch.dict(
+                os.environ, {"SGLANG_TORCH_PROFILER_DIR": tmp}, clear=False
+            ):
+                self.assertEqual(
+                    putils.graph_capture_profile_dir(),
+                    os.path.join(tmp, "graph_capture_profile"),
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py

@@ -0,0 +1,162 @@
+"""Unit tests for ``FullCudaGraphBackend.capture_one`` profiling hooks — CPU-only.
+
+These cover the changes from the "cuda graph profile traces" PR that wire the
+runner's torch profiler into the capture loop:
+
+  * When profiling is disabled, ``capture_one`` runs exactly two warmups + one
+    capture and never touches a profiler (behavior-identical to before the PR).
+  * When the runner exposes an active ``_profiler`` (per-bs capture profiling,
+    ``--enable-profile-cuda-graph`` + ``SGLANG_GRAPH_BATCH_CAPTURE``),
+    ``capture_one`` calls ``profiler.step()`` past the two warmups and once after
+    the capture (schedule ``wait=2, warmup=0, active=1``). The captured forward is
+    NOT wrapped in a ``record_function``; per-bs trace naming is handled by the
+    profiler's ``on_trace_ready`` callback instead.
+  * The ``getattr`` guards mean a runner that sets the flag but has no
+    ``_profiler`` attribute degrades gracefully (no stepping, no crash).
+
+The real capture path needs CUDA (``torch.cuda.CUDAGraph`` + device graph
+context), so those are mocked; the logic under test (call counts, ordering,
+profiler stepping) is pure-Python and runs on CPU.
+"""
+
+import contextlib
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
+from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
+    FullCudaGraphBackend,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+# Sentinel: distinguishes "runner has no _profiler attribute" from
+# "_profiler is None" in the test fixtures.
+_UNSET = object()
+
+
+class _FakeGraphCtx:
+    """Stand-in for ``device_module.graph(...)`` — a no-op context manager."""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+
+def _make_backend(runner):
+    """Build a ``FullCudaGraphBackend`` without running ``__init__`` (which would
+    touch CUDA), wiring just the attributes ``capture_one`` reads."""
+    backend = FullCudaGraphBackend.__new__(FullCudaGraphBackend)
+    backend._graphs = {}
+    backend._outputs = {}
+    backend._pool = None
+    backend._capture_stream = None
+    backend._memory_saver_adapter = None
+    backend._cuda_graph_runner = runner
+    backend._device_module = runner.device_module
+    backend._tp_group = runner.model_runner.tp_group
+    return backend
+
+
+def _make_runner(*, enable_profile, profiler, num_tokens_per_bs=1, mode_name="DECODE"):
+    device_module = SimpleNamespace(
+        synchronize=mock.Mock(name="synchronize"),
+        graph=mock.Mock(name="graph", side_effect=lambda **kw: _FakeGraphCtx()),
+    )
+    tp_group = SimpleNamespace(barrier=mock.Mock(name="barrier"))
+    runner = SimpleNamespace(
+        device_module=device_module,
+        model_runner=SimpleNamespace(tp_group=tp_group),
+        num_tokens_per_bs=num_tokens_per_bs,
+        capture_forward_mode=SimpleNamespace(name=mode_name),
+        enable_profile_cuda_graph=enable_profile,
+    )
+    if profiler is not _UNSET:
+        runner._profiler = profiler
+    return runner
+
+
+class TestCaptureOneNoProfiling(CustomTestCase):
+    def test_runs_two_warmups_and_capture_without_stepping(self):
+        runner = _make_runner(enable_profile=False, profiler=None)
+        backend = _make_backend(runner)
+
+        sentinel_out = object()
+        forward_fn = mock.Mock(return_value=sentinel_out)
+        post_warmup_hook = mock.Mock()
+        shape_key = ShapeKey(size=4)
+
+        with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"):
+            backend.capture_one(
+                shape_key, forward_fn, post_warmup_hook=post_warmup_hook
+            )
+
+        # 2 warmups + 1 capture.
+        self.assertEqual(forward_fn.call_count, 3)
+        # post_warmup_hook only runs in the two warmup iterations.
+        self.assertEqual(post_warmup_hook.call_count, 2)
+        # Graph + output are recorded against the shape key.
+        self.assertEqual(backend._graphs[shape_key], "GRAPH")
+        self.assertIs(backend._outputs[shape_key], sentinel_out)
+
+    def test_enable_flag_set_but_no_profiler_attr_does_not_step(self):
+        # The runner advertises the flag but never created a profiler; the
+        # getattr guard must keep capture_one on the non-profiling path.
+        runner = _make_runner(enable_profile=True, profiler=_UNSET)
+        backend = _make_backend(runner)
+        self.assertFalse(hasattr(runner, "_profiler"))
+
+        forward_fn = mock.Mock(return_value=object())
+        with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"):
+            backend.capture_one(ShapeKey(size=2), forward_fn)
+
+        self.assertEqual(forward_fn.call_count, 3)
+
+
+class TestCaptureOneWithProfiling(CustomTestCase):
+    def _run(self, *, size, num_tokens_per_bs, mode_name):
+        profiler = SimpleNamespace(step=mock.Mock(name="step"))
+        runner = _make_runner(
+            enable_profile=True,
+            profiler=profiler,
+            num_tokens_per_bs=num_tokens_per_bs,
+            mode_name=mode_name,
+        )
+        backend = _make_backend(runner)
+
+        forward_fn = mock.Mock(return_value=object())
+        rf_names = []
+
+        def _fake_record_function(name):
+            rf_names.append(name)
+            return contextlib.nullcontext()
+
+        with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"), mock.patch(
+            "torch.profiler.record_function", side_effect=_fake_record_function
+        ):
+            backend.capture_one(ShapeKey(size=size), forward_fn)
+
+        return profiler, forward_fn, rf_names
+
+    def test_steps_twice_in_warmup_and_once_after_capture(self):
+        profiler, forward_fn, _ = self._run(
+            size=4, num_tokens_per_bs=1, mode_name="DECODE"
+        )
+        # Schedule wait=2 + active=1 => one step per warmup (x2) + one post-capture.
+        self.assertEqual(profiler.step.call_count, 3)
+        self.assertEqual(forward_fn.call_count, 3)
+
+    def test_capture_not_wrapped_in_record_function(self):
+        # The capture forward is no longer wrapped in a record_function; per-bs
+        # trace naming is handled by the profiler's on_trace_ready callback.
+        _, _, rf_names = self._run(size=4, num_tokens_per_bs=1, mode_name="DECODE")
+        self.assertEqual(rf_names, [])
+
+
+if __name__ == "__main__":
+    unittest.main()