update fused_moe,py

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -644,7 +644,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             expert_ids,
             num_tokens_post_padded,
             B.size(1) if not use_nn_moe else B.size(2),
-            B.size(1),
+            A.size(1),
             EM,
             num_tokens,
             A.stride(0),
@@ -1081,7 +1081,7 @@ def inplace_fused_experts(
                        use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,
                        use_mxfp4_w4a4, per_channel_quant, global_num_experts,
                        expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
-                       a2_scale, block_shape, w1_bias, w2_bias)
+                       a2_scale, block_shape, w1_bias, w2_bias, use_nn_moe)
 
 
 def inplace_fused_experts_fake(hidden_states: torch.Tensor,
@@ -1108,7 +1108,8 @@ def inplace_fused_experts_fake(hidden_states: torch.Tensor,
                                a2_scale: Optional[torch.Tensor] = None,
                                block_shape: Optional[list[int]] = None,
                                w1_bias: Optional[torch.Tensor] = None,
-                               w2_bias: Optional[torch.Tensor] = None) -> None:
+                               w2_bias: Optional[torch.Tensor] = None,
+                               use_nn_moe: Optional[bool] = False) -> None:
     pass