chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU] NPU quantization refactoring & more quantization formats support (#14504)

Browse Source 
Co-authored-by: TamirBaydasov <mr.jeijy@gmail.com>
Co-authored-by: Tamir Baydasov <41994229+TamirBaydasov@users.noreply.github.com>
Co-authored-by: Савкин Артем <savkinartem@MacBook-Air-Viktoria.local>
Co-authored-by: Edward Shogulin <edward.shogulin@gmail.com>
This commit is contained in:
Артем Савкин
2026-01-14 23:25:15 +03:00
committed by GitHub
parent f091858304
commit 424a380077
30 changed files with 1962 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/srt/configs/model_config.py
View File

@@ -17,6 +17,7 @@ import logging
import math
import os
from enum import Enum, IntEnum, auto
from pathlib import Path
from typing import Any, List, Optional, Set, Union
import torch
@@ -632,6 +633,18 @@ class ModelConfig:
quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
return quant_cfg
def _find_quant_modelslim_config(self):
quant_config_file = Path(self.model_path, "quant_model_description.json")
quant_cfg = None
if quant_config_file.is_file():
with open(quant_config_file) as f:
quant_cfg = json.load(f)
# This field is required for flagless model loading but is not present in
# modelslim model description, so we're adding it here manually.
quant_cfg["quant_method"] = "modelslim"
return quant_cfg
def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]:
"""Parse ModelOpt quantization config and return the appropriate quant_method."""
json_quant_configs = quant_config_dict["quantization"]
@@ -744,6 +757,7 @@ class ModelConfig:
"w4afp8",
"petit_nvfp4",
"quark",
"modelslim",
]
compatible_quantization_methods = {
"modelopt_fp8": ["modelopt"],
@@ -755,8 +769,19 @@ class ModelConfig:
if self.quantization is not None:
self.quantization = self.quantization.lower()
# Parse quantization method from the HF model config, if available.
quant_cfg = self._parse_quant_hf_config()
# Parse quantization method from the HF and ModelSlim model config, if available.
# Only one function should return config, other should return None.
cfg_list = []
cfg_list.append(self._parse_quant_hf_config())
cfg_list.append(self._find_quant_modelslim_config())
# Filter out None values
cfg_list = [item for item in cfg_list if item is not None]
if len(cfg_list) > 1:
raise ValueError(
"Config list contains configs from 2 methods, must be only 1"
)
quant_cfg = cfg_list[0] if cfg_list else None
if quant_cfg is not None:
quant_method = quant_cfg.get(

587
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
View File

@@ -1,18 +1,17 @@
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import numpy as np
import torch
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
def npu_fused_experts(
@@ -140,79 +139,18 @@ def npu_fused_moe_without_routing_weights_bf16(
return hidden_states
class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
class _NPUFusedMoEMethodBase(FusedMoEMethodBase):
def create_weights(
def __init__(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
quant_config: Optional["QuantizationConfig"] = None,
):
self.quant_config = quant_config
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
def release_weight_cache(self, weight: torch.Tensor):
def _release_weight_cache(self, weight: torch.Tensor):
# .contiguous() introduces additional memory overhead and needs to be released using resize_(0)
origin_weight = weight.data.transpose(1, 2)
new_weight = origin_weight.contiguous()
@@ -220,10 +158,10 @@ class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
return new_weight
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_data = self.release_weight_cache(layer.w13_weight.data)
weight_data = self._release_weight_cache(layer.w13_weight.data)
layer.w13_weight = torch.nn.Parameter(weight_data, requires_grad=False)
weight_data = self.release_weight_cache(layer.w2_weight.data)
weight_data = self._release_weight_cache(layer.w2_weight.data)
layer.w2_weight = torch.nn.Parameter(weight_data, requires_grad=False)
layer.w13_weight_scale = torch.nn.Parameter(
@@ -233,21 +171,21 @@ class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
layer.w2_weight_scale = torch.nn.Parameter(
layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False
)
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
)
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
)
# Compressed-tensors format doesn't have this field
if hasattr(layer, "w13_weight_offset"):
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.squeeze(-1).contiguous(),
requires_grad=False,
)
if hasattr(layer, "w2_weight_offset"):
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.squeeze(-1).contiguous(),
requires_grad=False,
)
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
@@ -321,237 +259,19 @@ class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
return hidden_states
class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
class NPUW4A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
def __init__(self, activation_use_clip: bool) -> None:
self.group_size = 0
self.tp_size = 1
self.activation_use_clip = activation_use_clip
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.is_per_channel_weight = self.group_size == 0
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# >> weight
w13_output_size = intermediate_size_per_partition
w2_output_size = hidden_size // 2
w13_weight = torch.nn.Parameter(
torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w2_output_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# >> scale
weight_scale_dtype = torch.int64 if self.activation_use_clip else torch.float32
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
1,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=weight_scale_dtype),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# >> offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
# >>> special param for w4a8
if self.activation_use_clip:
self._init_activation_clip_params(
layer,
num_experts,
hidden_size,
intermediate_size_per_partition,
extra_weight_attrs,
)
else:
self._init_extra_scale_params(
layer,
num_experts,
hidden_size,
intermediate_size_per_partition,
extra_weight_attrs,
)
def _init_activation_clip_params(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
extra_weight_attrs: dict,
) -> None:
"""
Initializes bias and alpha parameters for quantization schemes that use activation clipping.
This helper registers `w13_bias`, `w2_bias`, and `w2_alpha`, which are required to
shift and scale the activations or outputs to compensate for the precision loss
introduced by clamping activations.
"""
w13_bias = torch.nn.Parameter(
torch.ones(
num_experts, 2 * intermediate_size_per_partition, dtype=torch.float
),
requires_grad=False,
)
layer.register_parameter("w13_bias", w13_bias)
set_weight_attrs(w13_bias, extra_weight_attrs)
w2_bias = torch.nn.Parameter(
torch.ones(num_experts, hidden_size, dtype=torch.float),
requires_grad=False,
)
layer.register_parameter("w2_bias", w2_bias)
set_weight_attrs(w2_bias, extra_weight_attrs)
w2_alpha = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float), requires_grad=False
)
layer.register_parameter("w2_alpha", w2_alpha)
set_weight_attrs(w2_alpha, extra_weight_attrs)
def _init_extra_scale_params(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
extra_weight_attrs: dict,
) -> None:
"""
Initializes additional scaling, offset, and bias parameters for quantization schemes without activation clipping.
This method registers the following parameters:
1. Scale Biases: `w13_scale_bias` and `w2_scale_bias`.
2. Secondary Quantization Params (initialized only for grouped quantization):
`w13_weight_scale_second`, `w13_weight_offset_second`,
`w2_weight_scale_second`, and `w2_weight_offset_second`.
"""
if not self.is_per_channel_weight:
w13_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second)
set_weight_attrs(w13_weight_scale_second, extra_weight_attrs)
w13_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter(
"w13_weight_offset_second", w13_weight_offset_second
)
set_weight_attrs(w13_weight_offset_second, extra_weight_attrs)
w2_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second)
set_weight_attrs(w2_weight_scale_second, extra_weight_attrs)
w2_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second)
set_weight_attrs(w2_weight_offset_second, extra_weight_attrs)
w13_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_scale_bias", w13_scale_bias)
set_weight_attrs(w13_scale_bias, extra_weight_attrs)
w2_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w2_scale_bias", w2_scale_bias)
set_weight_attrs(w2_scale_bias, extra_weight_attrs)
def process_scale(self, weight: torch.Tensor, scale, per_group_scale):
def _process_scale(
self, weight: torch.Tensor, scale, per_group_scale, is_per_channel_weight
):
scale = scale.transpose(1, 2).contiguous()
if self.is_per_channel_weight:
if is_per_channel_weight:
scale_np = scale.cpu().numpy()
scale_np.dtype = np.uint32
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
return scale_uint64_tensor, None
per_group_scale = per_group_scale.transpose(1, 2).contiguous()
group_num, k, n = weight.shape
# the weight of the new version is reduced by half by pack n, so it needs to be restored
@@ -576,7 +296,7 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
sscale_uint64_tensor = sscale_uint64_tensor.npu()
return sscale_uint64_tensor, bias
def update_bias(self, layer, w13_bias, w2_bias):
def _update_bias(self, layer, w13_bias, w2_bias):
layer.w13_scale_bias.data = (
layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
)
@@ -584,16 +304,18 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
)
def pack_to_int32(self, weight: torch.Tensor):
def _pack_to_int32(self, weight: torch.Tensor):
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
assert (
weight.shape[-1] % 4 == 0
), "the last dim of weight needs to be divided by 4"
return weight.view(torch.int32).contiguous()
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if not self.activation_use_clip:
self._process_weights_without_clip(layer)
def process_weights_after_loading(
self, layer: torch.nn.Module, is_per_channel_weight, activation_use_clip
) -> None:
if not activation_use_clip:
self._process_weights_without_clip(layer, is_per_channel_weight)
else:
self._process_weights_with_clip(layer)
@@ -607,10 +329,12 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data)
layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data)
layer.w13_weight.data = self._pack_to_int32(layer.w13_weight.data)
layer.w2_weight.data = self._pack_to_int32(layer.w2_weight.data)
def _process_weights_without_clip(self, layer: torch.nn.Module) -> None:
def _process_weights_without_clip(
self, layer: torch.nn.Module, is_per_channel_weight
) -> None:
w13_weight_scale_second = (
layer.w13_weight_scale_second.data
if hasattr(layer, "w13_weight_scale_second")
@@ -621,11 +345,17 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
if hasattr(layer, "w2_weight_scale_second")
else None
)
layer.w13_weight_scale.data, w13_bias = self.process_scale(
layer.w13_weight, layer.w13_weight_scale.data, w13_weight_scale_second
layer.w13_weight_scale.data, w13_bias = self._process_scale(
layer.w13_weight,
layer.w13_weight_scale.data,
w13_weight_scale_second,
is_per_channel_weight,
)
layer.w2_weight_scale.data, w2_bias = self.process_scale(
layer.w2_weight, layer.w2_weight_scale.data, w2_weight_scale_second
layer.w2_weight_scale.data, w2_bias = self._process_scale(
layer.w2_weight,
layer.w2_weight_scale.data,
w2_weight_scale_second,
is_per_channel_weight,
)
if hasattr(layer, "w13_weight_scale_second"):
# scale_second is no longer used, release this part of the memory
@@ -634,7 +364,7 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
del layer.w13_weight_offset_second
del layer.w2_weight_offset_second
self.update_bias(layer, w13_bias, w2_bias)
self._update_bias(layer, w13_bias, w2_bias)
def _process_weights_with_clip(self, layer: torch.nn.Module) -> None:
w13_weight_scale = (
@@ -650,21 +380,90 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
layer.w13_scale_bias = layer.w13_bias
layer.w2_scale_bias = layer.w2_bias
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
# FIXME W4A8 only support with deepep
raise NotImplementedError(
f"W4A8 only support with deepep for now, please enable --moe-a2a-backend deepep"
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
top_k = topk_ids.shape[1]
group_list_type = 1
original_shape = hidden_states.shape
topk_weights = topk_weights
num_tokens = hidden_states.shape[:-1].numel()
first_expert_idx = 0
last_expert_idx = layer.num_experts
global_num_experts = layer.num_experts
sorted_hidden_states, expanded_row_idx, expert_tokens, pertoken_scale = (
torch.ops.npu.npu_moe_init_routing_v2(
hidden_states,
topk_ids,
active_num=num_tokens * top_k,
expert_num=global_num_experts,
expert_tokens_num_type=1,
expert_tokens_num_flag=True,
active_expert_range=[first_expert_idx, last_expert_idx],
quant_mode=1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
bias1 = [layer.w13_scale_bias]
bias2 = [layer.w2_scale_bias]
w1_scale = [layer.w13_weight_scale]
w2_scale = [layer.w2_weight_scale]
_output_dtype = torch.bfloat16
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[sorted_hidden_states],
weight=[layer.w13_weight],
scale=w1_scale,
bias=bias1,
per_token_scale=[pertoken_scale],
group_list=expert_tokens,
split_item=2,
group_type=0,
group_list_type=group_list_type,
output_dtype=_output_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
output = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight],
scale=w2_scale,
bias=bias2,
per_token_scale=[swiglu_out_scale],
group_list=expert_tokens,
split_item=2,
group_type=0,
group_list_type=group_list_type,
output_dtype=_output_dtype,
)[0]
assert original_shape is not None
final_hidden_states = torch.ops.npu.npu_moe_token_unpermute(
permuted_tokens=output,
sorted_indices=torch.abs(expanded_row_idx),
probs=topk_weights,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return StandardCombineInput(hidden_states=final_hidden_states)
def apply_without_routing_weights(
self,
layer,
@@ -709,118 +508,9 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
return hidden_states
class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase):
class NPUW4A16Int4DynamicMoEMethod(_NPUFusedMoEMethodBase):
def __init__(self, quantization_config) -> None:
self.pack_factor = 8 # weight dtype is int4, but use int32 to create
target = (
"MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear"
)
if target in quantization_config.target_scheme_map:
self.group_size = quantization_config.target_scheme_map[target][
"weights"
].group_size
else:
self.group_size = 128
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
if (
extra_weight_attrs.get(
"intermediate_size_full", intermediate_size_per_partition
)
// intermediate_size_per_partition
> 1
):
quant_method = FusedMoeWeightScaleSupported.GROUP.value
else:
quant_method = FusedMoeWeightScaleSupported.CHANNEL.value
extra_weight_attrs.update({"quant_method": quant_method})
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
weight_scale_dtype = torch.bfloat16
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def pack_to_int32(self, weight: torch.Tensor):
def _pack_to_int32(self, weight: torch.Tensor):
assert weight.dim() == 3
if weight.dtype == torch.int32:
# pack 8 int4 to int32, we use a int32 to represent a int4
@@ -841,7 +531,7 @@ class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase):
raise ValueError(f"{weight.dtype=} is not supported !")
return new_weight
def unpack_from_int32(
def _unpack_from_int32(
self,
value: torch.Tensor,
num_bits: int,
@@ -926,31 +616,26 @@ class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase):
# w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous()
# w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous()
unpacked_w13_weight = (
self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4)
self._unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4)
.view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1)
.transpose(1, 2)
.contiguous()
.int()
)
unpacked_w2_weight = (
self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4)
self._unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4)
.view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1)
.transpose(1, 2)
.contiguous()
.int()
)
w13_weight = self.pack_to_int32(unpacked_w13_weight)
w2_weight = self.pack_to_int32(unpacked_w2_weight)
w13_weight = self._pack_to_int32(unpacked_w13_weight)
w2_weight = self._pack_to_int32(unpacked_w2_weight)
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,

