sglang 0.2.13__py3-none-any.whl → 0.2.14__py3-none-any.whl

sglang/srt/models/mixtral_quant.py

@@ -45,6 +45,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -333,6 +334,7 @@ class QuantMixtralForCausalLM(nn.Module):
         self.model = MixtralModel(config, quant_config=quant_config)
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
 
     @torch.no_grad()
     def forward(
@@ -343,9 +345,11 @@ class QuantMixtralForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/qwen.py

@@ -39,6 +39,7 @@ from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -251,6 +252,7 @@ class QWenLMHeadModel(nn.Module):
         vocab_size = ((config.vocab_size + 63) // 64) * 64
         self.lm_head = ParallelLMHead(vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
 
     @torch.no_grad()
     def forward(
@@ -260,10 +262,11 @@ class QWenLMHeadModel(nn.Module):
         input_metadata: InputMetadata,
     ):
         hidden_states = self.transformer(input_ids, positions, input_metadata)
-        next_tokens = self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
-        return next_tokens
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/qwen2.py

@@ -38,7 +38,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.pooler import Pooler, PoolingType
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 Qwen2Config = None
@@ -275,6 +277,8 @@ class Qwen2ForCausalLM(nn.Module):
         self.model = Qwen2Model(config, quant_config=quant_config)
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
+        self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
 
     @torch.no_grad()
     def forward(
@@ -283,11 +287,17 @@ class Qwen2ForCausalLM(nn.Module):
         positions: torch.Tensor,
         input_metadata: InputMetadata,
         input_embeds: torch.Tensor = None,
+        get_embedding: bool = False,
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
-        return self.logits_processor(
-            input_ids, hidden_states, self.lm_head.weight, input_metadata
-        )
+        if not get_embedding:
+            logits_output = self.logits_processor(
+                input_ids, hidden_states, self.lm_head.weight, input_metadata
+            )
+            sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+            return sample_output, logits_output
+        else:
+            return self.pooler(hidden_states, input_metadata)
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/qwen2_moe.py

@@ -35,10 +35,8 @@ from vllm.model_executor.layers.linear import (
     ReplicatedLinear,
     RowParallelLinear,
 )
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.sampler import Sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -49,6 +47,7 @@ from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -366,6 +365,7 @@ class Qwen2MoeForCausalLM(nn.Module):
             config.vocab_size, config.hidden_size, quant_config=quant_config
         )
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
 
     @torch.no_grad()
     def forward(
@@ -376,20 +376,11 @@ class Qwen2MoeForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
-
-    def compute_logits(
-        self,
-        input_ids: torch.Tensor,
-        hidden_states: torch.Tensor,
-        input_metadata: InputMetadata,
-    ) -> torch.Tensor:
-        logits = self.logits_processor(
-            input_ids, hidden_states, self.lm_head.weight, input_metadata
-        )
-        return logits
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/stablelm.py

@@ -40,6 +40,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -249,6 +250,7 @@ class StableLmForCausalLM(nn.Module):
         self.model = StableLMEpochModel(config, quant_config=quant_config)
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
 
     @torch.no_grad()
     def forward(
@@ -259,9 +261,11 @@ class StableLmForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/openai_api/adapter.py

@@ -17,6 +17,7 @@ limitations under the License.
 
 import asyncio
 import json
+import logging
 import os
 import time
 import uuid
@@ -64,6 +65,8 @@ from sglang.srt.openai_api.protocol import (
     UsageInfo,
 )
 
+logger = logging.getLogger(__name__)
+
 chat_template_name = None
 
 
@@ -120,7 +123,7 @@ def create_streaming_error_response(
 def load_chat_template_for_openai_api(tokenizer_manager, chat_template_arg):
     global chat_template_name
 
-    print(f"Use chat template: {chat_template_arg}")
+    logger.info(f"Use chat template: {chat_template_arg}")
     if not chat_template_exists(chat_template_arg):
         if not os.path.exists(chat_template_arg):
             raise RuntimeError(
@@ -276,6 +279,12 @@ async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRe
             request_data = json.loads(line)
             file_request_list.append(request_data)
             body = request_data["body"]
+
+            # Although streaming is supported for standalone completions, it is not supported in
+            # batch mode (multiple completions in single request).
+            if body.get("stream", False):
+                raise ValueError("Streaming requests are not supported in batch mode")
+
             if end_point == "/v1/chat/completions":
                 all_requests.append(ChatCompletionRequest(**body))
             elif end_point == "/v1/completions":
@@ -346,7 +355,7 @@ async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRe
         }
 
     except Exception as e:
-        print("error in SGLang:", e)
+        logger.error("error in SGLang:", e)
         # Update batch status to "failed"
         retrieve_batch = batch_storage[batch_id]
         retrieve_batch.status = "failed"
@@ -383,20 +392,33 @@ async def v1_retrieve_file_content(file_id: str):
     return StreamingResponse(iter_file(), media_type="application/octet-stream")
 
 
-def v1_generate_request(all_requests):
+def v1_generate_request(all_requests: List[CompletionRequest]):
     prompts = []
     sampling_params_list = []
     return_logprobs = []
+    logprob_start_lens = []
     top_logprobs_nums = []
-    first_prompt_type = type(all_requests[0].prompt)
 
+    # NOTE: with openai API, the prompt's logprobs are always not computed
+    first_prompt_type = type(all_requests[0].prompt)
     for request in all_requests:
-        prompt = request.prompt
         assert (
-            type(prompt) == first_prompt_type
+            type(request.prompt) == first_prompt_type
         ), "All prompts must be of the same type in file input settings"
-        prompts.append(prompt)
+        if len(all_requests) > 1 and request.n > 1:
+            raise ValueError(
+                "Parallel sampling is not supported for completions from files"
+            )
+        if request.echo and request.logprobs:
+            logger.warning(
+                "Echo is not compatible with logprobs. "
+                "To compute logprobs of input prompt, please use SGLang /request API."
+            )
+
+    for request in all_requests:
+        prompts.append(request.prompt)
         return_logprobs.append(request.logprobs is not None and request.logprobs > 0)
+        logprob_start_lens.append(-1)
         top_logprobs_nums.append(
             request.logprobs if request.logprobs is not None else 0
         )
@@ -416,14 +438,11 @@ def v1_generate_request(all_requests):
                 "ignore_eos": request.ignore_eos,
             }
         )
-        if len(all_requests) > 1 and request.n > 1:
-            raise ValueError(
-                "Parallel sampling is not supported for completions from files"
-            )
 
     if len(all_requests) == 1:
         prompt = prompts[0]
         sampling_params_list = sampling_params_list[0]
+        logprob_start_lens = logprob_start_lens[0]
         return_logprobs = return_logprobs[0]
         top_logprobs_nums = top_logprobs_nums[0]
         if isinstance(prompt, str) or isinstance(prompt[0], str):
@@ -441,6 +460,7 @@ def v1_generate_request(all_requests):
         sampling_params=sampling_params_list,
         return_logprob=return_logprobs,
         top_logprobs_num=top_logprobs_nums,
+        logprob_start_len=logprob_start_lens,
         return_text_in_logprobs=True,
         stream=all_requests[0].stream,
     )
@@ -580,27 +600,45 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
     if adapted_request.stream:
 
         async def generate_stream_resp():
-            stream_buffer = ""
-            n_prev_token = 0
+            stream_buffers = {}
+            n_prev_tokens = {}
+            prompt_tokens = {}
+            completion_tokens = {}
             try:
                 async for content in tokenizer_manager.generate_request(
                     adapted_request, raw_request
                 ):
+                    index = content["index"]
+
+                    stream_buffer = stream_buffers.get(index, "")
+                    n_prev_token = n_prev_tokens.get(index, 0)
+
                     text = content["text"]
-                    prompt_tokens = content["meta_info"]["prompt_tokens"]
-                    completion_tokens = content["meta_info"]["completion_tokens"]
+                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
+                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
 
                     if not stream_buffer:  # The first chunk
                         if request.echo:
                             if isinstance(request.prompt, str):
                                 # for the case of single str prompts
                                 prompts = request.prompt
-                            elif isinstance(request.prompt, list) and isinstance(
-                                request.prompt[0], int
-                            ):
-                                prompts = tokenizer_manager.tokenizer.decode(
-                                    request.prompt, skip_special_tokens=True
-                                )
+                            elif isinstance(request.prompt, list):
+                                if isinstance(request.prompt[0], str):
+                                    # for the case of multiple str prompts
+                                    prompts = request.prompt[index // request.n]
+                                elif isinstance(request.prompt[0], int):
+                                    # for the case of single token ids prompt
+                                    prompts = tokenizer_manager.tokenizer.decode(
+                                        request.prompt, skip_special_tokens=True
+                                    )
+                                elif isinstance(request.prompt[0], list) and isinstance(
+                                    request.prompt[0][0], int
+                                ):
+                                    # for the case of multiple token ids prompts
+                                    prompts = tokenizer_manager.tokenizer.decode(
+                                        request.prompt[index // request.n],
+                                        skip_special_tokens=True,
+                                    )
 
                             # Prepend prompt in response text.
                             text = prompts + text
@@ -637,7 +675,7 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
                     delta = text[len(stream_buffer) :]
                     stream_buffer = stream_buffer + delta
                     choice_data = CompletionResponseStreamChoice(
-                        index=0,
+                        index=index,
                         text=delta,
                         logprobs=logprobs,
                         finish_reason=format_finish_reason(
@@ -650,12 +688,24 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
                         choices=[choice_data],
                         model=request.model,
                     )
+
+                    stream_buffers[index] = stream_buffer
+                    n_prev_tokens[index] = n_prev_token
+
                     yield f"data: {chunk.model_dump_json()}\n\n"
                 if request.stream_options and request.stream_options.include_usage:
+                    total_prompt_tokens = sum(
+                        tokens
+                        for i, tokens in prompt_tokens.items()
+                        if i % request.n == 0
+                    )
+                    total_completion_tokens = sum(
+                        tokens for tokens in completion_tokens.values()
+                    )
                     usage = UsageInfo(
-                        prompt_tokens=prompt_tokens,
-                        completion_tokens=completion_tokens,
-                        total_tokens=prompt_tokens + completion_tokens,
+                        prompt_tokens=total_prompt_tokens,
+                        completion_tokens=total_completion_tokens,
+                        total_tokens=total_prompt_tokens + total_completion_tokens,
                     )
 
                     final_usage_chunk = CompletionStreamResponse(
@@ -694,12 +744,18 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
     return response
 
 
-def v1_chat_generate_request(all_requests, tokenizer_manager):
+def v1_chat_generate_request(
+    all_requests: List[ChatCompletionRequest], tokenizer_manager
+):
     input_ids = []
     sampling_params_list = []
     image_data_list = []
     return_logprobs = []
+    logprob_start_lens = []
     top_logprobs_nums = []
+
+    # NOTE: with openai API, the prompt's logprobs are always not computed
+
     for request in all_requests:
         # Prep the data needed for the underlying GenerateReqInput:
         #  - prompt: The full prompt string.
@@ -732,6 +788,7 @@ def v1_chat_generate_request(all_requests, tokenizer_manager):
             image_data = None
         input_ids.append(prompt_ids)
         return_logprobs.append(request.logprobs)
+        logprob_start_lens.append(-1)
         top_logprobs_nums.append(request.top_logprobs)
         sampling_params_list.append(
             {
@@ -758,17 +815,20 @@ def v1_chat_generate_request(all_requests, tokenizer_manager):
         sampling_params_list = sampling_params_list[0]
         image_data = image_data_list[0]
         return_logprobs = return_logprobs[0]
+        logprob_start_lens = logprob_start_lens[0]
         top_logprobs_nums = top_logprobs_nums[0]
     else:
         if isinstance(input_ids[0], str):
             prompt_kwargs = {"text": input_ids}
         else:
             prompt_kwargs = {"input_ids": input_ids}
+
     adapted_request = GenerateReqInput(
         **prompt_kwargs,
         image_data=image_data,
         sampling_params=sampling_params_list,
         return_logprob=return_logprobs,
+        logprob_start_len=logprob_start_lens,
         top_logprobs_num=top_logprobs_nums,
         stream=all_requests[0].stream,
         return_text_in_logprobs=True,
@@ -892,16 +952,23 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
     if adapted_request.stream:
 
         async def generate_stream_resp():
-            is_first = True
-
-            stream_buffer = ""
-            n_prev_token = 0
+            is_firsts = {}
+            stream_buffers = {}
+            n_prev_tokens = {}
+            prompt_tokens = {}
+            completion_tokens = {}
             try:
                 async for content in tokenizer_manager.generate_request(
                     adapted_request, raw_request
                 ):
-                    prompt_tokens = content["meta_info"]["prompt_tokens"]
-                    completion_tokens = content["meta_info"]["completion_tokens"]
+                    index = content["index"]
+
+                    is_first = is_firsts.get(index, True)
+                    stream_buffer = stream_buffers.get(index, "")
+                    n_prev_token = n_prev_tokens.get(index, 0)
+
+                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
+                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
                     if request.logprobs:
                         logprobs = to_openai_style_logprobs(
                             output_token_logprobs=content["meta_info"][
@@ -951,7 +1018,7 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
                         # First chunk with role
                         is_first = False
                         choice_data = ChatCompletionResponseStreamChoice(
-                            index=0,
+                            index=index,
                             delta=DeltaMessage(role="assistant"),
                             finish_reason=format_finish_reason(
                                 content["meta_info"]["finish_reason"]
@@ -969,7 +1036,7 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
                     delta = text[len(stream_buffer) :]
                     stream_buffer = stream_buffer + delta
                     choice_data = ChatCompletionResponseStreamChoice(
-                        index=0,
+                        index=index,
                         delta=DeltaMessage(content=delta),
                         finish_reason=format_finish_reason(
                             content["meta_info"]["finish_reason"]
@@ -981,12 +1048,25 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
                         choices=[choice_data],
                         model=request.model,
                     )
+
+                    is_firsts[index] = is_first
+                    stream_buffers[index] = stream_buffer
+                    n_prev_tokens[index] = n_prev_token
+
                     yield f"data: {chunk.model_dump_json()}\n\n"
                 if request.stream_options and request.stream_options.include_usage:
+                    total_prompt_tokens = sum(
+                        tokens
+                        for i, tokens in prompt_tokens.items()
+                        if i % request.n == 0
+                    )
+                    total_completion_tokens = sum(
+                        tokens for tokens in completion_tokens.values()
+                    )
                     usage = UsageInfo(
-                        prompt_tokens=prompt_tokens,
-                        completion_tokens=completion_tokens,
-                        total_tokens=prompt_tokens + completion_tokens,
+                        prompt_tokens=total_prompt_tokens,
+                        completion_tokens=total_completion_tokens,
+                        total_tokens=total_prompt_tokens + total_completion_tokens,
                     )
 
                     final_usage_chunk = ChatCompletionStreamResponse(