From bff489284b50a57e16c697dfe6cc36f0d2dad1c6 Mon Sep 17 00:00:00 2001 From: yuyu5333 <77156718+yuyu5333@users.noreply.github.com> Date: Fri, 17 Jul 2026 05:24:34 +0800 Subject: [PATCH] [Feature] Support DeepSeek-V4 Wint4Abf16 and Win4Afp8. (#25763) Co-authored-by: zekai --- .../csrc/gemm/per_tensor_quant_fp8.cuh | 47 ++- .../sglang/jit_kernel/per_tensor_quant_fp8.py | 32 ++ .../sglang/kernels/ops/moe/ep_moe_kernels.py | 62 ++++ .../sglang/srt/layers/moe/cutlass_w4a8_moe.py | 24 +- .../compressed_tensors/compressed_tensors.py | 52 ++- .../compressed_tensors/schemes/__init__.py | 2 + .../compressed_tensors_w4a8_fp8_moe.py | 323 ++++++++++++++++++ .../quantization/compressed_tensors/utils.py | 1 + .../deepseek_common/deepseek_weight_loader.py | 4 +- .../srt/models/deepseek_common/utils.py | 22 ++ python/sglang/srt/models/deepseek_v2.py | 5 +- python/sglang/srt/models/deepseek_v4.py | 21 +- 12 files changed, 561 insertions(+), 34 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py diff --git a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh index 17c5de888..651e18c28 100644 --- a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh @@ -94,9 +94,14 @@ __global__ void per_tensor_quant_fp8_kernel( } } -template -void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { +template +void per_tensor_quant_fp8_impl( + tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { using namespace host; + static_assert( + !(kIsStatic && kSkipQuant), + "kIsStatic+kSkipQuant=no work. Use per_tensor_absmax_fp8 for absmax-only " + "or per_tensor_quant_fp8 for static-scale quant."); auto device = SymbolicDevice{}; auto N = SymbolicSize{"num_elements"}; @@ -106,10 +111,12 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu .with_dtype() .with_device(device) .verify(input); - TensorMatcher({N}) // - .with_dtype() - .with_device(device) - .verify(output_q); + if constexpr (!kSkipQuant) { + TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(output_q); + } TensorMatcher({1}) // .with_dtype() .with_device(device) @@ -128,12 +135,28 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu static_cast(num_elements)); } - LaunchKernel(num_blocks, kBlockSize, device.unwrap())( - per_tensor_quant_fp8_kernel, - static_cast(input.data_ptr()), - static_cast(output_q.data_ptr()), - static_cast(output_s.data_ptr()), - static_cast(num_elements)); + if constexpr (!kSkipQuant) { + LaunchKernel(num_blocks, kBlockSize, device.unwrap())( + per_tensor_quant_fp8_kernel, + static_cast(input.data_ptr()), + static_cast(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + static_cast(num_elements)); + } +} + +template +void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { + per_tensor_quant_fp8_impl(input, output_q, output_s); +} + +// output_q is statically unused when kSkipQuant=true; reuse output_s as a zero-cost placeholder. +template +void per_tensor_absmax_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_s) { + per_tensor_quant_fp8_impl< + /*kIsStatic=*/false, + /*kSkipQuant=*/true, + DType>(input, output_s, output_s); } } // namespace diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py index 9225aa45d..40aee2607 100644 --- a/python/sglang/jit_kernel/per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/per_tensor_quant_fp8.py @@ -43,3 +43,35 @@ def per_tensor_quant_fp8( """ module = _jit_per_tensor_quant_fp8_module(is_static, input.dtype) module.per_tensor_quant_fp8(input.view(-1), output_q.view(-1), output_s.view(-1)) + + +@cache_once +def _jit_per_tensor_absmax_fp8_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "per_tensor_absmax_fp8", + *args, + cuda_files=["gemm/per_tensor_quant_fp8.cuh"], + cuda_wrappers=[("per_tensor_absmax_fp8", f"per_tensor_absmax_fp8<{args}>")], + ) + + +@register_custom_op( + op_name="per_tensor_absmax_fp8", + mutates_args=["output_s"], +) +def per_tensor_absmax_fp8( + input: torch.Tensor, + output_s: torch.Tensor, +) -> None: + """Compute scale = max(abs(input)) / fp8_e4m3_max via atomic-max reduction. + + The caller must zero-initialise ``output_s`` before the call (the kernel + uses ``atomic_max`` across blocks, so starting from 0 is required). + + Args: + input: Input tensor (float16, bfloat16, or float32). Any shape. + output_s: Pre-allocated float32 tensor of shape (1,), zero-initialised. + """ + module = _jit_per_tensor_absmax_fp8_module(input.dtype) + module.per_tensor_absmax_fp8(input.view(-1), output_s.view(-1)) diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 780ba0e52..cd54fc0e2 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -699,6 +699,68 @@ def silu_and_mul_masked_fwd( return output +@triton.jit +def silu_mul_dynamic_scale_triton_kernel_for_cutlass_moe( + input_ptr, + scale_ptr, + num_tokens_tensor_ptr, + intermediate_size, + fp8_max, + BLOCK_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + num_tokens = tl.load(num_tokens_tensor_ptr) + numel = num_tokens * intermediate_size + gate_ptr = input_ptr + up_ptr = input_ptr + intermediate_size + + start_idx = tl.program_id(0) * BLOCK_SIZE + step = tl.num_programs(0) * BLOCK_SIZE + absmax = 0.0 + + for id in tl.range(start_idx, numel, step, num_stages=NUM_STAGES): + ids = id + tl.arange(0, BLOCK_SIZE) + token_ids = ids // intermediate_size + mask = ids < numel + + offs = ids + token_ids * intermediate_size + gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) + up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + gate_up = gate / (1 + tl.exp(-gate)) * up + absmax = tl.maximum(absmax, tl.max(tl.abs(gate_up))) + + absmax = tl.maximum(absmax, 1e-10) + tl.atomic_max(scale_ptr, absmax / fp8_max) + + +def silu_mul_dynamic_tensorwise_quant_for_cutlass_moe( + input: torch.Tensor, + output: torch.Tensor, + scale: torch.Tensor, + num_tokens_tensor: torch.Tensor, + expected_num_tokens: int, + intermediate_size: int, +): + grid, block_dim = _get_launch_config_1d( + input.device, expected_num_tokens * intermediate_size + ) + scale.zero_() + fp8_max = torch.finfo(torch.float8_e4m3fn).max + + silu_mul_dynamic_scale_triton_kernel_for_cutlass_moe[grid]( + input_ptr=input, + scale_ptr=scale, + num_tokens_tensor_ptr=num_tokens_tensor, + intermediate_size=intermediate_size, + fp8_max=fp8_max, + BLOCK_SIZE=block_dim, + NUM_STAGES=3, + ) + silu_mul_static_tensorwise_quant_for_cutlass_moe( + input, output, scale, num_tokens_tensor, expected_num_tokens, intermediate_size + ) + + @triton.jit def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe( input_ptr, diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 4209ca094..ee6332367 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -22,7 +22,10 @@ if _is_cuda: else: from sgl_kernel import silu_and_mul -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.jit_kernel.per_tensor_quant_fp8 import ( + per_tensor_absmax_fp8, + per_tensor_quant_fp8, +) from sglang.kernels.ops.moe.ep_moe_kernels import ( cutlass_w4_run_moe_ep_preproess, deepep_ll_get_cutlass_w4a8_moe_mm_data, @@ -33,6 +36,7 @@ from sglang.kernels.ops.moe.ep_moe_kernels import ( post_reorder_for_cutlass_moe, pre_reorder_for_cutlass_moe, silu_and_mul_masked_post_per_tensor_quant_fwd, + silu_mul_dynamic_tensorwise_quant_for_cutlass_moe, silu_mul_static_tensorwise_quant_for_cutlass_moe, ) @@ -137,6 +141,11 @@ def cutlass_w4a8_moe( dtype=torch.float8_e4m3fn, ) + # TODO: fuse per_tensor_absmax_fp8 and pre_reorder_for_cutlass_moe + if a1_scale is None: + a1_scale = torch.zeros(1, dtype=torch.float32, device=device) + per_tensor_absmax_fp8(a, a1_scale) + pre_reorder_for_cutlass_moe( a, gateup_input, @@ -188,9 +197,16 @@ def cutlass_w4a8_moe( intermediate_q = torch.empty( (m * topk, n), dtype=torch.float8_e4m3fn, device=device ) - silu_mul_static_tensorwise_quant_for_cutlass_moe( - c1, intermediate_q, a2_scale.float(), expert_offsets[-1:], m * topk, n - ) + + if a2_scale is None: + a2_scale = torch.zeros(1, dtype=torch.float32, device=device) + silu_mul_dynamic_tensorwise_quant_for_cutlass_moe( + c1, intermediate_q, a2_scale, expert_offsets[-1:], m * topk, n + ) + else: + silu_mul_static_tensorwise_quant_for_cutlass_moe( + c1, intermediate_q, a2_scale.float(), expert_offsets[-1:], m * topk, n + ) cutlass_w4a8_moe_mm( c2, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 0495e951f..f9d3e2ea0 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMxInt4MoE, CompressedTensorsW4A4Fp4, CompressedTensorsW4A4Nvfp4MoE, + CompressedTensorsW4AFP8MoE, CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Fp8MoE, CompressedTensorsW8A8Int8, @@ -306,15 +307,16 @@ class CompressedTensorsConfig(QuantizationConfig): 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 + # When activation quant format is set but no + # input_activations provided: valid for w8a16fp8 (FLOAT + # weights) and pack-quantized without activation quant + # (INT weights, W4A16 style). if not input_activations: - assert ( - target_scheme_map[target]["weights"].type - == QuantizationType.FLOAT - ) + weight_type = target_scheme_map[target]["weights"].type + if weight_type == QuantizationType.INT: + pass + else: + assert weight_type == QuantizationType.FLOAT else: target_scheme_map[target]["input_activations"] = ( QuantizationArgs.model_validate( # noqa: E501 @@ -365,6 +367,33 @@ class CompressedTensorsConfig(QuantizationConfig): and is_dynamic ) + def _is_wint4afp8(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + """Detect W4AFP8: packed INT4 weights + 8-bit dynamic per-token activations.""" + if weight_quant is None or input_quant is None: + return False + return ( + self.quant_format == CompressionFormat.pack_quantized.value + and weight_quant.num_bits == 4 + and weight_quant.type == QuantizationType.INT + and weight_quant.symmetric + and not weight_quant.dynamic + and input_quant.num_bits == 8 + and input_quant.type in [QuantizationType.FLOAT, QuantizationType.INT] + and input_quant.dynamic # currently not support static input scales + ) + + def _is_wint4abf16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: + """Detect W4A16: packed INT4 weights with no activation quantization (activations stay BF16).""" + if weight_quant is None or input_quant is not None: + return False + return ( + self.quant_format == CompressionFormat.pack_quantized.value + and weight_quant.num_bits == 4 + and weight_quant.type == QuantizationType.INT + and weight_quant.symmetric + and not weight_quant.dynamic + ) + def _is_static_tensor_w8a8( self, weight_quant: BaseModel, input_quant: BaseModel ) -> bool: @@ -721,6 +750,13 @@ class CompressedTensorsConfig(QuantizationConfig): raise NotImplementedError( f"The W8A8Int8 Fused MoE scheme is implemented only for NPU for now." ) + elif self._is_wint4afp8(weight_quant, input_quant): + # On NPU prefer the dedicated NPU W4A8Int8 path when activations are INT8. + if _is_npu and self._is_dynamic_token_w4a8(weight_quant, input_quant): + logger.info_once("Using NPUCompressedTensorsW4A8Int8DynamicMoE") + return NPUCompressedTensorsW4A8Int8DynamicMoE(self) + logger.info_once("Using CompressedTensorsW4AFP8MoE") + return CompressedTensorsW4AFP8MoE(self, weight_quant, input_quant) elif self._is_dynamic_token_w4a8(weight_quant, input_quant): if _is_npu: logger.info_once("Using NPUCompressedTensorsW4A8Int8DynamicMoE") diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py index 8f67e5ba5..2ee711799 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py @@ -7,6 +7,7 @@ from .compressed_tensors_scheme import ( from .compressed_tensors_w4a4_mxint4_moe import CompressedTensorsMxInt4MoE from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4 from .compressed_tensors_w4a4_nvfp4_moe import CompressedTensorsW4A4Nvfp4MoE +from .compressed_tensors_w4a8_fp8_moe import CompressedTensorsW4AFP8MoE from .compressed_tensors_w4a8_int8_moe import NPUCompressedTensorsW4A8Int8DynamicMoE from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8 from .compressed_tensors_w8a8_fp8_moe import CompressedTensorsW8A8Fp8MoE @@ -41,4 +42,5 @@ __all__ = [ "CompressedTensorsW4A4Nvfp4MoE", "NPUCompressedTensorsW4A8Int8DynamicMoE", "CompressedTensorsMxInt4MoE", + "CompressedTensorsW4AFP8MoE", ] diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py new file mode 100644 index 000000000..ef8d56abe --- /dev/null +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py @@ -0,0 +1,323 @@ +"""W4AFP8 MoE scheme: INT4 group-quantized weights + 8-bit dynamic activations. + +Loads INT4 weights from compressed-tensors pack-quantized format, +converts to CUTLASS W4A8 layout, and runs CUTLASS grouped GEMM +with dynamic 8-bit (FP8 or INT8) activation quantization. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch +from compressed_tensors import CompressionFormat + +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsMoEScheme, +) +from sglang.srt.layers.quantization.w4afp8 import interleave_scales +from sglang.srt.utils import set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + ) + +logger = logging.getLogger(__name__) + +__all__ = ["CompressedTensorsW4AFP8MoE"] + + +def _unpack_repack_int32_to_cutlass_int8( + weight_packed: torch.Tensor, num_bits: int +) -> torch.Tensor: + """Convert compressed-tensors pack_to_int32 format to CUTLASS int8-packed format. + + pack_to_int32 stores 8 unsigned-offset int4 values per int32. + CUTLASS expects pairs of signed int4 values packed into int8 + (low nibble = even index, high nibble = odd index, two's complement). + + Args: + weight_packed: [E, N, K // pack_factor] int32 (pack_factor = 32 // num_bits) + num_bits: quantization bit width (e.g. 4) + + Returns: + [E, N, K // 2] int8 in CUTLASS layout + """ + pack_factor = 32 // num_bits + mask = (1 << num_bits) - 1 + offset = 1 << (num_bits - 1) + pair_factor = pack_factor // 2 + + # Repack directly into CUTLASS int8 without materializing full unpacked int32. + # This reduces peak memory from O(pack_factor) large int32 buffers to O(1) temporaries. + out = torch.empty( + (*weight_packed.shape[:-1], weight_packed.shape[-1], pair_factor), + dtype=torch.int8, + device=weight_packed.device, + ) + for pair_idx in range(pair_factor): + low_shift = num_bits * (2 * pair_idx) + high_shift = low_shift + num_bits + + low_nibbles = ((weight_packed >> low_shift) & mask) - offset + high_nibbles = ((weight_packed >> high_shift) & mask) - offset + out[..., pair_idx] = ((high_nibbles << 4) | (low_nibbles & 0x0F)).to(torch.int8) + + return out.flatten(-2).contiguous() + + +class CompressedTensorsW4AFP8MoE(CompressedTensorsMoEScheme): + """W4AFP8 MoE: INT4 weights (pack-quantized) + dynamic per-token 8-bit activations, + using CUTLASS W4A8 grouped GEMM kernel.""" + + def __init__( + self, + quant_config: CompressedTensorsConfig, + weight_quant, + input_quant, + ): + self.quant_config = quant_config + config = self.quant_config.target_scheme_map["Linear"].get("weights") + self.num_bits = config.num_bits + self.packed_factor = 32 // config.num_bits + self.group_size = config.group_size + self.weight_quant = weight_quant + self.input_quant = input_quant + + assert config.symmetric, "Only symmetric quantization is supported" + assert ( + self.quant_config.quant_format == CompressionFormat.pack_quantized.value + ), f"W4AFP8MoE requires pack-quantized format, got {self.quant_config.quant_format}" + + @classmethod + def get_min_capability(cls) -> int: + return 90 + + 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 + + # Weights in checkpoint (non-transposed) layout: [E, N, K // pack_factor] + # This matches the pack-quantized checkpoint format directly. + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.packed_factor, + 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, + hidden_size, + intermediate_size_per_partition // self.packed_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # Scales: [E, N, K // group_size] + num_groups_w13 = hidden_size // self.group_size + num_groups_w2 = intermediate_size_per_partition // self.group_size + + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + + w13_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + num_groups_w13, + dtype=torch.float32, + ), + 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.zeros( + num_experts, + hidden_size, + num_groups_w2, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + + # Placeholder params to accept checkpoint tensors that we don't use. + # Without these, the weight loader warns "not found in params_dict". + for name, shape in [ + ("w13_weight_shape", (num_experts, 2)), + ("w2_weight_shape", (num_experts, 2)), + ]: + p = torch.nn.Parameter( + torch.empty(shape, dtype=torch.int32), requires_grad=False + ) + layer.register_parameter(name, p) + set_weight_attrs(p, extra_weight_attrs) + + self._init_cutlass_buffers( + num_experts, + hidden_size, + intermediate_size_per_partition, + layer.w13_weight_packed.device, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Convert pack-quantized INT32 weights to CUTLASS INT8-packed format. + + Unpacks weights from compressed-tensors pack_to_int32 format and repacks + into CUTLASS int8 layout, then interleaves scales. + INT32 [E, N, K//8] → unpack signed int4 → repack INT8 [E, N, K//2] + Scales [E, N, K//gs] → interleaved bfloat16 + """ + if getattr(layer, "is_w4afp8_converted", False): + return + + dtype = torch.bfloat16 + device = layer.w2_weight_packed.device + + # TODO: currently only support per tensor quant. + layer.a13_scale = None + layer.a2_scale = None + + w13 = _unpack_repack_int32_to_cutlass_int8( + layer.w13_weight_packed.data, self.num_bits + ) + layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False) + + w2 = _unpack_repack_int32_to_cutlass_int8( + layer.w2_weight_packed.data, self.num_bits + ) + layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False) + + w13_weight_scale = layer.w13_weight_scale.to(dtype) + w13_weight_scale = interleave_scales(w13_weight_scale) + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + + w2_weight_scale = layer.w2_weight_scale.to(dtype) + w2_weight_scale = interleave_scales(w2_weight_scale) + layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) + + layer.is_w4afp8_converted = True + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def _init_cutlass_buffers( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + device: torch.device, + ): + """Pre-allocate stride and workspace tensors for CUTLASS grouped GEMM.""" + self.a_strides1 = torch.full( + (num_experts, 3), hidden_size, device=device, dtype=torch.int64 + ) + self.c_strides1 = torch.full( + (num_experts, 3), + 2 * intermediate_size, + device=device, + dtype=torch.int64, + ) + self.a_strides2 = torch.full( + (num_experts, 3), intermediate_size, device=device, dtype=torch.int64 + ) + self.c_strides2 = torch.full( + (num_experts, 3), hidden_size, device=device, dtype=torch.int64 + ) + self.b_strides1 = self.a_strides1 + self.s_strides13 = self.c_strides1 + self.b_strides2 = self.a_strides2 + self.s_strides2 = self.c_strides2 + + self.expert_offsets = torch.empty( + num_experts + 1, dtype=torch.int32, device=device + ) + self.problem_sizes1 = torch.empty( + (num_experts, 3), dtype=torch.int32, device=device + ) + self.problem_sizes2 = torch.empty( + (num_experts, 3), dtype=torch.int32, device=device + ) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + """Run the forward pass for W4AFP8 MoE.""" + return self.apply_weights(layer, dispatch_output) + + def apply_weights( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + topk_weights, topk_ids, _ = topk_output + + # TODO: currently, group_size is hardcoded to 128 in the cutlass_w4a8_moe kernel. + output = cutlass_w4a8_moe( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + layer.w13_weight_scale, + layer.w2_weight_scale, + topk_weights, + topk_ids, + self.a_strides1, + self.b_strides1, + self.c_strides1, + self.a_strides2, + self.b_strides2, + self.c_strides2, + self.s_strides13, + self.s_strides2, + self.expert_offsets, + self.problem_sizes1, + self.problem_sizes2, + layer.a13_scale, + layer.a2_scale, + routed_scaling_factor=self.moe_runner_config.routed_scaling_factor or 1.0, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py index e7d6db63d..d6c2ca3d2 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py @@ -16,6 +16,7 @@ def is_activation_quantization_format(format: str) -> bool: CompressionFormat.int_quantized.value, CompressionFormat.float_quantized.value, CompressionFormat.nvfp4_pack_quantized.value, + CompressionFormat.pack_quantized.value, ] return format in _ACTIVATION_QUANTIZATION_FORMATS diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index ad83266b7..165119e2d 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -58,6 +58,7 @@ from sglang.srt.models.deepseek_common.utils import ( _use_aiter_gfx95, awq_dequantize_func, enable_nextn_moe_bf16_cast_to_fp8, + is_wint4afp8_or_wint4a16_config, ) from sglang.srt.utils import bind_or_assign, get_bool_env_var, log_info_on_rank0 @@ -182,7 +183,7 @@ class DeepseekV2WeightLoaderMixin: # Params for special naming rules in mixed-precision models, for example: # model.layers.xx.mlp.experts.xx.w1.input_scale. For details, # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main. - if self.quant_config and self.quant_config.get_name() == "w4afp8": + if is_wint4afp8_or_wint4a16_config(self.quant_config): expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping( num_experts=self.config.n_routed_experts ) @@ -203,6 +204,7 @@ class DeepseekV2WeightLoaderMixin: futures = [] params_dict = dict(self.named_parameters()) weight_names = [] + for name, loaded_weight in weights: use_async_loading = should_async_load(loaded_weight) layer_id = get_layer_id(name) diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 8db799f1f..eef411711 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -112,6 +112,28 @@ def enable_nextn_moe_bf16_cast_to_fp8( ) +def is_wint4afp8_or_wint4a16_config( + quant_config: Optional[QuantizationConfig], +) -> bool: + if quant_config is None: + return False + if quant_config.get_name() == "w4afp8": + return True + + from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + ) + + if not isinstance(quant_config, CompressedTensorsConfig): + return False + linear_scheme = quant_config.target_scheme_map.get("Linear", {}) + weight_quant = linear_scheme.get("weights") + input_quant = linear_scheme.get("input_activations") + return quant_config._is_wint4afp8( + weight_quant, input_quant + ) or quant_config._is_wint4abf16(weight_quant, input_quant) + + def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: if scale <= 1: return 1.0 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c0087d34a..a086013be 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -179,6 +179,7 @@ from sglang.srt.models.deepseek_common.utils import ( _use_aiter, _use_aiter_bpreshuffle_gfx95, _use_aiter_gfx95, + is_wint4afp8_or_wint4a16_config, ) from sglang.srt.runtime_context import ( get_flags, @@ -2808,8 +2809,8 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) " "can use shared experts fusion optimization under expert parallelism." ) - elif self.quant_config and self.quant_config.get_name() == "w4afp8": - disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts." + elif is_wint4afp8_or_wint4a16_config(self.quant_config): + disable_reason = "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts." if disable_reason is not None: from sglang.srt.arg_groups.overrides import declare_load_time_override diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index d17daba3b..264806123 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -125,7 +125,10 @@ from sglang.srt.models.dbrx import ReplicatedLinear from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import ( try_fused_hc_post_pre, ) -from sglang.srt.models.deepseek_common.utils import _use_aiter_bpreshuffle_gfx95 +from sglang.srt.models.deepseek_common.utils import ( + _use_aiter_bpreshuffle_gfx95, + is_wint4afp8_or_wint4a16_config, +) from sglang.srt.models.deepseek_v2 import ( ParallelLMHead, _is_cuda, @@ -2743,7 +2746,7 @@ class DeepseekV4ForCausalLM(nn.Module): num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, ) - if self.quant_config and self.quant_config.get_name() == "w4afp8": + if is_wint4afp8_or_wint4a16_config(self.quant_config): expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping( num_experts=self.config.n_routed_experts ) @@ -2879,16 +2882,18 @@ class DeepseekV4ForCausalLM(nn.Module): loaded_params.add(name) break else: + skip_unmaterialized_expert_param = False for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue if _is_npu: name = name.replace("weight_packed", "weight") - name = name.replace(weight_name, param_name) - if name not in params_dict: + resolved_name = name.replace(weight_name, param_name) + if resolved_name not in params_dict: + skip_unmaterialized_expert_param = True continue - param = params_dict[name] + param = params_dict[resolved_name] weight_loader = param.weight_loader maybe_executor_submit( executor=executor, @@ -2898,16 +2903,18 @@ class DeepseekV4ForCausalLM(nn.Module): func_args=( param, loaded_weight, - name, + resolved_name, ), func_kwargs={ "shard_id": shard_id, "expert_id": expert_id, }, ) - loaded_params.add(name) + loaded_params.add(resolved_name) break else: + if skip_unmaterialized_expert_param: + continue if name.endswith(".bias") and name not in params_dict: continue if (