184
python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py
View File

@@ -1,13 +1,8 @@
from typing import TYPE_CHECKING, List, Optional
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.parameter import (
ChannelQuantScaleParameter,
ModelWeightParameter,
PerTensorScaleParameter,
)
from sglang.srt.layers.quantization.base_config import LinearMethodBase
if TYPE_CHECKING:
@@ -20,82 +15,33 @@ class _NPULinearMethodBase(LinearMethodBase):
self,
quant_config: Optional["QuantizationConfig"] = None,
):
super().__init__()
self.quant_config = quant_config
class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
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,
):
weight_loader = extra_weight_attrs.get("weight_loader")
output_size_per_partition = sum(output_partition_sizes)
def process_weights_after_loading(self, layer: torch.nn.Module):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight.data = npu_format_cast(layer.weight.data)
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
layer.weight_scale.data = layer.weight_scale.data.flatten()
# Compressed-tensors format doesn't have this field
if hasattr(layer, "weight_offset"):
layer.weight_offset.data = layer.weight_offset.data.flatten()
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
expanding_factor = layer.weight.data.shape[0]
layer.aclnn_input_scale = torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.register_parameter("weight_scale", weight_scale)
weight_offset = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.register_parameter("weight_offset", weight_offset)
input_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
layer.aclnn_input_offset = torch.nn.Parameter(
layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
input_scale.ignore_warning = True
layer.register_parameter("input_scale", input_scale)
input_offset = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
)
input_offset.ignore_warning = True
layer.register_parameter("input_offset", input_offset)
quant_bias = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=torch.int32),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("quant_bias", quant_bias)
if params_dtype == torch.bfloat16:
deq_scale_dtype = torch.float32
elif params_dtype == torch.float16:
deq_scale_dtype = torch.int64
else:
raise ValueError(f"Unsupported params_dtype: {params_dtype}")
deq_scale = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=deq_scale_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("deq_scale", deq_scale)
def apply(
self,
@@ -129,66 +75,17 @@ class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
output_dtype=original_dtype,
)
class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight.data = npu_format_cast(layer.weight.data)
layer.weight_scale.data = torch.flatten(layer.weight_scale.data)
layer.weight_offset.data = torch.flatten(layer.weight_offset.data)
expanding_factor = layer.weight.data.shape[0]
layer.aclnn_input_scale = torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.aclnn_input_offset = torch.nn.Parameter(
layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
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,
):
weight_loader = extra_weight_attrs.get("weight_loader")
output_size_per_partition = sum(output_partition_sizes)
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
weight_offset = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_offset", weight_offset)
layer.weight_scale.data = layer.weight_scale.data.flatten()
# Compressed-tensors format doesn't have this field
if hasattr(layer, "weight_offset"):
layer.weight_offset.data = layer.weight_offset.data.flatten()
def apply(
self,
@@ -207,9 +104,34 @@ class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
output_dtype=original_dtype,
)
def process_weights_after_loading(self, layer: torch.nn.Module):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight.data = npu_format_cast(layer.weight.data)
class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
def process_weights_after_loading(self, layer):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight_scale.data = layer.weight_scale.data.flatten()
layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32)
layer.weight_offset.data = layer.weight_offset.data.flatten()
layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack(
layer.weight.data.to(torch.int32)
)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
tp_rank: Optional[int] = 0,
) -> torch.Tensor:
original_dtype = x.dtype
quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(
x, dst_type=torch.quint4x2
)
return torch.ops.npu.npu_quant_matmul(
quant_out,
layer.weight,
layer.weight_scale,
pertoken_scale=dynamic_scale,
bias=bias,
output_dtype=original_dtype,
)

