diff --git a/arch/arm64/kvm/hisilicon/hisi_virt.c b/arch/arm64/kvm/hisilicon/hisi_virt.c
index d108b0ba31469f583c31a73ed006dbc3d6c611c3..647fa62f2a4e765239f0736e272d06407b42836c 100644
--- a/arch/arm64/kvm/hisilicon/hisi_virt.c
+++ b/arch/arm64/kvm/hisilicon/hisi_virt.c
@@ -486,7 +486,9 @@ static void kvm_update_vm_lsudvmbm_hip12(struct kvm *kvm)
 
 	if (nr_dies == 1) {
 		val = DVMBM_RANGE_ONE_DIE << DVMBM_RANGE_SHIFT	|
-		      vm_aff3s[0] << DVMBM_DIE1_VDIE_SHIFT_HIP12;
+		      (vm_aff3s[0] & MPIDR_AFF3_DIE_ID_MASK) << DVMBM_DIE1_VDIE_SHIFT_HIP12 |
+		      ((vm_aff3s[0] & MPIDR_AFF3_SOCKET_ID_MASK) >> MPIDR_AFF3_SOCKET_ID_SHIFT)
+		            << DVMBM_DIE1_SOCKET_SHIFT_HIP12;
 
 		/* fulfill bits [11:6] */
 		for_each_cpu(cpu, kvm->arch.sched_cpus) {
@@ -502,8 +504,12 @@ static void kvm_update_vm_lsudvmbm_hip12(struct kvm *kvm)
 	/* nr_dies == 2 */
 	val = DVMBM_RANGE_TWO_DIES << DVMBM_RANGE_SHIFT	|
 	      DVMBM_GRAN_CLUSTER << DVMBM_GRAN_SHIFT	|
-	      vm_aff3s[0] << DVMBM_DIE1_VDIE_SHIFT_HIP12    |
-	      vm_aff3s[1] << DVMBM_DIE2_VDIE_SHIFT_HIP12;
+	      (vm_aff3s[0] & MPIDR_AFF3_DIE_ID_MASK) << DVMBM_DIE1_VDIE_SHIFT_HIP12 |
+	      ((vm_aff3s[0] & MPIDR_AFF3_SOCKET_ID_MASK) >> MPIDR_AFF3_SOCKET_ID_SHIFT)
+	            << DVMBM_DIE1_SOCKET_SHIFT_HIP12 |
+	      (vm_aff3s[1] & MPIDR_AFF3_DIE_ID_MASK) << DVMBM_DIE2_VDIE_SHIFT_HIP12 |
+	      ((vm_aff3s[1] & MPIDR_AFF3_SOCKET_ID_MASK) >> MPIDR_AFF3_SOCKET_ID_SHIFT)
+	            << DVMBM_DIE2_SOCKET_SHIFT_HIP12;
 
 	/* and fulfill bits [11:0] */
 	for_each_cpu(cpu, kvm->arch.sched_cpus) {
diff --git a/arch/arm64/kvm/hisilicon/hisi_virt.h b/arch/arm64/kvm/hisilicon/hisi_virt.h
index 85dccafde8a63c3cd3f1db7595470c9566190a90..3ae2cb9e7890a258989b51cb91de10d3d46c670d 100644
--- a/arch/arm64/kvm/hisilicon/hisi_virt.h
+++ b/arch/arm64/kvm/hisilicon/hisi_virt.h
@@ -70,11 +70,16 @@ enum hisi_cpu_type {
 #define DVMBM_MAX_DIES			32
 
 /* HIP12 */
+#define DVMBM_DIE1_SOCKET_SHIFT_HIP12		59
+#define DVMBM_DIE2_SOCKET_SHIFT_HIP12		55
 #define DVMBM_DIE1_VDIE_SHIFT_HIP12		57
 #define DVMBM_DIE2_VDIE_SHIFT_HIP12		53
 #define DVMBM_DIE1_CLUSTER_SHIFT_HIP12		6
 #define DVMBM_DIE2_CLUSTER_SHIFT_HIP12		0
 #define DVMBM_MAX_DIES_HIP12				8
+#define MPIDR_AFF3_DIE_ID_MASK				0x7
+#define MPIDR_AFF3_SOCKET_ID_MASK			0x18
+#define MPIDR_AFF3_SOCKET_ID_SHIFT			3
 
 void probe_hisi_cpu_type(void);
 bool hisi_ncsnp_supported(void);