sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/speculative/eagle_utils.py

@@ -4,14 +4,128 @@ from typing import List, Optional
 
 import torch
 
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip, is_npu
 
-if is_cuda() or is_hip():
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+
+if _is_cuda or _is_hip:
     from sgl_kernel import (
         build_tree_kernel_efficient as sgl_build_tree_kernel_efficient,
     )
 
 
+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],
@@ -114,20 +228,41 @@ def build_tree_kernel_efficient(
             (bs * num_verify_tokens,), device=device, dtype=torch.long
         )
 
-    sgl_build_tree_kernel_efficient(
-        parent_list,
-        top_scores_index,
-        seq_lens,
-        tree_mask,
-        positions,
-        retrive_index,
-        retrive_next_token,
-        retrive_next_sibling,
-        topk,
-        spec_steps,
-        num_verify_tokens,
-        tree_mask_mode,
-    )
+    if _is_npu:
+        (
+            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,
+            num_verify_tokens,
+            tree_mask_mode,
+            bs,
+        )
+    else:
+        sgl_build_tree_kernel_efficient(
+            parent_list,
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrive_index,
+            retrive_next_token,
+            retrive_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
     return (
         tree_mask,
         positions,
@@ -136,3 +271,113 @@ def build_tree_kernel_efficient(
         retrive_next_sibling,
         draft_tokens,
     )
+
+
+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,
+    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,
+):
+    if _is_cuda or _is_hip:
+        from sgl_kernel import verify_tree_greedy
+
+        verify_tree_greedy(
+            predicts=predicts,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=accept_token_num,  # mutable
+            candidates=candidates,
+            retrive_index=retrive_index,
+            retrive_next_token=retrive_next_token,
+            retrive_next_sibling=retrive_next_sibling,
+            target_predict=target_predict,
+        )
+
+    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
+            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

sglang/srt/speculative/eagle_worker.py

@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple
 import torch
 
 from sglang.srt.distributed import get_tp_group
+from sglang.srt.layers.dp_attention import get_attention_tp_group
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -52,9 +53,12 @@ from sglang.srt.utils import (
     get_available_gpu_memory,
     get_bool_env_var,
     is_cuda,
+    is_npu,
     next_power_of_2,
 )
 
+_is_npu = is_npu()
+
 if is_cuda():
     from sgl_kernel import segment_packbits  # noqa: F401
 
@@ -117,7 +121,11 @@ class EAGLEWorker(TpModelWorker):
             self.hot_token_id = None
 
         # Init draft worker
-        with empty_context():
+        if server_args.enable_dp_attention and self.speculative_algorithm.is_eagle3():
+            ctx = draft_tp_context(get_attention_tp_group())
+        else:
+            ctx = empty_context()
+        with ctx:
             super().__init__(
                 server_args=server_args,
                 gpu_id=gpu_id,
@@ -200,7 +208,7 @@ class EAGLEWorker(TpModelWorker):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if self.server_args.disable_cuda_graph:
+        if self.server_args.disable_cuda_graph or _is_npu:
             return
 
         # Capture draft
@@ -940,7 +948,7 @@ class EAGLEWorker(TpModelWorker):
         draft_input.hidden_states = logits_output.hidden_states
 
 
-@torch.compile(dynamic=True)
+@torch.compile(dynamic=True, disable=_is_npu)
 def get_last_loc_large_page_size_top_k_1(
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,

sglang/srt/speculative/eagle_worker_v2.py

@@ -4,7 +4,6 @@ import time
 from typing import List, Optional, Tuple
 
 import torch
-from torch.cuda import Stream as CudaStream
 
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -38,18 +37,21 @@ from sglang.srt.utils.common import (
     empty_context,
     fast_topk,
     get_available_gpu_memory,
+    is_npu,
     next_power_of_2,
 )
 
+_is_npu = is_npu()
+
 logger = logging.getLogger(__name__)
 
 
 def _get_plan_stream(
     device: str,
-) -> Tuple[Optional[CudaStream], contextlib.AbstractContextManager]:
+) -> Tuple[any, contextlib.AbstractContextManager]:
     if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
-        plan_stream: CudaStream = torch.get_device_module(device).Stream()
-        plan_stream_ctx = torch.cuda.stream(plan_stream)
+        plan_stream = torch.get_device_module(device).Stream()
+        plan_stream_ctx = torch.get_device_module(device).stream(plan_stream)
         return plan_stream, plan_stream_ctx
     else:
         return None, contextlib.nullcontext()
@@ -206,7 +208,7 @@ class EagleDraftWorker(BaseDraftWorker):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if self.server_args.disable_cuda_graph:
+        if self.server_args.disable_cuda_graph or _is_npu:
             return
 
         # Capture draft
@@ -456,7 +458,9 @@ class EagleDraftWorker(BaseDraftWorker):
             )
 
         if self.plan_stream:
-            torch.cuda.current_stream().wait_stream(self.plan_stream)
+            torch.get_device_module(self.device).current_stream().wait_stream(
+                self.plan_stream
+            )
 
         # Run draft extend batch in the main compute stream
         draft_logits_output = self.draft_runner.model.forward(
@@ -577,7 +581,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
         # Since batch.seq_lens is allocated in another stream, we need
         # record_stream() to prevent pytorch gc and reuse the gpu memory
         # while forward_stream is still running.
-        batch.seq_lens.record_stream(torch.cuda.current_stream())
+        batch.seq_lens.record_stream(
+            torch.get_device_module(self.device).current_stream()
+        )
 
         # Parse args
         verify_input: EagleVerifyInput = batch.spec_info
@@ -596,7 +602,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
 
         # Correct some buffers due to the overlap plan
         if self.plan_stream:
-            torch.cuda.current_stream().wait_stream(self.plan_stream)
+            torch.get_device_module().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
@@ -628,7 +634,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
             accept_index,
         ) = verify_input.sample(batch, logits_output)
         new_seq_lens = batch.seq_lens + accept_length
-        verify_done = torch.cuda.Event()
+        verify_done = torch.get_device_module(self.device).Event()
         verify_done.record()
 
         all_verified_id = predict[accept_index]