8
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -7,9 +7,6 @@ import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A16Int4DynamicMoEMethod,
)
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.moe import (
get_deepep_mode,
@@ -27,6 +24,9 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
)
from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import (
NPUCompressedTensorsW4A16Int4DynamicMoEMethod,
)
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
@@ -374,7 +374,7 @@ class DeepEPMoE(FusedMoE):
else:
input_quant = get_bool_env_var("DEEP_NORMAL_MODE_USE_INT8_QUANT")
if not input_quant and not isinstance(
self.quant_method, NPUW4A16Int4DynamicMoEMethod
self.quant_method, NPUCompressedTensorsW4A16Int4DynamicMoEMethod
):
hidden_states, hidden_states_scale = torch_npu.npu_dynamic_quant(
hidden_states

11
python/sglang/srt/layers/quantization/__init__.py
View File

@@ -31,6 +31,7 @@ from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp8Config,
)
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config
from sglang.srt.layers.quantization.petit import PetitNvFp4Config
@@ -69,6 +70,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
"fbgemm_fp8": FBGEMMFp8Config,
"quark": QuarkConfig,
"auto-round": AutoRoundConfig,
"modelslim": ModelSlimConfig,
"quark_int4fp8_moe": QuarkInt4Fp8Config,
}
@@ -80,15 +82,6 @@ if is_cuda() or (_is_mxfp_supported and is_hip()):
}
)
if is_npu():
from sglang.srt.hardware_backend.npu.quantization.modelslim import ModelSlimConfig
BASE_QUANTIZATION_METHODS.update(
{
"modelslim": ModelSlimConfig,
}
)
QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS}

8
python/sglang/srt/layers/quantization/awq.py
View File

