From 1c8551169d90ac6563a86d5cff4a04b0e589dc9f Mon Sep 17 00:00:00 2001
From: Lucia Fang <116399278+luccafong@users.noreply.github.com>
Date: Thu, 18 Jun 2026 18:51:14 -0700
Subject: [PATCH] Add opt-in CUDA-graph capture-trace export (#28551)

---
 python/sglang/srt/environ.py                  |  1 +
 .../runner/decode_cuda_graph_runner.py        | 11 +++++++++
 python/sglang/srt/utils/profile_utils.py      | 24 +++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index c154e0438..6b85f2f98 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -251,6 +251,7 @@ class Envs:
         False, deprecated_name="SGLANG_OPERATIONS_ENABLE_PROFILE"
     )
     SGLANG_RECORD_STEP_TIME = EnvBool(False)
+    SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE = EnvBool(False)
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_DEBUG_REVERT_PR = EnvInt(0)
diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
index 72c59193f..fb75707e4 100644
--- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
@@ -101,6 +101,7 @@ from sglang.srt.utils import (
     require_mlp_sync,
     require_mlp_tp_gather,
 )
+from sglang.srt.utils.profile_utils import export_cuda_graph_capture_trace
 
 try:
     from kt_kernel import KTMoEWrapper
@@ -604,6 +605,16 @@ 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.
+        export_cuda_graph_capture_trace(
+            prof_context,
+            runner_name=type(self).__name__,
+            tp_rank=get_tensor_model_parallel_rank(),
+        )
+
     def capture_prepare(
         self,
         size: int,
diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py
index 7b6bcac31..84e8db0b9 100644
--- a/python/sglang/srt/utils/profile_utils.py
+++ b/python/sglang/srt/utils/profile_utils.py
@@ -9,6 +9,7 @@ from typing import Callable, Dict, List, Optional
 import torch
 
 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 ForwardMode
 from sglang.srt.server_args import get_global_server_args
@@ -29,6 +30,29 @@ if _is_npu:
 logger = logging.getLogger(__name__)
 
 
+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.
+
+    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.
+    """
+    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"
+    )
+    os.makedirs(output_dir, exist_ok=True)
+    path = os.path.join(
+        output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz"
+    )
+    prof_context.export_chrome_trace(path)
+    logger.info(f"CUDA graph capture trace saved to: {path}")
+
+
 class ProfileManager:
     def __init__(self, ps: ParallelState, cpu_group):
         self.stage_based_trigger = _StageBasedTrigger(