diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 3ab58a354..cf5203f35 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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 ( diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index 4f828d355..bf0f6bb1e 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -56,6 +56,7 @@ _NON_PERSISTENT_BUFFER_PATTERNS = ( "inv_freq", "freqs_cis", "_weight_fp32", + "expert_mask_gpu", )