@@ -628,8 +628,8 @@ class AWQLinearAscendMethod(AWQLinearMethod):
qzeros_tmp = -(qzeros_tmp - 8)
qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
layer.qzeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
layer.qweight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
def apply(
self,
@@ -637,9 +637,9 @@ class AWQLinearAscendMethod(AWQLinearMethod):
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
qweight = layer.qweight
qweight = layer.weight
scales = layer.scales
qzeros = layer.qzeros
qzeros = layer.zeros
pack_factor = self.quant_config.pack_factor
out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
reshaped_x = x.reshape(-1, x.shape[-1])

4
python/sglang/srt/layers/quantization/base_config.py
View File

@@ -17,7 +17,6 @@ if TYPE_CHECKING:
class QuantizeMethodBase(ABC):
"""Base class for different quantized methods."""
@abstractmethod
def create_weights(
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
):
@@ -44,7 +43,6 @@ class QuantizeMethodBase(ABC):
class LinearMethodBase(QuantizeMethodBase):
"""Base class for different (maybe quantized) linear methods."""
@abstractmethod
def create_weights(
self,
layer: torch.nn.Module,
@@ -84,7 +82,6 @@ class LinearMethodBase(QuantizeMethodBase):
class FusedMoEMethodBase(QuantizeMethodBase):
@abstractmethod
def create_weights(
self,
layer: torch.nn.Module,
@@ -96,7 +93,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
):
raise NotImplementedError
@abstractmethod
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):

88
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
View File

@@ -45,6 +45,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsW8A8Int8,
CompressedTensorsW8A16Fp8,
CompressedTensorsWNA16,
NPUCompressedTensorsW8A8Int8,
)
from sglang.srt.layers.quantization.compressed_tensors.utils import (
find_matched_target,
@@ -53,6 +54,10 @@ from sglang.srt.layers.quantization.compressed_tensors.utils import (
)
from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.utils import is_cuda, is_npu
_is_cuda = is_cuda()
_is_npu = is_npu()
if TYPE_CHECKING:
from sglang.srt.models.utils import WeightsMapper
@@ -305,6 +310,28 @@ class CompressedTensorsConfig(QuantizationConfig):
else:
return False
def _is_dynamic_token_w4a8(
self, weight_quant: BaseModel, input_quant: BaseModel
) -> bool:
is_weight_4_bits = weight_quant.num_bits == 4
is_activation_8_bits = input_quant.num_bits == 8
weight_strategy = (
weight_quant.strategy == QuantizationStrategy.GROUP.value
or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
)
is_token = (
weight_strategy and input_quant.strategy == QuantizationStrategy.TOKEN.value
)
is_dynamic = not weight_quant.dynamic and input_quant.dynamic
return (
is_weight_4_bits
and is_activation_8_bits
and is_token
and weight_quant.symmetric
and is_dynamic
)
def _is_static_tensor_w8a8(
self, weight_quant: BaseModel, input_quant: BaseModel
) -> bool:
@@ -444,6 +471,29 @@ class CompressedTensorsConfig(QuantizationConfig):
return is_channel_group and input_quant_none and is_symmetric and is_static
def _is_dynamic_token_w4(
self, weight_quant: BaseModel, input_quant: BaseModel
) -> bool:
is_w4 = weight_quant.num_bits == 4
weight_strategy = (
weight_quant.strategy == QuantizationStrategy.TENSOR.value
or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
or weight_quant.strategy == QuantizationStrategy.GROUP.value
)
if input_quant is not None:
is_token = (
weight_strategy
and input_quant.strategy == QuantizationStrategy.TOKEN.value
)
is_dynamic = not weight_quant.dynamic and input_quant.dynamic
else:
is_token = weight_strategy
is_dynamic = not weight_quant.dynamic
# Both symmetric and asymmetric input quantization supported.
# Only symmetric weight quantization supported.
return is_w4 and weight_quant.symmetric and is_token and is_dynamic
def _get_scheme_from_parts(
self, weight_quant: BaseModel, input_quant: BaseModel
) -> CompressedTensorsScheme:
@@ -505,18 +555,32 @@ class CompressedTensorsConfig(QuantizationConfig):
)
if self._is_static_tensor_w8a8(weight_quant, input_quant):
return CompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=True,
input_symmetric=input_quant.symmetric,
)
if not _is_npu:
return CompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=True,
input_symmetric=input_quant.symmetric,
)
else:
return NPUCompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=True,
input_symmetric=input_quant.symmetric,
)
if self._is_dynamic_token_w8a8(weight_quant, input_quant):
return CompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=False,
input_symmetric=input_quant.symmetric,
)
if not _is_npu:
return CompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=False,
input_symmetric=input_quant.symmetric,
)
else:
return NPUCompressedTensorsW8A8Int8(
strategy=weight_quant.strategy,
is_static_input_scheme=False,
input_symmetric=input_quant.symmetric,
)
raise NotImplementedError("No compressed-tensors compatible scheme was found.")
@@ -594,7 +658,9 @@ class CompressedTensorsConfig(QuantizationConfig):
# Raise error if device does not support the scheme
# (e.g. fp8 needs ada lovelace)
self._check_scheme_supported(scheme.get_min_capability())
# Note: NPU devices do not support min_capability function
if not _is_npu:
self._check_scheme_supported(scheme.get_min_capability())
logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)
return scheme

569
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
View File

