diff --git a/drivers/md/dm-table.c b/drivers/md/dm-table.c
index 7b22cc8d30f445526ed83a7730e77b6d3645a4eb..504e79bc3a556a8d94fa547185f23f170daecbf5 100644
--- a/drivers/md/dm-table.c
+++ b/drivers/md/dm-table.c
@@ -1823,6 +1823,8 @@ void dm_table_set_restrictions(struct dm_table *t, struct request_queue *q,
 	}
 	blk_queue_write_cache(q, wc, fua);
 
+	if (dm_table_supports_dax(t))
+		queue_flag_set_unlocked(QUEUE_FLAG_DAX, q);
 	if (dm_table_supports_dax_write_cache(t))
 		dax_write_cache(t->md->dax_dev, true);
 
diff --git a/drivers/md/dm.c b/drivers/md/dm.c
index 4e7682afebfa9be03f4ece3db7eb39e12da05e01..308d178fff738cbd772dc347241096fc73453ae9 100644
--- a/drivers/md/dm.c
+++ b/drivers/md/dm.c
@@ -2034,9 +2034,6 @@ int dm_setup_md_queue(struct mapped_device *md, struct dm_table *t)
 	case DM_TYPE_DAX_BIO_BASED:
 		dm_init_normal_md_queue(md);
 		blk_queue_make_request(md->queue, dm_make_request);
-
-		if (type == DM_TYPE_DAX_BIO_BASED)
-			queue_flag_set_unlocked(QUEUE_FLAG_DAX, md->queue);
 		break;
 	case DM_TYPE_NONE:
 		WARN_ON_ONCE(true);