[infer] Fix per_token_group_quant and support per_tensor_quant (#10359)

* support per_tensor_quant and per_token_group_quant

* check

* add per_token_group_quant, per_tensor_quant_fp8 into pybind_ops_list

Author: ckl117

csrc/gpu/cpp_extensions.cu

@@ -220,12 +220,14 @@ paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num
                                     const paddle::optional<paddle::Tensor>& output_padding_offset,
                                     int max_input_length);
 
-std::vector<paddle::Tensor> GroupQuant(const paddle::Tensor& x,
+std::vector<paddle::Tensor> PerTokenGroupQuant(const paddle::Tensor& x,
                                         const int group_size,
                                         const bool transpose_scale,
                                         const float quant_max_bound,
                                         const float quant_min_bound);
 
+std::vector<paddle::Tensor> PerTensorQuantFp8(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale);
+
 std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
                                                const paddle::Tensor& cum_offsets,
                                                const paddle::Tensor& token_num,
@@ -295,7 +297,8 @@ PYBIND11_MODULE(paddlenlp_ops, m) {
   m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
   m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
   m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
-  m.def("f_group_quant", &GroupQuant, "GroupQuant");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
   m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
   m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
   m.def("f_get_output", &GetOutput, "GetOutput");
@@ -324,7 +327,8 @@ PYBIND11_MODULE(paddlenlp_ops_80, m) {
   m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
   m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
   m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
-  m.def("f_group_quant", &GroupQuant, "GroupQuant");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
   m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
   m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
   m.def("f_get_output", &GetOutput, "GetOutput");
@@ -352,7 +356,8 @@ PYBIND11_MODULE(paddlenlp_ops_90, m) {
   m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
   m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
   m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
-  m.def("f_group_quant", &GroupQuant, "GroupQuant");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
   m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
   m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
   m.def("f_get_output", &GetOutput, "GetOutput");

csrc/gpu/group_quant.cu

@@ -162,6 +162,13 @@ class PDTraits<paddle::DataType::FLOAT8_E4M3FN> {
   typedef paddle::float8_e4m3fn data_t;
 };
 
+template <>
+class PDTraits<paddle::DataType::INT8> {
+public:
+  typedef int8_t DataType;
+  typedef int8_t data_t;
+};
+
 template <typename T, int Size>
 struct alignas(sizeof(T) * Size) AlignedVector {
   T val[Size];
@@ -245,3 +252,35 @@ inline bool GetMlaUseTensorcore() {
   const bool mla_use_tensorcore = flags_mla_use_tensorcore && enable_mla_tensorcore;
   return mla_use_tensorcore;
 }
+
+__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+    float old;
+    old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+                        : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
+    return old;
+}
+
+__device__ __forceinline__ float warpReduceMax(float max_value) {
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 16));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 8));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 4));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 2));
+    max_value = fmaxf(max_value, __shfl_xor_sync(0xffffffff, max_value, 1));
+    return max_value;
+}
+
+__device__ __forceinline__ float blockReduceMax(float max_value) {
+    static __shared__ float warpLevelMaxs[32];
+    const int laneId = threadIdx.x & 0x1f;;
+    const int warpId = threadIdx.x >> 5;
+
+    max_value = warpReduceMax(max_value);
+
+    if (laneId == 0) warpLevelMaxs[warpId] = max_value;
+        __syncthreads();
+
+    max_value = (threadIdx.x < blockDim.x / 32) ? warpLevelMaxs[laneId] : 0;
+    if (warpId == 0) max_value = warpReduceMax(max_value);
+
+    return max_value;
+}