apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2026-04-23 00:17:25 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity

[Feature] Unify fp8 block_wise quant ops (#5991)

Browse Source 
* quant stash

* blockwise_quant

* precommit

* rm tensor.cut

* tp ok

* add swiglu

* rm outdate code

* fix activate ut

* change baseline

* fix baseline error
This commit is contained in:
fxyfxy777
2026-01-15 21:50:37 +08:00
committed by GitHub
parent d38cd8b40b
commit 4c92035f2d
17 changed files with 55 additions and 571 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_ops/gpu_ops/per_token_quant_fp8.cu
-323
View File
@@ -16,313 +16,6 @@
constexpr float epsilon = 1e-10;
template <typename T>
__global__ void quant_per_token_per_block(
const T *input,
phi::dtype::float8_e4m3fn *quanted_res,
float *quanted_scale,
const int token_num,
const int hidden_size,
const int hidden_size_scale,
const bool use_finegrained_range) {
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int num_warp = blockDim.x / 32;
static constexpr int NUM_PER_THREADS = 128 / 32; // 4
static constexpr float MAX_VALUE = 448.f;
// Note(ZKK) use ceil_div!!
const int end_iter = (hidden_size + 127) / 128; // warp_iter_num
AlignedVector<T, NUM_PER_THREADS> load_vec;
AlignedVector<float, NUM_PER_THREADS> load_vec_float;
AlignedVector<phi::dtype::float8_e4m3fn, NUM_PER_THREADS> res_vec;
for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
const T *input_now = input + token_idx * hidden_size;
phi::dtype::float8_e4m3fn *quanted_res_now =
quanted_res + token_idx * hidden_size;
float *quanted_scale_now = quanted_scale + token_idx * hidden_size_scale;
// deal a block per warp
for (int iter = warp_id; iter < end_iter; iter += num_warp) {
const int start_offset = iter * 128;
const bool is_valid_data =
start_offset + lane_id * NUM_PER_THREADS < hidden_size;
if (is_valid_data) {
Load<T, NUM_PER_THREADS>(
input_now + start_offset + lane_id * NUM_PER_THREADS, &load_vec);
} else {
#pragma unroll
for (int vid = 0; vid < NUM_PER_THREADS; vid++) load_vec[vid] = T(0.f);
}
// get max value per thread
float max_value_thread = -5e4;
#pragma unroll
for (int vid = 0; vid < NUM_PER_THREADS; vid++) {
load_vec_float[vid] = static_cast<float>(load_vec[vid]);
max_value_thread = max(abs(load_vec_float[vid]), max_value_thread);
}
// get max value per warp
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 16),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 8),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 4),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 2),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 1),
max_value_thread);
// broadcast max_value
max_value_thread = __shfl_sync(0xFFFFFFFF, max_value_thread, 0);
max_value_thread = max(max_value_thread, epsilon);
if (use_finegrained_range) {
max_value_thread *= 7.0f;
}
float scale_to_store = max_value_thread / MAX_VALUE;
// quant
#pragma unroll
for (int vid = 0; vid < NUM_PER_THREADS; vid++) {
res_vec[vid] = static_cast<phi::dtype::float8_e4m3fn>(
load_vec_float[vid] * MAX_VALUE / max_value_thread);
}
// store
if (is_valid_data)
Store<phi::dtype::float8_e4m3fn, NUM_PER_THREADS>(
res_vec,
quanted_res_now + start_offset + lane_id * NUM_PER_THREADS);
if (lane_id == 0) {
quanted_scale_now[iter] = scale_to_store;
}
}
}
}
std::vector<paddle::Tensor> PerTokenQuant(paddle::Tensor &input,
const int block_size) {
auto input_dim = input.dims();
const int token_num = input_dim[0];
const int hidden_size = input_dim[1];
// Note(ZKK) here we use ceil_dive to support 4.5T runing on 8 GPUS
// where moe_intermediate_size is 448, can not be divided by 128.
const int hidden_size_scale = (hidden_size + block_size - 1) / block_size;
auto quanted_x = GetEmptyTensor(
{token_num, hidden_size}, paddle::DataType::FLOAT8_E4M3FN, input.place());
auto quanted_scale = GetEmptyTensor(
{token_num, hidden_size_scale}, paddle::DataType::FLOAT32, input.place());
const int gridx = min(132 * 8, token_num);
const int blockx = min(1024, hidden_size / 128 * 32);
bool use_finegrained_range = false;
char *env_var = getenv("PER_TOKEN_QUANT_FP8_USE_FINEGRAINED_RANGE");
if (env_var) {
use_finegrained_range = static_cast<bool>(std::stoi(env_var));
}
switch (input.dtype()) {
case paddle::DataType::BFLOAT16:
quant_per_token_per_block<<<gridx, blockx, 0, input.stream()>>>(
input.data<paddle::bfloat16>(),
quanted_x.data<phi::dtype::float8_e4m3fn>(),
quanted_scale.data<float>(),
token_num,
hidden_size,
hidden_size_scale,
use_finegrained_range);
break;
case paddle::DataType::FLOAT16:
quant_per_token_per_block<<<gridx, blockx, 0, input.stream()>>>(
input.data<paddle::float16>(),
quanted_x.data<phi::dtype::float8_e4m3fn>(),
quanted_scale.data<float>(),
token_num,
hidden_size,
hidden_size_scale,
use_finegrained_range);
break;
default:
PD_THROW("Unsupported data type for PerTokenQuant");
}
return {quanted_x, quanted_scale};
}
std::vector<std::vector<int64_t>> PerTokenQuantInferShape(
std::vector<int64_t> input_shape, const int block_size) {
const int token_num = input_shape[0];
const int hidden_size = input_shape[1];
const int hidden_size_scale = (hidden_size + block_size - 1) / block_size;
return {{token_num, hidden_size}, {token_num, hidden_size_scale}};
}
std::vector<paddle::DataType> PerTokenQuantInferDtype(
paddle::DataType input_dtype, const int block_size) {
return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
}
template <typename T>
__global__ void quant_per_token_per_block_padding(
const T *input,
phi::dtype::float8_e4m3fn *quanted_res,
float *quanted_scale,
const int token_num,
const int padded_token_num,
const int hidden_size,
const int hidden_size_scale,
const bool use_finegrained_range) {
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int num_warp = blockDim.x / 32;
static constexpr int NUM_PER_THREADS = 128 / 32; // 4
static constexpr float MAX_VALUE = 448.f;
const int end_iter = hidden_size / 128; // warp_iter_num
AlignedVector<T, NUM_PER_THREADS> load_vec;
AlignedVector<float, NUM_PER_THREADS> load_vec_float;
AlignedVector<phi::dtype::float8_e4m3fn, NUM_PER_THREADS> res_vec;
for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
const T *input_now = input + token_idx * hidden_size;
phi::dtype::float8_e4m3fn *quanted_res_now =
quanted_res + token_idx * hidden_size;
// deal a block per warp
for (int iter = warp_id; iter < end_iter; iter += num_warp) {
float *quanted_scale_now =
quanted_scale + iter * padded_token_num + token_idx;
const int start_offset = iter * 128;
Load<T, NUM_PER_THREADS>(
input_now + start_offset + lane_id * NUM_PER_THREADS, &load_vec);
// get max value per thread
float max_value_thread = -5e4;
#pragma unroll
for (int vid = 0; vid < NUM_PER_THREADS; vid++) {
load_vec_float[vid] = static_cast<float>(load_vec[vid]);
max_value_thread = max(abs(load_vec_float[vid]), max_value_thread);
}
// get max value per warp
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 16),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 8),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 4),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 2),
max_value_thread);
max_value_thread = max(__shfl_down_sync(0xffffffff, max_value_thread, 1),
max_value_thread);
// broadcast max_value
max_value_thread = __shfl_sync(0xFFFFFFFF, max_value_thread, 0);
max_value_thread = max(max_value_thread, epsilon);
if (use_finegrained_range) {
max_value_thread *= 7.0f;
}
float scale_to_store = max_value_thread / MAX_VALUE;
// quant
#pragma unroll
for (int vid = 0; vid < NUM_PER_THREADS; vid++) {
res_vec[vid] = static_cast<phi::dtype::float8_e4m3fn>(
load_vec_float[vid] * MAX_VALUE / max_value_thread);
}
// store
Store<phi::dtype::float8_e4m3fn, NUM_PER_THREADS>(
res_vec, quanted_res_now + start_offset + lane_id * NUM_PER_THREADS);
if (lane_id == 0) {
*quanted_scale_now = scale_to_store;
}
}
}
}
std::vector<paddle::Tensor> PerTokenQuantPadding(paddle::Tensor &input,
const int block_size) {
using ScaleDtype = float;
auto input_dim = input.dims();
const int token_num = input_dim[0];
const int hidden_size = input_dim[1];
PADDLE_ENFORCE(block_size == 128, "now only support block_size = 128");
PADDLE_ENFORCE(hidden_size % 128 == 0,
"hidden_size must be divisible by 128");
const int hidden_size_scale = hidden_size / block_size;
auto quanted_x = GetEmptyTensor(
{token_num, hidden_size}, paddle::DataType::FLOAT8_E4M3FN, input.place());
const int tma_alignment_bytes = 16;
const int tma_alignment_elements = tma_alignment_bytes / sizeof(ScaleDtype);
const int padded_token_num =
((token_num + tma_alignment_elements - 1) / tma_alignment_elements) *
tma_alignment_elements;
auto quanted_scale = GetEmptyTensor({padded_token_num, hidden_size_scale},
{1, padded_token_num},
paddle::DataType::FLOAT32,
input.place());
const int gridx = min(132 * 8, token_num);
const int blockx = min(1024, hidden_size / 128 * 32);
bool use_finegrained_range = false;
char *env_var = getenv("PER_TOKEN_QUANT_FP8_USE_FINEGRAINED_RANGE");
if (env_var) {
use_finegrained_range = static_cast<bool>(std::stoi(env_var));
}
switch (input.dtype()) {
case paddle::DataType::BFLOAT16:
quant_per_token_per_block_padding<<<gridx, blockx, 0, input.stream()>>>(
input.data<paddle::bfloat16>(),
quanted_x.data<phi::dtype::float8_e4m3fn>(),
quanted_scale.data<ScaleDtype>(),
token_num,
padded_token_num,
hidden_size,
hidden_size_scale,
use_finegrained_range);
break;
case paddle::DataType::FLOAT16:
quant_per_token_per_block_padding<<<gridx, blockx, 0, input.stream()>>>(
input.data<paddle::float16>(),
quanted_x.data<phi::dtype::float8_e4m3fn>(),
quanted_scale.data<ScaleDtype>(),
token_num,
padded_token_num,
hidden_size,
hidden_size_scale,
use_finegrained_range);
break;
default:
PD_THROW("Unsupported data type for PerTokenQuant");
}
return {quanted_x, quanted_scale};
}
std::vector<std::vector<int64_t>> PerTokenQuantPaddingInferShape(
std::vector<int64_t> input_shape, const int block_size) {
using ScaleDtype = float;
const int token_num = input_shape[0];
const int hidden_size = input_shape[1];
const int hidden_size_scale = hidden_size / block_size;
const int tma_alignment_bytes = 16;
const int tma_alignment_elements = tma_alignment_bytes / sizeof(ScaleDtype);
const int padded_token_num =
((token_num + tma_alignment_elements - 1) / tma_alignment_elements) *
tma_alignment_elements;
return {{token_num, hidden_size}, {padded_token_num, hidden_size_scale}};
}
std::vector<paddle::DataType> PerTokenQuantPaddingInferDtype(
paddle::DataType input_dtype) {
return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
}
template <typename T>
__global__ void masked_quant_per_token_per_block(
const T *input,
@@ -472,22 +165,6 @@ std::vector<paddle::Tensor> MaskedPerTokenQuant(
return {quanted_x, quanted_scale};
}
PD_BUILD_STATIC_OP(per_token_quant)
.Inputs({"input"})
.Outputs({"output", "output_scale"})
.Attrs({"block_size: int"})
.SetKernelFn(PD_KERNEL(PerTokenQuant))
.SetInferShapeFn(PD_INFER_SHAPE(PerTokenQuantInferShape))
.SetInferDtypeFn(PD_INFER_DTYPE(PerTokenQuantInferDtype));
PD_BUILD_STATIC_OP(per_token_quant_padding)
.Inputs({"input"})
.Outputs({"output", "output_scale"})
.Attrs({"block_size: int"})
.SetKernelFn(PD_KERNEL(PerTokenQuantPadding))
.SetInferShapeFn(PD_INFER_SHAPE(PerTokenQuantPaddingInferShape))
.SetInferDtypeFn(PD_INFER_DTYPE(PerTokenQuantPaddingInferDtype));
PD_BUILD_STATIC_OP(masked_per_token_quant)
.Inputs({"input", "recv_expert_count"})
.Outputs({"output", "output_scale"})