From dd20826e0a87fbf9b775727276a341f52b6d967a Mon Sep 17 00:00:00 2001 From: Zhiyao Jiang Date: Mon, 10 Aug 2026 22:59:15 -0400 Subject: [PATCH] [AMD] Preserve the AITER expert mask across torch_memory_saver pause/resume (#34220) --- python/sglang/srt/layers/moe/fused_moe_triton/layer.py | 5 +++++ python/sglang/srt/utils/weight_checker.py | 1 + 2 files changed, 6 insertions(+) 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", )