[MoE Refactor] [NPU] Refactor Ascend MoE implementation to reduce code duplication and align with community design (#25663)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Артем Савкин
2026-07-15 14:59:42 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent c9b17403e7
commit 8ed82afcc8
54 changed files with 2851 additions and 2375 deletions
+2
View File
@@ -577,6 +577,8 @@ class Envs:
# Master switch for the experimental TRT-LLM LoRA fast path; when OFF (default) every
# fine-grained opt switch reads False, keeping non-experimental paths byte-identical.
SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False)
# Enable int4x2 weights loading
SGLANG_NPU_W4A4_NEW_PACKING = EnvBool(False)
# Quantize x to int8 in the dispatch operator
DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) # This argument is deprecated
SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1)
@@ -0,0 +1,183 @@
from abc import ABC, abstractmethod
from typing import Any, Optional, Tuple
import torch
import torch.nn.functional as F
from sglang.srt.distributed.communication_op import (
tensor_model_parallel_all_gather,
)
from sglang.srt.layers.activation import GeluAndMul
from sglang.srt.runtime_context import get_parallel
# =============================================================================
# Abstract base for all activation variants
# =============================================================================
class BaseActivation(ABC):
@abstractmethod
def _apply_activation(
self, *args, **kwargs
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: ...
# =============================================================================
# Concrete activation implementations (unchanged except removed 8.)
# =============================================================================
class NPUSwiglu(BaseActivation):
def _apply_activation(self, hidden_states: torch.Tensor):
return torch.ops.npu.npu_swiglu(hidden_states), None
class NPUSwigluQuant(BaseActivation):
def _apply_activation(self, hidden_states: torch.Tensor):
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant(
hidden_states,
quant_mode=1,
activate_left=True,
)
return hidden_states, swiglu_out_scale
class NPUSwigluQuantWithScales(BaseActivation):
def _apply_activation(
self,
hidden_states: torch.Tensor,
weight_scale: torch.Tensor,
activation_scale: torch.Tensor,
group_index: torch.Tensor,
bias: Optional[torch.Tensor] = None,
quant_scale: Optional[torch.Tensor] = None,
quant_offset: Optional[torch.Tensor] = None,
):
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant(
x=hidden_states,
weight_scale=weight_scale,
activation_scale=activation_scale,
bias=bias,
quant_scale=quant_scale,
quant_offset=quant_offset,
group_index=group_index,
activate_left=True,
quant_mode=1,
)
return hidden_states, swiglu_out_scale
class NPUSwigluDeepEPKernel(BaseActivation):
def __init__(self, need_quant: bool = True):
from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant
self._kernel = swiglu_quant
self.need_quant = need_quant
def _apply_activation(
self,
hidden_states: torch.Tensor,
group_list: torch.Tensor,
group_list_type: int,
):
hidden_states, per_token_scale = self._kernel(
hidden_states, group_list, group_list_type, need_quant=self.need_quant
)
if self.need_quant:
return hidden_states, per_token_scale
return hidden_states, None
class NPUGeluAndMul(BaseActivation):
def __init__(self):
self._gelu = GeluAndMul()
def _apply_activation(self, hidden_states: torch.Tensor):
return self._gelu(hidden_states), None
class NPUSwigluOAI(BaseActivation):
def __init__(self, moe_runner_config=None):
from sgl_kernel_npu.activation.swiglu_oai import swiglu_oai_triton
self._kernel = swiglu_oai_triton
self._moe_runner_config = moe_runner_config
def _apply_activation(self, hidden_states: torch.Tensor):
# hidden_states is the output of the grouped matmul with shape
# [num_tokens, 2 * inter]. The old swiglu_oai kernel derived the
# gate_up dimension from layer.w13_weight.shape[2], which now fails
# because w13_weight is stored un-transposed. Instead we pass
# the gate_up dimension explicitly from the tensor itself.
alpha = 1.0
clamp = None
if self._moe_runner_config is not None:
alpha = getattr(self._moe_runner_config, "gemm1_alpha", 1.0)
clamp = getattr(self._moe_runner_config, "gemm1_clamp_limit", None)
output = self._kernel(
hidden_states,
hidden_states.shape[-1], # gate_up dim = 2 * inter
alpha,
clamp,
)
return output, None
class NPUSwigluStepAndMul(BaseActivation):
def __init__(self, clamp_limit: Optional[float] = None):
self._clamp_limit = clamp_limit
def _apply_activation(self, hidden_states: torch.Tensor):
if self._clamp_limit is not None:
return self._swiglustep_and_mul(hidden_states, self._clamp_limit), None
return torch.ops.npu.npu_swiglu(hidden_states), None
@staticmethod
def _swiglustep_and_mul(x: torch.Tensor, limit: float = 7.0) -> torch.Tensor:
gate, up = x.chunk(2, dim=-1)
gate = F.silu(gate).clamp(max=limit)
up = up.clamp(min=-limit, max=limit)
return gate * up
# =============================================================================
# Generic TP all‑gather wrapper – used by the runner when needed
# =============================================================================
class AllGatherActivationWrapper(BaseActivation):
"""
Wraps any activation and adds an all‑gather along `dim` if TP > 1.
This allows the runner to stay TP‑agnostic: the wrapper is applied
transparently at construction time.
"""
def __init__(self, inner: BaseActivation, dim: int = -1):
self.inner = inner
self.dim = dim
def _apply_activation(self, *args, **kwargs):
out, scale = self.inner._apply_activation(*args, **kwargs)
if get_parallel().tp_size > 1:
out = tensor_model_parallel_all_gather(out, dim=self.dim)
return out, scale
# =============================================================================
# Factory (unchanged, returns *base* activations)
# =============================================================================
def get_swiglu_variant(method: str, **kwargs: Any) -> BaseActivation:
variants: dict[str, type[BaseActivation]] = {
"standard": NPUSwiglu,
"dequant_swiglu_quant": NPUSwigluQuant,
"dequant_swiglu_quant_with_scales": NPUSwigluQuantWithScales,
"swiglu_quant_deepep_kernel": NPUSwigluDeepEPKernel,
"gelu_and_mul": NPUGeluAndMul,
}
if method == "swiglu_oai":
# The OAI variant now uses the triton kernel that derives the gate_up
# dimension from the tensor itself. No extra parameters are needed.
return NPUSwigluOAI()
if method == "swiglustep_and_mul":
clamp_limit = kwargs.pop("clamp_limit", None)
return NPUSwigluStepAndMul(clamp_limit=clamp_limit)
if method not in variants:
raise ValueError(f"Unknown SwiGLU variant: {method}")
return variants[method]()
@@ -0,0 +1,99 @@
"""
NPU MoE finalize routing components.
These classes reassemble expert outputs into the original token order
after the expert computation. A generic TP‑all‑gather wrapper is provided
to transparently gather the hidden dimension when needed (e.g. GGUF with
full weights).
"""
from abc import ABC, abstractmethod
import torch
from sglang.srt.distributed.communication_op import (
tensor_model_parallel_all_gather,
)
from sglang.srt.runtime_context import get_parallel
class BaseFinalizeRouting(ABC):
@abstractmethod
def _finalize_routing(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
expanded_row_idx: torch.Tensor,
topk_ids: torch.Tensor,
) -> torch.Tensor: ...
# ---------------------------------------------------------------------------
# Concrete implementations (unchanged)
# ---------------------------------------------------------------------------
class NPUFinalizeRouting(BaseFinalizeRouting):
def __init__(self, drop_pad_mode: int = 0):
self.drop_pad_mode = drop_pad_mode
def _finalize_routing(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
expanded_row_idx: torch.Tensor,
topk_ids: torch.Tensor,
) -> torch.Tensor:
return torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
drop_pad_mode=self.drop_pad_mode,
)
class NPUMoETokenUnpermute(BaseFinalizeRouting):
def _finalize_routing(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
expanded_row_idx: torch.Tensor,
topk_ids: torch.Tensor,
) -> torch.Tensor:
return torch.ops.npu.npu_moe_token_unpermute(
permuted_tokens=hidden_states,
sorted_indices=expanded_row_idx.abs(),
probs=topk_weights,
)
# ---------------------------------------------------------------------------
# Generic TP‑all‑gather wrapper – transparently adds communication
# ---------------------------------------------------------------------------
class AllGatherFinalizeRoutingWrapper(BaseFinalizeRouting):
"""
Wraps any finalize routing and performs an all‑gather along `dim`
after the routing if tensor‑parallelism is active.
This keeps the runner / permute hooks free of TP logic.
"""
def __init__(self, inner: BaseFinalizeRouting, dim: int = -1):
self.inner = inner
self.dim = dim
def _finalize_routing(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
expanded_row_idx: torch.Tensor,
topk_ids: torch.Tensor,
) -> torch.Tensor:
out = self.inner._finalize_routing(
hidden_states, topk_weights, expanded_row_idx, topk_ids
)
if get_parallel().tp_size > 1:
out = tensor_model_parallel_all_gather(out, dim=self.dim)
return out
@@ -14,9 +14,10 @@ import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.utils import FusedMoEMode, npu_format_cast
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
from sglang.srt.layers.moe.utils import DeepEPMode
from sglang.srt.runtime_context import get_server_args
if TYPE_CHECKING:
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
@@ -56,7 +57,7 @@ def forward_fuseep(
envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
),
num_experts=layer.num_experts,
fuse_mode=envs.SGLANG_NPU_FUSED_MOE_MODE.get(),
fuse_mode=get_server_args().fuseep_mode,
)
return hidden_states
@@ -120,52 +121,58 @@ def _scale_from_float_to_int64(scale: torch.Tensor) -> torch.nn.Parameter:
return torch.nn.Parameter(converted, requires_grad=False)
def process_fuseep_weights(layer: torch.nn.Module) -> None:
"""Apply the Ascend FuseEP-specific weight layout.
def process_fuseep_weights(layer: torch.nn.Module, weight_prefix: str) -> None:
"""Apply the Ascend FuseEP-specific weight layout for a single weight group.
Replaces NPU quant_method weight layouts with the form required by the
fused_deep_moe op. Invoked from NPU ``process_weights_after_loading``
when ``--moe-a2a-backend ascend_fuseep`` is set.
Invoked by ``maybe_apply_fuseep_weights`` for both ``"w13"`` and ``"w2"``.
"""
if envs.SGLANG_NPU_FUSED_MOE_MODE.get() == FusedMoEMode.DISPATCH_FFN_COMBINE.value:
w13_weight = _release_weight_cache(layer.w13_weight)
layer.w13_weight.data = npu_format_cast(w13_weight)
w2_weight = _release_weight_cache(layer.w2_weight)
layer.w2_weight.data = npu_format_cast(w2_weight)
if get_server_args().fuseep_mode == 1:
# -- The fused MoE optimization mode "1": dispatch_gmm_combine_decode --
if weight_prefix == "w13":
cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu()
layer.w13_weight.data = _reshape_w13_weight(cpu_w13, -1).npu()
w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous()
w13_scale = _permute_w13_weight_scale(w13_scale, 128)
layer.w13_weight_scale = torch.nn.Parameter(
w13_scale.to(torch.float32), requires_grad=False
)
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
else: # weight_prefix == "w2"
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
layer.w2_weight_scale = torch.nn.Parameter(
w2_scale.to(torch.float32), requires_grad=False
)
elif get_server_args().fuseep_mode == 2:
# -- The fused MoE optimization mode "2": dispatch_ffn_combine --
if weight_prefix == "w13":
w13_weight = _release_weight_cache(layer.w13_weight)
layer.w13_weight.data = npu_format_cast(w13_weight)
layer.w13_weight_scale.data = layer.w13_weight_scale.data.view(
layer.w13_weight_scale.data.shape[0], -1
)
layer.w13_weight_scale = _scale_from_float_to_int64(
layer.w13_weight_scale.data
)
else: # weight_prefix == "w2"
w2_weight = _release_weight_cache(layer.w2_weight)
layer.w2_weight.data = npu_format_cast(w2_weight)
w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
layer.w2_weight_scale = torch.nn.Parameter(
w2_scale.to(torch.float32), requires_grad=False
)
layer.w2_weight_scale = _scale_from_float_to_int64(
layer.w2_weight_scale.data
)
layer.w13_weight_scale.data = layer.w13_weight_scale.data.view(
layer.w13_weight_scale.data.shape[0], -1
)
w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
layer.w2_weight_scale = torch.nn.Parameter(
w2_scale.to(torch.float32), requires_grad=False
)
layer.w13_weight_scale = _scale_from_float_to_int64(layer.w13_weight_scale.data)
layer.w2_weight_scale = _scale_from_float_to_int64(layer.w2_weight_scale.data)
else:
cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu()
layer.w13_weight.data = _reshape_w13_weight(cpu_w13, -1).npu()
w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous()
w13_scale = _permute_w13_weight_scale(w13_scale, 128)
layer.w13_weight_scale = torch.nn.Parameter(
w13_scale.to(torch.float32), 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)
w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
layer.w2_weight_scale = torch.nn.Parameter(
w2_scale.to(torch.float32), requires_grad=False
)
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,
# -- offsets (exist or not, same logic for both prefixes) ---------------
offset_attr = f"{weight_prefix}_weight_offset"
if hasattr(layer, offset_attr):
setattr(
layer,
offset_attr,
torch.nn.Parameter(
getattr(layer, offset_attr).data.squeeze(-1).contiguous(),
requires_grad=False,
),
)
@@ -0,0 +1,72 @@
"""
Hidden state quantization utilities for NPU MoE.
Each class quantises hidden states and returns a (quantized_tensor, scale) tuple.
For static quantization the scale is ``None``.
"""
from abc import ABC, abstractmethod
from typing import Optional, Tuple
import torch
class BaseHiddenStatesQuant(ABC):
"""Abstract base for NPU hidden state quantisation."""
def __init__(self, quant_dtype: torch.dtype) -> None:
self.quant_dtype = quant_dtype
@abstractmethod
def __call__(
self, hidden_states: torch.Tensor, **kwargs
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: ...
class HiddenStatesDynamicQuant(BaseHiddenStatesQuant):
"""
Dynamic per‑token quantisation of hidden states.
Returns ``(quantized_hidden_states, per‑token_scale)``.
"""
def __call__(
self, hidden_states: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
quantized, scale = torch.ops.npu.npu_dynamic_quant(
hidden_states, dst_type=self.quant_dtype
)
return quantized, scale
class HiddenStatesStaticQuant(BaseHiddenStatesQuant):
"""
Static quantisation using pre‑computed layer‑specific scales and offsets.
The ``layer`` argument must expose ``aclnn_input_scale_reciprocal`` and
``aclnn_input_offset``. Returns ``(quantized_hidden_states, None)``.
"""
def __call__(
self,
hidden_states: torch.Tensor,
layer: torch.nn.Module,
) -> Tuple[torch.Tensor, None]:
# Optional defensive check (as suggested in the review)
if not hasattr(layer, "aclnn_input_scale_reciprocal") or not hasattr(
layer, "aclnn_input_offset"
):
raise AttributeError(
"Static quantisation requires layer attributes "
"'aclnn_input_scale_reciprocal' and 'aclnn_input_offset'."
)
quantized = torch.ops.npu.npu_quantize(
hidden_states,
layer.aclnn_input_scale_reciprocal,
layer.aclnn_input_offset,
self.quant_dtype,
-1,
False,
)
return quantized, None
@@ -0,0 +1,129 @@
"""
NPU MoE init routing components.
Prepare token routing before expert computation. Two API versions are provided:
- v1: legacy routing using ``npu_moe_init_routing``.
- v2: improved routing using ``npu_moe_init_routing_v2``.
"""
from abc import ABC, abstractmethod
from typing import Optional, Tuple
import torch
class BaseInitRouting(ABC):
"""Abstract base for NPU MoE init routing."""
@abstractmethod
def _init_routing(
self,
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
num_experts: int,
top_k: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: ...
class NPUMoEInitRouting_v1(BaseInitRouting):
"""
NPU MoE init routing (v1 API).
Uses ``npu_moe_init_routing`` with a manually constructed ``row_idx`` tensor.
"""
def _init_routing(
self,
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
num_experts: int,
top_k: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
num_tokens = hidden_states.shape[0]
row_idx_len = num_tokens * top_k
row_idx = (
torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_ids.device)
.view(topk_ids.shape[1], -1)
.permute(1, 0)
.contiguous()
)
hidden_states, expanded_row_idx, expanded_expert_idx = (
torch.ops.npu.npu_moe_init_routing(
hidden_states,
row_idx=row_idx,
expert_idx=topk_ids,
active_num=num_tokens,
)
)
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
)
expert_tokens = expert_tokens.to(torch.int64)
return hidden_states, expanded_row_idx, expert_tokens, None
class NPUMoEInitRouting_v2(BaseInitRouting):
"""
NPU MoE init routing (v2 API).
Uses ``npu_moe_init_routing_v2``, which integrates expert token counting.
"""
def __init__(self, quant_mode: int = -1):
self.quant_mode = quant_mode
def _init_routing(
self,
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
num_experts: int,
top_k: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
num_tokens = hidden_states.shape[0]
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=num_experts,
expert_tokens_num_type=1,
expert_tokens_num_flag=True,
active_expert_range=[0, num_experts],
quant_mode=self.quant_mode,
)
)
if self.quant_mode == -1:
pertoken_scale = None
expert_tokens = expert_tokens.to(torch.int64)
return hidden_states, expanded_row_idx, expert_tokens, pertoken_scale
class NPUMoEInitRouting_Quant(BaseInitRouting):
"""
NPU MoE init routing (Quant API).
Uses ``npu_moe_init_routing_quant``, which integrates expert token counting.
"""
def _init_routing(
self,
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
num_experts: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
num_tokens = hidden_states.shape[0]
hidden_states, expanded_row_idx, expert_tokens, _, pertoken_scale = (
torch.ops.npu.npu_moe_init_routing_quant(
hidden_states,
topk_ids,
active_num=num_tokens * topk_ids.shape[1],
expert_num=num_experts,
expert_tokens_num_mode=1,
expert_tokens_before_capacity_flag=False,
quant_mode=1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
return hidden_states, expanded_row_idx, expert_tokens, pertoken_scale
@@ -0,0 +1,49 @@
from abc import ABC, abstractmethod
import torch
class BaseMatmul(ABC):
@abstractmethod
def forward(
self,
layer: torch.nn.Module,
weight_prefix: str,
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
output_dtype: torch.dtype,
group_list_type: int,
transposed: bool,
**scale_args,
) -> torch.Tensor:
pass
class GroupedMatmul(BaseMatmul):
def forward(
self,
layer: torch.nn.Module,
weight_prefix: str,
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
output_dtype: torch.dtype,
group_list_type: int,
transposed: bool,
**scale_args,
) -> torch.Tensor:
# Access the weight attribute directly from the layer
weight = getattr(layer, f"{weight_prefix}_weight", None)
if weight is None:
raise AttributeError(
f"Weight attribute '{weight_prefix}_weight' not found in layer"
)
return torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[weight] if transposed else [weight.transpose(1, 2)],
**scale_args,
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=expert_tokens,
output_dtype=output_dtype,
)[0]
@@ -3,45 +3,139 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
import torch.nn.functional as F
import torch_npu
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A16Int4DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUWNA16Int4MoEMethod,
)
from sglang.srt.layers.quantization.utils import replace_parameter
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
import torch_npu
class AWQAscendLinearKernel:
def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Keep scales as (groups, N) – NPU kernel expects this layout
layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
qweight_tmp = torch.zeros_like(layer.qweight.data)
qzeros_tmp = layer.qzeros.data
qzeros_list = []
raw_qweight = layer.qweight.data # (K, N // pack_factor)
raw_qzeros = layer.qzeros.data # (groups, N // pack_factor)
pack_factor = self.quant_config.pack_factor
# shifts control which 4-bit nibble we extract from each packed byte:
# byte = [nibble_7 | nibble_6 | ... | nibble_0]
# shift = 4*i gives the i-th nibble's bit offset.
shifts = [0, 4, 1, 5, 2, 6, 3, 7]
for i in range(0, self.quant_config.pack_factor):
shift_num = shifts[i] * 4
qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF)
qweight_tmp.bitwise_or_(
((layer.qweight.data >> shift_num) & 0xF) << (4 * i)
K = raw_qweight.shape[0]
N = raw_qweight.shape[1] * pack_factor
num_groups = layer.scales.shape[0]
if K % num_groups != 0:
raise RuntimeError(f"K={K} not divisible by scale groups {num_groups}")
group_size = K // num_groups
# NPU fast-path constraint:
# The NPU's `npu_weight_quant_batchmatmul` kernel requires group_size
# to be a multiple of 32 and at least 32, but less than K (otherwise
# per-tensor scaling would apply, which is a different code path).
# This aligns with the NPU's SIMD vectorization width (32 elements)
# and ensures efficient memory access patterns.
is_support_npu_quant_mm = (group_size == 0) or (
group_size % 32 == 0 and 32 <= group_size < K
)
if is_support_npu_quant_mm:
# ----- NPU fast path: unsigned weight + raw zero point -----
# The NPU kernel expects:
# 1. qweight: packed unsigned 4-bit values (no XOR)
# 2. zeros: raw zero-point values (not dequantized)
#
# Step 1: Pack weight as unsigned nibbles (NO XOR).
# We extract each 4-bit nibble from the original packed tensor
# and repack them into a new tensor where each byte contains
# two 4-bit values in the order expected by the NPU kernel.
qweight_tmp = torch.zeros_like(raw_qweight)
qzeros_list = []
for i in range(pack_factor):
shift_num = shifts[i] * 4
qzeros_list.append((raw_qzeros.reshape(-1, 1) >> shift_num) & 0xF)
qweight_tmp.bitwise_or_(
((layer.qweight.data >> shift_num) & 0xF) << (4 * i)
)
# Step 2: XOR with 0x88888888 to convert from signed to unsigned
# representation. The original weights are stored as signed int4
# (values -8..7). XOR with 0x8 flips the sign bit, mapping
# -8 → 0, -7 → 1, ..., 7 → 15. This yields the unsigned
# representation the NPU kernel expects.
#
# Mathematical formula:
# unsigned_val = signed_val ^ 0x8 (for each 4-bit nibble)
# Since we pack two nibbles per byte, we XOR the whole byte
# with 0x88 to flip both sign bits simultaneously.
qweight_tmp.bitwise_xor_(
0x88888888
) # 0x88 per byte = flip sign bit of both nibbles
# Step 3: Convert zero points from signed to unsigned.
# The zero points are stored as signed int4 (-8..7).
# We convert them to unsigned (0..15) by subtracting 8,
# then negate to get the raw zero-point value expected by the NPU.
# unsigned_zero = signed_zero + 8
# raw_zero = -unsigned_zero
qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(raw_qzeros.shape[0], -1)
qzeros_tmp = -(qzeros_tmp - 8) # convert signed → unsigned → negated
qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
layer.use_npu_matmul = True
layer.npu_group_size = group_size
else:
# ----- Fallback: asymmetric dequantisation on CPU/NPU via standard linear -----
# When group_size doesn't meet the NPU constraint, we fall back to
# a standard dequantisation + FP16 linear. This is gives memory overhead but correct
# for all group_size values.
weight_u8 = torch.zeros((K, N), dtype=torch.int8, device=raw_qweight.device)
zeros_u8 = torch.zeros(
(num_groups, N), dtype=torch.int8, device=raw_qzeros.device
)
qweight_tmp.bitwise_xor_(0x88888888)
for i in range(pack_factor):
shift = shifts[i] * 4
nib_w = (raw_qweight >> shift) & 0xF
weight_u8[:, i::pack_factor] = nib_w.to(torch.int8)
nib_z = (raw_qzeros >> shift) & 0xF
zeros_u8[:, i::pack_factor] = nib_z.to(torch.int8)
qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1)
qzeros_tmp = -(qzeros_tmp - 8)
qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
# Dequantize: weight_fp = (weight_u8 - zeros) * scales
if group_size > 0:
zeros_exp = zeros_u8.repeat_interleave(group_size, dim=0)
scales_exp = layer.scales.data.repeat_interleave(group_size, dim=0)
else:
zeros_exp = zeros_u8
scales_exp = layer.scales.data
layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
weight_float = (weight_u8.float() - zeros_exp.float()) * scales_exp.float()
weight_float = weight_float.t().contiguous().to(torch.bfloat16)
layer.register_parameter(
"weight", torch.nn.Parameter(weight_float, requires_grad=False)
)
delattr(layer, "scales")
layer.use_npu_matmul = False
# Clean original packed tensors to free memory
for attr in ("qweight", "qzeros"):
if hasattr(layer, attr):
delattr(layer, attr)
def apply(
self,
@@ -49,32 +143,41 @@ class AWQAscendLinearKernel:
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
qweight = layer.weight
scales = layer.scales
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])
pack_factor = self.quant_config.pack_factor
if bias is not None and bias.dtype == torch.bfloat16:
bias = bias.float()
if layer.use_npu_matmul:
qweight = layer.weight # (K, N//pack) int32, unsigned
scales = layer.scales # (groups, N)
offset = layer.zeros # (groups, N) raw zero point
out = torch_npu.npu_weight_quant_batchmatmul(
reshaped_x,
qweight,
antiquant_scale=scales,
antiquant_offset=qzeros,
antiquant_group_size=self.quant_config.group_size,
bias=bias,
)
out_shape = x.shape[:-1] + (qweight.shape[1] * pack_factor,)
if bias is not None and bias.dtype == torch.bfloat16:
bias = bias.float()
return out.reshape(out_shape)
# NPU-accelerated quantized matmul.
# The kernel internally does:
# out = (x @ qweight_dequantized) + bias
# where qweight_dequantized = (qweight_unsigned - offset) * scales
# with group-wise scaling applied.
out = torch_npu.npu_weight_quant_batchmatmul(
reshaped_x,
qweight,
antiquant_scale=scales,
antiquant_offset=offset, # raw zero point
antiquant_group_size=layer.npu_group_size,
bias=bias,
)
return out.reshape(out_shape)
else:
return F.linear(x, layer.weight, bias)
class AWQAscendMoEKernel:
def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config
self.kernel = NPUW4A16Int4DynamicMoEMethod()
self.w13_kernel = NPUWNA16Int4MoEMethod()
self.w2_kernel = NPUWNA16Int4MoEMethod()
@staticmethod
def _register_or_replace_parameter(
@@ -87,88 +190,69 @@ class AWQAscendMoEKernel:
name, torch.nn.Parameter(tensor, requires_grad=False)
)
def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor:
num_experts, input_size, _ = qweight.shape
unpacked_weight = (
self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4)
.view(num_experts, input_size, -1)
.transpose(1, 2)
.contiguous()
.int()
)
return self.kernel._pack_to_int32(unpacked_weight)
def _convert_awq_qzeros_to_npu_offset(
self, qzeros: torch.Tensor, dtype: torch.dtype
) -> torch.Tensor:
num_experts, num_groups, _ = qzeros.shape
offset = (
-self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4)
.view(num_experts, num_groups, -1)
.transpose(1, 2)
.contiguous()
)
return offset.to(dtype)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self._register_or_replace_parameter(
layer,
"w13_weight",
self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data),
)
self._register_or_replace_parameter(
layer,
"w2_weight",
self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data),
)
self._register_or_replace_parameter(
layer,
"w13_weight_scale",
layer.w13_scales.data.transpose(1, 2).contiguous(),
)
self._register_or_replace_parameter(
layer,
"w2_weight_scale",
layer.w2_scales.data.transpose(1, 2).contiguous(),
)
self._register_or_replace_parameter(
layer,
"w13_weight_offset",
self._convert_awq_qzeros_to_npu_offset(
layer.w13_qzeros.data, layer.w13_scales.data.dtype
),
)
self._register_or_replace_parameter(
layer,
"w2_weight_offset",
self._convert_awq_qzeros_to_npu_offset(
layer.w2_qzeros.data, layer.w2_scales.data.dtype
),
)
w13_qweight_tmp = torch.zeros_like(layer.w13_qweight.data)
w2_qweight_tmp = torch.zeros_like(layer.w2_qweight.data)
w13_qzeros_list = []
w2_qzeros_list = []
self.kernel.process_weights_after_loading(layer)
# shifts control which 4-bit nibble we extract from each packed byte.
# For AWQ with pack_factor=8, each byte contains 8 nibbles (4-bit values).
# shifts = [0,4,1,5,2,6,3,7] extracts nibbles in the order:
# nibble_0, nibble_1, nibble_2, ..., nibble_7
# but interleaved to match the NPU kernel's expected layout.
shifts = [0, 4, 1, 5, 2, 6, 3, 7]
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> torch.Tensor:
return self.kernel.apply(layer, dispatch_output)
for i in range(self.quant_config.pack_factor):
shift_num = shifts[i] * 4
w13_qzeros_list.append(
(layer.w13_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF
)
w2_qzeros_list.append(
(layer.w2_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF
)
w13_qweight_tmp.bitwise_or_(
((layer.w13_qweight.data >> shift_num) * (2 ** (4 * i)))
& (0xF << (4 * i))
)
w2_qweight_tmp.bitwise_or_(
((layer.w2_qweight.data >> shift_num) * (2 ** (4 * i)))
& (0xF << (4 * i))
)
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,
# XOR with 0x88888888 converts signed int4 to unsigned int4.
# Each byte contains two 4-bit values, so 0x88 flips the sign bit
# of both nibbles simultaneously.
#
# signed_val: -8 -7 -6 -5 -4 -3 -2 -1 0 1 2 3 4 5 6 7
# unsigned_val: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# signed_val ^ 0x8 maps -8→0, -7→1, ..., 7→15.
w13_qweight_tmp.bitwise_xor_(0x88888888)
w2_qweight_tmp.bitwise_xor_(0x88888888)
# Convert zero points: signed int4 → unsigned → negated.
# The NPU kernel expects raw zero-point values (not dequantized).
w13_qzeros_tmp = torch.cat(w13_qzeros_list, dim=-1).reshape(
layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1
)
w13_qzeros_tmp = -(w13_qzeros_tmp - 8) # signed → unsigned → negated
w13_qzeros_tmp = w13_qzeros_tmp.to(layer.w13_scales.data.dtype)
w2_qzeros_tmp = torch.cat(w2_qzeros_list, dim=-1).reshape(
layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1
)
w2_qzeros_tmp = -(w2_qzeros_tmp - 8)
w2_qzeros_tmp = w2_qzeros_tmp.to(layer.w2_scales.data.dtype)
layer.register_parameter(
"w13_qzeros", torch.nn.Parameter(w13_qzeros_tmp, requires_grad=False)
)
layer.register_parameter(
"w13_qweight", torch.nn.Parameter(w13_qweight_tmp, requires_grad=False)
)
layer.register_parameter(
"w2_qzeros", torch.nn.Parameter(w2_qzeros_tmp, requires_grad=False)
)
layer.register_parameter(
"w2_qweight", torch.nn.Parameter(w2_qweight_tmp, requires_grad=False)
)
File diff suppressed because it is too large Load Diff
@@ -3,17 +3,14 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
import torch_npu
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
npu_fused_experts,
)
if TYPE_CHECKING:
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
from sglang.srt.layers.quantization.base_config import QuantizationConfig
import logging
logger = logging.getLogger(__name__)
def unpack_from_int32(
weight: torch.Tensor,
@@ -90,7 +87,7 @@ class GPTQLinearAscendKernel:
# for 4bit case we need to pack 4bit weight to int32 to save memory
layer.qweight = torch.nn.Parameter(
torch_npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)),
torch.ops.npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)),
requires_grad=False,
)
@@ -115,7 +112,7 @@ class GPTQLinearAscendKernel:
else:
out_shape = x.shape[:-1] + (qweight.shape[-1],)
out = torch_npu.npu_weight_quant_batchmatmul(
out = torch.ops.npu.npu_weight_quant_batchmatmul(
reshaped_x,
qweight,
antiquant_scale=scales,
@@ -131,17 +128,9 @@ class GPTQMoEAscendKernel:
def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config
self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
self.moe_runner_config: Optional[MoeRunnerConfig] = None
def create_moe_runner(
self,
layer: torch.nn.Module,
moe_runner_config: MoeRunnerConfig,
**extra_weight_attrs,
):
self.moe_runner_config = moe_runner_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# ----- zero‑points (unchanged) -----
w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape(
-1, layer.w13_qzeros.shape[-1]
)
@@ -174,6 +163,7 @@ class GPTQMoEAscendKernel:
if not self.use_v2_format:
layer.w2_qzeros += 1
# ----- w13 -----
w13_qweight_2d = (
layer.w13_qweight.data.transpose(-1, -2)
.contiguous()
@@ -185,35 +175,61 @@ class GPTQMoEAscendKernel:
if self.quant_config.weight_bits == 4:
group_size = self.quant_config.group_size
scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1)
k_shard_w13 = w13_qweight_tmp.shape[1]
neg_mask = scale_expanded < 0
# Check if the scales are compatible (expanded size must equal K_shard)
if layer.w13_scales.shape[1] * group_size != k_shard_w13:
logger.warning_once(
f"w13 scales expanded size {layer.w13_scales.shape[1] * group_size} "
f"does not match K_shard {k_shard_w13}. Skipping negative-scale correction."
f"This may break the accuracy, please try another TP-size or use DeepEP."
)
# pack directly
layer.w13_qweight = torch.nn.Parameter(
torch.ops.npu.npu_convert_weight_to_int4pack(
w13_qweight_tmp.reshape(
layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w13_qweight.shape[2])
.to(torch.int32)
)
.reshape(
layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1
)
.contiguous(),
requires_grad=False,
)
else:
scale_expanded = layer.w13_scales.data.repeat_interleave(
group_size, dim=1
)
neg_mask = scale_expanded < 0
if neg_mask.any():
neg_mask = neg_mask.transpose(-1, -2)
neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape)
w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask]
if w13_qweight_tmp.max() > 7:
w13_qweight_tmp.clamp_(max=7)
layer.w13_scales.data.abs_()
if neg_mask.any():
neg_mask = neg_mask.transpose(-1, -2)
neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape)
w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask]
if w13_qweight_tmp.max() > 7:
w13_qweight_tmp.clamp_(max=7)
layer.w13_scales.data.abs_()
layer.w13_qweight = torch.nn.Parameter(
torch_npu.npu_convert_weight_to_int4pack(
w13_qweight_tmp.reshape(
layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
layer.w13_qweight = torch.nn.Parameter(
torch.ops.npu.npu_convert_weight_to_int4pack(
w13_qweight_tmp.reshape(
layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w13_qweight.shape[2])
.to(torch.int32)
)
.reshape(
layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w13_qweight.shape[2])
.to(torch.int32)
.contiguous(),
requires_grad=False,
)
.reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1)
.contiguous(),
requires_grad=False,
)
# use int8 to store weight by default
else:
layer.w13_qweight = torch.nn.Parameter(
w13_qweight_tmp.reshape(
@@ -224,6 +240,7 @@ class GPTQMoEAscendKernel:
requires_grad=False,
)
# ----- w2 -----
w2_qweight_2d = (
layer.w2_qweight.data.transpose(-1, -2)
.contiguous()
@@ -235,35 +252,61 @@ class GPTQMoEAscendKernel:
if self.quant_config.weight_bits == 4:
group_size = self.quant_config.group_size
scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1)
k_shard_w2 = w2_qweight_tmp.shape[1]
neg_mask = scale_expanded < 0
# Check if the scales are compatible
if layer.w2_scales.shape[1] * group_size != k_shard_w2:
logger.warning_once(
f"w2 scales expanded size {layer.w2_scales.shape[1] * group_size} "
f"does not match K_shard {k_shard_w2}. Skipping negative-scale correction."
f"This may break the accuracy, please try another TP-size or use DeepEP."
)
# pack directly
layer.w2_qweight = torch.nn.Parameter(
torch.ops.npu.npu_convert_weight_to_int4pack(
w2_qweight_tmp.reshape(
layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w2_qweight.shape[2])
.to(torch.int32)
)
.reshape(
layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1
)
.contiguous(),
requires_grad=False,
)
else:
scale_expanded = layer.w2_scales.data.repeat_interleave(
group_size, dim=1
)
neg_mask = scale_expanded < 0
if neg_mask.any():
neg_mask = neg_mask.transpose(-1, -2)
neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape)
w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask]
if w2_qweight_tmp.max() > 7:
w2_qweight_tmp.clamp_(max=7)
layer.w2_scales.data.abs_()
if neg_mask.any():
neg_mask = neg_mask.transpose(-1, -2)
neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape)
w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask]
if w2_qweight_tmp.max() > 7:
w2_qweight_tmp.clamp_(max=7)
layer.w2_scales.data.abs_()
layer.w2_qweight = torch.nn.Parameter(
torch_npu.npu_convert_weight_to_int4pack(
w2_qweight_tmp.reshape(
layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
layer.w2_qweight = torch.nn.Parameter(
torch.ops.npu.npu_convert_weight_to_int4pack(
w2_qweight_tmp.reshape(
layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w2_qweight.shape[2])
.to(torch.int32)
)
.reshape(
layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1
)
.transpose(-1, -2)
.contiguous()
.reshape(-1, layer.w2_qweight.shape[2])
.to(torch.int32)
.contiguous(),
requires_grad=False,
)
.reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1)
.contiguous(),
requires_grad=False,
)
# use int8 to store weight by default
else:
layer.w2_qweight = torch.nn.Parameter(
w2_qweight_tmp.reshape(
@@ -273,43 +316,3 @@ class GPTQMoEAscendKernel:
.contiguous(),
requires_grad=False,
)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> torch.Tensor:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
assert self.moe_runner_config is not None, (
"moe_runner_config is not set. "
"Did you forget to call create_weights/create_moe_runner?"
)
assert self.moe_runner_config.activation in ("silu", "swiglu"), (
f"Only SiLU/Swiglu activation is supported, "
f"got {self.moe_runner_config.activation!r}."
)
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(x.dtype)
output = npu_fused_experts(
hidden_states=x,
w13=layer.w13_qweight,
w13_scale=layer.w13_scales,
w13_offset=layer.w13_qzeros,
w2=layer.w2_qweight,
w2_scale=layer.w2_scales,
w2_offset=layer.w2_qzeros,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=topk_ids.shape[1],
use_wna16=True,
)
return StandardCombineInput(hidden_states=output)
@@ -10,6 +10,8 @@ from sglang.srt.layers.quantization.base_config import LinearMethodBase
if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
MXFP8_BLOCK_SIZE = 32
@@ -313,9 +315,12 @@ class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
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)
)
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
layer.weight.data = layer.weight.data.view(torch.int32).contiguous()
else:
layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack(
layer.weight.data.to(torch.int32)
)
def apply(
self,
@@ -0,0 +1,724 @@
from typing import TYPE_CHECKING, Any, Dict, Optional
import numpy as np
import torch
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
import logging
from sglang.srt.hardware_backend.npu.moe.hidden_states_quant import (
HiddenStatesDynamicQuant,
)
from sglang.srt.hardware_backend.npu.moe.matmul import GroupedMatmul
logger = logging.getLogger(__name__)
# DEPRECATED METHOD
# TODO: Remove in future realeses
def fused_moe_npu(
x,
w1,
w2,
topk_output,
moe_runner_config,
):
logger.warning_once(
f"The fused_moe_npu method deprecated and will be removed in future releases"
)
topk_weights, topk_ids, _ = topk_output
original_dtype = x.dtype
num_tokens = x.shape[0]
topk_weights = topk_weights.to(x.dtype)
topk_ids = topk_ids.to(torch.int32)
num_experts = w1.shape[0]
top_k = topk_weights.shape[-1]
row_idx_len = num_tokens * top_k
row_idx = (
torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device)
.view(top_k, -1)
.permute(1, 0)
.contiguous()
)
hidden_states, expanded_row_idx, expanded_expert_idx = (
torch.ops.npu.npu_moe_init_routing(
x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
)
)
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
)
expert_tokens = expert_tokens.to(torch.int64)
# gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w1.permute(0, 2, 1)],
bias=None,
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# act_fn:
if moe_runner_config.activation == "silu":
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
else:
from sglang.srt.layers.activation import GeluAndMul
hidden_states = GeluAndMul()(hidden_states)
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w2.permute(0, 2, 1)],
bias=None,
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
)
return final_hidden_states
class _NPUMoEMethodBase(FusedMoEMethodBase):
"""Base class for NPU MoE methods with common helpers."""
def __init__(
self,
quant_config: Optional["QuantizationConfig"] = None,
):
super().__init__()
self.quant_config = quant_config
@staticmethod
def _set_dispatcher_output_dtype(layer: torch.nn.Module, dtype) -> None:
"""Set dispatcher output dtype if the layer has a dispatcher."""
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": dtype})
@staticmethod
def _validate_weight_prefix(layer: torch.nn.Module, weight_prefix: str) -> None:
"""Ensure the required attributes exist on the layer for the given prefix."""
required = [f"{weight_prefix}_weight"]
for attr in required:
if not hasattr(layer, attr):
raise AttributeError(
f"Layer {layer} is missing required attribute '{attr}' for "
f"weight_prefix='{weight_prefix}'"
)
@staticmethod
def _get_bias_args(
quant_info: "AscendQuantInfo", weight_prefix: str
) -> Dict[str, Any]:
bias = getattr(quant_info, f"{weight_prefix}_scale_bias", None)
if bias is None:
bias = getattr(quant_info, f"{weight_prefix}_weight_bias", None)
return {"bias": [bias]} if bias is not None else {}
# ---------------------------------------------------------------------------
# NPUW4A4Int4DynamicMoEMethod
# ---------------------------------------------------------------------------
class NPUW4A4Int4MoEMethod(_NPUMoEMethodBase):
"""W4A4 dynamic MoE – weights are int4, activations are int4."""
def __init__(self):
super().__init__(quant_config=None)
self.matmul = GroupedMatmul()
self.hidden_states_quantizer = HiddenStatesDynamicQuant(
quant_dtype=torch.quint4x2
)
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
) -> None:
self._validate_weight_prefix(layer, weight_prefix)
# Process scale
scale: torch.Tensor = getattr(layer, f"{weight_prefix}_weight_scale")
scale_np = scale.data.cpu().contiguous().numpy()
scale_np.dtype = np.uint32
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
processed_scale = torch.nn.Parameter(
scale_uint64_tensor.squeeze(-1), requires_grad=False
)
setattr(layer, f"{weight_prefix}_weight_scale", processed_scale)
# Process offset
offset: Optional[torch.Tensor] = getattr(
layer, f"{weight_prefix}_weight_offset", None
)
if offset is not None:
processed_offset = torch.nn.Parameter(
offset.data.squeeze(-1), requires_grad=False
)
setattr(layer, f"{weight_prefix}_weight_offset", processed_offset)
# Process weight
weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight")
if not envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
weight.data = self._w4a4_pack_int4(weight.data)
weight.data = weight.data.transpose(-2, -1).contiguous()
weight.data = npu_format_cast(weight.data)
weight.data = self._pack_to_int32(weight.data)
# Set DeepEP dispatcher output dtype
if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "bf16")
def _pack_int4(self, weight) -> torch.Tensor:
"""
Pack int4 weight to int8 weight
@param weight: torch.Tensor, int4 weight
@return: torch.Tensor, int8 weight
"""
weight = weight.to(torch.int8)
e = 0 # number of experts
if len(weight.shape) == 2:
k, n = weight.shape
elif len(weight.shape) == 3:
e, k, n = weight.shape
n_new = n // 2 + n % 2
if n_new != n // 2:
raise AssertionError("n dimension should be even")
weight = weight.reshape(-1, 2)
weight0 = weight[:, :1]
weight1 = weight[:, 1:]
weight1_4 = torch.bitwise_left_shift(weight1, 4)
weight2_4 = weight0 & 0b00001111
weight_add = torch.bitwise_or(weight1_4, weight2_4)
if e == 0:
weight_res = weight_add.reshape(k, n_new)
else:
weight_res = weight_add.reshape(e, k, n_new)
return weight_res
def _w4a4_pack_int4(self, save_quant_weight):
"""
Pack int4 weight to int8 weight
@param save_quant_weight: torch.Tensor, int4 weight
@return: torch.Tensor, int8 weight
"""
weight = save_quant_weight.transpose(-1, -2).contiguous()
packed_weight_tensor = self._pack_int4(weight)
packed_weight_tensor = packed_weight_tensor.transpose(-1, -2).contiguous()
return packed_weight_tensor
def _pack_to_int32(self, weight: torch.Tensor):
# pack 4 int8(int4*2) to int32
return weight.contiguous().view(torch.int32)
def apply(
self,
quant_info: "AscendQuantInfo",
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
pertoken_scale: torch.Tensor,
output_dtype: torch.dtype,
weight_prefix: str,
group_list_type,
) -> torch.Tensor:
scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None)
if pertoken_scale is None:
hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__(
hidden_states
)
scale_args: Dict[str, Any] = {
"scale": [scale],
"per_token_scale": [pertoken_scale],
}
scale_args.update(self._get_bias_args(quant_info, weight_prefix))
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens,
output_dtype,
group_list_type=group_list_type,
transposed=True,
**scale_args,
)
# ---------------------------------------------------------------------------
# NPUW8A8Int8MoEMethod
# ---------------------------------------------------------------------------
class NPUW8A8Int8MoEMethod(_NPUMoEMethodBase):
"""W8A8 MoE – weights are int8, activations in int8."""
def __init__(self):
super().__init__(quant_config=None)
self.matmul = GroupedMatmul()
self.hidden_states_quantizer = HiddenStatesDynamicQuant(quant_dtype=torch.int8)
@staticmethod
def maybe_process_fuseep_weights(layer: torch.nn.Module) -> bool:
"""Apply the FuseEP weight layout if --moe-a2a-backend is ascend_fuseep.
Returns True when the FuseEP layout was (or has already been) applied,
so that the caller can skip its own ``process_weights_after_loading`` body.
"""
from sglang.srt.layers.moe import get_moe_a2a_backend
if not get_moe_a2a_backend().is_ascend_fuseep():
return False
# Guard against double processing when called for multiple prefixes.
if getattr(layer, "_fuseep_weights_processed", False):
return True
from sglang.srt.hardware_backend.npu.moe.fuseep import process_fuseep_weights
for prefix in ("w13", "w2"):
process_fuseep_weights(layer, prefix)
layer._fuseep_weights_processed = True
return True
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
) -> None:
# If the FuseEP weight layout is used, process weights via
# maybe_apply_fuseep_weights and skip the rest of this method.
if self.maybe_process_fuseep_weights(layer):
return
self._validate_weight_prefix(layer, weight_prefix)
# Process scale
scale: torch.Tensor = getattr(layer, f"{weight_prefix}_weight_scale")
processed_scale = torch.nn.Parameter(
scale.data.squeeze(-1).to(dtype=torch.bfloat16), requires_grad=False
)
setattr(layer, f"{weight_prefix}_weight_scale", processed_scale)
# Process offset
offset: Optional[torch.Tensor] = getattr(
layer, f"{weight_prefix}_weight_offset", None
)
if offset is not None:
processed_offset = torch.nn.Parameter(
offset.data.squeeze(-1), requires_grad=False
)
setattr(layer, f"{weight_prefix}_weight_offset", processed_offset)
# Process weight
weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight")
weight.data = npu_format_cast(weight.data.transpose(1, 2))
setattr(
layer,
f"{weight_prefix}_weight",
torch.nn.Parameter(weight, requires_grad=False),
)
# Set dispatcher output dtype
if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "int8")
def apply(
self,
quant_info: "AscendQuantInfo",
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
pertoken_scale: torch.Tensor,
output_dtype: torch.dtype,
weight_prefix: str,
group_list_type,
) -> torch.Tensor:
scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None)
if pertoken_scale is None:
hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__(
hidden_states
)
scale_args: Dict[str, Any] = {
"scale": [scale],
"per_token_scale": [pertoken_scale],
}
scale_args.update(self._get_bias_args(quant_info, weight_prefix))
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens,
output_dtype,
group_list_type=group_list_type,
transposed=True,
**scale_args,
)
# ---------------------------------------------------------------------------
# NPUW4A8Int8MoEMethod
# ---------------------------------------------------------------------------
class NPUW4A8Int8MoEMethod(_NPUMoEMethodBase):
"""W4A8 MoE – weights are int4, activations quantized to int8."""
def __init__(
self,
quant_config: Optional["QuantizationConfig"] = None,
is_per_channel_weight: bool = False,
activation_use_clip: bool = False,
):
super().__init__(quant_config)
self.is_per_channel_weight = is_per_channel_weight
self.activation_use_clip = activation_use_clip
self.matmul = GroupedMatmul()
self.hidden_states_quantizer = HiddenStatesDynamicQuant(quant_dtype=torch.int8)
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
) -> None:
self._validate_weight_prefix(layer, weight_prefix)
# Process scale (and bias if needed)
scale = getattr(layer, f"{weight_prefix}_weight_scale")
scale_second = getattr(layer, f"{weight_prefix}_weight_scale_second", None)
bias = getattr(layer, f"{weight_prefix}_bias", None)
if not self.activation_use_clip:
# Process scale according to per-channel or per-group
processed_scale = self._process_scale(
getattr(layer, f"{weight_prefix}_weight"),
scale,
scale_second,
self.is_per_channel_weight,
)
setattr(
layer,
f"{weight_prefix}_weight_scale",
torch.nn.Parameter(processed_scale.squeeze(-1), requires_grad=False),
)
if scale_second is not None:
delattr(layer, f"{weight_prefix}_weight_scale_second")
delattr(layer, f"{weight_prefix}_weight_offset_second")
else:
# With clip: simple squeeze + unsqueeze
processed_scale = scale.data.squeeze(-1).unsqueeze(1).contiguous()
setattr(
layer,
f"{weight_prefix}_weight_scale",
torch.nn.Parameter(processed_scale, requires_grad=False),
)
if bias is not None:
setattr(
layer,
f"{weight_prefix}_scale_bias",
torch.nn.Parameter(
bias.data.transpose(1, 2).sum(dim=1).contiguous(),
requires_grad=False,
),
)
# Process weight
weight = getattr(layer, f"{weight_prefix}_weight")
weight.data = npu_format_cast(weight.data.transpose(1, 2))
weight.data = self._pack_to_int32(weight.data)
setattr(
layer,
f"{weight_prefix}_weight",
torch.nn.Parameter(weight, requires_grad=False),
)
# Set dispatcher output dtype
if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "int8")
def _process_scale(
self,
weight: torch.Tensor,
scale: torch.Tensor,
per_group_scale: Optional[torch.Tensor],
is_per_channel: bool,
) -> torch.Tensor:
scale = scale.transpose(1, 2).contiguous()
if is_per_channel:
scale_np = scale.cpu().contiguous().numpy()
scale_np.dtype = np.uint32
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
return scale_uint64_tensor
# Per‑group: multiply channel and group scales, then pack into uint64
per_group_scale = per_group_scale.transpose(1, 2).contiguous()
group_num, k, n = weight.shape
n = n * 2 # packed weight halves the column dimension
per_group_scale = per_group_scale.reshape(group_num, -1, n)
group_num, quantgroup_num, n = per_group_scale.shape
scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32)
scale_fp32_np = scale_fp32.cpu().numpy()
scale_fp32_np.dtype = np.uint32
sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32)
sscale_uint64[..., ::2] = scale_fp32_np
sscale_uint64_tensor = (
torch.from_numpy(sscale_uint64.view(np.int64).copy())
.reshape(group_num, quantgroup_num, n)
.npu()
)
return sscale_uint64_tensor
def _pack_to_int32(self, weight: torch.Tensor) -> torch.Tensor:
# pack 4 int8 (representing 8 int4) into int32
assert weight.shape[-1] % 4 == 0, (
f"Last dimension of weight must be divisible by 4 for int8→int32 packing, "
f"got shape {weight.shape}"
)
return weight.contiguous().view(torch.int32)
def apply(
self,
quant_info: "AscendQuantInfo",
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
pertoken_scale: torch.Tensor,
output_dtype: torch.dtype,
weight_prefix: str,
group_list_type,
) -> torch.Tensor:
scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None)
if pertoken_scale is None:
hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__(
hidden_states
)
scale_args: Dict[str, Any] = {
"scale": [scale],
"per_token_scale": [pertoken_scale],
}
scale_args.update(self._get_bias_args(quant_info, weight_prefix))
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens,
output_dtype,
group_list_type=group_list_type,
transposed=True,
**scale_args,
)
# ---------------------------------------------------------------------------
# NPUWNA16Int4MoEMethod
# ---------------------------------------------------------------------------
class NPUWNA16Int4MoEMethod(_NPUMoEMethodBase):
"""W4A16 MoE – weights are int4, activations stay in BF16."""
def __init__(self):
super().__init__(quant_config=None)
self.matmul = GroupedMatmul()
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
) -> None:
self._validate_weight_prefix(layer, weight_prefix)
# Process scale
scale = getattr(layer, f"{weight_prefix}_weight_scale") # shape [E, N, 1]
scale = scale.data.transpose(-1, -2).contiguous() # [E, N, 1] -> [E, 1, N]
setattr(
layer,
f"{weight_prefix}_weight_scale",
torch.nn.Parameter(scale, requires_grad=False),
)
# Process offset
offset = getattr(layer, f"{weight_prefix}_weight_offset", None)
if offset is not None:
offset = offset.data.transpose(-1, -2).contiguous()
setattr(
layer,
f"{weight_prefix}_weight_offset",
torch.nn.Parameter(offset, requires_grad=False),
)
# Process weight: unpack, transpose, repack
weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight")
unpacked_weight = (
self._unpack_from_int32(weight.data.flatten(0, 1), 4)
.view(weight.shape[0], weight.shape[1], -1)
.transpose(1, 2)
.int()
)
weight.data = self._pack_to_int32(unpacked_weight)
setattr(
layer,
f"{weight_prefix}_weight",
torch.nn.Parameter(weight, requires_grad=False),
)
# Set dispatcher output dtype
if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "bf16")
def _pack_to_int32(self, weight: torch.Tensor) -> torch.Tensor:
assert weight.dim() == 3
if weight.dtype == torch.int32:
assert weight.shape[-1] % 8 == 0, (
f"Last dimension of int32 weight must be divisible by 8 for int4 packing, "
f"got {weight.shape}"
)
new_weight = torch.ops.npu.npu_convert_weight_to_int4pack(
weight.flatten(0, 1)
)
new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1)
elif weight.dtype == torch.int8:
assert weight.shape[-1] % 4 == 0, (
f"Last dimension of int8 weight must be divisible by 4 for int32 packing, "
f"got {weight.shape}"
)
new_weight = weight.contiguous().view(torch.int32)
else:
raise ValueError(f"Unsupported weight dtype for packing: {weight.dtype}")
return new_weight.contiguous()
def _unpack_from_int32(
self,
value: torch.Tensor,
num_bits: int,
shape: Optional[torch.Size] = None,
packed_dim: int = 1,
) -> torch.Tensor:
"""
Unpacks a tensor of packed int32 weights into individual int8s,
maintaining the original bit range.
"""
if value.dtype is not torch.int32:
raise ValueError(
f"Expected {torch.int32} but got {value.dtype}, Aborting unpack."
)
if num_bits > 8:
raise ValueError("Unpacking is only supported for less than 8 bits")
pack_factor = 32 // num_bits
mask = (1 << num_bits) - 1
if packed_dim == 1:
unpacked = torch.zeros(
(value.shape[0], value.shape[1] * pack_factor),
device=value.device,
dtype=torch.int32,
)
for i in range(pack_factor):
unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask
if shape is not None:
original_row_size = int(shape[1])
unpacked = unpacked[:, :original_row_size]
else:
unpacked = torch.zeros(
(value.shape[0] * pack_factor, value.shape[1]),
device=value.device,
dtype=torch.int32,
)
for i in range(pack_factor):
unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask
if shape is not None:
original_row_size = int(shape[0])
unpacked = unpacked[:original_row_size, :]
offset = pow(2, num_bits) // 2
unpacked = (unpacked - offset).to(torch.int8)
return unpacked
def apply(
self,
quant_info: "AscendQuantInfo",
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
pertoken_scale: torch.Tensor, # not used, but kept for interface consistency
output_dtype: torch.dtype,
weight_prefix: str,
group_list_type,
) -> torch.Tensor:
scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None)
offset = getattr(quant_info, f"{weight_prefix}_weight_offset", None)
scale_args: Dict[str, Any] = {
"antiquant_scale": [scale],
"antiquant_offset": [offset] if offset is not None else [],
}
scale_args.update(self._get_bias_args(quant_info, weight_prefix))
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens,
output_dtype,
group_list_type=group_list_type,
transposed=True,
**scale_args,
)
# ---------------------------------------------------------------------------
# NPUWUnquantMoEMethod
# ---------------------------------------------------------------------------
class NPUUnquantMoEMethod(_NPUMoEMethodBase):
"""Unquant MoE – all computations in BF16, no quantization."""
def __init__(self):
super().__init__(quant_config=None)
self.matmul = GroupedMatmul()
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
) -> None:
self._validate_weight_prefix(layer, weight_prefix)
weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight")
weight.data = npu_format_cast(weight)
setattr(
layer,
f"{weight_prefix}_weight",
torch.nn.Parameter(weight, requires_grad=False),
)
if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "bf16")
def apply(
self,
quant_info: "AscendQuantInfo",
hidden_states: torch.Tensor,
expert_tokens: torch.Tensor,
pertoken_scale: torch.Tensor, # ignored
output_dtype: torch.dtype,
weight_prefix: str,
group_list_type,
) -> torch.Tensor:
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens,
output_dtype,
group_list_type=group_list_type,
transposed=False,
**self._get_bias_args(quant_info, weight_prefix),
)
@@ -24,11 +24,6 @@ class NPUACLFormat(IntEnum):
ACL_FORMAT_FRACTAL_NZ = 29
class FusedMoEMode(IntEnum):
FUSED_DEEP_MOE = 1
DISPATCH_FFN_COMBINE = 2
def _call_once(fn: Callable):
@functools.wraps(fn)
+9
View File
@@ -24,6 +24,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
is_allocation_symmetric,
)
@@ -670,6 +671,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
shard_size = loaded_weight.shape[output_dim]
shard_offset = loaded_weight.shape[output_dim] * loaded_shard_id
# Needed for experimental ModelSlim W4A4 int4x2 packing support
# TODO: remove env variable once new packing is fully released
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
pack_factor = getattr(param, "pack_factor", None)
if pack_factor is not None:
shard_size = shard_size // pack_factor
shard_offset = shard_offset // pack_factor
param_data = param_data.narrow(output_dim, shard_offset, shard_size)
start_idx = self.tp_rank * shard_size
@@ -285,9 +285,4 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]):
or get_moe_a2a_backend().is_nixl()
):
return DeepEPMoE
if get_moe_a2a_backend().is_ascend_fuseep():
# ascend_fuseep bypasses dispatch/combine inside FusedMoE.forward
# (see forward_fuseep in hardware_backend/npu/moe/fuseep.py).
return FusedMoE
return FusedMoE
@@ -34,6 +34,9 @@ from sglang.srt.layers.moe.kt_ep_wrapper import (
create_kt_config_from_server_args,
)
from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
from sglang.srt.layers.moe.token_dispatcher.ascend_tp import (
AscendTPDispatcher,
)
from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.token_dispatcher.standard import (
@@ -100,7 +103,9 @@ def _get_deepep_comm_group(a2a_backend):
def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
a2a_backend = get_moe_a2a_backend()
if (
if a2a_backend.is_none() and is_npu():
return AscendTPDispatcher(moe_runner_config)
elif (
a2a_backend.is_none()
or a2a_backend.is_megamoe()
or a2a_backend.is_ascend_fuseep()
@@ -0,0 +1,310 @@
"""Ascend MoE runner backend with NPU‑specific ops."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional
import torch
from sglang.srt.hardware_backend.npu.moe.activation import (
AllGatherActivationWrapper,
NPUGeluAndMul,
NPUSwiglu,
NPUSwigluDeepEPKernel,
NPUSwigluOAI,
NPUSwigluQuant,
NPUSwigluStepAndMul,
)
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW4A8Int8MoEMethod,
NPUW8A8Int8MoEMethod,
)
from sglang.srt.layers.moe.moe_runner.base import (
MoeQuantInfo,
MoeRunnerConfig,
MoeRunnerCore,
RunnerInput,
RunnerOutput,
register_post_permute,
register_pre_permute,
)
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher.deepep import (
DeepEPLLCombineInput,
DeepEPLLDispatchOutput,
DeepEPNormalCombineInput,
DeepEPNormalDispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.ascend_tp import (
AscendTPDispatchOutput,
AscendTPCombineInput,
)
from sglang.srt.layers.moe.utils import (
MoeRunnerBackend,
get_moe_a2a_backend,
)
# ---------------------------------------------------------------------------
# Runner IO dataclasses
# ---------------------------------------------------------------------------
@dataclass
class AscendRunnerInput(RunnerInput):
"""Input bundle for the NPU runner."""
hidden_states: torch.Tensor
hidden_states_scale: Optional[torch.Tensor] # None for unquant
expert_tokens: torch.Tensor
group_list_type: int # 0 or 1 (passed to NPU ops)
@property
def runner_backend(self) -> MoeRunnerBackend:
return MoeRunnerBackend.ASCEND
@dataclass
class AscendRunnerOutput(RunnerOutput):
"""Output bundle from the NPU runner."""
hidden_states: torch.Tensor
@property
def runner_backend(self) -> MoeRunnerBackend:
return MoeRunnerBackend.ASCEND
# ---------------------------------------------------------------------------
# Main runner core
# ---------------------------------------------------------------------------
class AscendRunnerCore(MoeRunnerCore):
runner_backend = MoeRunnerBackend.ASCEND
def __init__(self, config: MoeRunnerConfig):
super().__init__(config)
kernel = config.layer.w2_kernel
if get_moe_a2a_backend().is_deepep():
# DeepEP path: use a unified kernel that decides quantisation
is_quant_kernel = isinstance(
kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod)
)
self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel)
else:
# Non‑DeepEP (ascend_tp) path
# 1. Choose the base activation according to the quant method
if isinstance(kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod)):
inner = NPUSwigluQuant()
else:
if config.activation == "npu_swiglu_oai":
# NPUSwigluOAI requires the runner config to pass
# gemm1_alpha and gemm1_clamp_limit to the triton kernel.
inner = NPUSwigluOAI(moe_runner_config=config)
elif config.activation == "silu":
if config.gemm1_clamp_limit is not None:
inner = NPUSwigluStepAndMul(
clamp_limit=config.gemm1_clamp_limit
)
else:
inner = NPUSwiglu()
else:
inner = NPUGeluAndMul()
# 2. If the quant method (GGUF) needs TP all‑gather, wrap the activation
if getattr(config, "use_tp_all_gather_activation", False):
self.activation = AllGatherActivationWrapper(inner, dim=-1)
else:
self.activation = inner
def run(
self,
runner_input: AscendRunnerInput,
quant_info: AscendQuantInfo,
running_state: dict,
hooks: Optional[Any] = None,
) -> AscendRunnerOutput:
"""
Execute the MoE layer using NPU‑specific grouped matmul ops.
"""
x = runner_input.hidden_states
original_dtype = torch.float16 if x.dtype == torch.float16 else torch.bfloat16
expert_tokens = runner_input.expert_tokens
group_list_type = runner_input.group_list_type
# --- w13 (gate & up) projection ---
hidden_states = self.config.layer.w13_kernel.apply(
quant_info,
x,
expert_tokens,
pertoken_scale=runner_input.hidden_states_scale,
output_dtype=original_dtype,
weight_prefix="w13",
group_list_type=group_list_type,
)
# --- Activation ---
# The DeepEP kernel expects extra dispatch metadata
if isinstance(self.activation, NPUSwigluDeepEPKernel):
hidden_states, pertoken_scale = self.activation._apply_activation(
hidden_states,
group_list=expert_tokens,
group_list_type=group_list_type,
)
else:
hidden_states, pertoken_scale = self.activation._apply_activation(
hidden_states
)
# --- w2 (down) projection ---
hidden_states = self.config.layer.w2_kernel.apply(
quant_info,
hidden_states,
expert_tokens,
pertoken_scale=pertoken_scale,
output_dtype=original_dtype,
weight_prefix="w2",
group_list_type=group_list_type,
)
return AscendRunnerOutput(hidden_states=hidden_states)
# ---------------------------------------------------------------------------
# QuantInfo
# ---------------------------------------------------------------------------
@dataclass
class AscendQuantInfo(MoeQuantInfo):
"""Quantization payload for Ascend."""
w13_weight: torch.Tensor
w2_weight: torch.Tensor
w13_weight_scale: Optional[torch.Tensor] = None
w2_weight_scale: Optional[torch.Tensor] = None
w13_weight_offset: Optional[torch.Tensor] = None
w2_weight_offset: Optional[torch.Tensor] = None
w13_weight_bias: Optional[torch.Tensor] = None
w2_weight_bias: Optional[torch.Tensor] = None
w13_scale_bias: Optional[torch.Tensor] = None
w2_scale_bias: Optional[torch.Tensor] = None
# ---------------------------------------------------------------------------
# Pre/Post permute hooks
# ---------------------------------------------------------------------------
@register_pre_permute("ascend_tp", "ascend")
def pre_permute_ascend_tp_to_ascend(
dispatch_output: AscendTPDispatchOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> AscendRunnerInput:
return AscendRunnerInput(
hidden_states=dispatch_output.hidden_states,
hidden_states_scale=dispatch_output.hidden_states_scale,
expert_tokens=dispatch_output.expert_tokens,
group_list_type=dispatch_output.group_list_type,
)
@register_pre_permute("deepep_normal", "ascend")
def pre_permute_deepep_normal_to_ascend(
dispatch_output: DeepEPNormalDispatchOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> AscendRunnerInput:
(
hidden_states,
hidden_states_scale,
topk_ids,
topk_weights,
num_recv_tokens_per_expert,
) = dispatch_output
group_list = torch.tensor(
num_recv_tokens_per_expert,
dtype=torch.int64,
device=hidden_states.device,
)
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
return AscendRunnerInput(
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
expert_tokens=group_list,
group_list_type=1,
)
@register_pre_permute("deepep_ll", "ascend")
def pre_permute_deepep_ll_to_ascend(
dispatch_output: DeepEPLLDispatchOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> AscendRunnerInput:
(
hidden_states,
hidden_states_scale,
topk_ids,
topk_weights,
group_list,
_,
) = dispatch_output
group_list = group_list.to(torch.int64)
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
return AscendRunnerInput(
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
expert_tokens=group_list,
group_list_type=1,
)
@register_post_permute("ascend", "ascend_tp")
def post_permute_ascend_to_ascend_tp(
runner_output: AscendRunnerOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> AscendTPCombineInput:
from sglang.srt.layers.moe.token_dispatcher.ascend_tp import AscendTPCombineInput
return AscendTPCombineInput(hidden_states=runner_output.hidden_states)
@register_post_permute("ascend", "deepep_normal")
def post_permute_ascend_to_deepep_normal(
runner_output: AscendRunnerOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepEPNormalCombineInput:
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput
return DeepEPNormalCombineInput(
hidden_states=runner_output.hidden_states,
topk_ids=running_state["topk_ids"],
topk_weights=running_state["topk_weights"],
)
@register_post_permute("ascend", "deepep_ll")
def post_permute_ascend_to_deepep_ll(
runner_output: AscendRunnerOutput,
quant_info: AscendQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepEPLLCombineInput:
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
return DeepEPLLCombineInput(
hidden_states=runner_output.hidden_states,
topk_ids=running_state["topk_ids"],
topk_weights=running_state["topk_weights"],
)
@@ -60,6 +60,8 @@ class MoeRunnerConfig:
# silu+is_gated swiglu path consumes it (interleaved -> swiglu_gpt_oss_*,
# otherwise chunk gate/up then apply alpha/limit).
gate_up_interleaved: bool = True
layer: Optional[torch.nn.Module] = None
use_tp_all_gather_activation: bool = False
@dataclass
@@ -39,6 +39,10 @@ class MoeRunner:
if runner_backend.is_triton():
self.runner_core = TritonRunnerCore(config)
elif runner_backend.is_ascend():
from sglang.srt.layers.moe.moe_runner.ascend import AscendRunnerCore
self.runner_core = AscendRunnerCore(config)
elif runner_backend.is_triton_kernels():
self.runner_core = TritonKernelsRunnerCore(config)
elif runner_backend.is_deep_gemm():
@@ -1,3 +1,8 @@
from sglang.srt.layers.moe.token_dispatcher.ascend_tp import (
AscendTPCombineInput,
AscendTPDispatcher,
AscendTPDispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.base import (
BaseDispatcher,
BaseDispatcherConfig,
@@ -74,4 +79,7 @@ __all__ = [
"DeepEPLLDispatchOutput",
"DeepEPLLCombineInput",
"DeepEPNormalCombineInput",
"AscendTPDispatcher",
"AscendTPDispatchOutput",
"AscendTPCombineInput",
]
@@ -0,0 +1,137 @@
from __future__ import annotations
from typing import NamedTuple, Optional
import torch
from sglang.srt.hardware_backend.npu.moe.finalize_routing import (
AllGatherFinalizeRoutingWrapper,
NPUFinalizeRouting,
)
from sglang.srt.hardware_backend.npu.moe.init_routing import (
NPUMoEInitRouting_v2,
)
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher.base import (
BaseDispatcher,
CombineInputFormat,
DispatchOutputFormat,
)
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.utils import (
DispatcherOutputDtype,
get_ascend_dispatcher_output_dtype,
)
from sglang.srt.runtime_context import get_parallel
class AscendTPDispatchOutput(NamedTuple):
hidden_states: torch.Tensor
hidden_states_scale: Optional[torch.Tensor]
topk_weights: torch.Tensor
topk_ids: torch.Tensor
expanded_row_idx: torch.Tensor
expert_tokens: torch.Tensor
group_list_type: int
@property
def format(self) -> DispatchOutputFormat:
return DispatchOutputFormat.ASCEND_TP
class AscendTPCombineInput(NamedTuple):
hidden_states: torch.Tensor
@property
def format(self) -> CombineInputFormat:
return CombineInputFormat.ASCEND_TP
class AscendTPDispatcher(BaseDispatcher):
def __init__(self, moe_runner_config: MoeRunnerConfig):
super().__init__()
self.num_experts = moe_runner_config.num_experts
self.top_k = moe_runner_config.top_k
self._dispatch_output: Optional[AscendTPDispatchOutput] = None
self.quant_config: Optional[dict] = None
# Initialise routing kernels with default (no quant config yet)
self.set_ascend_dispatcher_output_dtype()
def set_quant_config(self, quant_config: dict) -> None:
self.quant_config = quant_config
self.set_ascend_dispatcher_output_dtype()
# If the quantisation is GGUF and TP is active, wrap the finalizer
# with an all‑gather so that the dispatcher stays completely clean.
if (
isinstance(self.quant_config, dict)
and self.quant_config.get("quant_type") == "gguf"
and get_parallel().tp_size > 1
):
self.finalize = AllGatherFinalizeRoutingWrapper(self.finalize, dim=-1)
def set_ascend_dispatcher_output_dtype(self) -> None:
"""Choose init & finalize routing kernels based on quant config."""
self.ascend_dispatcher_output_dtype = get_ascend_dispatcher_output_dtype(self)
if self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.BF16:
self.init = NPUMoEInitRouting_v2(quant_mode=-1)
self.finalize = NPUFinalizeRouting(drop_pad_mode=2)
self.group_list_type = 1
elif self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.INT8:
self.init = NPUMoEInitRouting_v2(quant_mode=1)
self.finalize = NPUFinalizeRouting(drop_pad_mode=2)
self.group_list_type = 1
else:
raise ValueError(
f"Unsupported ascend_dispatcher_output_dtype: {self.ascend_dispatcher_output_dtype}"
)
def dispatch(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> AscendTPDispatchOutput:
topk_weights, topk_ids, _ = topk_output
topk_weights = topk_weights.to(hidden_states.dtype)
topk_ids = topk_ids.to(torch.int32)
(
permuted_hidden_states,
expanded_row_idx,
expert_tokens,
hidden_states_scale,
) = self.init._init_routing(
hidden_states,
topk_ids,
self.num_experts,
self.top_k,
)
self._dispatch_output = AscendTPDispatchOutput(
hidden_states=permuted_hidden_states,
hidden_states_scale=hidden_states_scale,
topk_weights=topk_weights,
topk_ids=topk_ids,
expanded_row_idx=expanded_row_idx,
expert_tokens=expert_tokens,
group_list_type=self.group_list_type,
)
return self._dispatch_output
def combine(self, combine_input: AscendTPCombineInput) -> torch.Tensor:
if self._dispatch_output is None:
raise RuntimeError("combine() called before dispatch()")
dispatch_out = self._dispatch_output
# The finalizer (possibly wrapped with TP all‑gather) does all the work.
final_hidden_states = self.finalize._finalize_routing(
combine_input.hidden_states,
topk_weights=dispatch_out.topk_weights,
expanded_row_idx=dispatch_out.expanded_row_idx,
topk_ids=dispatch_out.topk_ids,
)
self._dispatch_output = None
return final_hidden_states
@@ -21,6 +21,8 @@ import torch
if TYPE_CHECKING:
from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs
from sglang.srt.layers.moe.token_dispatcher import (
AscendTPCombineInput,
AscendTPDispatchOutput,
DeepEPLLCombineInput,
DeepEPLLDispatchOutput,
DeepEPNormalCombineInput,
@@ -133,6 +135,12 @@ class DispatchOutputChecker:
) -> TypeGuard[StandardDispatchOutput]:
return dispatch_output.format.is_standard()
@staticmethod
def format_is_ascend_tp(
dispatch_output: DispatchOutput,
) -> TypeGuard[AscendTPDispatchOutput]:
return dispatch_output.format.is_ascend_tp()
@staticmethod
def format_is_deepep_normal(
dispatch_output: DispatchOutput,
@@ -164,10 +172,14 @@ class DispatchOutputFormat(Enum):
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer"
ASCEND_TP = "ascend_tp"
def is_standard(self) -> bool:
return self == DispatchOutputFormat.STANDARD
def is_ascend_tp(self) -> bool:
return self == DispatchOutputFormat.ASCEND_TP
def is_deepep_normal(self) -> bool:
return self == DispatchOutputFormat.DEEPEP_NORMAL
@@ -204,6 +216,12 @@ class CombineInputChecker:
) -> TypeGuard[StandardCombineInput]:
return combine_input.format == CombineInputFormat.STANDARD
@staticmethod
def format_is_ascend_tp(
combine_input: CombineInput,
) -> TypeGuard[AscendTPCombineInput]:
return combine_input.format == CombineInputFormat.ASCEND_TP
@staticmethod
def format_is_deepep_normal(
combine_input: CombineInput,
@@ -237,6 +255,7 @@ class CombineInputFormat(Enum):
DEEPEP_NORMAL = "deepep_normal"
DEEPEP_LL = "deepep_ll"
FLASHINFER = "flashinfer"
ASCEND_TP = "ascend_tp"
@runtime_checkable
@@ -22,7 +22,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.utils import (
DeepEPMode,
DeepEPOutputDtype,
DispatcherOutputDtype,
get_deepep_config,
get_deepep_output_dtype,
is_tbo_enabled,
@@ -422,22 +422,22 @@ class _DeepEPDispatcherImplBase:
# Configuration mapping for each dtype
config_map = {
DeepEPOutputDtype.BF16: {
DispatcherOutputDtype.BF16: {
"use_fp8": False,
"use_nvfp4": False,
},
DeepEPOutputDtype.FP8: {
DispatcherOutputDtype.FP8: {
"use_fp8": True,
"use_nvfp4": False,
},
# Needed for Ascend A2/A3 NPU case,
# despite the use_fp8 flag,
# quantization will be performed in int8
DeepEPOutputDtype.INT8: {
DispatcherOutputDtype.INT8: {
"use_fp8": True,
"use_nvfp4": False,
},
DeepEPOutputDtype.NVFP4: {
DispatcherOutputDtype.NVFP4: {
"use_fp8": False,
"use_nvfp4": True,
},
@@ -458,23 +458,23 @@ class _DeepEPDispatcherImplBase:
def _validate_and_adjust_dtype(self) -> None:
"""Validate dtype against hardware and adjust if necessary."""
if _is_npu:
if self.deepep_output_dtype == DeepEPOutputDtype.FP8:
if self.deepep_output_dtype == DispatcherOutputDtype.FP8:
logger.warning_once(
"Ascend A2/A3 NPU does not support fp8 "
"deepep_dispatcher_output_dtype, switching to int8..."
)
self.deepep_output_dtype = DeepEPOutputDtype.INT8
elif self.deepep_output_dtype == DeepEPOutputDtype.NVFP4:
self.deepep_output_dtype = DispatcherOutputDtype.INT8
elif self.deepep_output_dtype == DispatcherOutputDtype.NVFP4:
raise RuntimeError(
"Ascend A2/A3 NPU does not support nvfp4 deepep_dispatcher_output_dtype."
)
else:
if self.deepep_output_dtype == DeepEPOutputDtype.INT8:
if self.deepep_output_dtype == DispatcherOutputDtype.INT8:
logger.warning_once(
"GPU does not support int8 "
"deepep_dispatcher_output_dtype, switching to fp8..."
)
self.deepep_output_dtype = DeepEPOutputDtype.FP8
self.deepep_output_dtype = DispatcherOutputDtype.FP8
# NVFP4 is supported on GPU, no adjustment needed
def _update_int8_quant_env(self) -> None:
+34 -9
View File
@@ -33,6 +33,7 @@ class MoeA2ABackend(Enum):
NIXL = "nixl"
MORI = "mori"
ASCEND_FUSEEP = "ascend_fuseep"
ASCEND_TP = "ascend_tp"
FLASHINFER = "flashinfer"
MEGAMOE = "megamoe"
CUSTOMIZED = "customized"
@@ -64,6 +65,9 @@ class MoeA2ABackend(Enum):
def is_ascend_fuseep(self):
return self == MoeA2ABackend.ASCEND_FUSEEP
def is_ascend_tp(self):
return self == MoeA2ABackend.ASCEND_TP
def is_mori(self):
return self == MoeA2ABackend.MORI
@@ -89,6 +93,7 @@ class MoeRunnerBackend(Enum):
DEEP_GEMM = "deep_gemm"
TRITON = "triton"
TRITON_KERNELS = "triton_kernel"
ASCEND = "ascend"
FLASHINFER_TRTLLM = "flashinfer_trtllm"
EXPERIMENTAL_SGL_TRTLLM = "experimental_sgl_trtllm"
FLASHINFER_TRTLLM_ROUTED = "flashinfer_trtllm_routed"
@@ -109,6 +114,9 @@ class MoeRunnerBackend(Enum):
def is_triton(self):
return self == MoeRunnerBackend.TRITON
def is_ascend(self):
return self == MoeRunnerBackend.ASCEND
def is_triton_kernels(self):
return self == MoeRunnerBackend.TRITON_KERNELS
@@ -179,7 +187,7 @@ class DeepEPMode(Enum):
return self == DeepEPMode.AUTO
class DeepEPOutputDtype(Enum):
class DispatcherOutputDtype(Enum):
"""
Describes the dispatch output data type for DeepEP.
@@ -195,7 +203,7 @@ class DeepEPOutputDtype(Enum):
NVFP4 = "nvfp4"
def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
def get_deepep_output_dtype(self) -> DispatcherOutputDtype:
"""
Automatically choose the dispatch output dtype for DeepEP.
@@ -212,7 +220,7 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
# 0. Parse server argument.
server_args = get_server_args()
if server_args and server_args.deepep_dispatcher_output_dtype != "auto":
return DeepEPOutputDtype(server_args.deepep_dispatcher_output_dtype)
return DispatcherOutputDtype(server_args.deepep_dispatcher_output_dtype)
# 1. Parse deprecated environment variables.
if envs.SGLANG_DEEPEP_BF16_DISPATCH.get():
@@ -221,18 +229,18 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
"and will be removed in future releases. Please use a new "
"`--deepep-dispatcher-output-dtype bf16` argument instead."
)
return DeepEPOutputDtype.BF16
return DispatcherOutputDtype.BF16
# 2. NVFP4 is detected inside dispatch_a / _dispatch_core via quant_config; no need to infer here.
if self.quant_config is not None:
input_global_scale = self.quant_config.get("input_global_scale", None)
if input_global_scale is not None:
return DeepEPOutputDtype.NVFP4
return DispatcherOutputDtype.NVFP4
# 3. Parse quant config to determine the output dtype of dispatcher
dispatcher_output_dtype = self.quant_config.get("dispatcher_output_dtype", None)
if dispatcher_output_dtype is not None:
return DeepEPOutputDtype(dispatcher_output_dtype)
return DispatcherOutputDtype(dispatcher_output_dtype)
# 4. flashinfer_cutedsl / cutlass / humming expects BF16 dispatch
if (
@@ -240,14 +248,31 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
or get_moe_runner_backend().is_cutlass()
or get_moe_runner_backend().is_humming()
):
return DeepEPOutputDtype.BF16
return DispatcherOutputDtype.BF16
# 5. Default on NPU → BF16
if _is_npu:
return DeepEPOutputDtype.BF16
return DispatcherOutputDtype.BF16
# 6. Default → FP8
return DeepEPOutputDtype.FP8
return DispatcherOutputDtype.FP8
def get_ascend_dispatcher_output_dtype(dispatcher):
"""
Automatically choose the dispatch output dtype for Ascend.
"""
# 1. Parse quant config to determine the output dtype of dispatcher
if dispatcher.quant_config is not None:
dispatcher_output_dtype = dispatcher.quant_config.get(
"dispatcher_output_dtype", None
)
if dispatcher_output_dtype is not None:
return DispatcherOutputDtype(dispatcher_output_dtype)
# 2. Ascend dispatch defaults to BF16
return DispatcherOutputDtype.BF16
def initialize_moe_config(server_args: ServerArgs):
@@ -6,17 +6,16 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.layers.linear import set_weight_attrs
from sglang.srt.layers.moe import (
MoeRunner,
MoeRunnerBackend,
MoeRunnerConfig,
get_moe_runner_backend,
)
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from .awq_scheme import AWQMoESchemeBase
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
from sglang.srt.layers.quantization.awq.awq import AWQConfig, AWQMarlinConfig
__all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"]
@@ -151,6 +150,43 @@ class AWQAscendMoEScheme(AWQMoEScheme):
return AWQAscendMoEKernel(quant_config)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
self,
layer: torch.nn.Module,
moe_runner_config: MoeRunnerConfig,
**extra_weight_attrs,
):
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUWNA16Int4MoEMethod,
)
self.moe_runner_config = moe_runner_config
layer.w13_kernel = NPUWNA16Int4MoEMethod()
layer.w2_kernel = NPUWNA16Int4MoEMethod()
moe_runner_config.layer = layer
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def apply_weights(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.moe_runner.ascend import (
AscendQuantInfo,
)
quant_info = AscendQuantInfo(
w13_weight=layer.w13_qweight,
w2_weight=layer.w2_qweight,
w13_weight_scale=layer.w13_scales,
w2_weight_scale=layer.w2_scales,
w13_weight_offset=layer.w13_qzeros,
w2_weight_offset=layer.w2_qzeros,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
@@ -1,15 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Optional
from typing import Optional
import torch
from sglang.srt.layers.moe import MoeRunnerConfig
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
__all__ = ["BaseLinearScheme", "BaseMoEScheme"]
@@ -66,12 +61,6 @@ class BaseMoEScheme(ABC):
"""
raise NotImplementedError
@abstractmethod
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
raise NotImplementedError
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module):
"""
@@ -79,21 +68,3 @@ class BaseMoEScheme(ABC):
needs to occur.
"""
raise NotImplementedError
@abstractmethod
def apply_weights(
self,
layer: torch.nn.Module,
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
@@ -5,10 +5,11 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int8DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW4A8Int8MoEMethod,
)
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsMoEScheme,
)
@@ -38,7 +39,14 @@ class NPUCompressedTensorsW4A8Int8DynamicMoE(CompressedTensorsMoEScheme):
.get("group_1", {})
.get("activation_use_clip", False)
)
self.kernel = NPUW4A8Int8DynamicMoEMethod()
self.w13_kernel = NPUW4A8Int8MoEMethod(
is_per_channel_weight=self.is_per_channel_weight,
activation_use_clip=self.activation_use_clip,
)
self.w2_kernel = NPUW4A8Int8MoEMethod(
is_per_channel_weight=self.is_per_channel_weight,
activation_use_clip=self.activation_use_clip,
)
def create_weights(
self,
@@ -257,37 +265,38 @@ class NPUCompressedTensorsW4A8Int8DynamicMoE(CompressedTensorsMoEScheme):
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
)
self.w13_kernel.process_weights_after_loading(layer, "w13")
self.w2_kernel.process_weights_after_loading(layer, "w2")
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
layer.w13_kernel = self.w13_kernel
layer.w2_kernel = self.w2_kernel
moe_runner_config.layer = layer
self.moe_runner_config = moe_runner_config
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def apply_weights(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
return self.kernel.apply(layer, dispatch_output)
def apply_weights_with_router_logits(
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,
quant_info = AscendQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
w13_weight_offset=layer.w13_weight_offset,
w2_weight_offset=layer.w2_weight_offset,
w13_scale_bias=layer.w13_scale_bias,
w2_scale_bias=layer.w2_scale_bias,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
@@ -6,10 +6,11 @@ from typing import TYPE_CHECKING
import torch
from compressed_tensors.quantization import QuantizationStrategy
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW8A8Int8DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW8A8Int8MoEMethod,
)
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsMoEScheme,
)
@@ -31,7 +32,8 @@ class NPUCompressedTensorsW8A8Int8DynamicMoE(CompressedTensorsMoEScheme):
def __init__(self, weight_quant, input_quant):
self.weight_quant = weight_quant
self.input_quant = input_quant
self.kernel = NPUW8A8Int8DynamicMoEMethod()
self.w13_kernel = NPUW8A8Int8MoEMethod()
self.w2_kernel = NPUW8A8Int8MoEMethod()
self.static_input_scales = not self.input_quant.dynamic
per_channel = (
@@ -118,37 +120,38 @@ class NPUCompressedTensorsW8A8Int8DynamicMoE(CompressedTensorsMoEScheme):
layer.w2_input_scale = None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
self.w13_kernel.process_weights_after_loading(layer, "w13")
self.w2_kernel.process_weights_after_loading(layer, "w2")
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
layer.w13_kernel = self.w13_kernel
layer.w2_kernel = self.w2_kernel
moe_runner_config.layer = layer
self.moe_runner_config = moe_runner_config
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def apply_weights(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
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,
):
# NPU MoE bypasses MoeRunner: expose the kernel's existing
# apply_without_routing_weights directly through the scheme.
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
quant_info = AscendQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
w13_weight_offset=layer.w13_weight_offset,
w2_weight_offset=layer.w2_weight_offset,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
@@ -11,10 +11,11 @@ from compressed_tensors import CompressionFormat
from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import (
gptq_marlin_moe_repack,
)
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A16Int4DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUWNA16Int4MoEMethod,
)
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
WNA16_SUPPORTED_BITS,
CompressedTensorsMoEScheme,
@@ -38,7 +39,6 @@ if TYPE_CHECKING:
CompressedTensorsConfig,
)
__all__ = [
"CompressedTensorsWNA16MoE",
"CompressedTensorsWNA16TritonMoE",
@@ -578,7 +578,8 @@ class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme):
else:
self.group_size = 128
self.kernel = NPUW4A16Int4DynamicMoEMethod()
self.w13_kernel = NPUWNA16Int4MoEMethod()
self.w2_kernel = NPUWNA16Int4MoEMethod()
# TODO: See if we can merge this method's logic
# with CompressedTensorsWNA16MoE. Need more models and tests.
@@ -693,35 +694,38 @@ class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme):
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
self.w13_kernel.process_weights_after_loading(layer, "w13")
self.w2_kernel.process_weights_after_loading(layer, "w2")
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
layer.w13_kernel = self.w13_kernel
layer.w2_kernel = self.w2_kernel
moe_runner_config.layer = layer
self.moe_runner_config = moe_runner_config
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def apply_weights(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
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,
quant_info = AscendQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
w13_weight_offset=layer.w13_weight_offset,
w2_weight_offset=layer.w2_weight_offset,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
+41 -112
View File
@@ -12,8 +12,13 @@ import torch
from gguf import GGMLQuantizationType as WeightType
from torch.nn.parameter import Parameter, UninitializedParameter
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUUnquantMoEMethod,
)
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
LinearMethodBase,
@@ -780,6 +785,8 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase):
def __init__(self, quant_config: GGUFConfig):
self.quant_config = quant_config
self.w13_kernel = NPUUnquantMoEMethod()
self.w2_kernel = NPUUnquantMoEMethod()
def create_weights(
self,
@@ -870,16 +877,18 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase):
torch.from_numpy(dequant_np)
.to(dtype=self.params_dtype, device=w13_qweight.device)
.reshape(rows, cols)
.transpose(-1, -2)
.contiguous()
)
w13_dequant_list.append(dequant)
w13_full = torch.stack(w13_dequant_list, dim=0)
layer.register_buffer("w13_dequant", w13_full, persistent=False)
layer.register_buffer(
"w13_dequant", npu_format_cast(w13_full), persistent=False
)
else:
layer.register_buffer("w13_dequant", w13_qweight.data, persistent=False)
layer.register_buffer(
"w13_dequant", npu_format_cast(w13_qweight.data), persistent=False
)
# Pre-dequantize w2 weights (down projection)
w2_qweight = layer.w2_qweight
@@ -901,137 +910,57 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase):
torch.from_numpy(dequant_np)
.to(dtype=self.params_dtype, device=w2_qweight.device)
.reshape(rows, cols)
.transpose(-1, -2)
.contiguous()
)
w2_dequant_list.append(dequant)
w2_full = torch.stack(w2_dequant_list, dim=0)
layer.register_buffer("w2_dequant", w2_full, persistent=False)
layer.register_buffer(
"w2_dequant", npu_format_cast(w2_full), persistent=False
)
else:
layer.register_buffer("w2_dequant", w2_qweight.data, persistent=False)
layer.register_buffer(
"w2_dequant", npu_format_cast(w2_qweight.data), persistent=False
)
if hasattr(layer, "w2_qweight"):
del layer.w2_qweight
if hasattr(layer, "w13_qweight"):
del layer.w13_qweight
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"quant_type": "gguf"})
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
layer.w13_kernel = self.w13_kernel
layer.w2_kernel = self.w2_kernel
moe_runner_config.layer = layer
moe_runner_config.use_tp_all_gather_activation = True
self.moe_runner_config = moe_runner_config
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
"""Apply MoE forward pass on NPU using npu_grouped_matmul for maximum performance."""
from sglang.srt.distributed.communication_op import (
tensor_model_parallel_all_gather,
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
quant_info = AscendQuantInfo(
w13_weight=layer.w13_dequant,
w2_weight=layer.w2_dequant,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
)
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
# Check if pre-dequantized weights are available
use_pre_dequant = hasattr(layer, "w13_dequant") and hasattr(layer, "w2_dequant")
if not use_pre_dequant:
raise RuntimeError(
"GGUF MoE on NPU requires pre-dequantization (FusedMoE fix). Please report if this occurs."
)
w13 = layer.w13_dequant
w2 = layer.w2_dequant
num_experts = w13.shape[0]
tp_size = getattr(layer, "moe_tp_size", 1)
original_dtype = x.dtype
num_tokens = x.shape[0]
top_k = topk_ids.shape[1]
# Ensure correct dtypes for NPU ops
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(x.dtype)
# MoE routing initialization - reorder tokens by expert
row_idx_len = num_tokens * top_k
row_idx = (
torch.arange(0, row_idx_len, dtype=torch.int32, device=x.device)
.view(top_k, -1)
.permute(1, 0)
.contiguous()
)
sorted_hidden_states, expanded_row_idx, expanded_expert_idx = (
torch.ops.npu.npu_moe_init_routing(
x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
)
)
# Compute tokens per expert
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
)
expert_tokens = expert_tokens.to(torch.int64)
w13_gmm = w13 # No transpose needed
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[sorted_hidden_states],
weight=[w13_gmm],
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# Activation (SwiGLU)
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
# TP all-gather for intermediate dimension if needed
if tp_size > 1:
hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1)
w2_gmm = w2
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w2_gmm],
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# Finalize routing - reorder back and apply weights
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
)
if tp_size > 1:
final_hidden_states = tensor_model_parallel_all_gather(
final_hidden_states, dim=-1
)
# Ensure output matches input dtype
final_hidden_states = final_hidden_states.to(dtype=original_dtype)
return StandardCombineInput(hidden_states=final_hidden_states)
return self.runner.run(dispatch_output, quant_info)
class GGUFEmbeddingAscendMethod(GGUFLinearAscendMethod):
@@ -5,13 +5,20 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUWNA16Int4MoEMethod,
)
from sglang.srt.layers.linear import set_weight_attrs
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from .gptq_scheme import GPTQMoESchemeBase
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig, GPTQMarlinConfig
__all__ = ["GPTQMoEAscendScheme", "GPTQMarlinMoEScheme"]
@@ -122,9 +129,19 @@ class GPTQMoEAscendScheme(GPTQMoESchemeBase):
set_weight_attrs(w2_qzeros, extra_weight_attrs)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
self,
layer: torch.nn.Module,
moe_runner_config: MoeRunnerConfig,
**extra_weight_attrs,
):
self.kernel.create_moe_runner(layer, moe_runner_config)
self.moe_runner_config = moe_runner_config
layer.w13_kernel = NPUWNA16Int4MoEMethod()
layer.w2_kernel = NPUWNA16Int4MoEMethod()
moe_runner_config.layer = layer
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
@@ -133,8 +150,22 @@ class GPTQMoEAscendScheme(GPTQMoESchemeBase):
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
):
return self.kernel.apply(layer, dispatch_output)
) -> CombineInput:
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
quant_info = AscendQuantInfo(
w13_weight=layer.w13_qweight,
w2_weight=layer.w2_qweight,
w13_weight_scale=layer.w13_scales,
w2_weight_scale=layer.w2_scales,
w13_weight_offset=layer.w13_qzeros,
w2_weight_offset=layer.w2_qzeros,
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
class GPTQMarlinMoEScheme(GPTQMoESchemeBase):
@@ -9,6 +9,8 @@ import torch
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
_NPULinearMethodBase,
)
from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
QuantizationConfig,
@@ -34,7 +36,6 @@ if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizeMethodBase
from sglang.srt.layers.quantization.modelslim.schemes import (
ModelSlimLinearScheme,
ModelSlimMoEScheme,
)
logger = logging.getLogger(__name__)
@@ -182,7 +183,14 @@ class ModelSlimConfig(QuantizationConfig):
return UnquantizedLinearMethod()
return ModelSlimLinearMethod(self)
elif isinstance(layer, FusedMoE):
layer.scheme = self.get_moe_scheme(layer, prefix)
moe_schemes = self.get_moe_scheme(layer, prefix)
if moe_schemes is None:
raise ValueError(f"No ModelSlim MoE scheme found for layer {prefix}")
layer.w13_scheme, layer.w2_scheme = moe_schemes
layer.w13_kernel, layer.w2_kernel = (
layer.w13_scheme.kernel,
layer.w2_scheme.kernel,
)
return ModelSlimFusedMoEMethod(self)
return None
@@ -219,29 +227,69 @@ class ModelSlimConfig(QuantizationConfig):
self,
layer: torch.nn.Module,
prefix: str,
) -> Optional[ModelSlimMoEScheme]:
):
moe_quant_schemes = [
("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE),
("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE),
("W8A8_DYNAMIC", ModelSlimW8A8Int8MoE),
]
moe_weight_suffixes = [".0.gate_proj.weight", ".0.w2.weight"]
quant_schemes = [
self.quant_description.get(prefix + suffix, "")
for suffix in moe_weight_suffixes
w13_keys = [
prefix + ".0.gate_proj.weight",
prefix + ".0.up_proj.weight",
]
w2_key = prefix + ".0.down_proj.weight"
w13_entries = {
key: self.quant_description[key]
for key in w13_keys
if key in self.quant_description
}
if not w13_entries or w2_key not in self.quant_description:
missing_groups = []
if not w13_entries:
missing_groups.append(f"W13 ({', '.join(w13_keys)})")
if w2_key not in self.quant_description:
missing_groups.append(f"W2 ({w2_key})")
raise ValueError(
f"Missing ModelSlim MoE quantization description for layer {prefix}: "
+ ", ".join(missing_groups)
)
for scheme_name, scheme_class in moe_quant_schemes:
if any(s == scheme_name for s in quant_schemes):
logger.info_once(f"Using {scheme_class.__name__}")
return scheme_class(self)
w13_names = list(w13_entries.values())
w2_name = self.quant_description[w2_key]
logger.warning(
f"Unsupported FusedMoe modelslim scheme: "
f"{quant_schemes} in layer: {prefix}"
)
return None
# For w13, gate_proj and up_proj must agree on the scheme
unique_w13 = set(w13_names)
if len(unique_w13) > 1:
raise ValueError(
f"Mismatched ModelSlim quantization for W13 in layer {prefix}: "
f"{w13_entries}"
)
w13_scheme_name = w13_names[0]
# Map scheme names to classes
scheme_map = dict(
moe_quant_schemes
) # dict: "W4A4_DYNAMIC" -> ModelSlimW4A4Int4MoE, etc.
# Instantiate the schemes
def instantiate(name, weight_group):
cls = scheme_map.get(name)
if cls is None:
logger.warning(f"Unsupported scheme '{name}' for layer {prefix}")
return None
return cls(self, weight_group)
w13_scheme = instantiate(w13_scheme_name, weight_group="w13")
w2_scheme = instantiate(w2_name, weight_group="w2")
if w13_scheme is None or w2_scheme is None:
raise ValueError(
f"Unsupported ModelSlim MoE schemes for layer {prefix}: "
f"gate/up={w13_names}, down_proj='{w2_name}'"
)
logger.info_once(f"Using {type(w13_scheme).__name__} for gate_up_proj")
logger.info_once(f"Using {type(w2_scheme).__name__} for down_proj")
return w13_scheme, w2_scheme
def is_layer_skipped(
self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
@@ -332,12 +380,19 @@ class ModelSlimLinearMethod(_NPULinearMethodBase):
class ModelSlimFusedMoEMethod(FusedMoEMethodBase):
"""
Fused MoE method for ModelSlim quantization on Ascend NPU.
Delegates routing, activation, and finalization to the modular NPU MoE
components introduced in the hardware backend refactoring.
"""
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)
layer.w13_scheme.process_weights_after_loading(layer)
layer.w2_scheme.process_weights_after_loading(layer)
def create_weights(
self,
@@ -353,50 +408,53 @@ class ModelSlimFusedMoEMethod(FusedMoEMethodBase):
the necessary parameters for the layer. See FusedMoEMethodBase for param
details
"""
layer.scheme.create_weights(
layer.w13_scheme.create_weights(
layer=layer,
num_experts=num_experts,
hidden_size=hidden_size,
intermediate_size_per_partition=intermediate_size_per_partition,
params_dtype=params_dtype,
weight_prefix="w13",
**extra_weight_attrs,
)
layer.w2_scheme.create_weights(
layer=layer,
num_experts=num_experts,
hidden_size=hidden_size,
intermediate_size_per_partition=intermediate_size_per_partition,
weight_prefix="w2",
**extra_weight_attrs,
)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
return layer.scheme.create_moe_runner(layer, moe_runner_config)
moe_runner_config.layer = layer
self.moe_runner_config = moe_runner_config
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
self.runner = MoeRunner(backend, moe_runner_config)
# ------------------------------------------------------------------
# Main apply()
# ------------------------------------------------------------------
def apply(
self,
layer: torch.nn.Module,
layer,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
"""
Use the output of create_weights and the ModelSlimMoEScheme
associated with the layer to apply the forward pass with the
layer input. See FusedMoEMethodBase for param details
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
"""
scheme = layer.scheme
if scheme is None:
raise ValueError("A scheme must be defined for each layer")
return scheme.apply_weights(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return layer.scheme.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
quant_info = AscendQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
w13_weight_offset=layer.w13_weight_offset,
w2_weight_offset=layer.w2_weight_offset,
w13_scale_bias=getattr(layer, "w13_scale_bias", None),
w2_scale_bias=getattr(layer, "w2_scale_bias", None),
w13_weight_bias=getattr(layer, "w13_weight_bias", None),
w2_weight_bias=getattr(layer, "w2_weight_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
@@ -3,16 +3,12 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from typing import TYPE_CHECKING, Optional
from typing import 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"]
@@ -76,26 +72,3 @@ class ModelSlimMoEScheme(BaseMoEScheme):
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
@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPU_W4A4DynamicLinearMethod,
)
@@ -29,6 +30,8 @@ class ModelSlimW4A4Int4(ModelSlimLinearScheme):
def get_weight(
input_size: int, output_size: int, params_dtype: torch.dtype
) -> Dict[str, Any]:
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
output_size = output_size // 2
params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)}
return params_dict
@@ -54,15 +57,25 @@ class ModelSlimW4A4Int4(ModelSlimLinearScheme):
) -> None:
output_size_per_partition = sum(output_partition_sizes)
weight_loader = extra_weight_attrs.get("weight_loader")
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
weight_output_size_per_partition = output_size_per_partition // 2
else:
weight_output_size_per_partition = output_size_per_partition
weight_dict = {
"weight": torch.empty(
output_size_per_partition, input_size_per_partition, dtype=torch.int8
weight_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})
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
set_weight_attrs(
param, {"input_dim": 1, "output_dim": 0, "pack_factor": 2}
)
else:
set_weight_attrs(param, {"input_dim": 1, "output_dim": 0})
layer.register_parameter(weight_name, param)
set_weight_attrs(param, extra_weight_attrs)
@@ -1,23 +1,17 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict
from typing import Any, Dict
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A4Int4DynamicMoEMethod,
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW4A4Int4MoEMethod,
)
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__ = [
@@ -26,14 +20,30 @@ __all__ = [
class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme):
"""
W4A4 integer MoE scheme that creates weights for either the
w13 (gate+up) or w2 (down) projection group.
Two instances of this class are used per MoE layer:
- weight_prefix="w13" → handles the fused gate_proj + up_proj weights
- weight_prefix="w2" → handles the down_proj weights
"""
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
weight_prefix: str, # "w13" or "w2"
group_size: int = 0,
) -> None:
self.quant_config = quant_config
self.kernel = NPUW4A4Int4DynamicMoEMethod()
self.kernel = NPUW4A4Int4MoEMethod()
self.weight_prefix = weight_prefix
self.group_size = group_size
self.is_per_channel_weight = group_size == 0
if weight_prefix not in ("w13", "w2"):
raise ValueError(
f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'"
)
def create_weights(
self,
@@ -41,103 +51,87 @@ class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme):
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
"""
Create and register weight, scale, and offset parameters for the layer.
Shape depends on the W4A4 packing environment flag and whether the weight
prefix is "w13" or "w2".
"""
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# --- compute shapes based on the packing path and prefix ---
if envs.SGLANG_NPU_W4A4_NEW_PACKING.get():
if self.weight_prefix == "w13":
out_features = intermediate_size_per_partition
in_features = hidden_size
else: # w2
out_features = hidden_size // 2
in_features = intermediate_size_per_partition
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=torch.int8,
),
weight_shape = (num_experts, out_features, in_features)
scale_shape = (num_experts, 2 * out_features, 1)
else:
if self.weight_prefix == "w13":
a_dim = 2 * intermediate_size_per_partition
b_dim = hidden_size
else: # w2
a_dim = hidden_size
b_dim = intermediate_size_per_partition
weight_shape = (num_experts, a_dim, b_dim)
scale_shape = (num_experts, a_dim, 1)
offset_shape = scale_shape # offset always matches scale
self._create_weight_params(
layer,
self.weight_prefix,
weight_shape,
scale_shape,
offset_shape,
extra_weight_attrs,
)
@staticmethod
def _create_weight_params(
layer: torch.nn.Module,
prefix: str,
weight_shape: tuple,
scale_shape: tuple,
offset_shape: tuple,
extra_weight_attrs: dict,
) -> None:
"""Helper that registers weight, scale, and offset as parameters."""
# Weight
weight = torch.nn.Parameter(
torch.empty(weight_shape, 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,
),
layer.register_parameter(f"{prefix}_weight", weight)
set_weight_attrs(weight, extra_weight_attrs)
# Scale
scale = torch.nn.Parameter(
torch.empty(scale_shape, dtype=torch.float32),
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
),
layer.register_parameter(f"{prefix}_weight_scale", scale)
set_weight_attrs(scale, extra_weight_attrs)
# Offset
offset = torch.nn.Parameter(
torch.empty(offset_shape, 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)
layer.register_parameter(f"{prefix}_weight_offset", offset)
set_weight_attrs(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,
)
"""
Delegate weight processing to the NPU kernel for the fixed weight group.
"""
self.kernel.process_weights_after_loading(layer, self.weight_prefix)
@@ -1,43 +1,55 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict
from typing import Any, Dict
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int8DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW4A8Int8MoEMethod,
)
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",
]
__all__ = ["ModelSlimW4A8Int8MoE"]
class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme):
"""
W4A8 MoE scheme for a single weight group (w13 or w2).
Two instances of this class are created per MoE layer:
- weight_prefix="w13" → handles gate + up projections
- weight_prefix="w2" → handles down projection
Configuration flags (``is_per_channel_weight``, ``activation_use_clip``)
are passed to the underlying NPU kernel.
"""
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
weight_prefix: str,
group_size: int = 0,
tp_size: int = 1,
activation_use_clip: bool = False,
) -> None:
if weight_prefix not in ("w13", "w2"):
raise ValueError(
f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'"
)
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()
self.weight_prefix = weight_prefix
self.group_size = group_size
self.tp_size = tp_size
self.is_per_channel_weight = group_size == 0
self.activation_use_clip = activation_use_clip
self.kernel = NPUW4A8Int8MoEMethod(
is_per_channel_weight=self.is_per_channel_weight,
activation_use_clip=self.activation_use_clip,
)
def create_weights(
self,
@@ -45,173 +57,89 @@ class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme):
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)
# Determine dimensions based on weight group
if self.weight_prefix == "w13":
out_features = intermediate_size_per_partition
in_features = hidden_size
bias_last_dim = 1
else: # w2
out_features = hidden_size // 2
in_features = intermediate_size_per_partition
bias_last_dim = 16 // self.tp_size
# >> scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
prefix = self.weight_prefix
# ---- weight ----
weight = torch.nn.Parameter(
torch.empty(num_experts, out_features, in_features, dtype=torch.int8),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
layer.register_parameter(f"{prefix}_weight", weight)
set_weight_attrs(weight, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
# ---- scale ----
scale = torch.nn.Parameter(
torch.empty(num_experts, 2 * out_features, 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)
layer.register_parameter(f"{prefix}_weight_scale", scale)
set_weight_attrs(scale, extra_weight_attrs)
# >> offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
# ---- offset ----
offset = torch.nn.Parameter(
torch.empty(num_experts, 2 * out_features, 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)
layer.register_parameter(f"{prefix}_weight_offset", offset)
set_weight_attrs(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
# ---- per‑group second scale/offset (when not per‑channel) ----
if not self.is_per_channel_weight:
w13_weight_scale_second = torch.nn.Parameter(
scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
2 * out_features,
in_features // 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)
layer.register_parameter(f"{prefix}_weight_scale_second", scale_second)
set_weight_attrs(scale_second, extra_weight_attrs)
w2_weight_scale_second = torch.nn.Parameter(
offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
2 * out_features,
in_features // 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)
layer.register_parameter(f"{prefix}_weight_offset_second", offset_second)
set_weight_attrs(offset_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(
# ---- bias for scale (activation clip path) ----
# This parameter is always created; the kernel uses it only when activation_use_clip is True.
scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
num_experts, 2 * out_features, bias_last_dim, 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)
layer.register_parameter(f"{prefix}_scale_bias", scale_bias)
set_weight_attrs(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,
)
"""Delegate weight processing to the kernel for the assigned weight group."""
self.kernel.process_weights_after_loading(layer, self.weight_prefix)
@@ -1,23 +1,16 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict
from typing import Any, Dict
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW8A8Int8DynamicMoEMethod,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW8A8Int8MoEMethod,
)
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__ = [
@@ -26,14 +19,27 @@ __all__ = [
class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme):
"""
W8A8 integer MoE scheme that creates weights for either the
w13 (gate+up) or w2 (down) projection group.
Two instances of this class are used per MoE layer:
- weight_prefix="w13" → handles the fused gate_proj + up_proj weights
- weight_prefix="w2" → handles the down_proj weights
"""
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
weight_prefix: str, # "w13" or "w2"
) -> None:
self.quant_config = quant_config
self.kernel = NPUW8A8Int8DynamicMoEMethod()
self.kernel = NPUW8A8Int8MoEMethod()
self.weight_prefix = weight_prefix
if weight_prefix not in ("w13", "w2"):
raise ValueError(
f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'"
)
def create_weights(
self,
@@ -41,7 +47,6 @@ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme):
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
@@ -51,89 +56,45 @@ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme):
{"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,
),
# Determine shape based on weight group
if self.weight_prefix == "w13":
a_dim = 2 * intermediate_size_per_partition
b_dim = hidden_size
else: # w2
a_dim = hidden_size
b_dim = intermediate_size_per_partition
prefix = self.weight_prefix
# Create and register weight
weight_name = f"{prefix}_weight"
weight = torch.nn.Parameter(
torch.empty(num_experts, a_dim, b_dim, 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,
),
layer.register_parameter(weight_name, weight)
set_weight_attrs(weight, extra_weight_attrs)
# Create and register scale
scale_name = f"{prefix}_weight_scale"
scale = torch.nn.Parameter(
torch.empty(num_experts, a_dim, 1, dtype=torch.float32),
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
),
layer.register_parameter(scale_name, scale)
set_weight_attrs(scale, extra_weight_attrs)
# Create and register offset
offset_name = f"{prefix}_weight_offset"
offset = torch.nn.Parameter(
torch.empty(num_experts, a_dim, 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)
layer.register_parameter(offset_name, offset)
set_weight_attrs(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,
)
"""
Delegate weight processing to the NPU kernel for the fixed weight group.
"""
self.kernel.process_weights_after_loading(layer, self.weight_prefix)
+13 -151
View File
@@ -49,6 +49,9 @@ if TYPE_CHECKING:
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUUnquantMoEMethod,
)
_is_cpu_amx_available = cpu_has_amx_support()
_is_hip = is_hip()
@@ -60,9 +63,6 @@ if _use_aiter:
from aiter.ops.shuffle import shuffle_weight
from aiter.tuned_gemm import tgemm
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
class Bf16GemmBackend(Enum):
AUTO = "auto"
@@ -403,9 +403,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
layer.num_local_experts, *new_shape_w2
)
if _is_npu:
for weight_name in ["w13_weight", "w2_weight"]:
weight = getattr(layer, weight_name)
weight.data = npu_format_cast(weight)
layer.w13_kernel.process_weights_after_loading(layer, "w13")
layer.w2_kernel.process_weights_after_loading(layer, "w2")
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
return
@@ -474,6 +475,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
backend = MoeRunnerBackend.DEEP_GEMM
elif self.use_triton_kernels:
backend = MoeRunnerBackend.TRITON_KERNELS
elif _is_npu:
layer.w13_kernel = NPUUnquantMoEMethod()
layer.w2_kernel = NPUUnquantMoEMethod()
moe_runner_config.layer = layer
backend = MoeRunnerBackend.ASCEND
else:
backend = MoeRunnerBackend.TRITON
self.runner = MoeRunner(backend, moe_runner_config)
@@ -721,151 +727,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
dispatch_output: DispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
if DispatchOutputChecker.format_is_deepep(dispatch_output):
return self._forward_npu_deepep(layer, dispatch_output)
# x.shape = [B*S, H]
x = dispatch_output.hidden_states
# topk_weights.shape = [B*S, K]; topk_ids.shape = [B*S, K]
topk_weights, topk_ids, _ = dispatch_output.topk_output
original_dtype = x.dtype
num_tokens = x.shape[0]
topk_weights = topk_weights.to(x.dtype)
topk_ids = topk_ids.to(torch.int32)
num_experts = layer.num_experts
top_k = layer.top_k or topk_ids.shape[1] # in case layer.top_k is not set
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(
x,
topk_ids,
active_num=num_tokens * top_k,
expert_num=num_experts,
expert_tokens_num_type=1,
expert_tokens_num_flag=True,
active_expert_range=[0, num_experts],
quant_mode=-1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
w13_bias = [layer.w13_weight_bias] if self.with_bias else None
w2_bias = [layer.w2_weight_bias] if self.with_bias else None
# gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight.transpose(1, 2)],
bias=w13_bias,
split_item=2,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# act_fn:
if self.moe_runner_config.activation == "npu_swiglu_oai":
from sgl_kernel_npu.activation.swiglu_oai import swiglu_oai_triton
# `hidden_states` is the gmm1 output of shape [num_tokens, 2 * inter].
# Pass the gate_up dim from the activation itself instead of letting
# swiglu_oai() derive it from layer.w13_weight.shape[2]: w13_weight is
# now stored un-transposed (transposed on the fly for the grouped
# matmuls above), so shape[2] is `hidden`, not the gate_up dim, which
# makes the kernel's view(-1, dim) reshape fail.
hidden_states = swiglu_oai_triton(
hidden_states,
hidden_states.shape[-1],
self.moe_runner_config.gemm1_alpha,
self.moe_runner_config.gemm1_clamp_limit,
)
elif self.moe_runner_config.activation == "silu":
if self.moe_runner_config.gemm1_clamp_limit is not None:
from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant
hidden_states, _ = swiglu_quant(
hidden_states,
group_list=expert_tokens,
group_list_type=1,
need_quant=False,
do_limit=True,
limit=self.moe_runner_config.gemm1_clamp_limit,
)
else:
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
else:
from sglang.srt.layers.activation import GeluAndMul
hidden_states = GeluAndMul()(hidden_states)
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight.transpose(1, 2)],
bias=w2_bias,
split_item=2,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
drop_pad_mode=2,
)
return StandardCombineInput(hidden_states=final_hidden_states)
def _forward_npu_deepep(
self,
layer: torch.nn.Module,
dispatch_output: DispatchOutput,
) -> CombineInput:
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
npu_fused_moe_without_routing_weights_bf16,
)
from sglang.srt.layers.moe.token_dispatcher import (
DeepEPLLCombineInput,
DeepEPNormalCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
# NOTE: Ascend's Dispatch & Combine does not support FP16
output_dtype = torch.bfloat16
group_list_type = 1
if DispatchOutputChecker.format_is_deepep_normal(dispatch_output):
hidden_states, _, _, _, num_recv_tokens_per_expert = dispatch_output
group_list = torch.tensor(
num_recv_tokens_per_expert,
dtype=torch.int64,
device=hidden_states.device,
)
combine_cls = DeepEPNormalCombineInput
else:
hidden_states, _, _, _, group_list, _ = dispatch_output
group_list = group_list.to(torch.int64)
combine_cls = DeepEPLLCombineInput
hidden_states = npu_fused_moe_without_routing_weights_bf16(
layer, hidden_states, group_list_type, group_list, output_dtype
)
return combine_cls(
hidden_states=hidden_states,
topk_ids=dispatch_output.topk_ids,
topk_weights=dispatch_output.topk_weights,
)
return self.runner.run(dispatch_output, layer)
def forward_tpu(self, *args, **kwargs) -> CombineInput:
raise NotImplementedError("The TPU backend currently does not support MoE.")
+1 -1
View File
@@ -62,7 +62,7 @@ from sglang.srt.utils import add_prefix, is_npu
_is_npu = is_npu()
if _is_npu:
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
fused_moe_npu as fused_moe,
)
+1 -1
View File
@@ -27,7 +27,7 @@ from sglang.srt.configs import DbrxConfig
from sglang.srt.distributed import (
tensor_model_parallel_all_reduce,
)
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
fused_moe_npu,
)
from sglang.srt.layers.linear import (
+1 -1
View File
@@ -59,7 +59,7 @@ if _is_cpu and _is_cpu_amx_available:
import sgl_kernel # noqa: F401
if _is_npu:
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
fused_moe_npu as fused_moe,
)
else:
+1 -1
View File
@@ -22,7 +22,7 @@ from transformers import PretrainedConfig
from sglang.srt.distributed import (
tensor_model_parallel_all_reduce,
)
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
fused_moe_npu,
)
from sglang.srt.layers.activation import SiluAndMul
+42 -11
View File
@@ -274,6 +274,7 @@ MOE_A2A_BACKEND_CHOICES = [
"ascend_fuseep",
"flashinfer",
"megamoe",
"ascend_tp",
]
MXFP8_MOE_RUNNER_BACKEND_CHOICES = [
@@ -1901,6 +1902,7 @@ class ServerArgs:
"ascend_fuseep",
"flashinfer",
"megamoe",
"ascend_tp",
],
Arg(
help="Choose the backend for MoE A2A.",
@@ -1924,6 +1926,10 @@ class ServerArgs:
Literal["auto", "normal", "low_latency"],
"Select the mode when enable DeepEP or MoriEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.",
] = "auto"
fuseep_mode: A[
Literal[1, 2],
"Select the mode when enable Ascend FuseEP MoE, 1 -> dispatch_gmm_combine_decode is executed;2 -> dispatch_ffn_combine is executed (support hybrid deployment when 2).",
] = 2
deepep_dispatcher_output_dtype: A[
Literal["auto", "bf16", "fp8", "int8", "nvfp4"],
"Select DeepEP dispatcher output dtype",
@@ -5558,20 +5564,17 @@ class ServerArgs:
f"Nixl MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
)
if a2a_backend == "ascend_fuseep":
if (
self.moe_a2a_backend == "none" and is_npu()
) or self.moe_a2a_backend == "ascend_tp":
# FIXME (OrangeRedeng): for some reasons if pass "ascend_tp" accuracy drops to zero
self.moe_a2a_backend = "none"
if self.moe_a2a_backend == "ascend_fuseep":
logger.warning(
f"Ascend fused EP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
)
fuse_mode = envs.SGLANG_NPU_FUSED_MOE_MODE.get()
if fuse_mode not in [1, 2]:
raise ValueError(
f"Wrong value of {fuse_mode=}, the NPU only support 1 or 2."
)
elif fuse_mode == 2:
assert (
resolved_view(self).quantization == "modelslim"
), "When fuse_mode is set to 2, the NPU supports only ModelSlim quantization."
if a2a_backend == "flashinfer":
if self.moe_a2a_backend == "flashinfer":
assert (
resolved_view(self).enable_dp_attention and self.dp_size == self.tp_size
), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention"
@@ -7684,6 +7687,32 @@ def get_global_server_args() -> ServerArgs:
return get_context().server_args
def _has_cli_arg(argv: List[str], flag: str) -> bool:
return any(arg == flag or arg.startswith(f"{flag}=") for arg in argv)
def _apply_fuseep_mode_env_compat(
raw_args: argparse.Namespace, argv: List[str]
) -> None:
if not envs.SGLANG_NPU_FUSED_MOE_MODE.is_set() or _has_cli_arg(
argv, "--fuseep-mode"
):
return
fuseep_mode = envs.SGLANG_NPU_FUSED_MOE_MODE.get()
if fuseep_mode not in (1, 2):
raise ValueError(
f"Wrong value of SGLANG_NPU_FUSED_MOE_MODE={fuseep_mode}, "
"the NPU only supports 1 or 2."
)
logger.warning(
"The env variable SGLANG_NPU_FUSED_MOE_MODE is deprecated and will be "
"removed in a future release. Please use --fuseep-mode instead."
)
raw_args.fuseep_mode = fuseep_mode
def prepare_server_args(argv: List[str]) -> ServerArgs:
"""
Prepare the server arguments from the command line arguments.
@@ -7718,6 +7747,8 @@ def prepare_server_args(argv: List[str]) -> ServerArgs:
force=True,
)
_apply_fuseep_mode_env_compat(raw_args, argv)
return ServerArgs.from_cli_args(raw_args)
@@ -77,6 +77,9 @@ DOTS_OCR_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "rednote-hilab/dots.ocr"
ECO_TECH_QWEN3_32B_W4A4_LAOS_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-32B-w4a4-LAOS"
)
ECO_TECH_QWEN3_30B_A3B_W4A4_LAOS_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS"
)
ERNIE_4_5_21B_A3B_PT_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "baidu/ERNIE-4.5-21B-A3B-PT"
)
@@ -301,9 +304,6 @@ DEEPSEEK_R1_0528_W4A8_PER_CHANNEL_WEIGHTS_PATH = os.path.join(
DEEPSEEK_R1_0528_W8A8_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "vllm-ascend/DeepSeek-R1-0528-W8A8"
)
QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS"
)
QWEN3_5_397B_W4A8_MODEL_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3.5-397B-A17B-w4a8-mtp"
)