[AMD] Preserve the AITER expert mask across torch_memory_saver pause/resume (#34220)

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -413,6 +413,11 @@ class FusedMoE(torch.nn.Module):
 
         self.quant_method.create_moe_runner(self, self.moe_runner_config)
         self.dispatcher = create_moe_dispatcher(self.moe_runner_config)
+        # Dispatchers are not nn.Modules, so they cannot register their own
+        # buffers; the AITER expert mask would not survive a memory-saver resume.
+        expert_mask = getattr(self.dispatcher, "expert_mask_gpu", None)
+        if expert_mask is not None:
+            self.register_buffer("expert_mask_gpu", expert_mask, persistent=False)
         self._use_ascend_fuseep = get_moe_a2a_backend().is_ascend_fuseep()
 
         if (

python/sglang/srt/utils/weight_checker.py

@@ -56,6 +56,7 @@ _NON_PERSISTENT_BUFFER_PATTERNS = (
     "inv_freq",
     "freqs_cis",
     "_weight_fp32",
+    "expert_mask_gpu",
 )