chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Feature] Support DeepSeek MTP on NPU (#11897)

Browse Source 
Co-authored-by: liupeng374 <liupeng374@huawei.com>
This commit is contained in:
Even Zhou
2025-10-30 21:56:56 +08:00
committed by GitHub
parent cafebef154
commit ce6b17c0f9
16 changed files with 850 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/speculative/draft_utils.py
View File

@@ -49,6 +49,7 @@ class DraftBackendFactory:
"trtllm_mha": self._create_trtllm_mha_decode_backend,
"trtllm_mla": self._create_trtllm_mla_decode_backend,
"nsa": self._create_nsa_decode_backend,
"ascend": self._create_ascend_decode_backend,
}
return self._create_backend(
@@ -72,6 +73,7 @@ class DraftBackendFactory:
"trtllm_mha": self._create_trtllm_mha_prefill_backend,
"trtllm_mla": self._create_trtllm_mla_prefill_backend,
"nsa": self._create_nsa_prefill_backend,
"ascend": self._create_ascend_prefill_backend,
}
backend_name = (
"decode_attention_backend"
@@ -173,6 +175,15 @@ class DraftBackendFactory:
self.draft_model_runner, self.topk, self.speculative_num_steps
)
def _create_ascend_decode_backend(self):
from sglang.srt.layers.attention.ascend_backend import (
AscendAttnMultiStepDraftBackend,
)
return AscendAttnMultiStepDraftBackend(
self.draft_model_runner, self.topk, self.speculative_num_steps
)
def _create_flashinfer_prefill_backend(self):
if not get_global_server_args().use_mla_backend:
from sglang.srt.layers.attention.flashinfer_backend import (
@@ -219,6 +230,11 @@ class DraftBackendFactory:
return TRTLLMMLABackend(self.draft_model_runner, skip_prefill=False)
def _create_ascend_prefill_backend(self):
from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend
return AscendAttnBackend(self.draft_model_runner)
def _create_flashmla_prefill_backend(self):
logger.warning(
"flashmla prefill backend is not yet supported for draft extend."

78
python/sglang/srt/speculative/eagle_info.py
View File

@@ -24,12 +24,13 @@ from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin,
)
from sglang.srt.speculative.eagle_utils import verify_tree_greedy_func
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_LEN,
TREE_SPEC_KERNEL_AVAILABLE,
align_evict_mask_to_page_size,
assign_req_to_token_pool,
assign_req_to_token_pool_func,
create_accept_length_filter,
create_extend_after_decode_spec_info,
filter_finished_cache_loc_kernel,
@@ -37,17 +38,16 @@ from sglang.srt.speculative.spec_utils import (
get_src_tgt_cache_loc,
get_target_cache_loc,
)
from sglang.srt.utils import is_cuda, is_hip, next_power_of_2
from sglang.srt.utils import is_cuda, is_npu, next_power_of_2
_is_npu = is_npu()
if is_cuda():
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
verify_tree_greedy,
)
elif is_hip():
from sgl_kernel import verify_tree_greedy
logger = logging.getLogger(__name__)
@@ -77,18 +77,22 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
@classmethod
def create_idle_input(cls, topk: int, spec_steps: int, num_verify_tokens: int):
if not _is_npu:
device = "cuda"
else:
device = "npu"
return cls(
draft_token=torch.empty((0,), dtype=torch.long, device="cuda"),
custom_mask=torch.full((0,), True, dtype=torch.bool, device="cuda"),
positions=torch.empty((0,), dtype=torch.int64, device="cuda"),
draft_token=torch.empty((0,), dtype=torch.long, device=device),
custom_mask=torch.full((0,), True, dtype=torch.bool, device=device),
positions=torch.empty((0,), dtype=torch.int64, device=device),
retrive_index=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
(0, num_verify_tokens), -1, dtype=torch.long, device=device
),
retrive_next_token=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
(0, num_verify_tokens), -1, dtype=torch.long, device=device
),
retrive_next_sibling=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
(0, num_verify_tokens), -1, dtype=torch.long, device=device
),
retrive_cum_len=None,
topk=topk,
@@ -134,14 +138,13 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
self.last_loc = last_loc
bs = batch.batch_size()
assign_req_to_token_pool[(bs,)](
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
end_offset,
batch.out_cache_loc,
batch.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
bs,
)
def generate_attn_arg_prefill(
@@ -151,16 +154,17 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
paged_kernel_lens_sum: int,
req_to_token: torch.Tensor,
):
device = req_pool_indices.device
batch_size = len(req_pool_indices)
qo_indptr = torch.arange(
0,
(1 + batch_size) * self.draft_token_num,
step=self.draft_token_num,
dtype=torch.int32,
device="cuda",
device=device,
)
cum_kv_seq_len = torch.zeros(
(batch_size + 1,), dtype=torch.int32, device="cuda"
(batch_size + 1,), dtype=torch.int32, device=device
)
paged_kernel_lens = paged_kernel_lens + self.draft_token_num
@@ -169,7 +173,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
kv_indices = torch.empty(
paged_kernel_lens_sum + self.draft_token_num * batch_size,
dtype=torch.int32,
device="cuda",
device=device,
)
create_flashinfer_kv_indices_triton[(batch_size,)](
req_to_token,
@@ -226,11 +230,11 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
predict = torch.empty(predict_shape, dtype=torch.int32, device="cuda")
predict = torch.empty(predict_shape, dtype=torch.int32, device=batch.device)
accept_index = torch.full(
(bs, self.spec_steps + 1), -1, dtype=torch.int32, device="cuda"
(bs, self.spec_steps + 1), -1, dtype=torch.int32, device=batch.device
)
accept_length = torch.empty((bs,), dtype=torch.int32, device="cuda")
accept_length = torch.empty((bs,), dtype=torch.int32, device=batch.device)
if bs != len(sampling_info):
sampling_info = copy.deepcopy(sampling_info)
@@ -254,7 +258,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
linear_penalty = torch.zeros(
(bs, logits_output.next_token_logits.shape[1]),
dtype=torch.float32,
device="cuda",
device=batch.device,
)
sampling_info.apply_logits_bias(linear_penalty)
logits_output.next_token_logits.add_(
@@ -276,11 +280,10 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
"Falling back to greedy verification."
)
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE:
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE or _is_npu:
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
verify_tree_greedy(
predict, accept_index, accept_length = verify_tree_greedy_func(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=accept_length, # mutable
@@ -289,7 +292,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
retrive_next_token=self.retrive_next_token,
retrive_next_sibling=self.retrive_next_sibling,
target_predict=target_predict,
topk=self.topk,
)
else:
# apply temperature and get target probs
expanded_temperature = torch.repeat_interleave(
@@ -315,14 +320,16 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
target_probs = target_probs.reshape(bs, self.draft_token_num, -1)
draft_probs = torch.zeros(
target_probs.shape, dtype=torch.float32, device="cuda"
target_probs.shape, dtype=torch.float32, device=batch.device
)
# coins for rejection sampling
coins = torch.rand_like(candidates, dtype=torch.float32, device="cuda")
coins = torch.rand_like(
candidates, dtype=torch.float32, device=batch.device
)
# coins for final sampling
coins_for_final_sampling = torch.rand(
(bs,), dtype=torch.float32, device="cuda"
(bs,), dtype=torch.float32, device=batch.device
)
tree_speculative_sampling_target_only(
predicts=predict, # mutable
@@ -468,14 +475,13 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
if not has_finished:
if page_size == 1 or self.topk == 1:
batch.out_cache_loc = batch.out_cache_loc[accept_index]
assign_req_to_token_pool[(bs,)](
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + accept_length + 1,
batch.out_cache_loc,
batch.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
bs,
)
else:
batch.out_cache_loc = tgt_cache_loc
@@ -501,14 +507,13 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
)
else:
if page_size == 1 or self.topk == 1:
assign_req_to_token_pool[(bs,)](
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + accept_length + 1,
batch.out_cache_loc[accept_index],
batch.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
bs,
)
batch.seq_lens.add_(accept_length + 1)
batch.seq_lens_cpu.add_(accept_length_cpu + 1)
@@ -695,17 +700,18 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
paged_kernel_lens_sum: int,
req_to_token: torch.Tensor,
):
device = req_pool_indices.device
bs = self.accept_length.numel()
qo_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device="cuda")
qo_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=device)
qo_indptr[1:] = torch.cumsum(self.accept_length, dim=0)
cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device="cuda")
cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=device)
cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
if paged_kernel_lens_sum is None:
paged_kernel_lens_sum = cum_kv_seq_len[-1]
kv_indices = torch.empty(
paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
paged_kernel_lens_sum, dtype=torch.int32, device=device
)
create_flashinfer_kv_indices_triton[(bs,)](

93
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -23,11 +23,16 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.eagle_utils import verify_tree_greedy_func
from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_LEN,
generate_simulated_accept_index,
)
from sglang.srt.utils.common import fast_topk, is_cuda, is_hip, next_power_of_2
from sglang.srt.utils.common import fast_topk, is_cuda, is_hip, is_npu, next_power_of_2
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
if TYPE_CHECKING:
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -41,11 +46,8 @@ if is_cuda():
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
verify_tree_greedy,
)
from sgl_kernel.top_k import fast_topk
elif is_hip():
from sgl_kernel import verify_tree_greedy
@triton.jit
@@ -78,7 +80,7 @@ def assign_draft_cache_locs_page_size_1(
@dataclass
class EagleDraftInputV2Mixin:
def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
from sglang.srt.speculative.spec_utils import assign_req_to_token_pool
from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
bs = batch.batch_size()
@@ -112,15 +114,15 @@ class EagleDraftInputV2Mixin:
extend_num_tokens,
)
assign_req_to_token_pool[(bs,)](
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
self.allocate_lens,
new_allocate_lens,
out_cache_loc,
batch.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
bs,
)
self.allocate_lens = new_allocate_lens
# FIXME(lsyin): make this sync optional
@@ -199,22 +201,16 @@ class EagleVerifyInputV2Mixin:
bs = len(batch.req_pool_indices)
batch.input_ids = self.draft_token
device = batch.input_ids.device
batch.out_cache_loc = torch.empty(
(bs * self.draft_token_num,),
dtype=torch.int64,
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,
)
assign_extend_cache_locs[(bs,)](
batch.req_pool_indices,
req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + self.draft_token_num,
batch.out_cache_loc,
req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
)
# Get a forward batch
batch.forward_mode = ForwardMode.TARGET_VERIFY
batch.capture_hidden_mode = CaptureHiddenMode.FULL
@@ -258,11 +254,10 @@ class EagleVerifyInputV2Mixin:
accept_length = torch.empty((bs,), dtype=torch.int32, device=device)
# Sample tokens
if sampling_info.is_all_greedy:
if sampling_info.is_all_greedy or _is_npu:
target_predict = torch.argmax(next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
verify_tree_greedy(
predict, accept_index, accept_length = verify_tree_greedy_func(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=accept_length, # mutable
@@ -271,6 +266,7 @@ class EagleVerifyInputV2Mixin:
retrive_next_token=self.retrive_next_token,
retrive_next_sibling=self.retrive_next_sibling,
target_predict=target_predict,
topk=self.topk,
)
else:
# Apply temperature and get target probs
@@ -338,7 +334,7 @@ class EagleVerifyInputV2Mixin:
return predict, accept_length, accept_index
@torch.compile(dynamic=True)
@torch.compile(dynamic=True, disable=_is_npu)
def select_top_k_tokens_tmp(
i: int,
topk_p: torch.Tensor,
@@ -456,3 +452,50 @@ def assign_extend_cache_locs(
tl.store(out_cache_ptr + save_offset, data, mask=mask)
load_offset += BLOCK_SIZE
save_offset += BLOCK_SIZE
def assign_extend_cache_locs_func(
req_pool_indices: torch.Tensor,
req_to_token: torch.Tensor,
start_offset: torch.Tensor,
end_offset: torch.Tensor,
batch_size: int,
draft_token_num: int,
device,
) -> torch.Tensor:
if _is_cuda or _is_hip:
out_cache_loc = torch.empty(
(batch_size * draft_token_num,),
dtype=torch.int64,
device=device,
)
assign_extend_cache_locs[(batch_size,)](
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
req_to_token.shape[1],
next_power_of_2(batch_size),
)
return out_cache_loc
elif _is_npu:
import sgl_kernel_npu # noqa: F401
out_cache_loc = torch.empty(
(batch_size * draft_token_num,),
dtype=torch.int32,
device=device,
)
torch.ops.npu.cache_loc_update(
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
)
out_cache_loc = out_cache_loc.to(dtype=torch.int64)
return out_cache_loc

277
python/sglang/srt/speculative/eagle_utils.py
View File

@@ -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

7
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -53,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
@@ -205,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
@@ -945,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,

24
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -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]

52
python/sglang/srt/speculative/spec_utils.py
View File

@@ -19,16 +19,22 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.environ import envs
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils import is_cuda, is_hip, is_npu, next_power_of_2
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
if TYPE_CHECKING:
from sglang.srt.speculative.eagle_info import EagleVerifyInput
if is_cuda():
if _is_cuda:
from sgl_kernel import fast_topk
elif is_hip():
elif _is_hip:
from sgl_kernel import fast_topk
else:
from sglang.srt.utils.common import fast_topk
logger = logging.getLogger(__name__)
@@ -39,7 +45,7 @@ SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get() # turn off if < 0
SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get()
TREE_TRAVERSE_TIME_THRESHOLD = 1 # TODO: set this properly
TREE_SPEC_KERNEL_AVAILABLE = is_cuda() # This kernel is only available for CUDA now
TREE_SPEC_KERNEL_AVAILABLE = _is_cuda # This kernel is only available for CUDA now
@triton.jit
@@ -103,6 +109,36 @@ def assign_req_to_token_pool(
load_offset += BLOCK_SIZE
def assign_req_to_token_pool_func(
req_pool_indices: torch.Tensor,
req_to_token: torch.Tensor,
start_offset: torch.Tensor,
end_offset: torch.Tensor,
out_cache_loc: torch.Tensor,
batch_size: int,
):
if _is_cuda or _is_hip:
assign_req_to_token_pool[(batch_size,)](
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
req_to_token.shape[1],
next_power_of_2(batch_size),
)
elif _is_npu:
import sgl_kernel_npu # noqa: F401
torch.ops.npu.cache_loc_assign(
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
)
@triton.jit
def assign_draft_cache_locs(
req_pool_indices,
@@ -331,7 +367,7 @@ def get_target_cache_loc(
)
@torch.compile(dynamic=True)
@torch.compile(dynamic=True, disable=_is_npu)
def get_src_tgt_cache_loc(
seq_lens: torch.Tensor,
out_cache_loc: torch.Tensor,
@@ -381,7 +417,7 @@ def filter_finished_cache_loc_kernel(
)
@torch.compile(dynamic=True)
@torch.compile(dynamic=True, disable=_is_npu)
def create_accept_length_filter(
accept_length: torch.Tensor,
unfinished_index_device: torch.Tensor,
@@ -395,7 +431,7 @@ def create_accept_length_filter(
return accept_length_filter
@torch.compile(dynamic=True)
@torch.compile(dynamic=True, disable=_is_npu)
def select_top_k_tokens(
i: int,
topk_p: torch.Tensor,
@@ -413,7 +449,7 @@ def select_top_k_tokens(
tree_info = (
topk_p.unsqueeze(1), # shape: (b, 1, topk)
topk_index, # shape: (b, topk)
torch.arange(-1, topk, dtype=torch.long, device="cuda")
torch.arange(-1, topk, dtype=torch.long, device=hidden_states.device)
.unsqueeze(0)
.repeat(topk_p.shape[0], 1), # shape: (b, topk + 1)
)