sglang quant module remove vllm dependency (#4507)

python/sglang/srt/layers/quantization/__init__.py

@@ -6,21 +6,41 @@ from copy import deepcopy
 from typing import Callable, Dict, Optional, Type, Union
 
 import torch
-from vllm.model_executor.layers.quantization.aqlm import AQLMConfig
-from vllm.model_executor.layers.quantization.awq import AWQConfig
-from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
-from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig
-from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
-    CompressedTensorsConfig,
-)
-from vllm.model_executor.layers.quantization.deepspeedfp import DeepSpeedFPConfig
-from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config
-from vllm.model_executor.layers.quantization.fbgemm_fp8 import FBGEMMFp8Config
-from vllm.model_executor.layers.quantization.gguf import GGUFConfig
-from vllm.model_executor.layers.quantization.gptq_marlin_24 import GPTQMarlin24Config
-from vllm.model_executor.layers.quantization.marlin import MarlinConfig
-from vllm.model_executor.layers.quantization.qqq import QQQConfig
-from vllm.model_executor.layers.quantization.tpu_int8 import Int8TpuConfig
+
+try:
+    from vllm.model_executor.layers.quantization.aqlm import AQLMConfig
+    from vllm.model_executor.layers.quantization.awq import AWQConfig
+    from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
+    from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig
+    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+        CompressedTensorsConfig,
+    )
+    from vllm.model_executor.layers.quantization.deepspeedfp import DeepSpeedFPConfig
+    from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config
+    from vllm.model_executor.layers.quantization.fbgemm_fp8 import FBGEMMFp8Config
+    from vllm.model_executor.layers.quantization.gguf import GGUFConfig
+    from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
+        GPTQMarlin24Config,
+    )
+    from vllm.model_executor.layers.quantization.marlin import MarlinConfig
+    from vllm.model_executor.layers.quantization.qqq import QQQConfig
+    from vllm.model_executor.layers.quantization.tpu_int8 import Int8TpuConfig
+
+    VLLM_AVAILABLE = True
+except ImportError:
+    VLLM_AVAILABLE = False
+
+    # Define empty classes as placeholders when vllm is not available
+    class DummyConfig:
+        pass
+
+    AQLMConfig = AWQConfig = AWQMarlinConfig = BitsAndBytesConfig = (
+        CompressedTensorsConfig
+    ) = DummyConfig
+    DeepSpeedFPConfig = ExpertsInt8Config = FBGEMMFp8Config = GGUFConfig = (
+        GPTQMarlin24Config
+    ) = DummyConfig
+    MarlinConfig = QQQConfig = Int8TpuConfig = DummyConfig
 
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config
@@ -30,29 +50,37 @@ from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8Config
 from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config
 from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
 
-QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
-    "aqlm": AQLMConfig,
-    "awq": AWQConfig,
-    "deepspeedfp": DeepSpeedFPConfig,
-    "tpu_int8": Int8TpuConfig,
+# Base quantization methods that don't depend on vllm
+BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
     "fp8": Fp8Config,
     "blockwise_int8": BlockInt8Config,
-    "fbgemm_fp8": FBGEMMFp8Config,
-    "marlin": MarlinConfig,
     "modelopt": ModelOptFp8Config,
-    "gguf": GGUFConfig,
-    "gptq_marlin_24": GPTQMarlin24Config,
     "gptq_marlin": GPTQMarlinConfig,
-    "awq_marlin": AWQMarlinConfig,
     "gptq": GPTQConfig,
-    "compressed-tensors": CompressedTensorsConfig,
-    "bitsandbytes": BitsAndBytesConfig,
-    "qqq": QQQConfig,
-    "experts_int8": ExpertsInt8Config,
     "w8a8_int8": W8A8Int8Config,
     "w8a8_fp8": W8A8Fp8Config,
 }
 
+# Add vllm-dependent methods if available
+QUANTIZATION_METHODS = BASE_QUANTIZATION_METHODS.copy()
+if VLLM_AVAILABLE:
+    VLLM_QUANTIZATION_METHODS = {
+        "aqlm": AQLMConfig,
+        "awq": AWQConfig,
+        "deepspeedfp": DeepSpeedFPConfig,
+        "tpu_int8": Int8TpuConfig,
+        "fbgemm_fp8": FBGEMMFp8Config,
+        "marlin": MarlinConfig,
+        "gguf": GGUFConfig,
+        "gptq_marlin_24": GPTQMarlin24Config,
+        "awq_marlin": AWQMarlinConfig,
+        "compressed-tensors": CompressedTensorsConfig,
+        "bitsandbytes": BitsAndBytesConfig,
+        "qqq": QQQConfig,
+        "experts_int8": ExpertsInt8Config,
+    }
+    QUANTIZATION_METHODS.update(VLLM_QUANTIZATION_METHODS)
+
 
 def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
     if quantization not in QUANTIZATION_METHODS:
