[Refactor] Generalize attention graph variants in the decode runner (#38993)

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -31,6 +31,7 @@ from sglang.srt.layers.attention.dsa.utils import (
     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,
@@ -398,20 +399,11 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
         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
+        # topk_transform selects every valid page slot when kv_len <= index_topk;
-        # indexer's logits GEMM + paged_mqa_logits + top-k are wasted work: a plain
+        # logits are unnecessary in that case.
-        # 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.
         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,41 +411,18 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
 
         # 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 on dsa_dual_graph in decode_cuda_graph_runner, which
-            # already prevents the CUDA capture path from setting a "dense"
-            # variant.
-            if not _is_hip:
-                return False
             if get_is_capture_mode():
-                # Under a captured decode cuda graph the taken branch is frozen at
+                # Graph replay freezes this branch; use the capture variant,
-                # capture time, so we must NOT branch on a runtime seq_len (also a
+                # not capture-time sequence lengths.
-                # 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.
                 from sglang.srt.model_executor.runner_utils.capture_mode import (
-                    get_capture_dsa_variant,
+                    get_capture_attention_variant,
                 )
 
-                variant = get_capture_dsa_variant()
+                # No variant means the full indexer path for any context length.
-                if variant == "dense":
+                return get_capture_attention_variant() == DSA_DENSE
-                    return True
+            # Eager k-only decode skip is validated on ROCm only.
-                if variant == "sparse":
+            if not _is_hip:
                 return False
-
-                # No dual-variant capture signal: default to the correct-for-all
-                # full-indexer (sparse) path.
-                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:

python/sglang/srt/layers/attention/graph_variants.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+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
+
+
+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: ClassVar[tuple[str, ...]]
+
+    def select(self, forward_batch: ForwardBatch) -> str:
+        """Select one of capture_labels for the batch."""
+        ...
+
+
+@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, ...]] = (DSA_DENSE, DSA_SPARSE)
+
+    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 DSA_SPARSE
+        return DSA_DENSE if max_kv_len <= self.index_topk else DSA_SPARSE
+
+
+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
+
+    if is_hip() and is_deepseek_dsa(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

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

@@ -49,6 +49,10 @@ from sglang.srt.layers.attention.base_attn_backend import (
     SharedReadEnds,
 )
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+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,
@@ -91,7 +95,7 @@ from sglang.srt.model_executor.runner_utils.buffers import (
     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 +117,6 @@ from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout
 from sglang.srt.utils import (
     empty_context,
     get_available_gpu_memory,
-    is_hip,
     require_attn_tp_gather,
     require_mlp_tp_gather,
 )
@@ -222,8 +225,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         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
@@ -254,37 +259,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         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 +301,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         elif self.is_dllm:
             self.capture_forward_mode = ForwardMode.DLLM_EXTEND
 
+        self.attention_graph_variants: Optional[AttentionGraphVariants] = (
+            create_attention_graph_variants(model_runner.model_config.hf_config)
+        )
+
         # --- bucket sizes ---------------------------------------------
         self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(
             model_runner, self.captured_req_width
@@ -562,40 +540,24 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         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]:
+    def _resolve_attention_variant(self, forward_batch: ForwardBatch) -> Optional[str]:
-        """Host dispatch: pick which pre-captured DSA decode graph to replay
+        variants = self.attention_graph_variants
-        from the batch-max kv_len. If any request has kv_len > index_topk
+        return variants.select(forward_batch) if variants is not None else None
-        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_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
@@ -705,8 +667,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             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=(
+            attention_variant=(
-                self._resolve_dsa_variant(forward_batch)
+                self._resolve_attention_variant(forward_batch)
                 if self.disable_padding
                 else None
             ),
@@ -1109,19 +1071,12 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         )
         lora_variants = (
             [("lora", True), ("nolora", False)]
-            if getattr(self, "record_nolora_graph", False)
+            if self.record_nolora_graph
             else [(None, None)]
         )
-        # DSA: capture a dense (k-only) and a sparse (full indexer) graph
+        variants = self.attention_graph_variants
-        # per bs bucket. Order: dense first so its (smaller) capture-time peak
+        attention_variants = (
-        # runs while the shared pool is fresh; sparse's peak subsumes it.
+            variants.capture_labels if variants is not None else (None,)
-        # 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]
         )
         for bs in capture_range:
             if get_parallel().tp_rank == 0:
@@ -1136,23 +1091,22 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
 
             for variant_label, _variant_has_lora in lora_variants:
                 _set_capture_lora_variant(variant_label)
-                for dsa_variant in dsa_variants:
+                for attention_variant in attention_variants:
-                    _set_capture_dsa_variant(dsa_variant)
+                    _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:
                         self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label
+                            bs,
+                            forward,
+                            stream_idx,
+                            variant_label,
+                            attention_variant,
                         )
-                        else:
+        _set_capture_attention_variant(None)
-                            self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label, dsa_variant
-                            )
-        _set_capture_dsa_variant(None)
 
     def capture_one_shape(
         self,
@@ -1160,7 +1114,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         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 +1203,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     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 +1273,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     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 +1393,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             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:

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 @@ from typing import Optional
 
 @dataclass(frozen=True)
 class ShapeKey:
-    """Identifies one captured CUDA-graph shape across all runners.
+    # Tokens for prefill/ragged verify; requests for ordinary decode.
-
-    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.
-    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.
-    """
-
     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
-    dsa_variant: Optional[str] = None
+    # Independent attention variant; None selects the default.
+    attention_variant: Optional[str] = None

python/sglang/srt/model_executor/runner_utils/capture_mode.py

@@ -37,12 +37,8 @@ is_capture_mode = False
 # 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
+# Attention execution variant active through metadata preparation and capture.
-# being captured. Read by the DSA indexer's capture-time skip-logits branch to
+_capture_attention_variant: Optional[str] = None
-# 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
 
 
 def get_is_capture_mode() -> bool:
@@ -72,15 +68,13 @@ def _set_capture_lora_variant(variant: Optional[str]) -> None:
     _capture_lora_variant = variant
 
 
-def get_capture_dsa_variant() -> Optional[str]:
+def get_capture_attention_variant() -> Optional[str]:
-    """Return the DSA decode variant being captured ("dense"/"sparse"), or None
+    return _capture_attention_variant
-    when dual-variant capture is not active."""
-    return _capture_dsa_variant
 
 
-def _set_capture_dsa_variant(variant: Optional[str]) -> None:
+def _set_capture_attention_variant(variant: Optional[str]) -> None:
-    global _capture_dsa_variant
+    global _capture_attention_variant
-    _capture_dsa_variant = variant
+    _capture_attention_variant = variant
 
 
 @contextmanager

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -140,6 +140,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
         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()
@@ -343,6 +344,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
         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

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -135,6 +135,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         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()
@@ -343,6 +344,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         bs = size
         buffers = self.buffers

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -111,6 +111,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
         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()
@@ -247,8 +248,9 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
         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

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -180,6 +180,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         # 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()
@@ -431,6 +432,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         forward: Callable,
         stream_idx: Optional[int] = None,
         variant_label: Optional[str] = None,
+        attention_variant: Optional[str] = None,
     ):
         bs = size
 

python/sglang/srt/speculative/uno_cuda_graph_runner.py

@@ -32,7 +32,6 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
         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 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
         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 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
             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 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
             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(
@@ -128,9 +125,8 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
         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:
             self._lora_state.capture_draft(size)
             try:
@@ -139,7 +135,7 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
                     forward,
                     stream_idx,
                     None,
-                    dsa_variant,
+                    attention_variant,
                 )
             finally:
                 self._lora_state.reset()
@@ -152,7 +148,7 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
                     forward,
                     stream_idx,
                     None,
-                    dsa_variant,
+                    attention_variant,
                 )
             finally:
                 self.model_runner.lora_manager.reset_lora_batch()
@@ -169,7 +165,7 @@ class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
             forward,
             stream_idx,
             variant_label,
-            dsa_variant,
+            attention_variant,
         )
         self._lora_state.reset()