[Feature] Spec-Overlap supporting DP-ATTN; PD-Disaggregation; npugraph mode (#12443)

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -89,7 +89,7 @@ class EAGLEDraftCudaGraphRunner:
             set_torch_compile_config()
 
         # Graph inputs
-        with torch.device("cuda"):
+        with torch.device(model_runner.device):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.seq_lens = torch.full(
@@ -158,13 +158,30 @@ class EAGLEDraftCudaGraphRunner:
 
         return is_bs_supported
 
+    def _create_graph(self):
+        return torch.cuda.CUDAGraph()
+
+    def _capture_init(self, run_once_fn):
+        for _ in range(2):
+            torch.cuda.synchronize()
+            self.model_runner.tp_group.barrier()
+            run_once_fn()
+
+    def _capture_graph(self, graph, pool, stream, run_once_fn):
+        with torch.cuda.graph(graph, pool=pool, stream=stream):
+            out = run_once_fn()
+        return out
+
+    def _replay(self, forward_batch: ForwardBatch):
+        self.graphs[self.bs].replay()
+
     def capture(self):
         CudaGraphRunner.capture(self)
 
     def capture_one_batch_size(
         self, num_seqs: int, forward: Callable, stream_idx: int = 0
     ):
-        graph = torch.cuda.CUDAGraph()
+        graph = self._create_graph()
         stream = self.stream
         num_tokens = num_seqs * self.num_tokens_per_bs
 
@@ -285,16 +302,10 @@ class EAGLEDraftCudaGraphRunner:
 
         self.deepep_adapter.capture(is_extend_in_batch=False)
 
-        for _ in range(2):
-            torch.cuda.synchronize()
-            self.model_runner.tp_group.barrier()
-
-            run_once()
-
-        with torch.cuda.graph(
-            graph, pool=get_global_graph_memory_pool(), stream=stream
-        ):
-            out = run_once()
+        self._capture_init(run_once)
+        out = self._capture_graph(
+            graph, get_global_graph_memory_pool(), stream, run_once
+        )
 
         set_global_graph_memory_pool(graph.pool())
         return graph, out
@@ -362,10 +373,12 @@ class EAGLEDraftCudaGraphRunner:
         self.model_runner.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
             forward_batch, bs
         )
+        self.raw_bs = raw_bs
+        self.bs = bs
         # TODO: The forward_batch.seq_len_sum might need to be updated to reflect the padding in the cuda graph
 
         # Replay
-        self.graphs[bs].replay()
+        self._replay(forward_batch)
         out = self.output_buffers[bs]
 
         if bs != raw_bs:

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -43,8 +43,10 @@ class EAGLEDraftExtendCudaGraphRunner:
         if not hasattr(eagle_worker, "model_runner"):
             # V2: EagleDraftWorker
             self.model_runner = model_runner = eagle_worker.draft_runner
+            self.forward_mode = ForwardMode.DRAFT_EXTEND_V2
         else:
             self.model_runner = model_runner = eagle_worker.model_runner
+            self.forward_mode = ForwardMode.DRAFT_EXTEND
 
         self.graphs = {}
         self.output_buffers = {}
@@ -86,7 +88,7 @@ class EAGLEDraftExtendCudaGraphRunner:
             set_torch_compile_config()
 
         # Graph inputs
-        with torch.device("cuda"):
+        with torch.device(model_runner.device):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.out_cache_loc = torch.ones((self.max_num_token,), dtype=torch.int64)
@@ -116,8 +118,12 @@ class EAGLEDraftExtendCudaGraphRunner:
                     (self.max_num_token, self.model_runner.model_config.hidden_size),
                     dtype=self.model_runner.dtype,
                 )
