minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Chengze Fan
2026-07-22 07:50:23 -07:00
committed by GitHub
parent e8e765b9d6
commit 40b2119b23
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/standard.py
+2 -2
View File
@@ -200,7 +200,7 @@ class StandardDispatcher(BaseDispatcher):
) )
if self.local_expert_mapping is not None and not self.skip_local_expert_mapping: 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.expert_mask_gpu = (
( (
(self.local_expert_mapping >= 0) (self.local_expert_mapping >= 0)
@@ -209,7 +209,7 @@ class StandardDispatcher(BaseDispatcher):
.to(torch.int32) .to(torch.int32)
.to(device="cuda") .to(device="cuda")
) )
else: elif not self.use_aiter_moe_runner:
if TopKOutputChecker.format_is_standard(topk_output): if TopKOutputChecker.format_is_standard(topk_output):
topk_output = topk_output._replace( topk_output = topk_output._replace(
topk_ids=self.local_expert_mapping[topk_output.topk_ids] topk_ids=self.local_expert_mapping[topk_output.topk_ids]