[NPU] Fix LLaDA2 MoE OOM after the FRACTAL_NZ cast, re-enabling the NZ speedup (#31772)

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 = {
+        # Lazy: get_moe_weights() snapshots x.data, and building the map here would
-            layer_id: layer.mlp.get_moe_weights()
+        # pin every expert weight's pre-process_weights_after_loading storage.
-            for layer_id, layer in enumerate(self.model.layers)
+        if not hasattr(self, "routed_experts_weights_of_layer"):
-            if not isinstance(layer, PPMissingLayer)
+            self.routed_experts_weights_of_layer = LazyValue(
-            and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock)
+                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):

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