minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-11 00:45:58 -07:00
committed by GitHub
parent f618022b73
commit 17fa5ad327
10 changed files with 135 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+11 -42
View File
@@ -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:
python/sglang/srt/layers/attention/graph_variants.py
+59
View File
@@ -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
+39 -85
View File
@@ -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:
python/sglang/srt/model_executor/runner/shape_key.py
+5 -16
View File
@@ -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
python/sglang/srt/model_executor/runner_utils/capture_mode.py
+7 -13
View File
@@ -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
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+2
View File
@@ -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
+2
View File
@@ -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
+3 -1
View File
@@ -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
+2
View File
@@ -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
+5 -9
View File
@@ -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()