@@ -157,25 +185,31 @@ def get_linear_quant_method(
 
 
 def gptq_get_quant_method(self, layer, prefix):
-    from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod
-    from vllm.model_executor.layers.quantization.gptq_marlin import (
-        GPTQMarlinLinearMethod,
-        GPTQMarlinMoEMethod,
-    )
+    if not VLLM_AVAILABLE:
+        return None
 
-    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
-
-    if isinstance(layer, FusedMoE):
-        return GPTQMarlinMoEMethod(self)
-
-    if isinstance(self, GPTQConfig):
-        return get_linear_quant_method(
-            self, layer, prefix=prefix, linear_method_cls=GPTQLinearMethod
-        )
-    elif isinstance(self, GPTQMarlinConfig):
-        return get_linear_quant_method(
-            self, layer, prefix=prefix, linear_method_cls=GPTQMarlinLinearMethod
+    try:
+        from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod
+        from vllm.model_executor.layers.quantization.gptq_marlin import (
+            GPTQMarlinLinearMethod,
+            GPTQMarlinMoEMethod,
         )
+
+        from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+
+        if isinstance(layer, FusedMoE):
+            return GPTQMarlinMoEMethod(self)
+
+        if isinstance(self, GPTQConfig):
+            return get_linear_quant_method(
+                self, layer, prefix=prefix, linear_method_cls=GPTQLinearMethod
+            )
+        elif isinstance(self, GPTQMarlinConfig):
+            return get_linear_quant_method(
+                self, layer, prefix=prefix, linear_method_cls=GPTQMarlinLinearMethod
+            )
+    except ImportError:
+        pass
     return None
 
 
@@ -187,33 +221,40 @@ def monkey_patch_isinstance_for_vllm_base_layer(reverse: bool = False):
     Patch isinstance so that the `get_quant_method` in vllm's QuantizationConfig
     can recognize sglang layers
     """
+    if not VLLM_AVAILABLE:
+        return
 
     if reverse:
         builtins.isinstance = original_isinstance
         return
 
-    from vllm.model_executor.layers.fused_moe import FusedMoE
-    from vllm.model_executor.layers.linear import LinearBase
-    from vllm.model_executor.layers.vocab_parallel_embedding import (
-        VocabParallelEmbedding,
-    )
+    try:
+        from vllm.model_executor.layers.fused_moe import FusedMoE
+        from vllm.model_executor.layers.linear import LinearBase
+        from vllm.model_executor.layers.vocab_parallel_embedding import (
+            VocabParallelEmbedding,
+        )
 
-    from sglang.srt.layers.linear import LinearBase as PatchedLinearBase
-    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE as PatchedFusedMoE
-    from sglang.srt.layers.vocab_parallel_embedding import (
-        VocabParallelEmbedding as PatchedVocabParallelEmbedding,
-    )
+        from sglang.srt.layers.linear import LinearBase as PatchedLinearBase
+        from sglang.srt.layers.moe.fused_moe_triton.layer import (
+            FusedMoE as PatchedFusedMoE,
+        )
+        from sglang.srt.layers.vocab_parallel_embedding import (
+            VocabParallelEmbedding as PatchedVocabParallelEmbedding,
+        )
 
-    def patched_isinstance(obj, classinfo):
-        if classinfo is LinearBase:
-            return original_isinstance(obj, PatchedLinearBase)
-        if classinfo is FusedMoE:
-            return original_isinstance(obj, PatchedFusedMoE)
-        if classinfo is VocabParallelEmbedding:
-            return original_isinstance(obj, PatchedVocabParallelEmbedding)
-        return original_isinstance(obj, classinfo)
+        def patched_isinstance(obj, classinfo):
+            if classinfo is LinearBase:
+                return original_isinstance(obj, PatchedLinearBase)
+            if classinfo is FusedMoE:
+                return original_isinstance(obj, PatchedFusedMoE)
+            if classinfo is VocabParallelEmbedding:
+                return original_isinstance(obj, PatchedVocabParallelEmbedding)
+            return original_isinstance(obj, classinfo)
 
-    builtins.isinstance = patched_isinstance
+        builtins.isinstance = patched_isinstance
+    except ImportError:
+        return
 
 
 def monkey_patch_moe_apply(class_obj: "FusedMoEMethodBase"):
@@ -221,72 +262,88 @@ def monkey_patch_moe_apply(class_obj: "FusedMoEMethodBase"):
     Monkey patch the apply function of vllm's FusedMoEMethodBase.
     Convert sglang arguments to vllm arguments.
     """
-    original_apply = class_obj.apply
-    sig = inspect.signature(original_apply)
-    param_names = list(sig.parameters.keys())
-    has_correction_bias = "e_score_correction_bias" in param_names
+    if not VLLM_AVAILABLE:
+        return
 
-    def new_apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        router_logits: torch.Tensor,
-        top_k: int,
-        renormalize: bool,
-        use_grouped_topk: bool,
-        topk_group: Optional[int] = None,
-        num_expert_group: Optional[int] = None,
-        custom_routing_function: Optional[Callable] = None,
-        correction_bias: Optional[torch.Tensor] = None,
-        activation: str = "silu",
-        inplace: bool = True,
-        no_combine: bool = False,
-    ):
-        assert activation == "silu"
-        assert inplace and not no_combine
+    try:
+        original_apply = class_obj.apply
+        sig = inspect.signature(original_apply)
+        param_names = list(sig.parameters.keys())
+        has_correction_bias = "e_score_correction_bias" in param_names
 
-        kwargs = {
-            "self": self,
-            "layer": layer,
-            "x": x,
-            "router_logits": router_logits,
-            "top_k": top_k,
-            "renormalize": renormalize,
-            "use_grouped_topk": use_grouped_topk,
-            "topk_group": topk_group,
-            "num_expert_group": num_expert_group,
-            "custom_routing_function": custom_routing_function,
-        }
-        if correction_bias is not None:
-            if not has_correction_bias:
-                raise ValueError(
-                    "Please increase the version of your vllm. Try `pip install vllm==0.7.2`"
-                )
-            kwargs["e_score_correction_bias"] = correction_bias
-        return original_apply(**kwargs)
+        def new_apply(
+            self,
+            layer: torch.nn.Module,
+            x: torch.Tensor,
+            router_logits: torch.Tensor,
+            top_k: int,
+            renormalize: bool,
+            use_grouped_topk: bool,
+            topk_group: Optional[int] = None,
+            num_expert_group: Optional[int] = None,
+            custom_routing_function: Optional[Callable] = None,
+            correction_bias: Optional[torch.Tensor] = None,
+            activation: str = "silu",
+            inplace: bool = True,
+            no_combine: bool = False,
+        ):
+            assert activation == "silu"
+            assert inplace and not no_combine
 
-    setattr(class_obj, "apply", new_apply)
+            kwargs = {
+                "self": self,
+                "layer": layer,
+                "x": x,
+                "router_logits": router_logits,
+                "top_k": top_k,
+                "renormalize": renormalize,
+                "use_grouped_topk": use_grouped_topk,
+                "topk_group": topk_group,
+                "num_expert_group": num_expert_group,
+                "custom_routing_function": custom_routing_function,
+            }
+            if correction_bias is not None:
+                if not has_correction_bias:
+                    raise ValueError(
+                        "Please increase the version of your vllm. Try `pip install vllm==0.7.2`"
+                    )
+                kwargs["e_score_correction_bias"] = correction_bias
+            return original_apply(**kwargs)
+
+        setattr(class_obj, "apply", new_apply)
+    except (ImportError, AttributeError):
+        return
 
 
 def monkey_patch_quant_configs():
     """Apply all monkey patches in one place."""
-    from vllm.model_executor.layers.quantization.awq_marlin import AWQMoEMethod
-    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (
-        CompressedTensorsW8A8Fp8MoEMethod,
-        CompressedTensorsWNA16MoEMethod,
-    )
-    from vllm.model_executor.layers.quantization.gptq_marlin import GPTQMarlinMoEMethod
+    if not VLLM_AVAILABLE:
+        return
 
-    setattr(GPTQMarlinConfig, "get_quant_method", gptq_get_quant_method)
-    setattr(GPTQConfig, "get_quant_method", gptq_get_quant_method)
+    try:
+        from vllm.model_executor.layers.quantization.awq_marlin import AWQMoEMethod
+        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (
+            CompressedTensorsW8A8Fp8MoEMethod,
+            CompressedTensorsWNA16MoEMethod,
+        )
+        from vllm.model_executor.layers.quantization.gptq_marlin import (
+            GPTQMarlinMoEMethod,
+        )
 
-    monkey_patch_moe_apply(AWQMoEMethod)
-    monkey_patch_moe_apply(GPTQMarlinMoEMethod)
-    monkey_patch_moe_apply(CompressedTensorsW8A8Fp8MoEMethod)
-    monkey_patch_moe_apply(CompressedTensorsWNA16MoEMethod)
+        setattr(GPTQMarlinConfig, "get_quant_method", gptq_get_quant_method)
+        setattr(GPTQConfig, "get_quant_method", gptq_get_quant_method)
+
+        monkey_patch_moe_apply(AWQMoEMethod)
+        monkey_patch_moe_apply(GPTQMarlinMoEMethod)
+        monkey_patch_moe_apply(CompressedTensorsW8A8Fp8MoEMethod)
+        monkey_patch_moe_apply(CompressedTensorsWNA16MoEMethod)
+    except ImportError:
+        return
 
 
-monkey_patch_quant_configs()
+# Only apply monkey patches if vllm is available
+if VLLM_AVAILABLE:
+    monkey_patch_quant_configs()
 
 
 __all__ = [