diff --git a/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp b/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
index ab19dddc33..f9a00b9ab7 100644
--- a/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
+++ b/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
@@ -361,10 +361,8 @@ TEST_F(KFDDBGTest, HitTrapOnWaveStartEndEvent) {
         for (int i = 0; i < 2; i++) {
             uint32_t enableMask = !!!(i % 2) ? KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_START :
                                                KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_END;
-            uint32_t supportedMask = enableMask;
-            debug->SetWaveLaunchOverride(KFD_DBG_TRAP_OVERRIDE_OR,
-                                                    &enableMask,
-                                                    &supportedMask);
+            uint32_t supportedMask = enableMask, reqMask = enableMask;
+            debug->SetWaveLaunchOverride(KFD_DBG_TRAP_OVERRIDE_OR, &reqMask, &supportedMask);
 
             if (!!!(supportedMask & enableMask)) {
                 EXPECT_SUCCESS(queue.Destroy());
@@ -374,7 +372,8 @@ TEST_F(KFDDBGTest, HitTrapOnWaveStartEndEvent) {
                 goto exit;
             }
 
-            ASSERT_EQ(enableMask, 0); // previous set mask
+	    // previous set mask
+            ASSERT_EQ(reqMask, !!!(i % 2) ? 0 : KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_START);
 
             Dispatch *dispatch;
             dispatch = new Dispatch(isaBuf);
diff --git a/libhsakmt/tests/kfdtest/src/ShaderStore.cpp b/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
index ed3c602f21..946adbc967 100644
--- a/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
+++ b/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
@@ -1130,6 +1130,11 @@ const char *TrapHandlerIsa =
         s_waitcnt lgkmcnt(0)
         s_mov_b32 m0, ttmp2
         v_mov_b32 v4, ttmp1
+        .if (.amdgcn.gfx_generation_number >= 12)
+            s_getreg_b32 ttmp14, hwreg(HW_REG_EXCP_FLAG_PRIV)
+            s_and_b32 ttmp2, ttmp14, 0x300
+            s_cbranch_scc1 RESTORE_AND_EXIT
+         .endif
         /* restore and increment program counter to skip shader trap jump*/
         s_add_u32 ttmp0, ttmp0, 4
         s_addc_u32 ttmp1, ttmp1, 0