diff --git a/components/esp_hw_support/include/esp_private/sleep_cache.h b/components/esp_hw_support/include/esp_private/sleep_cache.h
index 3079b2ad42a..40eab5c76af 100644
--- a/components/esp_hw_support/include/esp_private/sleep_cache.h
+++ b/components/esp_hw_support/include/esp_private/sleep_cache.h
@@ -8,6 +8,7 @@
 
 #include <stdint.h>
 #include "sdkconfig.h"
+#include "soc/soc_caps.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -23,7 +24,7 @@ void sleep_cache_suspend(void);
  */
 void sleep_cache_resume(void);
 
-#if !CONFIG_SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE && CONFIG_SPIRAM
+#if SOC_PM_CPU_RETENTION_BY_SW && !SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE && CONFIG_SPIRAM
 /**
  * @brief Cache writeback before sleep when CPU/cache may power down (SPIRAM / PAU paths; Kconfig-driven).
  * @param sleep_flags Same flags as light sleep, e.g. @c PMU_SLEEP_PD_CPU when relevant.
diff --git a/components/esp_hw_support/sleep_cache.c b/components/esp_hw_support/sleep_cache.c
index 7d695cd2cd0..0ef117801ee 100644
--- a/components/esp_hw_support/sleep_cache.c
+++ b/components/esp_hw_support/sleep_cache.c
@@ -46,36 +46,22 @@ void sleep_cache_resume(void)
     }
 }
 
-#if !CONFIG_SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE && CONFIG_SPIRAM
+#if SOC_PM_CPU_RETENTION_BY_SW && !SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE && CONFIG_SPIRAM
+#if CONFIG_IDF_TARGET_ESP32S31
+#define DCACHE_WRITEBACK_MAP CACHE_MAP_L1_DCACHE
+#else
+#define DCACHE_WRITEBACK_MAP
+#endif
 void sleep_cache_safe_writeback(uint32_t sleep_flags)
 {
-#if CONFIG_PM_POWER_DOWN_CPU_IN_LIGHT_SLEEP && !SOC_PM_CACHE_RETENTION_BY_PAU && SOC_EXT_MEM_CACHE_TAG_IN_CPU_DOMAIN
-#if SOC_PMU_SUPPORTED
-    /* When SPIRAM is using, we need to use Cache_WriteBack_All to protect SPIRAM data
-        because the cache powers down when we power down the CPU */
     if (sleep_flags & PMU_SLEEP_PD_CPU) {
         if (s_cache_suspend_cnt) {
             spi_flash_restore_cache(esp_cpu_get_core_id(), s_cache_state);
         }
-        Cache_WriteBack_All();
+        Cache_WriteBack_All(DCACHE_WRITEBACK_MAP);
         if (s_cache_suspend_cnt) {
             spi_flash_disable_cache(esp_cpu_get_core_id(), &s_cache_state);
         }
     }
-#else
-    (void)sleep_flags;
-#endif
-#elif CONFIG_PM_POWER_DOWN_PERIPHERAL_IN_LIGHT_SLEEP && SOC_PM_CACHE_RETENTION_BY_PAU
-    (void)sleep_flags;
-    if (s_cache_suspend_cnt) {
-        spi_flash_restore_cache(esp_cpu_get_core_id(), s_cache_state);
-    }
-    Cache_WriteBack_All(CACHE_MAP_L1_DCACHE);
-    if (s_cache_suspend_cnt) {
-        spi_flash_disable_cache(esp_cpu_get_core_id(), &s_cache_state);
-    }
-#else
-    (void)sleep_flags;
-#endif
 }
 #endif
diff --git a/components/esp_hw_support/sleep_modes.c b/components/esp_hw_support/sleep_modes.c
index 435f1b3276d..f7bc70eb8ec 100644
--- a/components/esp_hw_support/sleep_modes.c
+++ b/components/esp_hw_support/sleep_modes.c
@@ -674,14 +674,24 @@ static SLEEP_FN_ATTR void misc_modules_sleep_prepare(uint32_t sleep_flags, bool
         mac_bb_power_down_cb_execute();
 #endif
 #if CONFIG_IDF_TARGET_ESP32
-    esp_sleep_gpio_pupd_config_workaround_apply();
+        esp_sleep_gpio_pupd_config_workaround_apply();
 #endif
-#if CONFIG_PM_POWER_DOWN_CPU_IN_LIGHT_SLEEP && SOC_PM_CPU_RETENTION_BY_RTCCNTL
+#if CONFIG_PM_POWER_DOWN_CPU_IN_LIGHT_SLEEP || CONFIG_PM_POWER_DOWN_PERIPHERAL_IN_LIGHT_SLEEP
+#if SOC_PM_CPU_RETENTION_BY_RTCCNTL
+        // Inside sleep_enable_cpu_retention, it will writeback the cache if the tag memory is power down with the CPU.
         sleep_enable_cpu_retention();
-#endif
-#if !SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE && CONFIG_SPIRAM
-    sleep_cache_safe_writeback(sleep_flags);
-#endif
+#elif SOC_PM_CPU_RETENTION_BY_SW && CONFIG_SPIRAM
+#if SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE
+        // For chips with SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE, if writeback is performed here,
+        // L1 dcache will still be dirty later, since the current function stack is in L2 MEM,
+        // so the writeback will be postponed to pmu_sleep.
+#else
+        // When SPIRAM is using, we need to writeback all dirty cache data to protect SPIRAM data conherence
+        // since the cache tag memory will be powered down with CPU.
+        sleep_cache_safe_writeback(sleep_flags);
+#endif // SOC_CACHE_INTERNAL_MEM_VIA_L1CACHE
+#endif // SOC_PM_CPU_RETENTION_BY_SW && CONFIG_SPIRAM
+#endif // CONFIG_PM_POWER_DOWN_CPU_IN_LIGHT_SLEEP || CONFIG_PM_POWER_DOWN_PERIPHERAL_IN_LIGHT_SLEEP
 #if ADC_LL_ANA_CALI_REG_PD_WORKAROUND
         adc_hal_i2c_saradc_reg_backup();
 #endif