diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c822f6820..3903f0be5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -279,6 +279,9 @@ class Envs: # Tool-Call behavior SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF) + # Ngram + SGLANG_NGRAM_FORCE_GREEDY_VERIFY = EnvBool(False) + # fmt: on diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index b0a5c8d99..5ba756aa3 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -15,6 +15,7 @@ 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 @@ -408,7 +409,9 @@ class NgramVerifyInput(SpecInput): ) # Sample tokens. Force greedy sampling on AMD - is_all_greedy = sampling_info.is_all_greedy + 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). " @@ -419,11 +422,7 @@ class NgramVerifyInput(SpecInput): self._greedy_verify(batch, logits_output) else: # NOTE: Compared with greedy_verify, the performance of _sampling_verify is relatively poor. - logger.warning( - "Currently Ngram speculative sampling forces to use greedy verification." - ) - self._greedy_verify(batch, logits_output) - # self._sampling_verify(batch, logits_output, sampling_info) + self._sampling_verify(batch, logits_output, sampling_info) self._fill_requests(batch, logits_output) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index b6a1bf554..d6ad689c2 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -5,6 +5,7 @@ import numpy as np import torch from sgl_kernel.speculative import reconstruct_indices_from_tree_mask +from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs from sglang.srt.managers.schedule_batch import ScheduleBatch @@ -15,11 +16,9 @@ from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache from sglang.srt.speculative.ngram_info import NgramVerifyInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.srt.utils import get_bool_env_var logger = logging.getLogger(__name__) -RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB") USE_FULL_MASK = True @@ -216,7 +215,7 @@ class NGRAMWorker: # acceptance indices are the indices in a "flattened" batch. # dividing it to num_draft_tokens will yield the actual batch index. temperatures = temperatures[accepted_indices // num_draft_tokens] - if RETURN_ORIGINAL_LOGPROB: + if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get(): logprobs = torch.nn.functional.log_softmax( logits_output.next_token_logits, dim=-1 )