From c0ed009f5b566be023661bd4e93065b8b4b8b31f Mon Sep 17 00:00:00 2001 From: wenxuewuhd Date: Tue, 21 Jul 2026 16:09:34 +0800 Subject: [PATCH] [NPU] Fix LLaDA2 MoE OOM after the FRACTAL_NZ cast, re-enabling the NZ speedup (#31772) --- python/sglang/srt/models/llada2.py | 18 ++++++++++++------ .../dllm/test_npu_llada2_mini.py | 5 ----- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index db97e2958..a6471de6b 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -83,6 +83,7 @@ from sglang.srt.runtime_context import ( get_stream, ) from sglang.srt.utils import ( + LazyValue, add_prefix, is_cuda, is_non_idle_and_non_empty, @@ -965,12 +966,17 @@ class LLaDA2MoeModelLM(nn.Module): ) weight_loader(param, loaded_weight) - self.routed_experts_weights_of_layer = { - layer_id: layer.mlp.get_moe_weights() - for layer_id, layer in enumerate(self.model.layers) - if not isinstance(layer, PPMissingLayer) - and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) - } + # Lazy: get_moe_weights() snapshots x.data, and building the map here would + # pin every expert weight's pre-process_weights_after_loading storage. + if not hasattr(self, "routed_experts_weights_of_layer"): + self.routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if not isinstance(layer, PPMissingLayer) + and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) + } + ) @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py index b8216893a..7550fd88d 100644 --- a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py +++ b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py @@ -1,4 +1,3 @@ -import os import unittest from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin @@ -31,10 +30,6 @@ class TestLLaDA2Mini(GSM8KAscendMixin, CustomTestCase): "LowConfidence", # TODO: Add dLLM configurations "--no-dllm-fdfo", # FDFO (PR #27551) halves single-batch speed on NPU; use sync mode ] - env = { - **os.environ, - "SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT": "1", # Need to avoid OOM issue - } accuracy = 0.88 output_throughput = 70