From 17fa5ad3276aec39f7626abcf7bbdf3dcff8297f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 11 Sep 2026 00:45:58 -0700 Subject: [PATCH] [Refactor] Generalize attention graph variants in the decode runner (#38993) --- .../srt/layers/attention/dsa/dsa_indexer.py | 53 ++------ .../srt/layers/attention/graph_variants.py | 59 +++++++++ .../runner/decode_cuda_graph_runner.py | 124 ++++++------------ .../srt/model_executor/runner/shape_key.py | 21 +-- .../runner_utils/capture_mode.py | 20 +-- .../eagle_draft_cuda_graph_runner.py | 2 + .../eagle_draft_extend_cuda_graph_runner.py | 2 + .../frozen_kv_mtp_cuda_graph_runner.py | 4 +- ...er_eagle_draft_extend_cuda_graph_runner.py | 2 + .../srt/speculative/uno_cuda_graph_runner.py | 14 +- 10 files changed, 135 insertions(+), 166 deletions(-) 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 6678e4b3a..26a457f62 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 @@ 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 - # 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,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 - # 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: 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 000000000..14ba77c7b --- /dev/null +++ b/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 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 51463f6e4..d5160d43b 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,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]: - """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 @@ 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=( - self._resolve_dsa_variant(forward_batch) + attention_variant=( + 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 - # 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,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: - _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, @@ -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: diff --git a/python/sglang/srt/model_executor/runner/shape_key.py b/python/sglang/srt/model_executor/runner/shape_key.py index f04d07b1e..7691b7e07 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 @@ from typing import Optional @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. - 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. - """ - + # 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 - dsa_variant: 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 f936cd0c9..a2195b9d0 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 @@ 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 -# 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,13 @@ 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 _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/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index cb603f1ff..0c6d80e3d 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 @@ 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 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 529867b00..4aeb6c5a3 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 @@ 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 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 aa9ff6612..c9a7c76cb 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 @@ 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 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 dc0d0230f..7a0a371fe 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 @@ 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 diff --git a/python/sglang/srt/speculative/uno_cuda_graph_runner.py b/python/sglang/srt/speculative/uno_cuda_graph_runner.py index 921353f89..bcd0f57fd 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 @@ 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()