wehub-resource-sync
282 lines
13 KiB
C++
282 lines
13 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/backends/xpu/xpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
|
|
#ifndef MAX_NUM_EXPERTS
|
|
#define MAX_NUM_EXPERTS 80
|
|
#endif
|
|
|
|
template <typename T, typename Context>
|
|
void dispatch_tokens_unzip_stable(const Context &dev_ctx,
|
|
const DenseTensor &X,
|
|
const DenseTensor &expert_routemap_topk,
|
|
const DenseTensor &expert_prob_topk,
|
|
const optional<DenseTensor> &XScale,
|
|
const DenseTensor &expert_offsets,
|
|
DenseTensor *X_unzipped,
|
|
DenseTensor *zipped_expertwise_rowmap,
|
|
DenseTensor *token_prob_unzipped,
|
|
DenseTensor *XScale_unzipped,
|
|
const int total_zipped_tokens_num,
|
|
const int token_length,
|
|
const int total_tokens_after_broadcast,
|
|
const int topk,
|
|
const int num_experts,
|
|
const int scale_length,
|
|
const bool do_gather) {
|
|
#define DTYPE_CASE(dtype, type) dtype == DataType::type
|
|
#define GET_DATA(tensor, type) tensor.data<type>()
|
|
#define GET_XPU_DATA(tensor, type, xpu_type) \
|
|
reinterpret_cast<const xpu_type *>(tensor.data<type>())
|
|
#define GET_PTR_XPU_DATA(tensor, type, xpu_type) \
|
|
reinterpret_cast<xpu_type *>(tensor->data<type>())
|
|
|
|
#define DISPATCH_CASE(TOKEN_T, PROB_T, INT_T, HAS_SCALE, DO_GATHER) \
|
|
using XPU_TOKEN_T = typename XPUTypeTrait<TOKEN_T>::Type; \
|
|
using XPU_PROB_T = typename XPUTypeTrait<PROB_T>::Type; \
|
|
using XPU_INT_T = typename XPUTypeTrait<INT_T>::Type; \
|
|
\
|
|
int r = xpu::moe_permute<XPU_TOKEN_T, XPU_INT_T, XPU_PROB_T>( \
|
|
dev_ctx.x_context(), \
|
|
reinterpret_cast<const XPU_TOKEN_T *>( \
|
|
X.data<TOKEN_T>()), /* hidden_states */ \
|
|
(XScale ? XScale.get_ptr()->data<float>() : nullptr), /* scale */ \
|
|
reinterpret_cast<const XPU_INT_T *>( \
|
|
expert_routemap_topk.data<INT_T>()), /* expert_routemap_topk */ \
|
|
reinterpret_cast<const XPU_PROB_T *>( \
|
|
expert_prob_topk.data<PROB_T>()), /* expert_prob_topk */ \
|
|
reinterpret_cast<const XPU_INT_T *>( \
|
|
expert_offsets.data<int>()), /* expert_base_offset */ \
|
|
reinterpret_cast<XPU_TOKEN_T *>( \
|
|
X_unzipped->data<TOKEN_T>()), /* hidden_states_unzipped */ \
|
|
reinterpret_cast<XPU_INT_T *>( \
|
|
zipped_expertwise_rowmap \
|
|
->data<INT_T>()), /* zipped_expertwise_rowmap */ \
|
|
reinterpret_cast<XPU_PROB_T *>( \
|
|
token_prob_unzipped->data<PROB_T>()), /* token_prob_unzipped */ \
|
|
XScale_unzipped->data<float>(), /* scale_unzipped */ \
|
|
static_cast<int64_t>(total_zipped_tokens_num), /* sequence_length */ \
|
|
static_cast<int64_t>(token_length), /* hidden_size */ \
|
|
static_cast<int64_t>( \
|
|
total_tokens_after_broadcast), /* total_tokens_after_broadcast */ \
|
|
static_cast<int64_t>(topk), /* topk */ \
|
|
static_cast<int64_t>(num_experts), /* num_experts */ \
|
|
128, /* num_scale */ \
|
|
DO_GATHER /* do_gather */ \
|
|
); \
|
|
\
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "moe_permute");
|
|
|
|
#define HANDLE_GATHER_CASE(TOKEN_T, PROB_T, INT_T, HAS_SCALE) \
|
|
if (do_gather) { \
|
|
DISPATCH_CASE(TOKEN_T, PROB_T, INT_T, HAS_SCALE, true) \
|
|
} else { \
|
|
DISPATCH_CASE(TOKEN_T, PROB_T, INT_T, HAS_SCALE, false) \
|
|
}
|
|
|
|
// HANDLE_GATHER_CASE(phi::float8_e4m3fn, PROB_T, INT_T, true)
|
|
#define HANDLE_TOKEN_TYPE(PROB_T, INT_T) \
|
|
if (DTYPE_CASE(X.dtype(), BFLOAT16)) { \
|
|
HANDLE_GATHER_CASE(phi::bfloat16, PROB_T, INT_T, false) \
|
|
} else if (DTYPE_CASE(X.dtype(), FLOAT8_E4M3FN)) { \
|
|
PADDLE_THROW(common::errors::Unimplemented( \
|
|
"moe_permute input only support bfloat16")); \
|
|
}
|
|
|
|
#define HANDLE_PROB_TYPE(INT_T) \
|
|
if (DTYPE_CASE(expert_prob_topk.dtype(), BFLOAT16)) { \
|
|
PADDLE_THROW(common::errors::Unimplemented( \
|
|
"moe_permute expert_prob_topk only support float32")); \
|
|
} else if (DTYPE_CASE(expert_prob_topk.dtype(), FLOAT32)) { \
|
|
HANDLE_TOKEN_TYPE(float, INT_T) \
|
|
}
|
|
|
|
if (DTYPE_CASE(zipped_expertwise_rowmap->dtype(), INT32)) {
|
|
HANDLE_PROB_TYPE(int)
|
|
}
|
|
|
|
#undef DTYPE_CASE
|
|
#undef GET_DATA
|
|
#undef GET_XPU_DATA
|
|
#undef GET_PTR_XPU_DATA
|
|
#undef DISPATCH_CASE
|
|
#undef HANDLE_EXPERT_CASE
|
|
#undef HANDLE_TOKEN_TYPE
|
|
#undef HANDLE_PROB_TYPE
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void MoePermuteKernel(const Context &dev_ctx,
|
|
const DenseTensor &X, // hidden_states
|
|
const optional<DenseTensor> &XScale,
|
|
const DenseTensor &expert_routemap_topk,
|
|
const DenseTensor &expert_prob_topk,
|
|
const int num_experts,
|
|
const std::vector<int> &tokens_per_expert,
|
|
const int padding_multiplex,
|
|
const bool do_gather,
|
|
const bool using_ue8m0_scale,
|
|
const bool return_expert_indices,
|
|
const int override_buffer_size,
|
|
DenseTensor *X_unzipped,
|
|
DenseTensor *zipped_expertwise_rowmap,
|
|
DenseTensor *token_prob_unzipped,
|
|
DenseTensor *XScale_unzipped,
|
|
DenseTensor *expert_indices) {
|
|
PADDLE_ENFORCE_EQ(
|
|
return_expert_indices,
|
|
false,
|
|
common::errors::Unimplemented("moe_permute on XPU does not support "
|
|
"return_expert_indices=true yet."));
|
|
PADDLE_ENFORCE_EQ(
|
|
override_buffer_size,
|
|
-1,
|
|
common::errors::Unimplemented(
|
|
"moe_permute on XPU does not support override_buffer_size yet."));
|
|
const int64_t rows = X.dims()[0];
|
|
const int64_t cols = X.dims()[1];
|
|
PADDLE_ENFORCE_LE(
|
|
rows,
|
|
std::numeric_limits<int32_t>::max(),
|
|
common::errors::InvalidArgument("X.dims()[0] should be less than "
|
|
"INT_MAX, received X.dims()[0]: (%ld)",
|
|
rows));
|
|
PADDLE_ENFORCE_LE(
|
|
cols,
|
|
std::numeric_limits<int32_t>::max(),
|
|
common::errors::InvalidArgument("X.dims()[1] should be less than "
|
|
"INT_MAX, received X.dims()[1]: (%ld)",
|
|
cols));
|
|
PADDLE_ENFORCE_LE(
|
|
num_experts,
|
|
MAX_NUM_EXPERTS,
|
|
common::errors::InvalidArgument(
|
|
"Currently we support no more than (%ld), received num_expert: "
|
|
"(%ld). Please check input "
|
|
"value.",
|
|
MAX_NUM_EXPERTS,
|
|
num_experts));
|
|
const int64_t quanted_cols = (XScale) ? XScale.get_ptr()->dims()[1] : 0;
|
|
PADDLE_ENFORCE_LE(
|
|
quanted_cols,
|
|
std::numeric_limits<int32_t>::max(),
|
|
common::errors::InvalidArgument("quanted_cols should be less than "
|
|
"INT_MAX, received quanted_cols: (%ld)",
|
|
quanted_cols));
|
|
|
|
// Expert base offset initialization, tensor numeric range [0, max_token_num]
|
|
int expert_offset[MAX_NUM_EXPERTS];
|
|
int tokens_cumulated = 0;
|
|
for (int i = 0; i < MAX_NUM_EXPERTS; i++) {
|
|
if (i < num_experts) {
|
|
expert_offset[i] = tokens_cumulated;
|
|
tokens_cumulated +=
|
|
((tokens_per_expert[i] + padding_multiplex - 1) / padding_multiplex) *
|
|
padding_multiplex;
|
|
} else {
|
|
expert_offset[i] = 0;
|
|
}
|
|
}
|
|
DenseTensor expert_offset_tensor;
|
|
expert_offset_tensor.Resize({MAX_NUM_EXPERTS});
|
|
dev_ctx.template Alloc<int>(&expert_offset_tensor);
|
|
PADDLE_ENFORCE_XPU_SUCCESS(
|
|
cudaMemcpyAsync(expert_offset_tensor.data<int>(),
|
|
expert_offset,
|
|
sizeof(int) * MAX_NUM_EXPERTS,
|
|
cudaMemcpyHostToDevice,
|
|
reinterpret_cast<cudaStream_t>(dev_ctx.stream())));
|
|
// ------------------- resource allocate -------------------------
|
|
const int output_rows = tokens_cumulated;
|
|
const int64_t topk = expert_routemap_topk.dims()[1];
|
|
PADDLE_ENFORCE_LE(
|
|
topk,
|
|
std::numeric_limits<int32_t>::max(),
|
|
common::errors::InvalidArgument(
|
|
"topk should be less than INT_MAX, received topk: (%ld)", topk));
|
|
|
|
dev_ctx.template Alloc<T>(X_unzipped);
|
|
dev_ctx.template Alloc<float>(XScale_unzipped);
|
|
dev_ctx.template Alloc<int>(zipped_expertwise_rowmap);
|
|
dev_ctx.template Alloc<float>(token_prob_unzipped);
|
|
auto X_unzipped_ptr = reinterpret_cast<void *>(X_unzipped->data<T>());
|
|
auto token_prob_unzipped_ptr =
|
|
reinterpret_cast<void *>(token_prob_unzipped->data<float>());
|
|
auto XScale_unzipped_ptr =
|
|
reinterpret_cast<void *>(XScale_unzipped->data<float>());
|
|
|
|
// -------- Memset all padding area to zero, with regard to do_gather
|
|
auto memset_invalid_rows =
|
|
[&](void *ptr, int64_t element_size, int64_t stride) {
|
|
for (int i = 0; i < num_experts; i++) {
|
|
int64_t next_expert_offset =
|
|
i < num_experts - 1 ? expert_offset[i + 1] : output_rows;
|
|
int64_t invalid_rows =
|
|
next_expert_offset - expert_offset[i] - tokens_per_expert[i];
|
|
int64_t cur_expert_end = expert_offset[i] + tokens_per_expert[i];
|
|
|
|
PADDLE_ENFORCE_XPU_SUCCESS(cudaMemsetAsync(
|
|
ptr + cur_expert_end * stride * element_size,
|
|
0,
|
|
element_size * invalid_rows * stride,
|
|
reinterpret_cast<cudaStream_t>(dev_ctx.stream())));
|
|
}
|
|
};
|
|
if (do_gather) { // no gather, no memset
|
|
memset_invalid_rows(X_unzipped_ptr, sizeof(T), cols);
|
|
if (XScale) {
|
|
memset_invalid_rows(XScale_unzipped_ptr, sizeof(float), quanted_cols);
|
|
}
|
|
}
|
|
// Probs will be memset to zero whatsoever
|
|
memset_invalid_rows(token_prob_unzipped_ptr, sizeof(float), 1);
|
|
|
|
// Handle 0-size input
|
|
if (X.numel() == 0) return;
|
|
|
|
// -------- Initialize semaphore for cumsum ---------------
|
|
dispatch_tokens_unzip_stable<T, Context>(dev_ctx,
|
|
X,
|
|
expert_routemap_topk,
|
|
expert_prob_topk,
|
|
XScale,
|
|
expert_offset_tensor,
|
|
X_unzipped,
|
|
zipped_expertwise_rowmap,
|
|
token_prob_unzipped,
|
|
XScale_unzipped,
|
|
static_cast<int>(rows),
|
|
static_cast<int>(cols),
|
|
static_cast<int>(output_rows),
|
|
static_cast<int>(topk),
|
|
num_experts,
|
|
static_cast<int>(quanted_cols),
|
|
do_gather);
|
|
}
|
|
#undef MAX_NUM_EXPERTS
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(moe_permute,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::MoePermuteKernel,
|
|
// phi::float8_e4m3fn,
|
|
phi::bfloat16) {}
|