sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/layers/sampler.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -11,8 +11,8 @@ from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import crash_on_warnings, get_bool_env_var, is_cuda
 
 if is_cuda():
@@ -27,13 +27,13 @@ if is_cuda():
 logger = logging.getLogger(__name__)
 
 SYNC_TOKEN_IDS_ACROSS_TP = get_bool_env_var("SYNC_TOKEN_IDS_ACROSS_TP")
-RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB")
+SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
 
 
 class Sampler(nn.Module):
     def __init__(self):
         super().__init__()
-        self.use_nan_detection = global_server_args_dict["enable_nan_detection"]
+        self.use_nan_detection = get_global_server_args().enable_nan_detection
         self.tp_sync_group = get_tp_group().device_group
 
         if is_dp_attention_enabled():
@@ -65,6 +65,7 @@ class Sampler(nn.Module):
         return_logprob: bool,
         top_logprobs_nums: List[int],
         token_ids_logprobs: List[List[int]],
+        positions: torch.Tensor,
     ):
         """Run a sampler & compute logprobs and update logits_output accordingly.
 
@@ -77,6 +78,8 @@ class Sampler(nn.Module):
             batch_next_token_ids: next token IDs. If set, skip sampling and only
                 compute output logprobs It is used for speculative decoding which
                 performs sampling in draft workers.
+            positions: The positions of the tokens in the sequence. Used for deterministic sampling
+                to get the unique seed for each position.
         """
         logits = logits_output.next_token_logits
 
@@ -88,10 +91,15 @@ class Sampler(nn.Module):
             batch_next_token_ids = torch.argmax(logits, -1)
             if return_logprob:
                 logprobs = torch.nn.functional.log_softmax(logits, dim=-1)
-
         else:
+            can_sample_directly_from_probs = (
+                not sampling_info.need_top_p_sampling
+                and not sampling_info.need_top_k_sampling
+                and not sampling_info.need_min_p_sampling
+            )
+
             # If requested, cache probabilities from original logits before temperature scaling.
-            if return_logprob and RETURN_ORIGINAL_LOGPROB:
+            if return_logprob and SGLANG_RETURN_ORIGINAL_LOGPROB:
                 probs_without_temp_scaling = torch.softmax(logits, dim=-1)
 
             # Post process logits
@@ -100,8 +108,15 @@ class Sampler(nn.Module):
             probs = logits
             del logits
 
-            if True:  # Keep this redundant check to simplify some internal code sync
-                if global_server_args_dict["sampling_backend"] == "flashinfer":
+            if can_sample_directly_from_probs:
+                # when we don't need top-k, top-p, or min-p sampling, we can directly sample from the probs
+                batch_next_token_ids = sampling_from_probs_torch(
+                    probs,
+                    sampling_seed=sampling_info.sampling_seed,
+                    positions=positions,
+                )
+            else:
+                if get_global_server_args().sampling_backend == "flashinfer":
                     if sampling_info.need_min_p_sampling:
                         probs = top_k_renorm_prob(probs, sampling_info.top_ks)
                         probs = top_p_renorm_prob(probs, sampling_info.top_ps)
@@ -116,7 +131,7 @@ class Sampler(nn.Module):
                             filter_apply_order="joint",
                             check_nan=self.use_nan_detection,
                         )
-                elif global_server_args_dict["sampling_backend"] == "pytorch":
+                elif get_global_server_args().sampling_backend == "pytorch":
                     # A slower fallback implementation with torch native operations.
                     batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_torch(
                         probs,
@@ -124,15 +139,17 @@ class Sampler(nn.Module):
                         sampling_info.top_ps,
                         sampling_info.min_ps,
                         sampling_info.need_min_p_sampling,
+                        sampling_info.sampling_seed,
+                        positions,
                     )
                 else:
                     raise ValueError(
-                        f"Invalid sampling backend: {global_server_args_dict['sampling_backend']}"
+                        f"Invalid sampling backend: {get_global_server_args().sampling_backend}"
                     )
 
             if return_logprob:
                 # clamp to avoid -inf
-                if RETURN_ORIGINAL_LOGPROB:
+                if SGLANG_RETURN_ORIGINAL_LOGPROB:
                     logprobs = torch.log(probs_without_temp_scaling).clamp(
                         min=torch.finfo(probs_without_temp_scaling.dtype).min
                     )
