from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch import torch_npu from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( npu_fused_experts, ) if TYPE_CHECKING: from sglang.srt.layers.moe import MoeRunnerConfig from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig def unpack_from_int32( weight: torch.Tensor, num_bits: int, packed_dim: int = 1, ) -> torch.Tensor: """ Unpacks quantized weights from int32 format back to original bits. :param weight: The packed int32 tensor containing quantized weights :param num_bits: The number of bits used for quantization (<= 8) :param packed_dim: Dimension along which weights are packed (0 or 1), defaults to 1 :return: Unpacked tensor with int8 dtype after applying offset correction """ assert ( weight.dtype == torch.int32 ), f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}." assert ( num_bits <= 8 ), f"Expecting `num_bits` should not be larger than 8 but got {num_bits}." pack_factor = 32 // num_bits mask = (1 << num_bits) - 1 if packed_dim == 1: unpacked_weight = torch.zeros( (weight.shape[0], weight.shape[1] * pack_factor), device=weight.device, dtype=torch.int32, ) for i in range(pack_factor): unpacked_weight[:, i::pack_factor] = (weight >> (num_bits * i)) & mask else: unpacked_weight = torch.zeros( (weight.shape[0] * pack_factor, weight.shape[1]), device=weight.device, dtype=torch.int32, ) for i in range(pack_factor): unpacked_weight[i::pack_factor, :] = (weight >> (num_bits * i)) & mask offset = pow(2, num_bits) // 2 unpacked_weight = (unpacked_weight - offset).to(torch.int8) return unpacked_weight class GPTQLinearAscendKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.qzeros = torch.nn.Parameter( unpack_from_int32( layer.qzeros.data.contiguous(), self.quant_config.weight_bits, packed_dim=1, ).to(layer.scales.dtype), requires_grad=False, ) if not self.use_v2_format: layer.qzeros += 1 qweight_tmp = unpack_from_int32( layer.qweight.data.contiguous(), self.quant_config.weight_bits, packed_dim=0 ) # use int8 to store weight by default if self.quant_config.weight_bits != 4: layer.qweight = torch.nn.Parameter( qweight_tmp, requires_grad=False, ) return # for 4bit case we need to pack 4bit weight to int32 to save memory layer.qweight = torch.nn.Parameter( torch_npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)), 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 reshaped_x = x.reshape(-1, x.shape[-1]) if bias is not None and bias.dtype == torch.bfloat16: bias = bias.float() # 4bit weight is packed to int32(8 x int4) if self.quant_config.weight_bits == 4: out_shape = x.shape[:-1] + (qweight.shape[-1] * 8,) else: out_shape = x.shape[:-1] + (qweight.shape[-1],) out = torch_npu.npu_weight_quant_batchmatmul( reshaped_x, qweight, antiquant_scale=scales, antiquant_offset=qzeros, antiquant_group_size=self.quant_config.group_size, bias=bias, ) return out.reshape(out_shape) class GPTQMoEAscendKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" self.moe_runner_config: Optional[MoeRunnerConfig] = None def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig, **extra_weight_attrs, ): self.moe_runner_config = moe_runner_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape( -1, layer.w13_qzeros.shape[-1] ) layer.w13_qzeros = torch.nn.Parameter( unpack_from_int32( w13_qzeros_2d, self.quant_config.weight_bits, packed_dim=1, ) .reshape(layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1) .to(layer.w13_scales.dtype), requires_grad=False, ) if not self.use_v2_format: layer.w13_qzeros += 1 w2_qzeros_2d = layer.w2_qzeros.data.contiguous().reshape( -1, layer.w2_qzeros.shape[-1] ) layer.w2_qzeros = torch.nn.Parameter( unpack_from_int32( w2_qzeros_2d, self.quant_config.weight_bits, packed_dim=1, ) .reshape(layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1) .to(layer.w2_scales.dtype), requires_grad=False, ) if not self.use_v2_format: layer.w2_qzeros += 1 w13_qweight_2d = ( layer.w13_qweight.data.transpose(-1, -2) .contiguous() .reshape(-1, layer.w13_qweight.shape[-2]) ) w13_qweight_tmp = unpack_from_int32( w13_qweight_2d, self.quant_config.weight_bits, packed_dim=1 ) if self.quant_config.weight_bits == 4: group_size = self.quant_config.group_size scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1) neg_mask = scale_expanded < 0 if neg_mask.any(): neg_mask = neg_mask.transpose(-1, -2) neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape) w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask] if w13_qweight_tmp.max() > 7: w13_qweight_tmp.clamp_(max=7) layer.w13_scales.data.abs_() layer.w13_qweight = torch.nn.Parameter( torch_npu.npu_convert_weight_to_int4pack( w13_qweight_tmp.reshape( layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 ) .transpose(-1, -2) .contiguous() .reshape(-1, layer.w13_qweight.shape[2]) .to(torch.int32) ) .reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1) .contiguous(), requires_grad=False, ) # use int8 to store weight by default else: layer.w13_qweight = torch.nn.Parameter( w13_qweight_tmp.reshape( layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 ) .transpose(-1, -2) .contiguous(), requires_grad=False, ) w2_qweight_2d = ( layer.w2_qweight.data.transpose(-1, -2) .contiguous() .reshape(-1, layer.w2_qweight.shape[-2]) ) w2_qweight_tmp = unpack_from_int32( w2_qweight_2d, self.quant_config.weight_bits, packed_dim=1 ) if self.quant_config.weight_bits == 4: group_size = self.quant_config.group_size scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1) neg_mask = scale_expanded < 0 if neg_mask.any(): neg_mask = neg_mask.transpose(-1, -2) neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape) w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask] if w2_qweight_tmp.max() > 7: w2_qweight_tmp.clamp_(max=7) layer.w2_scales.data.abs_() layer.w2_qweight = torch.nn.Parameter( torch_npu.npu_convert_weight_to_int4pack( w2_qweight_tmp.reshape( layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 ) .transpose(-1, -2) .contiguous() .reshape(-1, layer.w2_qweight.shape[2]) .to(torch.int32) ) .reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1) .contiguous(), requires_grad=False, ) # use int8 to store weight by default else: layer.w2_qweight = torch.nn.Parameter( w2_qweight_tmp.reshape( layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 ) .transpose(-1, -2) .contiguous(), requires_grad=False, ) def apply( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput assert self.moe_runner_config is not None, ( "moe_runner_config is not set. " "Did you forget to call create_weights/create_moe_runner?" ) assert self.moe_runner_config.activation in ("silu", "swiglu"), ( f"Only SiLU/Swiglu activation is supported, " f"got {self.moe_runner_config.activation!r}." ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output topk_weights, topk_ids, _ = topk_output topk_ids = topk_ids.to(torch.int32) topk_weights = topk_weights.to(x.dtype) output = npu_fused_experts( hidden_states=x, w13=layer.w13_qweight, w13_scale=layer.w13_scales, w13_offset=layer.w13_qzeros, w2=layer.w2_qweight, w2_scale=layer.w2_scales, w2_offset=layer.w2_qzeros, topk_weights=topk_weights, topk_ids=topk_ids, top_k=topk_ids.shape[1], use_wna16=True, ) return StandardCombineInput(hidden_states=output)