update to comply with main

Author: dsikka

vllm/model_executor/layers/fused_moe/layer.py

@@ -28,8 +28,8 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
 
     @abstractmethod
     def apply(self, layer: torch.nn.Module, x: torch.Tensor,
-              topk_weights: torch.Tensor,
-              topk_ids: torch.Tensor) -> torch.Tensor:
+              topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+              **kwargs) -> torch.Tensor:
         raise NotImplementedError
 
 
@@ -59,21 +59,18 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
         set_weight_attrs(w2_weight, extra_weight_attrs)
 
     def apply(self, layer: torch.nn.Module, x: torch.Tensor,
-              topk_weights: torch.Tensor,
-              topk_ids: torch.Tensor) -> torch.Tensor:
+              topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+              **kwargs) -> torch.Tensor:
 
         return self.forward(x=x,
                             layer=layer,
                             topk_weights=topk_weights,
-                            topk_ids=topk_ids)
+                            topk_ids=topk_ids,
+                            **kwargs)
 
-    def forward_cuda(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-    ) -> torch.Tensor:
+    def forward_cuda(self, layer: torch.nn.Module, x: torch.Tensor,
+                     topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                     **kwargs) -> torch.Tensor:
         return fused_experts(hidden_states=x,
                              w1=layer.w13_weight,
                              w2=layer.w2_weight,
@@ -85,17 +82,11 @@ def forward_cpu(self, *args, **kwargs):
         raise NotImplementedError(
             "The CPU backend currently does not support MoE.")
 
-    def forward_tpu(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-    ) -> torch.Tensor:
-
-        #assert not use_grouped_topk
-        #assert num_expert_group is None
-        #assert topk_group is None
+    def forward_tpu(self, layer: torch.nn.Module, x: torch.Tensor,
+                    topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                    use_grouped_topk: bool) -> torch.Tensor:
+
+        assert not use_grouped_topk
         return fused_experts(hidden_states=x,
                              w1=layer.w13_weight,
                              w2=layer.w2_weight,
@@ -294,7 +285,8 @@ def forward(self, hidden_states: torch.Tensor,
             layer=self,
             x=hidden_states,
             topk_weights=topk_weights,
-            topk_ids=topk_ids)
+            topk_ids=topk_ids,
+            use_grouped_topk=self.use_grouped_topk)
 
         # Optionally reduce.
         if self.reduce_results and self.tp_size > 1:

vllm/model_executor/layers/quantization/awq.py

@@ -280,8 +280,8 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
             })
 
     def apply(self, layer: torch.nn.Module, x: torch.Tensor,
-              topk_weights: torch.Tensor,
-              topk_ids: torch.Tensor) -> torch.Tensor:
+              topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+              **kwargs) -> torch.Tensor:
 
         return fused_experts_awq(hidden_states=x,
                                  w1=layer.w13_qweight,

vllm/model_executor/layers/quantization/fp8.py

@@ -408,8 +408,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
             return
 
     def apply(self, layer: torch.nn.Module, x: torch.Tensor,
-              topk_weights: torch.Tensor,
-              topk_ids: torch.Tensor) -> torch.Tensor:
+              topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+              **kwargs) -> torch.Tensor:
 
         return fused_experts(x,
                              layer.w13_weight,

vllm/model_executor/models/deepseek_v2.py

@@ -115,7 +115,7 @@ def __init__(
                                 hidden_size=config.hidden_size,
                                 intermediate_size=config.moe_intermediate_size,
                                 reduce_results=False,
-                                renormalize=False,
+                                renormalize=True,
                                 quant_config=quant_config,
                                 use_grouped_topk=True,
                                 num_expert_group=config.n_group,