sglang 0.4.6.post3__py3-none-any.whl → 0.4.6.post5__py3-none-any.whl

sglang/srt/layers/quantization/qoq.py

@@ -0,0 +1,244 @@
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.srt.distributed import get_tensor_model_parallel_world_size
+from sglang.srt.layers.linear import LinearMethodBase
+from sglang.srt.layers.parameter import (
+    ChannelQuantScaleParameter,
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+)
+from sglang.srt.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
+from sglang.srt.utils import is_cuda
+
+_is_cuda = is_cuda()
+if _is_cuda:
+    from sgl_kernel import qserve_w4a8_per_chn_gemm, qserve_w4a8_per_group_gemm
+
+
+QoQ_SUPPORTED_WEIGHT_BITS = [4]
+QoQ_SUPPORTED_GROUP_SIZES = [-1, 128]
+
+
+class QoQConfig(QuantizationConfig):
+    """Config class for QoQ Quantization.
+
+    - Weight: static, per-channel/group, asymmetric
+    - Activation: dynamic, per-token, symmetric
+
+    Reference: https://arxiv.org/abs/2405.04532
+    https://github.com/mit-han-lab/omniserve
+    """
+
+    def __init__(self, weight_bits: int, group_size: int) -> None:
+        self.weight_bits = weight_bits
+        self.group_size = group_size
+
+        # Verify
+        if self.weight_bits not in QoQ_SUPPORTED_WEIGHT_BITS:
+            raise ValueError(
+                f"QoQ does not support weight_bits = {self.weight_bits}. "
+                f"Only weight_bits = {QoQ_SUPPORTED_WEIGHT_BITS} "
+                "are supported."
+            )
+        if self.group_size not in QoQ_SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"QoQ does not support group_size = {self.group_size}. "
+                f"Only group_sizes = {QoQ_SUPPORTED_GROUP_SIZES} "
+                "are supported."
+            )
+
+        # 4 bits packed into 8 bit datatype.
+        self.pack_factor = 8 // self.weight_bits
+
+    def __repr__(self) -> str:
+        return "QoQConfig(weight_bits={}, group_size={})".format(
+            self.weight_bits, self.group_size
+        )
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def get_name(self) -> str:
+        return "qoq"
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        """List of filenames to search for in the model directory."""
+        return [
+            "quant_config.json",
+            "quantize_config.json",
+        ]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "QoQConfig":
+        weight_bits = cls.get_from_keys(config, ["wbits"])
+        group_size = cls.get_from_keys(config, ["group_size"])
+        return cls(weight_bits, group_size)
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> Optional["QuantizeMethodBase"]:
+        from sglang.srt.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            return QoQLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class QoQLinearMethod(LinearMethodBase):
+    """Linear method for QoQ.
+
+    Args:
+        quant_config: The QoQ quantization config.
+    """
+
+    def __init__(self, quant_config: QoQConfig):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        # Validate output_size_per_partition
+        output_size_per_partition = sum(output_partition_sizes)
+        if output_size_per_partition % 32 != 0:
+            raise ValueError(
+                f"Weight output_size_per_partition = "
+                f"{output_size_per_partition} is not divisible by 32."
+            )
+
+        # Validate input_size_per_partition
+        if input_size_per_partition % self.quant_config.pack_factor != 0:
+            raise ValueError(
+                f"Weight input_size_per_partition = "
+                f"{input_size_per_partition} is not divisible by "
+                f"pack_factor = {self.quant_config.pack_factor}."
+            )
+        if (
+            self.quant_config.group_size != -1
+            and input_size_per_partition % self.quant_config.group_size != 0
+        ):
+            raise ValueError(
+                f"Weight input_size_per_partition = "
+                f"{input_size_per_partition} is not divisible by "
+                f"group_size = {self.quant_config.group_size}."
+            )
+
+        qweight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("qweight", qweight)
+
+        s1_scales = ChannelQuantScaleParameter(
+            data=torch.empty(output_size_per_partition, dtype=torch.float16),
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("s1_scales", s1_scales)
+
+        if self.quant_config.group_size == -1:
+            s1_szeros = ChannelQuantScaleParameter(
+                data=torch.empty(output_size_per_partition, dtype=torch.float16),
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("s1_szeros", s1_szeros)
+        else:
+            s2_scales = GroupQuantScaleParameter(
+                data=torch.empty(
+                    (
+                        input_size_per_partition // self.quant_config.group_size,
+                        output_size_per_partition,
+                    ),
+                    dtype=torch.int8,
+                ),
+                input_dim=0,
+                output_dim=1,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("s2_scales", s2_scales)
+
+            s2_zeros = GroupQuantScaleParameter(
+                data=torch.empty(
+                    (
+                        input_size_per_partition // self.quant_config.group_size,
+                        output_size_per_partition,
+                    ),
+                    dtype=torch.int8,
+                ),
+                input_dim=0,
+                output_dim=1,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("s2_zeros", s2_zeros)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.s1_scales = Parameter(layer.s1_scales.data, requires_grad=False)
+        if self.quant_config.group_size == -1:
+            layer.s1_szeros = Parameter(layer.s1_szeros.data, requires_grad=False)
+        else:
+            layer.s2_scales = Parameter(layer.s2_scales.data, requires_grad=False)
+            layer.s2_zeros = Parameter(layer.s2_zeros.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        assert x.dtype == torch.float16, "QoQ only supports float16 input now"
+        if self.quant_config.group_size == -1:
+            x_q, x_scale, x_sum = per_token_quant_int8(
+                x, scale_dtype=x.dtype, cal_sum=True
+            )
+            out = qserve_w4a8_per_chn_gemm(
+                x_q, layer.qweight, layer.s1_scales, x_scale, layer.s1_szeros, x_sum
+            )
+        else:
+            x_q, x_scale = per_token_quant_int8(x, scale_dtype=x.dtype)
+            out = qserve_w4a8_per_group_gemm(
+                x_q,
+                layer.qweight,
+                layer.s2_zeros,
+                layer.s2_scales,
+                layer.s1_scales,
+                x_scale,
+            )
+        if bias is not None:
+            out = out + bias
+        return out

sglang/srt/layers/sampler.py

@@ -239,10 +239,6 @@ def top_p_normalize_probs_torch(
 
 
 def get_top_logprobs(logprobs: torch.Tensor, top_logprobs_nums: List[int]):
-    assert len(top_logprobs_nums) == logprobs.shape[0], (
-        len(top_logprobs_nums),
-        logprobs.shape[0],
-    )
     max_k = max(top_logprobs_nums)
     ret = logprobs.topk(max_k, dim=1)
     values = ret.values.tolist()

sglang/srt/layers/vocab_parallel_embedding.py

@@ -13,6 +13,7 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
 from sglang.srt.layers.parameter import BasevLLMParameter
 from sglang.srt.layers.quantization.base_config import (
     QuantizationConfig,
@@ -214,12 +215,14 @@ class VocabParallelEmbedding(torch.nn.Module):
         self,
         num_embeddings: int,
         embedding_dim: int,
+        *,
         params_dtype: Optional[torch.dtype] = None,
         org_num_embeddings: Optional[int] = None,
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         enable_tp: bool = True,
+        use_attn_tp_group: bool = False,
         use_presharded_weights: bool = False,
     ):
         super().__init__()
@@ -227,9 +230,14 @@ class VocabParallelEmbedding(torch.nn.Module):
 
         self.enable_tp = enable_tp
         if self.enable_tp:
-            tp_rank = get_tensor_model_parallel_rank()
-            self.tp_size = get_tensor_model_parallel_world_size()
+            if use_attn_tp_group:
+                tp_rank = get_attention_tp_rank()
+                self.tp_size = get_attention_tp_size()
+            else:
+                tp_rank = get_tensor_model_parallel_rank()
+                self.tp_size = get_tensor_model_parallel_world_size()
         else:
+            assert use_attn_tp_group is False
             tp_rank = 0
             self.tp_size = 1
 
@@ -519,22 +527,25 @@ class ParallelLMHead(VocabParallelEmbedding):
         self,
         num_embeddings: int,
         embedding_dim: int,
+        *,
         bias: bool = False,
         params_dtype: Optional[torch.dtype] = None,
         org_num_embeddings: Optional[int] = None,
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        use_attn_tp_group: bool = False,
         use_presharded_weights: bool = False,
     ):
         super().__init__(
             num_embeddings,
             embedding_dim,
-            params_dtype,
-            org_num_embeddings,
-            padding_size,
-            quant_config,
-            prefix,
+            params_dtype=params_dtype,
+            org_num_embeddings=org_num_embeddings,
+            padding_size=padding_size,
+            quant_config=quant_config,
+            prefix=prefix,
+            use_attn_tp_group=use_attn_tp_group,
             use_presharded_weights=use_presharded_weights,
         )
         self.quant_config = quant_config

sglang/srt/lora/lora_manager.py

@@ -100,7 +100,7 @@ class LoRAManager:
             self.configs[name] = LoRAConfig(path)
             self.hf_target_names.update(self.configs[name].target_modules)
 
-        # Target lora weight names for lora_a and lora_b modules repectively.
+        # Target lora weight names for lora_a and lora_b modules respectively.
         # e.g., {("qkv_proj", "q_proj"), ("qkv_proj", "kv_proj")}
         self.lora_weight_names: Set[Tuple[str]] = set(
             [get_stacked_name(module) for module in self.hf_target_names]
@@ -170,9 +170,7 @@ class LoRAManager:
                 dim=0,
                 out=self.cuda_graph_batch_info.seg_indptr[1 : bs + 1],
             )
-            self.cuda_graph_batch_info.max_len = int(
-                torch.max(self.cuda_graph_batch_info.seg_lens[:bs])
-            )
+            self.cuda_graph_batch_info.max_len = 1
 
             for i, lora_path in enumerate(forward_batch.lora_paths):
                 self.cuda_graph_batch_info.weight_indices[i] = (

sglang/srt/lora/mem_pool.py

@@ -50,15 +50,15 @@ class LoRAMemoryPool:
         self.uid_to_buffer_id: Dict[Optional[str], int] = {}
 
         # Buffer idx -> lora uid in memory pool
-        # All uids are initalized as empty strings for empty buffer slots
-        # Here we don't initalize to None since None is a valid uid
+        # All uids are initialized as empty strings for empty buffer slots
+        # Here we don't initialize to None since None is a valid uid
         self.buffer_id_to_uid: List[Optional[str]] = [""] * self.max_loras_per_batch
 
     def get_lora_A_shape(
         self, module_name: str, base_model: torch.nn.Module
     ) -> Tuple[int]:
         """
-        Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+        Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
         """
         input_dim, _ = get_hidden_dim(module_name, self.base_hf_config, base_model)
         c = get_stacked_multiply(module_name)
@@ -75,7 +75,7 @@ class LoRAMemoryPool:
         self, module_name: str, base_model: torch.nn.Module
     ) -> Tuple[int]:
         """
-        Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+        Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
         """
         _, output_dim = get_hidden_dim(module_name, self.base_hf_config, base_model)
         c = get_stacked_multiply(module_name)

sglang/srt/lora/triton_ops/gate_up_lora_b.py

@@ -77,7 +77,7 @@ def _gate_up_lora_b_kernel(
         k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
     )
 
-    # Iteate to compute the block in output matrix
+    # Iterate to compute the block in output matrix
     partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_K)):
         x_tile = tl.load(

sglang/srt/lora/triton_ops/qkv_lora_b.py

@@ -79,7 +79,7 @@ def _qkv_lora_b_kernel(
         k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
     )
 
-    # Iteate to compute the block in output matrix
+    # Iterate to compute the block in output matrix
     partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_K)):
         x_tile = tl.load(

sglang/srt/lora/triton_ops/sgemm_lora_a.py

@@ -67,7 +67,7 @@ def _sgemm_lora_a_kernel(
         k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
     )
 
-    # Iteate to compute the block in output matrix
+    # Iterate to compute the block in output matrix
     partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_K)):
         x_tile = tl.load(

sglang/srt/lora/triton_ops/sgemm_lora_b.py

@@ -69,7 +69,7 @@ def _sgemm_lora_b_kernel(
         k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
     )
 
-    # Iteate to compute the block in output matrix
+    # Iterate to compute the block in output matrix
     partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_K)):
         x_tile = tl.load(

sglang/srt/lora/utils.py

@@ -79,7 +79,7 @@ def get_hidden_dim(
     module_name: str, config: AutoConfig, base_model: torch.nn.Module
 ) -> Tuple[int]:
     """
-    Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+    Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
     """
 
     if hasattr(base_model, "get_hidden_dim"):

sglang/srt/managers/data_parallel_controller.py

@@ -17,13 +17,13 @@ import logging
 import multiprocessing as mp
 import signal
 import threading
+import time
 from enum import Enum, auto
 
 import psutil
 import setproctitle
 import zmq
 
-from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.managers.io_struct import (
     TokenizedEmbeddingReqInput,
@@ -158,7 +158,7 @@ class DataParallelController:
         # This thread cannot be closed because otherwise the `kill_itself_when_parent_died`
         # function in scheduler.py will kill the scheduler.
         while True:
-            pass
+            time.sleep(30 * 24 * 3600)
 
     def launch_dp_attention_schedulers(self, server_args, port_args):
         self.launch_tensor_parallel_group(server_args, port_args, 0, None)
@@ -210,7 +210,7 @@ class DataParallelController:
                     )
                     # compute zmq ports for this dp rank
                     rank_port_args = PortArgs.init_new(server_args, dp_rank)
-                    # Data parallelism resues the tensor parallelism group,
+                    # Data parallelism reuses the tensor parallelism group,
                     # so all dp ranks should use the same nccl port.
                     rank_port_args.nccl_port = port_args.nccl_port