static void ufs_mtk_vreg_set_lpm(struct ufs_hba *hba, bool lpm)
 {
-       if (!hba->vreg_info.vccq2 || !hba->vreg_info.vcc)
+       struct ufs_vreg *vccqx = NULL;
+
+       if (!hba->vreg_info.vccq && !hba->vreg_info.vccq2)
+               return;
+
+       /* Skip if VCC is assumed always-on */
+       if (!hba->vreg_info.vcc)
                return;
 
        /* Bypass LPM when device is still active */
        if (lpm && hba->vreg_info.vcc->enabled)
                return;
 
-       if (lpm)
-               regulator_set_mode(hba->vreg_info.vccq2->reg,
-                                  REGULATOR_MODE_IDLE);
+       if (hba->vreg_info.vccq)
+               vccqx = hba->vreg_info.vccq;
        else
-               regulator_set_mode(hba->vreg_info.vccq2->reg,
-                                  REGULATOR_MODE_NORMAL);
+               vccqx = hba->vreg_info.vccq2;
+
+       regulator_set_mode(vccqx->reg,
+                          lpm ? REGULATOR_MODE_IDLE : REGULATOR_MODE_NORMAL);
 }
 
 static void ufs_mtk_auto_hibern8_disable(struct ufs_hba *hba)