@@ -15,6 +15,11 @@ from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int8DynamicMoEMethod,
NPUW4A16Int4DynamicMoEMethod,
NPUW8A8Int8DynamicMoEMethod,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
@@ -43,6 +48,7 @@ from sglang.srt.utils import (
get_bool_env_var,
is_cuda,
is_hip,
is_npu,
next_power_of_2,
set_weight_attrs,
)
@@ -58,6 +64,7 @@ if TYPE_CHECKING:
)
_is_hip = is_hip()
_is_npu = is_npu()
_is_cuda = is_cuda()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -79,8 +86,11 @@ class GPTQMarlinState(Enum):
__all__ = [
"CompressedTensorsMoEMethod",
"CompressedTensorsW4A4Nvfp4MoEMethod",
"NPUCompressedTensorsW4A8Int8DynamicMoEMethod",
"CompressedTensorsW8A8Fp8MoEMethod",
"NPUCompressedTensorsW8A8Int8MoEMethod",
"CompressedTensorsWNA16MoEMethod",
"NPUCompressedTensorsW4A16Int4DynamicMoEMethod",
]
@@ -103,14 +113,40 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
input_quant = quant_config.target_scheme_map["Linear"].get("input_activations")
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
return CompressedTensorsWNA16MoEMethod(quant_config)
if not _is_npu:
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
return CompressedTensorsWNA16MoEMethod(quant_config)
else:
if (
quant_config._is_dynamic_token_w4(weight_quant, input_quant)
and input_quant is None
):
logger.info_once(
"Using NPUCompressedTensorsW4A16Int4DynamicMoEMethod"
)
return NPUCompressedTensorsW4A16Int4DynamicMoEMethod(quant_config)
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
logger.info_once("Using CompressedTensorsW4A4Nvfp4MoEMethod")
return CompressedTensorsW4A4Nvfp4MoEMethod(quant_config)
elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
logger.info_once("Using CompressedTensorsW8A8Fp8MoEMethod")
return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
if _is_npu:
logger.info_once("Using NPUCompressedTensorsW8A8Int8DynamicMoEMethod")
return NPUCompressedTensorsW8A8Int8DynamicMoEMethod(quant_config)
else:
raise NotImplementedError(
f"The W8A8Int8 Fused MoE scheme is implemented only for NPU for now."
)
elif quant_config._is_dynamic_token_w4a8(weight_quant, input_quant):
if _is_npu:
logger.info_once("Using NPUCompressedTensorsW4A8Int8DynamicMoEMethod")
return NPUCompressedTensorsW4A8Int8DynamicMoEMethod(quant_config)
else:
raise NotImplementedError(
f"The W4A8Int8 Fused MoE scheme is implemented only for NPU for now."
)
else:
raise RuntimeError(
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
@@ -853,6 +889,117 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
return self.runner.run(dispatch_output, quant_info)
class NPUCompressedTensorsW8A8Int8DynamicMoEMethod(CompressedTensorsMoEMethod):
def __init__(self, quant_config: CompressedTensorsConfig):
self.quant_config = quant_config
self.weight_quant = self.quant_config.target_scheme_map["Linear"].get("weights")
self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
"input_activations"
)
self.kernel = NPUW8A8Int8DynamicMoEMethod()
self.static_input_scales = not self.input_quant.dynamic
per_channel = (
self.weight_quant.strategy == QuantizationStrategy.CHANNEL
and self.input_quant.strategy == QuantizationStrategy.TOKEN
)
if not per_channel:
raise ValueError(
"For INT8 Fused MoE layers, we require channelwise, "
"dynamic per token quantization. Found "
f"{self.weight_quant}, {self.input_quant}"
)
self.static_input_scales = not self.input_quant.dynamic
if self.static_input_scales:
raise ValueError(
"For INT8 Fused MoE layers, we require channelwise, "
"dynamic per token quantization. Found static input scales."
)
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
params_dtype = torch.int8
# WEIGHTS
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# WEIGHT_SCALES
assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
w2_weight_scale = torch.nn.Parameter(
torch.ones(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
# Add PER-CHANNEL quantization for FusedMoE.weight_loader.
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# INPUT_SCALES
assert not self.static_input_scales
layer.w13_input_scale = None
layer.w2_input_scale = None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
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:
return self.kernel.apply(layer, dispatch_output)
class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
def __init__(self, quant_config: CompressedTensorsConfig, num_gpu_experts=-1):
@@ -1200,3 +1347,421 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
routed_scaling_factor=self.moe_runner_config.routed_scaling_factor,
)
return StandardCombineInput(hidden_states=output)
class NPUCompressedTensorsW4A8Int8DynamicMoEMethod(CompressedTensorsMoEMethod):
### TODO: Get rid of code duplication with python/sglang/srt/modelslim/modelslim_moe.py @OrangeRedeng @TamirBaydasov
def __init__(self, quantization_config) -> None:
self.group_size = 0
self.is_per_channel_weight = self.group_size == 0
self.tp_size = 1
self.activation_use_clip = (
self.quantization_config.get("config_groups", {})
.get("group_1", {})
.get("activation_use_clip", False)
)
self.kernel = NPUW4A8Int8DynamicMoEMethod()
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# >> weight
w13_output_size = intermediate_size_per_partition
w2_output_size = hidden_size // 2
w13_weight = torch.nn.Parameter(
torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w2_output_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# >> scale
weight_scale_dtype = torch.int64 if self.activation_use_clip else torch.float32
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
1,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=weight_scale_dtype),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# >> offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
# >>> special param for w4a8
if self.activation_use_clip:
self._init_activation_clip_params(
layer,
num_experts,
hidden_size,
intermediate_size_per_partition,
extra_weight_attrs,
)
else:
self._init_extra_scale_params(
layer,
num_experts,
hidden_size,
intermediate_size_per_partition,
extra_weight_attrs,
)
def _init_activation_clip_params(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
extra_weight_attrs: dict,
) -> None:
"""
Initializes bias and alpha parameters for quantization schemes that use activation clipping.
This helper registers `w13_bias`, `w2_bias`, and `w2_alpha`, which are required to
shift and scale the activations or outputs to compensate for the precision loss
introduced by clamping activations.
"""
w13_bias = torch.nn.Parameter(
torch.ones(
num_experts, 2 * intermediate_size_per_partition, dtype=torch.float
),
requires_grad=False,
)
layer.register_parameter("w13_bias", w13_bias)
set_weight_attrs(w13_bias, extra_weight_attrs)
w2_bias = torch.nn.Parameter(
torch.ones(num_experts, hidden_size, dtype=torch.float),
requires_grad=False,
)
layer.register_parameter("w2_bias", w2_bias)
set_weight_attrs(w2_bias, extra_weight_attrs)
w2_alpha = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float), requires_grad=False
)
layer.register_parameter("w2_alpha", w2_alpha)
set_weight_attrs(w2_alpha, extra_weight_attrs)
def _init_extra_scale_params(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
extra_weight_attrs: dict,
) -> None:
"""
Initializes additional scaling, offset, and bias parameters for quantization schemes without activation clipping.
This method registers the following parameters:
1. Scale Biases: `w13_scale_bias` and `w2_scale_bias`.
2. Secondary Quantization Params (initialized only for grouped quantization):
`w13_weight_scale_second`, `w13_weight_offset_second`,
`w2_weight_scale_second`, and `w2_weight_offset_second`.
"""
if not self.is_per_channel_weight:
w13_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second)
set_weight_attrs(w13_weight_scale_second, extra_weight_attrs)
w13_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter(
"w13_weight_offset_second", w13_weight_offset_second
)
set_weight_attrs(w13_weight_offset_second, extra_weight_attrs)
w2_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second)
set_weight_attrs(w2_weight_scale_second, extra_weight_attrs)
w2_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second)
set_weight_attrs(w2_weight_offset_second, extra_weight_attrs)
w13_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_scale_bias", w13_scale_bias)
set_weight_attrs(w13_scale_bias, extra_weight_attrs)
w2_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w2_scale_bias", w2_scale_bias)
set_weight_attrs(w2_scale_bias, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(
layer, self.is_per_channel_weight, self.activation_use_clip
)
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:
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)
class NPUCompressedTensorsW4A16Int4DynamicMoEMethod(CompressedTensorsMoEMethod):
def __init__(self, quantization_config) -> None:
self.pack_factor = 8 # weight dtype is int4, but use int32 to create
target = (
"MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear"
)
if target in quantization_config.target_scheme_map:
self.group_size = quantization_config.target_scheme_map[target][
"weights"
].group_size
else:
self.group_size = 128
self.kernel = NPUW4A16Int4DynamicMoEMethod()
# TODO: See if we can merge this method's logic
# with CompressedTensorsWNA16MoEMethod. Need more models and tests.
# @OrangeRedeng @TamirBaydasov
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
if (
extra_weight_attrs.get(
"intermediate_size_full", intermediate_size_per_partition
)
// intermediate_size_per_partition
> 1
):
quant_method = FusedMoeWeightScaleSupported.GROUP.value
else:
quant_method = FusedMoeWeightScaleSupported.CHANNEL.value
extra_weight_attrs.update({"quant_method": quant_method})
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
weight_scale_dtype = torch.bfloat16
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
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:
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)

6
python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py
View File

