--- a/drivers/mtd/mtk-snand/mtk-snand-mtd.c
+++ b/drivers/mtd/mtk-snand/mtk-snand-mtd.c
@@ -16,6 +16,7 @@
 #include <linux/dma-mapping.h>
 #include <linux/wait.h>
 #include <linux/mtd/mtd.h>
+#include <linux/mtd/mtk_bmt.h>
 #include <linux/mtd/partitions.h>
 #include <linux/of_platform.h>
 
@@ -612,6 +613,8 @@ static int mtk_snand_probe(struct platfo
 	mtd->_block_isbad = mtk_snand_mtd_block_isbad;
 	mtd->_block_markbad = mtk_snand_mtd_block_markbad;
 
+	mtk_bmt_attach(mtd);
+
 	ret = mtd_device_register(mtd, NULL, 0);
 	if (ret) {
 		dev_err(msm->pdev.dev, "failed to register mtd partition\n");
@@ -623,6 +626,7 @@ static int mtk_snand_probe(struct platfo
 	return 0;
 
 errout4:
+	mtk_bmt_detach(mtd);
 	devm_kfree(msm->pdev.dev, msm->page_cache);
 
 errout3:
@@ -650,6 +654,8 @@ static int mtk_snand_remove(struct platf
 	if (ret)
 		return ret;
 
+	mtk_bmt_detach(mtd);
+
 	mtk_snand_cleanup(msm->snf);
 
 	if (msm->irq >= 0)