minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU]Refactor weight processing and add NPUSwigluLimit activation (#38420)

Browse Source 
Co-authored-by: AndyLi429 <AndyLi429@noreply.gitcode.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
AndyLi429
2026-09-17 16:40:32 +08:00
committed by GitHub
co-authored by AndyLi429 Even Zhou
parent 2d08cc5ede
commit e970453b43
6 changed files with 365 additions and 892 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/moe/activation.py
+29
View File
@@ -201,6 +201,35 @@ class NPUSwigluStepAndMul(BaseActivation):
return gate * up
class NPUSwigluMxfp8Quant(BaseActivation):
"""DeepSeek-V4 grouped SwiGLU with MXFP8 requantization for GMM2."""
def __init__(self, limit: float):
self._limit = float(limit)
def _apply_activation(
self,
hidden_states: torch.Tensor,
group_list: torch.Tensor,
group_list_type: int,
):
# The op sums the group list as per-expert counts and has no cumulative layout;
# a cusum list passed through would silently process the wrong rows.
if group_list_type != 1:
raise ValueError(
"swiglu_group_quant takes a per-expert count group list, got "
f"group_list_type={group_list_type}"
)
out, scale, _ = torch.ops.npu.swiglu_group_quant(
x=hidden_states,
group_index=group_list,
quant_mode=2, # MX: one e8m0 scale per 32-element block
group_list_type=0, # sglang numbers the count layout 1, the op numbers it 0
clamp_value=self._limit,
)
return out, scale
# =============================================================================
# Generic TP all‑gather wrapper – used by the runner when needed
# =============================================================================
python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py
+66 -504
View File
@@ -1,69 +1,28 @@
"""MXFP4 routed-expert MoE method for Ascend A5 (Ascend 950).
DeepSeek-V4's FP4 expert checkpoint stores block-32 MXFP4 weights with E8M0
scales. This module wires those weights to the A5 grouped-matmul kernels, both
for the plain (init-routing) path and for the DeepEP dispatch path.
scales. This module adapts those checkpoint weights to the shared Ascend MoE
runner and A5 grouped-matmul kernels.
"""
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING
import torch
from sgl_kernel_npu.activation.swiglu_mxfp8_quant import swiglu_quant
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
_get_float4_e2m1fn_x2_dtype,
_get_float8_e8m0fnu_dtype,
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUW4A8MXFP4MoEMethod,
prepare_w4a8_mxfp_weight,
)
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
from sglang.srt.layers.moe.token_dispatcher import DispatchOutput
# MXFP4 group size, fixed at 32 by the msmodelslim export format.
MXFP4_BLOCK_SIZE = 32
def _configure_dsv4_deepep_dispatcher(layer: torch.nn.Module) -> None:
"""Select the DSV4 FP4 DeepEP wire format without changing other MoEs."""
dispatcher = getattr(layer, "dispatcher", None)
if dispatcher is None:
return
# This method is only instantiated for DSV4 FP4 experts on A5 today, but
# retain the former BF16 setting if that selection changes in the future.
if not is_npu_arch35():
dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
return
# Import lazily to avoid importing the MoE backend during quant method
# module initialization.
from sglang.srt.layers.moe import get_moe_a2a_backend
if not get_moe_a2a_backend().is_deepep():
dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
return
low_latency_dtype = envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.get()
if low_latency_dtype not in {"mxfp8", "bf16"}:
raise ValueError(
"SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE must be one of "
"'mxfp8' or 'bf16' for A5 DSV4 DeepEP low-latency dispatch; "
f"got {low_latency_dtype!r}."
)
# The concrete dispatcher selects one mode-specific value. Normal (prefill)
# remains BF16, while low-latency (decode) defaults to MXFP8.
dispatcher.set_quant_config(
{
"normal_dispatcher_output_dtype": "bf16",
"low_latency_dispatcher_output_dtype": low_latency_dtype,
}
)
def _wrap_mxfp4_scale_weight_loader(weight_loader):
def load_scale(param, loaded_weight, *args, **kwargs):
if param.dtype == torch.uint8 and loaded_weight.dtype == torch.float8_e8m0fnu:
@@ -73,17 +32,25 @@ def _wrap_mxfp4_scale_weight_loader(weight_loader):
return load_scale
class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase):
class NPUW4A8MXFP4FusedMoEMethod(FusedMoEMethodBase):
"""DeepSeek-V4 routed experts on Ascend A5: W4A8 MXFP weights.
Delegates nothing to ``fp8_method`` except the shared runner config; it is
held so the FP8 method sees the same ``moe_runner_config`` the layer built.
The checkpoint-specific loading remains here while execution is delegated
to the shared Ascend MoE runner.
"""
def __init__(self, fp8_method, prefix: str = ""):
self._fp8 = fp8_method
def __init__(self, prefix: str = ""):
self.prefix = prefix
self.moe_runner_config = None
# ``None`` selects the full MX dynamic-quant defaults used by the
# original DeepSeek-V4 path; the shared ModelSlim path keeps its
# historical explicit ``dst_type`` behavior.
# TODO: Fuse DeepSeek-V4 W13 GMM + SwiGLU + MXFP8 requant with
# npu_grouped_matmul_swiglu_quant_v2 once it accepts swiglu_limit.
# V4 sets swiglu_limit=10.0; the current fused op implements only
# standard SwiGLU and would skip the required gate/up clamps.
self.w13_kernel = NPUW4A8MXFP4MoEMethod(dynamic_quant_kwargs=None)
self.w2_kernel = NPUW4A8MXFP4MoEMethod(dynamic_quant_kwargs=None)
self.runner = None
def create_weights(
self,
@@ -153,11 +120,27 @@ class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase):
set_weight_attrs(w2_weight_scale, scale_attrs)
def create_moe_runner(self, layer: torch.nn.Module, moe_runner_config):
self.moe_runner_config = moe_runner_config
self._fp8.moe_runner_config = moe_runner_config
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.utils import (
MoeRunnerBackend,
get_moe_runner_backend,
)
backend = get_moe_runner_backend()
if backend.is_auto():
backend = MoeRunnerBackend.ASCEND
if not backend.is_ascend():
raise ValueError(
f"NPU W4A8 MXFP4 requires the Ascend MoE runner, got {backend.value}"
)
layer.w13_kernel = self.w13_kernel
layer.w2_kernel = self.w2_kernel
moe_runner_config.layer = layer
self.runner = MoeRunner(backend, moe_runner_config)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast
from sglang.srt.hardware_backend.npu.utils import NPUACLFormat
if layer.w13_weight_scale_inv.data.max() == 0:
raise RuntimeError(
@@ -172,459 +155,38 @@ class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase):
"not match w2_weight_scale_inv."
)
nz_kwargs = {
"customize_dtype": torch.float8_e4m3fn,
"input_dtype": _get_float4_e2m1fn_x2_dtype(),
}
nz_format = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ
layer.w13_weight.data = npu_format_cast(
layer.w13_weight.data.view(torch.uint8), nz_format, **nz_kwargs
).transpose(1, 2)
layer.w2_weight.data = npu_format_cast(
layer.w2_weight.data.view(torch.uint8), nz_format, **nz_kwargs
).transpose(1, 2)
layer.w13_weight_scale_inv = torch.nn.Parameter(
_reshape_mxfp4_scale_for_npu(layer.w13_weight_scale_inv.data),
requires_grad=False,
layer.w13_weight.data, w13_scale = prepare_w4a8_mxfp_weight(
layer.w13_weight.data.view(torch.uint8),
layer.w13_weight_scale_inv.data,
npu_format=nz_format,
)
layer.w2_weight_scale_inv = torch.nn.Parameter(
_reshape_mxfp4_scale_for_npu(layer.w2_weight_scale_inv.data),
requires_grad=False,
layer.w2_weight.data, w2_scale = prepare_w4a8_mxfp_weight(
layer.w2_weight.data.view(torch.uint8),
layer.w2_weight_scale_inv.data,
npu_format=nz_format,
)
layer.w13_weight_scale_inv = torch.nn.Parameter(w13_scale, requires_grad=False)
layer.w2_weight_scale_inv = torch.nn.Parameter(w2_scale, requires_grad=False)
_configure_dsv4_deepep_dispatcher(layer)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> "CombineInput":
combine_input = npu_apply_w4a8_mxfp_moe_deepep(layer, dispatch_output)
if combine_input is not None:
return combine_input
combine_input = npu_apply_w4a4_mxfp_moe_ascend_tp(layer, dispatch_output)
if combine_input is not None:
return combine_input
# Standard dispatch. Unreachable on NPU today — create_moe_dispatcher
# picks AscendTPDispatcher whenever is_npu() and no a2a backend is set —
# but kept so this method is not silently wrong if that changes.
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
hidden_states = dispatch_output.hidden_states
topk_weights, topk_ids, _ = dispatch_output.topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(hidden_states.dtype)
moe_runner_config = layer.moe_runner_config
output = npu_fused_experts_w4a4_mxfp(
hidden_states,
layer.w13_weight,
layer.w13_weight_scale_inv,
layer.w2_weight,
layer.w2_weight_scale_inv,
topk_weights,
topk_ids,
moe_runner_config.top_k,
swiglu_limit=moe_runner_config.swiglu_limit,
)
return StandardCombineInput(hidden_states=output)
def _reshape_mxfp4_scale_for_npu(scale: torch.Tensor) -> torch.Tensor:
"""``[E, N, K/32] -> [E, K/64, N, 2]``, the packed-pair layout the GMM wants."""
if scale.dim() != 3:
return scale
num_experts, n, k32 = scale.shape
if k32 % 2 != 0:
raise ValueError(
"MXFP4 scale K dimension must be divisible by 2 for the "
f"[E, K/64, N, 2] layout, got {tuple(scale.shape)}."
)
return scale.view(num_experts, n, k32 // 2, 2).transpose(1, 2)
def _apply_swiglu_limit_npu(
gate_up: torch.Tensor, swiglu_limit: Optional[float]
) -> None:
"""Clamp the SwiGLU input in place before ``npu_swiglu`` (DeepSeek-V4).
gate (first half) <= limit; up (second half) in
[-limit, limit]. ``chunk`` returns views, so the in-place clamps mutate
``gate_up`` directly. No-op when ``swiglu_limit`` is unset or <= 0.
"""
if swiglu_limit is None or swiglu_limit <= 0:
return
gate, up = gate_up.chunk(2, dim=-1)
gate.clamp_(max=swiglu_limit)
up.clamp_(min=-swiglu_limit, max=swiglu_limit)
def npu_fused_experts_w4a4_mxfp(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_weight_scale_inv: torch.Tensor,
w2: torch.Tensor,
w2_weight_scale_inv: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
swiglu_limit: Optional[float] = None,
**kwargs,
):
if torch.npu.is_current_stream_capturing():
return npu_fused_experts_w4a4_mxfp_decode(
hidden_states=hidden_states,
w13=w13,
w13_weight_scale_inv=w13_weight_scale_inv,
w2=w2,
w2_weight_scale_inv=w2_weight_scale_inv,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=top_k,
swiglu_limit=swiglu_limit,
**kwargs,
)
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
if len(original_shape) == 3:
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
num_tokens = hidden_states.shape[0]
num_experts = w13.shape[0]
row_idx = (
torch.arange(
0, num_tokens * top_k, 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(
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
).to(torch.int64)
# npu_moe_init_routing pads its output to the worst case; rows past the last
# expert boundary hold garbage and must not reach finalize_routing.
row_ids = torch.arange(
hidden_states.shape[0], device=hidden_states.device, dtype=torch.int64
)
valid_mask_2d = (row_ids < expert_tokens[-1]).unsqueeze(1)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=None,
weight=w13,
weight_scale=w13_weight_scale_inv,
group_list_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)
assert swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=expert_tokens,
group_list_type=0,
need_quant=True,
do_limit=True,
limit=swiglu_limit,
)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=w2,
weight_scale=w2_weight_scale_inv,
group_list_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)
hidden_states = hidden_states * valid_mask_2d.to(hidden_states.dtype)
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 len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_fused_experts_w4a4_mxfp_decode(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_weight_scale_inv: torch.Tensor,
w2: torch.Tensor,
w2_weight_scale_inv: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
swiglu_limit: Optional[float] = None,
**kwargs,
):
"""Graph-capturable variant: routing v2 + token_unpermute, no host syncs."""
num_tokens = hidden_states.shape[:-1].numel()
global_num_experts = w13.shape[0]
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
group_list_type = 1
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(
hidden_states,
topk_ids,
active_num=num_tokens * top_k,
expert_num=global_num_experts,
expert_tokens_num_type=group_list_type,
expert_tokens_num_flag=True,
active_expert_range=[0, global_num_experts],
quant_mode=-1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=None,
weight=w13,
weight_scale=w13_weight_scale_inv,
group_list_type=group_list_type,
group_list=expert_tokens,
output_dtype=original_dtype,
)
assert swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=expert_tokens,
group_list_type=group_list_type,
need_quant=True,
do_limit=True,
limit=swiglu_limit,
)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=w2,
weight_scale=w2_weight_scale_inv,
group_list_type=group_list_type,
group_list=expert_tokens,
output_dtype=original_dtype,
)
final_hidden_states = torch.ops.npu.npu_moe_token_unpermute(
permuted_tokens=hidden_states,
sorted_indices=torch.abs(expanded_row_idx),
probs=topk_weights,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_apply_w4a4_mxfp_moe_ascend_tp(
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> Optional["CombineInput"]:
"""Ascend TP path. Returns ``None`` when the dispatch is not an Ascend TP one.
AscendTPDispatcher already ran npu_moe_init_routing_v2 on dispatch and runs
npu_moe_finalize_routing (with topk_weights) on combine, so this only owns
the grouped-matmul chain in between — no permute, no routing-weight apply.
"""
from sglang.srt.layers.moe.token_dispatcher import AscendTPCombineInput
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
if not DispatchOutputChecker.format_is_ascend_tp(dispatch_output):
return None
hidden_states = npu_apply_without_routing_weights_w4a4_mxfp(
layer,
dispatch_output.hidden_states,
dispatch_output.hidden_states_scale,
group_list_type=dispatch_output.group_list_type,
group_list=dispatch_output.expert_tokens,
output_dtype=torch.bfloat16,
)
return AscendTPCombineInput(hidden_states=hidden_states)
def npu_apply_w4a8_mxfp_moe_deepep(
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> Optional["CombineInput"]:
"""DeepEP path. Returns ``None`` when the dispatch is not a DeepEP one."""
from sglang.srt.layers.moe.token_dispatcher import (
DeepEPLLCombineInput,
DeepEPNormalCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
if not dispatch_output.format.is_deepep():
return None
if DispatchOutputChecker.format_is_deepep_normal(dispatch_output):
hidden_states, hidden_states_scale, _, _, 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, hidden_states_scale, _, _, group_list, _ = dispatch_output
group_list = group_list.to(torch.int64)
combine_cls = DeepEPLLCombineInput
hidden_states = npu_apply_without_routing_weights_w4a4_mxfp(
layer,
hidden_states,
hidden_states_scale,
group_list_type=1,
group_list=group_list,
output_dtype=torch.bfloat16,
)
return combine_cls(
hidden_states=hidden_states,
topk_ids=dispatch_output.topk_ids,
topk_weights=dispatch_output.topk_weights,
)
def npu_apply_without_routing_weights_w4a4_mxfp(
layer,
hidden_states,
hidden_states_scale,
*,
group_list_type,
group_list,
output_dtype,
):
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=layer.w13_weight,
weight_scale=layer.w13_weight_scale_inv,
group_list_type=group_list_type,
group_list=group_list,
output_dtype=output_dtype,
)
assert layer.moe_runner_config.swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=group_list,
group_list_type=group_list_type,
need_quant=True,
do_limit=True,
limit=layer.moe_runner_config.swiglu_limit,
)
return w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=layer.w2_weight,
weight_scale=layer.w2_weight_scale_inv,
group_list_type=group_list_type,
group_list=group_list,
output_dtype=output_dtype,
)
def _pair_pack_mxfp_act_scale(
scale: torch.Tensor, input_shape: Optional[tuple[int, int]] = None
) -> torch.Tensor:
"""Adapt MXFP activation scales to the A5 GMM ``[M, K/64, 2]`` layout.
Low-latency DeepEP MXFP8 returns a flat E8M0 scale buffer, one byte for
every 32 activation elements. The grouped-matmul kernel expects those
bytes paired on the final dimension instead.
"""
if scale.ndim == 1:
if input_shape is None or len(input_shape) != 2:
raise ValueError(
"A flat MXFP activation scale requires its two-dimensional "
"activation input shape."
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config(
{
"normal_dispatcher_output_dtype": "bf16",
"low_latency_dispatcher_output_dtype": "mxfp8",
}
)
num_tokens, hidden_size = input_shape
if hidden_size % (2 * MXFP4_BLOCK_SIZE) != 0:
raise ValueError(
"MXFP activation hidden size must be divisible by "
f"{2 * MXFP4_BLOCK_SIZE}; got {hidden_size}."
)
expected_num_scales = num_tokens * (hidden_size // MXFP4_BLOCK_SIZE)
if scale.numel() != expected_num_scales:
raise ValueError(
"Invalid flat MXFP activation scale length: expected "
f"{expected_num_scales} for input shape {input_shape}, got "
f"{scale.numel()}."
)
scale = scale.reshape(num_tokens, hidden_size // MXFP4_BLOCK_SIZE)
# ``[M, K/32] -> [M, K/64, 2]`` MX per-token scale layout for the A5 GMM.
if scale.ndim != 2:
return scale
if scale.shape[-1] % 2 != 0:
raise ValueError(f"Invalid MXFP per-token scale shape: {tuple(scale.shape)}")
return scale.reshape(scale.shape[0], scale.shape[1] // 2, 2)
def apply(self, layer: torch.nn.Module, dispatch_output: "DispatchOutput"):
from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
if self.runner is None:
raise RuntimeError("The NPU FP4 MoE runner has not been initialized")
def w4a8_mxfp_gmm(
*,
input: torch.Tensor,
input_scale: Optional[torch.Tensor],
weight: torch.Tensor,
weight_scale: torch.Tensor,
group_list_type: int,
group_list: torch.Tensor,
output_dtype: torch.dtype,
scale_alg=None,
) -> torch.Tensor:
"""FP4 weight x FP8-e4m3 activation (the checkpoint's W4A8_MXFP scheme).
W4A8MXFP GMM call: FP8 ``x_dtype``, FP4
``weight_dtype``, and the weight block scales fed through ``antiquant_scale``
with ``scale=None`` — the ``scale=`` + ``scale_dtype=`` form belongs to
W4A4_MXFP4 and dequantizes differently.
"""
group_list = group_list.to(torch.int64)
if input_scale is None:
x, x_scale = torch.ops.npu.npu_dynamic_mx_quant(
input,
axis=1,
round_mode="rint",
dst_type=torch.float8_e4m3fn,
block_size=MXFP4_BLOCK_SIZE,
scale_alg=scale_alg,
quant_info = AscendQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale_inv,
w2_weight_scale=layer.w2_weight_scale_inv,
)
else:
x, x_scale = input, input_scale
return torch.ops.npu.npu_grouped_matmul(
[x],
[weight],
scale=None,
antiquant_scale=[weight_scale],
scale_dtype=None,
per_token_scale=[
_pair_pack_mxfp_act_scale(x_scale, input_shape=tuple(x.shape))
],
split_item=2,
group_type=0,
group_list=group_list,
group_list_type=group_list_type,
output_dtype=output_dtype,
x_dtype=torch.float8_e4m3fn,
weight_dtype=_get_float4_e2m1fn_x2_dtype(),
per_token_scale_dtype=_get_float8_e8m0fnu_dtype(),
)[0]
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py
+107 -43
View File
@@ -9,8 +9,8 @@ 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
from sglang.srt.layers.quantization.base_config import QuantizationConfig
import logging
@@ -28,6 +28,7 @@ logger = logging.getLogger(__name__)
_E8M0_DTYPE = None
_DEFAULT_DYNAMIC_QUANT = object()
def _require_e8m0_dtype():
@@ -65,6 +66,93 @@ def _require_e8m0_dtype():
return _E8M0_DTYPE
def reshape_w4a8_mxfp_weight_scale_for_npu(scale: torch.Tensor) -> torch.Tensor:
"""Pack MXFP4 scales from ``[E, N, K/32]`` to the A5 GMM layout."""
if scale.dim() != 3:
return scale
num_experts, n, k32 = scale.shape
if k32 % 2 != 0:
raise ValueError(
"MXFP4 scale K dimension must be divisible by 2 for the "
f"[E, K/64, N, 2] layout, got {tuple(scale.shape)}."
)
return scale.view(num_experts, n, k32 // 2, 2).transpose(1, 2)
def prepare_w4a8_mxfp_weight(
weight: torch.Tensor,
weight_scale: torch.Tensor,
*,
npu_format=None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Convert W4A8 MXFP weights and scales to the shared A5 GMM layout."""
cast_args = {
"customize_dtype": torch.float8_e4m3fn,
"input_dtype": _get_float4_e2m1fn_x2_dtype(),
}
if npu_format is None:
weight = npu_format_cast(weight, **cast_args)
else:
weight = npu_format_cast(weight, npu_format, **cast_args)
return weight.transpose(-1, -2), reshape_w4a8_mxfp_weight_scale_for_npu(
weight_scale
)
def _pair_pack_mxfp_act_scale(scale: torch.Tensor) -> torch.Tensor:
"""Pack MXFP activation scales from ``[M, K/32]`` to A5 GMM layout."""
if scale.ndim != 2:
return scale
if scale.shape[-1] % 2 != 0:
raise ValueError(f"Invalid MXFP per-token scale shape: {tuple(scale.shape)}")
return scale.reshape(scale.shape[0], scale.shape[1] // 2, 2)
def w4a8_mxfp_gmm(
*,
input: torch.Tensor,
input_scale: Optional[torch.Tensor],
weight: torch.Tensor,
weight_scale: torch.Tensor,
group_list_type: int,
group_list: torch.Tensor,
output_dtype: torch.dtype,
scale_alg=None,
dynamic_quant_kwargs: Optional[Dict[str, Any]] = None,
) -> torch.Tensor:
"""Run the shared A5 FP4-weight × FP8-activation grouped matmul."""
group_list = group_list.to(torch.int64)
if input_scale is None:
if dynamic_quant_kwargs is None:
dynamic_quant_kwargs = {
"axis": 1,
"round_mode": "rint",
"dst_type": torch.float8_e4m3fn,
"block_size": 32,
"scale_alg": scale_alg,
}
x, x_scale = torch.ops.npu.npu_dynamic_mx_quant(input, **dynamic_quant_kwargs)
else:
x, x_scale = input, input_scale
return torch.ops.npu.npu_grouped_matmul(
[x],
[weight],
scale=None,
antiquant_scale=[weight_scale],
scale_dtype=None,
per_token_scale=[_pair_pack_mxfp_act_scale(x_scale)],
split_item=2,
group_type=0,
group_list=group_list,
group_list_type=group_list_type,
output_dtype=output_dtype,
x_dtype=torch.float8_e4m3fn,
weight_dtype=_get_float4_e2m1fn_x2_dtype(),
per_token_scale_dtype=_require_e8m0_dtype(),
)[0]
# DEPRECATED METHOD
# TODO: Remove in future realeses
def fused_moe_npu(
@@ -191,12 +279,9 @@ class _NPUMoEMethodBase(FusedMoEMethodBase):
class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
"""ModelSlim W4A8 MoE with packed MXFP4 weights and MXFP8 activations."""
def __init__(self):
def __init__(self, dynamic_quant_kwargs=_DEFAULT_DYNAMIC_QUANT):
super().__init__(quant_config=None)
self.matmul = GroupedMatmul()
self.hidden_states_quantizer = HiddenStatesDynamicQuant(
quant_dtype=torch.float8_e4m3fn
)
self.dynamic_quant_kwargs = dynamic_quant_kwargs
def process_weights_after_loading(
self, layer: torch.nn.Module, weight_prefix: str
@@ -208,20 +293,10 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.")
weight = getattr(layer, f"{weight_prefix}_weight")
weight.data = npu_format_cast(
weight.data,
customize_dtype=torch.float8_e4m3fn,
input_dtype=fp4_dtype,
).transpose(-1, -2)
weight_scale = getattr(layer, f"{weight_prefix}_weight_scale")
scale = weight_scale.data.reshape(
weight_scale.shape[0],
weight_scale.shape[1],
weight_scale.shape[2] // 2,
2,
).transpose(1, 2)
weight_scale.data = scale
weight.data, weight_scale.data = prepare_w4a8_mxfp_weight(
weight.data, weight_scale.data
)
# The refactored NPU dispatchers currently support BF16 and INT8.
# Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM.
@@ -238,35 +313,24 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
weight_prefix: str,
group_list_type: int,
) -> torch.Tensor:
fp4_dtype = _get_float4_e2m1fn_x2_dtype()
if fp4_dtype is None:
raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.")
e8m0_dtype = _require_e8m0_dtype()
if pertoken_scale is None:
hidden_states, pertoken_scale = self.hidden_states_quantizer(hidden_states)
elif pertoken_scale is not None:
if pertoken_scale is not None:
pertoken_scale = pertoken_scale.reshape(
hidden_states.shape[0], hidden_states.shape[1] // 64, 2
)
return self.matmul.forward(
quant_info,
weight_prefix,
hidden_states,
expert_tokens.to(torch.int64),
output_dtype,
dynamic_quant_kwargs = self.dynamic_quant_kwargs
if dynamic_quant_kwargs is _DEFAULT_DYNAMIC_QUANT:
dynamic_quant_kwargs = {"dst_type": torch.float8_e4m3fn}
return w4a8_mxfp_gmm(
input=hidden_states,
input_scale=pertoken_scale,
weight=getattr(quant_info, f"{weight_prefix}_weight"),
weight_scale=getattr(quant_info, f"{weight_prefix}_weight_scale"),
group_list_type=group_list_type,
transposed=True,
scale=None,
scale_dtype=None,
per_token_scale=[pertoken_scale],
antiquant_scale=[
getattr(quant_info, f"{weight_prefix}_weight_scale", None)
],
x_dtype=torch.float8_e4m3fn,
weight_dtype=fp4_dtype,
per_token_scale_dtype=e8m0_dtype,
group_list=expert_tokens,
output_dtype=output_dtype,
dynamic_quant_kwargs=dynamic_quant_kwargs,
)
python/sglang/srt/layers/moe/moe_runner/ascend.py
+9 -1
View File
@@ -13,6 +13,7 @@ from sglang.srt.hardware_backend.npu.moe.activation import (
NPUSitu,
NPUSwiglu,
NPUSwigluDeepEPKernel,
NPUSwigluMxfp8Quant,
NPUSwigluOAI,
NPUSwigluQuant,
NPUSwigluStepAndMul,
@@ -20,6 +21,7 @@ from sglang.srt.hardware_backend.npu.moe.activation import (
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
NPUMXFP8MoEMethod,
NPUW4A8Int8MoEMethod,
NPUW4A8MXFP4MoEMethod,
NPUW8A8Int8MoEMethod,
)
from sglang.srt.layers.moe.moe_runner.base import (
@@ -97,6 +99,12 @@ class AscendRunnerCore(MoeRunnerCore):
# both dispatchers: ascend_tp gets its activation quant fused into
# routing, DeepEP dispatches bf16 and gmm1 quantises it itself.
self.activation = None
elif (
isinstance(kernel, NPUW4A8MXFP4MoEMethod)
and config.swiglu_limit is not None
and config.swiglu_limit > 0
):
self.activation = NPUSwigluMxfp8Quant(config.swiglu_limit)
elif get_moe_a2a_backend().is_deepep():
# DeepEP path: use a unified kernel that decides quantisation
is_quant_kernel = isinstance(
@@ -186,7 +194,7 @@ class AscendRunnerCore(MoeRunnerCore):
# Grouped-row activations require dispatch metadata.
if isinstance(
self.activation,
(NPUSwigluDeepEPKernel, NPUSitu),
(NPUSwigluDeepEPKernel, NPUSitu, NPUSwigluMxfp8Quant),
):
hidden_states, pertoken_scale = self.activation._apply_activation(
hidden_states,
python/sglang/srt/layers/quantization/fp8.py
+2 -2
View File
@@ -421,10 +421,10 @@ class Fp8Config(QuantizationConfig):
and self.is_dsv4_fp4_experts
):
from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import (
NPUW4A4Fp4MoEMethod,
NPUW4A8MXFP4FusedMoEMethod,
)
return NPUW4A4Fp4MoEMethod(fp8_method, prefix=prefix)
return NPUW4A8MXFP4FusedMoEMethod(prefix=prefix)
if self.is_fp4_experts and get_moe_runner_backend().is_marlin():
from sglang.srt.layers.quantization.mxfp4_marlin_moe import (