-
-            self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
+            self.seq_len_fill_value = (
+                self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
+            )
+            self.seq_lens = torch.full(
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
+            )
             self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
             self.accept_length = torch.full(
                 (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32
@@ -153,7 +159,14 @@ class EAGLEDraftExtendCudaGraphRunner:
                 vocab_size = self.model_runner.model_config.vocab_size
 
             self.next_token_logits_buffer = torch.zeros(
-                (self.max_bs, vocab_size),
+                (
+                    (
+                        self.max_bs * self.num_tokens_per_bs
+                        if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2
+                        else self.max_bs
+                    ),
+                    vocab_size,
+                ),
                 dtype=torch.float,
             )
 
@@ -187,11 +200,28 @@ class EAGLEDraftExtendCudaGraphRunner:
 
         return is_bs_supported
 
+    def _create_graph(self):
+        return torch.cuda.CUDAGraph()
+
+    def _capture_init(self, run_once_fn):
+        for _ in range(2):
+            torch.cuda.synchronize()
+            self.model_runner.tp_group.barrier()
+            run_once_fn()
+
+    def _capture_graph(self, graph, pool, stream, run_once_fn):
+        with torch.cuda.graph(graph, pool=pool, stream=stream):
+            out = run_once_fn()
+        return out
+
+    def _replay(self, forward_batch: ForwardBatch):
+        self.graphs[self.bs].replay()
+
     def capture(self):
         CudaGraphRunner.capture(self)
 
     def capture_one_batch_size(self, bs: int, forward: Callable, stream_idx: int = 0):
-        graph = torch.cuda.CUDAGraph()
+        graph = self._create_graph()
         stream = self.stream
         num_tokens = bs * self.num_tokens_per_bs
 
@@ -207,7 +237,9 @@ class EAGLEDraftExtendCudaGraphRunner:
         positions = self.positions[:num_tokens]
         mrope_positions = self.mrope_positions[:, :num_tokens]
         hidden_states = self.hidden_states[:num_tokens]
-        next_token_logits_buffer = self.next_token_logits_buffer[:bs]
+        next_token_logits_buffer = self.next_token_logits_buffer[
+            : bs if self.forward_mode == ForwardMode.DRAFT_EXTEND else num_tokens
+        ]
 
         if self.require_mlp_tp_gather:
             self.global_num_tokens_gpu.copy_(
@@ -254,7 +286,7 @@ class EAGLEDraftExtendCudaGraphRunner:
 
         # Forward batch
         forward_batch = ForwardBatch(
-            forward_mode=ForwardMode.DRAFT_EXTEND,
+            forward_mode=self.forward_mode,
             batch_size=bs,
             input_ids=input_ids,
             req_pool_indices=req_pool_indices,
@@ -287,7 +319,7 @@ class EAGLEDraftExtendCudaGraphRunner:
             req_pool_indices=req_pool_indices,
             seq_lens=seq_lens,
             encoder_lens=None,
-            forward_mode=ForwardMode.DRAFT_EXTEND,
+            forward_mode=self.forward_mode,
             spec_info=spec_info,
         )
 
@@ -318,16 +350,11 @@ class EAGLEDraftExtendCudaGraphRunner:
             forward_batch.spec_info.hidden_states = hidden_states_backup
             return ret
 
-        for _ in range(2):
-            torch.cuda.synchronize()
-            self.model_runner.tp_group.barrier()
+        self._capture_init(run_once)
 
-            run_once()
-
-        with torch.cuda.graph(
-            graph, pool=get_global_graph_memory_pool(), stream=stream
-        ):
-            out = run_once()
+        out = self._capture_graph(
+            graph, get_global_graph_memory_pool(), stream, run_once
+        )
 
         set_global_graph_memory_pool(graph.pool())
         return graph, out
@@ -399,21 +426,32 @@ class EAGLEDraftExtendCudaGraphRunner:
             seq_lens_sum=forward_batch.seq_lens_sum
             + (bs - raw_bs) * self.seq_len_fill_value,
             encoder_lens=None,
-            forward_mode=ForwardMode.DRAFT_EXTEND,
+            forward_mode=self.forward_mode,
             spec_info=forward_batch.spec_info,
             seq_lens_cpu=self.seq_lens_cpu,
         )
 
         # Replay
-        self.graphs[bs].replay()
+        self.raw_bs = raw_bs
+        self.bs = bs
+        self._replay(forward_batch)
         out = self.output_buffers[bs]
-        if bs != raw_bs:
+
+        if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2:
+            # DRAFT_EXTEND_V2: all tokens calculations whether accepted or not.
+            unpadding_bs = num_tokens
+        elif bs != raw_bs:
             forward_batch.spec_info.accept_length = self.accept_length[:raw_bs]
+            unpadding_bs = raw_bs
+        else:
+            unpadding_bs = None
+
+        if unpadding_bs is not None:
             out_copy = out
             out = LogitsProcessorOutput(
-                next_token_logits=out.next_token_logits[:raw_bs],
-                hidden_states=out.hidden_states[:raw_bs],
+                next_token_logits=out.next_token_logits[:unpadding_bs],
+                hidden_states=out.hidden_states[:unpadding_bs],
             )
-            out.topk_p = out_copy.topk_p[:raw_bs]
-            out.topk_index = out_copy.topk_index[:raw_bs]
+            out.topk_p = out_copy.topk_p[:unpadding_bs]
+            out.topk_index = out_copy.topk_index[:unpadding_bs]
         return out

python/sglang/srt/speculative/eagle_draft_extend_npu_graph_runner.py

@@ -0,0 +1,68 @@
+# Copyright 2024-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# 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_nsa
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
+    EAGLEDraftExtendCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker import EAGLEWorker
+
+
+class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
+    def __init__(self, eagle_worker: EAGLEWorker):
+        super().__init__(eagle_worker)
+
+    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_update(self, seq_lens):
+        self.graphs[self.bs].update(
+            cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
+        )
+
+    def _replay(self, forward_batch: ForwardBatch):
+        if not is_deepseek_nsa(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()
+        else:
+            self.graphs[self.bs].replay()

python/sglang/srt/speculative/eagle_draft_npu_graph_runner.py

@@ -0,0 +1,81 @@
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# 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
+
+import torch
+
+from sglang.srt.configs.model_config import is_deepseek_nsa
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+    EAGLEDraftCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker import EAGLEWorker
+
+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: EAGLEWorker):
+        super().__init__(eagle_worker)
+
+    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_update(self, seq_lens):
+        self.graphs[self.bs].update(
+            cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
+        )
+
+    def _replay(self, forward_batch: ForwardBatch):
+        if not is_deepseek_nsa(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()
+        else:
+            self.graphs[self.bs].replay()

python/sglang/srt/speculative/eagle_info.py

@@ -665,6 +665,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
             topk_p=torch.empty((0, topk), device=device, dtype=torch.float32),
             topk_index=torch.empty((0, topk), device=device, dtype=torch.int64),
             capture_hidden_mode=capture_hidden_mode,
+            allocate_lens=torch.empty((0,), device=device, dtype=torch.int32),
+            new_seq_lens=torch.empty((0,), device=device, dtype=torch.int32),
             accept_length=torch.empty((0,), device=device, dtype=torch.int32),
             accept_length_cpu=[],
         )

python/sglang/srt/speculative/eagle_info_v2.py

@@ -142,26 +142,29 @@ class EagleDraftInputV2Mixin:
         topk: int,
         num_steps: int,
     ):
-        bs = len(batch.seq_lens)
+        if not batch.forward_mode.is_idle():
+            bs = len(batch.seq_lens)
 
-        # Assign cache locations
-        batch.out_cache_loc = torch.empty(
-            (bs * topk * num_steps,),
-            dtype=torch.int64,
-            device=batch.input_ids.device,
-        )
-        # FIXME(lsyin): align with the default code path
-        assign_draft_cache_locs_page_size_1[(bs,)](
-            batch.req_pool_indices,
-            req_to_token_pool.req_to_token,
-            batch.seq_lens,
-            batch.out_cache_loc,
-            req_to_token_pool.req_to_token.shape[1],
-            topk,
-            num_steps,
-        )
+            # Assign cache locations
+            batch.out_cache_loc = torch.empty(
+                (bs * topk * num_steps,),
+                dtype=torch.int64,
+                device=batch.input_ids.device,
+            )
+            # FIXME(lsyin): align with the default code path
+            assign_draft_cache_locs_page_size_1[(bs,)](
+                batch.req_pool_indices,
+                req_to_token_pool.req_to_token,
+                batch.seq_lens,
+                batch.out_cache_loc,
+                req_to_token_pool.req_to_token.shape[1],
+                topk,
+                num_steps,
+            )
 
         # Get a forward batch
+        self.num_tokens_per_batch = topk
+        self.num_tokens_for_logprob_per_batch = topk
         batch.capture_hidden_mode = CaptureHiddenMode.LAST
         self.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
         forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
@@ -174,6 +177,7 @@ class EagleDraftInputV2Mixin:
         predict: torch.Tensor,
         num_draft_tokens: int,
         draft_model_runner: Any,
+        cuda_graph_runner: Any,
     ):
         seq_lens_cpu_ = batch.seq_lens_cpu
         extend_num_tokens = len(batch.seq_lens) * num_draft_tokens
