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

[NPU] NZ for non-quantized MOE, Qwen3 MOE double memory consumption fix (#15904)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Артем Савкин
2026-01-29 00:55:08 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
parent 1953efb60e
commit b77b0ffd60
7 changed files with 111 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
+6 -21
View File
@@ -150,42 +150,27 @@ class _NPUFusedMoEMethodBase(FusedMoEMethodBase):
class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
def _release_weight_cache(self, weight: torch.Tensor):
# .contiguous() introduces additional memory overhead and needs to be released using resize_(0)
origin_weight = weight.data.transpose(1, 2)
new_weight = origin_weight.contiguous()
origin_weight.untyped_storage().resize_(0)
return new_weight
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_data = self._release_weight_cache(layer.w13_weight.data)
layer.w13_weight = torch.nn.Parameter(weight_data, requires_grad=False)
weight_data = self._release_weight_cache(layer.w2_weight.data)
layer.w2_weight = torch.nn.Parameter(weight_data, requires_grad=False)
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data.transpose(1, 2))
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data.transpose(1, 2))
layer.w13_weight_scale = torch.nn.Parameter(
layer.w13_weight_scale.data.squeeze(-1).contiguous().to(torch.float32),
requires_grad=False,
layer.w13_weight_scale.data.squeeze(-1), requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(
layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False
layer.w2_weight_scale.data.squeeze(-1), requires_grad=False
)
# Compressed-tensors format doesn't have this field
if hasattr(layer, "w13_weight_offset"):
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.squeeze(-1).contiguous(),
layer.w13_weight_offset.data.squeeze(-1),
requires_grad=False,
)
if hasattr(layer, "w2_weight_offset"):
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.squeeze(-1).contiguous(),
layer.w2_weight_offset.data.squeeze(-1),
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)
def apply(
self,
layer,