diff --git a/kernel/ethosu_capabilities.h b/kernel/ethosu_capabilities.h
index 3a79dc5..30bf475 100644
--- a/kernel/ethosu_capabilities.h
+++ b/kernel/ethosu_capabilities.h
@@ -42,7 +42,7 @@
  * struct ethosu_capabilities - Capabilities internal struct
  */
 struct ethosu_capabilities {
-	struct ethosu_device                   *edev;
+	struct device                          *dev;
 	struct completion                      done;
 	struct ethosu_uapi_device_capabilities *uapi;
 	struct ethosu_mailbox_msg              msg;
@@ -53,10 +53,11 @@
  * Functions
  ****************************************************************************/
 
-int ethosu_capabilities_request(struct ethosu_device *edev,
+int ethosu_capabilities_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
 				struct ethosu_uapi_device_capabilities *uapi);
 
-void ethosu_capability_rsp(struct ethosu_device *edev,
+void ethosu_capability_rsp(struct ethosu_mailbox *mailbox,
 			   int msg_id,
 			   struct ethosu_core_msg_capabilities_rsp *rsp);