diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
index cd68540017001..bd63aedd3a25a 100644
--- a/paddle/phi/backends/xpu/xpu2_op_list.cc
+++ b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -227,6 +227,12 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::INT8,
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32})},
+      {"flatten_with_xshape",
+       XPUKernelSet({phi::DataType::INT64,
+                     phi::DataType::INT32,
+                     phi::DataType::INT8,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::FLOAT32})},
       {"flatten_grad",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,