diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index 51f3101dd9ab..b4209ee80b57 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -392,8 +392,8 @@ struct pkvm_iommu_driver {
 	atomic_t state;
 };
 
-int pkvm_iommu_driver_init(struct pkvm_iommu_driver *drv, void *data, size_t size);
-int pkvm_iommu_register(struct device *dev, struct pkvm_iommu_driver *drv,
+int pkvm_iommu_driver_init(u64 drv, void *data, size_t size);
+int pkvm_iommu_register(struct device *dev, u64 drv,
 			phys_addr_t pa, size_t size, struct device *parent);
 int pkvm_iommu_suspend(struct device *dev);
 int pkvm_iommu_resume(struct device *dev);
diff --git a/arch/arm64/kvm/hyp/nvhe/hyp-main.c b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
index 040b98c0eebd..4e59ed5ae591 100644
--- a/arch/arm64/kvm/hyp/nvhe/hyp-main.c
+++ b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
@@ -1131,8 +1131,6 @@ static void handle___pkvm_iommu_driver_init(struct kvm_cpu_context *host_ctxt)
 	DECLARE_REG(void *, data, host_ctxt, 2);
 	DECLARE_REG(size_t, size, host_ctxt, 3);
 
-	/* __pkvm_iommu_driver_init expects hyp_va as it can be called from EL2 as a function. */
-	drv = kern_hyp_va(drv);
 	data = kern_hyp_va(data);
 
 	cpu_reg(host_ctxt, 1) = __pkvm_iommu_driver_init(drv, data, size);
@@ -1148,8 +1146,6 @@ static void handle___pkvm_iommu_register(struct kvm_cpu_context *host_ctxt)
 	DECLARE_REG(void *, mem, host_ctxt, 6);
 	DECLARE_REG(size_t, mem_size, host_ctxt, 7);
 
-	/* drv_id is the hyp address of the driver. */
-	drv_id = kern_hyp_va(drv_id);
 	cpu_reg(host_ctxt, 1) = __pkvm_iommu_register(dev_id, drv_id, dev_pa,
 						      dev_size, parent_id,
 						      mem, mem_size);
diff --git a/arch/arm64/kvm/iommu.c b/arch/arm64/kvm/iommu.c
index f827edb92eb6..6ea342c034d4 100644
--- a/arch/arm64/kvm/iommu.c
+++ b/arch/arm64/kvm/iommu.c
@@ -12,12 +12,12 @@ static unsigned long dev_to_id(struct device *dev)
 	return (unsigned long)dev;
 }
 
-int pkvm_iommu_driver_init(struct pkvm_iommu_driver *drv, void *data, size_t size)
+int pkvm_iommu_driver_init(u64 drv, void *data, size_t size)
 {
 	return kvm_call_hyp_nvhe(__pkvm_iommu_driver_init, drv, data, size);
 }
 
-int pkvm_iommu_register(struct device *dev, struct pkvm_iommu_driver *drv, phys_addr_t pa,
+int pkvm_iommu_register(struct device *dev, u64 drv, phys_addr_t pa,
 			size_t size, struct device *parent)
 {
 	void *mem;