@@ -3,7 +3,10 @@
from .compressed_tensors_scheme import CompressedTensorsScheme
from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
from .compressed_tensors_w8a8_int8 import (
CompressedTensorsW8A8Int8,
NPUCompressedTensorsW8A8Int8,
)
from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8
from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS, CompressedTensorsWNA16
@@ -12,6 +15,7 @@ __all__ = [
"CompressedTensorsW8A8Fp8",
"CompressedTensorsW8A16Fp8",
"CompressedTensorsW8A8Int8",
"NPUCompressedTensorsW8A8Int8",
"CompressedTensorsWNA16",
"WNA16_SUPPORTED_BITS",
"CompressedTensorsW4A4Fp4",

138
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
View File

@@ -7,6 +7,9 @@ import torch
from compressed_tensors.quantization import QuantizationStrategy
from torch.nn import Parameter
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUW8A8Int8DynamicLinearMethod,
)
from sglang.srt.layers.parameter import (
ChannelQuantScaleParameter,
ModelWeightParameter,
@@ -33,6 +36,61 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
self.is_static_input_scheme = is_static_input_scheme
self.input_symmetric = input_symmetric
def create_weights(
self,
layer: torch.nn.Module,
output_partition_sizes: list[int],
input_size_per_partition: int,
params_dtype: torch.dtype,
weight_loader: Callable,
**kwargs,
):
output_size_per_partition = sum(output_partition_sizes)
layer.logical_widths = output_partition_sizes
# WEIGHT
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition, input_size_per_partition, dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
# WEIGHT SCALE
if self.strategy == QuantizationStrategy.CHANNEL:
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
output_dim=0,
weight_loader=weight_loader,
)
else:
assert self.strategy == QuantizationStrategy.TENSOR
weight_scale = PerTensorScaleParameter(
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
# INPUT SCALE
if self.is_static_input_scheme:
input_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
)
layer.register_parameter("input_scale", input_scale)
if not self.input_symmetric:
# Note: compressed-tensors stores the zp using the same dtype
# as the weights
# AZP loaded as int8 but used as int32
input_zero_point = PerTensorScaleParameter(
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
)
layer.register_parameter("input_zero_point", input_zero_point)
@classmethod
def get_min_capability(cls) -> int:
# ampere and up
@@ -107,61 +165,6 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
else:
layer.azp_adj = None
def create_weights(
self,
layer: torch.nn.Module,
output_partition_sizes: list[int],
input_size_per_partition: int,
params_dtype: torch.dtype,
weight_loader: Callable,
**kwargs,
):
output_size_per_partition = sum(output_partition_sizes)
layer.logical_widths = output_partition_sizes
# WEIGHT
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition, input_size_per_partition, dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
# WEIGHT SCALE
if self.strategy == QuantizationStrategy.CHANNEL:
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
output_dim=0,
weight_loader=weight_loader,
)
else:
assert self.strategy == QuantizationStrategy.TENSOR
weight_scale = PerTensorScaleParameter(
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
# INPUT SCALE
if self.is_static_input_scheme:
input_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
)
layer.register_parameter("input_scale", input_scale)
if not self.input_symmetric:
# Note: compressed-tensors stores the zp using the same dtype
# as the weights
# AZP loaded as int8 but used as int32
input_zero_point = PerTensorScaleParameter(
data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
)
layer.register_parameter("input_zero_point", input_zero_point)
def apply_weights(
self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
) -> torch.Tensor:
@@ -171,3 +174,28 @@ class CompressedTensorsW8A8Int8(CompressedTensorsScheme):
return int8_scaled_mm(
x_q, layer.weight, x_scale, layer.weight_scale, out_dtype=x.dtype, bias=bias
)
class NPUCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
def __init__(
self, strategy: str, is_static_input_scheme: bool, input_symmetric: bool
):
super().__init__(strategy, is_static_input_scheme, input_symmetric)
# TODO: Currently, NPU kernel for static quant requires quant_bias field,
# which can't be replicated in compressed-tensors.
if self.is_static_input_scheme:
raise NotImplementedError(
"Static compressed-tensors scheme is not yet supported on NPU."
)
self.kernel = NPUW8A8Int8DynamicLinearMethod()
@classmethod
def get_min_capability(cls) -> int:
return NotImplementedError
def process_weights_after_loading(self, layer):
return self.kernel.process_weights_after_loading(layer)
def apply_weights(self, layer, x, bias):
return self.kernel.apply(layer, x, bias)

14
python/sglang/srt/layers/quantization/modelslim/README.md Normal file
View File

