sglang 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

sglang/srt/models/roberta.py

@@ -6,7 +6,7 @@ from typing import Iterable, Optional, Tuple
 import torch
 from torch import nn
 
-from sglang.srt.layers.pooler import Pooler, PoolingType
+from sglang.srt.layers.pooler import CrossEncodingPooler, Pooler, PoolingType
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -16,6 +16,23 @@ from sglang.srt.models.bert import BertEncoder
 RobertaConfig = None
 
 
+# Adapted from transformers
+class RobertaClassificationHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, config: RobertaConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+
+    def forward(self, features, **kwargs):
+        x = features[0, :]  # take <s> token (equiv. to [CLS])
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.out_proj(x)
+        return x
+
+
 class RobertaEmbedding(nn.Module):
 
     def __init__(self, config: RobertaConfig):
@@ -51,8 +68,7 @@ class RobertaEmbedding(nn.Module):
         input_ids: torch.Tensor,
         seq_lens: torch.Tensor,
         position_ids: torch.Tensor,
-        inputs_embeds=None,
-        token_type_ids: Optional[torch.Tensor] = None,
+        forward_batch: ForwardBatch,
     ) -> torch.Tensor:
         input_shape = input_ids.size()
         inputs_embeds = self.word_embeddings(input_ids)
@@ -82,6 +98,8 @@ class RobertaEmbedding(nn.Module):
 
         # Position embeddings.
         position_embeddings = self.position_embeddings(position_ids)
