From 185543b2ad7407d0f05b97b9c6af9b2ec4d771ed Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 20:41:24 -0700
Subject: [PATCH 1/8] Separate attention graph variants from decode runner
 mechanics

---
 .../srt/layers/attention/dsa/dsa_indexer.py   |   6 +-
 .../layers/attention/dsa/indexer_capture.py   |  39 +++++++
 .../srt/layers/attention/graph_variants.py    |  27 +++++
 .../runner/decode_cuda_graph_runner.py        | 109 +++++-------------
 .../srt/model_executor/runner/shape_key.py    |   8 +-
 .../runner_utils/capture_mode.py              |  21 ++--
 .../srt/speculative/uno_cuda_graph_runner.py  |   8 +-
 .../runner/test_decode_cuda_graph_runner.py   |  60 +++++++++-
 8 files changed, 176 insertions(+), 102 deletions(-)
 create mode 100644 python/sglang/srt/layers/attention/dsa/indexer_capture.py
 create mode 100644 python/sglang/srt/layers/attention/graph_variants.py

diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index 6678e4b3abb6..ef5d079c5327 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -424,7 +424,7 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
             # this common code keeps the original behavior (decode never skips
             # the indexer, i.e. always runs the full logits path) because the
             # decode k-only path has not been validated on CUDA yet. Mirrors the
-            # is_hip() gate on dsa_dual_graph in decode_cuda_graph_runner, which
+            # is_hip() gate in create_attention_graph_variants, which
             # already prevents the CUDA capture path from setting a "dense"
             # variant.
             if not _is_hip:
@@ -440,10 +440,10 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
                 # at replay. The capture-variant signal tells us which one to
                 # bake in.
                 from sglang.srt.model_executor.runner_utils.capture_mode import (
-                    get_capture_dsa_variant,
+                    get_capture_attention_variant,
                 )
 
-                variant = get_capture_dsa_variant()
+                variant = get_capture_attention_variant()
                 if variant == "dense":
                     return True
                 if variant == "sparse":
diff --git a/python/sglang/srt/layers/attention/dsa/indexer_capture.py b/python/sglang/srt/layers/attention/dsa/indexer_capture.py
new file mode 100644
index 000000000000..0f8a1a99eefb
--- /dev/null
+++ b/python/sglang/srt/layers/attention/dsa/indexer_capture.py
@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class DsaGraphVariants:
+    index_topk: int
+    # Dense comes first: the sparse capture peak subsumes its shared-pool storage.
+    capture_labels: ClassVar[tuple[str, ...]] = ("dense", "sparse")
+
+    def __post_init__(self):
+        logger.info(
+            "[dense-decode] DSA dual-graph enabled: capturing "
+            "dense (k-only) + sparse (full indexer) decode graphs; "
+            "dispatch on max_kv_len vs index_topk=%d.",
+            self.index_topk,
+        )
+
+    def select(self, forward_batch: ForwardBatch) -> str:
+        seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
+        if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0:
+            # Host-side mirror (maintained incrementally for plain decode) — no
+            # d2h sync needed.
+            max_kv_len = int(seq_lens_cpu.max().item())
+        elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0:
+            # Fallback: a single scalar reduction d2h (cheap, per-step).
+            max_kv_len = int(forward_batch.seq_lens.max().item())
+        else:
+            # No length info: be safe and use the correct-for-all sparse graph.
+            return "sparse"
+        return "dense" if max_kv_len <= self.index_topk else "sparse"
diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py
new file mode 100644
index 000000000000..d08e4f6c8f20
--- /dev/null
+++ b/python/sglang/srt/layers/attention/graph_variants.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Protocol
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+
+class AttentionGraphVariants(Protocol):
+    # Capture order is significant when variants share a graph memory pool.
+    capture_labels: tuple[str, ...]
+
+    def select(self, forward_batch: ForwardBatch) -> str:
+        """Select one of capture_labels from host-side batch metadata."""
+        ...
+
+
+def create_attention_graph_variants(
+    hf_config, forward_mode: ForwardMode
+) -> Optional[AttentionGraphVariants]:
+    from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
+    from sglang.srt.layers.attention.dsa.indexer_capture import DsaGraphVariants
+    from sglang.srt.utils import is_hip
+
+    if is_hip() and is_deepseek_dsa(hf_config):
+        return DsaGraphVariants(get_dsa_index_topk(hf_config))
+    return None
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 51463f6e4cfc..3b60412e6655 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
@@ -49,6 +49,7 @@
     SharedReadEnds,
 )
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+from sglang.srt.layers.attention.graph_variants import create_attention_graph_variants
 from sglang.srt.layers.cp.utils import is_mla_cp_enabled
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
@@ -91,7 +92,7 @@
     DecodeInputBuffers,
 )
 from sglang.srt.model_executor.runner_utils.capture_mode import (
-    _set_capture_dsa_variant,
+    _set_capture_attention_variant,
     _set_capture_lora_variant,
     model_capture_mode,
 )
@@ -113,7 +114,6 @@
 from sglang.srt.utils import (
     empty_context,
     get_available_gpu_memory,
-    is_hip,
     require_attn_tp_gather,
     require_mlp_tp_gather,
 )
