From 40b2119b23e49be767da1f9f73746ac8e158dae5 Mon Sep 17 00:00:00 2001 From: Chengze Fan Date: Wed, 22 Jul 2026 07:50:23 -0700 Subject: [PATCH] [AMD] Cache AITER expert mask across decode (#31889) --- python/sglang/srt/layers/moe/token_dispatcher/standard.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 244f19e11..009221228 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -200,7 +200,7 @@ class StandardDispatcher(BaseDispatcher): ) if self.local_expert_mapping is not None and not self.skip_local_expert_mapping: - if self.use_aiter_moe_runner: + if self.use_aiter_moe_runner and self.expert_mask_gpu is None: self.expert_mask_gpu = ( ( (self.local_expert_mapping >= 0) @@ -209,7 +209,7 @@ class StandardDispatcher(BaseDispatcher): .to(torch.int32) .to(device="cuda") ) - else: + elif not self.use_aiter_moe_runner: if TopKOutputChecker.format_is_standard(topk_output): topk_output = topk_output._replace( topk_ids=self.local_expert_mapping[topk_output.topk_ids]