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

[Feature] [Ngram spec] Support ngram spec v2 (#17260)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Ratish P <114130421+Ratish1@users.noreply.github.com>
This commit is contained in:
Siyuan Chen
2026-06-10 02:46:00 -07:00
committed by GitHub
co-authored by hnyls2002 Ratish P
parent 255843d454
commit 111009ea54
13 changed files with 323 additions and 502 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+6 -2
View File
@@ -399,7 +399,6 @@ def _handle_ngram(server_args: "ServerArgs") -> None:
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
)
server_args.disable_overlap_schedule = True
server_args.enable_mixed_chunk = False
server_args.speculative_eagle_topk = server_args.speculative_ngram_max_bfs_breadth
if server_args.speculative_num_draft_tokens is None:
@@ -408,6 +407,11 @@ def _handle_ngram(server_args: "ServerArgs") -> None:
"speculative_num_draft_tokens is set to 12 by default for ngram speculative decoding. "
"You can override this by explicitly setting --speculative-num-draft-tokens."
)
if server_args.speculative_num_steps is None:
server_args.speculative_num_steps = (
server_args.speculative_num_draft_tokens
// server_args.speculative_eagle_topk
)
if server_args.speculative_ngram_external_corpus_path is not None:
if server_args.speculative_ngram_external_sam_budget <= 0:
raise ValueError(
@@ -428,7 +432,7 @@ def _handle_ngram(server_args: "ServerArgs") -> None:
f"speculative_num_draft_tokens - 1 ({server_args.speculative_num_draft_tokens - 1})."
)
logger.warning(
"The overlap scheduler and mixed chunked prefill are disabled because of "
"The mixed chunked prefill are disabled because of "
"using ngram speculative decoding."
)
python/sglang/srt/managers/overlap_utils.py
+6
View File
@@ -171,6 +171,9 @@ class FutureMap:
)
def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
if self.spec_algo.is_ngram():
# FIXME: remove once precomputed draft is supported.
return
draft_input: EagleDraftInput = batch.spec_info
if draft_input is None:
# FIXME(lsyin): only prefill; not compatible with mixed mode
@@ -256,6 +259,9 @@ class FutureMap:
future_indices: torch.Tensor,
payload: Union[torch.Tensor, EagleDraftInput],
) -> None:
if self.spec_algo.is_ngram():
# FIXME: remove once precomputed draft is supported.
return
indices = future_indices
if indices.shape[0] == 0:
# DP idle: payload is empty stub; lazy-init shape peek would IndexError.
python/sglang/srt/managers/schedule_batch.py
+1
View File
@@ -2789,6 +2789,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
return_logprob=self.return_logprob,
decoding_reqs=self.decoding_reqs,
spec_algorithm=self.spec_algorithm,
spec_info=self.spec_info,
global_num_tokens=self.global_num_tokens,
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
python/sglang/srt/mem_cache/common.py
+5 -1
View File
@@ -169,7 +169,11 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
spec_tokens = server_args.max_speculative_num_draft_tokens
page_size = server_args.page_size
if page_size == 1 or spec_topk == 1:
# NGRAM drafts are a flat candidate list written to contiguous slots after
# seq_lens (no per-topk page duplication), so the flat formula applies at
# any page_size.
is_ngram = (server_args.speculative_algorithm or "").upper() == "NGRAM"
if page_size == 1 or spec_topk == 1 or is_ngram:
return max(spec_steps * spec_topk, spec_tokens)
else:
# page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
python/sglang/srt/model_executor/model_runner.py
+1 -1
View File
@@ -2723,7 +2723,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
spec_info = NgramVerifyInput(
draft_token=None,
tree_mask=buffers.custom_mask,
custom_mask=buffers.custom_mask,
positions=None,
retrieve_index=None,
retrieve_next_token=None,
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+1 -1
View File
@@ -1155,7 +1155,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
spec_info = NgramVerifyInput(
draft_token=None,
tree_mask=self.buffers.custom_mask,
custom_mask=self.buffers.custom_mask,
positions=None,
retrieve_index=None,
retrieve_next_token=None,
python/sglang/srt/speculative/eagle_info_v2.py
+15 -4
View File
@@ -493,8 +493,17 @@ class EagleVerifyInputV2Mixin:
candidates = self.draft_token.reshape(bs, self.draft_token_num)
predict_shape = list(next_token_logits.shape)[:-1]
predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
# Longest root-to-leaf chain of the verify tree, incl. the root; bounds
# the accept_index row width. EAGLE trees are depth-bounded by the draft
# loop (spec_steps + 1); NGRAM trees are node-budgeted with no depth cap
# (a single corpus match can chain all draft_token_num nodes).
max_tree_depth = (
self.draft_token_num
if batch.spec_algorithm.is_ngram()
else self.spec_steps + 1
)
accept_index = torch.full(
(bs, self.spec_steps + 1), -1, dtype=torch.int32, device=device
(bs, max_tree_depth), -1, dtype=torch.int32, device=device
)
num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
@@ -511,7 +520,7 @@ class EagleVerifyInputV2Mixin:
retrieve_next_token=self.retrieve_next_token,
retrieve_next_sibling=self.retrieve_next_sibling,
target_predict=target_predict,
topk=self.topk,
topk=-1 if batch.spec_algorithm.is_ngram() else self.topk,
)
else:
# Apply temperature and get target probs
@@ -580,14 +589,16 @@ class EagleVerifyInputV2Mixin:
tp_group.broadcast(num_correct_drafts, src=0)
if SIMULATE_ACC_LEN > 0:
# Do simulation
# Do simulation. The helper builds (and returns) a replacement
# accept_index of width spec_steps + 1, so pass max_tree_depth - 1
# to keep the simulated width identical to the real one.
accept_index = generate_simulated_accept_index(
accept_index=accept_index,
predict=predict, # mutable
num_correct_drafts=num_correct_drafts, # mutable
simulate_acc_len=SIMULATE_ACC_LEN,
bs=bs,
spec_steps=self.spec_steps,
spec_steps=max_tree_depth - 1,
)
# `num_correct_drafts` stays drafts-only inside this function; the returned
python/sglang/srt/speculative/ngram_info.py
+70 -410
View File
@@ -1,123 +1,57 @@
from __future__ import annotations
import copy
import logging
from typing import Optional, Tuple
import torch
import triton
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.server_args import get_global_server_args
logger = logging.getLogger(__name__)
from dataclasses import dataclass
import torch.nn.functional as F
from sglang.srt.environ import envs
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import apply_custom_logit_processor
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
alloc_token_slots,
get_last_loc,
from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin,
)
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
from sglang.srt.speculative.spec_utils import (
TREE_SPEC_KERNEL_AVAILABLE,
assign_req_to_token_pool,
get_src_tgt_cache_loc,
get_target_cache_loc,
)
from sglang.srt.utils import is_cuda, is_hip, is_musa, next_power_of_2
if is_cuda() or is_musa():
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
@dataclass
class NgramVerifyInput(SpecInput):
class NgramVerifyInput(SpecInput, EagleDraftInputV2Mixin, EagleVerifyInputV2Mixin):
def __init__(
self,
draft_token: torch.Tensor,
tree_mask: torch.Tensor,
positions: torch.Tensor,
retrieve_index: torch.Tensor,
retrieve_next_token: torch.Tensor,
retrieve_next_sibling: torch.Tensor,
draft_token_num: int,
draft_token: torch.Tensor = None,
custom_mask: torch.Tensor = None,
positions: torch.Tensor = None,
retrieve_index: torch.Tensor = None,
retrieve_next_token: torch.Tensor = None,
retrieve_next_sibling: torch.Tensor = None,
draft_token_num: int = None,
grammar: BaseGrammarObject = None,
future_indices: Optional[torch.Tensor] = None,
new_seq_lens: Optional[torch.Tensor] = None,
accept_tokens: Optional[torch.Tensor] = None,
accept_lens: Optional[torch.Tensor] = None,
):
super().__init__(SpecInputType.NGRAM_VERIFY)
self.draft_token = draft_token
self.custom_mask = tree_mask
self.custom_mask = custom_mask
self.positions = positions
self.retrieve_index = retrieve_index
self.retrieve_next_token = retrieve_next_token
self.retrieve_next_sibling = retrieve_next_sibling
self.draft_token_num = draft_token_num
self.device = self.custom_mask.device
self.grammar = grammar
# Inputs for V2 overlap worker
self.future_indices = future_indices
self.new_seq_lens = new_seq_lens
self.accept_tokens = accept_tokens
self.accept_lens = accept_lens
self.device = (
custom_mask.device if custom_mask is not None else new_seq_lens.device
)
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
return self.draft_token_num, self.draft_token_num
def prepare_for_verify(self, batch: ScheduleBatch, page_size: int):
if batch.forward_mode.is_idle():
return
batch.input_ids = self.draft_token
if page_size == 1:
batch.out_cache_loc = alloc_token_slots(
batch.tree_cache,
len(batch.input_ids),
)
end_offset = batch.seq_lens + self.draft_token_num
else:
# TODO(lsyin): add prefix lens cpu here to support page size > 1
prefix_lens = batch.seq_lens
prefix_lens_cpu = batch.seq_lens_cpu
end_offset = prefix_lens + self.draft_token_num
end_offset_cpu = prefix_lens_cpu + self.draft_token_num
last_loc = get_last_loc(
batch.req_to_token_pool.req_to_token,
batch.req_pool_indices,
prefix_lens,
)
batch.out_cache_loc = alloc_paged_token_slots_extend(
batch.tree_cache,
prefix_lens,
prefix_lens_cpu,
end_offset,
end_offset_cpu,
last_loc,
len(batch.input_ids),
)
bs = batch.batch_size()
assign_req_to_token_pool[(bs,)](
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],
triton.next_power_of_2(bs),
)
def generate_attn_arg_prefill(
self,
req_pool_indices: torch.Tensor,
@@ -138,7 +72,9 @@ class NgramVerifyInput(SpecInput):
)
kv_indices = torch.empty(
cum_kv_seq_len[-1], dtype=torch.int32, device=self.device
paged_kernel_lens_sum + self.draft_token_num * bs,
dtype=torch.int32,
device=self.device,
)
create_flashinfer_kv_indices_triton[(bs,)](
@@ -150,328 +86,52 @@ class NgramVerifyInput(SpecInput):
kv_indices,
req_to_token.size(1),
)
return kv_indices, cum_kv_seq_len, self.qo_indptr, self.custom_mask
def _fill_requests(
self,
batch: ScheduleBatch,
logits_output: torch.Tensor,
):
accept_index_cpu = self.accept_indices.tolist()
predict_cpu = self.predict.tolist()
has_finished = False
think_end_id = batch.model_config.think_end_id
# Iterate every accepted token and check if req has finished after append the token
# should be checked BEFORE free kv cache slots
for i, (req, accept_index_row) in enumerate(zip(batch.reqs, accept_index_cpu)):
for j, idx in enumerate(accept_index_row):
if idx == -1:
break
id = predict_cpu[idx]
req.output_ids.append(id)
if req.require_reasoning and think_end_id is not None:
req.update_reasoning_tokens(id, think_end_id)
req.update_finish_state()
if req.finished():
has_finished = True
# set all tokens after finished token to -1 and break
self.accept_indices[i, j + 1 :] = -1
break
else:
if req.grammar is not None:
try:
req.grammar.accept_token(id)
except ValueError as e:
logger.info(
f"{i=}, {req=}\n"
f"{self.accept_indices=}\n"
f"{self.predict=}\n"
)
raise e
req.spec_verify_ct += 1
num_correct_drafts_this_req = (
sum(1 for idx in accept_index_row if idx != -1) - 1
)
req.spec_num_correct_drafts += num_correct_drafts_this_req
req.update_spec_correct_drafts_histogram(num_correct_drafts_this_req)
if has_finished:
self.num_correct_drafts = (self.accept_indices != -1).sum(dim=1) - 1
self.accept_indices = self.accept_indices[self.accept_indices != -1]
logits_output.next_token_logits = logits_output.next_token_logits[
self.accept_indices
]
if logits_output.hidden_states:
logits_output.hidden_states = logits_output.hidden_states[
self.accept_indices
]
self.accept_tokens = self.predict[self.accept_indices]
def _free_cache(
self,
batch: ScheduleBatch,
page_size: int,
num_correct_drafts_cpu: torch.Tensor,
):
bs = batch.batch_size()
# Free the KV cache for unaccepted tokens
if page_size == 1:
# TODO: boolean array index leads to a device sync. Remove it.
evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool)
evict_mask[self.accept_indices] = False
batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
batch.out_cache_loc = batch.out_cache_loc[self.accept_indices]
else:
# Shift the accepted tokens to the beginning.
# Only evict the last part
src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc(
batch.seq_lens,
batch.out_cache_loc,
self.accept_indices,
self.num_correct_drafts,
self.draft_token_num,
page_size,
)
to_free_slots = torch.empty(
(to_free_num_slots.sum().item(),),
dtype=torch.int64,
device=to_free_num_slots.device,
)
# out_cache_loc: [0 1 2, 3 4 5, 6 7 8]
# accept_index: [0 -1 2, 3 4 -1, 6 -1 -1]
# tgt_cache_loc: [0 1 , 3 4 , 6 ]
# to_free_slots: [ 2, 5, 7 8]
# to_free_slots also needs to be page-aligned without the first partial page
#
# split each row of out_cache_loc into two parts.
# 1. the first part goes to tgt_cache_loc. length = num_correct_drafts[i] + 1
# 2. the second part goes to to_free_slots.
get_target_cache_loc[(bs,)](
tgt_cache_loc,
to_free_slots,
self.num_correct_drafts,
to_free_num_slots,
batch.out_cache_loc,
self.draft_token_num,
next_power_of_2(self.draft_token_num),
next_power_of_2(bs),
)
# Free the kv cache
batch.token_to_kv_pool_allocator.free(to_free_slots)
# Copy the kv cache
batch.token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
tgt_cache_loc, src_cache_loc
)
batch.out_cache_loc = tgt_cache_loc
num_correct_drafts_list = num_correct_drafts_cpu.tolist()
for i, req in enumerate(batch.reqs):
req.kv_committed_len += num_correct_drafts_list[i] + 1
req.kv_allocated_len = req.kv_committed_len
assign_req_to_token_pool[(bs,)](
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + self.num_accept_tokens,
batch.out_cache_loc,
batch.req_to_token_pool.req_to_token.shape[1],
triton.next_power_of_2(bs),
# Pad custom_mask when CUDA graph pads batch size beyond the actual number of requests.
mask_numel = (
paged_kernel_lens_sum * self.draft_token_num
+ (self.draft_token_num**2) * bs
)
def _greedy_verify(
self,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
):
bs = batch.batch_size()
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
candidates = self.draft_token.reshape(bs, self.draft_token_num)
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accept_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
)
self.num_correct_drafts = torch.empty(
(bs,), dtype=torch.int32, device=self.device
)
verify_tree_greedy(
predicts=self.predict, # mutable
accept_index=self.accept_indices, # mutable
accept_token_num=self.num_correct_drafts, # mutable
candidates=candidates,
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
retrive_index=self.retrieve_index,
retrive_next_token=self.retrieve_next_token,
retrive_next_sibling=self.retrieve_next_sibling,
target_predict=target_predict,
)
def _sampling_verify(
self,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo,
):
bs = batch.batch_size()
candidates = self.draft_token.reshape(bs, self.draft_token_num)
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accept_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
)
self.num_correct_drafts = torch.empty(
(bs,), dtype=torch.int32, device=self.device
)
# apply temperature and get target probs
expanded_temperature = torch.repeat_interleave(
sampling_info.temperatures, self.draft_token_num, dim=0
) # (bs * draft_token_num, 1)
target_probs = F.softmax(
logits_output.next_token_logits / expanded_temperature, dim=-1
) # (bs * draft_token_num, vocab_size)
# NOTE: The test shows that top_p_renorm_prob and top_k_renorm_prob are the key factors
# contributing to the poor performance of _sampling_verify.
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(sampling_info.top_ks, self.draft_token_num, dim=0),
) # (bs * draft_token_num, vocab_size)
if sampling_info.need_top_p_sampling:
# logger.info("Using top-p sampling in speculative decoding verification.")
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, self.draft_token_num, dim=0
),
custom_mask = self.custom_mask
if custom_mask.numel() < mask_numel:
custom_mask = torch.cat(
[
custom_mask,
torch.full(
(mask_numel - custom_mask.numel(),),
True,
dtype=torch.bool,
device=self.device,
),
],
dim=0,
)
target_probs = target_probs.reshape(bs, self.draft_token_num, -1)
draft_probs = torch.zeros(
target_probs.shape, dtype=torch.float32, device=self.device
)
# coins for rejection sampling
coins = torch.rand_like(candidates, dtype=torch.float32, device=self.device)
# coins for final sampling
coins_for_final_sampling = torch.rand(
(bs,), dtype=torch.float32, device=self.device
)
tree_speculative_sampling_target_only(
predicts=self.predict, # mutable
accept_index=self.accept_indices, # mutable
accept_token_num=self.num_correct_drafts, # mutable
candidates=candidates.to(torch.int64),
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
retrive_index=self.retrieve_index.to(torch.int64),
retrive_next_token=self.retrieve_next_token.to(torch.int64),
retrive_next_sibling=self.retrieve_next_sibling.to(torch.int64),
uniform_samples=coins,
uniform_samples_for_final_sampling=coins_for_final_sampling,
target_probs=target_probs,
draft_probs=draft_probs,
threshold_single=get_global_server_args().speculative_accept_threshold_single,
threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
deterministic=True,
)
def verify(
self,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
page_size: int,
vocab_mask: Optional[torch.Tensor] = None, # For grammar
) -> torch.Tensor:
bs = self.retrieve_index.shape[0]
sampling_info = batch.sampling_info
if bs != len(sampling_info):
sampling_info = copy.deepcopy(sampling_info)
# NOTE: retrieve_index are the indices of the requests that are kept.
sampling_info.filter_batch(
self.retrieve_index.tolist(), self.retrieve_index
)
# Apply the custom logit processors if registered in the sampling info.
if sampling_info.has_custom_logit_processor:
apply_custom_logit_processor(
logits_output.next_token_logits,
sampling_info,
num_tokens_in_batch=self.draft_token_num,
)
# Apply penalty
if (
sampling_info.penalizer_orchestrator.is_required
or sampling_info.logit_bias is not None
):
# This is a relaxed version of penalties for speculative decoding.
sampling_info.penalizer_orchestrator.apply(
logits_output.next_token_logits, repeat=self.draft_token_num
)
if sampling_info.logit_bias is not None:
logits_output.next_token_logits.add_(
torch.repeat_interleave(
sampling_info.logit_bias, self.draft_token_num, dim=0
)
)
# Apply grammar mask
if vocab_mask is not None:
assert self.grammar is not None
self.grammar.apply_vocab_mask(
logits=logits_output.next_token_logits, vocab_mask=vocab_mask
)
# Sample tokens. Force greedy sampling on AMD
is_all_greedy = (
sampling_info.is_all_greedy or envs.SGLANG_NGRAM_FORCE_GREEDY_VERIFY.get()
)
if (not is_all_greedy) and (not TREE_SPEC_KERNEL_AVAILABLE):
logger.warning(
"Tree speculative sampling kernel unavailable (likely AMD/HIP build). "
"Falling back to greedy verification."
)
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE:
self._greedy_verify(batch, logits_output)
else:
# NOTE: Compared with greedy_verify, the performance of _sampling_verify is relatively poor.
self._sampling_verify(batch, logits_output, sampling_info)
self._fill_requests(batch, logits_output)
# Sync the bonus-included view after the kernel + `_fill_requests`
# finalize `num_correct_drafts`.
self.num_accept_tokens = self.num_correct_drafts + 1
num_correct_drafts_cpu = self.num_correct_drafts.cpu()
num_accept_tokens_cpu = num_correct_drafts_cpu + 1
num_correct_drafts = num_correct_drafts_cpu.sum().item()
self._free_cache(batch, page_size, num_correct_drafts_cpu)
batch.seq_lens.add_(self.num_accept_tokens)
batch.seq_lens_cpu.add_(num_accept_tokens_cpu)
# Keep seq_lens_sum in sync; attn backends size kv_indices from it.
batch.seq_lens_sum += int(num_accept_tokens_cpu.sum())
return logits_output, self.accept_tokens, num_correct_drafts
return kv_indices, cum_kv_seq_len, self.qo_indptr, custom_mask
def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True):
pass
if self.future_indices is not None:
self.future_indices = self.future_indices[new_indices]
if self.new_seq_lens is not None:
self.new_seq_lens = self.new_seq_lens[new_indices]
self.accept_tokens = self.accept_tokens.reshape(-1, self.draft_token_num)[
new_indices, :
]
self.accept_tokens = self.accept_tokens.flatten()
self.accept_lens = self.accept_lens[new_indices]
def merge_batch(self, spec_info: NgramVerifyInput):
pass
if self.future_indices is not None:
assert spec_info.future_indices is not None
self.future_indices = torch.cat(
(self.future_indices, spec_info.future_indices), dim=0
)
if self.new_seq_lens is not None:
assert spec_info.new_seq_lens is not None
self.new_seq_lens = torch.cat(
(self.new_seq_lens, spec_info.new_seq_lens), dim=0
)
self.accept_tokens = torch.cat(
(self.accept_tokens, spec_info.accept_tokens), dim=0
)
self.accept_lens = torch.cat((self.accept_lens, spec_info.accept_lens), dim=0)
python/sglang/srt/speculative/ngram_worker.py
+201 -71
View File
@@ -5,18 +5,24 @@ import numpy as np
import torch
from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1
from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.req_time_stats import set_time_batch
from sglang.srt.observability.trace import get_global_tracing_enabled
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
from sglang.srt.speculative.ngram_info import NgramVerifyInput
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import generate_token_bitmask
from sglang.srt.speculative.spec_utils import (
generate_token_bitmask,
move_accept_tokens_to_target_kvcache,
record_stream_for_v2_verify,
)
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_extend_cache_locs_func as assign_extend_cache_locs_func,
)
from sglang.srt.utils.async_probe import maybe_detect_inf, maybe_detect_nan
logger = logging.getLogger(__name__)
@@ -37,18 +43,31 @@ class NGRAMWorker:
nccl_port: int,
target_worker: TpModelWorker,
):
self.server_args = server_args
self.enable_overlap = not server_args.disable_overlap_schedule
self.target_worker = target_worker
self.model_runner = target_worker.model_runner
self.tp_rank = tp_rank
self.page_size = server_args.page_size
self.draft_token_num: int = server_args.speculative_num_draft_tokens
self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
self.max_batch_size = target_worker.max_running_requests
self.device = f"cuda:{gpu_id}" if gpu_id >= 0 else "cuda"
self._init_preallocated_tensors()
self.adaptive_controller = None
# rids of the last decode batch; used to erase corpus match state for
# requests that left the batch (see forward_batch_generation).
self._prev_decode_rids: set = set()
self.ngram_corpus = NgramCorpus(
min_bfs_breadth=server_args.speculative_ngram_min_bfs_breadth,
max_bfs_breadth=server_args.speculative_ngram_max_bfs_breadth,
@@ -82,6 +101,7 @@ class NGRAMWorker:
def clear_cache_pool(self):
self.ngram_corpus.reset()
self._prev_decode_rids = set()
def update_weights_from_tensor(self, recv_req):
# NGRAM has no draft weights of its own — the n-gram corpus is a CPU
@@ -162,22 +182,61 @@ class NGRAMWorker:
self.tree_mask[: bs * self.draft_token_num * self.draft_token_num]
)
def on_verify_complete_cpu(
self, num_correct_drafts_per_req: list[int], batch_size: int = 0
) -> None:
# Signature must match BaseSpecWorker.on_verify_complete_cpu; the
# result processor calls it with batch_size as a keyword argument.
if self.adaptive_controller is not None:
self.adaptive_controller.on_verify_complete(num_correct_drafts_per_req)
def _prepare_draft_tokens(
self, batch: ScheduleBatch
) -> tuple[np.ndarray, np.ndarray]:
bs = batch.batch_size()
bs = len(batch.reqs)
stride = self.draft_token_num
prev_token_ids, prev_accept_lens = (
batch.spec_info.accept_tokens,
batch.spec_info.accept_lens,
)
if not prev_token_ids.is_cpu:
prev_token_ids = prev_token_ids.cpu()
prev_accept_lens = prev_accept_lens.cpu()
# Worker-level staging: written here at draft prep, consumed by
# _update_ngram_corpus after verify within the same forward call.
self.prev_token_ids = prev_token_ids.tolist()
self.prev_accept_lens = prev_accept_lens.tolist()
self.ngram_corpus.synchronize()
req_ids = []
batch_tokens = []
total_lens = []
assert len(batch.reqs) == len(self.prev_accept_lens)
# Overlap mode processes results one iteration behind, so the last
# round's accepted tokens are not yet in req.output_ids and must be
# spliced in from spec_info. Sync mode and grammar batches process
# results before the next draft prep, so output_ids is already
# complete and splicing would duplicate the tail.
use_prev_tokens = self.enable_overlap and not batch.has_grammar
i = 0
for req in batch.reqs:
prev_tokens = (
self.prev_token_ids[i * stride : i * stride + self.prev_accept_lens[i]]
if use_prev_tokens
else []
)
check_token = self._efficient_concat_last_n(
req.origin_input_ids, req.output_ids, self.max_trie_depth
list(req.origin_input_ids),
list(req.output_ids[-self.max_trie_depth :]) + prev_tokens,
self.max_trie_depth,
)
req_ids.append(req.rid)
batch_tokens.append(check_token)
total_lens.append(len(req.origin_input_ids) + len(req.output_ids))
i += 1
total_lens.append(
len(req.origin_input_ids) + len(req.output_ids) + len(prev_tokens)
)
req_drafts, mask = self.ngram_corpus.batch_get(
req_ids, batch_tokens, total_lens
)
@@ -190,10 +249,13 @@ class NGRAMWorker:
return req_drafts, mask
def _prepare_for_speculative_decoding(self, batch: ScheduleBatch):
if batch.forward_mode.is_extend():
# Decode-only: extend goes through the plain target forward, and an
# IDLE batch must keep its forward_mode instead of being rewritten to
# TARGET_VERIFY below (relevant once DP attention support lands).
if not batch.forward_mode.is_decode():
return
bs = batch.batch_size()
bs = len(batch.reqs)
retrieve_index = self.retrieve_indexes_batch[bs]
retrieve_next_token = self.retrieve_next_token_batch[bs]
@@ -206,6 +268,7 @@ class NGRAMWorker:
tree_mask.copy_(torch.from_numpy(mask), non_blocking=True)
draft_tokens.copy_(torch.from_numpy(req_drafts), non_blocking=True)
# generate positions and some indices using tree_mask
reconstruct_indices_from_tree_mask(
tree_mask,
batch.seq_lens,
@@ -221,70 +284,99 @@ class NGRAMWorker:
# Testing shows about 8% performance improvement (the effect is roughly proportional to batch size).
if USE_FULL_MASK:
tree_mask = []
mask = mask.reshape(
batch.batch_size(), self.draft_token_num, self.draft_token_num
)
for i, req in enumerate(batch.reqs):
seq_len = len(req.origin_input_ids) + len(req.output_ids)
req_mask = torch.ones((self.draft_token_num, seq_len - 1)).cuda()
mask = mask.reshape(bs, self.draft_token_num, self.draft_token_num)
# TODO(siyuan): the for loop here leads to significant overhead in large batch size. Can be written into a kernel.
for i in range(bs):
seq_len = batch.seq_lens_cpu[i]
req_mask = torch.ones(
(self.draft_token_num, seq_len), device=self.device
)
req_mask = torch.cat(
(req_mask, torch.from_numpy(mask[i]).cuda()), dim=1
(
req_mask,
torch.from_numpy(mask[i]).to(
device=self.device, non_blocking=True
),
),
dim=1,
).to(torch.bool)
tree_mask.append(req_mask.flatten())
tree_mask = torch.cat(tree_mask, dim=0)
batch.spec_algorithm = SpeculativeAlgorithm.NGRAM
batch.forward_mode = ForwardMode.TARGET_VERIFY
batch.input_ids = draft_tokens
batch.out_cache_loc = assign_extend_cache_locs_func(
req_pool_indices=batch.req_pool_indices,
req_to_token=batch.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=self.device,
)
batch.spec_info = NgramVerifyInput(
draft_tokens,
tree_mask,
positions,
retrieve_index,
retrieve_next_token,
retrieve_next_sibling,
self.draft_token_num,
draft_token=draft_tokens,
custom_mask=tree_mask,
positions=positions,
retrieve_index=retrieve_index,
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
draft_token_num=self.draft_token_num,
)
batch.spec_info.prepare_for_verify(batch, self.page_size)
def _update_ngram_corpus(self, batch: ScheduleBatch):
batch_tokens = []
i, stride = 0, self.draft_token_num
# Same splice condition as _prepare_draft_tokens: only overlap mode
# has accepted tokens missing from req.output_ids.
use_prev_tokens = self.enable_overlap and not batch.has_grammar
for req in batch.reqs:
# FIXME: Whether to insert 'extend' into the cache or not, after testing,
# there is not much difference, so we will not insert it for now.
# if batch.forward_mode.is_extend():
# put_ids = req.origin_input_ids + req.output_ids
# else:
prev_tokens = (
self.prev_token_ids[i * stride : i * stride + self.prev_accept_lens[i]]
if use_prev_tokens
else []
)
put_ids = self._efficient_concat_last_n(
req.origin_input_ids, req.output_ids, self.max_trie_depth
list(req.origin_input_ids),
list(req.output_ids[-self.max_trie_depth :]) + prev_tokens,
self.max_trie_depth,
)
batch_tokens.append(put_ids)
i += 1
self.ngram_corpus.batch_put(batch_tokens)
def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
def forward_batch_generation(
self, batch: ScheduleBatch, on_publish=None
) -> GenerationBatchResult:
fwd_stream = torch.get_device_module(self.device).current_stream()
record_stream_for_v2_verify(batch, None, fwd_stream)
bs = len(batch.reqs)
set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
self._prepare_for_speculative_decoding(batch)
set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
spec_info = batch.spec_info
num_correct_drafts = 0
accept_lens = None
num_correct_drafts_per_req_cpu = None
verify_input: NgramVerifyInput = batch.spec_info
accept_lens = torch.ones(bs, dtype=torch.int32, device=self.device)
if batch.forward_mode.is_target_verify():
# Prepare grammar data on CPU if needed
if batch.has_grammar:
retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
draft_tokens_cpu = spec_info.draft_token.view(
spec_info.retrieve_next_token.shape
retrieve_next_token_cpu = verify_input.retrieve_next_token.cpu()
retrieve_next_sibling_cpu = verify_input.retrieve_next_sibling.cpu()
draft_tokens_cpu = verify_input.draft_token.view(
verify_input.retrieve_next_token.shape
).cpu()
set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
batch_result = self.target_worker.forward_batch_generation(
batch, is_verify=True
)
logits_output, can_run_cuda_graph = (
batch_result.logits_output,
batch_result.can_run_cuda_graph,
@@ -311,55 +403,93 @@ class NGRAMWorker:
# and will be applied to produce wrong results
batch.sampling_info.vocab_mask = None
logits_output, next_token_ids, num_correct_drafts = verify_input.verify(
batch, logits_output, self.page_size, vocab_mask
# Sample
maybe_detect_nan(
logits_output.next_token_logits, "verify: target model logits"
)
num_correct_drafts_per_req_cpu = (
verify_input.num_correct_drafts.cpu().tolist()
maybe_detect_inf(
logits_output.next_token_logits, "verify: target model logits"
)
(
predict,
accept_lens,
accept_index,
) = verify_input.sample(batch, logits_output, vocab_mask)
new_seq_lens = batch.seq_lens + accept_lens
accept_tokens = predict[accept_index].flatten()
next_token_ids = accept_tokens
if get_global_tracing_enabled():
for idx, req in enumerate(batch.reqs):
num_correct_drafts = (
verify_input.num_correct_drafts[idx].item()
if verify_input.num_correct_drafts is not None
else 0
)
req.time_stats.set_spec_verify_end_time(
num_correct_drafts=num_correct_drafts
)
# Store accept_lens (with bonus) for per-request metrics; downstream
# subtracts 1 to recover drafts-only counts.
accept_lens = verify_input.num_accept_tokens
# The KV mover expects drafts-only counts. NGRAM's
# accept_lens includes the bonus token, matching scheduler output.
num_correct_drafts_per_req = accept_lens - 1
move_accept_tokens_to_target_kvcache(
batch,
accept_index,
num_correct_drafts_per_req,
self.token_to_kv_pool_allocator,
)
if batch.return_logprob:
add_output_logprobs_for_spec_v1(batch, verify_input, logits_output)
# The last arg is the accept_index row width minus 1. NGRAM's
# accept_index is (bs, draft_token_num) -- the tree depth is not
# bounded by spec_steps like EAGLE's (bs, spec_steps + 1).
compute_spec_v2_logprobs(
batch,
logits_output,
predict,
accept_index,
self.draft_token_num - 1,
)
if on_publish is not None:
on_publish(new_seq_lens)
self._update_ngram_corpus(batch)
# Clean up per-request match state for finished/retracted requests.
# State entries are created in _prepare_draft_tokens and cleaned here.
# If a request is removed without passing through verify, the entry
# persists until reset(); this is acceptable because MatchState is small.
finished_req_ids = []
for req in batch.reqs:
if req.finished() or req.is_retracted:
finished_req_ids.append(req.rid)
if finished_req_ids:
self.ngram_corpus.erase_match_state(finished_req_ids)
# Erase match state of requests that left the decode batch.
# req.finished() is unusable here: under overlap it flips at result
# processing, one iteration after the request left the batch.
# The last batch's entries persist while idle (bounded, small).
cur_rids = {req.rid for req in batch.reqs}
departed_rids = self._prev_decode_rids - cur_rids
if departed_rids:
self.ngram_corpus.erase_match_state(list(departed_rids))
self._prev_decode_rids = cur_rids
batch.forward_mode = ForwardMode.DECODE
else:
batch_result = self.target_worker.forward_batch_generation(batch)
logits_output, next_token_ids, can_run_cuda_graph = (
logits_output, predict, can_run_cuda_graph = (
batch_result.logits_output,
batch_result.next_token_ids,
batch_result.can_run_cuda_graph,
)
new_seq_lens = batch.seq_lens.clone()
accept_tokens = torch.zeros(
bs, self.draft_token_num, dtype=torch.int32, device=self.device
)
accept_tokens[:, 0] = predict
accept_tokens = accept_tokens.flatten()
next_token_ids = predict
if on_publish is not None:
on_publish(new_seq_lens)
# Construct the next draft input
next_draft_input = NgramVerifyInput(
draft_token_num=self.draft_token_num,
new_seq_lens=new_seq_lens,
accept_tokens=accept_tokens,
accept_lens=accept_lens,
)
return GenerationBatchResult(
logits_output=logits_output,
next_token_ids=next_token_ids,
num_correct_drafts=num_correct_drafts,
num_correct_drafts_per_req_cpu=num_correct_drafts_per_req_cpu,
can_run_cuda_graph=can_run_cuda_graph,
accept_lens=accept_lens,
# Consumed by the non-overlap V2 scheduler branch to advance
# batch.seq_lens after the isolation restore; overlap mode relays
# it via on_publish instead.
new_seq_lens=new_seq_lens,
next_draft_input=next_draft_input,
speculative_num_draft_tokens=self.speculative_num_draft_tokens,
)
python/sglang/srt/speculative/spec_info.py
+1 -6
View File
@@ -146,7 +146,7 @@ class SpeculativeAlgorithm(Enum):
return None
def supports_spec_v2(self) -> bool:
return self.is_eagle() or self.is_standalone()
return self.is_eagle() or self.is_standalone() or self.is_ngram()
def get_num_tokens_per_bs_for_target_verify(
self, num_draft_tokens: int, is_draft_worker: bool
@@ -205,11 +205,6 @@ class SpeculativeAlgorithm(Enum):
return StandaloneWorkerV2
elif self.is_ngram():
if enable_overlap:
raise ValueError(
f"Speculative algorithm {self.name} does not support overlap worker creation."
)
from sglang.srt.speculative.ngram_worker import NGRAMWorker
return NGRAMWorker
python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py
+1 -1
View File
@@ -311,7 +311,7 @@ def _make_spec_verify_input(
if spec_kind == "ngram":
return NgramVerifyInput(
draft_token=batch.input_ids,
tree_mask=custom_mask,
custom_mask=custom_mask,
positions=batch.positions,
retrieve_index=retrieve_index,
retrieve_next_token=retrieve_next_token,