diff --git a/drivers/gpu/msm/a6xx_reg.h b/drivers/gpu/msm/a6xx_reg.h
index b3b7f2e5ce1b..6a72de66bbb6 100644
--- a/drivers/gpu/msm/a6xx_reg.h
+++ b/drivers/gpu/msm/a6xx_reg.h
@@ -989,6 +989,7 @@
 #define A6XX_GMU_ALWAYS_ON_COUNTER_L		0x1F888
 #define A6XX_GMU_ALWAYS_ON_COUNTER_H		0x1F889
 #define A6XX_GMU_GMU_PWR_COL_KEEPALIVE		0x1F8C3
+#define A6XX_GMU_PWR_COL_PREEMPT_KEEPALIVE      0x1F8C4
 #define A6XX_GMU_HFI_CTRL_STATUS		0x1F980
 #define A6XX_GMU_HFI_VERSION_INFO		0x1F981
 #define A6XX_GMU_HFI_SFR_ADDR			0x1F982
diff --git a/drivers/gpu/msm/adreno_a6xx_preempt.c b/drivers/gpu/msm/adreno_a6xx_preempt.c
index da8851edb4a2..7a52c786174e 100644
--- a/drivers/gpu/msm/adreno_a6xx_preempt.c
+++ b/drivers/gpu/msm/adreno_a6xx_preempt.c
@@ -70,6 +70,30 @@ static void _update_wptr(struct adreno_device *adreno_dev, bool reset_timer)
 	}
 }
 
+static void _power_collapse_set(struct adreno_device *adreno_dev, bool val)
+{
+	struct kgsl_device *device = KGSL_DEVICE(adreno_dev);
+
+	if (!gmu_core_isenabled(device))
+		return;
+
+	if (val) {
+		if (adreno_is_a660(adreno_dev))
+			gmu_core_regwrite(device,
+				 A6XX_GMU_PWR_COL_PREEMPT_KEEPALIVE, 0x1);
+		else
+			gmu_core_regrmw(device,
+				 A6XX_GMU_AO_SPARE_CNTL, 0x0, 0x2);
+	} else {
+		if (adreno_is_a660(adreno_dev))
+			gmu_core_regwrite(device,
+				 A6XX_GMU_PWR_COL_PREEMPT_KEEPALIVE, 0x0);
+		else
+			gmu_core_regrmw(device,
+				 A6XX_GMU_AO_SPARE_CNTL, 0x2, 0x0);
+	}
+}
+
 static void _a6xx_preemption_done(struct adreno_device *adreno_dev)
 {
 	struct kgsl_device *device = KGSL_DEVICE(adreno_dev);
@@ -287,8 +311,7 @@ void a6xx_preemption_trigger(struct adreno_device *adreno_dev)
 	 * free when the GPU is already powered on, whereas an OOB requires an
 	 * unconditional handshake with the GMU.
 	 */
-	if (gmu_core_isenabled(device))
-		gmu_core_regrmw(device, A6XX_GMU_AO_SPARE_CNTL, 0x0, 0x2);
+	_power_collapse_set(adreno_dev, true);
 
 	/*
 	 * Fenced writes on this path will make sure the GPU is woken up
@@ -384,8 +407,8 @@ err:
 	adreno_set_preempt_state(adreno_dev, ADRENO_PREEMPT_NONE);
 	adreno_dispatcher_schedule(device);
 	/* Clear the keep alive */
-	if (gmu_core_isenabled(device))
-		gmu_core_regrmw(device, A6XX_GMU_AO_SPARE_CNTL, 0x2, 0x0);
+	_power_collapse_set(adreno_dev, false);
+
 }
 
 void a6xx_preemption_callback(struct adreno_device *adreno_dev, int bit)
@@ -421,9 +444,7 @@ void a6xx_preemption_callback(struct adreno_device *adreno_dev, int bit)
 	 * We can now safely clear the preemption keepalive bit, allowing
 	 * power collapse to resume its regular activity.
 	 */
-	if (gmu_core_isenabled(KGSL_DEVICE(adreno_dev)))
-		gmu_core_regrmw(KGSL_DEVICE(adreno_dev),
-				A6XX_GMU_AO_SPARE_CNTL, 0x2, 0x0);
+	_power_collapse_set(adreno_dev, false);
 
 	del_timer(&adreno_dev->preempt.timer);