f4e2264a83
Copied from upstream_ref (NOT rewritten — exact upstream code):
ixformer C++ API (the authoritative header):
include/ixformer.h — ixformer::infer namespace: topk_softmax,
moe_compute_token_index_api, moe_w16a16_group_gemm, moe_expand_input,
moe_output_reduce_sum, silu_and_mul, rms_norm, xllm_paged_attention, etc.
include/ilu_ops_api.h — xllm::kernel::ilu namespace: moe_active_topk,
moe_gen_idx, moe_expand_input, group_gemm, moe_combine_result,
batch_prefill, batch_decode, rms_norm, matmul, act_and_mul, etc.
ILU kernel wrappers (call ixformer::infer directly):
csrc/ilu_kernel_fused_moe.cpp — topk routing + gen_idx + expand + combine
csrc/ilu_kernel_group_gemm.cpp — batched expert GEMM
csrc/ilu_kernel_{activation,norm,rope,matmul,attention}.cpp
ILU layer implementations (full pipeline):
csrc/ilu_layer_fused_moe.{cpp,h} — 797 lines, the complete MoE pipeline
that competitor 168 ran as corex_moe.py
csrc/ilu_layer_attention.{cpp,h} — prefill/decode attention dispatch
CUDA MoE kernels (from xllm + ds_vllm):
csrc/moe/moe_topk_softmax_kernels.cuh — CUB BlockReduce + warp topk
csrc/moe/moe_topk_sigmoid_kernels.cuh — sigmoid scoring variant
csrc/moe/moe_topk.cuh + moe_fused_topk.cu — entry points
csrc/moe/moeTopKFuncs.cuh — TRT-LLM derived vllm-compatible topk
csrc/moe/moe_ops.h + moe_align_sum_kernels.cu — alignment kernels
Common layer headers:
csrc/common_fused_moe{,_base}.h + common_moe_fused_topk.{cpp,h}
798 lines
32 KiB
C++
798 lines
32 KiB
C++
/* Copyright 2026 The xLLM 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
|
|
|
|
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
|
|
|
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 "fused_moe.h"
|
|
|
|
#include <glog/logging.h>
|
|
|
|
#include <iomanip>
|
|
|
|
#include "common/global_flags.h"
|
|
#include "framework/parallel_state/parallel_state.h"
|
|
#include "kernels/ops_api.h"
|
|
#include "layers/common/dp_utils.h"
|
|
#include "util/utils.h"
|
|
|
|
namespace {
|
|
|
|
int32_t get_dtype_size(torch::ScalarType dtype) {
|
|
return static_cast<int32_t>(torch::elementSize(dtype));
|
|
}
|
|
|
|
} // namespace
|
|
|
|
namespace xllm {
|
|
namespace layer {
|
|
|
|
FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
|
|
const FusedMoEArgs& moe_args,
|
|
const QuantArgs& quant_args,
|
|
const ParallelArgs& parallel_args,
|
|
const torch::TensorOptions& options)
|
|
: num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
|
|
topk_(model_args.num_experts_per_tok()),
|
|
num_expert_group_(model_args.n_group()),
|
|
topk_group_(model_args.topk_group()),
|
|
route_scale_(model_args.routed_scaling_factor()),
|
|
hidden_size_(model_args.hidden_size()),
|
|
n_shared_experts_(model_args.n_shared_experts()),
|
|
is_gated_(moe_args.is_gated),
|
|
renormalize_(model_args.norm_topk_prob() ? 1 : 0),
|
|
hidden_act_(model_args.hidden_act()),
|
|
scoring_func_(model_args.scoring_func()),
|
|
quant_args_(quant_args),
|
|
parallel_args_(parallel_args),
|
|
options_(options),
|
|
device_(options.device()) {
|
|
const int64_t num_experts = num_total_experts_;
|
|
const int64_t intermediate_size =
|
|
static_cast<int64_t>(model_args.moe_intermediate_size());
|
|
const std::string& topk_method = model_args.topk_method();
|
|
int64_t ep_size = parallel_args.ep_size();
|
|
int64_t ep_rank = 0;
|
|
tp_pg_ = parallel_args.tp_group_;
|
|
if (ep_size > 1) {
|
|
ep_rank = parallel_args.moe_ep_group_->rank();
|
|
tp_pg_ = parallel_args.moe_tp_group_;
|
|
}
|
|
|
|
// smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
|
|
// supported
|
|
if (!quant_args.quant_method().empty()) {
|
|
if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
|
|
!quant_args.activation_dynamic()) {
|
|
LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
|
|
"quant_method is set. "
|
|
<< "Got quant_method=" << quant_args.quant_method()
|
|
<< ", bits=" << quant_args.bits()
|
|
<< ", activation_dynamic=" << quant_args.activation_dynamic();
|
|
}
|
|
// If confirmed as smoothquant w8a8, set is_smoothquant_ to true
|
|
is_smoothquant_ = true;
|
|
} else {
|
|
is_smoothquant_ = false;
|
|
}
|
|
|
|
// Deep EP initialization check
|
|
enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1;
|
|
if (enable_deep_ep_) {
|
|
// for now, we only implement the deep ep for decode stage.
|
|
// so we will assume the max_token_num is limited to max_batch_size * (1+K)
|
|
// K is the number of speculative tokens.
|
|
int64_t dispatch_token_size;
|
|
if (quant_args.quant_method() == "smoothquant") {
|
|
// float32 is for the scale of the quantized input
|
|
dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
|
|
get_dtype_size(torch::kFloat32);
|
|
} else {
|
|
dispatch_token_size =
|
|
hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
|
|
}
|
|
torch::ScalarType combine_dtype = options_.dtype().toScalarType();
|
|
int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
|
|
// Ensure calculation base is at least ep_size
|
|
int64_t effective_seqs =
|
|
std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size);
|
|
// NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size,
|
|
// regardless of the dp size. To ensure robust scheduling and account
|
|
// for the worst-case scenario, we must guarantee that each rank is capable
|
|
// of handling the maximum possible number of tokens. Therefore, we define
|
|
// max_num_tokens_per_rank as the full maximum value, without dividing by
|
|
// either the rank count or the dp size.
|
|
int64_t max_num_tokens_per_rank =
|
|
(1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
|
|
|
|
// make sure that all layers share the same deep ep instance
|
|
// so that the memory footprint is minimized
|
|
deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
|
|
combine_token_size,
|
|
max_num_tokens_per_rank,
|
|
num_experts,
|
|
parallel_args,
|
|
options_);
|
|
|
|
// obtain the buffer and parameters of deep ep
|
|
deep_ep_buffer_ = deep_ep_->get_buffer();
|
|
deep_ep_params_ = deep_ep_->get_params();
|
|
|
|
// intermediate buffer that can be initialized once
|
|
// we place these tensor here in order to speed up forward pass
|
|
int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
|
|
int64_t token_bytes = is_smoothquant_
|
|
? get_dtype_size(torch::kInt8)
|
|
: get_dtype_size(options_.dtype().toScalarType());
|
|
token_bytes = token_bytes * hidden_size_;
|
|
int64_t head_size = n_tokens_recv * token_bytes;
|
|
dispatch_recv_token_tensor_head_ =
|
|
deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
|
|
.view({n_tokens_recv, token_bytes});
|
|
// input scale in smoothquant
|
|
if (is_smoothquant_) {
|
|
int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
|
|
dispatch_recv_token_tensor_tail_ =
|
|
deep_ep_buffer_.combine_send_token_tensor
|
|
.narrow(0, head_size, tail_size)
|
|
.view({n_tokens_recv, -1});
|
|
}
|
|
}
|
|
|
|
// calculate the number of experts per rank
|
|
num_experts_per_rank_ = num_experts / ep_size;
|
|
start_expert_id_ = ep_rank * num_experts_per_rank_;
|
|
|
|
if (topk_method == "noaux_tc") {
|
|
e_score_correction_bias_ = register_parameter(
|
|
"e_score_correction_bias", torch::empty({num_experts}, options), false);
|
|
}
|
|
|
|
gate_ = register_module(
|
|
"gate_proj",
|
|
ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
|
|
if (n_shared_experts_ > 0) {
|
|
ProcessGroup* shared_expert_pg;
|
|
if (parallel_args_.ep_size() > 1) {
|
|
// we use tp=1 for shared experts computation in deep ep mode
|
|
CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
|
|
<< "Models with shared experts only support ep_size equal to "
|
|
"world size for now.";
|
|
shared_expert_pg = parallel_args.moe_tp_group_;
|
|
} else {
|
|
shared_expert_pg = parallel_args.process_group_;
|
|
}
|
|
// The shared experts computation can proceed in parallel with the
|
|
// final communication step during the MoE computation, as long as it
|
|
// remains independent of any communication operations. For optimal
|
|
// performance, ensure that the shared experts layer on each rank always
|
|
// maintains its own unique weights.
|
|
shared_experts_ =
|
|
register_module("shared_experts",
|
|
DenseMLP(hidden_size_,
|
|
intermediate_size * n_shared_experts_,
|
|
is_gated_,
|
|
false,
|
|
hidden_act_,
|
|
/*enable_result_reduction=*/true,
|
|
quant_args,
|
|
shared_expert_pg,
|
|
options));
|
|
}
|
|
|
|
// create weight buffer
|
|
const int64_t world_size = tp_pg_->world_size();
|
|
int64_t local_intermediate_size = intermediate_size / world_size;
|
|
if (is_smoothquant_) {
|
|
auto quant_option = options_.dtype(torch::kInt8);
|
|
auto fp_option = options_.dtype(torch::kFloat32);
|
|
w13_ = register_parameter(
|
|
"w13",
|
|
torch::empty(
|
|
{num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
|
|
quant_option),
|
|
false);
|
|
w13_scale_ = register_parameter(
|
|
"w13_scale",
|
|
torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
|
|
fp_option),
|
|
false);
|
|
// Note: We do not check enable_deep_ep_ here, since smooth quantization
|
|
// information may be needed even when deep EP mode is disabled. This allows
|
|
// retrieving quantization parameters for any subset of experts as required.
|
|
input_smooth_ = register_parameter(
|
|
"input_smooth",
|
|
torch::empty({num_total_experts_, hidden_size_}, fp_option),
|
|
false);
|
|
w2_ = register_parameter(
|
|
"w2",
|
|
torch::empty(
|
|
{num_experts_per_rank_, hidden_size_, local_intermediate_size},
|
|
quant_option),
|
|
false);
|
|
w2_scale_ = register_parameter(
|
|
"w2_scale",
|
|
torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
|
|
false);
|
|
act_smooth_ = register_parameter(
|
|
"act_smooth",
|
|
torch::empty({num_experts_per_rank_, local_intermediate_size},
|
|
fp_option),
|
|
false);
|
|
|
|
} else {
|
|
w13_ = register_parameter(
|
|
"w13",
|
|
torch::empty(
|
|
{num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
|
|
options_),
|
|
false);
|
|
w2_ = register_parameter(
|
|
"w2",
|
|
torch::empty(
|
|
{num_experts_per_rank_, hidden_size_, local_intermediate_size},
|
|
options_),
|
|
false);
|
|
}
|
|
}
|
|
|
|
torch::Tensor FusedMoEImpl::create_group_gemm_output(
|
|
const torch::Tensor& a,
|
|
const torch::Tensor& b,
|
|
const torch::Tensor& group_list,
|
|
torch::ScalarType dtype,
|
|
torch::Tensor& workspace) {
|
|
// unify shape logic: define the target shape once.
|
|
bool is_3d_weight = (b.dim() != 2);
|
|
int64_t num_tokens = a.size(0);
|
|
int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
|
|
|
|
std::vector<int64_t> output_shape;
|
|
int64_t required_elements = num_tokens * out_dim;
|
|
|
|
if (is_3d_weight) {
|
|
output_shape = {num_tokens, out_dim};
|
|
} else {
|
|
output_shape = {group_list.size(0), num_tokens, out_dim};
|
|
required_elements *= group_list.size(0);
|
|
}
|
|
|
|
auto options = a.options().dtype(dtype);
|
|
|
|
// non-smoothquant: direct allocation
|
|
if (!is_smoothquant_) {
|
|
return torch::empty(output_shape, options);
|
|
}
|
|
|
|
// smoothquant: managed workspace logic
|
|
if (!workspace.defined()) {
|
|
// Lazy initialization: allocate max buffer for the lifecycle
|
|
// Note: accessing class members w13_ and w2_ directly for context
|
|
int64_t max_width = std::max(w13_.size(1), w2_.size(1));
|
|
workspace = torch::empty({num_tokens * max_width}, options);
|
|
}
|
|
|
|
// view construction
|
|
CHECK(workspace.numel() >= required_elements)
|
|
<< "FusedMoE Workspace too small! Alloc: " << workspace.numel()
|
|
<< ", Req: " << required_elements;
|
|
|
|
// utilize the pre-calculated output_shape
|
|
return workspace.slice(0, 0, required_elements).view(output_shape);
|
|
}
|
|
|
|
torch::Tensor FusedMoEImpl::select_experts(
|
|
const torch::Tensor& hidden_states_2d,
|
|
const torch::Tensor& router_logits_2d,
|
|
SelectedExpertInfo& selected_expert_info,
|
|
bool enable_all2all_communication) {
|
|
// prepare the parameters for select_experts
|
|
std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
|
|
if (e_score_correction_bias_.defined()) {
|
|
e_score_correction_bias = e_score_correction_bias_;
|
|
}
|
|
int64_t expert_size = w13_.size(0);
|
|
|
|
// Step 1: apply softmax topk or sigmoid topk / routing logic
|
|
torch::Tensor reduce_weight;
|
|
torch::Tensor expert_id;
|
|
{
|
|
xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
|
|
moe_active_topk_params.input = router_logits_2d;
|
|
moe_active_topk_params.topk = topk_;
|
|
moe_active_topk_params.num_expert_group = num_expert_group_;
|
|
moe_active_topk_params.topk_group = topk_group_;
|
|
moe_active_topk_params.normalize = renormalize_;
|
|
moe_active_topk_params.normed_by = "topk_logit";
|
|
moe_active_topk_params.scoring_func = scoring_func_;
|
|
moe_active_topk_params.route_scale = route_scale_;
|
|
moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
|
|
std::tie(reduce_weight, expert_id) =
|
|
xllm::kernel::moe_active_topk(moe_active_topk_params);
|
|
}
|
|
|
|
// Step 2: generate expert ids
|
|
torch::Tensor gather_idx;
|
|
torch::Tensor combine_idx;
|
|
torch::Tensor token_count;
|
|
std::optional<torch::Tensor> cusum_token_count;
|
|
{
|
|
xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
|
|
moe_gen_idx_params.expert_id = expert_id;
|
|
moe_gen_idx_params.expert_num = num_total_experts_;
|
|
std::vector<torch::Tensor> output_vec =
|
|
xllm::kernel::moe_gen_idx(moe_gen_idx_params);
|
|
gather_idx = output_vec[0];
|
|
combine_idx = output_vec[1];
|
|
token_count = output_vec[2];
|
|
// during all2all communication, we do not need cusum_token_count in the
|
|
// following computation
|
|
if (enable_all2all_communication) {
|
|
cusum_token_count = std::nullopt;
|
|
} else {
|
|
cusum_token_count = output_vec[3];
|
|
}
|
|
}
|
|
|
|
// Step 3: expand and quantize input if needed
|
|
torch::Tensor expand_hidden_states;
|
|
torch::Tensor hidden_states_scale;
|
|
torch::Tensor token_count_slice;
|
|
// all2all related variables
|
|
torch::Tensor dispatch_send_token_tensor;
|
|
// in all2all, the input is scattered, so there is no need to slice the token
|
|
// count, and we can use the dispatch buffer directly
|
|
if (enable_all2all_communication) {
|
|
token_count_slice = token_count;
|
|
int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
|
|
int64_t dispatch_bytes =
|
|
num_token_expand * deep_ep_params_.dispatch_token_size;
|
|
dispatch_send_token_tensor =
|
|
deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
|
|
.view({num_token_expand, deep_ep_params_.dispatch_token_size});
|
|
} else {
|
|
token_count_slice =
|
|
token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
|
|
}
|
|
|
|
if (is_smoothquant_) {
|
|
xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
|
|
scaled_quantize_params.x = hidden_states_2d;
|
|
// use dispatch_send_token_tensor buffer for input
|
|
// to reduce memory footprint
|
|
if (enable_all2all_communication) {
|
|
scaled_quantize_params.smooth = input_smooth_;
|
|
scaled_quantize_params.output =
|
|
dispatch_send_token_tensor.slice(1, 0, hidden_size_);
|
|
} else {
|
|
scaled_quantize_params.smooth = input_smooth_.slice(
|
|
0, start_expert_id_, start_expert_id_ + expert_size);
|
|
scaled_quantize_params.gather_index_start_position =
|
|
cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
|
|
}
|
|
scaled_quantize_params.token_count = token_count_slice;
|
|
scaled_quantize_params.gather_index = gather_idx;
|
|
scaled_quantize_params.act_mode = "none";
|
|
scaled_quantize_params.active_coef = 1.0;
|
|
scaled_quantize_params.is_gated = false;
|
|
scaled_quantize_params.quant_type = torch::kChar;
|
|
std::tie(expand_hidden_states, hidden_states_scale) =
|
|
xllm::kernel::scaled_quantize(scaled_quantize_params);
|
|
if (enable_all2all_communication) {
|
|
// since view_as_dtype has not supported stride yet,
|
|
// we need to copy the scale output to the dispatch buffer
|
|
torch::Tensor dispatch_scale_slice =
|
|
dispatch_send_token_tensor.slice(1, hidden_size_);
|
|
torch::Tensor hidden_states_scale_bytes =
|
|
view_as_dtype(hidden_states_scale, torch::kInt8)
|
|
.view_as(dispatch_scale_slice);
|
|
dispatch_scale_slice.copy_(hidden_states_scale_bytes);
|
|
}
|
|
} else {
|
|
xllm::kernel::MoeExpandInputParams moe_expand_input_params;
|
|
moe_expand_input_params.input = hidden_states_2d;
|
|
moe_expand_input_params.gather_index = gather_idx;
|
|
moe_expand_input_params.combine_idx = combine_idx;
|
|
moe_expand_input_params.topk = topk_;
|
|
expand_hidden_states =
|
|
xllm::kernel::moe_expand_input(moe_expand_input_params);
|
|
if (enable_all2all_communication) {
|
|
// use copy to place the output inside the dispatch buffer
|
|
torch::Tensor dispatch_tensor =
|
|
view_as_dtype(expand_hidden_states, torch::kChar);
|
|
dispatch_send_token_tensor.copy_(dispatch_tensor);
|
|
}
|
|
}
|
|
|
|
// collect the selected tensor
|
|
selected_expert_info.reduce_weight = reduce_weight;
|
|
selected_expert_info.combine_idx = combine_idx;
|
|
selected_expert_info.token_count_slice = token_count_slice;
|
|
selected_expert_info.cusum_token_count = cusum_token_count;
|
|
if (is_smoothquant_) {
|
|
selected_expert_info.input_scale = hidden_states_scale;
|
|
}
|
|
|
|
return expand_hidden_states;
|
|
}
|
|
|
|
torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
|
|
const torch::Tensor& router_logits,
|
|
bool enable_all2all_communication) {
|
|
if (!stream_initialized_) {
|
|
// update device record
|
|
device_ = xllm::Device(hidden_states.device());
|
|
|
|
// acquire streams from the pool again
|
|
routed_stream_ = device_.get_stream_from_pool();
|
|
shared_stream_ = device_.get_stream_from_pool();
|
|
stream_initialized_ = true;
|
|
}
|
|
|
|
std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
|
|
if (e_score_correction_bias_.defined()) {
|
|
e_score_correction_bias = e_score_correction_bias_;
|
|
}
|
|
|
|
// prepare the parameters for MoE computation
|
|
torch::Tensor shared_expert_output;
|
|
torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
|
|
torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
|
|
torch::Tensor hidden_states_2d =
|
|
hidden_states.reshape({-1, hidden_states.size(-1)});
|
|
torch::Tensor router_logits_2d =
|
|
router_logits.reshape({-1, router_logits.size(-1)});
|
|
int64_t group_gemm_max_dim = enable_all2all_communication
|
|
? deep_ep_params_.max_num_tokens_recv / topk_
|
|
: hidden_states_2d.size(0);
|
|
int64_t expert_size = w13_.size(0);
|
|
|
|
// Step 1-3: select experts
|
|
SelectedExpertInfo selected_expert_info;
|
|
torch::Tensor expand_hidden_states =
|
|
select_experts(hidden_states_2d,
|
|
router_logits_2d,
|
|
selected_expert_info,
|
|
enable_all2all_communication);
|
|
|
|
// Communciation Step 1: Dipatch
|
|
// intermediate outputs that are used both in dispatch and combine
|
|
torch::Tensor gather_by_rank_index;
|
|
torch::Tensor token_sum;
|
|
if (enable_all2all_communication) {
|
|
int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
|
|
|
|
// 1. Dispatch Step: Generate layout and send data
|
|
deep_ep_->dispatch_step(dispatch_token_num,
|
|
selected_expert_info.token_count_slice);
|
|
|
|
// 2. Process Result: Generate indices and unpack to computation buffer
|
|
// use the buffer during initialization for the output
|
|
expand_hidden_states = dispatch_recv_token_tensor_head_;
|
|
std::optional<torch::Tensor> output_tail = std::nullopt;
|
|
if (is_smoothquant_) {
|
|
output_tail = dispatch_recv_token_tensor_tail_;
|
|
// update selected_expert_info with the tail (input scale)
|
|
selected_expert_info.input_scale = output_tail;
|
|
}
|
|
|
|
DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
|
|
num_experts_per_rank_, expand_hidden_states, output_tail);
|
|
|
|
// Extract metadata for subsequent steps
|
|
gather_by_rank_index = deep_ep_meta.gather_rank_index;
|
|
selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
|
|
token_sum = deep_ep_meta.token_sum;
|
|
}
|
|
|
|
// common gemm workspace for reduce memory footprint
|
|
torch::Tensor gemm_workspace;
|
|
|
|
// Step 4: group gemm 1
|
|
torch::Tensor gemm1_out =
|
|
create_group_gemm_output(expand_hidden_states,
|
|
w13_,
|
|
selected_expert_info.token_count_slice,
|
|
hidden_states_dtype,
|
|
gemm_workspace);
|
|
// ensure the lifespan of these parameters via brace
|
|
{
|
|
xllm::kernel::GroupGemmParams group_gemm_params;
|
|
torch::ScalarType a_dtype =
|
|
is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
|
|
group_gemm_params.a =
|
|
view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
|
|
group_gemm_params.b = w13_;
|
|
group_gemm_params.token_count =
|
|
selected_expert_info.token_count_slice.to("cpu");
|
|
if (is_smoothquant_) {
|
|
torch::Tensor a_scale =
|
|
selected_expert_info.input_scale.value().flatten();
|
|
selected_expert_info.input_scale =
|
|
view_as_dtype(a_scale, torch::kFloat32);
|
|
group_gemm_params.a_scale = selected_expert_info.input_scale;
|
|
group_gemm_params.b_scale = w13_scale_;
|
|
}
|
|
group_gemm_params.max_dim = group_gemm_max_dim;
|
|
group_gemm_params.trans_a = false;
|
|
group_gemm_params.trans_b = true;
|
|
group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
|
|
group_gemm_params.output = gemm1_out;
|
|
group_gemm_params.combine_idx = std::nullopt;
|
|
gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
|
|
}
|
|
|
|
// Step 5: activation or scaled quantization(fused with activation)
|
|
torch::Tensor act_out;
|
|
torch::Tensor act_out_scale;
|
|
if (is_smoothquant_) {
|
|
int64_t slice_dim = gemm1_out.size(1);
|
|
if (is_gated_) slice_dim /= 2;
|
|
// slice operation is a view, does not take up extra memory, but points to
|
|
// the same memory
|
|
act_out = expand_hidden_states.slice(1, 0, slice_dim);
|
|
act_out_scale =
|
|
selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
|
|
// call scaled quantization kernel (also fused with activation)
|
|
xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
|
|
scaled_quantize_params.x = gemm1_out;
|
|
scaled_quantize_params.smooth = act_smooth_;
|
|
scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
|
|
scaled_quantize_params.output = act_out;
|
|
scaled_quantize_params.output_scale = act_out_scale;
|
|
scaled_quantize_params.act_mode = hidden_act_;
|
|
scaled_quantize_params.active_coef = 1.0;
|
|
scaled_quantize_params.is_gated = is_gated_;
|
|
scaled_quantize_params.quant_type = torch::kChar;
|
|
std::tie(act_out, act_out_scale) =
|
|
xllm::kernel::scaled_quantize(scaled_quantize_params);
|
|
} else {
|
|
act_out = is_gated_
|
|
? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
|
|
: gemm1_out;
|
|
// call activation kernel
|
|
xllm::kernel::ActivationParams activation_params;
|
|
activation_params.input = gemm1_out;
|
|
activation_params.output = act_out;
|
|
activation_params.cusum_token_count =
|
|
selected_expert_info.cusum_token_count;
|
|
activation_params.act_mode = hidden_act_;
|
|
activation_params.is_gated = is_gated_;
|
|
activation_params.start_expert_id = start_expert_id_;
|
|
activation_params.expert_size = expert_size;
|
|
xllm::kernel::active(activation_params);
|
|
}
|
|
|
|
// Step 6: group gemm 2
|
|
torch::Tensor gemm2_out =
|
|
create_group_gemm_output(act_out,
|
|
w2_,
|
|
selected_expert_info.token_count_slice,
|
|
hidden_states_dtype,
|
|
gemm_workspace);
|
|
// ensure the lifespan of these parameters via brace
|
|
{
|
|
xllm::kernel::GroupGemmParams group_gemm_params;
|
|
group_gemm_params.a = act_out;
|
|
group_gemm_params.b = w2_;
|
|
group_gemm_params.token_count =
|
|
selected_expert_info.token_count_slice.to("cpu");
|
|
if (is_smoothquant_) {
|
|
group_gemm_params.a_scale = act_out_scale;
|
|
group_gemm_params.b_scale = w2_scale_;
|
|
}
|
|
group_gemm_params.max_dim = group_gemm_max_dim;
|
|
group_gemm_params.trans_a = false;
|
|
group_gemm_params.trans_b = true;
|
|
group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
|
|
group_gemm_params.output = gemm2_out;
|
|
group_gemm_params.combine_idx = selected_expert_info.combine_idx;
|
|
gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
|
|
}
|
|
|
|
// Communciation Step 2: Combine
|
|
if (enable_all2all_communication) {
|
|
int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
|
|
// Delegate pack, layout generation and combine to DeepEP
|
|
torch::Tensor combine_send_layout =
|
|
deep_ep_->combine_step_pack(gemm2_out,
|
|
gather_by_rank_index,
|
|
token_sum,
|
|
hidden_size_,
|
|
hidden_states_dtype);
|
|
|
|
// create a wait event for the current stream to finish computation
|
|
auto current_stream = device_.current_stream();
|
|
routed_stream_->wait_stream(*current_stream);
|
|
// pure communciation kernel: dispatch
|
|
{
|
|
torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
|
|
gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
|
|
num_token_expand,
|
|
hidden_size_,
|
|
hidden_states_dtype);
|
|
}
|
|
|
|
// pure computation kernel: shared experts
|
|
if (n_shared_experts_ > 0) {
|
|
shared_stream_->wait_stream(*current_stream);
|
|
torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
|
|
shared_expert_output = shared_experts_(hidden_states);
|
|
}
|
|
|
|
// join for parallelization
|
|
current_stream->wait_stream(*routed_stream_);
|
|
if (n_shared_experts_ > 0) {
|
|
current_stream->wait_stream(*shared_stream_);
|
|
}
|
|
}
|
|
|
|
// After group gemm is finished, some tensors are no
|
|
// longer needed. We must explicitly release the memory.
|
|
expand_hidden_states = torch::Tensor();
|
|
selected_expert_info.input_scale = std::nullopt;
|
|
act_out = torch::Tensor();
|
|
|
|
// Step 7: combine the intermediate results and get the final hidden states
|
|
torch::Tensor final_hidden_states;
|
|
// ensure the lifespan of these parameters via brace
|
|
{
|
|
xllm::kernel::MoeCombineResultParams moe_combine_result_params;
|
|
moe_combine_result_params.input = gemm2_out;
|
|
moe_combine_result_params.reduce_weight =
|
|
selected_expert_info.reduce_weight;
|
|
moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
|
|
moe_combine_result_params.cusum_token_count =
|
|
selected_expert_info.cusum_token_count;
|
|
moe_combine_result_params.start_expert_id = start_expert_id_;
|
|
moe_combine_result_params.expert_size = expert_size;
|
|
moe_combine_result_params.bias = std::nullopt;
|
|
// if all2all communication is enabled and shared output is provided,
|
|
// we will fused the add up to combine result
|
|
if (enable_all2all_communication && n_shared_experts_ > 0) {
|
|
moe_combine_result_params.residual =
|
|
shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
|
|
}
|
|
|
|
final_hidden_states =
|
|
xllm::kernel::moe_combine_result(moe_combine_result_params);
|
|
}
|
|
|
|
// reshape the final hidden states to the original shape
|
|
final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
|
|
|
|
if (enable_all2all_communication) {
|
|
return final_hidden_states;
|
|
}
|
|
|
|
// Communciation Step 3: AllReduce for non-all2all communication
|
|
// shared experts can be parallelized with the final communication step
|
|
// during moe computation.
|
|
auto current_stream = device_.current_stream();
|
|
routed_stream_->wait_stream(*current_stream);
|
|
{
|
|
torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
|
|
if (tp_pg_->world_size() > 1) {
|
|
final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
|
|
}
|
|
if (parallel_args_.ep_size() > 1) {
|
|
final_hidden_states = parallel_state::reduce(
|
|
final_hidden_states, parallel_args_.moe_ep_group_);
|
|
}
|
|
}
|
|
|
|
if (n_shared_experts_ > 0) {
|
|
shared_stream_->wait_stream(*current_stream);
|
|
torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
|
|
// for non all2all, we compute the shared experts parallelized with the
|
|
// final communication step
|
|
shared_expert_output = shared_experts_(hidden_states);
|
|
shared_expert_output =
|
|
shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
|
|
}
|
|
|
|
// join for parallelization
|
|
current_stream->wait_stream(*routed_stream_);
|
|
if (n_shared_experts_ > 0) {
|
|
current_stream->wait_stream(*shared_stream_);
|
|
final_hidden_states += shared_expert_output;
|
|
}
|
|
|
|
return final_hidden_states;
|
|
}
|
|
|
|
torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
|
|
const ModelInputParams& input_params) {
|
|
// we only support all2all communication for decode stage for now
|
|
bool enable_all2all_communication =
|
|
enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(),
|
|
input_params.dp_is_decode.end(),
|
|
[](int32_t val) { return val == 1; });
|
|
|
|
bool is_dp_ep_parallel =
|
|
parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
|
|
// during all2all communication, the output has been
|
|
// gathered and sliced by dispatch and combine steps,
|
|
// so we do not need to gather input and slice output again
|
|
bool need_gather_and_slice =
|
|
is_dp_ep_parallel && !enable_all2all_communication;
|
|
|
|
auto input = hidden_states;
|
|
if (need_gather_and_slice) {
|
|
input = parallel_state::gather(input,
|
|
parallel_args_.dp_local_process_group_,
|
|
input_params.dp_global_token_nums);
|
|
}
|
|
// MoE Gate
|
|
auto router_logits = gate_(input);
|
|
|
|
// MoE Experts
|
|
auto output =
|
|
forward_experts(input, router_logits, enable_all2all_communication);
|
|
|
|
if (need_gather_and_slice) {
|
|
output = get_dp_local_slice(output, input_params, parallel_args_);
|
|
}
|
|
|
|
return output;
|
|
}
|
|
|
|
void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
|
|
if (e_score_correction_bias_.defined() &&
|
|
!e_score_correction_bias_is_loaded_) {
|
|
LOAD_WEIGHT(e_score_correction_bias);
|
|
}
|
|
}
|
|
|
|
void FusedMoEImpl::load_experts(const StateDict& state_dict) {
|
|
const int64_t rank = tp_pg_->rank();
|
|
const int64_t world_size = tp_pg_->world_size();
|
|
const int64_t start_expert_id = start_expert_id_;
|
|
const int64_t num_experts_per_rank = num_experts_per_rank_;
|
|
const int64_t num_total_experts = num_total_experts_;
|
|
std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
|
|
if (is_smoothquant_) {
|
|
LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
|
|
LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
|
|
// When supporting DeepEP All2All mode,
|
|
// we need to load the complete set of expert weights corresponding to
|
|
// "up_proj.smooth". Note that even if deep EP mode is not enabled, it
|
|
// remains possible to retrieve the smooth quantization information for a
|
|
// subset of experts. Therefore, we intentionally do not check whether
|
|
// deep_ep_ is enabled in this case.
|
|
LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
|
|
LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
|
|
LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
|
|
LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
|
|
} else {
|
|
LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
|
|
LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
|
|
}
|
|
}
|
|
|
|
void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
|
|
if (state_dict.size() == 0) {
|
|
return;
|
|
}
|
|
|
|
if (n_shared_experts_ > 0) {
|
|
shared_experts_->load_state_dict(
|
|
state_dict.get_dict_with_prefix("shared_experts."));
|
|
}
|
|
gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
|
|
load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
|
|
load_experts(state_dict.get_dict_with_prefix("experts."));
|
|
}
|
|
|
|
} // namespace layer
|
|
} // namespace xllm
|