[Refactor] Generalize attention graph variants in the decode runner

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,
@@ -398,62 +399,30 @@ 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
             return int(fb.seq_lens_cpu.max().item()) <= self.index_topk
 
         # 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
-                # 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_dsa_variant,
+                    get_capture_attention_variant,
                 )
 
-                variant = get_capture_dsa_variant()
-                if variant == "dense":
-                    return True
-                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 get_capture_attention_variant() == DSA_DENSE
+            # Eager k-only decode skip is validated on ROCm only.
+            if not _is_hip:
                 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 @@
     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 @@
     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.utils import (
     empty_context,
     get_available_gpu_memory,
-    is_hip,
     require_attn_tp_gather,
     require_mlp_tp_gather,
 )
@@ -222,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
@@ -254,37 +259,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 +301,10 @@ def __init__(
         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 @@ 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 = 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
@@ -705,8 +667,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
             ),
@@ -1109,19 +1071,12 @@ 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)]
         )
-        # 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]
+        variants = self.attention_graph_variants
+        attention_variants = (
+            variants.capture_labels if variants is not None else (None,)
         )
         for bs in capture_range:
             if get_parallel().tp_rank == 0:
@@ -1136,31 +1091,30 @@ 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:
-                            self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label
-                            )
-                        else:
-                            self.capture_one_shape(
-                                bs, forward, stream_idx, variant_label, dsa_variant
-                            )
-        _set_capture_dsa_variant(None)
+                        self.capture_one_shape(
+                            bs,
+                            forward,
+                            stream_idx,
+                            variant_label,
+                            attention_variant,
+                        )
+        _set_capture_attention_variant(None)
 
     def capture_one_shape(
         self,
         size: int,
         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 @@ 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 +1273,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 +1393,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: