[9/N] MoE Refactor: cleanup dispatcher interfaces (#11847)

This commit is contained in:
Cheng Wan
2025-10-20 10:11:46 -07:00
committed by GitHub
parent da5bde4d16
commit bfc3b3f786
24 changed files with 394 additions and 428 deletions
+3 -2
View File
@@ -365,9 +365,10 @@ class TopK(CustomOp):
def empty_topk_output(self, device: torch.device) -> TopKOutput:
topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
topk_idx = torch.full((0, topk), -1, dtype=torch.int32, device=device)
topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device)
# FIXME: router_logits should be of size (0, num_experts)
router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
return StandardTopKOutput(topk_weights, topk_idx, router_logits)
return StandardTopKOutput(topk_weights, topk_ids, router_logits)
# ------------------------------- TopK implementation -------------------------------------