@@ -187,9 +191,15 @@ class EagleDraftInputV2Mixin:
         batch.extend_prefix_lens = seq_lens_cpu_.tolist()
         batch.extend_num_tokens = extend_num_tokens
         batch.capture_hidden_mode = CaptureHiddenMode.FULL
-        batch.forward_mode = ForwardMode.DRAFT_EXTEND_V2
+        batch.forward_mode = (
+            ForwardMode.IDLE
+            if batch.forward_mode.is_idle()
+            else ForwardMode.DRAFT_EXTEND_V2
+        )
         forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
-        draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
+        can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
+        if not batch.forward_mode.is_idle() and not can_cuda_graph:
+            draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
         return forward_batch
 
 
@@ -201,22 +211,27 @@ class EagleVerifyInputV2Mixin:
         batch: ModelWorkerBatch,
         target_worker: TpModelWorker,
     ):
-        # Assign cache locations
-        bs = len(batch.req_pool_indices)
-        batch.input_ids = self.draft_token
-        device = batch.input_ids.device
-        batch.out_cache_loc = assign_extend_cache_locs_func(
-            req_pool_indices=batch.req_pool_indices,
-            req_to_token=req_to_token_pool.req_to_token,
-            start_offset=batch.seq_lens,
-            end_offset=batch.seq_lens + self.draft_token_num,
-            batch_size=bs,
-            draft_token_num=self.draft_token_num,
-            device=device,
-        )
+        if not batch.forward_mode.is_idle():
+            # Assign cache locations
+            bs = len(batch.req_pool_indices)
+            batch.input_ids = self.draft_token
+            device = batch.input_ids.device
+            batch.out_cache_loc = assign_extend_cache_locs_func(
+                req_pool_indices=batch.req_pool_indices,
+                req_to_token=req_to_token_pool.req_to_token,
+                start_offset=batch.seq_lens,
+                end_offset=batch.seq_lens + self.draft_token_num,
+                batch_size=bs,
+                draft_token_num=self.draft_token_num,
+                device=device,
+            )
 
         # Get a forward batch
