[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:
co-authored by
ronnie_zheng
parent
c9b17403e7
commit
8ed82afcc8
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
+34
-25
@@ -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)
|
||||
|
||||
+28
-25
@@ -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)
|
||||
|
||||
+28
-24
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+91
-97
@@ -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)
|
||||
|
||||
+78
-150
@@ -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)
|
||||
|
||||
+53
-92
@@ -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)
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user