diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 7f6f6a010..4e22d1f83 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -1,5 +1,5 @@ import logging -from typing import List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import torch import torch.distributed as dist @@ -31,6 +31,8 @@ logger = logging.getLogger(__name__) SYNC_TOKEN_IDS_ACROSS_TP = get_bool_env_var("SYNC_TOKEN_IDS_ACROSS_TP") SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") +_CUSTOM_SAMPLER_FACTORIES: Dict[str, Callable[[], "Sampler"]] = {} +_BUILT_IN_SAMPLING_BACKENDS = {"flashinfer", "pytorch", "ascend"} class Sampler(nn.Module): @@ -268,6 +270,42 @@ class Sampler(nn.Module): ) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs) +def register_sampler_backend(backend: str, factory: Callable[[], "Sampler"]) -> None: + """Register a custom sampler factory for a backend string.""" + + if not backend: + raise ValueError("backend must be a non-empty string") + + from sglang.srt.server_args import SAMPLING_BACKEND_CHOICES + + if backend in _CUSTOM_SAMPLER_FACTORIES: + logger.warning("Overriding existing sampler factory for backend '%s'", backend) + SAMPLING_BACKEND_CHOICES.add(backend) + _CUSTOM_SAMPLER_FACTORIES[backend] = factory + + +def create_sampler(backend: Optional[str] = None) -> "Sampler": + """Create a sampler honoring custom backend registrations.""" + + server_args = get_global_server_args() + backend = backend or (server_args.sampling_backend if server_args else None) + + if backend in _CUSTOM_SAMPLER_FACTORIES: + sampler = _CUSTOM_SAMPLER_FACTORIES[backend]() + if not isinstance(sampler, Sampler): + raise TypeError( + f"Custom sampler factory for backend '{backend}' must return a Sampler" + ) + return sampler + + if backend is None or backend in _BUILT_IN_SAMPLING_BACKENDS: + return Sampler() + + raise ValueError( + f"Unknown sampling backend '{backend}'. Register it via register_sampler_backend()." + ) + + def top_k_top_p_min_p_sampling_from_probs_torch( probs: torch.Tensor, top_ks: torch.Tensor, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 0fb96a700..41044f06c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -99,7 +99,7 @@ from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype -from sglang.srt.layers.sampler import Sampler +from sglang.srt.layers.sampler import create_sampler from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model from sglang.srt.lora.lora_manager import LoRAManager from sglang.srt.lora.lora_registry import LoRARef @@ -451,7 +451,7 @@ class ModelRunner: else None ) # Load the model - self.sampler = Sampler() + self.sampler = create_sampler() self.load_model() if ( diff --git a/python/sglang/srt/models/hunyuan.py b/python/sglang/srt/models/hunyuan.py index 7c6fd9e48..87f162867 100644 --- a/python/sglang/srt/models/hunyuan.py +++ b/python/sglang/srt/models/hunyuan.py @@ -40,7 +40,7 @@ from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.sampler import Sampler +from sglang.srt.layers.sampler import create_sampler from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -603,7 +603,7 @@ class HunYuanMoEV1ForCausalLM(nn.Module): logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor(config, logit_scale=logit_scale) - self.sampler = Sampler() + self.sampler = create_sampler() def forward( self, diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index cff9419b7..20f020baa 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -89,8 +89,15 @@ class SamplingBatchInfo: ) sampling_seed = ( torch.tensor( - [r.sampling_params.sampling_seed for r in reqs], - dtype=torch.int32, + [ + ( + r.sampling_params.sampling_seed + if r.sampling_params.sampling_seed is not None + else 42 + ) + for r in reqs + ], + dtype=torch.int64, device=device, ) if enable_deterministic @@ -173,8 +180,13 @@ class SamplingBatchInfo: device=device, logit_bias=logit_bias, ) + ret.adjusted_from_schedule_batch(batch, vocab_size) return ret + # placeholder for override + def adjusted_from_schedule_batch(self, batch: ScheduleBatch, vocab_size: int): + pass + def __len__(self): return len(self.temperatures) diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 34e5252c8..2acf565b5 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -58,7 +58,7 @@ class SamplingParams: custom_params: Optional[Dict[str, Any]] = None, stream_interval: Optional[int] = None, logit_bias: Optional[Dict[str, float]] = None, - sampling_seed: int = 42, + sampling_seed: Optional[int] = None, ) -> None: self.max_new_tokens = max_new_tokens self.stop_strs = stop @@ -146,8 +146,6 @@ class SamplingParams: f"logit_bias must has keys in [0, {vocab_size - 1}], got " f"{token_id}." ) - if self.sampling_seed is None: - raise ValueError("sampling_seed should not be None") grammars = [ self.json_schema, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7e8648031..e54f1c9a2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -70,6 +70,7 @@ logger = logging.getLogger(__name__) # Define constants +SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"} LOAD_FORMAT_CHOICES = [ "auto", "pt", @@ -3220,7 +3221,7 @@ class ServerArgs: parser.add_argument( "--sampling-backend", type=str, - choices=["flashinfer", "pytorch", "ascend"], + choices=SAMPLING_BACKEND_CHOICES, default=ServerArgs.sampling_backend, help="Choose the kernels for sampling layers.", )