diff --git a/target/i386/kvm/tdx.c b/target/i386/kvm/tdx.c
index 0041be589c21..60c70b6bed79 100644
--- a/target/i386/kvm/tdx.c
+++ b/target/i386/kvm/tdx.c
@@ -714,11 +714,7 @@ static void tdx_finalize_vm(Notifier *notifier, void *unused)
     tdx_post_init_vcpus();
 
     for_each_tdx_fw_entry(tdvf, entry) {
-        struct kvm_tdx_init_mem_region mem_region = {
-            .source_addr = (__u64)entry->mem_ptr,
-            .gpa = entry->address,
-            .nr_pages = entry->size / 4096,
-        };
+        struct kvm_memory_mapping mapping;
 
         r = kvm_set_memory_attributes_private(entry->address, entry->size);
         if (r < 0) {
@@ -726,14 +722,31 @@ static void tdx_finalize_vm(Notifier *notifier, void *unused)
              exit(1);
         }
 
-        __u32 flags = entry->attributes & TDVF_SECTION_ATTRIBUTES_MR_EXTEND ?
-                      KVM_TDX_MEASURE_MEMORY_REGION : 0;
-
-        trace_kvm_tdx_init_mem_region(entry->type, entry->attributes, mem_region.source_addr, mem_region.gpa, mem_region.nr_pages);
-        r = tdx_vm_ioctl(KVM_TDX_INIT_MEM_REGION, flags, &mem_region);
+        mapping = (struct kvm_memory_mapping) {
+            .base_gfn = entry->address / 4096,
+            .nr_pages = entry->size / 4096,
+            .source = (__u64)entry->mem_ptr,
+        };
+        do {
+            r = kvm_vcpu_ioctl(first_cpu, KVM_MEMORY_MAPPING, &mapping);
+        } while (r == -EAGAIN);
         if (r < 0) {
-             error_report("KVM_TDX_INIT_MEM_REGION failed %s", strerror(-r));
-             exit(1);
+            error_report("KVM_MEMORY_MAPPING failed %s", strerror(-r));
+            exit(1);
+        }
+
+        if (entry->attributes & TDVF_SECTION_ATTRIBUTES_MR_EXTEND) {
+            mapping = (struct kvm_memory_mapping) {
+                .base_gfn = entry->address / 4096,
+                .nr_pages = entry->size / 4096,
+            };
+            do {
+                r = tdx_vm_ioctl(KVM_TDX_EXTEND_MEMORY, 0, &mapping);
+            } while (r == -EAGAIN);
+            if (r < 0) {
+                error_report("KVM_TDX_EXTEND_MEMORY failed %s", strerror(-r));
+                exit(1);
+            }
         }
 
         if (entry->type == TDVF_SECTION_TYPE_TD_HOB ||
--

Required changes to tdx kselftest.

diff --git a/tools/testing/selftests/kvm/lib/x86_64/tdx/tdx_util.c b/tools/testing/selftests/kvm/lib/x86_64/tdx/tdx_util.c
index edc1b227a014..a2f4921c416f 100644
--- a/tools/testing/selftests/kvm/lib/x86_64/tdx/tdx_util.c
+++ b/tools/testing/selftests/kvm/lib/x86_64/tdx/tdx_util.c
@@ -23,7 +23,7 @@ static char *tdx_cmd_str[] = {
 	"KVM_TDX_CAPABILITIES",
 	"KVM_TDX_INIT_VM",
 	"KVM_TDX_INIT_VCPU",
-	"KVM_TDX_INIT_MEM_REGION",
+	"KVM_TDX_EXTEND_MEMORY",
 	"KVM_TDX_FINALIZE_VM",
 	"KVM_TDX_RELEASE_VM"
 };
@@ -179,19 +179,29 @@ static void tdx_td_vcpu_init(struct kvm_vcpu *vcpu)
 static void tdx_init_mem_region(struct kvm_vm *vm, void *source_pages,
 				uint64_t gpa, uint64_t size)
 {
-	struct kvm_tdx_init_mem_region mem_region = {
-		.source_addr = (uint64_t)source_pages,
-		.gpa = gpa,
+	struct kvm_vcpu * vcpu = list_first_entry(&vm->vcpus,
+						  struct kvm_vcpu, list);
+	struct kvm_memory_mapping mapping = {
+		.base_gfn = gpa / PAGE_SIZE,
 		.nr_pages = size / PAGE_SIZE,
+		.source = (uint64_t)source_pages,
 	};
-	uint32_t metadata = KVM_TDX_MEASURE_MEMORY_REGION;
+	int r;
 
-	TEST_ASSERT((mem_region.nr_pages > 0) &&
-			    ((mem_region.nr_pages * PAGE_SIZE) == size),
+	TEST_ASSERT((mapping.nr_pages > 0) &&
+		    ((mapping.nr_pages * PAGE_SIZE) == size),
 		    "Cannot add partial pages to the guest memory.\n");
 	TEST_ASSERT(((uint64_t)source_pages & (PAGE_SIZE - 1)) == 0,
 		    "Source memory buffer is not page aligned\n");
-	tdx_ioctl(vm->fd, KVM_TDX_INIT_MEM_REGION, metadata, &mem_region);
+	r = ioctl(vcpu->fd, KVM_MEMORY_MAPPING, &mapping);
+	TEST_ASSERT(r == 0, "KVM_MEMORY_MAPPING failed: %d %d",
+		    r, errno);
+
+	mapping = (struct kvm_memory_mapping) {
+		.base_gfn = gpa / PAGE_SIZE,
+		.nr_pages = size / PAGE_SIZE,
+	};
+	tdx_ioctl(vm->fd, KVM_TDX_EXTEND_MEMORY, 0, &mapping);
 }
 
 static void tdx_td_finalizemr(struct kvm_vm *vm)
@@ -231,7 +241,7 @@ static void tdx_enable_capabilities(struct kvm_vm *vm)
 		      KVM_X2APIC_API_USE_32BIT_IDS |
 			      KVM_X2APIC_API_DISABLE_BROADCAST_QUIRK);
 	vm_enable_cap(vm, KVM_CAP_SPLIT_IRQCHIP, 24);
-	vm_enable_cap(vm, KVM_CAP_MAX_VCPUS, 1024);
+	vm_enable_cap(vm, KVM_CAP_MAX_VCPUS, 512);
 }
 
 static void tdx_configure_memory_encryption(struct kvm_vm *vm)