@@ -0,0 +1,14 @@
Quantization [ModelSlim](https://gitcode.com/Ascend/msit) module.
`--quantization modelslim` flag introduced. To load already quantized models, simply load the model weights. For models quantized with ModelSlim, there's no need to add `--quantization modelslim` argument when starting the engine. The quantization method will be automatically parsed from the downloaded `quant_model_description.json` config.
ModelSlim was developed in the format of compressed_tensors and includes support for various quantization schemes, such as:
- [x] W4A4 dynamic linear
- [x] W8A8 static linear
- [x] W8A8 dynamic linear
- [x] W4A8 dynamic MOE
- [x] W8A8 dynamic MOE
Also ModelSlim module include:
- [x] Automated config detection for modelslim format (without the need to specify --quantization modelslim flag)
- [x] Unit-tests for w4a4 modelslim, w8a8 modelslim

179
python/sglang/srt/hardware_backend/npu/quantization/modelslim.py → python/sglang/srt/layers/quantization/modelslim/modelslim.py
View File

@@ -1,31 +1,30 @@
from __future__ import annotations
import logging
from types import MappingProxyType
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast
import torch
from compressed_tensors.quantization import QuantizationStrategy
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int4DynamicMoEMethod,
NPUW4A16Int4DynamicMoEMethod,
NPUW8A8Int8DynamicMoEMethod,
)
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUW8A8Int8DynamicLinearMethod,
NPUW8A8Int8LinearMethod,
_NPULinearMethodBase,
)
from sglang.srt.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsConfig,
)
from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
from sglang.srt.layers.quantization.modelslim.modelslim_moe import ModelSlimMoEMethod
from sglang.srt.layers.quantization.modelslim.schemes import (
ModelSlimScheme,
ModelSlimW4A4Int4,
ModelSlimW8A8Int8,
)
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.utils import apply_module_patch
logger = logging.getLogger(__name__)
# func refers to RMSNorm.__init__
def npu_wrapper_rmsnorm_init(func):
@@ -74,33 +73,12 @@ class ModelSlimConfig(QuantizationConfig):
def __init__(self, quant_config: Dict[str, Any] = {}):
super().__init__()
self.quant_description = quant_config
self.is_dynamic = quant_config.get("is_dynamic", False)
self.is_moe_w4_dynamic = False
ignore = cast(List[str], quant_config.get("ignore", []))
self.ignore = ignore if ignore is not None else []
packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
self.packed_modules_mapping = (
packed_modules_mapping if packed_modules_mapping is not None else {}
)
self.activation_use_clip = (
self.quant_description.get("config_groups", {})
.get("group_1", {})
.get("activation_use_clip", False)
)
self.target_scheme_map = (
CompressedTensorsConfig._quantization_scheme_map_from_config(
config=quant_config
)
)
target = "MoEGMM" if "MoEGMM" in self.target_scheme_map else "Linear"
target_scheme = self.target_scheme_map.get(target, None)
if target_scheme is None:
self.is_moe_w4_dynamic = False
else:
weight_quant = target_scheme.get("weights")
input_quant = target_scheme.get("input_activations")
self.is_moe_w4_dynamic = self.is_dynamic_token_w4(weight_quant, input_quant)
self.is_moe_input_quant = input_quant
for name in self.quant_description.keys():
if "norm.bias" in name:
@@ -115,6 +93,9 @@ class ModelSlimConfig(QuantizationConfig):
[npu_wrapper_rmsnorm_forward],
)
def get_linear_method(self) -> ModelSlimLinearMethod:
return ModelSlimLinearMethod(self)
@classmethod
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
return [torch.int8, torch.float16, torch.bfloat16]
@@ -124,7 +105,7 @@ class ModelSlimConfig(QuantizationConfig):
return 0
@classmethod
def get_name(self) -> str:
def get_name(cls) -> str:
return "modelslim"
@classmethod
@@ -163,37 +144,47 @@ class ModelSlimConfig(QuantizationConfig):
prefix_in_quant_config = prefix.replace(
proj_name, packed_modules_mapping_subset[proj_name][0]
)
self.is_dynamic = (
self.quant_description.get(prefix_in_quant_config + ".weight", "")
== "W8A8_DYNAMIC"
or self.quant_description.get("quant_method", "")
== "modelslim" # TODO: This path is for compress-tensor config，needs refactor @zhengdqin
)
if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
return UnquantizedLinearMethod()
return (
NPUW8A8Int8DynamicLinearMethod(self)
if self.is_dynamic
else NPUW8A8Int8LinearMethod(self)
)
scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config)
layer.scheme = scheme
return ModelSlimLinearMethod(self)
elif isinstance(layer, FusedMoE):
prefix_in_quant_config = prefix + ".0.down_proj.weight"
is_moe_w4a8_dynamic = (
self.quant_description.get(prefix_in_quant_config, "STATIC")
== "W4A8_DYNAMIC"
)
if (
self.is_moe_w4_dynamic and self.is_moe_input_quant is not None
) or is_moe_w4a8_dynamic:
return NPUW4A8Int4DynamicMoEMethod(
activation_use_clip=self.activation_use_clip
)
elif self.is_moe_w4_dynamic and self.is_moe_input_quant is None:
return NPUW4A16Int4DynamicMoEMethod(self)
else:
return NPUW8A8Int8DynamicMoEMethod()
return ModelSlimMoEMethod.get_moe_method(self, layer, prefix)
return None
def _get_scheme_from_parts(
self,
layer_name: str,
) -> ModelSlimScheme:
quant_type = self.quant_description.get(layer_name + ".weight", "")
if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8":
return ModelSlimW8A8Int8(
quant_config=self.quant_description, prefix=layer_name
)
elif quant_type == "W4A4_DYNAMIC":
return ModelSlimW4A4Int4(
quant_config=self.quant_description, prefix=layer_name
)
raise NotImplementedError("No modelslim compatible scheme was found.")
def get_scheme(
self, layer: torch.nn.Module, layer_name: Optional[str] = None
) -> Optional[ModelSlimScheme]:
"""
get_scheme method adjusted for modelslim, taken from
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
"""
scheme = self._get_scheme_from_parts(
layer_name=layer_name,
)
# Ascend doesn't support device capability
logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)
return scheme
def is_layer_skipped(
self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
):
@@ -228,23 +219,55 @@ class ModelSlimConfig(QuantizationConfig):
def get_scaled_act_names(self) -> List[str]:
return []
def is_dynamic_token_w4(self, weight_quant, input_quant) -> bool:
is_w4 = weight_quant.num_bits == 4
weight_strategy = (
weight_quant.strategy == QuantizationStrategy.TENSOR.value
or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
or weight_quant.strategy == QuantizationStrategy.GROUP.value
)
if input_quant is not None:
is_token = (
weight_strategy
and input_quant.strategy == QuantizationStrategy.TOKEN.value
)
is_dynamic = not weight_quant.dynamic and input_quant.dynamic
else:
is_token = weight_strategy
is_dynamic = not weight_quant.dynamic
# Both symmetric and asymmetric input quantization supported.
# Only symmetric weight quantization supported.
return is_w4 and weight_quant.symmetric and is_token and is_dynamic
class ModelSlimLinearMethod(_NPULinearMethodBase):
def __init__(self, quantization_config: ModelSlimConfig):
self.quantization_config = quantization_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.scheme.process_weights_after_loading(layer)
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,
):
"""
Use the ModelSlimScheme associated with each layer to create
the necessary parameters for the layer. See LinearMethodBase for param
details
"""
weight_loader = extra_weight_attrs.get("weight_loader")
layer.scheme.create_weights(
layer=layer,
input_size=input_size,
input_size_per_partition=input_size_per_partition,
output_partition_sizes=output_partition_sizes,
output_size=output_size,
params_dtype=params_dtype,
weight_loader=weight_loader,
)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
):
"""
Use the output of create_weights and the CompressedTensorsScheme
associated with the layer to apply the forward pass with the
layer input. See LinearMethodBase for param details
"""
scheme = layer.scheme
if scheme is None:
raise ValueError("A scheme must be defined for each layer")
return scheme.apply_weights(layer, x, bias=bias)

377
python/sglang/srt/layers/quantization/modelslim/modelslim_moe.py Normal file
View File

@@ -0,0 +1,377 @@
# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int8DynamicMoEMethod,
NPUW8A8Int8DynamicMoEMethod,
)
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
logger = logging.getLogger(__name__)
__all__ = [
"ModelSlimMoEMethod",
"ModelSlimW4A8Int8MoE",
"ModelSlimW8A8Int8MoE",
]
class ModelSlimMoEMethod(FusedMoEMethodBase):
def __new__(cls, *args, **kwargs):
if cls is ModelSlimMoEMethod:
return super().__new__(cls)
return super().__new__(cls)
@staticmethod
def get_moe_method(
quant_config: ModelSlimConfig,
layer: torch.nn.Module,
prefix: str,
) -> "ModelSlimMoEMethod":
# TODO: @dsikka: refactor this to use schemes as other kernels
# are supported + check if the layer is being ignored.
prefix_in_quant_config = prefix + ".0.down_proj.weight"
is_moe_w4a8_dynamic = (
quant_config.quant_description.get(prefix_in_quant_config, "STATIC")
== "W4A8_DYNAMIC"
)
is_moe_w8a8_dynamic = (
quant_config.quant_description.get(prefix_in_quant_config, "STATIC")
== "W8A8_DYNAMIC"
)
if is_moe_w4a8_dynamic:
logger.info_once("Using ModelSlimW4A8Int8MoE")
return ModelSlimW4A8Int8MoE(quant_config)
elif is_moe_w8a8_dynamic:
logger.info_once("Using ModelSlimW8A8Int8MoE")
return ModelSlimW8A8Int8MoE(quant_config)
else:
logger.warning(
f"Unsupported FusedMoe modelslim scheme: \
{quant_config.quant_description.get(prefix_in_quant_config.strip())} \
in layer: {prefix}"
)
return None
class ModelSlimW4A8Int8MoE(ModelSlimMoEMethod):
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
self.quant_config = quant_config
self.group_size = 0
self.is_per_channel_weight = self.group_size == 0
self.tp_size = 1
self.activation_use_clip = False
self.kernel = NPUW4A8Int8DynamicMoEMethod()
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.is_per_channel_weight = self.group_size == 0
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# >> weight
w13_output_size = intermediate_size_per_partition
w2_output_size = hidden_size // 2
w13_weight = torch.nn.Parameter(
torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w2_output_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# >> scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# >> offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
# >>> special param for w4a8
if not self.is_per_channel_weight:
w13_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second)
set_weight_attrs(w13_weight_scale_second, extra_weight_attrs)
w13_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter(
"w13_weight_offset_second", w13_weight_offset_second
)
set_weight_attrs(w13_weight_offset_second, extra_weight_attrs)
w2_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second)
set_weight_attrs(w2_weight_scale_second, extra_weight_attrs)
w2_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second)
set_weight_attrs(w2_weight_offset_second, extra_weight_attrs)
w13_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_scale_bias", w13_scale_bias)
set_weight_attrs(w13_scale_bias, extra_weight_attrs)
w2_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w2_scale_bias", w2_scale_bias)
set_weight_attrs(w2_scale_bias, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(
layer, self.is_per_channel_weight, self.activation_use_clip
)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
# FIXME W4A8 without EP can give 0 accuracy
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)
class ModelSlimW8A8Int8MoE(ModelSlimMoEMethod):
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
self.quant_config = quant_config
self.kernel = NPUW8A8Int8DynamicMoEMethod()
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,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)

11
python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py Normal file
View File

@@ -0,0 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
from .modelslim_scheme import ModelSlimScheme
from .modelslim_w4a4_int4 import ModelSlimW4A4Int4
from .modelslim_w8a8_int8 import ModelSlimW8A8Int8
__all__ = [
"ModelSlimScheme",
"ModelSlimW8A8Int8",
"ModelSlimW4A4Int4",
]

48
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py Normal file
View File

@@ -0,0 +1,48 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
from abc import ABC, abstractmethod
from typing import Optional
import torch
__all__ = ["ModelSlimScheme"]
class ModelSlimScheme(ABC):
"""
Abstract class used to describe the weight creation and forward pass
of different quantization schemes supported by CompressedTensors.
"""
@abstractmethod
def create_weights(self, *args, **kwargs):
"""
Weight creation for the particular scheme. Inputs to this function
"""
raise NotImplementedError
@abstractmethod
def apply_weights(
self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
):
"""
Run the forward pass for the particular scheme. This is where
scheme-specific dequant/quant steps/kernels should be applied.
:param layer: torch.nn.Module with the registered weights and
other parameters relevant to the particular scheme.
:param x: input to the layer
:param bias: bias parameter
"""
raise NotImplementedError
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module):
"""
Called after weight loading is complete for any cleanup that
needs to occur.
"""
raise NotImplementedError

99
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py Normal file
View File

@@ -0,0 +1,99 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Dict, List, Optional
import torch
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPU_W4A4DynamicLinearMethod,
)
from sglang.srt.layers.parameter import PerTensorScaleParameter
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimScheme
from sglang.srt.utils import set_weight_attrs
class ModelSlimW4A4Int4(ModelSlimScheme):
def __init__(
self,
quant_config: Dict[str, any],
prefix: str,
):
self.quant_config = quant_config
self.is_dynamic = self.quant_config[prefix + ".weight"] == "W4A4_DYNAMIC"
self.kernel = NPU_W4A4DynamicLinearMethod()
@staticmethod
def get_weight(
input_size: int, output_size: int, params_dtype: torch.dtype
) -> Dict[str, Any]:
params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)}
return params_dict
@staticmethod
def get_perchannel_param(
output_size: int,
params_dtype: torch.dtype,
) -> Dict[str, Any]:
params_dict = {}
params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype)
params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype)
return params_dict
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,
) -> None:
output_size_per_partition = sum(output_partition_sizes)
weight_loader = extra_weight_attrs.get("weight_loader")
weight_dict = {
"weight": torch.empty(
output_size_per_partition, input_size_per_partition, dtype=torch.int8
)
}
for weight_name, weight_param in weight_dict.items():
param = torch.nn.Parameter(weight_param, requires_grad=False)
set_weight_attrs(param, {"input_dim": 1, "output_dim": 0})
layer.register_parameter(weight_name, param)
set_weight_attrs(param, extra_weight_attrs)
pertensor_dict = {}
for pertensor_name, pertensor_param in pertensor_dict.items():
param = PerTensorScaleParameter(
data=pertensor_param, weight_loader=weight_loader
)
# disable warning
param.ignore_warning = True
layer.register_parameter(pertensor_name, param)
perchannel_dict = {}
perchannel_dict["weight_scale"] = torch.empty(
output_size_per_partition, 1, dtype=params_dtype
)
perchannel_dict["weight_offset"] = torch.empty(
output_size_per_partition, 1, dtype=params_dtype
)
for perchannel_name, perchannel_param in perchannel_dict.items():
param = torch.nn.Parameter(perchannel_param, requires_grad=False)
set_weight_attrs(param, {"output_dim": 0})
layer.register_parameter(perchannel_name, param)
set_weight_attrs(param, extra_weight_attrs)
def process_weights_after_loading(self, layer):
self.kernel.process_weights_after_loading(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
return self.kernel.apply(layer, x, bias)

117
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py Normal file
View File

@@ -0,0 +1,117 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
from typing import Dict, List, Optional
import torch
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUW8A8Int8DynamicLinearMethod,
NPUW8A8Int8LinearMethod,
)
from sglang.srt.layers.parameter import (
ChannelQuantScaleParameter,
ModelWeightParameter,
PerTensorScaleParameter,
)
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimScheme
class ModelSlimW8A8Int8(ModelSlimScheme):
def __init__(
self,
quant_config: Dict[str, any],
prefix: str,
):
self.quant_config = quant_config
self.is_dynamic = (
self.quant_config.get(prefix + ".weight", "") == "W8A8_DYNAMIC"
)
if self.is_dynamic:
self.kernel = NPUW8A8Int8DynamicLinearMethod()
else:
self.kernel = NPUW8A8Int8LinearMethod()
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,
):
weight_loader = extra_weight_attrs.get("weight_loader")
output_size_per_partition = sum(output_partition_sizes)
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
weight_offset = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_offset", weight_offset)
if not self.is_dynamic:
input_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
)
input_scale.ignore_warning = True
layer.register_parameter("input_scale", input_scale)
input_offset = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
)
input_offset.ignore_warning = True
layer.register_parameter("input_offset", input_offset)
quant_bias = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=torch.int32),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("quant_bias", quant_bias)
if params_dtype == torch.bfloat16:
deq_scale_dtype = torch.float32
elif params_dtype == torch.float16:
deq_scale_dtype = torch.int64
else:
raise ValueError(f"Unsupported params_dtype: {params_dtype}")
deq_scale = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=deq_scale_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("deq_scale", deq_scale)
def process_weights_after_loading(self, layer: torch.nn.Module):
self.kernel.process_weights_after_loading(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
return self.kernel.apply(layer, x, bias)