[AMD] Resolving bug in per batch size graph capture due to multiple capture runs per batch size

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

@@ -835,35 +835,76 @@ def _graph_batch_capture_active(self) -> bool:
             and not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get()
         )
 
+    def _set_profile_trace_tag(
+        self,
+        bs: int,
+        stream_idx: Optional[int],
+        variant_label: Optional[str],
+        attention_variant: Optional[str],
+    ) -> None:
+        """Name the trace for one capture. bs alone is not unique: the capture
+        loop multiplies it by the lora and attention variants (DSA captures
+        dense + sparse per bs). Absent variants are omitted so single-variant
+        runners keep the historical `bs_<n>` stem.
+        """
+        if getattr(self, "_profiler", None) is None:
+            return
+        parts = [f"bs_{bs}"]
+        if attention_variant is not None:
+            parts.append(attention_variant)
+        if variant_label is not None:
+            parts.append(variant_label)
+        if stream_idx is not None:
+            parts.append(f"stream_{stream_idx}")
+        self._profile_trace_tag = "_".join(parts)
+
+    def _profile_runner_name(self) -> str:
+        """Runner name used in capture-trace filenames. DSPARK / DFLASH draft
+        workers capture their verify graphs with this base class too, at the
+        same shapes as the target, so the name alone would collide; prefix it.
+        Draft-specific subclasses (EAGLEDraft*, Uno*, ...) are already unique.
+        """
+        name = type(self).__name__
+        if type(self) is DecodeCudaGraphRunner and self.model_runner.is_draft_worker:
+            return f"Draft{name}"
+        return name
+
     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.
+            # SGLANG_GRAPH_BATCH_CAPTURE: the profiler is stepped once per
+            # captured shape (see FullCudaGraphBackend.capture_one), so
+            # on_trace_ready writes one chrome trace per shape.
             rank = get_parallel().tp_rank
-            runner_name = type(self).__name__
+            runner_name = self._profile_runner_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
+            self._profile_trace_tag = None
 
             def on_trace_ready(prof):
