diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/conversation.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/conversation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6cba339ba9014f7275707427d6e647161985f79 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/conversation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/function_call_parser.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/function_call_parser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7db4378ad9f4075c0dde9dd84dffc65e778ec9b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/function_call_parser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/mm_utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/mm_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ecd54bad4e5b5199c6c312866a24b53fa24f604 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/mm_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/warmup.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/warmup.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..563626558a7b79c85d91d845ac1dd6480f50e74b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/warmup.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/base_attn_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/base_attn_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efad841e4685523a6cab5a9e35c4e0b2740488e6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/base_attn_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/cutlass_mla_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/cutlass_mla_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..951bf4cea9eb31ac60a3776f662d73cbee0d4af4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/cutlass_mla_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/double_sparsity_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/double_sparsity_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3adfbcabb3a36f138766cd916895130aaf0126fe Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/double_sparsity_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashattention_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashattention_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..507bf549fd9c3b0873cc8bf83fabfc12abbaf8bf Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashattention_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c86ce7bd55a5d5e3587ae3cb10ec92a659a54275 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_mla_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_mla_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a431d2ea90ad6a3e32951e64b7fc5620650d8c39 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashinfer_mla_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashmla_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashmla_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0cd4174b043d4c3336443e49f2c5c6fedbdb8635 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/flashmla_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/torch_native_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/torch_native_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22263b3ec0ad381ee76e870777916e12e00f4fb9 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/torch_native_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/triton_backend.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/triton_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9b313a822ca0c8b63d473d75278c96ebb8395b4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/triton_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..669988a4dc74288aa44a272416dea477d44edb78 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/vision.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/vision.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93c1d6f638d51164382075aa1a7441e2c49d2fc6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/__pycache__/vision.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/decode_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/decode_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..539816d879631a82aede4c63d7fc0dde2edb0b0a Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/decode_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/double_sparsity_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/double_sparsity_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56584f48ed5494f2e02b4cd8af52c358fdf670db Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/double_sparsity_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/extend_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/extend_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9cd6f327b7f836acf4a3ca04b3a55382ea6f468a Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/extend_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/prefill_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/prefill_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb238d19728d227aac8e4513408af4059e23d18a Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/prefill_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/rocm_mla_decode_rope.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/rocm_mla_decode_rope.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2505cdac7479dbc015441933fdd15a6c7bb611a Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/__pycache__/rocm_mla_decode_rope.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__init__.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6b8719bfa40dddcd5d91853670c801b7cbf72428 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__init__.py @@ -0,0 +1,334 @@ +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/v0.5.5/vllm/model_executor/layers/quantization/__init__.py +import builtins +import inspect +import re +from copy import deepcopy +from typing import Callable, Dict, Optional, Type, Union + +import torch + +try: + from vllm.model_executor.layers.quantization.aqlm import AQLMConfig + from vllm.model_executor.layers.quantization.awq_marlin import ( + AWQMarlinConfig, + AWQMoEMethod, + ) + from vllm.model_executor.layers.quantization.bitsandbytes import BitsAndBytesConfig + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( + CompressedTensorsW8A8Fp8MoEMethod, + CompressedTensorsWNA16MoEMethod, + ) + 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 import GPTQLinearMethod + from vllm.model_executor.layers.quantization.gptq_marlin import ( + GPTQMarlinLinearMethod, + GPTQMarlinMoEMethod, + ) + 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: + def override_quantization_method(self, *args, **kwargs): + return None + + AQLMConfig = AWQMarlinConfig = BitsAndBytesConfig = CompressedTensorsConfig = ( + DeepSpeedFPConfig + ) = ExpertsInt8Config = FBGEMMFp8Config = GGUFConfig = GPTQMarlin24Config = ( + MarlinConfig + ) = QQQConfig = Int8TpuConfig = DummyConfig + + +from sglang.srt.layers.linear import LinearBase, UnquantizedLinearMethod +from sglang.srt.layers.quantization.awq import AWQConfig +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + ModelOptFp8Config, +) +from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config +from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config +from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config + +# Base quantization methods that don't depend on vllm +BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { + "fp8": Fp8Config, + "blockwise_int8": BlockInt8Config, + "modelopt": ModelOptFp8Config, + "modelopt_fp4": ModelOptFp4Config, + "w8a8_int8": W8A8Int8Config, + "w8a8_fp8": W8A8Fp8Config, + "moe_wna16": MoeWNA16Config, + "compressed-tensors": CompressedTensorsConfig, +} + +# VLLM-dependent quantization methods +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, + "bitsandbytes": BitsAndBytesConfig, + "qqq": QQQConfig, + "experts_int8": ExpertsInt8Config, + "gptq_marlin": GPTQMarlinConfig, + "gptq": GPTQConfig, +} + +QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS, **VLLM_QUANTIZATION_METHODS} + + +def get_quantization_config(quantization: str) -> Type[QuantizationConfig]: + if quantization not in QUANTIZATION_METHODS: + raise ValueError( + f"Invalid quantization method: {quantization}. " + f"Available methods: {list(QUANTIZATION_METHODS.keys())}" + ) + if quantization in VLLM_QUANTIZATION_METHODS and not VLLM_AVAILABLE: + raise ValueError( + f"{quantization} quantization requires some operators from vllm. " + "Pleaes install vllm by `pip install vllm==0.7.2`" + ) + + return QUANTIZATION_METHODS[quantization] + + +# Match dynamic rules with module name (prefix) and override quantize +# config if module (prefix) matches a rule +def override_config(config: QuantizationConfig, prefix: str): + weight_bits = get_dynamic_override(config, prefix, "bits", config.weight_bits) + if isinstance(weight_bits, int): + config.weight_bits = weight_bits + group_size = get_dynamic_override(config, prefix, "group_size", config.group_size) + if isinstance(group_size, int): + config.group_size = group_size + desc_act = get_dynamic_override(config, prefix, "desc_act", config.desc_act) + if isinstance(desc_act, bool): + config.desc_act = desc_act + + config.pack_factor = 32 // config.weight_bits # packed into int32 + if config.get_name() == "gptq_marlin": + is_sym = get_dynamic_override(config, prefix, "sym", config.is_sym) + if isinstance(is_sym, bool): + config.is_sym = is_sym + + if (config.weight_bits, config.is_sym) not in config.TYPE_MAP: + raise ValueError( + "Unsupported quantization config: " + f"bits={config.weight_bits}, sym={config.is_sym}" + ) + + config.quant_type = config.TYPE_MAP[(config.weight_bits, config.is_sym)] + elif config.get_name() == "gptq": + if config.weight_bits not in [2, 3, 4, 8]: + raise ValueError( + "Currently, only 2/3/4/8-bit weight quantization is " + f"supported for GPTQ, but got {config.weight_bits} bits." + ) + + +def get_dynamic_override( + config: QuantizationConfig, + layer_name: str, + key: Optional[str] = None, + default_value: Union[int, bool, None] = None, +) -> Union[Dict, int, bool, None]: + for pattern, pattern_dict in config.dynamic.items(): + # Negative match: matched modules are excluded from quantized init + if pattern.startswith("-:"): + if re.match(pattern.removeprefix("-:"), layer_name): + return False + # Positive match: matched modules have quant properties overrides + # base quant config + elif re.match(pattern.removeprefix("+:"), layer_name): + if key is None: + return pattern_dict + else: + return pattern_dict.get(key, default_value) + return default_value + + +def get_linear_quant_method( + config: QuantizationConfig, + layer: torch.nn.Module, + prefix: str, + linear_method_cls: type, +): + # Move import here to avoid circular import. This is only used in monkey patching + # of vllm's QuantizationConfig. + from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, + ) + + cloned_config = deepcopy(config) + parallel_lm_head_quantized = ( + isinstance(layer, ParallelLMHead) and cloned_config.lm_head_quantized + ) + + if isinstance(layer, LinearBase) or parallel_lm_head_quantized: + # False = skip module, None = no override, else = Positive match + if ( + get_dynamic_override( # noqa: E712 + cloned_config, layer_name=prefix # noqa: E712 + ) + == False + ): # noqa: E712 + if parallel_lm_head_quantized: + return UnquantizedEmbeddingMethod() + return UnquantizedLinearMethod() + + if prefix: + # Dynamic per module/layer rules may override base config + override_config(cloned_config, prefix=prefix) + + return linear_method_cls(cloned_config) + return None + + +def gptq_get_quant_method(self, layer, prefix): + 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 + ) + return None + + +original_isinstance = builtins.isinstance + + +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, + ) + + 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) + + builtins.isinstance = patched_isinstance + + +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 + + 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", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ): + assert activation == "silu" + assert inplace and not no_combine + + 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) + + +def monkey_patch_quant_configs(): + """Apply all monkey patches in one place.""" + 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) + + +# Only apply monkey patches if vllm is available +if VLLM_AVAILABLE: + monkey_patch_quant_configs() diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/awq.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/awq.py new file mode 100644 index 0000000000000000000000000000000000000000..9f14ac4c1cac52262993a3ab0258e470fa7b45fc --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/awq.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +import logging +from typing import Any, Dict, List, Optional + +import torch + +from sglang.srt.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import awq_dequantize + +logger = logging.getLogger(__name__) + + +def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]): + return any(module_name in prefix for module_name in modules_to_not_convert) + + +class AWQConfig(QuantizationConfig): + """Config class for AWQ. + + Reference: https://arxiv.org/abs/2306.00978 + """ + + def __init__( + self, + weight_bits: int, + group_size: int, + zero_point: bool, + modules_to_not_convert: Optional[List[str]] = None, + ) -> None: + super().__init__() + self.weight_bits = weight_bits + self.group_size = group_size + self.zero_point = zero_point + self.modules_to_not_convert = modules_to_not_convert or [] + + if self.weight_bits != 4: + raise ValueError( + "Currently, only 4-bit weight quantization is supported for " + f"AWQ, but got {self.weight_bits} bits." + ) + self.pack_factor = 32 // self.weight_bits + + def __repr__(self) -> str: + return ( + f"AWQConfig(weight_bits={self.weight_bits}, " + f"group_size={self.group_size}, " + f"zero_point={self.zero_point}, " + f"modules_to_not_convert={self.modules_to_not_convert})" + ) + + def get_scaled_act_names(self) -> List[str]: + return [] + + def get_name(self) -> str: + return "awq" + + def get_supported_act_dtypes(self) -> List[torch.dtype]: + return [torch.half] + + @classmethod + def get_min_capability(cls) -> int: + # The AWQ kernel only supports Turing or newer GPUs. + return 75 + + @staticmethod + def get_config_filenames() -> List[str]: + return [ + "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq + # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq + "quantize_config.json", + ] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "AWQConfig": + weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) + group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) + zero_point = cls.get_from_keys(config, ["zero_point"]) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls(weight_bits, group_size, zero_point, modules_to_not_convert) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["LinearMethodBase"]: + + if isinstance(layer, LinearBase): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + return AWQLinearMethod(self) + return None + + +class AWQLinearMethod(LinearMethodBase): + """Linear method for AWQ. + + Args: + quant_config: The AWQ quantization config. + """ + + def __init__(self, quant_config: AWQConfig): + 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, + ): + if input_size_per_partition % self.quant_config.group_size != 0: + raise ValueError( + "The input size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % self.quant_config.pack_factor != 0: + raise ValueError( + "The output size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + weight_loader = extra_weight_attrs.get("weight_loader") + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + qzeros = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + scales = GroupQuantScaleParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qweight = layer.qweight + scales = layer.scales + qzeros = layer.qzeros + pack_factor = self.quant_config.pack_factor + out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) + reshaped_x = x.reshape(-1, x.shape[-1]) + + out = awq_dequantize(qweight, scales, qzeros) + out = torch.matmul(reshaped_x, out) + + if bias is not None: + out.add_(bias) + return out.reshape(out_shape) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/base_config.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/base_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6058702c9a10dd2e3c161c4753fa6dd73bf7cf79 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/base_config.py @@ -0,0 +1,140 @@ +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/v0.5.5/vllm/model_executor/layers/quantization/base_config.py + +import inspect +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Type + +import torch +from torch import nn + + +class QuantizeMethodBase(ABC): + """Base class for different quantized methods.""" + + @abstractmethod + def create_weights( + self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs + ): + """Create weights for a layer. + + The weights will be set as attributes of the layer.""" + raise NotImplementedError + + @abstractmethod + def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: + """Apply the weights in layer to the input tensor. + + Expects create_weights to have been called before on the layer.""" + raise NotImplementedError + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Process the weight after loading. + + This can be used for example, to transpose weights for computation. + """ + return + + +class QuantizationConfig(ABC): + """Base class for quantization configs.""" + + def __init__(self): + super().__init__() + # mapping is updated by models as they initialize + self.packed_modules_mapping: Dict[str, List[str]] = dict() + + @abstractmethod + def get_name(self) -> str: + """Name of the quantization method.""" + raise NotImplementedError + + @abstractmethod + def get_supported_act_dtypes(self) -> List[torch.dtype]: + """List of supported activation dtypes.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """Minimum GPU capability to support the quantization method. + + E.g., 70 for Volta, 75 for Turing, 80 for Ampere. + This requirement is due to the custom CUDA kernels used by the + quantization method. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_config_filenames() -> List[str]: + """List of filenames to search for in the model directory.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def from_config(cls, config: Dict[str, Any]) -> "QuantizationConfig": + """Create a config class from the model's quantization config.""" + raise NotImplementedError + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + """ + Detects if this quantization method can support a given checkpoint + format by overriding the user specified quantization method -- + this method should only be overwritten by subclasses in exceptional + circumstances + """ + return None + + @staticmethod + def get_from_keys(config: Dict[str, Any], keys: List[str]) -> Any: + """Get a value from the model's quantization config.""" + for key in keys: + if key in config: + return config[key] + raise ValueError( + f"Cannot find any of {keys} in the model's " "quantization config." + ) + + @staticmethod + def get_from_keys_or(config: Dict[str, Any], keys: List[str], default: Any) -> Any: + """Get a optional value from the model's quantization config.""" + try: + return QuantizationConfig.get_from_keys(config, keys) + except ValueError: + return default + + @abstractmethod + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + """Get the quantize method to use for the quantized layer. + + Args: + layer: The layer for the quant method. + prefix: The full name of the layer in the state dict + Returns: + The quantize method. None if the given layer doesn't support quant + method. + """ + raise NotImplementedError + + @abstractmethod + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError + + +def method_has_implemented_embedding(method_class: Type[QuantizeMethodBase]) -> bool: + """ + Not all quant methods have embedding implemented, so we need to check that + it exists for our given method. We check this by making sure the function + has been changed from the base implementation. + """ + base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None) + class_embedding = inspect.getattr_static(method_class, "embedding", None) + + return class_embedding is not None and class_embedding is not base_embedding diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/blockwise_int8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/blockwise_int8.py new file mode 100644 index 0000000000000000000000000000000000000000..25c91da6eb20b4c7a303dbf854c0d52a5a087c63 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/blockwise_int8.py @@ -0,0 +1,412 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/quantization/fp8.py + +import logging +from typing import Any, Callable, Dict, List, Optional + +import torch +from torch.nn import Module + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from sglang.srt.layers.parameter import BlockQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.int8_utils import apply_w8a8_block_int8_linear +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import set_weight_attrs + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class BlockInt8Config(QuantizationConfig): + """Config class for INT8.""" + + def __init__( + self, + is_checkpoint_int8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: Optional[List[str]] = None, + weight_block_size: List[int] = None, + ) -> None: + self.is_checkpoint_int8_serialized = is_checkpoint_int8_serialized + if is_checkpoint_int8_serialized: + logger.warning( + "Detected int8 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + if weight_block_size is not None: + if not is_checkpoint_int8_serialized: + raise ValueError( + f"The block-wise quantization only supports int8-serialized checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." + ) + if activation_scheme != "dynamic": + raise ValueError( + f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + + @classmethod + def get_name(cls) -> str: + return "blockwise_int8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "BlockInt8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_int8_serialized = "int8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + return cls( + is_checkpoint_int8_serialized=is_checkpoint_int8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return BlockInt8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return BlockInt8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class BlockInt8LinearMethod(LinearMethodBase): + """Linear method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic activation scale. + + Limitations: + Only support block-wise int8 quantization and int8 checkpoint + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: BlockInt8Config): + self.quant_config = quant_config + assert self.quant_config.weight_block_size is not None + assert self.quant_config.is_checkpoint_int8_serialized + + 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, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + tp_size = get_tensor_model_parallel_world_size() + + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + # Required by collum parallel or enabling merged weights + if (tp_size > 1 and output_size // output_size_per_partition == tp_size) or len( + output_partition_sizes + ) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight_dtype = ( + torch.int8 + if self.quant_config.is_checkpoint_int8_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, input_size_per_partition, dtype=weight_dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + + scale = BlockQuantScaleParameter( + data=torch.empty( + (output_size_per_partition + block_n - 1) // block_n, + (input_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale_inv", scale) + + # INPUT ACTIVATION SCALE + assert self.quant_config.activation_scheme == "dynamic" + layer.register_parameter("input_scale", None) + + def process_weights_after_loading(self, layer: Module) -> None: + # Block quant doesn't need to process weights after loading + # Use torch Parameter to avoid cuda graph capturing issue + layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) + layer.weight_scale_inv = torch.nn.Parameter( + layer.weight_scale_inv.data, requires_grad=False + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_w8a8_block_int8_linear( + input=x, + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=None, + bias=bias, + ) + + +class BlockInt8MoEMethod: + """MoE method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic activation scale. + + Limitations: + Only support block-wise int8 quantization and int8 checkpoint + + Args: + quant_config: The quantization config. + """ + + def __new__(cls, *args, **kwargs): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoEMethodBase + + if not hasattr(cls, "_initialized"): + original_init = cls.__init__ + new_cls = type( + cls.__name__, + (FusedMoEMethodBase,), + { + "__init__": original_init, + **{k: v for k, v in cls.__dict__.items() if k != "__dict__"}, + }, + ) + obj = super(new_cls, new_cls).__new__(new_cls) + obj.__init__(*args, **kwargs) + return obj + return super().__new__(cls) + + def __init__(self, quant_config): + self.quant_config = quant_config + assert self.quant_config.weight_block_size is not None + assert self.quant_config.is_checkpoint_int8_serialized + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + if self.quant_config.is_checkpoint_int8_serialized: + params_dtype = torch.int8 + tp_size = get_tensor_model_parallel_world_size() + + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + # Required by collum parallel or enabling merged weights + if intermediate_size % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1: + # Required by row parallel + if intermediate_size % block_k != 0: + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size, hidden_size, dtype=params_dtype + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, hidden_size, intermediate_size, dtype=params_dtype + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * ((intermediate_size + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + assert self.quant_config.activation_scheme == "dynamic" + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: Module) -> None: + # Block quant doesn't need to process weights after loading + return + + def 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", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + # Expert selection + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + # Expert fusion with INT8 quantization + return fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + use_int8_w8a8=True, + w1_scale=(layer.w13_weight_scale_inv), + w2_scale=(layer.w2_weight_scale_inv), + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.quant_config.weight_block_size, + no_combine=no_combine, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__init__.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab8ea54c57db6e197be18b9770e2df079cabba2e Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a863d319301be19a8a704d141269a1ac6b77cd5 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors_moe.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef34efd07bbac33c146f5ab0805ee028d72aac5f Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/compressed_tensors_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44863a18f6753e3d76cb479e0b48082f67fd4352 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..7ce89345fd6b18e51ec722d7b3489bca5a067b29 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -0,0 +1,664 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +import logging +from contextlib import suppress +from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, cast + +import torch +from compressed_tensors.config import ( + CompressionFormat, + SparsityCompressionConfig, + SparsityStructure, +) +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) +from pydantic import BaseModel + +from sglang.srt.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + CompressedTensorsMoEMethod, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, + CompressedTensorsW8A8Fp8, + CompressedTensorsW8A16Fp8, +) +from sglang.srt.layers.quantization.compressed_tensors.utils import ( + find_matched_target, + is_activation_quantization_format, + should_ignore_layer, +) + +try: + import vllm + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +logger = logging.getLogger(__name__) + +__all__ = ["CompressedTensorsLinearMethod"] + +SPARSITY_CONFIG_NAME: Literal["sparsity_config"] = "sparsity_config" +QUANTIZATION_SCHEME_MAP_TYPE = Dict[str, Optional[Dict[str, QuantizationArgs]]] + + +class DeviceCapability(NamedTuple): + major: int + minor: int + + def as_version_str(self) -> str: + return f"{self.major}.{self.minor}" + + def to_int(self) -> int: + """ + Express device capability as an integer ````. + + It is assumed that the minor version is always a single digit. + """ + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class CompressedTensorsConfig(QuantizationConfig): + + def __init__( + self, + target_scheme_map: Dict[str, Any], + ignore: List[str], + quant_format: str, + sparsity_scheme_map: Dict[str, SparsityCompressionConfig], + sparsity_ignore_list: List[str], + kv_cache_scheme: Optional[Dict[str, Any]] = None, + config: Optional[Dict[str, Any]] = None, + packed_modules_mapping: Dict[str, List[str]] = {}, + ): + super().__init__() + self.ignore = ignore + self.quant_format = quant_format + # Map from [target -> scheme] + self.target_scheme_map = target_scheme_map + self.kv_cache_scheme = kv_cache_scheme + self.sparsity_scheme_map = sparsity_scheme_map + self.sparsity_ignore_list = sparsity_ignore_list + self.config = config + self.packed_modules_mapping = packed_modules_mapping + + def get_linear_method(self) -> "CompressedTensorsLinearMethod": + return CompressedTensorsLinearMethod(self) + + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def get_name(self) -> str: + return "compressed_tensors" + + def get_scaled_act_names(self) -> List[str]: + return [] + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + + # Check if the layer is skipped for quantization. + # TODO (@robertgshaw2): support module names + if should_ignore_layer( + prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping + ): + return UnquantizedLinearMethod() + if isinstance(layer, LinearBase): + scheme = self.get_scheme(layer=layer, layer_name=prefix) + if scheme is None: + return UnquantizedLinearMethod() + layer.scheme = scheme + return CompressedTensorsLinearMethod(self) + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, FusedMoE): + return CompressedTensorsMoEMethod.get_moe_method(self) + return None + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "CompressedTensorsConfig": + ignore: List[str] = cast(List[str], config.get("ignore", [])) + quant_format = cast(str, config.get("format")) + target_scheme_map = cls._quantization_scheme_map_from_config(config=config) + sparsity_scheme_map, sparsity_ignore_list = cls._parse_sparsity_config( + config=config + ) + packed_modules_mapping = config.get("packed_modules_mapping", {}) + + return cls( + target_scheme_map=target_scheme_map, + ignore=ignore, + quant_format=quant_format, + sparsity_scheme_map=sparsity_scheme_map, + sparsity_ignore_list=sparsity_ignore_list, + config=config, + packed_modules_mapping=packed_modules_mapping, + ) + + @classmethod + def _parse_sparsity_config( + cls, config: Dict[str, Any] + ) -> Tuple[Dict[str, SparsityCompressionConfig], List[str]]: + """ + :param config: The `quantization_config` dictionary from config.json + :return: A tuple with two elements + 1. A dictionary mapping target layer names to their corresponding + sparsity_config + 2. A list of layer names to ignore for sparsity + """ + if not (sparsity_config := config.get(SPARSITY_CONFIG_NAME)): + return dict(), [] + + sparsity_config = SparsityCompressionConfig.model_validate(sparsity_config) + sparse_scheme_map: Dict[str, SparsityCompressionConfig] = { + target: sparsity_config for target in sparsity_config.targets or list() + } + sparsity_ignore_list = sparsity_config.ignore or list() + return sparse_scheme_map, sparsity_ignore_list + + @classmethod + def _quantization_scheme_map_from_config( + cls, config: Dict[str, Any] + ) -> QUANTIZATION_SCHEME_MAP_TYPE: + """ + :param config: The `quantization_config` dictionary from config.json + :return: A dictionary mapping target layer names to their corresponding + quantization_args for weights and input activations + """ + target_scheme_map: Dict[str, Any] = dict() + quant_format = cast(str, config.get("format")) + + # The quant_config has multiple config_groups, each containing + # an input_activations key with details about how the activations are + # quantized, a weights key indicating how the weights are quantized, + # and a list of targets under the `targets` key, dictating which + # layers are impacted by the quantization details. The quantization + # details follow the structure defined by the QuantizationArgs + # pydantic model, which is used to verify the structure of the + # quant_config and also store the details for later use. + + config_groups = config.get("config_groups", dict()) + for _, quant_config in config_groups.items(): + targets = quant_config.get("targets") + for target in targets: + target_scheme_map[target] = {} + target_scheme_map[target]["weights"] = QuantizationArgs.model_validate( + quant_config.get("weights") + ) + + target_scheme_map[target]["input_activations"] = None + if is_activation_quantization_format(quant_format): + input_activations = quant_config.get("input_activations") + # The only case where we have activation quant supported + # but no input_activations provided in the config + # should be w8a16fp8 w8a16fp8 can also run for cases where + # there is an input_quant but it is ignored + if not input_activations: + assert ( + target_scheme_map[target]["weights"].type + == QuantizationType.FLOAT + ) + else: + target_scheme_map[target]["input_activations"] = ( + QuantizationArgs.model_validate( # noqa: E501 + quant_config.get("input_activations") + ) + ) + return target_scheme_map + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool: + capability_tuple = DeviceCapability(*torch.cuda.get_device_capability()) + + if capability_tuple is not None: + capability = capability_tuple.to_int() + supported = capability >= min_capability + if error and not supported: + raise RuntimeError( + "Quantization scheme is not supported for ", + f"the current GPU. Min capability: {min_capability}. ", + f"Current capability: {capability}.", + ) + return supported + else: + return False + + def _is_static_tensor_w8a8( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8 + weight_strategy = ( + weight_quant.strategy == QuantizationStrategy.TENSOR.value + or weight_quant.strategy == QuantizationStrategy.CHANNEL.value + ) + is_tensor = ( + weight_strategy + and input_quant.strategy == QuantizationStrategy.TENSOR.value + ) + is_static = not weight_quant.dynamic and not input_quant.dynamic + + # Both symmetric and asymmetric input quantization supported. + # Only symmetric weight quantization supported. + return is_8_bits and is_tensor and weight_quant.symmetric and is_static + + def _is_dynamic_token_w8a8( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8 + weight_strategy = ( + weight_quant.strategy == QuantizationStrategy.TENSOR.value + or weight_quant.strategy == QuantizationStrategy.CHANNEL.value + ) + is_token = ( + weight_strategy and input_quant.strategy == QuantizationStrategy.TOKEN.value + ) + is_dynamic = not weight_quant.dynamic and input_quant.dynamic + + # Both symmetric and asymmetric input quantization supported. + # Only symmetric weight quantization supported. + return is_8_bits and is_token and weight_quant.symmetric and is_dynamic + + def _is_fp8_w8a8(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + # Confirm weights and activations quantized. + if weight_quant is None or input_quant is None: + return False + + # Confirm weight scheme is supported. + is_floating_point = ( + weight_quant.type == QuantizationType.FLOAT + and input_quant.type == QuantizationType.FLOAT + ) + is_symmetric_weight = weight_quant.symmetric + is_static_weight = not weight_quant.dynamic + is_per_tensor_or_channel_weight = weight_quant.strategy in [ + QuantizationStrategy.TENSOR, + QuantizationStrategy.CHANNEL, + ] + if not ( + is_floating_point + and is_symmetric_weight + and is_static_weight + and is_per_tensor_or_channel_weight + ): + return False + + # Dynamic quantization is always supported if weights supported. + if input_quant.dynamic: + return True + + # Confirm activation scheme is supported. + is_symmetric_activation = input_quant.symmetric + is_per_tensor_activation = input_quant.strategy == QuantizationStrategy.TENSOR + return is_symmetric_activation and is_per_tensor_activation + + def _is_fp8_w8a16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + # Confirm weights quantized. + if weight_quant is None: + return False + + # Confirm we have floating points. + if weight_quant.type != QuantizationType.FLOAT: + return False + + # Confirm weight scheme is supported. + is_symmetric_weight = weight_quant.symmetric + is_static_weight = not weight_quant.dynamic + is_per_tensor_or_channel_weight = weight_quant.strategy in [ + QuantizationStrategy.TENSOR, + QuantizationStrategy.CHANNEL, + ] + if not ( + is_symmetric_weight + and is_static_weight # noqa: SIM103 + and is_per_tensor_or_channel_weight + ): + return False + + # All conditions satisfied. + return True + + def _is_wNa16_group_channel( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> bool: + input_quant_none = input_quant is None + is_symmetric = weight_quant.symmetric + is_channel_group = ( + weight_quant.strategy == QuantizationStrategy.CHANNEL.value + or weight_quant.strategy == QuantizationStrategy.GROUP.value + ) + is_static = not weight_quant.dynamic + + return is_channel_group and input_quant_none and is_symmetric and is_static + + def _get_scheme_from_parts( + self, weight_quant: BaseModel, input_quant: BaseModel + ) -> "CompressedTensorsScheme": + + # Detect If Mixed Precision + if self._is_wNa16_group_channel(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsW4A16Sparse24 and CompressedTensorsWNA16, please install vllm" + ) + if ( + self.quant_format == CompressionFormat.marlin_24.value + and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS + ): + return CompressedTensorsW4A16Sparse24( + strategy=weight_quant.strategy, + num_bits=weight_quant.num_bits, + group_size=weight_quant.group_size, + ) + if ( + self.quant_format == CompressionFormat.pack_quantized.value + and weight_quant.num_bits in WNA16_SUPPORTED_BITS + ): + return CompressedTensorsWNA16( + num_bits=weight_quant.num_bits, + strategy=weight_quant.strategy, + group_size=weight_quant.group_size, + actorder=weight_quant.actorder, + ) + + if is_activation_quantization_format(self.quant_format): + if self._is_fp8_w8a8(weight_quant, input_quant): + is_fp8_w8a8_supported = self._check_scheme_supported( + CompressedTensorsW8A8Fp8.get_min_capability(), error=False + ) + if is_fp8_w8a8_supported: + return CompressedTensorsW8A8Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=( + input_quant and not input_quant.dynamic + ), + ) + else: + # note: input_quant will be present for converted models; + # will be ignored during inference post loading + return CompressedTensorsW8A16Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=not input_quant.dynamic, + ) + + # note: input_quant can be None + if self._is_fp8_w8a16(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsW8A16Fp8, please install vllm" + ) + is_static_input_scheme = input_quant and not input_quant.dynamic + return CompressedTensorsW8A16Fp8( + strategy=weight_quant.strategy, + is_static_input_scheme=is_static_input_scheme, + ) + + if self._is_static_tensor_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Int8( + strategy=weight_quant.strategy, + is_static_input_scheme=True, + input_symmetric=input_quant.symmetric, + ) + + if self._is_dynamic_token_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Int8( + strategy=weight_quant.strategy, + is_static_input_scheme=False, + input_symmetric=input_quant.symmetric, + ) + + raise NotImplementedError("No compressed-tensors compatible scheme was found.") + + def get_scheme( + self, layer: torch.nn.Module, layer_name: Optional[str] = None + ) -> Optional["CompressedTensorsScheme"]: + """ + compressed-tensors supports non uniform in the following way: + + targets of config_groups: There can be N config_groups which each + have a quantization scheme. Each config_group has a list of targets + which can be a full layer_name, a regex for a layer_name, or + an nn.Module name. + + Detect whether a layer_name is found in any target and + use the quantization scheme corresponding to the matched target + to select the CompressedTensorsScheme used for infernece. + """ + + # Find the "target" in the compressed-tensors config + # that our layer conforms to. + # TODO (@robertgshaw): add compressed-tensors as dep + # so we do not have to re-write these functions + # need to make accelerate optional in ct to do this + + # Will be empty for models with only sparsity + weight_quant = input_quant = None + if self.target_scheme_map: + matched_target = find_matched_target( + layer_name=layer_name, + module=layer, + targets=self.target_scheme_map.keys(), + fused_mapping=self.packed_modules_mapping, + ) + + scheme_dict = self.target_scheme_map[matched_target] + weight_quant = scheme_dict.get("weights") + input_quant = scheme_dict.get("input_activations") + + # Find the sparsity scheme of the layer + # assume that fused layers inerhit first component's sparsity scheme + sparsity_targets = self.sparsity_scheme_map.keys() - set( + self.sparsity_ignore_list + ) + sparsity_scheme: Optional[SparsityCompressionConfig] = None + with suppress(ValueError): + matched_target = find_matched_target( + layer_name=layer_name, + module=layer, + targets=sparsity_targets, + fused_mapping=self.packed_modules_mapping, + ) + sparsity_scheme = self.sparsity_scheme_map[matched_target] + + if self.supports_cutlass_24( + weight_quant=weight_quant, + input_quant=input_quant, + sparsity_scheme=sparsity_scheme, + ): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensors24, please install vllm" + ) + # Have a valid sparsity scheme + # Validate layer is supported by Cutlass 2:4 Kernel + model_compression_config = ( + None + if sparsity_scheme is None or sparsity_scheme.format == "dense" + else self.config + ) + + scheme = CompressedTensors24( + quantized=weight_quant is not None or input_quant is not None, + weight_quant=weight_quant, + input_quant=input_quant, + model_compression_config=model_compression_config, + ) + elif weight_quant is None: + logger.warning_once( + "Acceleration for non-quantized schemes is " + "not supported by Compressed Tensors. " + "Falling back to UnquantizedLinearMethod" + ) + return None + + else: + # Find the quant_scheme + scheme = self._get_scheme_from_parts( # type: ignore + weight_quant=weight_quant, + input_quant=input_quant, + ) + + # Raise error if device does not support the scheme + # (e.g. fp8 needs ada lovelace) + self._check_scheme_supported(scheme.get_min_capability()) + logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) + return scheme + + def get_cache_scale(self, name: str) -> Optional[str]: + """ + Check whether the param name matches the format for k/v cache scales + in compressed-tensors. If this is the case, return its equivalent + param name expected by vLLM + + :param name: param name + :return: matching param name for KV cache scale in vLLM + """ + if name.endswith(".output_scale") and ".k_proj" in name: + return name.replace(".k_proj.output_scale", ".attn.k_scale") + if name.endswith(".output_scale") and ".v_proj" in name: + return name.replace(".v_proj.output_scale", ".attn.v_scale") + # If no matches, return None + return None + + @staticmethod + def supports_cutlass_24( + weight_quant: Optional[QuantizationArgs], + input_quant: Optional[QuantizationArgs], + sparsity_scheme: Optional[SparsityCompressionConfig] = None, + ) -> bool: + """ + Check if the layer is supported by the Cutlass 2:4 Kernel + Conditions: + - Overarching condition: Sparsity Structure is 2:4 + - Unquantized cases are supported + - Weight only quantization is not-supported + - Supported weight quantization strategies are TENSOR and CHANNEL + - Supported input quantization strategies are TENSOR and TOKEN + - Only 8 bit quantization is supported + + :return: True if the layer is supported by the Cutlass 2:4 Kernel + False otherwise + """ + if sparsity_scheme is None: + return False + + is_valid_sparsity_structure: bool = ( + sparsity_scheme.sparsity_structure == SparsityStructure.TWO_FOUR.value + ) + + valid_compressors = { + CompressionFormat.dense.value, + CompressionFormat.sparse_24_bitmask.value, + } + + is_valid_sparsity = ( + is_valid_sparsity_structure and sparsity_scheme.format in valid_compressors + ) + + if not is_valid_sparsity: + return False + + # Unquantized cases are supported + if weight_quant is None and input_quant is None: + return True + + # Weight only quantization is not-supported + if weight_quant is not None and input_quant is None: + return False + + supported_weight_quant_strategies = [ + QuantizationStrategy.TENSOR.value, + QuantizationStrategy.CHANNEL.value, + ] + + assert weight_quant is not None + assert input_quant is not None + if weight_quant.strategy not in supported_weight_quant_strategies: + return False + + supported_input_quant_strategies = [ + QuantizationStrategy.TENSOR.value, + QuantizationStrategy.TOKEN.value, + ] + + if input_quant.strategy not in supported_input_quant_strategies: + return False + + return weight_quant.num_bits == input_quant.num_bits == 8 + + +class CompressedTensorsLinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: CompressedTensorsConfig): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + 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, + ): + """ + Use the CompressedTensorsScheme associated with each layer to create + the necessary parameters for the layer. See LinearMethodBase for param + details + """ + weight_loader = extra_weight_attrs.get("weight_loader") + layer.scheme.create_weights( + layer=layer, + input_size=input_size, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + output_size=output_size, + params_dtype=params_dtype, + weight_loader=weight_loader, + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + """ + Use the output of create_weights and the CompressedTensorsScheme + associated with the layer to apply the forward pass with the + layer input. See LinearMethodBase for param details + + """ + + scheme = layer.scheme + if scheme is None: + raise ValueError("A scheme must be defined for each layer") + return scheme.apply_weights(layer, x, bias=bias) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..67496b14b59413981d824feeea9944a6480cb17d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -0,0 +1,677 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +import enum +import logging +from enum import Enum +from typing import Callable, List, Optional + +import torch +from compressed_tensors import CompressionFormat +from compressed_tensors.quantization import QuantizationStrategy + +from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant +from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz +from sglang.srt.layers.quantization.utils import ( + all_close_1d, + is_cuda, + is_fp8_fnuz, + per_tensor_dequantize, + replace_parameter, +) +from sglang.srt.utils import set_weight_attrs + +_is_cuda = is_cuda() + +if not _is_cuda: + from vllm import _custom_ops as vllm_ops + from vllm._custom_ops import scaled_fp8_quant + +try: + import vllm + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +logger = logging.getLogger(__name__) + + +class GPTQMarlinState(Enum): + REPACK = enum.auto() + READY = enum.auto() + + +__all__ = [ + "CompressedTensorsMoEMethod", + "CompressedTensorsW8A8Fp8MoEMethod", + "CompressedTensorsWNA16MoEMethod", +] + + +class CompressedTensorsMoEMethod: + def __new__(cls, *args, **kwargs): + if cls is CompressedTensorsMoEMethod: + return super().__new__(cls) + return super().__new__(cls) + + @staticmethod + def get_moe_method( + quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501 + ) -> "CompressedTensorsMoEMethod": + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + weight_quant = quant_config.target_scheme_map["Linear"].get("weights") + input_quant = quant_config.target_scheme_map["Linear"].get("input_activations") + + if quant_config._is_wNa16_group_channel(weight_quant, input_quant): + if not VLLM_AVAILABLE: + raise ImportError( + "vllm is not installed, to use CompressedTensorsWNA16MoEMethod, please install vllm." + ) + return CompressedTensorsWNA16MoEMethod(quant_config) + elif quant_config._is_fp8_w8a8(weight_quant, input_quant): + return CompressedTensorsW8A8Fp8MoEMethod(quant_config) + else: + raise RuntimeError( + f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}" + ) + + +class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): + + def __init__( + self, quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501 + ): + self.quant_config = quant_config + self.weight_quant = self.quant_config.target_scheme_map["Linear"].get("weights") + self.input_quant = self.quant_config.target_scheme_map["Linear"].get( + "input_activations" + ) + + self.static_input_scales = not self.input_quant.dynamic + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + params_dtype = torch.float8_e4m3fn + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + # per-tensor quantization + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + weight_quant_method = FusedMoeWeightScaleSupported.TENSOR.value + elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + 1, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + weight_quant_method = FusedMoeWeightScaleSupported.CHANNEL.value + else: + raise ValueError( + f"Unsupported weight quantization strategy: {self.weight_quant.strategy}" + ) + + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update({"quant_method": weight_quant_method}) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.static_input_scales: + assert ( + self.input_quant.strategy == QuantizationStrategy.TENSOR + ), "Only per-tensor quantization is supported for static input scales" + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Fp8 moe kernels require a single activation scale. + # We take the max of all the scales in case they differ. + if self.static_input_scales: + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + logger.warning( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer." + ) + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + if is_fp8_fnuz(): + # Normalize the weights and scales + w13_weight, w13_weight_scale, w13_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale + ) + ) + w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( + layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + if w13_input_scale is not None: + layer.w13_input_scale = torch.nn.Parameter( + w13_input_scale, requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + if w2_input_scale is not None: + layer.w2_input_scale = torch.nn.Parameter( + w2_input_scale, requires_grad=False + ) + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + # Fp8 moe kernel needs single weight scale for w13 per expert. + # We take the max then dequant and requant each expert. + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.local_num_experts): + start = 0 + for shard_id in range(2): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter( + max_w13_scales, requires_grad=False + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + use_grouped_topk: bool = False, + topk_group: Optional[int] = None, + num_expert_group: Optional[int] = None, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + custom_routing_function: Optional[Callable] = None, + scoring_func: str = "softmax", + correction_bias: Optional[torch.Tensor] = None, + activation: str = "silu", + inplace: bool = True, + no_combine: bool = False, + apply_router_weight_on_input: bool = False, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.fused_moe_triton import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + return fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + activation=activation, + use_fp8_w8a8=True, + per_channel_quant=self.weight_quant.strategy + == QuantizationStrategy.CHANNEL, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + +class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): + + def __init__( + self, quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501 + ): + self.quant_config = quant_config + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + config = self.quant_config.target_scheme_map["Linear"].get("weights") + self.num_bits = config.num_bits + self.packed_factor = 32 // config.num_bits + self.strategy = config.strategy + self.group_size = config.group_size + self.actorder = config.actorder + assert config.symmetric, "Only symmetric quantization is supported for MoE" + + if not ( + self.quant_config.quant_format == CompressionFormat.pack_quantized.value + and self.num_bits in WNA16_SUPPORTED_BITS + ): + raise ValueError( + "For Fused MoE layers, only ", + f"{CompressionFormat.pack_quantized.value} ", + "is supported for the following bits: ", + f"{WNA16_SUPPORTED_BITS}", + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + assert ( + params_dtype == torch.float16 + ), "float16 is required for MoE compressed models. Set dtype=torch.float16" # noqa: E501 + + intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") + + # Will transpose the loaded weight along the + # intermediate and hidden dim sizes. Will + # shard for TP along the transposed dims + extra_weight_attrs.update( + {"is_transposed": True, "quant_method": self.strategy} + ) + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // self.packed_factor, + 2 * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # In the case where we have actorder/g_idx, + # we do not partition the w2 scales + load_full_w2 = self.actorder and self.group_size != -1 + w2_scales_size = ( + intermediate_size_full if load_full_w2 else intermediate_size_per_partition + ) + + self.is_k_full = (not self.actorder) or ( + intermediate_size_per_partition == intermediate_size_full + ) + + if self.strategy == "channel": + num_groups_w2 = num_groups_w13 = 1 + self.group_size = -1 + else: + num_groups_w2 = w2_scales_size // self.group_size + num_groups_w13 = hidden_size // self.group_size + + w13_scale = torch.nn.Parameter( + torch.ones( + num_experts, + num_groups_w13, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_scale) + set_weight_attrs(w13_scale, extra_weight_attrs) + + w2_scale = torch.nn.Parameter( + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + layer.a13_scale = None + layer.a2_scale = None + layer.marlin_state = GPTQMarlinState.REPACK + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + + def replace_tensor(name, new_t): + # It is important to use resize_() here since it ensures + # the same buffer is reused + getattr(layer, name).resize_(new_t.shape) + getattr(layer, name).copy_(new_t) + del new_t + + def get_scale_perms(num_bits: int): + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single: List[int] = [] + for i in range(4): + scale_perm_single.extend( + [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]] + ) + return scale_perm, scale_perm_single + + def marlin_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int, num_bits: int + ): + scale_perm, scale_perm_single = get_scale_perms(num_bits) + if group_size < size_k and group_size != -1: + s = s.reshape((-1, len(scale_perm)))[:, scale_perm] + else: + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + s = s.reshape((-1, size_n)).contiguous() + return s + + def marlin_moe_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int, num_bits: int + ): + num_experts = s.shape[0] + output = torch.empty( + (num_experts, s.shape[1], s.shape[2]), device=s.device, dtype=s.dtype + ) + for e in range(num_experts): + output[e] = marlin_permute_scales( + s[e], size_k, size_n, group_size, num_bits + ) + return output + + size_k2 = layer.w2_weight_packed.shape[2] + size_k13 = layer.w13_weight_packed.shape[2] + + num_experts = layer.w13_weight_g_idx.shape[0] + device = layer.w13_weight_g_idx.device + + # when running models with grouped act order, + # resort to g_idx values provided in checkpoint + if self.actorder == "group": + w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx) + w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx) + w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx) + w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx) + + for e in range(num_experts): + w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to( + torch.int32 + ) + w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to( + torch.int32 + ) + w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][ + w13_g_idx_sort_indices[e] + ] + w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]] + + replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx) + replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx) + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + + else: + layer.w13_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + marlin_w13_qweight = vllm_ops.gptq_marlin_moe_repack( + layer.w13_weight_packed, + layer.w13_g_idx_sort_indices, + layer.w13_weight_packed.shape[1] * self.packed_factor, + layer.w13_weight_packed.shape[2], + self.num_bits, + ) + replace_tensor("w13_weight_packed", marlin_w13_qweight) + marlin_w2_qweight = vllm_ops.gptq_marlin_moe_repack( + layer.w2_weight_packed, + layer.w2_g_idx_sort_indices, + layer.w2_weight_packed.shape[1] * self.packed_factor, + layer.w2_weight_packed.shape[2], + self.num_bits, + ) + replace_tensor("w2_weight_packed", marlin_w2_qweight) + # Repack scales + marlin_w13_scales = marlin_moe_permute_scales( + layer.w13_weight_scale, + size_k13, + layer.w13_weight_scale.shape[2], + self.group_size, + self.num_bits, + ) + replace_tensor("w13_weight_scale", marlin_w13_scales) + marlin_w2_scales = marlin_moe_permute_scales( + layer.w2_weight_scale, + layer.w2_weight_scale.shape[1] + * (self.group_size if self.group_size != -1 else self.packed_factor), + size_k2, + self.group_size, + self.num_bits, + ) + replace_tensor("w2_weight_scale", marlin_w2_scales) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + use_grouped_topk: bool = False, + topk_group: Optional[int] = None, + num_expert_group: Optional[int] = None, + global_num_experts: int = -1, + expert_map: Optional[torch.Tensor] = None, + custom_routing_function: Optional[Callable] = None, + scoring_func: str = "softmax", + correction_bias: Optional[torch.Tensor] = None, + activation: str = "silu", + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.topk import select_experts + + assert activation == "silu", "Only SiLU activation is supported." + if expert_map is not None: + raise NotImplementedError( + "Expert Parallelism is not supported for " "fused Marlin MoE method." + ) + + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + return torch.ops.vllm.fused_marlin_moe( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + layer.w13_weight_scale, + layer.w2_weight_scale, + router_logits, + topk_weights, + topk_ids, + g_idx1=layer.w13_weight_g_idx, + g_idx2=layer.w2_weight_g_idx, + sort_indices1=layer.w13_g_idx_sort_indices, + sort_indices2=layer.w2_g_idx_sort_indices, + num_bits=self.num_bits, + is_k_full=self.is_k_full, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c94575316758e25f99a26f828170ed9974286768 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .compressed_tensors_scheme import CompressedTensorsScheme +from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8 +from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8 + +__all__ = [ + "CompressedTensorsScheme", + "CompressedTensorsW8A8Fp8", + "CompressedTensorsW8A16Fp8", +] diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a16_fp8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a16_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdf8381e75694bd72d7da3afd31305fe84b27545 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a16_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a8_fp8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a8_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b379d07e70925a731cf6ebd0324ae7698adc1345 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/__pycache__/compressed_tensors_w8a8_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py new file mode 100644 index 0000000000000000000000000000000000000000..3795d0a54ebe25a77921fffa70adf44892375150 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py @@ -0,0 +1,56 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from typing import Optional + +import torch + +__all__ = ["CompressedTensorsScheme"] + + +class CompressedTensorsScheme(ABC): + """ + Abstract class used to describe the weight creation and forward pass + of different quantization schemes supported by CompressedTensors. + """ + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """ + Get minimum device capability. + """ + raise NotImplementedError + + @abstractmethod + def create_weights(self, *args, **kwargs): + """ + Weight creation for the particular scheme. Inputs to this function + + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + """ + Run the forward pass for the particular scheme. This is where + scheme-specific dequant/quant steps/kernels should be applied. + + :param layer: torch.nn.Module with the registered weights and + other parameters relevant to the particular scheme. + :param x: input to the layer + :param bias: bias parameter + + """ + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + """ + Called after weight loading is complete for any cleanup that + needs to occur. + """ + raise NotImplementedError diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..fa7d77f2832bc0a058de4ea3a56cf44ecc4e6254 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -0,0 +1,153 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable, List, Optional + +import torch +from compressed_tensors.quantization import QuantizationStrategy + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from sglang.srt.layers.quantization.utils import convert_to_channelwise + +try: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, + ) + + MARLIN_FP8_AVAILABLE = True +except ImportError: + MARLIN_FP8_AVAILABLE = False + + def apply_fp8_marlin_linear(*args, **kwargs): + raise ImportError("vllm is not installed") + + def prepare_fp8_layer_for_marlin(*args, **kwargs): + raise ImportError("vllm is not installed") + + +__all__ = ["CompressedTensorsW8A16Fp8"] + +SUPPORTED_STRATEGIES = [QuantizationStrategy.CHANNEL, QuantizationStrategy.TENSOR] + + +class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): + def __init__(self, strategy: str, is_static_input_scheme: bool): + self.strategy = strategy + self.is_static_input_scheme = is_static_input_scheme + + if not MARLIN_FP8_AVAILABLE: + raise ImportError( + "vllm is not installed. To use CompressedTensorsW8A16Fp8, please install vllm" + ) + + @classmethod + def get_min_capability(cls) -> int: + # ampere and up + return 80 + + # W8A8-Fp8 kernels support only per-tensor and per-channel cases. + # So if we have a fused module (QKV, MLP) with per tensor scales, + # we expand each scale to its shard's channels. + def process_weights_after_loading(self, layer) -> None: + if self.strategy == QuantizationStrategy.TENSOR: + ws_channelwise = convert_to_channelwise( + layer.weight_scale, layer.logical_widths + ) + layer.weight_scale = torch.nn.Parameter(ws_channelwise, requires_grad=False) + else: + # required by torch.compile to be torch.nn.Parameter + layer.weight_scale = torch.nn.Parameter( + layer.weight_scale.data, requires_grad=False + ) + + # Weights must be transposed for marlin + layer.weight = torch.nn.Parameter(layer.weight.t(), requires_grad=False) + + if self.is_static_input_scheme: + # required by torch.compile to be torch.nn.Parameter + layer.input_scale = torch.nn.Parameter( + layer.input_scale.data, requires_grad=False + ) + prepare_fp8_layer_for_marlin(layer, strategy="channel") + + def create_weights( + self, + layer: torch.nn.Module, + input_size: int, + output_partition_sizes: List[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + if self.strategy == QuantizationStrategy.CHANNEL: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + elif self.strategy == QuantizationStrategy.TENSOR: + weight_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + else: + raise ValueError( + f"Unsupported weight strategy={self.strategy}, " + f"supported strategies are {SUPPORTED_STRATEGIES}" + ) + + weight_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", weight_scale) + + # INPUT SCALE (to deal with converted checkpoints) + if self.is_static_input_scheme: + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("input_scale", input_scale) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..ed13195c52dbf9eedd258b99c7e221a46170bade --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -0,0 +1,159 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable, List, Optional + +import torch +from compressed_tensors.quantization import QuantizationStrategy +from torch.nn import Parameter + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.utils import is_fp8_fnuz, requantize_with_max_scale + +__all__ = ["CompressedTensorsW8A8Fp8"] + + +class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): + + def __init__(self, strategy: str, is_static_input_scheme: bool): + self.strategy = strategy + self.is_static_input_scheme = is_static_input_scheme + + @classmethod + def get_min_capability(cls) -> int: + # lovelace and up + return 89 + + def process_weights_after_loading(self, layer) -> None: + # If per tensor, when we have a fused module (e.g. QKV) with per + # tensor scales (thus N scales being passed to the kernel), + # requantize so we can always run per tensor + if self.strategy == QuantizationStrategy.TENSOR: + max_w_scale, weight = requantize_with_max_scale( + weight=layer.weight, + weight_scale=layer.weight_scale, + logical_widths=layer.logical_widths, + ) + + if is_fp8_fnuz(): + input_scale = getattr(layer, "input_scale", None) + + weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=max_w_scale, input_scale=input_scale + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + + # If channelwise, scales are already lined up, so just transpose. + elif self.strategy == QuantizationStrategy.CHANNEL: + weight = layer.weight + + if is_fp8_fnuz(): + input_scale = getattr(layer, "input_scale", None) + + weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, + weight_scale=layer.weight_scale, + input_scale=input_scale, + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + else: + weight_scale = layer.weight_scale.data + + layer.weight = Parameter(weight.t(), requires_grad=False) + # required by torch.compile to be torch.nn.Parameter + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + else: + raise ValueError(f"Unknown quantization strategy {self.strategy}") + + # INPUT SCALE + if self.is_static_input_scheme and hasattr(layer, "input_scale"): + layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + else: + layer.input_scale = None + + def create_weights( + self, + layer: torch.nn.Module, + output_partition_sizes: List[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + + # WEIGHT + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + # TODO: update create_xxx_parameter functions to return + # the newly added parameters + if self.strategy == QuantizationStrategy.CHANNEL: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + else: + assert self.strategy == QuantizationStrategy.TENSOR + weight_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + # min requirement for fp8 kernels + weight_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", weight_scale) + + # INPUT SCALE + if self.is_static_input_scheme: + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + input_scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("input_scale", input_scale) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + use_per_token_if_dynamic=True, + compressed_tensor_quant=True, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/utils.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ddefa5ea3de5f9fdfc82417802f064ed250a89e3 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/compressed_tensors/utils.py @@ -0,0 +1,218 @@ +# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors +# SPDX-License-Identifier: Apache-2.0 + +import re +from types import MappingProxyType +from typing import Iterable, List, Mapping, Optional + +from compressed_tensors import CompressionFormat +from torch.nn import Module + + +def is_activation_quantization_format(format: str) -> bool: + _ACTIVATION_QUANTIZATION_FORMATS = [ + CompressionFormat.naive_quantized.value, + CompressionFormat.int_quantized.value, + CompressionFormat.float_quantized.value, + ] + return format in _ACTIVATION_QUANTIZATION_FORMATS + + +def should_ignore_layer( + layer_name: Optional[str], + ignore: Iterable[str] = tuple(), + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> bool: + if layer_name is None: + return False + + # layer_name = model.layers.0.self_attn.qkv_proj + # proj_name = qkv_proj + proj_name = layer_name.split(".")[-1] + + # Fused layers like gate_up_proj or qkv_proj will not be fused + # in the safetensors checkpoint. So, we convert the name + # from the fused version to unfused + check to make sure that + # each shard of the fused layer has the same scheme. + if proj_name in fused_mapping and layer_name not in ignore: + shard_proj_names = fused_mapping[proj_name] + + # Convert fused_name --> [shard_names] + shard_names = [ + layer_name.replace(proj_name, shard_proj_name) + for shard_proj_name in shard_proj_names + ] + + # Layer should be ignored if shards are ignored. + should_ignore_layer = None + for shard_name in shard_names: + should_ignore_shard = check_equal_or_regex_match( + layer_name=shard_name, targets=ignore + ) + + # If shard_idx=0, set layer ignore to match shard. + if should_ignore_layer is None: + should_ignore_layer = should_ignore_shard + + # If shard_idx=1+ confirm scheme matches prior shards. + elif should_ignore_shard != should_ignore_layer: + raise ValueError( + f"Found a different quantization schemes for " + f"{shard_proj_names} in {layer_name}. vLLM " + "requires all to use the same scheme." + ) + + # Unfused layers like down_proj and o_proj will match + # the safetensors checkpoint already. + else: + should_ignore_layer = check_equal_or_regex_match( + layer_name=layer_name, targets=ignore + ) + + assert should_ignore_layer is not None + return should_ignore_layer + + +def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool: + """ + Checks whether a layer_name is exactly equal or a regex match for + if target starts with 're:' to any target in list. + """ + for target in targets: + if _is_equal_or_regex_match(layer_name, target): + return True + return False + + +def find_matched_target( + layer_name: Optional[str], + module: Module, + targets: Iterable[str], + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> str: + """ + Helper function to look up which "target" in the compressed-tensors + config that a layer corresponds to. + + Recall that a compressed-tensors configs has a concept of + config_groups, where each layer can be quantized with with a different + scheme. + + targets in each config_group will be a list of either layer names + (or regexes corresponding to layer names) or names of torch Modules. + + First, we try to match the layer_name with a target + Second, we try to match the module's name with a target + Third, we try to map the layer_name to a list of fused module names. + *All* component module names must match in order for a match to be + successful. A successful match returns the first component target + + :param layer_name: layer name + :param module: torch.nn.Module + :param targets: list of targets to match the layer against + :param fused_mapping: map from fused layer names to its components + :param fused_strategy: either "all" or "any". If using "all", fused + layers match if "all" of its components match + """ + + if layer_name is None: + layer_name = "" + + matched_target = ( + _find_first_match(layer_name, targets) + or _find_first_match(module.__class__.__name__, targets, True) + or _match_fused_layer(layer_name, targets, fused_mapping) + ) + + if matched_target is None: + raise ValueError( + f"Unable to find matching target for {layer_name} in the " + "compressed-tensors config." + ) + + return matched_target + + +def _find_first_match( + value: str, targets: Iterable[str], check_contains: bool = False +) -> Optional[str]: + """ + Returns first element of target that matches value either + exactly or as a regex after 're:'. If check_contains is set to True, + additionally checks if the target string is contained within the value. + + :param value: string to compare the list of targets against + :param targets: list of targets to match the layer against + :param check_contains: whether or not to do a substring match + """ + + for target in targets: + if _is_equal_or_regex_match(value, target, check_contains=check_contains): + return target + return None + + +def _is_equal_or_regex_match( + value: str, target: str, check_contains: bool = False +) -> bool: + """ + Checks whether a value is exactly equal or a regex match for target + if target starts with 're:'. If check_contains is set to True, + additionally checks if the target string is contained within the value. + """ + + if target.startswith("re:"): + pattern = target[3:] + if re.match(pattern, value): + return True + elif check_contains: + if target.lower() in value.lower(): + return True + elif target == value: + return True + return False + + +def _match_fused_layer( + layer_name: str, + target_layers: Iterable[str], + fused_mapping: Mapping[str, List[str]], +) -> Optional[str]: + """ + Match a fused layer name to its corresponding individual layer in + target_layers. Returns first value in fused_mapping which matches targets + + Implements an "all" matching strategy where a fused layer matches iff + "all" of its components match + + :param layer_name: layer name + :param target_layers: list of targets to match the layer against + :param fused_mapping: map from fused layer names to its components + + Examples: + layer_name = "model.layers.0.self_attn.qkv_proj" + target_layers = ["model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.k_proj", + "model.layers.0.self_attn.v_proj"] + """ + # find layer_name in mapping + fused = next((key for key in fused_mapping if layer_name.endswith(key)), None) + if fused is None: + return None + + # expand path of unfused components + unfused_paths = [ + layer_name.replace(fused, unfused) for unfused in fused_mapping[fused] + ] + + # for each unfused component, find a match in targets + unfused_matches: List[Optional[str]] = [] + for unfused in unfused_paths: + for target in target_layers: + if _is_equal_or_regex_match(unfused, target): + unfused_matches.append(target) + break + else: + unfused_matches.append(None) + + return unfused_matches[0] if all(unfused_matches) else None diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/deep_gemm.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/deep_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..08ba0b9f97c306cb2dfae50c6a47a33a36383ddd --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/deep_gemm.py @@ -0,0 +1,385 @@ +import logging +import os +from contextlib import contextmanager +from dataclasses import dataclass +from enum import IntEnum, auto +from typing import Callable, Dict, List, Optional, Tuple + +import torch +from tqdm.contrib.concurrent import thread_map + +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import get_bool_env_var, get_device_sm, get_int_env_var, is_cuda + +_ENABLE_JIT_DEEPGEMM = False +if is_cuda(): + import deep_gemm + from deep_gemm import get_num_sms + from deep_gemm.jit_kernels.gemm import get_best_configs + from deep_gemm.jit_kernels.gemm import includes as deep_gemm_includes + from deep_gemm.jit_kernels.gemm import template as deep_gemm_gemm_template + from deep_gemm.jit_kernels.m_grouped_gemm import ( + template as deep_gemm_grouped_gemm_template, + ) + from deep_gemm.jit_kernels.tuner import jit_tuner + + sm_version = get_device_sm() + if sm_version == 90: + if get_bool_env_var("SGL_ENABLE_JIT_DEEPGEMM", default="true"): + _ENABLE_JIT_DEEPGEMM = True + +logger = logging.getLogger(__name__) + +_BUILTIN_M_LIST = list(range(1, 1024 * 16 + 1)) +_ENABLE_JIT_DEEPGEMM_PRECOMPILE = get_bool_env_var( + "SGL_JIT_DEEPGEMM_PRECOMPILE", "true" +) +_DO_COMPILE_ALL = True +_IS_FIRST_RANK_ON_NODE = get_bool_env_var("SGL_IS_FIRST_RANK_ON_NODE", "true") +_COMPILE_WORKERS = get_int_env_var("SGL_JIT_DEEPGEMM_COMPILE_WORKERS", 4) +_IN_PRECOMPILE_STAGE = get_bool_env_var("SGL_IN_DEEPGEMM_PRECOMPILE_STAGE", "false") + +# Force redirect deep_gemm cache_dir +os.environ["DG_CACHE_DIR"] = os.getenv( + "SGL_DG_CACHE_DIR", os.path.expanduser("~") + "/.cache/deep_gemm" +) + + +def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): + global _BUILTIN_M_LIST + global _DO_COMPILE_ALL + global _IS_FIRST_RANK_ON_NODE + + # Generate m_max + m_max = 1024 * 16 + if server_args.chunked_prefill_size < 1: + m_max = 1024 * 64 + elif server_args.chunked_prefill_size > 8192: + m_max = server_args.chunked_prefill_size * 2 + m_max = min(1024 * 128, m_max) + _BUILTIN_M_LIST = list(range(1, m_max + 1)) + + _IS_FIRST_RANK_ON_NODE = ServerArgs.base_gpu_id == gpu_id + + # Check if is the first rank on node. + # Default each rank will try compile all Ms to + # load all symbols at the launch stages. + # Avoid loading symbols at the serving stages. + _DO_COMPILE_ALL = _IS_FIRST_RANK_ON_NODE or not _IN_PRECOMPILE_STAGE + + +class DeepGemmKernelType(IntEnum): + GROUPED_GEMM_NT_F8F8BF16_MASKED = auto() + GROUPED_GEMM_NT_F8F8BF16_CONTIG = auto() + GEMM_NT_F8F8BF16 = auto() + + +@dataclass +class DeepGemmKernelHelper: + name: str + compile_func: Callable[ + [ + int, + int, + int, + Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]], + ], + None, + ] + configure_func: Callable[ + [int, int, int, int, int], + Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]], + ] + + +_INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict() + + +def _compile_warning_1(): + if not _IN_PRECOMPILE_STAGE and _IS_FIRST_RANK_ON_NODE: + logger.warning( + "Entering DeepGEMM JIT Pre-Complie session. " + "And it may takes a long time(Typically 10-20 mins) " + "if you have not run `sglang.compile_deep_gemm`. " + "Recommand to run `sglang.compile_deep_gemm` with same args as `sglang.launch_server`" + " for pre-compilation to reduce the overhead if you have not run it before. " + "For example: " + "`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`" + ) + + +def _compile_warning_2(): + logger.warning( + "Entering DeepGEMM JIT Single Kernel Complie session. " + "And it will makes inference throughput becomes flaky. " + "Please run `sglang.compile_deep_gemm` with same args as `sglang.launch_server`" + " for pre-compilation to solve this issue. " + "For example: " + "`python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code`" + ) + + +def _compile_grouped_gemm_nt_f8f8bf16_masked_one( + n: int, + k: int, + num_groups: int, + config: Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]], +) -> None: + # Auto-tuning with compilation + global deep_gemm_includes, deep_gemm_grouped_gemm_template + _, block_m, block_n, num_stages, tma_multicast_config, smem_config = config + _ = jit_tuner.compile_and_tune( + name="m_grouped_gemm_fp8_fp8_bf16_nt", + keys={ + "N": n, + "K": k, + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "SWIZZLE_D_MODE": smem_config[1], + "BLOCK_N_PADDING": smem_config[2], + "NUM_GROUPS": num_groups, + "NUM_STAGES": num_stages, + "NUM_TMA_MULTICAST": tma_multicast_config[0], + "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1], + "GEMM_TYPE": "GroupedMasked", + }, + space=(), + includes=deep_gemm_includes, + arg_defs=( + ("lhs", torch.float8_e4m3fn), + ("lhs_scales", torch.float), + ("rhs", torch.float8_e4m3fn), + ("rhs_scales", torch.float), + ("out", torch.bfloat16), + ("grouped_layout", torch.int32), + ("m", int), + ("stream", torch.cuda.Stream), + ("num_sms", int), + ("smem_size", int), + ), + template=deep_gemm_grouped_gemm_template, + args=[], + ) + + +def _compile_grouped_gemm_nt_f8f8bf16_contig_one( + n: int, + k: int, + num_groups: int, + config: Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]], +) -> None: + global deep_gemm_includes, deep_gemm_grouped_gemm_template + _, block_m, block_n, num_stages, tma_multicast_config, smem_config = config + _ = jit_tuner.compile_and_tune( + name="m_grouped_gemm_fp8_fp8_bf16_nt", + keys={ + "N": n, + "K": k, + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "SWIZZLE_D_MODE": smem_config[1], + "BLOCK_N_PADDING": smem_config[2], + "NUM_GROUPS": num_groups, + "NUM_STAGES": num_stages, + "NUM_TMA_MULTICAST": tma_multicast_config[0], + "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1], + "GEMM_TYPE": "GroupedContiguous", + }, + space=(), + includes=deep_gemm_includes, + arg_defs=( + ("lhs", torch.float8_e4m3fn), + ("lhs_scales", torch.float), + ("rhs", torch.float8_e4m3fn), + ("rhs_scales", torch.float), + ("out", torch.bfloat16), + ("grouped_layout", torch.int32), + ("m", int), + ("num_groups", int), + ("stream", torch.cuda.Stream), + ("num_sms", int), + ("smem_size", int), + ), + template=deep_gemm_grouped_gemm_template, + args=[], + ) + + +def _compile_gemm_nt_f8f8bf16_one( + n: int, + k: int, + _: int, # _ is a dummy parameter to align with other interfaces + config: Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]], +) -> None: + global deep_gemm_includes, deep_gemm_gemm_template + _, block_m, block_n, num_stages, tma_multicast_config, smem_config = config + _ = jit_tuner.compile_and_tune( + name="gemm_fp8_fp8_bf16_nt", + keys={ + "N": n, + "K": k, + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "SWIZZLE_D_MODE": smem_config[1], + "BLOCK_N_PADDING": smem_config[2], + "NUM_STAGES": num_stages, + "NUM_TMA_MULTICAST": tma_multicast_config[0], + "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1], + }, + space=(), + includes=deep_gemm_includes, + arg_defs=( + ("lhs", torch.float8_e4m3fn), + ("lhs_scales", torch.float), + ("rhs", torch.float8_e4m3fn), + ("rhs_scales", torch.float), + ("out", torch.bfloat16), + ("m", int), + ("stream", torch.cuda.Stream), + ("num_sms", int), + ("smem_size", int), + ), + template=deep_gemm_gemm_template, + args=[], + ) + + +_KERNEL_HELPER_DICT: Dict[DeepGemmKernelType, DeepGemmKernelHelper] = { + DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: DeepGemmKernelHelper( + name="m_grouped_gemm_fp8_fp8_bf16_nt_masked", + compile_func=_compile_grouped_gemm_nt_f8f8bf16_masked_one, + configure_func=lambda m, n, k, num_groups, num_sms: get_best_configs( + m, n, k, num_groups, num_sms, is_grouped_masked=True + ), + ), + DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: DeepGemmKernelHelper( + name="m_grouped_gemm_fp8_fp8_bf16_nt_contiguous", + compile_func=_compile_grouped_gemm_nt_f8f8bf16_contig_one, + configure_func=lambda m, n, k, _, num_sms: get_best_configs( + m, n, k, 1, num_sms, is_grouped_contiguous=True + ), + ), + DeepGemmKernelType.GEMM_NT_F8F8BF16: DeepGemmKernelHelper( + name="gemm_fp8_fp8_bf16_nt", + compile_func=_compile_gemm_nt_f8f8bf16_one, + configure_func=lambda m, n, k, _, num_sms: get_best_configs( + m, n, k, 1, num_sms + ), + ), +} + + +def _maybe_compile_deep_gemm_one_type_all( + kernel_type: DeepGemmKernelType, + n: int, + k: int, + num_groups: int, + m_list: Optional[List[int]] = None, +) -> None: + + global _INITIALIZATION_DICT + global _BUILTIN_M_LIST + + query_key = (kernel_type, n, k, num_groups) + if ( + _ENABLE_JIT_DEEPGEMM_PRECOMPILE + and _DO_COMPILE_ALL + and _INITIALIZATION_DICT.get(query_key) is None + ): + _INITIALIZATION_DICT[query_key] = True + + kernel_helper = _KERNEL_HELPER_DICT[kernel_type] + _compile_warning_1() + logger.info( + f"Try DeepGEMM JIT Compiling for " + f"<{kernel_helper.name}> N={n}, K={k}, num_groups={num_groups} with all Ms." + f"{' It only takes a litte time(Typically 1 sec) if you have run `sglang.compile_deep_gemm`. ' if not _IN_PRECOMPILE_STAGE else ''}" + ) + + # NOTE(alcanderian): get_num_sms should be change when 2-batch-overlap is introduced + num_sms = get_num_sms() + collected_configs = set() + for m in m_list if m_list is not None else _BUILTIN_M_LIST: + # Put config into set to get unique configs and reduce cases to be compiled + collected_configs.add( + kernel_helper.configure_func(m, n, k, num_groups, num_sms) + ) + compile_func = lambda config: kernel_helper.compile_func( + n, k, num_groups, config + ) + thread_map(compile_func, collected_configs, max_workers=_COMPILE_WORKERS) + + +def grouped_gemm_nt_f8f8bf16_masked( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + masked_m: torch.Tensor, + expected_m: int, +): + num_groups, _, k = lhs[0].shape + _, n, _ = rhs[0].shape + + kernel_type = DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED + _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups) + + with _log_jit_build(expected_m, n, k, kernel_type): + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked( + lhs, rhs, out, masked_m, expected_m + ) + + +def grouped_gemm_nt_f8f8bf16_contig( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + m_indices: torch.Tensor, +): + m, k = lhs[0].shape + num_groups, n, _ = rhs[0].shape + + kernel_type = DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG + _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups) + + with _log_jit_build(m, n, k, kernel_type): + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs, rhs, out, m_indices) + + +def gemm_nt_f8f8bf16( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, +): + m, k = lhs[0].shape + n, _ = rhs[0].shape + + kernel_type = DeepGemmKernelType.GEMM_NT_F8F8BF16 + _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, 1) + + with _log_jit_build(m, n, k, kernel_type): + deep_gemm.gemm_fp8_fp8_bf16_nt(lhs, rhs, out) + + +@contextmanager +def _log_jit_build(M: int, N: int, K: int, kernel_type: DeepGemmKernelType): + if _IN_PRECOMPILE_STAGE: + yield + return + + from deep_gemm.jit.runtime import RuntimeCache + + origin_func = RuntimeCache.__getitem__ + + def __patched_func(self, *args, **kwargs): + ret = origin_func(self, *args, **kwargs) + if ret is None: + kernel_helper = _KERNEL_HELPER_DICT[kernel_type] + _compile_warning_2() + logger.warning( + f"DeepGEMM JIT Compiling for <{kernel_helper.name}> M={M}, N={N}, K={K}. Please wait." + ) + return ret + + RuntimeCache.__getitem__ = __patched_func + yield + RuntimeCache.__getitem__ = origin_func diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..2a5de2f21bd1ec19f64b0ad4f7bc1a36baaeb41b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8.py @@ -0,0 +1,997 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/quantization/fp8.py + +import logging +from typing import Any, Callable, Dict, List, Optional + +import torch +import torch.nn.functional as F +from torch.nn import Module +from torch.nn.parameter import Parameter + +try: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, + ) + + MARLIN_FP8_AVAILABLE = True +except ImportError: + MARLIN_FP8_AVAILABLE = False + + def dummy_func(*args, **kwargs): + raise ImportError( + "marlin FP8 requires some operators from vllm. Please install vllm." + ) + + apply_fp8_marlin_linear = prepare_fp8_layer_for_marlin = dummy_func + + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from sglang.srt.layers.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import ( + per_token_group_quant_fp8, + scaled_fp8_quant, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + apply_w8a8_block_fp8_linear, + cutlass_fp8_supported, + input_to_float8, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.utils import ( + all_close_1d, + convert_to_channelwise, + is_layer_skipped, + per_tensor_dequantize, + requantize_with_max_scale, +) +from sglang.srt.utils import ( + get_bool_env_var, + is_cuda, + is_hip, + print_warning_once, + set_weight_attrs, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() + +if _is_hip: + from aiter import ActivationType, QuantType + from aiter.fused_moe_bf16_asm import asm_moe, ck_moe_2stages + from aiter.ops.shuffle import shuffle_weight + +if not _is_cuda: + from vllm._custom_ops import scaled_fp8_quant + + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class Fp8Config(QuantizationConfig): + """Config class for FP8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: Optional[List[str]] = None, + weight_block_size: List[int] = None, + ) -> None: + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + if is_checkpoint_fp8_serialized: + logger.warning( + "Detected fp8 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + f"The block-wise quantization only supports fp8-serialized checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." + ) + if activation_scheme != "dynamic": + raise ValueError( + f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + + @classmethod + def get_name(cls) -> str: + return "fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "Fp8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return Fp8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return Fp8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class Fp8LinearMethod(LinearMethodBase): + """Linear method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Limitations: + 1. Only support per-tensor quantization due to torch._scaled_mm support. + 2. Only support float8_e4m3fn data type due to the limitation of + torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + # For GPUs that lack FP8 hardware support, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + self.use_marlin = ( + get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") and MARLIN_FP8_AVAILABLE + ) + # Disable marlin for ROCm + if _is_hip: + self.use_marlin = False + + self.block_quant = self.quant_config.weight_block_size is not None + if self.block_quant: + # Marlin doesn't support block-wise fp8 + self.use_marlin = False + + 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, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + tp_size = get_tensor_model_parallel_world_size() + if self.block_quant: + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + # Required by collum parallel or enabling merged weights + if ( + tp_size > 1 and output_size // output_size_per_partition == tp_size + ) or len(output_partition_sizes) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, input_size_per_partition, dtype=weight_dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # If checkpoint is serialized fp8, load them. + # Otherwise, wait until process_weights_after_loading. + if self.quant_config.is_checkpoint_fp8_serialized: + # WEIGHT SCALE + if self.block_quant: + assert self.quant_config.activation_scheme == "dynamic" + scale = BlockQuantScaleParameter( + data=torch.empty( + (output_size_per_partition + block_n - 1) // block_n, + (input_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale_inv", scale) + else: + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", scale) + + # INPUT ACTIVATION SCALE + if self.quant_config.activation_scheme == "static": + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("input_scale", scale) + else: + layer.register_parameter("input_scale", None) + + def process_weights_after_loading(self, layer: Module) -> None: + # Block quant doesn't need to process weights after loading + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_hip: + # activation_scheme: dynamic + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.weight, + weight_scale=layer.weight_scale_inv, + input_scale=None, + ) + layer.weight = torch.nn.Parameter(weight, requires_grad=False) + layer.weight_scale_inv = torch.nn.Parameter( + weight_scale, requires_grad=False + ) + layer.input_scale = None + else: + layer.weight = torch.nn.Parameter( + layer.weight.data, requires_grad=False + ) + layer.weight_scale_inv = torch.nn.Parameter( + layer.weight_scale_inv.data, requires_grad=False + ) + return + + layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) + + # If checkpoint not serialized fp8, quantize the weights. + if not self.quant_config.is_checkpoint_fp8_serialized: + if self.cutlass_fp8_supported or self.use_marlin: + # apply per-channel quantization default, as cutlass sgl-kernel and marlin only support per-channel scale + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + else: + # per-tensor quantization + qweight, weight_scale = input_to_float8(layer.weight) + + # Update the layer with the new values. + layer.weight = Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + layer.input_scale = None + + # If checkpoint is fp8, handle that there are N scales for N + # shards in a fused module + else: + layer.weight_scale = torch.nn.Parameter( + layer.weight_scale.data, requires_grad=False + ) + if self.quant_config.activation_scheme == "static": + layer.input_scale = torch.nn.Parameter( + layer.input_scale.data, requires_grad=False + ) + + # cutlass sgl-kernel and marlin only support per-channel scale + if self.cutlass_fp8_supported or self.use_marlin: + weight = layer.weight + weight_scale = convert_to_channelwise( + layer.weight_scale, layer.logical_widths + ) + else: + # Dequant -> Quant with max scale so we can run per tensor. + weight = layer.weight + weight_scale = layer.weight_scale + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_hip: + weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, + weight_scale=weight_scale, + input_scale=layer.input_scale, + ) + if input_scale is not None: + layer.input_scale = Parameter(input_scale, requires_grad=False) + + weight_scale, weight = requantize_with_max_scale( + weight=weight, + weight_scale=weight_scale, + logical_widths=layer.logical_widths, + ) + + # Update layer with new values. + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + if self.quant_config.activation_scheme == "static": + layer.input_scale = Parameter( + layer.input_scale.max(), requires_grad=False + ) + + if self.use_marlin: + prepare_fp8_layer_for_marlin(layer) + # Activations not quantized for marlin. + del layer.input_scale + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + if self.use_marlin: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + if self.block_quant: + return apply_w8a8_block_fp8_linear( + input=x, + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=None, + bias=bias, + ) + + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + use_per_token_if_dynamic=False, + ) + + +class Fp8MoEMethod: + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Args: + quant_config: The quantization config. + """ + + def __new__(cls, *args, **kwargs): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoEMethodBase + + if not hasattr(cls, "_initialized"): + original_init = cls.__init__ + new_cls = type( + cls.__name__, + (FusedMoEMethodBase,), + { + "__init__": original_init, + **{k: v for k, v in cls.__dict__.items() if k != "__dict__"}, + }, + ) + obj = super(new_cls, new_cls).__new__(new_cls) + obj.__init__(*args, **kwargs) + return obj + return super().__new__(cls) + + def __init__(self, quant_config): + self.quant_config = quant_config + self.block_quant = self.quant_config.weight_block_size is not None + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + if self.quant_config.is_checkpoint_fp8_serialized: + params_dtype = ( + torch.uint32 + if get_bool_env_var("SGLANG_INT4_WEIGHT") + else torch.float8_e4m3fn + ) + tp_size = get_tensor_model_parallel_world_size() + if self.block_quant: + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + # Required by collum parallel or enabling merged weights + if intermediate_size % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1: + # Required by row parallel + if intermediate_size % block_k != 0: + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + # INT4 MoE weight - INT32 packed + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size, + hidden_size // 8, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, hidden_size, intermediate_size // 8, dtype=params_dtype + ), + requires_grad=False, + ) + else: + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size, hidden_size, dtype=params_dtype + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, hidden_size, intermediate_size, dtype=params_dtype + ), + requires_grad=False, + ) + + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + if self.block_quant: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * ((intermediate_size + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + assert self.quant_config.activation_scheme == "dynamic" + else: + # Allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + if ( + _is_hip + ): # and get_bool_env_var("SGLANG_AITER_MOE"): TODO: add check back after triton kernel + # ROCm - using column scaling, duplicate scaling numbers in case per tensor scaling + w13_weight_scale1 = torch.nn.Parameter( + torch.ones(num_experts, 2 * intermediate_size, dtype=torch.float32), + requires_grad=False, + ) + w2_weight_scale1 = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale1", w13_weight_scale1) + layer.register_parameter("w2_weight_scale1", w2_weight_scale1) + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + if self.block_quant + else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + # If loading fp8 checkpoint, pass the weight loaders. + # If loading an fp16 checkpoint, do not (we will quantize in + # process_weights_after_loading() + if self.quant_config.is_checkpoint_fp8_serialized: + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + set_weight_attrs(w13_weight_scale1, extra_weight_attrs) + set_weight_attrs(w2_weight_scale1, extra_weight_attrs) + + # INPUT_SCALES + if self.quant_config.activation_scheme == "static": + if not self.quant_config.is_checkpoint_fp8_serialized: + raise ValueError( + "Found static activation scheme for checkpoint that " + "was not serialized fp8." + ) + + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: Module) -> None: + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): + self.process_weights_hip_int4(layer) + return + + # Block quant doesn't need to process weights after loading + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_hip: + # activation_scheme: dynamic + w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w13_weight, + weight_scale=layer.w13_weight_scale_inv, + input_scale=None, + ) + w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w2_weight, + weight_scale=layer.w2_weight_scale_inv, + input_scale=None, + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale_inv = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + layer.w13_input_scale = None + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale_inv = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + layer.w2_input_scale = None + + if get_bool_env_var("SGLANG_AITER_MOE"): + # Pre-shuffle weights + layer.w13_weight.data = shuffle_weight( + layer.w13_weight.contiguous(), (16, 16) + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight.contiguous(), (16, 16) + ) + return + + # If checkpoint is fp16 or bfloat16, quantize in place. + if not self.quant_config.is_checkpoint_fp8_serialized: + # If ROCm, use float8_e4m3fnuz instead (MI300x HW) + fp8_dtype = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype) + w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype) + + # Re-initialize w13_scale because we directly quantize + # merged w13 weights and generate a single scaling factor. + layer.w13_weight_scale = torch.nn.Parameter( + torch.ones( + layer.num_experts, dtype=torch.float32, device=w13_weight.device + ), + requires_grad=False, + ) + for expert in range(layer.num_experts): + w13_weight[expert, :, :], layer.w13_weight_scale[expert] = ( + scaled_fp8_quant(layer.w13_weight.data[expert, :, :]) + ) + w2_weight[expert, :, :], layer.w2_weight_scale[expert] = ( + scaled_fp8_quant(layer.w2_weight.data[expert, :, :]) + ) + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + + if _is_hip: + self.process_weights_hip_scale_padding(layer) + return + + # If checkpoint is fp8, we need to handle that the + # MoE kernels require single activation scale and single weight + # scale for w13 per expert. + else: + # Fp8 moe kernels require a single activation scale. + # We take the max of all the scales in case they differ. + if self.quant_config.activation_scheme == "static": + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + print_warning_once( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer. " + ) + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_hip: + # Normalize the weights and scales + w13_weight, w13_weight_scale, w13_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale + ) + ) + w2_weight, w2_weight_scale, w2_input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale + ) + ) + # Reset the parameter + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + if w13_input_scale is not None: + layer.w13_input_scale = torch.nn.Parameter( + w13_input_scale, requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter( + w2_weight_scale, requires_grad=False + ) + if w2_input_scale is not None: + layer.w2_input_scale = torch.nn.Parameter( + w2_input_scale, requires_grad=False + ) + # Fp8 moe kernel needs single weight scale for w13 per expert. + # We take the max then dequant and requant each expert. + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.num_experts): + start = 0 + for shard_id in range(2): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter( + max_w13_scales, requires_grad=False + ) + + if _is_hip: + self.process_weights_hip_scale_padding(layer) + return + + def process_weights_hip_int4(self, layer: Module): + # TODO: and get_bool_env_var("SGLANG_AITER_MOE"): add after triton kernel added + # INT4-FP8 (INT4 MoE Weight, FP8 Compute) + # Weight Permutation + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + + # INT4-FP8 : offset INT4 w13_weight_scale1 to single w13_weight_scale + # Fp8 moe kernel needs single fp8 w13_weight_scale for w13 per expert. + # We won't do requant each expert's fp8 weight (not direct available), + # instead we adjust half of INT4 w13_weight_scale1 numbers + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + for expert_id in range(layer.num_experts): + start = 0 + max_w13_scale_fp8 = max_w13_scales[expert_id] + for shard_id in range(2): + if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8: + int4_rescale = ( + layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8 + ) + layer.w13_weight_scale1[expert_id][ + start : start + shard_size + ] *= int4_rescale + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False) + + # special hack to asm_moe, which takes (weight_scale1 * weight_scale) as post GEMM scaling + # optimal design - shall apply per-column weight_scale1 before GEMM, and weight_scale post + for expert_id in range(layer.num_experts): + layer.w13_weight_scale1[expert_id] *= max_w13_scales[expert_id] + layer.w2_weight_scale1[expert_id] *= layer.w2_weight_scale[expert_id] + + def process_weights_hip_scale_padding(self, layer: Module): + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + padding_size, # Avoid circular import + ) + + if get_bool_env_var("SGLANG_AITER_MOE"): + layer.w13_weight = torch.nn.Parameter( + shuffle_weight(layer.w13_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + shuffle_weight(layer.w2_weight.data, (16, 16)), + requires_grad=False, + ) + torch.cuda.empty_cache() + # ROCm (SGLANG_AITER_MOE): using column-wise scaling + layer.w13_weight_scale1 *= layer.w13_weight_scale.unsqueeze(-1) + layer.w2_weight_scale1 *= layer.w2_weight_scale.unsqueeze(-1) + elif get_bool_env_var("SGLANG_MOE_PADDING"): + # If ROCm, apply weight padding (min. Mem channel contention) only if set + layer.w13_weight = torch.nn.Parameter( + F.pad(layer.w13_weight.data, (0, padding_size), "constant", 0), + requires_grad=False, + ) + torch.cuda.empty_cache() + layer.w2_weight = torch.nn.Parameter( + F.pad(layer.w2_weight.data, (0, padding_size), "constant", 0), + requires_grad=False, + ) + torch.cuda.empty_cache() + + def 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", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + # Expert selection + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + if _is_hip: + if get_bool_env_var("SGLANG_INT4_WEIGHT"): + # TODO: add triton kernel and add check get_bool_env_var("SGLANG_AITER_MOE") + assert not no_combine, f"{no_combine=} is not supported." + return ck_moe_2stages( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + QuantType.per_Token, + layer.w13_weight_scale1, + layer.w2_weight_scale1, + activation=( + ActivationType.Silu + if activation == "silu" + else ActivationType.Gelu + ), + ) + + if get_bool_env_var("SGLANG_AITER_MOE"): + assert not no_combine, f"{no_combine=} is not supported." + if self.block_quant: + # TODO(SGLANG_AITER_MOE): FP8 block_quant only supports 'silu' for the time-being. + assert ( + activation == "silu" + ), f"SGLANG_AITER_MOE: FP8 bloack_quant {activation=} will be supported later, unset SGLANG_AITER_MOE" + return asm_moe( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + layer.w13_weight_scale_inv, + layer.w2_weight_scale_inv, + block_shape=tuple(self.quant_config.weight_block_size), + expert_mask=None, + ) + else: + return ck_moe_2stages( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + QuantType.per_Token, + layer.w13_weight_scale1, + layer.w2_weight_scale1, + activation=( + ActivationType.Silu + if activation == "silu" + else ActivationType.Gelu + ), + ) + + # Expert fusion with FP8 quantization + return fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace and not no_combine, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + use_fp8_w8a8=True, + w1_scale=( + layer.w13_weight_scale_inv + if self.block_quant + else layer.w13_weight_scale + ), + w2_scale=( + layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale + ), + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.quant_config.weight_block_size, + no_combine=no_combine, + ) + + +class Fp8KVCacheMethod(BaseKVCacheMethod): + """ + Supports loading kv-cache scaling factors from FP8 checkpoints. + """ + + def __init__(self, quant_config: Fp8Config): + super().__init__(quant_config) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_kernel.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..45157527ebf23ce0d6fe50a38677bcc5d2b705c2 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_kernel.py @@ -0,0 +1,1068 @@ +# Copyright 2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization.deep_gemm import _ENABLE_JIT_DEEPGEMM +from sglang.srt.utils import ( + direct_register_custom_op, + get_device_core_count, + get_device_name, + is_cuda, + is_hip, + supports_custom_op, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_fp8_type = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn +if _is_hip: + fp8_max = 224.0 +else: + fp8_max = torch.finfo(_fp8_type).max +fp8_min = -fp8_max + +if _is_cuda: + from sgl_kernel import ( + sgl_per_tensor_quant_fp8, + sgl_per_token_group_quant_fp8, + sgl_per_token_quant_fp8, + ) + + from sglang.srt.layers.quantization.deep_gemm import ( + gemm_nt_f8f8bf16 as deep_gemm_gemm_nt_f8f8bf16, + ) + +logger = logging.getLogger(__name__) + +if supports_custom_op(): + + def deep_gemm_fp8_fp8_bf16_nt( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + ) -> None: + deep_gemm_gemm_nt_f8f8bf16((A, As), (B, Bs), C) + + def deep_gemm_fp8_fp8_bf16_nt_fake( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + ) -> None: + return + + direct_register_custom_op( + op_name="deep_gemm_fp8_fp8_bf16_nt", + op_func=deep_gemm_fp8_fp8_bf16_nt, + mutates_args=["C"], + fake_impl=deep_gemm_fp8_fp8_bf16_nt_fake, + ) + + +@triton.jit +def _per_token_group_quant_fp8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + y_stride, + # Collums of input + N, + # Avoid to divide zero + eps, + # Information for float8 + fp8_min, + fp8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group quantization on a + tensor. + + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + y_s_inv = 1.0 / y_s + y_q = tl.clamp(y * y_s_inv, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +@triton.jit +def _per_token_group_quant_fp8_colmajor( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + # Num columns of y + y_num_columns, + # Stride from one column to the next of y_s + y_s_col_stride, + # Avoid to divide zero + eps, + # Information for float8 + fp8_min, + fp8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * group_size + y_q_ptr += g_id * group_size + + # Convert g_id the flattened block coordinate to 2D so we can index + # into the output y_scales matrix + blocks_per_row = y_num_columns // group_size + scale_col = g_id % blocks_per_row + scale_row = g_id // blocks_per_row + y_s_ptr += scale_col * y_s_col_stride + scale_row + + cols = tl.arange(0, BLOCK) # group_size <= BLOCK + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_fp8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed float8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + x_q = torch.empty_like(x, device=x.device, dtype=_fp8_type) + M = x.numel() // group_size + N = group_size + if column_major_scales: + if scale_tma_aligned: + # aligned to 4 * sizeof(float) + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = torch.empty( + x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + else: + x_s = torch.empty( + (x.shape[-1] // group_size,) + x.shape[:-1], + device=x.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + if column_major_scales: + _per_token_group_quant_fp8_colmajor[(M,)]( + x, + x_q, + x_s, + group_size, + x.shape[1], + x_s.stride(1), + eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + else: + _per_token_group_quant_fp8[(M,)]( + x, + x_q, + x_s, + group_size, + N, + eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + return x_q, x_s + + +def sglang_per_token_group_quant_fp8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + column_major_scales: bool = False, + scale_tma_aligned: bool = False, +): + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + x_q = torch.empty_like(x, device=x.device, dtype=_fp8_type) + if column_major_scales: + if scale_tma_aligned: + # aligned to 4 * sizeof(float) + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = torch.empty( + x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + else: + x_s = torch.empty( + (x.shape[-1] // group_size,) + x.shape[:-1], + device=x.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, eps, fp8_min, fp8_max) + + return x_q, x_s + + +def sglang_per_token_quant_fp8( + x: torch.Tensor, + dtype: torch.dtype = _fp8_type, +): + assert x.is_contiguous(), "`x` is not contiguous" + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + x_s = torch.empty( + x.shape[0], + 1, + device=x.device, + dtype=torch.float32, + ) + + sgl_per_token_quant_fp8(x, x_q, x_s) + + return x_q, x_s + + +@triton.jit +def _static_quant_fp8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + y_s_repeat_ptr, + # Stride of input + y_stride, + # Collums of input + N, + # Information for float8 + fp8_min, + fp8_max, + # Meta-parameters + BLOCK: tl.constexpr, + REPEAT_SCALE: tl.constexpr, +): + """A Triton-accelerated function to perform quantization using the given scale on a + tensor + + This function converts the tensor values into float8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + if REPEAT_SCALE: + y_s_repeat_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + y_s = tl.load(y_s_ptr).to(tl.float32) + y_s_inv = 1.0 / y_s + y_q = tl.clamp(y * y_s_inv, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + if REPEAT_SCALE: + tl.store(y_s_repeat_ptr, y_s) + + +def static_quant_fp8( + x: torch.Tensor, + x_s: torch.Tensor, + repeat_scale: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform static quantization using the given scale on an input tensor `x`. + + It converts the tensor values into signed float8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + x_s: The quantization scale. + repeat_scale: Whether to broadcast per-tensor scale to per-channel scale. + dtype: The dype of output tensor. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert x.is_contiguous(), "`x` is not contiguous" + assert x_s.numel() == 1, "only supports per-tensor scale" + + x_q = torch.empty_like(x, device=x.device, dtype=_fp8_type) + M = x.numel() // x.shape[-1] + N = x.shape[-1] + if repeat_scale: + x_s_repeat = torch.empty( + (M, 1), + device=x.device, + dtype=torch.float32, + ) + else: + x_s_repeat = None + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _static_quant_fp8[(M,)]( + x, + x_q, + x_s, + x_s_repeat, + N, + N, + fp8_min=fp8_min, + fp8_max=fp8_max, + BLOCK=BLOCK, + REPEAT_SCALE=repeat_scale, + num_warps=num_warps, + num_stages=num_stages, + ) + x_s = x_s_repeat if repeat_scale else x_s + return x_q, x_s + + +@triton.jit +def _w8a8_block_fp8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and store the result in output + tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.jit +def _w8a8_block_fp8_matmul_unrolledx4( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and store the result in output + tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + # manually unroll to 4 iterations + UNROLL_FACTOR = 4 + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * UNROLL_FACTOR)): + # 1st iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = (k * UNROLL_FACTOR) * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # 2nd iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # 3rd iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # 4th iteration + a = tl.load( + a_ptrs, + mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K, + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K, + other=0.0, + ) + + k_start = k_start + BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@functools.lru_cache +def get_w8a8_block_fp8_configs( + N: int, K: int, block_n: int, block_k: int +) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the w8a8 block fp8 kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + + # First look up if an optimized configuration is available in the configs + # directory + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info( + "Using configuration from %s for W8A8 Block FP8 kernel.", + config_file_path, + ) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + logger.warning( + ( + "Using default W8A8 Block FP8 kernel config. Performance might be sub-optimal! " + "Config file not found at %s" + ), + config_file_path, + ) + return None + + +def w8a8_block_fp8_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Default config + # Block-wise quant: BLOCK_SIZE_K must be divisable by block_size[1] + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_size[0], + "BLOCK_SIZE_K": block_size[1], + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + + def grid(META): + return ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + # Use manually unrolledx4 kernel on AMD GPU when the grid size is small. + # Empirical testing shows the sweet spot lies when it's less than the # of + # compute units available on the device. + num_workgroups = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv( + N, config["BLOCK_SIZE_N"] + ) + + # deepgemm only support bf16 + if C.dtype == torch.bfloat16 and _ENABLE_JIT_DEEPGEMM: + if supports_custom_op(): + torch.ops.sglang.deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C) + else: + deep_gemm_gemm_nt_f8f8bf16((A, As), (B, Bs), C) + else: + kernel = ( + _w8a8_block_fp8_matmul_unrolledx4 + if (_is_hip == True and num_workgroups <= get_device_core_count()) + else _w8a8_block_fp8_matmul + ) + + kernel[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C + + +@triton.jit +def _per_tensor_quant_mla_fp8_stage1( + x_ptr, + x_s_ptr, + head_size, + x_stride_h, + x_stride_s, + eps, + fp8_max, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + head_id = tl.program_id(1) + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < head_size + + x_ptr += head_id * x_stride_h + seq_id * x_stride_s + x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(x)), eps) + + tl.atomic_max(x_s_ptr, _absmax / fp8_max) + + +@triton.jit +def _per_tensor_quant_mla_fp8_stage2( + x_ptr, + x_s_ptr, + x_q_ptr, + num_seq, + head_size, + x_stride_h, + x_stride_s, + fp8_min, + fp8_max, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + head_id = tl.program_id(1) + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < head_size + + x_s = tl.load(x_s_ptr) + x_s_inv = 1.0 / x_s + + x_ptr += head_id * x_stride_h + seq_id * x_stride_s + x_q_ptr += head_id * num_seq * head_size + seq_id * head_size + + x = tl.load(x_ptr + offset, mask=mask, other=0.0).to(tl.float32) + x_q = tl.clamp(x * x_s_inv, fp8_min, fp8_max).to(x_q_ptr.dtype.element_ty) + tl.store(x_q_ptr + offset, x_q, mask=mask) + + +def per_tensor_quant_mla_fp8( + x: torch.Tensor, x_s_out: torch.Tensor, eps: float = 1e-12 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function quantizes input values to float8 values with tensor-wise quantization + and specialized for mla absorbed case. + """ + assert x.dim() == 3, "`x` is not a 3d-tensor" + assert ( + x_s_out.shape == (1,) + and x_s_out.dtype == torch.float32 + and x_s_out.device == x.device + ) + + x_q = x.new_empty(x.size(), dtype=_fp8_type) + + num_head, num_seq, head_size = x.shape + BLOCK_SIZE = triton.next_power_of_2(head_size) + grid = (num_seq, num_head) + + _per_tensor_quant_mla_fp8_stage1[grid]( + x, + x_s_out, + head_size, + x.stride(0), + x.stride(1), + eps, + fp8_max, + BLOCK_SIZE, + ) + _per_tensor_quant_mla_fp8_stage2[grid]( + x, + x_s_out, + x_q, + num_seq, + head_size, + x.stride(0), + x.stride(1), + fp8_min, + fp8_max, + BLOCK_SIZE, + ) + + return x_q, x_s_out + + +@triton.jit +def _per_token_group_quant_mla_deep_gemm_masked_fp8( + y_ptr, + y_q_ptr, + y_s_ptr, + masked_m_ptr, + group_size, + y_stride_b, + y_stride_t, + y_q_stride_b, + y_q_stride_t, + y_s_stride_b, + y_s_stride_g, + eps, + fp8_min, + fp8_max, + NUM_GROUP: tl.constexpr, + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor for deep_gemm grouped_gemm_masked. + This function converts the tensor values into float8 values. + y and y_q: (b, t, k) + y_s: (b, k//group_size, t) + """ + t_id = tl.program_id(0) + b_id = tl.program_id(1) + + y_ptr += b_id * y_stride_b + t_id * y_stride_t + y_q_ptr += b_id * y_q_stride_b + t_id * y_q_stride_t + y_s_ptr += b_id * y_s_stride_b + t_id + + if t_id == 0: + tl.store(masked_m_ptr + b_id, tl.num_programs(0)) + + cols = tl.arange(0, BLOCK) # group_size <= BLOCK + mask = cols < group_size + + for gid in range(NUM_GROUP): + y = tl.load(y_ptr + gid * group_size + cols, mask=mask, other=0.0).to( + tl.float32 + ) + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + gid * group_size + cols, y_q, mask=mask) + tl.store(y_s_ptr + gid * y_s_stride_g, y_s) + + +def per_tensor_quant_mla_deep_gemm_masked_fp8( + x: torch.Tensor, + group_size: int = 128, + eps: float = 1e-12, + dtype: torch.dtype = torch.float8_e4m3fn, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function quantizes input values to float8 values with per-token-group-quantization + for deep_gemm grouped_gemm_masked and specialized for mla absorbed case. + """ + assert x.dim() == 3, "`x` is not a 3d-tensor" + + finfo = torch.finfo(dtype) + fp8_max = finfo.max + if _is_hip: + dtype = torch.float8_e4m3fnuz + fp8_max = 224.0 + + b, m, k = x.shape + aligned_m = (m + 255) // 256 * 256 # 256 is the max block_m of the gemm kernel + num_tiles_k = k // group_size + assert num_tiles_k * group_size == k, f"k % {group_size} must be zero" + + x_q = x.new_empty((b, aligned_m, k), dtype=dtype) + x_s = x.new_empty((b, num_tiles_k, aligned_m), dtype=torch.float32) + masked_m = x.new_empty((b,), dtype=torch.int32) + + BLOCK_SIZE = triton.next_power_of_2(group_size) + grid = (m, b) + + _per_token_group_quant_mla_deep_gemm_masked_fp8[grid]( + x, + x_q, + x_s, + masked_m, + group_size, + x.stride(0), + x.stride(1), + x_q.stride(0), + x_q.stride(1), + x_s.stride(0), + x_s.stride(1), + eps, + -fp8_max, + fp8_max, + num_tiles_k, + BLOCK_SIZE, + ) + + return x_q, x_s.transpose(1, 2), masked_m, m, aligned_m + + +def scaled_fp8_quant( + input: torch.Tensor, + scale: Optional[torch.Tensor] = None, + num_token_padding: Optional[int] = None, + use_per_token_if_dynamic: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP8 (8-bit floating point) format. + + Args: + input (torch.Tensor): Input tensor to be quantized + scale (Optional[torch.Tensor]): Pre-computed scaling factor for static quantization. + If None, scales will be computed dynamically. + num_token_padding (Optional[int]): If specified, pad the first dimension + of the output to at least this value. + use_per_token_if_dynamic (bool): When using dynamic scaling (scale=None), + determines the quantization granularity: + - True: compute scale per token + - False: compute single scale per tensor + + Returns: + Tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - quantized_tensor: The FP8 quantized version of input + - scale_tensor: The scaling factors used for quantization + + Raises: + AssertionError: If input is not 2D or if static scale's numel != 1 + """ + assert input.ndim == 2, f"Expected 2D input tensor, got {input.ndim}D" + shape = input.shape + out_dtype = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + if num_token_padding: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + output = torch.empty(shape, device=input.device, dtype=out_dtype) + + if scale is None: + # Dynamic scaling + if use_per_token_if_dynamic: + scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32) + sgl_per_token_quant_fp8(input, output, scale) + else: + scale = torch.zeros(1, device=input.device, dtype=torch.float32) + sgl_per_tensor_quant_fp8( + input, output, scale, is_static=False + ) # False for dynamic + else: + # Static scaling + assert scale.numel() == 1, f"Expected scalar scale, got numel={scale.numel()}" + sgl_per_tensor_quant_fp8( + input, output, scale, is_static=True + ) # True for static + + return output, scale diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_utils.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c137d3ad14eafb5d881682b85262540d64596c9a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/fp8_utils.py @@ -0,0 +1,514 @@ +import os +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 + +try: + from vllm import _custom_ops as ops + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + +from sglang.srt.layers.quantization.deep_gemm import _ENABLE_JIT_DEEPGEMM +from sglang.srt.layers.quantization.fp8_kernel import ( + per_token_group_quant_fp8, + scaled_fp8_quant, + sglang_per_token_quant_fp8, + static_quant_fp8, + w8a8_block_fp8_matmul, +) +from sglang.srt.utils import ( + get_bool_env_var, + get_cuda_version, + get_device_capability, + is_cuda, + is_hip, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() + +if _is_hip and get_bool_env_var("SGLANG_AITER_MOE"): + from aiter import gemm_a8w8_blockscale + +if _is_cuda: + from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm + +use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL") + +# Input scaling factors are no longer optional in _scaled_mm starting +# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale +TORCH_DEVICE_IDENTITY = None + +_TORCH_VERSION = torch.__version__.split("+")[0] +try: + _TORCH_VERSION_TUPLE = tuple(map(int, _TORCH_VERSION.split(".")[:3])) +except ValueError: + _TORCH_VERSION_TUPLE = (0, 0, 0) + +# The condition to determine if it is on a platform that supports +# torch._scaled_mm rowwise feature. +# The condition is determined once as the operations +# are time consuming. +USE_ROWWISE_TORCH_SCALED_MM = ( + _is_hip and get_device_capability() >= (9, 4) and _TORCH_VERSION_TUPLE >= (2, 7, 0) +) + + +def cutlass_fp8_supported(): + if not _is_cuda: + return False + major, minor = get_device_capability() + cuda_version = get_cuda_version() + if major >= 9: + return cuda_version >= (12, 0) + elif major == 8 and minor == 9: + return cuda_version >= (12, 4) + return False + + +def normalize_e4m3fn_to_e4m3fnuz( + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + assert weight.dtype == torch.float8_e4m3fn + # The bits pattern 10000000(-128) represents zero in e4m3fn + # but NaN in e4m3fnuz. So here we set it to 0. + # https://onnx.ai/onnx/technical/float8.html + weight_as_int8 = weight.view(torch.int8) + ROCM_FP8_NAN_AS_INT = -128 + weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0 + weight = weight_as_int8.view(torch.float8_e4m3fnuz) + + # For the same bits representation, e4m3fnuz value is half of + # the e4m3fn value, so we should double the scaling factor to + # get the same dequantized value. + # https://onnx.ai/onnx/technical/float8.html + weight_scale = weight_scale * 2.0 + if input_scale is not None: + input_scale = input_scale * 2.0 + return weight, weight_scale, input_scale + + +def cutlass_block_fp8_supported() -> bool: + if not get_bool_env_var("SUPPORT_CUTLASS_BLOCK_FP8"): + return False + if _is_cuda: + major, minor = torch.cuda.get_device_capability() + sm_version = major * 10 + minor + cuda_version = tuple(map(int, torch.version.cuda.split("."))) + if cuda_version >= (12, 0) and sm_version >= 90: + return True + return False + + +CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported() + + +def apply_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + # TODO: add more robust shape check here + shape_supported_by_cutlass = ( + weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0 + ) + if CUTLASS_BLOCK_FP8_SUPPORTED and shape_supported_by_cutlass: + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=True + ) + output = fp8_blockwise_scaled_mm( + q_input, weight.T, x_scale, weight_scale.T, out_dtype=input.dtype + ) + elif _is_hip and get_bool_env_var("SGLANG_AITER_MOE"): + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=False + ) + output = torch.zeros( + [q_input.shape[0], weight.shape[0]], + dtype=input.dtype, + device=q_input.device, + ) + gemm_a8w8_blockscale(q_input, weight, x_scale, weight_scale, output) + else: + if _ENABLE_JIT_DEEPGEMM: + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, + block_size[1], + column_major_scales=True, + scale_tma_aligned=True, + ) + else: + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=False + ) + output = w8a8_block_fp8_matmul( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype + ) + + if bias is not None: + output = output + bias + return output.to(dtype=input.dtype).view(*output_shape) + + +def input_to_float8( + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to float8 values with tensor-wise quantization.""" + finfo = torch.finfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).float().clamp(min=1e-12) + fp8_max = finfo.max + if _is_hip: + dtype = torch.float8_e4m3fnuz + fp8_max = 224.0 + scale = fp8_max / amax + x_scl_sat = (x.float() * scale).clamp(min=-fp8_max, max=fp8_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_quant_to_tensor_quant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function converts block-wise quantization to tensor-wise quantization. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The outputs are tensor-wise quantization tensor and tensor-wise quantization scale. + Note only float8 is supported for now. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + x_dq_block_tiles = [ + [ + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] + for i in range(k_tiles) + ] + for j in range(n_tiles) + ] + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block_tiles[j][i][:, :] = x_dq_block_tiles[j][i] * x_s[j][i] + + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_block) + if _is_cuda + else input_to_float8(x_dq_block, dtype=x_q_block.dtype) + ) + return x_q_tensor, scale + + +def channel_quant_to_tensor_quant( + x_q_channel: torch.Tensor, + x_s: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + x_dq_channel = x_q_channel.to(torch.float32) * x_s + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_channel) + if _is_cuda + else input_to_float8(x_dq_channel, dtype=x_q_channel.dtype) + ) + return x_q_tensor, scale + + +def _process_scaled_mm_output(output, input_2d_shape, output_shape): + if type(output) is tuple and len(output) == 2: + output = output[0] + return torch.narrow(output, 0, 0, input_2d_shape[0]).view(*output_shape) + + +def _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d_shape, + output_shape, + bias, + input_dtype, +): + global TORCH_DEVICE_IDENTITY + if TORCH_DEVICE_IDENTITY is None: + TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32, device=weight.device) + + output = torch._scaled_mm( + qinput, + weight, + scale_a=TORCH_DEVICE_IDENTITY, + scale_b=TORCH_DEVICE_IDENTITY, + out_dtype=torch.float32, + ) + + output = _process_scaled_mm_output(output, input_2d_shape, output_shape) + x_scale = torch.narrow(x_scale, 0, 0, input_2d_shape[0]) + + output = output * x_scale * weight_scale.t() + if bias is not None: + output = output + bias + return output.to(dtype=input_dtype) + + +def apply_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + input_scale_ub: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + cutlass_fp8_supported: bool = cutlass_fp8_supported(), + use_per_token_if_dynamic: bool = False, + pad_output: Optional[bool] = None, + compressed_tensor_quant: bool = False, +) -> torch.Tensor: + # Note: we pad the input because torch._scaled_mm is more performant + # for matrices with batch dimension > 16. + # This could change in the future. + # We also don't pad when using torch.compile, + # as it breaks with dynamic shapes. + if pad_output is None: + pad_output = not get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE") + output_padding = 17 if pad_output else None + + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[1]] + + if compressed_tensor_quant: + # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A + # for sgl-kernel fp8_scaled_mm, it support per channel W now + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: + qinput, x_scale = scaled_fp8_quant( + input_2d, + input_scale, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + + # Fused GEMM_DQ + if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel: + # Fall back to vllm cutlass w8a8 fp8 kernel + output = ops.cutlass_scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + else: + assert ( + weight_scale.numel() == weight.shape[1] + ), "cutlass w8a8 fp8 sgl-kernel only supports per-channel scale" + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) + + # torch.scaled_mm supports per tensor weights + activations only + # so fallback to naive if per channel or per token + else: + # Maybe apply padding to output, see comment in __init__ + qinput, x_scale = ( + scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=output_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + if _is_cuda + else ops.scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=output_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + ) + + per_tensor_weights = weight_scale.numel() == 1 + per_tensor_activations = x_scale.numel() == 1 + + if per_tensor_weights and per_tensor_activations: + # Fused GEMM_DQ + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + elif ( + use_per_token_if_dynamic + and not per_tensor_weights + and not per_tensor_activations + and USE_ROWWISE_TORCH_SCALED_MM + ): + # For now validated on ROCm platform + # fp8 rowwise scaling in torch._scaled_mm is introduced in + # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt + # and ROCm 6.3, which only exists in torch 2.7 and above. + # For CUDA platform please validate if the + # torch._scaled_mm support rowwise scaled GEMM + # Fused GEMM_DQ Rowwise GEMM + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale.t(), + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + else: + # Fallback for channelwise case, where we use unfused DQ + # due to limitations with scaled_mm + + # Symmetric quantized GEMM by definition computes the following: + # C = (s_x * X) (s_w * W) + bias + # This is equivalent to dequantizing the weights and activations + # before applying a GEMM. + # + # In order to compute quantized operands, a quantized kernel + # will rewrite the above like so: + # C = s_w * s_x * (X * W) + bias + # + # For the scaled_mm fallback case, we break this down, since it + # does not support s_w being a vector. + return _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d.shape, + output_shape, + bias, + input.dtype, + ) + else: + # cutlass w8a8 fp8 sgl-kernel only supports per-token scale + if input_scale is not None: + assert input_scale.numel() == 1 + # broadcast per-tensor scale to per-token scale when supporting cutlass + qinput, x_scale = static_quant_fp8( + input_2d, input_scale, repeat_scale=cutlass_fp8_supported + ) + else: + # default use per-token quantization if dynamic + if _is_cuda: + qinput, x_scale = sglang_per_token_quant_fp8(input_2d) + else: + # TODO(kkhuang): temporarily enforce per-tensor activation scaling if weight is per-tensor scaling + # final solution should be: 1. add support to per-tensor activation scaling. + # 2. solve the torch.compile error from weight_scale.numel() == 1 and x_scale.numel() > 1 (below line#308) + if _is_hip and weight_scale.numel() == 1: + qinput, x_scale = ops.scaled_fp8_quant( + input_2d, + input_scale, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + else: + qinput, x_scale = per_token_group_quant_fp8( + input_2d, group_size=input_2d.shape[1] + ) + + if cutlass_fp8_supported: + try: + if VLLM_AVAILABLE and use_vllm_cutlass_w8a8_fp8_kernel: + # Fall back to vllm cutlass w8a8 fp8 kernel + output = ops.cutlass_scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + else: + assert ( + weight_scale.numel() == weight.shape[1] + ), "cutlass w8a8 fp8 sgl-kernel only supports per-channel scale" + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) + except (ImportError, NameError, AttributeError): + pass + + # torch.scaled_mm supports per tensor weights + activations only + # so fallback to naive if per channel or per token + per_tensor_weights = weight_scale.numel() == 1 + per_tensor_activations = x_scale.numel() == 1 + + if per_tensor_weights and per_tensor_activations: + # Fused GEMM_DQ + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + else: + # Fallback for channelwise case, where we use unfused DQ + # due to limitations with scaled_mm + + # Symmetric quantized GEMM by definition computes the following: + # C = (s_x * X) (s_w * W) + bias + # This is equivalent to dequantizing the weights and activations + # before applying a GEMM. + # + # In order to compute quantized operands, a quantized kernel + # will rewrite the above like so: + # C = s_w * s_x * (X * W) + bias + # + # For the scaled_mm fallback case, we break this down, since it + # does not support s_w being a vector. + return _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d.shape, + output_shape, + bias, + input.dtype, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/gptq.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/gptq.py new file mode 100644 index 0000000000000000000000000000000000000000..af917a26df300d72dbe5c5011d09512609b9be52 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/gptq.py @@ -0,0 +1,439 @@ +import logging +from fractions import Fraction +from typing import Any, Dict, List, Optional, Union + +import torch + +from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() + +try: + from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase + from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod + from vllm.model_executor.layers.quantization.gptq_marlin import ( + GPTQMarlinLinearMethod, + GPTQMarlinMoEMethod, + ) + from vllm.model_executor.layers.quantization.marlin import MarlinLinearMethod + from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_marlin_supported, + ) + from vllm.scalar_type import scalar_types + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + + GPTQLinearMethod = MarlinLinearMethod = QuantizeMethodBase = Any + + class scalar_types: + uint4b8 = "uint4b8" + uint8b128 = "uint8b128" + + +logger = logging.getLogger(__name__) + + +def check_marlin_format(hf_quant_cfg: Dict[str, Any]) -> bool: + # compat: gptqmodel and autogptq (eol) main use checkpoint_format: str + # compat: autogptq <=0.7.1 is_marlin_format: bool + return hf_quant_cfg.get("checkpoint_format") == "marlin" or hf_quant_cfg.get( + "is_marlin_format", False + ) + + +class GPTQConfig(QuantizationConfig): + """Config class for GPTQ. + + Reference: https://arxiv.org/abs/2210.17323 + """ + + def __init__( + self, + weight_bits: int, + group_size: int, + desc_act: bool, + lm_head_quantized: bool, + dynamic: Dict[str, Dict[str, Union[int, bool]]], + ) -> None: + # GPTQModel use `dynamic` config property to allow per module + # quantization config so each module can be individually optimized. + # Format is Dict[str, Dict] where key is a regex string that can + # perform both positive ("+:" prefixed) or negative ("-:" prefixed) + # matching of a module. + # Default to positive match, override base quant config mode, if no + # prefix is used. Value is in dict format of field key and override + # value. + # Negative matching will skip quantization init for this module + # entirely: + # non-quantized inference. More details and quantization examples can be + # found at: https://github.com/ModelCloud/GPTQModel + # Example: + # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9 + # # last 1/4 of the layers 16-21 has 8bit and group_size 64 + # dynamic = { + # #`.*\.` matches the layers_node prefix + # # positive match layer 10-15 + # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,}, + # # positive match layer 16-21 + # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,}, + # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers + # } + super().__init__() + self.dynamic = dynamic + + self.weight_bits = weight_bits + self.group_size = group_size + self.desc_act = desc_act + self.lm_head_quantized = lm_head_quantized + self.pack_factor = Fraction(32, self.weight_bits) + if self.weight_bits not in [2, 3, 4, 8]: + raise ValueError( + "Currently, only 2/3/4/8-bit weight quantization is " + f"supported for GPTQ, but got {self.weight_bits} bits." + ) + + def __repr__(self) -> str: + return ( + f"GPTQConfig(weight_bits={self.weight_bits}, " + f"group_size={self.group_size}, " + f"desc_act={self.desc_act})," + f"lm_head_quantized={self.lm_head_quantized}), " + f"dynamic={self.dynamic}" + ) + + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError + + @classmethod + def get_name(cls) -> str: + return "gptq" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half] + + @classmethod + # Need to figure it out + def get_min_capability(cls) -> int: + return 60 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "GPTQConfig": + dynamic = cls.get_from_keys_or(config, ["dynamic"], default={}) + dynamic = {} if dynamic is None else dynamic + + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + desc_act = cls.get_from_keys(config, ["desc_act"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls(weight_bits, group_size, desc_act, lm_head_quantized, dynamic) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[GPTQLinearMethod]: + # Delay the import to avoid circular dependency + from sglang.srt.layers.quantization import get_linear_quant_method + + return get_linear_quant_method(self, layer, prefix, GPTQLinearMethod) + + +class GPTQMarlinConfig(QuantizationConfig): + """Config class for GPTQ Marlin""" + + # (num_bits, is_sym) -> quant_type + TYPE_MAP = { + (4, True): scalar_types.uint4b8, + (8, True): scalar_types.uint8b128, + } + + def __init__( + self, + weight_bits: int, + group_size: int, + desc_act: bool, + is_sym: bool, + lm_head_quantized: bool, + dynamic: Dict[str, Dict[str, Union[int, bool]]], + full_config: Dict[str, Any], + ) -> None: + super().__init__() + if desc_act and group_size == -1: + # In this case, act_order == True is the same as act_order == False + # (since we have only one group per output channel) + desc_act = False + + # GPTQModel use `dynamic` config property to allow per module + # quantization config so each module can be individually optimized. + # Format is Dict[str, Dict] where key is a regex string that can + # perform both positive ("+:" prefixed) or negative ("-:" prefixed) + # matching of a module. + # Default to positive match, override base quant config mode, if no + # prefix is used. Value is in dict format of field key and override + # value. + # Negative matching will skip quantization init for this module + # entirely: + # non-quantized inference. More details and quantization examples can be + # found at: https://github.com/ModelCloud/GPTQModel + # Example: + # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9 + # # last 1/4 of the layers 16-21 has 8bit and group_size 64 + # dynamic = { + # #`.*\.` matches the layers_node prefix + # # positive match layer 10-15 + # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,}, + # # positive match layer 16-21 + # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,}, + # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers + # } + self.dynamic = dynamic + + self.weight_bits = weight_bits + self.is_sym = is_sym + + self.pack_factor = 32 // weight_bits # packed into int32 + self.group_size = group_size + self.desc_act = desc_act + self.lm_head_quantized = lm_head_quantized + self.full_config = full_config + + if (weight_bits, is_sym) not in self.TYPE_MAP: + raise ValueError( + "Unsupported quantization config: " f"bits={weight_bits}, sym={is_sym}" + ) + + # (num_bits, is_sym) -> quant_type + self.quant_type = self.TYPE_MAP[(weight_bits, is_sym)] + + def __repr__(self) -> str: + return ( + f"GPTQMarlinConfig(quant_type={self.quant_type}, " + f"group_size={self.group_size}, " + f"desc_act={self.desc_act}, " + f"lm_head_quantized={self.lm_head_quantized}), " + f"dynamic={self.dynamic}" + ) + + def get_scaled_act_names(self) -> List[str]: + """Returns the activation function names that should be post-scaled. + + For now, this is only used by AWQ. + """ + raise NotImplementedError + + @classmethod + def get_name(cls) -> str: + return "gptq_marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "GPTQMarlinConfig": + dynamic = cls.get_from_keys_or(config, ["dynamic"], default={}) + dynamic = {} if dynamic is None else dynamic + + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + desc_act = cls.get_from_keys(config, ["desc_act"]) + is_sym = cls.get_from_keys(config, ["sym"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls( + weight_bits, + group_size, + desc_act, + is_sym, + lm_head_quantized, + dynamic, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + is_marlin_format = check_marlin_format(hf_quant_cfg) + + can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg) + + is_valid_user_quant = ( + user_quant is None or user_quant == "marlin" or user_quant == "gptq_marlin" + ) + + if not is_marlin_format and can_convert and is_valid_user_quant: + msg = ( + "The model is convertible to {} during runtime." + " Using {} kernel.".format(cls.get_name(), cls.get_name()) + ) + logger.info(msg) + return cls.get_name() + + if not is_marlin_format and can_convert and user_quant == "gptq": + logger.info( + "Detected that the model can run with gptq_marlin" + ", however you specified quantization=gptq explicitly," + " so forcing gptq. Use quantization=gptq_marlin for" + " faster inference" + ) + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + # Delay the import to avoid circular dependency + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.quantization import get_linear_quant_method + + if isinstance(layer, FusedMoE): + return GPTQMarlinMoEMethod(self) + # TODO: re-enable after SGLang syncs with vllm >= 0.7.3 + # if layer.num_experts > 32: + # # For MoEs with many experts the moe_wna16 kernel is faster + # return MoeWNA16Config.from_config(self.full_config).get_quant_method( + # layer, prefix + # ) + # else: + # return GPTQMarlinMoEMethod(self) + return get_linear_quant_method(self, layer, prefix, GPTQMarlinLinearMethod) + + @classmethod + def is_gptq_marlin_compatible(cls, quant_config: Dict[str, Any]): + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + group_size = quant_config.get("group_size") + sym = quant_config.get("sym") + desc_act = quant_config.get("desc_act") + + if not _is_cuda: + return False + + if quant_method != "gptq": + return False + + # Marlin conversion is only valid if required properties are found + if num_bits is None or group_size is None or sym is None or desc_act is None: + return False + + if (num_bits, sym) not in cls.TYPE_MAP: + return False + + return check_marlin_supported( + quant_type=cls.TYPE_MAP[(num_bits, sym)], group_size=group_size + ) + + +class MarlinConfig(QuantizationConfig): + """Config class for Marlin. + + Reference: https://github.com/IST-DASLab/marlin/tree/master + """ + + def __init__( + self, + group_size: int, + lm_head_quantized: bool, + ) -> None: + # Group size for the quantization. + self.group_size = group_size + self.lm_head_quantized = lm_head_quantized + if self.group_size != 128 and self.group_size != -1: + raise ValueError( + "Currently, only group size 128 and -1 (channelwise) " + "is supported for Marlin, but got group_size of " + f"{self.group_size}" + ) + + # 4 Bits packed into 32 bit datatype. + self.pack_factor = 32 // 4 + + # Tile size used by marlin kernels. + self.tile_size = 16 + + # Min out_features dim + self.min_n_threads = 64 + + # Min in_features dim + self.min_k_threads = 128 + + # Max parallel problems to solve at once (improves large + # batch performance) + self.max_parallel = 16 + + # Permutation length used by the marlin kernels. + self.perm_len = 1024 + + def __repr__(self) -> str: + return ( + f"MarlinConfig(group_size={self.group_size}, " + f"lm_head_quantized={self.lm_head_quantized})" + ) + + @classmethod + def get_name(cls) -> str: + return "marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half] + + @classmethod + # Need to figure it out + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "MarlinConfig": + group_size = cls.get_from_keys(config, ["group_size"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + return cls(group_size, lm_head_quantized) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + is_marlin_format = check_marlin_format(hf_quant_cfg) + + is_valid_user_quant = ( + user_quant is None or user_quant == "gptq" or user_quant == "marlin" + ) + + if is_marlin_format and is_valid_user_quant: + msg = "The model is serialized in {} format. Using {} kernel.".format( + cls.get_name(), cls.get_name() + ) + logger.info(msg) + return cls.get_name() + + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[MarlinLinearMethod]: + # Delay the import to avoid circular dependency + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if isinstance(layer, LinearBase) or ( + isinstance(layer, ParallelLMHead) and self.lm_head_quantized + ): + return MarlinLinearMethod(self) + return None diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_kernel.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..22f1c5069378e79075c71e93276abdd1033e25a0 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_kernel.py @@ -0,0 +1,412 @@ +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import get_device_name, is_cuda + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import sgl_per_token_group_quant_int8 + +logger = logging.getLogger(__name__) + + +@triton.jit +def _per_token_quant_int8( + x_ptr, + xq_ptr, + scale_ptr, + stride_x, + stride_xq, + N, + BLOCK: tl.constexpr, +): + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 + row_id = tl.program_id(0) + + cols = tl.arange(0, BLOCK) + mask = cols < N + + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) + scale_x = absmax / 127 + x_q = x * (127 / absmax) + x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8) + + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) + tl.store(scale_ptr + row_id, scale_x) + + +def per_token_quant_int8(x): + M = x.numel() // x.shape[-1] + N = x.shape[-1] + x_q = torch.empty_like(x, device=x.device, dtype=torch.int8) + scales = torch.empty(x.shape[:-1] + (1,), device=x.device, dtype=torch.float32) + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + + assert x.is_contiguous() + _per_token_quant_int8[(M,)]( + x, + x_q, + scales, + stride_x=x.stride(-2), + stride_xq=x_q.stride(-2), + N=N, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + + return x_q, scales + + +@triton.jit +def _per_token_group_quant_int8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + y_stride, + # Collums of input + N, + # Avoid to divide zero + eps, + # Information for int8 + int8_min, + int8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group quantization on a + tensor. + + This function converts the tensor values into int8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / int8_max + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed int8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + dtype: The dype of output tensor. Note that only `torch.int8` is supported for now. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. + """ + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + M = x.numel() // group_size + N = group_size + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_int8[(M,)]( + x, + x_q, + x_s, + group_size, + N, + eps, + int8_min=int8_min, + int8_max=int8_max, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + return x_q, x_s + + +def sglang_per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +): + assert ( + x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + + sgl_per_token_group_quant_int8(x, x_q, x_s, group_size, eps, int8_min, int8_max) + + return x_q, x_s + + +@triton.jit +def _w8a8_block_int8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and store the result in output + tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@functools.lru_cache +def get_w8a8_block_int8_configs( + N: int, K: int, block_n: int, block_k: int +) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the w8a8 block fp8 kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + + # First look up if an optimized configuration is available in the configs + # directory + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info( + "Using configuration from %s for W8A8 Block INT8 kernel.", + config_file_path, + ) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + logger.warning( + ( + "Using default W8A8 Block INT8 kernel config. Performance might be sub-optimal! " + "Config file not found at %s" + ), + config_file_path, + ) + return None + + +def w8a8_block_int8_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1]) + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Default config + # Block-wise quant: BLOCK_SIZE_K must be divisable by block_size[1] + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_size[0], + "BLOCK_SIZE_K": block_size[1], + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + + def grid(META): + return ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + _w8a8_block_int8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_utils.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a7aa6d825f2a2dfe4bc71cf923f288bf0f7f6f8 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/int8_utils.py @@ -0,0 +1,73 @@ +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.layers.quantization.int8_kernel import ( + per_token_group_quant_int8, + w8a8_block_int8_matmul, +) + + +def apply_w8a8_block_int8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) + output = w8a8_block_int8_matmul( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype + ) + + if bias is not None: + output = output + bias + return output.to(dtype=input.dtype).view(*output_shape) + + +def input_to_int8( + x: torch.Tensor, dtype: torch.dtype = torch.int8 +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to int8 values with tensor-wise quantization.""" + iinfo = torch.iinfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + int8_min, int8_max = iinfo.min, iinfo.max + scale = int8_max / amax + x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> torch.Tensor: + """This function conducts block-wise dequantization. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The outputs are dequantized tensor. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] *= x_s[j][i] + + return x_dq_block diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/kv_cache.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/kv_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..da6d91a9b040dfe8037f1f07bca632cec826f975 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/kv_cache.py @@ -0,0 +1,89 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/kv_cache.py + +import logging + +import torch + +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.utils import is_hip + +_is_hip = is_hip() + +logger = logging.getLogger(__name__) + + +class BaseKVCacheMethod(QuantizeMethodBase): + """ + Quant method that adds `k_scale` and `v_scale` attributes to the + Attention layer to support loading those scaling factors from checkpoints. + The k/v_scale will be used to: + - quantize k/v_cache entries before saving them to the cache + - dequantize k/v_cache entries before fetching them from the cache + + :param quant_config: the appropriate QuantizationConfig + """ + + def __init__(self, quant_config: QuantizationConfig): + self.quant_config = quant_config + + def create_weights(self, layer: torch.nn.Module): + """ + Create "weight" (aka k_scale and v_scale) for an attention layer. + """ + # Initialize the KV cache scales to -1.0, which is an invalid value. + # If the k/v_scale appears in the checkpoint, it will be + # overwritten when loading weights. + layer.k_scale = torch.nn.Parameter( + torch.tensor(-1.0, dtype=torch.float32), requires_grad=False + ) + layer.v_scale = torch.nn.Parameter( + torch.tensor(-1.0, dtype=torch.float32), requires_grad=False + ) + + @classmethod + def is_fp8_fnuz(cls) -> bool: + # only device 0 is checked, this assumes MI300 platforms are homogeneous + return "gfx94" in torch.cuda.get_device_properties(0).gcnArchName + + def apply(self, layer: torch.nn.Module) -> torch.Tensor: + raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.") + + def process_weights_after_loading(self, layer: RadixAttention) -> None: + if layer.k_scale > 0.0 and layer.v_scale > 0.0: + # We prefer to use separate k_scale and v_scale if present + k_scale = layer.k_scale.to("cpu").tolist() + v_scale = layer.v_scale.to("cpu").tolist() + if _is_hip and self.is_fp8_fnuz(): + k_scale *= 2 + v_scale *= 2 + elif layer.k_scale < 0.0 and layer.v_scale < 0.0: + # If no scales were loaded (both scales are invalid negative + # values), use the default value of 1.0 + k_scale = 1.0 + v_scale = 1.0 + else: + # If we find a single kv_scale in the checkpoint, we remap + # kv_scale to k_scale during weight loading, and duplicate + # k_scale to v_scale here + assert layer.k_scale > 0.0 + scale_to_duplicate = max(layer.k_scale, layer.v_scale) + k_scale = scale_to_duplicate.to("cpu").tolist() + v_scale = scale_to_duplicate.to("cpu").tolist() + if _is_hip and self.is_fp8_fnuz(): + k_scale *= 2 + v_scale *= 2 + + if not isinstance(k_scale, float) or not isinstance(v_scale, float): + raise ValueError( + "Only support per-tensor scaling factor " "for fp8 KV cache" + ) + + # These are used in the final Attention.forward() + layer.k_scale.copy_(k_scale) + layer.v_scale.copy_(v_scale) + layer.k_scale_float = k_scale + layer.v_scale_float = v_scale diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/modelopt_quant.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/modelopt_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..aebc452443696a7dc4477fb32aded1f6dab65e82 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/modelopt_quant.py @@ -0,0 +1,463 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py + +import logging +from typing import Any, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.linear import LinearBase, LinearMethodBase +from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, +) +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.utils import ( + convert_to_channelwise, + requantize_with_max_scale, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.utils import is_cuda + +if is_cuda(): + from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant + +# Initialize logger for the module +logger = logging.getLogger(__name__) + +# Supported activation schemes for the current configuration +ACTIVATION_SCHEMES = ["static"] + + +class ModelOptFp8Config(QuantizationConfig): + """Configuration for ModelOpt FP8 quantization, including serialization and compatibility checks.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + kv_cache_quant_method: Optional[str] = None, + exclude_modules: Optional[List[str]] = None, + ) -> None: + """ + Args: + is_checkpoint_fp8_serialized (bool): Indicates if the checkpoint uses serialized FP8 format. + """ + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + self.kv_cache_quant_method = kv_cache_quant_method + self.exclude_modules = exclude_modules + if is_checkpoint_fp8_serialized: + logger.warning( + "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change." + ) + + @classmethod + def get_name(cls) -> str: + return "modelopt" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 89 # Minimum hardware capability (e.g., Hopper GPUs). + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config": + quant_method = cls.get_from_keys(config, ["quantization"]).get("quant_algo") + kv_cache_quant_method = cls.get_from_keys(config, ["quantization"]).get( + "kv_cache_quant_algo" + ) + exclude_modules = cls.get_from_keys(config, ["quantization"]).get( + "exclude_modules" + ) + + if "FP8" not in quant_method: + raise ValueError( + "ModelOpt only supports static FP8 quantization in SGLang. " + "Check the `hf_quant_config.json` file for your model's configuration." + ) + + return cls( + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=kv_cache_quant_method, + exclude_modules=exclude_modules, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + if self.exclude_modules and any( + module in prefix for module in self.exclude_modules + ): + return None + + if isinstance(layer, LinearBase): + return ModelOptFp8LinearMethod(self) + if self.kv_cache_quant_method and isinstance(layer, RadixAttention): + return ModelOptFp8KVCacheMethod(self) + + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class ModelOptFp8LinearMethod(LinearMethodBase): + """Linear method for ModelOpt static FP8 quantization. + + Supports loading FP8 checkpoints with static weight and activation scales. + Future support may include dynamic scales. + + **Limitations**: + 1. Only supports per-tensor quantization due to `torch._scaled_mm` limitations. + 2. Only supports the `float8_e4m3fn` data type. + + Args: + quant_config (ModelOptFp8Config): The ModelOpt quantization configuration. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + super().__init__() + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + """Creates and registers weights, weight scales, and input scales for FP8 quantization.""" + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + # Set layer attributes + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + + # Register weight + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + if self.quant_config.is_checkpoint_fp8_serialized: + # Register weight and input scales + for scale_name in ["weight_scale", "input_scale"]: + layer.register_parameter( + scale_name, + PerTensorScaleParameter( + data=torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Requantizes weights after loading using the maximum scale.""" + max_w_scale, quantized_weight = requantize_with_max_scale( + layer.weight, layer.weight_scale, layer.logical_widths + ) + layer.weight = Parameter(quantized_weight.t(), requires_grad=False) + # cutlass sgl-kernel only supports per-channel scale + if self.cutlass_fp8_supported: + max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Applies FP8 linear transformation.""" + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) + + +class ModelOptFp8KVCacheMethod(BaseKVCacheMethod): + """ + Handles loading FP8 kv-cache scaling factors from modelopt quantized checkpoints. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + super().__init__(quant_config) + + +class ModelOptFp4Config(QuantizationConfig): + """Config class for FP4.""" + + def __init__( + self, + is_checkpoint_nvfp4_serialized: bool = False, + kv_cache_quant_algo: str = None, + group_size: int = None, + exclude_modules: List[str] = None, + ) -> None: + self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + if is_checkpoint_nvfp4_serialized: + logger.warning( + "Detected nvfp4 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + self.group_size = group_size + self.kv_cache_quant_algo = kv_cache_quant_algo + self.exclude_modules = exclude_modules + + @classmethod + def get_name(cls) -> str: + return "modelopt_fp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half, torch.float8_e4m3fn] + + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp4Config": + quant_config = cls.get_from_keys(config, ["quantization"]) + quant_method = quant_config["quant_algo"] + if not quant_method in ["FP8", "NVFP4"]: + raise ValueError( + f"ModelOpt currently only supports: FP8, NVFP4" + " quantizations in sglang. Please check the " + "`hf_quant_config.json` file for your model's " + "quant configuration." + ) + is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method + kv_cache_quant_algo = quant_config["kv_cache_quant_algo"] + group_size = quant_config["group_size"] + exclude_modules = quant_config["exclude_modules"] + if not (group_size and kv_cache_quant_algo and exclude_modules): + raise ValueError( + "NVFP4 quantization requires group size and " + "kv_cache_quant_algo specified in " + "hf_quant_config.json" + ) + return cls( + is_checkpoint_nvfp4_serialized, + kv_cache_quant_algo, + group_size, + exclude_modules, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + if self.exclude_modules and any( + module in prefix for module in self.exclude_modules + ): + return None + + if isinstance(layer, LinearBase): + return ModelOptFp4LinearMethod(self) + if self.kv_cache_quant_algo and isinstance(layer, RadixAttention): + return ModelOptFp8KVCacheMethod(self) + + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class ModelOptFp4LinearMethod(LinearMethodBase): + """Linear method for NVFP4. + Supports loading NVFP4 checkpoints with the following structure: + + |Tensor Name | datatype | shape | + |----------------------------------------------------| + |input_scale | torch.float32 | scalar | + |weight | NVFP4(SE2M1) | [1, X, y/2] | + |weight_scale | FP8-E4M3 | [X, Y] | + |weight_scale_2 | torch.float32 | scalar | + + The weights are quantized per block of 16 elements. + Args: quant_config: The ModelOpt quantization config. + """ + + def __init__(self, quant_config: ModelOptFp4Config): + 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, + ): + del input_size, output_size + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "NVFP4 quantization was selected, " + " dynamic quantization is not supported." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + if input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when in features size is " "not multiple of 16" + ) + + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_nvfp4_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + # 2 fp4 data is packed in one uint8 in the input dimension + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + layer.register_parameter("input_scale", input_scale) + + weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale_2", weight_scale_2) + + weight_scale = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.group_size, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + input_scale_2 = layer.input_scale.max().to(torch.float32) + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) + layer.input_scale = Parameter(input_scale_2, requires_grad=False) + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) + layer.alpha = Parameter( + layer.input_scale * layer.weight_scale_2, requires_grad=False + ) + + # Pad and blockwise interleave weight_scale + scales = layer.weight_scale + scale_ndim = scales.ndim + if scale_ndim == 2: + scales = scales.unsqueeze(0) + assert scales.ndim == 3 + B, M, K = scales.shape + round_up_multiple = lambda x, m: (x + m - 1) // m * m + M_padded = round_up_multiple(M, 128) + K_padded = round_up_multiple(K, 4) + padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype) + padded_scales[:B, :M, :K] = scales + batches, rows, cols = padded_scales.shape + assert rows % 128 == 0 + assert cols % 4 == 0 + padded_scales = padded_scales.reshape(batches, rows // 128, 4, 32, cols // 4, 4) + padded_scales = padded_scales.permute((0, 1, 4, 3, 2, 5)) + padded_scales = padded_scales.contiguous().cuda() + padded_scales = ( + padded_scales.reshape(M, K) + if scale_ndim == 2 + else padded_scales.reshape(B, M, K) + ) + layer.weight_scale_interleaved = Parameter(padded_scales, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + output_dtype = x.dtype + x_m, _ = x.shape + w_n, _ = layer.weight.shape + output_shape = [x_m, w_n] + + # Quantize BF16 or FP16 to (FP4 and interleaved block scale) + x_fp4, x_scale_interleaved = scaled_fp4_quant(x, 1 / layer.input_scale) + + assert x_fp4.dtype == torch.uint8 + assert x_scale_interleaved.dtype == torch.float8_e4m3fn + assert layer.weight.dtype == torch.uint8 + assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn + assert layer.alpha.dtype == torch.float32 + + out = cutlass_scaled_fp4_mm( + x_fp4, + layer.weight, + x_scale_interleaved, + layer.weight_scale_interleaved, + layer.alpha, + output_dtype, + ) + if bias is not None: + out = out + bias + return out.view(*output_shape) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/moe_wna16.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/moe_wna16.py new file mode 100644 index 0000000000000000000000000000000000000000..122f310dc8a62b82a655257c3ff1d06d0c4de7f7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/moe_wna16.py @@ -0,0 +1,505 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/moe_wna16.py + +import logging +from typing import Any, Callable, Dict, List, Optional + +import torch + +from sglang.srt.distributed import get_tensor_model_parallel_rank +from sglang.srt.distributed.parallel_state import get_tp_group +from sglang.srt.layers.linear import LinearBase, UnquantizedLinearMethod +from sglang.srt.layers.quantization.awq import AWQConfig +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig +from sglang.srt.utils import get_device_capability, set_weight_attrs + +logger = logging.getLogger(__name__) + + +class MoeWNA16Config(QuantizationConfig): + """Config class for MOE WNA16 (W8A16/W4A16) quantization.""" + + def __init__( + self, + linear_quant_method: str, + weight_bits: int, + group_size: int, + has_zp: bool, + lm_head_quantized: bool, + modules_to_not_convert: Optional[List[str]], + full_config: Dict[str, Any], + ) -> None: + super().__init__() + self.weight_bits = weight_bits + self.group_size = group_size + self.has_zp = has_zp + self.bit8_pack_factor = 8 // self.weight_bits + self.lm_head_quantized = lm_head_quantized + self.linear_quant_method = linear_quant_method + self.full_config = full_config + self.use_marlin = False + # Avoid circular import + + if self.linear_quant_method == "gptq": + self.use_marlin = GPTQMarlinConfig.is_gptq_marlin_compatible(full_config) + elif self.linear_quant_method == "awq": + capability_tuple = get_device_capability() + device_capability = ( + -1 + if capability_tuple is None + else capability_tuple[0] * 10 + capability_tuple[1] + ) + awq_min_capability = AWQConfig.get_min_capability() + if device_capability < awq_min_capability: + raise ValueError( + "The quantization method moe_wna16 + awq is not supported " + "for the current GPU. " + f"Minimum capability: {awq_min_capability}. " + f"Current capability: {device_capability}." + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + + if modules_to_not_convert is None: + self.modules_to_not_convert = [] + else: + self.modules_to_not_convert = modules_to_not_convert + + @classmethod + def get_name(cls) -> str: + return "moe_wna16" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["quantize_config.json"] + + def get_scaled_act_names(self) -> List[str]: + raise NotImplementedError + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "MoeWNA16Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + if quant_method == "gptq": + has_zp = not cls.get_from_keys(config, ["sym"]) + modules_to_not_convert = [] + elif quant_method == "awq": + has_zp = cls.get_from_keys(config, ["zero_point"]) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + + return cls( + quant_method, + weight_bits, + group_size, + has_zp, + lm_head_quantized, + modules_to_not_convert, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + can_convert = cls.is_moe_wna16_compatible(hf_quant_cfg) + if can_convert and user_quant == "moe_wna16": + return cls.get_name() + return None + + @classmethod + def is_moe_wna16_compatible(cls, quant_config: Dict[str, Any]): + # Extract data from quant config. + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + desc_act = quant_config.get("desc_act") + + capability_tuple = get_device_capability() + device_capability = ( + -1 + if capability_tuple is None + else capability_tuple[0] * 10 + capability_tuple[1] + ) + # Avoid circular import + awq_min_capability = AWQConfig.get_min_capability() + + gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8] + awq_compatible = ( + quant_method == "awq" + and num_bits == 4 + and device_capability >= awq_min_capability + ) + + return gptq_compatible or awq_compatible + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + # avoid circular import + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + + if is_layer_skipped_quant(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + elif isinstance(layer, LinearBase): + + if self.linear_quant_method == "gptq": + if self.use_marlin: + return GPTQMarlinConfig.from_config( + self.full_config + ).get_quant_method(layer, prefix) + else: + return GPTQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + elif self.linear_quant_method == "awq": + return AWQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + else: + raise ValueError("moe_wna16 only support gptq and awq.") + elif isinstance(layer, FusedMoE): + return MoeWNA16Method(self) + return None + + +def is_layer_skipped_quant(prefix: str, modules_to_not_convert: List[str]): + return any(module_name in prefix for module_name in modules_to_not_convert) + + +class MoeWNA16Method: + """Linear method for MOE WNA16 (W8A16/W4A16) quantization. + + Args: + quant_config: The MOE WNA16 (W8A16/W4A16) quantization config. + """ + + def __new__(cls, *args, **kwargs): + # avoid circular import + from sglang.srt.layers.moe.fused_moe_triton import FusedMoEMethodBase + + if not hasattr(cls, "_initialized"): + original_init = cls.__init__ + new_cls = type( + cls.__name__, + (FusedMoEMethodBase,), + { + "__init__": original_init, + **{k: v for k, v in cls.__dict__.items() if k != "__dict__"}, + }, + ) + obj = super(new_cls, new_cls).__new__(new_cls) + obj.__init__(*args, **kwargs) + return obj + return super().__new__(cls) + + def __init__(self, quant_config: MoeWNA16Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + layer.quant_config = self.quant_config + bit8_pack_factor = self.quant_config.bit8_pack_factor + group_size = self.quant_config.group_size + group_size_div_factor = 1 + + # make intermediate_size and hidden_size diviable by group_size + # we reduce the group size to ensure that + # and we would repeat the loaded_weight later + while intermediate_size_per_partition % group_size or hidden_size % group_size: + group_size = group_size // 2 + group_size_div_factor *= 2 + assert group_size >= 32 + layer.group_size = group_size + layer.group_size_div_factor = group_size_div_factor + + strategy = FusedMoeWeightScaleSupported.GROUP.value + extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False}) + + assert "weight_loader" in extra_weight_attrs + weight_loader = extra_weight_attrs["weight_loader"] + wrapped_weight_loader = MoeWNA16Method.get_weight_loader(layer, weight_loader) + extra_weight_attrs["weight_loader"] = wrapped_weight_loader + + # Fused gate_up_proj (column parallel) + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qweight", w13_qweight) + set_weight_attrs(w13_qweight, extra_weight_attrs) + + # down_proj (row parallel) + w2_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // bit8_pack_factor, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + + w13_scales = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_scales", w13_scales) + set_weight_attrs(w13_scales, extra_weight_attrs) + + w2_scales = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // group_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_scales", w2_scales) + set_weight_attrs(w2_scales, extra_weight_attrs) + + if self.quant_config.has_zp: + w13_qzeros = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition // bit8_pack_factor, + hidden_size // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qzeros", w13_qzeros) + set_weight_attrs(w13_qzeros, extra_weight_attrs) + + w2_qzeros = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size // bit8_pack_factor, + intermediate_size_per_partition // group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, extra_weight_attrs) + + if self.quant_config.linear_quant_method == "gptq": + # some param are unused, but we need to init them in order to + # load weights + invalid_param_keys = ["w13_g_idx", "w2_g_idx"] + if not self.quant_config.has_zp: + invalid_param_keys += ["w13_qzeros", "w2_qzeros"] + for key in invalid_param_keys: + param = torch.nn.Parameter( + torch.empty((0,), dtype=torch.int32), requires_grad=False + ) + layer.register_parameter(key, param) + set_weight_attrs(param, extra_weight_attrs) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + use_grouped_topk: bool = False, + 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", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + # avoid circular import + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + assert activation == "silu", "Only SiLU activation is supported." + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + top_k=top_k, + use_grouped_topk=use_grouped_topk, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + weight_bits = self.quant_config.weight_bits + has_zp = self.quant_config.has_zp + + return fused_experts( + x, + layer.w13_qweight, + layer.w2_qweight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + apply_router_weight_on_input=apply_router_weight_on_input, + use_int4_w4a16=weight_bits == 4, + use_int8_w8a16=weight_bits == 8, + w1_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w1_zp=layer.w13_qzeros if has_zp else None, + w2_zp=layer.w2_qzeros if has_zp else None, + block_shape=[0, layer.group_size], + no_combine=no_combine, + ) + + @staticmethod + def get_weight_loader(layer, weight_loader): + + def convert_awq_tensor(tensor, tensor_type): + # convert awq qweight/qzeros to a standard format (assume int4) + # qweight: (k, n // pack_factor_bit32) -> (n, k // pack_factor_bit8) + # qzeros: (k // group_size, n // pack_factor_bit32) -> + # (n // pack_factor_bit8, k // group_size) + # pack_factor_bit32 = 32 // weight_bits + # pack_factor_bit8 = 8 // weight_bits + + # 0. suppose origin shape (a, b), dtype int32 + # 1. convert to uint8, shape (a, b) -> (a, 4 * b) + size0 = tensor.size(0) + tensor = tensor.view(torch.uint8) + + # 2. unpack to uint4 (only when weight_bits == 4) + # shape (a, 4 * b) -> (a, 4 * b, 2) + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) + tensor = (tensor[:, :, None] >> shifter) & 0xF + + # 3. change order, see + # https://github.com/casper-hansen/AutoAWQ/blob/v0.2.8/awq/utils/quant_utils.py + # shape -> (a, 4 * b * pack_factor_bit8) + reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7] + tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order] + tensor = tensor.view(size0, -1) + + # 4. transpose, shape -> (4 * b * pack_factor_bit8, a) + tensor = tensor.T.contiguous() + + # 5. repack (only when weight_bits == 4) + # qweight shape -> (4 * b * pack_factor_bit8, a // pack_factor_bit8) + # qzeros shape -> (4 * b, a) + + if tensor_type == "qweight": + tensor = tensor[:, 1::2] * 16 + tensor[:, ::2] + elif tensor_type == "qzeros": + tensor = tensor[1::2, :] * 16 + tensor[::2, :] + return tensor + + def convert_gptq_int4_qzeros(tensor): + tensor = tensor.view(torch.uint8) + shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device) + tensor = (tensor[:, :, None] >> shifter) & 0xF + tensor = tensor + 1 + tensor = tensor[:, :, 0] + tensor[:, :, 1] * 16 + return tensor + + def moe_wna16_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ): + if "g_idx" in weight_name: + return + if not layer.quant_config.has_zp and "qzeros" in weight_name: + return + + device = get_tp_group().device + tp_rank = get_tensor_model_parallel_rank() + loaded_weight = loaded_weight.to(device) + shard_size = layer.intermediate_size_per_partition + + # convert gptq and awq weight to a standard format + if layer.quant_config.linear_quant_method == "awq": + assert layer.quant_config.weight_bits == 4 + if "weight" in weight_name: + loaded_weight = convert_awq_tensor(loaded_weight, "qweight") + elif "zeros" in weight_name: + loaded_weight = convert_awq_tensor(loaded_weight, "qzeros") + else: + loaded_weight = loaded_weight.T + elif layer.quant_config.linear_quant_method == "gptq": + assert layer.quant_config.weight_bits in [4, 8] + if "weight" in weight_name: + loaded_weight = loaded_weight.T.contiguous().view(torch.uint8) + elif "zeros" in weight_name: + # add 1 to gptq qzeros to align with awq + loaded_weight = loaded_weight.view(torch.uint8) + if layer.quant_config.weight_bits == 4: + loaded_weight = convert_gptq_int4_qzeros(loaded_weight).T + else: + loaded_weight = loaded_weight.T + 1 + else: + loaded_weight = loaded_weight.T + + # repeat the qzeros/scales to fit new group size + if ( + layer.group_size_div_factor > 1 + and "qzeros" in weight_name + or "scales" in weight_name + ): + loaded_weight = loaded_weight.repeat_interleave( + layer.group_size_div_factor, 1 + ) + + if "w13_qzeros" in weight_name: + tensor = loaded_weight.view(layer.tp_size, -1, loaded_weight.size(1))[ + tp_rank + ] + if shard_id == "w1": + param.data[expert_id, : shard_size // 2] = tensor + else: + param.data[expert_id, shard_size // 2 :] = tensor + elif "w2_qzeros" in weight_name: + param.data[expert_id] = loaded_weight.view( + loaded_weight.size(0), layer.tp_size, -1 + )[:, tp_rank] + else: + weight_loader(param, loaded_weight, weight_name, shard_id, expert_id) + + return moe_wna16_weight_loader diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/utils.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fdea997e359b4aae00f0b63052d074b2aa7caad2 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/utils.py @@ -0,0 +1,147 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/quant_utils.py + +from types import MappingProxyType +from typing import List, Mapping, Tuple, Union + +import torch + +from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() + +if not _is_cuda: + from vllm._custom_ops import scaled_fp8_quant + + +def is_fp8_fnuz() -> bool: + # only device 0 is checked, this assumes MI300 platforms are homogeneous + return "gfx94" in torch.cuda.get_device_properties(0).gcnArchName + + +def is_layer_skipped( + prefix: str, + ignored_layers: List[str], + fused_mapping: Mapping[str, List[str]] = MappingProxyType({}), +) -> bool: + # prefix: model.layers.0.self_attn.q_proj + # proj_name: q_proj + proj_name = prefix.split(".")[-1] + + # Fused layers like gate_up_proj or qkv_proj will not be fused + # in the safetensors checkpoint. So, we convert the name + # from the fused version to unfused + check to make sure that + # each shard of the fused layer has the same scheme. + if proj_name in fused_mapping: + shard_prefixes = [ + prefix.replace(proj_name, shard_proj_name) + for shard_proj_name in fused_mapping[proj_name] + ] + + is_skipped = None + for shard_prefix in shard_prefixes: + is_shard_skipped = shard_prefix in ignored_layers + + if is_skipped is None: + is_skipped = is_shard_skipped + elif is_shard_skipped != is_skipped: + raise ValueError( + f"Detected some but not all shards of {prefix} " + "are quantized. All shards of fused layers " + "to have the same precision." + ) + else: + is_skipped = prefix in ignored_layers + + assert is_skipped is not None + return is_skipped + + +def per_tensor_dequantize( + tensor: torch.Tensor, inv_scale: Union[float, torch.Tensor] +) -> torch.Tensor: + fake_qweight = tensor.to(torch.float16) + dq_weight = fake_qweight * inv_scale + return dq_weight + + +def all_close_1d(x: torch.Tensor) -> bool: + assert len(x.shape) == 1 + return all(torch.allclose(x[0], x[i]) for i in range(x.shape[0])) + + +def convert_to_channelwise( + weight_scale: torch.Tensor, logical_widths: List[int] +) -> Tuple[torch.Tensor, torch.Tensor]: + # Create channelwise buffer + weight_scale_channel = torch.empty( + (sum(logical_widths), 1), dtype=torch.float32, device=weight_scale.device + ) + + # Handle scalar tensor case: broadcast same scale to all channels + if weight_scale.dim() == 0: + weight_scale_channel.fill_(weight_scale.item()) + return weight_scale_channel + + # Expand each scale to match the size of each logical matrix. + start = 0 + for idx, logical_width in enumerate(logical_widths): + end = start + logical_width + weight_scale_channel[start:end, :] = weight_scale[idx] + start = end + + return weight_scale_channel + + +def requantize_with_max_scale( + weight: torch.Tensor, weight_scale: torch.Tensor, logical_widths: List[int] +) -> Tuple[torch.Tensor, torch.Tensor]: + # Max scale to be used for requanitzation. + max_w_scale = weight_scale.max() + + # QKV / MLP is fused in the on disk checkpoint if any of the + # weight scales are still set to the default since we initialize + # N weight scales for N shards but we only load 1 weight scale + # from disk in this case. Skip requantization in this case (since) + # we already are quantized with the single scale. + # * Sample Model: nm-testing/Phi-3-mini-128k-instruct-FP8 + unfused_module_in_checkpoint = ( + weight_scale[-1] > torch.finfo(torch.float8_e4m3fn).min + ) + + # If unfused checkpoint, need requanize with the single scale. + if unfused_module_in_checkpoint: + start = 0 + for idx, logical_width in enumerate(logical_widths): + end = start + logical_width + weight_dq = per_tensor_dequantize(weight[start:end, :], weight_scale[idx]) + weight[start:end, :], _ = scaled_fp8_quant(weight_dq, max_w_scale) + start = end + + return max_w_scale, weight + + +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/layer_utils.py +# Newly generated tensors need to replace existing tensors that are +# already registered as parameters by vLLM (and won't be freed) +def replace_parameter( + mod: torch.nn.Module, name: str, new: Union[torch.Tensor, torch.nn.Parameter] +) -> None: + + old = getattr(mod, name) + if ( + type(old) is type(new) + and old.dtype == new.dtype + and old.untyped_storage().nbytes() == new.untyped_storage().nbytes() + ): + # If we can just update in-place to avoid re-registering + # can be faster if the underlying storage is the same + update_tensor_inplace(old, new) + else: + # Fallback re-register parameter, convert to Parameter if necessary + # this not only ensures we don't register a tensor as a parameter, but + # also ensures that all parameter subclasses get re-registered as + # parameters for `torch.compile` compatibility + if not isinstance(new, torch.nn.Parameter): + new = torch.nn.Parameter(new, requires_grad=False) + mod.register_parameter(name, torch.nn.Parameter(new, requires_grad=False)) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_fp8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..48cf5db34e64bbc648988cae5f8952324ae3a95b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_fp8.py @@ -0,0 +1,331 @@ +from typing import Any, Callable, Dict, List, Optional + +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.linear import LinearMethodBase +from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8 +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, + input_to_float8, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.utils import is_hip, set_weight_attrs + +_is_hip = is_hip() + + +class W8A8Fp8Config(QuantizationConfig): + """Config class for W8A8 FP8 Quantization. + + Weight Quantization: + - Method: Static quantization + - Granularity: Per-channel + - Type: Symmetric + + Activation Quantization: + - Method: Dynamic quantization + - Granularity: Per-token + - Type: Symmetric + + Note: + - For models without offline quantization, weights will be quantized during model loading + - If CUTLASS is supported: Per-channel weight quantization is used + - If CUTLASS is not supported: Falls back to per-tensor weight quantization + """ + + def __init__(self, is_checkpoint_fp8_serialized: bool = False): + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 89 + + @classmethod + def get_name(self) -> str: + return "w8a8_fp8" + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "W8A8Fp8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = ( + "compressed-tensors" in quant_method or "w8a8_fp8" in quant_method + ) + return cls(is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + return W8A8Fp8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return W8A8FP8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class W8A8Fp8LinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: W8A8Fp8Config): + self.cutlass_fp8_supported = cutlass_fp8_supported() + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight + + if self.quantization_config.is_checkpoint_fp8_serialized: + weight_scale = layer.weight_scale.detach() + # If checkpoint offline quantized with w8a8_fp8, load the weight and weight_scale directly. + if _is_hip: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=weight_scale + ) + + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + else: + # If checkpoint not offline quantized, quantize the weights with per-channel quantization. + if self.cutlass_fp8_supported: + # if cutlass supported, we use cutlass_scaled_mm + # which requires per-channel quantization on weight + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + if _is_hip: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=weight_scale + ) + else: + # if cutlass not supported, we fall back to use torch._scaled_mm + # which requires per tensor quantization on weight + fp8_dtype = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + qweight, weight_scale = input_to_float8(layer.weight, dtype=fp8_dtype) + + # Update the layer with the new values. + layer.weight = Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + layer.input_scale = None + + 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_dtype = ( + torch.float8_e4m3fn + if self.quantization_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + weight_loader = extra_weight_attrs.get("weight_loader") + self.logical_widths = output_partition_sizes + + weight = ModelWeightParameter( + data=torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + if self.quantization_config.is_checkpoint_fp8_serialized: + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + else: + layer.weight_scale = None + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + return apply_fp8_linear( + x, + layer.weight, + layer.weight_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) + + +class W8A8FP8MoEMethod: + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + Args: + quant_config: The quantization config. + """ + + def __new__(cls, *args, **kwargs): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoEMethodBase + + if not hasattr(cls, "_initialized"): + original_init = cls.__init__ + new_cls = type( + cls.__name__, + (FusedMoEMethodBase,), + { + "__init__": original_init, + **{k: v for k, v in cls.__dict__.items() if k != "__dict__"}, + }, + ) + obj = super(new_cls, new_cls).__new__(new_cls) + obj.__init__(*args, **kwargs) + return obj + return super().__new__(cls) + + def __init__(self, quant_config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + fp8_dtype = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size, hidden_size, dtype=fp8_dtype + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, intermediate_size, dtype=fp8_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2 * intermediate_size, 1, dtype=torch.float32), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w13_input_scale = None + layer.register_parameter("w13_input_scale", w13_input_scale) + + w2_input_scale = None + layer.register_parameter("w2_input_scale", w2_input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False) + layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False) + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data, requires_grad=False + ) + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data, requires_grad=False + ) + + def 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, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + # Expert selection + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + return fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + activation=activation, + use_fp8_w8a8=True, + per_channel_quant=True, + w1_scale=(layer.w13_weight_scale), + w2_scale=(layer.w2_weight_scale), + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + no_combine=no_combine, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_int8.py b/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_int8.py new file mode 100644 index 0000000000000000000000000000000000000000..8af2bb7020182349a73dc712b5a03b6fc0c3f83a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/w8a8_int8.py @@ -0,0 +1,269 @@ +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, 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, set_weight_attrs + +_is_cuda = is_cuda() +if _is_cuda: + from sgl_kernel import int8_scaled_mm + + +class W8A8Int8Config(QuantizationConfig): + """Config class for W8A8 Int8 Quantization. + + - Weight: static, per-channel, symmetric + - Activation: dynamic, per-token, symmetric + """ + + def __init__(self): + pass + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_name(self) -> str: + return "w8a8_int8" + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "W8A8Int8Config": + return cls() + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + return W8A8Int8LinearMethod(self) + elif isinstance(layer, FusedMoE): + return W8A8Int8MoEMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class W8A8Int8LinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: W8A8Int8Config): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight = Parameter(layer.weight.t(), requires_grad=False) + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + 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") + self.logical_widths = output_partition_sizes + + weight = ModelWeightParameter( + data=torch.empty( + sum(output_partition_sizes), input_size_per_partition, dtype=torch.int8 + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + x_q, x_scale = per_token_quant_int8(x) + + return int8_scaled_mm( + x_q, layer.weight, x_scale, layer.weight_scale, out_dtype=x.dtype, bias=bias + ) + + +class W8A8Int8MoEMethod: + """MoE method for INT8. + Supports loading INT8 checkpoints with static weight scale and + dynamic/static activation scale. + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + Args: + quant_config: The quantization config. + """ + + def __new__(cls, *args, **kwargs): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoEMethodBase + + if not hasattr(cls, "_initialized"): + original_init = cls.__init__ + new_cls = type( + cls.__name__, + (FusedMoEMethodBase,), + { + "__init__": original_init, + **{k: v for k, v in cls.__dict__.items() if k != "__dict__"}, + }, + ) + obj = super(new_cls, new_cls).__new__(new_cls) + obj.__init__(*args, **kwargs) + return obj + return super().__new__(cls) + + def __init__(self, quant_config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + tp_size = get_tensor_model_parallel_world_size() + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size, hidden_size, dtype=torch.int8 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, intermediate_size, dtype=torch.int8), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2 * intermediate_size, 1, dtype=torch.float32), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + w13_input_scale = None + layer.register_parameter("w13_input_scale", w13_input_scale) + + w2_input_scale = None + layer.register_parameter("w2_input_scale", w2_input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False) + layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False) + layer.w13_weight_scale = Parameter( + layer.w13_weight_scale.data, requires_grad=False + ) + layer.w2_weight_scale = Parameter( + layer.w2_weight_scale.data, requires_grad=False + ) + + def 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", + apply_router_weight_on_input: bool = False, + inplace: bool = True, + no_combine: bool = False, + routed_scaling_factor: Optional[float] = None, + ) -> torch.Tensor: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts + from sglang.srt.layers.moe.topk import select_experts + + # Expert selection + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=custom_routing_function, + correction_bias=correction_bias, + routed_scaling_factor=routed_scaling_factor, + ) + + return fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + use_int8_w8a8=True, + per_channel_quant=True, + w1_scale=(layer.w13_weight_scale), + w2_scale=(layer.w2_weight_scale), + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + no_combine=no_combine, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/cache_controller.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/cache_controller.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f324e694b5609e50636c1f3dcd3a7164d38bd397 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/cache_controller.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/configure_logging.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/configure_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bffa520507e9b641df493760d3fc9cf997050bb0 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/configure_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/data_parallel_controller.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/data_parallel_controller.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0df3fb53fdbea344410772df33a09a16f2fc2bb Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/data_parallel_controller.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/detokenizer_manager.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/detokenizer_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2988831acd572076559d382a17494d17be214c1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/detokenizer_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/expert_distribution.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/expert_distribution.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fc9cd29aea5125033e6c5c6380cee759e94cd34 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/expert_distribution.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/io_struct.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/io_struct.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9633ef640e30d3f175dd791281e96202725fe581 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/io_struct.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/mm_utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/mm_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a258d7c4c08959a2f47c896377e23412852fce06 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/mm_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/multimodal_processor.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/multimodal_processor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d18d486d11943da738d6b4411ae6b7bf8d36f0ce Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/multimodal_processor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_batch.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93396c59d1e3ba3d6ed787955d5d9ae47258c926 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_policy.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16dbf6b5c354005c1c53b4df4e14b775d90ddc33 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/schedule_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f1b1127707256faa2c77a525ad65646eebb16cc Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler_output_processor_mixin.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler_output_processor_mixin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b4b455f3f1d80ca9833dff2fda35e1e55b8cfba Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/scheduler_output_processor_mixin.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/session_controller.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/session_controller.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2238bfcb11320480e996406929ca9fa76ae2fded Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/session_controller.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tokenizer_manager.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tokenizer_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d466385a071d389995b143f4d7430d4c7bb83e8 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tokenizer_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2710b4dc82ce2a71672f823a390be56f5f119ec9 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker_overlap_thread.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker_overlap_thread.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..021441c55569d536de27984ee7f8a6e0bcf4a06b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/tp_worker_overlap_thread.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e55058bd139025fa6c6c865367cf21c0a162c2e Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/managers/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/managers/cache_controller.py b/lib/python3.12/site-packages/sglang/srt/managers/cache_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..8cb237cf171f2034b2f7a90500450a148f041b50 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/cache_controller.py @@ -0,0 +1,524 @@ +from __future__ import annotations + +""" +Copyright 2023-2025 SGLang Team +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import concurrent.futures +import logging +import math +import threading +from queue import Empty, Full, PriorityQueue, Queue +from typing import List, Optional + +import torch + +from sglang.srt.mem_cache.memory_pool import HostKVCache, TokenToKVPoolAllocator + +logger = logging.getLogger(__name__) + + +class LayerDoneCounter: + def __init__(self, num_layers): + self.counter = num_layers + self.condition = threading.Condition() + + def increment(self): + with self.condition: + self.counter += 1 + self.condition.notify_all() + + def wait_until(self, threshold): + with self.condition: + while self.counter <= threshold: + self.condition.wait() + + def reset(self): + with self.condition: + self.counter = 0 + + +class CacheOperation: + + counter = 0 + + def __init__( + self, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + node_id: int, + priority: Optional[int] = None, + ): + self.host_indices = host_indices + self.device_indices = device_indices + self.node_ids = [node_id] + self.data = None + + self.id = CacheOperation.counter + CacheOperation.counter += 1 + # default priority is the order of creation + self.priority = priority if priority is not None else self.id + + def merge(self, other: "CacheOperation") -> None: + # multiple operations can be merged into a single operation for batch processing + self.host_indices = torch.cat([self.host_indices, other.host_indices]) + self.device_indices = torch.cat([self.device_indices, other.device_indices]) + self.priority = min(self.priority, other.priority) + self.node_ids.extend(other.node_ids) + + def split(self, factor) -> List["CacheOperation"]: + # split an operation into smaller operations to reduce the size of intermediate buffers + if factor <= 1: + return [self] + + chunk_size = math.ceil(len(self.host_indices) / factor) + split_ops = [] + for i in range(0, len(self.host_indices), chunk_size): + split_ops.append( + CacheOperation( + host_indices=self.host_indices[i : i + chunk_size], + device_indices=self.device_indices[i : i + chunk_size], + node_id=0, + ) + ) + # Inherit the node_ids on the final chunk + if split_ops: + split_ops[-1].node_ids = self.node_ids + + return split_ops + + def __lt__(self, other: "CacheOperation"): + return self.priority < other.priority + + +class TransferBuffer: + """ + Overlapping buffer preparation and transfer operations to improve throughput. + """ + + def __init__( + self, stop_event, buffer_count: int = 3, max_buffer_size: int = 1000 + ) -> None: + self.stop_event = stop_event + self.buffers = Queue(maxsize=buffer_count) + # todo: adjust the buffer size based on throughput profile of the system + self.max_buffer_size = max_buffer_size + + def full(self) -> bool: + return self.buffers.full() + + def empty(self) -> bool: + return self.buffers.empty() + + def put(self, item, block=True, timeout=1) -> None: + while not self.stop_event.is_set(): + try: + self.buffers.put(item, block=block, timeout=timeout) + break + except Full: + if not block: + break + continue + except Exception as e: + logger.error(e) + + def get(self, block=True, timeout=1) -> Optional[CacheOperation]: + try: + return self.buffers.get(block=block, timeout=timeout) + except Empty: + return None + except Exception as e: + logger.error(e) + + def clear(self): + self.buffers.queue.clear() + + +class HiCacheController: + + def __init__( + self, + token_to_kv_pool_allocator: TokenToKVPoolAllocator, + mem_pool_host: HostKVCache, + page_size: int, + load_cache_event: threading.Event = None, + write_policy: str = "write_through_selective", + ): + self.mem_pool_device_allocator = token_to_kv_pool_allocator + self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache() + self.mem_pool_host = mem_pool_host + self.write_policy = write_policy + self.page_size = page_size + + self.load_cache_event = load_cache_event + self.layer_done_counter = LayerDoneCounter(self.mem_pool_device.layer_num) + self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) + + if write_policy not in [ + "write_through", + "write_through_selective", + "write_back", + ]: + raise ValueError(f"Invalid write policy: {write_policy}") + + self.write_queue = PriorityQueue() + self.load_queue = PriorityQueue() + + self.ack_write_queue = Queue() + self.ack_load_queue = Queue() + + self.stop_event = threading.Event() + self.write_buffer = TransferBuffer(self.stop_event) + self.load_buffer = TransferBuffer( + self.stop_event, buffer_count=10, max_buffer_size=100 + ) + + self.write_stream = torch.cuda.Stream() + self.load_stream = torch.cuda.Stream() + + self.write_thread = threading.Thread( + target=( + self.write_thread_func_buffer + if self.page_size == 1 + else self.write_thread_func_direct + ), + daemon=True, + ) + self.load_thread = threading.Thread( + target=self.load_thread_func_layer_by_layer, daemon=True + ) + self.write_thread.start() + self.load_thread.start() + + def reset(self): + self.stop_event.set() + self.write_thread.join() + self.load_thread.join() + + self.write_queue.queue.clear() + self.load_queue.queue.clear() + self.write_buffer.clear() + self.load_buffer.clear() + self.ack_write_queue.queue.clear() + self.ack_load_queue.queue.clear() + + self.write_thread = threading.Thread( + target=( + self.write_thread_func_buffer + if self.page_size == 1 + else self.write_thread_func_direct + ), + daemon=True, + ) + self.load_thread = threading.Thread( + target=self.load_thread_func_layer_by_layer, daemon=True + ) + self.stop_event.clear() + self.write_thread.start() + self.load_thread.start() + + def write( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = 0, + ) -> Optional[torch.Tensor]: + """ + Back up KV caches from device memory to host memory. + """ + host_indices = self.mem_pool_host.alloc(len(device_indices)) + if host_indices is None: + return None + self.mem_pool_host.protect_write(host_indices) + self.write_queue.put( + CacheOperation(host_indices, device_indices, node_id, priority) + ) + return host_indices + + def load( + self, + host_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = 0, + ) -> Optional[torch.Tensor]: + """ + Load KV caches from host memory to device memory. + """ + device_indices = self.mem_pool_device_allocator.alloc(len(host_indices)) + if device_indices is None: + return None + self.mem_pool_host.protect_load(host_indices) + # to ensure the device indices are ready before accessed by another CUDA stream + torch.cuda.current_stream().synchronize() + self.load_queue.put( + CacheOperation(host_indices, device_indices, node_id, priority) + ) + return device_indices + + def write_thread_func_direct(self): + """ + Directly write through KV caches to host memory without buffering. + """ + with torch.cuda.stream(self.write_stream): + while not self.stop_event.is_set(): + try: + operation = self.write_queue.get(block=True, timeout=1) + self.mem_pool_host.write_page_all_layers( + operation.host_indices, + operation.device_indices, + self.mem_pool_device, + ) + self.write_stream.synchronize() + self.mem_pool_host.complete_io(operation.host_indices) + for node_id in operation.node_ids: + if node_id != 0: + self.ack_write_queue.put(node_id) + except Empty: + continue + except Exception as e: + logger.error(e) + + def load_thread_func_direct(self): + """ + Directly load KV caches from host memory to device memory without buffering. + """ + with torch.cuda.stream(self.load_stream): + while not self.stop_event.is_set(): + try: + operation = self.load_queue.get(block=True, timeout=1) + # time.sleep(18e-6 * len(operation.host_indices)) + operation.data = self.mem_pool_host.get_flat_data( + operation.host_indices + ) + self.mem_pool_device.transfer( + operation.device_indices, operation.data + ) + self.mem_pool_host.complete_io(operation.host_indices) + for node_id in operation.node_ids: + if node_id != 0: + self.ack_load_queue.put(node_id) + except Empty: + continue + except Exception as e: + logger.error(e) + + def load_thread_func_layer_by_layer(self): + """ + Load KV caches from host memory to device memory layer by layer. + """ + with torch.cuda.stream(self.load_stream): + while not self.stop_event.is_set(): + self.load_cache_event.wait(timeout=1) + if not self.load_cache_event.is_set(): + continue + self.load_cache_event.clear() + + batch_operation = None + while self.load_queue.qsize() > 0: + op = self.load_queue.get(block=True) + if batch_operation is None: + batch_operation = op + else: + batch_operation.merge(op) + if batch_operation is None: + continue + + self.layer_done_counter.reset() + for i in range(self.mem_pool_host.layer_num): + if self.page_size == 1: + flat_data = self.mem_pool_host.get_flat_data_by_layer( + batch_operation.host_indices, i + ) + self.mem_pool_device.transfer_per_layer( + batch_operation.device_indices, flat_data, i + ) + else: + self.mem_pool_host.load_page_per_layer( + batch_operation.host_indices, + batch_operation.device_indices, + self.mem_pool_device, + i, + ) + self.load_stream.synchronize() + self.layer_done_counter.increment() + + self.mem_pool_host.complete_io(batch_operation.host_indices) + for node_id in batch_operation.node_ids: + if node_id != 0: + self.ack_load_queue.put(node_id) + + def write_aux_func(self, no_wait=False): + """ + Auxiliary function to prepare the buffer for write operations. + """ + + def _to_op(op_): + assert op_.device_indices.is_cuda, "Device indices should be on GPU" + op_.data = self.mem_pool_device.get_flat_data(op_.device_indices).to( + self.mem_pool_host.device + ) + self.write_buffer.put(op_) + return op_ + + buffer = None + with torch.cuda.stream(self.write_stream): + while not self.stop_event.is_set(): + try: + operation = self.write_queue.get(block=True, timeout=1) + factor = ( + len(operation.device_indices) + // self.write_buffer.max_buffer_size + ) + + if factor >= 1: + if buffer is not None: + _to_op(buffer) + buffer = None + + if factor < 2: + _to_op(operation) + else: + split_ops = operation.split(factor) + for op_ in split_ops: + _to_op(op_) + continue + + if buffer is None: + buffer = operation + else: + buffer.merge(operation) + if ( + no_wait + or len(buffer.host_indices) >= self.write_buffer.max_buffer_size + or self.write_queue.empty() + or self.write_buffer.empty() + ): + _to_op(buffer) + buffer = None + except Empty: + continue + except Exception as e: + logger.error(e) + + def load_aux_func(self): + """ + Auxiliary function to prepare the buffer for load operations. + """ + + def _pin_op(op_, put=True): + op_.data = ( + self.mem_pool_host.get_flat_data(op_.host_indices) + .contiguous() + .pin_memory() + ) + if put: + self.load_buffer.put(op_) + return op_ + + buffer = None + while not self.stop_event.is_set(): + try: + operation = self.load_queue.get(block=True, timeout=1) + factor = len(operation.host_indices) // self.load_buffer.max_buffer_size + + if factor >= 1: + if buffer is not None: + _pin_op(buffer) + buffer = None + + if factor < 2: + _pin_op(operation) + else: + split_ops = operation.split(factor) + split_args = [(op_, True) for op_ in split_ops[:-1]] + split_args.append((split_ops[-1], False)) + # Spawn threads to pin each op concurrently + with concurrent.futures.ThreadPoolExecutor() as executor: + pinned_ops = list( + executor.map( + lambda x: _pin_op(x[0], put=x[1]), split_args + ) + ) + # preserve the order of last op to ensure correct ack + self.load_buffer.put(pinned_ops[-1]) + continue + + if buffer is None: + buffer = operation + else: + buffer.merge(operation) + if ( + len(buffer.host_indices) >= self.load_buffer.max_buffer_size + or self.load_queue.empty() + or self.load_buffer.empty() + ): + _pin_op(buffer) + buffer = None + except Empty: + continue + except Exception as e: + logger.error(e) + + def write_thread_func_buffer(self): + aux_thread = threading.Thread(target=self.write_aux_func, daemon=True) + aux_thread.start() + + while not self.stop_event.is_set(): + operation = self.write_buffer.get() + if operation is None: + continue + self.mem_pool_host.assign_flat_data(operation.host_indices, operation.data) + self.mem_pool_host.complete_io(operation.host_indices) + for node_id in operation.node_ids: + if node_id != 0: + self.ack_write_queue.put(node_id) + aux_thread.join() + + def load_thread_func_buffer(self): + aux_thread = threading.Thread(target=self.load_aux_func, daemon=True) + aux_thread.start() + + with torch.cuda.stream(self.load_stream): + while not self.stop_event.is_set(): + operation = self.load_buffer.get() + if operation is None: + continue + self.mem_pool_device.transfer(operation.device_indices, operation.data) + self.mem_pool_host.complete_io(operation.host_indices) + for node_id in operation.node_ids: + if node_id != 0: + self.ack_load_queue.put(node_id) + aux_thread.join() + + def evict_device( + self, device_indices: torch.Tensor, host_indices: torch.Tensor + ) -> int: + if self.mem_pool_host.is_synced(host_indices): + self.mem_pool_device_allocator.free(device_indices) + self.mem_pool_host.update_backup(host_indices) + return len(device_indices) + else: + raise ValueError( + f"Inconsistent states: {self.mem_pool_host.get_state(host_indices)}" + ) + + def evict_host(self, host_indices: torch.Tensor, backup_only: bool = True) -> int: + if not backup_only: + raise ValueError("Other eviction policies are not supported yet.") + + if self.mem_pool_host.is_backup(host_indices): + self.mem_pool_host.free(host_indices) + return len(host_indices) + else: + raise ValueError( + f"Inconsistent states: {self.mem_pool_host.get_state(host_indices)}" + ) diff --git a/lib/python3.12/site-packages/sglang/srt/managers/configure_logging.py b/lib/python3.12/site-packages/sglang/srt/managers/configure_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..d331990ff76bafeb6265187e4735de9db165b1e6 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/configure_logging.py @@ -0,0 +1,47 @@ +""" +Copyright 2023-2025 SGLang Team +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +""" +Configure the logging settings of a server. + +Usage: +python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000 +""" + +import argparse + +import requests + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--url", type=str, default="http://localhost:30000") + parser.add_argument("--log-requests", action="store_true") + parser.add_argument("--log-requests-level", type=int, default=2) + parser.add_argument( + "--dump-requests-folder", type=str, default="/tmp/sglang_request_dump" + ) + parser.add_argument("--dump-requests-threshold", type=int, default=1000) + args = parser.parse_args() + + response = requests.post( + args.url + "/configure_logging", + json={ + "log_requests": args.log_requests, + "log_requests_level": args.log_requests_level, # Log full requests + "dump_requests_folder": args.dump_requests_folder, + "dump_requests_threshold": args.dump_requests_threshold, + }, + ) + assert response.status_code == 200 diff --git a/lib/python3.12/site-packages/sglang/srt/managers/data_parallel_controller.py b/lib/python3.12/site-packages/sglang/srt/managers/data_parallel_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..ce921988b9b34e12ca35b0917ecf4ed18307b3ce --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/data_parallel_controller.py @@ -0,0 +1,293 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""A controller that dispatches requests to multiple data parallel workers.""" + +import logging +import multiprocessing as mp +import signal +import threading +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, + TokenizedGenerateReqInput, +) +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.scheduler import run_scheduler_process +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter +from sglang.srt.utils import bind_port, configure_logger, get_zmq_socket +from sglang.utils import get_exception_traceback + +logger = logging.getLogger(__name__) + + +class LoadBalanceMethod(Enum): + """Load balance method.""" + + ROUND_ROBIN = auto() + SHORTEST_QUEUE = auto() + + @classmethod + def from_str(cls, method: str): + method = method.upper() + try: + return cls[method] + except KeyError as exc: + raise ValueError(f"Invalid load balance method: {method}") from exc + + +class DataParallelController: + """A controller that dispatches requests to multiple data parallel workers.""" + + def __init__(self, server_args: ServerArgs, port_args: PortArgs) -> None: + # Parse args + self.max_total_num_tokens = None + self.server_args = server_args + self.port_args = port_args + self.load_balance_method = LoadBalanceMethod.from_str( + server_args.load_balance_method + ) + + # Init inter-process communication + self.context = zmq.Context(1 + server_args.dp_size) + if server_args.node_rank == 0: + self.recv_from_tokenizer = get_zmq_socket( + self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False + ) + + # Dispatch method + self.round_robin_counter = 0 + dispatch_lookup = { + LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler, + LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler, + } + self.dispatching = dispatch_lookup[self.load_balance_method] + + # Launch data parallel workers + self.scheduler_procs = [] + self.workers = [None] * server_args.dp_size + + if server_args.enable_dp_attention: + dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args) + self.control_message_step = server_args.tp_size + else: + dp_port_args = self.launch_dp_schedulers(server_args, port_args) + self.control_message_step = 1 + + # Only node rank 0 runs the real data parallel controller that dispatches the requests. + if server_args.node_rank == 0: + for dp_rank in range(server_args.dp_size): + self.workers[dp_rank] = get_zmq_socket( + self.context, + zmq.PUSH, + dp_port_args[dp_rank].scheduler_input_ipc_name, + True, + ) + + self.max_req_input_len = None + + def launch_dp_schedulers(self, server_args, port_args): + base_gpu_id = 0 + + threads = [] + sockets = [] + dp_port_args = [] + ready_events = [] + for dp_rank in range(server_args.dp_size): + tmp_port_args = PortArgs.init_new(server_args) + tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name + tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name + dp_port_args.append(tmp_port_args) + + # This port is checked free in PortArgs.init_new. + # We hold it first so that the next dp worker gets a different port + sockets.append(bind_port(tmp_port_args.nccl_port)) + + ready_event = threading.Event() + ready_events.append(ready_event) + + # Create a thread for each worker + thread = threading.Thread( + target=self.launch_tensor_parallel_group_thread, + args=(server_args, tmp_port_args, base_gpu_id, dp_rank, ready_event), + ) + threads.append(thread) + base_gpu_id += server_args.tp_size * server_args.gpu_id_step + + # Free all sockets before starting the threads to launch TP workers + for sock in sockets: + sock.close() + + # Start all threads + for thread in threads: + thread.start() + for event in ready_events: + event.wait() + + return dp_port_args + + def launch_tensor_parallel_group_thread( + self, + server_args: ServerArgs, + port_args: PortArgs, + base_gpu_id: int, + dp_rank: int, + ready_event: threading.Event, + ): + self.launch_tensor_parallel_group(server_args, port_args, base_gpu_id, dp_rank) + ready_event.set() + + # This thread cannot be closed because otherwise the `kill_itself_when_parent_died` + # function in scheduler.py will kill the scheduler. + while True: + pass + + def launch_dp_attention_schedulers(self, server_args, port_args): + self.launch_tensor_parallel_group(server_args, port_args, 0, None) + dp_port_args = [] + for dp_rank in range(server_args.dp_size): + dp_port_args.append(PortArgs.init_new(server_args, dp_rank)) + return dp_port_args + + def launch_tensor_parallel_group( + self, + server_args: ServerArgs, + port_args: PortArgs, + base_gpu_id: int, + dp_rank: int, + ): + if not server_args.enable_dp_attention: + logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.") + + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=server_args.enable_memory_saver + ) + + # Launch tensor parallel scheduler processes + scheduler_pipe_readers = [] + tp_size_per_node = server_args.tp_size // server_args.nnodes + tp_rank_range = range( + tp_size_per_node * server_args.node_rank, + tp_size_per_node * (server_args.node_rank + 1), + ) + for tp_rank in tp_rank_range: + rank_port_args = port_args + + if server_args.enable_dp_attention: + # dp attention has different sharding logic + _, _, dp_rank = compute_dp_attention_world_info( + server_args.enable_dp_attention, + tp_rank, + server_args.tp_size, + server_args.dp_size, + ) + # compute zmq ports for this dp rank + rank_port_args = PortArgs.init_new(server_args, dp_rank) + # Data parallelism resues the tensor parallelism group, + # so all dp ranks should use the same nccl port. + rank_port_args.nccl_port = port_args.nccl_port + + reader, writer = mp.Pipe(duplex=False) + gpu_id = ( + server_args.base_gpu_id + + base_gpu_id + + (tp_rank % tp_size_per_node) * server_args.gpu_id_step + ) + proc = mp.Process( + target=run_scheduler_process, + args=(server_args, rank_port_args, gpu_id, tp_rank, dp_rank, writer), + ) + with memory_saver_adapter.configure_subprocess(): + proc.start() + self.scheduler_procs.append(proc) + scheduler_pipe_readers.append(reader) + + # Wait for model to finish loading + scheduler_info = [] + for i in range(len(scheduler_pipe_readers)): + scheduler_info.append(scheduler_pipe_readers[i].recv()) + + self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"] + self.max_req_input_len = scheduler_info[0]["max_req_input_len"] + + def round_robin_scheduler(self, req: Req): + if self.server_args.disaggregation_mode == "null": + self.workers[self.round_robin_counter].send_pyobj(req) + self.round_robin_counter = (self.round_robin_counter + 1) % len( + self.workers + ) + else: + self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req) + + def shortest_queue_scheduler(self, input_requests): + raise NotImplementedError() + + def event_loop(self): + while True: + while True: + try: + recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) + except zmq.ZMQError: + break + + if isinstance( + recv_req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + ), + ): + self.dispatching(recv_req) + else: + # Send other control messages to first worker of tp group + for worker in self.workers[:: self.control_message_step]: + worker.send_pyobj(recv_req) + + +def run_data_parallel_controller_process( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, +): + setproctitle.setproctitle("sglang::data_parallel_controller") + configure_logger(server_args) + parent_process = psutil.Process().parent() + + try: + controller = DataParallelController(server_args, port_args) + pipe_writer.send( + { + "status": "ready", + "max_total_num_tokens": controller.max_total_num_tokens, + "max_req_input_len": controller.max_req_input_len, + } + ) + if server_args.node_rank == 0: + controller.event_loop() + for proc in controller.scheduler_procs: + proc.join() + logger.error( + f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}" + ) + except Exception: + traceback = get_exception_traceback() + logger.error(f"DataParallelController hit an exception: {traceback}") + parent_process.send_signal(signal.SIGQUIT) diff --git a/lib/python3.12/site-packages/sglang/srt/managers/detokenizer_manager.py b/lib/python3.12/site-packages/sglang/srt/managers/detokenizer_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..aeae266eba0741c476f04649dbd15910d3fe7ea1 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/detokenizer_manager.py @@ -0,0 +1,266 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""DetokenizerManager is a process that detokenizes the token ids.""" + +import dataclasses +import logging +import os +import signal +from collections import OrderedDict +from typing import Dict, List, Union + +import psutil +import setproctitle +import zmq + +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.srt.managers.io_struct import ( + BatchEmbeddingOut, + BatchMultimodalDecodeReq, + BatchStrOut, + BatchTokenIDOut, +) +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.utils import ( + configure_logger, + get_zmq_socket, + kill_itself_when_parent_died, +) +from sglang.utils import ( + TypeBasedDispatcher, + find_printable_text, + get_exception_traceback, +) + +logger = logging.getLogger(__name__) + +# Maximum number of request states that detokenizer can hold. When exceeded, +# oldest request states will be evicted. Default: 65536 (1<<16). +# For more details, see: https://github.com/sgl-project/sglang/issues/2812 +# Use power of 2 values for better memory allocation. +DETOKENIZER_MAX_STATES = int(os.environ.get("SGLANG_DETOKENIZER_MAX_STATES", 1 << 16)) + + +@dataclasses.dataclass +class DecodeStatus: + """Store the status of incremental decoding.""" + + decoded_text: str + decode_ids: List[int] + surr_offset: int + read_offset: int + + +class DetokenizerManager: + """DetokenizerManager is a process that detokenizes the token ids.""" + + def __init__( + self, + server_args: ServerArgs, + port_args: PortArgs, + ): + # Init inter-process communication + context = zmq.Context(2) + self.recv_from_scheduler = get_zmq_socket( + context, zmq.PULL, port_args.detokenizer_ipc_name, True + ) + self.send_to_tokenizer = get_zmq_socket( + context, zmq.PUSH, port_args.tokenizer_ipc_name, False + ) + + if server_args.skip_tokenizer_init: + self.tokenizer = None + else: + self.tokenizer = get_tokenizer( + server_args.tokenizer_path, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ) + + self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES) + self.is_dummy = server_args.load_format == "dummy" + + self._request_dispatcher = TypeBasedDispatcher( + [ + (BatchEmbeddingOut, self.handle_batch_embedding_out), + (BatchTokenIDOut, self.handle_batch_token_id_out), + (BatchMultimodalDecodeReq, self.handle_multimodal_decode_req), + ] + ) + + def event_loop(self): + """The event loop that handles requests""" + while True: + recv_obj = self.recv_from_scheduler.recv_pyobj() + output = self._request_dispatcher(recv_obj) + self.send_to_tokenizer.send_pyobj(output) + + def trim_matched_stop( + self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool + ): + if no_stop_trim or not finished_reason: + return output + + matched = finished_reason.get("matched", None) + if not matched: + return output + + # TODO(lmzheng): handle the case where multiple stop strs are hit + + # Trim stop str. + if isinstance(matched, str) and isinstance(output, str): + pos = output.find(matched) + return output[:pos] if pos != -1 else output + + # Trim stop token. + if isinstance(matched, int) and isinstance(output, list): + assert len(output) > 0 + return output[:-1] + return output + + def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut): + # If it is embedding model, no detokenization is needed. + return recv_obj + + def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut): + bs = len(recv_obj.rids) + + # Initialize decode status + read_ids, surr_ids = [], [] + for i in range(bs): + rid = recv_obj.rids[i] + if rid not in self.decode_status: + s = DecodeStatus( + decoded_text=recv_obj.decoded_texts[i], + decode_ids=recv_obj.decode_ids[i], + surr_offset=0, + read_offset=recv_obj.read_offsets[i], + ) + self.decode_status[rid] = s + else: + s = self.decode_status[rid] + s.decode_ids = recv_obj.decode_ids[i] + + read_ids.append( + self.trim_matched_stop( + s.decode_ids[s.surr_offset :], + recv_obj.finished_reasons[i], + recv_obj.no_stop_trim[i], + ) + ) + surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset]) + + # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request + surr_texts = self.tokenizer.batch_decode( + surr_ids, + skip_special_tokens=recv_obj.skip_special_tokens[0], + spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0], + ) + read_texts = self.tokenizer.batch_decode( + read_ids, + skip_special_tokens=recv_obj.skip_special_tokens[0], + spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0], + ) + + # Incremental decoding + output_strs = [] + for i in range(bs): + try: + s = self.decode_status[recv_obj.rids[i]] + except KeyError: + raise RuntimeError( + f"Decode status not found for request {recv_obj.rids[i]}. " + "It may be due to the request being evicted from the decode status due to memory pressure. " + "Please increase the maximum number of requests by setting " + "the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. " + f"The current value is {DETOKENIZER_MAX_STATES}. " + "For more details, see: https://github.com/sgl-project/sglang/issues/2812" + ) + new_text = read_texts[i][len(surr_texts[i]) :] + if recv_obj.finished_reasons[i] is None: + # Streaming chunk: update the decode status + if len(new_text) > 0 and not new_text.endswith("�"): + s.decoded_text = s.decoded_text + new_text + s.surr_offset = s.read_offset + s.read_offset = len(s.decode_ids) + new_text = "" + else: + new_text = find_printable_text(new_text) + + output_strs.append( + self.trim_matched_stop( + s.decoded_text + new_text, + recv_obj.finished_reasons[i], + recv_obj.no_stop_trim[i], + ) + ) + + return BatchStrOut( + rids=recv_obj.rids, + finished_reasons=recv_obj.finished_reasons, + output_strs=output_strs, + output_ids=None, + prompt_tokens=recv_obj.prompt_tokens, + completion_tokens=recv_obj.completion_tokens, + cached_tokens=recv_obj.cached_tokens, + spec_verify_ct=recv_obj.spec_verify_ct, + input_token_logprobs_val=recv_obj.input_token_logprobs_val, + input_token_logprobs_idx=recv_obj.input_token_logprobs_idx, + output_token_logprobs_val=recv_obj.output_token_logprobs_val, + output_token_logprobs_idx=recv_obj.output_token_logprobs_idx, + input_top_logprobs_val=recv_obj.input_top_logprobs_val, + input_top_logprobs_idx=recv_obj.input_top_logprobs_idx, + output_top_logprobs_val=recv_obj.output_top_logprobs_val, + output_top_logprobs_idx=recv_obj.output_top_logprobs_idx, + input_token_ids_logprobs_val=recv_obj.input_token_ids_logprobs_val, + input_token_ids_logprobs_idx=recv_obj.input_token_ids_logprobs_idx, + output_token_ids_logprobs_val=recv_obj.output_token_ids_logprobs_val, + output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx, + output_hidden_states=recv_obj.output_hidden_states, + ) + + def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq): + raise NotImplementedError() + + +class LimitedCapacityDict(OrderedDict): + def __init__(self, capacity: int, *args, **kwargs): + super().__init__(*args, **kwargs) + self.capacity = capacity + + def __setitem__(self, key, value): + if len(self) >= self.capacity: + # Remove the oldest element (first item in the dict) + self.popitem(last=False) + # Set the new item + super().__setitem__(key, value) + + +def run_detokenizer_process( + server_args: ServerArgs, + port_args: PortArgs, +): + kill_itself_when_parent_died() + setproctitle.setproctitle("sglang::detokenizer") + configure_logger(server_args) + parent_process = psutil.Process().parent() + + try: + manager = DetokenizerManager(server_args, port_args) + manager.event_loop() + except Exception: + traceback = get_exception_traceback() + logger.error(f"DetokenizerManager hit an exception: {traceback}") + parent_process.send_signal(signal.SIGQUIT) diff --git a/lib/python3.12/site-packages/sglang/srt/managers/expert_distribution.py b/lib/python3.12/site-packages/sglang/srt/managers/expert_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..226256ed2186f941f0ac1f2f2d334d3f322e6f25 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/expert_distribution.py @@ -0,0 +1,81 @@ +import json +import logging +import time +from collections import defaultdict +from typing import Dict, List, Tuple + +import torch + +logger = logging.getLogger(__name__) + + +# global expert distribution recording +class ExpertDistributionRecorder: + # This class is a singleton class + def __new__(cls): + if not hasattr(cls, "instance"): + cls.instance = super(ExpertDistributionRecorder, cls).__new__(cls) + return cls.instance + + def __init__(self): + # the length of the dictionary is the number of layers + # the length of the list is the number of tokens + # the length of the tuple is topk's k value + self._expert_distribution_record: Dict[int, List[Tuple[int]]] = defaultdict( + list + ) + self._record = False + self._current_layer_id = "UNKNOWN" + + def set_current_layer(self, layer_idx): + self._current_layer_id = layer_idx + + def record_new_token(self, topk_ids): + if not self._record: + return + topk_ids_list = topk_ids.to("cpu", non_blocking=True).numpy().tolist() + torch.cuda.synchronize() + for i in topk_ids_list: + self._expert_distribution_record[self._current_layer_id].append(tuple(i)) + + def reset(self): + """Reset the expert distribution recorder.""" + logger.info("Resetting expert distribution record...") + self._record = False + self._expert_distribution_record.clear() + self._current_layer_id = "UNKNOWN" + + def start_record(self): + """Start recording the expert distribution. Reset the recorder and set the recording flag to True.""" + if self._record == True: + logger.warning( + "SGLang server is already recording expert ids. Did you forget to dump the expert ids recorded so far by sending requests to the `/stop_expert_distribution_record` and `/dump_expert_distribution_record` endpoints?" + ) + self.reset() + self._record = True + + def stop_record(self): + """Stop recording the expert distribution. Set the recording flag to False.""" + if self._record == False: + logger.warning( + "SGLang server has not been recording expert ids. Did you forget to start recording by sending request to the `/start_expert_distribution_record` endpoint?" + ) + self._record = False + + def dump_record(self): + """Dump the expert distribution record to a file. Reset the recorder after dumping.""" + results = {} + for layer_idx, layer_record in self._expert_distribution_record.items(): + results[layer_idx] = defaultdict(int) + for token_record in layer_record: + for expert_idx in token_record: + results[layer_idx][expert_idx] += 1 + with open( + f"expert_distribution_rank{torch.distributed.get_rank()}_timestamp{time.time()}.csv", + "w", + ) as fd: + fd.write("layer_id,expert_id,count\n") + for layer_idx, layer_results in results.items(): + for expert_idx, count in layer_results.items(): + fd.write(f"{layer_idx},{expert_idx},{count}\n") + self.reset() diff --git a/lib/python3.12/site-packages/sglang/srt/managers/io_struct.py b/lib/python3.12/site-packages/sglang/srt/managers/io_struct.py new file mode 100644 index 0000000000000000000000000000000000000000..174656b2dcb5474a69e6206eb1d3984e960d7a5b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/io_struct.py @@ -0,0 +1,939 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" +The definition of objects transfered between different +processes (TokenizerManager, DetokenizerManager, Controller). +""" + +import copy +import uuid +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union + +# handle serialization of Image for pydantic +if TYPE_CHECKING: + from PIL.Image import Image +else: + Image = Any + +from sglang.srt.managers.schedule_batch import BaseFinishReason +from sglang.srt.sampling.sampling_params import SamplingParams + + +@dataclass +class SessionParams: + id: Optional[str] = None + rid: Optional[str] = None + offset: Optional[int] = None + replace: Optional[bool] = None + + +@dataclass +class GenerateReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: Optional[Union[List[str], str]] = None + # The token ids for text; one can specify either text or input_ids + input_ids: Optional[Union[List[List[int]], List[int]]] = None + # The embeddings for input_ids; one can specify either text or input_ids or input_embeds. + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[ + Union[List[List[Union[Image, str]]], List[Union[Image, str]], Union[Image, str]] + ] = None + # The audio input. Like image data, it can be a file name, a url, or base64 encoded string. + audio_data: Optional[Union[List[str], str]] = None + # The sampling_params. See descriptions below. + sampling_params: Optional[Union[List[Dict], Dict]] = None + # The request id. + rid: Optional[Union[List[str], str]] = None + # Whether to return logprobs. + return_logprob: Optional[Union[List[bool], bool]] = None + # If return logprobs, the start location in the prompt for returning logprobs. + # By default, this value is "-1", which means it will only return logprobs for output tokens. + logprob_start_len: Optional[Union[List[int], int]] = None + # If return logprobs, the number of top logprobs to return at each position. + top_logprobs_num: Optional[Union[List[int], int]] = None + # If return logprobs, the token ids to return logprob for. + token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None + # Whether to detokenize tokens in text in the returned logprobs. + return_text_in_logprobs: bool = False + # Whether to stream output. + stream: bool = False + # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) + log_metrics: bool = True + + # The modalities of the image data [image, multi-images, video] + modalities: Optional[List[str]] = None + # LoRA related + lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None + + # Session info for continual prompting + session_params: Optional[Union[List[Dict], Dict]] = None + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None + + # Whether to return hidden states + return_hidden_states: bool = False + + # For disaggregated inference + bootstrap_host: Optional[Union[List[str], str]] = None + bootstrap_port: Optional[Union[List[int], int]] = None + bootstrap_room: Optional[Union[List[int], int]] = None + + def normalize_batch_and_arguments(self): + """ + Normalize the batch size and arguments for the request. + + This method resolves various input formats and ensures all parameters + are properly formatted as either single values or batches depending on the input. + It also handles parallel sampling expansion and sets default values for + unspecified parameters. + + Raises: + ValueError: If inputs are not properly specified (e.g., none or all of + text, input_ids, input_embeds are provided) + """ + self._validate_inputs() + self._determine_batch_size() + self._handle_parallel_sampling() + + if self.is_single: + self._normalize_single_inputs() + else: + self._normalize_batch_inputs() + + self._validate_session_params() + + def _validate_inputs(self): + """Validate that the input configuration is valid.""" + if ( + self.text is None and self.input_ids is None and self.input_embeds is None + ) or ( + self.text is not None + and self.input_ids is not None + and self.input_embeds is not None + ): + raise ValueError( + "Either text, input_ids or input_embeds should be provided." + ) + + def _determine_batch_size(self): + """Determine if this is a single example or a batch and the batch size.""" + if self.text is not None: + if isinstance(self.text, str): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.text) + self.input_embeds = None + elif self.input_ids is not None: + if len(self.input_ids) == 0: + raise ValueError("input_ids cannot be empty.") + if isinstance(self.input_ids[0], int): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.input_ids) + self.input_embeds = None + else: + if isinstance(self.input_embeds[0][0], float): + self.is_single = True + self.batch_size = 1 + else: + self.is_single = False + self.batch_size = len(self.input_embeds) + + def _handle_parallel_sampling(self): + """Handle parallel sampling parameters and adjust batch size if needed.""" + # Determine parallel sample count + if self.sampling_params is None: + self.parallel_sample_num = 1 + elif isinstance(self.sampling_params, dict): + self.parallel_sample_num = self.sampling_params.get("n", 1) + else: # isinstance(self.sampling_params, list): + self.parallel_sample_num = self.sampling_params[0].get("n", 1) + for sampling_params in self.sampling_params: + if self.parallel_sample_num != sampling_params.get("n", 1): + raise ValueError( + "The parallel_sample_num should be the same for all samples in sample params." + ) + + # If using parallel sampling with a single example, convert to batch + if self.parallel_sample_num > 1 and self.is_single: + self.is_single = False + if self.text is not None: + self.text = [self.text] + if self.input_ids is not None: + self.input_ids = [self.input_ids] + + def _normalize_single_inputs(self): + """Normalize inputs for a single example.""" + if self.sampling_params is None: + self.sampling_params = {} + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.return_logprob is None: + self.return_logprob = False + if self.logprob_start_len is None: + self.logprob_start_len = -1 + if self.top_logprobs_num is None: + self.top_logprobs_num = 0 + if not self.token_ids_logprob: # covers both None and [] + self.token_ids_logprob = None + + def _normalize_batch_inputs(self): + """Normalize inputs for a batch of examples, including parallel sampling expansion.""" + # Calculate expanded batch size + if self.parallel_sample_num == 1: + num = self.batch_size + else: + # Expand parallel_sample_num + num = self.batch_size * self.parallel_sample_num + + # Expand input based on type + self._expand_inputs(num) + self._normalize_lora_paths(num) + self._normalize_image_data(num) + self._normalize_audio_data(num) + self._normalize_sampling_params(num) + self._normalize_rid(num) + self._normalize_logprob_params(num) + self._normalize_custom_logit_processor(num) + + def _expand_inputs(self, num): + """Expand the main inputs (text, input_ids, input_embeds) for parallel sampling.""" + if self.text is not None: + if not isinstance(self.text, list): + raise ValueError("Text should be a list for batch processing.") + self.text = self.text * self.parallel_sample_num + elif self.input_ids is not None: + if not isinstance(self.input_ids, list) or not isinstance( + self.input_ids[0], list + ): + raise ValueError( + "input_ids should be a list of lists for batch processing." + ) + self.input_ids = self.input_ids * self.parallel_sample_num + elif self.input_embeds is not None: + if not isinstance(self.input_embeds, list): + raise ValueError("input_embeds should be a list for batch processing.") + self.input_embeds = self.input_embeds * self.parallel_sample_num + + def _normalize_lora_paths(self, num): + """Normalize LoRA paths for batch processing.""" + if self.lora_path is not None: + if isinstance(self.lora_path, str): + self.lora_path = [self.lora_path] * num + elif isinstance(self.lora_path, list): + self.lora_path = self.lora_path * self.parallel_sample_num + else: + raise ValueError("lora_path should be a list or a string.") + + def _normalize_image_data(self, num): + """Normalize image data for batch processing.""" + if self.image_data is None: + self.image_data = [None] * num + elif not isinstance(self.image_data, list): + # Single image, convert to list of single-image lists + self.image_data = [[self.image_data]] * num + self.modalities = ["image"] * num + elif isinstance(self.image_data, list): + if len(self.image_data) != self.batch_size: + raise ValueError( + "The length of image_data should be equal to the batch size." + ) + + self.modalities = [] + if len(self.image_data) > 0 and isinstance(self.image_data[0], list): + # Already a list of lists, keep as is + for i in range(len(self.image_data)): + if self.image_data[i] is None or self.image_data[i] == [None]: + self.modalities.append(None) + elif len(self.image_data[i]) == 1: + self.modalities.append("image") + elif len(self.image_data[i]) > 1: + self.modalities.append("multi-images") + # Expand parallel_sample_num + self.image_data = self.image_data * self.parallel_sample_num + self.modalities = self.modalities * self.parallel_sample_num + else: + # List of images for a batch, wrap each in a list + wrapped_images = [[img] for img in self.image_data] + # Expand for parallel sampling + self.image_data = wrapped_images * self.parallel_sample_num + self.modalities = ["image"] * num + + def _normalize_audio_data(self, num): + """Normalize audio data for batch processing.""" + if self.audio_data is None: + self.audio_data = [None] * num + elif not isinstance(self.audio_data, list): + self.audio_data = [self.audio_data] * num + elif isinstance(self.audio_data, list): + self.audio_data = self.audio_data * self.parallel_sample_num + + def _normalize_sampling_params(self, num): + """Normalize sampling parameters for batch processing.""" + if self.sampling_params is None: + self.sampling_params = [{}] * num + elif isinstance(self.sampling_params, dict): + self.sampling_params = [self.sampling_params] * num + else: # Already a list + self.sampling_params = self.sampling_params * self.parallel_sample_num + + def _normalize_rid(self, num): + """Normalize request IDs for batch processing.""" + if self.rid is None: + self.rid = [uuid.uuid4().hex for _ in range(num)] + elif not isinstance(self.rid, list): + raise ValueError("The rid should be a list for batch processing.") + + def _normalize_logprob_params(self, num): + """Normalize logprob-related parameters for batch processing.""" + + # Helper function to normalize a parameter + def normalize_param(param, default_value, param_name): + if param is None: + return [default_value] * num + elif not isinstance(param, list): + return [param] * num + else: + if self.parallel_sample_num > 1: + raise ValueError( + f"Cannot use list {param_name} with parallel_sample_num > 1" + ) + return param + + # Normalize each logprob parameter + self.return_logprob = normalize_param( + self.return_logprob, False, "return_logprob" + ) + self.logprob_start_len = normalize_param( + self.logprob_start_len, -1, "logprob_start_len" + ) + self.top_logprobs_num = normalize_param( + self.top_logprobs_num, 0, "top_logprobs_num" + ) + + # Handle token_ids_logprob specially due to its nested structure + if not self.token_ids_logprob: # covers both None and [] + self.token_ids_logprob = [None] * num + elif not isinstance(self.token_ids_logprob, list): + self.token_ids_logprob = [[self.token_ids_logprob] for _ in range(num)] + elif not isinstance(self.token_ids_logprob[0], list): + self.token_ids_logprob = [ + copy.deepcopy(self.token_ids_logprob) for _ in range(num) + ] + elif self.parallel_sample_num > 1: + raise ValueError( + "Cannot use list token_ids_logprob with parallel_sample_num > 1" + ) + + def _normalize_custom_logit_processor(self, num): + """Normalize custom logit processor for batch processing.""" + if self.custom_logit_processor is None: + self.custom_logit_processor = [None] * num + elif not isinstance(self.custom_logit_processor, list): + self.custom_logit_processor = [self.custom_logit_processor] * num + elif self.parallel_sample_num > 1: + raise ValueError( + "Cannot use list custom_logit_processor with parallel_sample_num > 1" + ) + + def _validate_session_params(self): + """Validate that session parameters are properly formatted.""" + if self.session_params is not None: + if not isinstance(self.session_params, dict) and not isinstance( + self.session_params[0], dict + ): + raise ValueError("Session params must be a dict or a list of dicts.") + + def regenerate_rid(self): + """Generate a new request ID and return it.""" + self.rid = uuid.uuid4().hex + return self.rid + + def __getitem__(self, i): + return GenerateReqInput( + text=self.text[i] if self.text is not None else None, + input_ids=self.input_ids[i] if self.input_ids is not None else None, + image_data=self.image_data[i], + audio_data=self.audio_data[i], + sampling_params=self.sampling_params[i], + rid=self.rid[i], + return_logprob=self.return_logprob[i], + logprob_start_len=self.logprob_start_len[i], + top_logprobs_num=self.top_logprobs_num[i], + token_ids_logprob=self.token_ids_logprob[i], + return_text_in_logprobs=self.return_text_in_logprobs, + stream=self.stream, + log_metrics=self.log_metrics, + modalities=self.modalities[i] if self.modalities else None, + lora_path=self.lora_path[i] if self.lora_path is not None else None, + custom_logit_processor=( + self.custom_logit_processor[i] + if self.custom_logit_processor is not None + else None + ), + return_hidden_states=self.return_hidden_states, + bootstrap_host=( + self.bootstrap_host[i] if self.bootstrap_host is not None else None + ), + bootstrap_port=( + self.bootstrap_port[i] if self.bootstrap_port is not None else None + ), + bootstrap_room=( + self.bootstrap_room[i] if self.bootstrap_room is not None else None + ), + ) + + +@dataclass +class TokenizedGenerateReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: List[int] + # The multimodal inputs + mm_inputs: dict + # The sampling parameters + sampling_params: SamplingParams + # Whether to return the logprobs + return_logprob: bool + # If return logprobs, the start location in the prompt for returning logprobs. + logprob_start_len: int + # If return logprobs, the number of top logprobs to return at each position. + top_logprobs_num: int + # If return logprobs, the token id to return logprob for + token_ids_logprob: List[int] + # Whether to stream output + stream: bool + + # LoRA related + lora_path: Optional[str] = None # None means just use the base model + # The input embeds + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + + # Session info for continual prompting + session_params: Optional[SessionParams] = None + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: Optional[str] = None + + # Whether to return hidden states + return_hidden_states: bool = False + + # For disaggregated inference + bootstrap_host: Optional[str] = None + bootstrap_port: Optional[int] = None + bootstrap_room: Optional[int] = None + + +@dataclass +class EmbeddingReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: Optional[Union[List[str], str]] = None + # The image input. It can be an image instance, file name, URL, or base64 encoded string. + # Can be formatted as: + # - Single image for a single request + # - List of images (one per request in a batch) + # - List of lists of images (multiple images per request) + # See also python/sglang/srt/utils.py:load_image for more details. + image_data: Optional[ + Union[List[List[Union[Image, str]]], List[Union[Image, str]], Union[Image, str]] + ] = None + # The audio input. Like image data, it can be a file name, a url, or base64 encoded string. + audio_data: Optional[Union[List[str], str]] = None + # The token ids for text; one can either specify text or input_ids. + input_ids: Optional[Union[List[List[int]], List[int]]] = None + # The request id. + rid: Optional[Union[List[str], str]] = None + # Dummy sampling params for compatibility + sampling_params: Union[List[Dict], Dict] = None + # Dummy input embeds for compatibility + input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) + log_metrics: bool = True + # The modalities of the image data [image, multi-images, video] + modalities: Optional[List[str]] = None + + def normalize_batch_and_arguments(self): + # at least one of text, input_ids, or image should be provided + if self.text is None and self.input_ids is None and self.image_data is None: + raise ValueError( + "At least one of text, input_ids, or image should be provided" + ) + + # text and input_ids cannot be provided at the same time + if self.text is not None and self.input_ids is not None: + raise ValueError("text and input_ids cannot be provided at the same time") + + # Derive the batch size + self.batch_size = 0 + self.is_single = True + + # check the batch size of text + if self.text is not None: + if isinstance(self.text, list): + self.batch_size += len(self.text) + else: + self.batch_size += 1 + + # check the batch size of input_ids + if self.input_ids is not None: + if isinstance(self.input_ids[0], list): + self.batch_size += len(self.input_ids) + else: + self.batch_size += 1 + + if self.batch_size > 1: + self.is_single = False + + # Fill in default arguments + if self.is_single: + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.sampling_params is None: + self.sampling_params = {} + self.sampling_params["max_new_tokens"] = 0 + else: + if self.rid is None: + self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)] + else: + assert isinstance(self.rid, list), "The rid should be a list." + + if self.sampling_params is None: + self.sampling_params = [{}] * self.batch_size + for i in range(self.batch_size): + self.sampling_params[i]["max_new_tokens"] = 0 + + def regenerate_rid(self): + self.rid = uuid.uuid4().hex + return self.rid + + def __getitem__(self, i): + return EmbeddingReqInput( + text=self.text[i] if self.text is not None else None, + input_ids=self.input_ids[i] if self.input_ids is not None else None, + image_data=self.image_data[i] if self.image_data is not None else None, + sampling_params=self.sampling_params[i], + rid=self.rid[i], + ) + + +@dataclass +class TokenizedEmbeddingReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: List[int] + # The image inputs + image_inputs: dict + # Dummy sampling params for compatibility + sampling_params: SamplingParams + + +@dataclass +class BatchTokenIDOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[BaseFinishReason] + # For incremental decoding + decoded_texts: List[str] + decode_ids: List[int] + read_offsets: List[int] + # Only used when `--skip-tokenizer-init` is on + output_ids: Optional[List[int]] + # Detokenization configs + skip_special_tokens: List[bool] + spaces_between_special_tokens: List[bool] + no_stop_trim: List[bool] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + spec_verify_ct: List[int] + + # Logprobs + input_token_logprobs_val: List[float] + input_token_logprobs_idx: List[int] + output_token_logprobs_val: List[float] + output_token_logprobs_idx: List[int] + input_top_logprobs_val: List[List] + input_top_logprobs_idx: List[List] + output_top_logprobs_val: List[List] + output_top_logprobs_idx: List[List] + input_token_ids_logprobs_val: List[List] + input_token_ids_logprobs_idx: List[List] + output_token_ids_logprobs_val: List[List] + output_token_ids_logprobs_idx: List[List] + + # Hidden states + output_hidden_states: List[List[float]] + + +@dataclass +class BatchMultimodalDecodeReq: + # The request id + rids: List[str] + finished_reasons: List[BaseFinishReason] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + + +@dataclass +class BatchStrOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[dict] + # The output decoded strings + output_strs: List[str] + # The token ids + output_ids: Optional[List[int]] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + spec_verify_ct: List[int] + + # Logprobs + input_token_logprobs_val: List[float] + input_token_logprobs_idx: List[int] + output_token_logprobs_val: List[float] + output_token_logprobs_idx: List[int] + input_top_logprobs_val: List[List] + input_top_logprobs_idx: List[List] + output_top_logprobs_val: List[List] + output_top_logprobs_idx: List[List] + input_token_ids_logprobs_val: List[List] + input_token_ids_logprobs_idx: List[List] + output_token_ids_logprobs_val: List[List] + output_token_ids_logprobs_idx: List[List] + + # Hidden states + output_hidden_states: List[List[float]] + + +@dataclass +class BatchMultimodalOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[dict] + # The outputs + outputs: List[List[Dict]] + + # Token counts + prompt_tokens: List[int] + completion_tokens: List[int] + cached_tokens: List[int] + + +@dataclass +class BatchEmbeddingOut: + # The request id + rids: List[str] + # The finish reason + finished_reasons: List[BaseFinishReason] + # The output embedding + embeddings: List[List[float]] + # Token counts + prompt_tokens: List[int] + cached_tokens: List[int] + + +@dataclass +class FlushCacheReqInput: + pass + + +@dataclass +class FlushCacheReqOutput: + success: bool + + +@dataclass +class UpdateWeightFromDiskReqInput: + # The model path with the new weights + model_path: str + # The format to load the weights + load_format: Optional[str] = None + + +@dataclass +class UpdateWeightFromDiskReqOutput: + success: bool + message: str + # Number of paused requests during weight sync. + num_paused_requests: Optional[int] = 0 + + +@dataclass +class UpdateWeightsFromDistributedReqInput: + name: str + dtype: str + shape: List[int] + + +@dataclass +class UpdateWeightsFromDistributedReqOutput: + success: bool + message: str + + +@dataclass +class UpdateWeightsFromTensorReqInput: + """Update model weights from tensor input. + + - Tensors are serialized for transmission + - Data is structured in JSON for easy transmission over HTTP + """ + + serialized_named_tensors: List[Union[str, bytes]] + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to flush the cache after updating weights + flush_cache: bool = True + + +@dataclass +class UpdateWeightsFromTensorReqOutput: + success: bool + message: str + + +@dataclass +class InitWeightsUpdateGroupReqInput: + # The master address + master_address: str + # The master port + master_port: int + # The rank offset + rank_offset: int + # The world size + world_size: int + # The group name + group_name: str = "weight_update_group" + # The backend + backend: str = "nccl" + + +@dataclass +class InitWeightsUpdateGroupReqOutput: + success: bool + message: str + + +@dataclass +class GetWeightsByNameReqInput: + name: str + truncate_size: int = 100 + + +@dataclass +class GetWeightsByNameReqOutput: + parameter: list + + +@dataclass +class ReleaseMemoryOccupationReqInput: + pass + + +@dataclass +class ReleaseMemoryOccupationReqOutput: + pass + + +@dataclass +class ResumeMemoryOccupationReqInput: + pass + + +@dataclass +class ResumeMemoryOccupationReqOutput: + pass + + +@dataclass +class AbortReq: + # The request id + rid: str + + +@dataclass +class GetInternalStateReq: + pass + + +@dataclass +class GetInternalStateReqOutput: + internal_state: Dict[Any, Any] + + +@dataclass +class SetInternalStateReq: + server_args: Dict[str, Any] + + +@dataclass +class SetInternalStateReqOutput: + updated: bool + server_args: Dict[str, Any] + + +@dataclass +class ProfileReqInput: + # The output directory + output_dir: Optional[str] = None + # If set, it profile as many as this number of steps. + # If it is set, profiling is automatically stopped after this step, and + # the caller doesn't need to run stop_profile. + num_steps: Optional[int] = None + activities: Optional[List[Literal["CPU", "GPU", "MEM", "CUDA_PROFILER"]]] = None + + +class ProfileReqType(Enum): + START_PROFILE = 1 + STOP_PROFILE = 2 + + +class ExpertDistributionReq(Enum): + START_RECORD = 1 + STOP_RECORD = 2 + DUMP_RECORD = 3 + + +@dataclass +class ExpertDistributionReqOutput: + pass + + +@dataclass +class ProfileReq: + type: ProfileReqType + output_dir: Optional[str] = None + num_steps: Optional[int] = None + activities: Optional[List[str]] = None + with_stack: Optional[bool] = None + record_shapes: Optional[bool] = None + profile_id: Optional[str] = None + + +@dataclass +class ProfileReqOutput: + success: bool + message: str + + +@dataclass +class ConfigureLoggingReq: + log_requests: Optional[bool] = None + log_requests_level: Optional[int] = None + dump_requests_folder: Optional[str] = None + dump_requests_threshold: Optional[int] = None + + +@dataclass +class OpenSessionReqInput: + capacity_of_str_len: int + session_id: Optional[str] = None + + +@dataclass +class CloseSessionReqInput: + session_id: str + + +@dataclass +class OpenSessionReqOutput: + session_id: Optional[str] + success: bool + + +@dataclass +class HealthCheckOutput: + pass + + +@dataclass +class Function: + description: Optional[str] = None + name: Optional[str] = None + parameters: Optional[object] = None + + +@dataclass +class Tool: + function: Function + type: Optional[str] = "function" + + +@dataclass +class ParseFunctionCallReq: + text: str # The text to parse. + tools: List[Tool] = field( + default_factory=list + ) # A list of available function tools (name, parameters, etc.). + tool_call_parser: Optional[str] = ( + None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. + ) + + +@dataclass +class SeparateReasoningReqInput: + text: str # The text to parse. + reasoning_parser: str # Specify the parser type, e.g., "deepseek-r1". + + +@dataclass +class VertexGenerateReqInput: + instances: List[dict] + parameters: Optional[dict] = None + + +@dataclass +class RpcReqInput: + method: str + parameters: Optional[Dict] = None + + +@dataclass +class RpcReqOutput: + success: bool + message: str diff --git a/lib/python3.12/site-packages/sglang/srt/managers/mm_utils.py b/lib/python3.12/site-packages/sglang/srt/managers/mm_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..025a3010bedec9a35b78d8668f1b646147937203 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/managers/mm_utils.py @@ -0,0 +1,478 @@ +""" +Multi-modality utils +""" + +import logging +from abc import abstractmethod +from typing import Callable, List, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, + global_server_args_dict, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils import flatten_nested_list, print_warning_once + +logger = logging.getLogger(__name__) + + +class MultiModalityDataPaddingPattern: + """ + Data tokens (like image tokens) often need special handling during padding + to maintain model compatibility. This class provides the interface for + implementing different padding strategies for data tokens + """ + + @abstractmethod + def pad_input_tokens( + self, input_ids: List[int], mm_inputs: MultimodalInputs + ) -> List[int]: + """ + Pad the input ids sequence containing data tokens, and replace them with pad_values + """ + pass + + +class MultiModalityDataPaddingPatternTokenPairs(MultiModalityDataPaddingPattern): + """In this pattern, data tokens should be enclosed by special token pairs (e.g. ..., data_token_pairs) + + This strategy should be applied when data content is marked by start/end token pairs in the input sequence. + """ + + def __init__(self, data_token_pairs: Optional[List[Tuple[int, int]]]) -> None: + self.data_token_id_pairs = data_token_pairs + + def pad_input_tokens( + self, input_ids: List[int], mm_inputs: MultimodalInputs + ) -> List[int]: + """ + This function will replace the data-tokens inbetween with pad_values accordingly + """ + pad_values = [item.pad_value for item in mm_inputs.mm_items] + data_token_pairs = self.data_token_id_pairs + mm_inputs.data_offsets = [] + if data_token_pairs is None: + data_token_pairs = [mm_inputs.im_start_id, mm_inputs.im_end_id] + if data_token_pairs is None: + print_warning_once( + "No data_token_pairs provided, RadixAttention might be influenced." + ) + return input_ids + start_token_ids = [s for s, _e in data_token_pairs] + end_tokens_ids = [e for _s, e in data_token_pairs] + + padded_ids = [] + last_idx = 0 + data_idx = -1 + + start_indices = [i for i, x in enumerate(input_ids) if x in start_token_ids] + end_indices = [i for i, x in enumerate(input_ids) if x in end_tokens_ids] + + if len(start_indices) != len(end_indices): + return input_ids + + for start_idx, end_idx in zip(start_indices, end_indices): + padded_ids.extend(input_ids[last_idx : start_idx + 1]) + + if input_ids[start_idx] in start_token_ids: + data_idx += 1 + mm_inputs.data_offsets += [start_idx] + + if data_idx >= len(pad_values): + data_idx = len(pad_values) - 1 + + num_tokens = end_idx - start_idx - 1 + pad_value = pad_values[data_idx] + padded_ids.extend([pad_value] * num_tokens) + + last_idx = end_idx + + padded_ids.extend(input_ids[last_idx:]) + + assert len(input_ids) == len(padded_ids), "Length validation fails" + return padded_ids + + +class MultiModalityDataPaddingPatternMultimodalTokens(MultiModalityDataPaddingPattern): + """In this pattern, data tokens should be represented as repetitions of a single token + e.g. ...., or