[9/N] MoE Refactor: cleanup dispatcher interfaces (#11847)
This commit is contained in:
@@ -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 -------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user