diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gfx_v9.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gfx_v9.c
index 2e0702985bd90..6b822a51e0393 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gfx_v9.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gfx_v9.c
@@ -864,6 +864,13 @@ uint32_t kgd_gfx_v9_enable_debug_trap(struct kgd_dev *kgd,
 	data = 0;
 	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_MASK), data);
 
+	data = 0;
+	data = REG_SET_FIELD(data, SPI_GDBG_TRAP_CONFIG,
+		VMID_SEL, 1<<vmid);
+	data = REG_SET_FIELD(data, SPI_GDBG_TRAP_CONFIG,
+		TRAP_EN, 1);
+	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_CONFIG), data);
+
 	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_WAVE_CNTL), orig_stall_vmid);
 
 	mutex_unlock(&adev->grbm_idx_mutex);
@@ -877,6 +884,8 @@ uint32_t kgd_gfx_v9_disable_debug_trap(struct kgd_dev *kgd)
 
 	mutex_lock(&adev->grbm_idx_mutex);
 
+	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_CONFIG), 0);
+
 	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_DATA0), 0);
 	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_DATA1), 0);
 
diff --git a/drivers/gpu/drm/amd/amdgpu/gfx_v9_0.c b/drivers/gpu/drm/amd/amdgpu/gfx_v9_0.c
index 5c8107dc8803e..1627507e7f406 100644
--- a/drivers/gpu/drm/amd/amdgpu/gfx_v9_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gfx_v9_0.c
@@ -2427,8 +2427,6 @@ static void gfx_v9_0_init_compute_vmid(struct amdgpu_device *adev)
 	int i;
 	uint32_t sh_mem_config;
 	uint32_t sh_mem_bases;
-	uint32_t trap_config_vmid_mask = 0;
-	uint32_t data;
 
 	/*
 	 * Configure apertures:
@@ -2448,18 +2446,9 @@ static void gfx_v9_0_init_compute_vmid(struct amdgpu_device *adev)
 		/* CP and shaders */
 		WREG32_SOC15_RLC(GC, 0, mmSH_MEM_CONFIG, sh_mem_config);
 		WREG32_SOC15_RLC(GC, 0, mmSH_MEM_BASES, sh_mem_bases);
-
-		/* Calculate trap config vmid mask */
-		trap_config_vmid_mask |= (1 << i);
 	}
 	soc15_grbm_select(adev, 0, 0, 0, 0);
 	mutex_unlock(&adev->srbm_mutex);
-	data = 0;
-	data = REG_SET_FIELD(data, SPI_GDBG_TRAP_CONFIG,
-			VMID_SEL, trap_config_vmid_mask);
-	data = REG_SET_FIELD(data, SPI_GDBG_TRAP_CONFIG,
-			TRAP_EN, 1);
-	WREG32(SOC15_REG_OFFSET(GC, 0, mmSPI_GDBG_TRAP_CONFIG), data);
 }
 
 static void gfx_v9_0_init_gds_vmid(struct amdgpu_device *adev)