From 111009ea542e3430d8eed6e8e99dd8558982699f Mon Sep 17 00:00:00 2001 From: Siyuan Chen <41201609+SYChen123@users.noreply.github.com> Date: Wed, 10 Jun 2026 17:46:00 +0800 Subject: [PATCH] [Feature] [Ngram spec] Support ngram spec v2 (#17260) Co-authored-by: hnyls2002 Co-authored-by: Ratish P <114130421+Ratish1@users.noreply.github.com> --- .../sglang/srt/arg_groups/speculative_hook.py | 8 +- python/sglang/srt/managers/overlap_utils.py | 6 + python/sglang/srt/managers/schedule_batch.py | 1 + python/sglang/srt/mem_cache/common.py | 6 +- .../sglang/srt/model_executor/model_runner.py | 2 +- .../runner/decode_cuda_graph_runner.py | 2 +- .../sglang/srt/speculative/eagle_info_v2.py | 19 +- python/sglang/srt/speculative/ngram_info.py | 480 +++--------------- python/sglang/srt/speculative/ngram_worker.py | 272 +++++++--- python/sglang/srt/speculative/spec_info.py | 7 +- .../speculative_target_verify_runner.py | 2 +- test/registered/spec/test_spec_ngram.py | 5 +- test/registered/spec/test_spec_ngram_extra.py | 15 +- 13 files changed, 323 insertions(+), 502 deletions(-) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 2ea956fa9..ee350ad33 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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." ) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index ec7ce2add..f67b661db 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -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. diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index acd61e279..933015b2c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index f201c2257..6fb43be0c 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ae9a411a2..38059cfe6 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 601274ee6..38636296f 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 0cb891dd9..5d9f0720d 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 28dba86e1..7874cdab1 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -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) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index d53522d6b..c1a47e1bf 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -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, ) diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 3b12140a0..b1c5b9372 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -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 diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py index 29ea566cb..6ce5c18b1 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py @@ -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, diff --git a/test/registered/spec/test_spec_ngram.py b/test/registered/spec/test_spec_ngram.py index 02b21928a..24f360afc 100644 --- a/test/registered/spec/test_spec_ngram.py +++ b/test/registered/spec/test_spec_ngram.py @@ -2,15 +2,16 @@ import unittest from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.spec_server_kits import SpecLogprobKit from sglang.test.server_fixtures.ngram_fixture import NgramServerBase # Per-commit: Paged backend only. # - FA3 base test archived to test/manual/spec/test_spec_ngram_fa3.py # - Triton + Flashinfer moved to test_spec_ngram_extra.py -register_cuda_ci(est_time=254, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=400, stage="base-b", runner_config="1-gpu-large") -class TestNgramSpeculativeDecodingPaged(NgramServerBase, GSM8KMixin): +class TestNgramSpeculativeDecodingPaged(NgramServerBase, GSM8KMixin, SpecLogprobKit): attention_backend = "flashinfer" extra_args = ["--page-size", "64"] diff --git a/test/registered/spec/test_spec_ngram_extra.py b/test/registered/spec/test_spec_ngram_extra.py index 8a8039a13..b0cb4d44b 100644 --- a/test/registered/spec/test_spec_ngram_extra.py +++ b/test/registered/spec/test_spec_ngram_extra.py @@ -6,15 +6,24 @@ from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.eval_accuracy_kit import GSM8KMixin from sglang.test.server_fixtures.ngram_fixture import NgramServerBase -# Extra: Triton + Flashinfer NGRAM backends. Sibling per-commit file -# (test_spec_ngram.py) keeps the Paged variant. -register_cuda_ci(est_time=254, stage="extra-a", runner_config="1-gpu-large") +# Extra: Triton + Flashinfer NGRAM backends + non-overlap (sync V2) variant. +# Sibling per-commit file (test_spec_ngram.py) keeps the Paged variant. +register_cuda_ci(est_time=400, stage="extra-a", runner_config="1-gpu-large") class TestNgramSpeculativeDecodingTriton(NgramServerBase, GSM8KMixin): attention_backend = "triton" +class TestNgramSpeculativeDecodingNoOverlap(NgramServerBase, GSM8KMixin): + """Non-overlap path: the scheduler drives the V2 worker synchronously + (seq_lens advance via GenerationBatchResult.new_seq_lens, accepted tokens + come from req.output_ids instead of the spec_info splice).""" + + attention_backend = "flashinfer" + extra_args = ["--disable-overlap-schedule"] + + class TestNgramSpeculativeDecodingFlashinfer(NgramServerBase, GSM8KMixin): attention_backend = "flashinfer" extra_args = ["--speculative-ngram-external-sam-budget", "8"]