-                bs = self._profile_bs_list[self._profile_bs_idx]
+                tag = self._profile_trace_tag
+                if tag is None:
+                    # Reachable only if a subclass overrides _capture_one_stream
+                    # without tagging; exporting would collide every flush on
+                    # one name.
+                    logger.warning(
+                        "%s: capture trace flushed with no shape tag; skipping "
+                        "export. A subclass likely overrides _capture_one_stream "
+                        "without calling _set_profile_trace_tag.",
+                        runner_name,
+                    )
+                    return
                 trace_file = os.path.join(
-                    trace_dir, f"{runner_name}_bs_{bs}_rank{rank}.json.gz"
+                    trace_dir, f"{runner_name}_{tag}_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
+                logger.info(f"Saved trace for {tag} to {trace_file}")
 
             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.
+                # (capture run); repeat=0 repeats the cycle so each captured
+                # shape gets its own trace.
                 schedule=torch.profiler.schedule(wait=2, warmup=0, active=1, repeat=0),
                 record_shapes=True,
                 with_stack=True,
@@ -905,7 +946,7 @@ def _post_process_after_profile(self, prof_context):
         # and the per-bs on_trace_ready handles export instead.
         export_cuda_graph_capture_trace(
             prof_context,
-            runner_name=type(self).__name__,
+            runner_name=self._profile_runner_name(),
         )
 
     def capture_prepare(
@@ -1173,6 +1214,9 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
                 _set_capture_lora_variant(variant_label)
                 for attention_variant in attention_variants:
                     _set_capture_attention_variant(attention_variant)
+                    self._set_profile_trace_tag(
+                        bs, stream_idx, variant_label, attention_variant
+                    )
                     with torch_compile_decoration.patch_model(
                         self.model_runner.model,
                         bs in self.compile_bs,

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

@@ -145,7 +145,7 @@ def capture_one(
         # 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.
+        # captured shape 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.

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

@@ -10,9 +10,10 @@
   * **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
+    ``on_trace_ready`` hook that writes one trace per captured shape to
     ``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile/`` named
-    ``{runner_name}_bs_{bs}_rank{rank}.json.gz``.
+    ``{runner_name}_{tag}_rank{rank}.json.gz``, where ``tag`` comes from
+    ``_set_profile_trace_tag`` (bs plus whichever variants are active).
   * **Precedence**: when both env vars are set, the original single-trace path
     wins (no per-bs schedule / dir / bookkeeping).
 
@@ -49,9 +50,25 @@ def _make_fake_self(capture_bs):
     fake_self._graph_batch_capture_active = (
         DecodeCudaGraphRunner._graph_batch_capture_active.__get__(fake_self)
     )
+    # _profiler is the gate _set_profile_trace_tag checks.
+    fake_self._profiler = mock.Mock()
+    fake_self._set_profile_trace_tag = (
+        DecodeCudaGraphRunner._set_profile_trace_tag.__get__(fake_self)
+    )
+    fake_self._profile_runner_name = DecodeCudaGraphRunner._profile_runner_name.__get__(
+        fake_self
+    )
     return fake_self
 
 
+def _make_runner(cls, *, is_draft_worker):
+    """Uninitialized runner of ``cls`` with only the model_runner flag the
+    trace-naming helper reads."""
+    runner = cls.__new__(cls)
+    runner.model_runner = SimpleNamespace(is_draft_worker=is_draft_worker)
+    return runner
+
+
 class TestInitProfileBatchMode(CustomTestCase):
     """SGLANG_GRAPH_BATCH_CAPTURE -> scheduled per-bs profiler."""
 
@@ -85,12 +102,10 @@ def test_creates_graph_capture_profile_dir(self):
             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):
+    def test_primes_empty_trace_tag(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)
+            self.assertIsNone(fake_self._profile_trace_tag)
 
 
 class TestInitProfileOriginalMode(CustomTestCase):
@@ -122,7 +137,7 @@ def _invoke_original(self, *, env):
             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"))
+            self.assertFalse(hasattr(fake_self, "_profile_trace_tag"))
 
     def test_no_flags(self):
         self._invoke_original(env={})
@@ -135,8 +150,11 @@ def test_both_flags_original_takes_precedence(self):
 
 
 class TestOnTraceReadyNaming(CustomTestCase):
-    def _build_on_trace_ready(self, *, capture_bs, rank, tmp):
+    @staticmethod
+    def _build_on_trace_ready(*, capture_bs, rank, tmp, runner_name=None):
         fake_self = _make_fake_self(capture_bs)
+        if runner_name is not None:
+            fake_self._profile_runner_name = lambda: runner_name
         with (
             mock.patch.dict(
                 os.environ,
@@ -155,23 +173,24 @@ def _build_on_trace_ready(self, *, capture_bs, rank, tmp):
         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):
+    @staticmethod
+    def _flush(on_trace_ready, exported):
+        prof = mock.Mock()
+        prof.export_chrome_trace.side_effect = exported.append
+        on_trace_ready(prof)
+
+    def test_single_variant_keeps_bs_only_naming(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
+                capture_bs=[1, 2, 4], 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")
-                )
+            for bs in (4, 2, 1):
+                fake_self._set_profile_trace_tag(bs, None, None, None)
+                self._flush(on_trace_ready, exported)
 
             self.assertEqual(
                 exported,
@@ -181,15 +200,60 @@ def test_exports_one_named_trace_per_bs_and_advances_index(self):
                     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_dual_attention_variants_name_both_traces(self):
+        """DSA models capture dense + sparse per bs, so two flushes per bs."""
+        with tempfile.TemporaryDirectory() as tmp:
+            fake_self, on_trace_ready = self._build_on_trace_ready(
+                capture_bs=[1, 2], rank=0, tmp=tmp
+            )
+            trace_dir = os.path.join(tmp, "graph_capture_profile")
+            runner = type(fake_self).__name__
+
+            exported = []
+            for bs in (2, 1):
+                for attention_variant in ("dense", "sparse"):
+                    fake_self._set_profile_trace_tag(bs, None, None, attention_variant)
+                    self._flush(on_trace_ready, exported)
+
+            self.assertEqual(
+                exported,
+                [
+                    os.path.join(trace_dir, f"{runner}_bs_2_dense_rank0.json.gz"),
+                    os.path.join(trace_dir, f"{runner}_bs_2_sparse_rank0.json.gz"),
+                    os.path.join(trace_dir, f"{runner}_bs_1_dense_rank0.json.gz"),
+                    os.path.join(trace_dir, f"{runner}_bs_1_sparse_rank0.json.gz"),
+                ],
+            )
+            self.assertEqual(len(set(exported)), len(exported))
+
+    def test_more_flushes_than_batch_sizes_does_not_raise(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            fake_self, on_trace_ready = self._build_on_trace_ready(
+                capture_bs=[1, 2], rank=0, tmp=tmp
+            )
+            exported = []
+            for i in range(6):
+                fake_self._set_profile_trace_tag(2, None, None, f"v{i}")
+                self._flush(on_trace_ready, exported)
+            self.assertEqual(len(exported), 6)
+
+    def test_unset_tag_skips_export(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            _, on_trace_ready = self._build_on_trace_ready(
+                capture_bs=[8], rank=0, tmp=tmp
+            )
+            prof = mock.Mock()
+            on_trace_ready(prof)  # tag never set by the capture loop
+            prof.export_chrome_trace.assert_not_called()
 
     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__
+            fake_self._set_profile_trace_tag(8, None, None, None)
             prof = mock.Mock()
             on_trace_ready(prof)
             prof.export_chrome_trace.assert_called_once_with(
@@ -199,6 +263,64 @@ def test_rank_in_trace_filename(self):
             )
 
 
+class TestProfileTraceTag(CustomTestCase):
+    """_set_profile_trace_tag composition."""
+
+    def test_bs_only_when_no_variants(self):
+        fake_self = _make_fake_self([1])
+        fake_self._set_profile_trace_tag(16, None, None, None)
+        self.assertEqual(fake_self._profile_trace_tag, "bs_16")
+
+    def test_appends_attention_lora_and_stream(self):
+        fake_self = _make_fake_self([1])
+        fake_self._set_profile_trace_tag(8, 1, "nolora", "sparse")
+        self.assertEqual(fake_self._profile_trace_tag, "bs_8_sparse_nolora_stream_1")
+
+    def test_noop_without_scheduled_profiler(self):
+        fake_self = _make_fake_self([1])
+        fake_self._profiler = None
+        fake_self._set_profile_trace_tag(8, None, None, "sparse")
+        self.assertFalse(hasattr(fake_self, "_profile_trace_tag"))
+
+
+class TestProfileRunnerName(CustomTestCase):
+    """DSPARK / DFLASH drafts capture with the base class at the target's
+    shapes, so their traces must not share the target's runner name."""
+
+    def test_target_keeps_class_name(self):
+        runner = _make_runner(DecodeCudaGraphRunner, is_draft_worker=False)
+        self.assertEqual(runner._profile_runner_name(), "DecodeCudaGraphRunner")
+
+    def test_base_class_draft_worker_is_prefixed(self):
+        runner = _make_runner(DecodeCudaGraphRunner, is_draft_worker=True)
+        self.assertEqual(runner._profile_runner_name(), "DraftDecodeCudaGraphRunner")
+
+    def test_draft_subclass_keeps_class_name(self):
+        class EAGLEDraftLikeRunner(DecodeCudaGraphRunner):
+            pass
+
+        runner = _make_runner(EAGLEDraftLikeRunner, is_draft_worker=True)
+        self.assertEqual(runner._profile_runner_name(), "EAGLEDraftLikeRunner")
+
+    def test_target_and_draft_traces_do_not_collide(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            exported = []
+            for is_draft_worker in (False, True):
+                name = _make_runner(
+                    DecodeCudaGraphRunner, is_draft_worker=is_draft_worker
+                )._profile_runner_name()
+                fake_self, on_trace_ready = (
+                    TestOnTraceReadyNaming._build_on_trace_ready(
+                        capture_bs=[1, 2], rank=0, tmp=tmp, runner_name=name
+                    )
+                )
+                for bs in (2, 1):
+                    fake_self._set_profile_trace_tag(bs, None, None, None)
+                    TestOnTraceReadyNaming._flush(on_trace_ready, exported)
+            self.assertEqual(len(exported), 4)
+            self.assertEqual(len(set(exported)), len(exported))
+
+
 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.