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

[sgl] reduce specdec cpu overhead (#23321)

Browse Source
This commit is contained in:
Tarushii Goel
2026-05-04 15:02:03 -07:00
committed by GitHub
parent 4743cf6051
commit d7c93e183b
6 changed files with 93 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/allocator_npu.py
+9 -5
View File
@@ -33,17 +33,21 @@ class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
extend_num_tokens: int,
num_new_pages: int = None,
):
if self.debug_mode:
assert torch.all(
(last_loc + 1) % self.page_size == prefix_lens % self.page_size
)
num_new_pages = (
(seq_lens + self.roundup) // self.page_size
- (prefix_lens + self.roundup) // self.page_size
).sum()
num_new_pages_item = num_new_pages.item()
if num_new_pages is None:
num_new_pages_tensor = (
(seq_lens + self.roundup) // self.page_size
- (prefix_lens + self.roundup) // self.page_size
).sum()
num_new_pages_item = num_new_pages_tensor.item()
else:
num_new_pages_item = num_new_pages
if self.need_sort and num_new_pages_item > len(self.free_pages):
self.merge_and_sort_free()
python/sglang/srt/mem_cache/allocator.py
+7 -5
View File
@@ -410,6 +410,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
extend_num_tokens: int,
num_new_pages: int = None,
):
if self.debug_mode:
assert torch.all(
@@ -439,11 +440,12 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
if self.debug_mode:
assert len(torch.unique(out_indices)) == len(out_indices)
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
prefix_lens=prefix_lens_cpu,
)
if num_new_pages is None:
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
prefix_lens=prefix_lens_cpu,
)
if num_new_pages > len(self.free_pages):
return None
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+1
View File
@@ -307,6 +307,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
seq_lens_cpu,
hisparse_last_loc,
len(logical_indices),
num_new_pages=num_new_pages,
)
assert (
hisparse_indices is not None
python/sglang/srt/mem_cache/swa_memory_pool.py
+2
View File
@@ -395,6 +395,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
seq_lens_cpu,
last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
prefix_lens,
@@ -403,6 +404,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
seq_lens_cpu,
swa_last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
python/sglang/srt/speculative/eagle_info_v2.py
+14 -12
View File
@@ -116,23 +116,25 @@ class EagleDraftInputV2Mixin:
)
page_size = batch.token_to_kv_pool_allocator.page_size
cur_kv_lens_cpu = []
nxt_kv_lens_cpu = []
num_needed_tokens = 0
alloc_len_per_decode = get_alloc_len_per_decode()
for r in batch.reqs:
# Over-allocation happens here
x = r.kv_committed_len + 2 * alloc_len_per_decode - r.kv_allocated_len
cur_kv_lens_cpu.append(r.kv_allocated_len)
nxt_kv_lens_cpu.append(r.kv_allocated_len + x)
num_needed_tokens += x
r.kv_allocated_len += x
double_alloc = alloc_len_per_decode + alloc_len_per_decode
cur_kv_lens = [0] * bs
nxt_kv_lens = [0] * bs
num_needed_tokens = 0
for i, r in enumerate(batch.reqs):
cur = r.kv_allocated_len
nxt = r.kv_committed_len + double_alloc
cur_kv_lens[i] = cur
nxt_kv_lens[i] = nxt
num_needed_tokens += nxt - cur
r.kv_allocated_len = nxt
r.decode_batch_idx += 1
# Pre-claim bonus slot here (like normal decode); resolve subtracts 1.
r.kv_committed_len += 1
cur_kv_lens_cpu = torch.tensor(cur_kv_lens_cpu, dtype=torch.int32, device="cpu")
nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens_cpu, dtype=torch.int32, device="cpu")
cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")
nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu")
if page_size == 1:
out_cache_loc = alloc_token_slots(batch.tree_cache, num_needed_tokens)
python/sglang/srt/speculative/spec_utils.py
+60 -39
View File
@@ -469,7 +469,63 @@ def create_num_accepted_drafts_filter(
return num_accepted_drafts_filter
def _select_top_k_tokens_first(
topk_p: torch.Tensor,
topk_index: torch.Tensor,
hidden_states: Optional[torch.Tensor],
topk: int,
):
input_ids = topk_index.flatten()
if hidden_states is not None:
hidden_states = hidden_states.repeat_interleave(topk, dim=0)
tree_info = (
topk_p.unsqueeze(1), # (b, 1, topk)
topk_index, # (b, topk)
torch.arange(-1, topk, dtype=torch.long, device=input_ids.device).expand(
topk_p.shape[0], -1
), # (b, topk + 1) — expand avoids the allocation of repeat
)
return input_ids, hidden_states, topk_p, tree_info
@torch.compile(dynamic=True, disable=_is_npu)
def _select_top_k_tokens_later(
i: int,
topk_p: torch.Tensor,
topk_index: torch.Tensor,
hidden_states: torch.Tensor,
scores: torch.Tensor,
topk: int,
):
topk_sq = topk * topk
expand_scores = scores.unsqueeze(2) * topk_p.view(-1, topk, topk)
# (b, topk, 1) * (b, topk, topk) -> (b, topk, topk)
topk_cs_p, topk_cs_index = fast_topk(
expand_scores.flatten(start_dim=1), topk, dim=-1
) # (b, topk)
topk_index = topk_index.view(-1, topk_sq)
input_ids = torch.gather(topk_index, 1, topk_cs_index).flatten()
if hidden_states.shape[0] > 0:
flat_cs = topk_cs_index.flatten()
batch_offsets = torch.arange(
0, hidden_states.shape[0], step=topk, device=flat_cs.device
)
selected_input_index = flat_cs // topk + batch_offsets.repeat_interleave(topk)
hidden_states = hidden_states[selected_input_index]
tree_info = (
expand_scores, # (b, topk, topk)
topk_index, # (b, topk * topk)
topk_cs_index + (topk_sq * (i - 1) + topk), # (b, topk)
)
return input_ids, hidden_states, topk_cs_p, tree_info
def select_top_k_tokens(
i: int,
topk_p: torch.Tensor,
@@ -479,45 +535,10 @@ def select_top_k_tokens(
topk: int,
):
if i == 0:
# The first step after extend
input_ids = topk_index.flatten()
if hidden_states is not None:
hidden_states = hidden_states.repeat_interleave(topk, dim=0)
scores = topk_p # shape: (b, topk)
tree_info = (
topk_p.unsqueeze(1), # shape: (b, 1, topk)
topk_index, # shape: (b, topk)
torch.arange(-1, topk, dtype=torch.long, device=input_ids.device)
.unsqueeze(0)
.repeat(topk_p.shape[0], 1), # shape: (b, topk + 1)
)
else:
# The later decode steps
expand_scores = torch.mul(
scores.unsqueeze(2), topk_p.reshape(-1, topk, topk)
) # (b, topk, 1) x (b, topk ,topk) -> (b, topk, topk)
topk_cs_p, topk_cs_index = fast_topk(
expand_scores.flatten(start_dim=1), topk, dim=-1
) # (b, topk)
scores = topk_cs_p # shape: (b, topk)
topk_index = topk_index.reshape(-1, topk**2)
input_ids = torch.gather(topk_index, index=topk_cs_index, dim=1).flatten()
if hidden_states.shape[0] > 0:
selected_input_index = topk_cs_index.flatten() // topk + torch.arange(
0, hidden_states.shape[0], step=topk, device=topk_index.device
).repeat_interleave(topk)
hidden_states = hidden_states[selected_input_index, :]
tree_info = (
expand_scores, # shape: (b, topk, topk)
topk_index, # shape: (b, topk * topk)
topk_cs_index + (topk**2 * (i - 1) + topk), # shape: (b, topk)
)
return input_ids, hidden_states, scores, tree_info
return _select_top_k_tokens_first(topk_p, topk_index, hidden_states, topk)
return _select_top_k_tokens_later(
i, topk_p, topk_index, hidden_states, scores, topk
)
def generate_simulated_accept_index(