[NPU]Refactor weight processing and add NPUSwigluLimit activation (#38420)
Co-authored-by: AndyLi429 <AndyLi429@noreply.gitcode.com> Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
AndyLi429
co-authored by
AndyLi429
Even Zhou
AndyLi429
Even Zhou
parent
2d08cc5ede
commit
e970453b43
@@ -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
|
||||
# =============================================================================
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
import inspect
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
@@ -18,23 +16,19 @@ register_npu_ci(est_time=1, suite="stage-a-unit-test-npu")
|
||||
# initialized, `_get_float8_e8m0fnu_dtype` not yet defined). Initializing the
|
||||
# package first mirrors how the engine loads quantization at model-config time.
|
||||
import sglang.srt.layers.quantization # noqa: F401
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.hardware_backend.npu.quantization import fp4_moe_methods
|
||||
from sglang.srt.hardware_backend.npu.moe.activation import NPUSwigluMxfp8Quant
|
||||
from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import (
|
||||
NPUW4A4Fp4MoEMethod,
|
||||
_apply_swiglu_limit_npu,
|
||||
_configure_dsv4_deepep_dispatcher,
|
||||
NPUW4A8MXFP4FusedMoEMethod,
|
||||
)
|
||||
from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
|
||||
_pair_pack_mxfp_act_scale,
|
||||
_reshape_mxfp4_scale_for_npu,
|
||||
npu_apply_without_routing_weights_w4a4_mxfp,
|
||||
prepare_w4a8_mxfp_weight,
|
||||
reshape_w4a8_mxfp_weight_scale_for_npu,
|
||||
w4a8_mxfp_gmm,
|
||||
)
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
from sglang.srt.layers.moe.token_dispatcher import deepep
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
|
||||
|
||||
_NOT_PASSED = object()
|
||||
|
||||
|
||||
class TestFP4MethodGate(unittest.TestCase):
|
||||
def test_pre_arch35_keeps_fp8_moe_method(self):
|
||||
@@ -52,25 +46,63 @@ class TestFP4MethodGate(unittest.TestCase):
|
||||
|
||||
self.assertIsInstance(method, Fp8MoEMethod)
|
||||
|
||||
def test_arch35_uses_ascend_runner_method(self):
|
||||
config = Fp8Config(is_fp4_experts=True)
|
||||
config.is_dsv4_fp4_experts = True
|
||||
layer = FusedMoE.__new__(FusedMoE)
|
||||
|
||||
class TestApplySwiGLULimitNpu(unittest.TestCase):
|
||||
def test_clamps_gate_and_up_asymmetrically(self):
|
||||
# DeepSeek-V4 clamps gate (first half) to <= limit but only the upper
|
||||
# bound, while up (second half) is clamped symmetrically to [-limit, limit].
|
||||
# A regression that swapped these would silently change expert activations.
|
||||
gate_up = torch.tensor([[8.0, -9.0, 9.0, -9.0]])
|
||||
_apply_swiglu_limit_npu(gate_up, 7.0)
|
||||
self.assertTrue(torch.equal(gate_up, torch.tensor([[7.0, -9.0, 7.0, -7.0]])))
|
||||
with (
|
||||
patch("sglang.srt.layers.quantization.fp8.is_npu", return_value=True),
|
||||
patch(
|
||||
"sglang.srt.layers.quantization.fp8.is_npu_arch35",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
method = config.get_quant_method(layer, "model.layers.0.experts")
|
||||
|
||||
def test_noop_when_limit_none(self):
|
||||
gate_up = torch.tensor([[8.0, -9.0]])
|
||||
_apply_swiglu_limit_npu(gate_up, None)
|
||||
self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]])))
|
||||
self.assertIsInstance(method, NPUW4A8MXFP4FusedMoEMethod)
|
||||
|
||||
def test_noop_when_limit_nonpositive(self):
|
||||
gate_up = torch.tensor([[8.0, -9.0]])
|
||||
_apply_swiglu_limit_npu(gate_up, 0.0)
|
||||
self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]])))
|
||||
|
||||
class TestNPUSwigluMxfp8Quant(unittest.TestCase):
|
||||
def test_translates_runner_conventions_onto_the_ascend_c_op(self):
|
||||
# The sgl-kernel-npu op numbers the group-list layouts the other way round
|
||||
# from sglang, so count layout arrives as 1 and must leave as 0. The clamp is
|
||||
# inert at 0.0, so a limit that missed clamp_value would silently disable it.
|
||||
activation = NPUSwigluMxfp8Quant(7.0)
|
||||
output = torch.empty(2, 4, dtype=torch.float8_e4m3fn)
|
||||
scale = torch.empty(2, 1, 2, dtype=torch.float8_e8m0fnu)
|
||||
group_list = torch.tensor([1, 1], dtype=torch.int64)
|
||||
hidden_states = torch.empty(2, 8)
|
||||
|
||||
with patch.object(
|
||||
torch.ops.npu,
|
||||
"swiglu_group_quant",
|
||||
return_value=(output, scale, None),
|
||||
create=True,
|
||||
) as kernel:
|
||||
actual_output, actual_scale = activation._apply_activation(
|
||||
hidden_states, group_list, group_list_type=1
|
||||
)
|
||||
|
||||
self.assertIs(actual_output, output)
|
||||
self.assertIs(actual_scale, scale)
|
||||
kwargs = kernel.call_args.kwargs
|
||||
self.assertIs(kwargs["x"], hidden_states)
|
||||
self.assertIs(kwargs["group_index"], group_list)
|
||||
self.assertEqual(kwargs["quant_mode"], 2)
|
||||
self.assertEqual(kwargs["group_list_type"], 0)
|
||||
self.assertEqual(kwargs["clamp_value"], 7.0)
|
||||
|
||||
def test_rejects_a_cumulative_group_list(self):
|
||||
# The op has no cusum layout and sums its group list as counts, so a cusum list
|
||||
# must fail here rather than derive a row count from the sum of prefix sums.
|
||||
activation = NPUSwigluMxfp8Quant(7.0)
|
||||
with self.assertRaises(ValueError):
|
||||
activation._apply_activation(
|
||||
torch.empty(2, 8),
|
||||
torch.tensor([1, 2], dtype=torch.int64),
|
||||
group_list_type=0,
|
||||
)
|
||||
|
||||
|
||||
class TestReshapeMxfp4ScaleForNpu(unittest.TestCase):
|
||||
@@ -78,13 +110,35 @@ class TestReshapeMxfp4ScaleForNpu(unittest.TestCase):
|
||||
# [E, N, K/32] -> [E, K/64, N, 2] is the packed-pair layout the GMM reads;
|
||||
# getting the transpose axis wrong silently dequantizes with the wrong scale.
|
||||
scale = torch.arange(8, dtype=torch.uint8).view(1, 2, 4)
|
||||
out = _reshape_mxfp4_scale_for_npu(scale)
|
||||
out = reshape_w4a8_mxfp_weight_scale_for_npu(scale)
|
||||
self.assertEqual(tuple(out.shape), (1, 2, 2, 2))
|
||||
self.assertTrue(torch.equal(out, scale.view(1, 2, 2, 2).transpose(1, 2)))
|
||||
|
||||
def test_rejects_odd_k_dim(self):
|
||||
with self.assertRaises(ValueError):
|
||||
_reshape_mxfp4_scale_for_npu(torch.zeros(1, 2, 3, dtype=torch.uint8))
|
||||
reshape_w4a8_mxfp_weight_scale_for_npu(
|
||||
torch.zeros(1, 2, 3, dtype=torch.uint8)
|
||||
)
|
||||
|
||||
|
||||
class TestPrepareW4A8MxfpWeight(unittest.TestCase):
|
||||
def test_uses_shared_weight_and_scale_layout(self):
|
||||
# A wrong transpose or scale packing makes both ModelSlim W4A8 and
|
||||
# DeepSeek-V4 W4A8 read different blocks from the same checkpoint.
|
||||
weight = torch.arange(16, dtype=torch.uint8).view(1, 2, 8)
|
||||
scale = torch.arange(8, dtype=torch.uint8).view(1, 2, 4)
|
||||
formatted_weight = torch.arange(16, dtype=torch.uint8).view(1, 2, 8)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.hardware_backend.npu.quantization.moe_methods.npu_format_cast",
|
||||
return_value=formatted_weight,
|
||||
):
|
||||
prepared_weight, prepared_scale = prepare_w4a8_mxfp_weight(weight, scale)
|
||||
|
||||
self.assertTrue(torch.equal(prepared_weight, formatted_weight.transpose(1, 2)))
|
||||
self.assertTrue(
|
||||
torch.equal(prepared_scale, scale.view(1, 2, 2, 2).transpose(1, 2))
|
||||
)
|
||||
|
||||
|
||||
class TestMxfp4ScaleWeightLoader(unittest.TestCase):
|
||||
@@ -95,7 +149,7 @@ class TestMxfp4ScaleWeightLoader(unittest.TestCase):
|
||||
loaded.append(loaded_weight.clone())
|
||||
|
||||
layer = torch.nn.Module()
|
||||
method = NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test")
|
||||
method = NPUW4A8MXFP4FusedMoEMethod(prefix="test")
|
||||
method.create_weights(
|
||||
layer,
|
||||
num_experts=1,
|
||||
@@ -131,264 +185,6 @@ class TestPairPackMxfpActScale(unittest.TestCase):
|
||||
with self.assertRaises(ValueError):
|
||||
_pair_pack_mxfp_act_scale(torch.zeros(2, 3))
|
||||
|
||||
def test_unflattens_low_latency_deepep_scale_as_view(self):
|
||||
# DeepEP returns one flat E8M0 scale per 32-element block. Passing
|
||||
# that flat buffer to GMM would use the wrong scale layout and either
|
||||
# fail or dequantize activations incorrectly.
|
||||
flat = torch.arange(4, dtype=torch.uint8)
|
||||
packed = _pair_pack_mxfp_act_scale(flat, input_shape=(2, 64))
|
||||
|
||||
self.assertEqual(tuple(packed.shape), (2, 1, 2))
|
||||
self.assertEqual(packed.data_ptr(), flat.data_ptr())
|
||||
self.assertTrue(torch.equal(packed, torch.tensor([[[0, 1]], [[2, 3]]])))
|
||||
|
||||
def test_rejects_low_latency_deepep_scale_with_wrong_length(self):
|
||||
with self.assertRaises(ValueError):
|
||||
_pair_pack_mxfp_act_scale(
|
||||
torch.zeros(3, dtype=torch.uint8), input_shape=(2, 64)
|
||||
)
|
||||
|
||||
|
||||
class TestDsv4DeepEPMxfp8DispatcherConfig(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _deepep_backend():
|
||||
return SimpleNamespace(is_deepep=lambda: True)
|
||||
|
||||
def test_a5_deepep_defaults_low_latency_dispatch_to_mxfp8(self):
|
||||
dispatcher = MagicMock()
|
||||
layer = SimpleNamespace(dispatcher=dispatcher)
|
||||
|
||||
with (
|
||||
patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True),
|
||||
patch(
|
||||
"sglang.srt.layers.moe.get_moe_a2a_backend",
|
||||
return_value=self._deepep_backend(),
|
||||
),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
_configure_dsv4_deepep_dispatcher(layer)
|
||||
|
||||
dispatcher.set_quant_config.assert_called_once_with(
|
||||
{
|
||||
"normal_dispatcher_output_dtype": "bf16",
|
||||
"low_latency_dispatcher_output_dtype": "mxfp8",
|
||||
}
|
||||
)
|
||||
|
||||
def test_non_deepep_ignores_the_low_latency_quant_environment(self):
|
||||
dispatcher = MagicMock()
|
||||
layer = SimpleNamespace(dispatcher=dispatcher)
|
||||
|
||||
with (
|
||||
patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True),
|
||||
patch(
|
||||
"sglang.srt.layers.moe.get_moe_a2a_backend",
|
||||
return_value=SimpleNamespace(is_deepep=lambda: False),
|
||||
),
|
||||
envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"),
|
||||
):
|
||||
_configure_dsv4_deepep_dispatcher(layer)
|
||||
|
||||
dispatcher.set_quant_config.assert_called_once_with(
|
||||
{"dispatcher_output_dtype": "bf16"}
|
||||
)
|
||||
|
||||
def test_a5_deepep_allows_bf16_low_latency_fallback(self):
|
||||
dispatcher = MagicMock()
|
||||
layer = SimpleNamespace(dispatcher=dispatcher)
|
||||
|
||||
with (
|
||||
patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True),
|
||||
patch(
|
||||
"sglang.srt.layers.moe.get_moe_a2a_backend",
|
||||
return_value=self._deepep_backend(),
|
||||
),
|
||||
envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("bf16"),
|
||||
):
|
||||
_configure_dsv4_deepep_dispatcher(layer)
|
||||
|
||||
dispatcher.set_quant_config.assert_called_once_with(
|
||||
{
|
||||
"normal_dispatcher_output_dtype": "bf16",
|
||||
"low_latency_dispatcher_output_dtype": "bf16",
|
||||
}
|
||||
)
|
||||
|
||||
def test_a5_deepep_rejects_an_invalid_low_latency_quant_mode(self):
|
||||
layer = SimpleNamespace(dispatcher=MagicMock())
|
||||
|
||||
with (
|
||||
patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True),
|
||||
patch(
|
||||
"sglang.srt.layers.moe.get_moe_a2a_backend",
|
||||
return_value=self._deepep_backend(),
|
||||
),
|
||||
envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"),
|
||||
self.assertRaisesRegex(ValueError, "SGLANG_NPU_DSV4"),
|
||||
):
|
||||
_configure_dsv4_deepep_dispatcher(layer)
|
||||
|
||||
def test_non_a5_ignores_the_low_latency_quant_environment(self):
|
||||
dispatcher = MagicMock()
|
||||
layer = SimpleNamespace(dispatcher=dispatcher)
|
||||
|
||||
with (
|
||||
patch.object(fp4_moe_methods, "is_npu_arch35", return_value=False),
|
||||
patch(
|
||||
"sglang.srt.layers.moe.get_moe_a2a_backend",
|
||||
return_value=self._deepep_backend(),
|
||||
),
|
||||
envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"),
|
||||
):
|
||||
_configure_dsv4_deepep_dispatcher(layer)
|
||||
|
||||
dispatcher.set_quant_config.assert_called_once_with(
|
||||
{"dispatcher_output_dtype": "bf16"}
|
||||
)
|
||||
|
||||
|
||||
class _LowLatencyBuffer:
|
||||
def __init__(self):
|
||||
self.kwargs = None
|
||||
|
||||
def low_latency_dispatch(
|
||||
self,
|
||||
hidden_states,
|
||||
topk_ids,
|
||||
num_max_dispatch_tokens_per_rank,
|
||||
num_experts,
|
||||
*,
|
||||
use_fp8,
|
||||
quant_mode=_NOT_PASSED,
|
||||
**kwargs,
|
||||
):
|
||||
self.kwargs = {"use_fp8": use_fp8, "quant_mode": quant_mode, **kwargs}
|
||||
return torch.empty(0), torch.empty(0), object(), object(), object()
|
||||
|
||||
|
||||
class _LegacyLowLatencyBuffer:
|
||||
def low_latency_dispatch(
|
||||
self,
|
||||
hidden_states,
|
||||
topk_ids,
|
||||
num_max_dispatch_tokens_per_rank,
|
||||
num_experts,
|
||||
*,
|
||||
use_fp8,
|
||||
**kwargs,
|
||||
):
|
||||
return torch.empty(0), torch.empty(0), object(), object(), object()
|
||||
|
||||
|
||||
class TestDeepEPLowLatencyMxfp8Dispatch(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _dispatcher(quant_mode, buffer):
|
||||
dispatcher = object.__new__(deepep._DeepEPDispatcherImplLowLatency)
|
||||
dispatcher.quant_config = {}
|
||||
dispatcher.use_fp8 = False
|
||||
dispatcher.use_nvfp4 = False
|
||||
dispatcher.low_latency_quant_mode = quant_mode
|
||||
dispatcher._low_latency_quant_mode_runtime_checked = False
|
||||
dispatcher.num_max_dispatch_tokens_per_rank = 2
|
||||
dispatcher.num_experts = 2
|
||||
dispatcher.return_recv_hook = False
|
||||
dispatcher._get_buffer = lambda: buffer
|
||||
return dispatcher
|
||||
|
||||
def test_mxfp8_passes_the_kernel_quant_mode(self):
|
||||
buffer = _LowLatencyBuffer()
|
||||
dispatcher = self._dispatcher("mx_fp8_e4m3", buffer)
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
patch.object(deepep, "_deepep_precompile_tp_barrier"),
|
||||
):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
self.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3")
|
||||
|
||||
def test_mxfp8_ops_strategy_uses_legacy_mxfp8_flags(self):
|
||||
buffer = _LowLatencyBuffer()
|
||||
dispatcher = self._dispatcher("mx_fp8_e4m3", buffer)
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"DEEP_USE_MODE": "ops"}, clear=True),
|
||||
patch.object(deepep, "_deepep_precompile_tp_barrier"),
|
||||
):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
self.assertTrue(buffer.kwargs["use_fp8"])
|
||||
self.assertTrue(buffer.kwargs["use_ue8m0"])
|
||||
self.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3")
|
||||
|
||||
def test_mxfp8_rejects_an_unsupported_low_latency_strategy(self):
|
||||
dispatcher = self._dispatcher("mx_fp8_e4m3", _LowLatencyBuffer())
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"DEEP_USE_MODE": "alltoall"}, clear=True),
|
||||
self.assertRaisesRegex(RuntimeError, "DEEP_USE_MODE"),
|
||||
):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
def test_mxfp8_checks_runtime_interface_once_per_dispatcher(self):
|
||||
buffer = _LowLatencyBuffer()
|
||||
dispatcher = self._dispatcher("mx_fp8_e4m3", buffer)
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
patch.object(deepep, "_deepep_precompile_tp_barrier"),
|
||||
patch.object(
|
||||
deepep.inspect, "signature", wraps=inspect.signature
|
||||
) as signature,
|
||||
):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
self.assertEqual(signature.call_count, 1)
|
||||
|
||||
def test_bf16_does_not_pass_a_quant_mode(self):
|
||||
buffer = _LowLatencyBuffer()
|
||||
dispatcher = self._dispatcher(None, buffer)
|
||||
|
||||
with patch.object(deepep, "_deepep_precompile_tp_barrier"):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
self.assertIs(buffer.kwargs["quant_mode"], _NOT_PASSED)
|
||||
|
||||
def test_mxfp8_rejects_legacy_runtime_without_quant_mode(self):
|
||||
dispatcher = self._dispatcher("mx_fp8_e4m3", _LegacyLowLatencyBuffer())
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "quant_mode"):
|
||||
dispatcher._dispatch_core(
|
||||
torch.zeros(1, 64),
|
||||
torch.zeros(1, 1, dtype=torch.int64),
|
||||
torch.ones(1, 1),
|
||||
)
|
||||
|
||||
|
||||
class TestW4A8MxfpGmmInputScale(unittest.TestCase):
|
||||
def setUp(self):
|
||||
@@ -431,28 +227,6 @@ class TestW4A8MxfpGmmInputScale(unittest.TestCase):
|
||||
self.assertEqual(call_kwargs["group_list"].dtype, torch.int64)
|
||||
self.assertTrue(torch.equal(call_kwargs["group_list"], self.group_list))
|
||||
|
||||
def test_flat_deepep_scale_skips_dynamic_quant_after_layout_adaptation(self):
|
||||
flat_scale = torch.arange(4, dtype=torch.uint8)
|
||||
expected = torch.randn(2, 32)
|
||||
with (
|
||||
patch.object(
|
||||
torch.ops.npu, "npu_dynamic_mx_quant", create=True
|
||||
) as dynamic_quant,
|
||||
patch.object(
|
||||
torch.ops.npu,
|
||||
"npu_grouped_matmul",
|
||||
return_value=[expected],
|
||||
create=True,
|
||||
) as grouped_matmul,
|
||||
):
|
||||
output = self._call_gmm(flat_scale)
|
||||
|
||||
dynamic_quant.assert_not_called()
|
||||
self.assertIs(output, expected)
|
||||
packed_scale = grouped_matmul.call_args.kwargs["per_token_scale"][0]
|
||||
self.assertEqual(tuple(packed_scale.shape), (2, 1, 2))
|
||||
self.assertEqual(packed_scale.data_ptr(), flat_scale.data_ptr())
|
||||
|
||||
def test_missing_scale_uses_dynamic_quant(self):
|
||||
quantized = torch.empty(2, 64, dtype=torch.float8_e4m3fn)
|
||||
quantized_scale = torch.ones(2, 1, 2)
|
||||
@@ -480,51 +254,55 @@ class TestW4A8MxfpGmmInputScale(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestW4A8MxfpGmmChain(unittest.TestCase):
|
||||
def test_passes_swiglu_limit_to_quant(self):
|
||||
gate_up = torch.randn(1, 64)
|
||||
activated = torch.randn(1, 32)
|
||||
activated_scale = torch.randn(1, 1)
|
||||
expected = torch.randn(1, 32)
|
||||
class TestRunnerDelegation(unittest.TestCase):
|
||||
def test_apply_delegates_with_dsv4_scale_names(self):
|
||||
method = NPUW4A8MXFP4FusedMoEMethod(prefix="test")
|
||||
expected = object()
|
||||
method.runner = MagicMock()
|
||||
method.runner.run.return_value = expected
|
||||
layer = SimpleNamespace(
|
||||
w13_weight=MagicMock(),
|
||||
w13_weight_scale_inv=MagicMock(),
|
||||
w2_weight=MagicMock(),
|
||||
w2_weight_scale_inv=MagicMock(),
|
||||
moe_runner_config=SimpleNamespace(swiglu_limit=7.0),
|
||||
)
|
||||
dispatch_output = object()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
fp4_moe_methods, "w4a8_mxfp_gmm", side_effect=[gate_up, expected]
|
||||
) as gmm,
|
||||
patch.object(
|
||||
fp4_moe_methods,
|
||||
"swiglu_quant",
|
||||
return_value=(activated, activated_scale),
|
||||
) as swiglu,
|
||||
):
|
||||
output = npu_apply_without_routing_weights_w4a4_mxfp(
|
||||
layer,
|
||||
torch.randn(1, 4),
|
||||
torch.ones(1, 1, 2),
|
||||
group_list_type=1,
|
||||
group_list=torch.tensor([1], dtype=torch.int64),
|
||||
output_dtype=torch.bfloat16,
|
||||
)
|
||||
output = method.apply(layer, dispatch_output)
|
||||
|
||||
self.assertIs(output, expected)
|
||||
self.assertTrue(torch.equal(swiglu.call_args.args[0], gate_up))
|
||||
self.assertTrue(swiglu.call_args.kwargs["do_limit"])
|
||||
self.assertEqual(swiglu.call_args.kwargs["limit"], 7.0)
|
||||
self.assertIs(gmm.call_args_list[1].kwargs["input"], activated)
|
||||
self.assertIs(gmm.call_args_list[1].kwargs["input_scale"], activated_scale)
|
||||
method.runner.run.assert_called_once()
|
||||
self.assertIs(method.runner.run.call_args.args[0], dispatch_output)
|
||||
quant_info = method.runner.run.call_args.args[1]
|
||||
self.assertIs(quant_info.w13_weight_scale, layer.w13_weight_scale_inv)
|
||||
self.assertIs(quant_info.w2_weight_scale, layer.w2_weight_scale_inv)
|
||||
|
||||
def test_create_runner_installs_internal_kernels_before_runner(self):
|
||||
method = NPUW4A8MXFP4FusedMoEMethod(prefix="test")
|
||||
layer = SimpleNamespace()
|
||||
config = SimpleNamespace(layer=None)
|
||||
backend = MagicMock()
|
||||
backend.is_auto.return_value = True
|
||||
|
||||
with (
|
||||
patch("sglang.srt.layers.moe.moe_runner.runner.MoeRunner") as moe_runner,
|
||||
patch(
|
||||
"sglang.srt.layers.moe.utils.get_moe_runner_backend",
|
||||
return_value=backend,
|
||||
),
|
||||
):
|
||||
method.create_moe_runner(layer, config)
|
||||
|
||||
self.assertIs(layer.w13_kernel, method.w13_kernel)
|
||||
self.assertIs(layer.w2_kernel, method.w2_kernel)
|
||||
self.assertIs(config.layer, layer)
|
||||
moe_runner.assert_called_once()
|
||||
|
||||
|
||||
class TestProcessWeightsAfterLoadingZeroScale(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _method():
|
||||
return NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test")
|
||||
return NPUW4A8MXFP4FusedMoEMethod(prefix="test")
|
||||
|
||||
def test_raises_when_w13_scales_never_loaded(self):
|
||||
# An all-zero scale is the signature of a checkpoint whose scale names
|
||||
@@ -553,6 +331,38 @@ class TestProcessWeightsAfterLoadingZeroScale(unittest.TestCase):
|
||||
with self.assertRaises(RuntimeError):
|
||||
self._method().process_weights_after_loading(layer)
|
||||
|
||||
def test_keeps_low_latency_dispatch_in_mxfp8(self):
|
||||
layer = SimpleNamespace(
|
||||
w13_weight=torch.nn.Parameter(
|
||||
torch.ones(1, 2, 32, dtype=torch.uint8), requires_grad=False
|
||||
),
|
||||
w13_weight_scale_inv=torch.nn.Parameter(
|
||||
torch.ones(1, 2, 1, dtype=torch.uint8), requires_grad=False
|
||||
),
|
||||
w2_weight=torch.nn.Parameter(
|
||||
torch.ones(1, 32, 1, dtype=torch.uint8), requires_grad=False
|
||||
),
|
||||
w2_weight_scale_inv=torch.nn.Parameter(
|
||||
torch.ones(1, 32, 1, dtype=torch.uint8), requires_grad=False
|
||||
),
|
||||
dispatcher=MagicMock(),
|
||||
)
|
||||
prepared_weight = torch.ones(1, 32, 2, dtype=torch.uint8)
|
||||
prepared_scale = torch.ones(1, 1, 2, 2, dtype=torch.uint8)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods.prepare_w4a8_mxfp_weight",
|
||||
return_value=(prepared_weight, prepared_scale),
|
||||
):
|
||||
self._method().process_weights_after_loading(layer)
|
||||
|
||||
layer.dispatcher.set_quant_config.assert_called_once_with(
|
||||
{
|
||||
"normal_dispatcher_output_dtype": "bf16",
|
||||
"low_latency_dispatcher_output_dtype": "mxfp8",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user