Add fp4 quantize before all-gather for Flashinfer cutlass MoE DP (max throughput) (#7667)

2025-08-15 22:08:11 -07:00
parent 87dab54824
commit eff4eb3fdd
16 changed files with 360 additions and 52 deletions
@@ -327,6 +327,13 @@ class TopK(CustomOp):
                expert_location_dispatch_info=expert_location_dispatch_info,
            )

+    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)
+        router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
+        return StandardTopKOutput(topk_weights, topk_idx, router_logits)
+

 # ------------------------------- TopK implementation -------------------------------------