[XPU] bfloat16 support for gather/gather_grad/scatter/scatter_grad (#71132)

houj04 · web-flow · commit 719d038dc3b5 · 2025-02-17T10:10:41.000+08:00
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20250204")
+  set(XPU_XHPC_BASE_DATE "dev/20250213")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.2.3") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -701,7 +701,14 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"floor", XPUKernelSet({phi::DataType::FLOAT32})},
       {"gather_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64,
+                     phi::DataType::BOOL})},
       {"gather_nd_grad",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
@@ -717,6 +724,9 @@ XPUOpMap& get_kl3_ops() {
       {"gather",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::BOOL})},
@@ -1154,9 +1164,13 @@ XPUOpMap& get_kl3_ops() {
       {"scatter",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"scatter_grad",
-       XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
+       XPUKernelSet({phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::FLOAT32})},
       {"scatter_nd_add",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/cpu/gather_grad_kernel.cc b/paddle/phi/kernels/cpu/gather_grad_kernel.cc
@@ -72,9 +72,12 @@ PD_REGISTER_KERNEL(gather_grad,
                    phi::GatherGradKernel,
                    float,
                    double,
-                   int,
                    uint8_t,
+                   int8_t,
+                   int16_t,
+                   int32_t,
                    int64_t,
+                   bool,
                    phi::dtype::bfloat16,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/gather_kernel.cc b/paddle/phi/kernels/cpu/gather_kernel.cc
@@ -70,9 +70,12 @@ PD_REGISTER_KERNEL(gather,
                    phi::GatherKernel,
                    float,
                    double,
-                   int,
                    uint8_t,
+                   int8_t,
+                   int16_t,
+                   int32_t,
                    int64_t,
+                   bool,
                    phi::dtype::bfloat16,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
@@ -358,11 +358,13 @@ void XPUFusedRotaryHalf(const Context& dev_ctx,
                         nullptr,
                         reinterpret_cast<const XPUSCType*>(sin_data),
                         reinterpret_cast<const XPUSCType*>(cos_data),
+                        nullptr,
                         reinterpret_cast<XPUType*>(out_q->data()),
                         nullptr,
                         {batch_size, seq_len, num_heads, head_dim},
                         {batch_size, seq_len, 1, head_dim},
                         {},
+                        0,
                         "BLHD",
                         -1,
                         10000.0f);
@@ -374,11 +376,13 @@ void XPUFusedRotaryHalf(const Context& dev_ctx,
                         reinterpret_cast<const XPUType*>(in_k->data()),
                         reinterpret_cast<const XPUSCType*>(sin_data),
                         reinterpret_cast<const XPUSCType*>(cos_data),
+                        nullptr,
                         reinterpret_cast<XPUType*>(out_q->data()),
                         reinterpret_cast<XPUType*>(out_k->data()),
                         {batch_size, seq_len, num_heads, head_dim},
                         {batch_size, seq_len, 1, head_dim},
                         {},
+                        0,
                         "BLHD",
                         num_heads_k,
                         10000.0f);
@@ -392,11 +396,13 @@ void XPUFusedRotaryHalf(const Context& dev_ctx,
                         nullptr,
                         reinterpret_cast<const XPUSCType*>(sin_data),
                         reinterpret_cast<const XPUSCType*>(cos_data),
+                        nullptr,
                         reinterpret_cast<XPUType*>(out_v->data()),
                         nullptr,
                         {batch_size, seq_len, num_heads_v, head_dim},
                         {batch_size, seq_len, 1, head_dim},
                         {},
+                        0,
                         "BLHD",
                         -1,
                         10000.0f);
diff --git a/paddle/phi/kernels/xpu/gather_grad_kernel.cc b/paddle/phi/kernels/xpu/gather_grad_kernel.cc
@@ -101,4 +101,10 @@ PD_REGISTER_KERNEL(gather_grad,
                    ALL_LAYOUT,
                    phi::GatherGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int8_t,
+                   int16_t,
+                   int32_t,
+                   int64_t,
+                   bool) {}
diff --git a/paddle/phi/kernels/xpu/gather_kernel.cc b/paddle/phi/kernels/xpu/gather_kernel.cc
@@ -88,6 +88,9 @@ PD_REGISTER_KERNEL(gather,
                    phi::GatherKernel,
                    float,
                    phi::dtype::float16,
-                   int,
+                   phi::dtype::bfloat16,
+                   int8_t,
+                   int16_t,
+                   int32_t,
                    int64_t,
                    bool) {}
diff --git a/paddle/phi/kernels/xpu/scatter_grad_kernel.cc b/paddle/phi/kernels/xpu/scatter_grad_kernel.cc
@@ -93,4 +93,5 @@ PD_REGISTER_KERNEL(scatter_grad,
                    ALL_LAYOUT,
                    phi::ScatterGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/scatter_kernel.cc b/paddle/phi/kernels/xpu/scatter_kernel.cc
@@ -119,6 +119,7 @@ PD_REGISTER_KERNEL(scatter,
                    ALL_LAYOUT,
                    phi::ScatterKernel,
                    float,
-                   int,
+                   int32_t,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_gather_op_xpu.py b/test/xpu/test_gather_op_xpu.py
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -45,7 +46,11 @@ def setUp(self):
             self.dtype = self.in_type
 
             self.init_config()
-            xnp = np.random.random(self.x_shape).astype(self.dtype)
+            xnp = np.random.random(self.x_shape).astype(
+                self.dtype if self.dtype != np.uint16 else np.float32
+            )
+            if self.dtype == np.uint16:
+                xnp = convert_float_to_uint16(xnp)
             self.inputs = {
                 'X': xnp,
                 'Index': np.array(self.index).astype(self.index_type),
diff --git a/test/xpu/test_scatter_op_xpu.py b/test/xpu/test_scatter_op_xpu.py
@@ -21,6 +21,7 @@
     get_xpu_op_support_types,
     type_dict_str_to_numpy,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -116,17 +117,29 @@ def setUp(self):
             )
 
             if not hasattr(self, 'init_ref_np'):
-                self.ref_np = np.ones((3, 50)).astype(self.dtype)
+                self.ref_np = np.ones((3, 50)).astype(
+                    self.dtype if self.dtype != np.uint16 else np.float32
+                )
                 self.index_np = np.array([1, 2]).astype(self.index_type)
-                self.updates_np = np.random.random((2, 50)).astype(self.dtype)
+                self.updates_np = np.random.random((2, 50)).astype(
+                    self.dtype if self.dtype != np.uint16 else np.float32
+                )
                 self.output_np = np.copy(self.ref_np)
                 self.output_np[self.index_np] = self.updates_np
             else:
-                self.ref_np = self.init_ref_np.astype(self.dtype)
+                self.ref_np = self.init_ref_np.astype(
+                    self.dtype if self.dtype != np.uint16 else np.float32
+                )
                 self.index_np = self.init_index_np.astype(self.index_type)
-                self.updates_np = self.init_updates_np.astype(self.dtype)
+                self.updates_np = self.init_updates_np.astype(
+                    self.dtype if self.dtype != np.uint16 else np.float32
+                )
                 self.output_np = self.init_output_np.astype(self.dtype)
 
+            if self.dtype == np.uint16:
+                self.ref_np = convert_float_to_uint16(self.ref_np)
+                self.updates_np = convert_float_to_uint16(self.updates_np)
+
             self.inputs = {
                 'X': self.ref_np,
                 'Ids': self.index_np,
