diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 8ef5f1e63..07f81e62c 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -1770,26 +1770,45 @@ class AscendAttnBackend(AttentionBackend): mask = self.mtp_mask sparse_mode = 4 if is_swa_layer else 3 - attn_output, _ = torch_npu.npu_fused_infer_attention_score_v2( - query, - k_cache, - v_cache, - block_table=block_table, - block_size=self.page_size, - num_query_heads=layer.tp_q_head_num, - num_key_value_heads=layer.tp_k_head_num, - input_layout="TND", - atten_mask=mask, - softmax_scale=layer.scaling, - actual_seq_qlen=actual_seq_lengths, - actual_seq_kvlen=actual_seq_lengths_kv, - sparse_mode=sparse_mode, - pre_tokens=( - layer.sliding_window_size if is_swa_layer else FULL_ATTENTION_WINDOW - ), - next_tokens=0 if is_swa_layer else FULL_ATTENTION_WINDOW, - learnable_sink=sinks, - ) + if self.is_hybrid_swa: + attn_output, _ = torch_npu.npu_fused_infer_attention_score_v2( + query, + k_cache, + v_cache, + block_table=block_table, + block_size=self.page_size, + num_query_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + atten_mask=mask, + softmax_scale=layer.scaling, + actual_seq_qlen=actual_seq_lengths, + actual_seq_kvlen=actual_seq_lengths_kv, + sparse_mode=sparse_mode, + pre_tokens=( + layer.sliding_window_size + if is_swa_layer + else FULL_ATTENTION_WINDOW + ), + next_tokens=0 if is_swa_layer else FULL_ATTENTION_WINDOW, + learnable_sink=sinks, + ) + else: + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + query, + k_cache, + v_cache, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + atten_mask=mask, + scale=layer.scaling, + actual_seq_lengths=actual_seq_lengths, + actual_seq_lengths_kv=actual_seq_lengths_kv, + sparse_mode=sparse_mode, + ) attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) if ( not self.graph_mode diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py index 22ef7e237..394ff813f 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py @@ -11,17 +11,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run the model with npu graph and torch.compile.""" from __future__ import annotations -import threading from typing import TYPE_CHECKING import torch from sglang.srt.configs.model_config import is_deepseek_dsa -from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, ) @@ -34,38 +31,19 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner): def __init__(self, eagle_worker: EagleDraftWorker): super().__init__(eagle_worker) - def _create_graph(self): - return torch.npu.NPUGraph() - def _cache_loc_dtype(self): return torch.int32 - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.npu.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.npu.graph( - graph, pool=pool, stream=stream, auto_dispatch_capture=True - ): - out = run_once_fn() - return out - - def _replay_update(self, seq_lens): - self.graphs[self.bs].update( - cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}] - ) - - def _replay(self, forward_batch: ForwardBatch): + def _replay_graph(self, shape_key, forward_batch): if not is_deepseek_dsa(self.model_runner.model_config.hf_config): seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * ( self.bs - self.raw_bs ) - thread = threading.Thread(target=self._replay_update, args=(seq_lens,)) - thread.start() - self.graphs[self.bs].replay() - thread.join() + return self.backend.replay_with_input_update( + shape_key, + seq_lens=seq_lens, + attr_name="actual_seq_lengths_kv", + attr_type=[], + ) else: - self.graphs[self.bs].replay() + return self.backend.replay(shape_key, forward_batch) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 044fc538e..450a81f3a 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py @@ -11,19 +11,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run the model with npu graph and torch.compile""" from __future__ import annotations -import logging -import threading from typing import TYPE_CHECKING, Dict, Union -import numpy as np import torch from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa -from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, ) @@ -31,25 +26,11 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( if TYPE_CHECKING: from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker -from sglang.srt.utils import is_npu - -logger = logging.getLogger(__name__) - -if is_npu(): - torch.cuda.CUDAGraph = torch.npu.NPUGraph - torch.cuda.synchronize = torch.npu.synchronize - torch.cuda.graph = torch.npu.graph - torch.cuda.stream = torch.npu.stream - torch.cuda.Stream = torch.npu.Stream - torch.cuda.current_stream = torch.npu.current_stream - class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): def __init__(self, eagle_worker: EagleDraftWorker): - super().__init__(eagle_worker) - self.update_attr_name = None - self.update_attr_type = None self._init_arch_map() + super().__init__(eagle_worker) def _init_arch_map(self): self.attr_name: Dict[str, str] = { @@ -61,21 +42,8 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): AttentionArch.MHA: torch.Tensor(), } - def _create_graph(self): - return torch.npu.NPUGraph() - - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.npu.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.npu.graph( - graph, pool=pool, stream=stream, auto_dispatch_capture=True - ): - out = run_once_fn() - return out + def _cache_loc_dtype(self): + return torch.int32 def _get_update_attr_name(self): return self.attr_name[AttentionArch.MLA] @@ -83,33 +51,19 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): def _get_update_attr_type(self): return self.attr_type[AttentionArch.MLA] - def _replay_update(self, seq_lens_list): - if isinstance(self.update_attr_type, torch.Tensor): - seq_lens = torch.from_numpy(np.array(seq_lens_list).astype(np.int32)) - - self.graphs[self.bs].update( - cpu_update_input=[ - {self.update_attr_name: seq_lens} for seq_lens in seq_lens_list - ] - ) - - def _replay(self, forward_batch: ForwardBatch): - self.update_attr_name = self._get_update_attr_name() - self.update_attr_type = self._get_update_attr_type() + def _replay_graph(self, shape_key, forward_batch): if not is_deepseek_dsa(self.model_runner.model_config.hf_config): seq_lens_for_each_draft_step = [] for speculative_step_id in range(self.speculative_num_steps - 1): - seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1 + seq_lens_cpu = ( + forward_batch.seq_lens_cpu[: self.raw_bs] + speculative_step_id + 1 + ) seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs) seq_lens_for_each_draft_step.append(seq_lens) - thread = threading.Thread( - target=self._replay_update, args=(seq_lens_for_each_draft_step,) + attr_name = self._get_update_attr_name() + cpu_update_input = [{attr_name: sl} for sl in seq_lens_for_each_draft_step] + return self.backend.replay_with_input_update( + shape_key, seq_lens=None, cpu_update_input=cpu_update_input ) - thread.start() - self.graphs[self.bs].replay() - thread.join() else: - self.graphs[self.bs].replay() - - def _cache_loc_dtype(self): - return torch.int32 + return self.backend.replay(shape_key, forward_batch) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py index 4103cebb6..ed388ad77 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py @@ -11,19 +11,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run the multi-layer eagle draft extend model with npu graph.""" from __future__ import annotations import logging -import threading import time from typing import TYPE_CHECKING, List, Optional import torch from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled -from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleDraftExtendCudaGraphRunner, MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, @@ -44,36 +41,16 @@ class MultiLayerEagleDraftExtendNpuGraphRunner( def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int): super().__init__(eagle_worker, step) - def _create_graph(self): - return torch.npu.NPUGraph() - - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.npu.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.npu.graph( - graph, - pool=pool, - stream=stream, - auto_dispatch_capture=True, - ): - out = run_once_fn() - return out - - def _replay(self, forward_batch: ForwardBatch): + def _replay_graph(self, shape_key, forward_batch): seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * ( self.bs - self.raw_bs ) - thread = threading.Thread( - target=self.graphs[self.bs].update, - kwargs={"cpu_update_input": [{"actual_seq_kvlen": seq_lens}]}, + return self.backend.replay_with_input_update( + shape_key, + seq_lens=seq_lens, + attr_name="actual_seq_kvlen", + attr_type=[], ) - thread.start() - self.graphs[self.bs].replay() - thread.join() class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner( diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py index db5bed56b..e16569b30 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py @@ -140,19 +140,29 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend): def replay_with_input_update( self, shape_key: Any, - seq_lens: list, - attr_name: str, - attr_type: Any, + seq_lens: Any, + attr_name: str = None, + attr_type: Any = None, + cpu_update_input: list = None, ) -> Any: """Rebind seq_lens on the recorded NPU graph in a background - thread, then replay. Used when the model is not deepseek-nsa.""" - if isinstance(attr_type, torch.Tensor): - seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32)) + thread, then replay. Used when the model is not deepseek-nsa. + + Two calling conventions: + 1. (legacy) seq_lens + attr_name + attr_type: + Constructs cpu_update_input=[{attr_name: seq_lens}] internally. + 2. cpu_update_input: A list of {attr_name: seq_lens} dicts, + one per speculative step. Used by EAGLE draft runners. + """ + if cpu_update_input is None: + if isinstance(attr_type, torch.Tensor): + seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32)) + cpu_update_input = [{attr_name: seq_lens}] graph = self._graphs[shape_key] def _update(): - graph.update(cpu_update_input=[{attr_name: seq_lens}]) + graph.update(cpu_update_input=cpu_update_input) thread = threading.Thread(target=_update) thread.start() 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 ee8842869..c3154a5f9 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -26,7 +26,7 @@ from sglang.srt.model_executor.runner import ( get_batch_sizes_to_capture, model_capture_mode, ) -from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend from sglang.srt.model_executor.runner_backend_utils import ( CUDA_GRAPH_CAPTURE_FAILED_MSG, ) @@ -226,11 +226,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): ) self.buffers.share_buffers() - # Backend (Full CUDA graph capture) - self.backend = FullCudaGraphBackend( - self, - enable_memory_saver=model_runner.server_args.enable_memory_saver, - ) + self.backend = resolve_decode_backend(self) # Capture try: @@ -241,6 +237,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _replay_graph(self, shape_key, forward_batch): + return self.backend.replay(shape_key, forward_batch) + # ----------------------------------------------------------------- # Helpers # ----------------------------------------------------------------- @@ -533,7 +532,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): else contextlib.nullcontext() ) with timer_ctx: - out = self.backend.replay(shape_key, forward_batch) + out = self._replay_graph(shape_key, forward_batch) if bs != raw_bs: out = self._postprocess_output_to_raw_bs(out, raw_bs) 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 6ed2a08ee..c0a149f74 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 @@ -26,7 +26,7 @@ from sglang.srt.model_executor.runner import ( get_batch_sizes_to_capture, model_capture_mode, ) -from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend from sglang.srt.model_executor.runner_backend_utils import ( CUDA_GRAPH_CAPTURE_FAILED_MSG, ) @@ -242,10 +242,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ) self.buffers.share_buffers() - self.backend = FullCudaGraphBackend( - self, - enable_memory_saver=model_runner.server_args.enable_memory_saver, - ) + self.backend = resolve_decode_backend(self) try: with model_capture_mode(): @@ -255,6 +252,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _replay_graph(self, shape_key, forward_batch): + return self.backend.replay(shape_key, forward_batch) + def _cache_loc_dtype(self): return torch.int64 @@ -556,7 +556,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): else contextlib.nullcontext() ) with timer_ctx: - out = self.backend.replay(shape_key, forward_batch) + out = self._replay_graph(shape_key, forward_batch) if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: unpadding_bs = num_tokens 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 bef89ace6..03d5380fc 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 @@ -50,7 +50,7 @@ from sglang.srt.model_executor.runner import ( get_batch_sizes_to_capture, model_capture_mode, ) -from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend from sglang.srt.model_executor.runner_backend_utils import ( CUDA_GRAPH_CAPTURE_FAILED_MSG, ) @@ -272,10 +272,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) - self.backend = FullCudaGraphBackend( - self, - enable_memory_saver=self.model_runner.server_args.enable_memory_saver, - ) + self.backend = resolve_decode_backend(self) try: with model_capture_mode(): @@ -285,6 +282,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _replay_graph(self, shape_key, forward_batch): + return self.backend.replay(shape_key, forward_batch) + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): return bs @@ -588,7 +588,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): self.raw_bs = raw_bs self.bs = bs shape_key = self._make_graph_key(bs) - out = self.backend.replay(shape_key, forward_batch) + out = self._replay_graph(shape_key, forward_batch) if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: unpadding_bs = num_tokens