From 67e9d287eea004a3f6fac67be9708d9fca3a79bb Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Thu, 13 Nov 2025 08:16:00 -0800 Subject: [PATCH] [Quantization] Support Quark Dense + MoE FP8 & FP8 PTPC (#10485) Co-authored-by: HAI Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com> --- python/sglang/srt/configs/model_config.py | 1 + .../srt/layers/quantization/__init__.py | 14 +- .../schemes/compressed_tensors_w8a8_fp8.py | 2 +- .../srt/layers/quantization/fp8_utils.py | 336 ++++++------------ .../srt/layers/quantization/quark/quark.py | 43 ++- .../layers/quantization/quark/quark_moe.py | 301 +++++++++++++++- .../quantization/quark/schemes/__init__.py | 3 +- .../quark/schemes/quark_w4a4_mxfp4.py | 11 +- .../quark/schemes/quark_w8a8_fp8.py | 186 ++++++++++ .../srt/layers/quantization/quark/utils.py | 12 +- 10 files changed, 666 insertions(+), 243 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index b51d19d26..2c12ad0f3 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -663,6 +663,7 @@ class ModelConfig: "qoq", "w4afp8", "petit_nvfp4", + "quark", ] compatible_quantization_methods = { "modelopt_fp8": ["modelopt"], diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 5693d62d4..87d6eb03d 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -35,6 +35,7 @@ from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config from sglang.srt.layers.quantization.petit import PetitNvFp4Config from sglang.srt.layers.quantization.qoq import QoQConfig +from sglang.srt.layers.quantization.quark.quark import QuarkConfig from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config @@ -65,23 +66,14 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { "w4afp8": W4AFp8Config, "petit_nvfp4": PetitNvFp4Config, "fbgemm_fp8": FBGEMMFp8Config, + "quark": QuarkConfig, "auto-round": AutoRoundConfig, } -if is_cuda(): +if is_cuda() or (_is_mxfp_supported and is_hip()): BASE_QUANTIZATION_METHODS.update( { - "quark": Mxfp4Config, - "mxfp4": Mxfp4Config, - } - ) -elif _is_mxfp_supported and is_hip(): - from sglang.srt.layers.quantization.quark.quark import QuarkConfig - - BASE_QUANTIZATION_METHODS.update( - { - "quark": QuarkConfig, "mxfp4": Mxfp4Config, } ) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index a157ebc3e..7ea1545a0 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -84,7 +84,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): if _use_aiter: layer.weight = Parameter( - shuffle_weight(weight, (16, 16)), requires_grad=False + shuffle_weight(weight, (16, 16)).t(), requires_grad=False ) else: layer.weight = Parameter(weight.t(), requires_grad=False) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ff2d3c2c0..ba589cc74 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -604,158 +604,16 @@ def apply_fp8_linear( 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 + # Maybe apply padding to output, see comment in __init__ + num_token_padding = output_padding 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" - - cutlass_compatible_b = ( - weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 - ) - if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: - # Massage the input to be 2D - qinput = qinput.view(-1, qinput.shape[-1]) - output = triton_scaled_mm( - qinput, weight, x_scale, weight_scale, input.dtype, bias - ) - else: - 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 or _use_aiter) - ): - # into this sector means use dynamic per-token-per-channel quant - # per-token scale quant for input matrix, every row(one token) have one scale factor - # per-channel scale quant for weight matrix, every col(one channel) have one scale factor - if _use_aiter: - # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype) - # XQ -> input tensor, shape = (m, k) - # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf - # x_scale -> input scale tensor, shape = (m, 1) - # w_scale -> weight scale tensor, shape = (n ,1) - # dtype -> output dtype - output = gemm_a8w8_bpreshuffle( - XQ=qinput, - WQ=weight, - x_scale=x_scale, - w_scale=weight_scale, - dtype=input.dtype, - ) - if bias is not None: - output += bias - return _process_scaled_mm_output( - output, input_2d.shape, [*input.shape[:-1], weight.shape[0]] - ) - else: - # 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, - ) + num_token_padding = None + qinput, x_scale = scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=num_token_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) else: # cutlass w8a8 fp8 sgl-kernel only supports per-token scale if input_scale is not None: @@ -783,53 +641,12 @@ def apply_fp8_linear( 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" - - cutlass_compatible_b = ( - weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 - ) - if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: - # Massage the input to be 2D - qinput = qinput.view(-1, qinput.shape[-1]) - output = triton_scaled_mm( - qinput, weight, x_scale, weight_scale, input.dtype, bias - ) - else: - 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( + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: + # 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 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, @@ -837,33 +654,112 @@ def apply_fp8_linear( 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 + cutlass_compatible_b = ( + weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 + ) + if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: + # Massage the input to be 2D + qinput = qinput.view(-1, qinput.shape[-1]) + output = triton_scaled_mm( + qinput, weight, x_scale, weight_scale, input.dtype, bias + ) + else: + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) - # 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( + # 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 ( + use_per_token_if_dynamic + and not per_tensor_weights + and not per_tensor_activations + and (USE_ROWWISE_TORCH_SCALED_MM or _use_aiter) + ): + # into this sector means use dynamic per-token-per-channel quant + # per-token scale quant for input matrix, every row(one token) have one scale factor + # per-channel scale quant for weight matrix, every col(one channel) have one scale factor + if _use_aiter: + # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype) + # XQ -> input tensor, shape = (m, k) + # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf + # x_scale -> input scale tensor, shape = (m, 1) + # w_scale -> weight scale tensor, shape = (n ,1) + # dtype -> output dtype + output = gemm_a8w8_bpreshuffle( + XQ=qinput, + WQ=weight.T, + x_scale=x_scale, + w_scale=weight_scale, + dtype=input.dtype, + ) + if bias is not None: + output += bias + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + else: + # 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, - x_scale, - weight_scale, - input_2d.shape, - output_shape, - bias, - input.dtype, + 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) + + 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) + + # 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, + ) def can_auto_enable_marlin_fp8() -> bool: diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index d0fbe74ef..37500e687 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -14,7 +14,11 @@ from sglang.srt.layers.quantization.base_config import ( # noqa: E501 ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod from sglang.srt.layers.quantization.quark.quark_moe import QuarkMoEMethod -from sglang.srt.layers.quantization.quark.schemes import QuarkScheme, QuarkW4A4MXFP4 +from sglang.srt.layers.quantization.quark.schemes import ( + QuarkScheme, + QuarkW4A4MXFP4, + QuarkW8A8Fp8, +) from sglang.srt.layers.quantization.quark.utils import deep_compare, should_ignore_layer from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.utils import get_device_capability @@ -173,6 +177,37 @@ class QuarkConfig(QuantizationConfig): else: return False + def _is_fp8_w8a8( + self, + weight_quant: Optional[dict[str, Any]], + input_quant: Optional[dict[str, Any]], + ) -> bool: + # Confirm weights and input quantized. + if weight_quant is None or input_quant is None: + return False + + # Confirm weight scheme is supported + is_fp8_dtype = ( + weight_quant.get("dtype") == "fp8_e4m3" + and input_quant.get("dtype") == "fp8_e4m3" + ) + is_static_weight = not weight_quant.get("is_dynamic") + is_per_tensor_or_channel_weight = weight_quant.get("qscheme") in [ + "per_tensor", + "per_channel", + ] + + if not (is_fp8_dtype and is_static_weight and is_per_tensor_or_channel_weight): + return False + + # Dynamic quantization is always supported if weights supported. + if input_quant.get("is_dynamic"): + return True + + # Confirm activation scheme is supported. + is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor" + return is_per_tensor_activation + def _is_mx_fp4( self, weight_quant: Optional[dict[str, Any]], @@ -281,6 +316,12 @@ class QuarkConfig(QuantizationConfig): if self._is_mx_fp4(weight_config, input_config): return QuarkW4A4MXFP4(weight_config, input_config) + if self._is_fp8_w8a8(weight_config, input_config): + is_fp8_w8a8_supported = self._check_scheme_supported( + QuarkW8A8Fp8.get_min_capability(), error=False + ) + if is_fp8_w8a8_supported: + return QuarkW8A8Fp8(weight_config, input_config) raise NotImplementedError( "No quark compatible scheme was found. " diff --git a/python/sglang/srt/layers/quantization/quark/quark_moe.py b/python/sglang/srt/layers/quantization/quark/quark_moe.py index 9607d392e..497e69b8e 100644 --- a/python/sglang/srt/layers/quantization/quark/quark_moe.py +++ b/python/sglang/srt/layers/quantization/quark/quark_moe.py @@ -6,13 +6,13 @@ import logging from typing import TYPE_CHECKING, Any import torch -from aiter import ActivationType, QuantType -from aiter.fused_moe import fused_moe -from aiter.ops.shuffle import shuffle_weight -from aiter.utility.fp4_utils import e8m0_shuffle -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, 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, per_tensor_dequantize from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs if TYPE_CHECKING: @@ -29,6 +29,17 @@ _is_shuffle_moe_mxfp4 = get_bool_env_var("AITER_MXFP4_MOE_SF") and _is_hip __all__ = ["QuarkMoEMethod", "QuarkW4A4MXFp4MoEMethod"] +_is_fp8_fnuz = is_fp8_fnuz() +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +if _use_aiter: + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe + from aiter.ops.shuffle import shuffle_weight + from aiter.utility.fp4_utils import e8m0_shuffle + + from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1 + OCP_MX_BLOCK_SIZE = 32 if TYPE_CHECKING: @@ -59,6 +70,8 @@ class QuarkMoEMethod(FusedMoEMethodBase): if quant_config._is_mx_fp4(weight_config, input_config): return QuarkW4A4MXFp4MoEMethod(weight_config, input_config) + elif quant_config._is_fp8_w8a8(weight_config, input_config): + return QuarkW8A8FP8MoEMethod(weight_config, input_config) else: raise RuntimeError("Unsupported FusedMoe scheme") @@ -224,3 +237,281 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod): doweight_stage1=False, ) return StandardCombineInput(hidden_states=output) + + +class QuarkW8A8FP8MoEMethod(QuarkMoEMethod): + + def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): + self.is_static_input_scheme: bool = False + self.input_qscheme = None + + if input_config is not None: + self.is_static_input_scheme = not input_config.get("is_dynamic") + self.input_qscheme = input_config.get("qscheme") + + self.input_per_token = ( + not self.is_static_input_scheme and self.input_qscheme == "per_channel" + ) + self.weight_qscheme = weight_config.get("qscheme") + self.is_weight_per_channel = self.weight_qscheme == "per_channel" + self.out_dtype = torch.get_default_dtype() + + @classmethod + def get_min_capability(cls) -> int: + # lovelace and up + return 89 + + 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_qscheme == "per_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_qscheme == "per_channel": + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + weight_quant_method = FusedMoeWeightScaleSupported.CHANNEL.value + else: + raise ValueError( + f"Unsupported weight quantization strategy: {self.weight_qscheme}." + ) + + 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.is_static_input_scheme: + assert ( + self.input_qscheme == "per_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.is_static_input_scheme: + 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_qscheme == "per_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.num_local_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 + ) + elif self.weight_qscheme == "per_channel": + layer.w13_weight_scale = torch.nn.Parameter( + layer.w13_weight_scale.unsqueeze(-1), requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + layer.w2_weight_scale.unsqueeze(-1), requires_grad=False + ) + else: + raise ValueError( + f"Unsupported weight quantization strategy: {self.weight_qscheme}." + ) + + if ( + _use_aiter + and self.is_weight_per_channel + and self.moe_runner_config.apply_router_weight_on_input + ): + with torch.no_grad(): + # Pre-shuffle weights + 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() + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + moe_runner_config = self.moe_runner_config + + if ( + _use_aiter + and self.is_weight_per_channel + and moe_runner_config.apply_router_weight_on_input + ): + topk_weights, topk_ids, _ = topk_output + output = rocm_fused_experts_tkw1( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=moe_runner_config.activation, + apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + use_fp8_w8a8=True, + per_channel_quant=self.is_weight_per_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, + ) + return StandardCombineInput(hidden_states=output) + else: + quant_info = TritonMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8_w8a8=True, + per_channel_quant=self.is_weight_per_channel, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/quark/schemes/__init__.py b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py index 91b08c512..91ceb6a1e 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py @@ -2,5 +2,6 @@ from .quark_scheme import QuarkScheme from .quark_w4a4_mxfp4 import QuarkW4A4MXFP4 +from .quark_w8a8_fp8 import QuarkW8A8Fp8 -__all__ = ["QuarkScheme", "QuarkW4A4MXFP4"] +__all__ = ["QuarkScheme", "QuarkW4A4MXFP4", "QuarkW8A8Fp8"] diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index a8322b496..ccb34f749 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py @@ -3,12 +3,17 @@ from typing import Any, Callable, Optional import torch -from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 -from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import gemm_afp4wfp4_pre_quant -from aiter.ops.triton.quant import dynamic_mxfp4_quant from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter from sglang.srt.layers.quantization.quark.schemes import QuarkScheme +from sglang.srt.utils import is_hip + +_is_hip = is_hip() +if _is_hip: + from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 + from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import gemm_afp4wfp4_pre_quant + from aiter.ops.triton.quant import dynamic_mxfp4_quant + __all__ = ["QuarkW4A4MXFP4"] diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py new file mode 100644 index 000000000..53001842a --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Callable, Optional, cast + +import torch +from torch.nn import Parameter + +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.quark.schemes import QuarkScheme +from sglang.srt.layers.quantization.utils import requantize_with_max_scale +from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs + +__all__ = ["QuarkW8A8Fp8"] + +_is_fp8_fnuz = is_fp8_fnuz() +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +if _use_aiter: + from aiter.ops.shuffle import shuffle_weight + + +class QuarkW8A8Fp8(QuarkScheme): + + def __init__( + self, weight_config: dict[str, Any], input_config: Optional[dict[str, Any]] + ): + self.cutlass_fp8_supported = cutlass_fp8_supported() + self.weight_qscheme = cast(str, weight_config.get("qscheme")) + self.is_static_input_scheme: bool = False + self.input_qscheme: Optional[str] = None + if input_config is not None: + self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic")) + self.input_qscheme = cast(str, input_config.get("qscheme")) + + self.per_token = ( + not self.is_static_input_scheme and self.input_qscheme == "per_channel" + ) + self.out_dtype = torch.get_default_dtype() + + @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.weight_qscheme == "per_tensor": + if _is_fp8_fnuz: + input_scale = getattr(layer, "input_scale", None) + weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.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: + max_w_scale = layer.weight_scale + weight = layer.weight + + max_w_scale, weight = requantize_with_max_scale( + weight=weight, + weight_scale=max_w_scale, + logical_widths=layer.logical_widths, + ) + + 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.weight_qscheme == "per_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 + if self.per_token: + weight_scale = weight_scale.view(-1, 1) + if _use_aiter: + layer.weight = Parameter( + shuffle_weight(weight, (16, 16)).t(), requires_grad=False + ) + else: + 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 scheme {self.weight_qscheme}") + + # INPUT SCALE + if self.is_static_input_scheme: + 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 + if self.weight_qscheme == "per_channel": + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + else: + assert self.weight_qscheme == "per_tensor" + weight_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + set_weight_attrs(weight_scale, {"needs_scalar_to_array": True}) + + # 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 + set_weight_attrs(input_scale, {"needs_scalar_to_array": True}) + 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( + x, + layer.weight, + layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + use_per_token_if_dynamic=self.per_token, + ) diff --git a/python/sglang/srt/layers/quantization/quark/utils.py b/python/sglang/srt/layers/quantization/quark/utils.py index eacbf3ba9..6f7d25c73 100644 --- a/python/sglang/srt/layers/quantization/quark/utils.py +++ b/python/sglang/srt/layers/quantization/quark/utils.py @@ -6,7 +6,17 @@ from types import MappingProxyType from typing import Any, Optional import torch -from aiter.ops.triton.quant import dynamic_mxfp4_quant + +try: + from aiter.ops.triton.quant import dynamic_mxfp4_quant +except ImportError as err: + + def raise_aiter_import_error(*args, **kwargs): + raise ImportError( + "Failed to import aiter. " "Make sure AITER is installed and accessible." + ) + + dynamic_mxfp4_quant = raise_aiter_import_error from torch import nn