[SPEC] fix: use effective max draft tokens for adaptive spec initiali… (#26354)
Co-authored-by: maodoudou168 <maodoudou168@users.noreply.github.com>
This commit is contained in:
@@ -297,7 +297,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
# so we need to reserve the space for the draft tokens.
|
# so we need to reserve the space for the draft tokens.
|
||||||
self.num_reserved_tokens = max(
|
self.num_reserved_tokens = max(
|
||||||
server_args.speculative_eagle_topk * server_args.speculative_num_steps,
|
server_args.speculative_eagle_topk * server_args.speculative_num_steps,
|
||||||
server_args.speculative_num_draft_tokens,
|
server_args.max_speculative_num_draft_tokens,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
self.num_reserved_tokens = 0
|
self.num_reserved_tokens = 0
|
||||||
|
|||||||
@@ -240,7 +240,7 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
|||||||
|
|
||||||
spec_steps = server_args.speculative_num_steps or 1
|
spec_steps = server_args.speculative_num_steps or 1
|
||||||
spec_topk = server_args.speculative_eagle_topk or 1
|
spec_topk = server_args.speculative_eagle_topk or 1
|
||||||
spec_tokens = server_args.speculative_num_draft_tokens
|
spec_tokens = server_args.max_speculative_num_draft_tokens
|
||||||
page_size = server_args.page_size
|
page_size = server_args.page_size
|
||||||
|
|
||||||
if page_size == 1 or spec_topk == 1:
|
if page_size == 1 or spec_topk == 1:
|
||||||
|
|||||||
@@ -285,9 +285,7 @@ class ModelRunnerKVCacheMixin:
|
|||||||
# Initialize req_to_token_pool
|
# Initialize req_to_token_pool
|
||||||
if self.req_to_token_pool is None:
|
if self.req_to_token_pool is None:
|
||||||
# FIXME(lsyin): this is the temporary fix for the context length issue when using speculative decoding
|
# FIXME(lsyin): this is the temporary fix for the context length issue when using speculative decoding
|
||||||
max_spec_draft_tokens = (
|
max_spec_draft_tokens = self.server_args.max_speculative_num_draft_tokens
|
||||||
self.server_args.effective_max_speculative_num_draft_tokens()
|
|
||||||
)
|
|
||||||
extra_max_context_len = 4
|
extra_max_context_len = 4
|
||||||
if max_spec_draft_tokens is not None:
|
if max_spec_draft_tokens is not None:
|
||||||
extra_max_context_len += max_spec_draft_tokens
|
extra_max_context_len += max_spec_draft_tokens
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import random
|
import random
|
||||||
import tempfile
|
import tempfile
|
||||||
|
from functools import cached_property
|
||||||
from typing import Any, Callable, Dict, List, Literal, Optional, Union
|
from typing import Any, Callable, Dict, List, Literal, Optional, Union
|
||||||
|
|
||||||
from sglang.srt.arg_groups.argparse_actions import (
|
from sglang.srt.arg_groups.argparse_actions import (
|
||||||
@@ -590,8 +591,6 @@ class ServerArgs:
|
|||||||
speculative_moe_runner_backend: Optional[str] = None
|
speculative_moe_runner_backend: Optional[str] = None
|
||||||
speculative_moe_a2a_backend: Optional[str] = None
|
speculative_moe_a2a_backend: Optional[str] = None
|
||||||
speculative_draft_model_quantization: Optional[str] = None
|
speculative_draft_model_quantization: Optional[str] = None
|
||||||
speculative_adaptive: bool = False
|
|
||||||
speculative_adaptive_config: Optional[str] = None
|
|
||||||
speculative_skip_dp_mlp_sync: bool = False
|
speculative_skip_dp_mlp_sync: bool = False
|
||||||
|
|
||||||
# Speculative decoding (ngram)
|
# Speculative decoding (ngram)
|
||||||
@@ -605,6 +604,10 @@ class ServerArgs:
|
|||||||
speculative_ngram_external_corpus_max_tokens: int = 10000000
|
speculative_ngram_external_corpus_max_tokens: int = 10000000
|
||||||
enable_multi_layer_eagle: bool = False
|
enable_multi_layer_eagle: bool = False
|
||||||
|
|
||||||
|
# Adaptive speculative decoding
|
||||||
|
speculative_adaptive: bool = False
|
||||||
|
speculative_adaptive_config: Optional[str] = None
|
||||||
|
|
||||||
# Expert parallelism
|
# Expert parallelism
|
||||||
ep_size: int = 1
|
ep_size: int = 1
|
||||||
moe_a2a_backend: Literal[
|
moe_a2a_backend: Literal[
|
||||||
@@ -7050,8 +7053,9 @@ class ServerArgs:
|
|||||||
def enable_mamba_extra_buffer(self) -> bool:
|
def enable_mamba_extra_buffer(self) -> bool:
|
||||||
return self.mamba_scheduler_strategy == "extra_buffer"
|
return self.mamba_scheduler_strategy == "extra_buffer"
|
||||||
|
|
||||||
def effective_max_speculative_num_draft_tokens(self) -> Optional[int]:
|
@cached_property
|
||||||
"""Return the maximum draft-token count runtime speculative decoding may use."""
|
def max_speculative_num_draft_tokens(self) -> Optional[int]:
|
||||||
|
"""Return the maximum draft-token count speculative decoding may use."""
|
||||||
if self.speculative_num_draft_tokens is None:
|
if self.speculative_num_draft_tokens is None:
|
||||||
return None
|
return None
|
||||||
if not self.speculative_adaptive:
|
if not self.speculative_adaptive:
|
||||||
|
|||||||
Reference in New Issue
Block a user