[AMD] Cache AITER expert mask across decode (#31889)

2026-07-22 07:50:23 -07:00
parent e8e765b9d6
commit 40b2119b23
1 changed files with 2 additions and 2 deletions
@@ -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]