diff --git a/tools/testing/selftests/kvm/lib/x86_64/processor.c b/tools/testing/selftests/kvm/lib/x86_64/processor.c
index 2a6e28c769f2f..9e848d43ed5d1 100644
--- a/tools/testing/selftests/kvm/lib/x86_64/processor.c
+++ b/tools/testing/selftests/kvm/lib/x86_64/processor.c
@@ -636,7 +636,7 @@ struct kvm_vcpu *vm_arch_vcpu_add(struct kvm_vm *vm, uint32_t vcpu_id,
 	vcpu_regs_set(vcpu, &regs);
 
 	/* Setup the MP state */
-	mp_state.mp_state = 0;
+	mp_state.mp_state = KVM_MP_STATE_RUNNABLE;
 	vcpu_mp_state_set(vcpu, &mp_state);
 
 	return vcpu;
@@ -662,7 +662,7 @@ struct kvm_vcpu *vm_vcpu_add_tdx(struct kvm_vm *vm, uint32_t vcpu_id)
 	initialize_td_vcpu(vcpu);
 
 	/* Setup the MP state */
-	mp_state.mp_state = 0;
+	mp_state.mp_state = KVM_MP_STATE_RUNNABLE;
 	vcpu_mp_state_set(vcpu, &mp_state);
 
 	return vcpu;
diff --git a/tools/testing/selftests/kvm/lib/x86_64/tdx.h b/tools/testing/selftests/kvm/lib/x86_64/tdx.h
index be8564f4672db..bfa3709a76e5f 100644
--- a/tools/testing/selftests/kvm/lib/x86_64/tdx.h
+++ b/tools/testing/selftests/kvm/lib/x86_64/tdx.h
@@ -119,6 +119,7 @@ struct page_table {
 void add_td_memory(struct kvm_vm *vm, void *source_page,
 		   uint64_t gpa, int size);
 void finalize_td_memory(struct kvm_vm *vm);
+void get_tdx_capabilities(struct kvm_vm *vm);
 void initialize_td(struct kvm_vm *vm);
 void initialize_td_with_attributes(struct kvm_vm *vm, uint64_t attributes);
 void initialize_td_vcpu(struct kvm_vcpu *vcpu);
diff --git a/tools/testing/selftests/kvm/lib/x86_64/tdx_lib.c b/tools/testing/selftests/kvm/lib/x86_64/tdx_lib.c
index 23893949c3a1c..9193e2f80b617 100644
--- a/tools/testing/selftests/kvm/lib/x86_64/tdx_lib.c
+++ b/tools/testing/selftests/kvm/lib/x86_64/tdx_lib.c
@@ -25,12 +25,11 @@ char *tdx_cmd_str[] = {
 #define XFEATURE_MASK_XTILECFG	(1 << XFEATURE_XTILECFG)
 #define XFEATURE_MASK_XTILEDATA	(1 << XFEATURE_XTILEDATA)
 #define XFEATURE_MASK_XTILE	(XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA)
+#define XFEATURE_MASK_CET	((1 << 11) | (1 << 12))
 
-
-static void tdx_ioctl(int fd, int ioctl_no, uint32_t flags, void *data)
+static int __tdx_ioctl(int fd, int ioctl_no, uint32_t flags, void *data)
 {
 	struct kvm_tdx_cmd tdx_cmd;
-	int r;
 
 	TEST_ASSERT(ioctl_no < TDX_MAX_CMD_STR, "Unknown TDX CMD : %d\n",
 		    ioctl_no);
@@ -39,7 +38,15 @@ static void tdx_ioctl(int fd, int ioctl_no, uint32_t flags, void *data)
 	tdx_cmd.id = ioctl_no;
 	tdx_cmd.flags = flags;
 	tdx_cmd.data = (uint64_t)data;
-	r = ioctl(fd, KVM_MEMORY_ENCRYPT_OP, &tdx_cmd);
+	return ioctl(fd, KVM_MEMORY_ENCRYPT_OP, &tdx_cmd);
+}
+
+
+static void tdx_ioctl(int fd, int ioctl_no, uint32_t flags, void *data)
+{
+	int r;
+
+	r = __tdx_ioctl(fd, ioctl_no, flags, data);
 	TEST_ASSERT(r == 0, "%s failed: %d  %d", tdx_cmd_str[ioctl_no], r,
 		    errno);
 }
@@ -65,18 +72,71 @@ static struct tdx_cpuid_data get_tdx_cpuid_data(struct kvm_vm *vm)
 	for (i = 0; i < KVM_MAX_CPUID_ENTRIES; i++) {
 		struct kvm_cpuid_entry2 *e = &cpuid_data.entries[i];
 
-		/* TDX doesn't support LBR and AMX features yet.
+		/* TDX doesn't support LBR yet.
 		 * Disable those bits from the XCR0 register.
 		 */
 		if (e->function == 0xd && (e->index == 0)) {
 			e->eax &= ~XFEATURE_MASK_LBR;
-			e->eax &= ~XFEATURE_MASK_XTILE;
+
+			/*
+			 * TDX modules requires both CET_{U, S} to be set even
+			 * if only one is supported.
+			 */
+			if (e->eax & XFEATURE_MASK_CET) {
+				e->eax |= XFEATURE_MASK_CET;
+			}
+			/*
+			 * TDX module requires both XTILE_{CFG, DATA} to be set.
+			 * Both bits are required for AMX to be functional.
+			 */
+			if ((e->eax & XFEATURE_MASK_XTILE) != XFEATURE_MASK_XTILE) {
+				e->eax &= ~XFEATURE_MASK_XTILE;
+			}
 		}
 	}
 
 	return cpuid_data;
 }
 
+/* Call KVM_TDX_CAPABILITIES for API test. The result isn't used. */
+void get_tdx_capabilities(struct kvm_vm *vm)
+{
+	int i;
+	int rc;
+	int nr_cpuid_configs = 8;
+	struct kvm_tdx_capabilities *tdx_cap = NULL;
+
+	while (true) {
+		tdx_cap = realloc(
+			tdx_cap, sizeof(*tdx_cap) +
+			nr_cpuid_configs * sizeof(*tdx_cap->cpuid_configs));
+		tdx_cap->nr_cpuid_configs = nr_cpuid_configs;
+		TEST_ASSERT(tdx_cap != NULL,
+			"Could not allocate memory for tdx capability "
+			"nr_cpuid_configs %d\n", nr_cpuid_configs);
+		rc = __tdx_ioctl(vm->fd, KVM_TDX_CAPABILITIES, 0, tdx_cap);
+		if (rc < 0 && errno == E2BIG) {
+			nr_cpuid_configs *= 2;
+			continue;
+		}
+		TEST_ASSERT(rc == 0, "%s failed: %d %d",
+			tdx_cmd_str[KVM_TDX_CAPABILITIES], rc, errno);
+		break;
+	}
+	pr_debug("tdx_cap: attrs: fixed0 0x%016llx fixed1 0x%016llx\n"
+		"tdx_cap: xfam fixed0 0x%016llx fixed1 0x%016llx\n",
+		tdx_cap->attrs_fixed0, tdx_cap->attrs_fixed1,
+		tdx_cap->xfam_fixed0, tdx_cap->xfam_fixed1);
+	for (i = 0; i < tdx_cap->nr_cpuid_configs; i++) {
+		const struct kvm_tdx_cpuid_config *config =
+			&tdx_cap->cpuid_configs[i];
+		pr_debug("cpuid config[%d]: leaf 0x%x sub_leaf 0x%x "
+			"eax 0x%08x ebx 0x%08x ecx 0x%08x edx 0x%08x\n",
+			i, config->leaf, config->sub_leaf,
+			config->eax, config->ebx, config->ecx, config->edx);
+	}
+}
+
 /*
  * Initialize a VM as a TD with attributes.
  *
diff --git a/tools/testing/selftests/kvm/x86_64/tdx_vm_tests.c b/tools/testing/selftests/kvm/x86_64/tdx_vm_tests.c
index a96abada54b63..b3f9e3fa41f4b 100644
--- a/tools/testing/selftests/kvm/x86_64/tdx_vm_tests.c
+++ b/tools/testing/selftests/kvm/x86_64/tdx_vm_tests.c
@@ -197,6 +197,9 @@ void verify_td_lifecycle(void)
 	/* Create a TD VM with no memory.*/
 	vm = vm_create_tdx();
 
+	/* Get TDX capabilities */
+	get_tdx_capabilities(vm);
+
 	/* Allocate TD guest memory and initialize the TD.*/
 	initialize_td(vm);