[NPU] Support GGUF quantization for Ascend NPU (dense + MoE) (#17883)
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
1874.
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user