-        batch.forward_mode = ForwardMode.TARGET_VERIFY
+        batch.forward_mode = (
+            ForwardMode.IDLE
+            if batch.forward_mode.is_idle()
+            else ForwardMode.TARGET_VERIFY
+        )
         batch.capture_hidden_mode = CaptureHiddenMode.FULL
         verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
 
@@ -228,9 +243,10 @@ class EagleVerifyInputV2Mixin:
         if can_run_cuda_graph:
             target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
         else:
-            target_worker.model_runner.attn_backend.init_forward_metadata(
-                verify_forward_batch
-            )
+            if not batch.forward_mode.is_idle():
+                target_worker.model_runner.attn_backend.init_forward_metadata(
+                    verify_forward_batch
+                )
 
         return verify_forward_batch, can_run_cuda_graph
 
@@ -243,6 +259,16 @@ class EagleVerifyInputV2Mixin:
         Verify and find accepted tokens based on logits output and batch
         (which contains spec decoding information).
         """
+        if batch.forward_mode.is_idle():
+            predict = torch.empty(0, dtype=torch.long, device=batch.input_ids.device)
+            accept_length = torch.empty(
+                0, dtype=torch.int32, device=batch.input_ids.device
+            )
+            accept_index = torch.empty(
+                0, dtype=torch.int32, device=batch.input_ids.device
+            )
+            return predict, accept_length, accept_index
+
         bs = len(batch.seq_lens)
         sampling_info = batch.sampling_info
         next_token_logits = logits_output.next_token_logits

python/sglang/srt/speculative/eagle_utils.py

@@ -16,116 +16,6 @@ if _is_cuda or _is_hip:
     )
 
 
-def build_tree_efficient_native(
-    parent_list: torch.Tensor,
-    selected_index: torch.Tensor,
-    verified_seq_len: torch.Tensor,
-    tree_mask: torch.Tensor,
-    retrive_index: torch.Tensor,
-    retrive_next_token: torch.Tensor,
-    retrive_next_sibling: torch.Tensor,
-    topk: int,
-    draft_token_num: int,
-    tree_mask_mode: int,
-    bs: int,
-):
-    # Generate batch and token index ranges
-    bs_range = torch.arange(bs, device=tree_mask.device).view(-1, 1)
-    draft_token_num_range = torch.arange(draft_token_num, device=tree_mask.device)
-
-    # Optimized common case for performance.
-    if draft_token_num == 2 and topk == 1 and tree_mask_mode == TreeMaskMode.FULL_MASK:
-        positions = verified_seq_len.repeat_interleave(draft_token_num)
-        positions = (positions.view(bs, -1) + draft_token_num_range).view(-1)
-
-        retrive_index[:] = bs_range * draft_token_num + draft_token_num_range
-        retrive_next_token[:, 0] = 1
-        retrive_next_token[:, 1] = -1
-        return (
-            positions,
-            retrive_index,
-            retrive_next_token,
-            retrive_next_sibling,
-            tree_mask,
-        )
-
-    # Precompute sequence tree indices
-    draft_token_num_range1 = torch.arange(draft_token_num - 1, device=tree_mask.device)
-    cum_seq_len = torch.cumsum(verified_seq_len * draft_token_num, dim=0)
-    cum_seq_len = torch.cat((torch.tensor([0], device=tree_mask.device), cum_seq_len))
-    cum_seq_len = cum_seq_len[:-1]
-    seq_tree_idx = (
-        draft_token_num * draft_token_num * torch.arange(bs, device=tree_mask.device)
-        + cum_seq_len
-    )
-
-    # Batch processing for tree mask
-    if tree_mask_mode == TreeMaskMode.FULL_MASK:
-        token_tree_base = (
-            seq_tree_idx.view(-1, 1)
-            + (verified_seq_len.view(-1, 1) + draft_token_num) * draft_token_num_range
-        )
-        token_tree_indices = token_tree_base + verified_seq_len.view(-1, 1) + 1
-    else:
-        token_tree_indices = (
-            bs_range * draft_token_num**2 + draft_token_num_range * draft_token_num + 1
-        )
-
-    tree_mask[token_tree_indices.flatten() - 1] = True
-    indices = token_tree_indices.unsqueeze(-1) + draft_token_num_range1.view(1, 1, -1)
-    tree_mask[indices.view(-1)] = False
-
-    positions = verified_seq_len.repeat_interleave(draft_token_num)
-    parent_tb_indices = selected_index // topk
-    retrive_index[:] = bs_range * draft_token_num + draft_token_num_range
-    tree_mask[token_tree_indices.view(-1, 1) + draft_token_num_range1] = True
-
-    for bid in range(bs):
-        for tid in range(draft_token_num):
-            position = 0
-            if tid == 0:
-                # Process root node
-                for i in range(draft_token_num - 1, 0, -1):
-                    parent_position = 0
-                    parent_tb_idx = parent_tb_indices[bid][i - 1]
-                    if parent_tb_idx > 0:
-                        parent_token_idx = parent_list[bid][parent_tb_idx]
-                        loop_num = draft_token_num - parent_position
-                        for _ in range(loop_num):
-                            if selected_index[bid][parent_position] == parent_token_idx:
-                                parent_position += 1
-                                break
-                            parent_position += 1
-                    if parent_position == draft_token_num:
-                        continue
-
-                    if retrive_next_token[bid][parent_position] != -1:
-                        retrive_next_sibling[bid][i] = retrive_next_token[bid][
-                            parent_position
-                        ]
-                    retrive_next_token[bid][parent_position] = i
-            else:
-                # Process no-root nodes
-                cur_position = tid - 1
-                while True:
-                    position += 1
-                    if cur_position >= draft_token_num:
-                        tree_mask[token_tree_indices + cur_position] = True
-                        parent_tb_idx = selected_index[bid][cur_position] // topk
-                    else:
-                        parent_tb_idx = parent_tb_indices[bid][cur_position]
-                    if parent_tb_idx == 0:
-                        break
-                    token_idx = parent_list[bid][parent_tb_idx]
-                    cur_position = 0
-                    for _ in range(draft_token_num):
-                        if selected_index[bid][cur_position] == token_idx:
-                            break
-                        cur_position += 1
-                positions[bid * draft_token_num + tid] += position
-    return positions, retrive_index, retrive_next_token, retrive_next_sibling, tree_mask
-
-
 def organize_draft_results(
     score_list: List[torch.Tensor],
     token_list: List[torch.Tensor],
@@ -229,24 +119,19 @@ def build_tree_kernel_efficient(
         )
 
     if _is_npu:
-        (
+        torch.ops.npu.build_tree_kernel_efficient(
+            parent_list.to(dtype=torch.int64),
+            top_scores_index,
+            seq_lens,
+            tree_mask,
             positions,
             retrive_index,
             retrive_next_token,
             retrive_next_sibling,
-            tree_mask,
-        ) = build_tree_efficient_native(
-            parent_list,
-            top_scores_index,
-            seq_lens,
-            tree_mask,
-            retrive_index,
-            retrive_next_token,
-            retrive_next_sibling,
             topk,
+            spec_steps,
             num_verify_tokens,
             tree_mask_mode,
-            bs,
         )
     else:
         sgl_build_tree_kernel_efficient(
@@ -273,75 +158,6 @@ def build_tree_kernel_efficient(
     )
 
 
-def verify_tree_greedy_native(
-    predicts: torch.Tensor,
-    accept_index: torch.Tensor,
-    accept_token_num: torch.Tensor,
-    candidates: torch.Tensor,
-    retrive_index: torch.Tensor,
-    retrive_next_token: torch.Tensor,
-    retrive_next_sibling: torch.Tensor,
-    target_predict: torch.Tensor,
-    topk: int = -1,
-):
-    batch_size, num_draft_tokens = candidates.shape
-
-    # Optimized common case for performance.
-    if num_draft_tokens == 2 and accept_index.shape[1] == 2 and topk == 1:
-        comparison_result = candidates[:, 1] == target_predict[:, 0]
-
-        predicts = target_predict.flatten()
-
-        accept_index = torch.arange(
-            0, num_draft_tokens * batch_size, device=candidates.device, dtype=torch.long
-        ).reshape(batch_size, num_draft_tokens)
-        comparison_result = comparison_result.to(torch.int64)
-        accept_index_mask = accept_index[:, 1] * comparison_result
-        accept_index[:, 1] = accept_index_mask - (1 - comparison_result)
-
-        accept_token_num = comparison_result.int()
-        return predicts, accept_index, accept_token_num
-
-    # BFS
-    for bx in range(batch_size):
-        cur_candidates = candidates[bx]
-        cur_retrive_index = retrive_index[bx]
-        cur_next_token = retrive_next_token[bx]
-        cur_next_sibling = retrive_next_sibling[bx]
-        cur_target = target_predict[bx]
-
-        last_accepted_idx = cur_retrive_index[0]
-        accept_index[bx, 0] = last_accepted_idx
-        num_accepted = 0
-        cur_node = 0
-
-        for _ in range(1, num_draft_tokens):
-            cur_node = cur_next_token[cur_node]
-            found = False
-            while cur_node != -1:
-                draft_idx = cur_retrive_index[cur_node]
-                draft_token = cur_candidates[cur_node]
-                target_token = cur_target[last_accepted_idx - num_draft_tokens * bx]
-
-                if draft_token == target_token:
-                    predicts[last_accepted_idx] = target_token
-                    num_accepted += 1
-                    accept_index[bx, num_accepted] = draft_idx
-                    last_accepted_idx = draft_idx
-                    found = True
-                    break
-                else:
-                    cur_node = cur_next_sibling[cur_node]
-            if not found:
-                break
-
-        accept_token_num[bx] = num_accepted
-        predicts[last_accepted_idx] = cur_target[
-            last_accepted_idx - num_draft_tokens * bx
-        ]
-    return predicts, accept_index, accept_token_num
-
-
 def verify_tree_greedy_func(
     predicts: torch.Tensor,
     accept_index: torch.Tensor,
@@ -368,16 +184,16 @@ def verify_tree_greedy_func(
         )
 
     elif _is_npu:
-        predicts, accept_index, accept_token_num = verify_tree_greedy_native(
-            predicts=predicts,  # mutable
-            accept_index=accept_index,  # mutable
-            accept_token_num=accept_token_num,  # mutable
+        from sgl_kernel_npu.sample.verify_tree_greedy import verify_tree_greedy
+
+        verify_tree_greedy(
+            predicts=predicts,
+            accept_index=accept_index,
+            accept_token_num=accept_token_num,
             candidates=candidates,
             retrive_index=retrive_index,
             retrive_next_token=retrive_next_token,
             retrive_next_sibling=retrive_next_sibling,
             target_predict=target_predict,
-            topk=topk,
         )
-
     return predicts, accept_index, accept_token_num

python/sglang/srt/speculative/eagle_worker.py

@@ -31,6 +31,7 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
 from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
     EAGLEDraftExtendCudaGraphRunner,
 )
+from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner
 from sglang.srt.speculative.eagle_info import (
     EagleDraftInput,
     EagleVerifyInput,
@@ -214,9 +215,13 @@ class EAGLEWorker(TpModelWorker):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if self.server_args.disable_cuda_graph or _is_npu:
+        if self.server_args.disable_cuda_graph:
             return
 
+        Device2DraftCudaGraphRunner = {
+            "npu": EAGLEDraftNpuGraphRunner,
+            "cuda": EAGLEDraftCudaGraphRunner,
+        }
         # Capture draft
         if self.speculative_num_steps > 1:
             tic = time.perf_counter()
@@ -224,14 +229,16 @@ class EAGLEWorker(TpModelWorker):
             logger.info(
                 f"Capture draft cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
             )
-            self.cuda_graph_runner = EAGLEDraftCudaGraphRunner(self)
+            self.cuda_graph_runner = Device2DraftCudaGraphRunner[
+                self.target_worker.device
+            ](self)
             after_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(
                 f"Capture draft cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
             )
 
         # Capture extend
-        if self.draft_extend_attn_backend:
+        if self.draft_extend_attn_backend and not _is_npu:
             tic = time.perf_counter()
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -20,6 +20,10 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
 from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
     EAGLEDraftExtendCudaGraphRunner,
 )
+from sglang.srt.speculative.eagle_draft_extend_npu_graph_runner import (
+    EAGLEDraftExtendNpuGraphRunner,
+)
+from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.eagle_info_v2 import (
     assign_extend_cache_locs,
@@ -211,9 +215,13 @@ class EagleDraftWorker(BaseDraftWorker):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if self.server_args.disable_cuda_graph or _is_npu:
+        if self.server_args.disable_cuda_graph:
             return
 
+        Device2DraftCudaGraphRunner = {
+            "npu": EAGLEDraftNpuGraphRunner,
+            "cuda": EAGLEDraftCudaGraphRunner,
+        }
         # Capture draft
         if self.speculative_num_steps > 1:
             tic = time.perf_counter()
@@ -221,22 +229,29 @@ class EagleDraftWorker(BaseDraftWorker):
             logger.info(
                 f"Capture draft cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
             )
-            self.cuda_graph_runner = EAGLEDraftCudaGraphRunner(self)
+            self.cuda_graph_runner = Device2DraftCudaGraphRunner[
+                self.target_worker.device
+            ](self)
             after_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(
                 f"Capture draft cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
             )
 
+        Device2ExtendCudaGraphRunner = {
+            "npu": EAGLEDraftExtendNpuGraphRunner,
+            "cuda": EAGLEDraftExtendCudaGraphRunner,
+        }
         # Capture extend
-        if self.draft_extend_attn_backend:
+        # FIXME cuda not support draft_extend capture
+        if self.draft_extend_attn_backend and _is_npu:
             tic = time.perf_counter()
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(
                 f"Capture draft extend cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
             )
-            self.cuda_graph_runner_for_draft_extend = EAGLEDraftExtendCudaGraphRunner(
-                self
-            )
+            self.cuda_graph_runner_for_draft_extend = Device2ExtendCudaGraphRunner[
+                self.target_worker.device
+            ](self)
             after_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(
                 f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
@@ -259,7 +274,10 @@ class EagleDraftWorker(BaseDraftWorker):
                 forward_batch,
             )
         else:
-            if self.speculative_num_steps > 1:
+            if (
+                not forward_batch.forward_mode.is_idle()
+                and self.speculative_num_steps > 1
+            ):
                 # Skip attention backend init for 1-step draft,
                 # `draft_forward` only does sample in this case.
                 self.draft_attn_backend.init_forward_metadata(forward_batch)
@@ -267,6 +285,13 @@ class EagleDraftWorker(BaseDraftWorker):
                 forward_batch
             )
 
+        if model_worker_batch.forward_mode.is_idle():
+            return EagleVerifyInput.create_idle_input(
+                self.topk,
+                self.speculative_num_steps,
+                self.speculative_num_draft_tokens,
+            )
+
         # Build tree mask
         # Directly write to cuda graph buffers for verify attn
         tree_mask_buf, position_buf = (
@@ -408,13 +433,14 @@ class EagleDraftWorker(BaseDraftWorker):
             next_token_ids: Next token ids generated from the target forward.
         """
         # Construct input_ids