+
+        token_type_ids = forward_batch.token_type_ids
         if token_type_ids is None:
             token_type_ids = torch.zeros(
                 input_shape, dtype=torch.long, device=inputs_embeds.device
@@ -93,20 +111,25 @@ class RobertaEmbedding(nn.Module):
         return embeddings
 
 
-class XLMRobertaModel(nn.Module):
+class XLMRobertaBaseModel(nn.Module):
     def __init__(
         self,
         *,
         config: RobertaConfig,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        add_pooling_layer: bool = False,
     ):
         super().__init__()
 
         self.config = config
         self.embeddings = RobertaEmbedding(config)
         self.encoder = BertEncoder(config=config, quant_config=quant_config, prefix="")
-        self.pooler = Pooler(pooling_type=PoolingType.CLS, normalize=True)
+        self.pooler = (
+            Pooler(pooling_type=PoolingType.CLS, normalize=True)
+            if add_pooling_layer
+            else None
+        )
 
     @torch.no_grad()
     def forward(
@@ -124,11 +147,12 @@ class XLMRobertaModel(nn.Module):
             input_ids=input_ids,
             position_ids=positions,
             seq_lens=forward_batch.seq_lens,
+            forward_batch=forward_batch,
         )
 
         hidden_states = self.encoder(hidden_states, forward_batch=forward_batch)
-        pooler_out = self.pooler(hidden_states, forward_batch)
-        return pooler_out
+
+        return hidden_states
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
@@ -141,7 +165,7 @@ class XLMRobertaModel(nn.Module):
         params_dict = dict(self.named_parameters())
         for name, loaded_weight in weights:
             name = name.replace("self", "self_attn")
-            if "pooler" in name:
+            if self.pooler is None and "pooler" in name:
                 continue
             for param_name, weight_name, shard_id in stacked_params_mapping:
 
@@ -175,4 +199,88 @@ def create_position_ids_from_input_ids(
     return incremental_indices.long() + padding_idx
 
 
-EntryClass = [XLMRobertaModel]
+class XLMRobertaModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        config: RobertaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.roberta = XLMRobertaBaseModel(
+            config=config, quant_config=quant_config, prefix=prefix
+        )
+        self.pooler = Pooler(pooling_type=PoolingType.CLS, normalize=True)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: torch.Tensor = None,
+        get_embedding: bool = False,
+    ) -> torch.Tensor:
+        hidden_states = self.roberta(
+            input_ids, positions, forward_batch, input_embeds, get_embedding
+        )
+        return self.pooler(hidden_states, forward_batch)
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        self.roberta.load_weights(weights)
+
+
+class XLMRobertaForSequenceClassification(nn.Module):
+    def __init__(
+        self,
+        *,
+        config: RobertaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.roberta = XLMRobertaBaseModel(
+            config=config, quant_config=quant_config, prefix=prefix
+        )
+        self.classifier = RobertaClassificationHead(config)
+        self.pooler = CrossEncodingPooler(config, self.classifier, self.roberta.pooler)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: torch.Tensor = None,
+        get_embedding: bool = True,
+    ) -> torch.Tensor:
+        assert (
+            get_embedding
+        ), "XLMRobertaForSequenceClassification is only used for rerank"
+
+        hidden_states = self.roberta(
+            input_ids, positions, forward_batch, input_embeds, get_embedding
+        )
+        return self.pooler(hidden_states, forward_batch)
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        self_weights = []
+
+        def weight_filter():
+            for name, weight in weights:
+                if name.startswith("roberta."):
+                    yield (name[len("roberta.") :], weight)
+                else:
+                    self_weights.append((name, weight))
+
+        self.roberta.load_weights(weight_filter())
+
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in self_weights:
+            if name.startswith("classifier"):
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+
+EntryClass = [XLMRobertaModel, XLMRobertaForSequenceClassification]

sglang/srt/models/vila.py

@@ -0,0 +1,305 @@
+import logging
+from typing import Any, Dict, Iterable, List, Optional, Tuple, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
+
+import sglang.srt.managers.mm_utils as mm_utils
+import sglang.srt.model_loader.weight_utils as weight_utils
+import sglang.srt.utils as utils
+from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
+from sglang.srt.layers.pooler import Pooler, PoolingType
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+
+logger = logging.getLogger(__name__)
+
+
+##### BEGIN COPY configuration.py #####
+
+
+class VILAConfig(PretrainedConfig):
+    # Class attributes.
+    model_type: str = "vila"
+    sub_configs: Dict[str, PretrainedConfig] = {
+        "text_config": Qwen2Config(),
+        "vision_config": SiglipVisionConfig(),
+    }
+    _auto_class: Optional[str] = "AutoConfig"
+
+    # Configuration for sub-modules.
+    text_config: Qwen2Config = Qwen2Config()
+    vision_config: SiglipVisionConfig = SiglipVisionConfig()
+
+    # Model configuration.
+    hidden_size: int
+    image_token_id: int
+    mm_hidden_size: int
+    mm_projector_type: str
+    mm_vision_select_feature: str
+    mm_vision_select_layer: int
+    video_token_id: int
+
+    def __init__(
+        self,
+        text_config: Optional[Dict[str, Any]] = None,
+        vision_config: Optional[Dict[str, Any]] = None,
+        *,
+        hidden_size: int = 1536,
+        image_token_id: int = 151649,
+        mm_hidden_size: int = 1152,
+        mm_projector_type: str = "mlp_downsample_3x3_fix",
+        mm_vision_select_feature: str = "cls_patch",
+        mm_vision_select_layer: int = -2,
+        video_token_id: int = 151650,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.text_config = Qwen2Config(**text_config) if text_config else Qwen2Config()
+        self.vision_config = (
+            SiglipVisionConfig(**vision_config)
+            if vision_config
+            else SiglipVisionConfig()
+        )
+
+        self.hidden_size = hidden_size
+        self.image_token_id = image_token_id
+        self.mm_hidden_size = mm_hidden_size
+        self.mm_projector_type = mm_projector_type
+        self.mm_vision_select_feature = mm_vision_select_feature
+        self.mm_vision_select_layer = mm_vision_select_layer
+        self.video_token_id = video_token_id
+
+
+##### END COPY configuration.py #####
+
+##### BEGIN COPY modeling_vila.py #####
+
+
+class DownSample3x3BlockFix(nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
+
+        Returns:
+            The output tensor of shape (batch_size, image_pad_len, mm_hidden_size * 9).
+        """
+
+        batch_size, sequence_length, hidden_size = x.shape
+
+        feat_size = int(sequence_length**0.5)
+        if feat_size**2 != sequence_length:
+            raise ValueError(
+                f"Cannot take square root: sequence_length {sequence_length} is not a perfect square"
+            )
+
+        features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
+
+        pad_after = (3 - feat_size % 3) % 3
+        if pad_after > 0:
+            features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
+            feat_size = feat_size + pad_after
+
+        features = features.reshape(
+            batch_size, feat_size // 3, 3, feat_size // 3, 3, hidden_size
+        )
+        features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
+        features = features.reshape(batch_size, -1, 9 * hidden_size)
+
+        return features
+
+
+class MultimodalProjector(nn.Module):
+    layers: nn.Sequential
+
+    def __init__(
+        self,
+        config: VILAConfig,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        if config.mm_projector_type == "mlp_downsample_3x3_fix":
+            self.layers = nn.Sequential(
+                DownSample3x3BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 9),
+                nn.Linear(
+                    config.mm_hidden_size * 9,
+                    config.mm_hidden_size * 3,
+                ),
+                nn.GELU(),
+                nn.LayerNorm(config.vision_config.hidden_size * 3),
+                nn.Linear(config.vision_config.hidden_size * 3, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported mm_projector_type: {config.mm_projector_type}"
+            )
+
+        self.layers.type(config.torch_dtype)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
+
+        Returns:
+            The output tensor of shape (batch_size, image_pad_len, hidden_size).
+        """
+
+        return self.layers(x.to(device=self.device, dtype=self.dtype))
+
+
+##### END COPY modeling_vila.py #####
+
+
+class VILAForConditionalGeneration(nn.Module):
+    config: VILAConfig
+    quant_config: Optional[QuantizationConfig]
+
+    logits_processor: LogitsProcessor
+    pooler: Pooler
+
+    llm: Qwen2ForCausalLM
+    mm_projector: MultimodalProjector
+    vision_tower: SiglipVisionModel
+
+    def __init__(
+        self,
+        config: VILAConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.logits_processor = LogitsProcessor(config)
+        self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+
+        self.llm = Qwen2ForCausalLM(
+            config=config.text_config,
+            quant_config=quant_config,
+            prefix=utils.add_prefix("llm", prefix),
+        )
+        self.mm_projector = MultimodalProjector(config)
+        self.vision_tower = SiglipVisionModel(config.vision_config)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.config.torch_dtype
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        positions: Tensor,
+        forward_batch: ForwardBatch,
+        get_embedding: bool = False,
+    ) -> LogitsProcessorOutput:
+        output = mm_utils.general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.llm,
+            image_data_embedding_func=self.get_image_feature,
+            get_embedding=get_embedding,
+            positions=positions,
+        )
+
+        return cast(LogitsProcessorOutput, output)
+
+    def get_image_feature(self, mm_input: List[MultimodalDataItem]) -> Tensor:
+        pixel_values = cast(Tensor, mm_input[0].pixel_values)
+
+        ##### BEGIN COPY modeling_vila.py #####
+
+        vision_tower_output: BaseModelOutputWithPooling = self.vision_tower.__call__(
+            pixel_values.to(
+                device=self.vision_tower.device, dtype=self.vision_tower.dtype
+            ),
+            output_hidden_states=True,
+        )
+
+        mm_projector_input = self._vision_tower_output_to_mm_projector_input(
+            vision_tower_output
+        )
+
+        image_embedding: Tensor = self.mm_projector.__call__(
+            mm_projector_input.to(
+                device=self.mm_projector.device, dtype=self.mm_projector.dtype
+            )
+        )
+
+        ##### END COPY modeling_vila.py #####
+
+        return image_embedding
+
+    def load_weights(self, weights: Iterable[Tuple[str, Tensor]]) -> None:
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in weights:
+            if name.startswith("llm."):
+                self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(
+                    param, "weight_loader", weight_utils.default_weight_loader
+                )
+                weight_loader(param, loaded_weight)
+
+    def pad_input_ids(
+        self,
+        input_ids: List[int],
+        image_inputs: MultimodalInputs,
+    ) -> List[int]:
+        pattern = MultiModalityDataPaddingPatternMultimodalTokens(
+            token_ids=[self.config.image_token_id],
+        )
+
+        return pattern.pad_input_tokens(input_ids, image_inputs)
+
+    ##### BEGIN COPY modeling_vila.py #####
+
+    def _vision_tower_output_to_mm_projector_input(
+        self,
+        vision_tower_output: BaseModelOutputWithPooling,
+    ) -> Tensor:
+        assert vision_tower_output.hidden_states is not None
+
+        selected_layer_hidden_states = vision_tower_output.hidden_states[
+            self.config.mm_vision_select_layer
+        ]
+
+        if self.config.mm_vision_select_feature == "cls_patch":
+            return selected_layer_hidden_states
+        else:
+            raise NotImplementedError(
+                f"Unsupported mm_vision_select_feature: {self.config.mm_vision_select_feature}"
+            )
+
+    ##### END COPY modeling_vila.py #####
+
+
+EntryClass = [VILAForConditionalGeneration]

sglang/srt/reasoning_parser.py

@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple, Type
 
 
 class StreamingParseResult:
@@ -32,17 +32,26 @@ class BaseReasoningFormatDetector:
         One-time parsing: Detects and parses reasoning sections in the provided text.
         Returns both reasoning content and normal text separately.
         """
-        text = text.replace(self.think_start_token, "").strip()
-        if self.think_end_token not in text:
+        in_reasoning = self._in_reasoning or text.startswith(self.think_start_token)
+
+        if not in_reasoning:
+            return StreamingParseResult(normal_text=text)
+
+        # The text is considered to be in a reasoning block.
+        processed_text = text.replace(self.think_start_token, "").strip()
+
+        if self.think_end_token not in processed_text:
             # Assume reasoning was truncated before `</think>` token
-            return StreamingParseResult(reasoning_text=text)
+            return StreamingParseResult(reasoning_text=processed_text)
 
         # Extract reasoning content
-        splits = text.split(self.think_end_token, maxsplit=1)
+        splits = processed_text.split(self.think_end_token, maxsplit=1)
         reasoning_text = splits[0]
-        text = splits[1].strip()
+        normal_text = splits[1].strip()
 
-        return StreamingParseResult(normal_text=text, reasoning_text=reasoning_text)
+        return StreamingParseResult(
+            normal_text=normal_text, reasoning_text=reasoning_text
+        )
 
     def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
         """
@@ -61,6 +70,7 @@ class BaseReasoningFormatDetector:
         if not self.stripped_think_start and self.think_start_token in current_text:
             current_text = current_text.replace(self.think_start_token, "")
             self.stripped_think_start = True
+            self._in_reasoning = True
 
         # Handle end of reasoning block
         if self._in_reasoning and self.think_end_token in current_text:
@@ -131,11 +141,11 @@ class Qwen3Detector(BaseReasoningFormatDetector):
     """
 
     def __init__(self, stream_reasoning: bool = True):
-        # Qwen3 is assumed to be reasoning until `</think>` token
+        # Qwen3 won't be in reasoning mode when user passes `enable_thinking=False`
         super().__init__(
             "<think>",
             "</think>",
-            force_reasoning=True,
+            force_reasoning=False,
             stream_reasoning=stream_reasoning,
         )
 
@@ -151,12 +161,12 @@ class ReasoningParser:
             If True, streams reasoning content as it arrives.
     """
 
-    DetectorMap: Dict[str, BaseReasoningFormatDetector] = {
+    DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = {
         "deepseek-r1": DeepSeekR1Detector,
         "qwen3": Qwen3Detector,
     }
 
-    def __init__(self, model_type: str = None, stream_reasoning: bool = True):
+    def __init__(self, model_type: Optional[str] = None, stream_reasoning: bool = True):
         if not model_type:
             raise ValueError("Model type must be specified")
 

sglang/srt/sampling/sampling_batch_info.py

@@ -10,6 +10,7 @@ import torch
 import sglang.srt.sampling.penaltylib as penaltylib
 from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
+from sglang.srt.utils import merge_bias_tensor
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -63,6 +64,9 @@ class SamplingBatchInfo:
     # Device
     device: str = "cuda"
 
+    # Handle logit bias
+    logit_bias: Optional[torch.Tensor] = None
+
     @classmethod
     def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
         reqs = batch.reqs
@@ -85,6 +89,14 @@ class SamplingBatchInfo:
             [r.sampling_params.min_p for r in reqs], dtype=torch.float
         ).to(device, non_blocking=True)
 
+        logit_bias = None
+        if any(r.sampling_params.logit_bias is not None for r in reqs):
+            logit_bias = torch.zeros(len(reqs), vocab_size, device=device)
+            for i, r in enumerate(reqs):
+                if r.sampling_params.logit_bias is not None:
+                    for key, value in r.sampling_params.logit_bias.items():
+                        logit_bias[i, int(key)] = value
+
         # Check if any request has custom logit processor
         has_custom_logit_processor = (
             batch.enable_custom_logit_processor  # check the flag first.
@@ -150,6 +162,7 @@ class SamplingBatchInfo:
             custom_params=custom_params,
             custom_logit_processor=merged_custom_logit_processor,
             device=device,
+            logit_bias=logit_bias,
         )
         return ret
 
@@ -206,6 +219,9 @@ class SamplingBatchInfo:
         if self.vocab_mask is not None:
             self.apply_mask_func(logits=logits, vocab_mask=self.vocab_mask)
 
+        if self.logit_bias is not None:
+            logits.add_(self.logit_bias)
+
     def filter_batch(self, keep_indices: List[int], keep_indices_device: torch.Tensor):
         self.penalizer_orchestrator.filter(keep_indices_device)
 
@@ -221,6 +237,9 @@ class SamplingBatchInfo:
             value = getattr(self, item, None)
             setattr(self, item, value[keep_indices_device])
 
+        if self.logit_bias is not None:
+            self.logit_bias = self.logit_bias[keep_indices_device]
+
     def _filter_batch_custom_logit_processor(
         self, keep_indices: List[int], keep_indices_device: torch.Tensor
     ):
@@ -321,3 +340,8 @@ class SamplingBatchInfo:
         self.need_top_p_sampling |= other.need_top_p_sampling
         self.need_top_k_sampling |= other.need_top_k_sampling
         self.need_min_p_sampling |= other.need_min_p_sampling
+
+        # Merge logit bias
+        self.logit_bias = merge_bias_tensor(
+            self.logit_bias, other.logit_bias, len(self), len(other), self.device, 0.0
+        )

sglang/srt/sampling/sampling_params.py

@@ -52,6 +52,7 @@ class SamplingParams:
         no_stop_trim: bool = False,
         custom_params: Optional[Dict[str, Any]] = None,
         stream_interval: Optional[int] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
     ) -> None:
         self.max_new_tokens = max_new_tokens
         self.stop_strs = stop
@@ -78,6 +79,7 @@ class SamplingParams:
         self.no_stop_trim = no_stop_trim
         self.custom_params = custom_params
         self.stream_interval = stream_interval
+        self.logit_bias = logit_bias
 
         # Process some special cases
         if 0 <= self.temperature < _SAMPLING_EPS: