diff --git a/drivers/gpu/drm/rockchip/rockchip_drm_vop.h b/drivers/gpu/drm/rockchip/rockchip_drm_vop.h
index d3e299d9d4b6..251f44e9f0e0 100644
--- a/drivers/gpu/drm/rockchip/rockchip_drm_vop.h
+++ b/drivers/gpu/drm/rockchip/rockchip_drm_vop.h
@@ -613,6 +613,8 @@ struct vop2_wb_regs {
 	struct vop_reg scale_x_factor;
 	struct vop_reg scale_x_en;
 	struct vop_reg scale_y_en;
+	struct vop_reg axi_yrgb_id;
+	struct vop_reg axi_uv_id;
 };
 
 struct vop_win_data {
diff --git a/drivers/gpu/drm/rockchip/rockchip_drm_vop2.c b/drivers/gpu/drm/rockchip/rockchip_drm_vop2.c
index 6fb9754412af..e2dc55d92c73 100644
--- a/drivers/gpu/drm/rockchip/rockchip_drm_vop2.c
+++ b/drivers/gpu/drm/rockchip/rockchip_drm_vop2.c
@@ -2076,6 +2076,7 @@ static void vop2_initial(struct drm_crtc *crtc)
 	struct vop2_video_port *vp = to_vop2_video_port(crtc);
 	struct vop2 *vop2 = vp->vop2;
 	uint32_t current_vp_id = vp->id;
+	struct vop2_wb *wb = &vop2->wb;
 	int ret;
 
 	if (vop2->enable_count == 0) {
@@ -2094,6 +2095,10 @@ static void vop2_initial(struct drm_crtc *crtc)
 
 		memcpy(vop2->regsbak, vop2->regs, vop2->len);
 
+		VOP_MODULE_SET(vop2, wb, axi_yrgb_id, 0xd);
+		VOP_MODULE_SET(vop2, wb, axi_uv_id, 0xe);
+		vop2_wb_cfg_done(vop2);
+
 		VOP_CTRL_SET(vop2, cfg_done_en, 1);
 		/*
 		 * Disable auto gating, this is a workaround to
diff --git a/drivers/gpu/drm/rockchip/rockchip_vop2_reg.c b/drivers/gpu/drm/rockchip/rockchip_vop2_reg.c
index 614bac88bb1a..706bd9ead607 100644
--- a/drivers/gpu/drm/rockchip/rockchip_vop2_reg.c
+++ b/drivers/gpu/drm/rockchip/rockchip_vop2_reg.c
@@ -410,6 +410,8 @@ static const struct vop2_wb_regs rk3568_vop_wb_regs = {
 	.r2y_en = VOP_REG(RK3568_WB_CTRL, 0x1, 5),
 	.scale_x_en = VOP_REG(RK3568_WB_CTRL, 0x1, 7),
 	.scale_y_en = VOP_REG(RK3568_WB_CTRL, 0x1, 8),
+	.axi_yrgb_id = VOP_REG(RK3568_WB_CTRL, 0xff, 19),
+	.axi_uv_id = VOP_REG(RK3568_WB_CTRL, 0x1f, 27),
 	.scale_x_factor = VOP_REG(RK3568_WB_XSCAL_FACTOR, 0x3fff, 16),
 	.yrgb_mst = VOP_REG(RK3568_WB_YRGB_MST, 0xffffffff, 0),
 	.uv_mst = VOP_REG(RK3568_WB_CBR_MST, 0xffffffff, 0),