--- a/target/linux/qualcommax/image/ipq807x.mk
+++ b/target/linux/qualcommax/image/ipq807x.mk
@@ -256,12 +256,7 @@ define Device/xiaomi_ax3600
 	PAGESIZE := 2048
 	DEVICE_DTS_CONFIG := config@ac04
 	SOC := ipq8071
-	KERNEL_SIZE := 36608k
 	DEVICE_PACKAGES := ipq-wifi-xiaomi_ax3600 kmod-ath10k-ct-smallbuffers ath10k-firmware-qca9887-ct
-ifneq ($(CONFIG_TARGET_ROOTFS_INITRAMFS),)
-	ARTIFACTS := initramfs-factory.ubi
-	ARTIFACT/initramfs-factory.ubi := append-image-stage initramfs-uImage.itb | ubinize-kernel
-endif
 endef
 TARGET_DEVICES += xiaomi_ax3600
 
@@ -274,13 +269,8 @@ define Device/xiaomi_ax9000
 	PAGESIZE := 2048
 	DEVICE_DTS_CONFIG := config@hk14
 	SOC := ipq8072
-	KERNEL_SIZE := 57344k
 	DEVICE_PACKAGES := ipq-wifi-xiaomi_ax9000 kmod-ath11k-pci ath11k-firmware-qcn9074 \
 		kmod-ath10k-ct ath10k-firmware-qca9887-ct
-ifneq ($(CONFIG_TARGET_ROOTFS_INITRAMFS),)
-	ARTIFACTS := initramfs-factory.ubi
-	ARTIFACT/initramfs-factory.ubi := append-image-stage initramfs-uImage.itb | ubinize-kernel
-endif
 endef
 TARGET_DEVICES += xiaomi_ax9000
 

--- a/target/linux/qualcommax/ipq807x/base-files/etc/init.d/bootcount
+++ b/target/linux/qualcommax/ipq807x/base-files/etc/init.d/bootcount
@@ -9,5 +9,13 @@ boot() {
 		# Unset changed flag after sysupgrade complete
 		fw_setenv changed
 	;;
+	redmi,ax6|\
+	xiaomi,ax3600|\
+	xiaomi,ax9000)
+		# OTA handling should not be used. Reset it just in case.
+		fw_setenv flag_ota_reboot 0
+		# Not strictly needed but useful to handle partition crash condition
+		fw_setenv flag_boot_success 1
+	;;
 	esac
 }
 

--- a/target/linux/qualcommax/ipq807x/base-files/lib/upgrade/platform.sh
+++ b/target/linux/qualcommax/ipq807x/base-files/lib/upgrade/platform.sh
@@ -83,20 +83,29 @@ platform_do_upgrade() {
 	redmi,ax6|\
 	xiaomi,ax3600|\
 	xiaomi,ax9000)
-		# Make sure that UART is enabled
-		fw_setenv boot_wait on
-		fw_setenv uart_en 1
+		part_num="$(fw_printenv -n flag_boot_rootfs)"
+		if [ "$part_num" -eq "1" ]; then
+			CI_UBIPART="rootfs_1"
+			target_num=1
+			# Reset fail flag for the current partition
+			# With both partition set to fail, the partition 2 (bit 1)
+			# is loaded
+			fw_setenv flag_try_sys2_failed 0
+		else
+			CI_UBIPART="rootfs"
+			target_num=0
+			# Reset fail flag for the current partition
+			# or uboot will skip the loading of this partition
+			fw_setenv flag_try_sys1_failed 0
+		fi
+
+		# Tell uboot to switch partition
+		fw_setenv flag_boot_rootfs "$target_num"
+		fw_setenv flag_last_success "$target_num"
 
-		# Enforce single partition.
-		fw_setenv flag_boot_rootfs 0
-		fw_setenv flag_last_success 0
-		fw_setenv flag_boot_success 1
-		fw_setenv flag_try_sys1_failed 8
-		fw_setenv flag_try_sys2_failed 8
+		# Reset success flag
+		fw_setenv flag_boot_success 0
 
-		# Kernel and rootfs are placed in 2 different UBI
-		CI_KERN_UBIPART="ubi_kernel"
-		CI_ROOT_UBIPART="rootfs"
 		nand_do_upgrade "$1"
 		;;
 	*)