-        pt = 0
-        for i, extend_len in enumerate(batch.extend_seq_lens):
-            input_ids = batch.input_ids[pt : pt + extend_len]
-            batch.input_ids[pt : pt + extend_len] = torch.cat(
-                (input_ids[1:], next_token_ids[i].reshape(1))
-            )
-            pt += extend_len
+        if not batch.forward_mode.is_idle():
+            pt = 0
+            for i, extend_len in enumerate(batch.extend_seq_lens):
+                input_ids = batch.input_ids[pt : pt + extend_len]
+                batch.input_ids[pt : pt + extend_len] = torch.cat(
+                    (input_ids[1:], next_token_ids[i].reshape(1))
+                )
+                pt += extend_len
 
         # Construct spec_info
         next_draft_input = EagleDraftInput(
@@ -422,7 +448,11 @@ class EagleDraftWorker(BaseDraftWorker):
             verified_id=next_token_ids,
             new_seq_lens=batch.seq_lens,
             allocate_lens=batch.seq_lens,
+            # draft mode is same with decode mode, only 1 num token per batch
+            num_tokens_per_batch=1,
+            num_tokens_for_logprob_per_batch=1,
         )
+
         batch.spec_info = next_draft_input
 
         # Run forward
@@ -443,6 +473,8 @@ class EagleDraftWorker(BaseDraftWorker):
         # Batch 2: Draft extend
         draft_input = EagleDraftInput(
             hidden_states=batch_result.logits_output.hidden_states,
+            num_tokens_per_batch=self.speculative_num_steps + 1,
+            num_tokens_for_logprob_per_batch=1,
         )
         select_index = (
             torch.arange(len(batch.seq_lens), device=self.device)
@@ -458,6 +490,7 @@ class EagleDraftWorker(BaseDraftWorker):
                 batch_result.next_token_ids,
                 self.speculative_num_draft_tokens,
                 self.draft_runner,
+                self.cuda_graph_runner_for_draft_extend,
             )
 
         if self.plan_stream:
