diff --git a/docs/advanced_features/quantization.md b/docs/advanced_features/quantization.md index 5c816953b..2d244ca2e 100644 --- a/docs/advanced_features/quantization.md +++ b/docs/advanced_features/quantization.md @@ -36,7 +36,7 @@ The following table summarizes quantization method support across NVIDIA and AMD | `quark_int4fp8_moe` | No | Yes | No | AMD-only; online INT4-to-FP8 MoE quantization (CDNA3/CDNA4) | | `awq_marlin` | Yes | No | No | Marlin kernels are CUDA-only | | `gptq_marlin` | Yes | No | No | Marlin kernels are CUDA-only | -| `gguf` | Yes | No | WIP | CUDA-only kernels in sgl-kernel | +| `gguf` | Yes | No | Yes | CUDA-only kernels in sgl-kernel; Pre-dequantized on Ascend | | `modelopt` / `modelopt_fp8` | Yes (Hopper/SM90+) | No | No | [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer); requires NVIDIA hardware | | `modelopt_fp4` | Yes (Blackwell/SM100+) | No | No | [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer); native FP4 on Blackwell (B200, GB200) | | `petit_nvfp4` | No | Yes (MI250/MI300X/MI325X) | No | Enables NVFP4 on ROCm via [Petit](https://github.com/causalflow-ai/petit-kernel); use `modelopt_fp4` on NVIDIA Blackwell. Auto-selected when loading NVFP4 models on AMD. See [LMSYS blog](https://lmsys.org/blog/2025-09-21-petit-amdgpu/) and [AMD ROCm blog](https://rocm.blogs.amd.com/artificial-intelligence/fp4-mixed-precision/README.html). | diff --git a/docs/platforms/ascend/ascend_npu_quantization.md b/docs/platforms/ascend/ascend_npu_quantization.md index ca579156c..2524d2e3c 100644 --- a/docs/platforms/ascend/ascend_npu_quantization.md +++ b/docs/platforms/ascend/ascend_npu_quantization.md @@ -48,5 +48,9 @@ Compressed-tensors (LLM Compressor) on Ascend support: | [W8A8 dynamic](https://github.com/sgl-project/sglang/pull/14504) | MoE | **** | **** | **TBD** | [GGUF on Ascend support](https://github.com/sgl-project/sglang/pull/17883) +| Quantization scheme | Layer type | A2 Supported | A3 Supported | A5 Supported | +|-----------------------------------------------------------|--------------------------|:----------------------------------------:|:----------------------------------------:|:-----------------------------------------:| +| [GGUF (all types)](https://github.com/sgl-project/sglang/pull/17883) | Linear | **** | **** | **TBD** | +| [GGUF (all types)](https://github.com/sgl-project/sglang/pull/17883) | MoE | **** | **** | **TBD** | -in progress +> Note: On Ascend, GGUF weights are pre-dequantized to FP16/BF16 during model loading to ensure optimal inference performance. This enables support for all GGUF quantization types (Q2_K, Q4_K_M, IQ4_XS, etc.) while maintaining high inference speed. diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index f0baf579e..26980e7de 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -391,7 +391,11 @@ class ColumnParallelLinear(LinearBase): # Materialize GGUF UninitializedParameter if is_gguf_weight and isinstance(param, UninitializedParameter): - param.materialize(loaded_weight.shape, dtype=loaded_weight.dtype) + weight_shape = list(loaded_weight.shape) + if output_dim is not None: + weight_shape[output_dim] = weight_shape[output_dim] // self.tp_size + param.materialize(tuple(weight_shape), dtype=loaded_weight.dtype) + param_data = param.data # bitsandbytes loads the weights of the specific portion # no need to narrow here diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 9128b4ed6..f34ffc10f 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -6,6 +6,7 @@ from functools import cached_property from typing import List, Optional, Tuple import torch +from torch.nn.parameter import UninitializedParameter from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher @@ -671,6 +672,55 @@ class FusedMoE(torch.nn.Module): expert_id=expert_id, ) + def _load_gguf_weight( + self, + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: str, + expert_id: int, + tp_rank: int, + ) -> bool: + """Handle GGUF weight loading. + + Args: + param: The parameter to load the weight into. + loaded_weight: The weight tensor to load. + shard_id: The shard ID (w1, w2, or w3). + expert_id: The expert ID. + tp_rank: The tensor parallel rank. + + Returns: + True if the weight was handled as a GGUF weight, False otherwise. + """ + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + + if is_gguf_weight_type: + # Store weight type for this expert + param.weight_type = loaded_weight.item() + return True + + if is_gguf_weight: + output_dim = getattr(param, "output_dim", None) + if self.moe_tp_size > 1: + if shard_id in ["w1", "w3", "w2"] and output_dim == 0: + shard_size = loaded_weight.size(0) // self.moe_tp_size + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow( + 0, start_idx, shard_size + ).clone() + + # Store in data_container with expert/shard info + if not hasattr(param, "expert_data_map"): + param.expert_data_map = {} + + key = (expert_id, shard_id) + param.expert_data_map[key] = loaded_weight + param.data_container.append(loaded_weight) + return True + + return False + def _weight_loader_impl( self, param: torch.nn.Parameter, @@ -681,6 +731,10 @@ class FusedMoE(torch.nn.Module): ) -> None: tp_rank = self.moe_tp_rank + # Special case for GGUF weights + if self._load_gguf_weight(param, loaded_weight, shard_id, expert_id, tp_rank): + return + # compressed-tensors checkpoints with packed weights are stored flipped # TODO (mgoin): check self.quant_method.quant_config.quant_format # against known CompressionFormat enum values that have this quality @@ -1143,6 +1197,61 @@ class FusedMoE(torch.nn.Module): self.down_gemm_overlap_args = None self.meta_overlap_args = None + def materialize_gguf_weights(self) -> None: + """Process weights after loading, especially for GGUF quantization. + + This materializes GGUF UninitializedParameters from their data_containers. + """ + + for name, param in list(self.named_parameters()): + is_gguf_weight = getattr(param, "is_gguf_weight", False) + + if is_gguf_weight and isinstance(param, UninitializedParameter): + data_container = getattr(param, "data_container", []) + expert_data_map = getattr(param, "expert_data_map", {}) + tensor_shape = getattr(param, "tensor_shape", None) + + if data_container and tensor_shape: + # Determine the structure from expert_data_map + num_experts = tensor_shape[0] + + # Collect weights by expert + expert_weights = {} + for (expert_id, shard_id), weight in expert_data_map.items(): + if expert_id not in expert_weights: + expert_weights[expert_id] = {} + expert_weights[expert_id][shard_id] = weight + + # Build the full tensor + if "w13" in name: + # w13 is gate+up fused + weight_list = [] + for e in range(num_experts): + if e in expert_weights: + w1 = expert_weights[e].get("w1") + w3 = expert_weights[e].get("w3") + + if w1 is not None and w3 is not None: + fused = torch.cat([w1, w3], dim=0) + weight_list.append(fused) + + if weight_list: + stacked = torch.stack(weight_list, dim=0) + param.materialize(stacked.shape, dtype=stacked.dtype) + param.data.copy_(stacked) + elif "w2" in name: + # w2 is down projection + weight_list = [] + for e in range(num_experts): + if e in expert_weights and "w2" in expert_weights[e]: + w2_weight = expert_weights[e]["w2"] + weight_list.append(w2_weight) + + if weight_list: + stacked = torch.stack(weight_list, dim=0) + param.materialize(stacked.shape, dtype=stacked.dtype) + param.data.copy_(stacked) + @register_custom_op(out_shape="hidden_states") def moe_forward_piecewise_cuda_graph_impl( diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 3bc7b18a7..b020921c4 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -20,7 +20,7 @@ from sglang.srt.layers.quantization.base_config import ( QuantizeMethodBase, ) from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.utils import is_cuda, is_hip, is_musa, is_xpu, set_weight_attrs +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, is_xpu, set_weight_attrs if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( @@ -32,6 +32,7 @@ _is_cuda = is_cuda() _is_hip = is_hip() _is_xpu = is_xpu() _is_musa = is_musa() +_is_npu = is_npu() if _is_cuda: from sgl_kernel import moe_align_block_size, moe_sum @@ -55,9 +56,11 @@ elif _is_musa: ggml_mul_mat_a8, ggml_mul_mat_vec_a8, ) +elif _is_npu: + from gguf import dequantize as gguf_dequantize else: if not _is_hip: - warnings.warn(f"Only CUDA and MUSA support GGUF quantization currently.") + warnings.warn(f"Only CUDA, MUSA and NPU support GGUF quantization currently.") logger = logging.getLogger(__name__) @@ -107,10 +110,16 @@ class GGUFConfig(QuantizationConfig): if isinstance(layer, LinearBase): if is_layer_skipped_gguf(prefix, self.modules_to_not_convert): return UnquantizedLinearMethod() + if _is_npu: + return GGUFLinearAscendMethod(self) return GGUFLinearMethod(self) elif isinstance(layer, VocabParallelEmbedding): + if _is_npu: + return GGUFEmbeddingAscendMethod(self) return GGUFEmbeddingMethod(self) elif isinstance(layer, FusedMoE): + if _is_npu: + return GGUFMoEAscendMethod(self) return GGUFMoEMethod(self) return None @@ -575,3 +584,457 @@ class GGUFEmbeddingMethod(GGUFLinearMethod): class GGUFUninitializedParameter(UninitializedParameter): cls_to_become = Parameter data_container: list[torch.Tensor] + + +# ============================================================================= +# NPU-specific implementations for Ascend hardware +# ============================================================================= +def ggml_dequantize_ascend( + qweight: torch.Tensor, + qweight_type: int, + rows: int, + cols: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Dequantize GGML quantized weights for NPU. + + Uses gguf library's reference implementation which supports all GGML formats + and is guaranteed to be correct. The dequantization runs on CPU during model + loading, then the dequantized weights are transferred to NPU for inference. + """ + + # Move to CPU for dequantization using gguf library + qweight_cpu = qweight.cpu().numpy() + + # Use gguf library's dequantize (supports all GGML formats) + dequant_np = gguf_dequantize(qweight_cpu, qweight_type) + + # Convert to torch and move to target device + result = torch.from_numpy(dequant_np).to(dtype=dtype, device=qweight.device) + result = result.reshape(rows, cols) + + return result + + +class GGUFLinearAscendMethod(LinearMethodBase): + """Linear method for GGUF on Ascend NPU.""" + + def __init__(self, quant_config: GGUFConfig): + 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, + ): + self.params_dtype = params_dtype + output_size_per_partition = sum(output_partition_sizes) + + tensor_shape = (output_size_per_partition, input_size_per_partition) + qweight = GGUFUninitializedParameter(requires_grad=False) + set_weight_attrs( + qweight, + { + "input_dim": 1, + "output_dim": 0, + "tensor_shape": tensor_shape, + "is_gguf_weight": True, + "data_container": [], + "shard_id": [], + "shard_id_map": {}, + }, + ) + set_weight_attrs(qweight, extra_weight_attrs) + layer.register_parameter("qweight", qweight) + + qweight_type = Parameter( + torch.empty(len(output_partition_sizes), dtype=torch.uint8), + requires_grad=False, + ) + set_weight_attrs( + qweight_type, + { + "is_gguf_weight_type": True, + "weight_type": 0, + "shard_weight_type": {}, + "ignore_warning": True, + }, + ) + set_weight_attrs(qweight_type, extra_weight_attrs) + layer.register_parameter("qweight_type", qweight_type) + + def process_weights_after_loading(self, layer: torch.nn.Module): + qweight_type = layer.qweight_type.weight_type + if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES): + raise ValueError( + f"Unsupported GGUF quantization type {WeightType(qweight_type)} in layer." + ) + self._create_padded_weight_param(layer) + # Pre-dequantize weights for faster inference + self._pre_dequantize_weights(layer) + + def _create_padded_weight_param(self, layer: torch.nn.Module): + """Create padded weight parameter for GGUF MergedLinear layer.""" + qweight = layer.qweight + shard_id_map = qweight.shard_id_map + shard_id = qweight.shard_id + if len(data_container := qweight.data_container) > 1: + dtype = {data.dtype for data in data_container} + assert len(dtype) == 1 + dtype = next(iter(dtype)) + padded_side = max(x.size(1) for x in data_container) + concat_side = sum(x.size(0) for x in data_container) + padded_data = torch.zeros( + (concat_side, padded_side), dtype=dtype, device=qweight.device + ) + shard_offset_map = dict[str, tuple[int, int, int]]() + for idx in shard_id: + id_in_container = shard_id_map[idx] + start = sum(x.size(0) for x in data_container[:id_in_container]) + end = start + data_container[id_in_container].size(0) + size = data_container[id_in_container].size(1) + padded_data[start:end, :size] = data_container[id_in_container] + shard_offset_map[idx] = (start, end, size) + qweight.data_container.clear() + padded_param = Parameter(padded_data, requires_grad=False) + set_weight_attrs(padded_param, vars(qweight)) + set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map}) + layer.register_parameter("qweight", padded_param) + + def _pre_dequantize_weights(self, layer: torch.nn.Module): + """Pre-dequantize GGML weights to FP16 for faster inference. + + This eliminates runtime dequantization overhead at the cost of more memory. + """ + qweight = layer.qweight + qweight_type = layer.qweight_type.weight_type + + if qweight_type in UNQUANTIZED_TYPES and qweight.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + layer.dequantized_weight = qweight + return + + shard_id = getattr(qweight, "shard_id", None) + has_shard_offset = hasattr(qweight, "shard_offset_map") + + if shard_id and has_shard_offset: + # Handle sharded weights (QKV merged) + shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id + dequant_shards = [] + for idx in shard_id: + start, end, offset = qweight.shard_offset_map[idx] + shard_qtype = layer.qweight_type.shard_weight_type[idx] + shard_data = qweight[start:end, :offset].contiguous() + + block_size, type_size = gguf.GGML_QUANT_SIZES[shard_qtype] + shape = ( + shard_data.shape[0], + shard_data.shape[1] // type_size * block_size, + ) + dequant = ggml_dequantize_ascend( + shard_data, shard_qtype, *shape, self.params_dtype + ) + dequant_shards.append(dequant) + + dequant_weight = torch.cat(dequant_shards, dim=0) + else: + # Handle single weight + block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type] + shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size) + dequant_weight = ggml_dequantize_ascend( + qweight, qweight_type, *shape, self.params_dtype + ) + + layer.dequantized_weight = dequant_weight + + if hasattr(layer, "qweight"): + del layer.qweight + if hasattr(layer, "qweight_type"): + del layer.qweight_type + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + # Use pre-dequantized weight (always available after process_weights_after_loading) + weight = layer.dequantized_weight + out = x @ weight.T + if bias is not None: + out.add_(bias) + return out + + +class GGUFMoEAscendMethod(FusedMoEMethodBase): + """MoE method for GGUF on Ascend NPU.""" + + def __init__(self, quant_config: GGUFConfig): + 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, + ): + tensor_shape = (num_experts, 2 * intermediate_size_per_partition, hidden_size) + w13_qweight = GGUFUninitializedParameter(requires_grad=False) + set_weight_attrs( + w13_qweight, + { + "input_dim": 1, + "output_dim": 0, + "tensor_shape": tensor_shape, + "is_gguf_weight": True, + "data_container": [], + }, + ) + set_weight_attrs(w13_qweight, extra_weight_attrs) + layer.register_parameter("w13_qweight", w13_qweight) + + w13_qweight_type = Parameter( + torch.empty(1, dtype=torch.uint8), requires_grad=False + ) + set_weight_attrs( + w13_qweight_type, + {"is_gguf_weight_type": True, "weight_type": 0, "ignore_warning": True}, + ) + set_weight_attrs(w13_qweight_type, extra_weight_attrs) + layer.register_parameter("w13_qweight_type", w13_qweight_type) + + tensor_shape = (num_experts, intermediate_size_per_partition, hidden_size) + w2_qweight = GGUFUninitializedParameter(requires_grad=False) + set_weight_attrs( + w2_qweight, + { + "input_dim": 1, + "output_dim": 0, + "tensor_shape": tensor_shape, + "is_gguf_weight": True, + "data_container": [], + }, + ) + set_weight_attrs(w2_qweight, extra_weight_attrs) + layer.register_parameter("w2_qweight", w2_qweight) + + w2_qweight_type = Parameter( + torch.empty(1, dtype=torch.uint8), requires_grad=False + ) + set_weight_attrs( + w2_qweight_type, + {"is_gguf_weight_type": True, "weight_type": 0, "ignore_warning": True}, + ) + set_weight_attrs(w2_qweight_type, extra_weight_attrs) + layer.register_parameter("w2_qweight_type", w2_qweight_type) + + # Store params_dtype for pre-dequantization + self.params_dtype = params_dtype + + def process_weights_after_loading(self, layer: torch.nn.Module): + """Pre-dequantize MoE weights to FP16 for faster inference.""" + + if hasattr(layer, "materialize_gguf_weights"): + layer.materialize_gguf_weights() + + # Check if weights are actually loaded (not still UninitializedParameter/empty) + w13_qweight = layer.w13_qweight + w13_qtype = layer.w13_qweight_type.weight_type + + # Pre-dequantize w13 weights (gate+up projections) + if w13_qtype not in UNQUANTIZED_TYPES: + num_experts = w13_qweight.shape[0] + w13_dequant_list = [] + + block_size, type_size = gguf.GGML_QUANT_SIZES[w13_qtype] + + for e in range(num_experts): + qweight_cpu = w13_qweight[e].cpu().numpy() + rows = w13_qweight[e].shape[0] + cols = w13_qweight[e].shape[1] // type_size * block_size + + dequant_np = gguf_dequantize(qweight_cpu.flatten(), w13_qtype) + dequant = ( + torch.from_numpy(dequant_np) + .to(dtype=self.params_dtype, device=w13_qweight.device) + .reshape(rows, cols) + .transpose(-1, -2) + .contiguous() + ) + w13_dequant_list.append(dequant) + + w13_full = torch.stack(w13_dequant_list, dim=0) + + layer.register_buffer("w13_dequant", w13_full, persistent=False) + else: + layer.register_buffer("w13_dequant", w13_qweight.data, persistent=False) + + # Pre-dequantize w2 weights (down projection) + w2_qweight = layer.w2_qweight + w2_qtype = layer.w2_qweight_type.weight_type + + if w2_qtype not in UNQUANTIZED_TYPES: + num_experts = w2_qweight.shape[0] + w2_dequant_list = [] + + block_size, type_size = gguf.GGML_QUANT_SIZES[w2_qtype] + + for e in range(num_experts): + qweight_cpu = w2_qweight[e].cpu().numpy() + rows = w2_qweight[e].shape[0] + cols = w2_qweight[e].shape[1] // type_size * block_size + + dequant_np = gguf_dequantize(qweight_cpu.flatten(), w2_qtype) + dequant = ( + torch.from_numpy(dequant_np) + .to(dtype=self.params_dtype, device=w2_qweight.device) + .reshape(rows, cols) + .transpose(-1, -2) + .contiguous() + ) + w2_dequant_list.append(dequant) + + w2_full = torch.stack(w2_dequant_list, dim=0) + + layer.register_buffer("w2_dequant", w2_full, persistent=False) + else: + layer.register_buffer("w2_dequant", w2_qweight.data, persistent=False) + + if hasattr(layer, "w2_qweight"): + del layer.w2_qweight + if hasattr(layer, "w13_qweight"): + del layer.w13_qweight + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + """Apply MoE forward pass on NPU using npu_grouped_matmul for maximum performance.""" + from sglang.srt.distributed.communication_op import ( + tensor_model_parallel_all_gather, + ) + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + topk_weights, topk_ids, _ = topk_output + + # Check if pre-dequantized weights are available + use_pre_dequant = hasattr(layer, "w13_dequant") and hasattr(layer, "w2_dequant") + + if not use_pre_dequant: + raise RuntimeError( + "GGUF MoE on NPU requires pre-dequantization (FusedMoE fix). Please report if this occurs." + ) + + w13 = layer.w13_dequant + w2 = layer.w2_dequant + + num_experts = w13.shape[0] + + tp_size = getattr(layer, "moe_tp_size", 1) + + original_dtype = x.dtype + num_tokens = x.shape[0] + top_k = topk_ids.shape[1] + + # Ensure correct dtypes for NPU ops + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(x.dtype) + + # MoE routing initialization - reorder tokens by expert + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=x.device) + .view(top_k, -1) + .permute(1, 0) + .contiguous() + ) + + sorted_hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch.ops.npu.npu_moe_init_routing( + x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens + ) + ) + + # Compute tokens per expert + expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + expert_tokens = expert_tokens.to(torch.int64) + + w13_gmm = w13 # No transpose needed + + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[sorted_hidden_states], + weight=[w13_gmm], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + # Activation (SwiGLU) + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + + # TP all-gather for intermediate dimension if needed + if tp_size > 1: + hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1) + + w2_gmm = w2 + + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w2_gmm], + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + # Finalize routing - reorder back and apply weights + final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + ) + + if tp_size > 1: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, dim=-1 + ) + + # Ensure output matches input dtype + final_hidden_states = final_hidden_states.to(dtype=original_dtype) + + return StandardCombineInput(hidden_states=final_hidden_states) + + +class GGUFEmbeddingAscendMethod(GGUFLinearAscendMethod): + """Embedding method for GGUF on Ascend NPU.""" + + def embedding(self, layer: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + return torch.embedding(layer.dequantized_weight, x) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 5cdb3d4fc..0d90c49cc 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -2027,6 +2027,8 @@ class GGUFModelLoader(BaseModelLoader): # hack: ggufs have a different name than transformers if model_type == "cohere": model_type = "command-r" + elif model_type == "qwen3_moe": + model_type = "qwen3moe" arch = None for key, value in gguf.MODEL_ARCH_NAMES.items(): if value == model_type: diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index ddb7a3c36..dc343bfdf 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1097,21 +1097,85 @@ def gguf_quant_weights_iterator( reader = gguf.GGUFReader(gguf_file) + # MoE expert weight name patterns + MOE_WEIGHT_PATTERNS = { + "ffn_gate_exps": "gate_proj", # gate projection + "ffn_up_exps": "up_proj", # up projection + "ffn_down_exps": "down_proj", # down projection + } + + # First pass: yield weight types for tensor in reader.tensors: - if tensor.name in gguf_to_hf_name_map: - weight_type = tensor.tensor_type - name = gguf_to_hf_name_map[tensor.name] + weight_type = tensor.tensor_type + tensor_name = tensor.name + + # Check if this is a MoE expert weight (packed format) + is_moe_weight = any( + pattern in tensor_name for pattern in MOE_WEIGHT_PATTERNS.keys() + ) + + if is_moe_weight: + # MoE weights need special handling - extract layer_id and weight type + # Format: blk.{layer_id}.ffn_gate_exps.weight + import re + + match = re.match(r"blk\.(\d+)\.(ffn_\w+_exps)\.weight", tensor_name) + if match: + layer_id = int(match.group(1)) + weight_pattern = match.group(2) + hf_weight_name = MOE_WEIGHT_PATTERNS.get(weight_pattern) + + if hf_weight_name and weight_type.name != "F32": + # Yield weight type for each expert + weight = tensor.data + num_experts = weight.shape[0] + for expert_id in range(num_experts): + hf_name = f"model.layers.{layer_id}.mlp.experts.{expert_id}.{hf_weight_name}.qweight_type" + yield hf_name, torch.tensor(weight_type) + elif tensor_name in gguf_to_hf_name_map: + # Normal weight handling + name = gguf_to_hf_name_map[tensor_name] if weight_type.name != "F32": weight_type_name = name.replace("weight", "qweight_type") - weight_type = torch.tensor(weight_type) - yield weight_type_name, weight_type + yield weight_type_name, torch.tensor(weight_type) + # Second pass: yield actual weights for tensor in reader.tensors: - if tensor.name in gguf_to_hf_name_map: - weight = tensor.data - weight_type = tensor.tensor_type - name = gguf_to_hf_name_map[tensor.name] + weight = tensor.data + weight_type = tensor.tensor_type + tensor_name = tensor.name + + # Check if this is a MoE expert weight (packed format) + is_moe_weight = any( + pattern in tensor_name for pattern in MOE_WEIGHT_PATTERNS.keys() + ) + + if is_moe_weight: + # MoE weights: split packed format into individual expert weights + import re + + match = re.match(r"blk\.(\d+)\.(ffn_\w+_exps)\.weight", tensor_name) + if match: + layer_id = int(match.group(1)) + weight_pattern = match.group(2) + hf_weight_name = MOE_WEIGHT_PATTERNS.get(weight_pattern) + + if hf_weight_name: + # Packed format: [num_experts, ...] + num_experts = weight.shape[0] + for expert_id in range(num_experts): + expert_weight = weight[expert_id] + + if weight_type.name != "F32": + hf_name = f"model.layers.{layer_id}.mlp.experts.{expert_id}.{hf_weight_name}.qweight" + else: + hf_name = f"model.layers.{layer_id}.mlp.experts.{expert_id}.{hf_weight_name}.weight" + + yield hf_name, torch.tensor(expert_weight) + elif tensor_name in gguf_to_hf_name_map: + # Normal weight handling + name = gguf_to_hf_name_map[tensor_name] if weight_type.name != "F32": name = name.replace("weight", "qweight") diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index f01632ac7..755fa70ac 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -719,6 +719,7 @@ class Qwen2MoeModel(nn.Module): config.vocab_size, config.hidden_size, use_attn_tp_group=is_dp_attention_enabled(), + quant_config=quant_config, prefix=add_prefix("embed_tokens", prefix), ) else: diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index b94ba030c..8d2fe00de 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -248,9 +248,15 @@ class Qwen3MoeSparseMoeBlock(nn.Module): f"the number of experts {config.num_experts}." ) + from sglang.srt.layers.quantization.gguf import GGUFConfig + + norm_topk_prob = getattr(config, "norm_topk_prob", True) + if isinstance(quant_config, GGUFConfig): + norm_topk_prob = False + self.topk = TopK( top_k=config.num_experts_per_tok, - renormalize=config.norm_topk_prob, + renormalize=norm_topk_prob, use_grouped_topk=False, layer_id=layer_id, ) diff --git a/python/sglang/test/ascend/test_ascend_utils.py b/python/sglang/test/ascend/test_ascend_utils.py index 681299a30..90212acc5 100644 --- a/python/sglang/test/ascend/test_ascend_utils.py +++ b/python/sglang/test/ascend/test_ascend_utils.py @@ -120,12 +120,18 @@ QWEN3_235B_A22B_W8A8_WEIGHTS_PATH = os.path.join( QWEN3_30B_A3B_GPTQ_2507_INT4_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Qwen/Qwen3-30B-A3B-GPTQ-Int4" ) +QWEN3_30B_A3B_GGUF_Q4_K_M_WEIGHTS_PATH = os.path.join( + MODEL_WEIGHTS_DIR, "Qwen/Qwen3-30B-A3B-GGUF/Qwen3-30B-A3B-Q4_K_M.gguf" +) QWEN3_30B_A3B_INSTRUCT_2507_INT4_AUTOROUND_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Intel/Qwen3-30B-A3B-Instruct-2507-int4-AutoRound" ) QWEN3_30B_A3B_INSTRUCT_2507_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Qwen/Qwen3-30B-A3B-Instruct-2507" ) +QWEN3_4B_GGUF_Q4_K_M_WEIGHTS_PATH = os.path.join( + MODEL_WEIGHTS_DIR, "Qwen/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf" +) QWEN3_8B_INT4_AUTOROUND_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Intel/Qwen3-8B-int4-AutoRound" ) diff --git a/test/registered/ascend/basic_function/quant/test_npu_gguf.py b/test/registered/ascend/basic_function/quant/test_npu_gguf.py new file mode 100644 index 000000000..c02412e01 --- /dev/null +++ b/test/registered/ascend/basic_function/quant/test_npu_gguf.py @@ -0,0 +1,78 @@ +import logging +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.ascend.test_ascend_utils import QWEN3_4B_GGUF_Q4_K_M_WEIGHTS_PATH +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_npu_ci(est_time=300, suite="nightly-1-npu-a3", nightly=True) + +logger = logging.getLogger(__name__) + +TEST_MODEL_MATRIX = { + QWEN3_4B_GGUF_Q4_K_M_WEIGHTS_PATH: { + "accuracy": 0.80, + }, +} + + +class TestAscendGGUF(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.models = TEST_MODEL_MATRIX.keys() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.url = urlparse(DEFAULT_URL_FOR_TEST) + cls.common_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.8, + "--attention-backend", + "ascend", + ] + + def test_a_gsm8k(self): + for model in self.models: + with self.subTest(model=model): + logger.info(f"##=== Testing accuracy: {model} ===##") + + process = popen_launch_server( + model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *self.common_args, + ], + ) + + try: + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=128, + host=f"http://{self.url.hostname}", + port=int(self.url.port), + ) + + metrics = run_eval_few_shot_gsm8k(args) + self.assertGreaterEqual( + metrics["accuracy"], + TEST_MODEL_MATRIX[model]["accuracy"], + ) + finally: + kill_process_tree(process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/ascend/basic_function/quant/test_npu_gguf_moe.py b/test/registered/ascend/basic_function/quant/test_npu_gguf_moe.py new file mode 100644 index 000000000..25f6a0d97 --- /dev/null +++ b/test/registered/ascend/basic_function/quant/test_npu_gguf_moe.py @@ -0,0 +1,82 @@ +import logging +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.ascend.test_ascend_utils import ( + QWEN3_30B_A3B_GGUF_Q4_K_M_WEIGHTS_PATH, +) +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_npu_ci(est_time=500, suite="nightly-2-npu-a3", nightly=True) + +logger = logging.getLogger(__name__) + +TEST_MODEL_MATRIX = { + QWEN3_30B_A3B_GGUF_Q4_K_M_WEIGHTS_PATH: { + "accuracy": 0.85, + }, +} + + +class TestAscendGGUFMoE(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.models = TEST_MODEL_MATRIX.keys() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.url = urlparse(DEFAULT_URL_FOR_TEST) + cls.common_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.8, + "--attention-backend", + "ascend", + "--tensor-parallel-size", + 2, + ] + + def test_a_gsm8k(self): + for model in self.models: + with self.subTest(model=model): + logger.info(f"##=== Testing accuracy: {model} ===##") + + process = popen_launch_server( + model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *self.common_args, + ], + ) + + try: + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=128, + host=f"http://{self.url.hostname}", + port=int(self.url.port), + ) + + metrics = run_eval_few_shot_gsm8k(args) + self.assertGreaterEqual( + metrics["accuracy"], + TEST_MODEL_MATRIX[model]["accuracy"], + ) + finally: + kill_process_tree(process.pid) + + +if __name__ == "__main__": + unittest.main()