@@ -189,6 +206,7 @@ class Sampler(nn.Module):
         Optimized for prefill-only scoring requests that need token probabilities
         but don't require next token generation.
         """
+
         if logits_output.next_token_logits is None:
             logger.warning("No logits available for logprob computation")
             return
@@ -230,8 +248,14 @@ def top_k_top_p_min_p_sampling_from_probs_torch(
     top_ps: torch.Tensor,
     min_ps: torch.Tensor,
     need_min_p_sampling: bool,
+    sampling_seed: Optional[torch.Tensor],
+    positions: torch.Tensor,
 ):
-    """A top-k, top-p and min-p sampling implementation with native pytorch operations."""
+    """
+    A top-k, top-p and min-p sampling implementation with native pytorch operations.
+    When sampling_seed is not None, deterministic inference will be enabled, it will sample
+    with the sampling_seed of each request.
+    """
     probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
     probs_sum = torch.cumsum(probs_sort, dim=-1)
     probs_sort[
@@ -243,18 +267,62 @@ def top_k_top_p_min_p_sampling_from_probs_torch(
     if need_min_p_sampling:
         min_p_thresholds = probs_sort[:, 0] * min_ps
         probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0
-
-    sampled_index = torch.multinomial(probs_sort, num_samples=1)
+    if sampling_seed is not None:
+        sampled_index = multinomial_with_seed(probs_sort, sampling_seed, positions)
+    else:
+        sampled_index = torch.multinomial(probs_sort, num_samples=1)
     # int32 range is enough to represent the token ids
     probs_idx = probs_idx.to(torch.int32)
     batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1)
     return batch_next_token_ids
 
 
-def sampling_from_probs_torch(probs: torch.Tensor):
+def multinomial_with_seed(
+    inputs: torch.Tensor, seed: torch.Tensor, positions: torch.Tensor
+) -> torch.Tensor:
+    """
+    Samples n elements from an input tensor `inputs` of shape (n, m) using
+    a unique random seed for each row. This is a deterministic batched alternative to
+    `torch.multinomial`.
+
+    Args:
+        inputs: A float tensor of shape (n, m) representing n categorical
+                distributions with m categories each. The values are treated
+                as weights and do not need to sum to 1.
+        seed:   An integer tensor of shape (n,) containing the random seed
+                for each corresponding row in `inputs`.
+        positions: The positions of the tokens in the sequence. Used for deterministic sampling
+                to get the unique seed for each position.
+
+    Returns:
+        A tensor of shape (n,) where the i-th element is an index sampled
+        from the distribution in `inputs[i]` using `seed[i]`.
+    """
+    n, m = inputs.shape
+    col_indices = torch.arange(m, device=inputs.device).unsqueeze(0)
+    step_seed = (seed * 19349663) ^ (positions * 73856093)
+    seed_expanded = step_seed.unsqueeze(-1)
+    hashed = (seed_expanded * 8589934591) ^ (col_indices * 479001599)
+    uniform_samples = (hashed % (2**24)).float() / (2**24)
+    epsilon = 1e-10
+    uniform_samples = uniform_samples.clamp(epsilon, 1.0 - epsilon)
+    gumbel_noise = -torch.log(-torch.log(uniform_samples))
+    log_probs = torch.log(inputs + epsilon)
+    perturbed_log_probs = log_probs + gumbel_noise
+    return torch.argmax(perturbed_log_probs, dim=1, keepdim=True)
+
+
+def sampling_from_probs_torch(
+    probs: torch.Tensor,
+    sampling_seed: Optional[torch.Tensor] = None,
+    positions: Optional[torch.Tensor] = None,
+):
     """A sampling implementation with native pytorch operations, without
     top-k, top-p, or min-p filtering."""
-    sampled_index = torch.multinomial(probs, num_samples=1)
+    if sampling_seed is not None:
+        sampled_index = multinomial_with_seed(probs, sampling_seed, positions)
+    else:
+        sampled_index = torch.multinomial(probs, num_samples=1)
     batch_next_token_ids = sampled_index.view(-1).to(torch.int32)
     return batch_next_token_ids
 

sglang/srt/layers/sparse_pooler.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+
+from sglang.srt.model_executor.model_runner import ForwardBatch
+
+
+@dataclass
+class SparseEmbeddingOutput:
+    embeddings: torch.Tensor  # [batch_size, vocab_size]
+
+
+class SparsePooler(nn.Module):
+    """A layer that pools hidden states into sparse vocabulary-space embeddings.
+
+    This layer does the following:
+    1. Applies a linear transformation + ReLU to get token-level weights
+    2. Maps these weights to vocabulary positions using token IDs
+    3. Aggregates weights for repeated tokens using max pooling
+    4. Returns sparse embeddings in vocabulary space
+
+    Attributes:
+        config: Model configuration containing vocab_size and hidden_size
+        sparse_linear: Linear layer for computing token weights
+        vocab_size: Size of vocabulary for output embeddings
+    """
+
+    def __init__(self, config: PretrainedConfig):
+        super().__init__()
+
+        # Validate required attributes
+        if not hasattr(config, "vocab_size"):
+            raise AttributeError(
+                f"Config {type(config)} missing required 'vocab_size' attribute"
+            )
+        if not hasattr(config, "hidden_size"):
+            raise AttributeError(
+                f"Config {type(config)} missing required 'hidden_size' attribute"
+            )
+
+        self.vocab_size = config.vocab_size
+        self.sparse_linear = nn.Linear(config.hidden_size, 1)
+        self._weights_loaded = False
+
+    def forward(
+        self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
+    ) -> SparseEmbeddingOutput:
+        """
+        Forward pass for sparse pooling.
+
+        Args:
+            hidden_states: Packed sequence hidden states [total_tokens, hidden_size]
+            forward_batch: Batch information with sequence lengths and input_ids
+
+        Returns:
+            SparseEmbeddingOutput with embeddings of shape [batch_size, vocab_size]
+        """
+        if not self._weights_loaded:
+            raise ValueError(
+                "Sparse pooling weights not loaded. Call load_weights() first"
+            )
+
+        # Apply sparse linear + ReLU to get token weights
+        token_weights = F.relu(self.sparse_linear(hidden_states)).squeeze(
+            -1
+        )  # [total_tokens]
+
+        # Create batch indices for packed sequences
+        batch_indices = torch.repeat_interleave(
+            torch.arange(
+                len(forward_batch.extend_seq_lens), device=hidden_states.device
+            ),
+            forward_batch.extend_seq_lens,
+        )
+
+        # Initialize sparse embedding output
+        sparse_embedding = torch.zeros(
+            len(forward_batch.extend_seq_lens),
+            self.vocab_size,
+            dtype=token_weights.dtype,
+            device=token_weights.device,
+        )
+
+        # Map to vocabulary space using scatter_reduce with amax
+        flat_indices = batch_indices * self.vocab_size + forward_batch.input_ids
+        sparse_embedding.view(-1).scatter_reduce_(
+            0, flat_indices, token_weights, reduce="amax"
+        )
+
+        return SparseEmbeddingOutput(embeddings=sparse_embedding)
+
+    def load_weights(self, state_dict: dict):
+        """Load weights from state dict (called by the model)."""
+        self.sparse_linear.load_state_dict(state_dict)
+        self._weights_loaded = True

sglang/srt/layers/utils.py

@@ -1,6 +1,5 @@
 import logging
 import re
-from functools import lru_cache
 
 import torch
 
@@ -15,6 +14,29 @@ def get_layer_id(weight_name):
     return None
 
 
+def pad_or_narrow_weight(
+    loaded_weight: torch.Tensor, input_dim: int, start_idx: int, shard_size: int
+) -> torch.Tensor:
+    # Padding with zeros for special case such as qwen2_5_VL's mlp which is not 8-aligned
+    valid_size = max(loaded_weight.shape[input_dim] - start_idx, 0)
+
+    if valid_size > 0:
+        loaded_slice = loaded_weight.narrow(input_dim, start_idx, valid_size)
+        pad_shape = list(loaded_weight.shape)
+        pad_shape[input_dim] = shard_size - valid_size
+        pad = torch.zeros(
+            pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+        )
+        return torch.cat([loaded_slice, pad], dim=input_dim)
+
+    # All padding
+    pad_shape = list(loaded_weight.shape)
+    pad_shape[input_dim] = shard_size
+    return torch.zeros(
+        pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+    )
+
+
 class PPMissingLayer(torch.nn.Identity):
     # Adapted from
     # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1

sglang/srt/layers/vocab_parallel_embedding.py

@@ -540,7 +540,10 @@ class ParallelLMHead(VocabParallelEmbedding):
 
         # We only support pack LMHead if it's not quantized.
         if _is_cpu and _is_cpu_amx_available:
-            if hasattr(self, "weight") and self.weight.dtype == torch.bfloat16:
+            if hasattr(self, "weight") and self.weight.dtype in [
+                torch.bfloat16,
+                torch.float16,
+            ]:
                 self.quant_method = PackWeightMethod(weight_names=["weight"])
 
         if bias:

sglang/srt/lora/backend/base_backend.py

@@ -143,10 +143,10 @@ def get_backend_from_name(name: str) -> BaseLoRABackend:
         from sglang.srt.lora.backend.triton_backend import TritonLoRABackend
 
         return TritonLoRABackend
-    # elif name == "csgmv":
-    #     from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
+    elif name == "csgmv":
+        from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
 
-    #     return ChunkedSgmvLoRABackend
+        return ChunkedSgmvLoRABackend
     elif name == "flashinfer":
         raise ValueError(
             "FlashInfer LoRA backend has been deprecated, please use `triton` instead."