@@ -466,9 +499,18 @@ class EagleDraftWorker(BaseDraftWorker):
             )
 
         # Run draft extend batch in the main compute stream
-        draft_logits_output = self.draft_runner.model.forward(
-            forward_batch.input_ids, forward_batch.positions, forward_batch
+        can_cuda_graph = (
+            self.cuda_graph_runner_for_draft_extend
+            and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch)
         )
+        if can_cuda_graph:
+            draft_logits_output = self.cuda_graph_runner_for_draft_extend.replay(
+                forward_batch
+            )
+        else:
+            draft_logits_output, _ = self.draft_runner.forward(
+                forward_batch, skip_attn_backend_init=True
+            )
 
         # Reorganize the spec info for the next batch
         draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[
@@ -551,16 +593,10 @@ class EAGLEWorkerV2(BaseSpecWorker):
         pass
 
     def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
-        if model_worker_batch.forward_mode.is_decode():
-            draft_input: EagleDraftInput = model_worker_batch.spec_info
-            assert draft_input.is_draft_input()
-            verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch)
-            assert verify_input.is_verify_input()
-            model_worker_batch.spec_info = verify_input
-            batch_output = self.verify(model_worker_batch, draft_input.allocate_lens)
-            self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output)
-            return batch_output
-        else:
+        if (
+            model_worker_batch.forward_mode.is_extend()
+            or model_worker_batch.is_extend_in_batch
+        ):
             # Target prefill
             model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
             batch_output = self.target_worker.forward_batch_generation(
@@ -575,6 +611,22 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 batch_output.next_token_ids,
             )
             return batch_output
+        else:
+            if model_worker_batch.spec_info is None:
+                model_worker_batch.spec_info = EagleDraftInput.create_idle_input(
+                    device=self.device,
+                    hidden_size=self.target_worker.model_config.hidden_size,
+                    dtype=self.target_worker.model_config.dtype,
+                    topk=self.topk,
+                    capture_hidden_mode=CaptureHiddenMode.LAST,
+                )
+            draft_input: EagleDraftInput = model_worker_batch.spec_info
+            verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch)
+            assert verify_input.is_verify_input()
+            model_worker_batch.spec_info = verify_input
+            batch_output = self.verify(model_worker_batch, draft_input.allocate_lens)
+            self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output)
+            return batch_output
 
     def verify(
         self,
@@ -605,7 +657,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
 
         # Correct some buffers due to the overlap plan
         if self.plan_stream:
-            torch.get_device_module().current_stream().wait_stream(self.plan_stream)
+            torch.get_device_module(self.device).current_stream().wait_stream(
+                self.plan_stream
+            )
 
             # Some values such as custom_mask and position depend on the output of draft,
             # so the previous plan step used the wrong values. Here, we need to run the related
@@ -640,14 +694,17 @@ class EAGLEWorkerV2(BaseSpecWorker):
         verify_done = torch.get_device_module(self.device).Event()
         verify_done.record()
 
-        all_verified_id = predict[accept_index]
-        verified_id = torch.empty_like(accept_length, dtype=torch.int32)
-        fill_new_verified_id[(bs,)](
-            all_verified_id,
-            accept_length,
-            verified_id,
-            self.speculative_num_draft_tokens,
-        )
+        if not batch.forward_mode.is_idle():
+            all_verified_id = predict[accept_index]
+            verified_id = torch.empty_like(accept_length, dtype=torch.int32)
+            fill_new_verified_id[(bs,)](
+                all_verified_id,
+                accept_length,
+                verified_id,
+                self.speculative_num_draft_tokens,
+            )
+        else:
+            verified_id = torch.empty((0,), device=self.device, dtype=torch.int32)
 
         # Construct the next draft input
         next_draft_input = EagleDraftInput(

python/sglang/srt/speculative/spec_utils.py

@@ -468,7 +468,7 @@ def select_top_k_tokens(
 
         if hidden_states.shape[0] > 0:
             selected_input_index = topk_cs_index.flatten() // topk + torch.arange(
-                0, hidden_states.shape[0], step=topk, device="cuda"
+                0, hidden_states.shape[0], step=topk, device=topk_index.device
             ).repeat_interleave(topk)
             hidden_states = hidden_states[selected_input_index, :]