[NPU] NZ for non-quantized MOE, Qwen3 MOE double memory consumption fix (#15904)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
gemini-code-assist[bot]
parent
1953efb60e
commit
b77b0ffd60
@@ -25,6 +25,7 @@ from sglang.srt.utils import (
|
||||
get_bool_env_var,
|
||||
is_cpu,
|
||||
is_hip,
|
||||
is_npu,
|
||||
next_power_of_2,
|
||||
set_weight_attrs,
|
||||
use_intel_amx_backend,
|
||||
@@ -40,6 +41,7 @@ if TYPE_CHECKING:
|
||||
_is_cpu_amx_available = cpu_has_amx_support()
|
||||
_is_hip = is_hip()
|
||||
_is_cpu = is_cpu()
|
||||
_is_npu = is_npu()
|
||||
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
||||
|
||||
if _use_aiter:
|
||||
@@ -47,6 +49,9 @@ if _use_aiter:
|
||||
from aiter.fused_moe import fused_moe
|
||||
from aiter.ops.shuffle import shuffle_weight
|
||||
|
||||
if _is_npu:
|
||||
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
|
||||
|
||||
try:
|
||||
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
|
||||
except ImportError:
|
||||
@@ -296,6 +301,14 @@ 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 = weight.data.transpose(1, 2)
|
||||
weight.data = npu_format_cast(
|
||||
weight.data,
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
def create_moe_runner(
|
||||
@@ -494,14 +507,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
|
||||
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
|
||||
if layer.w13_weight.shape[-1] == layer.hidden_size:
|
||||
w13 = layer.w13_weight.transpose(1, 2)
|
||||
w2 = layer.w2_weight.transpose(1, 2)
|
||||
|
||||
# gmm1: gate_up_proj
|
||||
hidden_states = torch_npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[w13],
|
||||
weight=[layer.w13_weight],
|
||||
bias=w13_bias,
|
||||
split_item=2,
|
||||
group_list_type=0,
|
||||
@@ -525,7 +535,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
|
||||
# gmm2: down_proj
|
||||
hidden_states = torch_npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[w2],
|
||||
weight=[layer.w2_weight],
|
||||
bias=w2_bias,
|
||||
split_item=2,
|
||||
group_list_type=0,
|
||||
|
||||
Reference in New Issue
Block a user