@@ -254,37 +254,6 @@ def __init__(
         self.speculative_algorithm = get_spec().speculative_algorithm
         self.enable_profile_cuda_graph = get_exec().graph.enable_profile_cuda_graph
 
-        # --- DSA dense-decode dual-graph -------------------------------
-        # Capture a "dense" (k-only, skip-indexer) and a "sparse" (full indexer)
-        # decode graph per bs bucket, and dispatch on max_kv_len vs index_topk at
-        # replay. Auto-enabled for DSA models (index_topk present in the HF
-        # config) — correct for mixed lengths since any request with
-        # kv_len > index_topk falls back to the sparse graph. Adds ~52 graphs and
-        # ~2x capture time.
-        #
-        # Scoped to HIP (AMD): the k-only dense-decode fast path has only been
-        # validated on MI355X. This is common (non-hardware-gated) code, so on
-        # CUDA we deliberately keep the original behavior (no dual-graph) to
-        # avoid silently changing the CUDA decode path for DSA models (e.g.
-        # DeepSeek-V3.2). CUDA can opt in later once validated there.
-        self.dsa_dual_graph = False
-        self.dsa_index_topk: Optional[int] = None
-        from sglang.srt.configs.model_config import (
-            get_dsa_index_topk,
-            is_deepseek_dsa,
-        )
-
-        hf_config = model_runner.model_config.hf_config
-        if is_hip() and is_deepseek_dsa(hf_config):
-            self.dsa_index_topk = get_dsa_index_topk(hf_config)
-            self.dsa_dual_graph = True
-            logger.info(
-                "[dense-decode] DSA dual-graph enabled: capturing "
-                "dense (k-only) + sparse (full indexer) decode graphs; "
-                "dispatch on max_kv_len vs index_topk=%d.",
-                self.dsa_index_topk,
-            )
-
         self.attn_tp_size = get_parallel().attn_tp_size
         self.attn_tp_rank = get_parallel().attn_tp_rank
         # True if a DSACPLayerCommunicator-style prefill-CP flavor is active
@@ -327,6 +296,10 @@ def __init__(
         elif self.is_dllm:
             self.capture_forward_mode = ForwardMode.DLLM_EXTEND
 
+        self.attention_graph_variants = create_attention_graph_variants(
+            model_runner.model_config.hf_config, self.capture_forward_mode
+        )
+
         # --- bucket sizes ---------------------------------------------
         self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(
             model_runner, self.captured_req_width
@@ -562,37 +535,21 @@ def _cache_loc_dtype(self):
         return torch.int64
 
     def _make_graph_key(
-        self, size, stream_idx=None, variant_label=None, dsa_variant=None
+        self, size, stream_idx=None, variant_label=None, attention_variant=None
     ):
         return ShapeKey(
             size=size,
             stream_idx=stream_idx,
             variant_label=variant_label,
-            dsa_variant=dsa_variant,
+            attention_variant=attention_variant,
         )
 
     def _capture_graph_size(self, *, bs: int, num_tokens: int) -> int:
         return num_tokens if self.ragged_verify_mode else bs
 
-    def _resolve_dsa_variant(self, forward_batch: ForwardBatch) -> Optional[str]:
-        """Host dispatch: pick which pre-captured DSA decode graph to replay
-        from the batch-max kv_len. If any request has kv_len > index_topk
-        the dense (k-only) graph would be wrong for it, so the whole batch uses
-        the sparse (full indexer) graph. Returns None when dual-graph is off."""
-        if not getattr(self, "dsa_dual_graph", False):
-            return None
-        seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
-        if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0:
-            # Host-side mirror (maintained incrementally for plain decode) — no
-            # d2h sync needed.
-            max_kv_len = int(seq_lens_cpu.max().item())
-        elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0:
-            # Fallback: a single scalar reduction d2h (cheap, per-step).
-            max_kv_len = int(forward_batch.seq_lens.max().item())
-        else:
-            # No length info: be safe and use the correct-for-all sparse graph.
-            return "sparse"
-        return "dense" if max_kv_len <= self.dsa_index_topk else "sparse"
+    def _resolve_attention_variant(self, forward_batch: ForwardBatch) -> Optional[str]:
+        variants = getattr(self, "attention_graph_variants", None)
+        return variants.select(forward_batch) if variants is not None else None
 
     def _resolve_lora_variant(self, forward_batch: ForwardBatch):
         if not getattr(self, "record_nolora_graph", False):
@@ -705,8 +662,8 @@ def can_run_graph(self, forward_batch: ForwardBatch):
             cuda_graph_bs,
             stream_idx=get_current_stream_idx() if self.enable_pdmux else None,
             variant_label=self._resolve_lora_variant(forward_batch),
-            dsa_variant=(
-                self._resolve_dsa_variant(forward_batch)
+            attention_variant=(
+                self._resolve_attention_variant(forward_batch)
                 if self.disable_padding
                 else None
             ),
@@ -1112,16 +1069,10 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
             if getattr(self, "record_nolora_graph", False)
             else [(None, None)]
         )
-        # DSA: capture a dense (k-only) and a sparse (full indexer) graph
-        # per bs bucket. Order: dense first so its (smaller) capture-time peak
-        # runs while the shared pool is fresh; sparse's peak subsumes it.
-        # getattr default: subclasses like EAGLEDraftCudaGraphRunner reuse this
-        # capture() but don't run DecodeCudaGraphRunner.__init__ (so they never
-        # set dsa_dual_graph) and override capture_one_shape with a signature that
-        # has no dsa_variant. Default to no dual-graph and, for the None variant,
-        # call capture_one_shape without the extra arg so those overrides work.
-        dsa_variants = (
-            ["dense", "sparse"] if getattr(self, "dsa_dual_graph", False) else [None]
+        # Draft runners can reuse capture without initializing attention variants.
+        variants = getattr(self, "attention_graph_variants", None)
+        attention_variants = (
+            variants.capture_labels if variants is not None else (None,)
         )
         for bs in capture_range:
             if get_parallel().tp_rank == 0:
@@ -1136,23 +1087,27 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
 
             for variant_label, _variant_has_lora in lora_variants:
                 _set_capture_lora_variant(variant_label)
-                for dsa_variant in dsa_variants:
-                    _set_capture_dsa_variant(dsa_variant)
+                for attention_variant in attention_variants:
+                    _set_capture_attention_variant(attention_variant)
                     with torch_compile_decoration.patch_model(
                         self.model_runner.model,
                         bs in self.compile_bs,
                         num_tokens=bs * self.captured_req_width,
                         tp_group=self.model_runner.tp_group,
                     ) as forward:
-                        if dsa_variant is None:
+                        if attention_variant is None:
                             self.capture_one_shape(
                                 bs, forward, stream_idx, variant_label
                             )
                         else:
                             self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label, dsa_variant
+                                bs,
+                                forward,
+                                stream_idx,
+                                variant_label,
+                                attention_variant,
                             )
-        _set_capture_dsa_variant(None)
+        _set_capture_attention_variant(None)
 
     def capture_one_shape(
         self,
@@ -1160,7 +1115,7 @@ def capture_one_shape(
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
-        dsa_variant: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         num_tokens = size * self.captured_req_width
         bs = self._ragged_capture_slots(num_tokens) if self.ragged_verify_mode else size
@@ -1249,7 +1204,7 @@ def run_once():
                     self._capture_graph_size(bs=bs, num_tokens=num_tokens),
                     stream_idx,
                     variant_label,
-                    dsa_variant,
+                    attention_variant,
                 )
                 # Adaptive runners may own a different backend than model_runner.
                 post_warmup_hook = getattr(
@@ -1319,10 +1274,10 @@ def load_batch(
                     forward_batch.input_embeds
                 )
             variant_label = self._resolve_lora_variant(forward_batch)
-            dsa_variant = self._resolve_dsa_variant(forward_batch)
+            attention_variant = self._resolve_attention_variant(forward_batch)
             stream_idx = get_current_stream_idx() if self.enable_pdmux else None
             self._replay_graph_key = self._make_graph_key(
-                graph_size_key, stream_idx, variant_label, dsa_variant
+                graph_size_key, stream_idx, variant_label, attention_variant
             )
             return
 
@@ -1439,10 +1394,10 @@ def load_batch(
             self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs)
 
         variant_label = self._resolve_lora_variant(forward_batch)
-        dsa_variant = self._resolve_dsa_variant(forward_batch)
+        attention_variant = self._resolve_attention_variant(forward_batch)
         stream_idx = get_current_stream_idx() if self.enable_pdmux else None
         self._replay_graph_key = self._make_graph_key(
-            graph_size_key, stream_idx, variant_label, dsa_variant
+            graph_size_key, stream_idx, variant_label, attention_variant
         )
 
     def _ragged_graph_num_tokens(self, total_verify_tokens: int) -> int:
diff --git a/python/sglang/srt/model_executor/runner/shape_key.py b/python/sglang/srt/model_executor/runner/shape_key.py
index f04d07b1e68a..93e850891c4f 100644
--- a/python/sglang/srt/model_executor/runner/shape_key.py
+++ b/python/sglang/srt/model_executor/runner/shape_key.py
@@ -30,12 +30,12 @@ class ShapeKey:
     variant_label: optional execution variant (for example, "lora",
         "nolora", or "chunked_prefix"), or None for runners that don't
         record per-variant graphs.
-    dsa_variant: DSA decode dual-graph variant ("dense" / "sparse"), or None
-        when DSA dual-graph capture is not enabled. Composes with variant_label
-        so LoRA and DSA variants can be captured independently.
+    attention_variant: attention execution variant, or None for the default path.
+        Composes with variant_label so attention and LoRA variants can be
+        captured independently.
     """
 
     size: int
     stream_idx: Optional[int] = None
     variant_label: Optional[str] = None
-    dsa_variant: Optional[str] = None
+    attention_variant: Optional[str] = None
diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py
index f936cd0c9ddf..03c19835ec18 100644
--- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py
+++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py
@@ -37,12 +37,8 @@
 # None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant.
 _capture_lora_variant: Optional[str] = None
 
-# When capturing dual DSA decode graphs (dense/sparse), tracks which variant is
-# being captured. Read by the DSA indexer's capture-time skip-logits branch to
-# force k-only ("dense") vs full indexer ("sparse").
-# None = not dual-capturing; the indexer then bakes in the full-indexer path,
-# which is correct for any kv_len.
-_capture_dsa_variant: Optional[str] = None
+# Attention execution variant active through metadata preparation and capture.
+_capture_attention_variant: Optional[str] = None
 
 
 def get_is_capture_mode() -> bool:
@@ -72,15 +68,14 @@ def _set_capture_lora_variant(variant: Optional[str]) -> None:
     _capture_lora_variant = variant
 
 
-def get_capture_dsa_variant() -> Optional[str]:
-    """Return the DSA decode variant being captured ("dense"/"sparse"), or None
-    when dual-variant capture is not active."""
-    return _capture_dsa_variant
+def get_capture_attention_variant() -> Optional[str]:
+    """Return the attention execution variant being captured, or None."""
+    return _capture_attention_variant
 
 
-def _set_capture_dsa_variant(variant: Optional[str]) -> None:
-    global _capture_dsa_variant
-    _capture_dsa_variant = variant
+def _set_capture_attention_variant(variant: Optional[str]) -> None:
+    global _capture_attention_variant
+    _capture_attention_variant = variant
 
 
 @contextmanager
diff --git a/python/sglang/srt/speculative/uno_cuda_graph_runner.py b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
index 921353f89736..3cfba6d54d85 100644
--- a/python/sglang/srt/speculative/uno_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
@@ -128,7 +128,7 @@ def capture_one_shape(
         forward,
         stream_idx=None,
         variant_label=None,
-        dsa_variant=None,
+        attention_variant=None,
     ):
         """capture one CUDA graph with/out UNO LoRA."""
         if self._tree_draft_mode:
@@ -139,7 +139,7 @@ def capture_one_shape(
                     forward,
                     stream_idx,
                     None,
-                    dsa_variant,
+                    attention_variant,
                 )
             finally:
                 self._lora_state.reset()
@@ -152,7 +152,7 @@ def capture_one_shape(
                     forward,
                     stream_idx,
                     None,
-                    dsa_variant,
+                    attention_variant,
                 )
             finally:
                 self.model_runner.lora_manager.reset_lora_batch()
@@ -169,7 +169,7 @@ def capture_one_shape(
             forward,
             stream_idx,
             variant_label,
-            dsa_variant,
+            attention_variant,
         )
         self._lora_state.reset()
 
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index 03990cd68283..70c319cbb0cd 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -1,4 +1,4 @@
-"""Unit tests for ``DecodeCudaGraphRunner`` capture-phase profiling — CPU-only.
+"""Unit tests for ``DecodeCudaGraphRunner`` capture profiling and variant dispatch — CPU-only.
 
 Two capture-trace modes plus their precedence:
 
@@ -25,13 +25,16 @@
 import os
 import tempfile
 import unittest
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest import mock
 
+from sglang.srt.layers.attention.dsa.indexer_capture import DsaGraphVariants
 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.model_executor.runner_utils import capture_mode
 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
@@ -293,5 +296,60 @@ def test_dir_helper_uses_profiler_dir(self):
                 )
 
 
+class TestCaptureVariants(CustomTestCase):
+    def test_capture_order_and_draft_runner_without_variants(self):
+        module = "sglang.srt.model_executor.runner.decode_cuda_graph_runner"
+        runner = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner)
+        runner.attention_graph_variants = DsaGraphVariants(16)
+        runner.record_nolora_graph = True
+        runner.capture_bs = [1, 8]
+        runner.compile_bs = []
+        runner.captured_req_width = 1
+        runner.model_runner = SimpleNamespace(
+            device="cpu", gpu_id=0, model=object(), tp_group=object()
+        )
+        captured = []
+
+        def capture(size, forward, stream_idx, lora, attention=None):
+            self.assertEqual(capture_mode.get_capture_attention_variant(), attention)
+            self.assertEqual(capture_mode.get_capture_lora_variant(), lora)
+            captured.append((size, lora, attention))
+
+        with (
+            mock.patch(f"{module}.get_available_gpu_memory", return_value=1),
+            mock.patch(
+                f"{module}.get_parallel", return_value=SimpleNamespace(tp_rank=1)
+            ),
+            mock.patch(
+                f"{module}.torch_compile_decoration.patch_model",
+                side_effect=lambda *a, **kw: nullcontext(None),
+            ),
+        ):
+            try:
+                runner.capture_one_shape = capture
+                runner._capture_one_stream()
+                self.assertEqual(
+                    captured,
+                    [
+                        (size, lora, attention)
+                        for size in (8, 1)
+                        for lora in ("lora", "nolora")
+                        for attention in ("dense", "sparse")
+                    ],
+                )
+                self.assertIsNone(capture_mode.get_capture_attention_variant())
+                del runner.attention_graph_variants
+                runner.record_nolora_graph = False
+                captured.clear()
+                runner.capture_one_shape = lambda size, forward, stream, lora: capture(
+                    size, forward, stream, lora
+                )
+                runner._capture_one_stream()
+                self.assertEqual(captured, [(8, None, None), (1, None, None)])
+            finally:
+                capture_mode._set_capture_lora_variant(None)
+                capture_mode._set_capture_attention_variant(None)
+
+
 if __name__ == "__main__":
     unittest.main()

From ccd43adccf7a900e0ac38c7b41158fdadf9a9f87 Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:08:01 -0700
Subject: [PATCH 2/8] Trim graph variant comments and internal docstrings

---
 .../srt/layers/attention/dsa/dsa_indexer.py   | 40 ++++---------------
 .../layers/attention/dsa/indexer_capture.py   |  3 +-
 .../srt/model_executor/runner/shape_key.py    | 19 ++-------
 .../runner_utils/capture_mode.py              |  1 -
 .../srt/speculative/uno_cuda_graph_runner.py  |  1 -
 .../runner/test_decode_cuda_graph_runner.py   | 24 -----------
 6 files changed, 12 insertions(+), 76 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index ef5d079c5327..7e674fc656d4 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -398,20 +398,11 @@ def _maybe_rotate(self, x: torch.Tensor) -> torch.Tensor:
         return x if self.use_dsa_indexer_fusion else rotate_activation(x)
 
     def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
-        # When kv_len <= index_topk the top-k selects ALL valid positions, so the
-        # indexer's logits GEMM + paged_mqa_logits + top-k are wasted work: a plain
-        # topk_transform(dummy_logits) already yields the correct "select-all"
-        # (physical page-slot) indices. Skipping the logits path is safe here.
-        #
-        # Prefill/extend: original fast path, all platforms.
-        # Decode: new here, and ROCm-only for now (see the _is_hip gate below).
-        # Under a captured decode cuda graph the chosen branch is frozen at
-        # capture time and would replay incorrectly for kv_len > index_topk, so
-        # the decode skip is not decided per-step during capture; it is driven by
-        # which graph variant is being captured instead.
+        # topk_transform selects every valid page slot when kv_len <= index_topk;
+        # logits are unnecessary in that case.
         fb = forward_batch
 
-        # Prefill/extend: original per-step gate (host sync on seq_lens_cpu is fine).
+        # Prefill/extend.
         if fb.forward_mode.is_extend_without_speculative():
             if fb.seq_lens_cpu is None or fb.seq_lens_cpu.numel() == 0:
                 return False
@@ -419,26 +410,12 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
 
         # Decode/idle.
         if fb.forward_mode.is_decode_or_idle():
-            # Decode k-only skip (both the captured dual-graph "dense" variant
-            # and the eager per-step skip below) is currently HIP-only. On CUDA
-            # this common code keeps the original behavior (decode never skips
-            # the indexer, i.e. always runs the full logits path) because the
-            # decode k-only path has not been validated on CUDA yet. Mirrors the
-            # is_hip() gate in create_attention_graph_variants, which
-            # already prevents the CUDA capture path from setting a "dense"
-            # variant.
+            # Match the HIP-only gate in create_attention_graph_variants.
             if not _is_hip:
                 return False
             if get_is_capture_mode():
-                # Under a captured decode cuda graph the taken branch is frozen at
-                # capture time, so we must NOT branch on a runtime seq_len (also a
-                # host sync would break capture). The chosen branch is instead
-                # driven by which graph variant is being captured.
-                #
-                # The decode runner captures a "dense" (k-only) and a "sparse"
-                # (full indexer) graph per bs bucket and dispatches on max_kv_len
-                # at replay. The capture-variant signal tells us which one to
-                # bake in.
+                # Graph replay freezes this branch; use the capture variant,
+                # not capture-time sequence lengths.
                 from sglang.srt.model_executor.runner_utils.capture_mode import (
                     get_capture_attention_variant,
                 )
@@ -449,11 +426,8 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
                 if variant == "sparse":
                     return False
 
-                # No dual-variant capture signal: default to the correct-for-all
-                # full-indexer (sparse) path.
+                # No variant means the full indexer path for any context length.
                 return False
-            # Eager decode: safe to check per-step (host sync OK); correct for both
-            # kv_len<=index_topk (k-only) and kv_len>index_topk (falls through).
             if fb.seq_lens_cpu is not None and fb.seq_lens_cpu.numel() > 0:
                 max_kv_len = int(fb.seq_lens_cpu.max().item())
             elif fb.seq_lens is not None and fb.seq_lens.numel() > 0:
diff --git a/python/sglang/srt/layers/attention/dsa/indexer_capture.py b/python/sglang/srt/layers/attention/dsa/indexer_capture.py
index 0f8a1a99eefb..72e63af2e0bc 100644
--- a/python/sglang/srt/layers/attention/dsa/indexer_capture.py
+++ b/python/sglang/srt/layers/attention/dsa/indexer_capture.py
@@ -27,8 +27,7 @@ def __post_init__(self):
     def select(self, forward_batch: ForwardBatch) -> str:
         seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
         if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0:
-            # Host-side mirror (maintained incrementally for plain decode) — no
-            # d2h sync needed.
+            # Plain decode maintains this host mirror without a D2H sync.
             max_kv_len = int(seq_lens_cpu.max().item())
         elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0:
             # Fallback: a single scalar reduction d2h (cheap, per-step).
diff --git a/python/sglang/srt/model_executor/runner/shape_key.py b/python/sglang/srt/model_executor/runner/shape_key.py
index 93e850891c4f..7691b7e0707e 100644
--- a/python/sglang/srt/model_executor/runner/shape_key.py
+++ b/python/sglang/srt/model_executor/runner/shape_key.py
@@ -11,7 +11,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""ShapeKey — typed identifier for one captured CUDA-graph shape."""
 
 from __future__ import annotations
 
@@ -21,21 +20,11 @@
 
 @dataclass(frozen=True)
 class ShapeKey:
-    """Identifies one captured CUDA-graph shape across all runners.
-
-    size: the per-phase capture size — what the runner iterates over.
-        - prefill: num_tokens
-        - decode:  bs
-    stream_idx:   pdmux stream index, or None for single-stream runners.
-    variant_label: optional execution variant (for example, "lora",
-        "nolora", or "chunked_prefix"), or None for runners that don't
-        record per-variant graphs.
-    attention_variant: attention execution variant, or None for the default path.
-        Composes with variant_label so attention and LoRA variants can be
-        captured independently.
-    """
-
+    # Tokens for prefill/ragged verify; requests for ordinary decode.
     size: int
+    # PDMux stream, or None for a single stream.
     stream_idx: Optional[int] = None
+    # LoRA or prefill-prefix variant; None selects the default.
     variant_label: Optional[str] = None
+    # Independent attention variant; None selects the default.
     attention_variant: Optional[str] = None
diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py
index 03c19835ec18..a2195b9d00cc 100644
--- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py
+++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py
@@ -69,7 +69,6 @@ def _set_capture_lora_variant(variant: Optional[str]) -> None:
 
 
 def get_capture_attention_variant() -> Optional[str]:
-    """Return the attention execution variant being captured, or None."""
     return _capture_attention_variant
 
 
diff --git a/python/sglang/srt/speculative/uno_cuda_graph_runner.py b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
index 3cfba6d54d85..159f5bc0a12d 100644
--- a/python/sglang/srt/speculative/uno_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
@@ -130,7 +130,6 @@ def capture_one_shape(
         variant_label=None,
         attention_variant=None,
     ):
-        """capture one CUDA graph with/out UNO LoRA."""
         if self._tree_draft_mode:
             self._lora_state.capture_draft(size)
             try:
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index 70c319cbb0cd..af72ab98f332 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -1,27 +1,3 @@
-"""Unit tests for ``DecodeCudaGraphRunner`` capture profiling and variant dispatch — 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 9dbd239883d1df65e382b740028fece716fea18c Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:22:13 -0700
Subject: [PATCH 3/8] Name DSA graph variants module by its role

---
 .../attention/dsa/{indexer_capture.py => graph_variants.py}     | 0
 python/sglang/srt/layers/attention/graph_variants.py            | 2 +-
 .../unit/model_executor/runner/test_decode_cuda_graph_runner.py | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename python/sglang/srt/layers/attention/dsa/{indexer_capture.py => graph_variants.py} (100%)

diff --git a/python/sglang/srt/layers/attention/dsa/indexer_capture.py b/python/sglang/srt/layers/attention/dsa/graph_variants.py
similarity index 100%
rename from python/sglang/srt/layers/attention/dsa/indexer_capture.py
rename to python/sglang/srt/layers/attention/dsa/graph_variants.py
diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py
index d08e4f6c8f20..88482950bbba 100644
--- a/python/sglang/srt/layers/attention/graph_variants.py
+++ b/python/sglang/srt/layers/attention/graph_variants.py
@@ -19,7 +19,7 @@ def create_attention_graph_variants(
     hf_config, forward_mode: ForwardMode
 ) -> Optional[AttentionGraphVariants]:
     from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
-    from sglang.srt.layers.attention.dsa.indexer_capture import DsaGraphVariants
+    from sglang.srt.layers.attention.dsa.graph_variants import DsaGraphVariants
     from sglang.srt.utils import is_hip
 
     if is_hip() and is_deepseek_dsa(hf_config):
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index af72ab98f332..20b6719d9737 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -5,7 +5,7 @@
 from types import SimpleNamespace
 from unittest import mock
 
-from sglang.srt.layers.attention.dsa.indexer_capture import DsaGraphVariants
+from sglang.srt.layers.attention.dsa.graph_variants import DsaGraphVariants
 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 3ac47910a46dd932eaef25f3fbefe4fb875dbdd5 Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:30:17 -0700
Subject: [PATCH 4/8] Consolidate graph variants and declare runner defaults

---
 .../layers/attention/dsa/graph_variants.py    | 38 -------------------
 .../srt/layers/attention/graph_variants.py    | 36 +++++++++++++++++-
 .../runner/decode_cuda_graph_runner.py        | 16 +++++---
 .../runner/test_decode_cuda_graph_runner.py   |  2 +-
 4 files changed, 46 insertions(+), 46 deletions(-)
 delete mode 100644 python/sglang/srt/layers/attention/dsa/graph_variants.py

diff --git a/python/sglang/srt/layers/attention/dsa/graph_variants.py b/python/sglang/srt/layers/attention/dsa/graph_variants.py
deleted file mode 100644
index 72e63af2e0bc..000000000000
--- a/python/sglang/srt/layers/attention/dsa/graph_variants.py
+++ /dev/null
@@ -1,38 +0,0 @@
-from __future__ import annotations
-
-import logging
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar
-
-if TYPE_CHECKING:
-    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass(frozen=True)
-class DsaGraphVariants:
-    index_topk: int
-    # Dense comes first: the sparse capture peak subsumes its shared-pool storage.
-    capture_labels: ClassVar[tuple[str, ...]] = ("dense", "sparse")
-
-    def __post_init__(self):
-        logger.info(
-            "[dense-decode] DSA dual-graph enabled: capturing "
-            "dense (k-only) + sparse (full indexer) decode graphs; "
-            "dispatch on max_kv_len vs index_topk=%d.",
-            self.index_topk,
-        )
-
-    def select(self, forward_batch: ForwardBatch) -> str:
-        seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
-        if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0:
-            # Plain decode maintains this host mirror without a D2H sync.
-            max_kv_len = int(seq_lens_cpu.max().item())
-        elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0:
-            # Fallback: a single scalar reduction d2h (cheap, per-step).
-            max_kv_len = int(forward_batch.seq_lens.max().item())
-        else:
-            # No length info: be safe and use the correct-for-all sparse graph.
-            return "sparse"
-        return "dense" if max_kv_len <= self.index_topk else "sparse"
diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py
index 88482950bbba..e949fdf2e7e7 100644
--- a/python/sglang/srt/layers/attention/graph_variants.py
+++ b/python/sglang/srt/layers/attention/graph_variants.py
@@ -1,11 +1,16 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, Protocol
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar, Optional, Protocol
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 
 
+logger = logging.getLogger(__name__)
+
+
 class AttentionGraphVariants(Protocol):
     # Capture order is significant when variants share a graph memory pool.
     capture_labels: tuple[str, ...]
@@ -15,11 +20,38 @@ def select(self, forward_batch: ForwardBatch) -> str:
         ...
 
 
+@dataclass(frozen=True)
+class DsaGraphVariants:
+    index_topk: int
+    # Dense comes first: the sparse capture peak subsumes its shared-pool storage.
+    capture_labels: ClassVar[tuple[str, ...]] = ("dense", "sparse")
+
+    def __post_init__(self):
+        logger.info(
+            "[dense-decode] DSA dual-graph enabled: capturing "
+            "dense (k-only) + sparse (full indexer) decode graphs; "
+            "dispatch on max_kv_len vs index_topk=%d.",
+            self.index_topk,
+        )
+
+    def select(self, forward_batch: ForwardBatch) -> str:
+        seq_lens_cpu = forward_batch.seq_lens_cpu
+        if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0:
+            # Plain decode maintains this host mirror without a D2H sync.
+            max_kv_len = int(seq_lens_cpu.max().item())
+        elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0:
+            # Fallback: a single scalar reduction d2h (cheap, per-step).
+            max_kv_len = int(forward_batch.seq_lens.max().item())
+        else:
+            # No length info: be safe and use the correct-for-all sparse graph.
+            return "sparse"
+        return "dense" if max_kv_len <= self.index_topk else "sparse"
+
+
 def create_attention_graph_variants(
     hf_config, forward_mode: ForwardMode
 ) -> Optional[AttentionGraphVariants]:
     from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
-    from sglang.srt.layers.attention.dsa.graph_variants import DsaGraphVariants
     from sglang.srt.utils import is_hip
 
     if is_hip() and is_deepseek_dsa(hf_config):
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 3b60412e6655..825b214a986e 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
@@ -49,7 +49,10 @@
     SharedReadEnds,
 )
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
-from sglang.srt.layers.attention.graph_variants import create_attention_graph_variants
+from sglang.srt.layers.attention.graph_variants import (
+    AttentionGraphVariants,
+    create_attention_graph_variants,
+)
 from sglang.srt.layers.cp.utils import is_mla_cp_enabled
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
@@ -215,6 +218,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
     pluggable self.backend that handles the actual capture/replay.
     """
 
+    attention_graph_variants: Optional[AttentionGraphVariants] = None
+    record_nolora_graph: bool = False
+
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -548,11 +554,11 @@ def _capture_graph_size(self, *, bs: int, num_tokens: int) -> int:
         return num_tokens if self.ragged_verify_mode else bs
 
     def _resolve_attention_variant(self, forward_batch: ForwardBatch) -> Optional[str]:
-        variants = getattr(self, "attention_graph_variants", None)
+        variants = self.attention_graph_variants
         return variants.select(forward_batch) if variants is not None else None
 
     def _resolve_lora_variant(self, forward_batch: ForwardBatch):
-        if not getattr(self, "record_nolora_graph", False):
+        if not self.record_nolora_graph:
             return None
         if forward_batch.lora_ids is not None and any(
             uid is not None for uid in forward_batch.lora_ids
@@ -1066,11 +1072,11 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
         )
         lora_variants = (
             [("lora", True), ("nolora", False)]
-            if getattr(self, "record_nolora_graph", False)
+            if self.record_nolora_graph
             else [(None, None)]
         )
         # Draft runners can reuse capture without initializing attention variants.
-        variants = getattr(self, "attention_graph_variants", None)
+        variants = self.attention_graph_variants
         attention_variants = (
             variants.capture_labels if variants is not None else (None,)
         )
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index 20b6719d9737..17ce50ae106f 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -5,7 +5,7 @@
 from types import SimpleNamespace
 from unittest import mock
 
-from sglang.srt.layers.attention.dsa.graph_variants import DsaGraphVariants
+from sglang.srt.layers.attention.graph_variants import DsaGraphVariants
 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 80a88fe32bc326fdefd73f7de1de055ba129c4b0 Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:38:12 -0700
Subject: [PATCH 5/8] Initialize graph variant state on runner instances

---
 .../runner/decode_cuda_graph_runner.py               | 12 ++++++------
 .../srt/speculative/eagle_draft_cuda_graph_runner.py |  1 +
 .../eagle_draft_extend_cuda_graph_runner.py          |  1 +
 .../speculative/frozen_kv_mtp_cuda_graph_runner.py   |  1 +
 ...lti_layer_eagle_draft_extend_cuda_graph_runner.py |  1 +
 .../sglang/srt/speculative/uno_cuda_graph_runner.py  |  5 +----
 .../runner/test_decode_cuda_graph_runner.py          |  2 +-
 7 files changed, 12 insertions(+), 11 deletions(-)

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 825b214a986e..5725fdc99f96 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
@@ -218,9 +218,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
     pluggable self.backend that handles the actual capture/replay.
     """
 
-    attention_graph_variants: Optional[AttentionGraphVariants] = None
-    record_nolora_graph: bool = False
-
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -228,8 +225,10 @@ def __init__(
         attn_backend=None,
         speculative_num_steps: Optional[int] = None,
         speculative_num_draft_tokens: Optional[int] = None,
+        record_nolora_graph: bool = False,
     ):
         super().__init__(model_runner)
+        self.record_nolora_graph = record_nolora_graph
 
         # In-graph metadata prep: shared buffers -> in-graph private data
         self.in_graph_metadata_prep_done: Optional[torch.cuda.Event] = None
@@ -302,8 +301,10 @@ def __init__(
         elif self.is_dllm:
             self.capture_forward_mode = ForwardMode.DLLM_EXTEND
 
-        self.attention_graph_variants = create_attention_graph_variants(
-            model_runner.model_config.hf_config, self.capture_forward_mode
+        self.attention_graph_variants: Optional[AttentionGraphVariants] = (
+            create_attention_graph_variants(
+                model_runner.model_config.hf_config, self.capture_forward_mode
+            )
         )
 
         # --- bucket sizes ---------------------------------------------
@@ -1075,7 +1076,6 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
             if self.record_nolora_graph
             else [(None, None)]
         )
-        # Draft runners can reuse capture without initializing attention variants.
         variants = self.attention_graph_variants
         attention_variants = (
             variants.capture_labels if variants is not None else (None,)
diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
index cb603f1ff094..46960d87a5db 100644
--- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
@@ -140,6 +140,7 @@ def __init__(
         self.compile_bs = []  # disables patch_model torch.compile wrapping
         self.enable_pdmux = False
         self.record_nolora_graph = False
+        self.attention_graph_variants = None
         self.is_dllm = False
 
         self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
index 529867b00549..03b28848e8bd 100644
--- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
@@ -135,6 +135,7 @@ def __init__(
         self.compile_bs = []
         self.enable_pdmux = False
         self.record_nolora_graph = False
+        self.attention_graph_variants = None
         self.is_dllm = False
 
         self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
index aa9ff66129cc..f0d5c26bdc4d 100644
--- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
@@ -111,6 +111,7 @@ def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPDraftWorker):
         self.compile_bs = []
         self.enable_pdmux = False
         self.record_nolora_graph = False
+        self.attention_graph_variants = None
         self.is_dllm = False
 
         self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
index dc0d0230f829..7547e975604d 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
@@ -180,6 +180,7 @@ def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int):
         # Disable parent paths that don't apply.
         self.compile_bs = []
         self.record_nolora_graph = False
+        self.attention_graph_variants = None
         self.is_dllm = False
 
         self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
diff --git a/python/sglang/srt/speculative/uno_cuda_graph_runner.py b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
index 159f5bc0a12d..bcd0f57fd546 100644
--- a/python/sglang/srt/speculative/uno_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/uno_cuda_graph_runner.py
@@ -32,7 +32,6 @@ def __init__(
         self._tree_draft_mode = tree_draft_width is not None
 
         if self._tree_draft_mode:
-            self.record_nolora_graph = False
             self._capture_spec_input_type = SpecInputType.UNO_DRAFT
             self._lora_state = UnoCudaGraphLoRAState(
                 model_runner.lora_manager,
@@ -52,7 +51,6 @@ def __init__(
         if self._tree_mode:
             # Capture exactly one base-model target graph.  The internal UNO
             # adapter is active only in the rejected F-wide draft phase.
-            self.record_nolora_graph = False
             model_runner.lora_manager.reset_lora_batch()
             kwargs.update(
                 attn_backend=model_runner.attn_backend,
@@ -64,7 +62,6 @@ def __init__(
             return
 
         forward_width = model_runner.decode_num_tokens_per_req()
-        self.record_nolora_graph = forward_width > 1
         self._capture_spec_input_type = SpecInputType.UNO_VERIFY
         self._lora_state = UnoCudaGraphLoRAState(
             model_runner.lora_manager,
@@ -72,7 +69,7 @@ def __init__(
             forward_width,
         )
         model_runner.lora_manager.reset_lora_batch()
-        super().__init__(model_runner, **kwargs)
+        super().__init__(model_runner, record_nolora_graph=forward_width > 1, **kwargs)
 
     def capture_prepare(self, size, stream_idx=None, num_tokens=None):
         forward_batch, attn_backend, pp_proxy_tensors = super().capture_prepare(
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index 17ce50ae106f..cee700635e21 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -314,7 +314,7 @@ def capture(size, forward, stream_idx, lora, attention=None):
                     ],
                 )
                 self.assertIsNone(capture_mode.get_capture_attention_variant())
-                del runner.attention_graph_variants
+                runner.attention_graph_variants = None
                 runner.record_nolora_graph = False
                 captured.clear()
                 runner.capture_one_shape = lambda size, forward, stream, lora: capture(

From 1c75d5c0ce9fb02b17836b6a6f8775bd7d7975d7 Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:47:27 -0700
Subject: [PATCH 6/8] Trim graph variant API and refactor test scope

---
 .../srt/layers/attention/graph_variants.py    |  8 +-
 .../runner/decode_cuda_graph_runner.py        |  4 +-
 .../runner/test_decode_cuda_graph_runner.py   | 82 ++++++-------------
 3 files changed, 28 insertions(+), 66 deletions(-)

diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py
index e949fdf2e7e7..de60599b683e 100644
--- a/python/sglang/srt/layers/attention/graph_variants.py
+++ b/python/sglang/srt/layers/attention/graph_variants.py
@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, ClassVar, Optional, Protocol
 
 if TYPE_CHECKING:
-    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 
 logger = logging.getLogger(__name__)
@@ -16,7 +16,7 @@ class AttentionGraphVariants(Protocol):
     capture_labels: tuple[str, ...]
 
     def select(self, forward_batch: ForwardBatch) -> str:
-        """Select one of capture_labels from host-side batch metadata."""
+        """Select one of capture_labels for the batch."""
         ...
 
 
@@ -48,9 +48,7 @@ def select(self, forward_batch: ForwardBatch) -> str:
         return "dense" if max_kv_len <= self.index_topk else "sparse"
 
 
-def create_attention_graph_variants(
-    hf_config, forward_mode: ForwardMode
-) -> Optional[AttentionGraphVariants]:
+def create_attention_graph_variants(hf_config) -> Optional[AttentionGraphVariants]:
     from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
     from sglang.srt.utils import is_hip
 
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 5725fdc99f96..880837bf91b4 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
@@ -302,9 +302,7 @@ def __init__(
             self.capture_forward_mode = ForwardMode.DLLM_EXTEND
 
         self.attention_graph_variants: Optional[AttentionGraphVariants] = (
-            create_attention_graph_variants(
-                model_runner.model_config.hf_config, self.capture_forward_mode
-            )
+            create_attention_graph_variants(model_runner.model_config.hf_config)
         )
 
         # --- bucket sizes ---------------------------------------------
diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
index cee700635e21..03990cd68283 100644
--- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
+++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py
@@ -1,16 +1,37 @@
+"""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 contextlib import nullcontext
 from types import SimpleNamespace
 from unittest import mock
 
-from sglang.srt.layers.attention.graph_variants import DsaGraphVariants
 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.model_executor.runner_utils import capture_mode
 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
@@ -272,60 +293,5 @@ def test_dir_helper_uses_profiler_dir(self):
                 )
 
 
-class TestCaptureVariants(CustomTestCase):
-    def test_capture_order_and_draft_runner_without_variants(self):
-        module = "sglang.srt.model_executor.runner.decode_cuda_graph_runner"
-        runner = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner)
-        runner.attention_graph_variants = DsaGraphVariants(16)
-        runner.record_nolora_graph = True
-        runner.capture_bs = [1, 8]
-        runner.compile_bs = []
-        runner.captured_req_width = 1
-        runner.model_runner = SimpleNamespace(
-            device="cpu", gpu_id=0, model=object(), tp_group=object()
-        )
-        captured = []
-
-        def capture(size, forward, stream_idx, lora, attention=None):
-            self.assertEqual(capture_mode.get_capture_attention_variant(), attention)
-            self.assertEqual(capture_mode.get_capture_lora_variant(), lora)
-            captured.append((size, lora, attention))
-
-        with (
-            mock.patch(f"{module}.get_available_gpu_memory", return_value=1),
-            mock.patch(
-                f"{module}.get_parallel", return_value=SimpleNamespace(tp_rank=1)
-            ),
-            mock.patch(
-                f"{module}.torch_compile_decoration.patch_model",
-                side_effect=lambda *a, **kw: nullcontext(None),
-            ),
-        ):
-            try:
-                runner.capture_one_shape = capture
-                runner._capture_one_stream()
-                self.assertEqual(
-                    captured,
-                    [
-                        (size, lora, attention)
-                        for size in (8, 1)
-                        for lora in ("lora", "nolora")
-                        for attention in ("dense", "sparse")
-                    ],
-                )
-                self.assertIsNone(capture_mode.get_capture_attention_variant())
-                runner.attention_graph_variants = None
-                runner.record_nolora_graph = False
-                captured.clear()
-                runner.capture_one_shape = lambda size, forward, stream, lora: capture(
-                    size, forward, stream, lora
-                )
-                runner._capture_one_stream()
-                self.assertEqual(captured, [(8, None, None), (1, None, None)])
-            finally:
-                capture_mode._set_capture_lora_variant(None)
-                capture_mode._set_capture_attention_variant(None)
-
-
 if __name__ == "__main__":
     unittest.main()

From 788843c55747a6fbc05e15a7ef92446415ac188d Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:53:57 -0700
Subject: [PATCH 7/8] Unify attention graph capture contracts and DSA labels

---
 .../srt/layers/attention/dsa/dsa_indexer.py   | 12 ++------
 .../srt/layers/attention/graph_variants.py    | 28 ++++++++++---------
 .../runner/decode_cuda_graph_runner.py        | 19 +++++--------
 .../eagle_draft_cuda_graph_runner.py          |  1 +
 .../eagle_draft_extend_cuda_graph_runner.py   |  1 +
 .../frozen_kv_mtp_cuda_graph_runner.py        |  3 +-
 ...er_eagle_draft_extend_cuda_graph_runner.py |  1 +
 7 files changed, 30 insertions(+), 35 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index 7e674fc656d4..8d2a5e9a0cac 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -31,6 +31,7 @@
     is_dsa_enable_prefill_cp,
     is_graph_dsa_split_op_surface,
 )
+from sglang.srt.layers.attention.graph_variants import DSA_DENSE
 from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
     is_in_breakable_cuda_graph,
@@ -410,9 +411,6 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
 
         # Decode/idle.
         if fb.forward_mode.is_decode_or_idle():
-            # Match the HIP-only gate in create_attention_graph_variants.
-            if not _is_hip:
-                return False
             if get_is_capture_mode():
                 # Graph replay freezes this branch; use the capture variant,
                 # not capture-time sequence lengths.
@@ -420,13 +418,9 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
                     get_capture_attention_variant,
                 )
 
-                variant = get_capture_attention_variant()
-                if variant == "dense":
-                    return True
-                if variant == "sparse":
-                    return False
-
                 # No variant means the full indexer path for any context length.
+                return get_capture_attention_variant() == DSA_DENSE
+            if not _is_hip:
                 return False
             if fb.seq_lens_cpu is not None and fb.seq_lens_cpu.numel() > 0:
                 max_kv_len = int(fb.seq_lens_cpu.max().item())
diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py
index de60599b683e..14ba77c7b32a 100644
--- a/python/sglang/srt/layers/attention/graph_variants.py
+++ b/python/sglang/srt/layers/attention/graph_variants.py
@@ -10,10 +10,13 @@
 
 logger = logging.getLogger(__name__)
 
+DSA_DENSE = "dense"
+DSA_SPARSE = "sparse"
+
 
 class AttentionGraphVariants(Protocol):
     # Capture order is significant when variants share a graph memory pool.
-    capture_labels: tuple[str, ...]
+    capture_labels: ClassVar[tuple[str, ...]]
 
     def select(self, forward_batch: ForwardBatch) -> str:
         """Select one of capture_labels for the batch."""
@@ -24,15 +27,7 @@ def select(self, forward_batch: ForwardBatch) -> str:
 class DsaGraphVariants:
     index_topk: int
     # Dense comes first: the sparse capture peak subsumes its shared-pool storage.
-    capture_labels: ClassVar[tuple[str, ...]] = ("dense", "sparse")
-
-    def __post_init__(self):
-        logger.info(
-            "[dense-decode] DSA dual-graph enabled: capturing "
-            "dense (k-only) + sparse (full indexer) decode graphs; "
-            "dispatch on max_kv_len vs index_topk=%d.",
-            self.index_topk,
-        )
+    capture_labels: ClassVar[tuple[str, ...]] = (DSA_DENSE, DSA_SPARSE)
 
     def select(self, forward_batch: ForwardBatch) -> str:
         seq_lens_cpu = forward_batch.seq_lens_cpu
@@ -44,8 +39,8 @@ def select(self, forward_batch: ForwardBatch) -> str:
             max_kv_len = int(forward_batch.seq_lens.max().item())
         else:
             # No length info: be safe and use the correct-for-all sparse graph.
-            return "sparse"
-        return "dense" if max_kv_len <= self.index_topk else "sparse"
+            return DSA_SPARSE
+        return DSA_DENSE if max_kv_len <= self.index_topk else DSA_SPARSE
 
 
 def create_attention_graph_variants(hf_config) -> Optional[AttentionGraphVariants]:
@@ -53,5 +48,12 @@ def create_attention_graph_variants(hf_config) -> Optional[AttentionGraphVariant
     from sglang.srt.utils import is_hip
 
     if is_hip() and is_deepseek_dsa(hf_config):
-        return DsaGraphVariants(get_dsa_index_topk(hf_config))
+        index_topk = get_dsa_index_topk(hf_config)
+        logger.info(
+            "[dense-decode] DSA dual-graph enabled: capturing "
+            "dense (k-only) + sparse (full indexer) decode graphs; "
+            "dispatch on max_kv_len vs index_topk=%d.",
+            index_topk,
+        )
+        return DsaGraphVariants(index_topk)
     return None
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 880837bf91b4..d5160d43b925 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
@@ -1099,18 +1099,13 @@ def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
                         num_tokens=bs * self.captured_req_width,
                         tp_group=self.model_runner.tp_group,
                     ) as forward:
-                        if attention_variant is None:
-                            self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label
-                            )
-                        else:
-                            self.capture_one_shape(
-                                bs,
-                                forward,
-                                stream_idx,
-                                variant_label,
-                                attention_variant,
-                            )
+                        self.capture_one_shape(
+                            bs,
+                            forward,
+                            stream_idx,
+                            variant_label,
+                            attention_variant,
+                        )
         _set_capture_attention_variant(None)
 
     def capture_one_shape(
diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
index 46960d87a5db..0c6d80e3dfdd 100644
--- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
@@ -344,6 +344,7 @@ def capture_one_shape(
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         num_seqs = size  # EAGLE legacy name
         buffers = self.buffers
diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
index 03b28848e8bd..4aeb6c5a3cae 100644
--- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
@@ -344,6 +344,7 @@ def capture_one_shape(
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         bs = size
         buffers = self.buffers
diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
index f0d5c26bdc4d..c9a7c76cb574 100644
--- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
@@ -248,8 +248,9 @@ def capture_one_shape(
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
-        del forward, stream_idx, variant_label
+        del forward, stream_idx, variant_label, attention_variant
         buffers = self.buffers
         request_bs = size
         expanded_bs = request_bs * self.captured_req_width
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
index 7547e975604d..7a0a371feecf 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
@@ -432,6 +432,7 @@ def capture_one_shape(
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         bs = size
 

From 34959046e0aed44f27f0ba4c8197a5e3086885cd Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Thu, 10 Sep 2026 21:56:29 -0700
Subject: [PATCH 8/8] Explain the ROCm gate for eager DSA decode skip

---
 python/sglang/srt/layers/attention/dsa/dsa_indexer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index 8d2a5e9a0cac..26a457f62df9 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -420,6 +420,7 @@ def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
 
                 # No variant means the full indexer path for any context length.
                 return get_capture_attention_variant() == DSA_DENSE
+            # Eager k-only decode skip is validated on ROCm only.
             if not _is_hip:
                 return False
             if fb.seq_lens_cpu is not None and fb.seq_lens_cpu.numel() > 0: