chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:38:16 +08:00
commit 94057c3d3e
7152 changed files with 2120455 additions and 0 deletions
@@ -0,0 +1,28 @@
# SPDX-License-Identifier: Apache-2.0
# NOTE: Import order is critical to avoid circular dependency.
# modelslim_mxfp8 imports ModelSlimLinearScheme from this package,
# so the base class must be imported first.
# isort: off
from .modelslim_scheme import ModelSlimLinearScheme, ModelSlimMoEScheme
from .modelslim_mxfp8 import ModelSlimMXFP8Scheme
from .modelslim_mxfp4_w4a8 import ModelSlimMXFP4W4A8Scheme
# isort: on
from .modelslim_w4a4_int4 import ModelSlimW4A4Int4
from .modelslim_w4a4_int4_moe import ModelSlimW4A4Int4MoE
from .modelslim_w4a8_int8_moe import ModelSlimW4A8Int8MoE
from .modelslim_w8a8_int8 import ModelSlimW8A8Int8
from .modelslim_w8a8_int8_moe import ModelSlimW8A8Int8MoE
__all__ = [
"ModelSlimLinearScheme",
"ModelSlimMoEScheme",
"ModelSlimMXFP8Scheme",
"ModelSlimMXFP4W4A8Scheme",
"ModelSlimW8A8Int8",
"ModelSlimW4A4Int4",
"ModelSlimW4A4Int4MoE",
"ModelSlimW4A8Int8MoE",
"ModelSlimW8A8Int8MoE",
]
@@ -0,0 +1,107 @@
"""ModelSlim W4A8_MXFP scheme for pre-quantized weight inference on Ascend NPU (SRT).
The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as **packed FP4**:
weight: uint8 (pack_fp4_to_uint8), shape [out, in//2], group_size=32
weight_scale: uint8 (UE8M0, +127 biased), shape [out, in//32]
(verified on ``Qwen3-8B-mxw4a8-pack-full`` and matching the msmodelslim exporter
``ascendv1.py:on_w4a8_mx_dynamic_per_block``). This is a true W4(weight) A8(activation)
scheme: weights are 4-bit FP4, activations are dynamically quantised to MXFP8.
This is NOT the same layout as ``W8A8_MXFP8`` (which stores float8_e4m3fn weights
of shape [out, in]) — so weight creation and the forward pass differ from MXFP8.
Weight post-processing and the matmul are delegated to ``NPUMXFP4W4A8OfflineLinearMethod``
(``self.kernel``), mirroring vllm-ascend's ``AscendW4A8MXFPDynamicLinearMethod``:
``npu_format_cast`` the packed FP4 to FRACTAL_NZ + transpose, then ``x2_dtype=
float4_e2m1fn_x2`` matmul with ``group_sizes=[0, 0, 32]``. Requires a recent
torch_npu for the FP4 matmul on Ascend 950/A5 (older builds reject the NZ weight) —
see ``NPUMXFP4W4A8OfflineLinearMethod`` for the version caveat.
"""
from typing import Dict, List, Optional
import torch
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUMXFP4W4A8OfflineLinearMethod,
)
from sglang.srt.layers.parameter import GroupQuantScaleParameter, ModelWeightParameter
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
# Fixed by the msmodelslim W4A8_MXFP export format (ascendv1.py sets group_size=32).
MXFP4_W4A8_BLOCK_SIZE = 32
# FP4 weights are bit-packed two-per-byte along the input (reduction) dim.
MXFP4_W4A8_PACK_FACTOR = 2
class ModelSlimMXFP4W4A8Scheme(ModelSlimLinearScheme):
"""W4A8_MXFP offline scheme — packed-FP4 weights, MXFP8 activations."""
def __init__(
self,
quant_config: Optional[Dict[str, any]] = None,
prefix: Optional[str] = None,
):
# quant_config / prefix accepted to match ModelSlimConfig.get_linear_scheme's
# dispatch signature; W4A8_MXFP needs no per-layer config beyond create_weights.
del quant_config, prefix
self.kernel = NPUMXFP4W4A8OfflineLinearMethod()
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)
# Packed-FP4 weight: uint8, shape [out, in//2] (two FP4 nibbles per byte
# along the input dim). input_dim=1 is the packed dim; TP row-parallel
# sharding narrows by self.data.shape[input_dim] (already halved), so a
# plain ModelWeightParameter shards correctly without packing metadata
# (FP4 packs the reduction dim only; the output dim stays unpacked).
weight = ModelWeightParameter(
data=torch.empty(
(
output_size_per_partition,
input_size_per_partition // MXFP4_W4A8_PACK_FACTOR,
),
dtype=torch.uint8,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
# UE8M0 block scales: uint8, shape [out, in//32]. Named "weight_scale" to
# match the checkpoint key; the kernel re-layouts it into weight_scale_inv
# during process_weights_after_loading.
scale_dim = input_size_per_partition // MXFP4_W4A8_BLOCK_SIZE
weight_scale = GroupQuantScaleParameter(
data=torch.empty(
(output_size_per_partition, scale_dim),
dtype=torch.uint8,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_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)
@@ -0,0 +1,89 @@
"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU (SRT).
Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights,
uint8 scales) and runs MXFP8 matmul at inference.
Following the modelslim-scheme convention (see ModelSlimW8A8Int8), this scheme
owns only the hardware-agnostic weight creation; weight post-processing and the
forward pass are delegated to an NPUMXFP8LinearMethod kernel (self.kernel). Its
process_weights_after_loading detects the pre-quantized float8_e4m3fn weight and
takes the offline (transpose-only) branch.
"""
from typing import Dict, List, Optional
import torch
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUMXFP8LinearMethod,
)
from sglang.srt.layers.parameter import GroupQuantScaleParameter, ModelWeightParameter
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
MXFP8_BLOCK_SIZE = 32
class ModelSlimMXFP8Scheme(ModelSlimLinearScheme):
def __init__(
self,
quant_config: Optional[Dict[str, any]] = None,
prefix: Optional[str] = None,
):
# quant_config / prefix are accepted to match the linear-scheme
# dispatch signature used by ModelSlimConfig.get_linear_scheme;
# MXFP8 needs no per-layer config beyond what create_weights derives.
del quant_config, prefix
self.kernel = NPUMXFP8LinearMethod()
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)
# msmodelslim exports weight as float8_e4m3fn, shape [out, in]
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition),
dtype=torch.float8_e4m3fn,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
# msmodelslim exports weight_scale as uint8, shape [out, in/32].
# NOTE: Named "weight_scale" (not "weight_scale_inv") to match the
# checkpoint key exported by msmodelslim; the kernel re-layouts it into
# weight_scale_inv during process_weights_after_loading.
scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE
weight_scale = GroupQuantScaleParameter(
data=torch.empty(
(output_size_per_partition, scale_dim),
dtype=torch.uint8,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_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)
@@ -0,0 +1,101 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
__all__ = ["ModelSlimLinearScheme", "ModelSlimMoEScheme"]
class ModelSlimLinearScheme(BaseLinearScheme):
"""
Abstract class used to describe the weight creation and forward pass
of different quantization schemes supported by ModelSlim.
"""
@abstractmethod
def create_weights(self, *args, **kwargs):
"""
Weight creation for the particular scheme. Inputs to this function
"""
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
@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
class ModelSlimMoEScheme(BaseMoEScheme):
"""
Abstract class used to describe the weight creation and forward pass
of different quantization schemes supported by ModelSlim.
"""
@abstractmethod
def create_weights(self, *args, **kwargs):
"""
Weight creation for the particular scheme. Inputs to this function
"""
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
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer,
dispatch_output: "StandardDispatchOutput",
):
"""
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
@@ -0,0 +1,100 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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 ModelSlimLinearScheme
from sglang.srt.utils import set_weight_attrs
class ModelSlimW4A4Int4(ModelSlimLinearScheme):
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)
@@ -0,0 +1,143 @@
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 (
NPUW4A4Int4DynamicMoEMethod,
)
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme
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,
)
logger = logging.getLogger(__name__)
__all__ = [
"ModelSlimW4A4Int4MoE",
]
class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme):
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
self.quant_config = quant_config
self.kernel = NPUW4A4Int4DynamicMoEMethod()
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_weights(
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,
):
logger.warning_once(
"Warning: Performance may be reduced, because DeepEP Dispatcher does not support 4-bit quantization, "
"switching to the bf16 dispatcher, quantization will be performed separately..."
)
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)
@@ -0,0 +1,217 @@
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,
)
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme
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,
)
logger = logging.getLogger(__name__)
__all__ = [
"ModelSlimW4A8Int8MoE",
]
class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme):
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_weights(
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,
)
@@ -0,0 +1,118 @@
# Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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 ModelSlimLinearScheme
class ModelSlimW8A8Int8(ModelSlimLinearScheme):
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)
@@ -0,0 +1,139 @@
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 (
NPUW8A8Int8DynamicMoEMethod,
)
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme
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,
)
logger = logging.getLogger(__name__)
__all__ = [
"ModelSlimW8A8Int8MoE",
]
class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme):
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_weights(
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,
)